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())