LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
SystemsFlashAttention

Block Pointers and Multi-Dim Support

Premium

Scale from single sequence to Batch/Head parallelism and simplify pointer math with block pointers.

Get code access

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:

Log in to continue reading

This is premium content. Please log in to access the full article.

  • 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

Stride rule: in row-major layout, each stride = product of all trailing dims. Last dim stride is 1.

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

PyTorch computes these automatically via Q.stride().

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:

systems/flash_attention/04_batch_head.py
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,

Grid mapping:

  • pid_b = tl.program_id(0) → batch index
  • pid_h = tl.program_id(1) → head index
  • pid_m = tl.program_id(2) → Q block index

Manual Pointer Offsets

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

From Naive Implementation to Auto-Tuning

Write your first Flash Attention kernel and use Auto-Tune for performance optimization.

Causal Masking Optimization

Implement causal attention for autoregressive models, achieving ~2x speedup by skipping the upper-triangular computation.

Table of Contents

From Single Sequence to Batch/Head Parallelism
4D Tensor Memory Layout
3D Grid Parallelism
Manual Pointer Offsets
Block Pointers: The Elegant Solution
Core API
Pointer Advances in the Loop
Full Comparison
Summary
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