h3-blackwell-runtime/PLAN.md

H3 Blackwell Runtime Plan

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 still blocked on locating/adding the standalone Sol-Attn source or package. It is not installed in the Spark image and is not present in this repository.
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.