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>
236 lines
9.1 KiB
Python
236 lines
9.1 KiB
Python
"""Replication (waste) accounting + optimization opportunities.
|
|
|
|
Two things this module produces:
|
|
|
|
1. `replication_report(cfg)` — where memory is being *duplicated* across
|
|
the deployment, per category (attention weights, FFN weights, KV cache
|
|
under replicate mode, DP replicas). Answers "if I lifted this
|
|
duplication, how much would I save?"
|
|
|
|
2. `optimization_hints(cfg)` — a list of actionable hints. Each hint has
|
|
a `category` (space | comm | compute), a `severity` (info | warn |
|
|
good), and a short human message. UI code renders these as bullets.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from .model_config import FullConfig
|
|
from .memory_layout import per_pe_weight_bytes, per_pe_kv_cache_bytes
|
|
|
|
|
|
# ── Replication accounting ────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class ReplicationEntry:
|
|
tensor: str
|
|
per_pe_bytes: int
|
|
replicated_across: str # e.g. "CP (x4 copies)"
|
|
copies: int # total # copies globally
|
|
wasted_bytes: int # (copies - 1) * per_pe_bytes (per PE view)
|
|
|
|
|
|
def _attn_weight_bytes_per_pe(cfg: FullConfig) -> int:
|
|
"""W_Q + W_K + W_V + W_O per PE, one layer."""
|
|
m = cfg.model
|
|
tp = cfg.topo.tp
|
|
hq_per_pe = cfg.h_q_per_pe
|
|
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
|
|
per_layer = (
|
|
m.hidden * hq_per_pe * m.d_head +
|
|
m.hidden * hkv_per_pe * m.d_head +
|
|
m.hidden * hkv_per_pe * m.d_head +
|
|
hq_per_pe * m.d_head * m.hidden
|
|
) * m.bytes_per_elem
|
|
layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
|
return int(per_layer * layers_per_stage)
|
|
|
|
|
|
def _ffn_weight_bytes_per_pe(cfg: FullConfig) -> int:
|
|
"""W_gate + W_up + W_down per PE, all local layers."""
|
|
m = cfg.model
|
|
ep = max(1, cfg.topo.ep)
|
|
ffn_div = cfg.ffn_shard_divisor * ep
|
|
per_layer = 3 * m.hidden * (m.ffn_dim // max(1, ffn_div)) * m.bytes_per_elem
|
|
layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
|
return int(per_layer * layers_per_stage)
|
|
|
|
|
|
def _kv_bytes_per_pe(cfg: FullConfig) -> int:
|
|
return per_pe_kv_cache_bytes(cfg)
|
|
|
|
|
|
def replication_report(cfg: FullConfig) -> list[ReplicationEntry]:
|
|
"""Enumerate replicated (duplicated) tensors and their waste."""
|
|
entries: list[ReplicationEntry] = []
|
|
topo = cfg.topo
|
|
|
|
# Attention weights are NOT sharded by CP - so each CP rank holds a copy.
|
|
attn_pe = _attn_weight_bytes_per_pe(cfg)
|
|
cp_copies = topo.cp
|
|
if cp_copies > 1:
|
|
entries.append(ReplicationEntry(
|
|
tensor="Attn weights (W_Q/W_K/W_V/W_O)",
|
|
per_pe_bytes=attn_pe,
|
|
replicated_across=f"CP (x{cp_copies} identical group copies)",
|
|
copies=cp_copies,
|
|
wasted_bytes=attn_pe * (cp_copies - 1),
|
|
))
|
|
|
|
# FFN weights: replicated across scope levels NOT included.
|
|
# Scope=TP -> replicated across CP AND DP.
|
|
# Scope=TP+CP -> replicated across DP only.
|
|
# Scope=TP+CP+DP -> no replication (perfect share).
|
|
ffn_pe = _ffn_weight_bytes_per_pe(cfg)
|
|
scope = topo.ffn_shard_scope
|
|
ffn_cp_copies = 1 if "CP" in scope else topo.cp
|
|
ffn_dp_copies = 1 if "DP" in scope else topo.dp
|
|
ffn_total_copies = ffn_cp_copies * ffn_dp_copies
|
|
if ffn_total_copies > 1:
|
|
reasons = []
|
|
if ffn_cp_copies > 1:
|
|
reasons.append(f"CP (x{ffn_cp_copies})")
|
|
if ffn_dp_copies > 1:
|
|
reasons.append(f"DP (x{ffn_dp_copies})")
|
|
entries.append(ReplicationEntry(
|
|
tensor=f"FFN weights (scope={scope})",
|
|
per_pe_bytes=ffn_pe,
|
|
replicated_across=" & ".join(reasons),
|
|
copies=ffn_total_copies,
|
|
wasted_bytes=ffn_pe * (ffn_total_copies - 1),
|
|
))
|
|
|
|
# KV cache: if kv_shard_mode='replicate' and TP > H_kv, each KV head
|
|
# is duplicated tp // h_kv times across TP ranks.
|
|
if (cfg.kv_replication_needed
|
|
and topo.kv_shard_mode == "replicate"):
|
|
kv_pe = _kv_bytes_per_pe(cfg)
|
|
rep = cfg.head_dim_split_factor # tp // h_kv
|
|
entries.append(ReplicationEntry(
|
|
tensor="KV cache (replicate mode)",
|
|
per_pe_bytes=kv_pe,
|
|
replicated_across=f"TP (x{rep} copies of each KV head)",
|
|
copies=rep,
|
|
wasted_bytes=kv_pe * (rep - 1),
|
|
))
|
|
|
|
# DP replicas duplicate the entire per-PE footprint (attn + ffn + KV).
|
|
if topo.dp > 1:
|
|
total_pe = attn_pe + ffn_pe + _kv_bytes_per_pe(cfg)
|
|
entries.append(ReplicationEntry(
|
|
tensor="Full per-PE footprint (attn + FFN + KV)",
|
|
per_pe_bytes=total_pe,
|
|
replicated_across=f"DP (x{topo.dp} model replicas)",
|
|
copies=topo.dp,
|
|
wasted_bytes=total_pe * (topo.dp - 1),
|
|
))
|
|
|
|
return entries
|
|
|
|
|
|
# ── Optimization hints ────────────────────────────────────────────
|
|
|
|
@dataclass
|
|
class Hint:
|
|
category: str # "space" | "comm" | "compute" | "layout"
|
|
severity: str # "info" | "warn" | "good"
|
|
message: str
|
|
|
|
|
|
def optimization_hints(cfg: FullConfig) -> list[Hint]:
|
|
hints: list[Hint] = []
|
|
m = cfg.model
|
|
topo = cfg.topo
|
|
|
|
# SPACE hints -------------------------------------------------
|
|
scope = topo.ffn_shard_scope
|
|
ffn_pe = _ffn_weight_bytes_per_pe(cfg)
|
|
if scope == "TP" and topo.cp > 1:
|
|
current = ffn_pe
|
|
new = current // topo.cp
|
|
hints.append(Hint(
|
|
"space", "warn",
|
|
f"FFN scope=TP replicates FFN weights across all {topo.cp} CP "
|
|
f"groups. Switching to **TP+CP** would drop per-PE FFN weight "
|
|
f"from {current/1e9:.2f} GB to {new/1e9:.2f} GB "
|
|
f"(saves {(current-new)/1e9:.2f} GB/PE) at the cost of one "
|
|
f"extra AllReduce per layer."
|
|
))
|
|
if scope in ("TP", "TP+CP") and topo.dp > 1:
|
|
base = ffn_pe if scope == "TP+CP" else ffn_pe // topo.cp
|
|
new = base // topo.dp
|
|
hints.append(Hint(
|
|
"space", "info",
|
|
f"With DP={topo.dp}, scope=**TP+CP+DP** would shard FFN across "
|
|
f"replicas too, dropping per-PE FFN from {base/1e9:.2f} GB to "
|
|
f"{new/1e9:.2f} GB (comes with an inter-replica AllReduce)."
|
|
))
|
|
if (cfg.kv_replication_needed
|
|
and topo.kv_shard_mode == "replicate"):
|
|
rep = cfg.head_dim_split_factor
|
|
kv_pe = _kv_bytes_per_pe(cfg)
|
|
hints.append(Hint(
|
|
"space", "warn",
|
|
f"KV mode=replicate holds each KV head {rep}x. Switching to "
|
|
f"**split** saves ~{kv_pe*(rep-1)/rep/1e9:.2f} GB KV per PE "
|
|
f"but adds a Score AllReduce over {rep} ranks per hop."
|
|
))
|
|
if topo.cp == 1 and topo.tp <= m.h_kv:
|
|
# Suggest CP if s_kv is large and KV dominates.
|
|
kv_pe = _kv_bytes_per_pe(cfg)
|
|
attn_pe = _attn_weight_bytes_per_pe(cfg)
|
|
if kv_pe > attn_pe:
|
|
hints.append(Hint(
|
|
"space", "info",
|
|
f"KV cache ({kv_pe/1e9:.2f} GB) dominates weights "
|
|
f"({attn_pe/1e9:.2f} GB) at S_kv={topo.s_kv:,}. Adding CP "
|
|
f"shards the sequence axis (S_local = S_kv/CP)."
|
|
))
|
|
|
|
# COMM hints --------------------------------------------------
|
|
if topo.cp_inter_sip_hops > 0:
|
|
hints.append(Hint(
|
|
"comm", "warn",
|
|
f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
|
|
f"at {cfg.machine.bw_intersip_gbs:.0f} GB/s vs "
|
|
f"{cfg.machine.bw_inter_gbs:.0f} GB/s intra-SIP - "
|
|
f"consider a smaller CP or a topology that keeps CP inside one SIP."
|
|
))
|
|
if topo.tp_spans_cubes > 1:
|
|
hints.append(Hint(
|
|
"comm", "warn",
|
|
f"TP={topo.tp} spans {topo.tp_spans_cubes} cubes, so the W_O "
|
|
f"AllReduce runs cross-cube "
|
|
f"({cfg.machine.bw_intra_gbs:.0f} -> "
|
|
f"{cfg.machine.bw_inter_gbs:.0f} GB/s). A TP that fits in one "
|
|
f"cube (TP<={topo.pes_per_cube_hw}) keeps it intra-cube."
|
|
))
|
|
if topo.sips_used > 1 and topo.sip_topology == "ring":
|
|
hints.append(Hint(
|
|
"layout", "info",
|
|
f"With {topo.sips_used} SIPs on a ring, worst-case inter-SIP "
|
|
f"distance is {topo.sips_used - 1} hops. Switching to **torus2d** "
|
|
f"cuts worst-case hops to ~sqrt({topo.sips_used})."
|
|
))
|
|
if topo.tp <= m.h_kv:
|
|
hints.append(Hint(
|
|
"comm", "good",
|
|
f"TP ({topo.tp}) <= H_kv ({m.h_kv}): each KV head lives on "
|
|
f"exactly one TP rank - no head-split AllReduce needed."
|
|
))
|
|
|
|
# COMPUTE hints -----------------------------------------------
|
|
if topo.mode == "decode" and topo.T_q == 1:
|
|
# Decode is memory-bound almost everywhere; note it.
|
|
hints.append(Hint(
|
|
"compute", "info",
|
|
"Decode (T_q=1): most GEMMs are memory-bound. FLOPs matter far "
|
|
"less than HBM BW here - budget attention around "
|
|
"weight-read cost, not compute peak."
|
|
))
|
|
|
|
return hints
|