diff --git a/tests/analytical_visualization/__init__.py b/tests/analytical_visualization/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py
new file mode 100644
index 0000000..d733c7b
--- /dev/null
+++ b/tests/analytical_visualization/app.py
@@ -0,0 +1,1284 @@
+"""Streamlit app: interactive analytical model for transformer inference
+on the SIP architecture.
+
+Run:
+ python -m streamlit run tests/analytical_visualization/app.py
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import streamlit as st
+
+from tests.analytical_visualization.model_config import (
+ FullConfig, ModelConfig, TopologyConfig, MachineParams,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+from tests.analytical_visualization.stage_latencies import all_stages, all_ffn_stages
+from tests.analytical_visualization.memory_layout import (
+ compute_memory,
+ attention_weight_rows,
+ ffn_weight_rows,
+ kv_cache_rows,
+ sum_bytes_all_layers,
+ total_weight_bytes_full_model,
+ total_kv_bytes_full_model,
+)
+from tests.analytical_visualization.topology_map import draw_topology
+from tests.analytical_visualization.pe_weight_layout import (
+ draw_pe_layout, _per_pe_bytes, _q_heads_for_pe, _kv_heads_for_pe,
+)
+from tests.analytical_visualization.tensor_sharding import draw_tensor_sharding
+from tests.analytical_visualization.pipeline_diagram import draw_pipeline
+from tests.analytical_visualization.optimization_report import (
+ replication_report, optimization_hints,
+)
+from tests.analytical_visualization.autosuggest import auto_suggest
+
+
+st.set_page_config(page_title="SIP Attention Analytical Model",
+ layout="wide")
+
+# Tighter spacing so more info fits per page.
+st.markdown(
+ """
+
+ """,
+ unsafe_allow_html=True,
+)
+
+
+def _pick_in(container, label: str, options, default, key: str, help: str = ""):
+ """Selectbox helper — returns the selected value from options."""
+ return container.selectbox(
+ label, options,
+ index=options.index(default) if default in options else 0,
+ key=key, help=help,
+ )
+
+
+def _pick(label, options, default, key, help=""):
+ return _pick_in(st.sidebar, label, options, default, key, help)
+
+
+# ── Sidebar ─ compact, guided layout ────────────────────────────
+st.sidebar.markdown("### SIP config")
+st.sidebar.caption(
+ "Pick a model, workload, parallelism plan, then tune hardware. "
+ "Every change re-runs the analytical model live."
+)
+
+with st.sidebar.expander("Model", expanded=True):
+ preset_names = list(PRESETS.keys())
+ preset_name = st.selectbox(
+ "Preset", preset_names,
+ index=preset_names.index("Qwen 3 8B"),
+ key="preset_select",
+ label_visibility="collapsed",
+ )
+ preset = PRESETS[preset_name]
+ st.caption(
+ f"{preset.family or '-'} - {preset.attn_type} - "
+ f"H_q={preset.model.h_q}, H_kv={preset.model.h_kv}, "
+ f"L={preset.model.layers}"
+ )
+ with st.expander("Override dims", expanded=False):
+ d1, d2 = st.columns(2)
+ with d1:
+ hidden_ = st.number_input("hidden", 128, 65536,
+ value=preset.model.hidden, step=128)
+ h_q_ = st.number_input("H_q", 1, 256, value=preset.model.h_q)
+ d_head_ = st.number_input("d_head", 32, 512,
+ value=preset.model.d_head, step=32)
+ layers_ = st.number_input("layers", 1, 200,
+ value=preset.model.layers)
+ with d2:
+ ffn_ = st.number_input("ffn_dim", 128, 131072,
+ value=preset.model.ffn_dim, step=128)
+ h_kv_ = st.number_input("H_kv", 1, 256, value=preset.model.h_kv)
+ bpe_ = st.selectbox("dtype bytes", [1, 2, 4], index=1)
+ model = ModelConfig(
+ name=preset.model.name, hidden=hidden_, ffn_dim=ffn_,
+ h_q=h_q_, h_kv=h_kv_, d_head=d_head_, layers=layers_,
+ bytes_per_elem=bpe_,
+ )
+
+
+with st.sidebar.expander("Workload", expanded=True):
+ w1, w2 = st.columns([2, 1])
+ with w1:
+ s_kv = st.selectbox(
+ "Context S_kv",
+ [1024, 2048, 4096, 8192, 16384, 32768, 65536,
+ 131072, 262144, 524288, 1048576],
+ index=10, key="s_kv", # default 1M
+ )
+ with w2:
+ mode = st.selectbox("Mode", ["decode", "prefill"],
+ index=0, key="mode")
+
+
+# Compute auto-suggest for defaults / apply button
+_default_machine = MachineParams()
+_default_suggestion = auto_suggest(model, _default_machine, s_kv, mode)
+
+def _snap(v: int, opts: tuple) -> int:
+ return min(opts, key=lambda o: abs(o - v))
+
+_CP_OPTS = (1, 2, 4, 8, 16, 32, 48, 64, 96)
+_TP_OPTS = (1, 2, 4, 8, 16, 32)
+_PP_OPTS = (1, 2, 4, 8, 16, 32)
+_DP_OPTS = (1, 2, 4, 8, 16)
+_EP_OPTS = (1, 2, 4, 8, 16, 32)
+
+# Auto-reset dropdowns when model preset changes.
+if st.session_state.get("_last_preset") != preset_name:
+ st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
+ st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
+ st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
+ st.session_state["dp"] = 1
+ st.session_state["ep"] = 1
+ st.session_state["_last_preset"] = preset_name
+
+with st.sidebar.expander("Parallelism", expanded=True):
+ st.caption(
+ f"Auto-suggest: CP={_default_suggestion.cp}, "
+ f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp}"
+ )
+ if st.button("Apply auto-suggest", width='stretch'):
+ st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
+ st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
+ st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
+ st.rerun()
+
+ # Sharding dims — arrange in 2x3 grid to save vertical space
+ p_row1 = st.columns(3)
+ with p_row1[0]:
+ cp = st.selectbox(
+ "CP", list(_CP_OPTS),
+ index=_CP_OPTS.index(st.session_state.get(
+ "cp", _snap(_default_suggestion.cp, _CP_OPTS))),
+ key="cp",
+ )
+ with p_row1[1]:
+ tp = st.selectbox(
+ "TP", list(_TP_OPTS),
+ index=_TP_OPTS.index(st.session_state.get(
+ "tp", _snap(_default_suggestion.tp, _TP_OPTS))),
+ key="tp",
+ )
+ with p_row1[2]:
+ pp = st.selectbox(
+ "PP", list(_PP_OPTS),
+ index=_PP_OPTS.index(st.session_state.get(
+ "pp", _snap(_default_suggestion.pp, _PP_OPTS))),
+ key="pp",
+ )
+ p_row2 = st.columns(2)
+ with p_row2[0]:
+ dp = st.selectbox(
+ "DP", list(_DP_OPTS),
+ index=_DP_OPTS.index(st.session_state.get("dp", 1)),
+ key="dp",
+ )
+ with p_row2[1]:
+ ep = st.selectbox(
+ "EP (MoE)", list(_EP_OPTS),
+ index=_EP_OPTS.index(st.session_state.get("ep", 1)),
+ key="ep",
+ )
+
+ st.markdown("**Placement (PE-level vs cube-level)**")
+ pl_col1, pl_col2 = st.columns(2)
+ with pl_col1:
+ tp_placement = st.radio(
+ "TP on", options=["pe", "cube"], index=0,
+ key="tp_placement", horizontal=True,
+ help=("pe: TP shares PEs of one cube (intra-cube AllReduce). "
+ "cube: each TP rank owns a whole cube (inter-cube AllReduce)."),
+ )
+ with pl_col2:
+ cp_placement = st.radio(
+ "CP on", options=["pe", "cube"], index=1,
+ key="cp_placement", horizontal=True,
+ help=("pe: CP ring runs among PEs of one cube (intra-cube). "
+ "cube: each CP rank owns cubes (inter-cube ring)."),
+ )
+
+ st.markdown("**Sharding modes**")
+ # FFN scope labels show the effective divisor (dynamic per CP/TP/DP).
+ _ffn_labels = {
+ "TP": f"TP only (div={max(1,tp)})",
+ "TP+CP": f"TP*CP (div={max(1,tp*cp)})",
+ "TP+CP+DP": f"TP*CP*DP (div={max(1,tp*cp*dp)})",
+ }
+ _ffn_choice = st.radio(
+ "FFN shard scope",
+ options=list(_ffn_labels.values()),
+ index=1, key="ffn_scope_label", horizontal=True,
+ help=("Divisor = number of ranks the FFN weights are split across. "
+ "Larger divisor = less memory/PE, more AllReduce cost."),
+ )
+ # Map label -> raw scope key.
+ ffn_shard_scope = {v: k for k, v in _ffn_labels.items()}[_ffn_choice]
+
+ sip_topology = st.radio(
+ "SIP interconnect",
+ options=["ring", "mesh2d", "torus2d"],
+ index=0, key="sip_topo", horizontal=True,
+ help=("ring: 1D chain + wrap. mesh2d: 2D grid, no wrap. "
+ "torus2d: 2D grid + wrap."),
+ )
+ cp_ring_variant = st.radio(
+ "CP ring: what rotates",
+ options=["kv", "qoml"], index=0, key="cp_ring_variant",
+ horizontal=True,
+ help=("kv: K,V shards rotate each hop (Q,O,m,l stay local). "
+ "Bytes/hop scales with S_local*H_kv - good for prefill.\n"
+ "qoml: Q + running (O,m,l) rotate (K,V stay local). "
+ "Bytes/hop scales with T_q*H_q - much cheaper for decode."),
+ )
+ kv_shard_mode = "split"
+ if tp > model.h_kv:
+ _rep_factor = max(1, tp // model.h_kv)
+ kv_shard_mode = st.radio(
+ f"KV mode (TP={tp} > H_kv={model.h_kv})",
+ options=["split", "replicate"], index=0,
+ key="kv_mode", horizontal=True,
+ help=(f"split: fractional heads (KV/PE shrinks by {tp/max(1,model.h_kv):.1f}x, "
+ f"adds Score AllReduce across {_rep_factor} ranks). "
+ f"replicate: whole heads per PE, duplicated {_rep_factor}x, "
+ f"no Score AllReduce."),
+ )
+
+
+with st.sidebar.expander("Hardware", expanded=False):
+ hw_tab_pe, hw_tab_net = st.tabs(["Per-PE", "Interconnect"])
+ with hw_tab_pe:
+ hp1, hp2 = st.columns(2)
+ with hp1:
+ pe_hbm_gb = st.selectbox(
+ "PE HBM (GB)", [3.0, 6.0, 12.0, 24.0, 48.0, 96.0],
+ index=1, key="pe_hbm",
+ )
+ bw_hbm = st.selectbox(
+ "HBM BW GB/s", [128.0, 256.0, 512.0, 1024.0, 2048.0],
+ index=1, key="bw_hbm",
+ )
+ with hp2:
+ peak_tflops = st.selectbox(
+ "TFLOPs/PE", [2.0, 4.0, 8.0, 16.0, 32.0, 64.0],
+ index=2, key="tflops",
+ )
+ compute_util = st.selectbox(
+ "Compute util", [0.3, 0.5, 0.7, 0.8, 0.9, 1.0],
+ index=3, key="util",
+ )
+ with hw_tab_net:
+ n1, n2 = st.columns(2)
+ with n1:
+ bw_intra = st.selectbox(
+ "intra-cube GB/s",
+ [128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0],
+ index=2, key="bw_intra",
+ )
+ bw_inter = st.selectbox(
+ "inter-cube GB/s",
+ [32.0, 64.0, 128.0, 256.0, 512.0, 900.0],
+ index=2, key="bw_inter",
+ )
+ bw_intersip = st.selectbox(
+ "inter-SIP GB/s",
+ [12.5, 25.0, 50.0, 100.0, 200.0, 400.0],
+ index=2, key="bw_intersip",
+ )
+ with n2:
+ alpha_intra = st.selectbox(
+ "a_intra ns", [5.0, 10.0, 20.0, 50.0, 100.0],
+ index=2, key="a_intra",
+ )
+ alpha_inter = st.selectbox(
+ "a_inter ns", [50.0, 100.0, 200.0, 500.0, 1000.0],
+ index=1, key="a_inter",
+ )
+ alpha_intersip = st.selectbox(
+ "a_intersip ns",
+ [500.0, 1000.0, 2000.0, 5000.0, 10000.0],
+ index=1, key="a_intersip",
+ )
+
+machine = MachineParams(
+ pe_hbm_gb=pe_hbm_gb, peak_tflops_f16=peak_tflops,
+ bw_hbm_gbs=bw_hbm, bw_intra_gbs=bw_intra,
+ bw_inter_gbs=bw_inter, bw_intersip_gbs=bw_intersip,
+ alpha_intra_ns=alpha_intra, alpha_inter_ns=alpha_inter,
+ alpha_intersip_ns=alpha_intersip,
+ compute_util=compute_util,
+)
+
+topo = TopologyConfig(cp=cp, tp=tp, pp=pp, dp=dp, ep=ep,
+ s_kv=s_kv, mode=mode,
+ kv_shard_mode=kv_shard_mode,
+ ffn_shard_scope=ffn_shard_scope,
+ sip_topology=sip_topology,
+ tp_placement=tp_placement,
+ cp_placement=cp_placement,
+ cp_ring_variant=cp_ring_variant)
+cfg = FullConfig(model=model, topo=topo, machine=machine)
+
+
+# ── Main pane ─────────────────────────────────────────────────────
+st.title(f"SIP Analytical Model — {model.name}")
+
+st.caption(
+ f"Family: **{preset.family or '-'}** | Attn: **{preset.attn_type}** | "
+ f"H_q={model.h_q}, H_kv={model.h_kv}, hidden={model.hidden}, "
+ f"ffn={model.ffn_dim}, layers={model.layers}"
+)
+if preset.note:
+ st.info(f"Note: {preset.note}")
+
+
+# Recompute auto-suggest with the *current* machine (in case user changed HBM etc.)
+suggestion = auto_suggest(model, machine, s_kv, mode)
+
+# Compact top summary: auto-suggest + current in one row of tight bullets.
+_sug_status = "OK" if suggestion.fits else "NO FIT"
+_sug_line = (
+ f"**Auto-suggest:** CP={suggestion.cp} - TP={suggestion.tp} - "
+ f"PP={suggestion.pp} - PEs={suggestion.pes_used} - SIPs={suggestion.sips_used} "
+ f"({_sug_status}, slack {suggestion.slack_gb:.2f} GB / "
+ f"{machine.pe_hbm_gb:.1f} GB budget)"
+)
+_cur_line = (
+ f"**Current:** CP={cp}@{cp_placement} - TP={tp}@{tp_placement} - "
+ f"PP={pp} - DP={dp} - EP={ep} - PEs={topo.total_pes} - "
+ f"Cubes={topo.cubes_used} ({topo.pes_per_cube_used}/cube live) - "
+ f"SIPs={topo.sips_used} - FFN={cfg.topo.ffn_shard_scope} - "
+ f"SIPnet={cfg.topo.sip_topology}"
+)
+sum_l, sum_r = st.columns([1, 1])
+with sum_l:
+ st.markdown(_sug_line)
+with sum_r:
+ st.markdown(_cur_line)
+
+# Consolidated warnings (only if any apply)
+_warnings = []
+if not topo.placement_valid:
+ _spill_cubes = (topo.intra_cube_dims + topo.pes_per_cube_hw - 1) // topo.pes_per_cube_hw
+ _warnings.append(
+ f"Intra-cube demand = {topo.intra_cube_dims} PEs > "
+ f"{topo.pes_per_cube_hw} PEs/cube -> spills to {_spill_cubes} cubes "
+ f"per group (AllReduce/ring drops from intra-cube "
+ f"{machine.bw_intra_gbs:.0f} to inter-cube "
+ f"{machine.bw_inter_gbs:.0f} GB/s)."
+ )
+if tp > model.h_kv:
+ # Compute per-PE KV bytes under each mode so the difference is visible.
+ from tests.analytical_visualization.memory_layout import per_pe_kv_cache_bytes
+ import copy as _copy
+ _cfg_split = FullConfig(model=model, topo=TopologyConfig(
+ **{**topo.__dict__, "kv_shard_mode": "split"}), machine=machine)
+ _cfg_rep = FullConfig(model=model, topo=TopologyConfig(
+ **{**topo.__dict__, "kv_shard_mode": "replicate"}), machine=machine)
+ _kv_split_gb = per_pe_kv_cache_bytes(_cfg_split) / 1e9
+ _kv_rep_gb = per_pe_kv_cache_bytes(_cfg_rep) / 1e9
+ _kv_now = _kv_split_gb if kv_shard_mode == "split" else _kv_rep_gb
+ _warnings.append(
+ f"KV mode: **{kv_shard_mode}** (TP={tp} > H_kv={model.h_kv}) - "
+ f"per-PE KV: split={_kv_split_gb:.2f} GB vs "
+ f"replicate={_kv_rep_gb:.2f} GB (currently {_kv_now:.2f} GB). "
+ f"Split adds one Score AllReduce over "
+ f"{max(1, tp // model.h_kv)} ranks per hop."
+ )
+if cfg.kv_replication_needed:
+ _warnings.append(f"Head-dim split: {cfg.head_dim_split_factor} ranks share one KV head")
+if topo.tp_spans_cubes > 1:
+ _warnings.append(f"TP spans {topo.tp_spans_cubes} cubes "
+ f"(AllReduce goes cross-cube)")
+if topo.cp_inter_sip_hops > 0:
+ _warnings.append(f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
+ f"at {machine.bw_intersip_gbs:.0f} GB/s")
+if _warnings:
+ st.warning(" - ".join(_warnings))
+
+
+# ── Tabs ─────────────────────────────────────────────────────────
+tab_layout, tab_memory, tab_stages, tab_compare = st.tabs([
+ "Physical layout", "Memory breakdown", "Per-stage latency",
+ "Save & compare",
+])
+
+
+# ── TAB 1: Physical layout ──────────────────────────────────────
+with tab_layout:
+ # Row A: pipeline diagram (color-coded) + per-stage table (formulas).
+ st.markdown("**Per-layer pipeline (color = dominant bound)**")
+ fig_pipe = draw_pipeline(cfg)
+ st.pyplot(fig_pipe, width='stretch')
+ plt.close(fig_pipe)
+
+ with st.expander("Per-stage latency (attention + FFN) with formulas",
+ expanded=True):
+ import html as _html
+ _stages_attn = all_stages(cfg)
+ _stages_ffn = all_ffn_stages(cfg)
+ _bound_color = {
+ "compute": "#d6eaf8",
+ "memory": "#fdebd0",
+ "comm": "#f5b7b1",
+ "trivial": "#eaeded",
+ }
+
+ def _vsplit(formula: str) -> str:
+ """Break 'X = Y = Z' into vertical `= X` / `= Y` / `= Z` lines.
+ If the string contains a 'PURPOSE:' preamble ending with a line
+ of '---', keep that preamble verbatim and only vsplit the tail.
+ """
+ if not formula or formula in ("-", "0"):
+ return formula
+ _preamble = ""
+ _body = formula
+ if "PURPOSE:" in formula and "\n---\n" in formula:
+ _preamble, _body = formula.split("\n---\n", 1)
+ _preamble = _preamble.rstrip() + "\n---"
+ pieces = [p.strip() for p in _body.split(" = ")]
+ _vsplit_body = "\n".join(f"= {p}" for p in pieces)
+ return f"{_preamble}\n{_vsplit_body}" if _preamble else _vsplit_body
+
+ # Prefill: per-hop ring is concurrent with S5/S6/S7 compute -
+ # annotate them. In decode the O/m/l all-reduce is already
+ # folded into the S8 row, so no separate annotation needed.
+ _per_hop_ring_us = 0.0
+ _ring_variant_desc = ""
+ if topo.mode == "prefill" and topo.cp > 1:
+ _c1 = next((_x for _x in _stages_attn
+ if _x.name.startswith("C1 ")), None)
+ if _c1 is not None and _c1.comm_s > 0:
+ _hops = max(1, topo.cp - 1)
+ _per_hop_ring_us = (_c1.comm_s * 1e6) / _hops
+ _ring_variant_desc = ("K/V" if topo.cp_ring_variant == "kv"
+ else "Q+O/m/l")
+
+ def _fmt_time(sec: float) -> str:
+ """Auto-scale time: ns for < 1 us, us for < 10 ms, else ms."""
+ if sec <= 0:
+ return "0"
+ ns = sec * 1e9
+ if ns < 1000:
+ return f"{ns:.1f} ns"
+ us = sec * 1e6
+ if us < 1000:
+ return f"{us:.2f} us"
+ ms = sec * 1e3
+ return f"{ms:.2f} ms"
+
+ def _combined_formula(s):
+ parts = []
+ if s.compute_s > 0 or (s.flops and s.flops > 0):
+ parts.append(
+ f"Compute:\n{_vsplit(s.flops_formula)}\n"
+ f"-> {_fmt_time(s.compute_s)}"
+ )
+ if s.memory_s > 0 or (s.mem_bytes and s.mem_bytes > 0):
+ parts.append(
+ f"Memory:\n{_vsplit(s.mem_formula)}\n"
+ f"-> {_fmt_time(s.memory_s)}"
+ )
+ if s.comm_s > 0 or (s.comm_bytes and s.comm_bytes > 0):
+ # If this is a concurrent-comm stage (e.g. C1 CP ring in
+ # prefill) that gets overlapped with compute, visible < raw.
+ # Show both so the '0' in the Time column makes sense.
+ _tail = f"-> raw comm: {_fmt_time(s.comm_s)}"
+ if s.visible_s < s.comm_s * 0.999:
+ _tail += (
+ f"\n-> visible: {_fmt_time(s.visible_s)} "
+ f"(concurrent S5-S7 compute hides "
+ f"{_fmt_time(s.comm_s - s.visible_s)} of it)"
+ )
+ parts.append(
+ f"Comm:\n{_vsplit(s.comm_formula)}\n{_tail}"
+ )
+ # Prefill: S5/S6/S7 run concurrently with the CP ring.
+ if (_per_hop_ring_us > 0
+ and s.name.startswith(("S5 ", "S6 ", "S7 "))):
+ # convert us to seconds for _fmt_time
+ _per_hop_s = _per_hop_ring_us / 1e6
+ parts.append(
+ f"Concurrent CP ring ({_ring_variant_desc}, per hop):\n"
+ f"~{_fmt_time(_per_hop_s)} in flight while this compute runs\n"
+ f"(aggregated in C1 row; overlap kept only excess vs compute)"
+ )
+ return "\n\n".join(parts) if parts else "-"
+
+ # Compact CSS: rows auto-size to text; tight padding; columns fit content.
+ _stage_table_css = """
+
+ """
+
+ def _render_table(stages, title):
+ rows_html = []
+ for s in stages:
+ bg = _bound_color.get(s.bound, "#eaeded")
+ formula = _html.escape(_combined_formula(s))
+ rows_html.append(
+ f''
+ f'| {_html.escape(s.name)} | '
+ f'{formula} | '
+ f'{s.bound} | '
+ f'{_fmt_time(s.visible_s)} | '
+ f'
'
+ )
+ html = (
+ _stage_table_css
+ + f'{_html.escape(title)}'
+ + ''
+ + '| Stage | Formula | '
+ 'Bound | Time |
|---|
'
+ + '' + ''.join(rows_html) + '
'
+ )
+ # Prefer st.html (reliable) with a markdown fallback.
+ if hasattr(st, "html"):
+ st.html(html)
+ else:
+ st.markdown(html, unsafe_allow_html=True)
+ _t = sum(s.visible_s for s in stages)
+ st.caption(
+ f"{title} total: **{_fmt_time(_t)}**. "
+ f"Compute time = FLOPs / ({cfg.machine.peak_tflops_f16:.1f} TFLOPs "
+ f"x {cfg.machine.compute_util:.0%} util). "
+ f"Memory time = bytes / {cfg.machine.bw_hbm_gbs:.0f} GB/s HBM. "
+ f"Bound = max(compute, memory, comm) - row color: "
+ f"blue=compute, orange=memory, red=comm, grey=trivial."
+ )
+
+ _render_table(_stages_attn, "Attention")
+ _render_table(_stages_ffn, "FFN")
+
+ with st.expander("Symbol glossary (what does T_q, d_h, div... mean?)",
+ expanded=False):
+ _glossary = [
+ ("Model dims", [
+ ("d", f"hidden dim (current: {model.hidden})"),
+ ("d_h", f"per-head dim (current: {model.d_head})"),
+ ("H_q", f"total query heads (current: {model.h_q})"),
+ ("H_kv", f"total KV heads (current: {model.h_kv})"),
+ ("ffn", f"FFN intermediate dim (current: {model.ffn_dim})"),
+ ("L", f"transformer layers (current: {model.layers})"),
+ ("b", f"bytes per element (dtype size) "
+ f"(current: {model.bytes_per_elem} -> "
+ f"{'fp8' if model.bytes_per_elem == 1 else 'bf16/fp16' if model.bytes_per_elem == 2 else 'fp32'})"),
+ ]),
+ ("Workload", [
+ ("S_kv", f"global KV cache length / context (current: {topo.s_kv:,})"),
+ ("S_local", f"per-CP-rank KV length = S_kv / CP "
+ f"(current: {topo.s_local:,})"),
+ ("T_q", f"query tokens processed this step "
+ f"(decode=1, prefill=S_local; current: {topo.T_q})"),
+ ("mode", f"decode | prefill (current: {topo.mode})"),
+ ]),
+ ("Parallelism (degrees)", [
+ ("TP", f"tensor parallelism (attn heads / FFN cols) "
+ f"(current: {topo.tp}, on {topo.tp_placement})"),
+ ("CP", f"context parallelism (sequence axis) "
+ f"(current: {topo.cp}, on {topo.cp_placement})"),
+ ("PP", f"pipeline parallelism (layer stages) (current: {topo.pp})"),
+ ("DP", f"data parallelism (model replicas) (current: {topo.dp})"),
+ ("EP", f"expert parallelism (MoE) (current: {topo.ep})"),
+ ("div", f"FFN shard divisor = product of dims in "
+ f"'FFN shard scope' (current: {cfg.ffn_shard_divisor})"),
+ ("split", f"# ranks sharing one KV head via head-dim split "
+ f"(TP/H_kv when TP > H_kv; current: "
+ f"{cfg.head_dim_split_factor if cfg.kv_replication_needed else 1})"),
+ ]),
+ ("Hardware / machine", [
+ ("BW_HBM", f"per-PE HBM bandwidth "
+ f"(current: {machine.bw_hbm_gbs:.0f} GB/s)"),
+ ("BW_intra", f"intra-cube PE<->PE bandwidth "
+ f"(current: {machine.bw_intra_gbs:.0f} GB/s)"),
+ ("BW_inter", f"inter-cube (D2D) bandwidth "
+ f"(current: {machine.bw_inter_gbs:.0f} GB/s)"),
+ ("BW_intersip", f"inter-SIP (C2C) bandwidth "
+ f"(current: {machine.bw_intersip_gbs:.0f} GB/s)"),
+ ("alpha_intra", f"per-hop latency intra-cube "
+ f"(current: {machine.alpha_intra_ns:.0f} ns)"),
+ ("alpha_inter", f"per-hop latency inter-cube "
+ f"(current: {machine.alpha_inter_ns:.0f} ns)"),
+ ("alpha_intersip", f"per-hop latency inter-SIP "
+ f"(current: {machine.alpha_intersip_ns:.0f} ns)"),
+ ("peak TFLOPs", f"per-PE peak compute "
+ f"(current: {machine.peak_tflops_f16:.1f} TFLOPs f16)"),
+ ("util", f"achievable compute utilization "
+ f"(current: {machine.compute_util:.0%})"),
+ ]),
+ ("Stage prefixes", [
+ ("S1 .. S10", "attention forward stages (RMSNorm, W_Q, W_K+W_V, "
+ "KV append, Q.K^T, softmax, P.V, merge, "
+ "normalize, W_O)"),
+ ("F1 .. F5", "FFN forward stages (RMSNorm, W_gate, W_up, "
+ "SwiGLU, W_down)"),
+ ("C1", "CP K/V ring (K,V passed hop-by-hop)"),
+ ("C2", "TP AllReduce on W_O output"),
+ ("C3", "Score AllReduce when TP > H_kv with head-dim split"),
+ ("CF1", "FFN AllReduce over the FFN shard scope"),
+ ]),
+ ("Notation in formulas", [
+ ("H_q / TP", "Q heads per PE (or per TP rank)"),
+ ("H_kv / TP", "KV heads per PE (fractional allowed in split mode)"),
+ ("ffn / div", "FFN cols/rows per PE after sharding"),
+ ("2 * A * B * C", "GEMM FLOPs = 2 * M * K * N (multiply + add per fma)"),
+ ("bytes * b", "byte count (b = bytes/element)"),
+ ("-> X us", "resulting latency after dividing by BW or peak"),
+ ]),
+ ]
+ for section, items in _glossary:
+ st.markdown(f"**{section}**")
+ _rows = "".join(
+ f'| {sym} | '
+ f'{desc} |
'
+ for sym, desc in items
+ )
+ st.html(
+ f''
+ )
+
+ st.divider()
+
+ # Row B: topology diagram (full width - it can be tall for multi-SIP)
+ with st.expander(
+ f"Physical layout - {topo.total_pes} PEs, {topo.cubes_used} cubes, "
+ f"{topo.sips_used} SIP(s), inter-SIP: {topo.sip_topology}",
+ expanded=True,
+ ):
+ fig_topo = draw_topology(cfg)
+ st.pyplot(fig_topo, width='stretch')
+ plt.close(fig_topo)
+ l1, l2, l3, l4 = st.columns(4)
+ l1.metric("Cubes", topo.cubes_used)
+ l2.metric("SIPs", topo.sips_used)
+ l3.metric("CP intra-SIP hops", topo.cp_intra_sip_hops)
+ l4.metric("CP inter-SIP hops", topo.cp_inter_sip_hops)
+
+ st.divider()
+
+ # Row C: two-column - LEFT per-PE memory (pie + summary), RIGHT tensor sharding
+ mem_layout = compute_memory(cfg)
+ col_mem, col_shard = st.columns([1, 1], gap="medium")
+
+ with col_mem:
+ st.markdown("**Per-PE memory footprint**")
+ labels = ["Weights", "KV cache", "Transient", "Slack"]
+ vals = [mem_layout.weights_bytes, mem_layout.kv_cache_bytes,
+ mem_layout.transient_bytes, mem_layout.slack_bytes]
+ colors = ["#3a86ff", "#8338ec", "#ffbe0b", "#adb5bd"]
+ fig_mem_l, ax_mem_l = plt.subplots(figsize=(4.2, 4.2))
+ ax_mem_l.pie(vals,
+ labels=[f"{l}\n{v/1e9:.2f} GB"
+ for l, v in zip(labels, vals)],
+ colors=colors, autopct="%1.1f%%", startangle=90,
+ textprops={"fontsize": 8})
+ ax_mem_l.set_title(f"Per-PE (budget: "
+ f"{mem_layout.budget_bytes/1e9:.1f} GB)",
+ fontsize=10)
+ st.pyplot(fig_mem_l, width='stretch')
+ plt.close(fig_mem_l)
+ if mem_layout.over_budget:
+ st.error(f"OVER BUDGET by "
+ f"{(mem_layout.used_bytes - mem_layout.budget_bytes)/1e9:.2f} GB")
+ else:
+ st.caption(
+ f"Weights {mem_layout.weights_bytes/1e9:.2f} GB - "
+ f"KV {mem_layout.kv_cache_bytes/1e9:.2f} GB "
+ f"(S_local={topo.s_local:,}) - "
+ f"Transient {mem_layout.transient_bytes/1e9:.2f} GB - "
+ f"Slack {mem_layout.slack_bytes/1e9:.2f} GB"
+ )
+
+ with col_shard:
+ st.markdown("**Tensor sharding (which dim is split)**")
+ st.caption(
+ "W_Q/W_K/W_V/W_gate/W_up: **output cols**. "
+ "W_O/W_down: **input rows**. "
+ "KV cache: **rows by CP** x **cols by TP**."
+ )
+ s_c1, s_c2 = st.columns([1, 2])
+ with s_c1:
+ my_pe_view = st.number_input("PE index",
+ min_value=0,
+ max_value=max(0, topo.tp - 1),
+ value=0, key="my_pe_view")
+ with s_c2:
+ show_physical = st.checkbox(
+ "Enlarge + show physical PE/cube mapping",
+ value=False, key="tsv_physical",
+ help=("Larger figure with each shard cell labelled by owning "
+ "cube/PE. Especially useful for KV cache when CP is "
+ "split across cubes and PEs."),
+ )
+ fig_shard = draw_tensor_sharding(cfg, my_pe=my_pe_view,
+ show_physical=show_physical)
+ st.pyplot(fig_shard, width='stretch')
+ plt.close(fig_shard)
+
+ st.divider()
+
+ # Row D: PE-level view (wide, full-width) inside an expander
+ with st.expander(
+ f"PE-level view of one CP group (TP={topo.tp}, "
+ f"H_q={model.h_q}, H_kv={model.h_kv})",
+ expanded=False,
+ ):
+ st.caption(
+ f"One CP rank's {topo.tp} PE(s). Attention weights replicated "
+ f"across CP groups; KV sharded by CP (S/{topo.cp} tokens) x TP heads; "
+ f"FFN by TP x EP."
+ )
+ fig_pe = draw_pe_layout(cfg)
+ st.pyplot(fig_pe, width='stretch')
+ plt.close(fig_pe)
+
+ # Row E: Replication + Optimization panel (side-by-side)
+ st.divider()
+ col_rep, col_opt = st.columns([1, 1], gap="medium")
+
+ with col_rep:
+ st.markdown("**Replicated / duplicated storage**")
+ rep_entries = replication_report(cfg)
+ if not rep_entries:
+ st.caption("No duplicated storage detected - every tensor is uniquely sharded.")
+ else:
+ rep_rows = []
+ total_waste = 0
+ for e in rep_entries:
+ rep_rows.append({
+ "Tensor": e.tensor,
+ "Per-PE (GB)": round(e.per_pe_bytes / 1e9, 3),
+ "Replicated": e.replicated_across,
+ "Copies": e.copies,
+ "Waste/PE (GB)": round(e.wasted_bytes / 1e9, 3),
+ })
+ total_waste += e.wasted_bytes
+ st.dataframe(pd.DataFrame(rep_rows), width='stretch',
+ hide_index=True)
+ st.caption(
+ f"Total per-PE duplicated bytes: "
+ f"**{total_waste/1e9:.2f} GB** "
+ f"(vs {mem_layout.budget_bytes/1e9:.1f} GB budget). "
+ f"Each 'copy' represents one identical block stored on "
+ f"another rank."
+ )
+
+ with col_opt:
+ st.markdown("**Optimization opportunities**")
+ hints = optimization_hints(cfg)
+ if not hints:
+ st.caption("No obvious improvements at this configuration.")
+ else:
+ _icons = {"warn": ":warning:", "info": ":bulb:", "good": ":white_check_mark:"}
+ _cats = {"space": "SPACE", "comm": "COMM",
+ "compute": "COMPUTE", "layout": "LAYOUT"}
+ for h in hints:
+ icon = _icons.get(h.severity, ":bulb:")
+ cat = _cats.get(h.category, h.category.upper())
+ st.markdown(f"{icon} **[{cat}]** {h.message}")
+
+ st.divider()
+
+ # Row F: Per-PE breakdown table (collapsible)
+ with st.expander("Per-PE breakdown table", expanded=False):
+ rows = []
+ for pe_id in range(topo.tp):
+ q_heads = _q_heads_for_pe(pe_id, topo.tp, model.h_q)
+ kv_heads, kv_note = _kv_heads_for_pe(
+ pe_id, topo.tp, model.h_kv, topo.kv_shard_mode,
+ )
+ bytes_ = _per_pe_bytes(cfg, pe_id)
+ weights_bytes = sum(v for k, v in bytes_.items()
+ if k not in ("KV cache", "Transient"))
+ rows.append({
+ "PE": pe_id,
+ "Q heads": (f"{q_heads[0]}-{q_heads[-1]}"
+ if len(q_heads) > 1 else str(q_heads[0])
+ if q_heads else "-"),
+ "KV heads": (str(kv_heads[0]) + (f" ({kv_note})" if kv_note else "")
+ if kv_heads else "-"),
+ "W_Q (MB)": round(bytes_["W_Q"] / 1e6, 2),
+ "W_K (MB)": round(bytes_["W_K"] / 1e6, 2),
+ "W_V (MB)": round(bytes_["W_V"] / 1e6, 2),
+ "W_O (MB)": round(bytes_["W_O"] / 1e6, 2),
+ "FFN (MB)": round((bytes_["W_gate"] + bytes_["W_up"]
+ + bytes_["W_down"]) / 1e6, 2),
+ "Weights total (GB)": round(weights_bytes / 1e9, 3),
+ "KV cache (GB)": round(bytes_["KV cache"] / 1e9, 3),
+ "Transient (MB)": round(bytes_["Transient"] / 1e6, 2),
+ })
+ st.dataframe(pd.DataFrame(rows), width='stretch', hide_index=True)
+
+
+# ── TAB 2: Memory breakdown ─────────────────────────────────────
+with tab_memory:
+ mem = compute_memory(cfg)
+ layers_per_stage = (model.layers + pp - 1) // pp
+
+ parallelism_str = (
+ f"CP={cp}, TP={tp}, PP={pp}, DP={dp}, EP={ep}"
+ f" | layers={model.layers}, layers/stage={layers_per_stage}"
+ f" | KV mode: {kv_shard_mode}"
+ )
+ st.caption(f"**Parallelism:** {parallelism_str}")
+
+ # ─── ATTENTION (per-layer) ─────────────────────────────────
+ st.subheader(f"Attention weights (per layer) — {parallelism_str}")
+ attn_rows = attention_weight_rows(cfg)
+ display_attn = [{k: v for k, v in r.items() if not k.startswith("_")}
+ for r in attn_rows]
+ st.dataframe(pd.DataFrame(display_attn), width='stretch', hide_index=True)
+
+ attn_bytes_1L_global = sum(r["_p_global"] for r in attn_rows) * model.bytes_per_elem
+ attn_bytes_1L_pe = sum(r["_p_per_pe"] for r in attn_rows) * model.bytes_per_elem
+ attn_bytes_all_global = attn_bytes_1L_global * model.layers
+ attn_bytes_all_pe = attn_bytes_1L_pe * layers_per_stage
+
+ ac1, ac2, ac3, ac4 = st.columns(4)
+ ac1.metric("Attn / layer (global)", f"{attn_bytes_1L_global/1e6:.2f} MB")
+ ac2.metric(f"Attn all {model.layers} layers", f"{attn_bytes_all_global/1e9:.3f} GB")
+ ac3.metric("Attn / layer (per PE)", f"{attn_bytes_1L_pe/1e6:.2f} MB")
+ ac4.metric(f"Attn per PE ({layers_per_stage} layers/stage)",
+ f"{attn_bytes_all_pe/1e9:.3f} GB")
+
+ # ─── FFN / MoE (per-layer) ─────────────────────────────────
+ is_moe = "MoE" in (preset.family + preset.note)
+ ffn_label = "FFN (activated experts approx)" if is_moe else "FFN"
+ st.subheader(f"{ffn_label} weights (per layer) — {parallelism_str}")
+ ffn_rows = ffn_weight_rows(cfg)
+ display_ffn = [{k: v for k, v in r.items() if not k.startswith("_")}
+ for r in ffn_rows]
+ st.dataframe(pd.DataFrame(display_ffn), width='stretch', hide_index=True)
+
+ ffn_bytes_1L_global = sum(r["_p_global"] for r in ffn_rows) * model.bytes_per_elem
+ ffn_bytes_1L_pe = sum(r["_p_per_pe"] for r in ffn_rows) * model.bytes_per_elem
+ ffn_bytes_all_global = ffn_bytes_1L_global * model.layers
+ ffn_bytes_all_pe = ffn_bytes_1L_pe * layers_per_stage
+
+ fc1, fc2, fc3, fc4 = st.columns(4)
+ fc1.metric("FFN / layer (global)", f"{ffn_bytes_1L_global/1e6:.2f} MB")
+ fc2.metric(f"FFN all {model.layers} layers", f"{ffn_bytes_all_global/1e9:.3f} GB")
+ fc3.metric("FFN / layer (per PE)", f"{ffn_bytes_1L_pe/1e6:.2f} MB")
+ fc4.metric(f"FFN per PE ({layers_per_stage} layers/stage)",
+ f"{ffn_bytes_all_pe/1e9:.3f} GB")
+
+ # ─── KV cache (per-layer) ─────────────────────────────────
+ st.subheader(f"KV cache (per layer, S_kv = {s_kv:,}) — {parallelism_str}")
+ kv_rows = kv_cache_rows(cfg)
+ display_kv = [{k: v for k, v in r.items() if not k.startswith("_")}
+ for r in kv_rows]
+ st.dataframe(pd.DataFrame(display_kv), width='stretch', hide_index=True)
+
+ kv_bytes_1L_global = sum(r["_p_global"] for r in kv_rows) * model.bytes_per_elem
+ kv_bytes_1L_pe = sum(r["_p_per_pe"] for r in kv_rows) * model.bytes_per_elem
+ kv_bytes_all_global = kv_bytes_1L_global * model.layers
+ kv_bytes_all_pe = kv_bytes_1L_pe * layers_per_stage
+
+ kc1, kc2, kc3, kc4 = st.columns(4)
+ kc1.metric("KV / layer (global)", f"{kv_bytes_1L_global/1e6:.2f} MB")
+ kc2.metric(f"KV all {model.layers} layers", f"{kv_bytes_all_global/1e9:.3f} GB")
+ kc3.metric("KV / layer (per PE)", f"{kv_bytes_1L_pe/1e6:.2f} MB")
+ kc4.metric(f"KV per PE ({layers_per_stage} layers/stage)",
+ f"{kv_bytes_all_pe/1e9:.3f} GB")
+
+ st.divider()
+
+ # ─── FULL-MODEL TOTALS ───────────────────────────────────
+ st.subheader("Full-model totals (all layers)")
+ total_w_global = attn_bytes_all_global + ffn_bytes_all_global
+ total_all_global = total_w_global + kv_bytes_all_global
+ total_per_pe = attn_bytes_all_pe + ffn_bytes_all_pe + kv_bytes_all_pe
+ tc1, tc2, tc3, tc4 = st.columns(4)
+ tc1.metric("Weights (all layers, global)", f"{total_w_global/1e9:.2f} GB")
+ tc2.metric("KV (all layers, global)", f"{kv_bytes_all_global/1e9:.2f} GB")
+ tc3.metric("Weights + KV (global)", f"{total_all_global/1e9:.2f} GB")
+ tc4.metric("Per-PE total (W+KV, layers/stage)",
+ f"{total_per_pe/1e9:.2f} GB")
+
+ # ─── Per-PE memory pie ───────────────────────────────────
+ st.subheader("Per-PE memory footprint (with transient + slack)")
+ pie_col, info_col = st.columns([1, 1])
+ with pie_col:
+ labels = ["Weights", "KV cache", "Transient", "Slack"]
+ vals = [mem.weights_bytes, mem.kv_cache_bytes,
+ mem.transient_bytes, mem.slack_bytes]
+ colors = ["#3a86ff", "#8338ec", "#ffbe0b", "#adb5bd"]
+
+ fig_mem, ax_mem = plt.subplots(figsize=(6, 6))
+ ax_mem.pie(vals,
+ labels=[f"{l}\n{v/1e9:.2f} GB" for l, v in zip(labels, vals)],
+ colors=colors, autopct="%1.1f%%", startangle=90)
+ ax_mem.set_title(f"Per-PE (budget: {mem.budget_bytes/1e9:.1f} GB)")
+ st.pyplot(fig_mem, width='stretch')
+ plt.close(fig_mem)
+ with info_col:
+ if mem.over_budget:
+ st.error(f"OVER BUDGET by "
+ f"{(mem.used_bytes - mem.budget_bytes)/1e9:.2f} GB")
+ else:
+ st.success(f"{mem.slack_bytes/1e9:.2f} GB slack")
+
+ st.markdown(
+ f"- Weights: **{mem.weights_bytes/1e9:.2f} GB** "
+ f"(layers/PP = {layers_per_stage})\n"
+ f"- KV cache: **{mem.kv_cache_bytes/1e9:.2f} GB** "
+ f"(S_local = {topo.s_local:,})\n"
+ f"- Transient: {mem.transient_bytes/1e9:.2f} GB\n"
+ f"- Slack: {mem.slack_bytes/1e9:.2f} GB"
+ )
+
+
+# ── TAB 3: Per-stage latency ────────────────────────────────────
+with tab_stages:
+ st.subheader("Per-stage attention latency")
+ stages = all_stages(cfg)
+
+ rows = []
+ for s in stages:
+ rows.append({
+ "Stage": s.name,
+ "Bound": s.bound,
+ "Compute (us)": round(s.compute_s * 1e6, 3),
+ "Memory (us)": round(s.memory_s * 1e6, 3),
+ "Comm (us)": round(s.comm_s * 1e6, 3),
+ "Visible (us)": round(s.visible_s * 1e6, 3),
+ "Formula": s.formula,
+ })
+ df = pd.DataFrame(rows)
+ st.dataframe(df, width='stretch', hide_index=True)
+
+ total_visible_s = sum(s.visible_s for s in stages)
+ total_layers = total_visible_s * model.layers
+ st.markdown(
+ f"**Per-layer attention (analytical): {total_visible_s*1e6:.2f} us** \n"
+ f"**Full-model attention ({model.layers} layers): "
+ f"{total_layers*1e6:.2f} us = {total_layers*1e3:.2f} ms**"
+ )
+
+ st.subheader("Stage cost breakdown")
+ fig_bar, ax_bar = plt.subplots(figsize=(12, 5))
+ names = [s.name for s in stages]
+ cmp_vals = [s.compute_s * 1e6 for s in stages]
+ mem_vals = [s.memory_s * 1e6 for s in stages]
+ comm_vals = [s.comm_s * 1e6 for s in stages]
+
+ x = range(len(names))
+ ax_bar.bar(x, cmp_vals, label="Compute", color="#3a86ff")
+ ax_bar.bar(x, mem_vals, bottom=cmp_vals, label="Memory", color="#ffbe0b")
+ ax_bar.bar(x, comm_vals,
+ bottom=[c + m for c, m in zip(cmp_vals, mem_vals)],
+ label="Comm", color="#d90429")
+ ax_bar.set_xticks(x)
+ ax_bar.set_xticklabels(names, rotation=45, ha="right", fontsize=8)
+ ax_bar.set_ylabel("Time (us)")
+ ax_bar.set_title(f"Stage breakdown - mode={mode}, T_q={topo.T_q}")
+ ax_bar.legend()
+ ax_bar.grid(axis="y", alpha=0.3)
+ plt.tight_layout()
+ st.pyplot(fig_bar, width='stretch')
+ plt.close(fig_bar)
+
+
+# ── TAB 4: Save & compare ───────────────────────────────────────
+def _snapshot_current_config():
+ """Capture a compact snapshot of the current cfg + key metrics."""
+ _stages_attn = all_stages(cfg)
+ _stages_ffn = all_ffn_stages(cfg)
+ attn_us = sum(s.visible_s for s in _stages_attn) * 1e6
+ ffn_us = sum(s.visible_s for s in _stages_ffn) * 1e6
+ mem = compute_memory(cfg)
+ # Per-stage latencies keyed by short prefix (S1, S2, ..., C1, F1, ..., CF1).
+ # Value = seconds (raw), so display can format uniformly across snapshots.
+ _by_prefix_attn = {s.name.split()[0]: s.visible_s for s in _stages_attn}
+ _by_prefix_ffn = {s.name.split()[0]: s.visible_s for s in _stages_ffn}
+ # Preserve the full descriptive name too so a hover/second column can
+ # show what the abbreviation means for this particular snapshot.
+ _fullname_attn = {s.name.split()[0]: s.name for s in _stages_attn}
+ _fullname_ffn = {s.name.split()[0]: s.name for s in _stages_ffn}
+ return {
+ "model": model.name,
+ "mode": topo.mode,
+ "s_kv": topo.s_kv,
+ "cp": topo.cp, "tp": topo.tp,
+ "pp": topo.pp, "dp": topo.dp, "ep": topo.ep,
+ "cp_placement": topo.cp_placement,
+ "tp_placement": topo.tp_placement,
+ "ffn_scope": topo.ffn_shard_scope,
+ "kv_mode": topo.kv_shard_mode,
+ "sip_topology": topo.sip_topology,
+ "cp_ring_variant": topo.cp_ring_variant,
+ "pes": topo.total_pes,
+ "cubes": topo.cubes_used,
+ "sips": topo.sips_used,
+ "pe_hbm_gb": machine.pe_hbm_gb,
+ "peak_tflops": machine.peak_tflops_f16,
+ "bw_hbm_gbs": machine.bw_hbm_gbs,
+ "bw_inter_gbs": machine.bw_inter_gbs,
+ "bw_intersip_gbs": machine.bw_intersip_gbs,
+ "weights_gb": round(mem.weights_bytes / 1e9, 3),
+ "kv_gb": round(mem.kv_cache_bytes / 1e9, 3),
+ "transient_gb": round(mem.transient_bytes / 1e9, 3),
+ "used_gb": round(mem.used_bytes / 1e9, 3),
+ "slack_gb": round(mem.slack_bytes / 1e9, 3),
+ "over_budget": mem.over_budget,
+ "attn_us": round(attn_us, 3),
+ "ffn_us": round(ffn_us, 3),
+ "layer_us": round(attn_us + ffn_us, 3),
+ "stages_attn": _by_prefix_attn,
+ "stages_ffn": _by_prefix_ffn,
+ "stage_names_attn": _fullname_attn,
+ "stage_names_ffn": _fullname_ffn,
+ }
+
+
+with tab_compare:
+ st.markdown(
+ "Save the current configuration under a name, then load it back "
+ "later or compare multiple saved configs side-by-side."
+ )
+ if "saved_cfgs" not in st.session_state:
+ st.session_state["saved_cfgs"] = {} # name -> snapshot dict
+
+ # Default name: next free config1, config2, ...
+ def _next_config_name():
+ i = 1
+ while f"config{i}" in st.session_state.get("saved_cfgs", {}):
+ i += 1
+ return f"config{i}"
+
+ sv_c1, sv_c2, sv_c3 = st.columns([3, 1, 1])
+ with sv_c1:
+ _save_name = st.text_input(
+ "Config name",
+ value=_next_config_name(),
+ key=f"cfg_save_name_{len(st.session_state.get('saved_cfgs', {}))}",
+ )
+ with sv_c2:
+ if st.button("Save current", width='stretch'):
+ if _save_name.strip():
+ st.session_state["saved_cfgs"][_save_name.strip()] = \
+ _snapshot_current_config()
+ st.success(f"Saved '{_save_name.strip()}'")
+ st.rerun()
+ else:
+ st.warning("Give it a name first.")
+ with sv_c3:
+ if st.button("Clear all", width='stretch'):
+ st.session_state["saved_cfgs"] = {}
+ st.rerun()
+
+ _saved_all = st.session_state.get("saved_cfgs", {})
+ if not _saved_all:
+ st.info("No saved configurations yet.")
+ else:
+ st.markdown(f"**{len(_saved_all)} saved configuration(s)** - "
+ f"pick which to compare below")
+
+ # Pick which configs to include in the comparison (default all)
+ _picked = st.multiselect(
+ "Compare",
+ options=list(_saved_all.keys()),
+ default=list(_saved_all.keys()),
+ key="compare_picked",
+ )
+ _saved = {n: _saved_all[n] for n in _picked}
+
+ # Delete individual
+ with st.expander("Manage (delete individual)", expanded=False):
+ for _n in list(_saved_all.keys()):
+ d1, d2, d3 = st.columns([5, 1, 1])
+ d1.text(_n)
+ if d2.button("Rename", key=f"rn_btn_{_n}"):
+ st.session_state[f"rn_active_{_n}"] = True
+ if d3.button("Delete", key=f"del_{_n}"):
+ del st.session_state["saved_cfgs"][_n]
+ st.rerun()
+ if st.session_state.get(f"rn_active_{_n}"):
+ new_name = st.text_input(
+ f"New name for '{_n}'", value=_n,
+ key=f"rn_txt_{_n}",
+ )
+ if st.button("Confirm rename", key=f"rn_ok_{_n}"):
+ if new_name and new_name != _n and new_name not in _saved_all:
+ st.session_state["saved_cfgs"][new_name] = \
+ st.session_state["saved_cfgs"].pop(_n)
+ st.session_state[f"rn_active_{_n}"] = False
+ st.rerun()
+
+ if not _saved:
+ st.info("No configurations selected. Tick some above to compare.")
+ st.stop()
+
+ # Comparison table: rows = metrics, columns = config names
+ _key_order = [
+ ("Model / workload", ["model", "mode", "s_kv"]),
+ ("Parallelism", ["cp", "tp", "pp", "dp", "ep",
+ "cp_placement", "tp_placement",
+ "ffn_scope", "kv_mode", "cp_ring_variant",
+ "sip_topology"]),
+ ("Physical", ["pes", "cubes", "sips"]),
+ ("Hardware knobs", ["pe_hbm_gb", "peak_tflops",
+ "bw_hbm_gbs", "bw_inter_gbs",
+ "bw_intersip_gbs"]),
+ ("Memory / PE (GB)", ["weights_gb", "kv_gb", "transient_gb",
+ "used_gb", "slack_gb", "over_budget"]),
+ ("Per-layer latency (us)", ["attn_us", "ffn_us", "layer_us"]),
+ ]
+
+ # Build one wide DataFrame - metrics x saved names
+ _names = list(_saved.keys())
+ _rows_out = []
+ for _sec, _keys in _key_order:
+ _rows_out.append({"metric": f"= {_sec}",
+ **{n: "" for n in _names}})
+ for k in _keys:
+ _rows_out.append({
+ "metric": k,
+ **{n: _saved[n].get(k, "-") for n in _names},
+ })
+ _df_cmp = pd.DataFrame(_rows_out)
+
+ # Highlight best (lowest) for latency + used_gb rows, worst (over_budget True) red.
+ def _cmp_style(row):
+ styles = [""] * len(row)
+ m = row["metric"]
+ vals = row[1:]
+ try:
+ nums = {k: float(v) for k, v in vals.items()
+ if isinstance(v, (int, float)) and not isinstance(v, bool)}
+ except Exception:
+ nums = {}
+ if m in ("layer_us", "attn_us", "ffn_us", "used_gb",
+ "weights_gb", "kv_gb"):
+ if nums:
+ best = min(nums, key=nums.get)
+ for i, (k, v) in enumerate(vals.items(), start=1):
+ if k == best:
+ styles[i] = "background-color: #d4edda; font-weight:600;"
+ if m == "slack_gb" and nums:
+ best = max(nums, key=nums.get)
+ for i, (k, v) in enumerate(vals.items(), start=1):
+ if k == best:
+ styles[i] = "background-color: #d4edda; font-weight:600;"
+ if m == "over_budget":
+ for i, (k, v) in enumerate(vals.items(), start=1):
+ if v is True or v == "True":
+ styles[i] = "background-color: #f8d7da; font-weight:600;"
+ if m.startswith("= "):
+ styles = ["background-color: #e9ecef; font-weight:600;"] * len(row)
+ return styles
+
+ _styled_cmp = _df_cmp.style.apply(_cmp_style, axis=1)
+ st.dataframe(_styled_cmp, width='stretch', hide_index=True,
+ row_height=28)
+ st.caption(
+ "Green cell = best value in that row (lowest latency / memory, "
+ "highest slack). Red = over-budget. Grey rows = section headers."
+ )
+
+ # ── Per-stage side-by-side comparison ────────────────────
+ _fmt = _fmt_time # reuse the ns/us/ms auto-formatter
+ _short_label = {
+ "S1": "S1 RMSNorm", "S2": "S2 W_Q", "S3": "S3 W_K+W_V",
+ "S4": "S4 KV append", "S5": "S5 Q.K^T", "S6": "S6 softmax",
+ "S7": "S7 P.V", "S8": "S8 merge (+ AR if decode)",
+ "S9": "S9 norm O/l", "S10": "S10 W_O",
+ "C1": "C1 CP ring (prefill only)",
+ "C2": "C2 TP AR (W_O)",
+ "C3": "C3 Score AR (head-split)",
+ "F1": "F1 RMSNorm", "F2": "F2 W_gate", "F3": "F3 W_up",
+ "F4": "F4 SwiGLU", "F5": "F5 W_down", "CF1": "CF1 FFN AR",
+ }
+
+ def _stage_rows_across(prefix_order, key):
+ """Build one row per stage prefix, with each saved config as a col."""
+ rows = []
+ for pfx in prefix_order:
+ row = {"stage": _short_label.get(pfx, pfx)}
+ sec_values = {}
+ for n, snap in _saved.items():
+ v = snap.get(key, {}).get(pfx)
+ if v is None:
+ row[n] = "-"
+ else:
+ row[n] = _fmt(v)
+ sec_values[n] = v
+ rows.append(row)
+ return rows
+
+ def _stage_style_row(row):
+ styles = [""] * len(row)
+ # Try to find min seconds — reconstruct from formatted strings
+ # is fragile; recompute from saved secs directly.
+ pfx = row["stage"].split()[0]
+ secs = {}
+ for n, snap in _saved.items():
+ key = "stages_attn" if pfx in _attn_prefixes else "stages_ffn"
+ v = snap.get(key, {}).get(pfx)
+ if v is not None and v > 0:
+ secs[n] = v
+ if secs:
+ best = min(secs, key=secs.get)
+ for i, name in enumerate(row.index[1:], start=1):
+ if name == best:
+ styles[i] = "background-color: #d4edda; font-weight:600;"
+ return styles
+
+ _attn_prefixes = ["S1", "S2", "S3", "S4", "S5", "S6", "S7",
+ "S8", "S9", "S10", "C1", "C2", "C3"]
+ _ffn_prefixes = ["F1", "F2", "F3", "F4", "F5", "CF1"]
+
+ st.markdown("**Attention per-stage latency (side-by-side)**")
+ _df_attn = pd.DataFrame(_stage_rows_across(_attn_prefixes, "stages_attn"))
+ st.dataframe(_df_attn.style.apply(_stage_style_row, axis=1),
+ width='stretch', hide_index=True, row_height=28)
+
+ st.markdown("**FFN per-stage latency (side-by-side)**")
+ _df_ffn = pd.DataFrame(_stage_rows_across(_ffn_prefixes, "stages_ffn"))
+ st.dataframe(_df_ffn.style.apply(_stage_style_row, axis=1),
+ width='stretch', hide_index=True, row_height=28)
+
+ st.caption(
+ "Times auto-scale (ns/us/ms). Green cell = fastest for that stage. "
+ "'-' means the stage is absent for that config (e.g. C1 in "
+ "decode, since the O/m/l all-reduce is folded into S8)."
+ )
diff --git a/tests/analytical_visualization/autosuggest.py b/tests/analytical_visualization/autosuggest.py
new file mode 100644
index 0000000..7d8f7c4
--- /dev/null
+++ b/tests/analytical_visualization/autosuggest.py
@@ -0,0 +1,106 @@
+"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.
+
+Strategy: iterate over candidate (CP, TP, PP) triples in order of
+increasing total PE count and return the first one that satisfies
+memory + kernel-support constraints with >10% slack. If none fits, return
+the best-effort configuration and flag over-budget.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Iterable
+
+from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
+from .memory_layout import compute_memory
+
+
+# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
+_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
+_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
+_PP_OPTIONS = (1, 2, 4, 8, 16, 32)
+
+
+@dataclass
+class Suggestion:
+ cp: int
+ tp: int
+ pp: int
+ weights_gb: float
+ kv_gb: float
+ transient_gb: float
+ slack_gb: float
+ fits: bool
+ pes_used: int
+ sips_used: int
+ reason: str = ""
+
+
+def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
+ """Yield (CP, TP, PP) in order of increasing PE count."""
+ triples: list[tuple[int, int, int]] = []
+ for tp in _TP_OPTIONS:
+ for cp in _CP_OPTIONS:
+ for pp in _PP_OPTIONS:
+ # PP must not exceed layer count.
+ if pp > model.layers:
+ continue
+ # Skip TP > H_q * some factor (unrealistic).
+ if tp > model.h_q * 4:
+ continue
+ triples.append((cp, tp, pp))
+ # Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
+ # (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
+ triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
+ for t in triples:
+ yield t
+
+
+def _score_candidate(cp: int, tp: int, pp: int,
+ model: ModelConfig, machine: MachineParams,
+ s_kv: int, mode: str,
+ slack_frac: float = 0.10) -> Suggestion:
+ topo = TopologyConfig(cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode)
+ cfg = FullConfig(model=model, topo=topo, machine=machine)
+ mem = compute_memory(cfg)
+
+ fits = (not mem.over_budget
+ and mem.slack_bytes >= slack_frac * mem.budget_bytes)
+ reason = ""
+ if mem.over_budget:
+ reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
+ f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
+ elif not fits:
+ reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
+ return Suggestion(
+ cp=cp, tp=tp, pp=pp,
+ weights_gb=mem.weights_bytes / 1e9,
+ kv_gb=mem.kv_cache_bytes / 1e9,
+ transient_gb=mem.transient_bytes / 1e9,
+ slack_gb=mem.slack_bytes / 1e9,
+ fits=fits,
+ pes_used=topo.total_pes,
+ sips_used=topo.sips_used,
+ reason=reason,
+ )
+
+
+def auto_suggest(model: ModelConfig, machine: MachineParams,
+ s_kv: int, mode: str = "decode",
+ slack_frac: float = 0.10) -> Suggestion:
+ """Return the smallest-PE-count (CP, TP, PP) that fits.
+
+ If no candidate fits, returns the best-effort one (highest slack,
+ even if negative) with fits=False.
+ """
+ best_fit: Suggestion | None = None
+ best_effort: Suggestion | None = None
+
+ for cp, tp, pp in _iter_candidates(model):
+ s = _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac)
+ if s.fits and best_fit is None:
+ best_fit = s
+ break # candidates are pre-sorted by PE count
+ if best_effort is None or s.slack_gb > best_effort.slack_gb:
+ best_effort = s
+
+ return best_fit or best_effort # type: ignore
diff --git a/tests/analytical_visualization/memory_layout.py b/tests/analytical_visualization/memory_layout.py
new file mode 100644
index 0000000..e7c6274
--- /dev/null
+++ b/tests/analytical_visualization/memory_layout.py
@@ -0,0 +1,225 @@
+"""Per-PE memory footprint: weights + KV cache + activations + slack.
+
+All formulas are per-PE, per-PP-stage. If PP > 1, each PE holds only
+its stage's layers (layers/PP).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import FullConfig
+
+
+@dataclass
+class MemoryBreakdown:
+ weights_bytes: int
+ kv_cache_bytes: int
+ transient_bytes: int
+ budget_bytes: int
+
+ @property
+ def used_bytes(self) -> int:
+ return self.weights_bytes + self.kv_cache_bytes + self.transient_bytes
+
+ @property
+ def slack_bytes(self) -> int:
+ return max(0, self.budget_bytes - self.used_bytes)
+
+ @property
+ def over_budget(self) -> bool:
+ return self.used_bytes > self.budget_bytes
+
+
+def per_pe_weight_bytes(cfg: FullConfig) -> int:
+ """Attention + FFN weights per PE, bf16. Divided by TP, PP, and EP.
+
+ EP divides the FFN (experts) across ranks; attention weights are
+ unaffected. When kv_shard_mode='replicate' and TP > H_kv, each KV
+ head is replicated (per-PE W_K/W_V size doesn't drop below one head).
+ """
+ m = cfg.model
+ tp = cfg.topo.tp
+ pp = cfg.topo.pp
+ ep = max(1, cfg.topo.ep)
+
+ hq_per_pe = cfg.h_q_per_pe
+ if cfg.topo.kv_shard_mode == "replicate":
+ # Each KV head held fully; replicated across ranks if TP > H_kv.
+ hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
+ else: # "split"
+ # Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
+ hkv_per_pe_bytes = m.h_kv / tp
+
+ per_layer_attn = (
+ m.hidden * hq_per_pe * m.d_head + # W_Q
+ m.hidden * hkv_per_pe_bytes * m.d_head + # W_K
+ m.hidden * hkv_per_pe_bytes * m.d_head + # W_V
+ hq_per_pe * m.d_head * m.hidden # W_O
+ )
+ # FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
+ # EP further divides (MoE experts).
+ ffn_div = cfg.ffn_shard_divisor * ep
+ per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
+ per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
+
+ layers_per_stage = (m.layers + pp - 1) // pp
+ return int(per_layer * layers_per_stage)
+
+
+def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
+ """K + V per PE, across all layers this stage holds.
+
+ - CP shards the sequence dim -> S_local = S_kv/CP tokens per PE.
+ - TP splits KV heads across ranks. If kv_shard_mode='replicate' and
+ TP > H_kv, each KV head is duplicated across TP/H_kv ranks so
+ per-PE storage doesn't shrink below 1 head.
+ - PP: only layers_per_stage layers stored per PE.
+ """
+ m = cfg.model
+ pp = cfg.topo.pp
+ tp = cfg.topo.tp
+ if cfg.topo.kv_shard_mode == "replicate":
+ hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
+ else:
+ hkv_per_pe_bytes = m.h_kv / tp
+ layers_per_stage = (m.layers + pp - 1) // pp
+ per_layer = 2 * cfg.topo.s_local * hkv_per_pe_bytes * m.d_head * m.bytes_per_elem
+ return int(per_layer * layers_per_stage)
+
+
+def per_pe_transient_bytes(cfg: FullConfig) -> int:
+ """Rough peak transient (activations, GEMM outputs) per PE."""
+ m = cfg.model
+ T_q = cfg.topo.T_q
+ hq_per_pe = cfg.h_q_per_pe
+
+ if cfg.topo.mode == "decode":
+ return 4 * (m.hidden + m.head_dim_total_q // cfg.topo.tp) * m.bytes_per_elem
+ else:
+ TILE = 1024
+ tile_score = hq_per_pe * T_q * TILE * m.bytes_per_elem
+ return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
+
+
+def compute_memory(cfg: FullConfig) -> MemoryBreakdown:
+ return MemoryBreakdown(
+ weights_bytes=per_pe_weight_bytes(cfg),
+ kv_cache_bytes=per_pe_kv_cache_bytes(cfg),
+ transient_bytes=per_pe_transient_bytes(cfg),
+ budget_bytes=cfg.machine.pe_budget_bytes,
+ )
+
+
+def _one_layer_row(name: str,
+ global_shape: tuple[int, int],
+ per_pe_shape: tuple[int, int],
+ bytes_per_elem: int) -> dict:
+ """Per-tensor row for ONE layer (both global and per-PE shard)."""
+ p_global = global_shape[0] * global_shape[1]
+ p_per_pe = per_pe_shape[0] * per_pe_shape[1]
+ return {
+ "Tensor": name,
+ "Global shape": f"({global_shape[0]}, {global_shape[1]})",
+ "Params/layer": f"{p_global/1e6:.2f} M",
+ "Bytes/layer (global)": f"{p_global * bytes_per_elem / 1e6:.2f} MB",
+ "Per-PE shape": f"({per_pe_shape[0]}, {per_pe_shape[1]})",
+ "Bytes/layer (per PE)": f"{p_per_pe * bytes_per_elem / 1e6:.2f} MB",
+ "_p_global": p_global,
+ "_p_per_pe": p_per_pe,
+ }
+
+
+def attention_weight_rows(cfg: FullConfig) -> list[dict]:
+ """Per-tensor rows for attention weights (one layer)."""
+ m = cfg.model
+ hq_per_pe = cfg.h_q_per_pe
+ hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
+ return [
+ _one_layer_row("W_Q", (m.hidden, m.h_q * m.d_head),
+ (m.hidden, hq_per_pe * m.d_head),
+ m.bytes_per_elem),
+ _one_layer_row("W_K", (m.hidden, m.h_kv * m.d_head),
+ (m.hidden, hkv_per_pe * m.d_head),
+ m.bytes_per_elem),
+ _one_layer_row("W_V", (m.hidden, m.h_kv * m.d_head),
+ (m.hidden, hkv_per_pe * m.d_head),
+ m.bytes_per_elem),
+ _one_layer_row("W_O", (m.h_q * m.d_head, m.hidden),
+ (hq_per_pe * m.d_head, m.hidden),
+ m.bytes_per_elem),
+ ]
+
+
+def ffn_weight_rows(cfg: FullConfig) -> list[dict]:
+ """Per-tensor rows for FFN weights (one layer, activated for MoE)."""
+ m = cfg.model
+ ffn_per_pe = m.ffn_dim // cfg.topo.tp
+ return [
+ _one_layer_row("W_gate", (m.hidden, m.ffn_dim),
+ (m.hidden, ffn_per_pe), m.bytes_per_elem),
+ _one_layer_row("W_up", (m.hidden, m.ffn_dim),
+ (m.hidden, ffn_per_pe), m.bytes_per_elem),
+ _one_layer_row("W_down", (m.ffn_dim, m.hidden),
+ (ffn_per_pe, m.hidden), m.bytes_per_elem),
+ ]
+
+
+def kv_cache_rows(cfg: FullConfig) -> list[dict]:
+ """Per-tensor row for K and V cache (one layer)."""
+ m = cfg.model
+ hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
+ b = m.bytes_per_elem
+ global_k = cfg.topo.s_kv * m.h_kv * m.d_head
+ per_pe_k = cfg.topo.s_local * hkv_per_pe * m.d_head
+ return [
+ {
+ "Tensor": "K cache",
+ "Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
+ "Params/layer": f"{global_k/1e6:.2f} M",
+ "Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
+ "Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
+ "Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
+ "_p_global": global_k,
+ "_p_per_pe": per_pe_k,
+ },
+ {
+ "Tensor": "V cache",
+ "Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
+ "Params/layer": f"{global_k/1e6:.2f} M",
+ "Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
+ "Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
+ "Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
+ "_p_global": global_k,
+ "_p_per_pe": per_pe_k,
+ },
+ ]
+
+
+def sum_bytes_all_layers(rows: list[dict], cfg: FullConfig,
+ per_pe: bool = False) -> int:
+ """Sum bytes across all rows × N layers (or layers/PP for per_pe)."""
+ b = cfg.model.bytes_per_elem
+ layers = cfg.model.layers
+ if per_pe:
+ layers = (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
+ key = "_p_per_pe" if per_pe else "_p_global"
+ return sum(r[key] for r in rows) * b * layers
+
+
+def total_weight_bytes_full_model(cfg: FullConfig) -> int:
+ """Total bytes of ALL weights (attention + FFN) for the full unsharded model."""
+ m = cfg.model
+ attn_per_layer = (
+ m.hidden * m.h_q * m.d_head + # W_Q
+ m.hidden * m.h_kv * m.d_head * 2 + # W_K, W_V
+ m.h_q * m.d_head * m.hidden # W_O
+ )
+ ffn_per_layer = 3 * m.hidden * m.ffn_dim
+ return (attn_per_layer + ffn_per_layer) * m.bytes_per_elem * m.layers
+
+
+def total_kv_bytes_full_model(cfg: FullConfig) -> int:
+ """Total bytes of KV cache for full unsharded model at cfg.topo.s_kv."""
+ m = cfg.model
+ per_layer = 2 * cfg.topo.s_kv * m.h_kv * m.d_head * m.bytes_per_elem
+ return per_layer * m.layers
diff --git a/tests/analytical_visualization/model_config.py b/tests/analytical_visualization/model_config.py
new file mode 100644
index 0000000..3bb7df5
--- /dev/null
+++ b/tests/analytical_visualization/model_config.py
@@ -0,0 +1,310 @@
+"""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)
diff --git a/tests/analytical_visualization/model_presets.py b/tests/analytical_visualization/model_presets.py
new file mode 100644
index 0000000..0214718
--- /dev/null
+++ b/tests/analytical_visualization/model_presets.py
@@ -0,0 +1,126 @@
+"""Model preset library for the analytical visualization tool.
+
+Comprehensive coverage of popular MHA, GQA, MQA, and MLA models across
+sizes. MoE models are approximated as dense with (ffn_dim = activated
+experts × per_expert_ffn) for compute-side estimates; the note field
+flags this. MLA models are approximated as GQA with matching H_kv until
+we add proper MLA support.
+
+Source of numeric params: model cards / config.json from HuggingFace.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import ModelConfig
+
+
+@dataclass
+class Preset:
+ label: str
+ model: ModelConfig
+ family: str = "" # e.g. "Llama 3", "Qwen 3"
+ attn_type: str = "" # "MHA" | "GQA" | "MQA" | "MLA"
+ note: str = ""
+
+
+def _mk(name: str, hidden: int, ffn: int, hq: int, hkv: int,
+ dh: int, layers: int, *, family: str = "", attn: str = "GQA",
+ note: str = "") -> Preset:
+ return Preset(
+ label=name,
+ model=ModelConfig(
+ name=name, hidden=hidden, ffn_dim=ffn,
+ h_q=hq, h_kv=hkv, d_head=dh, layers=layers,
+ ),
+ family=family, attn_type=attn, note=note,
+ )
+
+
+PRESETS: dict[str, Preset] = {
+ # ── MHA (H_q == H_kv) ─────────────────────────────────────────
+ "GPT-2 (117M)": _mk("GPT-2 117M", 768, 3072, 12, 12, 64, 12, family="GPT-2", attn="MHA"),
+ "GPT-2 XL (1.5B)": _mk("GPT-2 XL 1.5B", 1600, 6400, 25, 25, 64, 48, family="GPT-2", attn="MHA"),
+ "Llama 2 7B": _mk("Llama 2 7B", 4096, 11008, 32, 32, 128, 32, family="Llama 2", attn="MHA"),
+ "Llama 2 13B": _mk("Llama 2 13B", 5120, 13824, 40, 40, 128, 40, family="Llama 2", attn="MHA"),
+ "Yi 6B": _mk("Yi 6B", 4096, 11008, 32, 4, 128, 32, family="Yi", attn="GQA"),
+ "Yi 9B": _mk("Yi 9B", 4096, 11008, 32, 4, 128, 48, family="Yi", attn="GQA"),
+ "OPT 30B": _mk("OPT 30B", 7168, 28672, 56, 56, 128, 48, family="OPT", attn="MHA"),
+
+ # ── GQA (H_q > H_kv > 1) ──────────────────────────────────────
+ "Qwen 2 0.5B": _mk("Qwen 2 0.5B", 896, 4864, 14, 2, 64, 24, family="Qwen 2"),
+ "Qwen 2 1.5B": _mk("Qwen 2 1.5B", 1536, 8960, 12, 2, 128, 28, family="Qwen 2"),
+ "Qwen 2 7B": _mk("Qwen 2 7B", 3584, 18944, 28, 4, 128, 28, family="Qwen 2"),
+ "Qwen 2 72B": _mk("Qwen 2 72B", 8192, 29568, 64, 8, 128, 80, family="Qwen 2"),
+
+ "Qwen 3 0.6B": _mk("Qwen 3 0.6B", 1024, 3072, 16, 8, 128, 28, family="Qwen 3"),
+ "Qwen 3 1.7B": _mk("Qwen 3 1.7B", 2048, 6144, 16, 8, 128, 28, family="Qwen 3"),
+ "Qwen 3 4B": _mk("Qwen 3 4B", 2560, 9728, 32, 8, 128, 36, family="Qwen 3"),
+ "Qwen 3 8B": _mk("Qwen 3 8B", 4096, 12288, 32, 8, 128, 36, family="Qwen 3"),
+ "Qwen 3 14B": _mk("Qwen 3 14B", 5120, 17408, 40, 8, 128, 40, family="Qwen 3"),
+ "Qwen 3 32B": _mk("Qwen 3 32B", 5120, 25600, 64, 8, 128, 64, family="Qwen 3"),
+
+ "Llama 3 8B": _mk("Llama 3 8B", 4096, 14336, 32, 8, 128, 32, family="Llama 3"),
+ "Llama 3 70B": _mk("Llama 3 70B", 8192, 28672, 64, 8, 128, 80, family="Llama 3"),
+ "Llama 3.1 8B": _mk("Llama 3.1 8B", 4096, 14336, 32, 8, 128, 32, family="Llama 3.1"),
+ "Llama 3.1 70B": _mk("Llama 3.1 70B", 8192, 28672, 64, 8, 128, 80, family="Llama 3.1"),
+ "Llama 3.1 405B": _mk("Llama 3.1 405B", 16384,53248, 128,8, 128, 126,family="Llama 3.1"),
+ "Llama 3.2 1B": _mk("Llama 3.2 1B", 2048, 8192, 32, 8, 64, 16, family="Llama 3.2"),
+ "Llama 3.2 3B": _mk("Llama 3.2 3B", 3072, 8192, 24, 8, 128, 28, family="Llama 3.2"),
+
+ "Mistral 7B": _mk("Mistral 7B v0.3", 4096, 14336, 32, 8, 128, 32, family="Mistral"),
+ "Mistral Nemo 12B": _mk("Mistral Nemo 12B", 5120, 14336, 32, 8, 128, 40, family="Mistral"),
+ "Mistral Small 22B": _mk("Mistral Small 22B",6144, 16384, 48, 8, 128, 56, family="Mistral"),
+ "Mistral Large 123B": _mk("Mistral Large 123B",12288,28672,96, 8, 128, 88, family="Mistral"),
+
+ "Gemma 2 2B": _mk("Gemma 2 2B", 2304, 9216, 8, 4, 256, 26, family="Gemma 2"),
+ "Gemma 2 9B": _mk("Gemma 2 9B", 3584, 14336, 16, 8, 256, 42, family="Gemma 2"),
+ "Gemma 2 27B": _mk("Gemma 2 27B", 4608, 36864, 32, 16, 128, 46, family="Gemma 2"),
+
+ "Phi 3 mini (3.8B)": _mk("Phi 3 mini 3.8B", 3072, 8192, 32, 32, 96, 32, family="Phi 3", attn="MHA"),
+ "Phi 3 small (7B)": _mk("Phi 3 small 7B", 4096, 14336, 32, 8, 128, 32, family="Phi 3"),
+ "Phi 3 medium (14B)": _mk("Phi 3 medium 14B", 5120, 17920, 40, 10, 128, 40, family="Phi 3"),
+
+ "Yi 34B": _mk("Yi 34B", 7168, 20480, 56, 8, 128, 60, family="Yi"),
+ "Command R+ 104B": _mk("Command R+ 104B", 12288,33792, 96, 8, 128, 64, family="Cohere"),
+
+ # ── MoE (activated-only approximation) ────────────────────────
+ "Mixtral 8x7B (MoE)": _mk("Mixtral 8x7B", 4096, 14336*2, 32, 8, 128, 32,
+ family="Mistral MoE",
+ note="MoE: 8 experts x 2 activated. FFN ~ 2 x per_expert."),
+ "Mixtral 8x22B (MoE)":_mk("Mixtral 8x22B", 6144, 16384*2, 48, 8, 128, 56,
+ family="Mistral MoE",
+ note="MoE: 8 experts x 2 activated."),
+ "Qwen 3 30B (MoE)": _mk("Qwen 3 30B (MoE)", 2048, 768*8, 32, 4, 128, 48,
+ family="Qwen 3 MoE",
+ note="128 experts x 8 activated per token."),
+ "Qwen 3 235B (MoE)": _mk("Qwen 3 235B (MoE)",4096, 1536*8, 64, 4, 128, 94,
+ family="Qwen 3 MoE",
+ note="128 experts x 8 activated. Dense-approx overstates weight."),
+ "DeepSeek V2 (dense-approx)":
+ _mk("DeepSeek V2", 5120, 12288, 128, 128, 128, 60,
+ family="DeepSeek", attn="MLA",
+ note="MLA compresses KV to ~576 B/token; d_c=512 latent. Dense-approx."),
+ "DeepSeek V3 (dense-approx)":
+ _mk("DeepSeek V3", 7168, 16384, 128, 128, 128, 61,
+ family="DeepSeek", attn="MLA",
+ note="MoE 256 x 8 activated + MLA. Both approximations."),
+ "Grok-1 314B (MoE)": _mk("Grok-1 314B", 6144, 32768*2, 48, 8, 128, 64,
+ family="xAI",
+ note="MoE 8 experts x 2 activated."),
+
+ # ── MQA (H_kv == 1) ───────────────────────────────────────────
+ "Falcon 7B (MQA)": _mk("Falcon 7B", 4544, 18176, 71, 1, 64, 32, family="Falcon", attn="MQA"),
+ "Falcon 40B (MQA)": _mk("Falcon 40B", 8192, 32768, 128,1, 64, 60, family="Falcon", attn="MQA"),
+
+ # ── Custom slot ──────────────────────────────────────────────
+ "Custom": Preset(label="Custom", model=ModelConfig(name="Custom")),
+}
+
+
+def preset_names_by_family() -> dict[str, list[str]]:
+ """Return preset names grouped by family (for a nested dropdown)."""
+ grouped: dict[str, list[str]] = {}
+ for name, p in PRESETS.items():
+ grouped.setdefault(p.family or "Other", []).append(name)
+ return grouped
diff --git a/tests/analytical_visualization/optimization_report.py b/tests/analytical_visualization/optimization_report.py
new file mode 100644
index 0000000..90fdf4a
--- /dev/null
+++ b/tests/analytical_visualization/optimization_report.py
@@ -0,0 +1,235 @@
+"""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
diff --git a/tests/analytical_visualization/pe_weight_layout.py b/tests/analytical_visualization/pe_weight_layout.py
new file mode 100644
index 0000000..f3d1d16
--- /dev/null
+++ b/tests/analytical_visualization/pe_weight_layout.py
@@ -0,0 +1,221 @@
+"""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
diff --git a/tests/analytical_visualization/pipeline_diagram.py b/tests/analytical_visualization/pipeline_diagram.py
new file mode 100644
index 0000000..4484ec1
--- /dev/null
+++ b/tests/analytical_visualization/pipeline_diagram.py
@@ -0,0 +1,300 @@
+"""Pipeline diagram: attention + FFN steps, colored by bound type.
+
+Each stage is drawn as a rounded box in a horizontal flow:
+ Attention: S1 -> S2 -> S3 -> ... -> S10 -> C2 (AllReduce)
+ FFN: F1 -> F2 -> F3 -> F4 -> F5 -> CF1 (AllReduce)
+
+CP ring communication (C1) and score AllReduce (C3) are drawn as
+"overhead" boxes attached above the QKt/PV region (they run
+concurrently with per-hop compute in the ring).
+
+Box fill color = the stage's bound classification:
+ compute-bound blue
+ memory-bound orange
+ comm-bound red
+ trivial grey
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+from .stage_latencies import (
+ all_stages, all_ffn_stages, StageCost,
+)
+
+
+BOUND_COLOR = {
+ "compute": "#4682b4", # blue
+ "memory": "#e37400", # orange
+ "comm": "#c62828", # red
+ "trivial": "#a9a9a9", # grey
+}
+BOUND_LABEL = {
+ "compute": "compute-bound",
+ "memory": "memory-bound",
+ "comm": "communication",
+ "trivial": "trivial / negligible",
+}
+
+
+def _fmt_us(sec: float) -> str:
+ us = sec * 1e6
+ if us < 1:
+ return f"{us*1e3:.1f}ns"
+ if us < 100:
+ return f"{us:.2f}us"
+ return f"{us:.0f}us"
+
+
+def _draw_block(ax, x, y, w, h, label, cost: StageCost, edge="#333"):
+ color = BOUND_COLOR.get(cost.bound, "#a9a9a9")
+ box = patches.FancyBboxPatch(
+ (x, y), w, h,
+ boxstyle="round,pad=0.05",
+ facecolor=color, edgecolor=edge, linewidth=1.2, alpha=0.85,
+ )
+ ax.add_patch(box)
+ # Two lines: stage name (bold) + visible latency (small)
+ ax.text(x + w / 2, y + h * 0.62, label,
+ ha="center", va="center",
+ fontsize=8.0, fontweight="bold", color="white")
+ ax.text(x + w / 2, y + h * 0.25, _fmt_us(cost.visible_s),
+ ha="center", va="center",
+ fontsize=7.0, color="white")
+
+
+_SHORT_NAMES = {
+ "S1": "RMSNorm",
+ "S2": "W_Q",
+ "S3": "W_K+W_V",
+ "S4": "KV append",
+ "S5": "Q.K^T",
+ "S6": "softmax",
+ "S7": "P.V",
+ "S8": "merge",
+ "S9": "norm O/l",
+ "S10": "W_O",
+ "C1": "CP ring",
+ "C2": "TP AR",
+ "C3": "Score AR",
+ "F1": "RMSNorm",
+ "F2": "W_gate",
+ "F3": "W_up",
+ "F4": "SwiGLU",
+ "F5": "W_down",
+ "CF1": "FFN AR",
+}
+
+
+def _short_name(name: str) -> str:
+ """Map 'S1 RMSNorm' -> 'RMSNorm'; fallback to first token."""
+ tok = name.split()[0]
+ return _SHORT_NAMES.get(tok, tok)
+
+
+def _draw_arrow(ax, x0, x1, y):
+ ax.annotate("", xy=(x1, y), xytext=(x0, y),
+ arrowprops=dict(arrowstyle="->", color="#333", lw=1.1))
+
+
+def _draw_ring_loop(ax, x_left, x_right, y_top, label: str, n_iters: int,
+ arc_lift: float = 0.9):
+ """Draw a curved back-arrow above [x_left..x_right] with a repeat label.
+
+ Visualizes 'repeat this block n_iters times' (like a loop in a diagram).
+ arc_lift: how high above y_top the arc sits.
+ """
+ color = "#5a189a"
+ lw = 1.6
+ arc_y = y_top + arc_lift
+ # Vertical tick up from the right edge
+ ax.plot([x_right, x_right], [y_top + 0.02, arc_y],
+ color=color, linewidth=lw)
+ # Horizontal top of the loop
+ ax.plot([x_right, x_left], [arc_y, arc_y],
+ color=color, linewidth=lw)
+ # Down-arrow into the left edge (a back-edge indicating the repeat)
+ ax.annotate("", xy=(x_left, y_top + 0.02), xytext=(x_left, arc_y),
+ arrowprops=dict(arrowstyle="->", color=color, lw=lw))
+ # Loop label centered above the arc
+ ax.text((x_left + x_right) / 2, arc_y + 0.08,
+ label,
+ ha="center", va="bottom",
+ fontsize=9, fontweight="bold", color=color)
+ # Iteration count as a badge on the right
+ ax.text(x_right + 0.08, arc_y - 0.03,
+ f"x {n_iters}",
+ ha="left", va="top",
+ fontsize=10, fontweight="bold", color=color,
+ bbox=dict(boxstyle="round,pad=0.15", facecolor="white",
+ edgecolor=color, linewidth=1.2, alpha=0.95))
+
+
+def _draw_row(ax, y_base, row_title, stages: list[StageCost],
+ x_start: float, box_w: float, box_h: float, gap: float):
+ """Draw one row of stage boxes with arrows between them."""
+ ax.text(x_start - 0.4, y_base + box_h / 2, row_title,
+ ha="right", va="center",
+ fontsize=11, fontweight="bold", color="#212529")
+ x = x_start
+ for i, st in enumerate(stages):
+ _draw_block(ax, x, y_base, box_w, box_h,
+ _short_name(st.name), st)
+ if i < len(stages) - 1:
+ _draw_arrow(ax, x + box_w + 0.02, x + box_w + gap - 0.02,
+ y_base + box_h / 2)
+ x += box_w + gap
+ return x - gap
+
+
+def _draw_overlap_boxes(ax, hidden_stages, x_left, x_right, y_top,
+ gap_between: float = 0.15):
+ """Draw dashed 'overlapped comm' boxes spanning x_left..x_right above y_top.
+
+ Boxes are widened to cover the horizontal range of the stages they
+ overlap with (e.g. C1 CP ring sits above S5-S7).
+ """
+ if not hidden_stages:
+ return
+ hy = y_top + 0.10
+ n = len(hidden_stages)
+ total_span = x_right - x_left
+ box_span = (total_span - (n - 1) * gap_between) / max(1, n)
+ box_h_local = 0.42
+ hx = x_left
+ for st in hidden_stages:
+ box = patches.FancyBboxPatch(
+ (hx, hy), box_span, box_h_local,
+ boxstyle="round,pad=0.03",
+ facecolor=BOUND_COLOR["comm"], edgecolor="#333",
+ linewidth=1.0, alpha=0.35, linestyle="--",
+ )
+ ax.add_patch(box)
+ ax.text(hx + box_span / 2, hy + box_h_local * 0.65,
+ f"{_short_name(st.name)} (concurrent)",
+ ha="center", va="center",
+ fontsize=7, color="#7a0e0e", fontweight="bold")
+ ax.text(hx + box_span / 2, hy + box_h_local * 0.2,
+ _fmt_us(st.visible_s),
+ ha="center", va="center",
+ fontsize=6.5, color="#7a0e0e")
+ hx += box_span + gap_between
+
+
+def draw_pipeline(cfg: FullConfig, ax=None):
+ """Draw attention + FFN pipeline. Boxes colored by bound type."""
+ attn = all_stages(cfg)
+ ffn = all_ffn_stages(cfg)
+
+ # Split attention: main stages, then C2 (TP AllReduce). C1/C3 = overlapped.
+ attn_main = [s for s in attn if not s.name.startswith("C")]
+ attn_c2 = next((s for s in attn if s.name.startswith("C2")), None)
+ attn_hidden = [s for s in attn
+ if s.name.startswith("C1") or s.name.startswith("C3")]
+ attn_hidden = [s for s in attn_hidden if s.visible_s > 0]
+ attn_row = attn_main + ([attn_c2] if attn_c2 and attn_c2.visible_s > 0 else [])
+
+ # FFN row: all F* + CF1 if non-trivial
+ ffn_main = [s for s in ffn if not s.name.startswith("CF")]
+ ffn_cf = next((s for s in ffn if s.name.startswith("CF1")), None)
+ ffn_row = ffn_main + ([ffn_cf] if ffn_cf and ffn_cf.visible_s > 0 else [])
+
+ n_max = max(len(attn_row), len(ffn_row))
+ box_w = 1.35
+ box_h = 0.78
+ gap = 0.18
+ x_start = 1.1
+
+ total_w = x_start + n_max * box_w + (n_max - 1) * gap + 0.5
+ total_h = 4.2
+
+ if ax is None:
+ fig, ax = plt.subplots(figsize=(max(11, total_w * 1.1),
+ max(4.5, total_h * 1.0)))
+ else:
+ fig = ax.figure
+
+ _draw_row(ax, y_base=2.6, row_title="Attention",
+ stages=attn_row, x_start=x_start,
+ box_w=box_w, box_h=box_h, gap=gap)
+
+ _draw_row(ax, y_base=0.6, row_title="FFN",
+ stages=ffn_row, x_start=x_start,
+ box_w=box_w, box_h=box_h, gap=gap)
+
+ # Ring loop indicator + concurrent-comm boxes anchored over S5..S8.
+ _idx_by_prefix = {s.name.split()[0]: i for i, s in enumerate(attn_row)}
+ _s5_idx = _idx_by_prefix.get("S5")
+ _s7_idx = _idx_by_prefix.get("S7")
+ _s8_idx = _idx_by_prefix.get("S8")
+
+ if cfg.topo.cp > 1 and cfg.topo.mode == "prefill" and _s5_idx is not None:
+ _right_idx = _s8_idx if _s8_idx is not None else _s7_idx
+ _lx = x_start + _s5_idx * (box_w + gap)
+ _rx = x_start + _right_idx * (box_w + gap) + box_w
+ _y_top = 2.6 + box_h
+ # Concurrent-comm boxes (C1 CP ring, C3 score AR) directly above S5-S7
+ _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
+ # Ring loop arc above those
+ _draw_ring_loop(
+ ax, _lx, _rx, _y_top,
+ label="ring attention loop (K/V ring per hop; S8 merges partial O,m,l)",
+ n_iters=cfg.topo.cp,
+ arc_lift=1.05,
+ )
+ elif cfg.topo.cp > 1 and cfg.topo.mode == "decode" and _s5_idx is not None:
+ _right_idx = _s7_idx if _s7_idx is not None else _s5_idx
+ _lx = x_start + _s5_idx * (box_w + gap)
+ _rx = x_start + _right_idx * (box_w + gap) + box_w
+ _y_top = 2.6 + box_h
+ # Any straggling concurrent comm (e.g. C3 if head-split active) above S5-S7
+ _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
+ _draw_ring_loop(
+ ax, _lx, _rx, _y_top,
+ label="decode: local pass only; S8 fires O/m/l all-reduce",
+ n_iters=1,
+ arc_lift=1.05 if attn_hidden else 0.85,
+ )
+ elif attn_hidden:
+ # Fallback: CP=1 but still have C3 or similar; anchor above S5-S7 if present
+ if _s5_idx is not None and _s7_idx is not None:
+ _lx = x_start + _s5_idx * (box_w + gap)
+ _rx = x_start + _s7_idx * (box_w + gap) + box_w
+ _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, 2.6 + box_h)
+
+ # Totals
+ attn_total = sum(s.visible_s for s in attn_row) \
+ + sum(s.visible_s for s in attn_hidden)
+ ffn_total = sum(s.visible_s for s in ffn_row)
+ ax.text(total_w - 0.3, 2.6 + box_h / 2 + 1.15,
+ f"Attn total: {_fmt_us(attn_total)}",
+ ha="right", va="bottom", fontsize=9,
+ fontweight="bold", color="#212529")
+ ax.text(total_w - 0.3, 0.6 + box_h + 0.15,
+ f"FFN total: {_fmt_us(ffn_total)}",
+ ha="right", va="bottom", fontsize=9,
+ fontweight="bold", color="#212529")
+
+ # Legend
+ handles = []
+ for k, lbl in BOUND_LABEL.items():
+ handles.append(patches.Patch(
+ facecolor=BOUND_COLOR[k], edgecolor="#333",
+ alpha=0.85, label=lbl,
+ ))
+ ax.legend(handles=handles, loc="upper center",
+ bbox_to_anchor=(0.5, -0.02),
+ ncol=4, fontsize=8, frameon=False)
+
+ ax.set_xlim(0, total_w + 0.6)
+ ax.set_ylim(0, 4.9)
+ ax.set_aspect("auto")
+ ax.axis("off")
+ ax.set_title(
+ f"Per-layer pipeline (one PE, {cfg.topo.mode}, T_q={cfg.topo.T_q}, "
+ f"S_local={cfg.topo.s_local}) - color = dominant bound",
+ fontsize=10, fontweight="bold",
+ )
+ return fig
diff --git a/tests/analytical_visualization/stage_latencies.py b/tests/analytical_visualization/stage_latencies.py
new file mode 100644
index 0000000..488e0bd
--- /dev/null
+++ b/tests/analytical_visualization/stage_latencies.py
@@ -0,0 +1,767 @@
+"""Per-stage cost formulas for attention forward pass.
+
+Each stage returns:
+ name, formula (str with the substituted numeric form),
+ compute_time, memory_time, comm_time, bound, visible_time.
+
+Bound is the max of the three; "hidden" comm can appear as 0 when
+overlapped with compute.
+
+All times are in seconds; convert to μs at display.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+from .model_config import FullConfig
+
+
+Bound = Literal["compute", "memory", "comm", "trivial"]
+
+
+@dataclass
+class StageCost:
+ name: str
+ formula: str
+ compute_s: float
+ memory_s: float
+ comm_s: float
+ bound: Bound
+ visible_s: float
+ hop_multiplier: int = 1 # some stages scale with N_cp hops
+ # Detailed breakdown for the per-stage table.
+ flops: int = 0 # total FLOPs (numeric)
+ mem_bytes: int = 0 # HBM bytes moved (numeric)
+ comm_bytes: int = 0 # comm bytes over the wire (numeric)
+ flops_formula: str = "" # human-readable FLOPs formula
+ mem_formula: str = "" # human-readable memory formula
+ comm_formula: str = "" # human-readable comm formula
+
+
+def _visible(compute: float, memory: float, comm: float) -> tuple[float, Bound]:
+ """Pick the dominant time as the visible stage cost."""
+ parts = {"compute": compute, "memory": memory, "comm": comm}
+ dominant = max(parts, key=parts.get) # type: ignore
+ return parts[dominant], dominant # type: ignore
+
+
+def stage_rmsnorm(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ bytes_ = T_q * d * b * 2 # load x + load weight
+ flops = 4 * T_q * d
+ mem_s = bytes_ / cfg.machine.bw_hbm
+ cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+ vis, bnd = _visible(cmp_s, mem_s, 0)
+ return StageCost(
+ name="S1 RMSNorm",
+ formula=f"bytes = {T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
+ compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+ bound=bnd, visible_s=vis,
+ flops=flops, mem_bytes=bytes_,
+ flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
+ mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
+ )
+
+
+def _gemm_time(flops: int, weight_bytes: int, cfg: FullConfig) -> tuple[float, float]:
+ cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+ mem_s = weight_bytes / cfg.machine.bw_hbm
+ return cmp_s, mem_s
+
+
+def stage_wq(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ hq_per_pe = cfg.h_q_per_pe
+ dh = cfg.model.d_head
+ flops = 2 * T_q * d * (hq_per_pe * dh)
+ weight_B = d * (hq_per_pe * dh) * b
+ cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+ vis, bnd = _visible(cmp_s, mem_s, 0)
+ return StageCost(
+ name="S2 W_Q GEMM",
+ formula=f"FLOPs = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
+ f"weight = {weight_B/1e6:.1f} MB",
+ compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+ bound=bnd, visible_s=vis,
+ flops=int(flops), mem_bytes=int(weight_B),
+ flops_formula=f"2*T_q*d*(H_q/TP*d_h) = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
+ mem_formula=f"d*(H_q/TP*d_h)*b = {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
+ )
+
+
+def stage_wkv(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
+ dh = cfg.model.d_head
+ flops_one = 2 * T_q * d * (hkv_per_pe * dh)
+ weight_B_one = d * (hkv_per_pe * dh) * b
+ flops = 2 * flops_one
+ weight_B = 2 * weight_B_one
+ cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+ vis, bnd = _visible(cmp_s, mem_s, 0)
+ return StageCost(
+ name="S3 W_K + W_V GEMM",
+ formula=f"FLOPs per proj = 2*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
+ compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+ bound=bnd, visible_s=vis,
+ flops=int(flops), mem_bytes=int(weight_B),
+ flops_formula=f"2*(2*T_q*d*(H_kv/TP*d_h)) = 2*(2*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
+ mem_formula=f"2*(d*(H_kv/TP*d_h)*b) = 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
+ )
+
+
+def stage_kv_append(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ b = cfg.model.bytes_per_elem
+ hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
+ dh = cfg.model.d_head
+ bytes_ = 2 * T_q * hkv_per_pe * dh * b
+ mem_s = bytes_ / cfg.machine.bw_hbm
+ return StageCost(
+ name="S4 KV cache append",
+ formula=f"bytes = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+ compute_s=0, memory_s=mem_s, comm_s=0,
+ bound="memory", visible_s=mem_s,
+ flops=0, mem_bytes=bytes_,
+ flops_formula="0 (no matmul, just cache write)",
+ mem_formula=f"2*T_q*(H_kv/TP)*d_h*b = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+ )
+
+
+def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
+ """Q·Kᵀ (and P·V has same FLOPs) per hop, per PE."""
+ T_q = cfg.topo.T_q
+ S_local = cfg.topo.s_local
+ dh = cfg.model.d_head
+ hq_per_pe = cfg.h_q_per_pe
+ flops = 2 * T_q * S_local * dh * hq_per_pe
+ cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+ formula = f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops:.2g} FLOPs/hop"
+ return cmp_s, formula
+
+
+def _cp_compute_passes(cfg: FullConfig) -> int:
+ """How many local S5/S6/S7 passes happen per token step under CP.
+ - decode: 1 (each PE computes once against its local K,V, then all-reduce O/m/l)
+ - prefill: CP (K/V or Q rotates through the ring; one local pass per hop)
+ """
+ if cfg.topo.cp <= 1:
+ return 1
+ return 1 if cfg.topo.mode == "decode" else cfg.topo.cp
+
+
+def stage_qkT(cfg: FullConfig) -> StageCost:
+ cmp_hop, _ = _per_hop_qkT_pv(cfg)
+ passes = _cp_compute_passes(cfg)
+ total = cmp_hop * passes
+ T_q = cfg.topo.T_q
+ S_local = cfg.topo.s_local
+ dh = cfg.model.d_head
+ hq_per_pe = cfg.h_q_per_pe
+ flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
+ total_flops = flops_per_hop * passes
+ _hop_word = "pass" if passes == 1 else "hops"
+ return StageCost(
+ name=f"S5 Q.K^T (x{passes} {_hop_word})",
+ formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
+ compute_s=total, memory_s=0, comm_s=0,
+ bound="compute", visible_s=total,
+ hop_multiplier=passes,
+ flops=int(total_flops), mem_bytes=0,
+ flops_formula=(
+ f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
+ f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
+ ),
+ mem_formula="0 (scores accumulated on-chip)",
+ )
+
+
+def stage_softmax(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ S_local = cfg.topo.s_local
+ hq_per_pe = cfg.h_q_per_pe
+ b = cfg.model.bytes_per_elem
+ elems = hq_per_pe * T_q * S_local
+ bytes_ = elems * b * 2
+ mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
+ passes = _cp_compute_passes(cfg)
+ total = mem_s_per_hop * passes
+ total_bytes = bytes_ * passes
+ _hop_word = "pass" if passes == 1 else "hops"
+ return StageCost(
+ name=f"S6 softmax (x{passes} {_hop_word})",
+ formula=f"elts/hop = {hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
+ compute_s=0, memory_s=total, comm_s=0,
+ bound="memory", visible_s=total,
+ hop_multiplier=passes,
+ flops=0, mem_bytes=int(total_bytes),
+ flops_formula="~O(elts) (negligible)",
+ mem_formula=(
+ f"{passes}*2*b*(H_q/TP)*T_q*S_local = "
+ f"{passes}*2*{b}*{hq_per_pe}*{T_q}*{S_local} = "
+ f"{total_bytes/1e6:.2f} MB"
+ ),
+ )
+
+
+def stage_pv(cfg: FullConfig) -> StageCost:
+ cmp_hop, _ = _per_hop_qkT_pv(cfg)
+ passes = _cp_compute_passes(cfg)
+ total = cmp_hop * passes
+ T_q = cfg.topo.T_q
+ S_local = cfg.topo.s_local
+ dh = cfg.model.d_head
+ hq_per_pe = cfg.h_q_per_pe
+ flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
+ total_flops = flops_per_hop * passes
+ _hop_word = "pass" if passes == 1 else "hops"
+ return StageCost(
+ name=f"S7 P.V (x{passes} {_hop_word})",
+ formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
+ compute_s=total, memory_s=0, comm_s=0,
+ bound="compute", visible_s=total,
+ hop_multiplier=passes,
+ flops=int(total_flops), mem_bytes=0,
+ flops_formula=(
+ f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
+ f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
+ ),
+ mem_formula="0 (accumulated on-chip)",
+ )
+
+
+def stage_merge(cfg: FullConfig) -> StageCost:
+ """S8: online-softmax merge of partial (o, m, l).
+ - prefill (K/V ring): merge happens in-place across CP hops; no comm here.
+ - decode: merge is the moment we all-reduce partial (o, m, l) across CP
+ ranks; comm is folded into this stage (no separate C1).
+ """
+ T_q = cfg.topo.T_q
+ hq_per_pe = cfg.h_q_per_pe
+ dh = cfg.model.d_head
+ b = cfg.model.bytes_per_elem
+ flops = 6 * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
+ cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+
+ comm_s = 0.0
+ comm_bytes = 0
+ comm_formula = ""
+ name_suffix = ""
+
+ if cfg.topo.mode == "decode" and cfg.topo.cp > 1:
+ cp = cfg.topo.cp
+ # (O + m + l) bytes per rank
+ M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+ if cfg.topo.cp_placement == "pe":
+ bw, alpha, tier = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
+ elif cfg.topo.sips_used > 1:
+ bw, alpha, tier = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "inter-SIP"
+ else:
+ bw, alpha, tier = cfg.machine.bw_inter, cfg.machine.alpha_inter, "inter-cube"
+ ar_bytes = 2 * (cp - 1) / cp * M
+ comm_s = ar_bytes / bw + 2 * (cp - 1) * alpha
+ comm_bytes = int(ar_bytes)
+ comm_formula = (
+ f"PURPOSE (folded into S8 in decode): each CP rank computed\n"
+ f"attention against its LOCAL slice of KV and holds a partial\n"
+ f"(O, m, l). This all-reduce merges those partials across all\n"
+ f"{cp} CP ranks (online-softmax combine). No separate C1 row\n"
+ f"in decode because this is the only CP comm and it happens\n"
+ f"exactly here, at the end of the local attention body.\n"
+ f"---\n"
+ f"AR of (O + m + l): 2*(CP-1)/CP * M "
+ f"= 2*({cp}-1)/{cp} * {M} B "
+ f"= {ar_bytes:.0f} B over {tier} ({bw/1e9:.0f} GB/s) "
+ f"+ 2*({cp}-1)*alpha"
+ )
+ name_suffix = f" + O/m/l AR (CP={cp} ranks, {tier})"
+
+ visible_s = max(cmp_s, comm_s)
+ if comm_s > cmp_s:
+ bound = "comm"
+ elif cmp_s > 0:
+ bound = "compute"
+ else:
+ bound = "trivial"
+
+ return StageCost(
+ name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}",
+ formula=f"~6*{T_q}*{hq_per_pe}*{dh}*(C-1) = {flops:.2g} FLOPs",
+ compute_s=cmp_s, memory_s=0, comm_s=comm_s,
+ bound=bound, visible_s=visible_s,
+ flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes,
+ flops_formula=(
+ f"6*T_q*(H_q/TP)*d_h*(CP-1) = "
+ f"6*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
+ ),
+ mem_formula="0 (in-register)",
+ comm_formula=comm_formula or "0",
+ )
+
+
+def stage_normalize(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ hq_per_pe = cfg.h_q_per_pe
+ dh = cfg.model.d_head
+ flops = T_q * hq_per_pe * dh
+ cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+ return StageCost(
+ name="S9 normalize O/l",
+ formula=f"{T_q}*{hq_per_pe}*{dh} = {flops} divisions",
+ compute_s=cmp_s, memory_s=0, comm_s=0,
+ bound="trivial", visible_s=cmp_s,
+ flops=int(flops), mem_bytes=0,
+ flops_formula=f"T_q*(H_q/TP)*d_h = {T_q}*{hq_per_pe}*{dh} = {flops}",
+ mem_formula="0",
+ )
+
+
+def stage_wo(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ hq_per_pe = cfg.h_q_per_pe
+ dh = cfg.model.d_head
+ flops = 2 * T_q * (hq_per_pe * dh) * d
+ weight_B = (hq_per_pe * dh) * d * b
+ cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+ vis, bnd = _visible(cmp_s, mem_s, 0)
+ return StageCost(
+ name="S10 W_O GEMM",
+ formula=f"FLOPs = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
+ f"weight = {weight_B/1e6:.1f} MB",
+ compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+ bound=bnd, visible_s=vis,
+ flops=int(flops), mem_bytes=int(weight_B),
+ flops_formula=f"2*T_q*(H_q/TP*d_h)*d = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
+ mem_formula=f"(H_q/TP*d_h)*d*b = {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
+ )
+
+
+def comm_cp_ring(cfg: FullConfig) -> StageCost:
+ """CP comm — depends on mode:
+
+ - decode: single O/m/l all-reduce over CP ranks AFTER local S5-S8 finish.
+ No per-hop ring; each PE computes attention against its local K,V once
+ and then contributes (o, m, l) to the reduce.
+ - prefill: per-hop ring during S5/S7 compute (either K/V or Q+O/m/l per
+ cp_ring_variant).
+
+ BW tier depends on cp_placement:
+ - cp_placement=pe: intra-cube BW for all hops
+ - cp_placement=cube: inter-cube BW; long rings cross SIP boundaries
+ """
+ if cfg.topo.cp <= 1:
+ return StageCost(
+ name="C1 CP comm", formula="C=1 -> no comm",
+ compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+ )
+ S_local = cfg.topo.s_local
+ hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
+ hq_per_pe = cfg.h_q_per_pe
+ T_q = cfg.topo.T_q
+ dh = cfg.model.d_head
+ b = cfg.model.bytes_per_elem
+
+ # ── Decode: single O/m/l all-reduce after local compute ──
+ if cfg.topo.mode == "decode":
+ cp = cfg.topo.cp
+ # per-rank O + m + l bytes
+ M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+ if cfg.topo.cp_placement == "pe":
+ bw = cfg.machine.bw_intra
+ alpha = cfg.machine.alpha_intra
+ tier = "intra-cube"
+ elif cfg.topo.sips_used > 1:
+ bw = cfg.machine.bw_intersip
+ alpha = cfg.machine.alpha_intersip
+ tier = "inter-SIP"
+ else:
+ bw = cfg.machine.bw_inter
+ alpha = cfg.machine.alpha_inter
+ tier = "inter-cube"
+ ar_bytes = 2 * (cp - 1) / cp * M
+ ar_time = ar_bytes / bw + 2 * (cp - 1) * alpha
+ return StageCost(
+ name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)",
+ formula=(f"decode: gather partial (O,m,l) once at end; "
+ f"M = T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
+ f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"),
+ compute_s=0, memory_s=0, comm_s=ar_time,
+ bound="comm", visible_s=ar_time,
+ flops=0, mem_bytes=0, comm_bytes=int(ar_bytes),
+ flops_formula="0",
+ mem_formula="0",
+ comm_formula=(
+ f"PURPOSE: In decode, each CP rank computed attention against\n"
+ f"its OWN slice of the KV cache. Each rank now holds a "
+ f"partial\n"
+ f"(O, m, l). This single all-reduce merges those partials "
+ f"across\n"
+ f"all {cp} CP ranks (using online-softmax math) to get the "
+ f"final O.\n"
+ f"---\n"
+ f"AR of (O + m + l): 2*(CP-1)/CP * M "
+ f"= 2*({cp}-1)/{cp} * {M} "
+ f"= {ar_bytes:.0f} B over {tier} at {bw/1e9:.0f} GB/s "
+ f"+ 2*({cp}-1)*alpha"
+ ),
+ )
+
+ # ── Prefill: per-hop ring (K/V or Q+O/m/l) concurrent with S5-S7 ──
+ if cfg.topo.cp_ring_variant == "qoml":
+ M_KV = (2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+ variant_desc = "Q+O/m/l ring"
+ formula_bytes = (
+ f"2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b "
+ f"= 2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b} "
+ f"= {M_KV} B/hop"
+ )
+ else:
+ M_KV = 2 * S_local * hkv_per_pe * dh * b
+ variant_desc = "K/V ring"
+ formula_bytes = (
+ f"2*S_local*(H_kv/TP)*d_h*b "
+ f"= 2*{S_local}*{hkv_per_pe}*{dh}*{b} "
+ f"= {M_KV/1e6:.3f} MB/hop"
+ )
+
+ if cfg.topo.cp_placement == "pe":
+ # All hops intra-cube (fastest).
+ intra_hops = cfg.topo.cp - 1
+ inter_hops = 0
+ intra_time = (intra_hops * M_KV / cfg.machine.bw_intra
+ + intra_hops * cfg.machine.alpha_intra)
+ inter_time = 0.0
+ else: # cp_placement == "cube"
+ intra_hops = cfg.topo.cp_intra_sip_hops
+ inter_hops = cfg.topo.cp_inter_sip_hops
+ intra_time = (intra_hops * M_KV / cfg.machine.bw_inter
+ + intra_hops * cfg.machine.alpha_inter)
+ inter_time = (inter_hops * M_KV / cfg.machine.bw_intersip
+ + inter_hops * cfg.machine.alpha_intersip)
+ comm_time = intra_time + inter_time
+
+ # Overlap check: per-hop compute (S5+S6+S7) hides intra-SIP hops well;
+ # inter-SIP hops usually dominate.
+ per_hop_cmp, _ = _per_hop_qkT_pv(cfg)
+ per_hop_mem_softmax = (cfg.h_q_per_pe * cfg.topo.T_q
+ * S_local * cfg.model.bytes_per_elem * 2
+ / cfg.machine.bw_hbm)
+ per_hop_compute = 2 * per_hop_cmp + per_hop_mem_softmax
+ per_intra_ring = M_KV / cfg.machine.bw_inter + cfg.machine.alpha_inter
+ per_inter_ring = M_KV / cfg.machine.bw_intersip + cfg.machine.alpha_intersip
+ visible_intra = intra_hops * max(0.0, per_intra_ring - per_hop_compute)
+ visible_inter = inter_hops * max(0.0, per_inter_ring - per_hop_compute)
+ visible_total = visible_intra + visible_inter
+
+ tier_desc = f"{intra_hops} intra-SIP + {inter_hops} inter-SIP hops"
+ total_ring_bytes = M_KV * (intra_hops + inter_hops)
+ _var_purpose = (
+ "K, V shards rotate between CP ranks each hop"
+ if cfg.topo.cp_ring_variant == "kv"
+ else "Q + running (O, m, l) rotate between CP ranks each hop"
+ )
+ return StageCost(
+ name=f"C1 CP {variant_desc} ({tier_desc})",
+ formula=f"{formula_bytes}; "
+ f"intra: {intra_hops}*M/BW + inter: {inter_hops}*M/BW",
+ compute_s=0, memory_s=0, comm_s=comm_time,
+ bound="comm", visible_s=visible_total,
+ flops=0, mem_bytes=0, comm_bytes=int(total_ring_bytes),
+ flops_formula="0",
+ mem_formula="0",
+ comm_formula=(
+ f"PURPOSE: CP shards the sequence axis. Each CP rank holds only\n"
+ f"1/{cfg.topo.cp} of the KV cache, so to compute full attention\n"
+ f"we must move data between CP ranks. Variant: {_var_purpose}.\n"
+ f"Runs concurrently with S5/S6/S7 - each hop, compute of the\n"
+ f"just-arrived shard overlaps with comm of the next shard.\n"
+ f"---\n"
+ f"M*(CP-1) with M = {formula_bytes}; "
+ f"total = {total_ring_bytes/1e6:.3f} MB over the ring"
+ ),
+ )
+
+
+def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
+ """TP AllReduce on W_O output."""
+ if cfg.topo.tp <= 1:
+ return StageCost(
+ name="C2 TP AllReduce W_O", formula="TP=1 -> no AR",
+ compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+ )
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ bytes_ = T_q * d * b
+ tp = cfg.topo.tp
+ tier = cfg.topo.tp_link_tier() # "intra" | "inter" | "intersip"
+ if tier == "intra":
+ bw, alpha, scope = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
+ elif tier == "inter":
+ bw, alpha, scope = cfg.machine.bw_inter, cfg.machine.alpha_inter, "cross-cube"
+ else:
+ bw, alpha, scope = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "cross-SIP"
+ comm_time = 2 * (tp - 1) / tp * bytes_ / bw + 2 * (tp - 1) * alpha
+ total_comm_bytes = int(2 * (tp - 1) / tp * bytes_)
+ return StageCost(
+ name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})",
+ formula=f"2*(TP-1)/TP * {T_q}*{d}*{b} B / BW + 2(TP-1)*alpha [{scope}]",
+ compute_s=0, memory_s=0, comm_s=comm_time,
+ bound="comm", visible_s=comm_time,
+ flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
+ flops_formula="0",
+ mem_formula="0",
+ comm_formula=(
+ f"PURPOSE: W_O is row-parallel across TP ranks (each rank holds\n"
+ f"1/{tp} of W_O's input rows). After the local W_O GEMM, each of\n"
+ f"the {tp} TP ranks holds a PARTIAL hidden vector (sum over its\n"
+ f"own Q-head slice only). This all-reduce sums those partials so\n"
+ f"every rank ends up with the full hidden vector for the next\n"
+ f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n"
+ f"---\n"
+ f"2*(TP-1)/TP * T_q*d*b = 2*({tp}-1)/{tp} * {T_q}*{d}*{b} "
+ f"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
+ f"{bw/1e9:.0f} GB/s"
+ ),
+ )
+
+
+def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
+ """Extra AllReduce on attention scores when TP > H_kv with head-dim split.
+
+ Ranks sharing a KV head have partial scores; must AllReduce across the
+ split_factor group to get final scores per hop.
+ """
+ if not cfg.kv_replication_needed or cfg.topo.kv_shard_mode != "split":
+ return StageCost(
+ name="C3 Score AllReduce (head-split)",
+ formula="TP <= H_kv or replicate mode: not needed",
+ compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+ )
+ split = cfg.head_dim_split_factor # ranks sharing one head
+ T_q = cfg.topo.T_q
+ S_local = cfg.topo.s_local
+ hq_per_pe = cfg.h_q_per_pe
+ b = cfg.model.bytes_per_elem
+ # score bytes per rank per hop
+ bytes_per_hop = hq_per_pe * T_q * S_local * b
+ # split-group AllReduce (intra-cube assumed if group fits)
+ bw = cfg.machine.bw_intra
+ alpha = cfg.machine.alpha_intra
+ per_hop = 2 * (split - 1) / split * bytes_per_hop / bw + 2 * (split - 1) * alpha
+ total = per_hop * cfg.topo.cp
+ total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp)
+ return StageCost(
+ name=f"C3 Score AllReduce ({split}-way, xCP hops)",
+ formula=f"2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} / BW + latency",
+ compute_s=0, memory_s=0, comm_s=total,
+ bound="comm", visible_s=total,
+ flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
+ flops_formula="0",
+ mem_formula="0",
+ comm_formula=(
+ f"PURPOSE: TP={cfg.topo.tp} > H_kv={cfg.model.h_kv}, so each KV\n"
+ f"head is split across {split} TP ranks along the head-dim (d_h\n"
+ f"/{split} per rank). Each rank's Q.K^T is therefore only a\n"
+ f"PARTIAL dot product; the {split} ranks sharing one KV head\n"
+ f"must AllReduce their partial scores to get the true score.\n"
+ f"This fires per hop of the ring, so {cfg.topo.cp}x per layer.\n"
+ f"To eliminate C3: set KV mode = 'replicate' (costs {split}x KV\n"
+ f"memory but no per-hop score AR), or lower TP so TP <= H_kv.\n"
+ f"---\n"
+ f"CP * 2*(split-1)/split * (H_q/TP)*T_q*S_local*b "
+ f"= {cfg.topo.cp} * 2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} "
+ f"= {total_comm_bytes/1e6:.2f} MB total"
+ ),
+ )
+
+
+def all_stages(cfg: FullConfig) -> list[StageCost]:
+ """Ordered per-layer attention stages.
+
+ In prefill with CP > 1, C1 (the CP ring) is inserted between S7 and S8
+ because it runs concurrently with S5-S7 per hop. Placing it there in the
+ stage list makes the per-stage table read left-to-right in the order
+ things actually happen. C3 (score AR) sits with C1 for the same reason.
+ C2 (TP AllReduce on W_O) sits right after S10.
+
+ In decode, C1's comm is folded into S8, so no separate C1 row.
+ """
+ stages = [
+ stage_rmsnorm(cfg),
+ stage_wq(cfg),
+ stage_wkv(cfg),
+ stage_kv_append(cfg),
+ stage_qkT(cfg),
+ stage_softmax(cfg),
+ stage_pv(cfg),
+ ]
+ # Concurrent comm (with S5-S7 in prefill) goes here, right after S7.
+ if cfg.topo.mode == "prefill" and cfg.topo.cp > 1:
+ stages.append(comm_cp_ring(cfg))
+ _c3 = comm_kv_split_allreduce(cfg)
+ if _c3.visible_s > 0 or _c3.comm_s > 0:
+ stages.append(_c3)
+ stages.extend([
+ stage_merge(cfg),
+ stage_normalize(cfg),
+ stage_wo(cfg),
+ ])
+ # C2 fires right after S10 (W_O produces per-TP-rank output; AR combines).
+ stages.append(comm_tp_allreduce(cfg))
+ return stages
+
+
+# ── FFN block stages (lightweight; per PE, per layer) ────────────
+def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ bytes_ = T_q * d * b * 2
+ flops = 4 * T_q * d
+ mem_s = bytes_ / cfg.machine.bw_hbm
+ return StageCost(
+ name="F1 RMSNorm (pre-FFN)",
+ formula=f"{T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
+ compute_s=0, memory_s=mem_s, comm_s=0,
+ bound="memory", visible_s=mem_s,
+ flops=flops, mem_bytes=bytes_,
+ flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
+ mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
+ )
+
+
+def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ flops = 2 * T_q * d * ffn_per_pe
+ weight_B = d * ffn_per_pe * b
+ cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+ vis, bnd = _visible(cmp_s, mem_s, 0)
+ return StageCost(
+ name=name,
+ formula=f"FLOPs = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
+ compute_s=cmp_s, memory_s=mem_s, comm_s=0,
+ bound=bnd, visible_s=vis,
+ flops=int(flops), mem_bytes=int(weight_B),
+ flops_formula=f"2*T_q*d*(ffn/div) = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
+ mem_formula=f"d*(ffn/div)*b = {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
+ )
+
+
+def stage_ffn_gate(cfg: FullConfig) -> StageCost:
+ ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+ return _ffn_gemm(cfg, "F2 W_gate GEMM", ffn_per_pe)
+
+
+def stage_ffn_up(cfg: FullConfig) -> StageCost:
+ ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+ return _ffn_gemm(cfg, "F3 W_up GEMM", ffn_per_pe)
+
+
+def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
+ """SwiGLU element-wise activation on the intermediate FFN tensor."""
+ T_q = cfg.topo.T_q
+ ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+ b = cfg.model.bytes_per_elem
+ bytes_ = 3 * T_q * ffn_per_pe * b
+ flops = 3 * T_q * ffn_per_pe # gate * silu(up) approximated as 3 flops/elt
+ mem_s = bytes_ / cfg.machine.bw_hbm
+ return StageCost(
+ name="F4 SwiGLU act",
+ formula=f"3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
+ compute_s=0, memory_s=mem_s, comm_s=0,
+ bound="memory", visible_s=mem_s,
+ flops=flops, mem_bytes=bytes_,
+ flops_formula=f"~3*T_q*(ffn/div) = 3*{T_q}*{ffn_per_pe} = {flops}",
+ mem_formula=f"3*T_q*(ffn/div)*b = 3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
+ )
+
+
+def stage_ffn_down(cfg: FullConfig) -> StageCost:
+ """W_down: (ffn_per_pe, hidden). Row-parallel over FFN scope."""
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+ flops = 2 * T_q * ffn_per_pe * d
+ weight_B = ffn_per_pe * d * b
+ cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+ vis, bnd = _visible(cmp_s, mem_s, 0)
+ return StageCost(
+ name="F5 W_down GEMM",
+ formula=f"FLOPs = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
+ compute_s=cmp_s, memory_s=mem_s, comm_s=0,
+ bound=bnd, visible_s=vis,
+ flops=int(flops), mem_bytes=int(weight_B),
+ flops_formula=f"2*T_q*(ffn/div)*d = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
+ mem_formula=f"(ffn/div)*d*b = {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
+ )
+
+
+def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
+ """AllReduce on FFN output across the FFN sharding scope."""
+ scope = cfg.topo.ffn_shard_scope
+ divisor = cfg.ffn_shard_divisor # TP or TP*CP or TP*CP*DP
+ if divisor <= 1:
+ return StageCost(
+ name="CF1 FFN AllReduce",
+ formula="FFN scope = 1: not needed",
+ compute_s=0, memory_s=0, comm_s=0,
+ bound="trivial", visible_s=0,
+ )
+ T_q = cfg.topo.T_q
+ d = cfg.model.hidden
+ b = cfg.model.bytes_per_elem
+ bytes_ = T_q * d * b
+ # Choose BW/alpha tier based on scope (rough): TP=intra-cube (or inter-cube),
+ # +CP=inter-SIP possible, +DP=inter-SIP always.
+ if "DP" in scope or "CP" in scope:
+ bw = cfg.machine.bw_intersip if cfg.topo.sips_used > 1 else cfg.machine.bw_inter
+ alpha = cfg.machine.alpha_intersip if cfg.topo.sips_used > 1 else cfg.machine.alpha_inter
+ else: # TP only
+ bw = cfg.machine.bw_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.bw_inter
+ alpha = cfg.machine.alpha_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.alpha_inter
+ comm_time = 2 * (divisor - 1) / divisor * bytes_ / bw + 2 * (divisor - 1) * alpha
+ total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_)
+ return StageCost(
+ name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})",
+ formula=f"2*({divisor}-1)/{divisor} * {T_q}*{d}*{b} B / BW + latency",
+ compute_s=0, memory_s=0, comm_s=comm_time,
+ bound="comm", visible_s=comm_time,
+ flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
+ flops_formula="0",
+ mem_formula="0",
+ comm_formula=(
+ f"PURPOSE: W_down is row-parallel across the {divisor} ranks in\n"
+ f"the FFN scope (scope={scope}). Each rank produced a PARTIAL\n"
+ f"hidden vector after W_down; this all-reduce sums them so every\n"
+ f"rank has the full FFN output for the next layer. Larger scope\n"
+ f"= less FFN weight memory per PE but bigger AR bill.\n"
+ f"---\n"
+ f"2*(div-1)/div * T_q*d*b = 2*({divisor}-1)/{divisor} * "
+ f"{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
+ f"{bw/1e9:.0f} GB/s (scope={scope})"
+ ),
+ )
+
+
+def all_ffn_stages(cfg: FullConfig) -> list[StageCost]:
+ return [
+ stage_ffn_rmsnorm(cfg),
+ stage_ffn_gate(cfg),
+ stage_ffn_up(cfg),
+ stage_ffn_swiglu(cfg),
+ stage_ffn_down(cfg),
+ comm_ffn_allreduce(cfg),
+ ]
diff --git a/tests/analytical_visualization/tensor_sharding.py b/tests/analytical_visualization/tensor_sharding.py
new file mode 100644
index 0000000..675ebcc
--- /dev/null
+++ b/tests/analytical_visualization/tensor_sharding.py
@@ -0,0 +1,330 @@
+"""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: (S_kv, H_kv*d_head), 2D shard ─
+ kv_row_splits = cp
+ kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
+ _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:,})",
+ 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="2D shard: seq axis (CP) x head axis (TP)",
+ 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")
+
+ ax.set_title(
+ f"Tensor sharding view (PE {my_pe}) | TP={tp}, CP={cp}, "
+ 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
diff --git a/tests/analytical_visualization/topology_map.py b/tests/analytical_visualization/topology_map.py
new file mode 100644
index 0000000..53a974c
--- /dev/null
+++ b/tests/analytical_visualization/topology_map.py
@@ -0,0 +1,611 @@
+"""Draw the SIP topology diagram, color-coded by parallelism group.
+
+Color scheme:
+ - CP ring: orange arrows (unchanged)
+ - Inter-SIP: red arrows across SIP boundaries
+ - PP stage: cube border color (each PP stage gets a distinct hue)
+ - TP group: PE fill color (each TP group within a cube gets one hue,
+ matches its cube's PP stage but slightly desaturated)
+ - EP group: a small badge on the cube (for MoE)
+ - DP replica: cubes of each DP replica get a hatched pattern
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.lines import Line2D
+
+from .model_config import FullConfig
+
+
+MESH_W = 4
+MESH_H = 4
+PE_COLS = 4
+PE_ROWS = 2
+
+
+def _group_label(cfg, group_idx: int) -> str:
+ """Label for a cube-level group under the current placement.
+ - cp_placement=cube -> CP rank (or CP/TP if both on cube)
+ - cp_placement=pe, tp_placement=cube -> TP rank
+ - both on pe -> single group ("PE")
+ """
+ parts = []
+ if cfg.topo.cp_placement == "cube":
+ cp_r = group_idx % max(1, cfg.topo.cp)
+ parts.append(f"CP{cp_r}")
+ if cfg.topo.tp_placement == "cube":
+ div = cfg.topo.cp if cfg.topo.cp_placement == "cube" else 1
+ tp_r = (group_idx // max(1, div)) % max(1, cfg.topo.tp)
+ parts.append(f"TP{tp_r}")
+ return "/".join(parts) if parts else "PE"
+
+# Fully-distinct palette for (PP stage x CP rank) — 24 unique colors.
+# Each cube's (pp, cp) pair maps to one entry; wraps if you exceed 24 groups.
+_GROUP_PALETTE = [
+ "#e6194b", "#3cb44b", "#ffe119", "#4363d8",
+ "#f58231", "#911eb4", "#42d4f4", "#f032e6",
+ "#bfef45", "#fabed4", "#469990", "#dcbeff",
+ "#9a6324", "#fffac8", "#800000", "#aaffc3",
+ "#808000", "#ffd8b1", "#000075", "#a9a9a9",
+ "#004d40", "#c62828", "#4527a0", "#00695c",
+]
+
+# Kept for backward-compat / PP-stage arrows
+_PP_HUES = _GROUP_PALETTE[:8]
+
+
+def _shade(hue_hex: str, t: float) -> str:
+ """Mix hue_hex toward white by t in [0, 1] (0 = original, 1 = white)."""
+ h = hue_hex.lstrip("#")
+ r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
+ rn = int(r + (255 - r) * t)
+ gn = int(g + (255 - g) * t)
+ bn = int(b + (255 - b) * t)
+ return f"#{rn:02x}{gn:02x}{bn:02x}"
+
+
+def _cp_color(pp_stage: int, cp_rank: int, cp_size: int) -> tuple[str, str]:
+ """Return (border_color, pe_fill_color) for a cube at (pp_stage, cp_rank).
+
+ Each unique (pp, cp) pair gets a **fully distinct color** from the
+ 24-entry palette. Border is the pure color; PE fill is a lightened
+ version so the cube border stays visible around the PE grid.
+ """
+ idx = pp_stage * max(1, cp_size) + cp_rank
+ color = _GROUP_PALETTE[idx % len(_GROUP_PALETTE)]
+ pe_fill = _shade(color, 0.30) # slightly lightened for PE fill
+ return color, pe_fill
+
+# Inactive (unused) colors.
+_INACTIVE_CUBE_FACE = "#f7f7f7"
+_INACTIVE_CUBE_EDGE = "#c8c8c8"
+_INACTIVE_PE = "#eaeaea"
+
+
+def _snake_path(n: int, mesh_w: int, mesh_h: int) -> list[int]:
+ path: list[int] = []
+ for r in range(mesh_h):
+ cols = range(mesh_w) if r % 2 == 0 else range(mesh_w - 1, -1, -1)
+ for c in cols:
+ path.append(r * mesh_w + c)
+ if len(path) == n:
+ return path
+ return path[:n]
+
+
+def _rect_shape(n: int, max_w: int, max_h: int) -> tuple[int, int]:
+ """Return (rows, cols) for n cubes packed as square as possible in max_wxmax_h."""
+ if n <= 0:
+ return (0, 0)
+ # Prefer near-square shapes; cap at mesh dims.
+ best = None
+ for cols in range(1, min(n, max_w) + 1):
+ rows = (n + cols - 1) // cols
+ if rows > max_h:
+ continue
+ # Score: penalise non-squareness + wasted cells
+ waste = rows * cols - n
+ aspect = abs(rows - cols)
+ score = aspect * 10 + waste
+ if best is None or score < best[0]:
+ best = (score, rows, cols)
+ if best is None:
+ # Fallback: single row
+ return (1, min(n, max_w))
+ return best[1], best[2]
+
+
+def _pack_groups_2d(items_per_group: int, n_groups: int,
+ mesh_w: int, mesh_h: int,
+ layout_mode: str = "compact") -> list[tuple[int, int]]:
+ """Return list of (row, col) positions for n_groups × items_per_group
+ cubes in a mesh_w x mesh_h mesh.
+
+ layout_mode:
+ - "compact": each group is a near-square rectangle; groups are also
+ tiled 2D as near-square. Ideal for 2x2, 2x4 etc.
+ - "linear": each group is a single row of cubes; groups tile down
+ as rows (original sequential layout).
+ """
+ if n_groups == 0 or items_per_group == 0:
+ return []
+
+ if layout_mode == "linear":
+ # Each group is a single horizontal row of cubes.
+ tp_rows, tp_cols = 1, min(items_per_group, mesh_w)
+ else: # compact
+ tp_rows, tp_cols = _rect_shape(items_per_group, mesh_w, mesh_h)
+ if tp_rows == 0:
+ return []
+
+ # Group grid: how many groups per row/col of GROUPS
+ max_group_cols = max(1, mesh_w // tp_cols)
+ max_group_rows = max(1, mesh_h // tp_rows)
+ if layout_mode == "linear":
+ # Row-major: fill each mesh row with groups, then wrap down.
+ g_cols = max_group_cols
+ g_rows = max_group_rows
+ else:
+ g_rows, g_cols = _rect_shape(n_groups, max_group_cols, max_group_rows)
+
+ positions: list[tuple[int, int]] = []
+ for g in range(n_groups):
+ gr = g // g_cols
+ gc = g % g_cols
+ base_row = gr * tp_rows
+ base_col = gc * tp_cols
+ for i in range(items_per_group):
+ r = i // tp_cols
+ c = i % tp_cols
+ row = base_row + r
+ col = base_col + c
+ if row >= mesh_h or col >= mesh_w:
+ fallback = (g * items_per_group + i)
+ row = fallback // mesh_w
+ col = fallback % mesh_w
+ positions.append((row, col))
+ return positions
+
+
+def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int,
+ sip_x0: float, sip_y0: float,
+ sip_width: float, sip_height: float,
+ cube_pp_cp: dict[int, tuple[int, int]],
+ ring_paths: dict[int, list[int]],
+ tp_group_cubes: dict[int, list[int]]):
+ """Draw one SIP with PP hue + CP shade + TP group boundaries."""
+ cube_size = min(sip_width / MESH_W, sip_height / MESH_H) * 0.85
+ cube_gap = cube_size * 0.15
+ total_w = MESH_W * (cube_size + cube_gap) - cube_gap
+ total_h = MESH_H * (cube_size + cube_gap) - cube_gap
+ ox = sip_x0 + (sip_width - total_w) / 2
+ oy = sip_y0 + (sip_height - total_h) / 2
+
+ sip_rect = patches.FancyBboxPatch(
+ (sip_x0, sip_y0), sip_width, sip_height,
+ boxstyle="round,pad=0.02", facecolor="#f8f9fa",
+ edgecolor="#212529", linewidth=1.5,
+ )
+ ax.add_patch(sip_rect)
+ ax.text(sip_x0 + sip_width / 2, sip_y0 + sip_height + 0.05,
+ f"SIP {sip_idx}", ha="center", va="bottom",
+ fontsize=10, fontweight="bold")
+
+ def _cube_center(cube_local: int) -> tuple[float, float]:
+ r = cube_local // MESH_W
+ c = cube_local % MESH_W
+ x = ox + c * (cube_size + cube_gap)
+ y = oy + (MESH_H - 1 - r) * (cube_size + cube_gap)
+ return x, y
+
+ for cube_local in range(MESH_W * MESH_H):
+ x, y = _cube_center(cube_local)
+ info = cube_pp_cp.get(cube_local, None)
+ is_used = info is not None
+
+ if is_used:
+ pp_stage, cp_rank = info
+ cube_edge, pe_fill = _cp_color(
+ pp_stage, cp_rank,
+ max(1, cfg.topo.inter_cube_dims),
+ )
+ cube_face = _shade(cube_edge, 0.85)
+ cube_lw = 2.0
+ else:
+ cube_edge = _INACTIVE_CUBE_EDGE
+ cube_face = _INACTIVE_CUBE_FACE
+ cube_lw = 0.6
+ pe_fill = _INACTIVE_PE
+
+ rect = patches.FancyBboxPatch(
+ (x, y), cube_size, cube_size,
+ boxstyle="round,pad=0.01",
+ facecolor=cube_face, edgecolor=cube_edge, linewidth=cube_lw,
+ )
+ ax.add_patch(rect)
+ if is_used:
+ pp_stage, cp_rank = info
+ # Small cube ID above
+ ax.text(x + cube_size / 2, y + cube_size + 0.01,
+ f"cube {cube_local}", ha="center", va="bottom",
+ fontsize=5.5, color=cube_edge, fontweight="bold")
+ # Group label (CP/TP depending on placement)
+ _lbl = _group_label(cfg, cp_rank)
+ ax.text(x + 0.03, y + cube_size - 0.03,
+ _lbl, ha="left", va="top",
+ fontsize=11, color=cube_edge, fontweight="bold",
+ bbox=dict(boxstyle="round,pad=0.15",
+ facecolor="white", edgecolor=cube_edge,
+ linewidth=1.0, alpha=0.9))
+ if cfg.topo.pp > 1:
+ ax.text(x + cube_size - 0.03, y + cube_size - 0.03,
+ f"PP{pp_stage}", ha="right", va="top",
+ fontsize=9, color=cube_edge, fontweight="bold",
+ bbox=dict(boxstyle="round,pad=0.1",
+ facecolor="white", edgecolor=cube_edge,
+ linewidth=0.8, alpha=0.9))
+ else:
+ ax.text(x + cube_size / 2, y + cube_size + 0.008,
+ f"{cube_local}", ha="center", va="bottom",
+ fontsize=6, color=cube_edge)
+
+ # PEs
+ pe_pad = cube_size * 0.08
+ pe_gap = cube_size * 0.02
+ inner_w = cube_size - 2 * pe_pad
+ inner_h = cube_size - 2 * pe_pad
+ pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
+ pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
+ pes_used = cfg.topo.pes_per_cube_used if is_used else 0
+ for pr in range(PE_ROWS):
+ for pc in range(PE_COLS):
+ pe_id = pr * PE_COLS + pc
+ px = x + pe_pad + pc * (pe_w + pe_gap)
+ py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
+ fill = pe_fill if pe_id < pes_used else _INACTIVE_PE
+ pe_rect = patches.Rectangle(
+ (px, py), pe_w, pe_h,
+ facecolor=fill, edgecolor="#666", linewidth=0.3,
+ )
+ ax.add_patch(pe_rect)
+
+ if is_used and cfg.topo.ep > 1:
+ ax.text(x + cube_size - 0.02, y + cube_size - 0.02,
+ f"EP{cfg.topo.ep}", ha="right", va="top",
+ fontsize=6, color="#5f0f40", fontweight="bold")
+
+ # (CP ring arrows removed per user request)
+
+ # Draw TP group bounding boxes (RED dashed) when TP > 8
+ for tp_g, cubes in tp_group_cubes.items():
+ if len(cubes) <= 1:
+ continue # TP fits in one cube, redundant with cube border
+ # Bounding box around all cubes in this TP group
+ xs, ys = [], []
+ for cube_local in cubes:
+ x, y = _cube_center(cube_local)
+ xs.extend([x, x + cube_size])
+ ys.extend([y, y + cube_size])
+ pad = 0.03
+ bbox_x = min(xs) - pad
+ bbox_y = min(ys) - pad
+ bbox_w = (max(xs) - min(xs)) + 2 * pad
+ bbox_h = (max(ys) - min(ys)) + 2 * pad
+ tp_rect = patches.Rectangle(
+ (bbox_x, bbox_y), bbox_w, bbox_h,
+ fill=False, edgecolor="#d90429", linewidth=1.8,
+ linestyle="--",
+ )
+ ax.add_patch(tp_rect)
+ ax.text(bbox_x + bbox_w / 2, bbox_y - 0.05,
+ f"TP group {tp_g}", ha="center", va="top",
+ color="#d90429", fontsize=7, fontweight="bold")
+
+
+def _sip_grid_layout(n_sips: int, topology: str) -> tuple[int, int]:
+ """(rows, cols) for arranging SIPs on the page.
+ - ring: horizontal chain (1 x N)
+ - mesh2d / torus2d: near-square grid
+ """
+ if n_sips <= 0:
+ return 0, 0
+ if topology == "ring" or n_sips <= 2:
+ return 1, n_sips
+ import math
+ cols = int(math.ceil(math.sqrt(n_sips)))
+ rows = (n_sips + cols - 1) // cols
+ return rows, cols
+
+
+def _draw_sip_links(ax, sip_xy: list[tuple[float, float]],
+ topology: str, sip_w: float, sip_h: float,
+ grid_rows: int, grid_cols: int):
+ """Draw inter-SIP interconnect lines.
+
+ Solid lines: adjacent (grid-neighbor) links.
+ Dashed lines: wrap-around links (ring, torus).
+ """
+ n = len(sip_xy)
+ if n <= 1:
+ return
+ color = "#c62828"
+ lw = 2.0
+
+ def center(i):
+ x, y = sip_xy[i]
+ return x + sip_w / 2, y + sip_h / 2
+
+ if topology == "ring":
+ # 1D chain (adjacent solid)
+ for i in range(n - 1):
+ x1, y1 = center(i)
+ x2, y2 = center(i + 1)
+ ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
+ color=color, linewidth=lw)
+ # wrap-around from last back to first (dashed arc below)
+ if n > 2:
+ x_first, y_first = center(0)
+ x_last, y_last = center(n - 1)
+ arc_y = min(y_first, y_last) - sip_h / 2 - 0.35
+ ax.plot([x_first, x_first], [y_first - sip_h / 2, arc_y],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([x_first, x_last], [arc_y, arc_y],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([x_last, x_last], [arc_y, y_last - sip_h / 2],
+ color=color, linewidth=lw, linestyle="--")
+ elif n == 2:
+ # 2 SIPs: single link
+ x1, y1 = center(0)
+ x2, y2 = center(1)
+ # wrap link: dashed loop below
+ arc_y = y1 - sip_h / 2 - 0.35
+ ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([x1, x2], [arc_y, arc_y],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
+ color=color, linewidth=lw, linestyle="--")
+ return
+
+ # mesh2d / torus2d: grid neighbours
+ def idx(r, c):
+ return r * grid_cols + c
+
+ for r in range(grid_rows):
+ for c in range(grid_cols):
+ i = idx(r, c)
+ if i >= n:
+ continue
+ # right
+ if c + 1 < grid_cols and idx(r, c + 1) < n:
+ x1, y1 = center(i)
+ x2, y2 = center(idx(r, c + 1))
+ ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
+ color=color, linewidth=lw)
+ # down (visually — grid row+1 is drawn below)
+ if r + 1 < grid_rows and idx(r + 1, c) < n:
+ x1, y1 = center(i)
+ x2, y2 = center(idx(r + 1, c))
+ ax.plot([x1, x2], [y1 - sip_h / 2, y2 + sip_h / 2],
+ color=color, linewidth=lw)
+
+ if topology == "torus2d":
+ # row wrap: last col -> first col in each row (dashed arc BELOW that row)
+ for r in range(grid_rows):
+ left_i = idx(r, 0)
+ right_i = idx(r, grid_cols - 1)
+ if left_i >= n or right_i >= n or left_i == right_i:
+ continue
+ x1, y1 = center(left_i)
+ x2, y2 = center(right_i)
+ arc_y = min(y1, y2) - sip_h / 2 - 0.35
+ ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([x1, x2], [arc_y, arc_y],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
+ color=color, linewidth=lw, linestyle="--")
+ # col wrap: last row -> first row in each col (dashed arc to the RIGHT)
+ for c in range(grid_cols):
+ top_i = idx(0, c)
+ bot_i = idx(grid_rows - 1, c)
+ if top_i >= n or bot_i >= n or top_i == bot_i:
+ continue
+ x1, y1 = center(top_i) # top row (higher y)
+ x2, y2 = center(bot_i) # bottom row (lower y)
+ arc_x = max(x1, x2) + sip_w / 2 + 0.35
+ ax.plot([x1 + sip_w / 2, arc_x], [y1, y1],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([arc_x, arc_x], [y1, y2],
+ color=color, linewidth=lw, linestyle="--")
+ ax.plot([arc_x, x2 + sip_w / 2], [y2, y2],
+ color=color, linewidth=lw, linestyle="--")
+
+
+def draw_topology(cfg: FullConfig, ax=None, layout_mode: str = "linear"):
+ """Draw one or more SIPs. Each (PP,CP) group gets a unique color.
+
+ SIP grid arrangement + inter-SIP links depend on cfg.topo.sip_topology:
+ "ring" -> 1D horizontal chain, wrap arc below
+ "mesh2d" -> near-square grid, no wrap
+ "torus2d" -> grid + row/col wrap arcs
+ """
+ sips_used = max(1, cfg.topo.sips_used)
+ total_pes = cfg.topo.total_pes
+ pes_per_stage = cfg.topo.pes_per_stage # CP*TP
+ n_replicas = cfg.topo.dp
+ n_pp_stages = cfg.topo.pp
+ pes_per_cube = cfg.topo.pes_per_cube_hw
+ # Placement-aware: how many cubes per "group" (was: cubes-per-TP).
+ # A group is one cube-level unit (one CP rank if cp on cube, or one TP
+ # rank if only tp on cube, etc.). Cubes per group == intra-cube spill.
+ cubes_per_group = max(
+ 1, (cfg.topo.intra_cube_dims + pes_per_cube - 1) // pes_per_cube
+ )
+ n_groups_per_stage = max(1, cfg.topo.inter_cube_dims)
+ cubes_per_stage = cfg.topo.cubes_per_stage
+ cubes_used = cfg.topo.cubes_used
+ cubes_per_sip = MESH_W * MESH_H
+
+ sip_topo = cfg.topo.sip_topology
+ grid_rows, grid_cols = _sip_grid_layout(sips_used, sip_topo)
+
+ sip_w = 4.0
+ sip_h = 4.0
+ sip_gap = 0.6
+
+ fig_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap + 1.5
+ fig_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap + 2.0
+ if ax is None:
+ fig, ax = plt.subplots(figsize=(max(6, fig_w), max(6, fig_h)))
+ else:
+ fig = ax.figure
+
+ # 2D packing: each group's cubes form a rectangle inside the SIP mesh.
+ # If more groups than fit in one SIP, spill to next SIP.
+ tp_rows, tp_cols = _rect_shape(cubes_per_group, MESH_W, MESH_H)
+ if tp_rows == 0:
+ tp_rows, tp_cols = 1, 1
+ groups_per_sip = max(1, (MESH_W // tp_cols) * (MESH_H // tp_rows))
+
+ cube_pp_cp_global: dict[int, tuple[int, int]] = {}
+ ring_globals_by_pp: dict[int, list[int]] = {}
+ tp_group_cubes_global: dict[int, list[int]] = {}
+
+ current_sip = 0
+ for rep in range(n_replicas):
+ for pp_s in range(n_pp_stages):
+ # Chunk cube-level groups into per-SIP batches.
+ for chunk_start in range(0, n_groups_per_stage, groups_per_sip):
+ chunk = min(groups_per_sip, n_groups_per_stage - chunk_start)
+ positions = _pack_groups_2d(cubes_per_group, chunk,
+ MESH_W, MESH_H,
+ layout_mode=layout_mode)
+ for local_gi in range(chunk):
+ group_idx = chunk_start + local_gi
+ tp_gid = ((rep * n_pp_stages) + pp_s) * n_groups_per_stage + group_idx
+ cubes_here: list[int] = []
+ for i in range(cubes_per_group):
+ idx_in_group = local_gi * cubes_per_group + i
+ r, c = positions[idx_in_group]
+ local_cube = r * MESH_W + c
+ gc = current_sip * cubes_per_sip + local_cube
+ cube_pp_cp_global[gc] = (pp_s, group_idx)
+ cubes_here.append(gc)
+ tp_group_cubes_global[tp_gid] = cubes_here
+ ring_globals_by_pp.setdefault(pp_s, []).append(cubes_here[0])
+ current_sip += 1 # next SIP for the next chunk (or (rep, pp))
+
+ # Partition (pp, cp) map by SIP using cube_positions_global
+ cubes_by_sip: list[dict[int, tuple[int, int]]] = [dict() for _ in range(sips_used)]
+ for gc, ppcp in cube_pp_cp_global.items():
+ sip_idx = gc // cubes_per_sip
+ local_cube = gc % cubes_per_sip
+ if sip_idx < sips_used:
+ cubes_by_sip[sip_idx][local_cube] = ppcp
+
+ # Per-SIP snake-path CP rings (organized by pp_stage)
+ ring_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
+ for pp_s, gc_list in ring_globals_by_pp.items():
+ by_sip_local: dict[int, list[int]] = {}
+ for gc in gc_list:
+ si = gc // cubes_per_sip
+ lc = gc % cubes_per_sip
+ by_sip_local.setdefault(si, []).append(lc)
+ for si, locals_list in by_sip_local.items():
+ if si < sips_used:
+ # Snake through the packed positions in order
+ ring_by_sip[si][pp_s] = locals_list
+
+ # TP group cubes by SIP
+ tp_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
+ for tp_g, gc_list in tp_group_cubes_global.items():
+ by_sip_local: dict[int, list[int]] = {}
+ for gc in gc_list:
+ si = gc // cubes_per_sip
+ lc = gc % cubes_per_sip
+ by_sip_local.setdefault(si, []).append(lc)
+ for si, locals_list in by_sip_local.items():
+ if si < sips_used:
+ tp_by_sip[si][tp_g] = locals_list
+
+ # Compute (x, y) top-left for each SIP based on grid layout.
+ # Row 0 is drawn at the TOP visually, so y decreases with row index.
+ sip_xy: list[tuple[float, float]] = []
+ for sip_idx in range(sips_used):
+ r = sip_idx // grid_cols
+ c = sip_idx % grid_cols
+ x = c * (sip_w + sip_gap)
+ y = (grid_rows - 1 - r) * (sip_h + sip_gap)
+ sip_xy.append((x, y))
+
+ for sip_idx in range(sips_used):
+ x0, y0 = sip_xy[sip_idx]
+ _draw_one_sip(
+ ax, cfg, sip_idx,
+ sip_x0=x0, sip_y0=y0,
+ sip_width=sip_w, sip_height=sip_h,
+ cube_pp_cp=cubes_by_sip[sip_idx],
+ ring_paths=ring_by_sip[sip_idx],
+ tp_group_cubes=tp_by_sip[sip_idx],
+ )
+
+ # Inter-SIP interconnect links per user-selected topology.
+ _draw_sip_links(ax, sip_xy, sip_topo, sip_w, sip_h, grid_rows, grid_cols)
+
+ total_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap
+ total_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap
+ # Extra room: bottom for wrap arcs, right for torus col-wrap arcs, top for SIP labels
+ ax.set_xlim(-0.3, total_w + 0.9)
+ ax.set_ylim(-0.9, total_h + 0.7)
+ ax.set_aspect("equal")
+ ax.axis("off")
+
+ # Legend — one entry per unique (PP, cube-group) + TP boundary + SIP links
+ legend_handles = []
+ for pp_s in range(n_pp_stages):
+ for gi in range(n_groups_per_stage):
+ color, _ = _cp_color(pp_s, gi, n_groups_per_stage)
+ legend_handles.append(Line2D(
+ [], [], color=color, marker="s", linestyle="",
+ markersize=10,
+ label=(f"PP{pp_s} " if n_pp_stages > 1 else "")
+ + _group_label(cfg, gi),
+ ))
+ if cfg.topo.tp > cfg.topo.pes_per_cube_hw:
+ legend_handles.append(Line2D(
+ [], [], color="#d90429", marker="s", linestyle="--",
+ markersize=10, markerfacecolor="none", markeredgewidth=1.5,
+ label="TP group (dashed red)",
+ ))
+ if cfg.topo.ep > 1:
+ legend_handles.append(Line2D(
+ [], [], color="#5f0f40", marker="s", linestyle="",
+ markersize=8, label=f"EP={cfg.topo.ep}",
+ ))
+ if sips_used > 1:
+ _topo_label = {
+ "ring": "Inter-SIP: ring (chain + wrap)",
+ "mesh2d": "Inter-SIP: 2D mesh (no wrap)",
+ "torus2d": "Inter-SIP: 2D torus (grid + wrap)",
+ }.get(sip_topo, f"Inter-SIP: {sip_topo}")
+ legend_handles.append(Line2D(
+ [], [], color="#c62828", linewidth=2.0,
+ label=_topo_label,
+ ))
+ # Legend below the figure; wrap ncols so it's readable
+ n_items = len(legend_handles)
+ ncol = min(6, max(3, n_items))
+ ax.legend(handles=legend_handles, loc="upper center",
+ bbox_to_anchor=(0.5, -0.02), ncol=ncol,
+ fontsize=8, frameon=False)
+
+ title = (
+ f"Layout: {total_pes} PEs = {n_replicas} DP replica(s) x "
+ f"{n_pp_stages} PP stage(s) x {cubes_per_stage} cubes (CP={cfg.topo.cp}) "
+ f"x TP={cfg.topo.tp} | cubes: {cubes_used}, SIPs: {sips_used}"
+ )
+ ax.set_title(title, fontsize=10)
+
+ return fig