[Common][PyTorch] Fuse the RHT into grouped NVFP4 quantize on non-SM100 architectures - #3317
[Common][PyTorch] Fuse the RHT into grouped NVFP4 quantize on non-SM100 architectures#3317davidkny22 wants to merge 3 commits into
Conversation
Outside the SM100 family split_quantize with RHT-enabled NVFP4 quantizers quantizes each split on its own, which runs two launches per split plus a separate Hadamard transform that materializes the transformed tensor to global memory. Complete the transposed output path in group_quantize_transpose with per-split direct stores, fold the 16-point random Hadamard transform into the columnwise read, and add nvte_group_quantize_with_colwise_rht so the dispatch can use one grouped launch pair per chunk. Chunk launch pairs by input bytes so the amax pass prefetches for the quantize pass. Drop the transposed staging buffers in shared memory, which the direct stores make dead. The transposed path requires 128-row-aligned splits and is reachable only through the columnwise RHT entry point, which refuses on the SM100 family, so generic grouped quantize keeps its existing refusal everywhere. Signed-off-by: David Kogan <davidkny22@gmail.com>
github-actions
Bot
added
the
community-contribution
Greptile SummaryThe PR fuses the columnwise random Hadamard transform into grouped NVFP4 quantization on supported non-SM100 architectures.
Confidence Score: 4/5The PR is not yet safe to merge because a later split requesting 2D quantization can still be routed through the first quantizer's non-2D grouped configuration. The added uniformity guard does not compare Files Needing Attention: transformer_engine/pytorch/csrc/extensions/cast.cpp Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["split_quantize with RHT"] --> B{"Non-SM100 fallback"}
B --> C{"Quantizers uniform and splits 128-row aligned?"}
C -->|No| D["Per-split quantize"]
C -->|Yes| E["Chunk splits by input bytes"]
E --> F["Grouped post-RHT amax"]
F --> G["Grouped rowwise quantize plus fused columnwise RHT"]
G --> H["Per-split NVFP4 outputs"]
Reviews (2): 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}
The grouped launch configures itself from quantizers.front() alone: with_post_rht_amax, rowwise_usage, columnwise_usage, row_scaled_nvfp4, nvfp4_4over6_mode, stochastic_rounding and the RHT sign mask. Grouped callers build one independent NVFP4Quantizer per split, and the Python-side cross-expert validator does not cover these fields for NVFP4Quantizer, so a heterogeneous list could silently take the grouped path and apply the first split's settings to the rest. Require every quantizer to agree with the first on these fields before taking the grouped path, mirroring the scaling-mode check already in this function. A mismatch falls through to the existing per-split loop, which reads each quantizer on its own. Added a launch-based regression test: two quantizers differing only in columnwise_usage do not produce byte-different rowwise output on the shapes tested, since rowwise computation does not depend on columnwise_usage and TE allocates every split's columnwise buffer regardless of that flag, so the test checks which kernel launches rather than the output bytes. Signed-off-by: David Kogan <davidkny22@gmail.com>
|
Confirmed and fixed in the latest commit. The exact mechanism Greptile named was not quite right (
Fix requires every quantizer in the list to agree with the first on the fields the grouped launch actually reads ( Added a regression test asserting on which kernel launches, not on output bytes, since a heterogeneous |
for more information, see https://pre-commit.ci
1 participant
Description
Follows #3265. That PR made
split_quantizewith RHT-enabled NVFP4 quantizers fall back to quantizing each split on its own outside the SM100 family, since the grouped Hadamard transform cast-fusion kernels are SM100 only. This makes that fallback fused.The fallback runs two launches per split plus a separate
HadamardTransformKernelthat writes the transformed tensor to global memory for the quantize pass to read back. At G=64 with 256-row splits that is 320 launches persplit_quantizecall and 11.125 bytes of traffic per input element.group_quantize_transposealready has most of what a fused version needs.MultiAmaxCastTransposeFusionArgscarriescolwise_amax_list,output_colwise_data_list,output_colwise_scale_inv_listandoutput_colwise_scale_stride, all marked "unused for rowwise only scaling", and the transposed store path is present but commented out withTODO(zhongbo): add back when transpose is supported. This completes that path with per-split direct stores instead of the shared memory staging it replaces, folds the 16-point random Hadamard transform into the columnwise read as an in-register butterfly, and addsnvte_group_quantize_with_colwise_rhtfor the PyTorch dispatch to call. The same G=64 case becomes 24 launches and 3.125 bytes per element.Direct stores make the transposed staging buffers dead, so they go, along with the commented-out stores that named them and two declarations (
tid_Y_t,out_mem_rowwise_data) that only fed those buffers. Dynamic shared memory drops from 25728 to 20608 bytes per 128-thread CTA, 3 resident CTAs per SM to 4.On GB10 (sm_121a) this is 2.22x median on
split_quantizeeager, 2.00x median of GPU work under graph capture, over 14 shapes. Atte.GroupedLinearlevel it is 1.00x to 1.17x of a forward and backward step, median 1.07x.The fused path is taken only when every split is a multiple of 128 rows. A split boundary is resolved once per 128-row chunk in the transposed direction, so a ragged split would write into the wrong buffer. Ragged cases decline to the existing per-split path. For MoE this means expert capacity has to be padded to 128 tokens to benefit.
The new launches are chunked by input bytes rather than by tensor count, at 8 MiB, so the amax pass prefetches for the quantize pass instead of being evicted from L2 before it reads the same data. Chunking by tensor count instead puts 128 MiB against a 24 MiB L2 at G=64, which costs most of the win.
Architecture behaviour
The transposed path is reachable only through the new entry point.
group_quantize_transposerefuses a columnwise output unlessnvfp4_colwise_rhtis set, and onlynvte_group_quantize_with_colwise_rhtsets it, so genericnvte_group_quantizekeeps today's refusal on every architecture. The new entry point additionally refuses onsm >= 100 && sm <= 110, the same bandsplit_quantize_nvfp4_impluses, so the SM100 family cannot reach the new code even through the C API.No existing entry point changes behaviour on any architecture.
Type of change
Changes
transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh: complete the transposed output path with per-split direct stores of columnwise data and scale factors, add the in-register 16-point RHT behind aCOLWISE_RHTtemplate switch, bind the columnwise pointers on the per-stage tensor switch, gate the transposed path on the columnwise RHT entry point and the architecture band, require 128-row-aligned splits, and drop the now-dead transposed staging buffers and the two declarations that only fed them.transformer_engine/common/cast/cast.cu:nvte_group_quantize_with_colwise_rht.transformer_engine/common/include/transformer_engine/cast.h: its declaration and docs.transformer_engine/common/common.h: two internal quantization config fields,nvfp4_colwise_rhtandnvfp4_rht_sign_mask_t, not exposed through the attribute API.transformer_engine/pytorch/csrc/extensions/cast.cpp: use the grouped path inside the existing non-SM100 branch when the quantizers carry post-RHT amax, both usages, no 2D quantization, no row-scaled NVFP4 and no 4over6, and every split is a multiple of 128 rows. Empty splits are filtered out of the grouped launch. Launch pairs chunk at 8 MiB of input.tests/pytorch/nvfp4/test_nvfp4_group_quantize.py: fused versus per-split equality across aligned, mixed, empty and ragged splits.Checklist:
No documentation changes needed beyond the header docs on the new entry point. On the last box: the added tests pass and nothing regresses, but
test_nvfp4_group_quantize.pyhas 120 pre-existing failures on main andtest_nvfp4_group_quantize_graph_safe.pyhas 528, unchanged by this PR. Numbers and cause below.Testing done
GB10 (DGX Spark), sm_121a, CUDA 13.0, driver 580.95.05. Built and gated at main @
d35eedf5f6d1cfba64b5e1dd4c8bfc3fc8214750, the #3265 merge commit, then re-verified asgit apply --checkclean againstaf1ed441255ec66e3555fd485dddb4344d7c495a, current main tip at submission, which adds one unrelated commit (#3276, NCCL EP) touching none of these six files. Built withNVTE_CUDA_ARCHS=121a. clang-format 18.1.6, cpplint 2.0.2 and black 24.4.2 clean on the changed files. No warning lines from any touched file, verified by rebuilding the affected translation units withninja -vrather than trusting thepip installlog, which hides per-file compiler output.Both arms live in one binary.
NVTE_NVFP4_DISABLE_GROUPED_RHTis read withstd::getenvat the dispatch on every call, so the A/B is one build and one process on identical operands. The new test uses it to compare the two routes, and it is the way out of the one-ulp difference below for a caller that needs output identical to the per-split path.Tests
test_rht_split_quantize_grouped_matches_unfusedcompares the two routes byte for byte over 12 cases: aligned, mixed-aligned, empty-front, empty-back, empty-mixed and two ragged split patterns, at two widths. It would have caught the empty-split bug I hit building this, where a leading zero-row split hashas_data()true with a null data pointer and aborted the grouped launch. It does not assert which route ran; the launch counts below do that.tests/pytorch/nvfp4/test_nvfp4_group_quantize.pygoes 120 failed, 453 passed, 270 skipped to 120 failed, 465 passed, 270 skipped, the 12 being the new cases.test_nvfp4_group_quantize_graph_safe.pyis 528 failed, 60 passed either way. Those failures are pre-existing on main: theoptimize_for_gemm=Truecases, where swizzled scale factor emission is gated on the same architecture band so outputs carry compact scale factors while the test swizzles the reference.Numerics
17 case classes compare the fused route against the per-split route byte for byte: aligned, ragged, mixed, empty-front, empty-middle, empty-back, and 64-way splits. Each case runs the fallback twice and the fused route once; the second fallback run is the control, since scale buffers are allocated at
roundup(rows, 128)and neither route writes the padding, so an uncontrolled comparison would report phantom mismatches that move with allocator state.16 of 17 are identical over the full buffer and over the defined region, on rowwise data, columnwise data, both scale planes and both amaxes.
One is not. A single byte differs,
columnwise_data[195, 78]of one split, reproducibly, with both arms individually deterministic. It is reduction order rather than a race.HadamardTransformKernelevaluates the 16x16 transform withmma_m16_n16_k16_b16_b16_b16_noacc, bf16 operands and fp32 accumulate; the fused path evaluates the same transform as a scalar fp32 butterfly. Both multiply by exactly representable values, so the only difference is the order of the fp32 additions in a 16-term reduction. For that element the input chunk spans 3.4375 down to 1.4e-07, the sum cancels heavily, the two orders land one bf16 ulp apart (-0.1000976562 against -0.1005859375), and the element sits on the FP4 rounding boundary for its block, so it moves one FP4 code. Block amax, both scale planes and the rowwise plane are unaffected.Rate over 1,073,741,824 columnwise elements across 5 shapes and 4 seeds each: zero further instances. Counting the one known case against everything compared, 1 in 1.1e9, bounded to one code step, no measurable norm or sign bias.
Identical on every tested input except that one-ulp reduction-order difference, which can move a single FP4 code. Not bit identical, and it cannot be while the two paths reduce in different orders.
Launch counts
Kernel launches per
split_quantizecall, from the profiler:Performance, split_quantize
Clocks locked at 3003 MHz, 9 windows, 20 iterations per window, medians, arms alternating window by window so they share clock history. Both arms also captured into CUDA graphs, which removes per-launch submission cost from both sides. Milliseconds per call.
Median 2.22x eager, range 1.38x to 2.95x. Median 2.00x of GPU work, range 1.21x to 2.92x. No shape regresses in either column. On a gate-only build, where both arms are the same fallback code, the same harness reads 1.00x median on 14 of 14 shapes, range 1.00x to 1.03x.
The gap between the columns is per-launch submission cost the fallback pays on 320 to 640 launches, removed by graph capture.
Performance, te.GroupedLinear
Forward and backward at the
te.GroupedLinearlevel, steady-state trainer settings (is_first_microbatch=False,fuse_wgrad_accumulation=True). Expert counts and per-expert token counts come from published MoE configurations attokens x top_k / experts. RecipeNVFP4BlockScaling(disable_stochastic_rounding=True), since the SR kernel uses acvt.rsinstruction sm_121 does not have. 7 windows, 5 iterations, medians. Milliseconds per step.Median 1.07x, range 1.00x to 1.17x. Noise floor on a gate-only build with both arms identical: 1.00x eager on all six shapes, 0.98x to 1.01x graph.
The ragged case is the
%128decline, and it reads 1.00x with identical launch counts on both arms. GPT-OSS declines in the forward for an unrelated reason:split_quantizealready refuses NVFP4 kernel fusion when the input's inner dimension is not a multiple of 128, and 2880 is not, so only its backward fuses.The share of the step is why the module number is smaller than the kernel number: on Mixtral fc1 the RHT quantize path is 25.1% of device time in the fallback and 14.2% after; on Qwen3 fc1 the GEMM is roughly 85% of the step and the RHT path 7.4%, which caps anything done here at about 1.08x on that shape.
Open questions
HadamardTransformKernelbit for bit? The one-ulp difference is inherent to evaluating the same transform in a different reduction order, so if exact match is required this needs a different approach.kChunkBytesbe derived from the device's L2 size rather than fixed at 8 MiB? I only have one L2 size to size it against.%128alignment requirement acceptable for how grouped quantize gets called in practice, or is ragged-split support worth pursuing?cc @zhongbozhu @Oleg-Goncharov