"""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 dataclasses import importlib import sys import matplotlib.pyplot as plt import pandas as pd import streamlit as st # Streamlit's hot-reload only re-runs THIS script; imported sub-modules stay # cached in sys.modules across reruns. When we edit auto_explore.py or # auto_hardware.py while the Streamlit process is alive, the app would # keep using the pre-edit versions until a full Ctrl+C + restart. Force # a reload of our own modules on every rerun so signature changes land # without a full restart. # Order matters: model_config first because everything else imports from it # (e.g. TopologyConfig). If we reload downstream modules while model_config # is stale, they end up holding the OLD dataclass definition and every # `cfg.topo.new_field` lookup blows up with AttributeError. for _mod_name in ( "tests.analytical_visualization.model_config", "tests.analytical_visualization.model_presets", "tests.analytical_visualization.autosuggest", "tests.analytical_visualization.memory_layout", "tests.analytical_visualization.stage_latencies", "tests.analytical_visualization.stage_shapes", "tests.analytical_visualization.chip_roofline", "tests.analytical_visualization.auto_explore", "tests.analytical_visualization.auto_hardware", "tests.analytical_visualization.pe_weight_layout", "tests.analytical_visualization.tensor_sharding", "tests.analytical_visualization.topology_map", "tests.analytical_visualization.pipeline_diagram", "tests.analytical_visualization.optimization_report", ): _m = sys.modules.get(_mod_name) if _m is not None: importlib.reload(_m) 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.stage_shapes import ( attn_stage_shape_rows, ffn_stage_shape_rows, ) from tests.analytical_visualization.chip_roofline import ( ai_sensitivity_curve, arithmetic_intensity, balance_context, bound_regime, critical_batch, good_batch, good_context, knee_batch, max_batch_within_slo, memory_budget_curve_vs_batch, memory_budget_curve_vs_skv, per_token_latency_curve, size_deployment, step_latency_curve, t_com, t_mem_long, t_mem_short, total_active_params, utilization_at, ) 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") b_batch = st.selectbox( "Batch B (concurrent requests)", [1, 2, 4, 8, 16, 32, 64, 128, 256], index=0, key="b_batch", help=("Number of concurrent requests processed in one decode/prefill " "step. KV cache scales linearly with B (each request holds its " "own slice). Compute + AR-comm bytes also scale with B, but " "weights are shared. Batching is the primary way to raise " "HBM-bandwidth efficiency for decode."), ) # Compute auto-suggest for defaults / apply button _default_machine = MachineParams() _default_suggestion = auto_suggest(model, _default_machine, s_kv, mode, b=b_batch) 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 # Apply the picked cp_placement so the sidebar reflects the packed # config (autosuggest may have chosen "pe" packing to minimize cubes). st.session_state["cp_placement"] = _default_suggestion.cp_placement st.session_state["_last_preset"] = preset_name with st.sidebar.expander("Parallelism", expanded=True): st.caption( f"Auto-suggest (memory-min): CP={_default_suggestion.cp}, " f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp} " f"→ **{_default_suggestion.cubes_used} cube(s)**, " f"{_default_suggestion.pes_used} PEs " f"(cp_placement={_default_suggestion.cp_placement})" ) # Three scope-aware memory-min buttons. Each runs auto_suggest with a # different (include_attention, include_ffn) filter so the memory-fit # constraint counts only the relevant block: # - Attn → attention weights + KV cache (no FFN) # - FFN/MoE → FFN weights only (no attn / no KV) # - Attn+FFN → everything (the traditional autosuggest) st.caption( "**Apply memory-min** — smallest deployment that fits the chosen " "block. Each button sizes for a different memory footprint." ) _sb_scope_meta = { "attn": (True, False, "Attn"), "ffn": (False, True, "FFN/MoE"), "full": (True, True, "Attn+FFN"), } _sb_cols = st.columns(3) _sb_clicks = {} for _sb_scope, _sb_col in zip(("attn", "ffn", "full"), _sb_cols): with _sb_col: _sb_clicks[_sb_scope] = st.button( _sb_scope_meta[_sb_scope][2], width='stretch', key=f"_sb_apply_memmin_{_sb_scope}", help=(f"Memory-min autosuggest — memory budget counts only " f"{_sb_scope_meta[_sb_scope][2]} weights" f"{' + KV cache' if _sb_scope == 'attn' else ''}" f". Picks the smallest deployment that fits."), ) for _sb_scope, _sb_click in _sb_clicks.items(): if not _sb_click: continue _sb_ia, _sb_ff, _sb_lbl = _sb_scope_meta[_sb_scope] with st.spinner(f"Sizing memory-min for {_sb_lbl}..."): _sb_sug = auto_suggest( model, _default_machine, s_kv=s_kv, mode=mode, b=b_batch, include_attention=_sb_ia, include_ffn=_sb_ff, ) if not _sb_sug.fits: st.warning( f"{_sb_lbl} memory-min: no config fits at 10% slack; " f"best-effort {_sb_sug.pes_used} PEs, {_sb_sug.cubes_used} " f"cubes ({_sb_sug.reason})." ) st.session_state["cp"] = _snap(_sb_sug.cp, _CP_OPTS) st.session_state["tp"] = _snap(_sb_sug.tp, _TP_OPTS) st.session_state["pp"] = _snap(_sb_sug.pp, _PP_OPTS) st.session_state["dp"] = 1 st.session_state["ep"] = 1 st.session_state["cp_placement"] = _sb_sug.cp_placement # auto_suggest only explores (cp, tp, pp, cp_placement); reset # every other TopologyConfig field to its dataclass default so # the applied topology matches what memory-min actually scored # (else stale tp_placement="cube" etc. explodes cubes_used). st.session_state["tp_placement"] = "pe" st.session_state["kv_mode"] = "split" st.session_state["cp_ring_variant"] = "kv" # ffn_scope_label is a dynamic string (embeds current tp/cp) — # dropping the key lets the radio fall back to index=1 which is # "TP+CP", the TopologyConfig default. st.session_state.pop("ffn_scope_label", None) st.session_state["_pl_active_scope"] = _sb_scope 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, b=b_batch, 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, b=b_batch) # 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 ───────────────────────────────────────────────────────── # One tab per auto feature; each has TWO sweep buttons (Attention only / # Attn + FFN/MoE) so the user can toggle scope without switching tabs. # Auto Suggest is renamed "Parallelism" since it only varies parallelism # knobs (hardware is held fixed at the sidebar values). (tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw, tab_roofline, tab_planning, tab_prules) = st.tabs([ "Physical layout", "Auto Suggest Parallelism", "Memory breakdown", "Per-stage latency", "Save & compare", "Auto Hardware", "Chip roofline & B*", "Capacity planning", "Parallelism rules", ]) # ── TAB 1: Physical layout ────────────────────────────────────── with tab_layout: # Quick shortcut: three buttons that pick the memory-optimal # (smallest-fit) Pareto config for a chosen scope and load it into # the sidebar. Layout re-renders immediately on the next Streamlit # run with the new config. from tests.analytical_visualization.auto_explore import ( run_auto_explore as _run_auto_explore_layout, ) st.markdown( "**Quick auto-suggest layout** — pick a scope; the memory-optimal " "(smallest-fit) Pareto config for that scope is loaded into the " "sidebar and this layout redraws around it." ) _pl1, _pl2, _pl3, _pl4 = st.columns([1, 1, 1, 1]) with _pl1: _pl_run_attn = st.button("Run sweep — Attention", type="primary", width='stretch', key="_pl_run_attn") with _pl2: _pl_run_ffn = st.button("Run sweep — FFN/MoE", type="primary", width='stretch', key="_pl_run_ffn") with _pl3: _pl_run_full = st.button("Run sweep — Attn + FFN/MoE", type="primary", width='stretch', key="_pl_run_full") _pl_scope_meta = { "attn": (True, False, "Attention only"), "ffn": (False, True, "FFN / MoE only"), "full": (True, True, "Attn + FFN/MoE"), } _pl_button_click = { "attn": _pl_run_attn, "ffn": _pl_run_ffn, "full": _pl_run_full, } for _pl_scope, _pl_clicked in _pl_button_click.items(): if not _pl_clicked: continue _pl_ia, _pl_ff, _pl_lbl = _pl_scope_meta[_pl_scope] with st.spinner(f"Finding memory-optimal parallelism ({_pl_lbl})..."): _pl_res = _run_auto_explore_layout( model, machine, s_kv=s_kv, mode=mode, include_attention=_pl_ia, include_ffn=_pl_ff, b=b_batch, ) if not _pl_res.pareto_scores: st.error( f"No feasible parallelism for {_pl_lbl} — every config " "exceeds per-PE HBM. Try raising `pe_hbm_gb` or reducing " "`s_kv`." ) else: # Smallest-fit Pareto point (see sidebar apply comment). _pl_best = min(_pl_res.pareto_scores, key=lambda s: (s.pes_used, s.hbm_utilization, s.total_latency_ns)) # Load the winning config into the sidebar's session state. st.session_state["cp"] = _pl_best.cp st.session_state["tp"] = _pl_best.tp st.session_state["pp"] = _pl_best.pp st.session_state["dp"] = _pl_best.dp st.session_state["tp_placement"] = _pl_best.tp_placement st.session_state["cp_placement"] = _pl_best.cp_placement st.session_state["cp_ring_variant"] = _pl_best.cp_ring_variant st.session_state["kv_mode"] = _pl_best.kv_shard_mode _pl_tp, _pl_cp, _pl_dp = _pl_best.tp, _pl_best.cp, _pl_best.dp _pl_ffn_map = { "TP": f"TP only (div={max(1, _pl_tp)})", "TP+CP": f"TP*CP (div={max(1, _pl_tp * _pl_cp)})", "TP+CP+DP": f"TP*CP*DP (div={max(1, _pl_tp * _pl_cp * _pl_dp)})", } st.session_state["ffn_scope_label"] = \ _pl_ffn_map[_pl_best.ffn_shard_scope] # Persist the chosen scope so subsequent layout renders filter # stage tables (attn/ffn/full). st.session_state["_pl_active_scope"] = _pl_scope st.success( f"Loaded memory-optimal {_pl_lbl} config: CP={_pl_best.cp} " f"TP={_pl_best.tp} PP={_pl_best.pp} DP={_pl_best.dp} → " f"{_pl_best.latency_ms:.2f} ms, {_pl_best.pes_used} PEs. " "Redrawing layout..." ) st.rerun() # Determine the effective scope for filtering the tables below. Default # to "full" so a fresh page load / sidebar-driven config shows everything. _pl_active_scope = st.session_state.get("_pl_active_scope", "full") _pl_show_attn = _pl_active_scope in ("attn", "full") _pl_show_ffn = _pl_active_scope in ("ffn", "full") _pl_active_label = _pl_scope_meta[_pl_active_scope][2] st.markdown(f"### Layout scope: **{_pl_active_label}**") st.divider() # 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." ) if _pl_show_attn: _render_table(_stages_attn, "Attention") if _pl_show_ffn: _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) # Row G: All tensor shapes (per PE) — one collapsed expander with # weight, KV, and per-stage shapes for anyone who wants them. with st.expander("All tensor shapes (per PE)", expanded=False): st.caption( "Every tensor this deployment holds, per PE. Weights + KV " "show global vs per-PE shape (one row per tensor, per layer). " "Per-stage shapes show input / weight / output tensors each " "stage of the forward pass operates on." ) _lp = (model.layers + topo.pp - 1) // topo.pp st.markdown(f"**Attention weights** (per layer, layers/stage = {_lp})") _attn_rows_layout = attention_weight_rows(cfg) _display_attn_layout = [ {k: v for k, v in r.items() if not k.startswith("_")} for r in _attn_rows_layout ] st.dataframe(pd.DataFrame(_display_attn_layout), width='stretch', hide_index=True) st.markdown(f"**FFN weights** (per layer, layers/stage = {_lp})") _ffn_rows_layout = ffn_weight_rows(cfg) _display_ffn_layout = [ {k: v for k, v in r.items() if not k.startswith("_")} for r in _ffn_rows_layout ] st.dataframe(pd.DataFrame(_display_ffn_layout), width='stretch', hide_index=True) st.markdown(f"**KV cache** (per layer, S_kv = {s_kv:,})") _kv_rows_layout = kv_cache_rows(cfg) _display_kv_layout = [ {k: v for k, v in r.items() if not k.startswith("_")} for r in _kv_rows_layout ] st.dataframe(pd.DataFrame(_display_kv_layout), width='stretch', hide_index=True) st.markdown("**Per-stage attention shapes**") st.dataframe(pd.DataFrame(attn_stage_shape_rows(cfg)), width='stretch', hide_index=True) st.markdown("**Per-stage FFN shapes**") st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)), width='stretch', hide_index=True) # Row H: Full-model latency (TTFT & TPOT) ────────────────────── st.divider() st.subheader("Full-model latency (TTFT & TPOT)") st.caption( "End-to-end per-token latency (TPOT) is one decode step summed " "across all layers. TTFT is one prefill pass (all layers) for a " "prompt of the length below. **Assumes PP=1** (no PP-bubble " "modeling); comm is the same as your Per-stage latency tab." ) _tt_cols = st.columns([1, 1, 2]) with _tt_cols[0]: _prompt_len = st.number_input( "Prompt length (for TTFT)", min_value=1, max_value=2_000_000, value=int(s_kv), step=128, key="_tt_prompt_len", help="T_q used for the prefill pass. Defaults to sidebar S_kv.", ) with _tt_cols[1]: _tt_n_out = st.number_input( "Output tokens (for end-to-end)", min_value=1, max_value=100_000, value=100, step=1, key="_tt_n_out_tokens", help="Number of tokens to generate. E2E = TTFT + (N-1) × TPOT.", ) # Build two temporary FullConfigs: one for decode (TPOT), one for # prefill (TTFT). Everything else — sharding, B, machine — stays # the same as the sidebar cfg. _tpot_topo = dataclasses.replace(topo, mode="decode", s_kv=int(s_kv)) _ttft_topo = dataclasses.replace(topo, mode="prefill", s_kv=int(_prompt_len)) _tpot_cfg = FullConfig(model=model, topo=_tpot_topo, machine=_default_machine) _ttft_cfg = FullConfig(model=model, topo=_ttft_topo, machine=_default_machine) _decode_attn_s = sum(s.visible_s for s in all_stages(_tpot_cfg)) _decode_ffn_s = sum(s.visible_s for s in all_ffn_stages(_tpot_cfg)) _prefill_attn_s = sum(s.visible_s for s in all_stages(_ttft_cfg)) _prefill_ffn_s = sum(s.visible_s for s in all_ffn_stages(_ttft_cfg)) _tpot_per_layer_s = _decode_attn_s + _decode_ffn_s _ttft_per_layer_s = _prefill_attn_s + _prefill_ffn_s _tpot_total_s = _tpot_per_layer_s * model.layers _ttft_total_s = _ttft_per_layer_s * model.layers _e2e_total_s = _ttft_total_s + (int(_tt_n_out) - 1) * _tpot_total_s _throughput_tokens_per_s = (max(1, int(b_batch)) / _tpot_total_s if _tpot_total_s > 0 else 0.0) _kpi = st.columns(4) _kpi[0].metric("TPOT (per token)", f"{_tpot_total_s * 1e3:.2f} ms", help="One decode-step time × all layers.") _kpi[1].metric(f"TTFT (prompt={_prompt_len:,})", f"{_ttft_total_s * 1e3:.2f} ms", help="One prefill pass × all layers.") _kpi[2].metric(f"E2E for {_tt_n_out} tokens", f"{_e2e_total_s * 1e3:.2f} ms", help="TTFT + (N-1) × TPOT.") _kpi[3].metric(f"Throughput / replica (B={b_batch})", f"{_throughput_tokens_per_s:.1f} tok/s", help="B ÷ TPOT — decode-side throughput this replica.") _tt_rows = [ {"Component": "Attention (decode, TPOT contribution)", "Per layer (µs)": round(_decode_attn_s * 1e6, 3), "All layers (ms)": round(_decode_attn_s * model.layers * 1e3, 3)}, {"Component": "FFN (decode, TPOT contribution)", "Per layer (µs)": round(_decode_ffn_s * 1e6, 3), "All layers (ms)": round(_decode_ffn_s * model.layers * 1e3, 3)}, {"Component": "Attention (prefill, TTFT contribution)", "Per layer (µs)": round(_prefill_attn_s * 1e6, 3), "All layers (ms)": round(_prefill_attn_s * model.layers * 1e3, 3)}, {"Component": "FFN (prefill, TTFT contribution)", "Per layer (µs)": round(_prefill_ffn_s * 1e6, 3), "All layers (ms)": round(_prefill_ffn_s * model.layers * 1e3, 3)}, {"Component": "TPOT total (per token)", "Per layer (µs)": round(_tpot_per_layer_s * 1e6, 3), "All layers (ms)": round(_tpot_total_s * 1e3, 3)}, {"Component": "TTFT total (prefill)", "Per layer (µs)": round(_ttft_per_layer_s * 1e6, 3), "All layers (ms)": round(_ttft_total_s * 1e3, 3)}, ] st.dataframe(pd.DataFrame(_tt_rows), width='stretch', hide_index=True) st.caption( f"Sidebar params in use: **{model.name}** on **{topo.cp}**×" f"**{topo.tp}**×**{topo.pp}** (CP·TP·PP), B=**{b_batch}**, " f"layers=**{model.layers}**. TTFT uses prefill-mode formulas " f"(compute-bound, T_q={_prompt_len:,}); TPOT uses decode-mode " f"(memory-bound, T_q=1). Cross-check per-stage breakdown on " f"the *Per-stage latency* tab." ) # ── 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)." ) # ── TAB (Auto Suggest Parallelism) — one tab, two sweep buttons ── def _render_auto_explore_tab(): """Render the Auto Suggest Parallelism tab body. Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its result under its own session_state key. Whichever button was last clicked drives the display below; both caches persist so re-clicking either without changing model/workload just re-shows the cached view. """ from tests.analytical_visualization.auto_explore import ( compute_parallelism_sensitivity, run_auto_explore, ) st.markdown( "Sweep every valid parallelism config (CP, TP, PP, DP, kv_shard, " "ffn_scope, tp_placement, cp_placement, cp_ring) for the currently-" "selected model + workload and extract the 3D Pareto frontier on " "(latency ↓, PEs ↓, efficiency ↑). Throughput = 1/latency for a " "single request so it collapses with latency and is shown for info only." ) st.caption( "**Attention only**: drop FFN / MoE from the summed latency — " "useful when isolating attention-kernel tuning (matches what our " "kernbench attention sim measures). \n" "**Attn + FFN/MoE**: sum every stage — the full per-token cost." ) st.markdown( f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}** | " f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**" ) _b1, _b2, _b3, _b4 = st.columns([1, 1, 1, 1]) with _b1: _run_attn = st.button("Run sweep — Attention", type="primary", width='stretch', key="_auto_run_attn") with _b2: _run_ffn = st.button("Run sweep — FFN/MoE", type="primary", width='stretch', key="_auto_run_ffn") with _b3: _run_full = st.button("Run sweep — Attn + FFN/MoE", type="primary", width='stretch', key="_auto_run_full") _scope_meta = { "attn": (True, False, "Attention only"), "ffn": (False, True, "FFN / MoE only"), "full": (True, True, "Attn + FFN/MoE"), } _button_click = {"attn": _run_attn, "ffn": _run_ffn, "full": _run_full} for _scope, _clicked in _button_click.items(): if not _clicked: continue _ia, _ff, _lbl = _scope_meta[_scope] with st.spinner(f"Sweeping ~28k configs ({_lbl})..."): _r = run_auto_explore( model, machine, s_kv=s_kv, mode=mode, include_attention=_ia, include_ffn=_ff, b=b_batch, ) st.session_state[f"_auto_explore_result_{_scope}"] = _r st.session_state[f"_auto_explore_ctx_{_scope}"] = ( model.name, s_kv, mode, machine.pe_hbm_gb, ) st.session_state["_auto_explore_active"] = _scope _active = st.session_state.get("_auto_explore_active") if _active is None: st.info( "Click one of the run buttons above to enumerate valid " "parallelism configurations and rank them on the Pareto " "frontier. Takes ~5–10 s per scope. All button results are " "cached so you can flip between scopes without re-running." ) return # Downstream code expects _suffix, _label, include_attention, include_ffn, # _res, _ctx_key. include_attention, include_ffn, _label = _scope_meta[_active] _suffix = f"_{_active}" _result_key = f"_auto_explore_result{_suffix}" _ctx_key = f"_auto_explore_ctx{_suffix}" _res = st.session_state.get(_result_key) if _res is None: st.info(f"Click **Run sweep — {_label.split(' ')[0]}** to compute this view.") return st.markdown(f"### Currently viewing: **{_label}**") # Warn if the cached result is stale relative to the current config. _cached_ctx = st.session_state.get(_ctx_key) _cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb) if _cached_ctx != _cur_ctx: st.warning( f"Cached {_label} result is from a different model/workload. " f"Click **Run sweep — {_label.split(' ')[0]}** to refresh." ) _m1, _m2, _m3, _m4 = st.columns(4) _m1.metric("Enumerated", f"{_res.total_enumerated:,}") _m2.metric("Feasible", f"{_res.total_feasible:,}") _m3.metric("Pareto configs", len(_res.pareto_scores)) if _res.pareto_scores: _best = min(_res.pareto_scores, key=lambda s: (s.total_latency_ns, s.pes_used, s.hbm_utilization)) _m4.metric("Best latency", f"{_best.latency_ms:.2f} ms") if not _res.pareto_scores: st.error( "No feasible configs — every parallelism choice exceeds " "per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or " "reducing `s_kv`." ) return # ── Pareto scatter: latency vs PEs, coloured by efficiency ─ import matplotlib.pyplot as _plt _fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5)) _feas = [s for s in _res.all_scores if s.fits_memory and s.placement_valid] _axs[0].scatter( [s.pes_used for s in _feas], [s.latency_ms for s in _feas], c="lightgrey", s=8, alpha=0.4, label="feasible", ) _pareto_sorted = sorted(_res.pareto_scores, key=lambda s: s.pes_used) _sc = _axs[0].scatter( [s.pes_used for s in _pareto_sorted], [s.latency_ms for s in _pareto_sorted], c=[s.efficiency_score for s in _pareto_sorted], cmap="viridis", s=60, edgecolors="black", linewidths=0.8, label="Pareto", ) _axs[0].plot( [s.pes_used for s in _pareto_sorted], [s.latency_ms for s in _pareto_sorted], "k--", linewidth=0.8, alpha=0.5, ) _axs[0].set_xlabel("PEs used"); _axs[0].set_ylabel("Latency (ms)") _axs[0].set_xscale("log", base=2) _axs[0].set_yscale("log") _axs[0].set_title("Pareto: latency vs PEs") _axs[0].grid(alpha=0.3) _axs[0].legend(loc="upper right", fontsize=8) _fig.colorbar(_sc, ax=_axs[0], label="efficiency") _pareto_by_lat = sorted(_res.pareto_scores, key=lambda s: s.total_latency_ns) _axs[1].scatter( [s.hbm_utilization * 100 for s in _pareto_by_lat], [s.latency_ms for s in _pareto_by_lat], c=[s.pes_used for s in _pareto_by_lat], cmap="plasma", s=60, edgecolors="black", linewidths=0.8, ) _axs[1].set_xlabel("HBM utilization (%)") _axs[1].set_ylabel("Latency (ms)") _axs[1].set_yscale("log") _axs[1].set_title("HBM usage vs latency (Pareto)") _axs[1].grid(alpha=0.3) _sc2 = _axs[1].collections[0] _fig.colorbar(_sc2, ax=_axs[1], label="PEs") _fig.tight_layout() st.pyplot(_fig, width='stretch') _plt.close(_fig) # ── Pareto table ──────────────────────────────────────────── st.markdown("**Pareto configurations** (sorted by latency)") _rows = [] for i, s in enumerate(_pareto_by_lat): _rows.append({ "#": i, "CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp, "kv": s.kv_shard_mode, "ffn": s.ffn_shard_scope, "tp_place": s.tp_placement, "cp_place": s.cp_placement, "cp_ring": s.cp_ring_variant, "lat (ms)": round(s.latency_ms, 3), "eff": round(s.efficiency_score, 4), "PEs": s.pes_used, "SIPs": s.sips_used, "HBM %": round(s.hbm_utilization * 100, 1), }) _df = pd.DataFrame(_rows) st.dataframe(_df, width='stretch', hide_index=True) # ── Top-N by memory (smallest → largest) ───────────────────── st.markdown( "**Top 10 configurations by memory footprint** " "(smallest HBM% first). Useful when memory pressure is the " "real binding constraint of your deployment." ) _pareto_by_mem = sorted( _res.pareto_scores, key=lambda s: (s.hbm_utilization, s.pes_used, s.total_latency_ns), )[:10] _mem_rows = [] for i, s in enumerate(_pareto_by_mem): _mem_rows.append({ "#": i, "HBM %": round(s.hbm_utilization * 100, 1), "weights (GB)": round(s.weights_gb, 2), "KV (GB)": round(s.kv_gb, 2), "PEs": s.pes_used, "SIPs": s.sips_used, "lat (ms)": round(s.latency_ms, 3), "CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp, "kv": s.kv_shard_mode, "ffn": s.ffn_shard_scope, "tp_place": s.tp_placement, "cp_place": s.cp_placement, }) _df_mem = pd.DataFrame(_mem_rows) st.dataframe(_df_mem, width='stretch', hide_index=True) # ── Parallelism sensitivity subplots ────────────────────── st.markdown( "**Parallelism sensitivity** — for the *baseline* config picked " "below, vary each knob one at a time (holding others fixed). " "Solid line = fits memory; red X = infeasible." ) _sens_c1, _sens_c2 = st.columns([1, 4]) with _sens_c1: _sens_row = st.number_input( "Baseline row #", min_value=0, max_value=len(_pareto_by_lat) - 1, value=0, step=1, key=f"_auto_sens_row{_suffix}", help="Which Pareto row is the sensitivity computed around?", ) _baseline_for_sens = _pareto_by_lat[int(_sens_row)] _sens_rows = compute_parallelism_sensitivity( _baseline_for_sens, model, machine, s_kv=s_kv, mode=mode, include_attention=include_attention, include_ffn=include_ffn, ) _fig_sens, _axes_sens = _plt.subplots(1, 5, figsize=(14, 3.2)) for _ax, _row in zip(_axes_sens, _sens_rows): _fit_vals, _fit_lats = [], [] _bad_vals = [] for _v, _lat_ns, _fit in zip(_row.values, _row.latencies_ns, _row.fits_flags): if _lat_ns != _lat_ns: continue if _fit: _fit_vals.append(_v) _fit_lats.append(_lat_ns / 1e6) else: _bad_vals.append((_v, _lat_ns / 1e6)) if _fit_vals: _ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50", linewidth=1.4, markersize=4) if _bad_vals: _ax.scatter( [v for v, _ in _bad_vals], [l for _, l in _bad_vals], marker="x", color="#c0392b", s=40, alpha=0.6, label="no fit", ) _ax.axvline(_row.baseline_value, color="#3498db", linestyle=":", linewidth=1.0) _ax.set_xscale("log") _ax.set_yscale("log") _ax.set_title( f"{_row.knob.upper()} (baseline={_row.baseline_value})", fontsize=10, ) _ax.set_xlabel("value") _ax.set_ylabel("lat (ms)") _ax.grid(alpha=0.3) if _bad_vals: _ax.legend(loc="best", fontsize=7) _fig_sens.tight_layout() st.pyplot(_fig_sens, width='stretch') _plt.close(_fig_sens) # ── Load-into-sidebar ─────────────────────────────────────── st.markdown("**Load a Pareto config into the main sidebar sliders**") _lc1, _lc2 = st.columns([1, 3]) with _lc1: _pick = st.number_input( "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1, value=0, step=1, key=f"_auto_pick_row{_suffix}", ) with _lc2: if st.button("Load into sidebar", type="secondary", key=f"_auto_load_btn{_suffix}"): _s = _pareto_by_lat[int(_pick)] st.session_state["cp"] = _s.cp st.session_state["tp"] = _s.tp st.session_state["pp"] = _s.pp st.session_state["dp"] = _s.dp st.session_state["tp_placement"] = _s.tp_placement st.session_state["cp_placement"] = _s.cp_placement st.session_state["cp_ring_variant"] = _s.cp_ring_variant st.session_state["kv_mode"] = _s.kv_shard_mode _tp, _cp, _dp = _s.tp, _s.cp, _s.dp _ffn_label_map = { "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)})", } st.session_state["ffn_scope_label"] = \ _ffn_label_map[_s.ffn_shard_scope] st.success( f"Loaded row {_pick}: CP={_s.cp} TP={_s.tp} PP={_s.pp} " f"DP={_s.dp} → latency {_s.latency_ms:.2f} ms, " f"{_s.pes_used} PEs. Flip to another tab to see the " f"full breakdown." ) st.rerun() with tab_auto: _render_auto_explore_tab() # ── TAB (Auto Hardware) — one tab, two sweep buttons ───────────── def _render_auto_hardware_tab(): """Render the Auto Hardware tab body. Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its result. Whichever button was last clicked drives the display; both caches persist for fast toggling.""" from tests.analytical_visualization.auto_hardware import ( _HW_KNOB_DEFAULTS, joint_explore, ) st.markdown( "Sweep hardware knobs (PE HBM, HBM BW, TFLOPs, PE↔PE / D2D / C2C " "interconnect BW) **together** with parallelism. For each hardware " "candidate, we pick the latency-minimum parallelism that fits, then " "Pareto-rank the joint (latency ↓, hardware cost ↓) space. Also " "computes per-knob sensitivity: which HW investment gives the " "biggest speedup?" ) st.caption( "**Attention only**: FFN / MoE stages are excluded from the summed " "latency (both in the joint search and sensitivity). \n" "**Attn + FFN/MoE**: full per-token cost." ) hx_l, hx_m = st.columns([1, 1]) with hx_l: st.markdown( f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}**" ) with hx_m: _depth = st.radio( "Sweep depth", options=["two_stage", "balanced", "coarse"], index=1, key="_hw_depth", horizontal=True, help=("two_stage: 1 HW candidate (defaults) + autosuggest " "parallelism, ~1s. Great for sensitivity alone.\n" "balanced (default): 64 HW × ~2k parallelism, ~10-20s.\n" "coarse: 729 HW × ~2k parallelism, ~2-5 min."), ) _bh1, _bh2, _bh3, _bh4 = st.columns([1, 1, 1, 1]) with _bh1: _run_hw_attn = st.button("Run joint sweep — Attention", type="primary", width='stretch', key="_hw_run_attn") with _bh2: _run_hw_ffn = st.button("Run joint sweep — FFN/MoE", type="primary", width='stretch', key="_hw_run_ffn") with _bh3: _run_hw_full = st.button("Run joint sweep — Attn + FFN/MoE", type="primary", width='stretch', key="_hw_run_full") _scope_meta_hw = { "attn": (True, False, "Attention only"), "ffn": (False, True, "FFN / MoE only"), "full": (True, True, "Attn + FFN/MoE"), } _button_click_hw = { "attn": _run_hw_attn, "ffn": _run_hw_ffn, "full": _run_hw_full, } for _scope, _clicked in _button_click_hw.items(): if not _clicked: continue _ia, _ff, _lbl = _scope_meta_hw[_scope] with st.spinner(f"Sweeping HW × parallelism ({_depth}, {_lbl})..."): _r = joint_explore(model, s_kv=s_kv, mode=mode, depth=_depth, include_attention=_ia, include_ffn=_ff, b=b_batch) st.session_state[f"_hw_result_{_scope}"] = _r st.session_state[f"_hw_ctx_{_scope}"] = ( model.name, s_kv, mode, _depth, ) st.session_state["_hw_active"] = _scope _active = st.session_state.get("_hw_active") if _active is None: st.info( "Click one of the run buttons to explore hardware × parallelism " "co-design for this model + workload. All button results are " "cached so you can flip between scopes without re-running." ) return include_attention, include_ffn, _label = _scope_meta_hw[_active] _suffix = f"_{_active}" _hw_result_key = f"_hw_result{_suffix}" _hw_ctx_key = f"_hw_ctx{_suffix}" _hw_res = st.session_state.get(_hw_result_key) if _hw_res is None: st.info(f"Click **Run joint sweep — {_label.split(' ')[0]}** to compute this view.") return st.markdown(f"### Currently viewing: **{_label}**") _cached_ctx = st.session_state.get(_hw_ctx_key) _cur_ctx = (model.name, s_kv, mode, _depth) if _cached_ctx != _cur_ctx: st.warning( f"Cached {_label} result is from a different model/workload/depth. " f"Click **Run joint sweep — {_label.split(' ')[0]}** to refresh." ) _hm1, _hm2, _hm3, _hm4 = st.columns(4) _hm1.metric("HW candidates", _hw_res.total_hw) _hm2.metric("Feasible joint", _hw_res.total_joint) _hm3.metric("Pareto configs", len(_hw_res.pareto_scores)) if _hw_res.pareto_scores: _best = min(_hw_res.pareto_scores, key=lambda s: (s.total_latency_ns, s.parallelism.pes_used, s.parallelism.hbm_utilization)) _hm4.metric("Best latency", f"{_best.latency_ms:.2f} ms") if not _hw_res.pareto_scores: st.error( "No feasible joint configurations — every HW+parallelism " "combo exceeds per-PE HBM at this workload." ) return if True: # ── Panel 1: latency vs cost Pareto ──────────────────────── import matplotlib.pyplot as _plt _fig1, _ax1 = _plt.subplots(figsize=(8, 4)) _feas = _hw_res.all_scores _ax1.scatter( [s.cost_score for s in _feas], [s.latency_ms for s in _feas], c="lightgrey", s=15, alpha=0.5, label="feasible", ) _pareto_by_cost = sorted(_hw_res.pareto_scores, key=lambda s: s.cost_score) _sc = _ax1.scatter( [s.cost_score for s in _pareto_by_cost], [s.latency_ms for s in _pareto_by_cost], c=[s.parallelism.pes_used for s in _pareto_by_cost], cmap="viridis", s=90, edgecolors="black", linewidths=1.0, label="Pareto", ) _ax1.plot( [s.cost_score for s in _pareto_by_cost], [s.latency_ms for s in _pareto_by_cost], "k--", linewidth=1.0, alpha=0.6, ) _ax1.set_xlabel("Hardware cost proxy (unitless — 6.0 = all defaults)") _ax1.set_ylabel("Latency (ms)") _ax1.set_yscale("log") _ax1.set_title("Optimal HW+parallelism: latency vs cost") _ax1.grid(alpha=0.3) _ax1.legend(loc="upper right", fontsize=9) _fig1.colorbar(_sc, ax=_ax1, label="PEs") _fig1.tight_layout() st.pyplot(_fig1, width='stretch') _plt.close(_fig1) # ── Panel 2: per-knob sensitivity ────────────────────────── st.markdown( "**Per-knob sensitivity** (doubling each HW knob from the " "latency-optimal baseline). Higher bar = more valuable " "investment for this workload." ) _fig2, _ax2 = _plt.subplots(figsize=(8, 3.5)) _knobs = [r.knob for r in _hw_res.sensitivity] _speedups = [r.rel_speedup * 100 for r in _hw_res.sensitivity] _colors = ["#2ecc71" if v > 0 else "#95a5a6" for v in _speedups] _ax2.barh(_knobs, _speedups, color=_colors, edgecolor="black", linewidth=0.5) _ax2.invert_yaxis() # biggest speedup on top _ax2.set_xlabel("Speedup when knob is doubled (%)") _ax2.set_title("Sensitivity — where to invest next") _ax2.grid(alpha=0.3, axis="x") for i, r in enumerate(_hw_res.sensitivity): _ax2.text( r.rel_speedup * 100 + 0.5, i, f"{r.baseline_value:g} → {r.doubled_value:g}", va="center", fontsize=8, ) _fig2.tight_layout() st.pyplot(_fig2, width='stretch') _plt.close(_fig2) # ── Panel 3: Pareto table (HW spec + parallelism) ────────── st.markdown("**Pareto joint configurations** (sorted by latency)") _pareto_by_lat = sorted(_hw_res.pareto_scores, key=lambda s: s.total_latency_ns) _rows = [] for i, s in enumerate(_pareto_by_lat): p = s.parallelism h = s.hardware _rows.append({ "#": i, "lat (ms)": round(s.latency_ms, 3), "cost": round(s.cost_score, 2), "PEs": p.pes_used, "CP": p.cp, "TP": p.tp, "PP": p.pp, "DP": p.dp, "kv": p.kv_shard_mode, "PE HBM": f"{h.pe_hbm_gb:g} GB", "HBM BW": f"{h.bw_hbm_gbs:g} GB/s", "TFLOPs": f"{h.peak_tflops_f16:g}", "PE↔PE": f"{h.bw_intra_gbs:g}", "D2D": f"{h.bw_inter_gbs:g}", "C2C": f"{h.bw_intersip_gbs:g}", }) _df_hw = pd.DataFrame(_rows) st.dataframe(_df_hw, width='stretch', hide_index=True) # ── Load hardware into sidebar ───────────────────────────── st.markdown( "**Load a Pareto HW config into the main sidebar** " "(overwrites the Hardware→Per-PE and Hardware→Interconnect " "sliders; also loads the paired parallelism)." ) _lch1, _lch2 = st.columns([1, 3]) with _lch1: _pick_hw = st.number_input( "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1, value=0, step=1, key=f"_hw_pick_row{_suffix}", ) with _lch2: if st.button("Load into sidebar", type="secondary", key=f"_hw_load_btn{_suffix}"): _s = _pareto_by_lat[int(_pick_hw)] _h = _s.hardware _p = _s.parallelism # HW knobs (only load if the value is in the selectbox # options; otherwise leave as-is). _apply_if_in = { "pe_hbm": ([3.0, 6.0, 12.0, 24.0, 48.0, 96.0], _h.pe_hbm_gb), "bw_hbm": ([128.0, 256.0, 512.0, 1024.0, 2048.0], _h.bw_hbm_gbs), "tflops": ([2.0, 4.0, 8.0, 16.0, 32.0, 64.0], _h.peak_tflops_f16), "bw_intra": ([128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0], _h.bw_intra_gbs), "bw_inter": ([32.0, 64.0, 128.0, 256.0, 512.0, 900.0], _h.bw_inter_gbs), "bw_intersip": ([12.5, 25.0, 50.0, 100.0, 200.0, 400.0], _h.bw_intersip_gbs), } for _key, (_opts, _val) in _apply_if_in.items(): if _val in _opts: st.session_state[_key] = _val # Parallelism knobs st.session_state["cp"] = _p.cp st.session_state["tp"] = _p.tp st.session_state["pp"] = _p.pp st.session_state["dp"] = _p.dp st.session_state["tp_placement"] = _p.tp_placement st.session_state["cp_placement"] = _p.cp_placement st.session_state["cp_ring_variant"] = _p.cp_ring_variant st.session_state["kv_mode"] = _p.kv_shard_mode # Ffn_scope_label reconstruction (same as auto_explore tab) _tp2, _cp2, _dp2 = _p.tp, _p.cp, _p.dp _ffn_label_map = { "TP": f"TP only (div={max(1, _tp2)})", "TP+CP": f"TP*CP (div={max(1, _tp2 * _cp2)})", "TP+CP+DP": f"TP*CP*DP (div={max(1, _tp2 * _cp2 * _dp2)})", } st.session_state["ffn_scope_label"] = \ _ffn_label_map[_p.ffn_shard_scope] st.success( f"Loaded row {_pick_hw}: latency {_s.latency_ms:.2f} " f"ms, cost {_s.cost_score:.2f}, " f"{_p.pes_used} PEs. Flip to another tab to see the " f"full breakdown." ) st.rerun() with tab_hw: _render_auto_hardware_tab() # ── TAB 7: Chip roofline & B* ──────────────────────────────────── with tab_roofline: st.subheader("Chip roofline & critical batch size (B*)") st.caption( "Back-of-envelope answer to 'is my chip well-matched to this " "model?' — see how big the batch must be before decode becomes " "compute-bound, and how big the context can grow before the KV " "bandwidth wall kills your utilization. All numbers per PE, " "peak roofline (no compute-util factor), no comm." ) _n_active = total_active_params(model) # ── Local knobs: S_kv, B, FLOPs, HBM BW ─────────────────────── # Override sidebar for exploring the plots without changing the # rest of the app's config. FLOPs/BW knobs drive ALL roofline # math and every plot/formula on this tab (memory-budget stacks # excluded — those are HBM-capacity, not BW). st.markdown("**Explore: local overrides for the plots below**") _kn_a, _kn_b, _kn_c, _kn_d = st.columns(4) with _kn_a: _skv_min = 128 _skv_max = 1_000_000 _rf_skv = st.slider( "S_kv (context length)", min_value=_skv_min, max_value=_skv_max, value=min(int(s_kv), _skv_max), step=_skv_min, key="_rf_skv_override", help="Local to this tab. Drives every plot's S_kv.", ) with _kn_b: _b_min = 1 _b_max = 256 _rf_b = st.slider( "B (batch size)", min_value=_b_min, max_value=_b_max, value=min(max(1, int(b_batch)), _b_max), key="_rf_b_override", help="Local to this tab. Marker on plots + PE-memory KV.", ) with _kn_c: _rf_flops_tflops = st.slider( "FLOPs (TFLOPS)", min_value=1.0, max_value=32.0, value=float(_default_machine.peak_tflops_f16), step=0.5, key="_rf_flops_tflops", help=("Chip peak TFLOPS for the AI-sensitivity plot. " "Higher FLOPs → higher AI and higher B*."), ) with _kn_d: _rf_bw_gbs = st.slider( "HBM BW (GB/s)", min_value=128.0, max_value=1024.0, value=float(_default_machine.bw_hbm_gbs), step=16.0, key="_rf_bw_gbs", help=("Chip HBM bandwidth for the AI-sensitivity plot. " "Higher BW → lower AI (more bytes to feed per FLOP)."), ) # Build the effective chip from the knobs; every roofline number # on this tab is derived from it. _scaled_machine = dataclasses.replace( _default_machine, peak_tflops_f16=_rf_flops_tflops, bw_hbm_gbs=_rf_bw_gbs, ) _ai = arithmetic_intensity(_scaled_machine) _b_star = critical_batch(_scaled_machine, model) _l_star = balance_context(_scaled_machine, model) _flops_mult_now = _rf_flops_tflops / _default_machine.peak_tflops_f16 _bw_mult_now = _rf_bw_gbs / _default_machine.bw_hbm_gbs _ai_base = arithmetic_intensity(_default_machine) _b_star_base = critical_batch(_default_machine, model) st.caption( f"**S_kv = {_rf_skv:,}**, **B = {_rf_b}**, " f"**FLOPs = {_rf_flops_tflops:.1f} TFLOPS** " f"(× {_flops_mult_now:.2f} of base), " f"**BW = {_rf_bw_gbs:.0f} GB/s** " f"(× {_bw_mult_now:.2f} of base). " f"Effective **AI = {_ai:.1f}** (base {_ai_base:.1f}), " f"**B\\* = {_b_star:.0f}** (base {_b_star_base:.0f}), " f"**L\\* = {_l_star:,.0f}** tokens → S_kv/L\\* = " f"**{_rf_skv/_l_star:.2f}**." ) # ── Headline plots: step latency and cost per token (=÷B) ───── _b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256] _step_pts = step_latency_curve(_scaled_machine, model, _b_range, s_kv=_rf_skv) _tok_pts = per_token_latency_curve(_scaled_machine, model, _b_range, s_kv=_rf_skv) _xs_head = [p.batch for p in _step_pts] _hcol1, _hcol2 = st.columns(2) # Plot A — Latency per decode step (undivided by B) with _hcol1: st.markdown("**Latency per decode step** (one forward pass)") _figA, _axA = plt.subplots(figsize=(6, 4.2)) _axA.plot(_xs_head, [p.weight_s * 1e3 for p in _step_pts], "^-", color="#ffbe0b", label="Weight fetch (flat in B)") _axA.plot(_xs_head, [p.compute_s * 1e3 for p in _step_pts], "o-", color="#3a86ff", label="Compute (linear ↑)") _axA.plot(_xs_head, [p.kv_s * 1e3 for p in _step_pts], "s-", color="#d90429", label="KV fetch (linear ↑)") _axA.plot(_xs_head, [p.total_s * 1e3 for p in _step_pts], "-", color="#212529", linewidth=2.5, label="Total") _axA.axvline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6, label=f"B = {_rf_b}") _axA.set_xscale("log", base=2) _axA.set_yscale("log") _axA.set_xlabel("Batch size B") _axA.set_ylabel("Step time (ms)") _axA.set_title("Step latency = weight + compute·B + KV·B") _axA.grid(True, which="both", alpha=0.3) _axA.legend(fontsize=8, loc="upper left") plt.tight_layout() st.pyplot(_figA, width='stretch') plt.close(_figA) # Plot B — Cost per token = latency ÷ B with _hcol2: st.markdown("**Cost per token** (= step latency ÷ B)") _figB, _axB = plt.subplots(figsize=(6, 4.2)) _axB.plot(_xs_head, [p.weight_s * 1e3 for p in _tok_pts], "^-", color="#ffbe0b", label="Weight fetch (shrinks 1/B)") _axB.plot(_xs_head, [p.compute_s * 1e3 for p in _tok_pts], "o--", color="#3a86ff", label="Compute (flat)") _axB.plot(_xs_head, [p.kv_s * 1e3 for p in _tok_pts], "s--", color="#d90429", label="KV fetch (flat)") _axB.plot(_xs_head, [p.total_s * 1e3 for p in _tok_pts], "-", color="#212529", linewidth=2.5, label="Total") _axB.axvline(_b_star, linestyle=":", color="#2e7d32", label=f"B* = {_b_star:.0f}") _axB.axvline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6, label=f"B = {_rf_b}") _axB.set_xscale("log", base=2) _axB.set_yscale("log") _axB.set_xlabel("Batch size B") _axB.set_ylabel("Per-token time (ms)") _axB.set_title("Cost/token = weight/B + compute + KV") _axB.grid(True, which="both", alpha=0.3) _axB.legend(fontsize=8, loc="upper right") plt.tight_layout() st.pyplot(_figB, width='stretch') plt.close(_figB) st.caption( "**Left (step latency)**: how long one forward pass takes. " "Grows with B because compute and KV per-sequence grow linearly; " "weight fetch is loaded once per step so it's flat. This is the " "**SLO / TTFT view** — bigger B → longer step.\n\n" "**Right (cost per token)**: step latency divided by B tokens " "produced. Weight fetch amortizes (shrinks 1/B); compute and KV " "per-sequence stay flat per token. This is the **efficiency / " "$-per-token view** — bigger B (up to ~2·B*) → cheaper per token.\n\n" "Same underlying decomposition, two divisors — the classic " "throughput ↔ latency tradeoff." ) st.divider() # ── PE memory budget (sharded, per-PE) ──────────────────────── st.markdown( "**PE memory budget — how batch and context fill the HBM**" ) _hbm_gb = _default_machine.pe_hbm_gb _skv_sweep = [128, 512, 2048, 8192, 32768, 131072, 524288, 1_000_000] _bud_skv = memory_budget_curve_vs_skv(cfg, _skv_sweep, batch=_rf_b) _b_sweep = [1, 2, 4, 8, 16, 32, 64, 128, 256] _bud_b = memory_budget_curve_vs_batch(cfg, _b_sweep, s_kv=_rf_skv) _mem_L, _mem_R = st.columns(2) with _mem_L: st.markdown(f"vs S_kv (at B = {_rf_b}) — sharded CP={topo.cp}, " f"TP={topo.tp}, PP={topo.pp}") _figM1, _axM1 = plt.subplots(figsize=(6, 4.2)) _xM = [p.axis_val for p in _bud_skv] _w = [p.weights_gb for p in _bud_skv] _k = [p.kv_gb for p in _bud_skv] _tr = [p.transient_gb for p in _bud_skv] _axM1.stackplot(_xM, _w, _k, _tr, labels=["Weights", "KV cache", "Transient"], colors=["#ffbe0b", "#d90429", "#3a86ff"], alpha=0.8) _axM1.axhline(_hbm_gb, color="#212529", linestyle="--", linewidth=1.5, label=f"HBM budget = {_hbm_gb} GB") _axM1.axvline(_rf_skv, color="#7b1fa2", linewidth=1.5, alpha=0.6, label=f"S_kv = {_rf_skv:,}") _axM1.set_xscale("log") _axM1.set_xlabel("S_kv (tokens)") _axM1.set_ylabel("Per-PE memory (GB)") _axM1.set_title("Weight is fixed; KV fills the budget as S_kv grows") _axM1.grid(True, which="both", alpha=0.3) _axM1.legend(fontsize=8, loc="upper left") plt.tight_layout() st.pyplot(_figM1, width='stretch') plt.close(_figM1) with _mem_R: st.markdown(f"vs B (at S_kv = {_rf_skv:,}) — same sharding") _figM2, _axM2 = plt.subplots(figsize=(6, 4.2)) _xM2 = [p.axis_val for p in _bud_b] _w2 = [p.weights_gb for p in _bud_b] _k2 = [p.kv_gb for p in _bud_b] _tr2 = [p.transient_gb for p in _bud_b] _axM2.stackplot(_xM2, _w2, _k2, _tr2, labels=["Weights", "KV cache", "Transient"], colors=["#ffbe0b", "#d90429", "#3a86ff"], alpha=0.8) _axM2.axhline(_hbm_gb, color="#212529", linestyle="--", linewidth=1.5, label=f"HBM budget = {_hbm_gb} GB") _axM2.axvline(_rf_b, color="#7b1fa2", linewidth=1.5, alpha=0.6, label=f"B = {_rf_b}") _axM2.set_xscale("log", base=2) _axM2.set_xlabel("Batch size B") _axM2.set_ylabel("Per-PE memory (GB)") _axM2.set_title("Weight fixed; KV grows linearly with B") _axM2.grid(True, which="both", alpha=0.3) _axM2.legend(fontsize=8, loc="upper left") plt.tight_layout() st.pyplot(_figM2, width='stretch') plt.close(_figM2) _bud_now = memory_budget_curve_vs_skv(cfg, [_rf_skv], batch=_rf_b)[0] st.caption( f"At **S_kv = {_rf_skv:,}, B = {_rf_b}**: weights = " f"**{_bud_now.weights_gb:.2f} GB**, KV = **{_bud_now.kv_gb:.2f} GB**, " f"transient = **{_bud_now.transient_gb:.2f} GB**, " f"used = **{_bud_now.used_gb:.2f} GB / {_hbm_gb:.1f} GB** " f"({'**OVER BUDGET**' if _bud_now.over_budget else f'free = {_bud_now.free_gb:.2f} GB'}). " "Weights are fixed by sharding (CP·TP·PP·EP splits them across " "PEs); KV per PE = B × S_kv × kv_bpt_per_PE." ) st.divider() # ── AI sensitivity to FLOPs / BW ────────────────────────────── st.markdown("**AI sensitivity — how scaling chip FLOPs or BW moves B\\***") _ai_mults = [0.25, 0.5, 1.0, 2.0, 4.0, 8.0] _ai_flops_pts = ai_sensitivity_curve(_default_machine, model, _ai_mults, axis="flops") _ai_bw_pts = ai_sensitivity_curve(_default_machine, model, _ai_mults, axis="bw") _figAI, (_axAI1, _axAI2) = plt.subplots(1, 2, figsize=(12, 4.2)) _axAI1.plot(_ai_mults, [p.ai for p in _ai_flops_pts], "o-", color="#3a86ff", linewidth=2, label="Scale FLOPs") _axAI1.plot(_ai_mults, [p.ai for p in _ai_bw_pts], "s-", color="#d90429", linewidth=2, label="Scale HBM BW") _axAI1.axhline(_ai_base, color="#2e7d32", linestyle=":", label=f"Base AI = {_ai_base:.1f}") _axAI1.axvline(1.0, color="#888", linestyle=":", alpha=0.5) _axAI1.set_xscale("log", base=2) _axAI1.set_yscale("log") _axAI1.set_xlabel("Multiplier on chip parameter") _axAI1.set_ylabel("AI (FLOPs/byte)") _axAI1.set_title("AI = C / W. FLOPs ↑ raises AI, BW ↑ lowers AI") _axAI1.grid(True, which="both", alpha=0.3) _axAI1.legend(fontsize=8, loc="best") _axAI2.plot(_ai_mults, [p.b_star for p in _ai_flops_pts], "o-", color="#3a86ff", linewidth=2, label="Scale FLOPs") _axAI2.plot(_ai_mults, [p.b_star for p in _ai_bw_pts], "s-", color="#d90429", linewidth=2, label="Scale HBM BW") _axAI2.axhline(_b_star_base, color="#2e7d32", linestyle=":", label=f"Base B* = {_b_star_base:.0f}") _axAI2.axvline(_flops_mult_now, color="#3a86ff", linewidth=1.2, alpha=0.5, label=f"FLOPs × {_flops_mult_now:.2f}") _axAI2.axvline(_bw_mult_now, color="#d90429", linewidth=1.2, alpha=0.5, label=f"BW × {_bw_mult_now:.2f}", linestyle="--") _axAI2.set_xscale("log", base=2) _axAI2.set_yscale("log") _axAI2.set_xlabel("Multiplier on chip parameter") _axAI2.set_ylabel("B* (critical batch)") _axAI2.set_title("B* = AI · b/2 — same shape as AI (for BF16)") _axAI2.grid(True, which="both", alpha=0.3) _axAI2.legend(fontsize=8, loc="best") plt.tight_layout() st.pyplot(_figAI, width='stretch') plt.close(_figAI) st.caption( f"**Left**: AI vs chip-parameter multiplier. Scaling **FLOPs** " f"raises AI linearly (more compute per byte). Scaling **HBM BW** " f"lowers AI (more bytes to feed per FLOP). " f"**Right**: B\\* moves in the same direction as AI. " f"**Your knob position**: FLOPs × {_flops_mult_now:.2f}, " f"BW × {_bw_mult_now:.2f} → **AI = {_ai:.2f}**, " f"**B\\* = {_b_star:.0f}** (base AI {_ai_base:.1f}, " f"B\\* {_b_star_base:.0f})." ) st.divider() _regime_now = bound_regime(_scaled_machine, model, batch=max(1, _rf_b), s_kv=_rf_skv) # ── KPI cards ───────────────────────────────────────────────── _r1, _r2, _r3, _r4 = st.columns(4) _r1.metric("AI (FLOPs/byte)", f"{_ai:.1f}", help="C/W. Peak FLOPs per byte of HBM bandwidth. " "Chip-only — no model-dependent factor.") _r2.metric("B* (dense)", f"{_b_star:.0f}", help="C*b/(2*W). Batch size where weight-fetch time " "equals compute time. Below this you're memory-bound.") _r3.metric("L* (tokens)", f"{_l_star:,.0f}", help="Context length where KV-read time equals compute " "time. Above this, no batch size gets you back to " "compute-bound.") _r4.metric(f"Regime @ (B={_rf_b}, S_kv={_rf_skv:,})", _regime_now.replace("-bound", ""), help="Which term dominates cost at the current (B, S_kv) " "you've selected via the sliders above.") # ── Recommended B and L (heuristics) ────────────────────────── _rec_b = good_batch(_scaled_machine, model, sparsity=1.0) _rec_l = good_context(_scaled_machine, model, s_kv=_rf_skv) _g1, _g2, _g3 = st.columns(3) _g1.metric("Good B (dense)", f"{_rec_b.effective:.0f}", help=_rec_b.reason) _g2.metric("Good L ceiling", f"{_rec_l.max_efficient:,.0f} tok", help="Stay at or below L* to keep decode compute-bound.") _g3.metric(f"Utilization @ S_kv={_rf_skv:,}", f"{_rec_l.utilization_at*100:.1f}%", help="Peak compute utilization at this context: " "1 / (1 + S_kv/L*).") # Optional MoE B*: only surface if the preset flags MoE. _is_moe = "MoE" in (preset.family + preset.note) if _is_moe: st.info( "This preset is dense-approximated as an MoE. Reiner Pope's " "'300 × sparsity' rule: for real sparsity k = N_total/N_active, " f"B*_moe ≈ {_b_star:.0f} × k. E.g. DeepSeek 32-of-256 experts " f"(k=8) → B* ≈ {_b_star*8:.0f}." ) st.caption( f"**Full-model attention+FFN params:** {_n_active/1e9:.2f} B. " f"**Effective chip (from knobs):** " f"{_scaled_machine.peak_tflops_f16:.1f} TFLOPs BF16, " f"{_scaled_machine.bw_hbm_gbs:.0f} GB/s HBM. " f"Base chip from sidebar Hardware: " f"{_default_machine.peak_tflops_f16:.0f} TFLOPs, " f"{_default_machine.bw_hbm_gbs:.0f} GB/s." ) st.divider() # ── Regime formulas (short vs long context) ─────────────────── st.markdown("**Cost-term formulas by regime**") _t_com_val = t_com(_scaled_machine, model) * 1e3 # ms _t_mem_s_at_bstar = t_mem_short(_scaled_machine, model, int(round(_b_star))) * 1e3 _t_mem_l_at_skv = t_mem_long(_scaled_machine, model, _rf_skv) * 1e3 _regime_rows = [ {"Term": "t_com (compute)", "Short context (S_kv < L*)": "2 · N / C", "Long context (S_kv > L*)": "2 · N / C (same)", "Value now (ms)": f"{_t_com_val:.3f}"}, {"Term": "t_mem — weight fetch", "Short context (S_kv < L*)": "N · b / (W · B)", "Long context (S_kv > L*)": "N · b / (W · B) — becomes small", "Value now (ms)": f"{_t_mem_s_at_bstar:.3f} at B = B*"}, {"Term": "t_mem — KV fetch", "Short context (S_kv < L*)": "S_kv · kv_bpt / W — small", "Long context (S_kv > L*)": "S_kv · kv_bpt / W — dominates", "Value now (ms)": f"{_t_mem_l_at_skv:.3f}"}, {"Term": "Bottleneck", "Short context (S_kv < L*)": "Weights (B < B*), else Compute (B ≥ B*)", "Long context (S_kv > L*)": "KV bandwidth (regardless of B)", "Value now (ms)": _regime_now.replace("-bound", "")}, ] st.dataframe(pd.DataFrame(_regime_rows), width='stretch', hide_index=True) st.caption( "N = active params, b = bytes/elem, W = HBM BW, C = peak FLOPs, " "kv_bpt = KV bytes per token = 2 · H_kv · d_head · b · layers. " "'Value now' uses the S_kv slider above and B = B\\*." ) st.divider() # ── Recommended B and L (recap) ─────────────────────────────── st.markdown("**How to pick a good B and a good S_kv**") st.markdown( f"- **Good batch**: `B_target = 2 · B* = {_rec_b.effective:.0f}`. " "Below B\\*: memory-bound (doubling B halves cost/token). " "Above 3·B\\*: diminishing returns; latency keeps rising. " "Also cap by your HBM budget: `B_max ≤ (HBM − weights_per_PE) / " "(S_kv · kv_bpt_per_PE)` — check the Memory Breakdown tab.\n" f"- **Good S_kv**: stay ≤ **L\\* = {_l_star:,.0f} tokens** to " "keep decode compute-bound (peak utilization). Past L\\*, " f"utilization = 1/(1 + S_kv/L\\*). At S_kv = {_rf_skv:,}, " f"you're at ~**{_rec_l.utilization_at*100:.1f}%** utilization." ) st.divider() # ── Interpretation table ─────────────────────────────────────── st.markdown("**Formulas & interpretation**") _C = _scaled_machine.peak_flops # FLOPs / s (from knobs) _W = _scaled_machine.bw_hbm # bytes / s (from knobs) _b_elem = model.bytes_per_elem # bytes / weight _kv_bpt_full = 2 * model.h_kv * model.d_head * model.bytes_per_elem * model.layers _knee_now = knee_batch(_scaled_machine, model, _rf_skv) if _knee_now is None: _sub_bknee = ( f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = " f"{_b_star:.0f} / (1 − {_rf_skv/_l_star:.3f}) [≤ 0]" ) _bknee_val = "no knee (S_kv ≥ L*)" else: _sub_bknee = ( f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = " f"{_b_star:.0f} / {1 - _rf_skv/_l_star:.3f} = " f"{_knee_now:.0f}" ) _bknee_val = f"{_knee_now:.0f}" _formula_rows = [ {"Symbol": "AI", "Formula": "C / W", "With numbers": f"{_C:.2e} / {_W:.2e}", "Meaning": "FLOPs per byte of HBM BW. Chip-only.", "Value": f"{_ai:.2f} FLOPs/byte"}, {"Symbol": "B*", "Formula": "C · b / (2 · W)", "With numbers": f"{_C:.2e} · {_b_elem} / (2 · {_W:.2e})", "Meaning": ("Batch where weight-fetch = compute. Below: " "memory-bound. Above: amortized."), "Value": f"{_b_star:.0f}"}, {"Symbol": "B*_moe", "Formula": "B* · (N_total / N_active)", "With numbers": (f"{_b_star:.0f} · 8 (k=8 example)" if _is_moe else "N/A (dense preset)"), "Meaning": "MoE fetches all experts; computes on active only.", "Value": (f"{_b_star * 8:.0f}" if _is_moe else "1× (dense)")}, {"Symbol": "L*", "Formula": "2 · N_active · W / (C · kv_bpt)", "With numbers": (f"2 · {_n_active:.2e} · {_W:.2e} / " f"({_C:.2e} · {_kv_bpt_full:,})"), "Meaning": ("Context where KV-read = compute. Above: " "KV-bound regardless of B."), "Value": f"{_l_star:,.0f} tokens"}, {"Symbol": "B_knee", "Formula": "B* / (1 − S_kv/L*)", "With numbers": _sub_bknee, "Meaning": ("Batch where cost curve bends onto its floor. " "Diverges at S_kv = L*."), "Value": _bknee_val}, ] st.dataframe(pd.DataFrame(_formula_rows), width='stretch', hide_index=True) # ── TAB 8: Capacity planning ───────────────────────────────────── with tab_planning: st.subheader("Capacity planning & standards") st.caption( "How hyperscalers decide GPU count from a workload, what SLO " "targets they aim for, and what deployment patterns they use. " "All numbers below use the sidebar model + the Chip Roofline " "tab's effective chip (from its FLOPs/BW knobs, or sidebar " "defaults if you haven't opened that tab)." ) # ── Section 1: three-axis sizing formula + calculator ───────── st.markdown("### 1. GPU count = max(A, B, C) × N_replicas") # Symbol legend — spell out every letter that appears in the formulas # below with a plain-English meaning + the value for the current # sidebar model + chip. Makes the terse formulas readable at a glance. _n_active_here = total_active_params(model) _hbm_pe_bytes_here = int(_default_machine.pe_hbm_gb * 1e9) _kv_bpt_here = 2 * model.h_kv * model.d_head * model.bytes_per_elem * model.layers _weight_bytes_here = _n_active_here * model.bytes_per_elem st.markdown("**Symbol legend** — what each letter in the formulas means") _legend_rows = [ {"Symbol": "N", "Plain English": "Total number of model parameters (attention " "weights + FFN weights, across ALL layers)", "Value now": f"{_n_active_here/1e9:.2f} billion params"}, {"Symbol": "b", "Plain English": "Bytes per parameter (2 for BF16, 1 for FP8, " "0.5 for INT4)", "Value now": f"{model.bytes_per_elem} bytes (BF16)"}, {"Symbol": "N · b", "Plain English": "Total bytes needed to store one full copy of " "the model weights", "Value now": f"{_weight_bytes_here/1e9:.2f} GB"}, {"Symbol": "HBM_per_PE", "Plain English": "Amount of high-bandwidth memory on ONE chip " "(GPU / PE)", "Value now": f"{_default_machine.pe_hbm_gb:.1f} GB"}, {"Symbol": "S_kv", "Plain English": "Context length in tokens (how much history " "one user's KV cache holds)", "Value now": f"{s_kv:,} tokens (sidebar)"}, {"Symbol": "kv_bpt", "Plain English": "KV cache bytes per token per user, summed " "across all layers " "(= 2 · H_kv · d_head · b · layers)", "Value now": f"{_kv_bpt_here:,} B/token " f"({_kv_bpt_here/1024:.0f} KB/token)"}, {"Symbol": "users_per_replica", "Plain English": "How many concurrent users one replica serves " "(= per-replica batch size)", "Value now": "(from calculator below)"}, {"Symbol": "TPOT SLO", "Plain English": "Latency budget per output token (Time Per " "Output Token — the SLA the operator commits to)", "Value now": "(from calculator below)"}, {"Symbol": "step_latency", "Plain English": "Time for one decode forward pass at a given B " "(weight fetch + compute + KV read)", "Value now": "See Chip Roofline tab"}, {"Symbol": "B_at_SLO", "Plain English": "Largest batch size that still finishes one " "decode step within the SLO budget", "Value now": "(from calculator below)"}, {"Symbol": "n_users", "Plain English": "Total concurrent users you need to serve " "across the whole deployment", "Value now": "(from calculator below)"}, {"Symbol": "N_replicas", "Plain English": "Number of complete model copies deployed " "(each is a self-contained instance)", "Value now": "(from calculator below)"}, {"Symbol": "⌈ x ⌉", "Plain English": "'Ceiling' — round UP to the next whole " "number (you can't have 3.2 GPUs)", "Value now": "e.g. ⌈ 3.2 ⌉ = 4"}, ] st.dataframe(pd.DataFrame(_legend_rows), width='stretch', hide_index=True) st.markdown("**The three axes** — plain-English formulas") _axes_rows = [ {"Axis": "A. Capacity floor", "In symbols": "⌈ N · b / HBM_per_PE ⌉", "In plain English": "⌈ (weight bytes) ÷ (HBM per chip) ⌉", "What it says": "You need at least this many chips just to " "hold ONE copy of the weights. Below this, " "the model literally doesn't fit.", "Value now": f"⌈ {_weight_bytes_here/1e9:.2f} GB / " f"{_default_machine.pe_hbm_gb:.1f} GB ⌉ = " f"⌈ {_weight_bytes_here/_hbm_pe_bytes_here:.2f} ⌉ = " f"**{-(-_weight_bytes_here // _hbm_pe_bytes_here)} PEs**"}, {"Axis": "B. KV headroom", "In symbols": "⌈ (N·b + users_per_replica · S_kv · kv_bpt) / HBM_per_PE ⌉", "In plain English": "⌈ (weight bytes + all users' KV cache) ÷ (HBM per chip) ⌉", "What it says": "Weights ALSO need room for every user's KV " "history. As users × context grows, KV eats " "the remaining HBM until you need more chips.", "Value now": "(depends on users + context — see calculator)"}, {"Axis": "C. Throughput SLO", "In symbols": "N_replicas = ⌈ n_users / B_at_SLO ⌉", "In plain English": "Replicas = ⌈ total users ÷ users each replica handles " "within the latency budget ⌉", "What it says": "Even if capacity fits, a single replica can " "only serve so many users before per-token " "latency blows the SLO. More users → more " "replicas.", "Value now": "(depends on SLO — see calculator)"}, ] st.dataframe(pd.DataFrame(_axes_rows), width='stretch', hide_index=True) st.markdown("#### Live calculator") _cp_a, _cp_b, _cp_c = st.columns(3) with _cp_a: _cp_users = st.number_input( "Concurrent users", min_value=1, max_value=100_000, value=100, step=1, key="_cp_n_users", help="Target simultaneous active decode sequences.", ) with _cp_b: _cp_ctx = st.number_input( "Avg context / user (tokens)", min_value=1, max_value=2_000_000, value=int(s_kv), step=128, key="_cp_avg_ctx", help="Mean S_kv across the active users.", ) with _cp_c: _cp_slo_ms = st.number_input( "TPOT SLO (ms/token)", min_value=1, max_value=1000, value=30, step=1, key="_cp_tpot_slo_ms", help=("Per-token latency target during decode. " "Tighter = smaller B per replica, more replicas."), ) # Use the base sidebar chip for the calculator so it's stable. _cp_machine = _default_machine _sizing = size_deployment( _cp_machine, model, n_users=int(_cp_users), avg_ctx=int(_cp_ctx), tpot_slo_s=_cp_slo_ms / 1000.0, ) _s1, _s2, _s3, _s4 = st.columns(4) _s1.metric("A. Capacity", f"{_sizing.pes_axis_a_capacity} PEs", help="Min PEs to hold weights alone.") _s2.metric("B. KV headroom", f"{_sizing.pes_axis_b_kv} PEs", help="Min PEs to hold weights + all replica KV.") _s3.metric("C. Throughput", (f"{_sizing.n_replicas} replicas" if _sizing.b_at_slo > 0 else "SLO INFEASIBLE"), help=(f"B_at_SLO = {_sizing.b_at_slo}; ceil(" f"{_cp_users}/{_sizing.users_per_replica}) " f"= {_sizing.n_replicas} replicas.")) _s4.metric("Total GPUs", f"{_sizing.total_pes}", delta=f"binds: {_sizing.binding_axis}", help="max(A, B) × N_replicas.") if _sizing.b_at_slo == 0: st.error( f"SLO infeasible: even B=1 step latency exceeds " f"{_cp_slo_ms} ms at S_kv = {_cp_ctx:,}. Loosen the SLO, " f"shorten context, or use a faster chip." ) else: _binding_hint = { "capacity": ("Weights dominate. Adding replicas or bigger-HBM " "chips buys the most headroom."), "kv": ("KV cache dominates. Shorter context / smaller B / " "more aggressive CP sharding help most."), "throughput": (f"Latency SLO forces more replicas — one " f"replica serves ~{_sizing.b_at_slo} users. " "Loosening TPOT or faster chip cuts count."), }[_sizing.binding_axis] st.caption( f"**Binding axis: {_sizing.binding_axis}.** {_binding_hint} " f"B_at_SLO = **{_sizing.b_at_slo}** / replica × " f"**{_sizing.n_replicas}** replicas = " f"**{_sizing.total_pes} GPUs**." ) st.divider() # ── Section 2: SLO targets by workload ───────────────────────── st.markdown("### 2. SLO standards — TTFT & TPOT by workload") st.caption( "TTFT = Time To First Token (dominated by prefill of the prompt). " "TPOT = Time Per Output Token (one decode step per token). " "Real workload targets vary widely; the table is a rough guide." ) _slo_rows = [ {"Use case": "Voice / realtime", "TTFT target": "200–300 ms (whole pipeline)", "TPOT / ITL target": "10–25 ms", "What binds": "TPOT — smooth speech needs tight per-token cadence"}, {"Use case": "Code completion", "TTFT target": "< 100 ms", "TPOT / ITL target": "mostly n/a (short response)", "What binds": "TTFT — user cursor is waiting"}, {"Use case": "Interactive chat", "TTFT target": "300 ms – 1 s", "TPOT / ITL target": "20–50 ms", "What binds": "Balanced; both matter to perceived latency"}, {"Use case": "Agentic / multi-step", "TTFT target": "as low as possible — compounds across steps", "TPOT / ITL target": "maximize throughput", "What binds": "Throughput — many sequential LLM calls"}, {"Use case": "Batch / offline", "TTFT target": "seconds to minutes", "TPOT / ITL target": "irrelevant", "What binds": "Cost/token — pack B to the ceiling"}, ] st.dataframe(pd.DataFrame(_slo_rows), width='stretch', hide_index=True) st.divider() # ── Section 3: Hybrid deployment layers (table form) ────────── st.markdown("### 3. Same setup or different for short vs long context?") st.caption( "**Hybrid: one elastic pool + length-tier routing on top + " "prefill/decode disaggregation at the frontier.**" ) _layers_rows = [ {"Layer": "1. Elastic pool", "Serves": "Mixed traffic (~90% of requests, chat/RAG/code)", "Key mechanism": "PagedAttention (KV in fixed-size pages) + " "continuous batching", "Admission rule": "Scheduler checks remaining KV page pool per " "iteration; admits if user's expected " "max_tokens fits", "Typical config": "Standard replica (e.g. TP=8, CP=1..4)"}, {"Layer": "2. Length-tier routing", "Serves": "The long-context tail (128k–1M+)", "Key mechanism": "API gateway inspects max context, dispatches " "to a dedicated pool", "Admission rule": "Standard tier for < 32-128k; long-context " "tier otherwise", "Typical config": ("Long-context replicas have more chips per " "instance + heavier CP sharding, lower B " "(e.g. TP=8, CP=32)")}, {"Layer": "3. Disaggregated prefill/decode", "Serves": "Frontier deployments (DistServe, Splitwise pattern)", "Key mechanism": "Prefill on FLOPs-heavy pool; decode on HBM-BW-" "heavy pool; KV moves over NVLink/RDMA", "Admission rule": "Prefill goes to a prefill GPU; KV then " "handed off to decode pool", "Typical config": ("Two GPU pools with different chip mixes " "optimized for prefill vs decode " "rooflines")}, ] st.dataframe(pd.DataFrame(_layers_rows), width='stretch', hide_index=True) st.divider() # ── Section 4: Practical rules of thumb by context regime ───── st.markdown("### 4. Practical rules of thumb (by context regime)") _rules_rows = [ {"Regime": "Short context (S_kv < L*)", "Batch strategy": "Pack as many users as HBM allows; aim for " "B ≈ 2·B*", "Utilization": "High (near compute floor)", "Cost / token": "Low — cheap tier", "Deployment": "Standard elastic pool"}, {"Regime": "Long context (S_kv > L*)", "Batch strategy": "Fewer users per replica; more HBM per user; " "lower B", "Utilization": "Drops as 1 / (1 + S_kv/L*)", "Cost / token": "Higher — accept lower utilization", "Deployment": "Length-tier replica with more CP sharding"}, {"Regime": "Extreme long context (S_kv ≫ L*)", "Batch strategy": "Dedicated pool, heavy CP; often " "disaggregated prefill", "Utilization": "Very low without sparse attention / MLA", "Cost / token": "Much higher — priced accordingly", "Deployment": "Long-context tier + disaggregated prefill/decode"}, ] st.dataframe(pd.DataFrame(_rules_rows), width='stretch', hide_index=True) st.divider() # ── Section 5: Sample deployment templates ──────────────────── st.markdown("### 5. Sample deployment templates (same model, different sharding)") st.caption( "One base model runs on all three configs. What changes is the " "sharding (CP/TP/PP) that shapes the effective 'small' vs " "'large' setup. The API gateway routes each request to the " "right tier based on its context length." ) _tmpl_rows = [ {"Tier": "Config_small", "CP": 1, "TP": 8, "PP": 1, "Total GPUs / replica": 8, "Max context": "up to 32k", "Typical B": 64, "Best for": "Chat, RAG, short prompts — high utilization"}, {"Tier": "Config_medium", "CP": 4, "TP": 8, "PP": 1, "Total GPUs / replica": 32, "Max context": "up to 128k", "Typical B": 32, "Best for": "Long-doc analysis, code review — standard tier"}, {"Tier": "Config_large", "CP": 32, "TP": 8, "PP": 1, "Total GPUs / replica": 256, "Max context": "up to 1M", "Typical B": 4, "Best for": "Whole-repo / long-video / massive-doc — premium tier"}, ] st.dataframe(pd.DataFrame(_tmpl_rows), width='stretch', hide_index=True) st.divider() # ── Section 6: Real-world provider examples ─────────────────── st.markdown("### 6. Actually-provisioned batch sizes (public examples)") st.caption( "None of these are officially published. The numbers are " "**backed out from public pricing** (long-context tiers cost " "2–5× per token vs standard) and stated latency SLAs. Treat " "as ballpark, not authoritative." ) _real_rows = [ {"Provider / tier": "Gemini 1.5 Pro (1M)", "Approx B / replica": "1–4", "Signal": "Pricing ~2× standard suggests low B"}, {"Provider / tier": "Gemini 1.5 Flash (1M)", "Approx B / replica": "2–8", "Signal": "Smaller model → more HBM headroom for KV"}, {"Provider / tier": "GPT-4.1 (1M)", "Approx B / replica": "1–4", "Signal": "Similar pricing structure to Gemini Pro"}, {"Provider / tier": "Claude 3.5 Sonnet (200k)", "Approx B / replica": "4–16", "Signal": "Shorter context = more headroom, higher B feasible"}, {"Provider / tier": "Enterprise 'dedicated' tiers", "Approx B / replica": "1", "Signal": "Whole replica per customer for latency guarantees"}, ] st.dataframe(pd.DataFrame(_real_rows), width='stretch', hide_index=True) st.divider() # ── Section 7: Binding-axis playbook ────────────────────────── st.markdown("### 7. What to do when each axis binds") _playbook_rows = [ {"Binding axis": "A. Capacity (weights)", "Symptom": "Weights alone are close to per-PE HBM budget", "First lever": "Increase TP or PP to shard weights across more PEs", "Second lever": "Move to a bigger-HBM chip; use lower-precision " "weights (FP8/INT4)", "Cost impact": "Sub-linear — you pay for more chips, but " "utilization stays high"}, {"Binding axis": "B. KV headroom", "Symptom": "Weights fit fine but users × context blows HBM", "First lever": "Increase CP (shard sequence dim), or reduce " "batch (accept fewer users)", "Second lever": "GQA / MQA / MLA to shrink kv_bpt; sparse " "attention; KV compression (INT4 KV)", "Cost impact": "Direct — long-context requests fundamentally " "cost more"}, {"Binding axis": "C. Throughput SLO", "Symptom": "Latency budget forces small B per replica", "First lever": "Add more replicas (linear scale)", "Second lever": "Loosen SLO if product allows; faster chip; " "reduce model size; speculative decoding", "Cost impact": "Linear in replica count"}, ] st.dataframe(pd.DataFrame(_playbook_rows), width='stretch', hide_index=True) # ── TAB 9: Parallelism rules ───────────────────────────────────── with tab_prules: st.subheader("Parallelism rules — heuristics, constraints, and honest caveats") st.info( "**This is directional guidance, not a production procedure.** " "Real hyperscalers combine the rules below with measured " "collectives, workload traces, and per-architecture search " "(GSPMD, Alpa, hand-tuned Megatron). Public disclosures " "(DeepSeek-V3 training: PP16 + EP64 + no TP; vLLM's " "PP-over-TP recommendation on non-NVLink GPUs) show that no " "universal 'start TP at 8' rule survives contact with real " "workloads. Use these tables as a starting point, then benchmark." ) # ── Section 1: core principle ───────────────────────────────── st.markdown("### 1. The core principle") st.markdown( "> **Rank the techniques by how much they communicate per unit " "of compute, and assign the chattiest ones to the fastest links.**\n\n" "Memory capacity is the feasibility filter (you fit or you don't). " "Once you fit, every remaining choice is about minimizing " "*exposed* communication (comm that couldn't be overlapped with " "compute)." ) st.divider() # ── Section 2: Parallelism technique comparison ────────────── st.markdown("### 2. Parallelism techniques — comm profile at a glance") _p_compare = [ {"Technique": "DP (Data Parallelism)", "Shards": "Nothing (or optimizer state w/ ZeRO)", "Comm pattern": "AllReduce (gradients)", "Frequency": "1× per training step", "Per-link vol as degree ↑": "Flat, but overlappable", "Hard ceiling": "Critical batch size"}, {"Technique": "TP (Tensor Parallelism)", "Shards": "Weights + activations + KV by head", "Comm pattern": "AllReduce", "Frequency": "2× per layer", "Per-link vol as degree ↑": "Flat (compute shrinks, comm doesn't)", "Hard ceiling": "≤ NVLink domain size; efficient ≤ n_kv_heads"}, {"Technique": "PP (Pipeline Parallelism)", "Shards": "Weights by layer + KV by layer", "Comm pattern": "Point-to-point (activation hand-off)", "Frequency": "Per stage boundary", "Per-link vol as degree ↑": "Negligible per link", "Hard ceiling": "n_layers"}, {"Technique": "CP (Context Parallelism)", "Shards": "Activations + KV by position", "Comm pattern": "Ring P2P (per hop)", "Frequency": "Per layer, overlappable", "Per-link vol as degree ↑": "Falls per hop (KV shard shrinks)", "Hard ceiling": "seq_len / block_size"}, {"Technique": "EP (Expert Parallelism, MoE only)", "Shards": "Expert weights", "Comm pattern": "All-to-all (dispatch + combine)", "Frequency": "2× per MoE layer", "Per-link vol as degree ↑": "Falls per peer (total flat)", "Hard ceiling": "n_experts"}, ] st.dataframe(pd.DataFrame(_p_compare), width='stretch', hide_index=True) st.caption( "The TP row is the one that decides most configs. Its " "AllReduce **volume doesn't shrink** as you add ranks — only " "compute does. That's why TP has a practical ceiling (~8) " "even without hitting a hard n_kv_heads limit." ) st.divider() # ── Section 3: dominant problem → first axis to try ────────── st.markdown("### 3. What's hurting → which parallelism to reach for") _p_problem = [ {"Dominant problem": "Weight memory doesn't fit", "First to try": "TP (within a fast domain)", "Second lever": "PP across nodes; FSDP/ZeRO in training", "Why": "TP shards weights + activations; PP shards by depth"}, {"Dominant problem": "Optimizer / gradient memory (training)", "First to try": "FSDP / ZeRO-2 or ZeRO-3", "Second lever": "Increase DP degree", "Why": "ZeRO shards optimizer state across DP ranks"}, {"Dominant problem": "Activation memory (long sequences)", "First to try": "CP / sequence parallelism", "Second lever": "Activation checkpointing", "Why": "Activations scale with seq_len; CP shards positions"}, {"Dominant problem": "KV cache of ONE long sequence", "First to try": "CP", "Second lever": "TP over KV heads; PP by layer", "Why": "DP can't split one sequence — CP shards positions"}, {"Dominant problem": "KV capacity for MANY sequences", "First to try": "DP replicas + request scheduling", "Second lever": "Attention-DP within EP groups", "Why": "Independent sequences shard by DP for free"}, {"Dominant problem": "MoE expert weight memory", "First to try": "EP", "Second lever": "Redundant copies of hot experts", "Why": "EP places different experts on different ranks"}, {"Dominant problem": "Single-request latency (TPOT)", "First to try": "TP (up to ~8)", "Second lever": "CP for very long prefill", "Why": "TP divides weight-fetch time linearly"}, {"Dominant problem": "Aggregate throughput (many users)", "First to try": "Outer DP", "Second lever": "Higher B per replica", "Why": "DP scales throughput linearly at zero comm cost"}, {"Dominant problem": "Model spans multiple nodes", "First to try": "TP within node, PP across nodes", "Second lever": "Consider CP if seq_len is also large", "Why": "Slow scale-out links favor PP's point-to-point"}, {"Dominant problem": "Expert layer dominates (MoE)", "First to try": "EP", "Second lever": "Attention-DP + EP combined", "Why": "vLLM supports DP-attn + EP-experts topology"}, {"Dominant problem": "Global batch already at critical size (training)", "First to try": "TP × PP × CP × EP", "Second lever": "Reduce DP further", "Why": "Additional DP past critical batch wastes samples"}, ] st.dataframe(pd.DataFrame(_p_problem), width='stretch', hide_index=True) st.divider() # ── Section 4: per-axis select/stop criteria ───────────────── st.markdown("### 4. When to add / when to stop — per axis") _p_criteria = [ {"Axis": "DP", "Add when": "More independent batches/requests are available", "Stop when": "Global batch exceeds critical batch, or replicas " "can't fit"}, {"Axis": "TP", "Add when": "Layers/weights need sharding or TPOT needs improvement", "Stop when": "Collectives dominate step time OR local GEMMs " "become too small OR TP > n_kv_heads (efficiency)"}, {"Axis": "PP", "Add when": "Model depth must be distributed, or TP would cross " "slow links", "Stop when": "Bubble fraction > ~10% OR stage imbalance grows"}, {"Axis": "CP", "Add when": "A single sequence is long enough to need position " "sharding", "Stop when": "Ring hops can't be hidden by compute OR local " "sequence blocks are too small for efficient kernels"}, {"Axis": "EP", "Add when": "MoE model + expert weights should be distributed", "Stop when": "Expert GEMMs become too small OR routing " "imbalance grows OR all-to-all latency dominates"}, {"Axis": "FSDP / ZeRO", "Add when": "Training-state memory is the dominant problem", "Stop when": "Parameter all-gather / reduce-scatter starts " "dominating step time"}, {"Axis": "Sequence parallelism", "Add when": "TP is used during training (typically bundled)", "Stop when": "Rarely disabled alone — coupled to TP support"}, ] st.dataframe(pd.DataFrame(_p_criteria), width='stretch', hide_index=True) st.divider() # ── Section 5: naive rules vs nuanced reality ──────────────── st.markdown("### 5. Common misconceptions vs reality") _p_miscon = [ {"Naive rule": "Memory is only a feasibility filter", "Reality": "Memory headroom is also a continuous performance " "variable. More free HBM = larger B, more prefix " "cache, higher throughput — even after weights fit."}, {"Naive rule": "TP has a hard ceiling at n_kv_heads", "Reality": "vLLM (and others) support KV-head replication when " "TP > n_kv_heads. It's an efficiency preference, " "not a correctness ceiling."}, {"Naive rule": "Always start TP at 8", "Reality": "Start with the smallest TP that fits memory + " "meets latency. DeepSeek-V3 inference uses TP=4 " "(and TP=1 for some MLPs) explicitly to limit " "collective overhead."}, {"Naive rule": "EP comm falls as EP grows", "Reality": "Per-peer traffic falls, total outgoing traffic is " "flat (~N·k·H). Large EP hurts via startup overhead, " "small local GEMMs, and imbalance."}, {"Naive rule": "Attention-DP shards KV 'for free'", "Reality": "Only across DIFFERENT sequences. One 1M-token " "sequence still needs CP/TP/PP to shard — DP can't " "split it. Attention weights are still replicated " "across DP ranks."}, {"Naive rule": "DP > CP always if you have many requests", "Reality": "For MANY short requests, yes. For ONE long " "sequence with strict prefill SLO, CP wins because " "DP can't help. Different dimensions solve " "different problems."}, {"Naive rule": "Avoid PP inside the domain — it's worthless", "Reality": "vLLM specifically recommends PP over TP on GPUs " "without NVLink. DeepSeek-V3 training used PP=16 " "and no TP. PP shines when TP collectives cost more " "than PP bubbles."}, {"Naive rule": "FSDP/ZeRO-3 is THE training substrate", "Reality": "DeepSeek-V3 used ZeRO-1 + PP16 + EP64 — no ZeRO-3. " "FSDP is one option; the right choice depends on " "the topology, batch size, and model shape."}, {"Naive rule": "Pipeline bubble = (P−1)/m", "Reality": "Actual bubble fraction = (P−1)/(m + P−1). The " "(P−1)/m form is an approximation valid only when " "m ≫ P. Interleaved / 1F1B / zero-bubble schedules " "have their own formulas."}, ] st.dataframe(pd.DataFrame(_p_miscon), width='stretch', hide_index=True) st.divider() # ── Section 6: inference decision procedure ────────────────── st.markdown("### 6. Inference decision procedure (in order)") _p_infer = [ {"Step": 1, "Action": "Feasibility check", "Criterion": "Compute weights + KV(B, S_kv). If it fits on " "one GPU, stop — use plain replication (DP only)."}, {"Step": 2, "Action": "Set TP from TPOT budget", "Criterion": "TP is the latency knob — divides weight-fetch " "time. Prefer TP=4–8 for latency-oriented; TP=1–2 " "for throughput-oriented. Try to keep TP ≤ n_kv_heads."}, {"Step": 3, "Action": "Fill remaining domain with EP (if MoE)", "Criterion": "EP is the cost knob for MoE. Cap at n_experts. " "Overlap dispatch/combine with a second micro-batch."}, {"Step": 4, "Action": "Attention-DP over EP groups", "Criterion": "Different attention replicas handle different " "sequence subsets → KV shards by sequence for " "free, no extra collective."}, {"Step": 5, "Action": "CP only if attention-DP can't help", "Criterion": "Trigger: one very long sequence or per-request " "KV exceeds one GPU. For many short concurrent " "requests, DP strictly wins."}, {"Step": 6, "Action": "PP across racks if the model still won't fit", "Criterion": "Only across slow scale-out boundaries. Inside a " "fast NVLink domain, PP is usually inferior to TP."}, {"Step": 7, "Action": "Replicate the whole unit as outer DP", "Criterion": "The throughput axis. Cap by traffic + critical " "batch (training) or by aggregate request " "arrival rate (inference)."}, ] st.dataframe(pd.DataFrame(_p_infer), width='stretch', hide_index=True) st.divider() # ── Section 7: concrete inference cases ────────────────────── st.markdown("### 7. Inference cases → recommended shape") _p_cases = [ {"Case": "A. Fits on one GPU (with KV headroom)", "Recommended": "TP=1, PP=1, CP=1; scale via outer DP", "Notes": "Benchmark TP=2/4 if TPOT is tight"}, {"Case": "B. Fits within one NVLink node", "Recommended": "TP = 2/4/8 (smallest that fits + meets SLO)", "Notes": "Consider (TP=2, PP=2) for very deep models"}, {"Case": "C. Spans multiple nodes", "Recommended": "TP = GPUs per fast domain, PP = " "number of nodes; outer DP for throughput", "Notes": "vLLM's standard multi-node recipe"}, {"Case": "D. One request with extreme context (≫ L*)", "Recommended": "TP × CP × PP (DP can't split one sequence)", "Notes": "Long-context tier; heavy CP dominates"}, {"Case": "E. Many moderate-length requests", "Recommended": "Minimum TP × PP per replica; maximize DP = " "total_GPUs / (TP × PP)", "Notes": "Throughput-optimized; the mainline serving case"}, ] st.dataframe(pd.DataFrame(_p_cases), width='stretch', hide_index=True) st.divider() # ── Section 8: training baseline recipe ────────────────────── st.markdown("### 8. Training baseline recipe (in order)") _p_train = [ {"Step": 1, "Action": "Start with DP (or sharded DP / FSDP)", "Rationale": "Simplest; shards optimizer state; scales at " "throughput cost only"}, {"Step": 2, "Action": "Add TP if large layers don't fit or GEMMs stay " "efficient after sharding", "Rationale": "TP divides weight memory + activation memory"}, {"Step": 3, "Action": "Add CP when sequence activations dominate", "Rationale": "Activation memory ~ seq_len × d; CP shards " "positions"}, {"Step": 4, "Action": "Add PP when depth still doesn't fit or scaling " "across nodes", "Rationale": "Depth sharding + slow-link tolerance"}, {"Step": 5, "Action": "Add EP for MoE expert weights", "Rationale": "Distributes expert weights; preserves tokens " "per expert"}, {"Step": 6, "Action": "Use remaining devices as outer DP", "Rationale": "Global batch = DP × micro-batch × grad-accum; " "cap by critical batch"}, ] st.dataframe(pd.DataFrame(_p_train), width='stretch', hide_index=True) st.caption( "Total GPUs = **DP × TP × PP × CP × EP**. Global batch = " "**DP × micro-batch × grad-accum**. The critical batch size " "(NOT bandwidth) is what forces model parallelism at very " "large scale — past that, additional DP wastes samples." ) st.divider() # ── Section 9: sanity checks for any config ────────────────── st.markdown("### 9. Sanity checks for any candidate config") _p_sanity = [ {"Check": "TP stays inside one NVLink / scale-up domain?", "Pass criterion": "Yes", "If not": "TP AllReduce crosses slow links → step time balloons"}, {"Check": "TP ≤ n_kv_heads?", "Pass criterion": "Yes (or KV replication is acceptable)", "If not": "Wastes KV BW to replication; usually not worth it"}, {"Check": "Any KV or weight state duplicated inside the " "scale-up domain?", "Pass criterion": "No", "If not": "You're wasting aggregate HBM bandwidth"}, {"Check": "Exposed (non-overlapped) comm fraction?", "Pass criterion": "≤ ~10% of step time", "If not": "Something is misplaced — swap axes across links"}, {"Check": "Pipeline bubble fraction (P−1)/(m+P−1)?", "Pass criterion": "≤ ~10%", "If not": "Need more microbatches or fewer stages"}, {"Check": "Predicted step time meets TPOT SLO?", "Pass criterion": "Yes", "If not": "Reduce B, add DP replicas, or accept slower SLO"}, {"Check": "B sits in the 2–3× B* sweet spot?", "Pass criterion": "Yes", "If not": "Below → underused compute; above → wasted latency"}, ] st.dataframe(pd.DataFrame(_p_sanity), width='stretch', hide_index=True) st.divider() # ── Section 10: honest caveats ──────────────────────────────── st.markdown("### 10. One-paragraph rule + caveats") st.markdown( "> **First use the smallest parallelism group that makes the " "model, training state, activations, and target KV capacity " "fit. Use TP mainly inside the fastest scale-up domain, PP to " "divide model depth — especially across slower boundaries — CP " "to divide an individual long sequence, and EP to distribute " "MoE experts while preserving enough tokens per expert. Use DP " "across independent batches or requests for throughput. Then " "benchmark several legal combinations using real kernel " "shapes, collective measurements, workload distributions, and " "end-to-end SLOs; choose based on exposed communication and " "cost per token, not communication volume alone.**" ) st.warning( "**Honest caveat.** The specific configurations frontier labs " "run are mostly unpublished, and they retune per architecture " "+ hardware generation. What's stable is the *ordering " "principle* and the *constraint list* — those follow from the " "hardware, not from anyone's cleverness. Public counterexamples " "(DeepSeek-V3, vLLM PP-on-non-NVLink) show that no single " "recipe fits all workloads." )