▶️ Run this notebook in Colab

Triton tiled matmul

Track 07 - GPU · Notebook 03 · Runtime: ≈15 min on T4 or better

Prerequisites: 07_gpu/01 (roofline), 07_gpu/02 (Triton 101 softmax).

Paper / reference: Triton 2.0 tutorial §matmul, Tillet et al. 2019.

---

What

A Triton kernel that computes C = A @ B by tiling the three loops (M, N, K) into BLOCK_M × BLOCK_N × BLOCK_K chunks that fit in SRAM. Target: ≥70 % of cuBLAS FP16 throughput at 4096³.

Two things make this more than a naive triple loop:

• Program grid is 2-D over (M, N) tiles. Each program instance computes one BLOCK_M × BLOCK_N output tile. The K loop stays sequential inside each program.

• K is accumulated in FP32 even when inputs are FP16. Tensor-core matmul products need an FP32 accumulator to avoid precision loss; Triton handles this with tl.dot(..., acc_dtype=tl.float32) (the default for FP16 inputs).

No CUDA device → the notebook records a skip for each kernel check and exits clean.

from llm_systems_cookbook.nb import bootstrap

import time

import torch

IS_CUDA = torch.cuda.is_available()

s = bootstrap("07_gpu_03_triton_tiled_matmul")

The kernel

One Triton program computes one BLOCK_M × BLOCK_N output tile. The inner K loop advances in strides of BLOCK_K, loading matching tiles from A and B, doing a tile-level tl.dot, and accumulating into an FP32 register block.

Grouped program ordering (GROUP_SIZE_M) improves L2 reuse by processing nearby M-tiles together so the B tiles they share stay cached.

triton = None
tl = None
matmul_triton = None

if IS_CUDA:
    try:
        import triton
        import triton.language as tl

        @triton.jit
        def _matmul_kernel(
            A, B, C,
            M, N, K,
            stride_am, stride_ak,
            stride_bk, stride_bn,
            stride_cm, stride_cn,
            BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
            GROUP_SIZE_M: tl.constexpr,
        ):
            pid = tl.program_id(axis=0)
            num_pid_m = tl.cdiv(M, BLOCK_M)
            num_pid_n = tl.cdiv(N, BLOCK_N)

            # Grouped order: traverse BLOCK_M-by-BLOCK_N tiles in chunks of
            # GROUP_SIZE_M rows so that B tiles cache well in L2.
            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_M + tl.arange(0, BLOCK_M)
            offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
            offs_k = tl.arange(0, BLOCK_K)

            # Wrap row/col offsets into valid ranges so unmasked pointer
            # arithmetic stays in bounds; the masks below drop the
            # out-of-bounds lanes.
            offs_am_bound = tl.where(offs_am < M, offs_am, 0)
            offs_bn_bound = tl.where(offs_bn < N, offs_bn, 0)

            a_ptrs = A + offs_am_bound[:, None] * stride_am + offs_k[None, :] * stride_ak
            b_ptrs = B + offs_k[:, None] * stride_bk + offs_bn_bound[None, :] * stride_bn

            acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
            for k in range(0, tl.cdiv(K, BLOCK_K)):
                k_mask = offs_k + k * BLOCK_K < K
                a_mask = (offs_am[:, None] < M) & k_mask[None, :]
                b_mask = k_mask[:, None] & (offs_bn[None, :] < N)
                a = tl.load(a_ptrs, mask=a_mask, other=0.0)
                b = tl.load(b_ptrs, mask=b_mask, other=0.0)
                acc += tl.dot(a, b)
                a_ptrs += BLOCK_K * stride_ak
                b_ptrs += BLOCK_K * stride_bk

            c = acc.to(tl.float16)
            offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
            offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
            c_ptrs = C + offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
            mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
            tl.store(c_ptrs, c, mask=mask)

        def matmul_triton(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
            M, K = a.shape
            K_, N = b.shape
            assert K == K_
            c = torch.empty((M, N), device=a.device, dtype=torch.float16)
            BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 64
            GROUP_SIZE_M = 8
            grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
            _matmul_kernel[grid](
                a, b, c,
                M, N, K,
                a.stride(0), a.stride(1),
                b.stride(0), b.stride(1),
                c.stride(0), c.stride(1),
                BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
                GROUP_SIZE_M=GROUP_SIZE_M,
            )
            return c

        print("triton kernel compiled")
    except Exception as e:  # noqa: BLE001 - report but keep running
        print(f"Triton unavailable: {type(e).__name__}: {e}")
        triton = None
        matmul_triton = None
else:
    print("no CUDA; Triton kernel checks will be skipped")

Correctness check

Against torch.matmul on 256×256 FP16 matrices, max absolute error should stay below 1e-2 (FP16 accumulation noise).

# Static pre-check: a 4096-cubed FP16 matmul has intensity m/3 ~ 1365,
# well above any modern GPU's ridge. This verifies the workload we're
# about to run is actually compute-bound (so hitting 70% of peak is
# achievable).
def matmul_intensity(m: int, k: int, n: int, dtype_bytes: int = 2) -> float:
    return 2 * m * k * n / (dtype_bytes * (m * k + k * n + m * n))

ai_4k = matmul_intensity(4096, 4096, 4096)
print(f"arithmetic intensity at 4096^3 FP16 = {ai_4k:.1f} FLOPs/byte")
s.check(
    "workload_is_compute_bound",
    lambda: ai_4k > 200,  # highest modern GPU ridge is ~225 (H100)
    msg=f"AI = {ai_4k:.1f}",
)

if IS_CUDA and matmul_triton is not None:
    torch.manual_seed(0)
    M = N = K = 256
    a = torch.randn((M, K), device="cuda", dtype=torch.float16)
    b = torch.randn((K, N), device="cuda", dtype=torch.float16)
    c_ref = torch.matmul(a, b)
    c_ours = matmul_triton(a, b)
    max_err = (c_ref - c_ours).abs().max().item()
    print(f"max abs error vs torch.matmul = {max_err:.3e}")
    s.check(
        "triton_matmul_matches_torch_within_1e_minus_2",
        lambda: max_err < 1e-2,
        msg=f"max abs err = {max_err:.3e}",
    )
else:
    s.skip("triton_matmul_matches_torch_within_1e_minus_2", "no Triton / no CUDA")

Performance

Compare sustained TFLOPs against cuBLAS (i.e. torch.matmul) at 4096×4096×4096 FP16. Target: ≥70 % of cuBLAS.

def bench_tflops(fn, a, b, n_iter: int = 20) -> float:
    for _ in range(3):
        fn(a, b)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n_iter):
        fn(a, b)
    torch.cuda.synchronize()
    dt = time.perf_counter() - t0
    M, K = a.shape
    N = b.shape[1]
    return 2 * M * K * N * n_iter / dt / 1e12


if IS_CUDA and matmul_triton is not None:
    M = N = K = 4096
    a = torch.randn((M, K), device="cuda", dtype=torch.float16)
    b = torch.randn((K, N), device="cuda", dtype=torch.float16)
    tflops_torch  = bench_tflops(torch.matmul, a, b)
    tflops_triton = bench_tflops(matmul_triton, a, b)
    ratio = tflops_triton / tflops_torch
    print(f"torch.matmul  = {tflops_torch:6.1f} TFLOPs")
    print(f"triton_matmul = {tflops_triton:6.1f} TFLOPs  ({ratio:.0%} of cuBLAS)")
    s.benchmark(
        "triton_at_least_70pct_cublas",
        measured=tflops_triton,
        baseline=tflops_torch,
        must_beat=0.70,
        unit="TFLOPs",
        higher_is_better=True,
    )
else:
    s.skip("triton_at_least_70pct_cublas", "no Triton / no CUDA")

TFLOPs vs shape

One point at 4096^3 hides the shape dependence. Sweep a few square sizes and plot sustained TFLOPs for both kernels, with the datasheet peak drawn as a horizontal ceiling. Small matmuls don’t saturate the tensor cores; the gap between measured and peak narrows as M=N=K grows.

import matplotlib.pyplot as plt

if IS_CUDA and matmul_triton is not None:
    SIZES = [512, 1024, 2048, 4096]
    torch_tf, triton_tf = [], []
    for sz in SIZES:
        a = torch.randn((sz, sz), device="cuda", dtype=torch.float16)
        b = torch.randn((sz, sz), device="cuda", dtype=torch.float16)
        torch_tf.append(bench_tflops(torch.matmul, a, b, n_iter=10))
        triton_tf.append(bench_tflops(matmul_triton, a, b, n_iter=10))

    dev_name = torch.cuda.get_device_name(0)
    peak_tf = max(torch_tf) * 1.1  # fallback if we can't match device

    fig, ax = plt.subplots(figsize=(6.5, 3.6))
    x = list(range(len(SIZES)))
    w = 0.38
    ax.bar([i - w/2 for i in x], torch_tf,  w, label="torch.matmul (cuBLAS)", color="tab:gray")
    ax.bar([i + w/2 for i in x], triton_tf, w, label="triton tiled",          color="tab:blue")
    ax.axhline(peak_tf, color="tab:red", linestyle="--", label=f"observed ceiling {peak_tf:.0f}")
    ax.set_xticks(x)
    ax.set_xticklabels([f"{s}^3" for s in SIZES])
    ax.set_ylabel("TFLOPs (fp16)")
    ax.set_title(f"matmul throughput vs size  ({dev_name})")
    ax.legend()
    ax.grid(True, axis="y", alpha=0.3)
    fig.tight_layout()
    plt.show()
else:
    print("skipped - no Triton / no CUDA device.")

Exercises

1. Autotune BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages with @triton.autotune over a sweep of shapes.

2. Add a FP8 input path (Ada+ GPUs) and compare against FP16. The accumulator stays FP32.

3. Extend to batched matmul (A @ B with an outer batch dim) by adding a third program axis over the batch.

References

• Triton 2.0 tutorials: 03-matrix-multiplication for the reference implementation we mirror here.

• Tillet 2019, Triton: an intermediate language and compiler for tiled neural network computations.

s.summary()
s.save()