224 lines
11 KiB
Markdown
224 lines
11 KiB
Markdown
# 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.
|