{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# sabench: does your long-context speedup actually work?\n", "\n", "Proxy metrics for sparse attention are easy to get and easy to be fooled by. A selector can look\n", "great by \"fraction of attention mass retained\" and lose on the number that actually matters.\n", "\n", "This notebook evaluates block-selection policies for sparse attention against the real thing:\n", "\n", "* **quality** = actual next-token negative log likelihood on a long document, with the model's\n", " attention genuinely restricted to what the policy selects\n", "* **speed** = `torch.cuda.Event` timings of the selection step itself, kept separate from the\n", " attention operation, which this harness does not claim to speed up\n", "\n", "**Stated before the run, because it is the likely outcome and it should not look like a\n", "disappointment later:** a hand-written gather-based sparse attention in PyTorch will probably be\n", "SLOWER on the clock than PyTorch's fused dense kernel at these context lengths. Dense attention is\n", "one of the most optimised routines on the machine and this is not a fused kernel. The theoretical\n", "win needs a real kernel and a context long enough that memory traffic dominates. So the timing\n", "section reports measured milliseconds AND the hardware-independent memory traffic, and treats the\n", "first as a fact about this implementation rather than about any one policy.\n", "\n", "The quality number has no such excuse attached to it. If perplexity blows up, the policy is wrong." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 1 - SUBSTRATE. Biggest model that fits, frozen.\n", "# ============================================================\n", "import os\n", "os.environ.setdefault(\"PYTORCH_CUDA_ALLOC_CONF\", \"expandable_segments:True\")\n", "import sys, math, time, json, random, gc\n", "import torch, torch.nn as nn, torch.nn.functional as F\n", "\n", "SEED = 7\n", "random.seed(SEED); torch.manual_seed(SEED)\n", "assert torch.cuda.is_available(), \"this notebook is a timing measurement, it needs a GPU\"\n", "NGPU = torch.cuda.device_count()\n", "# Kaggle accepted machineShape \"NvidiaTeslaT4x2\" and then allocated a single P100. sm_60 has\n", "# no kernels in this torch build, so every CUDA op dies with \"no kernel image is available\"\n", "# a hundred lines into the run. Fail here instead, with a message that says what happened.\n", "_cap = torch.cuda.get_device_capability(0)\n", "if _cap[0] < 7:\n", " raise SystemExit(\n", " f\"got {torch.cuda.get_device_name(0)} (sm_{_cap[0]}{_cap[1]}). This torch build has no \"\n", " f\"kernels below sm_70 and fp16 has no tensor cores on it either, so a timing \"\n", " f\"measurement here would be meaningless even if it ran. Re-push with \"\n", " f\"machineShape NvidiaTeslaT4.\")\n", "TOTAL_VRAM = sum(torch.cuda.get_device_properties(i).total_memory\n", " for i in range(NGPU)) / 1e9\n", "print(f\"[env] torch {torch.__version__} | {NGPU} x \"\n", " f\"{torch.cuda.get_device_name(0)} | {TOTAL_VRAM:.1f} GB total\")\n", "\n", "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "# fp16 weight cost, plus headroom for a long KV cache and activations. Quantisation is\n", "# deliberately NOT used: it would distort exactly the memory-traffic measurement this\n", "# notebook exists to take.\n", "# Qwen2.5-7B in pure fp16 emits NaN logits BEFORE anything of ours is attached: the family\n", "# has large activation outliers and is meant to run in bfloat16. An earlier version of this\n", "# harness verified a patch against that broken reference and reported \"MISMATCH\" as though\n", "# the patch were at fault. So the model itself is health-checked FIRST, and the ladder walks\n", "# dtypes as well as sizes until the stock model produces finite logits.\n", "#\n", "# Quantisation stays banned: it would distort the memory-traffic measurement this exists for.\n", "LADDER = [(\"Qwen/Qwen2.5-7B\", torch.bfloat16, 15.2),\n", " (\"Qwen/Qwen2.5-7B\", torch.float16, 15.2),\n", " (\"Qwen/Qwen2.5-3B\", torch.bfloat16, 6.2),\n", " (\"Qwen/Qwen2.5-1.5B\", torch.bfloat16, 3.1),\n", " (\"Qwen/Qwen2.5-1.5B\", torch.float32, 6.2)]\n", "\n", "def stock_is_healthy(m):\n", " \"\"\"Does the UNPATCHED model produce finite logits? Everything downstream compares\n", " against this, so if it is NaN there is nothing to compare against.\"\"\"\n", " with torch.no_grad():\n", " x = torch.randint(0, 1000, (1, 256), device=next(m.parameters()).device)\n", " lg = m(x).logits.float()\n", " return bool(torch.isfinite(lg).all()), float(lg.abs().max())\n", "\n", "MODEL_ID = model = tok = DTYPE = None\n", "for mid, dt, need in LADDER:\n", " if need + 4.0 > TOTAL_VRAM:\n", " print(f\"[substrate] skip {mid} {dt}, needs ~{need:.1f} GB + headroom\")\n", " continue\n", " try:\n", " print(f\"[substrate] trying {mid} in {str(dt).split('.')[-1]} ...\", flush=True)\n", " tok = AutoTokenizer.from_pretrained(mid)\n", " m = AutoModelForCausalLM.from_pretrained(\n", " mid, dtype=dt, attn_implementation=\"eager\",\n", " device_map=\"auto\" if NGPU > 1 else None)\n", " if NGPU == 1:\n", " m = m.cuda()\n", " m.eval()\n", " ok, mag = stock_is_healthy(m)\n", " print(f\" stock health check: finite={ok}, max|logit|={mag:.1f}\")\n", " if not ok:\n", " print(f\" REJECTED, the stock model itself emits NaN in this dtype\")\n", " del m; gc.collect(); torch.cuda.empty_cache(); continue\n", " model, MODEL_ID, DTYPE = m, mid, dt\n", " break\n", " except Exception as e:\n", " print(f\" failed: {type(e).__name__}: {str(e)[:200]}\")\n", " gc.collect(); torch.cuda.empty_cache()\n", "if model is None:\n", " raise SystemExit(\"no model produced finite logits. Nothing can be measured.\")\n", "print(f\"[substrate] using {MODEL_ID} in {str(DTYPE).split('.')[-1]}\")\n", "\n", "model.eval()\n", "for p in model.parameters():\n", " p.requires_grad_(False)\n", "CFG = model.config\n", "N_LAYER, N_HEAD = CFG.num_hidden_layers, CFG.num_attention_heads\n", "N_KV = getattr(CFG, \"num_key_value_heads\", N_HEAD)\n", "D_HEAD = CFG.hidden_size // N_HEAD\n", "GROUP = N_HEAD // N_KV\n", "DEV = next(model.parameters()).device\n", "KV_BYTES_PER_TOK = 2 * N_LAYER * N_KV * D_HEAD * 2\n", "print(f\"\\n[substrate] {MODEL_ID} FROZEN\")\n", "print(f\" {sum(p.numel() for p in model.parameters())/1e9:.2f}B params | {N_LAYER} layers | \"\n", " f\"{N_HEAD} q heads | {N_KV} kv heads | head dim {D_HEAD}\")\n", "print(f\" KV cache costs {KV_BYTES_PER_TOK/1024:.1f} KiB per token of context\")" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 2 - CONSTANTS\n", "# ============================================================\n", "# Context is bounded by the MASK, not by the KV cache. Supplying an explicit float mask\n", "# makes SDPA fall back to its math backend, which expands the (1,1,T,T) mask across all\n", "# heads: 28 x 8192^2 x 2 bytes = 7.0 GiB, which is exactly what an earlier version of this\n", "# harness failed to allocate. At 4,096 the same term is 1.87 GiB and fits. Going longer needs\n", "# a block-sparse kernel that consumes the block mask directly instead of an expanded dense\n", "# one, which is the same missing kernel the timing section already refuses to claim around.\n", "CTX = int(os.environ.get(\"SAB_CTX\", \"4096\")) # document length the model must condition on\n", "EVAL_TAIL = 1024 # NLL is scored on the LAST this-many tokens, where long range matters\n", "BLOCK = int(os.environ.get(\"SAB_BLOCK\", \"32\")) # KV block, the unit selected\n", "TOP_M = int(os.environ.get(\"SAB_TOPM\", \"16\")) # blocks attended per query\n", "LOCAL_W = 2 # blocks nearest the query ALWAYS kept. 4 was 128 tokens,\n", " # enough local context that the floor nearly tied dense, as every sparse method does\n", "SINK = 2 # first blocks always kept. StreamingLLM: dropping it wrecks softmax\n", "CHUNK = 256 # query chunk size, keeps the score matrix bounded\n", "SKIP_LAYERS = 2 # first layers stay dense. Standard practice, they are near-local anyway\n", "\n", "print(f\"[cfg] {MODEL_ID} | ctx {CTX:,} | eval tail {EVAL_TAIL:,} | block {BLOCK} | \"\n", " f\"top-m {TOP_M} (+{LOCAL_W} local +{SINK} sink)\")\n", "_dens = (TOP_M + LOCAL_W + SINK) * BLOCK\n", "print(f\"[cfg] attends {_dens:,} of {CTX:,} tokens = {_dens/CTX*100:.1f}% density \"\n", " f\"(top-m {TOP_M*BLOCK} retrieved + {(LOCAL_W+SINK)*BLOCK} local and sink)\")\n", "print(f\"[cfg] KV cache at this context: {CTX*KV_BYTES_PER_TOK/1e9:.2f} GB\")" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 3 - A LONG REAL DOCUMENT\n", "# ============================================================\n", "import glob\n", "def gather(min_chars=12_000_000):\n", " out, n = [], 0\n", " for root in (\"/usr/local/lib/python3.11\", \"/usr/lib/python3\", \"/opt/conda/lib/python3.11\",\n", " \"/usr/local/lib/python3.12\", \"/kaggle/input\"):\n", " for pat in (\"**/*.py\", \"**/*.md\"):\n", " for fp in glob.iglob(os.path.join(root, pat), recursive=True):\n", " try:\n", " if os.path.getsize(fp) < 4000:\n", " continue\n", " out.append(open(fp, errors=\"ignore\").read()); n += len(out[-1])\n", " except Exception:\n", " pass\n", " if n >= min_chars:\n", " return out\n", " return out\n", "\n", "CH = gather()\n", "random.Random(SEED).shuffle(CH)\n", "ids = []\n", "for c in CH:\n", " ids.extend(tok(c, add_special_tokens=False).input_ids)\n", " if len(ids) > CTX * 40:\n", " break\n", "ALL = torch.tensor(ids, dtype=torch.long)\n", "print(f\"[corpus] {len(CH)} files -> {len(ALL):,} tokens\")\n", "if len(ALL) < CTX * 2:\n", " raise SystemExit(f\"corpus too short for ctx {CTX}\")\n", "\n", "def doc(seed):\n", " g = random.Random(seed)\n", " s = g.randrange(0, len(ALL) - CTX - 1)\n", " return ALL[s:s + CTX]" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 4 - THE PATCH. Inject a mask, do not reimplement attention.\n", "# ============================================================\n", "# An earlier version of this harness reimplemented q/k/v projection, RoPE, GQA repeat and the\n", "# output reshape from scratch, and came out 0.515 relative against the real model. Rather than\n", "# hunt that bug, the approach is replaced with one that cannot have it: HF's own attention runs\n", "# untouched, and the ONLY thing injected is an additive mask that sends unselected blocks to\n", "# -inf. Eager attention already adds `attention_mask` to the scores, so this rides the\n", "# existing code path.\n", "#\n", "# The consequence worth stating: with an all-zero mask the patched model is not merely close\n", "# to the original, it is running the identical computation. The verification below should\n", "# therefore be exactly 0, and anything else means the injection itself is broken.\n", "\n", "def rotate_half(x):\n", " a, b = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]\n", " return torch.cat((-b, a), dim=-1)\n", "\n", "\n", "def apply_rope(x, cos, sin):\n", " return x * cos + rotate_half(x) * sin\n", "\n", "\n", "MASKS = {} # layer -> additive (1, 1, T, T) or absent for dense\n", "CAPTURE = {} # layer -> (q, k) pre-RoPE, for the selectors\n", "CAPTURE_ON = False\n", "ROPE_FEATURES = True # selection features must be post-rotation: attention is computed on\n", " # rotated vectors, so selecting on unrotated ones hides position\n", " # from every selector and unfairly hurts the learned ones most\n", "NEG = torch.finfo(torch.float32).min / 4 # survives a cast to bf16\n", "\n", "def wrap_layer(li, sa):\n", " orig = sa.forward\n", " def fwd(hidden_states, attention_mask=None, **kw):\n", " if CAPTURE_ON:\n", " T = hidden_states.shape[1]\n", " _q = sa.q_proj(hidden_states)[0].detach().view(T, N_HEAD, D_HEAD)\n", " _k = sa.k_proj(hidden_states)[0].detach().view(T, N_KV, D_HEAD)\n", " # position_embeddings arrives inside **kw, not as a named parameter\n", " _pe = kw.get(\"position_embeddings\")\n", " if ROPE_FEATURES and _pe is not None:\n", " _cos, _sin = _pe\n", " _cos, _sin = _cos[0].unsqueeze(1), _sin[0].unsqueeze(1) # (T, 1, d)\n", " _q, _k = apply_rope(_q, _cos, _sin), apply_rope(_k, _cos, _sin)\n", " CAPTURE[li] = (_q.reshape(T, -1), _k.reshape(T, -1))\n", " m = MASKS.get(li)\n", " if m is not None:\n", " am = m.to(hidden_states.device, dtype=hidden_states.dtype)\n", " attention_mask = am if attention_mask is None else attention_mask + am\n", " return orig(hidden_states, attention_mask=attention_mask, **kw)\n", " return fwd, orig\n", "\n", "LAYERS_MOD = model.model.layers\n", "ORIG = {}\n", "for i, lyr in enumerate(LAYERS_MOD):\n", " f, o = wrap_layer(i, lyr.self_attn)\n", " ORIG[i] = o\n", " lyr.self_attn.forward = f\n", "\n", "_vd = doc(1)[:512].unsqueeze(0).to(DEV)\n", "MASKS.clear()\n", "with torch.no_grad():\n", " _lp = model(_vd).logits.float()\n", " for i, lyr in enumerate(LAYERS_MOD):\n", " lyr.self_attn.forward = ORIG[i]\n", " _lo = model(_vd).logits.float()\n", " for i, lyr in enumerate(LAYERS_MOD):\n", " lyr.self_attn.forward = wrap_layer(i, lyr.self_attn)[0]\n", "PATCH_ERR = float((_lp - _lo).abs().max())\n", "_ref = float(_lo.abs().max())\n", "_rel = PATCH_ERR / _ref\n", "_ok = bool(_rel < 1e-6) # bool() of a NaN comparison is False, which is what we want\n", "print(f\"[verify] max|patched - original| = {PATCH_ERR:.4e} on logits of magnitude {_ref:.1f}\")\n", "print(f\"[verify] relative {_rel:.2e} -> {'IDENTICAL computation' if _ok else 'MISMATCH or NaN'}\")\n", "if not _ok:\n", " raise SystemExit(\"the replacement attention does not reproduce the model. Everything \"\n", " \"downstream would be measuring a different network.\")\n", "del _lp, _lo, _vd\n", "gc.collect(); torch.cuda.empty_cache()" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 4b - VERIFY THE ROTATION. Do not assume it.\n", "# ============================================================\n", "# The captured features are now claimed to be post-RoPE. If that rotation is wrong every\n", "# selector that uses these features trains or scores on nonsense and the run says nothing.\n", "# So reconstruct attention from the captured q and k alone and require it to reproduce the\n", "# model's own attention.\n", "\n", "_rd = doc(11)[:384].unsqueeze(0).to(DEV)\n", "MASKS.clear(); CAPTURE.clear(); CAPTURE_ON = True\n", "with torch.no_grad():\n", " _ro = model(_rd, output_attentions=True)\n", "CAPTURE_ON = False\n", "_RL = N_LAYER // 2\n", "_rq, _rk = CAPTURE[_RL]\n", "_T4 = _rd.shape[1]\n", "with torch.no_grad():\n", " _qh = _rq.view(_T4, N_HEAD, D_HEAD).permute(1, 0, 2).float()\n", " _kh = _rk.view(_T4, N_KV, D_HEAD).permute(1, 0, 2).float()\n", " _kh = _kh.repeat_interleave(GROUP, dim=0)\n", " _sc = torch.matmul(_qh, _kh.transpose(1, 2)) / math.sqrt(D_HEAD)\n", " _ar = torch.arange(_T4, device=_sc.device)\n", " _sc = _sc.masked_fill((_ar.unsqueeze(0) > _ar.unsqueeze(1)).unsqueeze(0), float(\"-inf\"))\n", " # the 7B is sharded, so layer _RL's captures sit on its own device while the collected\n", " # attentions come back on cuda:0. Compare on cpu.\n", " _mine = F.softmax(_sc, dim=-1).mean(0).cpu()\n", " _theirs = _ro.attentions[_RL][0].float().mean(0).cpu()\n", "_ROPE_ERR = float((_mine - _theirs).abs().max())\n", "print(f\"[verify] attention rebuilt from captured features vs the model's own: \"\n", " f\"max abs diff {_ROPE_ERR:.3e} -> \"\n", " f\"{'rotation CONFIRMED' if bool(_ROPE_ERR < 5e-3) else 'ROTATION WRONG'}\")\n", "if not bool(_ROPE_ERR < 5e-3):\n", " raise SystemExit(\"the captured features do not reproduce the model's attention, so any \"\n", " \"selector built on them would be scoring the wrong geometry.\")\n", "del _rd, _ro, _rq, _rk, _mine, _theirs\n", "CAPTURE.clear(); gc.collect(); torch.cuda.empty_cache()" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 5 - ATTENTION BACKEND SWITCH\n", "# ============================================================\n", "# Selection uses PRE-RoPE-model, POST-RoPE-hook q and k throughout: what a selector needs is a\n", "# consistent representation, and taking the projections straight off the hooks removes an\n", "# entire class of rotation bugs of exactly the kind that just cost this harness two runs.\n", "\n", "def qgroup(q, T):\n", " \"\"\"Average q heads within each GQA group, matching the head-averaged attention labels.\"\"\"\n", " return q.view(T, N_KV, GROUP, D_HEAD).mean(2).reshape(T, -1)\n", "\n", "def use_sdpa():\n", " \"\"\"Switch to SDPA for the scoring passes.\n", "\n", " Eager attention forms a (heads, T, T) fp32 score tensor: 28 x 4096^2 x 4 = 1.75 GiB per\n", " layer, which OOMs a 7B that already fills both cards. SDPA consumes the identical\n", " additive float mask and never materialises it.\"\"\"\n", " try:\n", " if hasattr(model, \"set_attn_implementation\"):\n", " model.set_attn_implementation(\"sdpa\")\n", " else:\n", " model.config._attn_implementation = \"sdpa\"\n", " for m in model.modules():\n", " if hasattr(m, \"config\"):\n", " m.config._attn_implementation = \"sdpa\"\n", " return True\n", " except Exception as e:\n", " print(f\"[attn] could not switch to sdpa: {e}\")\n", " return False\n", "\n", "SDPA = use_sdpa()\n", "print(f\"[attn] scoring passes use {'sdpa' if SDPA else 'eager'}\")\n", "# The backend changed, so the dense arm is no longer the thing that was verified in cell 4.\n", "# Re-verify: an empty mask under SDPA must still reproduce the model.\n", "_vd2 = doc(3)[:512].unsqueeze(0).to(DEV)\n", "MASKS.clear()\n", "with torch.no_grad():\n", " _a = model(_vd2).logits.float()\n", " for i, lyr in enumerate(LAYERS_MOD):\n", " lyr.self_attn.forward = ORIG[i]\n", " _b = model(_vd2).logits.float()\n", " for i, lyr in enumerate(LAYERS_MOD):\n", " lyr.self_attn.forward = wrap_layer(i, lyr.self_attn)[0]\n", "_rel2 = float((_a - _b).abs().max()) / float(_b.abs().max())\n", "print(f\"[verify] after backend switch, relative {_rel2:.2e} -> \"\n", " f\"{'still identical' if bool(_rel2 < 1e-5) else 'MISMATCH'}\")\n", "if not bool(_rel2 < 1e-5):\n", " raise SystemExit(\"the mask injection does not survive the sdpa switch\")\n", "del _vd2, _a, _b\n", "gc.collect(); torch.cuda.empty_cache()" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 6 - SELECTORS -> a full additive mask\n", "# ============================================================\n", "# Every arm keeps the same free blocks (local window + attention sink) and the same TOP_M\n", "# budget, so the arms differ ONLY in which distant blocks they choose.\n", "\n", "def build_mask(scores, nblk, T, dev):\n", " \"\"\"scores (nblk_q, nblk) -> additive (1, 1, T, T). Rows are query BLOCKS, which is the\n", " granularity every practical sparse-attention scheme selects at.\"\"\"\n", " qb = torch.arange(nblk, device=dev).unsqueeze(1)\n", " kb = torch.arange(nblk, device=dev).unsqueeze(0)\n", " legal = kb <= qb\n", " free = ((kb <= qb) & (kb > qb - LOCAL_W)) | (kb < SINK)\n", " sc = scores.masked_fill(~legal, float(\"-inf\"))\n", " keep = torch.zeros_like(legal)\n", " keep.scatter_(1, sc.topk(min(TOP_M, nblk), dim=1).indices, True)\n", " keep = (keep & legal) | free\n", " m = torch.where(keep, 0.0, NEG)\n", " full = m.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)[:T, :T]\n", " # Token-level causality. Block granularity is not enough: supplying any mask turns off\n", " # is_causal, so this mask must carry the causal structure itself or the model reads the\n", " # future within its own block.\n", " ar = torch.arange(T, device=dev)\n", " full = full.masked_fill(ar.unsqueeze(0) > ar.unsqueeze(1), NEG)\n", " return full.unsqueeze(0).unsqueeze(0), float(keep.float().sum(1).mean())\n", "\n", "def sel_meanpool(li, q, k, T, dev):\n", " nblk = T // BLOCK\n", " kb = k[:nblk * BLOCK].view(nblk, BLOCK, -1).mean(1).float()\n", " qg = qgroup(q.float(), T)[:nblk * BLOCK].view(nblk, BLOCK, -1).mean(1)\n", " return build_mask(qg @ kb.T, nblk, T, dev)\n", "\n", "def sel_local(li, q, k, T, dev):\n", " nblk = T // BLOCK\n", " return build_mask(torch.full((nblk, nblk), -1e9, device=dev), nblk, T, dev)\n", "\n", "def sel_oracle(li, q, k, T, dev):\n", " \"\"\"THE CEILING. Picks the genuinely best blocks using real attention.\n", "\n", " No selector can beat this at a given density, so the gap between oracle and dense tells\n", " you how much the block-granularity budget costs at all, and the gap between a method and\n", " oracle tells you how much is left for a better method to win. Computed chunked from the\n", " captured post-RoPE features, which are verified to reproduce the model's own attention.\"\"\"\n", " nblk = T // BLOCK\n", " qh = q.view(T, N_HEAD, D_HEAD).permute(1, 0, 2).float()\n", " kh = k.view(T, N_KV, D_HEAD).permute(1, 0, 2).float().repeat_interleave(GROUP, dim=0)\n", " sc = torch.zeros(nblk, nblk, device=dev)\n", " ar = torch.arange(T, device=dev)\n", " with torch.no_grad():\n", " for b in range(nblk):\n", " qs, qe = b * BLOCK, (b + 1) * BLOCK\n", " s_ = torch.matmul(qh[:, qs:qe], kh.transpose(1, 2)) / math.sqrt(D_HEAD)\n", " s_ = s_.masked_fill((ar.unsqueeze(0) > ar[qs:qe].unsqueeze(1)).unsqueeze(0),\n", " float(\"-inf\"))\n", " pr = F.softmax(s_, dim=-1).mean(0).sum(0) # (T,) mass per key\n", " sc[b] = pr.view(nblk, BLOCK).sum(1)\n", " return build_mask(sc, nblk, T, dev)\n", "\n", "\n", "def sel_random(li, q, k, T, dev):\n", " \"\"\"THE NULL. Any method that cannot beat this is not selecting, it is guessing.\"\"\"\n", " nblk = T // BLOCK\n", " g = torch.Generator(device=\"cpu\").manual_seed(1234 + li)\n", " return build_mask(torch.rand(nblk, nblk, generator=g).to(dev), nblk, T, dev)\n", "\n", "\n", "# ============================================================\n", "# =============== ADD YOUR IDEA HERE =======================\n", "# ============================================================\n", "# A selector is one function with this signature:\n", "#\n", "# f(layer_idx, q, k, T, device) -> (additive_mask, mean_blocks_kept)\n", "#\n", "# q (T, N_HEAD * D_HEAD) post-RoPE queries, verified against the model's own attention\n", "# k (T, N_KV * D_HEAD) post-RoPE keys\n", "# returns whatever build_mask(scores, nblk, T, dev) gives you, where `scores` is\n", "# (query_block, key_block) and higher means more worth reading.\n", "#\n", "# build_mask handles the top-m cut, the local window, the sink and TOKEN-LEVEL CAUSALITY.\n", "# Do not hand-roll that last part: a block-granularity mask lets queries read their own\n", "# block's future and produces arms that beat full attention, which is how a silent bug\n", "# looked like a breakthrough here once before.\n", "#\n", "# Register it below and it is automatically run against every baseline and every guard.\n", "\n", "SEL = {\"oracle\": sel_oracle, \"meanpool\": sel_meanpool,\n", " \"random\": sel_random, \"local\": sel_local}\n", "print(\"[arms] dense, oracle ceiling, meanpool incumbent, random null, local floor\")\n", "\n", "# STRONGER CHECK. An all-selected mask must reproduce dense. This exercises the mask code\n", "# path the sparse arms actually use, which the empty-mask check above does NOT, because an\n", "# empty mask means no mask at all and quietly keeps is_causal=True.\n", "# The verification sequence must be short enough that top-m covers EVERY legal block,\n", "# or 'select everything' silently selects a subset and the check fails on its own\n", "# construction rather than on the mask. nblk <= TOP_M guarantees it at any lever value;\n", "# an earlier version of this harness hit exactly this when TOP_M dropped from 16 to 8\n", "# against a 16-block probe.\n", "_vd3 = doc(5)[:TOP_M * BLOCK].unsqueeze(0).to(DEV)\n", "MASKS.clear()\n", "with torch.no_grad():\n", " _dense_ref = model(_vd3).logits.float()\n", " _T3 = TOP_M * BLOCK\n", " _nb3 = _T3 // BLOCK\n", " _allsel = torch.zeros(_nb3, _nb3, device=DEV)\n", " for li in range(SKIP_LAYERS, N_LAYER):\n", " MASKS[li] = build_mask(_allsel + 1e9, _nb3, _T3, DEV)[0]\n", " _full_ref = model(_vd3).logits.float()\n", "MASKS.clear()\n", "_rel3 = float((_dense_ref - _full_ref).abs().max()) / float(_dense_ref.abs().max())\n", "print(f\"[verify] all-selected mask vs dense, relative {_rel3:.2e} -> \"\n", " f\"{'the mask path is causal and faithful' if bool(_rel3 < 2e-2) else 'MISMATCH'}\")\n", "if not bool(_rel3 < 2e-2):\n", " raise SystemExit(\"selecting EVERY block does not reproduce dense attention, so the mask \"\n", " \"itself is wrong. An earlier run had every sparse arm beating dense, \"\n", " \"which is exactly what a non-causal mask looks like.\")\n", "del _vd3, _dense_ref, _full_ref\n", "gc.collect(); torch.cuda.empty_cache()" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 7 - PERPLEXITY. The number with no excuse attached.\n", "# ============================================================\n", "@torch.no_grad()\n", "def nll(arm, d):\n", " global CAPTURE_ON\n", " MASKS.clear()\n", " if arm != \"dense\":\n", " CAPTURE.clear(); CAPTURE_ON = True\n", " model(d)\n", " CAPTURE_ON = False\n", " T = d.shape[1]\n", " kept = []\n", " for li in range(SKIP_LAYERS, N_LAYER):\n", " q, k = CAPTURE[li]\n", " m, kb = SEL[arm](li, q, k, T, q.device)\n", " MASKS[li] = m\n", " kept.append(kb)\n", " CAPTURE.clear()\n", " dens = sum(kept) / len(kept) * BLOCK / T\n", " else:\n", " dens = 1.0\n", " gc.collect(); torch.cuda.empty_cache()\n", " lg = model(d).logits[0, -EVAL_TAIL - 1:-1].float()\n", " v = F.cross_entropy(lg, d[0, -EVAL_TAIL:]).item()\n", " MASKS.clear(); gc.collect(); torch.cuda.empty_cache()\n", " return v, dens\n", "\n", "N_DOCS = int(os.environ.get(\"SAB_DOCS\", \"3\"))\n", "print(f\"[mem] mask term at this context: \"\n", " f\"{N_HEAD * CTX * CTX * 2 / 1e9:.2f} GB per layer, expanded across heads\")\n", "ARMS = [\"dense\", \"oracle\", \"meanpool\", \"random\", \"local\"]\n", "PPL, DENSITY = {a: [] for a in ARMS}, {a: [] for a in ARMS}\n", "for di in range(N_DOCS):\n", " d = doc(9000 + di).unsqueeze(0).to(DEV)\n", " for a in ARMS:\n", " t0 = time.time()\n", " v, dn = nll(a, d)\n", " PPL[a].append(v); DENSITY[a].append(dn)\n", " print(f\" doc {di} {a:9s} nll {v:.4f} density {dn*100:5.1f}% ({time.time()-t0:.0f}s)\",\n", " flush=True)\n", " del d; gc.collect(); torch.cuda.empty_cache()\n", "\n", "MEAN = {a: sum(v) / len(v) for a, v in PPL.items()}\n", "DENS = {a: sum(v) / len(v) for a, v in DENSITY.items()}\n", "print(\"\\n\" + \"=\" * 86)\n", "print(f\"PERPLEXITY. {MODEL_ID}, {str(DTYPE).split('.')[-1]}, ctx {CTX:,}, \"\n", " f\"scored on the last {EVAL_TAIL:,} tokens, {N_DOCS} docs\")\n", "print(\"=\" * 86)\n", "print(f\"{'arm':<11} {'nll':>8} {'ppl':>9} {'vs dense':>10} {'density':>9} selection cost\")\n", "for a in ARMS:\n", " cost = {\"dense\": \"reads everything\", \"oracle\": \"CEILING, cheats\", \"meanpool\": \"O(N/B)\",\n", " \"random\": \"NULL, guesses\", \"local\": \"none\"}[a]\n", " print(f\"{a:<11} {MEAN[a]:8.4f} {math.exp(MEAN[a]):9.3f} {MEAN[a]-MEAN['dense']:+10.4f} \"\n", " f\"{DENS[a]*100:8.1f}% {cost}\")" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "# ============================================================\n", "# CELL 8 - THE STOPWATCH, honestly scoped\n", "# ============================================================\n", "# Two separate things, kept separate on purpose.\n", "# 1. SELECTOR cost: how long the method takes to DECIDE what to read. This is a like-for\n", "# like comparison against the incumbent baseline.\n", "# 2. MEMORY TRAFFIC: KV reads per generated token. Hardware independent, and what a fused\n", "# kernel would actually realise. It is NOT a measured speedup and is not presented as one.\n", "# End-to-end wall clock is deliberately NOT claimed: masking does not skip the reads, so this\n", "# harness cannot demonstrate a speedup even in principle. Only a real kernel can.\n", "\n", "global CAPTURE_ON\n", "d = doc(777).unsqueeze(0).to(DEV)\n", "MASKS.clear(); CAPTURE.clear(); CAPTURE_ON = True\n", "with torch.no_grad():\n", " model(d)\n", "CAPTURE_ON = False\n", "LI = N_LAYER // 2\n", "q, k = CAPTURE[LI]\n", "T = d.shape[1]\n", "\n", "def bench(fn, reps=5):\n", " torch.cuda.synchronize()\n", " for _ in range(2):\n", " fn()\n", " torch.cuda.synchronize()\n", " e0, e1 = torch.cuda.Event(True), torch.cuda.Event(True)\n", " e0.record()\n", " for _ in range(reps):\n", " fn()\n", " e1.record(); torch.cuda.synchronize()\n", " return e0.elapsed_time(e1) / reps\n", "\n", "SEL_MS = {a: bench(lambda a=a: SEL[a](LI, q, k, T, q.device)) for a in (\"meanpool\",)}\n", "CAPTURE.clear(); del d; gc.collect(); torch.cuda.empty_cache()\n", "\n", "nblk = CTX // BLOCK\n", "TRAFFIC = {\"dense\": CTX,\n", " \"meanpool\": nblk // 2 + (TOP_M + LOCAL_W + SINK) * BLOCK,\n", " \"local\": (LOCAL_W + SINK) * BLOCK}\n", "print(\"\\n\" + \"=\" * 86)\n", "print(f\"SELECTION COST, layer {LI}, {CTX:,} keys. Measured with cuda events.\")\n", "print(\"=\" * 86)\n", "print(f\"{'arm':<11} {'select ms':>11} {'KV reads/tok':>14} {'traffic vs dense':>18}\")\n", "print(f\"{'dense':<11} {0.0:11.2f} {TRAFFIC['dense']:14,} {1.0:17.1f}x\")\n", "for a in (\"meanpool\",):\n", " print(f\"{a:<11} {SEL_MS[a]:11.2f} {TRAFFIC[a]:14,} {CTX/TRAFFIC[a]:17.1f}x\")\n", "\n", "print(\"\\n\" + \"=\" * 86)\n", "print(\"LEADERBOARD, share of recoverable long-range signal\")\n", "print(\"=\" * 86)\n", "_floor, _top = MEAN[\"local\"], MEAN[\"dense\"]\n", "_span = _floor - _top\n", "print(f\" full attention {_top:.4f} nats, local-only floor {_floor:.4f}, \"\n", " f\"so {_span:.4f} nats are on the table\")\n", "if _span < 0.05:\n", " print(\"\\n RUN INVALID. The local-only floor is within 0.05 nats of full attention, so this\")\n", " print(\" document needs no long-range retrieval and NO arm was actually tested. Nothing\")\n", " print(\" above means anything. Use a longer context or a retrieval-heavy task.\")\n", "else:\n", " print(f\"\\n{'arm':<11} {'recovers':>10} bar\")\n", " for a_ in sorted([x for x in ARMS if x not in (\"dense\", \"local\")],\n", " key=lambda x: -(1 - (MEAN[x] - _top) / _span)):\n", " r = 1 - (MEAN[a_] - _top) / _span\n", " print(f\"{a_:<11} {r*100:9.1f}% {'#' * max(0, min(40, int(r * 40)))}\")\n", " _orc = 1 - (MEAN[\"oracle\"] - _top) / _span\n", " _rnd = 1 - (MEAN[\"random\"] - _top) / _span\n", " print(f\"\\n READING IT:\")\n", " print(f\" oracle {_orc*100:.1f}% is the CEILING at this density. A method near it is\")\n", " print(f\" near-optimal and the remaining loss is the block budget, not the selector.\")\n", " print(f\" random {_rnd*100:.1f}% is the NULL. Anything at or below it is not selecting.\")\n", " if _orc - _rnd < 0.15:\n", " print(\" WARNING: oracle and random are within 15 points, so block choice barely\")\n", " print(\" matters at this density and the test has little power. Lower TOP_M.\")\n", " _best = max([x for x in ARMS if x not in (\"dense\", \"local\", \"oracle\")],\n", " key=lambda x: -(MEAN[x] - _top))\n", " print(f\" best real method: {_best}, \"\n", " f\"{(1-(MEAN[_best]-_top)/_span)*100:.1f}% of the recoverable signal, \"\n", " f\"{(_orc-(1-(MEAN[_best]-_top)/_span))*100:.1f} points below the ceiling\")\n", "\n", "print(\"\\n\" + \"=\" * 86)\n", "print(\"SELECTION COST. Measured. Traffic is ANALYTIC and is not a measured speedup.\")\n", "print(\"=\" * 86)\n", "print(f\"{'arm':<11} {'select ms':>11} {'KV reads/tok':>14} {'traffic vs dense':>18}\")\n", "for a_ in (\"meanpool\",):\n", " print(f\"{a_:<11} {SEL_MS[a_]:11.2f} {TRAFFIC[a_]:14,} {CTX/TRAFFIC[a_]:17.1f}x\")\n", "print(\"\\n No end-to-end speedup is claimed. Masking does not skip reads, so this harness\")\n", "print(\" cannot show one even in principle. Only a fused block-sparse kernel could.\")\n", "\n", "with open(\"/kaggle/working/sabench.json\", \"w\") as f:\n", " json.dump({\"model\": MODEL_ID, \"rope_features\": ROPE_FEATURES, \"rope_err\": _ROPE_ERR, \"dtype\": str(DTYPE), \"ctx\": CTX, \"eval_tail\": EVAL_TAIL,\n", " \"block\": BLOCK, \"top_m\": TOP_M, \"patch_rel_err\": _rel, \"nll_mean\": MEAN,\n", " \"nll_all\": PPL, \"density\": DENS, \"select_ms\": SEL_MS, \"traffic\": TRAFFIC,\n", " \"n_layer\": N_LAYER, \"kv_bytes_per_tok\": KV_BYTES_PER_TOK}, f, indent=1)\n", "print(\"\\n[sabench] wrote sabench.json\")" ], "outputs": [], "execution_count": null } ], "metadata": { "kernelspec": { "name": "python3", "language": "python", "display_name": "Python 3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }