Support blocked MXFP8 scales

Author: drisspg

test/test_mxfp8_tma.py

@@ -10,10 +10,12 @@
 
 try:
     from transformer_nuggets.cute import (
+        BlockScaleLayout,
         GridScheduler,
         MXFP8_TMA_PROFILE_TAGS,
         get_mxfp8_tma_gemv,
         mxfp8_tma_gemv,
+        mxfp8_tma_scaled_mm,
         select_mxfp8_tma_compute_warps,
     )
     from transformer_nuggets.cute.mxfp8_tma import app
@@ -42,6 +44,27 @@ def test_select_mxfp8_tma_compute_warps(k, block_n, expected_sm100):
     assert select_mxfp8_tma_compute_warps(k, block_n) == expected
 
 
+def swizzle_mxfp8_scales(scales: torch.Tensor) -> torch.Tensor:
+    """Convert natural E8M0 scales to canonical SWIZZLE_32_4_4 storage."""
+    rows, cols = scales.shape
+    padded_rows = ((rows + 127) // 128) * 128
+    padded_cols = ((cols + 3) // 4) * 4
+    padded = torch.zeros(
+        (padded_rows, padded_cols),
+        dtype=scales.dtype,
+        device=scales.device,
+    )
+    padded[:rows, :cols] = scales
+    return (
+        padded.view(padded_rows // 128, 128, padded_cols // 4, 4)
+        .permute(0, 2, 1, 3)
+        .reshape(-1, 4, 32, 4)
+        .transpose(1, 2)
+        .reshape(-1)
+        .contiguous()
+    )
+
+
 def dequantize_mxfp8(value: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
     """Dequantize raw MXFP8 storage to float32."""
     expanded_scale = scale.view(torch.float8_e8m0fnu).float().repeat_interleave(32, dim=1)
@@ -81,6 +104,65 @@ def test_mxfp8_tma_gemv_matches_reference(k, block_n, num_stages, num_compute_wa
     torch.testing.assert_close(actual, expected, atol=1.0, rtol=0.05)
 
 
+@pytest.mark.parametrize(("block_n", "num_compute_warps"), [(8, 4), (6, 2)])
+def test_mxfp8_tma_swizzled_scales_match_raw(block_n, num_compute_warps):
+    """Read canonical blocked E8M0 scales across row-layout boundaries."""
+    n, k = 264, 2048
+    q_input, input_scale = quantize_mxfp8(torch.randn((1, k), dtype=torch.bfloat16, device="cuda"))
+    weight, weight_scale = quantize_mxfp8(torch.randn((n, k), dtype=torch.bfloat16, device="cuda"))
+    expected = mxfp8_tma_gemv(
+        q_input,
+        weight,
+        input_scale,
+        weight_scale,
+        block_n=block_n,
+        num_compute_warps=num_compute_warps,
+    )
+    actual = mxfp8_tma_gemv(
+        q_input,
+        weight,
+        swizzle_mxfp8_scales(input_scale),
+        swizzle_mxfp8_scales(weight_scale),
+        block_n=block_n,
+        num_compute_warps=num_compute_warps,
+        block_scale_layout=BlockScaleLayout.SWIZZLE_32_4_4,
+    )
+    torch.cuda.synchronize()
+    torch.testing.assert_close(actual, expected, atol=1.0, rtol=0.05)
+
+
+def test_mxfp8_tma_scaled_mm_adapter_matches_scaled_mm():
+    """Match scaled_mm's blocked E8M0 scale and transposed-weight contract."""
+    from torch.nn.functional import ScalingType, SwizzleType, scaled_mm
+
+    n, k = 128, 2048
+    q_input, input_scale = quantize_mxfp8(torch.randn((1, k), dtype=torch.bfloat16, device="cuda"))
+    weight, weight_scale = quantize_mxfp8(torch.randn((n, k), dtype=torch.bfloat16, device="cuda"))
+    input_scale = swizzle_mxfp8_scales(input_scale).view(torch.float8_e8m0fnu)
+    weight_scale = swizzle_mxfp8_scales(weight_scale).view(torch.float8_e8m0fnu)
+    actual = mxfp8_tma_scaled_mm(
+        q_input,
+        weight.t(),
+        input_scale,
+        weight_scale,
+        block_n=8,
+        num_compute_warps=4,
+    )
+    expected = scaled_mm(
+        q_input,
+        weight.t(),
+        scale_a=[input_scale],
+        scale_recipe_a=[ScalingType.BlockWise1x32],
+        swizzle_a=[SwizzleType.SWIZZLE_32_4_4],
+        scale_b=[weight_scale],
+        scale_recipe_b=[ScalingType.BlockWise1x32],
+        swizzle_b=[SwizzleType.SWIZZLE_32_4_4],
+        output_dtype=torch.bfloat16,
+    )
+    torch.cuda.synchronize()
+    torch.testing.assert_close(actual, expected, atol=1.0, rtol=0.05)
+
+
 def test_mxfp8_tma_gemv_persistent_grid_matches_reference():
     """Reuse a bounded physical CTA grid across all logical output tiles."""
     k = 2048

transformer_nuggets/cute/__init__.py

@@ -15,6 +15,7 @@
 
 _BLOCKSCALED_TMA_EXPORTS = {
     "DEFAULT_PERSISTENT_CTAS_PER_SM",
+    "BlockScaleLayout",
     "GridScheduler",
     "ProfileTag",
 }
@@ -24,6 +25,7 @@
     "Mxfp8TmaGemv",
     "get_mxfp8_tma_gemv",
     "mxfp8_tma_gemv",
+    "mxfp8_tma_scaled_mm",
     "select_mxfp8_tma_compute_warps",
 }
 

transformer_nuggets/cute/blockscaled_tma.py

@@ -8,11 +8,19 @@
 import cutlass.pipeline as pipeline
 from cuda.bindings import driver as cuda
 from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blockscaled_layout as blockscaled_utils
 
 from transformer_nuggets.cute.base import CuteOp
 from transformer_nuggets.cute.profiler.ops import profile_region
 
 
+class BlockScaleLayout(str, Enum):
+    """Describe caller-owned block-scale storage."""
+
+    RAW = "raw"
+    SWIZZLE_32_4_4 = "swizzle_32_4_4"
+
+
 class GridScheduler(str, Enum):
     """Map logical output tiles onto the launched CTA grid."""
 
@@ -127,6 +135,14 @@ def __init__(
         self.max_profile_events_per_cta = 2 + 3 * self.num_k_tiles
         self.num_profile_units = self.num_tiles
 
+    @cute.jit
+    def make_blocked_scale_layout(self, scale_vector_size: cutlass.Constexpr):
+        """Map logical matrix coordinates onto canonical blocked scale storage."""
+        return blockscaled_utils.tile_atom_to_shape_SF(
+            (((self.n + 127) // 128) * 128, self.k, 1),
+            scale_vector_size,
+        )
+
     def x_smem_layout(self):
         """Return the staged layout for one physical input tile."""
         return cute.make_ordered_layout(