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>
222 lines
8.5 KiB
Python
222 lines
8.5 KiB
Python
"""Draw a PE-level view of one CP group (one TP group of cubes).
|
||
|
||
Shows how weights + KV cache are distributed across the PEs of ONE CP
|
||
rank. Attention weights are sharded by head (H_q, H_kv split across TP);
|
||
FFN is sharded by dim (ffn_dim / TP / EP); KV cache is sharded by CP
|
||
(sequence) × TP (head).
|
||
|
||
If TP > 8 the group spans multiple cubes — layout draws all of them.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import matplotlib.pyplot as plt
|
||
import matplotlib.patches as patches
|
||
|
||
from .model_config import FullConfig
|
||
|
||
|
||
PE_COLS = 4
|
||
PE_ROWS = 2
|
||
PES_PER_CUBE = PE_COLS * PE_ROWS
|
||
|
||
|
||
def layers_per_stage_val(cfg: FullConfig) -> int:
|
||
return (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
||
|
||
|
||
def _q_heads_for_pe(pe_id_in_group: int, tp: int, h_q: int) -> list[int]:
|
||
"""Return the list of Q-head indices assigned to this PE."""
|
||
heads_per_pe = h_q / tp
|
||
start = int(round(pe_id_in_group * heads_per_pe))
|
||
end = int(round((pe_id_in_group + 1) * heads_per_pe))
|
||
return list(range(start, end))
|
||
|
||
|
||
def _kv_heads_for_pe(pe_id_in_group: int, tp: int, h_kv: int,
|
||
kv_shard_mode: str) -> tuple[list[int], str]:
|
||
"""KV heads assigned to this PE and a note about replication/split."""
|
||
if h_kv >= tp:
|
||
heads_per_pe = h_kv // tp
|
||
start = pe_id_in_group * heads_per_pe
|
||
end = start + heads_per_pe
|
||
return list(range(start, end)), ""
|
||
# TP > H_kv
|
||
if kv_shard_mode == "replicate":
|
||
rep_factor = tp // h_kv
|
||
head_idx = pe_id_in_group // rep_factor
|
||
return [head_idx], f"replicated x{rep_factor}"
|
||
# split mode
|
||
split_factor = tp // h_kv
|
||
head_idx = pe_id_in_group // split_factor
|
||
part = pe_id_in_group % split_factor
|
||
return [head_idx], f"head-split ({part+1}/{split_factor})"
|
||
|
||
|
||
def _per_pe_bytes(cfg: FullConfig, pe_id_in_group: int) -> dict:
|
||
"""Return per-tensor bytes for one PE within a TP group."""
|
||
m = cfg.model
|
||
tp = cfg.topo.tp
|
||
pp = cfg.topo.pp
|
||
ep = max(1, cfg.topo.ep)
|
||
b = m.bytes_per_elem
|
||
layers_per_stage = (m.layers + pp - 1) // pp
|
||
|
||
hq_per_pe = m.h_q / tp
|
||
if cfg.topo.kv_shard_mode == "replicate":
|
||
hkv_per_pe = max(1.0, m.h_kv / tp)
|
||
else:
|
||
hkv_per_pe = m.h_kv / tp
|
||
ffn_div = cfg.ffn_shard_divisor * ep
|
||
ffn_per_pe = m.ffn_dim // ffn_div
|
||
|
||
per_layer = {
|
||
"W_Q": int(m.hidden * hq_per_pe * m.d_head * b),
|
||
"W_K": int(m.hidden * hkv_per_pe * m.d_head * b),
|
||
"W_V": int(m.hidden * hkv_per_pe * m.d_head * b),
|
||
"W_O": int(hq_per_pe * m.d_head * m.hidden * b),
|
||
"W_gate": int(m.hidden * ffn_per_pe * b),
|
||
"W_up": int(m.hidden * ffn_per_pe * b),
|
||
"W_down": int(ffn_per_pe * m.hidden * b),
|
||
}
|
||
kv_per_layer = int(2 * cfg.topo.s_local * hkv_per_pe * m.d_head * b)
|
||
|
||
all_layers = {k: v * layers_per_stage for k, v in per_layer.items()}
|
||
all_layers["KV cache"] = kv_per_layer * layers_per_stage
|
||
# transient estimate
|
||
if cfg.topo.mode == "decode":
|
||
all_layers["Transient"] = int(4 * (m.hidden + hq_per_pe * m.d_head) * b)
|
||
else:
|
||
TILE = 1024
|
||
all_layers["Transient"] = int(2 * hq_per_pe * cfg.topo.T_q * TILE * b
|
||
+ cfg.topo.T_q * m.hidden * b)
|
||
return all_layers
|
||
|
||
|
||
def draw_pe_layout(cfg: FullConfig, ax=None):
|
||
"""Draw one CP group's PEs with per-tensor weight + KV breakdown.
|
||
|
||
Each PE shows: PE id, Q heads, KV heads, W_Q / W_K / W_V / W_O in
|
||
MB, FFN (gate+up+down) in MB, KV cache in GB, total in GB.
|
||
|
||
TP=8 -> 1 cube. TP=16 -> 2 cubes. TP=32 -> 4 cubes.
|
||
Under placement, n_cubes reflects cubes-per-group (may include CP spill).
|
||
"""
|
||
tp = cfg.topo.tp
|
||
# cubes per one cube-level group (= intra-cube dims / PEs-per-cube, rounded up).
|
||
n_cubes = max(1, (cfg.topo.intra_cube_dims + PES_PER_CUBE - 1) // PES_PER_CUBE)
|
||
|
||
if ax is None:
|
||
fig, ax = plt.subplots(figsize=(7.5 * n_cubes, 8.5))
|
||
else:
|
||
fig = ax.figure
|
||
|
||
cube_w = 7.0
|
||
cube_h = 6.0
|
||
cube_gap = 0.6
|
||
|
||
for cube_idx in range(n_cubes):
|
||
x0 = cube_idx * (cube_w + cube_gap)
|
||
y0 = 0
|
||
|
||
rect = patches.FancyBboxPatch(
|
||
(x0, y0), cube_w, cube_h,
|
||
boxstyle="round,pad=0.05",
|
||
facecolor="#f8f9fa", edgecolor="#212529", linewidth=1.5,
|
||
)
|
||
ax.add_patch(rect)
|
||
ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
|
||
f"Cube {cube_idx} (PEs {cube_idx*8}-{cube_idx*8+7})",
|
||
ha="center", va="bottom", fontsize=11, fontweight="bold")
|
||
|
||
pe_pad_x = 0.15
|
||
pe_pad_y = 0.2
|
||
pe_gap = 0.10
|
||
inner_w = cube_w - 2 * pe_pad_x
|
||
inner_h = cube_h - 2 * pe_pad_y
|
||
pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
|
||
pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
|
||
|
||
for pr in range(PE_ROWS):
|
||
for pc in range(PE_COLS):
|
||
pe_local = pr * PE_COLS + pc
|
||
pe_id_in_group = cube_idx * PES_PER_CUBE + pe_local
|
||
if pe_id_in_group >= cfg.topo.intra_cube_dims:
|
||
continue
|
||
px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
|
||
py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
|
||
|
||
pe_rect = patches.Rectangle(
|
||
(px, py), pe_w, pe_h,
|
||
facecolor="#f0f7ff", edgecolor="#3a86ff", linewidth=1.0,
|
||
)
|
||
ax.add_patch(pe_rect)
|
||
|
||
q_heads = _q_heads_for_pe(pe_id_in_group, tp, cfg.model.h_q)
|
||
kv_heads, kv_note = _kv_heads_for_pe(
|
||
pe_id_in_group, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
|
||
)
|
||
bytes_ = _per_pe_bytes(cfg, pe_id_in_group)
|
||
weights_gb = sum(v for k, v in bytes_.items()
|
||
if k not in ("KV cache", "Transient")) / 1e9
|
||
kv_gb = bytes_["KV cache"] / 1e9
|
||
trans_mb = bytes_["Transient"] / 1e6
|
||
total_gb = sum(bytes_.values()) / 1e9
|
||
|
||
q_str = (f"{q_heads[0]}-{q_heads[-1]}"
|
||
if len(q_heads) > 1 else str(q_heads[0])
|
||
if q_heads else "-")
|
||
kv_str = str(kv_heads[0]) if kv_heads else "-"
|
||
if kv_note:
|
||
kv_str += f"({kv_note})"
|
||
|
||
ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
|
||
+ bytes_["W_down"]) / 1e6
|
||
# Header (PE id and heads) bigger, then per-tensor breakdown
|
||
header = (f"PE {pe_id_in_group}\n"
|
||
f"Q heads: {q_str}\n"
|
||
f"KV head: {kv_str}")
|
||
ax.text(px + 0.08, py + pe_h - 0.05, header,
|
||
ha="left", va="top",
|
||
fontsize=7.0, fontweight="bold",
|
||
family="monospace", color="#0d47a1")
|
||
|
||
# Per-tensor bytes with explicit division divisors.
|
||
# Divisor is always the PRODUCT of the parallel dims listed
|
||
# (TP*CP means TP times CP, i.e. multiplication).
|
||
tp_d = cfg.topo.tp
|
||
ffn_div = cfg.ffn_shard_divisor
|
||
ffn_scope = cfg.topo.ffn_shard_scope.replace("+", "*")
|
||
kv_div = cfg.topo.cp * cfg.topo.tp
|
||
L = layers_per_stage_val(cfg)
|
||
detail = (
|
||
f"weights (all {L} layers):\n"
|
||
f" W_Q /TP={tp_d} : {bytes_['W_Q']/1e6:7.1f}MB\n"
|
||
f" W_K /TP={tp_d} : {bytes_['W_K']/1e6:7.1f}MB\n"
|
||
f" W_V /TP={tp_d} : {bytes_['W_V']/1e6:7.1f}MB\n"
|
||
f" W_O /TP={tp_d} : {bytes_['W_O']/1e6:7.1f}MB\n"
|
||
f" FFN /({ffn_scope})={ffn_div}: {ffn_mb:7.1f}MB\n"
|
||
f"runtime:\n"
|
||
f" KV /(CP*TP)={kv_div}: {kv_gb*1000:7.1f}MB\n"
|
||
f" trans : {trans_mb:7.1f}MB\n"
|
||
f"W total: {weights_gb*1000:7.1f}MB\n"
|
||
f"TOTAL : {total_gb*1000:7.1f}MB"
|
||
)
|
||
ax.text(px + pe_w / 2, py + 0.06, detail,
|
||
ha="center", va="bottom",
|
||
fontsize=5.8, family="monospace", color="#1a1a1a")
|
||
|
||
total_w = n_cubes * cube_w + (n_cubes - 1) * cube_gap
|
||
ax.set_xlim(-0.2, total_w + 0.2)
|
||
ax.set_ylim(-0.3, cube_h + 0.6)
|
||
ax.set_aspect("equal")
|
||
ax.axis("off")
|
||
|
||
title = (
|
||
f"Per-PE layout of one CP group "
|
||
f"(TP={tp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
|
||
f" | KV mode: {cfg.topo.kv_shard_mode}"
|
||
)
|
||
ax.set_title(title, fontsize=11, fontweight="bold")
|
||
|
||
return fig
|