Research map / Architectures and efficiency / Attention and transformers
Bifurcated Attention for Single-Context Large-Batch Sampling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes context-aware bifurcated attention, which splits incremental-decoding attention into shared-context and per-sample KV parts to cut KV memory IO during single-context batch sampling, combined with multi-query/multi-group attention.
Problem
Incremental decoding in single-context batch sampling is memory-IO bound because the shared context KV cache is loaded once per batch index, causing latency to grow with batch size and context length.
Approach
Split the attention computation during incremental decoding into a context part (shared KV loaded once, broadcast across the batch) and a decoding part (per-sample KV), joined by concatenation/summation. This yields identical outputs and FLOPs but reduces KV memory IO by up to b-fold when context dominates, and it composes with multi-query attention.
Key results
- For CodeGen 16B multi-head with 2k context, bifurcated attention raises batch size from 5 to 128, increasing pass@k from 59.0% to 84.6%.
- Multi-head with bifurcated attention at batch 128 and context 10000 gives roughly 4x lower incremental decoding latency.
- A multi-query model needs about 10% more parameters than a multi-head model to match its capability (scaling-law comparison from 125M to 13B).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use bifurcated attention instead of SDPA | 7B multi-head model, compiled, batch size 16 · 8k context | per-token latency (ms): 26.19 -> 12.60 | Helped |
| Use bifurcated attention instead of SDPA | 7B multi-head model, compiled, batch size 16 · 16k context | per-token latency (ms): OOM -> 18.46 | Helped |
| Use bifurcated attention instead of SDPA | 7B multi-head model, batch size 1 · 32k context | per-token latency (ms): 44.94 -> 39.97 | Helped |
| Use bifurcated attention instead of SDPA | 7B multi-head model, compiled, batch size 2 · 32k context | per-token latency (ms): OOM -> 29.34 | Helped |
| Increase the number of attention groups (ablation) | Generalized multi-query attention · Validation loss versus model size | validation loss | Helped |
| Increase the number of attention groups (ablation) | Generalized multi-query attention · Multilingual HumanEval and MBXP, 13 programming languages | pass rate | Helped |
| Increase multi-query model size by 1.1 times (ablation) | Multi-query model compared with multi-head · Capability-matched model comparison | model size: 1.104 times | Helped |
Surprising findings
- With bifurcated attention, multi-head attention rivals multi-query in latency for moderate batch sizes (up to 64), removing the need to train a new multi-query model.
Limitations
- Benefit is limited to single-context batch sampling; not applicable to batch inference with distinct contexts.
- Multi-query models require ~10% larger size to match multi-head capability.
- Speculative decoding and incremental decoding latency still dominated by model loading for small contexts.
Experimental setup
Datasets: MBPP HumanEval MBXP
Benchmarks: MBPP Multilingual HumanEval MBXP
Metrics: per-step latency pass@n pass@top3 validation loss latency (ms)
Baselines: multi-head attention multi-query attention SDPA kernel
Builds on: multi-query attention generalized multi-query attention GPT-Fast
Closest papers by meaning
- CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.87)
- ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction (NeurIPS 2024 · similarity 0.87)
- MagicPIG: LSH Sampling for Efficient LLM Generation (ICLR 2025 · similarity 0.87)
- Attention-aware Inference Optimizations for Large Vision-Language Models with Memory-efficient Decoding (CVPR 2026 · similarity 0.86)
- QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.86)
- FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.