Choosing the right attention configuration

The attention mechanism determines three things simultaneously: model quality, inference latency, and serving cost. The wrong choice wastes GPU memory on KV cache entries you don't need, or starves the model of attention patterns it requires. This lesson builds a decision framework mapping workload constraints to concrete attention configurations, with worked memory budgets.

The decision variables

Every deployment has five constraints in tension:

Model size — parameter count determines base quality and weight memory
Target context length — how many tokens the model must process per request
Batch size — concurrent requests sharing GPU resources
GPU memory budget — total available VRAM across your serving GPUs
Latency requirements — time-to-first-token (TTFT) and decode throughput (tokens/sec)

The attention configuration interacts with all five. MHA vs GQA vs MQA determines KV cache size per request. Sparse vs full determines compute per layer. The combination constrains how many concurrent requests fit in memory and how fast each one generates.

MHA: maximum quality, maximum cost

Multi-Head Attention stores separate K and V projections for each head. A model with 32 heads, head_dim 128, and 32 layers stores:

\text{KV} \text{cache} \text{per} \text{token} = 2 \times n_{\text{heads}} \times \text{head}\_\text{dim} \times n_{\text{layers}} \times \text{bytes}\_\text{per}\_\text{param}

= 2 \times 32 \times 128 \times 32 \times 2 = 524, 288 \text{bytes} = 0.5 \text{MB} \text{per} \text{token}(\text{float16})

At 128K context: $128000 \times 0.5 \text{MB} = 64 \text{GB}$ of KV cache per request. A single request consumes an entire A100-80GB after model weights.

Use MHA when:

Training from scratch and maximizing quality per FLOP (research, foundation model pre-training)
Small models (< 3B params) where KV cache is negligible relative to available memory
Short context (< 4K) where even full MHA KV cache is manageable
Batch size 1 serving where memory is not shared

GQA: the production default

Grouped-Query Attention (Ainslie et al. 2023) shares K and V heads across groups. With 32 query heads and 8 KV groups (GQA-8, the Llama 3 configuration), the KV cache drops to 8/32 = 25% of MHA. The actual overhead:

\text{KV} \text{cache} \text{per} \text{token}(\text{GQA} - 8) = 2 \times 8 \times 128 \times 32 \times 2 = 131, 072 \text{bytes} = 0.125 \text{MB} \text{per} \text{token}

At 128K context: $128000 \times 0.125 \text{MB} = 16 \text{GB}$ per request. Four concurrent requests fit alongside a quantized 70B model on 4×A100-80GB.

The quality cost is minimal: Llama 3 70B (GQA-8) matches or exceeds Llama 2 70B (full MHA) on all standard benchmarks — the saved KV cache parameters are redistributed as additional capacity in the feedforward layers (Touvron et al. 2023).

Use GQA when:

Serving any model larger than 7B in production
Context lengths exceed 8K tokens
You need concurrent request serving (batch size > 1)
You want the best quality/memory tradeoff available today

Llama 3 (8B, 70B, 405B), Mistral 7B/8x7B, Gemma 2, Command R+ — all use GQA with 8 KV groups. This is the settled industry default as of 2026.

MQA: extreme compression

Multi-Query Attention (Shazeer 2019) uses a single KV head shared across all query heads. KV cache drops to 1/32 of MHA:

\text{KV} \text{cache} \text{per} \text{token}(\text{MQA}) = 2 \times 1 \times 128 \times 32 \times 2 = 16, 384 \text{bytes} = 0.016 \text{MB} \text{per} \text{token}

At 128K context: $128000 \times 0.016 \text{MB} = 2 \text{GB}$ per request. On a single A10-24GB, after loading a 7B model (int4, 3.5 GB), you have ~18 GB free — enough for 9 concurrent requests at 128K context, or hundreds at shorter contexts.

The quality cost is measurable: MQA models lose 0.5–1.5% on standard benchmarks relative to GQA-8 at the same parameter count (Ainslie et al. 2023, Table 2). For latency-critical applications where the quality gap is acceptable, MQA enables dramatically higher throughput.

Use MQA when:

Edge/mobile deployment with severe memory limits (phones, embedded GPUs)
Real-time systems requiring sub-50ms TTFT
Maximum concurrent connections matter more than peak quality (chatbots serving thousands of users on limited GPUs)

Falcon 40B, PaLM, and StarCoder use MQA. It's less common in 2026 because GQA captures most of the benefit with less quality regression.

MLA: latent compression

Multi-head Latent Attention (DeepSeek-V2, Liu et al. 2024) compresses the KV cache into a low-rank latent representation. Instead of caching full K and V vectors, it caches a compressed latent c_t of dimension d_c (typically 512), from which K and V are reconstructed on the fly during attention:

$c_{t} = W_{\text{dkv}} @ x_{t}$ (compress to d_c dimensions)

$k_{t} = W_{\text{uk}} @ c_{t}$ (reconstruct keys)

$v_{t} = W_{\text{uv}} @ c_{t}$ (reconstruct values)

KV cache per token: $d_{c} \times 2 \text{bytes} = 512 \times 2 = 1 \text{KB}$ per layer. For DeepSeek-V2's 60 layers: $60 \times 1 \text{KB} = 60 \text{KB} \text{per} \text{token}$ . At 128K context: $128000 \times 60 \text{KB} = 7.5 \text{GB}$ per request — compared to 64 GB for a comparable MHA model.

The cost: additional compute to decompress c_t into K and V at each layer during generation. This is a matrix multiply of size d_c × d_model per layer per token — roughly 2x the per-token compute of standard GQA. The tradeoff favors MLA when KV cache is the bottleneck (large batches, long contexts).

Use MLA when:

Serving very large batch sizes where KV cache dominates memory
Designing new architectures from scratch (not widely adoptable as a retrofit)
DeepSeek-family models (V2, V3, R1) — currently the only production models using MLA

FlashAttention: always on

FlashAttention (Dao et al. 2022, Dao 2023) is not an attention variant — it's a memory-efficient kernel that computes exact standard attention. By tiling the computation and keeping intermediate results in SRAM (on-chip shared memory) rather than writing the N×N attention matrix to HBM (GPU main memory), it achieves:

2–4x wall-clock speedup over naive PyTorch attention
O(N) memory instead of O(N²) for the attention computation
Exact results — numerically identical to standard attention (up to floating-point ordering)

There is no reason not to use FlashAttention. All major serving frameworks (vLLM, TensorRT-LLM, SGLang) use it by default. PyTorch 2.0+ includes it as torch.nn.functional.scaled_dot_product_attention with automatic backend selection.

GPU memory budget: worked examples

Example A: Llama 3 70B at 128K context on 4×A100-80GB

Total memory available: 4 × 80 = 320 GB

Model weights (int4 quantization): 70B params × 0.5 bytes = 35 GB. With tensor parallelism across 4 GPUs: ~9 GB per GPU, 35 GB total.

KV cache per request (GQA-8):

Llama 3 70B: 80 layers, 8 KV heads, head_dim 128
Per token: $2 \times 8 \times 128 \times 80 \times 2 = 327, 680 \text{bytes} ≈ 0.31 \text{MB}$
At 128K context: $128000 \times 0.31 \text{MB} = 40 \text{GB} \text{per} \text{request}$

Available for KV cache: 320 - 35 (weights) - 15 (framework overhead, activations) = 270 GB

Max concurrent requests: $\frac{270}{40} = 6 \text{requests}$

With KV cache quantization (int8): cache drops to 20 GB per request → $\frac{270}{20} = 13 \text{concurrent} \text{requests}$

With KV cache quantization (int4, via KIVI): cache drops to 10 GB per request → $\frac{270}{10} = 27 \text{concurrent} \text{requests}$

Example B: Mistral 7B real-time chat on a single A10-24GB

Model weights (int4): 7B × 0.5 = 3.5 GB

KV cache per request (GQA-8):

Mistral 7B: 32 layers, 8 KV heads, head_dim 128
Per token: $2 \times 8 \times 128 \times 32 \times 2 = 131, 072 \text{bytes} ≈ 0.125 \text{MB}$
At 8K context (typical chat): $8000 \times 0.125 \text{MB} = 1.0 \text{GB} \text{per} \text{request}$
With sliding window (W = 4096): cache capped at $4096 \times 0.125 \text{MB} = 0.5 \text{GB} \text{per} \text{request}$

Available for KV cache: 24 - 3.5 - 2 (overhead) = 18.5 GB

Max concurrent requests (8K, sliding window): $\frac{18.5}{0.5} = 37 \text{requests}$

Max concurrent requests (8K, full cache): $\frac{18.5}{1.0} = 18 \text{requests}$

The sliding window doubles serving capacity at this configuration.

Example C: Batch processing with Llama 3 8B on A100-40GB

Model weights (float16): 8B × 2 = 16 GB

KV cache per request (GQA-8):

Llama 3 8B: 32 layers, 8 KV heads, head_dim 128
Per token: $2 \times 8 \times 128 \times 32 \times 2 = 131, 072 \text{bytes} ≈ 0.125 \text{MB}$
At 1K context (short batch jobs): $1000 \times 0.125 \text{MB} = 125 \text{MB} \text{per} \text{request}$

Available for KV cache: 40 - 16 - 2 = 22 GB

Max concurrent requests: $22000 \frac{\text{MB}}{125} \text{MB} = 176 \text{requests}$

At short contexts, the KV cache is small and the bottleneck shifts to compute (processing 176 prefills simultaneously saturates the GPU's FLOP budget long before memory fills).

Building a serving capacity calculator

pythonData structures

from dataclasses import dataclass

@dataclass
class ModelConfig:
    name: str
    params_billions: float
    n_layers: int
    n_kv_heads: int
    head_dim: int
    max_context: int

@dataclass
class GPUConfig:
    name: str
    memory_gb: float
    count: int = 1

@dataclass
class ServingConfig:
    weight_quantization: str = "float16"  # "float16", "int8", "int4"
    kv_quantization: str = "float16"      # "float16", "int8", "int4"
    target_context: int = 8192
    sliding_window: int = None            # None = full attention
    overhead_gb: float = 2.0              # framework, activations, etc.


QUANT_BYTES = {"float16": 2, "int8": 1, "int4": 0.5}

pythonComputation

def compute_serving_capacity(
    model: ModelConfig,
    gpu: GPUConfig,
    serving: ServingConfig
) -> dict:
    total_memory_gb = gpu.memory_gb * gpu.count

    # Model weight memory
    bytes_per_param = QUANT_BYTES[serving.weight_quantization]
    weight_memory_gb = model.params_billions * 1e9 * bytes_per_param / 1e9

    # KV cache per token
    kv_bytes_per_param = QUANT_BYTES[serving.kv_quantization]
    kv_per_token_bytes = (
        2 * model.n_kv_heads * model.head_dim * model.n_layers * kv_bytes_per_param
    )

    # Effective context length (capped by sliding window if set)
    effective_context = serving.target_context
    if serving.sliding_window is not None:
        effective_context = min(serving.target_context, serving.sliding_window)

    # KV cache per request
    kv_per_request_gb = effective_context * kv_per_token_bytes / 1e9

    # Available memory for KV cache
    available_gb = total_memory_gb - weight_memory_gb - serving.overhead_gb

    if available_gb <= 0:
        return {"error": "Model weights exceed available memory"}

    max_concurrent = int(available_gb / kv_per_request_gb) if kv_per_request_gb > 0 else 0

    return {
        "model": model.name,
        "gpu": f"{gpu.count}x {gpu.name} ({total_memory_gb:.0f} GB)",
        "weight_memory_gb": round(weight_memory_gb, 1),
        "kv_per_token_bytes": int(kv_per_token_bytes),
        "effective_context": effective_context,
        "kv_per_request_gb": round(kv_per_request_gb, 2),
        "available_for_kv_gb": round(available_gb, 1),
        "max_concurrent_requests": max_concurrent,
    }

pythonExample configurations

# Example configurations
llama3_70b = ModelConfig("Llama 3 70B", 70, 80, 8, 128, 128000)
llama3_8b = ModelConfig("Llama 3 8B", 8, 32, 8, 128, 128000)
mistral_7b = ModelConfig("Mistral 7B", 7.2, 32, 8, 128, 32768)

# Scenario A: Llama 3 70B, 128K context, 4x A100-80GB
result_a = compute_serving_capacity(
    llama3_70b,
    GPUConfig("A100-80GB", 80, count=4),
    ServingConfig(weight_quantization="int4", target_context=128000, overhead_gb=15)
)

# Scenario B: Mistral 7B, 8K context with sliding window, single A10
result_b = compute_serving_capacity(
    mistral_7b,
    GPUConfig("A10-24GB", 24),
    ServingConfig(weight_quantization="int4", target_context=8192, sliding_window=4096)
)

# Scenario C: Llama 3 8B, 1K batch context, A100-40GB
result_c = compute_serving_capacity(
    llama3_8b,
    GPUConfig("A100-40GB", 40),
    ServingConfig(weight_quantization="float16", target_context=1024)
)

for label, result in [("A", result_a), ("B", result_b), ("C", result_c)]:
    print(f"\n{'='*50}")
    print(f"Scenario {label}: {result['model']}")
    print(f"  GPU: {result['gpu']}")
    print(f"  Weights: {result['weight_memory_gb']} GB")
    print(f"  KV/token: {result['kv_per_token_bytes']} bytes")
    print(f"  Effective context: {result['effective_context']:,} tokens")
    print(f"  KV/request: {result['kv_per_request_gb']} GB")
    print(f"  Available for KV: {result['available_for_kv_gb']} GB")
    print(f"  Max concurrent requests: {result['max_concurrent_requests']}")

Running this produces:

Scenario A: Llama 3 70B
  GPU: 4x A100-80GB (320 GB)
  Weights: 35.0 GB
  KV/token: 131072 bytes
  Effective context: 128,000 tokens
  KV/request: 16.78 GB
  Available for KV: 270.0 GB
  Max concurrent requests: 16

Scenario B: Mistral 7B
  GPU: 1x A10-24GB (24 GB)
  Weights: 3.6 GB
  KV/token: 131072 bytes
  Effective context: 4,096 tokens
  KV/request: 0.54 GB
  Available for KV: 18.4 GB
  Max concurrent requests: 34

Scenario C: Llama 3 8B
  GPU: 1x A100-40GB (40 GB)
  Weights: 16.0 GB
  KV/token: 131072 bytes
  Effective context: 1,024 tokens
  KV/request: 0.13 GB
  Available for KV: 22.0 GB
  Max concurrent requests: 169

Decision matrix

For context < 8K tokens:

GQA is the default. KV cache is small at these lengths — the attention variant barely matters. Optimize model quality.
FlashAttention handles the compute efficiently. No need for sparse patterns.
Capacity is usually compute-bound, not memory-bound.

For context 8K–32K tokens:

GQA-8 remains the sweet spot. KV cache starts mattering for concurrent serving.
Consider KV cache quantization (int8) if batch size needs exceed available memory.
Sliding window viable if the task is chat/code where locality holds.

For context 32K–128K tokens:

KV cache dominates memory. GQA-8 + int8 KV quantization is the minimum viable config.
Sliding window (Mistral-style) or sparse attention recommended unless tasks require full retrieval across the context.
MLA (DeepSeek-style) provides the best cache compression but requires a model trained with it.

For context > 128K tokens:

Ring Attention for distributed exact attention (training and long-context inference).
Hybrid Mamba-attention (Jamba) for single-node serving — Mamba layers handle most positions cheaply, attention layers provide long-range precision.
Pure Mamba/SSM for streaming use cases where constant-memory generation matters more than exact retrieval.

For maximum throughput (tokens/sec/GPU):

MQA or aggressive GQA (4 groups) to minimize memory bandwidth for KV reads during decode.
Speculative decoding on top of any attention variant for 2–3x decode speedup.
Continuous batching (vLLM, SGLang) to keep the GPU saturated by dynamically interleaving prefill and decode across requests.

The overarching principle: attention configuration is a memory allocation decision. Every byte used for KV cache is a byte unavailable for concurrent requests. Quantize the cache (int8 is nearly lossless for KV values), compress it (GQA, MLA), or bound it (sliding window) — but the goal is always the same: serve more users at acceptable quality within the fixed memory budget of your hardware.

← Previous