[Common] row-scaled nvfp4 path: fuse row/col amax into a single TMA-tiled kernel - #3454
[Common] row-scaled nvfp4 path: fuse row/col amax into a single TMA-tiled kernel#3454cael-ling wants to merge 2 commits into
Conversation
The row-scaled path ran two amax kernels; the columnwise one read global memory column-major (uncoalesced) and dominated runtime. Compute both directions in one kernel that streams 128x128 chunks through shared memory via TMA and reduces columns from SMEM. Gated by fused_amax_supported (BF16, 128-aligned dims); other cases keep the two-kernel path. Set NVTE_NVFP4_FUSED_AMAX=0 to force fallback. Signed-off-by: Cael Ling <caell@nvidia.com>
github-actions
Bot
added
the
community-contribution
for more information, see https://pre-commit.ci
Greptile SummaryThe PR adds an SM100 TMA-tiled kernel that computes rowwise and columnwise NVFP4 amax metadata in one pass, with runtime fallback and an environment kill switch.
Confidence Score: 4/5The no-op state-preservation regression should be fixed before merging because skipped quantization calls can erase both amax buffers. The fused wrapper unconditionally queues zeroing of the amax outputs, while its kernel subsequently returns on the noop flag, leaving previously preserved quantization metadata cleared. Files Needing Attention: transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
Q[Row-scaled NVFP4 quantization] --> C{Columnwise output requested?}
C -- No --> R[Standalone rowwise amax]
C -- Yes --> S{BF16, 128-aligned, buffers present, fusion enabled?}
S -- No --> F[Standalone rowwise and columnwise amax]
S -- Yes --> Z[Clear both amax buffers]
Z --> K[Fused TMA row and column reduction]
K --> N{noop equals 1?}
N -- Yes --> E[Kernel returns; cleared metadata remains]
N -- No --> A[Atomically publish row and column maxima]
Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| NVTE_CHECK(output->amax.numel() == rows, "Fused rowwise amax must have ", rows, | ||
| " entries, got ", output->amax.shape, "."); | ||
| NVTE_CHECK_CUDA(cudaMemsetAsync(row_amax_ptr, 0, rows * sizeof(float), stream)); | ||
| } | ||
| if (do_col) { | ||
| NVTE_CHECK(output->columnwise_amax.numel() == cols, "Fused columnwise amax must have ", cols, | ||
| " entries, got ", output->columnwise_amax.shape, "."); | ||
| NVTE_CHECK_CUDA(cudaMemsetAsync(col_amax_ptr, 0, cols * sizeof(float), stream)); | ||
| } |
There was a problem hiding this comment.
When an eligible fused quantization call has noop[0] == 1, this wrapper clears both amax buffers before the kernel returns without writing, causing skipped graph-replay or update calls to destroy the previously preserved quantization metadata.
Knowledge Base Used: Native GEMM and quantization kernels
1 participant
Description
The NVFP4 row-scaled path that was originally proposed in #2931 computes per-row and per-column amax with two separate kernels. This PR fuses both directions into a single kernel that streams 128x128 chunks through shared memory via TMA (coalesced loads) and does the column reduction from SMEM. Amax is an exact max reduction, so results are byte-identical to the two-kernel path. The fused path is used only when the quantize call needs both directions (rowwise + columnwise amax) on a BF16 input with 128-aligned dims; the kernel then produces both amaxes in one pass. Any other case keeps the original two kernels.
NVTE_NVFP4_FUSED_AMAX=0forces the fallback at runtime.Type of change
Changes
compute_fused_amax_kernel(rowwise + columnwise) and its host wrappersfused_amax_supported/compute_fused_amaxinquantize_transpose_nvfp4.cuh.dispatch/quantize.cuh(fwd and bwd) when supported, else fall back to the standalone amax kernels.NVTE_NVFP4_FUSED_AMAXkill switch (default on).Performance
Full-quantize median latency, fused (fused amax + cast) vs (row-wise & columnwise amax + cast), the cast kernel is byte-identical so the delta is the amax step:
Reproduce
Single Blackwell (SM100) GPU.
pytest tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py -k "row_scaled and both_directions"passes; rowwise/columnwise amax and
qx/qx_tmatch the reference exactly.NVTE_NVFP4_FUSED_AMAXenv var (default enabled): set to0to disable the fused path at runtime and fall back to the two standalone amax kernels, no rebuild needed.Checklist: