eval/scripts/decode_savings.py

"""Decode-time savings from block-sparse decode, Qwen3-14B — WEIGHTS EXCLUDED.

Per the agreed accounting we drop the fixed 28 GB weight read (HBM) and the fixed FFN/linear FLOPs, and
report only the subsystem block-sparsity touches: KV-cache reads (HBM) and attention FLOPs. Both scale
with the number of ATTENDED context tokens, so they save the same fraction (and it's batch-independent).

Sparse attends ~k blocks' content + one summary per block:  attended ≈ k*block + (S/block)*summ.
As S grows this is ~flat vs dense's O(S), so savings grow — but the per-block resident summary caps it:
attended/S -> summ/block, so max saving ≈ block/summ.
"""
KVtok = 2*40*8*128*2            # Qwen3-14B: KV bytes/token/seq = 163,840
BLOCK, K, SUMM = 200, 2, 8      # measured recipe: ~200-tok blocks, k=2 active, 8-token summary/block


def attended(S):
    n = max(1, S // BLOCK)
    return K*BLOCK + n*SUMM


print(f"Qwen3-14B decode savings (WEIGHTS EXCLUDED; KV-read HBM & attention-FLOPs; batch-independent %)")
print(f"block={BLOCK} k={K} summary={SUMM}  -> savings cap ≈ block/summary = {BLOCK/SUMM:.0f}x\n")
print(f"{'ctx S':>8} | {'attended d→s':>16} | {'KV read/seq d→s':>20} | {'saved':>6} | {'factor':>7}")
for S in [2048, 4096, 8192, 32768, 131072, 524288]:
    a = attended(S); kd, ks = KVtok*S/1e6, KVtok*a/1e6
    print(f"{S:>8} | {S:>6} → {a:<7} | {kd:>7.1f} → {ks:>6.2f} MB | {(1-a/S)*100:>5.0f}% | {S/a:>5.0f}x")
print("\nBoth KV-read HBM and attention FLOPs drop by this same factor (4x @2k ... ~25x long-context).")
print("Accuracy is retained (sparse ~= dense F1). This is a KV-bandwidth optimization for long-context")
print("batched serving; to raise the cap, use fewer summary tokens (trades off selection quality).")