Block Pointers and Multi-Dim Support

In the previous chapters we built a functional, autotuned Flash Attention kernel. Our parallelism was along the sequence dimension: each program handles one Q block, iterating over all K/V blocks.

But real Transformer inputs are (Batch, Head, SeqLen, Dim). Batch and Head are independent, so we should parallelize them too.

We solve two problems:

1. Multi-dimensional parallelism: scale from SeqLen to Batch × Head × SeqLen
2. Pointer management: use block pointers to simplify 4D addressing

From Single Sequence to Batch/Head Parallelism

4D Tensor Memory Layout

When input is (B, H, N, D), pointer math gets complex. The GPU memory is still 1D contiguous (see Tensor Layout).

Example: (2, 4, 8, 64) (2 batches, 4 heads, seq length 8, dim 64):

Logical view: Q[batch, head, seq, dim]  →  Q[2, 4, 8, 64]
Physical storage: 1D array, total 2 × 4 × 8 × 64 = 4096 elements

Strides tell how many elements to skip per dimension:

• stride_b = 4 × 8 × 64 = 2048 — next batch
• stride_h = 8 × 64 = 512 — next head
• stride_m = 64 — next seq row
• stride_d = 1 — next column

Shape:   (B,     H,     N,    D )
         (2,     4,     8,    64)
          ↓      ↓      ↓     ↓
Stride:  H×N×D   N×D    D     1
         2048    512    64    1

To access submatrix (pid_b, pid_h), offset by pid_b * stride_b + pid_h * stride_h.

3D Grid Parallelism

In 04_batch_head.py, we use a 3D grid:

def call_flash_attention(Q, K, V):
    B, H, N, D = Q.shape
    O = torch.empty_like(Q)

    # 3D grid: (Batch, Head, SeqBlocks)
    grid = lambda META: (B, H, triton.cdiv(N, META["BLOCK_M"]))

    flash_attention[grid](
        Q, K, V, O,
        N, D,
        Q.stride(0), Q.stride(1), Q.stride(2),  # stride_b, stride_h, stride_m
        K.stride(0), K.stride(1), K.stride(2),
        V.stride(0), V.stride(1), V.stride(2),
        O.stride(0), O.stride(1), O.stride(2),
    )
    return O

Manual Pointer Offsets

Inside the kernel, we must manually offset by batch/head: