diff --git a/tests/end_to_end/test_gluon.py b/tests/end_to_end/test_gluon.py index 8869bf6c..20fea324 100644 --- a/tests/end_to_end/test_gluon.py +++ b/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), diff --git a/triton_viz/core/simulation/gluon.py b/triton_viz/core/simulation/gluon.py index 4960bdb4..f658194c 100644 --- a/triton_viz/core/simulation/gluon.py +++ b/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