Fix NVFP4 benchmark scaling and explicit CuTe exports

Author: drisspg

benchmarks/tma_m1_scaled_mm_heatmap.py

@@ -208,18 +208,23 @@ def scaled_mm_mxfp8(
 
 
 def scaled_mm_nvfp4(
-    mat_a: torch.Tensor, mat_b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+    scale_a: torch.Tensor,
+    scale_b: torch.Tensor,
+    global_scale_a: torch.Tensor,
+    global_scale_b: torch.Tensor,
 ):
-    """Run the cuBLASLt NVFP4 block-scaled GEMV baseline."""
+    """Run the cuBLASLt NVFP4 block- and tensor-scaled GEMV baseline."""
     return scaled_mm(
         mat_a,
         mat_b,
-        scale_a=[scale_a],
-        scale_recipe_a=[ScalingType.BlockWise1x16],
-        swizzle_a=[SwizzleType.SWIZZLE_32_4_4],
-        scale_b=[scale_b],
-        scale_recipe_b=[ScalingType.BlockWise1x16],
-        swizzle_b=[SwizzleType.SWIZZLE_32_4_4],
+        scale_a=[scale_a, global_scale_a],
+        scale_recipe_a=[ScalingType.BlockWise1x16, ScalingType.TensorWise],
+        swizzle_a=[SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE],
+        scale_b=[scale_b, global_scale_b],
+        scale_recipe_b=[ScalingType.BlockWise1x16, ScalingType.TensorWise],
+        swizzle_b=[SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE],
         output_dtype=torch.bfloat16,
     )
 
@@ -286,6 +291,8 @@ def run_mxfp8(n: int, k: int, rounds: int, iterations: int, seed: int) -> Benchm
 def run_nvfp4(n: int, k: int, rounds: int, iterations: int, seed: int) -> BenchmarkResult:
     """Benchmark the NVFP4 TMA specialization against the matching scaled_mm contract."""
     mat_a, mat_b, scale_a, scale_b = make_nvfp4_case(n, k, seed)
+    global_scale_a = torch.tensor([1.5], dtype=torch.float32, device="cuda")
+    global_scale_b = torch.tensor([0.75], dtype=torch.float32, device="cuda")
     config = select_nvfp4_config(n, k, mat_a.device)
     output = torch.empty((1, n), dtype=torch.bfloat16, device="cuda")
     partial_output = (
@@ -304,10 +311,19 @@ def run_nvfp4(n: int, k: int, rounds: int, iterations: int, seed: int) -> Benchm
         grid_scheduler=config.grid_scheduler,
         split_k=config.split_k,
         stage_weight_scales=config.stage_weight_scales,
+        global_scale_a=global_scale_a,
+        global_scale_b=global_scale_b,
         output=output,
         partial_output=partial_output,
     )
-    baseline = lambda: scaled_mm_nvfp4(mat_a, mat_b, scale_a, scale_b)
+    baseline = lambda: scaled_mm_nvfp4(
+        mat_a,
+        mat_b,
+        scale_a,
+        scale_b,
+        global_scale_a,
+        global_scale_b,
+    )
     torch.testing.assert_close(tma(), baseline(), atol=2.0, rtol=0.05)
     torch.cuda.synchronize()
     tma_us, scaled_mm_us = median_interleaved_latency(tma, baseline, rounds, iterations)

transformer_nuggets/cute/__init__.py

@@ -11,60 +11,35 @@
 from transformer_nuggets.cute.element_wise import ElementwiseOp, elementwise_op
 from transformer_nuggets.cute.utils import visualize_tv_layout
 from transformer_nuggets.cute import profiler
-
-
-_BLOCKSCALED_TMA_EXPORTS = {
-    "DEFAULT_PERSISTENT_CTAS_PER_SM",
-    "BlockScaleLayout",
-    "GridScheduler",
-    "ProfileTag",
-}
-
-_MXFP8_TMA_EXPORTS = {
-    "MXFP8_TMA_PROFILE_TAGS",
-    "Mxfp8TmaGemv",
-    "get_mxfp8_tma_gemv",
-    "mxfp8_tma_gemv",
-    "mxfp8_tma_scaled_mm",
-    "select_mxfp8_tma_compute_warps",
-}
-
-_NVFP4_TMA_EXPORTS = {
-    "NVFP4_TMA_PROFILE_TAGS",
-    "Nvfp4TmaGemv",
-    "get_nvfp4_tma_gemv",
-    "nvfp4_tma_gemv",
-    "nvfp4_tma_scaled_mm",
-    "select_nvfp4_tma_compute_warps",
-    "select_nvfp4_tma_config",
-    "select_nvfp4_tma_split_k",
-    "select_nvfp4_tma_stage_weight_scales",
-}
-
-_SYMMETRIC_MEMORY_EXPORTS = {
-    "compile_symmetric_memory_all_reduce",
-    "init_torchrun_process_group",
-    "run_symmetric_memory_all_reduce_example",
-    "symmetric_memory_all_reduce",
-    "symmetric_memory_peer_tensors",
-}
-
-
-def __getattr__(name):
-    if name in _BLOCKSCALED_TMA_EXPORTS:
-        from transformer_nuggets.cute import blockscaled_tma
-
-        return getattr(blockscaled_tma, name)
-    if name in _MXFP8_TMA_EXPORTS:
-        from transformer_nuggets.cute import mxfp8_tma
-
-        return getattr(mxfp8_tma, name)
-    if name in _NVFP4_TMA_EXPORTS:
-        from transformer_nuggets.cute import nvfp4_tma
-
-        return getattr(nvfp4_tma, name)
-    if name in _SYMMETRIC_MEMORY_EXPORTS:
-        from transformer_nuggets.cute import symmetric_memory
-
-        return getattr(symmetric_memory, name)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+from transformer_nuggets.cute.blockscaled_tma import (
+    DEFAULT_PERSISTENT_CTAS_PER_SM,
+    BlockScaleLayout,
+    GridScheduler,
+    ProfileTag,
+)
+from transformer_nuggets.cute.mxfp8_tma import (
+    MXFP8_TMA_PROFILE_TAGS,
+    Mxfp8TmaGemv,
+    get_mxfp8_tma_gemv,
+    mxfp8_tma_gemv,
+    mxfp8_tma_scaled_mm,
+    select_mxfp8_tma_compute_warps,
+)
+from transformer_nuggets.cute.nvfp4_tma import (
+    NVFP4_TMA_PROFILE_TAGS,
+    Nvfp4TmaGemv,
+    get_nvfp4_tma_gemv,
+    nvfp4_tma_gemv,
+    nvfp4_tma_scaled_mm,
+    select_nvfp4_tma_compute_warps,
+    select_nvfp4_tma_config,
+    select_nvfp4_tma_split_k,
+    select_nvfp4_tma_stage_weight_scales,
+)
+from transformer_nuggets.cute.symmetric_memory import (
+    compile_symmetric_memory_all_reduce,
+    init_torchrun_process_group,
+    run_symmetric_memory_all_reduce_example,
+    symmetric_memory_all_reduce,
+    symmetric_memory_peer_tensors,
+)