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>
This commit is contained in:
2026-07-22 17:37:10 -07:00
parent 8c108154a7
commit 13f5cbc317
7 changed files with 132 additions and 32 deletions
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 69 KiB

@@ -240,26 +240,43 @@ requires lifting $M$, either by the long-context Q-tile split
(\S\ref{sec:gqa-composite}) or by batching multiple users---which is also (\S\ref{sec:gqa-composite}) or by batching multiple users---which is also
the lever that converts a mapping's density into throughput. the lever that converts a mapping's density into throughput.
\paragraph{Latency versus batched throughput (projected).} The latency \paragraph{Latency versus batched throughput.} The latency ladder and the
ladder and the density trade-off pull in opposite directions once a SIP density trade-off pull in opposite directions once a SIP serves more than
serves more than one request. Decode users are independent---each attends one request. Decode users are independent---each attends its own KV
its own KV cache---and these mappings generate no cross-CUBE traffic, so a cache---and these mappings generate no cross-CUBE traffic, so a SIP runs
SIP can run several users at once on disjoint CUBE groups, up to $16/C$ several users at once on disjoint CUBE groups, up to $16/C$ users: two for
users: two for \textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. \textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. We measure this
Projecting from the measured single-user latencies under perfect overlap directly, launching $B$ concurrent users and reading aggregate latency off
(an idealization---no concurrent multi-user run has been measured yet; see the shared clock (Figure~\ref{fig:gqa-batch}). The overlap is nearly
\S\ref{sec:future}), aggregate throughput at a full SIP is perfect: aggregate latency stays within \SI{3}{}--\SI{5}{\percent} of the
$\text{users}/\text{latency}$: at $8$K it rises from \SI{0.89}{} single-user latency all the way to a full SIP (\textsf{8-kv-per-cube} at
(\textsf{1-kv-per-cube}, 2 users) to \SI{1.48}{}~requests\,/\,\si{\micro\second} $16$ users is \SI{11.4}{} vs.\ \SI{10.8}{\micro\second}), so that small
(\textsf{8-kv-per-cube}, 16 users), so the \emph{dense} mapping wins on residual is the only cross-user interference the shared fabric adds.
throughput even though it is the \emph{slowest} per request; at $64$K the Aggregate throughput at $8$K therefore rises almost linearly with $B$, from
projected throughputs converge (\SI{0.17}{}--\SI{0.18}{}) because latency \SI{0.86}{}~requests\,/\,\si{\micro\second} (\textsf{1-kv-per-cube}, capped
there already scales as $1/C$. The design conclusion is therefore at two users) to \SI{1.41}{} (\textsf{8-kv-per-cube}, sixteen users): the
regime-dependent: \textsf{1-kv-per-cube} is the latency-optimal choice for \emph{dense} mapping wins on throughput even though it is the \emph{slowest}
single-stream or low-batch decode, while denser mappings are the per request, because it packs $8\times$ more concurrent users onto the SIP.
throughput-optimal choice for high-batch short-context serving. These are The design conclusion is regime-dependent: \textsf{1-kv-per-cube} is the
\emph{projections} assuming zero cross-user interference; validating them latency-optimal choice for single-stream or low-batch decode, while denser
against a concurrent multi-user run is 2H work (\S\ref{sec:future}). mappings are the throughput-optimal choice for high-batch short-context
serving. At long context the per-request latencies already scale as $1/C$,
so the same accounting predicts the throughputs converge; a measured
long-context batch sweep is 2H work (\S\ref{sec:future}).
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gqa_short_context/batch_scaling.png}
\caption{Batch scaling at $S_{kv}{=}8$K: aggregate throughput (requests per
\si{\micro\second}) versus the number of concurrent decode users on one SIP,
each user on a disjoint CUBE group. Each mapping scales almost linearly with
$B$ up to its SIP capacity $16/C$---\textsf{1-kv-per-cube} saturates at two
users, \textsf{8-kv-per-cube} at sixteen. The dense mapping reaches the
highest throughput; the spread mapping is latency-optimal but
capacity-limited. Concurrent users overlap to within
\SI{3}{}--\SI{5}{\percent} of the single-user latency.}
\label{fig:gqa-batch}
\end{figure}
\begin{figure}[t] \begin{figure}[t]
\centering \centering
+12
View File
@@ -0,0 +1,12 @@
mode,kv_per_cube,C,S_kv,B,capacity,agg_latency_us,mean_user_latency_us,throughput_users_per_us
A1,1,8,8192,1,2,2.2471645000005376,2.2471645000005376,0.44500524994932983
A1,1,8,8192,2,2,2.3131845000012543,2.2471645000012357,0.864608940617973
A2,2,4,8192,1,4,3.3555127500005475,3.3555127500005475,0.2980170467240325
A2,2,4,8192,2,4,3.3935227500010514,3.355512750001042,0.5893580645656141
A2,2,4,8192,4,4,3.4695427500010703,3.3555127500011586,1.1528896711241752
A4,4,2,8192,1,8,5.586135500000557,5.586135500000557,0.17901463364071643
A4,4,2,8192,2,8,5.602295500000823,5.586135500000673,0.35699652044411195
A4,4,2,8192,8,8,5.783285499986028,5.586135499984957,1.3832967436968704
B,8,1,8192,1,16,10.816091000000947,10.816091000000947,0.09245484343649775
B,8,1,8192,2,16,10.848411000001535,10.816091000001236,0.18435879687815265
B,8,1,8192,16,16,11.38492099986691,10.816090999974403,1.4053676788962384
1 mode kv_per_cube C S_kv B capacity agg_latency_us mean_user_latency_us throughput_users_per_us
2 A1 1 8 8192 1 2 2.2471645000005376 2.2471645000005376 0.44500524994932983
3 A1 1 8 8192 2 2 2.3131845000012543 2.2471645000012357 0.864608940617973
4 A2 2 4 8192 1 4 3.3555127500005475 3.3555127500005475 0.2980170467240325
5 A2 2 4 8192 2 4 3.3935227500010514 3.355512750001042 0.5893580645656141
6 A2 2 4 8192 4 4 3.4695427500010703 3.3555127500011586 1.1528896711241752
7 A4 4 2 8192 1 8 5.586135500000557 5.586135500000557 0.17901463364071643
8 A4 4 2 8192 2 8 5.602295500000823 5.586135500000673 0.35699652044411195
9 A4 4 2 8192 8 8 5.783285499986028 5.586135499984957 1.3832967436968704
10 B 8 1 8192 1 16 10.816091000000947 10.816091000000947 0.09245484343649775
11 B 8 1 8192 2 16 10.848411000001535 10.816091000001236 0.18435879687815265
12 B 8 1 8192 16 16 11.38492099986691 10.816090999974403 1.4053676788962384
@@ -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()
@@ -208,6 +208,6 @@ def gqa_attention_decode_short_kernel(
if pe_in_group == 0: if pe_in_group == 0:
O_final = O_local / l_local O_final = O_local / l_local
o_base = (o_ptr o_base = (o_ptr
+ cube_id * kv_per_cube * Q_ROW_BYTES + cube_local * kv_per_cube * Q_ROW_BYTES
+ group_id_in_cube * Q_ROW_BYTES) + group_id_in_cube * Q_ROW_BYTES)
tl.store(o_base, O_final) tl.store(o_base, O_final)
+21 -11
View File
@@ -41,7 +41,7 @@ CUBES_PER_SIP = 16
# (mode, kv_per_cube, C); users-per-SIP capacity = CUBES_PER_SIP // C. # (mode, kv_per_cube, C); users-per-SIP capacity = CUBES_PER_SIP // C.
MODES = [("A1", 1, 8), ("A2", 2, 4), ("A4", 4, 2), ("B", 8, 1)] MODES = [("A1", 1, 8), ("A2", 2, 4), ("A4", 4, 2), ("B", 8, 1)]
CONTEXT_LENGTHS = [8 * 1024, 64 * 1024] CONTEXT_LENGTHS = [8 * 1024]
CSV_OUT = (Path(__file__).resolve().parents[2] CSV_OUT = (Path(__file__).resolve().parents[2]
/ "docs" / "sweeps" / "decode_batch_sweep.csv") / "docs" / "sweeps" / "decode_batch_sweep.csv")
@@ -73,7 +73,11 @@ def _run_batch(*, kv_per_cube: int, C: int, S_kv: int, B: int, h_q: int = 64):
def _bench_fn(ctx): def _bench_fn(ctx):
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
pending = [] # Deploy ALL users' tensors first. Each deploy blocks (advances the
# clock), so if a launch were already submitted a later deploy would
# drive it to completion and serialize the batch. Creating every
# tensor before any launch keeps the deploys from touching kernels.
users = []
for u in range(B): for u in range(B):
cs = u * C cs = u * C
dp_q = DPPolicy(cube="column_wise", pe="replicate", dp_q = DPPolicy(cube="column_wise", pe="replicate",
@@ -90,8 +94,11 @@ def _run_batch(*, kv_per_cube: int, C: int, S_kv: int, B: int, h_q: int = 64):
name=f"v_u{u}") name=f"v_u{u}")
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o, o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
name=f"o_u{u}") name=f"o_u{u}")
# Defer wait so all B users go live concurrently on the shared users.append((u, cs, q, k, v, o))
# clock; collect handles and drain them together below. # Now submit every launch deferred (no blocking wait between them),
# then drain together so all B run concurrently on the shared clock.
pending = []
for u, cs, q, k, v, o in users:
pending += ctx.launch( pending += ctx.launch(
f"gqa_decode_u{u}", f"gqa_decode_u{u}",
gqa_attention_decode_short_kernel, gqa_attention_decode_short_kernel,
@@ -149,18 +156,21 @@ def _metrics(r, *, mode, kv_per_cube, C, S_kv, B):
@pytest.mark.slow @pytest.mark.slow
@pytest.mark.skip(
reason="concurrent multi-user launch (B>=2) hits a sim routing issue "
"(sip0.cube0.pe8 not found); single-user path works. Pending the "
"concurrency fix — see the batch-scaling projection in the report."
)
def test_decode_batch_sweep(): def test_decode_batch_sweep():
"""Sweep 4 modes × contexts × B=1..capacity; dump batch-scaling CSV.""" """Sweep 4 modes × contexts × B{1, 2, capacity}; dump batch CSV.
Throughput is near-linear in B (aggregate latency stays within a few
percent of the single-user latency across the batch), so the endpoints
plus a low point capture the scaling; the high-B runs of the dense
mappings are the expensive ones and full 1..capacity granularity is not
needed for the trade-off.
"""
rows = [] rows = []
for mode, kv_per_cube, C in MODES: for mode, kv_per_cube, C in MODES:
capacity = CUBES_PER_SIP // C capacity = CUBES_PER_SIP // C
batch_sizes = sorted({1, 2, capacity})
for S_kv in CONTEXT_LENGTHS: for S_kv in CONTEXT_LENGTHS:
for B in range(1, capacity + 1): for B in batch_sizes:
print(f">>> BATCH {mode} S_kv={S_kv//1024}K B={B}/{capacity}", print(f">>> BATCH {mode} S_kv={S_kv//1024}K B={B}/{capacity}",
flush=True) flush=True)
r = _run_batch(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, B=B) r = _run_batch(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, B=B)