hook up real nvfp4 grouped_gemm

test/prototype/mx_formats/test_inference_workflow.py

@@ -268,17 +268,6 @@ def test_narrow_similar_to_vllm(self):
         )
         self._test_narrow_similar_to_vllm(config)
 
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.skipif(
-        not torch_version_at_least("2.8.0"),
-        reason="torch.compile requires PyTorch 2.8+",
-    )
-    def test_nvfp4_quantize_3d_param_similar_to_vllm(self):
-        config = NVFP4WeightOnlyConfig(
-            use_dynamic_per_tensor_scale=False,
-        )
-        self._test_quantize_3d_param_similar_to_vllm(config)
-
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 @pytest.mark.skipif(
@@ -452,16 +441,22 @@ def forward(self, x, offs):
 def test_grouped_mm_nvfp4():
     """Smoke test: torch._grouped_mm forward pass with bfloat16 inputs."""
     E, K, N = 4, 128, 256
-    m_per_group = [32, 96, 16, 112]
+    m_per_group = [1, 3, 4, 16]
     total_m = sum(m_per_group)
 
     device = "cuda"
     dtype = torch.bfloat16
 
     model_ref = GroupedMMModel(E, K, N, device=device, dtype=dtype)
+
+    # make the future nvfp4 weight scales interesting
+    model_ref.weight[0, :, :] *= 10.0
+    model_ref.weight[-1, :, :] *= 1e-3
+
     model = copy.deepcopy(model_ref)
 
     x = torch.randn(total_m, K, device=device, dtype=dtype)
+
     offs = torch.tensor(
         [sum(m_per_group[: i + 1]) for i in range(E)],
         device=device,
@@ -483,15 +478,14 @@ def test_grouped_mm_nvfp4():
     assert isinstance(model.weight, NVFP4Tensor), (
         f"Expected NVFP4Tensor weight, got {type(model.weight)}"
     )
+    assert model.weight.per_tensor_scale.shape == (E, 1, 1)
     w_sqnr = compute_error(model_ref.weight, model.weight.dequantize())
     assert w_sqnr > 18.0
 
     ww = model.weight
     wwt = ww.transpose(-2, -1)
     assert tuple(wwt.shape) == (ww.shape[0], ww.shape[2], ww.shape[1])
 
-    # For now, this is emulated. In the near future we'll hook up
-    # a real nvfp4 grouped gemm.
     y = model(x, offs)
     y_sqnr = compute_error(y_ref, y)
-    assert y_sqnr > 18.0
+    assert y_sqnr > 15.0

torchao/prototype/mx_formats/inference_workflow.py

@@ -264,6 +264,7 @@ def _nvfp4_inference_linear_transform(
         return NVFP4ObservedLinear.from_float(module)
 
     elif step == QuantizationStep.CONVERT or step == "convert":
+        assert len(weight.shape) == 2, "3D weights not yet supported here"
         if not isinstance(module, NVFP4ObservedLinear):
             return module
 
@@ -314,8 +315,17 @@ def _nvfp4_inference_linear_transform(
 
         per_tensor_scale = None
         if config.use_dynamic_per_tensor_scale:
-            tensor_amax = torch.max(torch.abs(weight))
-            per_tensor_scale = per_tensor_amax_to_scale(tensor_amax)
+            if len(weight.shape) == 2:
+                tensor_amax = torch.max(torch.abs(weight))
+                per_tensor_scale = per_tensor_amax_to_scale(tensor_amax)
+            else:
+                assert len(weight.shape) == 3, f"unsupported {weight.shape=}"
+                tensor_amax = torch.amax(torch.abs(weight), dim=(1, 2))
+                per_tensor_scale = per_tensor_amax_to_scale(tensor_amax)
+                # 1D -> 3D
+                per_tensor_scale = per_tensor_scale.view(
+                    per_tensor_scale.shape[0], 1, 1
+                )
 
         act_quant_kwargs = QuantizeTensorToNVFP4Kwargs(
             use_dynamic_per_tensor_scale=config.use_dynamic_per_tensor_scale,
@@ -383,6 +393,7 @@ def _nvfp4_weight_only_linear_transform(
     """Quantization handler for NVFP4WeightOnlyConfig"""
     weight = module.weight
 
+    assert len(weight.shape) == 2, "3D weights not yet supported in this workflow"
     if weight.shape[-2] % 16 != 0 or weight.shape[-1] % 16 != 0:
         raise RuntimeError(
             f"NVFP4 only supports weight shape with last 2 dims divisible by 16, got {weight.shape}"

torchao/prototype/mx_formats/kernels.py

@@ -1249,3 +1249,87 @@ def _(x, global_scale=None):
     scales = scales.view(*orig_leading_dims, -1, padded_cols)
     xq = xq.view(*orig_leading_dims, -1, N // 2)
     return scales, xq
+
+
+def mslk_calculate_group_max(x: torch.Tensor, m_sizes: torch.Tensor) -> torch.Tensor:
+    """Compute per-expert activation global scale (encoding convention).
+
+    Args:
+        x: [M, K] concatenated activation tensor (bf16/fp16).
+        m_sizes: [E] int64 tensor of rows per expert.
+
+    Returns:
+        Per-expert global scale in encoding convention (448 * FP4_MAX / amax),
+        shape [E], dtype float32.
+    """
+    return _mslk_calculate_group_max_custom_op(x, m_sizes)
+
+
+@torch.library.custom_op("ao::mslk_calculate_group_max", mutates_args=())
+def _mslk_calculate_group_max_custom_op(
+    x: torch.Tensor, m_sizes: torch.Tensor
+) -> torch.Tensor:
+    assert _mslk_available, (
+        "mslk is required for calculate_group_max. "
+        "Install from https://github.com/meta-pytorch/MSLK"
+    )
+    from mslk.quantize.triton.fp4_quantize import calculate_group_max
+
+    global_scale, _ = calculate_group_max(x, m_sizes)
+    return global_scale
+
+
+@_mslk_calculate_group_max_custom_op.register_fake
+def _(x, m_sizes):
+    E = m_sizes.shape[0]
+    return x.new_empty(E, dtype=torch.float32)
+
+
+def mslk_quantize_nvfp4_stacked(
+    m_sizes: torch.Tensor,
+    x: torch.Tensor,
+    global_scale: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize concatenated MoE activations to NVFP4 with per-expert global scales.
+
+    Args:
+        m_sizes: [E] int64 tensor of rows per expert.
+        x: [M, K] concatenated activation tensor (bf16/fp16).
+        global_scale: [E] fp32 per-expert global scales in encoding convention.
+
+    Returns:
+        Tuple of (xq, scale):
+            xq: [M, K//2] float4_e2m1fn_x2 packed FP4 data.
+            scale: Padded+swizzled float8_e4m3fn block scales.
+    """
+    return _mslk_quantize_nvfp4_stacked_custom_op(m_sizes, x, global_scale)
+
+
+@torch.library.custom_op("ao::mslk_quantize_nvfp4_stacked", mutates_args=())
+def _mslk_quantize_nvfp4_stacked_custom_op(
+    m_sizes: torch.Tensor,
+    x: torch.Tensor,
+    global_scale: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert _mslk_available, (
+        "mslk is required for nvfp4_quantize_stacked. "
+        "Install from https://github.com/meta-pytorch/MSLK"
+    )
+    from mslk.quantize.triton.fp4_quantize import nvfp4_quantize_stacked
+
+    xq, scale = nvfp4_quantize_stacked(m_sizes, x, global_scale)
+    return xq, scale
+
+
+@_mslk_quantize_nvfp4_stacked_custom_op.register_fake
+def _(m_sizes, x, global_scale):
+    M, K = x.shape[0], x.shape[1]
+    num_segments = m_sizes.shape[0]
+    # Upper-bound on padded total rows (each segment can add at most 127 padding rows)
+    padded_total_M_ub = M + num_segments * 127
+    num_scales_per_row = K // 16
+    n_col_blocks = triton.cdiv(num_scales_per_row, 4)
+    padded_cols = n_col_blocks * 4
+    xq = x.new_empty(M, K // 2, dtype=torch.float4_e2m1fn_x2)
+    scale = x.new_empty(padded_total_M_ub, padded_cols, dtype=torch.float8_e4m3fn)
+    return xq, scale

torchao/prototype/mx_formats/nvfp4_tensor.py

@@ -15,7 +15,9 @@
 from torchao.prototype.mx_formats.kernels import (
     f4_unpacked_to_f32,
     f32_to_f4_unpacked,
+    mslk_calculate_group_max,
     mslk_quantize_nvfp4,
+    mslk_quantize_nvfp4_stacked,
     pack_uint4,
     unpack_uint4,
 )
@@ -246,6 +248,7 @@ def get_hp_scales(self) -> torch.Tensor:
             scale_e4m3 = from_blocked(
                 scale_e4m3, math.prod(leading_dims) * M, K // self.block_size
             )
+            scale_e4m3 = scale_e4m3.view(*leading_dims, M, K // self.block_size)
 
         return (
             scale_e4m3.to(torch.float)
@@ -419,7 +422,6 @@ def nvfp4_t(func, types, args, kwargs):
 @implements([aten.transpose.int])
 def nvfp4_transpose(func, types, args, kwargs):
     old, dim0, dim1 = args
-    _assert_no_per_expert_scale(old)
     assert len(old.shape) == 3, f"unsupported rank {len(old.shape)}"
     valid_3d_dims = ((1, 2), (2, 1), (-1, -2), (-2, -1))
     assert (dim0, dim1) in valid_3d_dims, f"transpose unsupported for {dim0=} {dim1=}"
@@ -680,13 +682,49 @@ def nvfp4_addmm(func, types, args, kwargs):
 
 @implements([aten._grouped_mm.default])
 def nvfp4_grouped_mm(func, types, args, kwargs):
+    from torch.nn.functional import ScalingType, SwizzleType, scaled_grouped_mm
+
     mat_a, mat_b = args[0], args[1]
-    # temporary: dequantize nvfp4 tensor to call original grouped_mm
-    # TODO(future PR): enable nvfp4 with per-expert outer scale (current code
-    # has a single outer scale across all experts).
-    # TODO(future PR): hook up real nvfp4 grouped_mm
-    mat_b_dq = mat_b.dequantize()
-    return func(mat_a, mat_b_dq, *args[2:], **kwargs)
+    offs = args[2] if len(args) > 2 else kwargs.get("offs", None)
+    assert offs is not None, "offs is required for nvfp4 grouped_mm"
+
+    # mat_b is a transposed NVFP4Tensor: the model stores weight as (E, N, K)
+    # and calls weight.transpose(-2, -1). Undo the transpose to get the
+    # original NVFP4Tensor so we can extract qdata/scale in the right layout.
+    assert isinstance(mat_b, NVFP4Tensor)
+    is_transposed = mat_b.qdata.stride(-2) < mat_b.qdata.stride(-1)
+    assert is_transposed, "unsupported"
+
+    E = offs.shape[0]
+    m_sizes = torch.diff(offs, prepend=offs.new_zeros(1)).to(torch.int64)
+
+    # mslk_calculate_group_max returns encoding scale (448 * FP4_MAX / amax)
+    # mslk_quantize_nvfp4_stacked needs encoding scale, scaled_grouped_mm
+    # needs decoding scale (1 / encoding_scale)
+    a_global_scale_enc = mslk_calculate_group_max(mat_a, m_sizes)
+    mat_a_qdata, mat_a_scale = mslk_quantize_nvfp4_stacked(
+        m_sizes, mat_a, a_global_scale_enc
+    )
+    a_global_scale_dec = 1.0 / a_global_scale_enc
+
+    # Flatten 3D scale (E, padded_N, padded_K//16) -> 2D (E, padded_N * padded_K//16)
+    # as required by _scaled_grouped_mm for 2D-3D grouped GEMM
+    mat_b_t_scale = mat_b.scale.transpose(-2, -1).flatten(1)
+    # [E, 1, 1] -> E
+    b_global_scale = mat_b.per_tensor_scale.view(E)
+
+    return scaled_grouped_mm(
+        mat_a_qdata,
+        mat_b.qdata.view(torch.float4_e2m1fn_x2),
+        scale_a=[mat_a_scale, a_global_scale_dec],
+        scale_recipe_a=[ScalingType.BlockWise1x16, ScalingType.TensorWise],
+        scale_b=[mat_b_t_scale, b_global_scale],
+        scale_recipe_b=[ScalingType.BlockWise1x16, ScalingType.TensorWise],
+        swizzle_a=SwizzleType.SWIZZLE_32_4_4,
+        swizzle_b=SwizzleType.SWIZZLE_32_4_4,
+        offs=offs,
+        output_dtype=mat_a.dtype,
+    )
 
 
 def per_tensor_amax_to_scale(amax: torch.Tensor) -> torch.Tensor: