"""Pictorial view of how each weight/KV tensor is sharded across ranks. Every tensor is drawn as a rectangle labelled with its (rows, cols) shape. Grid lines show shard boundaries; one shard (this PE's share) is highlighted so it's obvious which dim the split is along. Sharding pattern: - W_Q (hidden, H_q*d_head) : split on **output cols** across TP - W_K (hidden, H_kv*d_head): split on **output cols** across TP - W_V (hidden, H_kv*d_head): split on **output cols** across TP - W_O (H_q*d_head, hidden) : split on **input rows** across TP - W_gate/up (hidden, ffn_dim): split on **output cols** across FFN scope - W_down (ffn_dim, hidden) : split on **input rows** across FFN scope - K/V cache (S_kv, H_kv*d_head): split on **rows (S axis)** by CP AND **cols (head axis)** by TP -> 2D shard grid """ from __future__ import annotations import matplotlib.pyplot as plt import matplotlib.patches as patches from .model_config import FullConfig _UNSHARDED_FACE = "#e9ecef" _UNSHARDED_EDGE = "#adb5bd" _MY_SHARD_FACE = "#3a86ff" _MY_SHARD_EDGE = "#0057b3" _TP_LINE_COLOR = "#d90429" # TP split lines _CP_LINE_COLOR = "#0f9d58" # CP split lines _FFN_LINE_COLOR = "#e37400" # FFN-scope split lines def _draw_tensor(ax, x, y, w, h, name, shape_str, row_splits: int = 1, col_splits: int = 1, my_row_shard: int = 0, my_col_shard: int = 0, row_line_color: str = _TP_LINE_COLOR, col_line_color: str = _TP_LINE_COLOR, row_label: str = "", col_label: str = "", note: str = "", cell_annot=None): """Draw a tensor with sharding grid. cell_annot: optional callable (row, col) -> str returning the physical label to render inside shard cell (r, c). If it returns "", cell is not annotated. """ """Draw a tensor as a rectangle with grid lines showing shards.""" # Base rectangle (unsharded background) rect = patches.Rectangle( (x, y), w, h, facecolor=_UNSHARDED_FACE, edgecolor=_UNSHARDED_EDGE, linewidth=1.0, ) ax.add_patch(rect) # Highlight this PE's shard shard_w = w / max(1, col_splits) shard_h = h / max(1, row_splits) sx = x + my_col_shard * shard_w sy = y + (row_splits - 1 - my_row_shard) * shard_h highlight = patches.Rectangle( (sx, sy), shard_w, shard_h, facecolor=_MY_SHARD_FACE, edgecolor=_MY_SHARD_EDGE, linewidth=1.5, alpha=0.8, ) ax.add_patch(highlight) # Horizontal grid lines (row splits) for r in range(1, row_splits): yy = y + r * shard_h ax.plot([x, x + w], [yy, yy], color=row_line_color, linewidth=1.2, linestyle="--") # Vertical grid lines (col splits) for c in range(1, col_splits): xx = x + c * shard_w ax.plot([xx, xx], [y, y + h], color=col_line_color, linewidth=1.2, linestyle="--") # Per-cell physical annotations (e.g. "cube3 PE5") if provided. if cell_annot is not None: for r in range(row_splits): for c in range(col_splits): text = cell_annot(r, c) if not text: continue cx = x + c * shard_w + shard_w / 2 cy = y + (row_splits - 1 - r) * shard_h + shard_h / 2 ax.text(cx, cy, text, ha="center", va="center", fontsize=6.0, color="#0d1b2a", alpha=0.85) # Stacked labels ABOVE rectangle: note (italic, top), shape, name (bold, closest) if note: ax.text(x + w / 2, y + h + 0.55, note, ha="center", va="bottom", fontsize=6.5, color="#555", style="italic") ax.text(x + w / 2, y + h + 0.30, shape_str, ha="center", va="bottom", fontsize=7, color="#666") ax.text(x + w / 2, y + h + 0.08, name, ha="center", va="bottom", fontsize=10, fontweight="bold") # Split labels BELOW rectangle (stacked if both present) if row_label and col_label: ax.text(x + w / 2, y - 0.10, row_label, ha="center", va="top", fontsize=6.5, color=row_line_color, fontweight="bold") ax.text(x + w / 2, y - 0.35, col_label, ha="center", va="top", fontsize=6.5, color=col_line_color, fontweight="bold") elif row_label: ax.text(x + w / 2, y - 0.10, row_label, ha="center", va="top", fontsize=6.5, color=row_line_color, fontweight="bold") elif col_label: ax.text(x + w / 2, y - 0.10, col_label, ha="center", va="top", fontsize=6.5, color=col_line_color, fontweight="bold") def _physical_label(cfg: FullConfig, dim: str, rank: int) -> str: """Where does rank `rank` of dim (`tp`|`cp`|`ffn`) live physically? Returns a short "cubeX PEy" style label based on the placement config. """ topo = cfg.topo if dim == "tp": if topo.tp_placement == "pe": # TP fills PEs within one cube (spilling if tp > 8) pe_in_cube = rank % topo.pes_per_cube_hw cube = rank // topo.pes_per_cube_hw return f"c{cube}p{pe_in_cube}" else: # tp on cube return f"cube{rank}" if dim == "cp": if topo.cp_placement == "pe": pe_in_cube = rank % topo.pes_per_cube_hw cube = rank // topo.pes_per_cube_hw return f"c{cube}p{pe_in_cube}" else: # cp on cube return f"cube{rank}" return "" def _kv_cell_annot(cfg: FullConfig): """Return an annotation function for KV cache cells: (cp_r, tp_r) -> label.""" def _ann(r, c): cp_lbl = _physical_label(cfg, "cp", r) tp_lbl = _physical_label(cfg, "tp", c) # Two lines if both are non-trivial; otherwise a single line if cp_lbl and tp_lbl: return f"{cp_lbl}\n{tp_lbl}" return cp_lbl or tp_lbl return _ann def _tp_cell_annot(cfg: FullConfig): """Col-parallel tensors (W_Q/W_K/W_V/W_gate/W_up): each col = one TP rank.""" def _ann(r, c): return _physical_label(cfg, "tp", c) return _ann def _tp_row_cell_annot(cfg: FullConfig): """Row-parallel tensors split by TP (W_O): each row = one TP rank.""" def _ann(r, c): return _physical_label(cfg, "tp", r) return _ann def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0, show_physical: bool = False): """Draw all tensors with their sharding overlays for PE `my_pe`. show_physical: annotate each shard cell with the owning cube/PE. """ m = cfg.model tp = cfg.topo.tp cp = cfg.topo.cp # PE's assigned shard index for each dim q_shard = my_pe % tp # W_Q col shard kv_shard_col = min(my_pe % tp, m.h_kv - 1) if m.h_kv >= tp else 0 ffn_shard = my_pe % cfg.ffn_shard_divisor kv_row_shard = 0 # per-CP; showing CP=0 view for clarity # Enlarged mode gives more room for physical annotations. _fig_w, _fig_h = (22, 13) if show_physical else (15, 9) _cell_w = 3.4 if show_physical else 2.6 _cell_h = 3.0 if show_physical else 2.2 _gap_x = 0.8 if show_physical else 0.55 _gap_y = 2.0 if show_physical else 1.7 if ax is None: fig, ax = plt.subplots(figsize=(_fig_w, _fig_h)) else: fig = ax.figure # Grid layout: 2 rows x 4 cols of tensor rects # Attention row: W_Q, W_K, W_V, W_O # FFN + KV row: W_gate, W_up, W_down, KV cache cell_w = _cell_w cell_h = _cell_h gap_x = _gap_x gap_y = _gap_y _kv_ann = _kv_cell_annot(cfg) if show_physical else None _tp_col_ann = _tp_cell_annot(cfg) if show_physical else None _tp_row_ann = _tp_row_cell_annot(cfg) if show_physical else None positions = [ (0, 1), # W_Q at row 0 (1, 1), # W_K (2, 1), # W_V (3, 1), # W_O (0, 0), # W_gate at row 1 (1, 0), # W_up (2, 0), # W_down (3, 0), # KV cache ] # ── attention tensors ──────────────────────── tp_txt = f"TP={tp}" ffn_txt = f"{cfg.topo.ffn_shard_scope}={cfg.ffn_shard_divisor}" cp_txt = f"CP={cp}" def px(col): return col * (cell_w + gap_x) def py(row): return row * (cell_h + gap_y) # W_Q: (hidden, H_q*d_head), split on COL by TP _draw_tensor(ax, px(0), py(1), cell_w, cell_h, "W_Q", f"({m.hidden}, {m.h_q * m.d_head})", row_splits=1, col_splits=tp, my_col_shard=q_shard, col_line_color=_TP_LINE_COLOR, col_label=f"cols split by {tp_txt}", note="col-parallel (output dim)", cell_annot=_tp_col_ann) # W_K, W_V: (hidden, H_kv*d_head) for (name, col_idx) in [("W_K", 1), ("W_V", 2)]: splits = min(tp, m.h_kv) if m.h_kv >= tp else tp _draw_tensor(ax, px(col_idx), py(1), cell_w, cell_h, name, f"({m.hidden}, {m.h_kv * m.d_head})", row_splits=1, col_splits=splits, my_col_shard=min(q_shard, splits - 1), col_line_color=_TP_LINE_COLOR, col_label=f"cols split by {tp_txt}", note="col-parallel (heads)", cell_annot=_tp_col_ann) # W_O: (H_q*d_head, hidden), split on ROW by TP _draw_tensor(ax, px(3), py(1), cell_w, cell_h, "W_O", f"({m.h_q * m.d_head}, {m.hidden})", row_splits=tp, col_splits=1, my_row_shard=q_shard, row_line_color=_TP_LINE_COLOR, row_label=f"rows split by {tp_txt}", note="row-parallel (input dim)", cell_annot=_tp_row_ann) # ── FFN tensors ────────────────────────────── # W_gate, W_up: (hidden, ffn_dim), split on COL by FFN scope for (name, col_idx) in [("W_gate", 0), ("W_up", 1)]: _draw_tensor(ax, px(col_idx), py(0), cell_w, cell_h, name, f"({m.hidden}, {m.ffn_dim})", row_splits=1, col_splits=cfg.ffn_shard_divisor, my_col_shard=ffn_shard, col_line_color=_FFN_LINE_COLOR, col_label=f"cols split by {ffn_txt}", note="col-parallel") # W_down: (ffn_dim, hidden), split on ROW by FFN scope _draw_tensor(ax, px(2), py(0), cell_w, cell_h, "W_down", f"({m.ffn_dim}, {m.hidden})", row_splits=cfg.ffn_shard_divisor, col_splits=1, my_row_shard=ffn_shard, row_line_color=_FFN_LINE_COLOR, row_label=f"rows split by {ffn_txt}", note="row-parallel") # ── KV cache: (B, S_kv, H_kv*d_head), 2D shard on S_kv x heads ─ # Batch (B) is a stacking dimension — each concurrent request holds # its own KV slice per PE, so total KV per PE = B * (per-slice size). kv_row_splits = cp kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp _B = max(1, cfg.topo.b) _batch_suffix = f", B={_B}" if _B > 1 else "" _batch_note = (f"; batch B={_B} => {_B}x KV bytes per PE" if _B > 1 else "") # Draw B-1 offset "shadow" rectangles behind the KV cache to give # a visual hint of the batch stacking dimension. Cap at 5 shadows # so a large B doesn't overwhelm the diagram. if _B > 1: _kv_shadow_x = px(3) _kv_shadow_y = py(0) _n_shadows = min(_B - 1, 5) _offset_step = 0.15 for _si in range(_n_shadows, 0, -1): _dx = _si * _offset_step _dy = _si * _offset_step _shadow = patches.Rectangle( (_kv_shadow_x + _dx, _kv_shadow_y + _dy), cell_w, cell_h, facecolor="#f5f5f5", edgecolor="#adb5bd", linewidth=0.8, linestyle="--", alpha=0.55, ) ax.add_patch(_shadow) _draw_tensor(ax, px(3), py(0), cell_w, cell_h, "KV cache (K and V)", f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,}" f"{_batch_suffix})", row_splits=kv_row_splits, col_splits=kv_col_splits if show_physical else 1, my_row_shard=kv_row_shard, row_line_color=_CP_LINE_COLOR, row_label=f"rows: S split by {cp_txt}", col_label=f"cols: heads split by {tp_txt}", col_line_color=_TP_LINE_COLOR, note=f"2D shard: seq axis (CP) x head axis (TP)" f"{_batch_note}", cell_annot=_kv_ann) # In non-physical mode, KV was drawn with col_splits=1; overlay TP col # splits as dotted lines to hint at the 2D nature. In physical mode # the _draw_tensor call already drew proper vertical dashes. if not show_physical: kv_x = px(3) kv_y = py(0) shard_col_w = cell_w / max(1, kv_col_splits) for c in range(1, kv_col_splits): xx = kv_x + c * shard_col_w ax.plot([xx, xx], [kv_y, kv_y + cell_h], color=_TP_LINE_COLOR, linewidth=1.2, linestyle=":") ax.set_xlim(-0.3, 4 * (cell_w + gap_x) - gap_x + 0.3) ax.set_ylim(-0.9, 2 * (cell_h + gap_y) - gap_y + 0.9) ax.set_aspect("auto") ax.axis("off") _B_title = max(1, cfg.topo.b) _b_title_suffix = f", B={_B_title}" if _B_title > 1 else "" ax.set_title( f"Tensor sharding view (PE {my_pe}) | TP={tp}, CP={cp}" f"{_b_title_suffix}, " f"FFN scope={cfg.topo.ffn_shard_scope} | " f"blue = this PE's shard", fontsize=11, fontweight="bold", ) # Legend for line colors ax.plot([], [], color=_TP_LINE_COLOR, linewidth=1.2, linestyle="--", label=f"TP split ({tp} ranks)") ax.plot([], [], color=_CP_LINE_COLOR, linewidth=1.2, linestyle="--", label=f"CP split ({cp} ranks, sequence dim)") ax.plot([], [], color=_FFN_LINE_COLOR, linewidth=1.2, linestyle="--", label=f"FFN scope split ({cfg.ffn_shard_divisor} ranks)") ax.legend(loc="lower center", bbox_to_anchor=(0.5, -0.05), ncol=3, fontsize=8, frameon=False) return fig