Algorithms, data structures, and the line of thinking

The whole skill is recognizing when a problem fits the shape Triton is good at — a regular tile, a local reduction, a cheap store — and when it doesn't, so you hand it to CUDA, CUB, cuBLAS, or PyTorch instead.

The fit map

Before any kernel, the interview answer is a triage. Memorize this table as the first thing you say:

1 · Row-wise reduction as a data-structure primitive

A row-major matrix A of shape M×N is already a data structure: the i-th row lives at A_ptr + i*stride_m, and the j-th element of that row at + j*stride_n. There is no pointer chasing — the address is arithmetic. That is why row reductions are the canonical Triton win: assign one program per row, load the row into a tile, reduce.

import triton
import triton.language as tl

@triton.jit
def row_max_kernel(A_ptr, out_ptr,
                   M, N,
                   stride_m, stride_n,
                   BLOCK_N: tl.constexpr):
    row = tl.program_id(0)                  # this program owns one row
    if row >= M:
        return

    # pointer tensor for the whole row (assumes N <= BLOCK_N)
    offs_n = tl.arange(0, BLOCK_N)
    ptrs = A_ptr + row * stride_m + offs_n * stride_n
    mask = offs_n < N

    # other=-inf so masked lanes never win the max
    vals = tl.load(ptrs, mask=mask, other=float("-inf"))

    m = tl.max(vals, axis=0)                # local reduction, no cross-program comm
    tl.store(out_ptr + row, m)

The grid is (M,). The crucial line is other=float("-inf"): masked lanes load a sentinel that can never be the maximum, so the reduction over BLOCK_N lanes equals the reduction over the real N elements. Sum reductions use other=0.0 for the same reason.

2 · Top-k for small k without sorting

You don't need a sort to find the largest two elements of a row. You need tl.max, tl.argmax, and a mask. The trick: find the max, record its index, overwrite that lane with -∞ using tl.where, find the max again.

@triton.jit
def top2_per_row_kernel(A_ptr,
                        v1_ptr, i1_ptr, v2_ptr, i2_ptr,
                        M, N, stride_m, stride_n,
                        BLOCK_N: tl.constexpr):
    row = tl.program_id(0)
    if row >= M:
        return

    offs = tl.arange(0, BLOCK_N)
    ptrs = A_ptr + row * stride_m + offs * stride_n
    mask = offs < N
    vals = tl.load(ptrs, mask=mask, other=float("-inf"))

    # --- first pick ---
    v1 = tl.max(vals, axis=0)
    i1 = tl.argmax(vals, axis=0)

    # mask out the winning lane, then pick again
    vals2 = tl.where(offs == i1, float("-inf"), vals)
    v2 = tl.max(vals2, axis=0)
    i2 = tl.argmax(vals2, axis=0)

    tl.store(v1_ptr + row, v1)
    tl.store(i1_ptr + row, i1)
    tl.store(v2_ptr + row, v2)
    tl.store(i2_ptr + row, i2)

Generalizing to top-k is a loop: each iteration does one tl.max/tl.argmax over the whole row and one tl.where to evict the chosen lane. That is O(k·N) work per row.

3 · Segmented reductions in a CSR-like layout

Compressed Sparse Row (CSR) is the classic irregular structure. The matrix is three arrays: row_ptr (length M+1, the start offset of each row), col_idx (column index of each nonzero), and vals (the nonzeros themselves). Row i's nonzeros are vals[row_ptr[i] : row_ptr[i+1]]. A sparse-matrix × dense-vector product (SpMV) is a segmented reduction: one dot product per row over a variable-length segment.

@triton.jit
def csr_spmv_kernel(row_ptr, col_idx, vals, x_ptr, y_ptr,
                    M, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    if row >= M:
        return

    start = tl.load(row_ptr + row)          # segment bounds for this row
    end   = tl.load(row_ptr + row + 1)

    # load the whole bounded segment at once (assumes nnz_per_row <= BLOCK)
    offs = tl.arange(0, BLOCK)
    j    = start + offs
    mask = j < end
    cols = tl.load(col_idx + j, mask=mask, other=0)
    a    = tl.load(vals    + j, mask=mask, other=0.0)
    xs   = tl.load(x_ptr + cols, mask=mask, other=0.0)  # gather x[col]

    acc = tl.sum(a * xs, axis=0)            # segmented reduction, local
    tl.store(y_ptr + row, acc)

The data structure became pointer arithmetic: start = row_ptr[i], j = start + arange(BLOCK), mask = j < end kills the lanes past the segment, and tl.load(x_ptr + cols) is a gather that translates column indices into dense addresses.

4 · Block-sparse attention metadata

Attention has a beautiful property: it is regular at the block level but sparse at the sequence level. For a query block, most key/value blocks contribute nothing (they're masked out by a sliding window, a causal mask, or a learned sparsity pattern). The irregularity is which K/V blocks matter — and that can be precomputed into metadata: for each query block, a list of the K/V block ids it must attend to.

So a 2-D grid (q_block, sparse_slot) covers exactly the work that survives the mask. Each program reads its target block id from the table and does a regular dense block matmul:

@triton.jit
def block_sparse_scores_kernel(Q_ptr, K_ptr, S_ptr,
                               kv_block_ids,        # metadata: [n_q_blocks, MAX_SLOTS]
                               n_slots_ptr,         # how many real slots per q_block
                               stride_qm, stride_qd,
                               stride_kn, stride_kd,
                               stride_sq, stride_ss, stride_sm, stride_sn,
                               D: tl.constexpr,
                               BLOCK_M: tl.constexpr,
                               BLOCK_N: tl.constexpr,
                               MAX_SLOTS: tl.constexpr):
    q_block = tl.program_id(0)
    slot    = tl.program_id(1)

    # skip slots beyond this query block's real fan-out
    if slot >= tl.load(n_slots_ptr + q_block):
        return

    # THE METADATA IS THE DATA STRUCTURE: one indirection picks the K/V block
    kv_block = tl.load(kv_block_ids + q_block * MAX_SLOTS + slot)

    offs_m = q_block * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = kv_block * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, D)

    # load Q block [BLOCK_M, D] and K block [BLOCK_N, D]
    q = tl.load(Q_ptr + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd)
    k = tl.load(K_ptr + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd)

    # regular dense block matmul: [BLOCK_M, D] x [D, BLOCK_N] -> [BLOCK_M, BLOCK_N]
    acc = tl.dot(q, tl.trans(k))            # tensor-core path

    offs_sm = tl.arange(0, BLOCK_M)
    offs_sn = tl.arange(0, BLOCK_N)
    s_ptrs = (S_ptr + q_block * stride_sq + slot * stride_ss
              + offs_sm[:, None] * stride_sm
              + offs_sn[None, :] * stride_sn)
    tl.store(s_ptrs, acc)

5 · Online (streaming) softmax — the FlashAttention recurrence

The block-sparse kernel above produced score blocks. But a full attention output needs softmax(QKᵀ)·V, and the whole point of FlashAttention is to never materialize the full score matrix. The insight: softmax can be computed incrementally as K/V blocks stream past, keeping only O(1) state per query row.

Per query row the running state is the running max m, the normalizer l, and the accumulator acc (a D-vector). When a new score block arrives:

# state carried across K/V blocks, per query row
#   m   : running max of scores seen so far   (init -inf)
#   l   : running sum of exp(score - m)        (init 0)
#   acc : running sum of p_i * v_i             (init 0, shape [D])

scores  = tl.dot(q, tl.trans(k))               # this block's [BLOCK_M, BLOCK_N] scores

m_new   = tl.maximum(m, tl.max(scores, axis=1))        # new running max
alpha   = tl.exp(m - m_new)                            # rescale factor for OLD state
p       = tl.exp(scores - m_new[:, None])              # new block, stable exponent

l_new   = l * alpha + tl.sum(p, axis=1)                # renormalize + add new mass
acc_new = acc * alpha[:, None] + tl.dot(p, v)          # rescale old acc + add new

m, l, acc = m_new, l_new, acc_new
# ... after the last block:  out = acc / l[:, None]

6 · Associative scan / cumsum inside a tile

Triton can do prefix operations within a tile. tl.cumsum(x, axis=0) gives a running sum along an axis, and tl.associative_scan(x, axis, combine_fn) generalizes it to any associative operator (running max, running product, segmented combinations). These are perfect for within-block running sums — e.g., computing offsets inside a compacted block, or a local prefix for a fused operation.

offs = tl.arange(0, BLOCK)
x    = tl.load(p + offs, mask=offs < N, other=0.0)
run  = tl.cumsum(x, axis=0)        # in-tile prefix sum, no cross-program coordination

7 · Seeded dropout with a Philox counter

Fused dropout needs per-element randomness that is (a) deterministic given a seed, so backward can regenerate the same mask without storing it, and (b) decorrelated across elements. Triton exposes tl.rand(seed, offset), a Philox-style counter-based generator: the random value is a pure function of the seed and the element's global offset — no per-element state, no RNG stream to thread through.

@triton.jit
def fused_dropout_kernel(x_ptr, out_ptr, N, p, seed,
                         BLOCK: tl.constexpr):
    pid  = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N

    x = tl.load(x_ptr + offs, mask=mask)
    r = tl.rand(seed, offs)                 # deterministic per-offset uniform in [0,1)
    keep = r > p                            # Bernoulli(1 - p)
    out  = tl.where(keep, x / (1.0 - p), 0.0)   # inverted dropout scaling

    tl.store(out_ptr + offs, out, mask=mask)

Because the mask is reproducible from (seed, offs), the backward pass recomputes it on the fly instead of reading a full mask tensor from HBM — another fusion win that removes a tensor-sized memory round-trip.

Where Triton is the wrong tool

Interview answer template

When asked "would you implement X in Triton?", walk this ladder out loud. It demonstrates exactly the judgment the question is probing:

1. Describe the data structure. Is it dense strides, CSR (row_ptr/col_idx/vals), block metadata (an index table of block ids), or a mask? Name it precisely.
2. Choose the program owner. What does one program_id own — a row, a block, a tile, a sparse slot? This decides the grid.
3. Show the pointer tensors. Write base + id*stride + arange(BLOCK)*stride and the mask. If you can't express the region as pointer arithmetic, that's the first red flag.
4. State the local primitive. Map / tl.sum / tl.max / tl.dot / online softmax update. This is the work that stays in registers.
5. Name the cross-tile issue. Does it need a second pass, an atomic, metadata scheduling, or a library fallback? Being honest here is the whole point.
6. Decide the tool. If steps 3–4 are clean and local, Triton. If step 5 forces global ordering, dynamic queues, or peak GEMM, hand it to CUDA / CUB / cuBLAS / PyTorch — and say why.