{"slug": "pivot-explained-from-paper-to-working-code-in-10-minutes", "title": "PIVOT Explained — From Paper to Working Code in 10 Minutes", "summary": "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.", "body_md": "You enabled sparse attention. Your model still chokes at 128K tokens. The indexer is why — and PIVOT fixes it without touching your weights.\n\nDynamic 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).\n\nExcept **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.\n\n**Observation 1**: Adjacent queries share ~90% of their top-k token selections — they process nearly identical context.\n\n**Observation 2**: Indexer scores are long-tailed — a proxy query produces a reliable candidate set.\n\n**PIVOT's algorithm:**\n\n```\ngroup = [q_i, q_{i+1}, ..., q_{i+g-1}]\nproxy_q = mean(group)\n\n# ONE scan instead of g scans\nscores = proxy_q · K[:i]        # O(L)\nC = top-K(scores)               # candidate set, K = 2 × top_k\n\n# per-query refine (PIVOT-Refine)\nfor q in group:\n    refine_scores = q · K[C]    # O(K), not O(L)\n    final_indices[q] = top-k(refine_scores)\n```\n\nIndexer cost: O(L²) → **O(L²/g)**. With g=8 that's 8× fewer full scans.\n\n``` python\nimport torch\nimport torch.nn.functional as F\n\ndef pivot_refine(\n    queries: torch.Tensor,   # (seq_len, num_heads, head_dim)\n    keys: torch.Tensor,\n    top_k: int = 64,\n    group_size: int = 8,\n    candidate_ratio: float = 2.0,\n) -> torch.Tensor:\n    \"\"\"\n    Returns sparse-attention token indices.\n    Shape: (seq_len, num_heads, top_k)\n    \"\"\"\n    seq_len, num_heads, head_dim = queries.shape\n    K = int(top_k * candidate_ratio)\n    indices = torch.zeros(seq_len, num_heads, top_k,\n                          dtype=torch.long, device=queries.device)\n\n    for h in range(num_heads):\n        q_h = queries[:, h, :]\n        k_h = keys[:, h, :]\n\n        for g_start in range(0, seq_len, group_size):\n            g_end = min(g_start + group_size, seq_len)\n\n            # Step 1: one proxy scan per group (not per query)\n            proxy = q_h[g_start:g_end].mean(dim=0)\n\n            for pos in range(g_start, g_end):\n                if pos == 0:\n                    continue\n                proxy_scores = proxy @ k_h[:pos].T\n                actual_K = min(K, pos)\n                _, C = torch.topk(proxy_scores, actual_K)\n\n                # Step 2: O(K) refine — cheap!\n                refine_scores = q_h[pos] @ k_h[C].T\n                actual_k = min(top_k, actual_K)\n                _, best = torch.topk(refine_scores, actual_k)\n                indices[pos, h, :actual_k] = C[best]\n\n    return indices\n\ndef pivot_reuse(\n    queries: torch.Tensor,\n    keys: torch.Tensor,\n    top_k: int = 64,\n    group_size: int = 8,\n    candidate_ratio: float = 2.0,\n) -> torch.Tensor:\n    \"\"\"Max-speed variant — all queries in group share indices.\"\"\"\n    seq_len, num_heads, head_dim = queries.shape\n    K = int(top_k * candidate_ratio)\n    indices = torch.zeros(seq_len, num_heads, top_k,\n                          dtype=torch.long, device=queries.device)\n\n    for h in range(num_heads):\n        q_h = queries[:, h, :]\n        k_h = keys[:, h, :]\n\n        for g_start in range(0, seq_len, group_size):\n            g_end = min(g_start + group_size, seq_len)\n            ref_pos = max(g_start, 1)\n            proxy = q_h[g_start:g_end].mean(dim=0)\n            proxy_scores = proxy @ k_h[:ref_pos].T\n            actual_K = min(K, ref_pos)\n            _, C = torch.topk(proxy_scores, actual_K)\n\n            for pos in range(g_start, g_end):\n                actual_k = min(top_k, actual_K)\n                indices[pos, h, :actual_k] = C[:actual_k]\n\n    return indices\n\n# --- Quick test ---\nif __name__ == \"__main__\":\n    S, H, D = 1024, 8, 64\n    q = torch.randn(S, H, D)\n    k = torch.randn(S, H, D)\n    v = torch.randn(S, H, D)\n\n    idx = pivot_refine(q, k, top_k=64, group_size=8)\n    print(f\"Index shape: {idx.shape}\")  # (1024, 8, 64)\n\n    out = torch.zeros_like(q)\n    scale = D ** -0.5\n    for pos in range(1, S):\n        for h in range(H):\n            sel = idx[pos, h]\n            w = F.softmax(q[pos, h] @ k[sel, h].T * scale, dim=-1)\n            out[pos, h] = w @ v[sel, h]\n    print(f\"Output shape: {out.shape}\")  # (1024, 8, 64)\n```\n\n⚠️ Reference implementation only. Production speedups need Triton/CUDA kernels. Official code not yet released (July 2026).\n\nTested on **DeepSeek-V3.2** and **GLM-5.1** — LongBench + RULER:\n\n| Method | Indexer Speed | E2E Latency | Accuracy |\n|---|---|---|---|\n| Dense DSA (baseline) | 1× | 1× | ✅ full |\n| PIVOT-Refine | ~3× faster | −28% | ✅ ≈ baseline |\nPIVOT-Reuse |\n4× faster |\n−37.5% |\n⚠️ minor drop |\n\nPIVOT-Refine = dense-indexer accuracy + 3× speed. Zero retraining.\n\n**DSA-only**: Works out of the box on DeepSeek-V3.2 / GLM-5.1. Standard full-attention models need DSA fine-tuning first.\n\n**Group size is manual**: No adaptive strategy in the paper — tune g per model/sequence length.\n\n**Batch inference untested**: Single-sequence results only.\n\n**Prefill vs. decode split missing**: The 1.6× E2E number doesn't break down by phase.\n\n*What's your experience with sparse attention in production? Drop a comment.*", "url": "https://wpnews.pro/news/pivot-explained-from-paper-to-working-code-in-10-minutes", "canonical_source": "https://dev.to/cofldus/pivot-explained-from-paper-to-working-code-in-10-minutes-5a4j", "published_at": "2026-07-31 00:50:49+00:00", "updated_at": "2026-07-31 01:30:01.672857+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-research", "developer-tools"], "entities": ["PIVOT"], "alternates": {"html": "https://wpnews.pro/news/pivot-explained-from-paper-to-working-code-in-10-minutes", "markdown": "https://wpnews.pro/news/pivot-explained-from-paper-to-working-code-in-10-minutes.md", "text": "https://wpnews.pro/news/pivot-explained-from-paper-to-working-code-in-10-minutes.txt", "jsonld": "https://wpnews.pro/news/pivot-explained-from-paper-to-working-code-in-10-minutes.jsonld"}}