Optimizations for MXFP8/NVFP4 dequantize kernels

[YigongQin]

Description

• Handle empty tensors in dequantize for CUDA graph compatibility
• Add swizzled scale support to the NVFP4 dequantize kernel, reusing the existing MXFP8 swizzle index computation
• Add C++ unit tests for both NVFP4 and MXFP8 dequantization (including swizzled scale variants)
• Fix to_cpu() and set_scale() in test infrastructure to correctly sync amax/scale for NVTE_NVFP4_1D_SCALING mode

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

Please list the changes introduced in this PR:

• Handle empty tensors in dequantize for CUDA graph compatibility — Early return when input has zero elements, avoiding kernel launches on empty tensors.
• Add GEMM-swizzled scale support to NVFP4 dequantize kernel — Template the kernel with WITH_GEMM_SWIZZLED_SCALES to support reading scales from swizzled layout, reusing the MXFP8 swizzle index computation.
• Add GEMM-swizzled scale support to MXFP8 dequantize kernel — Extend the MXFP8 dequantize kernel to handle swizzled scale inputs.
• Add C++ unit tests for NVFP4 dequantization — 21 tests for compact scales + 21 tests for swizzled scales, covering multiple sizes and output dtypes (fp32, bf16, fp16).
• Add C++ unit tests for MXFP8 dequantization with swizzled scales — New swizzled test suite for MXFP8.
• Fix to_cpu() to sync amax/scale for NVFP4 tensors — Previously only synced for NVTE_DELAYED_TENSOR_SCALING, causing the CPU reference to use stale amax=0.
• Fix set_scale() to work for NVFP4 tensors — Same condition fix, enabling the scale to be properly uploaded to GPU before quantization.
• Fix swizzled test ordering — Move from_cpu() before the FP4 data copy to prevent from_cpu() from overwriting the copied data with zeros.

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

[zianglih]

The following relevant unit tests passed on SM100 (with the drop optimize_for_gemm = False changes):

python3 -m pytest --tb=auto tests/pytorch/test_backward_override.py
python3 -m pytest --tb=auto tests/pytorch/test_sanity.py
python3 -m pytest --tb=auto tests/pytorch/test_cpu_offloading.py
PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto tests/pytorch/test_cuda_graphs.py
NVTE_TEST_NVINSPECT_ENABLED=1 NVTE_TEST_NVINSPECT_CONFIG_FILE=tests/pytorch/debug/test_configs/dummy_feature.yaml NVTE_TEST_NVINSPECT_FEATURE_DIRS=transformer_engine/debug/features PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto tests/pytorch/test_sanity.py

[zianglih]

After this PR, fwd is around 3%-4% faster for DeepSeek shape MoE:

# With the optimization
NVTE_BACKWARD_OVERRIDE=dequantized python benchmarks/linear/benchmark_grouped_linear.py --recipe mxfp8 --fwd-only
       m     k     n recipe  num_gemms  grouped_fwd_time_ms
0  16384  7168  2048  mxfp8          4             0.272829
1  32768  7168  2048  mxfp8          4             0.509788
2  65536  7168  2048  mxfp8          4             0.948633
3  98304  7168  2048  mxfp8          4             1.391146
0  16384  7168  2048  mxfp8          8             0.303238
1  32768  7168  2048  mxfp8          8             0.533896
2  65536  7168  2048  mxfp8          8             1.003446
3  98304  7168  2048  mxfp8          8             1.470030

# Without the optimization
git restore --source 77b8681de5cf -- transformer_engine/pytorch/module
NVTE_BACKWARD_OVERRIDE=dequantized python benchmarks/linear/benchmark_grouped_linear.py --recipe mxfp8 --fwd-only
       m     k     n recipe  num_gemms  grouped_fwd_time_ms
0  16384  7168  2048  mxfp8          4             0.282720
1  32768  7168  2048  mxfp8          4             0.526736
2  65536  7168  2048  mxfp8          4             0.982166
3  98304  7168  2048  mxfp8          4             1.451485
0  16384  7168  2048  mxfp8          8             0.313753
1  32768  7168  2048  mxfp8          8             0.551043
2  65536  7168  2048  mxfp8          8             1.040773
3  98304  7168  2048  mxfp8          8             1.527951

[greptile-apps]

Greptile Summary

This PR extends the MXFP8 and NVFP4 dequantize kernels to support GEMM-swizzled scale layouts by templating both kernels on WITH_GEMM_SWIZZLED_SCALES and reusing the existing gemm_swizzled_scale_idx helper, adds an empty-tensor early return in the dispatch layer for CUDA graph compatibility, removes corresponding Python workarounds that are now superseded, and adds C++ unit tests for both quantization formats across compact and swizzled paths.

Confidence Score: 5/5

Safe to merge; no P0/P1 issues found — all findings are P2 suggestions.

The kernel index math for both rowwise and colwise swizzled paths was verified to be within bounds. The empty-tensor guard is correctly placed at the dispatch level. Python workaround removals are consistent with the new dequantize capability. One P2 asymmetry exists in basic_linear.py where the removed condition was broader than the analogous sibling modules, but it does not affect current behavior.

transformer_engine/pytorch/ops/basic/basic_linear.py — removed condition was broader (backward_override is not None) than linear.py/layernorm_linear.py (== "dequantized").

Important Files Changed

Filename	Overview
transformer_engine/common/cast/dispatch/dequantize.cuh	Adds early-return guard for empty tensors (numel == 0), enabling CUDA graph compatibility. Clean and straightforward.
transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh	Adds WITH_GEMM_SWIZZLED_SCALES template parameter and num_scale_tiles_X to kernel; swizzle index reuses gemm_swizzled_scale_idx with transposed args for colwise. Removes prior assertion blocking swizzled input. Logic appears correct.
transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh	Adds WITH_GEMM_SWIZZLED_SCALES template and num_scale_tiles_X = DIVUP(Mread, 4) for swizzle; reuses mxfp8 swizzle helper. Index math checks out.
tests/cpp/operator/test_dequantize_nvfp4.cu	New test file with compact and swizzled NVFP4 dequantize tests. CPU reference computed from quantized data; swizzled test validates against compact path. Guards for Blackwell CC and FP4_TYPE_SUPPORTED are present.
tests/cpp/operator/test_dequantize_mxfp8.cu	Adds swizzled-scale test suite for MXFP8 (performTest_x1_swizzled) and zero-row test cases. Follows correct sequencing: copy data, swizzle scales, dequantize, compare.
transformer_engine/pytorch/ops/basic/basic_linear.py	Removes optimize_for_gemm=False block; the removed condition was broader (backward_override is not None) than the analogous linear.py/layernorm_linear.py condition (== "dequantized").
transformer_engine/pytorch/module/grouped_linear.py	Removes empty-tensor workaround in backward (now handled in C++ dispatch) and removes optimize_for_gemm=False guard for dequantized backward.
transformer_engine/pytorch/module/linear.py	Removes optimize_for_gemm=False for dequantized backward; now valid since dequantize supports swizzled scales.
transformer_engine/pytorch/module/layernorm_linear.py	Identical optimize_for_gemm=False removal as linear.py. Clean.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[nvte_dequantize] --> B[dequantize_helper dispatch]
    B --> C{numel == 0?}
    C -->|yes| D[early return — CUDA graph safe]
    C -->|no| E{scaling_mode}
    E -->|NVTE_MXFP8_1D_SCALING| F[mxfp8::dequantize]
    E -->|NVTE_NVFP4_1D_SCALING| G[nvfp4::dequantize]
    F --> H{with_gemm_swizzled_scales?}
    G --> I{with_gemm_swizzled_scales?}
    H -->|true| J[dequantize_mxfp8_kernel WITH_GEMM_SWIZZLED_SCALES=true]
    H -->|false| K[dequantize_mxfp8_kernel WITH_GEMM_SWIZZLED_SCALES=false]
    I -->|true| L[dequantize_fp4_kernel WITH_GEMM_SWIZZLED_SCALES=true]
    I -->|false| M[dequantize_fp4_kernel WITH_GEMM_SWIZZLED_SCALES=false]

Loading

_{Reviews (14): Last reviewed commit: "Remove unnecessary scale from NVFP4 C++ ..." | Re-trigger Greptile}

[zhongbozhu]

There is one edge case:

For MXFP8, When the input shape is like 64x64, it will produce scaling factor shape 64x2, but then zero padded to 128x4. We should be able to inject some very large random values in the padded region during malloc (because we don't use torch.zeros to malloc but torch.empty), and detect whether dequantize results is affected. If things work as expected, this line will be triggered

TransformerEngine/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh

Line 759 in a073ad5

// Zero out swizzled scales if padding is needed

and the dequantize numerics won't be affected.

For NVFP4, I think we optimize for GEMM (or swizzle fusion) is actually not enabled, same for the zero-out edge case handling logic?

TransformerEngine/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh

Line 1189 in a073ad5

NVTE_CHECK(!output->with_gemm_swizzled_scales, "Output must have scales in compact format.");

So there shouldn't be any unswizzle logic needed here?

[zhongbozhu]

For NVFP4, I believe currently only device-init grouped quantize with RHT has the swizzle fusion feature, so the scaling factor zero-out is the job of the dedicated swizzle kernel. So if we dequantize + unswizzle for NVFP4, the unswizzle logic might not be correct.

[YigongQin]

For both MXFP8 and NVFP4, the unit test logic is: 1. generate compact scales (or from quantization); 2. call nvte_swizzle_scaling_factors to swizzle compact scales; 3. compare results of nvte_dequantize with compact scales and swizzled scales. Quantize with swizzle fusion is never enabled for both MXFP8 and NVFP4

[timmoon10]

I'm of two minds about this:

• Logically, GEMM-optimized data is not guaranteed to support anything except GEMMs. Even if MXFP8 and NVFP4 dequant happens to support them, these are custom optimizations. Future recipes can not be expected to support dequantizing GEMM-optimzied data by default.
• It's a little pedantic to have edge-case logic that won't be triggered by any of our existing use-cases. Given how subtle this is, I worry about it becoming stale and distracting.

I think for now, this change is fine. However, if we encounter problems in a future recipe, we should reimplement it properly:

# LOGICALLY WRONG!
# Fails if we add a new recipe
if recipe.backward_override == "dequantized" and recipe.future_recipe():
    input_quantizer.optimize_for_gemm = False

# LOGICALLY RIGHT!
# Automatically handles new recipes
if recipe.backward_override == "dequantized" and not (
    recipe.float8_per_tensor_scaling()
    or recipe.float8_block_scaling()
    or recipe.mxfp8()
    or recipe.nvfp4()
):
    input_quantizer.optimize_for_gemm = False

CC @ptrendx @ksivaman @zhongbozhu

[zhongbozhu]

/te-ci pytorch L1

[timmoon10]

/te-ci core

[timmoon10]

LGTM, pending CI. These kernels will be very useful.

[timmoon10]

/te-ci