How Much Dense Attention is Necessary? Oracle-Guided Sparse Prefill for Full/GQA Layers in Hybrid Long-Context Models
About
Long-context prefill remains expensive because full/GQA layers still score the historical sequence, even in hybrid models with local, sparse, linear, or recurrent components. We study how much dense attention is needed to preserve task-level behavior under explicit support granularity and top-k budgets. We introduce an attention-mass top-k oracle for existing GQA checkpoints: for each layer and query position, it computes dense attention, selects head-averaged token support, and recomputes attention only on that support. The oracle is a diagnostic reference, not a deployable accelerator, and separates sparse-budget feasibility from indexer error and runtime realization effects. On Qwen-family retrieval-heavy evaluations, the longest per-query oracle rows stay within 1 point of dense, and a Qwen3.5-9B RULER-style sweep from 4K to 100K stays within 0.48 points. Guided by the oracle, we derive a head-collapsed auxiliary indexer trained by KL distillation from dense attention-mass distributions while keeping the backbone frozen. With separately distilled Qwen3.5-0.8B and Qwen3.5-9B indexers, the reported 16K/32K validation macro gaps are +2.04 and +1.13 points, treated as quality preservation rather than improvement; fused selection-block-shared support can introduce a larger realization gap. Preliminary single-card TTFT measurements show distilled-indexer sparse serving speedups of 1.71x for Qwen3.5-0.8B on NPU and 1.93x for Qwen3.5-9B on GPU against its dense FlashAttention-2 baseline. Additional random-init stress rows reach 3.44x, indicating sparse-runtime headroom but not validated output quality. This first release separates oracle feasibility, distilled-indexer quality, and runtime headroom, leaving a fully matched quality-latency frontier to future work.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Time to First Token (TTFT) measurement | Context sequences 32K-256K vLLM 0.18.0 / 0.20.0 (online-serving) | TTFT (Dense, s)2.944 | 7 | |
| Long-context document understanding | MMLongBench-Doc 16K context slice | Score48.57 | 4 | |
| Long-context document understanding | MMLongBench-Doc 32K context slice | MMLongBench-Doc (32K) Score45 | 4 | |
| Long-context Language Understanding | LongBench 16K context slice v2 | Score57.89 | 4 | |
| Long-context Language Understanding | LongBench 32K context slice v2 | Score56.86 | 4 | |
| Long-context retrieval and reasoning | RULER 16K context slice | Score90 | 4 | |
| Long-context retrieval and reasoning | RULER 32K context slice | Score95.83 | 4 | |
| Long-context video understanding | VideoMME 16K context slice | Score66 | 4 | |
| Long-context video understanding | VideoMME 32K context slice | Score70 | 4 |