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 bug where optimize_for_gemm (weight preswizzling) was unconditionally enabled for all non-primary-FP8 weights, causing stale pointer issues in CUDA graphs when the fused quantize+swizzle kernel is unavailable (e.g. non-128-aligned shapes, 1-D NVFP4 quantization). The fix introduces _enable_weight_preswizzle in the base module and calls it during forward() once the weight tensor dimensions are available, correctly gating the optimization on architecture, shape alignment, and quantizer configuration.

• base.py: Adds _enable_weight_preswizzle, which returns True only when MXFP8 is used, or when NVFP4 with 2-D quantization is used on a supported GPU with 128-row/col alignment (or 64/128 for RHT). The architecture bound is now an open-ended >= (10, 0), fixing a previous maintainability concern.
• linear.py, layernorm_linear.py, layernorm_mlp.py, grouped_linear.py: Each module's forward() now calls _enable_weight_preswizzle to set optimize_for_gemm after the weight tensor is available, replacing the unconditional True assignment that was previously in _get_weight_quantizers().
• Tests: New test test_weight_optimize_for_gemm_disabled_without_swizzle_fusion covers all module types with a non-128-aligned shape; existing caching and primary-FP8-weight tests are extended to cover both 1-D and 2-D NVFP4 recipe variants.

Confidence Score: 5/5

Safe to merge — the change correctly narrows a previously unconditional flag to only the cases where the fused swizzle kernel is actually supported, with no regressions visible in the logic.

The logic in _enable_weight_preswizzle directly matches the conditions under which the fused quantize+swizzle kernel is known to work (architecture, shape alignment, 2-D quantization mode). All four module types apply the check consistently, and the new test covers both the eligible and non-eligible shape cases across all module types including LayerNormMLP.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Adds _enable_weight_preswizzle with correct shape/arch/recipe gating; open-ended >= (10, 0) arch bound avoids future-GPU maintenance footgun.
transformer_engine/pytorch/module/linear.py	Moves optimize_for_gemm assignment from _get_weight_quantizers to forward() where weight shape is available; correctly gated on weight_quantizer is not None and not debug.
transformer_engine/pytorch/module/layernorm_linear.py	Same pattern as linear.py; correctly calls _enable_weight_preswizzle after weight tensor retrieval in forward().
transformer_engine/pytorch/module/layernorm_mlp.py	Handles both fc1 and fc2 quantizers independently with individual None guards, correctly fixing the bug for both MLP sub-GEMMs.
transformer_engine/pytorch/module/grouped_linear.py	Computes eligibility once from expert 0 then broadcasts to all quantizers; valid since all experts share shape and recipe, with weight_quantizers[0] is not None guard.
tests/pytorch/test_weight_swizzle_in_layers.py	Expands recipe coverage to split nvfp4 into 1-D and 2-D variants, adds 1056-shape non-alignment test across all four module types, and parametrizes the caching-consistency test over both shape variants.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["forward() called"] --> B["_get_weight_quantizers()"]
    B --> C{"weight_quantizer not None and not debug?"}
    C -- No --> G["skip (no preswizzle)"]
    C -- Yes --> D["_enable_weight_preswizzle(quantizer, weight)"]
    D --> E{"primary_weights_in_fp8?"}
    E -- Yes --> F["return False"]
    E -- No --> H{"MXFP8Quantizer?"}
    H -- Yes --> I["return True"]
    H -- No --> J{"NVFP4Quantizer?"}
    J -- No --> K["return False"]
    J -- Yes --> L["compute rows, cols; arch_supported = capability >= (10,0)"]
    L --> M{"with_rht?"}
    M -- Yes --> N["arch_supported & rows%64==0 & cols%128==0"]
    M -- No --> O["arch_supported & with_2d_quantization & not row_scaled & not 4over6 & rows%128==0 & cols%128==0"]
    N --> P["set optimize_for_gemm"]
    O --> P
    I --> P
    F --> G
    K --> G
    P --> Q["GEMM / kernel execution"]

Loading

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

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

/te-ci pytorch

[zhongbozhu]

LGTM, please check those CI issues and see if they are not relevant

[timmoon10]

Overall ok, although I have a nit on how the swizzling logic is implemented.

[timmoon10]

Nit: This implementation will become awkward if we ever add another format with non-trivial logic. I'd suggest something like:

# MXFP8 supports fused quantize-swizzle
if isinstance(quantizer, MXFP8Quantizer):
    return True

# NVFP4
if isinstance(quantizer, NVFP4Quantizer):
    ...

# Default to unswizzled scales
return False

[vthumbe1503]

/te-ci pytorch