Day 2: Transformer Attention as a Systems Diagram

Day 2: Transformer Attention as a Systems Diagram

Attention is usually taught as matrix multiplies in a paper. In production it is a data path: Q/K/V projections, KV cache growth per token, batch packing, and accelerator memory ceilings. If you design agents without that diagram, you mis-price latency, oversubscribe concurrency, and blame “the model” for queueing problems.

⚡ TL;DR: Prefill is compute-heavy; decode is memory-bandwidth-heavy. KV cache grows with layers × heads × sequence × dtype. Batching raises tokens/sec until you hit memory or latency SLOs. Architect prompts and tools to shrink prefill and reuse cacheable prefixes.

Attention as a systems block diagram

At decode step t, the serving engine reads cached keys/values for tokens 1..t-1 and computes attention for the new token. Split the path:

  • Prefill (first forward on the prompt): highly parallel over prompt tokens; builds the initial KV cache; dominates time-to-first-token (TTFT) for long prompts.
  • Decode (each new token): mostly sequential; bottleneck is loading the growing KV cache from high-bandwidth memory.
User + tool + RAG tokens ──► Prefill ──► KV cache (HBM/DRAM)
                                   │
                                   ▼
                        Decode loop (token t) ──► sample ──► stream
                                   │
                                   └── read KV[1..t-1] every step

❌ Treating “200k context supported” as free once the SDK accepts the HTTP request. Prefill time and KV memory scale with prompt length and concurrent sessions.

Instrument TTFT and inter-token latency separately. A system can look “fast” on average tokens/sec while interactive agents feel dead because TTFT spiked on fat tool schemas.

KV cache: the hidden capacity limit

Order-of-magnitude planning (always verify against your engine and model card):

# ✅ Capacity sketch for capacity planning meetings
def kv_gib(layers, kv_heads, seq, head_dim, dtype_bytes=2, batch=1):
    # K and V per layer; multiply by batch for concurrent sequences
    return batch * layers * 2 * kv_heads * seq * head_dim * dtype_bytes / (1024**3)

print(kv_gib(layers=80, kv_heads=8, seq=32_000, head_dim=128, batch=16), "GiB")

Implications that change architecture reviews:

  • Long system prompts duplicated per concurrent session burn memory linearly with concurrency.
  • Prompt caching (Bedrock and others) amortizes stable prefixes when bytes match exactly.
  • Multi-turn agents that re-send giant OpenAPI blobs every hop pay prefill tax repeatedly.

Shrink what must live in the KV path: retrieve docs instead of pasting manuals; keep tool schemas minimal; pin a stable system prefix.

Batching without lying to product

Continuous batching improves accelerator utilization by interleaving decode steps across requests. Product still has a p95 latency SLO. Design queues explicitly:

Knob Effect
Max concurrent sequences KV memory bound
Max batch tokens Prefill spikes
Priority / preemption Interactive vs offline eval
Speculative decoding Lower TTFT if draft accept rate is high
# ✅ Separate interactive agent traffic from offline eval jobs
ROUTING = {
    "interactive": {"queue": "low-latency", "max_input_tokens": 24_000},
    "eval_batch": {"queue": "throughput", "max_input_tokens": 128_000},
}

❌ One shared queue where a 100k RAG evaluation starves coding-agent TTFT during business hours.

Design rules derived from the data path

  • Keep stable prefixes identical byte-for-byte to maximize cache hits.
  • Move bulky reference material out of the prompt into retrieval (Day 5+).
  • Prefer streaming decode for UX; never block the UI on full completion.
  • Cap agent tool loops — each tool result re-enters prefill/decode economics.
  • For Bedrock on-demand vs provisioned throughput, match the traffic shape: spiky interactive vs steady batch.

When someone proposes “just stuff the whole monorepo into context,” answer with a KV and TTFT sketch, not a vibe.

Closing checklist

  • [ ] Draw prefill vs decode for your serving stack on a whiteboard
  • [ ] Estimate KV memory at p95 prompt length × target concurrency
  • [ ] Split interactive vs batch inference queues
  • [ ] Make system+tool prefixes stable for prompt caching
  • [ ] Track TTFT and decode tok/s as separate SLIs with alerts
  • [ ] Cap max input tokens per traffic class in the gateway

Worked example: why your agent feels slow

A 12k-token system+tools prefix with cold cache might cost 1.5–4s of prefill on a busy endpoint before the first visible token. Ten concurrent sessions with 16k effective KV each can exhaust memory and force queueing that looks like “model stupidity.”

Interactive SLO: TTFT p95 < 1.2s, tok/s p50 > 40
If TTFT regresses: shrink prefix, enable prompt cache, split queues
If tok/s regresses: check batch contention / throttling, not the prompt author

Draw the request path on a wiki page: API Gateway → queue → worker → engine. Label where KV lives and who owns max concurrency. That diagram prevents capacity guesswork during launches.

Failure modes to watch

  • Shared queues mixing eval jobs with IDE agents.
  • Unstable prefixes (timestamps in system prompts) destroying cache hit rates.
  • Oversized tool JSON re-sent every hop of a ReAct loop.
  • Ignoring decode memory when planning “just raise max_tokens.”

Field notes from production

Provisioned throughput decisions should start from p95 concurrent sequences × KV footprint, not average QPS. Interactive coding agents are bursty at the top of the hour. Keep a kill switch that sheds batch eval traffic when interactive TTFT SLO burns. Document the exact prompt-cache key inputs (system hash, tool schema hash) in the runbook.

Implementation sketch

# Implementation sketch: classify traffic before enqueue
def enqueue(req):
    q = "interactive" if req.feature == "ide_agent" else "batch"
    if q == "interactive" and req.input_tokens > 24_000:
        req = req.truncate_rag_to(12_000)
    return queues[q].put(req)

Series navigation

← Day 1 · Day 3 →

Last updated September 11, 2026


Discover more from CheatCoders

Subscribe to get the latest posts sent to your email.

Comments

No comments yet. Why don’t you start the discussion?

Leave a Reply