Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| a4a2683aad | |||
| faf011dae5 | |||
| 1dade267ff |
|
After Width: | Height: | Size: 136 KiB |
|
Before Width: | Height: | Size: 230 KiB After Width: | Height: | Size: 305 KiB |
|
Before Width: | Height: | Size: 260 KiB After Width: | Height: | Size: 381 KiB |
@@ -145,9 +145,9 @@ with credit return, splitting the control plane into PE\_IPCQ and the
|
|||||||
data plane into PE\_DMA, with head updates riding the payload and tail
|
data plane into PE\_DMA, with head updates riding the payload and tail
|
||||||
updates riding a 16\,B side-channel credit (\S\ref{sec:allreduce}).
|
updates riding a 16\,B side-channel credit (\S\ref{sec:allreduce}).
|
||||||
|
|
||||||
\begin{figure}[t]
|
\begin{figure*}[t]
|
||||||
\centering
|
\centering
|
||||||
\includegraphics[width=\linewidth]{ipcq_alternatives_architecture_stacked.png}
|
\includegraphics[width=0.78\linewidth]{ipcq_alternatives_architecture_flow.png}
|
||||||
\caption{Per-send data and control flow for the four PE-to-PE
|
\caption{Per-send data and control flow for the four PE-to-PE
|
||||||
signalling mechanisms (sender\,$\rightarrow$\,NoC\,$\rightarrow$\,receiver).
|
signalling mechanisms (sender\,$\rightarrow$\,NoC\,$\rightarrow$\,receiver).
|
||||||
Doorbell and RDMA-CQ each issue two fabric transactions (payload then
|
Doorbell and RDMA-CQ each issue two fabric transactions (payload then
|
||||||
@@ -158,7 +158,7 @@ the payload flit train and returns the tail credit on a side channel,
|
|||||||
so a send is one MMIO write and a receive is a flip-flop read. This is
|
so a send is one MMIO write and a receive is a flip-flop read. This is
|
||||||
a \emph{design schematic}, not a measured comparison.}
|
a \emph{design schematic}, not a measured comparison.}
|
||||||
\label{fig:ipcq-arch}
|
\label{fig:ipcq-arch}
|
||||||
\end{figure}
|
\end{figure*}
|
||||||
|
|
||||||
\begin{figure}[t]
|
\begin{figure}[t]
|
||||||
\centering
|
\centering
|
||||||
|
|||||||
@@ -284,7 +284,9 @@ is unchanged across all three forms: decode is bound by streaming the KV
|
|||||||
cache out of HBM, so the command form does not move the critical path.
|
cache out of HBM, so the command form does not move the critical path.
|
||||||
|
|
||||||
That juxtaposition is the point. The composite command is not a latency
|
That juxtaposition is the point. The composite command is not a latency
|
||||||
optimization for this memory-bound decode; it is a \emph{CPU-issue}
|
optimization for this memory-bound decode \emph{at the 64-way production scale} (a
|
||||||
|
single-rank caveat follows, Figure~\ref{fig:gqa-decode-stream}); it is a
|
||||||
|
\emph{CPU-issue}
|
||||||
optimization. Its value is removing the per-tile dispatch work that would
|
optimization. Its value is removing the per-tile dispatch work that would
|
||||||
otherwise grow without bound as context grows, freeing PE\_CPU to run
|
otherwise grow without bound as context grows, freeing PE\_CPU to run
|
||||||
ahead and keep the engines fed---which is exactly what lets the
|
ahead and keep the engines fed---which is exactly what lets the
|
||||||
@@ -295,8 +297,49 @@ memory-bound path its marginal cost over the plain GEMM composite is small
|
|||||||
(98 vs.\ 94 commands), and like the plain composite it keeps the issued
|
(98 vs.\ 94 commands), and like the plain composite it keeps the issued
|
||||||
count flat as context scales.
|
count flat as context scales.
|
||||||
|
|
||||||
\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
|
\paragraph{Isolating the rank: the masked streaming win.} That neutrality
|
||||||
form is latency-neutral---is a property of its regime, not of the
|
is a property of the \emph{full 64-way} critical path, not of the local
|
||||||
|
attention: at production scale the inter-CUBE $(m,\ell,O)$ reduce tail and
|
||||||
|
shared-HBM contention set the wall clock, so a faster local attention does
|
||||||
|
not surface. Stripping those away---a single rank, no cross-CUBE reduce,
|
||||||
|
swept over the per-rank context $S_{kv}$ (so $S_{kv}{=}16$\,K here is the
|
||||||
|
per-PE load of a 1M-token, 64-way-sharded decode)---exposes the local
|
||||||
|
attention directly (Figure~\ref{fig:gqa-decode-stream}), and the composite
|
||||||
|
\emph{does} win, by \SI{25}{}--\SI{28}{\percent}. The reason is the
|
||||||
|
memory-bound mirror of prefill: its scheduler-streamed concurrent per-tile
|
||||||
|
DMAs keep the HBM pipeline full and reach \SI{233}{\giga\byte\per\second}
|
||||||
|
---\SI{91}{\percent} of the per-rank \SI{256}{\giga\byte\per\second}
|
||||||
|
roofline---whereas the primitive kernel's blocking \textsf{tl.dot}
|
||||||
|
serializes one tile DMA at a time and plateaus at
|
||||||
|
\SI{166}{\giga\byte\per\second}. So even for memory-bound decode the
|
||||||
|
composite is not \emph{only} a CPU-issue optimization---it also extracts
|
||||||
|
bandwidth---but that latency benefit materializes only when the local
|
||||||
|
attention is on the critical path, which at 64-way production scale it is
|
||||||
|
not.
|
||||||
|
|
||||||
|
\begin{figure}[t]
|
||||||
|
\centering
|
||||||
|
\includegraphics[width=\linewidth]{gqa_decode_streaming.png}
|
||||||
|
\caption{Single-rank memory-bound decode ($T_q{=}1$, $M{=}8$), three
|
||||||
|
command forms, swept over per-rank context. \emph{Left:} end-to-end
|
||||||
|
latency---the composite forms run \SI{25}{}--\SI{28}{\percent} below the
|
||||||
|
primitive, a gap that widens with context. \emph{Right:} achieved HBM
|
||||||
|
bandwidth against the per-rank \SI{256}{\giga\byte\per\second} roofline.
|
||||||
|
The primitive's blocking load$\rightarrow$dot serializes the KV stream and
|
||||||
|
plateaus at \SI{166}{\giga\byte\per\second}; the composite forms pipeline
|
||||||
|
concurrent per-tile DMAs through the scheduler and reach
|
||||||
|
\SI{233}{\giga\byte\per\second}. This is the memory-bound mirror of the
|
||||||
|
prefill result (Figure~\ref{fig:gqa-prefill-cb}): there the composite
|
||||||
|
approaches the MAC roofline, here the bandwidth roofline. Capped at
|
||||||
|
$16$\,K---the plain composite materializes the full $(M,S_{kv})$ scores in
|
||||||
|
TCM, so beyond that only the \textsf{softmax\_merge} recipe, which tiles
|
||||||
|
the softmax, stays within scratch.}
|
||||||
|
\label{fig:gqa-decode-stream}
|
||||||
|
\end{figure}
|
||||||
|
|
||||||
|
\paragraph{The compute-bound mirror: prefill.} Decode's \emph{production}
|
||||||
|
verdict---command form is latency-neutral at 64-way scale---is a property
|
||||||
|
of its regime, not of the
|
||||||
composite command. A decode step has $T_q{=}1$, so its score and context
|
composite command. A decode step has $T_q{=}1$, so its score and context
|
||||||
products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
|
products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
|
||||||
the kernel is bound by streaming the KV cache. Prefill is the opposite
|
the kernel is bound by streaming the KV cache. Prefill is the opposite
|
||||||
|
|||||||
@@ -0,0 +1,204 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""SCRATCH EXPERIMENT (not production; do not commit).
|
||||||
|
|
||||||
|
Question: does charging the *primitive* decode kernel for per-HW-tile
|
||||||
|
(16x16x16) CPU dispatch flip the "composite gives no decode-latency
|
||||||
|
benefit" conclusion?
|
||||||
|
|
||||||
|
We monkeypatch TLContext.dot so that, in the primitive kernel, every
|
||||||
|
tl.dot whose (M,K,N) exceeds the HW GEMM tile (mac_m/mac_k/mac_n) is
|
||||||
|
split by the CPU into ceil(M/mac_m)*ceil(K/mac_k)*ceil(N/mac_n)
|
||||||
|
HW-tile-sized GemmCmds. Each tile GemmCmd is emitted through the normal
|
||||||
|
_emit() path, so it (a) charges PE_CPU dispatch overhead via
|
||||||
|
_charge_dispatch (PeCpuOverheadCmd), and (b) blocks like a normal
|
||||||
|
single-op GemmCmd on PE_GEMM at the cycle-accurate ceil-product latency.
|
||||||
|
|
||||||
|
We also inject mac_m/mac_k/mac_n into every pe_gemm topology node so BOTH
|
||||||
|
the primitive-tiled and the composite variants run on the *same*
|
||||||
|
cycle-accurate engine (fair comparison). Composite is left untouched:
|
||||||
|
the CPU emits ONE CompositeCmd, and PE_SCHEDULER tiles internally (no
|
||||||
|
per-HW-tile CPU dispatch).
|
||||||
|
|
||||||
|
Data correctness:
|
||||||
|
Inputs are ctx.zeros (q/k/v), so every matmul result is zeros and the
|
||||||
|
DataExecutor replay is trivial. To keep replay numerically correct
|
||||||
|
regardless, exactly ONE emitted tile per dot carries the *real* full
|
||||||
|
operands+output handles (so the DataExecutor computes the true (M,N)
|
||||||
|
result via the recorded handle shapes), while its timing fields
|
||||||
|
(m,k,n) are the HW-tile size so the engine charges exactly one tile of
|
||||||
|
cycle time. The remaining n_tiles-1 emitted tiles are timing-only
|
||||||
|
GemmCmds (16x16x16) writing to throwaway scratch. Net: n_tiles tiles
|
||||||
|
of engine time + n_tiles dispatch charges, and a correct final output.
|
||||||
|
|
||||||
|
Engine mode: enable_data=True (same as the production sweep's
|
||||||
|
_engine_latency_ns), op_log end-to-end latency.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from math import ceil
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# --- HW GEMM tile under test -------------------------------------------------
|
||||||
|
MAC_M = 16
|
||||||
|
MAC_K = 16
|
||||||
|
MAC_N = 16
|
||||||
|
# Legacy alt to also try: (8, 16, 32)
|
||||||
|
|
||||||
|
S_KV_LATENCY = (8192, 32_768, 65_536, 131_072)
|
||||||
|
|
||||||
|
ROOT = Path(__file__).resolve().parent
|
||||||
|
SWEEP_JSON = (
|
||||||
|
ROOT / "src" / "kernbench" / "benches" / "1H_milestone_output"
|
||||||
|
/ "gqa" / "long_ctx" / "sweep_decode_composite.json"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# --- mac-dim topology override -----------------------------------------------
|
||||||
|
def _topo_with_mac(mac_m: int, mac_k: int, mac_n: int):
|
||||||
|
"""Compiled topology with mac dims injected into every pe_gemm node."""
|
||||||
|
from kernbench.topology.builder import resolve_topology
|
||||||
|
|
||||||
|
handle = resolve_topology("topology.yaml")
|
||||||
|
g = handle.topology_obj
|
||||||
|
n = 0
|
||||||
|
for node in g.nodes.values():
|
||||||
|
if node.kind == "pe_gemm":
|
||||||
|
node.attrs["mac_m"] = mac_m
|
||||||
|
node.attrs["mac_k"] = mac_k
|
||||||
|
node.attrs["mac_n"] = mac_n
|
||||||
|
n += 1
|
||||||
|
print(f" injected mac=({mac_m},{mac_k},{mac_n}) into {n} pe_gemm nodes")
|
||||||
|
return handle
|
||||||
|
|
||||||
|
|
||||||
|
# --- tiling monkeypatch for TLContext.dot ------------------------------------
|
||||||
|
def _make_tiled_dot(orig_dot, mac_m: int, mac_k: int, mac_n: int):
|
||||||
|
from kernbench.common.pe_commands import GemmCmd
|
||||||
|
|
||||||
|
def tiled_dot(self, a, b):
|
||||||
|
if len(a.shape) < 2 or len(b.shape) < 2:
|
||||||
|
return orig_dot(self, a, b)
|
||||||
|
m, k = a.shape[-2], a.shape[-1]
|
||||||
|
k2, n = b.shape[-2], b.shape[-1]
|
||||||
|
if k != k2:
|
||||||
|
raise ValueError(f"dot shape mismatch: a.K={k} != b.K={k2}")
|
||||||
|
|
||||||
|
n_tiles = ceil(m / mac_m) * ceil(k / mac_k) * ceil(n / mac_n)
|
||||||
|
|
||||||
|
out_shape = (*a.shape[:-2], m, n)
|
||||||
|
out = self._make_compute_out(shape=out_shape, dtype=a.dtype)
|
||||||
|
self._await_pending(a, b)
|
||||||
|
|
||||||
|
if n_tiles <= 1:
|
||||||
|
self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n))
|
||||||
|
return out
|
||||||
|
|
||||||
|
# One real-data tile: full handles (so DataExecutor computes the
|
||||||
|
# true result), but timing fields = HW tile (one tile of cycles).
|
||||||
|
self._emit(GemmCmd(a=a, b=b, out=out, m=mac_m, k=mac_k, n=mac_n))
|
||||||
|
# Remaining timing-only tiles: throwaway scratch, 16x16x16.
|
||||||
|
scratch = self._make_compute_out(shape=(mac_m, mac_n), dtype=a.dtype)
|
||||||
|
for _ in range(n_tiles - 1):
|
||||||
|
self._emit(GemmCmd(a=a, b=b, out=scratch,
|
||||||
|
m=mac_m, k=mac_k, n=mac_n))
|
||||||
|
return out
|
||||||
|
|
||||||
|
return tiled_dot
|
||||||
|
|
||||||
|
|
||||||
|
# --- latency runner (replicates sweep's _engine_latency_ns) ------------------
|
||||||
|
def _engine_latency_ns(variant: str, S_kv: int, topo) -> float:
|
||||||
|
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import ( # noqa: E501
|
||||||
|
_end_to_end_ns, _run_panel_fn,
|
||||||
|
)
|
||||||
|
from kernbench.runtime_api.bench_runner import run_bench
|
||||||
|
from kernbench.runtime_api.types import resolve_device
|
||||||
|
from kernbench.sim_engine.engine import GraphEngine
|
||||||
|
|
||||||
|
result = run_bench(
|
||||||
|
topology=topo, bench_fn=_run_panel_fn(variant, S_kv),
|
||||||
|
device=resolve_device(None),
|
||||||
|
engine_factory=lambda t, d: GraphEngine(
|
||||||
|
getattr(t, "topology_obj", t), enable_data=True,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
if not result.completion.ok:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"{variant}@{S_kv} failed: {result.completion}"
|
||||||
|
)
|
||||||
|
return _end_to_end_ns(result.engine.op_log)
|
||||||
|
|
||||||
|
|
||||||
|
def _emit_dispatch(variant: str, S_kv: int) -> int:
|
||||||
|
"""PE_CPU command count at the center rank (cube 6, pe 0)."""
|
||||||
|
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import ( # noqa: E501
|
||||||
|
_emit_dispatch as prod_emit,
|
||||||
|
)
|
||||||
|
return prod_emit(variant, S_kv)[0]
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
import kernbench.triton_emu.tl_context as tlc
|
||||||
|
|
||||||
|
# Baseline (A): primitive UNTILED latencies from the production sweep
|
||||||
|
# (mac=0 / TFLOPS model). Read straight off the committed sweep JSON.
|
||||||
|
sweep = json.loads(SWEEP_JSON.read_text())
|
||||||
|
base_A = {}
|
||||||
|
for r in sweep["rows"]:
|
||||||
|
if r["variant"] == "primitive" and r["latency_ns"] is not None:
|
||||||
|
base_A[r["S_kv"]] = r["latency_ns"]
|
||||||
|
|
||||||
|
print(f"== mac tile = ({MAC_M},{MAC_K},{MAC_N}) ==")
|
||||||
|
|
||||||
|
# --- command-count sanity (emit-time, mac-independent) ---------------
|
||||||
|
orig_dot = tlc.TLContext.dot
|
||||||
|
print("\n[dispatch counts @ S_kv=131072]")
|
||||||
|
n_prim_untiled = _emit_dispatch("primitive", 131072)
|
||||||
|
tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
|
||||||
|
try:
|
||||||
|
n_prim_tiled = _emit_dispatch("primitive", 131072)
|
||||||
|
n_comp = None
|
||||||
|
finally:
|
||||||
|
tlc.TLContext.dot = orig_dot
|
||||||
|
n_comp = _emit_dispatch("composite", 131072)
|
||||||
|
print(f" primitive UNTILED PE_CPU cmds : {n_prim_untiled}")
|
||||||
|
print(f" primitive TILED PE_CPU cmds : {n_prim_tiled} "
|
||||||
|
f"(x{n_prim_tiled / max(n_prim_untiled,1):.0f})")
|
||||||
|
print(f" composite PE_CPU cmds : {n_comp}")
|
||||||
|
|
||||||
|
# --- latency sweep ----------------------------------------------------
|
||||||
|
rows = []
|
||||||
|
topo = _topo_with_mac(MAC_M, MAC_K, MAC_N)
|
||||||
|
|
||||||
|
for S_kv in S_KV_LATENCY:
|
||||||
|
A = base_A.get(S_kv)
|
||||||
|
|
||||||
|
# (C) composite on the mac engine (untouched dot path)
|
||||||
|
C = _engine_latency_ns("composite", S_kv, topo)
|
||||||
|
|
||||||
|
# (B) primitive TILED on the mac engine
|
||||||
|
tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
|
||||||
|
try:
|
||||||
|
B = _engine_latency_ns("primitive", S_kv, topo)
|
||||||
|
finally:
|
||||||
|
tlc.TLContext.dot = orig_dot
|
||||||
|
|
||||||
|
gap_pct = (B - C) / C * 100.0 if C else float("nan")
|
||||||
|
rows.append((S_kv, A, B, C, gap_pct))
|
||||||
|
print(f" S_kv={S_kv:>7}: A(untiled)={A!s:>12} "
|
||||||
|
f"B(tiled)={B:12.2f} C(comp)={C:12.2f} (B-C)/C={gap_pct:+6.1f}%")
|
||||||
|
|
||||||
|
# --- final table ------------------------------------------------------
|
||||||
|
print("\n==================== RESULT TABLE ====================")
|
||||||
|
print(f"{'S_kv':>8} | {'A untiled(ns)':>14} | {'B tiled(ns)':>14} | "
|
||||||
|
f"{'C comp(ns)':>14} | {'(B-C)/C':>9}")
|
||||||
|
print("-" * 72)
|
||||||
|
for S_kv, A, B, C, gap in rows:
|
||||||
|
a_s = f"{A:.2f}" if A is not None else "n/a"
|
||||||
|
print(f"{S_kv:>8} | {a_s:>14} | {B:>14.2f} | {C:>14.2f} | "
|
||||||
|
f"{gap:>+8.1f}%")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,124 @@
|
|||||||
|
"""Comparative figure for the memory-bound decode-streaming composite study.
|
||||||
|
|
||||||
|
Reads sweep_decode_streaming.json (emitted by milestone-1h-gqa, sweep
|
||||||
|
``decode_streaming``) and writes one two-panel PNG:
|
||||||
|
|
||||||
|
gqa_decode_streaming.png
|
||||||
|
Left — end-to-end single-rank decode latency (µs) vs per-rank context.
|
||||||
|
Right — achieved HBM bandwidth (GB/s) vs context, against the
|
||||||
|
256 GB/s per-rank roofline.
|
||||||
|
|
||||||
|
The memory-bound mirror of the compute-bound prefill figure. With T_q=1
|
||||||
|
the GEMMs are skinny (M=8) and the kernel is bound by streaming the KV
|
||||||
|
cache. Isolating a single rank (no inter-CUBE reduce) reveals what the
|
||||||
|
64-way Case-6 decode masks: the composite command still wins, not by
|
||||||
|
feeding the MAC array but by keeping the DMA pipeline full — its
|
||||||
|
scheduler-streamed concurrent tile DMAs extract ~230 GB/s (near the
|
||||||
|
256 GB/s roofline) while the primitive kernel's blocking tl.dot serializes
|
||||||
|
one tile DMA at a time and plateaus at ~166 GB/s. That bandwidth gap is a
|
||||||
|
~25-28 % latency win that grows nowhere near prefill's compute-bound
|
||||||
|
margin but is decidedly not zero.
|
||||||
|
|
||||||
|
Run (after the bench):
|
||||||
|
GQA_1H_RUN=1 GQA_1H_SWEEPS=decode_streaming python -m kernbench.cli.main \\
|
||||||
|
run --bench milestone-1h-gqa --topology topology.yaml
|
||||||
|
python scripts/paper/paper_plot_gqa_decode_streaming.py
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import matplotlib
|
||||||
|
|
||||||
|
matplotlib.use("Agg")
|
||||||
|
import matplotlib.pyplot as plt # noqa: E402
|
||||||
|
|
||||||
|
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||||
|
_FIG_DIR = (
|
||||||
|
_REPO_ROOT / "src" / "kernbench" / "benches"
|
||||||
|
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||||
|
)
|
||||||
|
_SWEEP_JSON = _FIG_DIR / "sweep_decode_streaming.json"
|
||||||
|
_PAPER_FIG_DIR = (
|
||||||
|
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Per-rank HBM roofline: 8 pseudo-channels × 32 GB/s (topology.yaml
|
||||||
|
# hbm_ctrl.num_pcs / pc_bw_gbs; = pe_dma_to_noc_bw_gbs).
|
||||||
|
_PEAK_HBM_GBS = 256.0
|
||||||
|
|
||||||
|
_VARIANT_STYLE = {
|
||||||
|
"primitive": ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
|
||||||
|
"composite": ("composite GEMM", "#3b6ea5", "s"),
|
||||||
|
"composite_extended": ("composite + softmax_merge", "#4f8a4f", "^"),
|
||||||
|
}
|
||||||
|
_ORDER = ("primitive", "composite", "composite_extended")
|
||||||
|
|
||||||
|
|
||||||
|
def _ctx_label(c: int) -> str:
|
||||||
|
return f"{c // 1024}K" if c >= 1024 else str(c)
|
||||||
|
|
||||||
|
|
||||||
|
def _series(rows, variant, key):
|
||||||
|
pts = sorted(((r["s_kv"], r[key]) for r in rows
|
||||||
|
if r["variant"] == variant), key=lambda t: t[0])
|
||||||
|
return [p[0] for p in pts], [p[1] for p in pts]
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
sweep = json.loads(_SWEEP_JSON.read_text())
|
||||||
|
rows = sweep["rows"]
|
||||||
|
ctxs = sweep["s_kv_points"]
|
||||||
|
|
||||||
|
fig, (ax_lat, ax_bw) = plt.subplots(1, 2, figsize=(13.0, 4.8))
|
||||||
|
|
||||||
|
for v in _ORDER:
|
||||||
|
label, color, marker = _VARIANT_STYLE[v]
|
||||||
|
xs, lat = _series(rows, v, "latency_ns")
|
||||||
|
ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
|
||||||
|
color=color, label=label, lw=2)
|
||||||
|
xs, bw = _series(rows, v, "achieved_bw_gbs")
|
||||||
|
ax_bw.plot(xs, bw, marker=marker, color=color, label=label, lw=2)
|
||||||
|
|
||||||
|
for ax in (ax_lat, ax_bw):
|
||||||
|
ax.set_xscale("log", base=2)
|
||||||
|
ax.set_xticks(ctxs)
|
||||||
|
ax.set_xticklabels([_ctx_label(c) for c in ctxs])
|
||||||
|
ax.set_xlabel(r"per-rank context length $S_{kv}$ ($T_q{=}1$)")
|
||||||
|
ax.grid(True, ls=":", alpha=0.5)
|
||||||
|
ax.legend(fontsize=9)
|
||||||
|
|
||||||
|
ax_lat.set_ylabel("end-to-end decode latency (µs)")
|
||||||
|
ax_lat.set_title("Single-rank memory-bound decode latency per command form")
|
||||||
|
|
||||||
|
ax_bw.set_ylabel("achieved HBM bandwidth (GB/s)")
|
||||||
|
ax_bw.set_title(
|
||||||
|
"HBM bandwidth — composite keeps the DMA pipe full; primitive plateaus"
|
||||||
|
)
|
||||||
|
ax_bw.axhline(_PEAK_HBM_GBS, color="#888", ls="--", lw=1, alpha=0.7)
|
||||||
|
ax_bw.text(ctxs[0], _PEAK_HBM_GBS - 8, "256 GB/s roofline",
|
||||||
|
fontsize=8, color="#555", va="top")
|
||||||
|
ax_bw.set_ylim(0, _PEAK_HBM_GBS * 1.08)
|
||||||
|
|
||||||
|
fig.suptitle(
|
||||||
|
"Memory-bound decode streaming — use of composite commands\n"
|
||||||
|
"single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
|
||||||
|
"{=}128$); $M{=}8$ skinny, KV-streaming-bound",
|
||||||
|
fontsize=11,
|
||||||
|
)
|
||||||
|
fig.tight_layout(rect=(0, 0, 1, 0.92))
|
||||||
|
|
||||||
|
out = _FIG_DIR / "gqa_decode_streaming.png"
|
||||||
|
fig.savefig(out, dpi=150)
|
||||||
|
plt.close(fig)
|
||||||
|
print(f"wrote {out}")
|
||||||
|
|
||||||
|
if _PAPER_FIG_DIR.is_dir():
|
||||||
|
dst = _PAPER_FIG_DIR / out.name
|
||||||
|
dst.write_bytes(out.read_bytes())
|
||||||
|
print(f"copied {dst}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -243,6 +243,12 @@ _FLOW_DATA = "#6699CC" # blue flit (data)
|
|||||||
_FLOW_CTRL = "#ED7D65" # orange flit (control)
|
_FLOW_CTRL = "#ED7D65" # orange flit (control)
|
||||||
_FLOW_CRED = "#7BB661" # green (credit / ack)
|
_FLOW_CRED = "#7BB661" # green (credit / ack)
|
||||||
|
|
||||||
|
# Font-size multiplier for the flow helpers. The stacked figure renders at
|
||||||
|
# full text width (a two-column figure*), so its fonts are scaled up here so
|
||||||
|
# they stay legible after the figure is fit to the page. Default 1.0 leaves
|
||||||
|
# the 2×2 flow figure unchanged.
|
||||||
|
_FS = 1.0
|
||||||
|
|
||||||
|
|
||||||
def _flow_swim(ax, x, y, w, h, label, sub=None):
|
def _flow_swim(ax, x, y, w, h, label, sub=None):
|
||||||
ax.add_patch(mpatches.FancyBboxPatch(
|
ax.add_patch(mpatches.FancyBboxPatch(
|
||||||
@@ -251,11 +257,11 @@ def _flow_swim(ax, x, y, w, h, label, sub=None):
|
|||||||
facecolor="white", edgecolor="#3a3a3a", linewidth=1.6))
|
facecolor="white", edgecolor="#3a3a3a", linewidth=1.6))
|
||||||
ax.text(x + w / 2, y + h / 2 + (0.20 if sub else 0),
|
ax.text(x + w / 2, y + h / 2 + (0.20 if sub else 0),
|
||||||
label, ha="center", va="center",
|
label, ha="center", va="center",
|
||||||
fontsize=10.5, fontweight="bold", color="#111")
|
fontsize=10.5 * _FS, fontweight="bold", color="#111")
|
||||||
if sub:
|
if sub:
|
||||||
ax.text(x + w / 2, y + h / 2 - 0.35, sub,
|
ax.text(x + w / 2, y + h / 2 - 0.35, sub,
|
||||||
ha="center", va="center",
|
ha="center", va="center",
|
||||||
fontsize=8, color="#555", fontstyle="italic")
|
fontsize=8 * _FS, color="#555", fontstyle="italic")
|
||||||
|
|
||||||
|
|
||||||
def _flow_flit(ax, x, y, color, letter="", size=0.32):
|
def _flow_flit(ax, x, y, color, letter="", size=0.32):
|
||||||
@@ -265,7 +271,7 @@ def _flow_flit(ax, x, y, color, letter="", size=0.32):
|
|||||||
facecolor=color, edgecolor="#222", linewidth=0.8))
|
facecolor=color, edgecolor="#222", linewidth=0.8))
|
||||||
if letter:
|
if letter:
|
||||||
ax.text(x, y, letter, ha="center", va="center",
|
ax.text(x, y, letter, ha="center", va="center",
|
||||||
fontsize=8.5, fontweight="bold", color="white")
|
fontsize=8.5 * _FS, fontweight="bold", color="white")
|
||||||
|
|
||||||
|
|
||||||
def _flow_wire(ax, x1, x2, y, color="#888"):
|
def _flow_wire(ax, x1, x2, y, color="#888"):
|
||||||
@@ -280,7 +286,7 @@ def _flow_dot(ax, x, y, color):
|
|||||||
|
|
||||||
def _flow_label(ax, x, y, text, color="#222", fontsize=9):
|
def _flow_label(ax, x, y, text, color="#222", fontsize=9):
|
||||||
ax.text(x, y, text, ha="center", va="center",
|
ax.text(x, y, text, ha="center", va="center",
|
||||||
fontsize=fontsize, color=color, fontweight="bold")
|
fontsize=fontsize * _FS, color=color, fontweight="bold")
|
||||||
|
|
||||||
|
|
||||||
def _flow_lead(ax, x_tail, y_tail, x_head, y_head, color="#777"):
|
def _flow_lead(ax, x_tail, y_tail, x_head, y_head, color="#777"):
|
||||||
@@ -297,7 +303,7 @@ def _flow_setup(ax, kind):
|
|||||||
facecolor="white", edgecolor=_COLOR[kind],
|
facecolor="white", edgecolor=_COLOR[kind],
|
||||||
linewidth=0.9, alpha=0.85, zorder=0))
|
linewidth=0.9, alpha=0.85, zorder=0))
|
||||||
ax.text(10, 0.65, _TITLE[kind],
|
ax.text(10, 0.65, _TITLE[kind],
|
||||||
ha="center", va="center", fontsize=11.5,
|
ha="center", va="center", fontsize=11.5 * _FS,
|
||||||
fontweight="bold",
|
fontweight="bold",
|
||||||
color=("#1B5E20" if kind == "ipcq" else "white"),
|
color=("#1B5E20" if kind == "ipcq" else "white"),
|
||||||
bbox=dict(
|
bbox=dict(
|
||||||
@@ -449,7 +455,7 @@ def _flow_legend(fig, leg_y=0.02):
|
|||||||
facecolor=color, edgecolor="#222", linewidth=0.6,
|
facecolor=color, edgecolor="#222", linewidth=0.6,
|
||||||
transform=fig.transFigure))
|
transform=fig.transFigure))
|
||||||
fig.text(lx + 0.025, y + 0.011, text,
|
fig.text(lx + 0.025, y + 0.011, text,
|
||||||
ha="left", va="center", fontsize=10,
|
ha="left", va="center", fontsize=10 * _FS,
|
||||||
fontweight="bold", color="#222")
|
fontweight="bold", color="#222")
|
||||||
lx += 0.32
|
lx += 0.32
|
||||||
|
|
||||||
@@ -461,21 +467,27 @@ def _flow_legend(fig, leg_y=0.02):
|
|||||||
def _plot_architecture_flow() -> Path:
|
def _plot_architecture_flow() -> Path:
|
||||||
"""4 panels arranged 2×2 — each a horizontal Sender → NoC → Receiver
|
"""4 panels arranged 2×2 — each a horizontal Sender → NoC → Receiver
|
||||||
flow with small flit-style packets, dot-marker annotations and a
|
flow with small flit-style packets, dot-marker annotations and a
|
||||||
bottom name plate. Aesthetic mirrors latency_model.png."""
|
bottom name plate. Aesthetic mirrors latency_model.png. Fonts scaled
|
||||||
fig, axes = plt.subplots(2, 2, figsize=(18.0, 11.5))
|
up via _FS so the figure stays legible when fit to the page."""
|
||||||
_draw_flow_case(axes[0, 0], "doorbell")
|
global _FS
|
||||||
_draw_flow_case(axes[0, 1], "hmq")
|
_FS = 1.5
|
||||||
_draw_flow_case(axes[1, 0], "rdma")
|
try:
|
||||||
_draw_flow_case(axes[1, 1], "ipcq")
|
fig, axes = plt.subplots(2, 2, figsize=(18.0, 11.5))
|
||||||
fig.suptitle(
|
_draw_flow_case(axes[0, 0], "doorbell")
|
||||||
"Per-send architecture = $\\Sigma$ control overhead + "
|
_draw_flow_case(axes[0, 1], "hmq")
|
||||||
"data-path DMA + receiver wake",
|
_draw_flow_case(axes[1, 0], "rdma")
|
||||||
fontsize=14.5, fontweight="bold", y=0.995)
|
_draw_flow_case(axes[1, 1], "ipcq")
|
||||||
_flow_legend(fig, leg_y=0.015)
|
fig.suptitle(
|
||||||
fig.tight_layout(rect=(0, 0.07, 1, 0.96))
|
"Per-send architecture = $\\Sigma$ control overhead + "
|
||||||
out = _OUT_DIR / "ipcq_alternatives_architecture_flow.png"
|
"data-path DMA + receiver wake",
|
||||||
fig.savefig(out, dpi=150, bbox_inches="tight")
|
fontsize=14.5 * _FS, fontweight="bold", y=0.995)
|
||||||
plt.close(fig)
|
_flow_legend(fig, leg_y=0.015)
|
||||||
|
fig.tight_layout(rect=(0, 0.07, 1, 0.96))
|
||||||
|
out = _OUT_DIR / "ipcq_alternatives_architecture_flow.png"
|
||||||
|
fig.savefig(out, dpi=150, bbox_inches="tight")
|
||||||
|
plt.close(fig)
|
||||||
|
finally:
|
||||||
|
_FS = 1.0
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
|
||||||
@@ -483,19 +495,25 @@ def _plot_architecture_flow() -> Path:
|
|||||||
def _plot_architecture_stacked() -> Path:
|
def _plot_architecture_stacked() -> Path:
|
||||||
"""Same per-case flow content as the 2×2 figure, but stacked one
|
"""Same per-case flow content as the 2×2 figure, but stacked one
|
||||||
case per row (4 rows × 1 column). Wider panels give each case more
|
case per row (4 rows × 1 column). Wider panels give each case more
|
||||||
horizontal room."""
|
horizontal room. Rendered at full text width (two-column figure*), so
|
||||||
fig, axes = plt.subplots(4, 1, figsize=(16.0, 19.0))
|
fonts are scaled up via _FS to stay legible after page fitting."""
|
||||||
for ax, kind in zip(axes, ["doorbell", "hmq", "rdma", "ipcq"]):
|
global _FS
|
||||||
_draw_flow_case(ax, kind)
|
_FS = 1.9
|
||||||
fig.suptitle(
|
try:
|
||||||
"Per-send architecture = $\\Sigma$ control overhead + "
|
fig, axes = plt.subplots(4, 1, figsize=(13.0, 15.0))
|
||||||
"data-path DMA + receiver wake",
|
for ax, kind in zip(axes, ["doorbell", "hmq", "rdma", "ipcq"]):
|
||||||
fontsize=14.5, fontweight="bold", y=0.997)
|
_draw_flow_case(ax, kind)
|
||||||
_flow_legend(fig, leg_y=0.010)
|
fig.suptitle(
|
||||||
fig.tight_layout(rect=(0, 0.05, 1, 0.97))
|
"Per-send architecture = $\\Sigma$ control overhead + "
|
||||||
out = _OUT_DIR / "ipcq_alternatives_architecture_stacked.png"
|
"data-path DMA + receiver wake",
|
||||||
fig.savefig(out, dpi=150, bbox_inches="tight")
|
fontsize=14.5 * _FS, fontweight="bold", y=0.997)
|
||||||
plt.close(fig)
|
_flow_legend(fig, leg_y=0.010)
|
||||||
|
fig.tight_layout(rect=(0, 0.06, 1, 0.97))
|
||||||
|
out = _OUT_DIR / "ipcq_alternatives_architecture_stacked.png"
|
||||||
|
fig.savefig(out, dpi=150, bbox_inches="tight")
|
||||||
|
plt.close(fig)
|
||||||
|
finally:
|
||||||
|
_FS = 1.0
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
|
Before Width: | Height: | Size: 253 KiB After Width: | Height: | Size: 252 KiB |
|
Before Width: | Height: | Size: 230 KiB After Width: | Height: | Size: 305 KiB |
|
Before Width: | Height: | Size: 260 KiB After Width: | Height: | Size: 381 KiB |
|
Before Width: | Height: | Size: 153 KiB After Width: | Height: | Size: 153 KiB |
|
Before Width: | Height: | Size: 189 KiB After Width: | Height: | Size: 189 KiB |
|
Before Width: | Height: | Size: 160 KiB After Width: | Height: | Size: 160 KiB |
|
Before Width: | Height: | Size: 87 KiB After Width: | Height: | Size: 86 KiB |
|
Before Width: | Height: | Size: 209 KiB After Width: | Height: | Size: 209 KiB |
|
After Width: | Height: | Size: 136 KiB |
@@ -0,0 +1,172 @@
|
|||||||
|
{
|
||||||
|
"version": 1,
|
||||||
|
"variants": [
|
||||||
|
"primitive",
|
||||||
|
"composite",
|
||||||
|
"composite_extended"
|
||||||
|
],
|
||||||
|
"s_kv_points": [
|
||||||
|
2048,
|
||||||
|
4096,
|
||||||
|
8192,
|
||||||
|
16384
|
||||||
|
],
|
||||||
|
"rows": [
|
||||||
|
{
|
||||||
|
"variant": "primitive",
|
||||||
|
"s_kv": 2048,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 6188.437999999816,
|
||||||
|
"gemm_busy_ns": 1048.576000000001,
|
||||||
|
"dma_busy_ns": 6322.0,
|
||||||
|
"kv_bytes": 1048576.0,
|
||||||
|
"achieved_bw_gbs": 169.44114169036374,
|
||||||
|
"achieved_tflops": 1.3555291335229098,
|
||||||
|
"mac_util": 0.16944114169036373,
|
||||||
|
"dma_occupancy": 1.021582505957107
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite",
|
||||||
|
"s_kv": 2048,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 4836.115999999989,
|
||||||
|
"gemm_busy_ns": 6629.632000000123,
|
||||||
|
"dma_busy_ns": 267480.720000013,
|
||||||
|
"kv_bytes": 1048576.0,
|
||||||
|
"achieved_bw_gbs": 216.82192900253062,
|
||||||
|
"achieved_tflops": 1.734575432020245,
|
||||||
|
"mac_util": 0.2168219290025306,
|
||||||
|
"dma_occupancy": 55.30899589671001
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite_extended",
|
||||||
|
"s_kv": 2048,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 4737.751999999986,
|
||||||
|
"gemm_busy_ns": 8481.408000000116,
|
||||||
|
"dma_busy_ns": 251398.7400000122,
|
||||||
|
"kv_bytes": 1048576.0,
|
||||||
|
"achieved_bw_gbs": 221.32353065335693,
|
||||||
|
"achieved_tflops": 1.7705882452268553,
|
||||||
|
"mac_util": 0.22132353065335691,
|
||||||
|
"dma_occupancy": 53.06287454472352
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "primitive",
|
||||||
|
"s_kv": 4096,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 12510.006000000496,
|
||||||
|
"gemm_busy_ns": 2097.152000000002,
|
||||||
|
"dma_busy_ns": 12602.0,
|
||||||
|
"kv_bytes": 2097152.0,
|
||||||
|
"achieved_bw_gbs": 167.6379691584414,
|
||||||
|
"achieved_tflops": 1.3411037532675312,
|
||||||
|
"mac_util": 0.1676379691584414,
|
||||||
|
"dma_occupancy": 1.0073536335633653
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite",
|
||||||
|
"s_kv": 4096,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 9360.883999999554,
|
||||||
|
"gemm_busy_ns": 19520.799999939867,
|
||||||
|
"dma_busy_ns": 1059043.5999999335,
|
||||||
|
"kv_bytes": 2097152.0,
|
||||||
|
"achieved_bw_gbs": 224.03354213128802,
|
||||||
|
"achieved_tflops": 1.792268337050304,
|
||||||
|
"mac_util": 0.224033542131288,
|
||||||
|
"dma_occupancy": 113.13499878857424
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite_extended",
|
||||||
|
"s_kv": 4096,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 9198.135999999562,
|
||||||
|
"gemm_busy_ns": 39531.583999941715,
|
||||||
|
"dma_busy_ns": 1026582.7399999356,
|
||||||
|
"kv_bytes": 2097152.0,
|
||||||
|
"achieved_bw_gbs": 227.99749862364504,
|
||||||
|
"achieved_tflops": 1.8239799889891604,
|
||||||
|
"mac_util": 0.22799749862364505,
|
||||||
|
"dma_occupancy": 111.60769312390951
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "primitive",
|
||||||
|
"s_kv": 8192,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 25153.141999997388,
|
||||||
|
"gemm_busy_ns": 4194.304000000004,
|
||||||
|
"dma_busy_ns": 25162.0,
|
||||||
|
"kv_bytes": 4194304.0,
|
||||||
|
"achieved_bw_gbs": 166.7506985807354,
|
||||||
|
"achieved_tflops": 1.334005588645883,
|
||||||
|
"mac_util": 0.16675069858073538,
|
||||||
|
"dma_occupancy": 1.0003521627637062
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite",
|
||||||
|
"s_kv": 8192,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 18419.187999999034,
|
||||||
|
"gemm_busy_ns": 64177.11999976149,
|
||||||
|
"dma_busy_ns": 4214541.840000685,
|
||||||
|
"kv_bytes": 4194304.0,
|
||||||
|
"achieved_bw_gbs": 227.713838416776,
|
||||||
|
"achieved_tflops": 1.821710707334208,
|
||||||
|
"mac_util": 0.227713838416776,
|
||||||
|
"dma_occupancy": 228.81257523409317
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite_extended",
|
||||||
|
"s_kv": 8192,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 18128.439999999027,
|
||||||
|
"gemm_busy_ns": 169650.68799976518,
|
||||||
|
"dma_busy_ns": 4149323.2200006745,
|
||||||
|
"kv_bytes": 4194304.0,
|
||||||
|
"achieved_bw_gbs": 231.365964197704,
|
||||||
|
"achieved_tflops": 1.850927713581632,
|
||||||
|
"mac_util": 0.231365964197704,
|
||||||
|
"dma_occupancy": 228.88473691067168
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "primitive",
|
||||||
|
"s_kv": 16384,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 50439.414000008954,
|
||||||
|
"gemm_busy_ns": 8388.607999999076,
|
||||||
|
"dma_busy_ns": 50282.0,
|
||||||
|
"kv_bytes": 8388608.0,
|
||||||
|
"achieved_bw_gbs": 166.31057609032712,
|
||||||
|
"achieved_tflops": 1.330484608722617,
|
||||||
|
"mac_util": 0.16631057609032712,
|
||||||
|
"dma_occupancy": 0.9968791469304357
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite",
|
||||||
|
"s_kv": 16384,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 36535.79600000568,
|
||||||
|
"gemm_busy_ns": 228978.46400288073,
|
||||||
|
"dma_busy_ns": 16815028.239995122,
|
||||||
|
"kv_bytes": 8388608.0,
|
||||||
|
"achieved_bw_gbs": 229.59970545047648,
|
||||||
|
"achieved_tflops": 1.8367976436038118,
|
||||||
|
"mac_util": 0.22959970545047648,
|
||||||
|
"dma_occupancy": 460.2343477063565
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"variant": "composite_extended",
|
||||||
|
"s_kv": 16384,
|
||||||
|
"M": 8,
|
||||||
|
"latency_ns": 35989.048000005685,
|
||||||
|
"gemm_busy_ns": 701990.3680028584,
|
||||||
|
"dma_busy_ns": 16684294.09999516,
|
||||||
|
"kv_bytes": 8388608.0,
|
||||||
|
"achieved_bw_gbs": 233.0877993771515,
|
||||||
|
"achieved_tflops": 1.864702395017212,
|
||||||
|
"mac_util": 0.2330877993771515,
|
||||||
|
"dma_occupancy": 463.5936493789034
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
@@ -0,0 +1,129 @@
|
|||||||
|
"""GQA decode kernel — Case 6, **primitive-TILED** (16×16×16 MAC blocking).
|
||||||
|
|
||||||
|
Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
|
||||||
|
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the only difference is
|
||||||
|
that each local-attention matmul is *hand-blocked into 16×16×16 GemmCmds*
|
||||||
|
(mac=16) instead of one coarse ``tl.dot`` per tile. This models a kernel
|
||||||
|
that issues the MAC-array fan-out from PE_CPU itself: the per-block GemmCmd
|
||||||
|
count is ``ceil(M/16)·ceil(K/16)·ceil(N/16)`` per matmul, charging the full
|
||||||
|
ADR-0064 dispatch cost for every block — the "dispatch explosion" the
|
||||||
|
composite form offloads to PE_SCHEDULER.
|
||||||
|
|
||||||
|
FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute of
|
||||||
|
its block; the blocks sum to the full matmul), so end-to-end compute time
|
||||||
|
is unchanged vs the coarse primitive — only the PE_CPU command count and
|
||||||
|
its dispatch cycles grow. Inputs are zero (decode bench convention), so the
|
||||||
|
blocked accumulation is identically zero; the kernel returns a single
|
||||||
|
zeroed ``(M, N)`` output handle that the downstream softmax consumes — no
|
||||||
|
per-block accumulation handle needed for this zero-input study.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from math import ceil
|
||||||
|
|
||||||
|
from kernbench.common.pe_commands import GemmCmd
|
||||||
|
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
|
||||||
|
_ROOT_CUBE,
|
||||||
|
_merge_running,
|
||||||
|
reduce_mlo,
|
||||||
|
)
|
||||||
|
|
||||||
|
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||||
|
MAC = 16 # 16×16×16 MAC-array blocking granularity.
|
||||||
|
|
||||||
|
|
||||||
|
def _blocked_dot(A, B, *, tl):
|
||||||
|
"""``A @ B`` issued as one GemmCmd per 16×16×16 block.
|
||||||
|
|
||||||
|
``A``:(M, K), ``B``:(K, N) → out:(M, N). Emits
|
||||||
|
``ceil(M/16)·ceil(K/16)·ceil(N/16)`` GemmCmds (each charged the full
|
||||||
|
ADR-0064 dispatch cost via ``tl.dot``'s emit path). **All blocks write
|
||||||
|
the same ``(M, N)`` output handle** (``out``), and that handle is
|
||||||
|
returned — so downstream softmax ops depend on it exactly like the
|
||||||
|
coarse ``tl.dot`` path (the engine tracks the producer by output handle
|
||||||
|
id, and ``GemmCmd`` is a blocking PE_CPU command, so the block GEMMs
|
||||||
|
serialize on the critical path ahead of the consuming ``tl.max``/
|
||||||
|
``tl.exp``). K is innermost so each (mi, ni) output tile accumulates
|
||||||
|
across the K blocks into the same handle.
|
||||||
|
"""
|
||||||
|
M, K = A.shape[-2], A.shape[-1]
|
||||||
|
K2, N = B.shape[-2], B.shape[-1]
|
||||||
|
assert K == K2, f"blocked_dot shape mismatch K={K} != {K2}"
|
||||||
|
out = tl._make_compute_out(shape=(M, N), dtype="f16")
|
||||||
|
# One GemmCmd per 16³ block, all writing the shared `out` handle. A/B
|
||||||
|
# reuse the full operand handles at block dims (m,k,n = 16, or the
|
||||||
|
# ragged tail) — operands are TCM-resident (pinned), so this charges
|
||||||
|
# dispatch + the block's TFLOPS-model compute. K innermost = accumulate
|
||||||
|
# the K blocks into each (mi, ni) output tile of the shared handle.
|
||||||
|
for _mi in range(ceil(M / MAC)):
|
||||||
|
bm = min(MAC, M - _mi * MAC)
|
||||||
|
for _ni in range(ceil(N / MAC)):
|
||||||
|
bn = min(MAC, N - _ni * MAC)
|
||||||
|
for _ki in range(ceil(K / MAC)):
|
||||||
|
bk = min(MAC, K - _ki * MAC)
|
||||||
|
tl._emit(GemmCmd(a=A, b=B, out=out, m=bm, k=bk, n=bn))
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled_kernel(
|
||||||
|
q_ptr: int,
|
||||||
|
k_ptr: int,
|
||||||
|
v_ptr: int,
|
||||||
|
o_ptr: int,
|
||||||
|
T_q: int,
|
||||||
|
S_kv: int,
|
||||||
|
h_q: int,
|
||||||
|
h_kv: int,
|
||||||
|
d_head: int,
|
||||||
|
C: int,
|
||||||
|
P: int,
|
||||||
|
*,
|
||||||
|
tl,
|
||||||
|
) -> None:
|
||||||
|
"""Case-6 decode, primitive-TILED (16×16×16 GemmCmd blocking)."""
|
||||||
|
G = h_q // h_kv
|
||||||
|
n_ranks = C * P
|
||||||
|
S_local = S_kv // n_ranks
|
||||||
|
pe_id = tl.program_id(axis=0)
|
||||||
|
cube_id = tl.program_id(axis=1)
|
||||||
|
|
||||||
|
Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
|
||||||
|
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||||
|
KV_ROW_BYTES = d_head * 2 # f16
|
||||||
|
|
||||||
|
tile_s0 = min(TILE_S_KV, S_local)
|
||||||
|
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
|
||||||
|
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
|
||||||
|
scores = _blocked_dot(Q, K_T, tl=tl)
|
||||||
|
m_local = tl.max(scores, axis=-1)
|
||||||
|
centered = scores - m_local
|
||||||
|
exp_scores = tl.exp(centered)
|
||||||
|
l_local = tl.sum(exp_scores, axis=-1)
|
||||||
|
O_local = _blocked_dot(exp_scores, V, tl=tl)
|
||||||
|
|
||||||
|
for tile_idx in range(1, n_tiles):
|
||||||
|
tile_start = tile_idx * TILE_S_KV
|
||||||
|
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||||
|
with tl.scratch_scope():
|
||||||
|
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||||
|
shape=(d_head, tile_s), dtype="f16")
|
||||||
|
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||||
|
shape=(tile_s, d_head), dtype="f16")
|
||||||
|
scores_t = _blocked_dot(Q, K_T_t, tl=tl)
|
||||||
|
m_tile = tl.max(scores_t, axis=-1)
|
||||||
|
centered_t = scores_t - m_tile
|
||||||
|
exp_scores_t = tl.exp(centered_t)
|
||||||
|
l_tile = tl.sum(exp_scores_t, axis=-1)
|
||||||
|
O_tile = _blocked_dot(exp_scores_t, V_t, tl=tl)
|
||||||
|
m_new, l_new, O_new = _merge_running(
|
||||||
|
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||||
|
)
|
||||||
|
tl.copy_to(m_local, m_new)
|
||||||
|
tl.copy_to(l_local, l_new)
|
||||||
|
tl.copy_to(O_local, O_new)
|
||||||
|
|
||||||
|
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
|
||||||
|
|
||||||
|
if pe_id == 0 and cube_id == _ROOT_CUBE:
|
||||||
|
O_final = O_local / l_local
|
||||||
|
tl.store(o_ptr, O_final)
|
||||||
@@ -0,0 +1,168 @@
|
|||||||
|
"""milestone-1h-gqa: memory-bound decode streaming composite-command study.
|
||||||
|
|
||||||
|
Single-rank companion to the compute-bound prefill study
|
||||||
|
(``gqa_prefill_compute_bound``). Same three command-form kernels
|
||||||
|
(``_gqa_prefill_compute_bound``), same single-rank (C=P=1) harness, but
|
||||||
|
driven with **T_q=1** (decode) and swept over a large S_kv. With T_q=1 the
|
||||||
|
score / context GEMMs are skinny (M = G·T_q = 8), so the MAC array is
|
||||||
|
barely fed and the kernel is **memory-bound** — streaming the KV cache out
|
||||||
|
of HBM dominates. This is the regime mirror of prefill: command form is
|
||||||
|
*latency-neutral* here, because the bottleneck is data movement, not
|
||||||
|
issue.
|
||||||
|
|
||||||
|
Running single-rank (no cross-CUBE (m,ℓ,O) reduce) isolates the
|
||||||
|
local-attention streaming vs. issue trade-off cleanly — unlike the 64-way
|
||||||
|
Case-6 decode sweep, whose small-S_kv end-to-end latency is masked by the
|
||||||
|
inter-CUBE reduce tail. S_kv here is the *per-rank* context, so S_kv=64K
|
||||||
|
single-rank is the per-PE load of a 4M-token, 64-way-sharded decode.
|
||||||
|
|
||||||
|
Records per (variant, S_kv) the end-to-end latency, GEMM/DMA engine busy
|
||||||
|
time, MAC utilization (achieved ÷ peak), and DMA occupancy (dma_busy ÷
|
||||||
|
e2e), so the comparative plot can show all three command forms landing on
|
||||||
|
the same latency curve while the MAC array stays floored and the DMA
|
||||||
|
channel stays saturated.
|
||||||
|
|
||||||
|
Runs in data mode (engine latency). Gated via the umbrella
|
||||||
|
``GQA_1H_SWEEPS=decode_streaming``.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
|
||||||
|
gqa_prefill_composite_ext_kernel,
|
||||||
|
gqa_prefill_composite_kernel,
|
||||||
|
gqa_prefill_primitive_kernel,
|
||||||
|
)
|
||||||
|
from kernbench.policy.placement.dp import DPPolicy
|
||||||
|
|
||||||
|
_OUTPUT_DIR = (
|
||||||
|
Path(__file__).resolve().parents[2]
|
||||||
|
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||||
|
)
|
||||||
|
_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_streaming.json"
|
||||||
|
|
||||||
|
# Same kernels as the prefill study — they are general attention kernels;
|
||||||
|
# the regime is set by T_q (=1 here → memory-bound).
|
||||||
|
_VARIANT_KERNELS = {
|
||||||
|
"primitive": gqa_prefill_primitive_kernel,
|
||||||
|
"composite": gqa_prefill_composite_kernel,
|
||||||
|
"composite_extended": gqa_prefill_composite_ext_kernel,
|
||||||
|
}
|
||||||
|
_VARIANTS = ("primitive", "composite", "composite_extended")
|
||||||
|
|
||||||
|
# Per-rank context length. T_q=1, so M = G·T_q = 8 (skinny, memory-bound)
|
||||||
|
# at every point. Capped at 16K: the *plain* composite materializes the
|
||||||
|
# full (M, S_kv) scores in TCM scratch (~48·S_kv B incl. the exp transients),
|
||||||
|
# which overflows the 1 MB kernel scratch beyond ~21K — itself a sign that
|
||||||
|
# the recipe (composite_extended), which tiles the softmax, is what long
|
||||||
|
# context actually needs. S_kv is per-rank, so 16K single-rank already
|
||||||
|
# equals the per-PE load of a 1M-token, 64-way-sharded decode.
|
||||||
|
_S_KV_POINTS = (2048, 4096, 8192, 16384)
|
||||||
|
|
||||||
|
_T_Q = 1
|
||||||
|
_H_Q, _H_KV, _D_HEAD = 8, 1, 128
|
||||||
|
_PEAK_TFLOPS = 8.0 # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)
|
||||||
|
|
||||||
|
|
||||||
|
def _run_panel_fn(variant: str, s_kv: int):
|
||||||
|
kernel = _VARIANT_KERNELS[variant]
|
||||||
|
panel = f"decode_stream_{variant}_s{s_kv}"
|
||||||
|
|
||||||
|
def _bench_fn(ctx):
|
||||||
|
dp = DPPolicy(cube="replicate", pe="replicate",
|
||||||
|
num_cubes=1, num_pes=1)
|
||||||
|
q = ctx.zeros((_T_Q, _H_Q * _D_HEAD),
|
||||||
|
dtype="f16", dp=dp, name=f"{panel}_q")
|
||||||
|
k = ctx.zeros((s_kv, _H_KV * _D_HEAD),
|
||||||
|
dtype="f16", dp=dp, name=f"{panel}_k")
|
||||||
|
v = ctx.zeros((s_kv, _H_KV * _D_HEAD),
|
||||||
|
dtype="f16", dp=dp, name=f"{panel}_v")
|
||||||
|
o = ctx.empty((_T_Q, _H_Q * _D_HEAD),
|
||||||
|
dtype="f16", dp=dp, name=f"{panel}_o")
|
||||||
|
ctx.launch(panel, kernel, q, k, v, o,
|
||||||
|
_T_Q, s_kv, _H_Q, _H_KV, _D_HEAD, 1, 1,
|
||||||
|
_auto_dim_remap=False)
|
||||||
|
|
||||||
|
return _bench_fn
|
||||||
|
|
||||||
|
|
||||||
|
def _end_to_end_ns(op_log) -> float:
|
||||||
|
if not op_log:
|
||||||
|
return 0.0
|
||||||
|
return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
|
||||||
|
|
||||||
|
|
||||||
|
def _engine_busy_ns(op_log, suffix: str) -> float:
|
||||||
|
return sum(r.t_end - r.t_start
|
||||||
|
for r in op_log if r.component_id.endswith("." + suffix))
|
||||||
|
|
||||||
|
|
||||||
|
def _run_panel(variant: str, s_kv: int, topology: str) -> dict:
|
||||||
|
from kernbench.runtime_api.bench_runner import run_bench
|
||||||
|
from kernbench.runtime_api.types import resolve_device
|
||||||
|
from kernbench.sim_engine.engine import GraphEngine
|
||||||
|
from kernbench.topology.builder import resolve_topology
|
||||||
|
|
||||||
|
topo = resolve_topology(topology)
|
||||||
|
result = run_bench(
|
||||||
|
topology=topo, bench_fn=_run_panel_fn(variant, s_kv),
|
||||||
|
device=resolve_device(None),
|
||||||
|
engine_factory=lambda t, d: GraphEngine(
|
||||||
|
getattr(t, "topology_obj", t), enable_data=True,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
if not result.completion.ok:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"gqa-decode-streaming {variant}@{s_kv} failed: {result.completion}"
|
||||||
|
)
|
||||||
|
op_log = result.engine.op_log
|
||||||
|
e2e = _end_to_end_ns(op_log)
|
||||||
|
gemm = _engine_busy_ns(op_log, "pe_gemm")
|
||||||
|
dma = _engine_busy_ns(op_log, "pe_dma")
|
||||||
|
G = _H_Q // _H_KV
|
||||||
|
M = G * _T_Q
|
||||||
|
# Useful attention flops (Q·Kᵀ + P·V), single rank.
|
||||||
|
useful_flops = 4.0 * M * _D_HEAD * s_kv
|
||||||
|
# KV-cache bytes streamed from HBM (K + V, f16) — the memory-bound
|
||||||
|
# denominator. achieved_bw = bytes / e2e (GB/s) measures how close the
|
||||||
|
# command form gets to the HBM roofline (the memory-bound mirror of
|
||||||
|
# MAC utilization in the compute-bound prefill study).
|
||||||
|
kv_bytes = 2.0 * s_kv * _D_HEAD * 2.0
|
||||||
|
return {
|
||||||
|
"variant": variant,
|
||||||
|
"s_kv": s_kv,
|
||||||
|
"M": M,
|
||||||
|
"latency_ns": e2e,
|
||||||
|
"gemm_busy_ns": gemm,
|
||||||
|
"dma_busy_ns": dma,
|
||||||
|
"kv_bytes": kv_bytes,
|
||||||
|
"achieved_bw_gbs": (kv_bytes / e2e) if e2e > 0 else 0.0,
|
||||||
|
"achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
|
||||||
|
"mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
|
||||||
|
"dma_occupancy": (dma / e2e) if e2e > 0 else 0.0,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def run_sweep(topology: str = "topology.yaml") -> int:
|
||||||
|
"""Drive all (variant, S_kv) decode-streaming panels; write sweep.json."""
|
||||||
|
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||||
|
rows = [
|
||||||
|
_run_panel(variant, s_kv, topology)
|
||||||
|
for s_kv in _S_KV_POINTS
|
||||||
|
for variant in _VARIANTS
|
||||||
|
]
|
||||||
|
sweep = {
|
||||||
|
"version": 1,
|
||||||
|
"variants": list(_VARIANTS),
|
||||||
|
"s_kv_points": list(_S_KV_POINTS),
|
||||||
|
"rows": rows,
|
||||||
|
}
|
||||||
|
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
|
||||||
|
print(f" gqa-decode-streaming: {len(rows)} rows -> {_SWEEP_JSON}")
|
||||||
|
return len(rows)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
run_sweep()
|
||||||
@@ -28,6 +28,9 @@ from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
|
|||||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
|
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
|
||||||
run_sweep as _run_composite_sweep,
|
run_sweep as _run_composite_sweep,
|
||||||
)
|
)
|
||||||
|
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_streaming import (
|
||||||
|
run_sweep as _run_decode_streaming_sweep,
|
||||||
|
)
|
||||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
|
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
|
||||||
run_sweep as _run_prefill_cb_sweep,
|
run_sweep as _run_prefill_cb_sweep,
|
||||||
)
|
)
|
||||||
@@ -66,6 +69,7 @@ def run(torch) -> None:
|
|||||||
"decode": _run_decode_sweep,
|
"decode": _run_decode_sweep,
|
||||||
"composite": _run_composite_sweep,
|
"composite": _run_composite_sweep,
|
||||||
"prefill_cb": _run_prefill_cb_sweep,
|
"prefill_cb": _run_prefill_cb_sweep,
|
||||||
|
"decode_streaming": _run_decode_streaming_sweep,
|
||||||
}
|
}
|
||||||
unknown = [s for s in sweeps if s not in runners]
|
unknown = [s for s in sweeps if s not in runners]
|
||||||
if unknown:
|
if unknown:
|
||||||
|
|||||||