Add support and tests for 3D MoE in MXFP8QTensor

Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>

Author: danisereb

modelopt/torch/quantization/qtensor/mxfp8_tensor.py

@@ -71,9 +71,12 @@ def get_weights_scaling_factor(cls, weight: torch.Tensor) -> torch.Tensor:
 
         Args:
             weight: The weight tensor to compute scale for. Must be at least 2D.
+                Supports 2D (out_dim, in_dim) and 3D MoE (num_experts, out_dim, in_dim).
 
         Returns:
             torch.Tensor: E8M0 scale as uint8 tensor with shape [..., out_dim, in_dim // 32].
+                For 2D input: (out_dim, in_dim // 32)
+                For 3D MoE input: (num_experts, out_dim, in_dim // 32)
         """
         assert weight.dim() >= 2, f"Weight must be at least 2D, got {weight.dim()}D"
 
@@ -83,7 +86,7 @@ def get_weights_scaling_factor(cls, weight: torch.Tensor) -> torch.Tensor:
             f"Weight inner dimension ({in_dim}) must be divisible by MXFP8 block size ({cls.BLOCK_SIZE})"
         )
 
-        # Compute amax per block (reduce_block_amax handles reshaping internally)
+        # Compute amax per block (reduce_block_amax handles N-dimensional tensors)
         amax = reduce_block_amax(weight, block_sizes={-1: cls.BLOCK_SIZE})
 
         # Compute E8M0 exponent and convert to biased uint8 (bias = 127)
@@ -102,11 +105,12 @@ def get_weights_scaling_factor_from_quantizer(
         with proper format conversion and shape correction.
 
         Args:
-            weight: The weight tensor.
+            weight: The weight tensor. Can be 2D (out_dim, in_dim) or
+                3D for MoE (num_experts, out_dim, in_dim).
             weight_quantizer: The weight quantizer with block_sizes and optional _scale.
 
         Returns:
-            torch.Tensor: E8M0 scale as uint8 tensor with shape [out_dim, in_dim // 32].
+            torch.Tensor: E8M0 scale as uint8 tensor with shape [..., out_dim, in_dim // 32].
         """
         assert hasattr(weight_quantizer, "block_sizes"), (
             "weight_quantizer must have 'block_sizes' attribute"
@@ -116,8 +120,11 @@ def get_weights_scaling_factor_from_quantizer(
         )
         assert weight.dim() >= 2, f"Weight must be at least 2D, got {weight.dim()}D"
 
-        out_dim, in_dim = weight.shape[-2], weight.shape[-1]
-        expected_shape = (out_dim, in_dim // cls.BLOCK_SIZE)
+        in_dim = weight.shape[-1]
+        # Expected scale shape: all dims except last, with last dim reduced by block size
+        # For 2D: (out_dim, in_dim // 32)
+        # For 3D MoE: (num_experts, out_dim, in_dim // 32)
+        expected_shape = (*weight.shape[:-1], in_dim // cls.BLOCK_SIZE)
 
         if hasattr(weight_quantizer, "_scale") and weight_quantizer._scale is not None:
             scale = weight_quantizer._scale
@@ -127,11 +134,16 @@ def get_weights_scaling_factor_from_quantizer(
             )
 
             # Reshape if needed (same number of elements but wrong shape)
-            if (
-                scale.shape != expected_shape
-                and scale.numel() == expected_shape[0] * expected_shape[1]
-            ):
-                scale = scale.reshape(expected_shape)
+            if scale.shape != expected_shape:
+                expected_numel = 1
+                for dim in expected_shape:
+                    expected_numel *= dim
+                if scale.numel() == expected_numel:
+                    scale = scale.reshape(expected_shape)
+
+            assert scale.shape == expected_shape, (
+                f"Scale shape {scale.shape} does not match expected shape {expected_shape}"
+            )
             return scale
 
         # No scale in quantizer, compute from weight

tests/gpu/torch/quantization/test_qtensor_cuda.py

@@ -850,3 +850,50 @@ def test_mxfp8_quantize_with_scale_asserts(self, device):
         scale = torch.randint(0, 255, (64, 1), dtype=torch.uint8, device=device)
         with pytest.raises(AssertionError, match="must be divisible by MXFP8 block size"):
             MXFP8QTensor.quantize_with_scale(weight_bad_dim, scale)
+
+    @pytest.mark.parametrize("device", ["cuda"])
+    def test_mxfp8_get_weights_scaling_factor_from_quantizer_3d_moe(self, device):
+        """Test get_weights_scaling_factor_from_quantizer handles 3D MoE tensors."""
+        input_shape = (4, 64, 128)  # (num_experts, out_dim, in_dim)
+        weight = torch.randn(input_shape, dtype=torch.float32, device=device)
+
+        class MockQuantizer:
+            block_sizes = {-1: MXFP8QTensor.BLOCK_SIZE}
+            _scale = None
+
+        quantizer = MockQuantizer()
+
+        # Test when _scale is None (should compute from weight)
+        scale = MXFP8QTensor.get_weights_scaling_factor_from_quantizer(weight, quantizer)
+
+        expected_shape = (
+            input_shape[0],
+            input_shape[1],
+            input_shape[2] // MXFP8QTensor.BLOCK_SIZE,
+        )
+        assert scale.shape == expected_shape
+
+        # Test when _scale is provided with correct 3D shape
+        quantizer._scale = torch.randint(0, 255, expected_shape, dtype=torch.uint8, device=device)
+        scale_from_quantizer = MXFP8QTensor.get_weights_scaling_factor_from_quantizer(
+            weight, quantizer
+        )
+        assert torch.equal(scale_from_quantizer, quantizer._scale)
+
+    @pytest.mark.parametrize("device", ["cuda"])
+    def test_mxfp8_get_weights_scaling_factor_from_quantizer_scale_shape_mismatch(self, device):
+        """Test get_weights_scaling_factor_from_quantizer raises assertion on shape mismatch."""
+        input_shape = (4, 64, 128)  # (num_experts, out_dim, in_dim)
+        weight = torch.randn(input_shape, dtype=torch.float32, device=device)
+
+        class MockQuantizer:
+            block_sizes = {-1: MXFP8QTensor.BLOCK_SIZE}
+            # Wrong shape: 2D instead of 3D (missing num_experts dimension)
+            _scale = torch.randint(
+                0, 255, (64, 4), dtype=torch.uint8, device=device
+            )
+
+        quantizer = MockQuantizer()
+
+        with pytest.raises(AssertionError, match="Scale shape .* does not match expected shape"):
+            MXFP8QTensor.get_weights_scaling_factor_from_quantizer(weight, quantizer)