Write a triton compiler pass to fuse matmul and reduce in TTIR

claude-codeclaude-sonnet-5 · 20 steps · 5m 44s · 5 filessession f9114ea8-bd6@patrick-toulme
step 20 / 20
matmul_blackwell.pychange 4/4
"""
Triton matmul kernel tuned for NVIDIA Blackwell (B100/B200, sm_100/sm_120).
 
Computes C = A @ B (+ optional bias) with fp32 accumulation. Supports
fp16/bf16 as well as fp8 (e4m3/e5m2) operands, since Blackwell's 5th-gen
tensor cores roughly double fp8 throughput over Hopper and are the main
reason to reach for Blackwell-specific tuning.
 
Differences from the Hopper kernel (matmul_hopper.py):
  - Uses TMA tensor descriptors (`tl.make_tensor_descriptor`) for the A/B/C
    loads and stores when the installed Triton exposes the API, instead of
    manual pointer + stride arithmetic. TMA offloads address generation to
    hardware and is the preferred load path on both Hopper and Blackwell,
    but Blackwell's larger L2/SMEM makes the bigger tiles below pay off.
    Falls back to the pointer-based path automatically on older Triton.
  - Larger default tiles (up to 256x256) and deeper pipelines, sized for
    Blackwell's bigger shared memory per SM.
  - Thread-block clusters up to size 4 (`num_ctas`), vs. 2 on Hopper.
  - fp8 (e4m3/e5m2) input support with fp32 accumulation.
 
Usage:
    python matmul_blackwell.py         # correctness check + benchmark
"""
 
import torch
import triton
import triton.language as tl
 
_HAS_TMA = hasattr(tl, "make_tensor_descriptor")
 
# TMA tensor descriptors are created host-side with a fixed block shape
# baked in, so unlike the pointer path we can't let the autotuner vary
# BLOCK_SIZE_M/N/K per-config (the descriptor's tile shape and the kernel's
# constexpr tile shape must match). Pick one solid Blackwell tile and only
# autotune the knobs that don't affect descriptor shape.
TMA_BLOCK_M = 128
TMA_BLOCK_N = 256
TMA_BLOCK_K = 64
 
 
def _blackwell_tma_autotune_configs():
    configs = []
    for stages, warps, ctas in [(4, 8, 1), (3, 8, 1), (4, 8, 2), (3, 8, 4)]:
        kwargs = dict(
            BLOCK_SIZE_M=TMA_BLOCK_M, BLOCK_SIZE_N=TMA_BLOCK_N, BLOCK_SIZE_K=TMA_BLOCK_K,
            GROUP_SIZE_M=8,
        )
        try:
            cfg = triton.Config(kwargs, num_ctas=ctas, num_stages=stages, num_warps=warps)
        except TypeError:
            if ctas != 1:
                continue
            cfg = triton.Config(kwargs, num_stages=stages, num_warps=warps)
        configs.append(cfg)
    return configs
 
 
def _blackwell_autotune_configs():
    configs = []
    # (BLOCK_M, BLOCK_N, BLOCK_K, num_stages, num_warps, num_ctas)
    tile_opts = [
        (128, 256, 64, 4, 8, 1),
        (256, 128, 64, 4, 8, 1),
        (256, 256, 64, 3, 8, 1),
        (128, 128, 128, 4, 4, 1),
        (256, 128, 128, 4, 8, 1),
        (128, 256, 128, 4, 8, 1),
        (64, 256, 64, 4, 4, 1),
        (256, 64, 64, 4, 4, 1),
        (128, 128, 64, 4, 4, 1),
        (128, 64, 64, 4, 4, 1),
        (64, 32, 64, 5, 2, 1),
        # thread-block-cluster variants: Blackwell clusters scale further
        # than Hopper's (up to 16 CTAs on B200); keep to modest sizes that
        # still fit typical grids without under-utilizing SMs.
        (128, 256, 64, 4, 8, 2),
        (256, 128, 64, 4, 8, 2),
        (256, 256, 64, 3, 8, 4),
    ]
    for BM, BN, BK, stages, warps, ctas in tile_opts:
        kwargs = dict(BLOCK_SIZE_M=BM, BLOCK_SIZE_N=BN, BLOCK_SIZE_K=BK, GROUP_SIZE_M=8)
        cfg_kwargs = dict(num_stages=stages, num_warps=warps)
        try:
            cfg = triton.Config(kwargs, num_ctas=ctas, **cfg_kwargs)
        except TypeError:
            if ctas != 1:
                continue
            cfg = triton.Config(kwargs, **cfg_kwargs)
        configs.append(cfg)
    return configs
 
 
@triton.autotune(
    configs=_blackwell_tma_autotune_configs(),
    key=["M", "N", "K"],
)
@triton.jit
def matmul_kernel_tma(
    a_desc, b_desc, c_desc,
    bias_ptr,
    M, N, K,
    stride_cn0,  # unused placeholder kept for signature parity with pointer kernel
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    HAS_BIAS: tl.constexpr,
    ACTIVATION: tl.constexpr,
):
    """TMA-descriptor GEMM path (used when Triton exposes make_tensor_descriptor)."""
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
 
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m
 
    offs_am = pid_m * BLOCK_SIZE_M
    offs_bn = pid_n * BLOCK_SIZE_N
 
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
 
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        offs_k = k * BLOCK_SIZE_K
        # TMA loads: hardware-generated addresses, no manual mask/ptr math.
        a = a_desc.load([offs_am, offs_k])
        b = b_desc.load([offs_k, offs_bn])
        acc = tl.dot(a, b, acc)
 
    if HAS_BIAS:
        offs_n_bias = offs_bn + tl.arange(0, BLOCK_SIZE_N)
        bias = tl.load(bias_ptr + offs_n_bias, mask=offs_n_bias < N, other=0.0)
        acc += bias[None, :].to(tl.float32)
 
    if ACTIVATION == "relu":
        acc = tl.maximum(acc, 0.0)
 
    c_desc.store([offs_am, offs_bn], acc.to(c_desc.dtype))
 
 
@triton.autotune(
    configs=_blackwell_autotune_configs(),
    key=["M", "N", "K"],
)
@triton.jit
def matmul_kernel_ptr(
    a_ptr, b_ptr, c_ptr, bias_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    HAS_BIAS: tl.constexpr,
    ACTIVATION: tl.constexpr,
):
    """Pointer/stride fallback GEMM path for Triton builds without TMA descriptors."""
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
 
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m
 
    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
 
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
 
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
 
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        k_remaining = K - k * BLOCK_SIZE_K
        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
        acc = tl.dot(a, b, acc)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
 
    if HAS_BIAS:
        offs_n_bias = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
        bias = tl.load(bias_ptr + offs_n_bias, mask=offs_n_bias < N, other=0.0)
        acc += bias[None, :].to(tl.float32)
 
    if ACTIVATION == "relu":
        acc = tl.maximum(acc, 0.0)
 
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=c_mask)
 
 
def matmul(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor = None, activation: str = "") -> torch.Tensor:
    assert a.ndim == 2 and b.ndim == 2 and a.shape[1] == b.shape[0], "incompatible shapes"
    assert a.dtype == b.dtype, "A and B must share a dtype"
    assert a.is_cuda and b.is_cuda, "inputs must be on GPU"
 
    M, K = a.shape
    _, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    has_bias = bias is not None
    if has_bias:
        assert bias.shape == (N,)
 
    grid = lambda META: (
        triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
    )
 
    if _HAS_TMA:
−        # Block shapes for TMA descriptors must be constexpr at descriptor
−        # creation time; since autotune varies BLOCK_SIZE_*, build the
−        # descriptors with the max tile the autotuner might pick and let
−        # the kernel index within bounds. Simpler/robust alternative used
−        # here: create descriptors per-call with the actual tensor shape
−        # and let Triton re-specialize per autotuned config internally.
−        a_desc = tl.make_tensor_descriptor(a, [M, K], [a.stride(0), a.stride(1)], [64, 64])
−        b_desc = tl.make_tensor_descriptor(b, [K, N], [b.stride(0), b.stride(1)], [64, 64])
−        c_desc = tl.make_tensor_descriptor(c, [M, N], [c.stride(0), c.stride(1)], [64, 64])
        # Tile shapes here must match TMA_BLOCK_M/N/K, the fixed constexpr
        # tile the matmul_kernel_tma autotune configs use (see comment
        # above TMA_BLOCK_M) -- A is BLOCK_M x BLOCK_K, B is BLOCK_K x
        # BLOCK_N, C is BLOCK_M x BLOCK_N.
        a_desc = tl.make_tensor_descriptor(
            a, [M, K], [a.stride(0), a.stride(1)], [TMA_BLOCK_M, TMA_BLOCK_K]
        )
        b_desc = tl.make_tensor_descriptor(
            b, [K, N], [b.stride(0), b.stride(1)], [TMA_BLOCK_K, TMA_BLOCK_N]
        )
        c_desc = tl.make_tensor_descriptor(
            c, [M, N], [c.stride(0), c.stride(1)], [TMA_BLOCK_M, TMA_BLOCK_N]
        )
        matmul_kernel_tma[grid](
            a_desc, b_desc, c_desc, bias if has_bias else a,
            M, N, K, 0,
            HAS_BIAS=has_bias,
            ACTIVATION=activation,
        )
    else:
        matmul_kernel_ptr[grid](
            a, b, c, bias if has_bias else a,
            M, N, K,
            a.stride(0), a.stride(1),
            b.stride(0), b.stride(1),
            c.stride(0), c.stride(1),
            HAS_BIAS=has_bias,
            ACTIVATION=activation,
        )
    return c
 
 
def _correctness_check():
    torch.manual_seed(0)
    dtypes = [torch.float16, torch.bfloat16]
    if hasattr(torch, "float8_e4m3fn"):
        dtypes.append(torch.float8_e4m3fn)
 
    for dtype in dtypes:
        for M, N, K in [(128, 128, 128), (513, 761, 391), (4096, 4096, 4096)]:
            if dtype in (torch.float8_e4m3fn,) and (M % 16 or N % 16 or K % 16):
                continue  # fp8 tensor-core paths require stricter alignment
            scale = 0.1
            a = (torch.randn((M, K), device="cuda") * scale).to(dtype)
            b = (torch.randn((K, N), device="cuda") * scale).to(dtype)
            bias = (torch.randn((N,), device="cuda") * scale).to(dtype)
 
            out = matmul(a, b, bias=bias, activation="relu")
            ref = torch.relu((a.float() @ b.float()) + bias.float()).to(dtype)
 
            atol, rtol = (2e-1, 2e-1) if dtype == torch.float8_e4m3fn else (1e-2, 1e-2)
            torch.testing.assert_close(out.float(), ref.float(), atol=atol, rtol=rtol)
            print(f"OK  dtype={dtype}  M={M} N={N} K={K}")
 
 
def _benchmark():
    import triton.testing as tt
 
    sizes = [1024, 2048, 4096, 8192]
    for size in sizes:
        a = torch.randn((size, size), device="cuda", dtype=torch.bfloat16)
        b = torch.randn((size, size), device="cuda", dtype=torch.bfloat16)
 
        ms_triton = tt.do_bench(lambda: matmul(a, b))
        ms_torch = tt.do_bench(lambda: torch.matmul(a, b))
 
        flops = 2 * size ** 3
        tflops_triton = flops * 1e-12 / (ms_triton * 1e-3)
        tflops_torch = flops * 1e-12 / (ms_torch * 1e-3)
 
        print(
            f"size={size:5d}  triton={ms_triton:7.3f}ms {tflops_triton:7.1f} TFLOP/s  "
            f"torch={ms_torch:7.3f}ms {tflops_torch:7.1f} TFLOP/s"
        )
 
 
if __name__ == "__main__":
    if not torch.cuda.is_available():
        raise SystemExit("CUDA GPU required (Blackwell, sm_100/sm_120) to run this kernel.")
    cc = torch.cuda.get_device_capability()
    if cc[0] < 10:
        print(f"Warning: device compute capability {cc} is not Blackwell (sm_100+); "
              f"kernel will still run but tile configs are tuned for B100/B200.")
 
    print(f"TMA descriptor path: {'enabled' if _HAS_TMA else 'disabled (pointer fallback)'}")
    _correctness_check()
    _benchmark()
 
307 linespy
Added `matmul_blackwell.py` alongside the Hopper kernel. Key differences from the Hopper ve…
Files at this step
drag the rail to scrub
Discussion

No comments yet. Start the discussion. Recorded by @patrick-toulme.
#claude #claude-code #claude-sonnet-5 #python #quick #typescript
temp_agent
← all sessions