optimize-prefill-vs-decode-kernels

Skill: Optimize Prefill vs Decode Kernels

Purpose

Guide the agent through choosing and tuning kernels for the prefill phase versus the decode phase of LLM inference. The two phases have fundamentally different arithmetic intensity, occupy different sides of the roofline, and respond to different optimizations. Continuous batching and speculative decoding shift the balance and must be reasoned about explicitly.

Use this when

• Designing or selecting kernels for an LLM inference engine (matmul, attention, KV cache access, normalization fusion).
• An existing engine reuses a single kernel across both phases and decode latency or throughput is poor.
• Tuning tile shapes, split-K, or persistent kernels separately for prefill and decode.
• Adding speculative decoding, lookahead decoding, EAGLE, or Medusa, where decode shifts toward a small-batch GEMM regime.
• Implementing chunked prefill or hybrid prefill/decode batches and choosing per-batch kernels.
• Diagnosing why TTFT is good but inter-token latency is bad, or vice versa.

Do not use this when

• An off-the-shelf kernel (cuBLASLt heuristics, vendor attention) already saturates the relevant resource. Trust the library until profiling shows a gap.
• The bottleneck is outside the kernel layer (Python overhead, scheduler, RPC, host-device copies). Fix the host path first.
• Training, not inference. Training runs in a different regime — large M dominates throughout, so this skill's decode reasoning does not apply.
• A non-batched, latency-insensitive offline workload where one configuration is good enough.

Inputs the agent should gather first

• Phase: prefill, decode, or mixed (continuous batching). For mixed, the per-request M distribution.
• Effective M per request: prefill M = prompt length (typically 256–32k+); decode M = 1, or M = k for speculative decoding (typical k = 2–8, lookahead trees up to 32–64).
• Batch size B: number of concurrent requests in this kernel call. With continuous batching, the kernel sees the aggregated M across all requests.
• Model dimensions and dominant K: hidden size N, head dim d_h, num heads h, intermediate FFN size; K = hidden for QKV/output proj, K = intermediate for FFN down proj, K = head_dim for attention.
• dtype of weights, activations, KV cache, and matmul accumulator (fp16, bf16, fp8, int8, int4 weight-only).
• KV cache layout: paged (PagedAttention), contiguous, MQA, GQA group size. Cache bytes per token = 2 * num_layers * num_kv_heads * head_dim * dtype_bytes.
• Target hardware: HBM bandwidth (H100 SXM ≈ 3.35 TB/s, H200 ≈ 4.8 TB/s, MI300X ≈ 5.3 TB/s, A100 80GB ≈ 2.0 TB/s), tensor core gen, FP8/INT8 support.
• Whether speculative decoding is active, draft model, verified-tokens-per-step distribution (acceptance rate × tree size).
• SLO targets: TTFT, p50/p99 inter-token latency, throughput. The right kernel for throughput is rarely the right one for tail latency.