scripts/sparse_attention_repro.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["torch", "transformers==4.49.0", "datasets"]
# ///
"""
Reconstruction / proxy reproduction for Claims 1 and 2.

TileSparse's own code, model, benchmark, and baseline implementations are not
available to us (see Conclusion page: OpenReview PDF gated behind Cloudflare
Turnstile, no arXiv preprint, no GitHub repo). This script implements our own
best-effort reconstruction of the two mechanisms the claims compare, from the
paper's title and the challenge brief alone:

  - "dynamic K-only sparse" baseline: for each query, keep only the top-K
    individual keys by exact QK score (oracle top-K), the standard form of
    methods like H2O / Quest / SnapKV. This reduces *memory bytes read* per
    query but does not respect tile/block structure, so under a tiled kernel
    it does not proportionally reduce *FLOPs* once query blocks (D>1, as with
    MTP) are batched -- a whole K-tile must still be computed if any one of
    its keys is selected by any query in the block.

  - "tile-sparse" reconstruction: keys are grouped into fixed-size tiles; for
    a query block of size D (proxy for the MTP degree), block-level
    importance is aggregated across the D queries and the whole tile is
    kept/dropped together. This directly cuts Q-tile x K-tile matmul FLOPs
    (matching the paper's "arithmetic-intensity-aware" framing) at a *matched
    compute budget* with the K-only baseline, rather than matching memory
    reads.

Both are evaluated as an oracle causal-attention mask (using exact QK scores
to pick the kept set) applied inside GPT-2's attention, patched via the
module-level `eager_attention_forward` function in transformers==4.49.0's
`modeling_gpt2` (GPT-2 does not call `GPT2Attention._attn` in this
transformers version -- it dispatches through `ALL_ATTENTION_FUNCTIONS`), and
scored as validation perplexity on WikiText-2.

SCOPE / LIMITATIONS (this is a `toy`-labeled proxy, not the paper's own
setup): (1) GPT-2 uses plain MHA, not MLA -- we could not integrate a real
MLA model's fused/absorbed attention path in the time available, so this
tests the K-only-vs-tile-sparse selection mechanism in isolation, not on an
MLA backbone; (2) "D" (query block) here groups teacher-forced positions
within one forward pass as a proxy for MTP-batched queries, not a real
speculative MTP head; (3) oracle top-K/top-tile selection (using exact scores)
upper-bounds what any real predictor-based selector (including the paper's)
could achieve, so absolute numbers are not directly comparable to the paper's
reported 40% / 99% / 40.8% figures -- only the *relative* K-only-vs-tile-sparse
comparison, at matched compute budget, is informative.
"""
import argparse
import json
import math

import torch
import torch.nn.functional as F
from transformers import GPT2LMHeadModel, GPT2TokenizerFast
from datasets import load_dataset

SPARSE_CFG = {"mode": "dense", "keep_frac": 1.0, "block_q": 1, "block_k": 64}


def patched_eager_attention_forward(module, query, key, value, attention_mask, head_mask=None, **kwargs):
    """Drop-in replacement for transformers.models.gpt2.modeling_gpt2.eager_attention_forward
    (same signature/return contract) that applies the sparsity mode in SPARSE_CFG."""
    attn_weights = torch.matmul(query, key.transpose(-1, -2))
    attn_weights = attn_weights / (value.size(-1) ** 0.5)

    q_len, k_len = query.size(-2), key.size(-2)
    causal_mask = torch.tril(torch.ones(q_len, k_len, dtype=torch.bool, device=query.device),
                              diagonal=k_len - q_len)
    attn_weights = attn_weights.masked_fill(~causal_mask, float("-inf"))

    mode = SPARSE_CFG["mode"]
    if mode != "dense":
        keep_mask = build_keep_mask(attn_weights, causal_mask, SPARSE_CFG)
        attn_weights = attn_weights.masked_fill(~keep_mask, float("-inf"))

    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask[:, :, :, :k_len]

    attn_weights = F.softmax(attn_weights, dim=-1)
    attn_weights = torch.nan_to_num(attn_weights)  # rows that got fully masked
    attn_weights = attn_weights.type(value.dtype)
    if head_mask is not None:
        attn_weights = attn_weights * head_mask
    attn_output = torch.matmul(attn_weights, value)
    attn_output = attn_output.transpose(1, 2)
    return attn_output, attn_weights


# Real compute cost under a tiled kernel: a Q-tile x K-tile block must be
# computed if ANY (query, key) pair inside it is kept -- so a method whose
# per-query keep-set is chosen independently per query (k_only) can force a
# much larger union of touched tiles across a D-query block than its nominal
# per-query keep_frac suggests. This is the metric that actually measures
# "compute budget" (FLOPs), as opposed to keep_frac which only measures
# per-query memory reads. Accumulated globally per eval config, reset by
# reset_tile_stats() before each config's document loop.
TILE_STATS = {"touched_sum": 0.0, "qblocks_sum": 0}


def reset_tile_stats():
    TILE_STATS["touched_sum"] = 0.0
    TILE_STATS["qblocks_sum"] = 0


def record_tiles_touched(keep_mask, bq, bk):
    q_len, k_len = keep_mask.shape[-2:]
    n_qb = math.ceil(q_len / bq)
    n_kb = math.ceil(k_len / bk)
    pad_q = n_qb * bq - q_len
    pad_k = n_kb * bk - k_len
    padded = F.pad(keep_mask, (0, pad_k, 0, pad_q), value=False)
    blocked = padded.view(*padded.shape[:-2], n_qb, bq, n_kb, bk)
    touched = blocked.any(dim=-1).any(dim=-2)  # (..., n_qb, n_kb): any kept pair in this tile
    touched_count = touched.sum(dim=-1).float()  # (..., n_qb): touched k-tiles per q-block
    TILE_STATS["touched_sum"] += touched_count.sum().item()
    TILE_STATS["qblocks_sum"] += touched_count.numel()


def build_keep_mask(attn_weights, causal_mask, cfg):
    keep_frac = cfg["keep_frac"]
    q_len, k_len = attn_weights.shape[-2:]

    if cfg["mode"] == "k_only":
        n_keep = max(1, int(round(keep_frac * k_len)))
        scores = attn_weights.masked_fill(~causal_mask, float("-inf"))
        topk = scores.topk(min(n_keep, k_len), dim=-1).indices
        keep = torch.zeros_like(attn_weights, dtype=torch.bool)
        keep.scatter_(-1, topk, True)
        keep = keep & causal_mask
        # k_only picks each query's keep-set independently; block_q/block_k here
        # are only used to measure the resulting union-of-tiles compute cost,
        # they do not affect the (per-query) attention pattern itself.
        record_tiles_touched(keep, cfg["block_q"], cfg["block_k"])
        return keep

    if cfg["mode"] == "tile_sparse":
        bq, bk = cfg["block_q"], cfg["block_k"]
        scores = attn_weights.masked_fill(~causal_mask, float("-inf"))
        n_qb = math.ceil(q_len / bq)
        n_kb = math.ceil(k_len / bk)
        pad_q = n_qb * bq - q_len
        pad_k = n_kb * bk - k_len
        padded = F.pad(scores, (0, pad_k, 0, pad_q), value=float("-inf"))
        blocked = padded.view(*padded.shape[:-2], n_qb, bq, n_kb, bk)
        block_importance = blocked.amax(dim=(-3, -1))  # aggregate over D queries and tile keys
        n_keep_blocks = max(1, int(round(keep_frac * n_kb)))
        top_blocks = block_importance.topk(min(n_keep_blocks, n_kb), dim=-1).indices
        block_keep = torch.zeros_like(block_importance, dtype=torch.bool)
        block_keep.scatter_(-1, top_blocks, True)
        # block_keep: (..., n_qb, n_kb) -> (..., n_qb, 1, n_kb, 1) -> broadcast to (..., n_qb, bq, n_kb, bk)
        keep = block_keep.unsqueeze(-2).unsqueeze(-1).expand(*block_keep.shape[:-1], bq, n_kb, bk)
        keep = keep.reshape(*keep.shape[:-4], n_qb * bq, n_kb * bk)[..., :q_len, :k_len]
        keep = keep & causal_mask
        record_tiles_touched(keep, bq, bk)
        return keep

    raise ValueError(cfg["mode"])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="gpt2")
    ap.add_argument("--n-docs", type=int, default=20)
    ap.add_argument("--max-len", type=int, default=1024)
    ap.add_argument("--out", default="outputs/sparse_repro_results.json")
    ap.add_argument("--quick", action="store_true", help="smoke test: 2 tiny docs, short seqs")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"device={device} model={args.model}")

    tokenizer = GPT2TokenizerFast.from_pretrained(args.model)
    model = GPT2LMHeadModel.from_pretrained(args.model, attn_implementation="eager").to(device)
    model.eval()
    import transformers.models.gpt2.modeling_gpt2 as gpt2_modeling
    gpt2_modeling.eager_attention_forward = patched_eager_attention_forward

    ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
    docs = [t for t in ds["text"] if len(t.strip()) > 200][: (2 if args.quick else args.n_docs)]
    max_len = 64 if args.quick else args.max_len

    # block_q=8 on every config (including k_only) is the MTP-degree proxy: it
    # does not change k_only's per-query attention pattern, only the grouping
    # used to measure the real union-of-tiles compute cost (see
    # record_tiles_touched) on an equal footing with tile_sparse.
    MTP_DEGREE = 8
    configs = [
        {"name": "dense", "mode": "dense", "keep_frac": 1.0, "block_q": MTP_DEGREE, "block_k": 64},
        {"name": "k_only_tight", "mode": "k_only", "keep_frac": 0.15, "block_q": MTP_DEGREE, "block_k": 64},
        {"name": "tile_sparse_tight", "mode": "tile_sparse", "keep_frac": 0.15, "block_q": MTP_DEGREE, "block_k": 64},
        {"name": "k_only_40pct_cut", "mode": "k_only", "keep_frac": 0.592, "block_q": MTP_DEGREE, "block_k": 64},
        {"name": "tile_sparse_40pct_cut", "mode": "tile_sparse", "keep_frac": 0.592, "block_q": MTP_DEGREE, "block_k": 64},
        # Matched-REAL-COMPUTE variants: tile_sparse's keep_frac raised so its
        # avg_tiles_touched_per_qblock matches k_only's *actual* union-of-tiles
        # cost at the corresponding tier (from a first CPU pass, see Claim 1/2
        # pages) -- the fair comparison the claims are actually about.
        {"name": "tile_sparse_tight_matched_compute", "mode": "tile_sparse", "keep_frac": 0.65, "block_q": MTP_DEGREE, "block_k": 64},
        {"name": "tile_sparse_40pct_cut_matched_compute", "mode": "tile_sparse", "keep_frac": 0.73, "block_q": MTP_DEGREE, "block_k": 64},
    ]
    if args.quick:
        configs = [c for c in configs if c["name"] in ("dense", "k_only_tight", "tile_sparse_tight")]

    results = {}
    for cfg in configs:
        SPARSE_CFG.update(cfg)
        reset_tile_stats()
        total_loss, total_tokens = 0.0, 0
        for doc in docs:
            enc = tokenizer(doc, return_tensors="pt", truncation=True, max_length=max_len)
            input_ids = enc.input_ids.to(device)
            if input_ids.size(1) < 8:
                continue
            with torch.no_grad():
                out = model(input_ids, labels=input_ids)
            n_tok = input_ids.size(1) - 1
            total_loss += out.loss.item() * n_tok
            total_tokens += n_tok
        avg_loss = total_loss / max(total_tokens, 1)
        ppl = math.exp(avg_loss)
        avg_tiles_touched = (TILE_STATS["touched_sum"] / TILE_STATS["qblocks_sum"]
                              if cfg["mode"] != "dense" and TILE_STATS["qblocks_sum"] else None)
        results[cfg["name"]] = {"loss": avg_loss, "ppl": ppl, "n_tokens": total_tokens,
                                 "avg_tiles_touched_per_qblock": avg_tiles_touched, "config": cfg}
        tt_str = f"{avg_tiles_touched:.2f}" if avg_tiles_touched is not None else "n/a"
        print(f"{cfg['name']:<24} keep_frac={cfg['keep_frac']:<6} loss={avg_loss:.4f} ppl={ppl:.3f}  avg_tiles_touched/qblock={tt_str}")

    dense_loss = results["dense"]["loss"]

    def summarize(tight_or_cut):
        k_res, t_res = results[f"k_only_{tight_or_cut}"], results[f"tile_sparse_{tight_or_cut}"]
        k_excess, t_excess = k_res["loss"] - dense_loss, t_res["loss"] - dense_loss
        rel_improve = (k_excess - t_excess) / k_excess if k_excess > 1e-9 else float("nan")
        k_tiles, t_tiles = k_res["avg_tiles_touched_per_qblock"], t_res["avg_tiles_touched_per_qblock"]
        return {
            "k_only_excess_loss": k_excess, "tile_sparse_excess_loss": t_excess,
            "tile_sparse_relative_improvement_over_k_only_AT_MATCHED_KEEP_FRAC": rel_improve,
            "k_only_avg_tiles_touched_per_qblock": k_tiles,
            "tile_sparse_avg_tiles_touched_per_qblock": t_tiles,
            "k_only_real_compute_vs_tile_sparse_x": (k_tiles / t_tiles) if (k_tiles and t_tiles) else None,
            "note": ("tile_sparse's keep_frac is a real tile-compute budget (avg_tiles_touched "
                     "matches keep_frac*n_kb by construction). k_only's keep_frac is a per-query "
                     "memory budget only -- k_only_avg_tiles_touched_per_qblock shows the REAL tile "
                     "count a tiled kernel must compute once independent per-query top-K choices "
                     "are unioned across an MTP query block; if that is larger than tile_sparse's, "
                     "the two configs are not at matched real compute despite equal keep_frac, and "
                     "the accuracy comparison above understates tile_sparse's advantage per unit "
                     "of actual FLOPs."),
        }

    def summarize_matched_compute(tight_or_cut):
        k_res = results[f"k_only_{tight_or_cut}"]
        t_res = results[f"tile_sparse_{tight_or_cut}_matched_compute"]
        k_excess, t_excess = k_res["loss"] - dense_loss, t_res["loss"] - dense_loss
        rel_improve = (k_excess - t_excess) / k_excess if abs(k_excess) > 1e-9 else float("nan")
        return {
            "k_only_avg_tiles_touched_per_qblock": k_res["avg_tiles_touched_per_qblock"],
            "tile_sparse_avg_tiles_touched_per_qblock": t_res["avg_tiles_touched_per_qblock"],
            "k_only_excess_loss": k_excess, "tile_sparse_excess_loss": t_excess,
            "tile_sparse_relative_improvement_over_k_only_AT_MATCHED_REAL_COMPUTE": rel_improve,
            "tile_sparse_accuracy_retention_vs_dense": math.exp(-t_excess),
        }

    if not args.quick:
        results["claim1_tight_budget_comparison"] = summarize("tight")
        results["claim1_matched_compute_comparison"] = summarize_matched_compute("tight")
        results["claim2_40pct_cut_comparison"] = {
            **summarize("40pct_cut"),
            "tile_sparse_accuracy_retention_vs_dense": math.exp(-(results["tile_sparse_40pct_cut"]["loss"] - dense_loss)),
        }
        results["claim2_matched_compute_comparison"] = summarize_matched_compute("40pct_cut")
        print("\nClaim 1 (tight budget, keep_frac=0.15, MATCHED keep_frac):", json.dumps(results["claim1_tight_budget_comparison"], indent=2))
        print("\nClaim 1 (tight budget, MATCHED REAL COMPUTE):", json.dumps(results["claim1_matched_compute_comparison"], indent=2))
        print("\nClaim 2 (40.8% compute cut, keep_frac=0.592, MATCHED keep_frac):", json.dumps(results["claim2_40pct_cut_comparison"], indent=2))
        print("\nClaim 2 (40.8% compute cut, MATCHED REAL COMPUTE):", json.dumps(results["claim2_matched_compute_comparison"], indent=2))

    import os
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved {args.out}")


if __name__ == "__main__":
    main()