h3-blackwell-runtime/PLAN.md

# H3 Blackwell Runtime Plan

Current implementation status is tracked in [`CURRENT_STATE.md`](CURRENT_STATE.md).
The dated handoffs below are retained as historical investigation records.

## Goal

Build a direct MiniMax H3 Ref2VA runtime for Blackwell and Grace Blackwell that consumes the current Comfy safetensors checkpoints while removing ComfyUI and Raylight from the denoising critical path.

The runtime must support one GPU first, then correct 2/4/6/8 GPU execution. It must retain the current NVFP4 model artifacts and use SageAttention3 where quality validation permits.

## Future LTX 2.5 Track

Add a separate LTX 2.5 direct-runtime adapter after H3 single-GPU parity is stable. Target the gated Lightricks `ltx-2.5-22b-distilled-transformer-nvfp4.safetensors` artifact (18.7 GB, release commit `dd53cc2cd45bbeaa3563dfb575cba3f49cf44761`).

- Keep LTX model loading, conditioning, scheduler, VAE, and validation isolated from H3; this is a second model family, not an H3 checkpoint variant.
- Inspect the safetensors header and published architecture/configuration before sharing H3 modules or kernels.
- Establish a LTX SDPA/Sage2 correctness baseline before evaluating Sage3, FlashAttention-4, Sol-Attn, cache methods, or distributed layouts.
- Respect the LTX 2 Community License Agreement and gated-access requirements; do not automate downloads without authorized access.

## Reference Baseline

The first acceptance target is the clean one-GPU ComfyUI baseline in `../h3-lab/h3-raylight-usp2-results.json`:

- RTX PRO 6000 Blackwell, 96 GB
- Ref2VA, 960x544, 124 frames, 24 fps
- 12 steps, `beta`, `res_multistep`, seed `440202`
- SageAttention3, 1 GPU
- ComfyUI execution time: `49.893s`

The direct runner must first match the model contract and output quality. Beating this timing comes after correctness is established.

## Architecture

1. Checkpoint adapter: read Comfy safetensors metadata, preserve packed low-precision weights and scale tensors, and map them into a canonical H3 state dictionary.
2. Conditioning service: execute and cache Qwen layer-50 embeddings, modality tags, and reference VAE latents once per request.
3. H3 denoiser: implement the packed Ref2VA DiT, 3-axis RoPE, AdaLN, dual audio/video schedule, and RES multistep solver without node-graph orchestration.
4. Kernel layer: retain the known-good NVFP4 linear path initially; add explicit SageAttention3 and CUDA-graph buckets after exact single-GPU output validation.
5. Distributed layer: use ragged Ulysses all-to-all for Q/K/V head exchange, Sage3 on full packed tokens per local head shard, then inverse exchange. Add tensor parallelism only after sequence parallel correctness is proven.

## Milestones

1. Inspect the actual local `pruned_nvfp4` checkpoint header and classify every tensor/scale layout.
2. Create a direct single-GPU denoiser step matching ComfyUI for a fixed captured payload.
3. Implement full single-GPU Ref2VA and compare per-step tensors plus final AV output against ComfyUI.
4. Apply SageAttention3 and CUDA graphs; benchmark against the 49.893s reference.
   - Optional backends and execution strategies to evaluate behind the same per-step quality gate: FlashAttention-4 (the Blackwell successor to Hopper-only FlashAttention-3), EasyCache/H3-Cache, Sol-Attn, and KJ exact memory-lifetime patches.
   - Keep backend selection explicit per run; retain only candidates that match the validated direct correctness path and improve the measured denoising bottleneck.
   - Current correctness baseline: SageAttention2 (`sage2`), which exactly matches the captured ComfyUI `--use-sage-attention` output. SDPA is a fallback; SageAttention3 remains experimental and must pass the same quality gate.
   - Sol-Attn and KJ Sage have prior H3 test evidence and are supported experimental candidates. Integrate each as an isolated standalone adapter, record the exact mode/version, and gate it against the Sage2 per-step reference before combining it with caches or other approximation strategies.
5. Implement ragged Ulysses Sage3 with transport-identity and distributed-versus-single-Sage3 tests.
6. Sweep Ulysses/tensor-parallel layouts on 2/4/6/8 GPUs in an NVLink/NVSwitch domain.

## Performance Backend Plan

Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Feature and performance work should proceed in this order:

1. Validate and benchmark the existing `sage3` backend against the same cat prompt, seed, dimensions, and FP16 VAE runtime path used for Sage2 parity.
   - First cat benchmark result: Sage3 runs successfully but is slower than Sage2 in this direct path. Sampling was `123.675s` versus Sage2 `114.414s`; warm after text conditioning was `158.111s` versus Sage2 `149.304s`. Same-seed MP4 frame diff versus Sage2 was mean `46.563`, max `255`, so keep Sage3 experimental pending human visual review and stricter tensor gates.
2. Build a persistent hot runtime service instead of measuring only process-per-run CLIs. A warm container must preload and retain Qwen, H3, video VAE, and audio VAE in GPU memory, then accept video jobs without model-load latency. Add explicit startup warmup, readiness reporting, request-level attention selection (`sage2`, `sdpa`, `sage3` initially), and timing fields that separate resident-model request latency from cold startup.
3. Add exact memory/lifetime optimizations next: `kj_head_sliced` and `kj_chunked_ffn`. These must preserve the validated direct outputs before being kept.
4. Evaluate prior H3-tested attention candidates as standalone adapters: `sol_attn` and `kj_sage`.
   - `kj_sage` is implemented as explicit SageAttention mode backends: `kj_sage_cuda`, `kj_sage_triton`, `kj_sage_fp8`, and `kj_sage_fp8pp`; all passed hot-runtime smoke tests.
   - `sol_attn` is implemented through the pinned ComfyUI Triton source vendored
     into the Spark image. The native QKV layout is deployed; the inactive fused
     layout remains a known repair item.
5. Evaluate approximate denoiser caches only after exact baselines are recorded: `easycache` and `h3_cache`.
   - Initial direct cache modes are implemented as opt-in approximate sampler modes. They reuse cached denoised deltas and report skipped-step stats; full-size quality/threshold sweeps are still required before using them for production output.
6. Keep every backend explicit per run, with separate quality and timing records for sampling, VAE, audio, and end-to-end output.

## Next Performance Plan: Blackwell NVFP4 GEMMs

The latest attention/cache sweeps show that attention backend swaps are not the main remaining speed lever. Full request times for `sage2`, `kj_sage_fp8`, `kj_sage_fp8pp`, `kj_head_sliced`, and `sol_attn` are close, while cache gains come from skipping denoiser calls and must remain quality-gated. The next exact performance target is therefore the dense transformer linear stack.

Target mixed-precision policy:

- Keep residual stream, RMSNorm, RoPE, modulation/AdaLN, residual adds, timestep embeddings, softmax, final layer, audio VAE, and video VAE in BF16/FP16.
- Keep attention accumulation/output in the best measured BF16/FP16 backend until a separate quality gate proves otherwise.
- Replace only the large transformer GEMMs first: QKV projection, attention output projection, MLP gate/up projection, and MLP down projection.
- Prioritize Blackwell-native NVFP4 activation x NVFP4 weight GEMMs with BF16 accumulation/output, avoiding unnecessary dequantize/contiguous boundaries.

Kernel candidates to test, in order:

1. Current `Nvfp4Linear`/Comfy Kitchen path as the measured baseline.
2. Transformer Engine NVFP4BlockScaling prototype to establish expected Blackwell FP4 behavior with a higher-level NVIDIA stack.
3. CUTLASS/CuTe block-scaled NVFP4 GEMM prototype for owned hot-path kernels, exact layouts, fused epilogues, and persistent-kernel experiments.
4. cuBLASLt NVFP4/grouped GEMM where CUDA exposes a supported operation on the current Blackwell target.

Current NVFP4 activation quantization status:

- `H3_NVFP4_SCALE_BACKEND=vortex` is parity-safe for BF16 H3 activations and is the useful optimization seam. The tuned scale kernel uses 128 threads / 256 blocks and reaches about 88-90% of DGX Spark's advertised LPDDR5X bandwidth on FC1/FC2 scale discovery.
- Standalone scale discovery is effectively done for Spark; remaining theoretical headroom is too small to justify more absmax-only tuning.
- `vortex_native` BF16->NVFP4 pack is parity-safe for real H3 FC1/FC2 activation shapes, including qdata, cuBLAS tiled block-scale layout, signed zero, tie-to-even FP4 rounding, and downstream `scaled_mm_nvfp4` linear output.
- Keep Comfy Kitchen as the default production packer. Pack-only sweep `8d90cec` shows native 512-thread pack is essentially tied on FC1 (`1.090 ms` vs CK `1.098 ms`) but still slower on FC2 (`2.977 ms` vs CK `2.886 ms`). Use `vortex_native` only as an experimental benchmark path until the wide-FC2 pack kernel is redesigned.

Validation and profiling sequence:

1. Use `tools/profile_h3_block.py` to measure one representative H3 block before writing kernels. Record QKV, RoPE/RMS, attention kernel, output projection, MLP fc1, activation, MLP fc2, modulation/gating, and total block time.
2. Repeat profiling for representative block indices, at least block `0`, `24`, and `49`, because token statistics and cache behavior can differ through depth.
3. Confirm which component dominates before implementing a kernel replacement.
4. Add candidate GEMM backends behind the existing `Nvfp4Linear` API so model code and correctness tests remain stable.
5. Gate each candidate by layer-level max/mean error, one-block output error, one denoiser-step tensor error, and finally full video quality.

DGX Spark caveat:

- GB10 reports SM121, not B200/GB200 SM100. CUTLASS, Transformer Engine, and cuBLASLt FP4 coverage must be probed independently on Spark before assuming B200 examples work unchanged.
- Keep separate kernel policy notes for GB10/Spark and B200/GB200-class Blackwell.

## Non-Negotiable Validation

- Never silently pad semantic H3 tokens for unmasked attention.
- Compare distributed output against the identical single-GPU Sage3 path before comparing to SDPA.
- Validate denoiser outputs at each scheduler step, not only encoded video.
- Record attention, GEMM, communication, VAE, and end-to-end timings separately.
- Treat SageAttention3 as an experimental quality-gated kernel for H3.

## 2026-08-15 Handoff

Current quality baseline:

- `sage2` is the correct visual path for the 960x544, 124-frame, 12-step cat prompt at seed `440407`.
- Fresh Spark hot-runtime Sage2/Vortex A/B asset: `\\192.168.1.162\StoryStudioAssets\H3-output\h3-blackwell-runtime\cat-sage2-vortex-5s-440407-ab.mp4`.
- `sol_attn` is faster but visually wrong for the same prompt/seed. Keep it experimental until localized and quality-gated.
- Sol-native BSHD layout is exact versus the old Sol path, so the visual issue is likely Sol's approximation/configuration, not the layout removal.

Tomorrow priority:

1. Diagnose `sol_attn` quality failure.
   - Save per-step latents for `sage2` and `sol_attn` with identical prompt, seed, dimensions, and sampler settings.
   - Measure per-step max/mean/RMSE/cosine divergence.
   - Try hybrid schedules: early steps `sage2`, later steps `sol_attn`; also try Sage2 every Nth step.
   - Sweep Sol params only after the divergence point is known: `H3_SOL_TAU`, `H3_SOL_THRESH_TYPE`, `H3_SOL_INT8_QK`, `H3_SOL_INT8_PV`.
2. Optimize the quality path (`sage2`).
   - Use `sage2` as visual/numerical reference.
   - Attack Q/K/V HND layout and Q/K RMS+RoPE first; a fused Q/K RMS+RoPE + Sage HND layout kernel is the most plausible exact win.
   - Keep Sage3 experimental; previous measurements were slower and visually/numerically different.
3. Add durable hot-runtime request logging.
   - Log `generate_start`, `generate_complete`, and `generate_failed` to stdout with output path, seed, attention, dimensions, and timings.
   - This prevents losing timing data when a client disconnects before reading the synchronous response.
4. Check LTX 2.5 downloads.
   - Detached Spark container: `ltx25-nvfp4-download`.
   - Expected files under `/home/daniel/aeon-spark-test/h3/comfy-models`:
     - `diffusion_models/ltx-2.5-22b-distilled-transformer-nvfp4.safetensors`
     - `text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors`
   - Official `Lightricks/LTX-2.5` does not list a matching NVFP4 text encoder; the int8 ConvRot Gemma4 encoder is the matching low-VRAM official text encoder.

## 2026-08-13 VAE Debug Handoff

Resolved on `vae-decode-optimization`: direct VAE temporal overlap constants now match upstream, audio decode/mux is implemented, and Comfy-equivalent FP16 video VAE is the default runtime path. The 960x544x124 cat benchmark now matches warm Comfy performance: Comfy `150.26s`, direct `149.304s` after text conditioning, direct VAE decode `25.085s`.

Current saved latent and comparison assets live under:

- Spark: `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime`
- Share: `\\192.168.1.162\StoryStudioAssets\H3-output\h3-blackwell-runtime`

Key assets:

- Good sampled latent: `direct-cat-house-backflip-disco-shades-960x544-5s-latent.pt`
- Known-good same-latent Comfy/upstream VAE decode: `direct-cat-house-backflip-disco-shades-960x544-5s-upstream-comfy-tiled-decode.mp4`
- Current direct standalone FP32/Sage VAE decode candidate: `direct-cat-house-backflip-disco-shades-960x544-5s-direct-fp32-sage-vae-decode.mp4`

What is proven:

- The latent is good. Same latent decoded through Comfy/upstream VAE is visually clean.
- The sampler/model path is already exact against Comfy free-run parity for the fixed baseline.
- Direct VAE BF16/checkpoint-dtype loading was wrong. Loading direct VAE weights as FP32 reduced isolated decoder clip drift from roughly `mean=1.5e-3, max=1e-1` to roughly `mean=1e-6, max=6e-5`.
- Individual direct tiled VAE clips match upstream closely (`mean` around `4e-7`).
- Full direct `decode_temporal()` still differs from upstream at global frames `17, 34, 51, 68, 85, 102`, i.e. temporal join boundaries.
- Temporal assembly tracer shows pre-blend current chunks and overlap tails each match upstream, but blended join output differs hugely (`mean` around `0.15`, max around `5`) when comparing direct blend result to upstream blend result using their respective near-identical inputs.
- Direct and upstream `blend()` return identical results on the exact same inputs, so the remaining issue is likely an input/aliasing/dtype/shape subtlety at the temporal join, not the blend formula itself.

Relevant debug tools committed:

- `tools/decode_video_latent.py`
- `tools/compare_frame_dirs.py`
- `tools/compare_vae_decoder_clip.py`
- `tools/compare_vae_full_decode.py`
- `tools/compare_vae_tiled_clip.py`
- `tools/compare_vae_temporal_assembly.py`

Next VAE debugging steps:

- In `tools/compare_vae_temporal_assembly.py`, compare `prev_d` vs `prev_u` and `part_d` vs `part_u` after casting both pairs to a shared dtype and before blending. The current stats say they are close, but the blend of respective inputs explodes, which suggests a subtle shape/stride/dim broadcasting mismatch.
- Log `shape`, `stride`, `dtype`, `is_contiguous`, and `storage_offset` for `prev_*`, `part_*`, blend weights, and slices at `chunk_1_blended_at_17`.
- Try forcing `prev_d`, `prev_u`, `part_d`, and `part_u` to `.contiguous()` immediately before temporal blend in both direct and tracer paths.
- If that fixes it, patch direct `decode_temporal()` only. If not, compare exact selected slices (`prev[..., -9:, :, :]`, `part[..., :9, :, :]`) elementwise before and after multiplying weights.

Backends/experiments to evaluate later:

- Attention backend sweeps after VAE quality is fixed: `sage2`, `sage3`, `sdpa`, plus future FlashAttention-4/Sol-Attn/KJ candidates.
- Spectrum MiniMax H3 repo to inspect later: https://github.com/xmarre/ComfyUI-Spectrum-MiniMax-H3

## H3 Audio VAE Notes

Comfy source references:

- `comfy/ldm/minimax/audio_vae.py`: `MiniMaxH3AudioVAE`
- `comfy_extras/nodes_minimax_h3.py`: AV latent creation and audio guide encoding
- `comfy_extras/nodes_lt_audio.py`: `LTXVAudioVAEDecode`, which also handles nested AV latents and calls `audio_vae.decode(audio_latent)`
- `comfy_extras/nodes_audio.py`: generic audio save/preview and normalization helpers

Local checkpoint:

- Spark: `/home/daniel/aeon-spark-test/h3/comfy-models/vae/minimax_h3_audio_vae_fp32.safetensors`
- Runtime mount: `/vae/minimax_h3_audio_vae_fp32.safetensors`
- Shape inspection: 917 keys; top-level names include `dec_in_proj.*`, `decoder.*`, `encoder.*`, `latents_mean`, `latents_std`, `mean_proj.*`, `logs_proj.*`, and `pre_block.*`.

Architecture summary:

- This is not similar to the video ViT3D VAE. It is a DAC-lineage waveform encoder plus BigVGAN decoder.
- Latents are normalized `[B, 32, 2, T]`: 32 channels, 2 stereo channels, 40 latent frames/sec.
- Decoder output is stereo waveform `[B, 2, L]` at 32 kHz.
- One audio latent frame equals 800 samples (`32000 / 40`).
- Comfy decode path for nested H3 AV latent is effectively: select nested audio tensor, call `audio_vae.decode(audio_latent)`, then return waveform with sample rate 32000. VideoHelperSuite/Comfy audio save nodes use ffmpeg or audio helpers to save/mux.

Direct runtime implementation path:

- Change `sample_video_res_multistep()` to return both final video and final audio latents, or add a `sample_av_res_multistep()` wrapper that preserves the existing video-only API.
- Port `MiniMaxH3AudioVAE` from Comfy into a standalone `audio_vae_decoder.py`, starting decoder-only if we only need generated audio.
- Load `/vae/minimax_h3_audio_vae_fp32.safetensors` strictly in FP32, mirroring the video VAE precision lesson.
- Add a latent-only decode tool that accepts saved audio latent `[1,32,2,T]`, writes WAV/FLAC at 32 kHz, and optionally muxes with MP4 using ffmpeg.
- Validate first by capturing or decoding the same final audio latent through Comfy and direct, comparing waveform tensors before muxing.