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>
311 lines
10 KiB
Python
311 lines
10 KiB
Python
"""Model + machine parameters for the analytical visualization tool.
|
||
|
||
All numbers are per-rank / per-PE unless noted. bytes are counted in
|
||
bytes (b), FLOPs in raw ops, bandwidth in bytes/second, latency in
|
||
seconds. Convert to μs / GB / GFLOP at display time only.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass, field
|
||
|
||
|
||
@dataclass
|
||
class ModelConfig:
|
||
"""Transformer model dimensions."""
|
||
name: str = "Qwen 3 8B"
|
||
hidden: int = 4096
|
||
ffn_dim: int = 12288
|
||
h_q: int = 32 # Q heads
|
||
h_kv: int = 8 # KV heads
|
||
d_head: int = 128
|
||
layers: int = 36
|
||
bytes_per_elem: int = 2 # bf16
|
||
|
||
@property
|
||
def group_size(self) -> int:
|
||
"""GQA group size = H_q / H_kv."""
|
||
return self.h_q // self.h_kv
|
||
|
||
@property
|
||
def head_dim_total_q(self) -> int:
|
||
return self.h_q * self.d_head
|
||
|
||
@property
|
||
def head_dim_total_kv(self) -> int:
|
||
return self.h_kv * self.d_head
|
||
|
||
|
||
@dataclass
|
||
class TopologyConfig:
|
||
"""Physical mapping — DP × PP × CP × TP (× EP for MoE)."""
|
||
cp: int = 4 # inter-cube ring size (sequence shard)
|
||
tp: int = 8 # PEs per cube (head/hidden shard)
|
||
pp: int = 1 # pipeline stages (layer shard)
|
||
dp: int = 1 # data-parallel replicas
|
||
ep: int = 1 # expert-parallel degree (MoE only)
|
||
s_kv: int = 4096
|
||
mode: str = "decode" # "decode" or "prefill"
|
||
kv_shard_mode: str = "split" # "split" or "replicate" (used when TP > H_kv)
|
||
# FFN sharding scope: default TP+CP (across cube group) trades a small
|
||
# extra AllReduce for lower per-PE weight bytes.
|
||
ffn_shard_scope: str = "TP+CP" # "TP" | "TP+CP" | "TP+CP+DP"
|
||
# Inter-SIP interconnect: how physical SIPs are wired together.
|
||
# "ring" (1D wrap), "mesh2d" (grid, no wrap), "torus2d" (grid + wrap).
|
||
sip_topology: str = "ring"
|
||
# Placement of parallelism dims on the hierarchy.
|
||
# "pe" -> dim lives on PEs within a single cube (intra-cube).
|
||
# "cube" -> dim lives across cubes (inter-cube).
|
||
tp_placement: str = "pe" # default: TP fills PEs of a cube
|
||
cp_placement: str = "cube" # default: CP fills cubes of a SIP
|
||
# CP ring variant: what gets passed around the CP ring each hop.
|
||
# "kv" -> K and V shards rotate (Q,O,m,l stay local). Bytes/hop scales
|
||
# with S_local * H_kv * d_h. Good for prefill (large T_q).
|
||
# "qoml" -> Q and running (O, m, l) rotate; K,V stay local. Bytes/hop
|
||
# scales with T_q * H_q * d_h. Cheap for decode (T_q=1).
|
||
cp_ring_variant: str = "kv"
|
||
pes_per_cube_hw: int = 8 # SIP hardware constant
|
||
cubes_per_sip_hw: int = 16
|
||
|
||
@property
|
||
def pes_per_stage(self) -> int:
|
||
"""PEs per pipeline stage (all stages have the same size)."""
|
||
return self.cp * self.tp
|
||
|
||
@property
|
||
def pes_per_replica(self) -> int:
|
||
"""PEs per one full model replica (CP*TP*PP)."""
|
||
return self.cp * self.tp * self.pp
|
||
|
||
@property
|
||
def total_pes(self) -> int:
|
||
return self.pes_per_replica * self.dp
|
||
|
||
# ── Placement-aware layout ────────────────────────────────
|
||
@property
|
||
def intra_cube_dims(self) -> int:
|
||
"""Product of dims placed on PE level (must fit in one cube, else spills)."""
|
||
n = 1
|
||
if self.tp_placement == "pe":
|
||
n *= self.tp
|
||
if self.cp_placement == "pe":
|
||
n *= self.cp
|
||
return n
|
||
|
||
@property
|
||
def inter_cube_dims(self) -> int:
|
||
"""Product of dims placed at cube level = distinct cube groups per stage."""
|
||
n = 1
|
||
if self.tp_placement == "cube":
|
||
n *= self.tp
|
||
if self.cp_placement == "cube":
|
||
n *= self.cp
|
||
return n
|
||
|
||
@property
|
||
def pes_per_cube_used(self) -> int:
|
||
"""How many PEs are live in each cube."""
|
||
return min(self.intra_cube_dims, self.pes_per_cube_hw)
|
||
|
||
@property
|
||
def cubes_per_stage(self) -> int:
|
||
"""Cubes needed per PP stage under current placement.
|
||
If intra-cube demand exceeds PEs/cube, the intra dim spills to cubes."""
|
||
spill = max(1, (self.intra_cube_dims + self.pes_per_cube_hw - 1)
|
||
// self.pes_per_cube_hw)
|
||
return self.inter_cube_dims * spill
|
||
|
||
@property
|
||
def cubes_used(self) -> int:
|
||
"""Physical cubes used across all PP stages and DP replicas."""
|
||
return self.cubes_per_stage * self.pp * self.dp
|
||
|
||
@property
|
||
def sips_used(self) -> int:
|
||
"""Number of SIPs used (each SIP = pes_per_cube_hw × cubes_per_sip_hw)."""
|
||
cubes_per_sip = self.cubes_per_sip_hw
|
||
return max(1, (self.cubes_used + cubes_per_sip - 1) // cubes_per_sip)
|
||
|
||
@property
|
||
def placement_valid(self) -> bool:
|
||
"""False when intra-cube demand exceeds one cube (dim spills to cubes)."""
|
||
return self.intra_cube_dims <= self.pes_per_cube_hw
|
||
|
||
@property
|
||
def tp_spans_cubes(self) -> int:
|
||
"""How many cubes one TP group spans (>=1). Depends on placement:
|
||
- tp_placement=pe: 1 if TP fits in a cube, else spill.
|
||
- tp_placement=cube: TP itself is the number of cubes per TP group.
|
||
"""
|
||
if self.tp_placement == "cube":
|
||
return max(1, self.tp)
|
||
return max(1, (self.tp + self.pes_per_cube_hw - 1) // self.pes_per_cube_hw)
|
||
|
||
@property
|
||
def cp_intra_sip_hops(self) -> int:
|
||
"""CP ring hops that stay within one SIP."""
|
||
if self.cp <= 1:
|
||
return 0
|
||
sips_in_ring = self.sips_used # one PP stage's CP ring may span SIPs
|
||
return max(0, (self.cp - 1) - (sips_in_ring - 1))
|
||
|
||
@property
|
||
def cp_inter_sip_hops(self) -> int:
|
||
"""CP ring hops that cross SIP boundaries."""
|
||
if self.cp <= 1:
|
||
return 0
|
||
# Approximation: one inter-SIP hop per SIP boundary in the ring.
|
||
# For CP=32 across 2 SIPs → 1 inter-SIP + 30 intra-SIP hops.
|
||
return max(0, self.sips_used - 1)
|
||
|
||
@property
|
||
def T_q(self) -> int:
|
||
return 1 if self.mode == "decode" else self.s_kv // self.cp
|
||
|
||
@property
|
||
def s_local(self) -> int:
|
||
return self.s_kv // self.cp
|
||
|
||
@property
|
||
def layers_per_stage(self) -> str:
|
||
"""Layers per PP stage. Requires model.layers to render — leave as fn."""
|
||
return "N_layers / PP"
|
||
|
||
def tp_link_tier(self) -> str:
|
||
"""Return which BW tier a TP AllReduce uses: 'intra' | 'inter' | 'intersip'.
|
||
- tp_placement=pe and TP fits in one cube -> 'intra'
|
||
- tp_placement=cube and TP fits in one SIP -> 'inter'
|
||
- otherwise -> 'intersip'
|
||
"""
|
||
if self.tp_placement == "pe":
|
||
if self.tp <= self.pes_per_cube_hw:
|
||
return "intra"
|
||
# spilled to multiple cubes: cross-cube
|
||
if self.tp_spans_cubes <= self.cubes_per_sip_hw:
|
||
return "inter"
|
||
return "intersip"
|
||
else: # tp_placement == "cube"
|
||
if self.tp <= self.cubes_per_sip_hw:
|
||
return "inter"
|
||
return "intersip"
|
||
|
||
def cp_link_tier(self) -> str:
|
||
"""Return BW tier for the CP ring: 'intra' | 'inter' | 'intersip'."""
|
||
if self.cp_placement == "pe":
|
||
return "intra"
|
||
# cp_placement == "cube"
|
||
if self.cp <= self.cubes_per_sip_hw:
|
||
return "inter"
|
||
return "intersip"
|
||
|
||
|
||
@dataclass
|
||
class MachineParams:
|
||
"""SIP hardware parameters — all user-adjustable via sliders.
|
||
|
||
Bandwidth tiers, from fastest to slowest:
|
||
HBM (per-PE local memory)
|
||
> intra-cube (PE↔PE on same die)
|
||
> inter-cube (UCIe D2D between dies on same SIP)
|
||
> inter-SIP (C2C / RDMA-class, across SIPs / servers)
|
||
"""
|
||
# Per-PE compute peak (TFLOPs f16).
|
||
peak_tflops_f16: float = 8.0
|
||
# Per-PE HBM budget (GB) — the memory ceiling per PE.
|
||
pe_hbm_gb: float = 6.0
|
||
# HBM per-PE bandwidth (GB/s).
|
||
bw_hbm_gbs: float = 256.0
|
||
# Intra-cube PE↔PE link BW (GB/s).
|
||
bw_intra_gbs: float = 512.0
|
||
# Inter-cube (UCIe D2D) BW (GB/s).
|
||
bw_inter_gbs: float = 128.0
|
||
# Inter-SIP (C2C / RDMA) BW (GB/s).
|
||
bw_intersip_gbs: float = 50.0
|
||
# Per-hop latencies (ns).
|
||
alpha_intra_ns: float = 20.0
|
||
alpha_inter_ns: float = 100.0
|
||
alpha_intersip_ns: float = 1000.0
|
||
# Achievable utilization for compute-bound stages (0-1).
|
||
compute_util: float = 0.8
|
||
|
||
@property
|
||
def peak_flops(self) -> float:
|
||
return self.peak_tflops_f16 * 1e12
|
||
|
||
@property
|
||
def bw_hbm(self) -> float:
|
||
return self.bw_hbm_gbs * 1e9
|
||
|
||
@property
|
||
def bw_intra(self) -> float:
|
||
return self.bw_intra_gbs * 1e9
|
||
|
||
@property
|
||
def bw_inter(self) -> float:
|
||
return self.bw_inter_gbs * 1e9
|
||
|
||
@property
|
||
def bw_intersip(self) -> float:
|
||
return self.bw_intersip_gbs * 1e9
|
||
|
||
@property
|
||
def alpha_intra(self) -> float:
|
||
return self.alpha_intra_ns * 1e-9
|
||
|
||
@property
|
||
def alpha_inter(self) -> float:
|
||
return self.alpha_inter_ns * 1e-9
|
||
|
||
@property
|
||
def alpha_intersip(self) -> float:
|
||
return self.alpha_intersip_ns * 1e-9
|
||
|
||
@property
|
||
def pe_budget_bytes(self) -> int:
|
||
return int(self.pe_hbm_gb * 1e9)
|
||
|
||
|
||
@dataclass
|
||
class FullConfig:
|
||
"""Bundled config for one analysis run."""
|
||
model: ModelConfig = field(default_factory=ModelConfig)
|
||
topo: TopologyConfig = field(default_factory=TopologyConfig)
|
||
machine: MachineParams = field(default_factory=MachineParams)
|
||
|
||
@property
|
||
def h_q_per_pe(self) -> int:
|
||
return max(1, self.model.h_q // self.topo.tp)
|
||
|
||
@property
|
||
def h_kv_per_pe(self) -> float:
|
||
"""Fractional if TP > H_kv (head-dim split needed)."""
|
||
return self.model.h_kv / self.topo.tp
|
||
|
||
@property
|
||
def d_per_pe(self) -> int:
|
||
return self.model.hidden // self.topo.tp
|
||
|
||
@property
|
||
def ffn_per_pe(self) -> int:
|
||
return self.model.ffn_dim // self.topo.tp
|
||
|
||
@property
|
||
def kv_replication_needed(self) -> bool:
|
||
"""True if TP > H_kv (each KV head shared across TP/H_kv ranks)."""
|
||
return self.topo.tp > self.model.h_kv
|
||
|
||
@property
|
||
def head_dim_split_factor(self) -> int:
|
||
"""Number of TP ranks sharing one head (via d_head split)."""
|
||
return max(1, self.topo.tp // self.model.h_kv)
|
||
|
||
@property
|
||
def ffn_shard_divisor(self) -> int:
|
||
"""How much the FFN dim is sharded across ranks."""
|
||
scope = self.topo.ffn_shard_scope
|
||
d = self.topo.tp
|
||
if "CP" in scope:
|
||
d *= self.topo.cp
|
||
if "DP" in scope:
|
||
d *= self.topo.dp
|
||
return max(1, d)
|