From 8f7695f909948f4aa86553fa3442c648530e2c1b Mon Sep 17 00:00:00 2001
From: Varun Thumbe <vthumbe@nvidia.com>
Date: Tue, 21 Jul 2026 00:24:54 +0000
Subject: [PATCH 1/7] optimize for gemm's conditional enablement

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
---
 .../pytorch/test_weight_swizzle_in_layers.py  | 56 +++++++++++++++++--
 .../pytorch/module/grouped_linear.py          | 51 ++++++++++++++---
 .../pytorch/module/layernorm_linear.py        | 43 +++++++++++---
 .../pytorch/module/layernorm_mlp.py           | 50 +++++++++++++----
 transformer_engine/pytorch/module/linear.py   | 43 +++++++++++---
 5 files changed, 203 insertions(+), 40 deletions(-)

diff --git a/tests/pytorch/test_weight_swizzle_in_layers.py b/tests/pytorch/test_weight_swizzle_in_layers.py
index b9f19f2fd9..af24e16163 100644
--- a/tests/pytorch/test_weight_swizzle_in_layers.py
+++ b/tests/pytorch/test_weight_swizzle_in_layers.py
@@ -112,11 +112,9 @@ def _run_step(module, x, is_first_microbatch, recipe, m_splits):
 @pytest.mark.parametrize("layer_type", _LAYER_TYPES)
 @pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES)
 def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls):
-    """Caching the quantized weight across microbatches must pre-swizzle its
-    scales: the weight quantizer enables ``optimize_for_gemm`` and the cached
-    workspace carries ``_with_gemm_swizzled_scales``.
-
-    Generic across every swizzling recipe (MXFP8, NVFP4) and module type.
+    """When direct swizzle fusion is supported, cached weights must pre-swizzle scales.
+    Generic across every swizzling recipe (MXFP8, NVFP4) and module type. Uses
+    128-aligned weight shapes so NVFP4 2D swizzle fusion is eligible.
     """
     torch.manual_seed(1234)
     device = "cuda"
@@ -190,6 +188,54 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls):
         ), "quantized_model_init weight parameter must not have GEMM-swizzled scales"
 
 
+@pytest.mark.parametrize("layer_type", ["Linear", "LayerNormLinear", "GroupedLinear"])
+@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4)
+def test_weight_optimize_for_gemm_disabled_without_swizzle_fusion(layer_type):
+    """NVFP4 weights that cannot use in-kernel swizzle fusion must keep compact cached scales."""
+    torch.manual_seed(1234)
+    device = "cuda"
+    # Valid for NVFP4 quantization, but not aligned to 128x128 swizzle tiles.
+    in_features, out_features = 1056, 1056
+    batch = 512
+    num_gemms = 2 if layer_type == "GroupedLinear" else 1
+    m_splits = _grouped_m_splits(layer_type, batch, num_gemms)
+    recipe = NVFP4BlockScaling(disable_stochastic_rounding=True)
+
+    module = _make_module(layer_type, in_features, out_features, device, num_gemms)
+    x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True)
+
+    weight_quantizers = _forward_backward(module, x, True, recipe, m_splits)
+
+    for weight_quantizer in weight_quantizers:
+        assert weight_quantizer is not None
+        assert weight_quantizer.optimize_for_gemm is False
+
+    for _, ws in module._fp8_workspaces.items():
+        assert getattr(ws, "_with_gemm_swizzled_scales", False) is False
+
+
+@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4)
+def test_weight_optimize_for_gemm_disabled_without_nvfp4_2d_quantization():
+    """NVFP4 weights with 2D quantization disabled cannot preswizzle at quantize time."""
+    torch.manual_seed(1234)
+    device = "cuda"
+    in_features, out_features = 1024, 1024
+    batch = 512
+    recipe = NVFP4BlockScaling(
+        disable_stochastic_rounding=True,
+        disable_2d_quantization=True,
+    )
+
+    module = te.Linear(in_features, out_features, bias=True, params_dtype=torch.bfloat16).to(device)
+    x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True)
+
+    weight_quantizers = _forward_backward(module, x, True, recipe, None)
+
+    assert weight_quantizers[0].optimize_for_gemm is False
+    for _, ws in module._fp8_workspaces.items():
+        assert getattr(ws, "_with_gemm_swizzled_scales", False) is False
+
+
 @pytest.mark.parametrize("layer_type", _LAYER_TYPES)
 @pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES)
 def test_weight_caching_matches_no_caching(layer_type, recipe_cls):
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index a65ee3b5c3..a25999d8b4 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -1836,6 +1836,14 @@ def forward(
                 grad_weight_quantizers,
                 grad_output_quantizers,
             ) = quantizers
+            if not debug:
+                for weight_quantizer, weight_tensor in zip(
+                    weight_quantizers, weight_tensors
+                ):
+                    if weight_quantizer is not None:
+                        self._configure_weight_quantizer_optimize_for_gemm(
+                            weight_quantizer, weight_tensor
+                        )
 
             if is_grad_enabled:
                 linear_fn = _GroupedLinear.apply
@@ -2014,18 +2022,45 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
             ]
             for i in range(self.num_gemms)
         ]
-        # Preswizzle the weights during quantization instead of lazily inside every GEMM.
-        # This wont work when primay weights are in fp8 because of 2 reasons
-        # 1. optimizer step updates would need to dequantize the weights. But swizzled weights
-        # currently dont support dequantization.
-        # 2. For FSDP2, quantized weight all-gather would need to be done in the
-        # unswizzled layout.
         for i in range(self.num_gemms):
             weight_quantizers[i].internal = not self.primary_weights_in_fp8
-            if not self.primary_weights_in_fp8:
-                weight_quantizers[i].optimize_for_gemm = True
         return weight_quantizers
 
+    def _configure_weight_quantizer_optimize_for_gemm(
+        self,
+        quantizer: Quantizer,
+        weight: torch.Tensor,
+    ) -> None:
+        """Configure preswizzling for the per-weight, single-tensor quantize kernel.
+
+        GroupedLinear uses ``group_quantize`` for activations, but its cached weights
+        are quantized individually by ``quantize_weight``. NVFP4 RHT eligibility
+        therefore uses the single-tensor kernel's 64-row alignment, not the grouped
+        kernel's 128-row alignment.
+        """
+        if self.primary_weights_in_fp8:
+            quantizer.optimize_for_gemm = False
+            return
+        if not isinstance(quantizer, NVFP4Quantizer):
+            quantizer.optimize_for_gemm = True
+            return
+
+        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+        capability = get_device_compute_capability()
+        arch_supported = (10, 0) <= capability <= (11, 0)
+        if quantizer.with_rht:
+            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
+        else:
+            enabled = (
+                arch_supported
+                and quantizer.with_2d_quantization
+                and not quantizer.row_scaled_nvfp4
+                and not quantizer.nvfp4_use_4over6
+                and rows % 128 == 0
+                and cols % 128 == 0
+            )
+        quantizer.optimize_for_gemm = enabled
+
     def _get_quantizers(self):
         weight_quantizers = self._get_weight_quantizers()
         input_quantizers, output_quantizers = (
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 43799b003c..0ee10a01d8 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -37,6 +37,7 @@
     cast_if_needed,
     clear_tensor_data,
     divide,
+    get_device_compute_capability,
     get_default_init_method,
     init_method_constant,
     nvtx_range_pop,
@@ -72,6 +73,7 @@
     prepare_for_saving,
     restore_from_func_ctx,
 )
+from ..tensor.nvfp4_tensor import NVFP4Quantizer
 from ...debug.pytorch.debug_state import TEDebugState
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..cpu_offload import (
@@ -1720,6 +1722,10 @@ def forward(
                 grad_weight_quantizer,
                 grad_output_quantizer,
             ) = quantizers
+            if weight_quantizer is not None and not debug:
+                self._configure_weight_quantizer_optimize_for_gemm(
+                    weight_quantizer, weight_tensor
+                )
 
             if is_grad_enabled:
                 fwd_fn = _LayerNormLinear.apply
@@ -1967,12 +1973,33 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
             return [None]
         weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
         weight_quantizer.internal = True
-        # Preswizzle the weights during quantization instead of lazily inside every GEMM.
-        # This wont work when primay weights are in fp8 because of 2 reasons
-        # 1. optimizer step updates would need to dequantize the weights. But swizzled weights
-        # currently dont support dequantization.
-        # 2. For FSDP2, quantized weight all-gather would need to be done in the
-        # unswizzled layout.
-        if not self.primary_weights_in_fp8:
-            weight_quantizer.optimize_for_gemm = True
         return [weight_quantizer]
+
+    def _configure_weight_quantizer_optimize_for_gemm(
+        self,
+        quantizer: Quantizer,
+        weight: torch.Tensor,
+    ) -> None:
+        """Configure preswizzling for the single-tensor weight quantize kernel."""
+        if self.primary_weights_in_fp8:
+            quantizer.optimize_for_gemm = False
+            return
+        if not isinstance(quantizer, NVFP4Quantizer):
+            quantizer.optimize_for_gemm = True
+            return
+
+        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+        capability = get_device_compute_capability()
+        arch_supported = (10, 0) <= capability <= (11, 0)
+        if quantizer.with_rht:
+            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
+        else:
+            enabled = (
+                arch_supported
+                and quantizer.with_2d_quantization
+                and not quantizer.row_scaled_nvfp4
+                and not quantizer.nvfp4_use_4over6
+                and rows % 128 == 0
+                and cols % 128 == 0
+            )
+        quantizer.optimize_for_gemm = enabled
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index ad629c3f11..29015647e6 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -46,6 +46,7 @@
     cast_if_needed,
     assert_dim_for_fp8_exec,
     clear_tensor_data,
+    get_device_compute_capability,
     requires_grad,
     needs_quantized_gemm,
     get_nvtx_range_context,
@@ -2337,6 +2338,15 @@ def forward(
             fc1_weight, fc2_weight = self._get_weight_tensors()
             fc1_bias = self.fc1_bias if self.use_bias else None
             fc2_bias = self.fc2_bias if self.use_bias else None
+            if not debug:
+                if fc1_weight_quantizer is not None:
+                    self._configure_weight_quantizer_optimize_for_gemm(
+                        fc1_weight_quantizer, fc1_weight
+                    )
+                if fc2_weight_quantizer is not None:
+                    self._configure_weight_quantizer_optimize_for_gemm(
+                        fc2_weight_quantizer, fc2_weight
+                    )
             if not self.fp8:
                 if isinstance(fc1_weight, Float8Tensor):
                     fc1_weight = fc1_weight.dequantize()
@@ -2713,19 +2723,37 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
         fc1_weight_quantizer.internal = True
         fc2_weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_WEIGHT]
         fc2_weight_quantizer.internal = True
-        # Weight scale factors must be GEMM-swizzled before cuBLAS/CUTLASS can
-        # consume them. Pre-swizzle once at quantize time (persisted on the cached
-        # workspace when the weight is cached) instead of lazily inside every GEMM.
-        # No-op for recipes whose scales don't need swizzling (e.g. per-tensor FP8).
-        # The exception is quantized_model_init, where the weight parameter is
-        # itself quantized and gets all-gathered (FSDP2) and optimizer-updated in
-        # its unswizzled layout; swizzling would break the all-gather and the
-        # dequantize-on-update, so leave the quantizer state untouched.
-        if not self.primary_weights_in_fp8:
-            fc1_weight_quantizer.optimize_for_gemm = True
-            fc2_weight_quantizer.optimize_for_gemm = True
         return [fc1_weight_quantizer, fc2_weight_quantizer]
 
+    def _configure_weight_quantizer_optimize_for_gemm(
+        self,
+        quantizer: Quantizer,
+        weight: torch.Tensor,
+    ) -> None:
+        """Configure preswizzling for the single-tensor weight quantize kernel."""
+        if self.primary_weights_in_fp8:
+            quantizer.optimize_for_gemm = False
+            return
+        if not isinstance(quantizer, NVFP4Quantizer):
+            quantizer.optimize_for_gemm = True
+            return
+
+        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+        capability = get_device_compute_capability()
+        arch_supported = (10, 0) <= capability <= (11, 0)
+        if quantizer.with_rht:
+            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
+        else:
+            enabled = (
+                arch_supported
+                and quantizer.with_2d_quantization
+                and not quantizer.row_scaled_nvfp4
+                and not quantizer.nvfp4_use_4over6
+                and rows % 128 == 0
+                and cols % 128 == 0
+            )
+        quantizer.optimize_for_gemm = enabled
+
     def backward_dw(self):
         """
         Execute the delayed weight gradient computation.
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 78a4d31852..016610754f 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -39,6 +39,7 @@
     init_method_constant,
     needs_quantized_gemm,
     assert_dim_for_fp8_exec,
+    get_device_compute_capability,
     nvtx_range_pop,
     nvtx_range_push,
     get_nvtx_range_context,
@@ -70,6 +71,7 @@
 )
 from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
+from ..tensor.nvfp4_tensor import NVFP4Quantizer
 from ..tensor.utils import clear_columnwise_cache, is_custom
 from ..export import is_in_onnx_export_mode, assert_warmed_up
 from ..cpu_offload import (
@@ -1868,6 +1870,10 @@ def forward(
                 grad_weight_quantizer,
                 grad_output_quantizer,
             ) = quantizers
+            if weight_quantizer is not None and not debug:
+                self._configure_weight_quantizer_optimize_for_gemm(
+                    weight_quantizer, weight_tensor
+                )
 
             if is_grad_enabled:
                 linear_fn = _Linear.apply
@@ -2155,12 +2161,33 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
             return [None]
         weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
         weight_quantizer.internal = True
-        # Preswizzle the weights during quantization instead of lazily inside every GEMM.
-        # This wont work when primay weights are in fp8 because of 2 reasons
-        # 1. optimizer step updates would need to dequantize the weights. But swizzled weights
-        # currently dont support dequantization.
-        # 2. For FSDP2, quantized weight all-gather would need to be done in the
-        # unswizzled layout.
-        if not self.primary_weights_in_fp8:
-            weight_quantizer.optimize_for_gemm = True
         return [weight_quantizer]
+
+    def _configure_weight_quantizer_optimize_for_gemm(
+        self,
+        quantizer: Quantizer,
+        weight: torch.Tensor,
+    ) -> None:
+        """Configure preswizzling for the single-tensor weight quantize kernel."""
+        if self.primary_weights_in_fp8:
+            quantizer.optimize_for_gemm = False
+            return
+        if not isinstance(quantizer, NVFP4Quantizer):
+            quantizer.optimize_for_gemm = True
+            return
+
+        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+        capability = get_device_compute_capability()
+        arch_supported = (10, 0) <= capability <= (11, 0)
+        if quantizer.with_rht:
+            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
+        else:
+            enabled = (
+                arch_supported
+                and quantizer.with_2d_quantization
+                and not quantizer.row_scaled_nvfp4
+                and not quantizer.nvfp4_use_4over6
+                and rows % 128 == 0
+                and cols % 128 == 0
+            )
+        quantizer.optimize_for_gemm = enabled

From 3b1a05e2e3869061c9c5605e79b5fad22f66f248 Mon Sep 17 00:00:00 2001
From: Varun Thumbe <vthumbe@nvidia.com>
Date: Tue, 21 Jul 2026 23:08:56 +0000
Subject: [PATCH 2/7] avoid code repeatition

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
---
 transformer_engine/pytorch/module/base.py     | 40 +++++++++++++++++++
 .../pytorch/module/grouped_linear.py          | 35 ----------------
 .../pytorch/module/layernorm_linear.py        | 31 --------------
 .../pytorch/module/layernorm_mlp.py           | 30 --------------
 transformer_engine/pytorch/module/linear.py   | 31 --------------
 5 files changed, 40 insertions(+), 127 deletions(-)

diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 5eab45b7ac..40d7eef91c 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -58,6 +58,7 @@
 from ..utils import (
     is_non_tn_fp8_gemm_supported,
     torch_get_autocast_gpu_dtype,
+    get_device_compute_capability,
     get_nvtx_range_context,
     nvtx_range_push,
     nvtx_range_pop,
@@ -1207,6 +1208,45 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
             f"{self.__class__.__name__} class does not implement _get_weight_quantizers function"
         )
 
+    def _configure_weight_quantizer_optimize_for_gemm(
+        self,
+        quantizer: Quantizer,
+        weight: torch.Tensor,
+    ) -> None:
+        """Configure preswizzling for the single-tensor weight quantize kernel.
+
+        When ``optimize_for_gemm`` is enabled, scale-factor swizzling is fused into
+        weight quantization instead of being done lazily inside every GEMM.
+        This is disabled when primary weights are already quantized (dequant and
+        FSDP2 all-gather expect the unswizzled layout). For NVFP4, it is only
+        enabled for shapes/architectures where the fused swizzle+quantize kernel
+        is supported. Weight quantization always uses the single-tensor kernel
+        (including GroupedLinear's cached weights), so NVFP4 RHT eligibility uses
+        that kernel's 64-row alignment.
+        """
+        if self.primary_weights_in_fp8:
+            quantizer.optimize_for_gemm = False
+            return
+        if not isinstance(quantizer, NVFP4Quantizer):
+            quantizer.optimize_for_gemm = True
+            return
+
+        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+        capability = get_device_compute_capability()
+        arch_supported = (10, 0) <= capability <= (11, 0)
+        if quantizer.with_rht:
+            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
+        else:
+            enabled = (
+                arch_supported
+                and quantizer.with_2d_quantization
+                and not quantizer.row_scaled_nvfp4
+                and not quantizer.nvfp4_use_4over6
+                and rows % 128 == 0
+                and cols % 128 == 0
+            )
+        quantizer.optimize_for_gemm = enabled
+
     def init_fp8_meta_tensors(self, recipe: Recipe) -> None:
         """Init scales and amaxes."""
         self.set_meta_tensor(True, recipe)
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index a25999d8b4..1eef11bd10 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -2026,41 +2026,6 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
             weight_quantizers[i].internal = not self.primary_weights_in_fp8
         return weight_quantizers
 
-    def _configure_weight_quantizer_optimize_for_gemm(
-        self,
-        quantizer: Quantizer,
-        weight: torch.Tensor,
-    ) -> None:
-        """Configure preswizzling for the per-weight, single-tensor quantize kernel.
-
-        GroupedLinear uses ``group_quantize`` for activations, but its cached weights
-        are quantized individually by ``quantize_weight``. NVFP4 RHT eligibility
-        therefore uses the single-tensor kernel's 64-row alignment, not the grouped
-        kernel's 128-row alignment.
-        """
-        if self.primary_weights_in_fp8:
-            quantizer.optimize_for_gemm = False
-            return
-        if not isinstance(quantizer, NVFP4Quantizer):
-            quantizer.optimize_for_gemm = True
-            return
-
-        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
-        capability = get_device_compute_capability()
-        arch_supported = (10, 0) <= capability <= (11, 0)
-        if quantizer.with_rht:
-            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
-        else:
-            enabled = (
-                arch_supported
-                and quantizer.with_2d_quantization
-                and not quantizer.row_scaled_nvfp4
-                and not quantizer.nvfp4_use_4over6
-                and rows % 128 == 0
-                and cols % 128 == 0
-            )
-        quantizer.optimize_for_gemm = enabled
-
     def _get_quantizers(self):
         weight_quantizers = self._get_weight_quantizers()
         input_quantizers, output_quantizers = (
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 0ee10a01d8..deb8282f1f 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -37,7 +37,6 @@
     cast_if_needed,
     clear_tensor_data,
     divide,
-    get_device_compute_capability,
     get_default_init_method,
     init_method_constant,
     nvtx_range_pop,
@@ -73,7 +72,6 @@
     prepare_for_saving,
     restore_from_func_ctx,
 )
-from ..tensor.nvfp4_tensor import NVFP4Quantizer
 from ...debug.pytorch.debug_state import TEDebugState
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..cpu_offload import (
@@ -1974,32 +1972,3 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
         weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
         weight_quantizer.internal = True
         return [weight_quantizer]
-
-    def _configure_weight_quantizer_optimize_for_gemm(
-        self,
-        quantizer: Quantizer,
-        weight: torch.Tensor,
-    ) -> None:
-        """Configure preswizzling for the single-tensor weight quantize kernel."""
-        if self.primary_weights_in_fp8:
-            quantizer.optimize_for_gemm = False
-            return
-        if not isinstance(quantizer, NVFP4Quantizer):
-            quantizer.optimize_for_gemm = True
-            return
-
-        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
-        capability = get_device_compute_capability()
-        arch_supported = (10, 0) <= capability <= (11, 0)
-        if quantizer.with_rht:
-            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
-        else:
-            enabled = (
-                arch_supported
-                and quantizer.with_2d_quantization
-                and not quantizer.row_scaled_nvfp4
-                and not quantizer.nvfp4_use_4over6
-                and rows % 128 == 0
-                and cols % 128 == 0
-            )
-        quantizer.optimize_for_gemm = enabled
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 29015647e6..177200b546 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -46,7 +46,6 @@
     cast_if_needed,
     assert_dim_for_fp8_exec,
     clear_tensor_data,
-    get_device_compute_capability,
     requires_grad,
     needs_quantized_gemm,
     get_nvtx_range_context,
@@ -2725,35 +2724,6 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
         fc2_weight_quantizer.internal = True
         return [fc1_weight_quantizer, fc2_weight_quantizer]
 
-    def _configure_weight_quantizer_optimize_for_gemm(
-        self,
-        quantizer: Quantizer,
-        weight: torch.Tensor,
-    ) -> None:
-        """Configure preswizzling for the single-tensor weight quantize kernel."""
-        if self.primary_weights_in_fp8:
-            quantizer.optimize_for_gemm = False
-            return
-        if not isinstance(quantizer, NVFP4Quantizer):
-            quantizer.optimize_for_gemm = True
-            return
-
-        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
-        capability = get_device_compute_capability()
-        arch_supported = (10, 0) <= capability <= (11, 0)
-        if quantizer.with_rht:
-            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
-        else:
-            enabled = (
-                arch_supported
-                and quantizer.with_2d_quantization
-                and not quantizer.row_scaled_nvfp4
-                and not quantizer.nvfp4_use_4over6
-                and rows % 128 == 0
-                and cols % 128 == 0
-            )
-        quantizer.optimize_for_gemm = enabled
-
     def backward_dw(self):
         """
         Execute the delayed weight gradient computation.
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 016610754f..732839b08c 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -39,7 +39,6 @@
     init_method_constant,
     needs_quantized_gemm,
     assert_dim_for_fp8_exec,
-    get_device_compute_capability,
     nvtx_range_pop,
     nvtx_range_push,
     get_nvtx_range_context,
@@ -71,7 +70,6 @@
 )
 from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
-from ..tensor.nvfp4_tensor import NVFP4Quantizer
 from ..tensor.utils import clear_columnwise_cache, is_custom
 from ..export import is_in_onnx_export_mode, assert_warmed_up
 from ..cpu_offload import (
@@ -2162,32 +2160,3 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
         weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT]
         weight_quantizer.internal = True
         return [weight_quantizer]
-
-    def _configure_weight_quantizer_optimize_for_gemm(
-        self,
-        quantizer: Quantizer,
-        weight: torch.Tensor,
-    ) -> None:
-        """Configure preswizzling for the single-tensor weight quantize kernel."""
-        if self.primary_weights_in_fp8:
-            quantizer.optimize_for_gemm = False
-            return
-        if not isinstance(quantizer, NVFP4Quantizer):
-            quantizer.optimize_for_gemm = True
-            return
-
-        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
-        capability = get_device_compute_capability()
-        arch_supported = (10, 0) <= capability <= (11, 0)
-        if quantizer.with_rht:
-            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
-        else:
-            enabled = (
-                arch_supported
-                and quantizer.with_2d_quantization
-                and not quantizer.row_scaled_nvfp4
-                and not quantizer.nvfp4_use_4over6
-                and rows % 128 == 0
-                and cols % 128 == 0
-            )
-        quantizer.optimize_for_gemm = enabled

From cc8ea83ad8423a1f4d409d1ca4a3b084cc7a408d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 21 Jul 2026 23:11:11 +0000
Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/module/grouped_linear.py   | 4 +---
 transformer_engine/pytorch/module/layernorm_linear.py | 4 +---
 transformer_engine/pytorch/module/linear.py           | 4 +---
 3 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 1eef11bd10..8054941be2 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -1837,9 +1837,7 @@ def forward(
                 grad_output_quantizers,
             ) = quantizers
             if not debug:
-                for weight_quantizer, weight_tensor in zip(
-                    weight_quantizers, weight_tensors
-                ):
+                for weight_quantizer, weight_tensor in zip(weight_quantizers, weight_tensors):
                     if weight_quantizer is not None:
                         self._configure_weight_quantizer_optimize_for_gemm(
                             weight_quantizer, weight_tensor
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index deb8282f1f..3b3533c47c 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -1721,9 +1721,7 @@ def forward(
                 grad_output_quantizer,
             ) = quantizers
             if weight_quantizer is not None and not debug:
-                self._configure_weight_quantizer_optimize_for_gemm(
-                    weight_quantizer, weight_tensor
-                )
+                self._configure_weight_quantizer_optimize_for_gemm(weight_quantizer, weight_tensor)
 
             if is_grad_enabled:
                 fwd_fn = _LayerNormLinear.apply
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 732839b08c..21f03ab960 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -1869,9 +1869,7 @@ def forward(
                 grad_output_quantizer,
             ) = quantizers
             if weight_quantizer is not None and not debug:
-                self._configure_weight_quantizer_optimize_for_gemm(
-                    weight_quantizer, weight_tensor
-                )
+                self._configure_weight_quantizer_optimize_for_gemm(weight_quantizer, weight_tensor)
 
             if is_grad_enabled:
                 linear_fn = _Linear.apply

From 3f8b2ed57df41e8f2f2d0a7fb3e4971accb1e4c6 Mon Sep 17 00:00:00 2001
From: Varun Thumbe <vthumbe@nvidia.com>
Date: Wed, 22 Jul 2026 00:17:10 +0000
Subject: [PATCH 4/7] address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
---
 .../pytorch/test_weight_swizzle_in_layers.py  | 78 +++++++++++--------
 transformer_engine/pytorch/module/base.py     | 52 ++++++-------
 .../pytorch/module/grouped_linear.py          | 15 ++--
 .../pytorch/module/layernorm_linear.py        |  2 +-
 .../pytorch/module/layernorm_mlp.py           |  8 +-
 transformer_engine/pytorch/module/linear.py   |  2 +-
 6 files changed, 82 insertions(+), 75 deletions(-)

diff --git a/tests/pytorch/test_weight_swizzle_in_layers.py b/tests/pytorch/test_weight_swizzle_in_layers.py
index af24e16163..8f44dccfd6 100644
--- a/tests/pytorch/test_weight_swizzle_in_layers.py
+++ b/tests/pytorch/test_weight_swizzle_in_layers.py
@@ -21,19 +21,25 @@
 # skipped individually when the hardware/recipe is unavailable.
 _SWIZZLING_RECIPES = [
     pytest.param(
-        MXFP8BlockScaling,
+        "mxfp8",
         marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8),
-        id="mxfp8",
     ),
     pytest.param(
-        NVFP4BlockScaling,
+        "nvfp4-2d",
+        marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4),
+    ),
+    pytest.param(
+        "nvfp4-1d",
         marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4),
-        id="nvfp4",
     ),
 ]
 
 _LAYER_TYPES = ["Linear", "LayerNormLinear", "LayerNormMLP", "GroupedLinear"]
 
+# 1024 is 128-aligned (NVFP4 swizzle fusion eligible). 1056 is valid MXFP8/NVFP4
+# (divisible by 32) but not 128-aligned, so NVFP4 keeps compact cached scales.
+_WEIGHT_SHAPES = [1024, 1056]
+
 
 def _make_module(layer_type, in_features, out_features, device, num_gemms=1):
     common = dict(bias=True, params_dtype=torch.bfloat16)
@@ -65,12 +71,17 @@ def _grouped_m_splits(layer_type, batch, num_gemms):
     return [batch // num_gemms] * num_gemms
 
 
-def _make_recipe(recipe_cls):
+def _make_recipe(recipe_name):
     """Instantiate a recipe with run-to-run nondeterminism disabled where it
     exists (NVFP4 stochastic rounding); MXFP8 has none."""
-    if recipe_cls is NVFP4BlockScaling:
-        return recipe_cls(disable_stochastic_rounding=True)
-    return recipe_cls()
+    if recipe_name == "mxfp8":
+        return MXFP8BlockScaling()
+    if recipe_name in ("nvfp4-2d", "nvfp4-1d"):
+        return NVFP4BlockScaling(
+            disable_stochastic_rounding=True,
+            disable_2d_quantization=recipe_name == "nvfp4-1d",
+        )
+    raise ValueError(f"unknown recipe {recipe_name}")
 
 
 def _forward_backward(module, x, is_first_microbatch, recipe, m_splits):
@@ -110,19 +121,16 @@ def _run_step(module, x, is_first_microbatch, recipe, m_splits):
 
 
 @pytest.mark.parametrize("layer_type", _LAYER_TYPES)
-@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES)
-def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls):
-    """When direct swizzle fusion is supported, cached weights must pre-swizzle scales.
-    Generic across every swizzling recipe (MXFP8, NVFP4) and module type. Uses
-    128-aligned weight shapes so NVFP4 2D swizzle fusion is eligible.
-    """
+@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES)
+def test_weight_swizzling_with_workspace_caching(layer_type, recipe_name):
+    """Cached weights use preswizzled scales only when direct fusion is supported."""
     torch.manual_seed(1234)
     device = "cuda"
     in_features, out_features = 1024, 1024
     batch = 512
     num_gemms = 2 if layer_type == "GroupedLinear" else 1
     m_splits = _grouped_m_splits(layer_type, batch, num_gemms)
-    recipe = recipe_cls()
+    recipe = _make_recipe(recipe_name)
 
     module = _make_module(layer_type, in_features, out_features, device, num_gemms)
     x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True)
@@ -130,11 +138,12 @@ def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls):
     # is_first_microbatch=True caches the quantized weight.
     weight_quantizers = _forward_backward(module, x, True, recipe, m_splits)
 
+    expect_preswizzle = recipe_name != "nvfp4-1d"
     for weight_quantizer in weight_quantizers:
         assert weight_quantizer is not None
         assert (
-            weight_quantizer.optimize_for_gemm is True
-        ), f"optimize_for_gemm must be enabled for cached {layer_type} weights"
+            weight_quantizer.optimize_for_gemm is expect_preswizzle
+        ), f"unexpected optimize_for_gemm for cached {layer_type} weights"
 
     workspaces = module._fp8_workspaces
     assert len(workspaces) == _expected_num_weights(
@@ -142,13 +151,13 @@ def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls):
     ), f"unexpected cached weight workspace count for {layer_type}: {len(workspaces)}"
     for name, ws in workspaces.items():
         assert (
-            getattr(ws, "_with_gemm_swizzled_scales", False) is True
-        ), f"cached weight workspace {name!r} scales were not pre-swizzled"
+            getattr(ws, "_with_gemm_swizzled_scales", False) is expect_preswizzle
+        ), f"unexpected scale layout for cached weight workspace {name!r}"
 
 
 @pytest.mark.parametrize("layer_type", _LAYER_TYPES)
-@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES)
-def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls):
+@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES)
+def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_name):
     """With quantized_model_init the weight parameter is itself quantized and is
     all-gathered (FSDP2) / optimizer-updated in its unswizzled layout, so the
     eager-swizzle optimization must stay off: the weight quantizer must keep
@@ -162,7 +171,7 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls):
     batch = 512
     num_gemms = 2 if layer_type == "GroupedLinear" else 1
     m_splits = _grouped_m_splits(layer_type, batch, num_gemms)
-    recipe = recipe_cls()
+    recipe = _make_recipe(recipe_name)
 
     with te.quantized_model_init(enabled=True, recipe=recipe):
         module = _make_module(layer_type, in_features, out_features, device, num_gemms)
@@ -188,7 +197,7 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls):
         ), "quantized_model_init weight parameter must not have GEMM-swizzled scales"
 
 
-@pytest.mark.parametrize("layer_type", ["Linear", "LayerNormLinear", "GroupedLinear"])
+@pytest.mark.parametrize("layer_type", _LAYER_TYPES)
 @pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4)
 def test_weight_optimize_for_gemm_disabled_without_swizzle_fusion(layer_type):
     """NVFP4 weights that cannot use in-kernel swizzle fusion must keep compact cached scales."""
@@ -214,31 +223,36 @@ def test_weight_optimize_for_gemm_disabled_without_swizzle_fusion(layer_type):
         assert getattr(ws, "_with_gemm_swizzled_scales", False) is False
 
 
+@pytest.mark.parametrize("layer_type", _LAYER_TYPES)
 @pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4)
-def test_weight_optimize_for_gemm_disabled_without_nvfp4_2d_quantization():
+def test_weight_optimize_for_gemm_disabled_without_nvfp4_2d_quantization(layer_type):
     """NVFP4 weights with 2D quantization disabled cannot preswizzle at quantize time."""
     torch.manual_seed(1234)
     device = "cuda"
     in_features, out_features = 1024, 1024
     batch = 512
+    num_gemms = 2 if layer_type == "GroupedLinear" else 1
+    m_splits = _grouped_m_splits(layer_type, batch, num_gemms)
     recipe = NVFP4BlockScaling(
         disable_stochastic_rounding=True,
         disable_2d_quantization=True,
     )
-
-    module = te.Linear(in_features, out_features, bias=True, params_dtype=torch.bfloat16).to(device)
+    module = _make_module(layer_type, in_features, out_features, device, num_gemms)
     x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True)
 
-    weight_quantizers = _forward_backward(module, x, True, recipe, None)
+    weight_quantizers = _forward_backward(module, x, True, recipe, m_splits)
 
-    assert weight_quantizers[0].optimize_for_gemm is False
+    for weight_quantizer in weight_quantizers:
+        assert weight_quantizer is not None
+        assert weight_quantizer.optimize_for_gemm is False
     for _, ws in module._fp8_workspaces.items():
         assert getattr(ws, "_with_gemm_swizzled_scales", False) is False
 
 
 @pytest.mark.parametrize("layer_type", _LAYER_TYPES)
-@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES)
-def test_weight_caching_matches_no_caching(layer_type, recipe_cls):
+@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES)
+@pytest.mark.parametrize("features", _WEIGHT_SHAPES)
+def test_weight_caching_matches_no_caching(layer_type, recipe_name, features):
     """Caching the quantized (pre-swizzled) weight across microbatches must be
     numerically identical to the uncached flow that re-quantizes the weight every
     microbatch. Verified per microbatch for fprop output, dgrad and wgrad, across
@@ -251,12 +265,12 @@ def test_weight_caching_matches_no_caching(layer_type, recipe_cls):
     """
     torch.manual_seed(1234)
     device = "cuda"
-    in_features, out_features = 1024, 1024
+    in_features, out_features = features, features
     batch = 512
     microbatches = 4
     num_gemms = 2 if layer_type == "GroupedLinear" else 1
     m_splits = _grouped_m_splits(layer_type, batch, num_gemms)
-    recipe = _make_recipe(recipe_cls)
+    recipe = _make_recipe(recipe_name)
 
     # Identical modules: one drives the cached path, one the uncached path.
     cached = _make_module(layer_type, in_features, out_features, device, num_gemms)
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 40d7eef91c..3076b44e84 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -1208,44 +1208,38 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
             f"{self.__class__.__name__} class does not implement _get_weight_quantizers function"
         )
 
-    def _configure_weight_quantizer_optimize_for_gemm(
+    def _enable_weight_preswizzle(
         self,
         quantizer: Quantizer,
         weight: torch.Tensor,
-    ) -> None:
-        """Configure preswizzling for the single-tensor weight quantize kernel.
-
-        When ``optimize_for_gemm`` is enabled, scale-factor swizzling is fused into
-        weight quantization instead of being done lazily inside every GEMM.
-        This is disabled when primary weights are already quantized (dequant and
-        FSDP2 all-gather expect the unswizzled layout). For NVFP4, it is only
-        enabled for shapes/architectures where the fused swizzle+quantize kernel
-        is supported. Weight quantization always uses the single-tensor kernel
-        (including GroupedLinear's cached weights), so NVFP4 RHT eligibility uses
-        that kernel's 64-row alignment.
+    ) -> bool:
+        """Whether to fuse scale-factor swizzling into weight quantization.
+
+        When enabled, scales are preswizzled during quantization instead of lazily
+        inside every GEMM. Disabled when primary weights are already quantized
+        (dequant and FSDP2 all-gather expect the unswizzled layout). For NVFP4,
+        enabled only for shapes/architectures where the fused swizzle+quantize
+        kernel is supported. Weight quantization always uses the single-tensor
+        kernel (including GroupedLinear's cached weights), so NVFP4 RHT
+        eligibility uses that kernel's 64-row alignment.
         """
         if self.primary_weights_in_fp8:
-            quantizer.optimize_for_gemm = False
-            return
+            return False
         if not isinstance(quantizer, NVFP4Quantizer):
-            quantizer.optimize_for_gemm = True
-            return
+            return True
 
         rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
-        capability = get_device_compute_capability()
-        arch_supported = (10, 0) <= capability <= (11, 0)
+        arch_supported = get_device_compute_capability() >= (10, 0)
         if quantizer.with_rht:
-            enabled = arch_supported and rows % 64 == 0 and cols % 128 == 0
-        else:
-            enabled = (
-                arch_supported
-                and quantizer.with_2d_quantization
-                and not quantizer.row_scaled_nvfp4
-                and not quantizer.nvfp4_use_4over6
-                and rows % 128 == 0
-                and cols % 128 == 0
-            )
-        quantizer.optimize_for_gemm = enabled
+            return arch_supported and rows % 64 == 0 and cols % 128 == 0
+        return (
+            arch_supported
+            and quantizer.with_2d_quantization
+            and not quantizer.row_scaled_nvfp4
+            and not quantizer.nvfp4_use_4over6
+            and rows % 128 == 0
+            and cols % 128 == 0
+        )
 
     def init_fp8_meta_tensors(self, recipe: Recipe) -> None:
         """Init scales and amaxes."""
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 1eef11bd10..0c2e953823 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -1836,14 +1836,13 @@ def forward(
                 grad_weight_quantizers,
                 grad_output_quantizers,
             ) = quantizers
-            if not debug:
-                for weight_quantizer, weight_tensor in zip(
-                    weight_quantizers, weight_tensors
-                ):
-                    if weight_quantizer is not None:
-                        self._configure_weight_quantizer_optimize_for_gemm(
-                            weight_quantizer, weight_tensor
-                        )
+            if not debug and weight_quantizers[0] is not None:
+                # Experts share shape and recipe settings: compute once and broadcast.
+                optimize_for_gemm = self._enable_weight_preswizzle(
+                    weight_quantizers[0], weight_tensors[0]
+                )
+                for q in weight_quantizers:
+                    q.optimize_for_gemm = optimize_for_gemm
 
             if is_grad_enabled:
                 linear_fn = _GroupedLinear.apply
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index deb8282f1f..796e371827 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -1721,7 +1721,7 @@ def forward(
                 grad_output_quantizer,
             ) = quantizers
             if weight_quantizer is not None and not debug:
-                self._configure_weight_quantizer_optimize_for_gemm(
+                weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
                     weight_quantizer, weight_tensor
                 )
 
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 177200b546..52efb8a480 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -2339,12 +2339,12 @@ def forward(
             fc2_bias = self.fc2_bias if self.use_bias else None
             if not debug:
                 if fc1_weight_quantizer is not None:
-                    self._configure_weight_quantizer_optimize_for_gemm(
-                        fc1_weight_quantizer, fc1_weight
+                    fc1_weight_quantizer.optimize_for_gemm = (
+                        self._enable_weight_preswizzle(fc1_weight_quantizer, fc1_weight)
                     )
                 if fc2_weight_quantizer is not None:
-                    self._configure_weight_quantizer_optimize_for_gemm(
-                        fc2_weight_quantizer, fc2_weight
+                    fc2_weight_quantizer.optimize_for_gemm = (
+                        self._enable_weight_preswizzle(fc2_weight_quantizer, fc2_weight)
                     )
             if not self.fp8:
                 if isinstance(fc1_weight, Float8Tensor):
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 732839b08c..bd0731aac6 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -1869,7 +1869,7 @@ def forward(
                 grad_output_quantizer,
             ) = quantizers
             if weight_quantizer is not None and not debug:
-                self._configure_weight_quantizer_optimize_for_gemm(
+                weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
                     weight_quantizer, weight_tensor
                 )
 

From ffbd9518abce967d7c60000a7060353396d6b800 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 22 Jul 2026 00:21:01 +0000
Subject: [PATCH 5/7] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/module/layernorm_mlp.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 52efb8a480..19e20d775c 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -2339,12 +2339,12 @@ def forward(
             fc2_bias = self.fc2_bias if self.use_bias else None
             if not debug:
                 if fc1_weight_quantizer is not None:
-                    fc1_weight_quantizer.optimize_for_gemm = (
-                        self._enable_weight_preswizzle(fc1_weight_quantizer, fc1_weight)
+                    fc1_weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
+                        fc1_weight_quantizer, fc1_weight
                     )
                 if fc2_weight_quantizer is not None:
-                    fc2_weight_quantizer.optimize_for_gemm = (
-                        self._enable_weight_preswizzle(fc2_weight_quantizer, fc2_weight)
+                    fc2_weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
+                        fc2_weight_quantizer, fc2_weight
                     )
             if not self.fp8:
                 if isinstance(fc1_weight, Float8Tensor):

From 33a2c8f684e789a6273c773282207396ce1b48f4 Mon Sep 17 00:00:00 2001
From: Varun Thumbe <vthumbe@nvidia.com>
Date: Wed, 22 Jul 2026 00:27:51 +0000
Subject: [PATCH 6/7] remove redundant test

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
---
 .../pytorch/test_weight_swizzle_in_layers.py  | 26 -------------------
 1 file changed, 26 deletions(-)

diff --git a/tests/pytorch/test_weight_swizzle_in_layers.py b/tests/pytorch/test_weight_swizzle_in_layers.py
index 8f44dccfd6..ff961c5083 100644
--- a/tests/pytorch/test_weight_swizzle_in_layers.py
+++ b/tests/pytorch/test_weight_swizzle_in_layers.py
@@ -223,32 +223,6 @@ def test_weight_optimize_for_gemm_disabled_without_swizzle_fusion(layer_type):
         assert getattr(ws, "_with_gemm_swizzled_scales", False) is False
 
 
-@pytest.mark.parametrize("layer_type", _LAYER_TYPES)
-@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4)
-def test_weight_optimize_for_gemm_disabled_without_nvfp4_2d_quantization(layer_type):
-    """NVFP4 weights with 2D quantization disabled cannot preswizzle at quantize time."""
-    torch.manual_seed(1234)
-    device = "cuda"
-    in_features, out_features = 1024, 1024
-    batch = 512
-    num_gemms = 2 if layer_type == "GroupedLinear" else 1
-    m_splits = _grouped_m_splits(layer_type, batch, num_gemms)
-    recipe = NVFP4BlockScaling(
-        disable_stochastic_rounding=True,
-        disable_2d_quantization=True,
-    )
-    module = _make_module(layer_type, in_features, out_features, device, num_gemms)
-    x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True)
-
-    weight_quantizers = _forward_backward(module, x, True, recipe, m_splits)
-
-    for weight_quantizer in weight_quantizers:
-        assert weight_quantizer is not None
-        assert weight_quantizer.optimize_for_gemm is False
-    for _, ws in module._fp8_workspaces.items():
-        assert getattr(ws, "_with_gemm_swizzled_scales", False) is False
-
-
 @pytest.mark.parametrize("layer_type", _LAYER_TYPES)
 @pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES)
 @pytest.mark.parametrize("features", _WEIGHT_SHAPES)

From 1f33eaeebf76408ff5f73dec2976e1332192c6c3 Mon Sep 17 00:00:00 2001
From: Varun Thumbe <vthumbe@nvidia.com>
Date: Thu, 23 Jul 2026 21:23:33 +0000
Subject: [PATCH 7/7] address review comment

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
---
 transformer_engine/pytorch/module/base.py | 29 ++++++++++++-----------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 3076b44e84..7c11e635c6 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -1225,21 +1225,22 @@ def _enable_weight_preswizzle(
         """
         if self.primary_weights_in_fp8:
             return False
-        if not isinstance(quantizer, NVFP4Quantizer):
+        if isinstance(quantizer, MXFP8Quantizer):
             return True
-
-        rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
-        arch_supported = get_device_compute_capability() >= (10, 0)
-        if quantizer.with_rht:
-            return arch_supported and rows % 64 == 0 and cols % 128 == 0
-        return (
-            arch_supported
-            and quantizer.with_2d_quantization
-            and not quantizer.row_scaled_nvfp4
-            and not quantizer.nvfp4_use_4over6
-            and rows % 128 == 0
-            and cols % 128 == 0
-        )
+        if isinstance(quantizer, NVFP4Quantizer):
+            rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+            arch_supported = get_device_compute_capability() >= (10, 0)
+            if quantizer.with_rht:
+                return arch_supported and rows % 64 == 0 and cols % 128 == 0
+            return (
+                arch_supported
+                and quantizer.with_2d_quantization
+                and not quantizer.row_scaled_nvfp4
+                and not quantizer.nvfp4_use_4over6
+                and rows % 128 == 0
+                and cols % 128 == 0
+            )
+        return False
 
     def init_fp8_meta_tensors(self, recipe: Recipe) -> None:
         """Init scales and amaxes."""