PIVOT Explained — From Paper to Working Code in 10 Minutes

A developer introduced PIVOT, an algorithm that reduces the indexer cost of dynamic sparse attention from O(L²) to O(L²/g) by grouping queries and using a proxy query for a single scan per group. The algorithm exploits the observation that adjacent queries share ~90% of their top-k token selections, enabling efficient long-context inference without modifying model weights.

You enabled sparse attention. Your model still chokes at 128K tokens. The indexer is why — and PIVOT fixes it without touching your weights. 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. 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 Indexer cost: O L² → O L²/g . With g=8 that's 8× fewer full scans. python 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 ⚠️ Reference implementation only. Production speedups need Triton/CUDA kernels. Official code not yet released July 2026 . 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 drop | PIVOT-Refine = dense-indexer accuracy + 3× speed. Zero retraining. 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. What's your experience with sparse attention in production? Drop a comment.