From 9fdde44922007c1a4a37564321bc555485899e9a Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Sun, 26 Jul 2026 22:11:33 -0700 Subject: [PATCH] analytical-viz: interactive Streamlit tool for SIP transformer analysis New tests/analytical_visualization/ module - an interactive dashboard for exploring memory / latency tradeoffs of transformer inference on the SIP architecture. Highlights: - 30+ model presets (Qwen, Llama 2/3/3.1, Mistral, Gemma 2, Phi 3, DeepSeek MLA, Mixtral/Qwen 3 MoE, Grok-1, ...) - Placement toggles: TP and CP each on PE-level vs cube-level - CP ring variant: K/V ring vs Q+O/m/l ring (prefill); in decode the O/m/l all-reduce is folded into S8 (no separate C1 row) - SIP interconnect: ring / mesh2d / torus2d with matching link drawing - Per-stage latency table with compute + memory + comm formulas, auto-scaled ns/us/ms, colored by dominant bound - Ring attention loop indicator on the pipeline diagram (purple arc over S5-S8 with 'xN hops' badge) - Tensor sharding view with optional physical PE/cube annotations - Replication-waste + optimization-hints panel - Save & compare configurations (config1, config2, ...): summary table plus side-by-side per-stage attention and FFN latency, best-in-row highlighting - Symbol glossary with current values for every symbol used in formulas Not tied to production sim_engine or runtime API; purely analytical tooling for design-space exploration. Co-Authored-By: Claude Opus 4.7 --- tests/analytical_visualization/__init__.py | 0 tests/analytical_visualization/app.py | 1284 +++++++++++++++++ tests/analytical_visualization/autosuggest.py | 106 ++ .../analytical_visualization/memory_layout.py | 225 +++ .../analytical_visualization/model_config.py | 310 ++++ .../analytical_visualization/model_presets.py | 126 ++ .../optimization_report.py | 235 +++ .../pe_weight_layout.py | 221 +++ .../pipeline_diagram.py | 300 ++++ .../stage_latencies.py | 767 ++++++++++ .../tensor_sharding.py | 330 +++++ .../analytical_visualization/topology_map.py | 611 ++++++++ 12 files changed, 4515 insertions(+) create mode 100644 tests/analytical_visualization/__init__.py create mode 100644 tests/analytical_visualization/app.py create mode 100644 tests/analytical_visualization/autosuggest.py create mode 100644 tests/analytical_visualization/memory_layout.py create mode 100644 tests/analytical_visualization/model_config.py create mode 100644 tests/analytical_visualization/model_presets.py create mode 100644 tests/analytical_visualization/optimization_report.py create mode 100644 tests/analytical_visualization/pe_weight_layout.py create mode 100644 tests/analytical_visualization/pipeline_diagram.py create mode 100644 tests/analytical_visualization/stage_latencies.py create mode 100644 tests/analytical_visualization/tensor_sharding.py create mode 100644 tests/analytical_visualization/topology_map.py 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)}' + + '' + + '' + '' + + '' + ''.join(rows_html) + '
StageFormulaBoundTime
' + ) + # 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'{_rows}
' + ) + + 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