Enable Weight Preswizzling only when Swizzle fusion is available

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,47 +121,43 @@ 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.
-    """
+@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)
 
     # 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(
         layer_type, num_gemms
     ), 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
@@ -164,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)
@@ -191,8 +198,35 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls):
 
 
 @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.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.parametrize("layer_type", _LAYER_TYPES)
+@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
@@ -205,12 +239,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)

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,40 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
             f"{self.__class__.__name__} class does not implement _get_weight_quantizers function"
         )
 
+    def _enable_weight_preswizzle(
+        self,
+        quantizer: Quantizer,
+        weight: torch.Tensor,
+    ) -> 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:
+            return False
+        if isinstance(quantizer, MXFP8Quantizer):
+            return True
+        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."""
         self.set_meta_tensor(True, recipe)

transformer_engine/pytorch/module/grouped_linear.py

@@ -1900,6 +1900,13 @@ def forward(
                 grad_weight_quantizers,
                 grad_output_quantizers,
             ) = quantizers
+            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
@@ -2078,16 +2085,8 @@ 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 _get_quantizers(self):

transformer_engine/pytorch/module/layernorm_linear.py

@@ -1747,6 +1747,10 @@ def forward(
                 grad_weight_quantizer,
                 grad_output_quantizer,
             ) = quantizers
+            if weight_quantizer is not None and not debug:
+                weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
+                    weight_quantizer, weight_tensor
+                )
 
             if is_grad_enabled:
                 fwd_fn = _LayerNormLinear.apply
@@ -1994,12 +1998,4 @@ 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]

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -2337,6 +2337,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:
+                    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
+                    )
             if not self.fp8:
                 if isinstance(fc1_weight, Float8Tensor):
                     fc1_weight = fc1_weight.dequantize()
@@ -2713,17 +2722,6 @@ 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 backward_dw(self):

transformer_engine/pytorch/module/linear.py

@@ -1918,6 +1918,10 @@ def forward(
                 grad_weight_quantizer,
                 grad_output_quantizer,
             ) = quantizers
+            if weight_quantizer is not None and not debug:
+                weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
+                    weight_quantizer, weight_tensor
+                )
 
             if is_grad_enabled:
                 linear_fn = _Linear.apply
@@ -2205,12 +2209,4 @@ 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]