The Formula

Attention(Q, K, V) =
  softmax(QK^T / √d_k) ⋅ V

• Q = what we’re looking for
• K = what each position advertises
• V = what each position actually contains
• √d_k = scaling factor (prevents softmax saturation)

Shape Trace

def attention(q, k, v):
    d_k = q.shape[-1]
    scores = torch.matmul(q, k.transpose(-2, -1))
    scores = scores / math.sqrt(d_k)
    weights = torch.softmax(scores, dim=-1)
    return torch.matmul(weights, v)

Multi-Head Attention

• Split hidden_dim across heads: d_head = 4096 / 32 = 128
• Each head learns different attention patterns
• Reshape: (batch, seq, 4096) → (batch, 32, seq, 128)

GQA: The Memory Trick

• 8 KV heads shared across 32 query heads
• 4:1 ratio — 4 query heads per KV head
• Uses repeat_interleave to expand K,V

KV Cache Memory Comparison

# SwiGLU forward pass
def forward(self, x):
    gate = F.silu(self.w1(x))  # silu = x * sigmoid(x)
    up   = self.w3(x)
    return self.w2(gate * up)

# Shapes for 7B model:
#   x:      (batch, seq, 4096)
#   w1(x):  (batch, seq, 11008)  gate projection
#   w3(x):  (batch, seq, 11008)  up projection
#   w2(...): (batch, seq, 4096)  down projection

Parameter Count Per Layer

32-layer model: 4.3B parameters in FFN alone

Prefill

Builds the initial KV cache.

Decode

Appends one K,V entry per step.

Without Cache

Every decode step recomputes K,V for all previous tokens.

O(n²) per token generated

With Cache

Store K,V once computed. Only compute for the new token.

O(n) per token generated

Per Token

Per Sequence

Total (batch)

of 80 GB H100 VRAM

GEMM — Prefill

C = A @ B    (M,K) × (K,N)

GEMV — Decode

y = A @ x    (M,K) × (K,)

Thread Hierarchy

H100: 132 SMs, each with CUDA cores + Tensor cores

Memory Hierarchy

// CUDA vector addition kernel
__global__ void vector_add(
    float* a, float* b,
    float* c, int n) {
  int idx = blockIdx.x * blockDim.x
           + threadIdx.x;
  if (idx < n) {
    c[idx] = a[idx] + b[idx];
  }
}

// Launch configuration
int block_size = 256;
int num_blocks = (n + 255) / 256;
vector_add<<<num_blocks, block_size>>>(
    d_a, d_b, d_c, n);

Thread Indexing

# Profile your inference script
$ nsys profile -o output python inference.py
$ nsys stats output.nsys-rep

# Label regions with NVTX
from torch.cuda import nvtx
with nvtx.range("prefill"):
    output = model(input_ids)
with nvtx.range("decode"):
    for i in range(100):
        output = model(token)

• Load from HBM to shared memory once, access many times
• Up to 164 KB per SM, 20-30 cycle latency vs 300 cycles for global
• 32 banks — stride of 32 causes bank conflicts (serialization)

4.5x faster bandwidth, 10-15x lower latency

Tiled Reduction Pattern

__global__ void sum_reduce(
    float* data, float* result, int n) {
  __shared__ float smem[256];
  int idx = blockIdx.x * blockDim.x
            + threadIdx.x;
  smem[threadIdx.x] =
    (idx < n) ? data[idx] : 0.0f;
  __syncthreads();
  for (int s = blockDim.x/2; s > 0; s >>= 1) {
    if (threadIdx.x < s)
      smem[threadIdx.x] += smem[threadIdx.x+s];
    __syncthreads();
  }
  if (threadIdx.x == 0)
    atomicAdd(result, smem[0]);
}

30x faster for matrix math

• Block-level programming (not per-thread like CUDA)
• Auto-handles coalescing, shared memory, tensor cores
• @triton.autotune searches for best block sizes
• 573 TFLOPS on 4096x4096 matmul (83% of cuBLAS)

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr, M, N, K,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr):
  pid_m = tl.program_id(0)
  pid_n = tl.program_id(1)
  offs_m = pid_m*BLOCK_M + tl.arange(0, BLOCK_M)
  offs_n = pid_n*BLOCK_N + tl.arange(0, BLOCK_N)
  acc = tl.zeros((BLOCK_M, BLOCK_N),
                dtype=tl.float32)
  for k in range(0, K, BLOCK_K):
    a = tl.load(a_ptr + ...)
    b = tl.load(b_ptr + ...)
    acc += tl.dot(a, b) # Tensor cores!
  tl.store(c_ptr + ..., acc)

Attention Matrix Size Scaling

GPU SRAM per SM: only ~200 KB

Standard Softmax (3 passes):

Online Softmax (1 pass):

def online_softmax(x):
  m, d = float('-inf'), 0.0
  for xi in x:
    m_new = max(m, xi)
    # Correction factor!
    d = d * math.exp(m - m_new) \
      + math.exp(xi - m_new)
    m = m_new
  return [math.exp(xi-m)/d for xi in x]

Step-Through

SRAM Budget Per Tile

(B=64, d=128, fp16)

FlashInfer: Inference-Specialized

• Variable sequence lengths (ragged tensors, no padding)
• Paged KV cache support
• Separate prefill vs decode kernels

GQA: 32 heads / 8 KV heads, seq=2048

Static: wait for all to finish

Continuous: check every step

OSDI 2022 — Yu et al.

Scheduler Strategies

Savings by Workload

def match_prefix(self, tokens):
  node, matched = self.root, 0
  for token in tokens:
    if token in node.children:
      node = node.children[token]
      matched += 1
    else:
      break
  return node, matched

Naive: Contiguous Allocation

Paged: Fixed-Size Blocks

Block Table

• No external fragmentation
• Copy-on-write for beam search
• Waste: at most 1 block per sequence

Monolithic Prefill

Chunked Prefill (512-token chunks)

Naive Padding Waste

Smart: Split Attention, Share FFN

Without CUDA Graph

With CUDA Graph

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
  output = model.decode_step(
    input, kv_cache)
# Inference loop
for token in range(max_tokens):
  g.replay()  # single launch!

How it works

• Router selects top_k experts per token
• Only activated experts run — sparse computation
• Output = weighted sum of selected expert outputs
• Router is just a linear projection: d_model → num_experts

Architecture Flow

# Router selects 2 of 8 experts
router_logits = Linear(x)  # 4096 → 8
weights, indices = topk(logits, 2)
output = 0.7 * expert_3(x) + 0.3 * expert_5(x)

SSE Streaming Format

// Server-Sent Events over HTTP
data: {"choices": [{"delta": {"content": "Hello"}}]}

data: {"choices": [{"delta": {"content": " world"}}]}

data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}

data: [DONE]

SSE over HTTP — no WebSockets needed. Works everywhere, reconnects gracefully.

Key Metrics (7B on H100)

FlashAttention-3 (NeurIPS 2024)

• Warp-specialized async: producer warps (data movement) + consumer warps (compute) overlap
• Interleaved matmul & softmax — both run simultaneously
• FP8 path with block quantization — 2.6x more accurate than naive FP8

Hardware Generations

Production Adoption

• Google Search — spec decoding powers AI Overviews, same quality, faster
• Vertex AI — EAGLE-3 integration for production inference
• Red Hat / vLLM — "Speculators" library, standardized spec decode

KV cache transfer via NVLink / RDMA

Why Disaggregate?

Prefill and decode have radically different compute profiles. Different hardware for different regimes.

Results

Production Deployments

• Meta, LinkedIn, Mistral, HuggingFace via vLLM
• NVIDIA Dynamo + NIXL for KV transfer
• Emerging: intra-GPU disaggregation (Nexus)

• Disaggregation: prefill is compute-bound, decode is memory-bound — different hardware for each
• Batch size: moves you along the roofline — bigger batches push toward compute-bound
• Speculative decoding: decode has spare compute, use it to verify draft tokens for free
• FlashAttention: attention was memory-bound — IO-aware tiling makes it compute-bound
• MoE: sparse compute per token, but same memory bandwidth cost — stays efficient on the roofline