PIVOT Explained — From Paper to Working Code in 10 Minutes

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)

scores = proxy_q · K[:i]        # O(L)
C = top-K(scores)               # candidate set, K = 2 × top_k

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.

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)

            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)

                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

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:

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.