write-triton-gemm-kernel

Skill: Write a Triton GEMM Kernel

Purpose

Guide the agent through implementing a correct, performant blocked matrix multiplication kernel in Triton. This covers tile assignment via program_id, pointer arithmetic for A/B/C tiles, accumulation with tl.dot, boundary masking for non-divisible shapes, swizzled tile ordering for L2 reuse, and autotuning for BLOCK_M/BLOCK_N/BLOCK_K/num_stages/num_warps.

---

Use this when

• You need a custom GEMM or batched GEMM that fuses an epilogue (bias add, activation, scaling, etc.) that cuBLAS or CUTLASS cannot express without a separate kernel.
• You need a GEMM on a dtype or layout combination that vendor libraries do not natively support efficiently (e.g., mixed-precision accumulation, custom quantized formats).
• You are building a research kernel and need full visibility into the tiling strategy.
• The matmul is not on the hot path and you want a single portable Triton implementation rather than a CUTLASS build dependency.

---

Do not use this when

• The operation is a standard fp16/bf16/fp32 GEMM with no epilogue fusion requirements. Use torch.compile, torch.mm, or cuBLAS directly — they will match or beat a hand-written Triton GEMM at most shapes.
• The required shapes are very small (M or N < 64). cuBLAS handles these with batched or grouped GEMM routines that are difficult to match in Triton.
• You need int8 or fp8 tensor core throughput with fused dequantization. Prefer CUTLASS or cuDNN unless you have a specific reason to own this kernel.
• Latency matters more than throughput and the problem is memory-bandwidth-bound at small batch. Profiling should drive this decision — do not assume Triton wins.

---

Inputs the agent should gather first

Before writing any code, confirm:

1. M, N, K — exact values or the range of values expected at runtime (static vs dynamic shapes).
2. Input dtype — fp16, bf16, fp32, or mixed (e.g., bf16 inputs, fp32 accumulation).
3. Layout of A and B — row-major or column-major. If transposed, clarify whether the caller passes the transpose or the kernel should handle it internally.
4. Output dtype — same as input or upcast.
5. Epilogue — plain C = A @ B, or is there a scaling factor alpha, bias addition, activation function, or in-place accumulation into an existing C?
6. Batch dimension — standard 2D matmul, batched (B, M, K) x (B, K, N), or broadcasted batch?
7. Hardware target — A100, H100, or other. This determines tensor core eligibility and the optimal pipeline depth.
8. Whether autotuning is allowed — production kernels that ship with a fixed config need to justify that choice; autotuned kernels need a representative benchmark shape.