"""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