ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa(decode): add d_head-TP kernels + measurement runner + figure generators

Browse Source 
Two new long-ctx decode attention kernels for the d_head-TP sharding
variants the 6-case chart predicts:

  · _gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py
        Cube-SP × PE-TP-d_head (Case 4 in chart). Per cube holds
        S_kv/C tokens of full d_head; per PE holds same tokens but
        only d_head/P dims. Partial Q·Kᵀ scores reduced intra-cube
        before softmax; outer (m,ℓ,O) merge two-phase (intra+inter).

  · _gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py
        Cube-TP-d_head × PE-SP (Case 5). Per cube holds full S_kv
        with only d_head/C dims; per PE holds S_kv/P of those dims.
        Partial scores reduced inter-cube (UCIe) before softmax.

Sweep dispatch (gqa_decode_long_ctx_4cases.py) extended with two
new panels so the milestone-1h-gqa sweep covers all 6 cases.

Smoke test scripts/verify_case4_dhead_tp.py runs Cases 4/5/6 at
S_kv=2K to validate the kernels load and execute end-to-end.

Plus the figure-generation toolchain that produced the committed
PNGs in the prior commit (dd3337f):

  · paper_plot_gqa_4cases_summary.py    - 3-panel summary +
        2-panel (analytical / paired-measured) chart generator.
        _plot_comm now takes mode="analytical" | "paired".
  · paper_plot_gqa_kv_sharding_diagram.py  - 6-case 2-D KV-tensor
        diagram + companion comparison-table PNG.
  · measure_gqa_decode_placement_comm.py   - runs all 6 kernels at
        S_kv=8K, sums actual IPCQ-copy bytes from engine.op_log,
        scales partial-score AR ×128 to S_kv=1M, writes
        gqa_3cases_measured_comm.json (committed).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-18 11:26:05 -07:00
parent dd3337f2e4
commit 84bb418e1e
8 changed files with 1770 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/paper/measure_gqa_decode_placement_comm.py
+362
View File
@@ -0,0 +1,362 @@
"""Measured comm overlay for all 6 GQA decode KV placements.
Runs the simulator for Cases 1-6 (the same 6 placements the chart in
paper_plot_gqa_4cases_summary.py covers analytically) at S_kv = 64 K,
sums actual IPCQ-copy bytes from the engine op_log, projects to
per-token (x80 layers), scales the partial-score-AR component of
Cases 4/5 from S_kv = 64 K -> S_kv = 1 M (linear in S_kv; other
cases are S_kv-independent), adds the constant Wo + FFN AR
(1.25 MB / token), and writes the result to JSON for
paper_plot_gqa_4cases_summary.py to overlay on the analytical bars.
Single layer of decode attention only — the projection × 80 takes
the per-layer measurement to a per-token total.
Usage:
python scripts/paper/measure_gqa_decode_placement_comm.py
"""
from __future__ import annotations
import json
import os
import sys
from pathlib import Path
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel as _case3_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel as _case1_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel as _case2_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel,
)
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy
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
_C = 8
_P = 8
_N_LAYERS = 80
_S_KV_MEAS = 8 * 1024 # per-run simulator S_kv (small — 1/128th of headline)
_S_KV_HEADLINE = 1 << 20 # 1 Mi tokens, the chart's headline S_kv
# Per-cube S_kv share for the d_head-TP partial-score AR cost.
# Case 4 (Cube-SP × PE-TP_dhead): per-cube = S_kv / C
# Case 5 (Cube-TP_dhead × PE-SP): per-cube = S_kv (KV replicated across
# cubes for the cube-axis d_head-TP), so partial-score AR scales
# with the full S_kv.
# Headline / measured per-cube ratios give the partial-score-AR scale-up
# from S_kv = 64 K to S_kv = 1 M. (m,ℓ,O) AR is S_kv-independent.
_PARTIAL_SCORE_SCALE = _S_KV_HEADLINE / _S_KV_MEAS # = 16
# Per-token Wo + FFN AR (constant across all cases, comes from the
# attn-output and FFN-down all-reduces NOT measured by the attention-
# only kernel run here).
_WO_PER_LAYER_BYTES = 8 * 1024
_FFN_PER_LAYER_BYTES = 8 * 1024
_WO_FFN_PER_TOKEN_BYTES = (
(_WO_PER_LAYER_BYTES + _FFN_PER_LAYER_BYTES) * _N_LAYERS
) # 1.25 MB
# Total PE count in one KV-head group — average per-PE comm = total / N.
_NUM_PES = _C * _P
# Total partial-score-AR slice produced by the attention compute when
# d_head is sharded. Used to split measured IPCQ traffic into the
# S_kv-scaling component (partial scores) vs the S_kv-independent
# component ((m,ℓ,O) merge). Same as the analytical formula in
# paper_plot_gqa_4cases_summary.py: h_q · S_q · per_cube_S_kv · 2 bytes.
_H_Q = 8
_S_Q = 1
_BYTES_PER_ELEM = 2
_PARAMS = dict(C=_C, P=_P, T_q=_S_Q, S_kv=_S_KV_MEAS,
d_head=128, h_q=_H_Q, h_kv=1)
def _bench_fn_case1(ctx):
"""Case 1: Cube-Repl x PE-repl (PE-TP doesn't shard KV)."""
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp = DPPolicy(cube="replicate", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp, name="q_c1")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp, name="k_c1")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp, name="v_c1")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp, name="o_c1")
ctx.launch("case1_repl_repl", _case1_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _bench_fn_case2(ctx):
"""Case 2: Cube-SP x PE-repl (PE-TP doesn't shard KV)."""
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="q_c2")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c2")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c2")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="o_c2")
ctx.launch("case2_sp_repl", _case2_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _bench_fn_case3(ctx):
"""Case 3: Cube-Repl x PE-SP."""
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="replicate", pe="row_wise",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="q_c3")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c3")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c3")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="o_c3")
ctx.launch("case3_repl_sp", _case3_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _bench_fn_case4(ctx):
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="column_wise",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="q_c4")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c4")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c4")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="o_c4")
ctx.launch("case4_dhead_tp", _case4_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _bench_fn_case5(ctx):
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_q = DPPolicy(cube="column_wise", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_q, name="q_c5")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c5")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c5")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_q, name="o_c5")
ctx.launch("case5_dhead_tp_inter", _case5_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _bench_fn_case6(ctx):
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="q_c6")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c6")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c6")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="o_c6")
ctx.launch("case6_sp_sp", _case6_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _sum_ipcq_bytes(op_log) -> int:
"""Sum nbytes across all ipcq_copy records."""
return sum(
r.params.get("nbytes", 0)
for r in op_log
if r.op_kind == "memory" and r.op_name == "ipcq_copy"
)
def _partial_score_slices(case: int) -> int:
"""Divisor that splits S_kv into partial-score tiles, per the
analytical model in paper_plot_gqa_4cases_summary.py:
partial_score_per_PE = h_q * S_q * (s_kv / slices) * bytes
Case 4 (Cube-SP x PE-TP-dhead): intra-cube AR over d_head shards
on PE axis -> partial tile per PE has per-cube S_kv = s_kv/C.
Case 5 (Cube-TP-dhead x PE-SP): inter-cube AR over d_head shards
on cube axis -> partial tile per PE has per-PE S_kv = s_kv/P.
Cases 1, 2, 3, 6: no partial-score AR (only (m,l,O) merge).
"""
if case == 4:
return _C
if case == 5:
return _P
return 0
def _split_attn_layer_bytes(case: int, total_ipcq_bytes: int,
s_kv: int) -> tuple[int, int]:
"""Split per-layer attention-time IPCQ bytes into:
(partial_score_component, mlo_component).
The partial-score component scales with s_kv (so it must be scaled
when projecting from the measure-time s_kv to the headline s_kv);
the (m,l,O) component is constant in s_kv.
Partial-score size is analytically known per-case (formula in
_partial_score_slices); the remainder is treated as (m,l,O) + any
other S_kv-independent overhead. Per-PE = total / NUM_PES.
"""
per_pe_total = total_ipcq_bytes // _NUM_PES
slices = _partial_score_slices(case)
if slices == 0:
# No partial-score AR for this case.
return 0, per_pe_total
partial_score_per_pe = (
_H_Q * _S_Q * (s_kv // slices) * _BYTES_PER_ELEM
)
partial_score_per_pe = min(partial_score_per_pe, per_pe_total)
mlo_per_pe = per_pe_total - partial_score_per_pe
return partial_score_per_pe, mlo_per_pe
_KERNELS = (
(1, "Case 1 (Cube-Repl x PE-repl)", _bench_fn_case1),
(2, "Case 2 (Cube-SP x PE-repl)", _bench_fn_case2),
(3, "Case 3 (Cube-Repl x PE-SP)", _bench_fn_case3),
(4, "Case 4 (Cube-SP x PE-TP d_head)", _bench_fn_case4),
(5, "Case 5 (Cube-TP d_head x PE-SP)", _bench_fn_case5),
(6, "Case 6 (Cube-SP x PE-SP) [*]", _bench_fn_case6),
)
def main() -> int:
topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
topo = resolve_topology(topology)
out: dict = {
"S_kv_measured": _S_KV_MEAS,
"S_kv_headline": _S_KV_HEADLINE,
"n_layers": _N_LAYERS,
"num_pes": _NUM_PES,
"wo_ffn_per_token_bytes": _WO_FFN_PER_TOKEN_BYTES,
"cases": {},
}
print(f"Measuring at S_kv={_S_KV_MEAS:,} ; scaling partial-score AR "
f"to S_kv={_S_KV_HEADLINE:,} (×{int(_PARTIAL_SCORE_SCALE)})")
print()
for case_id, label, bench_fn in _KERNELS:
try:
res = run_bench(
topology=topo, bench_fn=bench_fn,
device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(
getattr(t, "topology_obj", t), enable_data=True,
),
)
except Exception as e:
print(f" {label:<42} FAIL: {type(e).__name__}: {e}")
return 1
if not res.completion.ok:
print(f" {label:<42} ENGINE FAIL: {res.completion}")
return 1
total_ipcq = _sum_ipcq_bytes(res.engine.op_log)
partial_pe, mlo_pe = _split_attn_layer_bytes(
case_id, total_ipcq, _S_KV_MEAS,
)
# Per-token attention-time comm at S_kv = 1 M:
# (partial_score_per_layer × scale + mlo_per_layer) × 80 layers
scaled_partial_per_token = (
partial_pe * int(_PARTIAL_SCORE_SCALE) * _N_LAYERS
)
mlo_per_token = mlo_pe * _N_LAYERS
attn_per_token = scaled_partial_per_token + mlo_per_token
total_per_token = attn_per_token + _WO_FFN_PER_TOKEN_BYTES
out["cases"][str(case_id)] = {
"label": label,
"total_ipcq_bytes_one_layer": total_ipcq,
"per_pe_partial_score_bytes_one_layer": partial_pe,
"per_pe_mlo_bytes_one_layer": mlo_pe,
"per_pe_attn_bytes_per_token_at_1M": attn_per_token,
"per_pe_total_bytes_per_token_at_1M": total_per_token,
}
print(f" {label:<42} "
f"ipcq_total={total_ipcq:>10,} "
f"per_pe_attn(1L)={(partial_pe + mlo_pe):>9,} "
f"per_pe_total/tok@1M={total_per_token / (1<<20):>7.2f} MB")
out_path = (
Path(__file__).resolve().parents[2]
/ "src" / "kernbench" / "benches"
/ "1H_milestone_output" / "gqa" / "long_ctx"
/ "gqa_3cases_measured_comm.json"
)
out_path.parent.mkdir(parents=True, exist_ok=True)
out_path.write_text(json.dumps(out, indent=2))
print()
print(f"wrote {out_path}")
return 0
if __name__ == "__main__":
sys.exit(main())
scripts/paper/paper_plot_gqa_4cases_summary.py
+151 -44
View File
@@ -114,6 +114,7 @@ Output PNG:
"""
from __future__ import annotations
import json
from pathlib import Path
import matplotlib.patches as mpatches
@@ -240,6 +241,22 @@ _OUT_DIR = (
/ "src" / "kernbench" / "benches"
/ "1H_milestone_output" / "gqa" / "long_ctx"
)
_MEASURED_JSON = _OUT_DIR / "gqa_3cases_measured_comm.json"
def _load_measured() -> dict[int, float] | None:
"""Load measured per-PE comm bytes (already scaled to S_kv=1M).
Returns {case_id: per_token_total_bytes} or None if JSON missing.
Produced by scripts/paper/measure_gqa_decode_placement_comm.py.
"""
if not _MEASURED_JSON.exists():
return None
data = json.loads(_MEASURED_JSON.read_text())
return {
int(cid): info["per_pe_total_bytes_per_token_at_1M"]
for cid, info in data["cases"].items()
}
# ── Formulae ────────────────────────────────────────────────────────
@@ -388,51 +405,114 @@ def _plot_memory(ax) -> None:
ax.legend(loc="upper right", fontsize=9)
def _plot_comm(ax) -> None:
wo_mb = []
ffn_mb = []
attn_mb = []
def _plot_comm(ax, *, mode: str = "analytical") -> None:
"""Per-PE comm panel.
mode = "analytical": single solid bars from per_token_bytes formula.
mode = "paired" : analytical (solid) + simulator-measured
(hatched) side-by-side per case, when the
measurement JSON is available.
"""
labels = [_CASE_LABEL[c] for c in _CASES]
wo_mb_list: list[float] = []
ffn_mb_list: list[float] = []
attn_mb: list[float] = []
for c in _CASES:
wo, ffn, attn = per_token_bytes(c, _HEADLINE_S_KV)
wo_mb.append(wo / (1 << 20))
ffn_mb.append(ffn / (1 << 20))
wo_mb_list.append(wo / (1 << 20))
ffn_mb_list.append(ffn / (1 << 20))
attn_mb.append(attn / (1 << 20))
labels = [_CASE_LABEL[c] for c in _CASES]
x = list(range(len(_CASES)))
measured = _load_measured() if mode == "paired" else None
paired = measured is not None
source_tag = "analytical (solid) vs simulator-measured (hatched)" \
if paired else "analytical"
ax.bar(x, wo_mb, color=_WO_COLOR, edgecolor="black")
ax.bar(x, ffn_mb, bottom=wo_mb, color=_FFN_COLOR, edgecolor="black")
bottoms_attn = [w + f for w, f in zip(wo_mb, ffn_mb)]
ax.bar(x, attn_mb, bottom=bottoms_attn,
n_cases = len(_CASES)
x = list(range(n_cases))
bar_w = 0.36 if paired else 0.65
x_ana = [xi - bar_w / 2 for xi in x] if paired else x
x_meas = [xi + bar_w / 2 for xi in x] if paired else None
# Analytical bars (solid).
ax.bar(x_ana, wo_mb_list, width=bar_w,
color=_WO_COLOR, edgecolor="black")
ax.bar(x_ana, ffn_mb_list, width=bar_w, bottom=wo_mb_list,
color=_FFN_COLOR, edgecolor="black")
bottoms_attn = [w + f for w, f in zip(wo_mb_list, ffn_mb_list)]
ax.bar(x_ana, attn_mb, width=bar_w, bottom=bottoms_attn,
color=_ATTN_COLOR, edgecolor="black")
# Measured bars (hatched) — same Wo+FFN base, attn from op_log.
meas_attn_mb: list[float] = []
if paired:
for i, c in enumerate(_CASES):
meas_total = measured.get(c, 0) / (1 << 20)
meas_attn_mb.append(
max(meas_total - wo_mb_list[i] - ffn_mb_list[i], 0.0))
ax.bar(x_meas, wo_mb_list, width=bar_w,
color=_WO_COLOR, edgecolor="black",
hatch="///", alpha=0.85)
ax.bar(x_meas, ffn_mb_list, width=bar_w, bottom=wo_mb_list,
color=_FFN_COLOR, edgecolor="black",
hatch="///", alpha=0.85)
ax.bar(x_meas, meas_attn_mb, width=bar_w, bottom=bottoms_attn,
color=_ATTN_COLOR, edgecolor="black",
hatch="///", alpha=0.85)
ax.set_xticks(x)
ax.set_xticklabels(labels, fontsize=9)
ax.set_ylabel("Comm per token per PE (MB, log)")
ax.set_yscale("log")
ax.set_title(
title = (
f"Per-PE communication at S_kv = {_HEADLINE_S_KV:,} tokens "
f"(decode S_q=1, B=1; {_N_LAYERS} layers)",
fontsize=11,
f"(decode S_q=1, B=1; {_N_LAYERS} layers) — {source_tag}"
)
if paired:
title += "\n(simulator measured at S_kv = 8K; " \
"partial-score AR scaled ×128 to S_kv = 1M)"
ax.set_title(title, fontsize=10)
ax.grid(axis="y", ls=":", alpha=0.5, which="both")
totals = [wo_mb[i] + ffn_mb[i] + attn_mb[i] for i in range(len(_CASES))]
ax.set_ylim(top=max(totals) * 10)
totals_ana = [wo_mb_list[i] + ffn_mb_list[i] + attn_mb[i]
for i in range(n_cases)]
ymax = max(totals_ana)
if paired:
ymax = max(ymax, max(
(measured.get(c, 0) / (1 << 20)) for c in _CASES))
ax.set_ylim(top=ymax * 22)
for i, c in enumerate(_CASES):
total_bytes = (wo_mb[i] + ffn_mb[i] + attn_mb[i]) * (1 << 20)
ax.text(x[i], totals[i] * 1.5, _fmt_bytes(total_bytes),
ha="center", fontsize=9, weight="bold")
ana_bytes = totals_ana[i] * (1 << 20)
if paired:
meas_bytes = measured.get(c, 0)
top_y = max(totals_ana[i], meas_bytes / (1 << 20))
label = (f"ana: {_fmt_bytes(ana_bytes)}\n"
f"sim: {_fmt_bytes(meas_bytes)}")
else:
top_y = totals_ana[i]
label = _fmt_bytes(ana_bytes)
ax.text(x[i], top_y * 1.5, label,
ha="center", va="bottom",
fontsize=8 if paired else 9,
weight="bold", linespacing=1.05)
# In-bar attention-time AR description (online-softmax (m,ℓ,O)
# merge, partial-score AR, intra/inter-cube). Placed on the
# analytical bar; with paired mode it lives on the solid bar so
# the hatched measured bar stays uncluttered.
if attn_mb[i] > 0:
mid = bottoms_attn[i] + attn_mb[i] / 2
ax.text(x[i], mid, _ATTN_DESC[c],
ha="center", va="center", fontsize=7,
ax.text(x_ana[i], mid, _ATTN_DESC[c],
ha="center", va="center",
fontsize=6 if paired else 7,
color="black", weight="bold")
else:
ax.text(x[i], totals[i] * 0.4, _ATTN_DESC[c],
ha="center", fontsize=7.5, color="grey", style="italic")
ax.text(x_ana[i], totals_ana[i] * 0.4, _ATTN_DESC[c],
ha="center",
fontsize=6.5 if paired else 7.5,
color="grey", style="italic")
legend_handles = [
mpatches.Patch(facecolor=_WO_COLOR, edgecolor="black",
@@ -442,40 +522,67 @@ def _plot_comm(ax) -> None:
mpatches.Patch(facecolor=_ATTN_COLOR, edgecolor="black",
label="Attn-time collective"),
]
if paired:
legend_handles.append(
mpatches.Patch(facecolor="white", edgecolor="black",
hatch="///", label="simulator-measured"))
ax.legend(handles=legend_handles, loc="upper right", fontsize=8,
framealpha=0.92)
def main() -> Path:
_OUT_DIR.mkdir(parents=True, exist_ok=True)
fig = plt.figure(figsize=(21.0, 6.0))
gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22)
ax_b = fig.add_subplot(gs[0, 0])
ax_m = fig.add_subplot(gs[0, 1])
ax_c = fig.add_subplot(gs[0, 2])
# (a) Per-PE HBM budget — standalone PNG.
fig_b, ax_b = plt.subplots(figsize=(4.0, 6.0))
_plot_budget(ax_b)
fig_b.tight_layout()
out_b = _OUT_DIR / "gqa_hbm_budget.png"
fig_b.savefig(out_b, dpi=150)
plt.close(fig_b)
print(f"wrote {out_b}")
# (b) Combined 3-panel summary — HBM budget + KV memory + comm.
fig = plt.figure(figsize=(22.0, 6.5))
gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22)
ax_b2 = fig.add_subplot(gs[0, 0])
ax_m = fig.add_subplot(gs[0, 1])
ax_c = fig.add_subplot(gs[0, 2])
_plot_budget(ax_b2)
_plot_memory(ax_m)
_plot_comm(ax_c)
fig.suptitle(
f"GQA per-PE memory + communication — LLaMA-3.1-70B "
f"single-KV-head group (C={_C}, P={_P}, {_N_LAYERS} layers, "
f"S_kv = 1 M, FP16)",
fontsize=12, y=0.985,
)
fig.text(
0.5, 0.94,
"Cases ordered by memory · Case 1 (40 GB, no sharding) → "
"Cases 2-3 (5 GB, 1-axis sharded) → Cases 4-6 (640 MB, 2-axis sharded) "
"· Case 6 ★ = lowest comm among memory-feasible cases",
ha="center", fontsize=9.5, color="#444",
)
fig.tight_layout(rect=(0, 0, 1, 0.92))
fig.tight_layout()
out = _OUT_DIR / "gqa_4cases_summary.png"
fig.savefig(out, dpi=150)
plt.close(fig)
print(f"wrote {out}")
# (c) 2-panel companion (analytical only).
fig2 = plt.figure(figsize=(18.0, 6.5))
gs2 = fig2.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18)
ax_m2 = fig2.add_subplot(gs2[0, 0])
ax_c2 = fig2.add_subplot(gs2[0, 1])
_plot_memory(ax_m2)
_plot_comm(ax_c2, mode="analytical")
fig2.tight_layout()
out2 = _OUT_DIR / "gqa_4cases_memory_comm_analytical.png"
fig2.savefig(out2, dpi=150)
plt.close(fig2)
print(f"wrote {out2}")
# (d) 2-panel companion — analytical vs simulator-measured paired.
fig3 = plt.figure(figsize=(19.0, 6.5))
gs3 = fig3.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18)
ax_m3 = fig3.add_subplot(gs3[0, 0])
ax_c3 = fig3.add_subplot(gs3[0, 1])
_plot_memory(ax_m3)
_plot_comm(ax_c3, mode="paired")
fig3.tight_layout()
out3 = _OUT_DIR / "gqa_4cases_memory_comm_paired.png"
fig3.savefig(out3, dpi=150)
plt.close(fig3)
print(f"wrote {out3}")
# Paper-ready table to stdout.
print()
print(f" {'Case':<27} {'KV/tok·PE':>12} {'KV @ 1M':>11} "
scripts/paper/paper_plot_gqa_kv_sharding_diagram.py
+323
View File
@@ -0,0 +1,323 @@
"""6-case KV-sharding tensor diagram (the slide-13 PNG export).
Flat 2-D rectangles, one per sharding case, with:
Y axis = S_kv (vertical) — Cube-SP / PE-SP slice it
X axis = d_head (horizontal) — Cube-TP-d_head / PE-TP-d_head slice it
Drops the batch axis entirely (decode: B = 1, T_q = 1). Same case set
and visual encoding as slide 13 of GQA_full_deck.pptx; matplotlib
renders it cleanly so the PNG sits next to the other GQA summary
artifacts in 1H_milestone_output/gqa/long_ctx/.
"""
from __future__ import annotations
from pathlib import Path
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
_C = 8
_P = 8
_GROUP_FILLS = [
"#A5D8FF", "#B2F2BB", "#FFD8A8", "#FFC9C9",
"#D0BFFF", "#99E9F2", "#FCC2D7", "#FFEC99",
]
_ACC = {
"red": "#E03131",
"orange": "#FD7E14",
"blue": "#1C7ED6",
"green": "#37B24D",
}
# (label, accent, kv, comm, overflow, encoding-flags, axis-spec)
# y_split = 8 horizontal Y bands (Cube-SP on S_kv)
# x_split = 8 vertical X bands (Cube-TP-d_head)
# pe_y = 7 fine horizontal dividers within each Y band
# pe_x = 7 fine vertical dividers within each X band
# axes = small annotation under the chip naming the axes
# that the cube/PE actually shard, so the reader can
# parse Case 5 (where cube colour fills run X instead
# of Y, breaking the visual symmetry of the rest).
_CASES = [
dict(label="Case 1\nCube-Repl / PE-repl", accent=_ACC["red"],
kv="40 GB", comm="1.2 MB", overflow=True,
y_split=False, x_split=False, pe_y=False, pe_x=False,
axes="Cube: replicated PE: replicated"),
dict(label="Case 2\nCube-SP / PE-repl", accent=_ACC["orange"],
kv="5 GB", comm="3.8 MB", overflow=True,
y_split=True, x_split=False, pe_y=False, pe_x=False,
axes="Cube → Y (S_kv) PE: replicated"),
dict(label="Case 3\nCube-Repl / PE-SP", accent=_ACC["orange"],
kv="5 GB", comm="3.8 MB", overflow=True,
y_split=False, x_split=False, pe_y=True, pe_x=False,
axes="Cube: replicated PE → Y (S_kv)"),
dict(label="Case 4\nCube-SP / PE-TP-d_head", accent=_ACC["blue"],
kv="640 MB", comm="166 MB", overflow=False,
y_split=True, x_split=False, pe_y=False, pe_x=True,
axes="Cube → Y (S_kv) PE → X (d_head)"),
dict(label="Case 5\nCube-TP-d_head / PE-SP", accent=_ACC["blue"],
kv="640 MB", comm="166 MB", overflow=False,
y_split=False, x_split=True, pe_y=True, pe_x=False,
axes="Cube → X (d_head) PE → Y (S_kv)"),
dict(label="Case 6 ★\nCube-SP / PE-SP", accent=_ACC["green"],
kv="640 MB", comm="6.2 MB", overflow=False,
y_split=True, x_split=False, pe_y=True, pe_x=False,
axes="Cube → Y (S_kv) PE → Y (S_kv)"),
]
_OUT_DIR = (
Path(__file__).resolve().parents[2]
/ "src" / "kernbench" / "benches"
/ "1H_milestone_output" / "gqa" / "long_ctx"
)
def _draw_panel(ax, cfg):
"""Draw one case's 2-D KV-tensor rectangle into a panel ax."""
ax.set_xlim(0, 1)
ax.set_ylim(1, 0) # Y points down (S_kv ↓)
ax.set_aspect("auto")
ax.set_xticks([])
ax.set_yticks([])
cube_repl = not cfg["y_split"] and not cfg["x_split"]
pe_repl = not cfg["pe_y"] and not cfg["pe_x"]
# Cube-level colour fill.
if cfg["y_split"] and not cfg["x_split"]:
# 8 horizontal Y bands.
for c in range(_C):
ax.add_patch(mpatches.Rectangle(
(0, c / _C), 1, 1 / _C,
facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
ax.text(0.04, c / _C + 0.5 / _C, f"C{c}",
ha="left", va="center", fontsize=8,
fontweight="bold", color="#333")
elif cfg["x_split"] and not cfg["y_split"]:
# 8 vertical X bands.
for c in range(_C):
ax.add_patch(mpatches.Rectangle(
(c / _C, 0), 1 / _C, 1,
facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
ax.text(c / _C + 0.5 / _C, 0.04, f"C{c}",
ha="center", va="top", fontsize=8,
fontweight="bold", color="#333")
else:
ax.add_patch(mpatches.Rectangle(
(0, 0), 1, 1,
facecolor="#F5F5F5", edgecolor="black", linewidth=0.8))
ax.text(0.5, 0.5, "× 8 cubes\nfull KV",
ha="center", va="center",
fontsize=10, fontweight="bold",
fontstyle="italic", color="#666")
# PE-level fine dividers — distinguished from cube boundaries by
# using a dashed style + slightly stronger contrast. This is what
# makes Case 5's PE-SP (horizontal lines across vertical cube
# bands) read as "different axis from the cubes" at a glance.
if cfg["pe_y"]:
outer = _C if cfg["y_split"] else 1
band = 1 / outer
for o in range(outer):
for p in range(1, _P):
y = o * band + band * p / _P
ax.axhline(y, color="#222", linewidth=0.8,
linestyle=(0, (3, 2)), alpha=0.75)
if cfg["pe_x"]:
outer = _C if cfg["x_split"] else 1
band = 1 / outer
for o in range(outer):
for p in range(1, _P):
x = o * band + band * p / _P
ax.axvline(x, color="#222", linewidth=0.8,
linestyle=(0, (3, 2)), alpha=0.75)
# Heavy outline on top.
ax.add_patch(mpatches.Rectangle(
(0, 0), 1, 1, facecolor="none",
edgecolor="black", linewidth=1.2))
# Replication badges — small text-only badges in the corners of
# the rectangle, no ghost-card stacking (which mis-reads as a
# larger enclosing tensor).
badges: list[str] = []
if cube_repl:
badges.append("× 8 cube copies")
if pe_repl and (cfg["y_split"] or cfg["x_split"]):
# Cube is sharded but PEs in each cube replicate that shard.
badges.append("× 8 PEs / cube replicate")
elif pe_repl and cube_repl:
# Both replicated — PE replication adds to the cube one.
badges.append("× 8 PEs / cube replicate")
if badges:
ax.text(0.98, 0.02, "\n".join(badges),
ha="right", va="top", fontsize=7,
fontweight="bold", color="#444",
fontstyle="italic",
bbox=dict(facecolor="white", edgecolor="#888",
boxstyle="round,pad=0.20", linewidth=0.5))
def _make_table_png() -> Path:
"""Slide-14 companion table: per-PE memory + comm for all 6 cases."""
headers = ["Case", "Sharding", "KV / PE", "Fit",
"Comm/tok\n(analytical)", "Comm/tok\n(measured)", "Notes"]
rows = [
("Case 1", "Cube-Repl · PE-repl", "40 GB", "✗",
"1.2 MB", "1.25 MB",
"no sharding — full KV on every PE"),
("Case 2", "Cube-SP · PE-repl", "5 GB", "✗",
"3.8 MB", "1.27 MB",
"cube-axis sharded only"),
("Case 3", "Cube-Repl · PE-SP", "5 GB", "✗",
"3.8 MB", "1.39 MB",
"PE-axis sharded only"),
("Case 4", "Cube-SP · PE-TP-d_head", "640 MB", "✓",
"166 MB", "162 MB",
"d_head split intra-cube — partial-score AR ∝ S_kv"),
("Case 5", "Cube-TP-d_head · PE-SP", "640 MB", "✓",
"166 MB", "162 MB",
"d_head split inter-cube — partial-score AR on UCIe"),
("Case 6 ★", "Cube-SP · PE-SP", "640 MB", "✓",
"6.2 MB", "1.41 MB",
"S_kv split on both axes — (m,ℓ,O) AR only, S_kv-indep."),
]
accents = [_ACC["red"], _ACC["orange"], _ACC["orange"],
_ACC["blue"], _ACC["blue"], _ACC["green"]]
fig, ax = plt.subplots(figsize=(20.0, 4.6))
ax.set_axis_off()
cell_data = [headers] + [list(r) for r in rows]
tbl = ax.table(cellText=cell_data,
colWidths=[0.07, 0.18, 0.08, 0.05, 0.12, 0.12, 0.38],
cellLoc="center", loc="center")
tbl.auto_set_font_size(False)
tbl.set_fontsize(11)
tbl.scale(1.0, 2.2)
n_cols = len(headers)
n_rows = len(rows) + 1 # +1 header
# Header styling.
for ci in range(n_cols):
cell = tbl[(0, ci)]
cell.set_facecolor("#1F4E79")
cell.set_text_props(color="white", weight="bold")
cell.set_edgecolor("#1F4E79")
# Body styling.
for ri, row in enumerate(rows, start=1):
is_pareto = row[0].endswith("★")
row_fill = "#E8F5E9" if is_pareto else (
"white" if ri % 2 == 1 else "#F5F5F7")
# Case-name cell uses accent.
case_cell = tbl[(ri, 0)]
case_cell.set_facecolor(accents[ri - 1])
case_cell.set_text_props(color="white", weight="bold")
# Remaining cells.
for ci in range(1, n_cols):
cell = tbl[(ri, ci)]
cell.set_facecolor(row_fill)
txt_kwargs = {"weight": "bold" if is_pareto else "normal",
"color": "#333"}
if ci == 2: # KV / PE
txt_kwargs["color"] = (
"#C62828" if row[3] == "✗" else "#2E7D32")
txt_kwargs["weight"] = "bold"
if ci == 3: # Fit
txt_kwargs["color"] = (
"#C62828" if row[3] == "✗" else "#2E7D32")
txt_kwargs["weight"] = "bold"
cell.set_text_props(**txt_kwargs)
# Last-column (Notes) cells left-aligned for readability.
tbl[(ri, n_cols - 1)].get_text().set_ha("left")
# Force left-align on the Notes header too.
tbl[(0, n_cols - 1)].get_text().set_ha("left")
fig.suptitle(
"GQA decode KV-sharding — per-PE memory & communication "
"(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, 80 layers)",
fontsize=12, y=0.97,
)
out = _OUT_DIR / "gqa_kv_sharding_6cases_table.png"
fig.savefig(out, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f"wrote {out}")
return out
def main() -> Path:
_OUT_DIR.mkdir(parents=True, exist_ok=True)
n = len(_CASES)
fig = plt.figure(figsize=(20.0, 7.0))
# Three rows per column: case chip · axis-spec annotation · rectangle.
gs = fig.add_gridspec(3, n,
height_ratios=[0.55, 0.32, 8.5],
hspace=0.05, wspace=0.20,
left=0.04, right=0.99,
top=0.93, bottom=0.06)
for i, cfg in enumerate(_CASES):
# Top: case chip header.
ax_chip = fig.add_subplot(gs[0, i])
ax_chip.set_xticks([])
ax_chip.set_yticks([])
for spine in ax_chip.spines.values():
spine.set_visible(False)
ax_chip.add_patch(mpatches.Rectangle(
(0, 0), 1, 1, transform=ax_chip.transAxes,
facecolor=cfg["accent"], edgecolor=cfg["accent"]))
ax_chip.text(0.5, 0.5, cfg["label"],
ha="center", va="center",
fontsize=10, fontweight="bold",
color="white")
# Middle: axis-spec annotation — names which axis the cube
# shards on and which axis the PE shards on (essential for
# parsing Case 5 where the cube colour fills run X instead
# of Y, breaking the visual symmetry of the rest).
ax_axes = fig.add_subplot(gs[1, i])
ax_axes.set_xticks([])
ax_axes.set_yticks([])
for spine in ax_axes.spines.values():
spine.set_visible(False)
ax_axes.add_patch(mpatches.Rectangle(
(0, 0), 1, 1, transform=ax_axes.transAxes,
facecolor="#F5F5F7", edgecolor="#CCCCCC",
linewidth=0.6))
ax_axes.text(0.5, 0.5, cfg["axes"],
ha="center", va="center",
fontsize=8.5, fontweight="bold",
color="#1F4E79")
# Bottom: the tensor rectangle.
ax = fig.add_subplot(gs[2, i])
_draw_panel(ax, cfg)
ax.set_xlabel("X : d_head = 128 →",
fontsize=9, fontweight="bold",
fontstyle="italic", color="#1F4E79")
ax.set_ylabel("Y : S_kv = 1 M ↓",
fontsize=9, fontweight="bold",
fontstyle="italic", color="#1F4E79")
fig.suptitle(
"GQA decode KV-tensor sharding — 6 cases · "
"LLaMA 70B GQA single KV-head group · "
"C = 8 cubes × P = 8 PEs · S_kv = 1 M, FP16, 80 layers",
fontsize=12, y=0.99,
)
out = _OUT_DIR / "gqa_kv_sharding_6cases_diagram.png"
fig.savefig(out, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f"wrote {out}")
# Companion table PNG (slide-14 export).
_make_table_png()
return out
if __name__ == "__main__":
main()
scripts/verify_case4_dhead_tp.py
+148
View File
@@ -0,0 +1,148 @@
"""Phase 1a smoke test for the Case 4 d_head-TP decode kernel.
Runs `gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel` at small
S_kv (2K) to verify:
1. it imports cleanly,
2. it runs without error under the same harness as Case 6,
3. captures op_log_summary (gemm/ipcq/dma counts),
4. compares vs analytical predictions and vs Case 6 (same memory tier).
Usage:
python scripts/verify_case4_dhead_tp.py
"""
from __future__ import annotations
import os
import sys
from pathlib import Path
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel,
)
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
_ccl_cfg, _summarize_op_log,
)
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy
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
# Small smoke-test params (S_kv=2K is enough to exercise tile-loop + AR).
_PARAMS = dict(C=8, P=8, T_q=1, S_kv=2_048,
d_head=128, h_q=8, h_kv=1)
def _bench_fn_case4(ctx):
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="column_wise",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="q_c4")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c4")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c4")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="o_c4")
ctx.launch("case4_dhead_tp", _case4_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _bench_fn_case5(ctx):
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_q = DPPolicy(cube="column_wise", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_q, name="q_c5")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c5")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c5")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_q, name="o_c5")
ctx.launch("case5_dhead_tp_inter", _case5_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def _bench_fn_case6(ctx):
p = _PARAMS
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="replicate",
num_cubes=p["C"], num_pes=p["P"])
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
num_cubes=p["C"], num_pes=p["P"])
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="q_c6")
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="k_c6")
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
dtype="f16", dp=dp_kv, name="v_c6")
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
dtype="f16", dp=dp_full, name="o_c6")
ctx.launch("case6_sp_sp", _case6_kernel,
q, k, v, o,
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
p["d_head"], p["C"], p["P"],
_auto_dim_remap=False)
def main() -> int:
topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
topo = resolve_topology(topology)
print(f"Smoke params: {_PARAMS}")
print()
for label, bench_fn in (
("Case 4 (Cube-SP × PE-TP d_head)", _bench_fn_case4),
("Case 5 (Cube-TP d_head × PE-SP)", _bench_fn_case5),
("Case 6 (Cube-SP × PE-SP S_kv)", _bench_fn_case6),
):
try:
res = run_bench(
topology=topo, bench_fn=bench_fn,
device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(
getattr(t, "topology_obj", t), enable_data=True,
),
)
except Exception as e:
print(f" {label:<42} FAIL: {type(e).__name__}: {e}")
continue
if not res.completion.ok:
print(f" {label:<42} ENGINE FAIL: {res.completion}")
continue
s = _summarize_op_log(res.engine.op_log)
lat = (res.engine.op_log[-1].t_end if res.engine.op_log else 0.0)
print(f" {label:<42} "
f"gemm={s['gemm_count']:>4} "
f"ipcq={s['ipcq_copy_count']:>4} "
f"dma_r={s['dma_read_count']:>4} "
f"dma_w={s['dma_write_count']:>3} "
f"latency={lat:.1f} ns "
f"(n_ops={len(res.engine.op_log)})")
return 0
if __name__ == "__main__":
sys.exit(main())