[Common] Optimize NVFP4 RHT cast-fusion kernels by CalebDu · Pull Request #3740 · NVIDIA/TransformerEngine

CalebDu · 2026-10-26T06:52:57Z

Description

Performance optimization of the NVFP4 RHT cast-fusion kernels on SM100 (dense, grouped, and graph-safe grouped quantization). Output numerics are unchanged.

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Remove register spills in the RHT cast-fusion kernels (column-quant epilogue processed in chunks; per-group tensors rebuilt per tile).
Use packed mul.rn.f32x2 (FMUL2) for scale computation and data scaling.
256-bit row-quant store layout.
Graph-safe grouped kernel: stage group metadata in shared memory and cache the current tensor range instead of re-reading global memory per tile.
Cheaper tile-scheduler index math.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

…, graph-safe metadata Grouped (graph-safe and non-graph-safe) and dense RHT cast-fusion kernels on SM100: - Row-quant store: port the CUTLASS 256-bit row store layout (S2RWarpLayout<_2,_16>, 64 elements per thread, SM100_STORE_256bit_CACHE_NOALLOCATION) to cut STG count. - Use packed mul.rn.f32x2 (FMUL2) for scale computation and data scaling via mul_f32x2_by_scalar in hadamard_transform_utils.cuh (bit-identical to scalar FMUL). - Graph-safe: stage group offsets / first dims in shared memory and cache the current tensor range in registers instead of re-reading global memory per tile. - Tile scheduler: 32-bit division instead of FastDivmodU64, per-work-unit const locals. - Col-quant epilogue: process the accumulator in 64-column TMEM chunks (load, fence, quantize per chunk) to bound register pressure; removes all LDL/STL in the previously spilling SR configurations. - Dense: pass scheduler constants (k_tile_max, tiles_in_m, tiles_in_n) as kernel params. Validated on GB300 (te-nt4 base): NVFP4 RHT/group-quantize test suites pass, dense outputs bit-identical to baseline for all row/col x SR x fast-math x swizzle configs, 0 LDL/STL in all dense and grouped instances. Example (ncu, 8 x [10240,10240], row+col, SR, FM=false): graph-safe 928.2 -> 717.0 us; dense 8192x10240 SR row+col FM=false 120.0 -> 87.6 us. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caleb Du <cadu@nvidia.com>

…kernels Keep only the group index and global scales live across k-tiles in the column-quant epilogue and rebuild the per-group output/scale tensors per tile. This lowers register pressure and removes the remaining register spills in the grouped (graph-safe and non-graph-safe) RHT cast-fusion kernels. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caleb Du <cadu@nvidia.com>

for more information, see https://pre-commit.ci

greptile-apps · 2026-10-26T06:59:07Z

Confidence Score: 4/5

[Medium impact] The changes are limited to SM100 NVFP4 RHT cast-fusion kernel paths.

The PR still needs the equal-shaped row-only quantization fix before merging.

Findings

Equal-shaped rowwise quantization fails ▶

Summary

This PR reorganizes the SM100 NVFP4 RHT cast-fusion kernels to reduce register use.

RHT cast-fusion kernels process quantization work in smaller chunks.
Graph-safe grouped quantization reuses staged group metadata.

No new findings were accepted.

_{Reviews (4) · Last reviewed commit: "Merge branch 'main' into caleb/fix_hadam..." · Reviewed by Greptile}

greptile-apps · 2026-10-26T06:59:10Z

+      if (i < static_cast<int>(num_tensors)) {
+        shared_storage.group_first_dims[i] = static_cast<int32_t>(__ldg(first_dims + i));


Equal-shaped rowwise quantization fails

The new staging loop reads first_dims even when only rowwise output is requested. For equal-shaped groups, GroupedTensor uses an empty first_dims, so this pointer is null. The old row-only path never read it; the new path hits an invalid device-memory read before quantization starts.

Load group_first_dims only when columnwise output needs it, and derive equal-group dimensions from the logical shape instead of reading the optional array.

Use 8 column-quant warps with lighter register allocation and batch Philox generation for the stochastic-rounding paths. Outputs are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caleb Du <cadu@nvidia.com>

…nels Adjust per-role register allocation for the column/row quantization warps to avoid register spills with older CUDA toolkits, and reduce the dense kernel scheduler pipeline depth. Outputs are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caleb Du <cadu@nvidia.com>

ksivaman · 2026-10-29T05:46:01Z

/te-ci

0z5a · 2026-10-31T19:29:34Z

I tested this PR on a B300 (sm_103). TE was built for 103a from main 63451d6b, once as-is and once with this PR (5e7536a1) applied, CUDA 13.3.

Numerics. For each path I dumped every output buffer of the NVFP4 quantization with RHT, rowwise + columnwise, from the same inputs and seeds. Each build gives byte-identical buffers on repeated runs.

path	shapes	without stochastic rounding	with stochastic rounding
dense (`NVFP4Quantizer(x)`)	(128, 128) to (8192, 7168), 6 shapes	identical	identical
grouped (`tex.split_quantize`)	4 split configurations, including an empty split	identical	rowwise data differs in the 3 configurations with 128-row-aligned splits, columnwise data identical
graph-safe grouped (`tex.group_quantize`)	3 configurations with 128-row-aligned splits	identical	rowwise data differs, columnwise data identical

The fourth grouped configuration (8 × 320 rows) is not 128-row aligned. It runs group_rht_gemm_device, which this PR doesn't change, so it serves as a control rather than as coverage. Its data matches; the rowwise scales differ only in the padding rows past row 320.

The amax of the empty split differs because neither build writes it. That predates this PR.

So with stochastic rounding, the grouped and graph-safe grouped kernels produce different rowwise bytes than main. The aggregate error statistics I measured are the same in both builds: rowwise dequantization rms error / rms x 0.1373–0.1374, and mean bias / |x| within ±1.3e-4, the same range as without stochastic rounding. I did not check that the random bits stay independent per element. Without stochastic rounding, every output is bit-identical.

tests/pytorch/nvfp4/ with this PR (the RHT exact, grouped, graph-safe grouped, SR and RHT swizzle-fusion files): 1555 passed, 0 failed, 274 skipped.

Kernel time (Nsight Compute gpu__time_duration, median of 3 launches, µs; cases matched to kernels by launch order):

case	main	this PR
dense 4096×2048	16.9	16.3
dense 8192×7168	45.0	39.2
dense 2560×2048	16.0	14.8
grouped, 8 × 1024 rows, N=4096	31.7	27.8
grouped, 6 splits (256/0/768/1280/512/128), N=7168	29.3	25.1
graph-safe grouped, 8 × 1024 rows, N=4096	39.3	27.4
graph-safe grouped, same 6 splits, N=7168	37.1	26.8
`group_rht_gemm_device` (not changed by this PR), 8 × 320 rows	17.2	17.2

Registers per thread go from 128 to 96 in the three changed kernels.

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Oct 26, 2026

Caleb Du and others added 2 commits October 25, 2026 23:54

CalebDu force-pushed the caleb/fix_hadamard branch from a15e37a to 3181072 Compare October 26, 2026 06:54

[pre-commit.ci] auto fixes from pre-commit.com hooks

5ae9f2e

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Oct 26, 2026

View reviewed changes

Caleb Du and others added 3 commits October 27, 2026 06:07

Merge branch 'main' into caleb/fix_hadamard

5e7536a

ksivaman self-assigned this Oct 29, 2026

Navigation Menu

Repository navigation

[Common] Optimize NVFP4 RHT cast-fusion kernels - #3740

[Common] Optimize NVFP4 RHT cast-fusion kernels#3740
CalebDu wants to merge 6 commits into
NVIDIA:mainfrom
CalebDu:caleb/fix_hadamard

CalebDu commented Oct 26, 2026

Uh oh!

greptile-apps Bot commented Oct 26, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot Oct 26, 2026

Uh oh!

ksivaman commented Oct 29, 2026

Uh oh!

0z5a commented Oct 31, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		if (i < static_cast<int>(num_tensors)) {
		shared_storage.group_first_dims[i] = static_cast<int32_t>(__ldg(first_dims + i));

Conversation

CalebDu commented Oct 26, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 4/5

Findings

Summary

Uh oh!

greptile-apps Bot Oct 26, 2026

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Oct 29, 2026

Uh oh!

0z5a commented Oct 31, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

greptile-apps Bot commented Oct 26, 2026 •

edited

Loading