《The Stack》 EP.43 2026 · APR · 20 42 MIN IMAGINARY

What Triton hides,
and what it exposes

Spread seven files on the table and read them line by line. When the same kernels move from CUDA to Triton — how does the code fold? What details disappear under the compiler, and what stays in your hands?

Host

Sam

"OK, so why is this the comfortable one?"

Guest

Jensen

Wasn't thrilled about it internally, but has come to accept it

CHAPTERS · 7 ACTS · 7 FILES42:08

01:2001Smallest Triton programsmoke_vector_add.py 06:4802One tl.sum replaces an entire warp shufflereduction.py 13:0203One program = one rowsoftmax.py 18:4104One thing where Triton beats CUDAmatmul.py 26:1005Flash Attention in 40 linesflash_attention.py 31:3506The magic of one constexprflash_attention_mha.py 37:2007Climbing up to torch.ops.*flash_attention_mha_op.py

00 · 00:28Cold Open

Seven files. One line at a time.

SAM

Last episode I asked "five years from now what else will you be optimizing besides matmul," and you punted to the next one.

JENSEN

(laughs) Still not telling. Something more fun instead. This time our friend rewrote the same kernels in Triton. Seven files — let's open them and read line by line.

FIG 0 · today's material · triton_kernels/7 files

01 · 01:20Smallest Triton program

One `program` = one block, not one thread.

CUDA's threadIdx.x didn't vanish — it just slipped under the compiler. Triton is block-level SPMD: the thread parallelism inside a block is decided by the compiler from num_warps.

smoke_vector_add.py · Tritonblock-level SPMD

@triton.jit
def vector_add_kernel(x_ptr, y_ptr, out_ptr, n,
                      BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

Same address math · CUDAthread-level

int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) out[idx] = x[idx] + y[idx];

SAM

Whatever I was doing with threadIdx.x in CUDA — where did it go in Triton?

JENSEN

Gone. More precisely, the compiler hid it. tl.arange(0, 1024) is a vector pointing to "all 1024 indices this block handles." Which thread handles which lane — Triton decides.

SAM

Sounds like SIMD.

JENSEN

Exactly that. "Block-level SPMD." The program is per block, not per thread. Thread parallelism inside is the compiler's problem.

Both patterns lower to the same PTX.
But Triton lets you think at the "index space" level.— Jensen, 04:12

FIG 1 · abstraction heightwhere does threadIdx.x live?

↳ Lesson 08 · Triton port

02 · 06:48One tl.sum replaces an entire warp shuffle

But it charges you three things.

FIG 2 · same reduction — CUDA vs Triton~15 lines → 2 lines

SAM

So it really is free?

JENSEN

(laughs) Not quite. Three costs.

① launch overhead

Reducing 67M elements — CUDA v4 = 1.039 ms, Triton = 1.097 ms (5% slower). Python → autotune cache → JIT cache → argument binding → cuLaunchKernel eats ~50–100 µs. For tiny kernels that overhead can exceed the compute time. → Launch 30 element-wise ops as separate Triton kernels and you're done.

② autotune footgun

Autotune runs configs sequentially against the same output buffer. Leftover stale partial sums from previous attempts mix into the result. Fix — reset_to_zero=["partial_ptr"]. It's faintly mentioned in the docs; miss it and you debug for hours.

③ autotune + 2-pass reduce idiom

Change BLOCK_SIZE and num_programs changes too. Size the partial buffer for the maximum case, then slice it to the prefix that matches the chosen config.

"Triton is high-level but thin."
Its abstraction is shallow, so the internals keep leaking out. Handling the leaks well is the value of a Triton engineer.— Jensen, 12:04

reduction.py · footgun remedyreset_to_zero

@triton.autotune(
    configs=AUTOTUNE_CONFIGS,
    key=["n_elements"],
    reset_to_zero=["partial_ptr"],  # ← matters
)

↳ Lesson 03 · Reduction

03 · 13:02One program = one row

Mask logic dissolves naturally into the data values.

SAM

So how do you handle rows with N=1000?

JENSEN

BLOCK_SIZE=1024, mask out the last 24. The trick is other=-float("inf"). If OOB lanes hold -inf, tl.max is unaffected and exp(-inf)=0 so sum doesn't get contributions either. Mask logic melts into the data values.

SAM

The autotune key is BLOCK_SIZE. Clever not to use N directly.

JENSEN

Right. With BLOCK_SIZE = _next_pow2(N), N=513–1024 all bucket into 1024. Cache-friendly autotune key design — the last skill you pick up when learning Triton.

FIG 3 · handling OOB lanes at N=1000mask = data value

softmax.py · the three core lines

offs = tl.arange(0, BLOCK_SIZE)       # 0..1023
mask = offs < n_cols                  # only first 1000 True
x = tl.load(in_row + offs, mask=mask,
            other=-float("inf"))       # OOB → -inf

↳ Lesson 05 · Softmax

04 · 18:41One thing where Triton beats CUDA

Grouped program-id swizzling — possible in CUDA, but nobody writes it.

Change only the order in which output C tiles are visited, and L2 reuse changes dramatically. Row-major eviction sweeps B's columns. Group-wise traversal lets the same B columns be reused.

FIG 4A · row-major · B columns evicted every rownaive

Tile 0→1→2→… same row of A, different column of B. If B doesn't fit in L2, it gets swept out.

FIG 4B · grouped · GROUP_SIZE_M=4swizzled

Tiles 0→1→2→3 reuse the same B column four times. L2 efficiency ↑.

SAM

You can write this in CUDA too. Why say CUDA doesn't have it?

JENSEN

You can, but nobody does. Three reasons — ① In CUDA, blockIdx.x is just linear hardware order. You have to write the math by hand at the top of the kernel. ② That math is painful to read. ③ Change GROUP_SIZE_M and you recompile. In Triton it's an autotune parameter and a standard idiom.

FIG 4C · measured — 4096³ matmulL4 sm_89

variant	TFLOPS	note
our CUDA v3 (FMA only)	3.9	register blocking
torch.matmul (cuBLAS + TF32)	25.8	years of NVIDIA tuning
Triton fp32	28.9	cuBLAS + 12%
our CUDA v4 (WMMA fp16)	18.5	hand-written mma
cuBLAS fp16	51.8	—
Triton fp16	54.0	cuBLAS + 4% · 40 lines

Twenty years of cuBLAS loses to 40 lines of Python.
Autotune explores the config space better than a human. A textbook case of measurement beating theory.— Jensen, 23:58

↳ Lesson 04 · Matmul Lesson 08 · Triton port

05 · 26:10Flash Attention in 40 lines

Half the code. And 6.1× faster.

FIG 5A · line countkernel body only

acc update · online + P@V in one line

acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)

In CUDA this logic stretches over 30+ lines. Complexity collapses when the abstraction is at the right height.

FIG 5B · perf — N=8192, seq · single headL4

impl	time (ms)	speedup
CUDA FA v1 (fp32)	3.045	1.00×
Triton FA (fp16)	0.496	6.14×

Where does the 6× come from

① tl.dot uses Tensor Cores (our CUDA v1 is fp32 FMA). ② Autotune sweeps 6 configs of (BLOCK_M, BLOCK_N, num_warps, num_stages) — sweeping that by hand in CUDA means 6 recompiles. ③ tl.trans, 2-D pointer broadcasts, swizzled smem layouts — all automatic.

↳ Lesson 06 · CUDA FA Lesson 08 · Triton port

06 · 31:35The magic of one constexpr

Causal mask isn't about "filling" — it's about "removing from the loop."

Filling with -inf still computes the full QKᵀ. FA-v2's real win comes from pulling the upper-triangle K tiles out of the iteration itself.

FIG 6A · causal loop · different end_n per pid_mhalf the tiles skipped

constexpr · two kernels JIT-compiledno runtime branch

def flash_attention_mha_fwd_kernel(...,
    IS_CAUSAL: tl.constexpr,     # ← key
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr):
    if IS_CAUSAL:
        end_n = tl.minimum(N, (pid_m+1) * BLOCK_M)
    else:
        end_n = N

→ Compile separate kernels for IS_CAUSAL=True and =False. The if doesn't exist at runtime. Equivalent to C++ template specialization.

FIG 6B · our Triton vs SDPA (cuDNN FA-2)LLaMA-7B shape

(B, H, N, d)	ours	SDPA	ratio
(1, 32, 2048, 128)	0.784	0.613	0.78×
(1, 32, 4096, 128)	2.964	2.559	0.86×
(16, 12, 512, 64)	0.249	0.282	1.13×

One 268-line file at 78–90% of cuDNN. What's missing — async copy, persistent kernel, warp specialization — all still experimental in Triton.

↳ Lesson 09 · MHA causal FA

07 · 37:20Climbing up to torch.ops.*

70 lines. The last glue that turns a Triton kernel into a component.

flash_attention_mha_op.py · 70 lines

@custom_op(
    "triton_training::flash_attention_mha",
    mutates_args=(),
    device_types="cuda",
)
def flash_attention_mha_op(q, k, v, is_causal=False):
    return triton_flash_attention_mha(q, k, v, is_causal=is_causal)

@flash_attention_mha_op.register_fake
def _fake(q, k, v, is_causal=False):
    return torch.empty_like(q)   # ← shape decl for Dynamo

Why register_fake matters

When torch.compile traces a model, it uses FakeTensors (shape, dtype, device — no data). Our Triton kernel can't run on those. Instead, we declare "this op's output shape is this" → Dynamo doesn't break the graph. Without it, fullgraph=True fails.

FIG 7 · what these 70 lines unlockvLLM pattern

SAM

So what Brian built is essentially a vLLM-style production op.

JENSEN

Nearly. What's missing is backward (autograd) and GQA support. Both have clear designs — next-lesson material. Bolt those two on and the next goal is porting vLLM's PagedAttention.

↳ Lesson 07 · custom_op Lesson 09 · MHA

What Triton hides · what it exposes

Not a "high-level DSL." A language that sits at exactly the right abstraction height. The five below go under the compiler; the five above stay in your hands.

HIDESunder the compiler

Thread-level parallelism — threadIdx.x is gone
Warp shuffle, smem tree reduction — tl.sum
Tensor Core instruction selection — tl.dot picks by dtype
Smem layout swizzling — bank conflicts dodged automatically
Launch config tuning — delegated to autotune

EXPOSESstill in your hands

Block size, grid structure — program_id, grid=lambda meta
Memory-hierarchy awareness — HBM patterns of tl.load(mask=...)
Compile-time vs runtime boundary — tl.constexpr
Autotune key design — too broad explodes; too narrow misses
Numerical behavior — online softmax, fp16 vs fp32 accumulator

CUDA is assembly, Triton is C, PyTorch is Python.
Write most of it in C; drop to assembly only for hot paths.— Sam & Jensen, 41:02

Four reasons you still need CUDA

When Triton hits a wall — new mma (Blackwell FP4), persistent kernels, async copy fine control — still CUTLASS/CUDA.
To debug, you need to read the PTX Triton emits — *.ptx under TRITON_CACHE_DIR.
vLLM, FlashAttention-3, Mamba kernels are still CUDA-based. Reading them needs CUDA as first language.
Chasing "why is this Triton slow" lands on bank conflict, register spill, occupancy. Answers live in CUDA concepts.

Code ref · triton_kernels/ all 7 files · L4 sm_89, CUDA 13.0, PyTorch 2.11, Triton 3.6.

← Ep.42 Why did your GPT end up using my GPU this way Back → Index · 11-post log

What Triton hides,and what it exposes

Seven files. One line at a time.

One program = one block, not one thread.

But it charges you three things.

Mask logic dissolves naturally into the data values.

Grouped program-id swizzling — possible in CUDA, but nobody writes it.

Half the code. And 6.1× faster.

Causal mask isn't about "filling" — it's about "removing from the loop."

70 lines. The last glue that turns a Triton kernel into a component.

What Triton hides · what it exposes

What Triton hides,
and what it exposes

One `program` = one block, not one thread.