Add Blackwell MXFP8 TMA GEMV

Author: drisspg

test/test_mxfp8_tma.py

@@ -0,0 +1,136 @@
+import pytest
+import torch
+
+
+if not torch.cuda.is_available():
+    pytest.skip("CUDA not available", allow_module_level=True)
+if torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}:
+    pytest.skip("MXFP8 TMA GEMV requires SM100 or SM103", allow_module_level=True)
+
+try:
+    from transformer_nuggets.cute import mxfp8_tma_gemv
+except ImportError:
+    pytest.skip("CuTe DSL not available", allow_module_level=True)
+
+
+def quantize_mxfp8(value: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize rows into E4M3 values with raw E8M0 block scales."""
+    blocks = value.float().reshape(value.shape[0], -1, 32)
+    exponent = torch.ceil(torch.log2(blocks.abs().amax(dim=-1) / 448.0)).clamp(-126, 127)
+    scale = torch.exp2(exponent).unsqueeze(-1)
+    quantized = (blocks / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
+    return quantized.reshape_as(value), (exponent + 127).to(torch.uint8)
+
+
+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)
+    return value.float() * expanded_scale
+
+
+@pytest.mark.parametrize(
+    ("k", "block_n", "num_stages"),
+    [(2048, 4, 2), (4096, 8, 3)],
+)
+def test_mxfp8_tma_gemv_matches_reference(k, block_n, num_stages):
+    """Match independently dequantized float32 matmul."""
+    torch.manual_seed(k)
+    q_input, input_scale = quantize_mxfp8(torch.randn((1, k), dtype=torch.bfloat16, device="cuda"))
+    weight, weight_scale = quantize_mxfp8(
+        torch.randn((128, k), dtype=torch.bfloat16, device="cuda")
+    )
+    expected = (
+        dequantize_mxfp8(q_input, input_scale) @ dequantize_mxfp8(weight, weight_scale).T
+    ).bfloat16()
+    output = torch.empty_like(expected)
+
+    actual = mxfp8_tma_gemv(
+        q_input,
+        weight,
+        input_scale,
+        weight_scale,
+        block_n=block_n,
+        num_stages=num_stages,
+        output=output,
+    )
+    torch.cuda.synchronize()
+
+    assert actual is output
+    torch.testing.assert_close(actual, expected, atol=1.0, rtol=0.05)
+
+
+def test_mxfp8_tma_gemv_cuda_graph_replay():
+    """Replay into caller-owned output without hidden allocation or copies."""
+    k = 2048
+    q_input, input_scale = quantize_mxfp8(torch.randn((1, k), dtype=torch.bfloat16, device="cuda"))
+    weight, weight_scale = quantize_mxfp8(
+        torch.randn((128, k), dtype=torch.bfloat16, device="cuda")
+    )
+    output = torch.empty((1, 128), dtype=torch.bfloat16, device="cuda")
+    mxfp8_tma_gemv(
+        q_input,
+        weight,
+        input_scale,
+        weight_scale,
+        block_n=4,
+        output=output,
+    )
+    graph = torch.cuda.CUDAGraph()
+
+    with torch.cuda.graph(graph):
+        mxfp8_tma_gemv(
+            q_input,
+            weight,
+            input_scale,
+            weight_scale,
+            block_n=4,
+            output=output,
+        )
+    graph.replay()
+    torch.cuda.synchronize()
+
+    expected = (
+        dequantize_mxfp8(q_input, input_scale) @ dequantize_mxfp8(weight, weight_scale).T
+    ).bfloat16()
+    torch.testing.assert_close(output, expected, atol=1.0, rtol=0.05)
+
+
+def test_mxfp8_tma_gemv_combines_cancelling_scales():
+    """Avoid an infinite intermediate when E8M0 scale exponents cancel."""
+    k = 2048
+    q_input = torch.ones((1, k), dtype=torch.float8_e4m3fn, device="cuda")
+    weight = torch.ones((128, k), dtype=torch.float8_e4m3fn, device="cuda")
+    input_scale = torch.full((1, k // 32), 254, dtype=torch.uint8, device="cuda")
+    weight_scale = torch.zeros((128, k // 32), dtype=torch.uint8, device="cuda")
+
+    actual = mxfp8_tma_gemv(
+        q_input,
+        weight,
+        input_scale,
+        weight_scale,
+        block_n=4,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(actual, torch.full_like(actual, k), rtol=0, atol=0)
+
+
+def test_mxfp8_tma_gemv_preserves_nan_scale():
+    """Decode the reserved E8M0 byte as NaN rather than infinity."""
+    k = 2048
+    q_input = torch.ones((1, k), dtype=torch.float8_e4m3fn, device="cuda")
+    weight = torch.ones((128, k), dtype=torch.float8_e4m3fn, device="cuda")
+    input_scale = torch.full((1, k // 32), 127, dtype=torch.uint8, device="cuda")
+    weight_scale = torch.full((128, k // 32), 127, dtype=torch.uint8, device="cuda")
+    input_scale[:, 0] = 0xFF
+
+    actual = mxfp8_tma_gemv(
+        q_input,
+        weight,
+        input_scale,
+        weight_scale,
+        block_n=4,
+    )
+    torch.cuda.synchronize()
+
+    assert torch.isnan(actual).all()

transformer_nuggets/cute/__init__.py

@@ -13,6 +13,12 @@
 from transformer_nuggets.cute import profiler
 
 
+_MXFP8_TMA_EXPORTS = {
+    "Mxfp8TmaGemv",
+    "get_mxfp8_tma_gemv",
+    "mxfp8_tma_gemv",
+}
+
 _SYMMETRIC_MEMORY_EXPORTS = {
     "compile_symmetric_memory_all_reduce",
     "init_torchrun_process_group",
@@ -23,6 +29,10 @@
 
 
 def __getattr__(name):
+    if name in _MXFP8_TMA_EXPORTS:
+        from transformer_nuggets.cute import mxfp8_tma
+
+        return getattr(mxfp8_tma, name)
     if name in _SYMMETRIC_MEMORY_EXPORTS:
         from transformer_nuggets.cute import symmetric_memory