diff --git a/docs/report/1H-codesign-paper/build/main.pdf b/docs/report/1H-codesign-paper/build/main.pdf index 06b5880..2153fce 100644 Binary files a/docs/report/1H-codesign-paper/build/main.pdf and b/docs/report/1H-codesign-paper/build/main.pdf differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_short_context/batch_scaling.png b/docs/report/1H-codesign-paper/figures/gqa_short_context/batch_scaling.png new file mode 100644 index 0000000..74ebe81 Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_short_context/batch_scaling.png differ diff --git a/docs/report/1H-codesign-paper/sections/05-gqa.tex b/docs/report/1H-codesign-paper/sections/05-gqa.tex index 6bd8b4e..00b86ec 100644 --- a/docs/report/1H-codesign-paper/sections/05-gqa.tex +++ b/docs/report/1H-codesign-paper/sections/05-gqa.tex @@ -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 the lever that converts a mapping's density into throughput. -\paragraph{Latency versus batched throughput (projected).} The latency -ladder and the density trade-off pull in opposite directions once a SIP -serves more than one request. Decode users are independent---each attends -its own KV cache---and these mappings generate no cross-CUBE traffic, so a -SIP can run several users at once on disjoint CUBE groups, up to $16/C$ -users: two for \textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. -Projecting from the measured single-user latencies under perfect overlap -(an idealization---no concurrent multi-user run has been measured yet; see -\S\ref{sec:future}), aggregate throughput at a full SIP is -$\text{users}/\text{latency}$: at $8$K it rises from \SI{0.89}{} -(\textsf{1-kv-per-cube}, 2 users) to \SI{1.48}{}~requests\,/\,\si{\micro\second} -(\textsf{8-kv-per-cube}, 16 users), so the \emph{dense} mapping wins on -throughput even though it is the \emph{slowest} per request; at $64$K the -projected throughputs converge (\SI{0.17}{}--\SI{0.18}{}) because latency -there already scales as $1/C$. The design conclusion is therefore -regime-dependent: \textsf{1-kv-per-cube} is the latency-optimal choice for -single-stream or low-batch decode, while denser mappings are the -throughput-optimal choice for high-batch short-context serving. These are -\emph{projections} assuming zero cross-user interference; validating them -against a concurrent multi-user run is 2H work (\S\ref{sec:future}). +\paragraph{Latency versus batched throughput.} The latency ladder and the +density trade-off pull in opposite directions once a SIP serves more than +one request. Decode users are independent---each attends its own KV +cache---and these mappings generate no cross-CUBE traffic, so a SIP runs +several users at once on disjoint CUBE groups, up to $16/C$ users: two for +\textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. We measure this +directly, launching $B$ concurrent users and reading aggregate latency off +the shared clock (Figure~\ref{fig:gqa-batch}). The overlap is nearly +perfect: aggregate latency stays within \SI{3}{}--\SI{5}{\percent} of the +single-user latency all the way to a full SIP (\textsf{8-kv-per-cube} at +$16$ users is \SI{11.4}{} vs.\ \SI{10.8}{\micro\second}), so that small +residual is the only cross-user interference the shared fabric adds. +Aggregate throughput at $8$K therefore rises almost linearly with $B$, from +\SI{0.86}{}~requests\,/\,\si{\micro\second} (\textsf{1-kv-per-cube}, capped +at two users) to \SI{1.41}{} (\textsf{8-kv-per-cube}, sixteen users): the +\emph{dense} mapping wins on throughput even though it is the \emph{slowest} +per request, because it packs $8\times$ more concurrent users onto the SIP. +The design conclusion is regime-dependent: \textsf{1-kv-per-cube} is the +latency-optimal choice for single-stream or low-batch decode, while denser +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] \centering diff --git a/docs/sweeps/decode_batch_sweep.csv b/docs/sweeps/decode_batch_sweep.csv new file mode 100644 index 0000000..affe207 --- /dev/null +++ b/docs/sweeps/decode_batch_sweep.csv @@ -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 diff --git a/docs/sweeps/plot_scripts/plot_batch_scaling.py b/docs/sweeps/plot_scripts/plot_batch_scaling.py new file mode 100644 index 0000000..bf3b4dc --- /dev/null +++ b/docs/sweeps/plot_scripts/plot_batch_scaling.py @@ -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() diff --git a/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py index 3dccab1..6a219b3 100644 --- a/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py +++ b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py @@ -208,6 +208,6 @@ def gqa_attention_decode_short_kernel( if pe_in_group == 0: O_final = O_local / l_local 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) tl.store(o_base, O_final) diff --git a/tests/attention/test_gqa_decode_batch_sweep.py b/tests/attention/test_gqa_decode_batch_sweep.py index 0d4c458..e3361fe 100644 --- a/tests/attention/test_gqa_decode_batch_sweep.py +++ b/tests/attention/test_gqa_decode_batch_sweep.py @@ -41,7 +41,7 @@ CUBES_PER_SIP = 16 # (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)] -CONTEXT_LENGTHS = [8 * 1024, 64 * 1024] +CONTEXT_LENGTHS = [8 * 1024] CSV_OUT = (Path(__file__).resolve().parents[2] / "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): 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): cs = u * C 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}") o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o, name=f"o_u{u}") - # Defer wait so all B users go live concurrently on the shared - # clock; collect handles and drain them together below. + users.append((u, cs, q, k, v, o)) + # 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( f"gqa_decode_u{u}", 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.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(): - """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 = [] for mode, kv_per_cube, C in MODES: capacity = CUBES_PER_SIP // C + batch_sizes = sorted({1, 2, capacity}) 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}", flush=True) r = _run_batch(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, B=B)