mukesh
84bb418e1e
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>
324 lines
13 KiB
Python
324 lines
13 KiB
Python
"""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()
|