2.2 KiB
NVFP4 Modulation Producer Fusion
Target
Fuse the two per-block AdaLN modulation operations into the exact NVFP4 activation scale and packing producer for attention QKV and MLP FC1. Retain the current Comfy Kitchen GEMMs and all BF16 numerical boundaries.
At the canonical 37,810-token shape, block 24 spends a median 10.47 ms on MSA modulation and 10.60 ms on MLP modulation. Production already mutates the two RMSNorm outputs in place, so this fusion targets memory traffic rather than a separate allocation.
Dataflow
Current:
- Read normalized BF16 input and write modulated BF16 output.
- Read modulated output to compute the global NVFP4 activation scale.
- Read modulated output again to produce packed E2M1 data and E4M3 scales.
- Execute the unchanged Comfy NVFP4 GEMM.
Prototype:
- Read normalized input, reproduce the reference modulation BF16 boundary, and reduce its absolute maximum without materializing the modulated tensor.
- Read normalized input again, reproduce the same BF16 values, and write exact QDATA/SFA directly into caller-owned buffers.
- Execute the unchanged Comfy NVFP4 GEMM.
Exactness Contract
- Preserve segment-specific shift and scale selection.
- Preserve
(x * (scale + 1)) + shiftoperation order and BF16 rounding. - Preserve Vortex's approximate-reciprocal NVFP4 encoding behavior.
- Produce byte-identical QDATA and E4M3 block scales.
- Produce bit-identical QKV/FC1 outputs before changing block dispatch.
Gate
- Adversarial modulation and quantization tests.
- Exact packed bytes and linear outputs for blocks 0, 24, and 49.
- Alternating module-forward benchmark in one loaded block.
- Full 50-block step and two-step trajectory only after all three blocks gain.
- Canonical 12-step trajectory before enabling the path by default.
Result
The producer is byte-exact for complete 37,810-row block 0, 24, and 49 inputs.
It takes 6.02 ms for block 24 QKV input. Integrated block medians improve by
0.28-0.79%, a warmed two-step trajectory improves by 0.52%, and the warmed
canonical 12-step run improves from 285.996 s to 284.408 s (0.56%). Video and
audio tensors are bit-identical. Spark deployments enable the path with
H3_NVFP4_MODULATE_FUSION=1.