gqa decode: fix cube_base output-store; measure batch scaling on one SIP

Concurrency fix (completes the cube_base change): the decode kernel's output store o_base still used the global cube_id while every load used the user-local cube_local. For a single user (cube_base=0) they coincide, so it was masked; a second batched user (cube_base>0) overshot its o shard into an unmapped VA, mis-decoded to a phantom sip0.cube0.pe8 and raised a RoutingError. Switch o_base to cube_local (one line); single-user results are byte-identical (decode smoke/correctness pass). Batch harness (un-skipped): create all users' tensors first, then submit all launches deferred and drain together, so deploys don't drive an already-submitted launch and serialize the batch. Concurrent users on disjoint CUBE groups now overlap to within 3-5% of the single-user latency. Swept B in {1,2,capacity} at 8K (high-B dense runs are the expensive ones; throughput is near-linear in B). Measured result: aggregate throughput at a full SIP rises from 0.86 (1-kv-per-cube, 2 users) to 1.41 requests/us (8-kv-per-cube, 16 users) — the dense mapping wins throughput, the spread mapping wins latency. S5.2 batch paragraph updated projected -> measured; add batch_scaling figure and plot_batch_scaling.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-22 17:37:10 -07:00
parent 8c108154a7
commit 13f5cbc317
7 changed files with 132 additions and 32 deletions
@@ -0,0 +1,61 @@
+"""Batch-scaling figure for concurrent decode users on one SIP.
+
+Reads ``docs/sweeps/decode_batch_sweep.csv`` (from
+``tests/attention/test_gqa_decode_batch_sweep.py``) and plots aggregate
+throughput (requests / us) versus batch size B for the four KV mappings,
+one panel per context length. Each mapping's curve runs up to its SIP
+capacity 16/C (A1=2 ... B=16 users).
+"""
+from pathlib import Path
+
+import csv
+
+import matplotlib.pyplot as plt
+
+ROOT = Path(__file__).resolve().parents[3]
+CSV = ROOT / "docs" / "sweeps" / "decode_batch_sweep.csv"
+OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
+       / "figures" / "gqa_short_context" / "batch_scaling.png")
+
+MODE_LABEL = {"A1": "1-kv-per-cube", "A2": "2-kv-per-cube",
+              "A4": "4-kv-per-cube", "B": "8-kv-per-cube"}
+MODE_COLOR = {"A1": "#1f77b4", "A2": "#2ca02c",
+              "A4": "#ffd000", "B": "#d62728"}
+MODES = ["A1", "A2", "A4", "B"]
+
+
+def main():
+    rows = list(csv.DictReader(CSV.open()))
+    contexts = sorted({int(r["S_kv"]) for r in rows})
+    fig, axes = plt.subplots(1, len(contexts), figsize=(6.5 * len(contexts), 4.5),
+                             squeeze=False)
+    for col, S in enumerate(contexts):
+        ax = axes[0][col]
+        for m in MODES:
+            pts = sorted(
+                ((int(r["B"]), float(r["throughput_users_per_us"]))
+                 for r in rows if r["mode"] == m and int(r["S_kv"]) == S),
+                key=lambda t: t[0],
+            )
+            if not pts:
+                continue
+            xs, ys = zip(*pts)
+            ax.plot(xs, ys, marker="o", color=MODE_COLOR[m],
+                    label=MODE_LABEL[m])
+            ax.annotate(f"{ys[-1]:.2f}", (xs[-1], ys[-1]),
+                        textcoords="offset points", xytext=(4, 4), fontsize=8)
+        ax.set_title(f"S_kv = {S // 1024}K")
+        ax.set_xlabel("batch size B (concurrent users)")
+        ax.set_ylabel("throughput (requests / µs)")
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=9)
+    fig.suptitle("Batch scaling: aggregate throughput vs concurrent users "
+                 "(one SIP)", fontsize=12, fontweight="bold")
+    fig.tight_layout()
+    OUT.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(OUT, dpi=140, bbox_inches="tight")
+    print(f"  ✓ {OUT}")
+
+
+if __name__ == "__main__":
+    main()