Triton Tensor Descriptor

Introduction

Triton tensor descriptors are used to describe the layout of tensors in memory, abstracting away manual 2D/N-D index and pointer arithmetic. They are also mandatory for TMA (Tensor Memory Access) operations on NVIDIA GPUs with SM90+ architecture, including Hopper and Blackwell GPUs.

In this blog post, I would like to quickly demonstrate how to use Triton host and device tensor descriptors for implementing GEMM operations.

Triton Tensor Descriptor

There are two types of Triton tensor descriptors: device tensor descriptors and host tensor descriptors. Device tensor descriptors are metadata objects that describe the layout of tensors in device memory, while host tensor descriptors describe the layout of tensors in host memory.

Triton now supports compiler-managed warp specialization through tl.range(..., warp_specialization=True). On Hopper (SM90) architecture, to use warp specialization, the tensor descriptors must be device tensor descriptors. On Blackwell (SM100) architecture or later, to use warp specialization, both device and host tensor descriptors can be used.

Triton tensor descriptors contain metadata about the base memory address, shape, strides, and block sizes of tensors on global memory. The following is an example of how to create and use Triton tensor descriptors in Python. No warp specialization is used in this example.

triton_tensor_descriptor.py

"""
Standalone Triton GEMM: Raw Pointers vs Device vs Host Tensor Descriptors (TMA)
Target: NVIDIA GPUs with TMA support (SM90+)
"""

import argparse
from dataclasses import dataclass
import torch
import triton
import triton.language as tl

try:
    from triton.tools.tensor_descriptor import TensorDescriptor
except ImportError:
    TensorDescriptor = None


@dataclass(frozen=True)
class GEMMConfig:
    block_m: int = 128
    block_n: int = 128
    block_k: int = 64
    num_warps: int = 4
    num_stages: int = 4


# ============================================================================
# Kernels
# ============================================================================


@triton.jit
def gemm_raw_kernel(
    a_ptr,
    b_ptr,
    c_ptr,
    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,
):
    pid_m, pid_n = tl.program_id(0), tl.program_id(1)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k0 in range(0, K, BLOCK_K):
        offs_k = k0 + tl.arange(0, BLOCK_K)
        a_ptrs = a_ptr + offs_m[:,
                                None] * stride_am + offs_k[None, :] * stride_ak
        b_ptrs = b_ptr + offs_k[:,
                                None] * stride_bk + offs_n[None, :] * stride_bn

        a = tl.load(a_ptrs,
                    mask=(offs_m[:, None] < M) & (offs_k[None, :] < K),
                    other=0.0)
        b = tl.load(b_ptrs,
                    mask=(offs_k[:, None] < K) & (offs_n[None, :] < N),
                    other=0.0)
        acc = tl.dot(a, b, acc)

    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs,
             acc.to(tl.float16),
             mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


@triton.jit
def gemm_device_descriptor_kernel(
    a_ptr,
    b_ptr,
    c_ptr,
    M,
    N,
    K,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    a_desc = tl.make_tensor_descriptor(a_ptr,
                                       shape=[M, K],
                                       strides=[K, 1],
                                       block_shape=[BLOCK_M, BLOCK_K])
    b_desc = tl.make_tensor_descriptor(b_ptr,
                                       shape=[K, N],
                                       strides=[N, 1],
                                       block_shape=[BLOCK_K, BLOCK_N])
    c_desc = tl.make_tensor_descriptor(c_ptr,
                                       shape=[M, N],
                                       strides=[N, 1],
                                       block_shape=[BLOCK_M, BLOCK_N])

    offs_m = tl.program_id(0) * BLOCK_M
    offs_n = tl.program_id(1) * BLOCK_N

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k0 in range(0, K, BLOCK_K):
        acc = tl.dot(a_desc.load([offs_m, k0]), b_desc.load([k0, offs_n]), acc)

    c_desc.store([offs_m, offs_n], acc.to(tl.float16))


@triton.jit
def gemm_host_descriptor_kernel(
    a_desc,
    b_desc,
    c_desc,
    M,
    N,
    K,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    offs_m = tl.program_id(0) * BLOCK_M
    offs_n = tl.program_id(1) * BLOCK_N

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k0 in range(0, K, BLOCK_K):
        acc = tl.dot(a_desc.load([offs_m, k0]), b_desc.load([k0, offs_n]), acc)

    c_desc.store([offs_m, offs_n], acc.to(tl.float16))


# ============================================================================
# Wrappers & Execution
# ============================================================================


def get_grid(M, N, cfg: GEMMConfig):
    return (triton.cdiv(M, cfg.block_m), triton.cdiv(N, cfg.block_n))


def run_raw(a, b, cfg: GEMMConfig):
    M, K = a.shape
    N = b.shape[1]
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    gemm_raw_kernel[get_grid(M, N, cfg)](
        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=cfg.block_m,
        BLOCK_N=cfg.block_n,
        BLOCK_K=cfg.block_k,
        num_warps=cfg.num_warps,
        num_stages=cfg.num_stages,
    )
    return c


def run_device_descriptor(a, b, cfg: GEMMConfig):
    M, K = a.shape
    N = b.shape[1]
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    gemm_device_descriptor_kernel[get_grid(M, N, cfg)](
        a,
        b,
        c,
        M,
        N,
        K,
        BLOCK_M=cfg.block_m,
        BLOCK_N=cfg.block_n,
        BLOCK_K=cfg.block_k,
        num_warps=cfg.num_warps,
        num_stages=cfg.num_stages,
    )
    return c


def run_host_descriptor(a, b, cfg: GEMMConfig):
    if TensorDescriptor is None:
        raise RuntimeError(
            "TensorDescriptor module is unavailable in this Triton build.")

    M, K = a.shape
    N = b.shape[1]
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    a_desc = TensorDescriptor.from_tensor(a, [cfg.block_m, cfg.block_k])
    b_desc = TensorDescriptor.from_tensor(b, [cfg.block_k, cfg.block_n])
    c_desc = TensorDescriptor.from_tensor(c, [cfg.block_m, cfg.block_n])

    gemm_host_descriptor_kernel[get_grid(M, N, cfg)](
        a_desc,
        b_desc,
        c_desc,
        M,
        N,
        K,
        BLOCK_M=cfg.block_m,
        BLOCK_N=cfg.block_n,
        BLOCK_K=cfg.block_k,
        num_warps=cfg.num_warps,
        num_stages=cfg.num_stages,
    )
    return c


def get_modes():
    return {
        "raw": run_raw,
        "device": run_device_descriptor,
        "host": run_host_descriptor,
    }


# ============================================================================
# Validation & Benchmarking
# ============================================================================


def check_correctness(shapes, modes, cfg: GEMMConfig):
    print("\n=== Correctness ===")
    torch.manual_seed(0)
    mode_map = get_modes()

    for M, N, K in shapes:
        print(f"\nShape M={M}, N={N}, K={K}")
        a = torch.randn((M, K), device="cuda", dtype=torch.float16)
        b = torch.randn((K, N), device="cuda", dtype=torch.float16)

        ref = torch.matmul(a, b)

        for mode in modes:
            out = mode_map[mode](a, b, cfg)
            torch.cuda.synchronize()

            out_f = out.float()
            ref_f = ref.float()
            diff = (out_f - ref_f).abs()

            max_abs = diff.max().item()
            cos_sim = torch.nn.functional.cosine_similarity(out_f.flatten(),
                                                            ref_f.flatten(),
                                                            dim=0).item()
            rel_norm_err = (torch.linalg.norm(out_f - ref_f) /
                            torch.linalg.norm(ref_f)).item()

            passed = cos_sim > 0.999 and rel_norm_err < 0.01
            status = "PASS" if passed else "FAIL"

            print(
                f"  {mode:>6}: {status} | max_abs={max_abs:.4g} | cos_sim={cos_sim:.6f} | rel_norm_err={rel_norm_err:.6f}"
            )
            if not passed:
                raise AssertionError(
                    f"Validation failed for mode '{mode}' with shape ({M}, {N}, {K})"
                )


def benchmark(shapes, modes, warmup_ms, rep_ms, cfg: GEMMConfig):
    print("\n=== Benchmark ===")
    mode_map = get_modes()

    for M, N, K in shapes:
        print(f"\nShape M={M}, N={N}, K={K}")
        a = torch.randn((M, K), device="cuda", dtype=torch.float16)
        b = torch.randn((K, N), device="cuda", dtype=torch.float16)

        results = {}
        for mode in modes:
            ms = triton.testing.do_bench(lambda m=mode: mode_map[m](a, b, cfg),
                                         warmup=warmup_ms,
                                         rep=rep_ms)
            results[mode] = ms

        raw_ms = results.get("raw", list(results.values())[0])

        print(f"{'Mode':>8} {'Time (ms)':>12} {'TFLOP/s':>12} {'Speedup':>12}")
        print("-" * 48)
        for mode, ms in results.items():
            tflops = (2.0 * M * N * K) / (ms * 1e9)
            speedup = raw_ms / ms
            print(f"{mode:>8} {ms:12.4f} {tflops:12.3f} {speedup:12.3f}x")


# ============================================================================
# Helpers & Entry
# ============================================================================


def parse_shape(s):
    try:
        dims = tuple(map(int, s.split(",")))
        if len(dims) == 3 and all(d > 0 for d in dims):
            return dims
    except ValueError:
        pass
    raise argparse.ArgumentTypeError(
        f"Shape must be 'M,N,K' with positive integers. Got: {s!r}")


def main():
    parser = argparse.ArgumentParser(description="Triton GEMM TMA Benchmarks")
    parser.add_argument("--mode",
                        choices=["raw", "device", "host", "all"],
                        default="all")
    parser.add_argument(
        "--shapes",
        nargs="+",
        type=parse_shape,
        default=[(512, 512, 512), (2048, 2048, 2048), (8192, 8192, 8192)],
    )
    parser.add_argument("--warmup-ms", type=int, default=100)
    parser.add_argument("--rep-ms", type=int, default=500)
    parser.add_argument("--no-check", action="store_true")

    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required.")

    major, minor = torch.cuda.get_device_capability()
    print(f"Triton: {triton.__version__}")
    print(f"PyTorch: {torch.__version__}")
    print(f"GPU: {torch.cuda.get_device_name()} (SM{major}{minor})")

    modes = ["raw", "device", "host"] if args.mode == "all" else [args.mode]

    if ("device" in modes or "host" in modes) and major < 9:
        raise SystemExit(
            f"Descriptor modes require SM90+ (Hopper/Blackwell). Detected SM{major}{minor}."
        )

    if "host" in modes:
        triton.set_allocator(lambda size, align, stream: torch.empty(
            size, dtype=torch.int8, device="cuda"))

    cfg = GEMMConfig()

    if not args.no_check:
        check_correctness(args.shapes, modes, cfg)

    benchmark(args.shapes, modes, args.warmup_ms, args.rep_ms, cfg)


if __name__ == "__main__":
    main()

With Triton tensor descriptors, the GEMM kernel implementations are simplified. Benchmarks were conducted on an NVIDIA RTX 5080 GPU using the nvcr.io/nvidia/pytorch:26.09-py3 PyTorch Docker image from NVIDIA NGC.

$ python triton_tensor_descriptor.py
Triton: 3.8.0
PyTorch: 2.14.0a0+b2c75dd062.nv26.09
GPU: NVIDIA GeForce RTX 5080 (SM120)

=== Correctness ===

Shape M=512, N=512, K=512
     raw: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000
  device: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000
    host: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000

Shape M=2048, N=2048, K=2048
     raw: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000
  device: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000
    host: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000

Shape M=8192, N=8192, K=8192
     raw: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000
  device: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000
    host: PASS | max_abs=0 | cos_sim=1.000000 | rel_norm_err=0.000000

=== Benchmark ===

Shape M=512, N=512, K=512
    Mode    Time (ms)      TFLOP/s      Speedup
------------------------------------------------
     raw       0.0191       14.076        1.000x
  device       0.0211       12.699        0.902x
    host       0.0193       13.893        0.987x

Shape M=2048, N=2048, K=2048
    Mode    Time (ms)      TFLOP/s      Speedup
------------------------------------------------
     raw       0.2366       72.617        1.000x
  device       0.2275       75.528        1.040x
    host       0.2206       77.882        1.073x

Shape M=8192, N=8192, K=8192
    Mode    Time (ms)      TFLOP/s      Speedup
------------------------------------------------
     raw      11.2056       98.121        1.000x
  device      10.5258      104.459        1.065x
    host      10.4400      105.317        1.073x

References

• Triton Tensor Descriptor - Triton Documentation
• Fused Attention Tutorial - Triton Documentation