[TEST] Cover Gluon Blackwell tcgen05 scaled ops

tests/end_to_end/test_gluon.py

@@ -82,6 +82,25 @@ def mangle(self):
         return f"TL{self.block[0]}x{self.block[1]}C{self.col_stride}{cga_layout_str}{two_ctas_str}TL"
 
 
+class _TensorMemoryScalesLayoutForCpu(blackwell.TensorMemoryScalesLayout):
+    def __init__(self, cga_layout=None):
+        object.__setattr__(
+            self,
+            "cga_layout",
+            [] if cga_layout is None else [list(basis) for basis in cga_layout],
+        )
+        object.__setattr__(self, "cta_split_num", None)
+
+    def _to_ir(self, builder):
+        return builder.get_tensor_memory_scales_layout(self.cga_layout)
+
+    def mangle(self):
+        cga_layout_str = "_".join(
+            "~".join(map(str, basis)) for basis in self.cga_layout
+        )
+        return f"TLS{cga_layout_str}TLS"
+
+
 @gluon.jit
 def _copy_scalar_kernel(in_ptr, out_ptr):
     value = gl.load(in_ptr)
@@ -761,6 +780,128 @@ def _run_tcgen05_small_mma_case(seed: int, lhs_in_tmem: bool, use_commit: bool):
     torch.testing.assert_close(out, a.float() @ b.float() + c, atol=1e-2, rtol=1e-2)
 
 
+@gluon.jit
+def _tcgen05_scaled_mma_kernel(
+    a_ptr,
+    b_ptr,
+    a_scale_ptr,
+    b_scale_ptr,
+    out_ptr,
+    M: gl.constexpr,
+    N: gl.constexpr,
+    K: gl.constexpr,
+    VEC: gl.constexpr,
+    A_TYPE: gl.constexpr,
+    B_TYPE: gl.constexpr,
+    num_warps: gl.constexpr,
+):
+    layout: gl.constexpr = gl.BlockedLayout(
+        [1, 1],
+        [1, 32],
+        [1, gl.num_warps()],
+        [1, 0],
+    )
+    offs_m = gl.arange(0, M, gl.SliceLayout(1, layout))
+    offs_n = gl.arange(0, N, gl.SliceLayout(0, layout))
+    offs_n_rows = gl.arange(0, N, gl.SliceLayout(1, layout))
+    offs_k = gl.arange(0, K, gl.SliceLayout(0, layout))
+    offs_k_rows = gl.arange(0, K, gl.SliceLayout(1, layout))
+    offs_scale = gl.arange(0, K // VEC, gl.SliceLayout(0, layout))
+
+    a = gl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :])
+    b = gl.load(b_ptr + offs_k_rows[:, None] * N + offs_n[None, :])
+    a_scales = gl.load(a_scale_ptr + offs_m[:, None] * (K // VEC) + offs_scale[None, :])
+    b_scales = gl.load(
+        b_scale_ptr + offs_n_rows[:, None] * (K // VEC) + offs_scale[None, :]
+    )
+
+    a_smem_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for(
+        [M, K],
+        a_ptr.dtype.element_ty,
+    )
+    b_smem_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for(
+        [K, N],
+        b_ptr.dtype.element_ty,
+    )
+    a_smem = gl.allocate_shared_memory(a_ptr.dtype.element_ty, [M, K], a_smem_layout)
+    b_smem = gl.allocate_shared_memory(b_ptr.dtype.element_ty, [K, N], b_smem_layout)
+    a_smem.store(a)
+    b_smem.store(b)
+
+    scale_layout: gl.constexpr = _TensorMemoryScalesLayoutForCpu([[1, 0], [0, 1]])
+    a_scale_tmem = blackwell.allocate_tensor_memory(
+        a_scale_ptr.dtype.element_ty,
+        [M, K // VEC],
+        scale_layout,
+    )
+    b_scale_tmem = blackwell.allocate_tensor_memory(
+        b_scale_ptr.dtype.element_ty,
+        [N, K // VEC],
+        scale_layout,
+    )
+    a_scale_tmem.store(a_scales)
+    b_scale_tmem.store(b_scales)
+
+    acc_layout: gl.constexpr = _TensorMemoryLayoutForCpu(
+        [M, N],
+        col_stride=1,
+        cga_layout=[[1, 0], [0, 1]],
+    )
+    acc = blackwell.allocate_tensor_memory(gl.float32, [M, N], acc_layout)
+    blackwell.tcgen05_mma_scaled(
+        a_smem,
+        b_smem,
+        acc,
+        a_scale_tmem,
+        b_scale_tmem,
+        A_TYPE,
+        B_TYPE,
+        use_acc=False,
+    )
+    result = acc.load(layout)
+    gl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], result)
+
+
+def _run_tcgen05_scaled_mma_case():
+    torch.manual_seed(7)
+    a = torch.randn(128, 4, dtype=torch.float32) / 4
+    b = torch.randn(4, 16, dtype=torch.float32) / 4
+    a_scale = (
+        torch.arange(a.shape[0] * 2, dtype=torch.uint8).reshape(a.shape[0], 2) % 3
+    ) + 126
+    b_scale = (
+        torch.arange(b.shape[1] * 2, dtype=torch.uint8).reshape(b.shape[1], 2) % 3
+    ) + 126
+    out = torch.empty((a.shape[0], b.shape[1]), dtype=torch.float32)
+    _run_gluon_on_cpu(
+        _tcgen05_scaled_mma_kernel,
+        (1,),
+        a,
+        b,
+        a_scale,
+        b_scale,
+        out,
+        a.shape[0],
+        b.shape[1],
+        a.shape[1],
+        2,
+        "e4m3",
+        "e4m3",
+        num_warps=4,
+    )
+
+    a_scale_float = torch.pow(2.0, a_scale.float() - 127.0).repeat_interleave(
+        2,
+        dim=1,
+    )
+    b_scale_float = torch.pow(2.0, b_scale.float() - 127.0).repeat_interleave(
+        2,
+        dim=1,
+    )
+    expected = (a * a_scale_float) @ (b * b_scale_float.T)
+    torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
+
+
 @gluon.jit
 def _tensor_memory_roundtrip_kernel(
     in_ptr,
@@ -1117,6 +1258,10 @@ def test_gluon_blackwell_tcgen05_runs_small_mma_with_lhs_tensor_memory_on_cpu():
     _run_tcgen05_small_mma_case(seed=6, lhs_in_tmem=True, use_commit=False)
 
 
+def test_gluon_blackwell_tcgen05_scaled_mma_uses_scale_groups_on_cpu():
+    _run_tcgen05_scaled_mma_case()
+
+
 def test_gluon_builder_preserves_tensor_memory_fp4_padding():
     layout = gluon_builder.get_tensor_memory_layout(
         (64, 64),

triton_viz/core/simulation/gluon.py

@@ -654,6 +654,8 @@ def _descriptor_from_host(value: Any) -> TensorDescriptor:
 
 def _implicit_gluon_cvt(name: str, value: Any, constexprs: set[str]) -> Any:
     if name in constexprs:
+        if isinstance(value, str):
+            return tl.constexpr(value)
         return value
     if isinstance(value, TensorDescriptor) or (
         type(value).__name__ in _HOST_TENSOR_DESCRIPTOR_TYPES
@@ -2628,7 +2630,7 @@ def __init__(self, fn: Callable, arg_names: list[str] | None = None) -> None:
         self.fn = fn
         signature = inspect.signature(fn)
         self.arg_names = arg_names or [v.name for v in signature.parameters.values()]
-        annotations = fn.__annotations__
+        annotations = getattr(fn, "fn", fn).__annotations__
         self.constexprs = {
             name
             for name in self.arg_names