# sabench: does your long-context speedup actually work?

An evaluation harness for sparse-attention / KV-selection methods on frozen open-weight LLMs. It
takes any block-selection policy, runs it against a **real frozen Qwen2.5-7B**, and tells you in
about 20 minutes whether it survives contact with real perplexity.

## Why not just measure attention recall

Because it can invert the answer. A common shortcut scores selection policies by "what fraction of
attention mass falls inside the blocks you selected." That proxy is fast to compute and easy to
be fooled by: two methods can trade places on the proxy and swap right back on actual next-token
loss. A metric that can flip which of two methods is better is not weak evidence. It is no
evidence.

**A proxy that flips which of two methods is better is not weak evidence. It is no evidence.**
This harness only reports the real thing: cross-entropy on held-out tokens, from a model whose
attention is genuinely restricted to what the policy selects.

## Usage

```bash
python3 assemble_sabench.py            # -> sabench.ipynb
# push to Kaggle with a GPU, internet ON, machineShape NvidiaTeslaT4
```

Levers, all environment variables, no code edits needed for a sweep:

```
SAB_CTX     context length            default 4096
SAB_TOPM    blocks selected per query default 16
SAB_BLOCK   KV block size             default 32
SAB_DOCS    documents scored          default 3
```

## Adding your idea

Write one function and register it. Everything else is automatic.

```python
def sel_mine(li, q, k, T, dev):
    # q  (T, N_HEAD * D_HEAD)  post-RoPE queries
    # k  (T, N_KV   * D_HEAD)  post-RoPE keys
    scores = ...                      # (query_block, key_block), higher = more worth reading
    return build_mask(scores, T // BLOCK, T, dev)

SEL["mine"] = sel_mine
ARMS = ["dense", "oracle", "meanpool", "mine", "random", "local"]
```

`build_mask` applies the top-m cut, the local window, the attention sink, **and token-level
causality**. Do not hand-roll that last part. See defect 4 below.

## The arms you are automatically measured against

| arm | what it is |
|---|---|
| `dense` | full attention. the target |
| `oracle` | picks the genuinely best blocks using real attention. **the ceiling** |
| `meanpool` | average each block's keys, dot with the query. Quest and MoBA style. **the incumbent** |
| `random` | random blocks. **the null** |
| `local` | recent window plus sink only. **the floor** |

Reading the leaderboard:

- **at or below `random`** means you are not selecting, you are guessing
- **near `oracle`** means you are near optimal and the remaining loss is the block budget, not
  your selector
- **below `meanpool`** means a two-line baseline with no training beats you, which is the bar that
  actually matters

## The guards, and why each one exists

Every one of these was added after it caught a real bug during development.

| guard | what it caught |
|---|---|
| stock model health check | Qwen2.5-7B in fp16 emits NaN before anything is attached. A patch was blamed for a broken reference. |
| `patched == original` at 0.0 | verifies the wrapper is the identical computation |
| re-verify after backend switch | eager to sdpa is a changed computation until proven otherwise |
| **all-selected mask == dense** | the empty-mask check passed at 0.0 while every arm was broken, because an empty mask takes a path the arms never use |
| rotation check | rebuilds attention from captured features and requires it to match the model's own, so no selector can be trained or scored on the wrong geometry |
| **RUN INVALID null** | if `local` is within 0.05 nats of `dense`, the document needed no long-range retrieval and nothing was tested |
| oracle-vs-random spread | if the ceiling and the null are within 15 points, block choice barely matters at this density and the test has no power |
| cell-order guard | static use-before-definition across notebook cells |
| banned patterns | quantisation (distorts memory traffic), `model.train()`, and the NaN-blind comparison form |

## Defects this harness was built out of, keep them in mind

1. **NaN-blind guard.** `if err >= tol: raise` never fires on NaN, because every comparison against
   NaN is False. A totally failed verification passed silently and printed a table of NaN under a
   confident verdict. Always test the positive condition: `if not (err < tol)`.
2. **Verifying against an unchecked reference.** A baseline model was itself broken, not the code
   under test. Health-check the reference FIRST.
3. **Reimplementing attention.** A from-scratch q/k/v + RoPE + GQA + reshape came out 0.515
   relative. Injecting a mask into the model's own attention cannot have those bugs and verifies at
   exactly 0. Prefer intercepting over reimplementing.
4. **A block-granularity mask is not causal.** Supplying any mask makes SDPA set `is_causal=False`
   and trust your mask. A block-level mask lets a query read its own block's future, up to
   `BLOCK-1` tokens ahead. This produced a run where every sparse arm BEAT full attention, which is
   impossible, and that impossibility is what exposed it.
5. **Too-small corpus.** A too-short corpus of site-packages text gave a dense perplexity and a
   local floor that nearly tied it. The null caught it.

## Cross-model reference results

Same harness, same corpus recipe, ~13% attention density, three base models:

| model | dense NLL | oracle | meanpool | random |
|---|---|---|---|---|
| Qwen2.5-7B (bf16) | 0.475 | 100% | 93% | 44% |
| Mistral-7B-v0.1 (bf16) | 0.521 | 97% | 92% | 55% |
| Qwen2.5-1.5B (bf16) | 0.653 | 100% | 91% | 45% |

Two findings are stable across families and sizes: near-perfect block selection matches full
attention at this density (the quality problem is selection, not the budget), and the
training-free mean-pool baseline already recovers >90% of the recoverable signal.

**Use base models, not chat-tuned ones.** A chat model evaluated on raw text can score
near-chance under dense attention, at which point every arm below it is noise. If your dense
NLL is implausibly high for the corpus, the substrate is the problem, not the selectors.

## What it cannot tell you

- **No end-to-end speedup.** Masking does not skip reads. Selection time is measured; memory
  traffic is analytic. Only a fused block-sparse kernel realises the traffic number, and that
  kernel is also what caps context here, since an explicit mask forces the dense fallback.
- **Nothing above ~4k context** for the same reason.
- Selection is per query BLOCK, and one index is shared across all heads.

## Reference numbers, Qwen2.5-7B at ctx 4096, ~12% density

| arm | nll | recovers |
|---|---|---|
| dense | 0.4754 | target |
| oracle | ~0.47 | ~100% |
| meanpool | ~0.49 | ~93% |
| random | ~0.55 | ~44% |
| local | dense + 0.618 nats | floor |

These are representative numbers from prior runs at this configuration; run the notebook yourself
to reproduce them on your own document and hardware.