mukesh
73e0b315fe
Removes the (measured) comm column, widens the Notes column with shorter wrapped text, scales the table figure down, and wraps the heading to two lines so it sits flush with the table. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
325 lines
13 KiB
Python
325 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)", "Notes"]
|
||
rows = [
|
||
("Case 1", "Cube-Repl · PE-repl", "40 GB", "✗",
|
||
"1.2 MB",
|
||
"no sharding —\nfull KV on every PE"),
|
||
("Case 2", "Cube-SP · PE-repl", "5 GB", "✗",
|
||
"3.8 MB",
|
||
"cube-axis\nsharded only"),
|
||
("Case 3", "Cube-Repl · PE-SP", "5 GB", "✗",
|
||
"3.8 MB",
|
||
"PE-axis\nsharded only"),
|
||
("Case 4", "Cube-SP · PE-TP-d_head", "640 MB", "✓",
|
||
"166 MB",
|
||
"d_head split intra-cube\npartial-score AR ∝ S_kv"),
|
||
("Case 5", "Cube-TP-d_head · PE-SP", "640 MB", "✓",
|
||
"166 MB",
|
||
"d_head split inter-cube\npartial-score AR on UCIe"),
|
||
("Case 6 ★", "Cube-SP · PE-SP", "640 MB", "✓",
|
||
"6.2 MB",
|
||
"S_kv split both axes\n(m,ℓ,O) AR only"),
|
||
]
|
||
accents = [_ACC["red"], _ACC["orange"], _ACC["orange"],
|
||
_ACC["blue"], _ACC["blue"], _ACC["green"]]
|
||
|
||
fig, ax = plt.subplots(figsize=(15.0, 5.0))
|
||
ax.set_axis_off()
|
||
|
||
cell_data = [headers] + [list(r) for r in rows]
|
||
tbl = ax.table(cellText=cell_data,
|
||
colWidths=[0.07, 0.20, 0.09, 0.05, 0.14, 0.28],
|
||
cellLoc="center", loc="center")
|
||
tbl.auto_set_font_size(False)
|
||
tbl.set_fontsize(10.5)
|
||
tbl.scale(1.0, 2.4)
|
||
|
||
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\n"
|
||
"(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, "
|
||
"80 layers)",
|
||
fontsize=11.5, y=0.94,
|
||
)
|
||
out = _OUT_DIR / "gqa_long_ctx_6cases_kv_sharding_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_long_ctx_6cases_kv_sharding_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()
|