scripts/sparse_attention_repro_mla.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["torch", "transformers>=4.56", "datasets", "accelerate"]
# ///
"""
Claims 1 & 2 reconstruction on a REAL MLA model: deepseek-ai/DeepSeek-V2-Lite
(16B total / 2.4B active MoE, native Multi-head Latent Attention). This is the
less-toy upgrade of scripts/sparse_attention_repro.py, which used GPT-2 (plain
MHA) as a substitute architecture. Same K-only-vs-tile-sparse reconstruction
and matched-real-compute methodology -- see that script's docstring for the
full mechanism explanation and the honesty caveats that still apply (no
TileSparse code/paper, oracle exact-score selection upper-bounds any real
predictor).

DeepseekV2Attention (transformers>=4.56, native support, no trust_remote_code
needed -- confirmed via scripts/_discover_deepseek_attn.py) materializes full
per-head query/key/value tensors (not the fused "absorbed" inference form)
and dispatches through the same module-level `eager_attention_forward`
pattern as Llama/GPT2, so the same patch technique applies. Its signature
differs from GPT2's: attention_mask arrives as an already-built additive
causal mask (0.0 = keep, large-negative = masked) rather than being built
inside the attention module, and scaling is passed in explicitly rather than
computed from head_dim.
"""
import argparse
import json
import math

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

MODEL_ID = "deepseek-ai/DeepSeek-V2-Lite"

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

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, pad_k = n_qb * bq - q_len, 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)
    touched_count = touched.sum(dim=-1).float()
    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
        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, n_kb = math.ceil(q_len / bq), math.ceil(k_len / bk)
        pad_q, pad_k = n_qb * bq - q_len, 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))
        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)
        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 patched_eager_attention_forward(module, query, key, value, attention_mask, scaling, dropout=0.0, **kwargs):
    """Drop-in replacement for DeepseekV2's module-level eager_attention_forward
    (transformers.models.deepseek_v2.modeling_deepseek_v2), same signature and
    return contract, applying the sparsity mode in SPARSE_CFG."""
    from transformers.models.deepseek_v2.modeling_deepseek_v2 import repeat_kv

    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    q_len, k_len = query.size(-2), key_states.size(-2)

    if attention_mask is not None:
        causal_mask = attention_mask[..., :q_len, :k_len] == 0
        attn_weights = attn_weights + attention_mask[..., :q_len, :k_len]
    else:
        causal_mask = torch.tril(torch.ones(q_len, k_len, dtype=torch.bool, device=query.device),
                                  diagonal=k_len - q_len)

    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"))

    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = torch.nan_to_num(attn_weights)
    attn_weights = F.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()
    return attn_output, attn_weights


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n-docs", type=int, default=8)
    ap.add_argument("--max-len", type=int, default=512)
    ap.add_argument("--out", default="outputs/sparse_repro_mla_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"
    dtype = torch.bfloat16 if device == "cuda" else torch.float32
    print(f"device={device} dtype={dtype} model={MODEL_ID}")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, dtype=dtype, attn_implementation="eager", device_map=device,
    )
    model.eval()
    print("attn class:", type(model.model.layers[0].self_attn))

    import transformers.models.deepseek_v2.modeling_deepseek_v2 as ds_modeling
    ds_modeling.eager_attention_forward = patched_eager_attention_forward

    ds = load_dataset("Salesforce/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

    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},
    ]
    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(tier):
        k_res, t_res = results[f"k_only_{tier}"], results[f"tile_sparse_{tier}"]
        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_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_res["avg_tiles_touched_per_qblock"],
                "tile_sparse_avg_tiles_touched_per_qblock": t_res["avg_tiles_touched_per_qblock"],
                "tile_sparse_accuracy_retention_vs_dense": math.exp(-t_excess)}

    if not args.quick:
        results["claim1_tight_budget_comparison"] = summarize("tight")
        results["claim2_40pct_cut_comparison"] = summarize("40pct_cut")
        print("\nClaim 1 (tight budget):", json.dumps(results["claim1_tight_budget_comparison"], indent=2))
        print("\nClaim 2 (40.8% compute cut):", json.dumps(results["claim2_40pct_cut_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()