"""Predict decode-speed ceilings from measured bandwidth, a model's bytes and a memory budget. Purpose: turn the bandwidth you just measured into a falsifiable prediction. Decode reads the active weights and the key-value cache once per token, so the token rate cannot exceed bandwidth divided by those bytes. Part 6 runs the models and shows how close the engines get. Platform: all (this is arithmetic; it needs no accelerator and no model download) Minimum memory: 8 GB Assumes: a readable copy of the course's models.json, whose path is passed on the command line; a bandwidth figure, either typed or read from the latest part-05/bandwidth-test line in the lab notebook Usage: python3 predict-decode.py --models models.json --labbook labbook.md [--bandwidth-gbps 250] [--budget-gb 22.5] [--split | --cpu-bandwidth-gbps 60] [--model-ids qwen3-8b,qwen3-30b-a3b,qwen3-32b] [--quants q4_k_m,bf16] [--context-tokens 4096] Method. 1 GB = 1e9 bytes throughout, the convention of the bandwidth scripts. bandwidth --bandwidth-gbps, or else the accelerator's "read" figure from the latest part-05/bandwidth-test line in --labbook; --split takes that line's CPU "read" figure as --cpu-bandwidth-gbps active GB file size x active / total parameters (the whole file for a dense model); Part 3 shows this is a few per cent high for dense Qwen3 files KV GB kv.bytesPerTokenFp16 x context tokens: an f16 cache, one sequence ceiling bandwidth / (active GB + KV GB), at 0 tokens and at --context-tokens fit with --budget-gb: "fits" when file + KV <= budget. Otherwise, with --cpu-bandwidth-gbps, "split": the cache and (budget - KV) GB of weights stay on the accelerator, the rest is read from system memory, and seconds per token = (active on accelerator + KV) / bandwidth + active in system memory / cpu bandwidth Without it, "over": the prediction assumes memory the device does not have. Every figure is a ceiling: sampling, kernel launches and the arithmetic itself are given no time, and a split also ignores the CPU's arithmetic on the layers it holds. """ import argparse import json import sys from datetime import datetime, timezone from pathlib import Path DEFAULT_IDS = "qwen3-8b,qwen3-30b-a3b,qwen3-32b" DEFAULT_QUANTS = "q4_k_m,bf16" QUANT_LABELS = {"q4_k_m": "Q4_K_M", "q8_0": "Q8_0", "q6_k": "Q6_K", "bf16": "BF16", "mxfp4": "MXFP4", "iq4_xs": "IQ4_XS"} def load_models(path): try: data = json.loads(Path(path).read_text(encoding="utf-8")) except (OSError, ValueError) as exc: sys.exit(f"predict-decode: cannot read {path}: {exc}") return {m["id"]: m for m in data.get("models", [])} def active_fraction(model): params = model.get("params", {}) total = float(params.get("totalB") or 0.0) active = float(params.get("activeB") or 0.0) if total <= 0 or active <= 0: return 1.0 return min(active / total, 1.0) def ceiling(size_gb, fraction, kv_gb, bandwidth, budget, cpu_bandwidth): """Return (fit label, tokens per second or None).""" active = size_gb * fraction if budget is None or size_gb + kv_gb <= budget: return ("fits" if budget is not None else "-"), bandwidth / (active + kv_gb) if cpu_bandwidth is None or budget <= kv_gb: return "over", None on_device = (budget - kv_gb) / size_gb # share of the file the accelerator keeps seconds = (active * on_device + kv_gb) / bandwidth + active * (1 - on_device) / cpu_bandwidth return "split", 1.0 / seconds def from_labbook(path): """Return the latest part-05/bandwidth-test record in the notebook, or exit.""" latest = None try: lines = Path(path).read_text(encoding="utf-8").splitlines() except OSError as exc: sys.exit(f"predict-decode: cannot read {path}: {exc}") for line in lines: if line.startswith("{") and '"part-05/bandwidth-test"' in line: try: latest = json.loads(line) except ValueError: continue if latest is None or "read" not in latest.get("gbps", {}).get(latest.get("device"), {}): sys.exit(f"predict-decode: no part-05/bandwidth-test line with a read figure in {path}; " f"run bandwidth-test.py --labbook {path} first, or pass --bandwidth-gbps") return latest def show(rate): return f"{rate:>9.1f}" if rate is not None else f"{'-':>9}" def main(): parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--bandwidth-gbps", type=float, default=None, help="the read figure bandwidth-test.py measured, in GB/s; " "default: the latest one in --labbook") parser.add_argument("--models", required=True, help="path to the course models.json") parser.add_argument("--model-ids", default=DEFAULT_IDS) parser.add_argument("--quants", default=DEFAULT_QUANTS) parser.add_argument("--context-tokens", type=int, default=4096, help="tokens already in the cache for the second ceiling column") parser.add_argument("--budget-gb", type=float, default=None, help="memory the accelerator may use for weights and cache, in GB") parser.add_argument("--cpu-bandwidth-gbps", type=float, default=None, help="the CPU read figure, to estimate models that must be split") parser.add_argument("--split", action="store_true", help="use the CPU read figure from the same notebook line for " "rows over budget (Track N)") parser.add_argument("--labbook", default=None, help="read the bandwidth from, and append one JSON line to, this file") args = parser.parse_args() source = "--bandwidth-gbps" if args.bandwidth_gbps is None or args.split: if not args.labbook: sys.exit("predict-decode: pass --bandwidth-gbps, or --labbook with a " "part-05/bandwidth-test line in it") record = from_labbook(args.labbook) if args.bandwidth_gbps is None: args.bandwidth_gbps = record["gbps"][record["device"]]["read"] source = (f"{record['device']} read figure, {record.get('size_mb')} MiB run of " f"{record.get('date', 'an undated run')}") if args.split and args.cpu_bandwidth_gbps is None: if record["device"] == "cpu" or "cpu" not in record["gbps"]: sys.exit("predict-decode: --split needs an accelerator run with a cpu figure") args.cpu_bandwidth_gbps = record["gbps"]["cpu"]["read"] if args.bandwidth_gbps <= 0 or args.context_tokens < 0: sys.exit("predict-decode: --bandwidth-gbps must be positive and --context-tokens " "not negative") catalogue = load_models(args.models) wanted = [i.strip() for i in args.model_ids.split(",") if i.strip()] quants = [q.strip() for q in args.quants.split(",") if q.strip()] ctx = args.context_tokens print(f"bandwidth: {args.bandwidth_gbps:.1f} GB/s, from {source}") if args.cpu_bandwidth_gbps: print(f"cpu bandwidth for split rows: {args.cpu_bandwidth_gbps:.1f} GB/s") print(f"budget: {args.budget_gb:.1f} GB" if args.budget_gb is not None else "budget: not given, so no fit column") print("ceiling = bandwidth / (active GB + KV GB); 1 GB = 1e9 bytes\n") header = (f"{'model':<15}{'arch':<6}{'format':<8}{'file GB':>8}{'active GB':>10}" f"{f'KV@{ctx}':>9}{'fit':>6}{'tok/s@0':>9}{f'tok/s@{ctx}':>11}") print(header) print("-" * len(header)) predictions = [] for model_id in wanted: model = catalogue.get(model_id) if model is None: print(f"{model_id:<15}not found in {args.models}") continue fraction = active_fraction(model) kv_bytes = float(model.get("kv", {}).get("bytesPerTokenFp16") or 0.0) kv_gb = kv_bytes * ctx / 1e9 sizes = model.get("sizesGB", {}) for quant in quants: size_gb = sizes.get(quant) if size_gb is None: print(f"{model_id:<15}{model.get('architecture', '?'):<6}{quant:<8}" f" not listed (available: {', '.join(sorted(sizes)) or 'none'})") continue size_gb = float(size_gb) fit0, rate0 = ceiling(size_gb, fraction, 0.0, args.bandwidth_gbps, args.budget_gb, args.cpu_bandwidth_gbps) fit, rate = ceiling(size_gb, fraction, kv_gb, args.bandwidth_gbps, args.budget_gb, args.cpu_bandwidth_gbps) print(f"{model_id:<15}{model.get('architecture', '?'):<6}" f"{QUANT_LABELS.get(quant, quant):<8}{size_gb:>8.1f}" f"{size_gb * fraction:>10.2f}{kv_gb:>9.2f}{fit:>6}{show(rate0)}" f"{show(rate):>11}") predictions.append({ "model": model_id, "quant": quant, "size_gb": size_gb, "active_gb": round(size_gb * fraction, 3), "kv_gb": round(kv_gb, 3), "fit": fit, "fit_at_zero_context": fit0, "predicted_tokens_per_second": round(rate0, 1) if rate0 else None, "predicted_tokens_per_second_at_context": round(rate, 1) if rate else None, }) if args.labbook and predictions: record = { "lab": "part-05/predict-decode", "date": datetime.now(timezone.utc).isoformat(timespec="seconds"), "bandwidth_gbps": args.bandwidth_gbps, "bandwidth_source": source, "cpu_bandwidth_gbps": args.cpu_bandwidth_gbps, "budget_gb": args.budget_gb, "context_tokens": ctx, "models_json": str(args.models), "method": "bandwidth / (file x active/total + kv bytes x context); split rows " "add the system-memory share at the cpu bandwidth", "predictions": predictions, } with Path(args.labbook).open("a", encoding="utf-8") as handle: handle.write(json.dumps(record, sort_keys=True) + "\n") print(f"\nrecorded in {args.labbook}") if __name__ == "__main__": main()