You enabled sparse attention. Your model still chokes at 128K tokens. The indexer is why — and PIVOT fixes it without touching your weights.
TL;DR
- Sparse attention’s indexer scores all L tokens per query → still O(L²)
- PIVOT groups nearby queries (which select ~90% overlapping top-k tokens), runs one proxy scan per group → O(L²/g)
- Result: 4× indexer speedup, 1.6× end-to-end latency reduction on DeepSeek-V3.2 and GLM-5.1
- Training-free: plug into existing DSA models at inference time
- Two modes: PIVOT-Reuse (max speed) and PIVOT-Refine (matches dense indexer accuracy)
The Problem
Dynamic Sparse Attention (DSA) should make long-context inference fast. Score all tokens → pick top-k → attend only to those k. Complexity drops from O(L²) to O(L·k).
Except scoring all tokens is itself O(L²). The “indexer” does a full O(L) scan per query position. With L queries, you’re back to O(L²). At 100K tokens, the indexer dominates latency. Sparse attention becomes a lie.
How It Works
Observation 1: Adjacent queries share ~90% of their top-k token selections — they process nearly identical context.
Observation 2: Indexer scores are long-tailed — a proxy query produces a reliable candidate set.
PIVOT’s algorithm:
group = [q_i, q_{i+1}, ..., q_{i+g-1}]
proxy_q = mean(group)
# ONE scan instead of g scans
scores = proxy_q · K[:i] # O(L)
C = top-K(scores) # candidate set, K = 2 × top_k
# per-query refine (PIVOT-Refine)
for q in group:
refine_scores = q · K[C] # O(K), not O(L)
final_indices[q] = top-k(refine_scores)
Enter fullscreen mode Exit fullscreen mode
Indexer cost: O(L²) → O(L²/g). With g=8 that’s 8× fewer full scans.
Show Me The Code
import torch
import torch.nn.functional as F
def pivot_refine(
queries: torch.Tensor, # (seq_len, num_heads, head_dim)
keys: torch.Tensor,
top_k: int = 64,
group_size: int = 8,
candidate_ratio: float = 2.0,
) -> torch.Tensor:
"""
Returns sparse-attention token indices.
Shape: (seq_len, num_heads, top_k)
"""
seq_len, num_heads, head_dim = queries.shape
K = int(top_k * candidate_ratio)
indices = torch.zeros(seq_len, num_heads, top_k,
dtype=torch.long, device=queries.device)
for h in range(num_heads):
q_h = queries[:, h, :]
k_h = keys[:, h, :]
for g_start in range(0, seq_len, group_size):
g_end = min(g_start + group_size, seq_len)
# Step 1: one proxy scan per group (not per query)
proxy = q_h[g_start:g_end].mean(dim=0)
for pos in range(g_start, g_end):
if pos == 0:
continue
proxy_scores = proxy @ k_h[:pos].T
actual_K = min(K, pos)
_, C = torch.topk(proxy_scores, actual_K)
# Step 2: O(K) refine — cheap!
refine_scores = q_h[pos] @ k_h[C].T
actual_k = min(top_k, actual_K)
_, best = torch.topk(refine_scores, actual_k)
indices[pos, h, :actual_k] = C[best]
return indices
def pivot_reuse(
queries: torch.Tensor,
keys: torch.Tensor,
top_k: int = 64,
group_size: int = 8,
candidate_ratio: float = 2.0,
) -> torch.Tensor:
"""Max-speed variant — all queries in group share indices."""
seq_len, num_heads, head_dim = queries.shape
K = int(top_k * candidate_ratio)
indices = torch.zeros(seq_len, num_heads, top_k,
dtype=torch.long, device=queries.device)
for h in range(num_heads):
q_h = queries[:, h, :]
k_h = keys[:, h, :]
for g_start in range(0, seq_len, group_size):
g_end = min(g_start + group_size, seq_len)
ref_pos = max(g_start, 1)
proxy = q_h[g_start:g_end].mean(dim=0)
proxy_scores = proxy @ k_h[:ref_pos].T
actual_K = min(K, ref_pos)
_, C = torch.topk(proxy_scores, actual_K)
for pos in range(g_start, g_end):
actual_k = min(top_k, actual_K)
indices[pos, h, :actual_k] = C[:actual_k]
return indices
# --- Quick test ---
if __name__ == "__main__":
S, H, D = 1024, 8, 64
q = torch.randn(S, H, D)
k = torch.randn(S, H, D)
v = torch.randn(S, H, D)
idx = pivot_refine(q, k, top_k=64, group_size=8)
print(f"Index shape: {idx.shape}") # (1024, 8, 64)
out = torch.zeros_like(q)
scale = D ** -0.5
for pos in range(1, S):
for h in range(H):
sel = idx[pos, h]
w = F.softmax(q[pos, h] @ k[sel, h].T * scale, dim=-1)
out[pos, h] = w @ v[sel, h]
print(f"Output shape: {out.shape}") # (1024, 8, 64)
Enter fullscreen mode Exit fullscreen mode
⚠️ Reference implementation only. Production speedups need Triton/CUDA kernels. Official code not yet released (July 2026).
Benchmark Results
Tested on DeepSeek-V3.2 and GLM-5.1 — LongBench + RULER:
Method Indexer Speed E2E Latency Accuracy Dense DSA (baseline) 1× 1× ✅ full PIVOT-Refine ~3× faster −28% ✅ ≈ baseline PIVOT-Reuse 4× faster −37.5% ⚠️ minor dropPIVOT-Refine = dense-indexer accuracy + 3× speed. Zero retraining.
Gotchas & Limitations
DSA-only: Works out of the box on DeepSeek-V3.2 / GLM-5.1. Standard full-attention models need DSA fine-tuning first.
Group size is manual: No adaptive strategy in the paper — tune g per model/sequence length.
Batch inference untested: Single-sequence results only.
Prefill vs. decode split missing: The 1.6× E2E number doesn’t break down by phase.
🚀 Try It Today
- Running DeepSeek-V3.2 via vLLM? Request PIVOT integration upstream.
- Prototype with the reference code above; profile vs. vanilla DSA on your target seq length.
- Stack with IndexCache (arXiv:2603.12201) for cross-layer index reuse on top.
- Paper: arXiv:2607.24593
Sources
- Hong Liu et al., “PIVOT: Efficient Query-Group Indexing for Token-Level Sparse Attention”, arXiv:2607.24593 — https://arxiv.org/abs/2607.24593
- “IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse”, arXiv:2603.12201 — https://arxiv.org/abs/2603.12201
- “MISA: Mixture of Indexer Sparse Attention”, arXiv:2605.07363 — https://arxiv.org/abs/2605.07363
- “FlashMemory-DeepSeek-V4: Lookahead Sparse Attention”, arXiv:2606.09079 — https://arxiv.org/abs/2606.09079
- DeepSeek-V3.2 vLLM blog — https://vllm.ai/blog/2025-09-29-deepseek-v3-2
What’s your experience with sparse attention in production? Drop a comment.
답글 남기기