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.

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.