h3-blackwell-runtime/NVFP4_STREAMING_DESIGN.md

NVFP4 Streaming Design

Decision

Comfy Kitchen 0.2.31 cannot consume incrementally produced NVFP4 tiles. Its Python path allocates complete activation QDATA, complete swizzled block scales, and a complete BF16 output before calling cublas_gemm_blockwise_fp4 with fixed full-tensor pointers. The registered operator has no tile callback, producer interface, output stride, user output, or custom epilogue argument.

True activation-materialization removal therefore requires an owned CUTLASS or CuTe DSL mainloop. Wrapping scaled_mm_nvfp4 cannot provide it.

Available Foundation

The Spark image contains:

• CUDA 13.0.88.
• NVIDIA CUTLASS DSL 4.6.2.
• CUDA Python bindings.
• Working SM121 block-scaled E2M1/E4M3 MMA support.

The official CUTLASS 4.6 SM120/SM121 persistent cooperative block-scaled GEMM was compiled and executed on GB10 with E2M1 A/B, E4M3 scales, FP32 accumulation, and BF16 output. A 1024x1024x1024 case passed its reference check at 16.9664 us and approximately 126.57 TFLOP/s.

The public kernel still takes complete global A, B, SFA, and SFB tensors and uses separate TMA descriptors for all four. CUTLASS does not expose a ready BF16-to-NVFP4 input producer callback.

Numerical Constraint

H3's activation quantization uses:

1. One exact global BF16 absolute maximum.
2. One FP32 tensor scale derived from that maximum.
3. One E4M3 block scale per 16 activation values.
4. E2M1 values quantized using both scales.

No activation tile can be packed reference-exactly until the full activation's global reduction is complete. The minimum exact design therefore has two phases:

1. Reduce the complete BF16 activation to one FP32 scale.
2. Produce packed tiles and consume them in block-scaled MMA.

A future cooperative mega-kernel could combine the phases around a grid-wide barrier, but that is not the first prototype. Programmatic dependent launch can overlap independent weight preparation but cannot transfer shared-memory tiles between kernels.

Reuse Constraint

The SM121 block-scaled kernel uses cluster shape 1x1x1; there is no TMA multicast path for sharing one packed activation tile among output-N CTAs. Replacing global packed A with an ordinary per-CTA software producer would make every output-N CTA reread BF16 A and repeat scale/pack work. BF16 A is materially larger than packed E2M1 plus E4M3 scales, so this can lose despite removing the initial packed-tensor write.

The owned kernel must therefore test an N-group schedule where one CTA retains one or more packed A K-tiles while accumulating multiple output-N tiles. The number of simultaneous N tiles is constrained by accumulator registers and shared memory. This reuse factor is a first-class tuning parameter.

Prototype Sequence

P0: Owned Prepacked Baseline

Port the official persistent cooperative SM120/SM121 CuTe DSL kernel behind an experimental Vortex entry point. Match Comfy Kitchen for the four H3 projection shapes using existing packed activation and weight tensors.

Acceptance:

• Exact BF16 output versus scaled_mm_nvfp4.
• H3 widths 5376, 7168, and 14336 plus QKV output width 21504.
• No regression beyond measurement noise before changing the A producer.

Current P0 result: real 128-row H3 tiles execute through the CuTe kernel. Raw block-scaled BF16 output is bit-exact for QKV, attention output, and FC1. P0 is not complete:

• The experimental epilogue now applies H3's FP32 global-scale product before BF16 conversion. QKV, attention output, and FC1 are bit-exact.
• FC2 at K=14336 differs even before global scaling (max_abs=8192, mean 0.015625). The correct alpha epilogue reduces final error to mean 2.42e-5, but bit equality still fails with max_abs=16. Its cuBLAS path uses a different reduction/Stream-K policy; CUTLASS K tiles 128 and 256 produce the same non-reference result.
• The stock CuTe Float32 output mode fails the official example's own GB10 reference check and cannot be used as an accumulation oracle.

Do not begin P1 until the FC2 reduction contract is exact or FC2 is deliberately kept on the existing cuBLAS fallback. See benchmarks/gb10-cute-p0-h3-summary.json.

Decision: FC2 remains explicitly on the existing cuBLAS path. P1 and later streaming work target QKV, attention output, and FC1 only. This preserves the reference FC2 Stream-K reduction order and therefore full-model bit parity while the owned producer-consumer path is developed independently.

P1: Single-Tile Software A Producer

Replace the A and SFA TMA loads for one fixed output tile with a software warp that loads BF16 A, generates E4M3 scales and E2M1 values into the existing staged shared-memory layouts, and commits a producer-consumer pipeline stage. Keep prepacked B/SFB and the existing MMA consumer and BF16 epilogue.

Acceptance:

• Packed bits and scale bytes agree with Comfy for every consumed tile.
• GEMM output is bit-exact for adversarial and randomized tile inputs.
• No global QDATA or SFA allocation in the captured kernel.

Current P1 checkpoint: the fixed 128-row producer-consumer path is exact. The existing 32-thread DMA warp loads four BF16 rows per lane, writes E2M1 and E4M3 directly into the staged sA/sSFA layouts, and lets retained B/SFB TMA completion publish each stage to the unchanged MMA consumer. The kernel takes BF16 A and one tensor scale; it does not take or allocate global activation QDATA or SFA.

Every 128-K tile has zero differing packed or scale bytes for real QKV, attention-output, and FC1 activations. Their complete 128-row GEMM outputs are also bit-exact after the alpha-before-BF16 epilogue. Exact Comfy compatibility requires PTX rcp.approx.ftz.f32, because Comfy Kitchen builds the quantizer with --use_fast_math; ordinary FP32 division changes E2M1 decisions at midpoints. CUTLASS DSL 4.6.2 also requires a static contiguous view for vector FP4 stores, as lowering the same store through a rank-2 dynamic layout aborts MLIR construction. Evidence is in benchmarks/gb10-cute-p1-stream-a-summary.json.

P1 is not deployment-complete. The next work is canonical row-count and padding support and an owned runtime module rather than validator-time patching of the CUTLASS example. FC2 continues to reject streaming and remains on Comfy/cuBLAS.

P2: N-Group Reuse

Accumulate multiple output-N tiles per produced A tile. Sweep N-group size and mainloop stages while recording registers, occupancy, achieved tensor throughput, BF16 bytes read, and total projection latency.

Reject the no-materialization design if repeated BF16 reads or reduced occupancy make it slower than the prepacked baseline. In that case, test a bounded global ring buffer of packed tiles as the producer-consumer compromise.

Current P2 result: reject the one-producer-per-output-N-CTA schedule. At 128 rows, streamed execution is 9.14x, 12.98x, and 12.44x slower than the complete scale+pack+GEMM reference for QKV, attention output, and FC1. Producer overhead scales at approximately 0.56-0.60 us per (N tile, K tile). Merely matching the removed activation quantization cost would require theoretical N reuse factors of 78, 28, and 102 respectively. Those factors are not practical through duplicate accumulator state, especially with the existing 232-register consumer requirement.

The next prototype is therefore a bounded global packed-tile ring or persistent work queue: produce each (M,K) A tile once, let all required N consumers reuse it, then recycle the slot. This deliberately restores bounded packed global traffic while avoiding complete activation materialization. See benchmarks/gb10-cute-p1-stream-a-timing-summary.json.

The bounded-ring capacity sweep is complete. A caller-owned _into CUDA producer writes exact QDATA and SFA into reusable buffers, and one full-activation scale is reused by every chunk. Capacities 128, 512, 1024, 2048, and 4096 were measured; 2048 rows is the selected point. It uses 6.19 MB for QKV/FC1 and 8.26 MB for attention output. Smaller chunks leave launch overhead, while 4096 rows regresses.

At 2048 rows, measured chunk latency improves by 10.7% for QKV, 14.4% for attention output, and 11.3% for FC1. Modeled canonical totals improve by 10.1-12.5% after charging one full-activation scale reduction. All packed bytes, scale bytes, and BF16 outputs are exact. QKV, attention output, and FC1 therefore advance to complete real-projection validation; FC2 remains on Comfy/cuBLAS. See benchmarks/gb10-cute-p2-ring-capacity-summary.json.

Complete 37,810-row projection parity now passes for QKV, attention output, and FC1 in blocks 0, 24, and 49. All 19 chunks, including the final 946-row chunk, are bit-exact with equal full checksums. QKV has a stable approximately 32 ms ring time and improves in every tested block, so it advances to opt-in resident runtime validation. Attention output timing is mixed. FC1 timing is unstable under duplicate-model unified-memory pressure, including one container stop, despite exact outputs. Do not approve those two roles from this harness. See benchmarks/gb10-cute-p2-ring-full-projection-summary.json.

The opt-in QKV-only integration passes numerical validation but fails the single-model performance gate. Alternating baseline/ring calls in one loaded block are bit-exact for blocks 0, 24, and 49, while the 2048-row schedule is 0.52-0.81% slower at block level. Capacity checks from 3072 through a full 37888-row workspace do not recover a gain; 4096 rows is closest at 0.24% slower. The reusable output also requires serialized model requests because the next projection reuses its storage after GPU work is enqueued. Keep this backend disabled. The next valid prototype must fuse chunk launches through a persistent work queue or replace the scheduler before repeating this block gate; trajectory validation is intentionally skipped. See benchmarks/gb10-cute-qkv-runtime-block-gate-summary.json.

P3: Full H3 Projection

Support canonical H3 M/N/K shapes, padding, current weight/block-scale layout, and token-row slicing. Integrate through Nvfp4Linear behind an opt-in feature flag.

P4: Exact Epilogue

After the owned GEMM is exact, add the explicit BF16 conversion, gate multiply, residual add, and final BF16 rounding sequence. Validate this independently from input streaming.

Validation Ladder

Every phase must pass, in order:

1. Adversarial scale, E4M3, E2M1, signed-zero, midpoint, NaN, and infinity cases.
2. Randomized packed-tile and GEMM comparisons.
3. Real QKV, output, FC1, and FC2 projections at blocks 0, 24, and 49.
4. Complete block equality.
5. All 50 blocks in a denoiser step.
6. Two-step video/audio trajectory equality.
7. Canonical 12-step video/audio checksum equality.
8. Full block Nsight recapture and peak-memory comparison.

No streamed path becomes a deployment default before all applicable exactness gates pass.