Enable Weight Preswizzling only when Swizzle fusion is available

[vthumbe1503]

Description

#3190 recently fixed a bug where enabled nvfp4 2d quantize swizzle fusion for the weights. Enabling this was needed, since without swizzle fusion, if we cache swizzle weights, it causes stale pointer issue with cuda graphs.

However, swizzle fusion isnt available in all nvfp4 quantization cases. And so this PR makes sure to disable optimize_for_gemm for the cases where the fusion isnt even available.

Fixes #3220 (issue)

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 fixes a CUDA graph stale pointer bug by conditionally enabling weight preswizzling only when the fused quantize+swizzle kernel is actually available. The previous code unconditionally set optimize_for_gemm=True for all non-primary NVFP4 weights, but the fused kernel has alignment and quantization-mode requirements that aren't always met.

• Introduces _enable_weight_preswizzle in the base module that checks architecture capability (SM ≥ 10.0), quantization mode (with_2d_quantization, with_rht), and weight shape alignment (128×128 for 2D, 64×128 for RHT) before enabling preswizzle.
• Moves the optimize_for_gemm assignment from _get_weight_quantizers() into the forward pass of all four module types (Linear, LayerNormLinear, LayerNormMLP, GroupedLinear), where the actual weight tensor is available for shape checks.
• Expands the test suite with a new test_weight_optimize_for_gemm_disabled_without_swizzle_fusion test, a separate nvfp4-1d recipe variant, and parametrization over non-128-aligned shapes (1056×1056) to validate both alignment branches.

Confidence Score: 5/5

Safe to merge — the change narrows the set of cases where preswizzling is enabled, directly fixing the CUDA graph stale pointer bug without regressing the aligned-shape fast path.

The logic in _enable_weight_preswizzle is a straightforward guard that defaults to the conservative (non-preswizzle) path when fusion requirements aren't met. All four module types are updated consistently. The test suite now covers both aligned and non-aligned shapes, 1D vs 2D NVFP4 quantization modes, and all layer types including LayerNormMLP.

No files require special attention. The most nuanced change is in base.py where the eligibility conditions are encoded, but these match the kernel requirements documented in the surrounding comments.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Adds _enable_weight_preswizzle to base module; correctly checks arch, quantization mode, and shape alignment before enabling preswizzle; architecture lower-bound is open-ended (≥ SM 10.0)
transformer_engine/pytorch/module/linear.py	Moves optimize_for_gemm assignment from _get_weight_quantizers into the forward pass; correctly placed before workspace retrieval; debug-mode guard is consistent with other modules
transformer_engine/pytorch/module/layernorm_linear.py	Same pattern as linear.py; removes unconditional optimize_for_gemm=True from _get_weight_quantizers and sets it conditionally in forward
transformer_engine/pytorch/module/layernorm_mlp.py	Handles both fc1 and fc2 weight quantizers independently; each gets its own _enable_weight_preswizzle call with its own weight tensor, correctly handling potentially different fc1/fc2 shapes
transformer_engine/pytorch/module/grouped_linear.py	Uses weight_quantizers[0] / weight_tensors[0] to compute once and broadcast to all experts; valid given experts share shape in standard MoE usage
tests/pytorch/test_weight_swizzle_in_layers.py	Adds nvfp4-1d recipe variant, 1056 non-aligned shape, and new test_weight_optimize_for_gemm_disabled_without_swizzle_fusion test covering all four layer types including LayerNormMLP

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[forward pass called] --> B{primary_weights_in_fp8?}
    B -- yes --> C[optimize_for_gemm = False]
    B -- no --> D{isinstance NVFP4Quantizer?}
    D -- no --> E[optimize_for_gemm = True\ne.g. MXFP8]
    D -- yes --> F{quantizer.with_rht?}
    F -- yes --> G{arch >= SM10.0 AND\nrows%64==0 AND cols%128==0?}
    G -- yes --> H[optimize_for_gemm = True]
    G -- no --> C
    F -- no --> I{arch >= SM10.0 AND\nwith_2d_quantization AND\nnot row_scaled AND\nnot 4over6 AND\nrows%128==0 AND cols%128==0?}
    I -- yes --> H
    I -- no --> C
    C --> J[weight quantized without preswizzle]
    E --> K[weight quantized with preswizzle]
    H --> K

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into optimize_for_ge..." | Re-trigger Greptile}

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

/te-ci pytorch