# Triton Kernel Authoring

> Triton Kernel Authoring

- Skill: `sanjeevrg89/triton-kernel-authoring` (Agent Skill, multi-file: 5 files)
- Install (CLI): `npx skillmds@latest add sanjeevrg89/triton-kernel-authoring`
- Raw SKILL.md: https://api.skillmd.com/api/skills/sanjeevrg89/triton-kernel-authoring/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: Coding & Dev Tools
- Author: sanjeevrg89 (https://skillmd.com/u/sanjeevrg89)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/sanjeevrg89/triton-kernel-authoring

---


# Triton Kernel Authoring

Apply the judgment of an engineer who writes production Triton kernels for a living: who starts from a
correct PyTorch reference, diffs every kernel with `torch.allclose` (including on non-power-of-2 and masked
edge sizes), autotunes before believing a number, and reports GB/s or TFLOP/s against the roofline — never
a bare "it's faster." This skill **produces the kernel**; profiling it is [[gpu-performance-engineering]]
and how it lowers is [[ml-compilers-codegen]].

## How to use this skill

1. **Read `triton-kernel-authoring-guide.md`** in this directory — the full reference (tile/SPMD model,
   the authoring procedure with correctness + perf gates, performance patterns, the canonical kernel
   patterns, anti-patterns, and the definition of done). Apply it to the task.
2. For full worked kernels to imitate — a masked fused vector op, an autotuned tiled matmul with `tl.dot`,
   and a numerically-stable fused softmax, each with its `torch.allclose` check and a `do_bench` stub —
   read **`examples.md`**.
3. Match the surrounding codebase's conventions (dtypes, layouts, how kernels are wrapped/launched); apply
   the correctness-first and benchmark-honestly rules regardless. **Never fabricate a `tl.*` name or a
   decorator flag — flag version-sensitive APIs "verify against current Triton docs."**

## Essentials (full detail in `triton-kernel-authoring-guide.md`)

- **Tile + SPMD, not threads.** You write **one program that computes one tile** of the output; the
  runtime launches a **grid** of them and the compiler owns warps/coalescing/SMEM/pipelining/MMA. You
  reason about blocks and pointer arithmetic; tile size / `num_warps` / `num_stages` are your perf knobs.
- **The primitives:** `@triton.jit`; `tl.program_id(axis)` for the tile index; `BLOCK_SIZE: tl.constexpr`
  (compile-time, usually powers of two); `tl.arange(0, BLOCK)` for local indices; pointer arithmetic with
  **strides** for multi-dim (and `[:, None]`/`[None, :]` broadcasts for 2-D tiles); `tl.load`/`tl.store`.
- **Mask every boundary.** `mask=offs < n` on loads and stores. A kernel without masks is correct only
  when every dim exactly divides its block — almost never. **Masking bugs are the #1 silent corruption.**
- **Launch with `triton.cdiv`.** `grid = lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),)`; the mask
  covers the partial last tile. Launch is async — **synchronize before timing.**
- **The procedure (in order, each a gate):** (1) write a **PyTorch reference** (oracle + baseline); (2)
  write the simplest correct masked kernel; (3) **validate with `torch.allclose`/`assert_close` on
  multiple sizes incl. non-power-of-2 and a non-dividing size**; (4) **`@triton.autotune`** over block
  sizes × `num_warps` × `num_stages` (re-check correctness after); (5) **benchmark vs the reference** with
  `triton.testing.do_bench` and report **GB/s** (memory-bound) or **TFLOP/s** (compute-bound) vs roofline.
- **`tl.dot` for matmul/tensor-cores; accumulate in fp32.** Even with fp16/bf16 inputs, `acc` is fp32;
  cast down at the end. fp16 accumulation loses precision over long K.
- **Numerical stability:** softmax subtracts `tl.max` before `tl.exp` (else `inf`/`NaN` on big logits);
  compute norm/softmax reductions in fp32.
- **Fusion wins on the memory wall.** The reason to write a kernel is usually to **fuse** a chain PyTorch
  runs as separate HBM round-trips — keep intermediates in registers/SMEM, write only the result.
  Epilogue fusion (bias+activation+cast into the matmul/reduction kernel) is the highest leverage.
- **Coalesce + tune.** Map the innermost offset to contiguous memory; sweep `num_warps`/`num_stages` via
  autotune (more isn't always better — large tiles × deep pipelines spill registers / exhaust SMEM).
- **Don't reinvent what's tuned.** A bare plain GEMM/elementwise that cuBLAS or `torch.compile`/Inductor
  already fuses well is not worth a hand kernel. Write one when you've shown bandwidth is left on the table.
- **Version-sensitive (verify current docs):** the block-pointer API (`tl.make_block_ptr`/`tl.advance`,
  `boundary_check`/`padding_option`), the exact `triton.Config`/`@triton.autotune` keyword set, `do_bench`
  signature, newer `tl.*` intrinsics, and non-NVIDIA backend warp/stage semantics. Never invent an API.

## Related skills

- `[[ml-compilers-codegen]]` — how your kernel **lowers**: TTIR→TTGIR (adds the layout/encoding deciding
  coalescing & MMA selection) →LLVM→PTX/SASS; read it to understand why tile shape/dtype are load-bearing.
- `[[gpu-performance-engineering]]` — how to **profile** what you wrote (roofline, Nsight Compute SOL,
  occupancy/coalescing/bank-conflict reads) and benchmark with statistical rigor. Pair with §5 here.
- `[[ml-frameworks]]` — PyTorch/JAX/XLA, the reference op you diff against, and how the kernel plugs into
  `torch.autograd`/`torch.library`/`torch.compile`.
- `[[inference-optimization]]` — where fused kernels (attention, norms, quantized matmul) pay off in
  decode/serving; the model-level memory-bound vs compute-bound reasoning.
- `[[ai-research-science]]` — when a novel op from a paper has no library kernel and you must author one.