Grouped Linear CPU Overhead Reductions By Removing Unnecessary Python Logic

tests/pytorch/test_fusible_ops.py

@@ -4,7 +4,9 @@
 
 from __future__ import annotations
 
+
 from collections.abc import Iterable, Sequence
+from unittest.mock import Mock
 import io
 import os
 import math
@@ -22,7 +24,9 @@
 from transformer_engine.pytorch.ops.basic.grouped_linear import (
     OUTPUT_BUFFER_KEY,
     GRAD_INPUT_BUFFER_KEY,
+    is_op_fuser_grouped_tensor_path_supported,
 )
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 from transformer_engine.pytorch.ops.fuser import OperationFuser
 from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV
 
@@ -5661,3 +5665,88 @@ def test_linear_inference_loop(
             quantization=quantization,
             recipe=recipe,
         )
+
+
+@pytest.mark.parametrize("split_alignment", [64, 128])
+def test_basic_grouped_pre_fuser_split_alignment(monkeypatch, split_alignment):
+    """Split fallback configures every expert before adjacent ops consume the quantizers."""
+    fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling()
+    if not is_op_fuser_grouped_tensor_path_supported(fp8_recipe, torch.bfloat16):
+        pytest.skip("Native MXFP8 grouped GEMM is unavailable.")
+    monkeypatch.delenv("NVTE_CUTEDSL_FUSED_GROUPED_MLP", raising=False)
+    FP8GlobalStateManager.reset()
+    op = te_ops.GroupedLinear(
+        3,
+        128,
+        128,
+        bias=False,
+        dtype=torch.bfloat16,
+        device="cuda",
+        split_alignment=split_alignment,
+    )
+    with te.autocast(enabled=True, recipe=fp8_recipe):
+        op.reset_recipe_state(recipe=fp8_recipe)
+        for quantizers in (op._input_quantizers, op._grad_output_quantizers):
+            for quantizer in quantizers:
+                quantizer.set_usage(rowwise=True, columnwise=False)
+        op.pre_fuser_forward(requires_grad=True)
+        for quantizers in (op._input_quantizers, op._grad_output_quantizers):
+            for index, quantizer in enumerate(quantizers):
+                assert quantizer.columnwise_usage == (split_alignment == 64 or index == 0)
+
+
+@pytest.mark.parametrize(
+    "fp8_recipe",
+    [
+        transformer_engine.common.recipe.MXFP8BlockScaling(),
+        transformer_engine.common.recipe.Float8CurrentScaling(),
+    ],
+)
+def test_grouped_discrete_weight_quantizers(fp8_recipe):
+    """Op usages may change between calls without disabling existing parameter usages."""
+    if not is_op_fuser_grouped_tensor_path_supported(fp8_recipe, torch.bfloat16):
+        pytest.skip("Native grouped GEMM is unavailable.")
+    FP8GlobalStateManager.reset()
+    op = te_ops.GroupedLinear(3, 128, 128, bias=False, dtype=torch.bfloat16, device="cuda")
+    with te.autocast(enabled=True, recipe=fp8_recipe):
+        op.reset_recipe_state(recipe=fp8_recipe)
+    quantizers = [op.get_quantizer("forward", 2 * i + 1) for i in range(op.num_groups)]
+    assert len({id(q) for q in quantizers}) == 3
+    for q in quantizers:
+        q.set_usage(rowwise=True, columnwise=True)
+    weights = [torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)]
+    reference = [q(w) for w, q in zip(weights, quantizers)]
+    weights[1] = reference[1]
+    param_quantizer = weights[1]._quantizer
+    assert param_quantizer is not quantizers[1]
+    # The data already has both layouts; its quantizer must adopt the op's requirement.
+    param_quantizer.set_usage(rowwise=True, columnwise=False)
+    for q in quantizers:
+        q.set_usage(rowwise=False, columnwise=False)
+
+    # Enable columnwise usage, stop requesting it, then request it again.
+    for columnwise_usage in (True, False, True):
+        quantizer_spies = [Mock(wraps=q) for q in quantizers]
+        outputs = op._get_discrete_weights_for_gemm(
+            weights,
+            quantizer_spies,
+            columnwise_usage=columnwise_usage,
+            with_quantized_compute=True,
+            dtype=torch.bfloat16,
+        )
+        for i in (0, 2):
+            quantizer_spies[i].assert_called_once_with(weights[i])
+        for spy, quantizer in zip(quantizer_spies, quantizers):
+            spy.set_usage.assert_called_once_with(rowwise=True, columnwise=columnwise_usage)
+            assert quantizer.rowwise_usage
+            assert quantizer.columnwise_usage == columnwise_usage
+        quantizer_spies[1].assert_not_called()
+        assert outputs[1] is weights[1]
+        assert weights[1]._quantizer is param_quantizer
+        assert param_quantizer.rowwise_usage
+        assert param_quantizer.columnwise_usage
+        assert weights[1].get_usages()["columnwise"]
+        for output, ref in zip(outputs, reference):
+            torch.testing.assert_close(output.dequantize(), ref.dequantize(), rtol=0, atol=0)
+        scales = [w._rowwise_scale_inv if fp8_recipe.mxfp8() else w._scale_inv for w in outputs]
+        assert len({s.data_ptr() for s in scales}) == 3