Stage compact NVFP4 weight scales with TMA

Author: drisspg

benchmarks/tma_m1_scaled_mm_heatmap.py

@@ -21,6 +21,7 @@
     select_mxfp8_tma_compute_warps,
     select_nvfp4_tma_config,
     select_nvfp4_tma_split_k,
+    select_nvfp4_tma_stage_weight_scales,
 )
 from transformer_nuggets.utils.benchmark import benchmark_cuda_function_in_microseconds
 
@@ -35,6 +36,7 @@ class KernelConfig:
     num_compute_warps: int
     grid_scheduler: GridScheduler
     split_k: int = 1
+    stage_weight_scales: bool = False
 
 
 @dataclass(frozen=True)
@@ -50,6 +52,7 @@ class BenchmarkResult:
     num_compute_warps: int | None
     grid_scheduler: str
     split_k: int
+    stage_weight_scales: bool
 
 
 def parse_sizes(value: str, dimension: str, multiple: int) -> list[int]:
@@ -168,12 +171,22 @@ def select_nvfp4_config(n: int, k: int, device: torch.device) -> KernelConfig:
     elif n >= 12288 and k == 12288:
         block_n, num_compute_warps = 16, 4
     split_k = select_nvfp4_tma_split_k(n, k, device)
+    stage_weight_scales = select_nvfp4_tma_stage_weight_scales(
+        n,
+        k,
+        block_n,
+        num_compute_warps,
+        GridScheduler.STATIC,
+        split_k,
+        device,
+    )
     return KernelConfig(
         block_n,
         num_stages,
         num_compute_warps,
         GridScheduler.STATIC,
         split_k,
+        stage_weight_scales,
     )
 
 
@@ -266,6 +279,7 @@ def run_mxfp8(n: int, k: int, rounds: int, iterations: int, seed: int) -> Benchm
         config.num_compute_warps,
         config.grid_scheduler.value,
         config.split_k,
+        config.stage_weight_scales,
     )
 
 
@@ -289,6 +303,7 @@ def run_nvfp4(n: int, k: int, rounds: int, iterations: int, seed: int) -> Benchm
         num_compute_warps=config.num_compute_warps,
         grid_scheduler=config.grid_scheduler,
         split_k=config.split_k,
+        stage_weight_scales=config.stage_weight_scales,
         output=output,
         partial_output=partial_output,
     )
@@ -308,6 +323,7 @@ def run_nvfp4(n: int, k: int, rounds: int, iterations: int, seed: int) -> Benchm
         config.num_compute_warps,
         config.grid_scheduler.value,
         config.split_k,
+        config.stage_weight_scales,
     )
 
 

test/test_nvfp4_tma.py

@@ -19,6 +19,7 @@
         select_nvfp4_tma_compute_warps,
         select_nvfp4_tma_config,
         select_nvfp4_tma_split_k,
+        select_nvfp4_tma_stage_weight_scales,
     )
     from transformer_nuggets.cute.profiler import profile_session
     from transformer_nuggets.cute.profiler.host import decode_events
@@ -54,6 +55,7 @@ def test_select_nvfp4_tma_compute_warps(k, block_n, expected):
         (14336, 16384, (8, 3, 4)),
         (24576, 24576, (8, 2, 4)),
         (32768, 32768, (8, 3, 4)),
+        (16384, 8192, (8, 2, 4)),
     ],
 )
 def test_select_nvfp4_tma_config(n, k, expected):
@@ -78,6 +80,38 @@ def test_select_nvfp4_tma_split_k(n, k, expected):
     assert select_nvfp4_tma_split_k(n, k) == expected
 
 
+@pytest.mark.parametrize(
+    ("n", "k", "block_n", "num_compute_warps", "grid_scheduler", "split_k", "expected"),
+    [
+        (16384, 6144, 8, 4, GridScheduler.STATIC, 1, True),
+        (32768, 8192, 8, 4, GridScheduler.STATIC, 1, True),
+        (14336, 8192, 8, 4, GridScheduler.STATIC, 1, False),
+        (16384, 12288, 8, 4, GridScheduler.STATIC, 1, False),
+        (16384, 8192, 16, 4, GridScheduler.STATIC, 1, False),
+        (16384, 8192, 8, 2, GridScheduler.STATIC, 1, False),
+        (16384, 8192, 8, 4, GridScheduler.PERSISTENT, 1, False),
+        (16384, 8192, 8, 4, GridScheduler.STATIC, 2, False),
+    ],
+)
+def test_select_nvfp4_tma_stage_weight_scales(
+    n, k, block_n, num_compute_warps, grid_scheduler, split_k, expected
+):
+    """Stage physical scale subsets only in the measured B200 regime."""
+    if torch.cuda.get_device_capability() != (10, 0):
+        expected = False
+    assert (
+        select_nvfp4_tma_stage_weight_scales(
+            n,
+            k,
+            block_n,
+            num_compute_warps,
+            grid_scheduler,
+            split_k,
+        )
+        is expected
+    )
+
+
 def pack_fp4(codes: torch.Tensor) -> torch.Tensor:
     """Pack low-nibble-first E2M1 codes into the PyTorch FP4 shell dtype."""
     packed = codes[:, 0::2] | (codes[:, 1::2] << 4)
@@ -178,6 +212,23 @@ def test_nvfp4_tma_scale_layout_and_paired_loads(block_n, num_compute_warps):
     torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
 
 
+def test_nvfp4_tma_compact_scale_tma_matches_blocked_layout():
+    """Match exact scales across physical 128-row atoms with compact TMA staging."""
+    q_input, weight, input_scale, weight_scale, expected, _ = make_case(264, 6144)
+    actual = nvfp4_tma_gemv(
+        q_input,
+        weight,
+        input_scale,
+        weight_scale,
+        block_n=8,
+        num_stages=3,
+        num_compute_warps=4,
+        stage_weight_scales=True,
+    )
+    torch.cuda.synchronize()
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
 def test_nvfp4_tma_fp16_partial_reduction_preserves_extreme_values():
     """Keep exact FP16 partial sums before accumulating them in FP32."""
     n, k = 128, 2048
@@ -374,8 +425,11 @@ def test_nvfp4_tma_split_k_accepts_scalar_aligned_workspace():
     torch.testing.assert_close(output, expected, atol=2.0, rtol=0.05)
 
 
-@pytest.mark.parametrize(("k", "split_k", "num_stages"), [(8192, 2, 2), (16384, 4, 3)])
-def test_nvfp4_tma_split_k_cuda_graph(k, split_k, num_stages):
+@pytest.mark.parametrize(
+    ("k", "split_k", "num_stages", "stage_weight_scales"),
+    [(8192, 2, 2, True), (16384, 4, 3, False)],
+)
+def test_nvfp4_tma_split_k_cuda_graph(k, split_k, num_stages, stage_weight_scales):
     """Reduce parallel K partitions in FP32 before the final BF16 conversion."""
     n = 128
     q_input, weight, input_scale, weight_scale, expected, _ = make_case(n, k)
@@ -386,6 +440,7 @@ def test_nvfp4_tma_split_k_cuda_graph(k, split_k, num_stages):
         "num_compute_warps": 4,
         "num_stages": num_stages,
         "split_k": split_k,
+        "stage_weight_scales": stage_weight_scales,
         "output": output,
         "partial_output": partial_output,
     }
@@ -398,6 +453,24 @@ def test_nvfp4_tma_split_k_cuda_graph(k, split_k, num_stages):
     torch.testing.assert_close(output, expected, atol=2.0, rtol=0.05)
 
 
+def test_nvfp4_tma_compact_scale_tma_persistent_reuse():
+    """Reuse compact scale stages across persistent output tiles."""
+    q_input, weight, input_scale, weight_scale, expected, _ = make_case(256, 2048)
+    actual = nvfp4_tma_gemv(
+        q_input,
+        weight,
+        input_scale,
+        weight_scale,
+        block_n=8,
+        num_compute_warps=4,
+        grid_scheduler=GridScheduler.PERSISTENT,
+        num_persistent_ctas=4,
+        stage_weight_scales=True,
+    )
+    torch.cuda.synchronize()
+    torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
+
+
 def test_nvfp4_tma_persistent_cuda_graph():
     """Replay persistent NVFP4 GEMV into caller-owned output."""
     q_input, weight, input_scale, weight_scale, expected, _ = make_case(256, 2048)

transformer_nuggets/cute/__init__.py

@@ -38,6 +38,7 @@
     "select_nvfp4_tma_compute_warps",
     "select_nvfp4_tma_config",
     "select_nvfp4_tma_split_k",
+    "select_nvfp4_tma_stage_weight_scales",
 }
 
 _SYMMETRIC_MEMORY_EXPORTS = {