[PyTorch] NVFP4: emit GEMM-swizzled scales in the non-RHT 2D quantize kernel

[cael-ling]

Description

Support the CUDA graph capture when enabling NVFP4 weight caching in #3093

Scope: the non-RHT 2D quantize path on 128-aligned shapes, which is the default cached-weight case. Other paths (1D weight scaling, non-128-aligned, etc) keep the compact + post-quantize swizzle fallback and are left for follow-ups.

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:

• Change A
• Change B

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

[greptile-apps]

Greptile Summary

This PR extends the NVFP4 2D quantize kernel to emit GEMM-swizzled scale factors directly (no standalone swizzle pass) for 128-aligned shapes on the non-RHT path, fixing CUDA graph capture issues when NVFP4 weight caching is enabled. It also corrects a latent bug where the backward-path fallback kernel hardcoded compact-scale output even when the tensor was pre-allocated expecting swizzled scales.

• New swizzle.cuh: adds gemm_swizzled_scale_idx device function — mathematically equivalent to the existing mxfp8::swizzle::gemm_swizzled_scale_idx, verified against the Python reference in nvfp4_utils.py.
• Kernel dispatch (quantize_transpose_nvfp4.cuh): adds WITH_GEMM_SWIZZLED_SCALES template parameter to the 2D kernel; NVTE_CHECKs guard that this mode is only used on the 2D path with 128-aligned dims; rowwise and columnwise scale writes branch on the template param.
• C++ quantizer (quantizer.cpp): nvfp4_emits_gemm_swizzled_scales helper unifies eligibility logic for both RHT and non-RHT 2D paths; both create_tensor and convert_and_update_tensor now agree on whether to set with_gemm_swizzled_scales.
• Fallback fix (quantize.cuh): forward and backward fallback kernel calls now propagate output_tensor->with_gemm_swizzled_scales instead of the previously hardcoded false.

Confidence Score: 5/5

This PR is safe to merge. The swizzle index formula is mathematically identical to the existing mxfp8 reference and the Python test utility, buffer sizes are unchanged for the gated 128-aligned shapes, and NVTE_CHECKs correctly constrain the new path to the 2D kernel on 128-aligned dims only.

The swizzle formula in the new swizzle.cuh expands to the same expression as mxfp8::swizzle::gemm_swizzled_scale_idx and the Python reference used by the fidelity tests. Namespace resolution is correct (function is in transformer_engine::dispatch::nvfp4::swizzle, accessed from the nvfp4 namespace as swizzle::). Buffer sizes for compact and swizzled layouts are equal for 128-aligned shapes, so no allocation size change is required. The fallback path fix in quantize.cuh (removing the hardcoded false for the backward path) is consistent with the forward path and with how the fallback kernel already supports swizzled scales. The two new test files provide direct fidelity, pointer-stability, and CUDA-graph-replay coverage for the fix.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/nvfp4/swizzle.cuh	New device helper for the GEMM-swizzled scale index; formula verified equivalent to mxfp8::swizzle::gemm_swizzled_scale_idx and the Python reference in nvfp4_utils.py.
transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh	Adds WITH_GEMM_SWIZZLED_SCALES template parameter; NVTE_CHECKs correctly enforce 2D-path and 128-aligned-shape preconditions; rowwise and columnwise scale write paths are correctly branched; col_length_t and SFs_per_row derivations are exact for 128-aligned dims.
transformer_engine/pytorch/csrc/quantizer.cpp	nvfp4_emits_gemm_swizzled_scales consolidates eligibility for RHT and non-RHT 2D paths; both create_tensor and convert_and_update_tensor updated consistently; is_eligible_for_2d_swizzle_fusion correctly mirrors the kernel's 128-alignment and sm_arch gate.
transformer_engine/common/cast/dispatch/quantize.cuh	Forward and backward fallback paths now correctly read with_gemm_swizzled_scales from the tensor instead of hardcoded false; the fallback kernel already supports swizzled scales via scale_factor_swizzled_offset.
transformer_engine/pytorch/csrc/common.h	Adds is_eligible_for_2d_swizzle_fusion public static declaration; clean addition consistent with existing is_eligible_for_rht_cast_fusion pattern.
tests/pytorch/nvfp4/test_nvfp4_2d_quantize_swizzle_fusion.py	New fidelity tests covering byte-equality of swizzled vs. compact+post-swizzle SF; shape-gate test distinguishes eligible from ineligible shapes; end-to-end test verifies both in-kernel and fallback paths yield swizzled SF.
tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py	New integration tests: pointer-stability, numerics, and CUDA-graph-replay tests for all major TE module types; the graph-replay test directly reproduces the original hang/corruption scenario.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["quantizer(x) called\noptimize_for_gemm=True\nwith_2d_quantization=True"] --> B["nvfp4_emits_gemm_swizzled_scales()"]
    B --> C{with_rht?}
    C -- Yes --> D["is_eligible_for_rht_cast_fusion(shape)"]
    C -- No --> E{with_2d_quantization\n&& !row_scaled\n&& !4over6?}
    E -- No --> F["with_gemm_swizzled_scales = false\n compact scales only"]
    E -- Yes --> G["is_eligible_for_2d_swizzle_fusion(shape)\n(rows%128==0 && cols%128==0 && sm100-110)"]
    D -- eligible --> H["with_gemm_swizzled_scales = true"]
    D -- ineligible --> F
    G -- eligible --> H
    G -- ineligible --> F
    H --> I["create_tensor allocates scale buffer"]
    F --> I
    I --> J["quantize_fwd_helper / quantize_bwd_helper"]
    J --> K{use_optimized_kernel?}
    K -- Yes --> L["nvfp4::quantize_transpose (2D kernel)"]
    K -- No --> M["quantize_transpose_vector_blockwise_fp4\nswizzled_scale=with_gemm_swizzled_scales"]
    L --> N{WITH_GEMM_SWIZZLED_SCALES?}
    N -- true --> O["gemm_swizzled_scale_idx()\nwrite SF at swizzled offset in-kernel"]
    N -- false --> P["compact SF write + post-quantize swizzle pass"]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["quantizer(x) called\noptimize_for_gemm=True\nwith_2d_quantization=True"] --> B["nvfp4_emits_gemm_swizzled_scales()"]
    B --> C{with_rht?}
    C -- Yes --> D["is_eligible_for_rht_cast_fusion(shape)"]
    C -- No --> E{with_2d_quantization\n&& !row_scaled\n&& !4over6?}
    E -- No --> F["with_gemm_swizzled_scales = false\n compact scales only"]
    E -- Yes --> G["is_eligible_for_2d_swizzle_fusion(shape)\n(rows%128==0 && cols%128==0 && sm100-110)"]
    D -- eligible --> H["with_gemm_swizzled_scales = true"]
    D -- ineligible --> F
    G -- eligible --> H
    G -- ineligible --> F
    H --> I["create_tensor allocates scale buffer"]
    F --> I
    I --> J["quantize_fwd_helper / quantize_bwd_helper"]
    J --> K{use_optimized_kernel?}
    K -- Yes --> L["nvfp4::quantize_transpose (2D kernel)"]
    K -- No --> M["quantize_transpose_vector_blockwise_fp4\nswizzled_scale=with_gemm_swizzled_scales"]
    L --> N{WITH_GEMM_SWIZZLED_SCALES?}
    N -- true --> O["gemm_swizzled_scale_idx()\nwrite SF at swizzled offset in-kernel"]
    N -- false --> P["compact SF write + post-quantize swizzle pass"]

Loading

_{Reviews (3): Last reviewed commit: "Merge branch 'main' into fix/nvfp4-weigh..." | Re-trigger Greptile}

[vthumbe1503]

LGTM. Thanks for the PR @cael-ling. Great work!

For more context, we enabled weight preswizzling in the previous PR by setting optimize_for_gemm in the weight quantizer to be True for unquantized primary weights.

However when weight caching + cuda graph is enabled, this led to a crash with nvfp4 recipe. Previously, NVFP4 quantization with weight preswizzling first quantizes the weights and then modifies the scales to point to a new swizzled buffer. This didnt work, since cuda graph capture on the cached weights wont be aware of the scale buffer being pointed to a new swizzled location.

Now, the preswizzling being done as part of the quantize kernel, means the quantizer kernel directly writes into the cached weight's swizzled buffer and the scale buffer is never modified

[vthumbe1503]

/te-ci