# PIVOT Explained — From Paper to Working Code in 10 Minutes

> Source: <https://dev.to/cofldus/pivot-explained-from-paper-to-working-code-in-10-minutes-5a4j>
> Published: 2026-07-31 00:50:49+00:00

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.*