mukesh
9fdde44922
New tests/analytical_visualization/ module - an interactive dashboard for exploring memory / latency tradeoffs of transformer inference on the SIP architecture. Highlights: - 30+ model presets (Qwen, Llama 2/3/3.1, Mistral, Gemma 2, Phi 3, DeepSeek MLA, Mixtral/Qwen 3 MoE, Grok-1, ...) - Placement toggles: TP and CP each on PE-level vs cube-level - CP ring variant: K/V ring vs Q+O/m/l ring (prefill); in decode the O/m/l all-reduce is folded into S8 (no separate C1 row) - SIP interconnect: ring / mesh2d / torus2d with matching link drawing - Per-stage latency table with compute + memory + comm formulas, auto-scaled ns/us/ms, colored by dominant bound - Ring attention loop indicator on the pipeline diagram (purple arc over S5-S8 with 'xN hops' badge) - Tensor sharding view with optional physical PE/cube annotations - Replication-waste + optimization-hints panel - Save & compare configurations (config1, config2, ...): summary table plus side-by-side per-stage attention and FFN latency, best-in-row highlighting - Symbol glossary with current values for every symbol used in formulas Not tied to production sim_engine or runtime API; purely analytical tooling for design-space exploration. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
301 lines
11 KiB
Python
301 lines
11 KiB
Python
"""Pipeline diagram: attention + FFN steps, colored by bound type.
|
|
|
|
Each stage is drawn as a rounded box in a horizontal flow:
|
|
Attention: S1 -> S2 -> S3 -> ... -> S10 -> C2 (AllReduce)
|
|
FFN: F1 -> F2 -> F3 -> F4 -> F5 -> CF1 (AllReduce)
|
|
|
|
CP ring communication (C1) and score AllReduce (C3) are drawn as
|
|
"overhead" boxes attached above the QKt/PV region (they run
|
|
concurrently with per-hop compute in the ring).
|
|
|
|
Box fill color = the stage's bound classification:
|
|
compute-bound blue
|
|
memory-bound orange
|
|
comm-bound red
|
|
trivial grey
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as patches
|
|
|
|
from .model_config import FullConfig
|
|
from .stage_latencies import (
|
|
all_stages, all_ffn_stages, StageCost,
|
|
)
|
|
|
|
|
|
BOUND_COLOR = {
|
|
"compute": "#4682b4", # blue
|
|
"memory": "#e37400", # orange
|
|
"comm": "#c62828", # red
|
|
"trivial": "#a9a9a9", # grey
|
|
}
|
|
BOUND_LABEL = {
|
|
"compute": "compute-bound",
|
|
"memory": "memory-bound",
|
|
"comm": "communication",
|
|
"trivial": "trivial / negligible",
|
|
}
|
|
|
|
|
|
def _fmt_us(sec: float) -> str:
|
|
us = sec * 1e6
|
|
if us < 1:
|
|
return f"{us*1e3:.1f}ns"
|
|
if us < 100:
|
|
return f"{us:.2f}us"
|
|
return f"{us:.0f}us"
|
|
|
|
|
|
def _draw_block(ax, x, y, w, h, label, cost: StageCost, edge="#333"):
|
|
color = BOUND_COLOR.get(cost.bound, "#a9a9a9")
|
|
box = patches.FancyBboxPatch(
|
|
(x, y), w, h,
|
|
boxstyle="round,pad=0.05",
|
|
facecolor=color, edgecolor=edge, linewidth=1.2, alpha=0.85,
|
|
)
|
|
ax.add_patch(box)
|
|
# Two lines: stage name (bold) + visible latency (small)
|
|
ax.text(x + w / 2, y + h * 0.62, label,
|
|
ha="center", va="center",
|
|
fontsize=8.0, fontweight="bold", color="white")
|
|
ax.text(x + w / 2, y + h * 0.25, _fmt_us(cost.visible_s),
|
|
ha="center", va="center",
|
|
fontsize=7.0, color="white")
|
|
|
|
|
|
_SHORT_NAMES = {
|
|
"S1": "RMSNorm",
|
|
"S2": "W_Q",
|
|
"S3": "W_K+W_V",
|
|
"S4": "KV append",
|
|
"S5": "Q.K^T",
|
|
"S6": "softmax",
|
|
"S7": "P.V",
|
|
"S8": "merge",
|
|
"S9": "norm O/l",
|
|
"S10": "W_O",
|
|
"C1": "CP ring",
|
|
"C2": "TP AR",
|
|
"C3": "Score AR",
|
|
"F1": "RMSNorm",
|
|
"F2": "W_gate",
|
|
"F3": "W_up",
|
|
"F4": "SwiGLU",
|
|
"F5": "W_down",
|
|
"CF1": "FFN AR",
|
|
}
|
|
|
|
|
|
def _short_name(name: str) -> str:
|
|
"""Map 'S1 RMSNorm' -> 'RMSNorm'; fallback to first token."""
|
|
tok = name.split()[0]
|
|
return _SHORT_NAMES.get(tok, tok)
|
|
|
|
|
|
def _draw_arrow(ax, x0, x1, y):
|
|
ax.annotate("", xy=(x1, y), xytext=(x0, y),
|
|
arrowprops=dict(arrowstyle="->", color="#333", lw=1.1))
|
|
|
|
|
|
def _draw_ring_loop(ax, x_left, x_right, y_top, label: str, n_iters: int,
|
|
arc_lift: float = 0.9):
|
|
"""Draw a curved back-arrow above [x_left..x_right] with a repeat label.
|
|
|
|
Visualizes 'repeat this block n_iters times' (like a loop in a diagram).
|
|
arc_lift: how high above y_top the arc sits.
|
|
"""
|
|
color = "#5a189a"
|
|
lw = 1.6
|
|
arc_y = y_top + arc_lift
|
|
# Vertical tick up from the right edge
|
|
ax.plot([x_right, x_right], [y_top + 0.02, arc_y],
|
|
color=color, linewidth=lw)
|
|
# Horizontal top of the loop
|
|
ax.plot([x_right, x_left], [arc_y, arc_y],
|
|
color=color, linewidth=lw)
|
|
# Down-arrow into the left edge (a back-edge indicating the repeat)
|
|
ax.annotate("", xy=(x_left, y_top + 0.02), xytext=(x_left, arc_y),
|
|
arrowprops=dict(arrowstyle="->", color=color, lw=lw))
|
|
# Loop label centered above the arc
|
|
ax.text((x_left + x_right) / 2, arc_y + 0.08,
|
|
label,
|
|
ha="center", va="bottom",
|
|
fontsize=9, fontweight="bold", color=color)
|
|
# Iteration count as a badge on the right
|
|
ax.text(x_right + 0.08, arc_y - 0.03,
|
|
f"x {n_iters}",
|
|
ha="left", va="top",
|
|
fontsize=10, fontweight="bold", color=color,
|
|
bbox=dict(boxstyle="round,pad=0.15", facecolor="white",
|
|
edgecolor=color, linewidth=1.2, alpha=0.95))
|
|
|
|
|
|
def _draw_row(ax, y_base, row_title, stages: list[StageCost],
|
|
x_start: float, box_w: float, box_h: float, gap: float):
|
|
"""Draw one row of stage boxes with arrows between them."""
|
|
ax.text(x_start - 0.4, y_base + box_h / 2, row_title,
|
|
ha="right", va="center",
|
|
fontsize=11, fontweight="bold", color="#212529")
|
|
x = x_start
|
|
for i, st in enumerate(stages):
|
|
_draw_block(ax, x, y_base, box_w, box_h,
|
|
_short_name(st.name), st)
|
|
if i < len(stages) - 1:
|
|
_draw_arrow(ax, x + box_w + 0.02, x + box_w + gap - 0.02,
|
|
y_base + box_h / 2)
|
|
x += box_w + gap
|
|
return x - gap
|
|
|
|
|
|
def _draw_overlap_boxes(ax, hidden_stages, x_left, x_right, y_top,
|
|
gap_between: float = 0.15):
|
|
"""Draw dashed 'overlapped comm' boxes spanning x_left..x_right above y_top.
|
|
|
|
Boxes are widened to cover the horizontal range of the stages they
|
|
overlap with (e.g. C1 CP ring sits above S5-S7).
|
|
"""
|
|
if not hidden_stages:
|
|
return
|
|
hy = y_top + 0.10
|
|
n = len(hidden_stages)
|
|
total_span = x_right - x_left
|
|
box_span = (total_span - (n - 1) * gap_between) / max(1, n)
|
|
box_h_local = 0.42
|
|
hx = x_left
|
|
for st in hidden_stages:
|
|
box = patches.FancyBboxPatch(
|
|
(hx, hy), box_span, box_h_local,
|
|
boxstyle="round,pad=0.03",
|
|
facecolor=BOUND_COLOR["comm"], edgecolor="#333",
|
|
linewidth=1.0, alpha=0.35, linestyle="--",
|
|
)
|
|
ax.add_patch(box)
|
|
ax.text(hx + box_span / 2, hy + box_h_local * 0.65,
|
|
f"{_short_name(st.name)} (concurrent)",
|
|
ha="center", va="center",
|
|
fontsize=7, color="#7a0e0e", fontweight="bold")
|
|
ax.text(hx + box_span / 2, hy + box_h_local * 0.2,
|
|
_fmt_us(st.visible_s),
|
|
ha="center", va="center",
|
|
fontsize=6.5, color="#7a0e0e")
|
|
hx += box_span + gap_between
|
|
|
|
|
|
def draw_pipeline(cfg: FullConfig, ax=None):
|
|
"""Draw attention + FFN pipeline. Boxes colored by bound type."""
|
|
attn = all_stages(cfg)
|
|
ffn = all_ffn_stages(cfg)
|
|
|
|
# Split attention: main stages, then C2 (TP AllReduce). C1/C3 = overlapped.
|
|
attn_main = [s for s in attn if not s.name.startswith("C")]
|
|
attn_c2 = next((s for s in attn if s.name.startswith("C2")), None)
|
|
attn_hidden = [s for s in attn
|
|
if s.name.startswith("C1") or s.name.startswith("C3")]
|
|
attn_hidden = [s for s in attn_hidden if s.visible_s > 0]
|
|
attn_row = attn_main + ([attn_c2] if attn_c2 and attn_c2.visible_s > 0 else [])
|
|
|
|
# FFN row: all F* + CF1 if non-trivial
|
|
ffn_main = [s for s in ffn if not s.name.startswith("CF")]
|
|
ffn_cf = next((s for s in ffn if s.name.startswith("CF1")), None)
|
|
ffn_row = ffn_main + ([ffn_cf] if ffn_cf and ffn_cf.visible_s > 0 else [])
|
|
|
|
n_max = max(len(attn_row), len(ffn_row))
|
|
box_w = 1.35
|
|
box_h = 0.78
|
|
gap = 0.18
|
|
x_start = 1.1
|
|
|
|
total_w = x_start + n_max * box_w + (n_max - 1) * gap + 0.5
|
|
total_h = 4.2
|
|
|
|
if ax is None:
|
|
fig, ax = plt.subplots(figsize=(max(11, total_w * 1.1),
|
|
max(4.5, total_h * 1.0)))
|
|
else:
|
|
fig = ax.figure
|
|
|
|
_draw_row(ax, y_base=2.6, row_title="Attention",
|
|
stages=attn_row, x_start=x_start,
|
|
box_w=box_w, box_h=box_h, gap=gap)
|
|
|
|
_draw_row(ax, y_base=0.6, row_title="FFN",
|
|
stages=ffn_row, x_start=x_start,
|
|
box_w=box_w, box_h=box_h, gap=gap)
|
|
|
|
# Ring loop indicator + concurrent-comm boxes anchored over S5..S8.
|
|
_idx_by_prefix = {s.name.split()[0]: i for i, s in enumerate(attn_row)}
|
|
_s5_idx = _idx_by_prefix.get("S5")
|
|
_s7_idx = _idx_by_prefix.get("S7")
|
|
_s8_idx = _idx_by_prefix.get("S8")
|
|
|
|
if cfg.topo.cp > 1 and cfg.topo.mode == "prefill" and _s5_idx is not None:
|
|
_right_idx = _s8_idx if _s8_idx is not None else _s7_idx
|
|
_lx = x_start + _s5_idx * (box_w + gap)
|
|
_rx = x_start + _right_idx * (box_w + gap) + box_w
|
|
_y_top = 2.6 + box_h
|
|
# Concurrent-comm boxes (C1 CP ring, C3 score AR) directly above S5-S7
|
|
_draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
|
|
# Ring loop arc above those
|
|
_draw_ring_loop(
|
|
ax, _lx, _rx, _y_top,
|
|
label="ring attention loop (K/V ring per hop; S8 merges partial O,m,l)",
|
|
n_iters=cfg.topo.cp,
|
|
arc_lift=1.05,
|
|
)
|
|
elif cfg.topo.cp > 1 and cfg.topo.mode == "decode" and _s5_idx is not None:
|
|
_right_idx = _s7_idx if _s7_idx is not None else _s5_idx
|
|
_lx = x_start + _s5_idx * (box_w + gap)
|
|
_rx = x_start + _right_idx * (box_w + gap) + box_w
|
|
_y_top = 2.6 + box_h
|
|
# Any straggling concurrent comm (e.g. C3 if head-split active) above S5-S7
|
|
_draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
|
|
_draw_ring_loop(
|
|
ax, _lx, _rx, _y_top,
|
|
label="decode: local pass only; S8 fires O/m/l all-reduce",
|
|
n_iters=1,
|
|
arc_lift=1.05 if attn_hidden else 0.85,
|
|
)
|
|
elif attn_hidden:
|
|
# Fallback: CP=1 but still have C3 or similar; anchor above S5-S7 if present
|
|
if _s5_idx is not None and _s7_idx is not None:
|
|
_lx = x_start + _s5_idx * (box_w + gap)
|
|
_rx = x_start + _s7_idx * (box_w + gap) + box_w
|
|
_draw_overlap_boxes(ax, attn_hidden, _lx, _rx, 2.6 + box_h)
|
|
|
|
# Totals
|
|
attn_total = sum(s.visible_s for s in attn_row) \
|
|
+ sum(s.visible_s for s in attn_hidden)
|
|
ffn_total = sum(s.visible_s for s in ffn_row)
|
|
ax.text(total_w - 0.3, 2.6 + box_h / 2 + 1.15,
|
|
f"Attn total: {_fmt_us(attn_total)}",
|
|
ha="right", va="bottom", fontsize=9,
|
|
fontweight="bold", color="#212529")
|
|
ax.text(total_w - 0.3, 0.6 + box_h + 0.15,
|
|
f"FFN total: {_fmt_us(ffn_total)}",
|
|
ha="right", va="bottom", fontsize=9,
|
|
fontweight="bold", color="#212529")
|
|
|
|
# Legend
|
|
handles = []
|
|
for k, lbl in BOUND_LABEL.items():
|
|
handles.append(patches.Patch(
|
|
facecolor=BOUND_COLOR[k], edgecolor="#333",
|
|
alpha=0.85, label=lbl,
|
|
))
|
|
ax.legend(handles=handles, loc="upper center",
|
|
bbox_to_anchor=(0.5, -0.02),
|
|
ncol=4, fontsize=8, frameon=False)
|
|
|
|
ax.set_xlim(0, total_w + 0.6)
|
|
ax.set_ylim(0, 4.9)
|
|
ax.set_aspect("auto")
|
|
ax.axis("off")
|
|
ax.set_title(
|
|
f"Per-layer pipeline (one PE, {cfg.topo.mode}, T_q={cfg.topo.T_q}, "
|
|
f"S_local={cfg.topo.s_local}) - color = dominant bound",
|
|
fontsize=10, fontweight="bold",
|
|
)
|
|
return fig
|