mukesh
9fdde44922
New tests/analytical_visualization/ module - an interactive dashboard for exploring memory / latency tradeoffs of transformer inference on the SIP architecture. Highlights: - 30+ model presets (Qwen, Llama 2/3/3.1, Mistral, Gemma 2, Phi 3, DeepSeek MLA, Mixtral/Qwen 3 MoE, Grok-1, ...) - Placement toggles: TP and CP each on PE-level vs cube-level - CP ring variant: K/V ring vs Q+O/m/l ring (prefill); in decode the O/m/l all-reduce is folded into S8 (no separate C1 row) - SIP interconnect: ring / mesh2d / torus2d with matching link drawing - Per-stage latency table with compute + memory + comm formulas, auto-scaled ns/us/ms, colored by dominant bound - Ring attention loop indicator on the pipeline diagram (purple arc over S5-S8 with 'xN hops' badge) - Tensor sharding view with optional physical PE/cube annotations - Replication-waste + optimization-hints panel - Save & compare configurations (config1, config2, ...): summary table plus side-by-side per-stage attention and FFN latency, best-in-row highlighting - Symbol glossary with current values for every symbol used in formulas Not tied to production sim_engine or runtime API; purely analytical tooling for design-space exploration. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1285 lines
57 KiB
Python
1285 lines
57 KiB
Python
"""Streamlit app: interactive analytical model for transformer inference
|
|
on the SIP architecture.
|
|
|
|
Run:
|
|
python -m streamlit run tests/analytical_visualization/app.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import matplotlib.pyplot as plt
|
|
import pandas as pd
|
|
import streamlit as st
|
|
|
|
from tests.analytical_visualization.model_config import (
|
|
FullConfig, ModelConfig, TopologyConfig, MachineParams,
|
|
)
|
|
from tests.analytical_visualization.model_presets import PRESETS
|
|
from tests.analytical_visualization.stage_latencies import all_stages, all_ffn_stages
|
|
from tests.analytical_visualization.memory_layout import (
|
|
compute_memory,
|
|
attention_weight_rows,
|
|
ffn_weight_rows,
|
|
kv_cache_rows,
|
|
sum_bytes_all_layers,
|
|
total_weight_bytes_full_model,
|
|
total_kv_bytes_full_model,
|
|
)
|
|
from tests.analytical_visualization.topology_map import draw_topology
|
|
from tests.analytical_visualization.pe_weight_layout import (
|
|
draw_pe_layout, _per_pe_bytes, _q_heads_for_pe, _kv_heads_for_pe,
|
|
)
|
|
from tests.analytical_visualization.tensor_sharding import draw_tensor_sharding
|
|
from tests.analytical_visualization.pipeline_diagram import draw_pipeline
|
|
from tests.analytical_visualization.optimization_report import (
|
|
replication_report, optimization_hints,
|
|
)
|
|
from tests.analytical_visualization.autosuggest import auto_suggest
|
|
|
|
|
|
st.set_page_config(page_title="SIP Attention Analytical Model",
|
|
layout="wide")
|
|
|
|
# Tighter spacing so more info fits per page.
|
|
st.markdown(
|
|
"""
|
|
<style>
|
|
.block-container { padding-top: 1.2rem; padding-bottom: 1rem;
|
|
padding-left: 1.5rem; padding-right: 1.5rem; }
|
|
div[data-testid="stMetricValue"] { font-size: 1.05rem; }
|
|
div[data-testid="stMetricLabel"] { font-size: 0.75rem; }
|
|
h1 { font-size: 1.6rem !important; margin-bottom: 0.2rem; }
|
|
h2 { font-size: 1.15rem !important; margin-top: 0.4rem; }
|
|
h3 { font-size: 1.0rem !important; margin-top: 0.3rem; }
|
|
div[data-testid="stExpander"] summary p { font-size: 0.85rem !important; }
|
|
.streamlit-expanderContent { padding-top: 0.3rem !important; }
|
|
section[data-testid="stSidebar"] .stSelectbox label,
|
|
section[data-testid="stSidebar"] .stRadio label,
|
|
section[data-testid="stSidebar"] .stNumberInput label {
|
|
font-size: 0.78rem; }
|
|
</style>
|
|
""",
|
|
unsafe_allow_html=True,
|
|
)
|
|
|
|
|
|
def _pick_in(container, label: str, options, default, key: str, help: str = ""):
|
|
"""Selectbox helper — returns the selected value from options."""
|
|
return container.selectbox(
|
|
label, options,
|
|
index=options.index(default) if default in options else 0,
|
|
key=key, help=help,
|
|
)
|
|
|
|
|
|
def _pick(label, options, default, key, help=""):
|
|
return _pick_in(st.sidebar, label, options, default, key, help)
|
|
|
|
|
|
# ── Sidebar ─ compact, guided layout ────────────────────────────
|
|
st.sidebar.markdown("### SIP config")
|
|
st.sidebar.caption(
|
|
"Pick a model, workload, parallelism plan, then tune hardware. "
|
|
"Every change re-runs the analytical model live."
|
|
)
|
|
|
|
with st.sidebar.expander("Model", expanded=True):
|
|
preset_names = list(PRESETS.keys())
|
|
preset_name = st.selectbox(
|
|
"Preset", preset_names,
|
|
index=preset_names.index("Qwen 3 8B"),
|
|
key="preset_select",
|
|
label_visibility="collapsed",
|
|
)
|
|
preset = PRESETS[preset_name]
|
|
st.caption(
|
|
f"{preset.family or '-'} - {preset.attn_type} - "
|
|
f"H_q={preset.model.h_q}, H_kv={preset.model.h_kv}, "
|
|
f"L={preset.model.layers}"
|
|
)
|
|
with st.expander("Override dims", expanded=False):
|
|
d1, d2 = st.columns(2)
|
|
with d1:
|
|
hidden_ = st.number_input("hidden", 128, 65536,
|
|
value=preset.model.hidden, step=128)
|
|
h_q_ = st.number_input("H_q", 1, 256, value=preset.model.h_q)
|
|
d_head_ = st.number_input("d_head", 32, 512,
|
|
value=preset.model.d_head, step=32)
|
|
layers_ = st.number_input("layers", 1, 200,
|
|
value=preset.model.layers)
|
|
with d2:
|
|
ffn_ = st.number_input("ffn_dim", 128, 131072,
|
|
value=preset.model.ffn_dim, step=128)
|
|
h_kv_ = st.number_input("H_kv", 1, 256, value=preset.model.h_kv)
|
|
bpe_ = st.selectbox("dtype bytes", [1, 2, 4], index=1)
|
|
model = ModelConfig(
|
|
name=preset.model.name, hidden=hidden_, ffn_dim=ffn_,
|
|
h_q=h_q_, h_kv=h_kv_, d_head=d_head_, layers=layers_,
|
|
bytes_per_elem=bpe_,
|
|
)
|
|
|
|
|
|
with st.sidebar.expander("Workload", expanded=True):
|
|
w1, w2 = st.columns([2, 1])
|
|
with w1:
|
|
s_kv = st.selectbox(
|
|
"Context S_kv",
|
|
[1024, 2048, 4096, 8192, 16384, 32768, 65536,
|
|
131072, 262144, 524288, 1048576],
|
|
index=10, key="s_kv", # default 1M
|
|
)
|
|
with w2:
|
|
mode = st.selectbox("Mode", ["decode", "prefill"],
|
|
index=0, key="mode")
|
|
|
|
|
|
# Compute auto-suggest for defaults / apply button
|
|
_default_machine = MachineParams()
|
|
_default_suggestion = auto_suggest(model, _default_machine, s_kv, mode)
|
|
|
|
def _snap(v: int, opts: tuple) -> int:
|
|
return min(opts, key=lambda o: abs(o - v))
|
|
|
|
_CP_OPTS = (1, 2, 4, 8, 16, 32, 48, 64, 96)
|
|
_TP_OPTS = (1, 2, 4, 8, 16, 32)
|
|
_PP_OPTS = (1, 2, 4, 8, 16, 32)
|
|
_DP_OPTS = (1, 2, 4, 8, 16)
|
|
_EP_OPTS = (1, 2, 4, 8, 16, 32)
|
|
|
|
# Auto-reset dropdowns when model preset changes.
|
|
if st.session_state.get("_last_preset") != preset_name:
|
|
st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
|
|
st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
|
|
st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
|
|
st.session_state["dp"] = 1
|
|
st.session_state["ep"] = 1
|
|
st.session_state["_last_preset"] = preset_name
|
|
|
|
with st.sidebar.expander("Parallelism", expanded=True):
|
|
st.caption(
|
|
f"Auto-suggest: CP={_default_suggestion.cp}, "
|
|
f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp}"
|
|
)
|
|
if st.button("Apply auto-suggest", width='stretch'):
|
|
st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
|
|
st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
|
|
st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
|
|
st.rerun()
|
|
|
|
# Sharding dims — arrange in 2x3 grid to save vertical space
|
|
p_row1 = st.columns(3)
|
|
with p_row1[0]:
|
|
cp = st.selectbox(
|
|
"CP", list(_CP_OPTS),
|
|
index=_CP_OPTS.index(st.session_state.get(
|
|
"cp", _snap(_default_suggestion.cp, _CP_OPTS))),
|
|
key="cp",
|
|
)
|
|
with p_row1[1]:
|
|
tp = st.selectbox(
|
|
"TP", list(_TP_OPTS),
|
|
index=_TP_OPTS.index(st.session_state.get(
|
|
"tp", _snap(_default_suggestion.tp, _TP_OPTS))),
|
|
key="tp",
|
|
)
|
|
with p_row1[2]:
|
|
pp = st.selectbox(
|
|
"PP", list(_PP_OPTS),
|
|
index=_PP_OPTS.index(st.session_state.get(
|
|
"pp", _snap(_default_suggestion.pp, _PP_OPTS))),
|
|
key="pp",
|
|
)
|
|
p_row2 = st.columns(2)
|
|
with p_row2[0]:
|
|
dp = st.selectbox(
|
|
"DP", list(_DP_OPTS),
|
|
index=_DP_OPTS.index(st.session_state.get("dp", 1)),
|
|
key="dp",
|
|
)
|
|
with p_row2[1]:
|
|
ep = st.selectbox(
|
|
"EP (MoE)", list(_EP_OPTS),
|
|
index=_EP_OPTS.index(st.session_state.get("ep", 1)),
|
|
key="ep",
|
|
)
|
|
|
|
st.markdown("**Placement (PE-level vs cube-level)**")
|
|
pl_col1, pl_col2 = st.columns(2)
|
|
with pl_col1:
|
|
tp_placement = st.radio(
|
|
"TP on", options=["pe", "cube"], index=0,
|
|
key="tp_placement", horizontal=True,
|
|
help=("pe: TP shares PEs of one cube (intra-cube AllReduce). "
|
|
"cube: each TP rank owns a whole cube (inter-cube AllReduce)."),
|
|
)
|
|
with pl_col2:
|
|
cp_placement = st.radio(
|
|
"CP on", options=["pe", "cube"], index=1,
|
|
key="cp_placement", horizontal=True,
|
|
help=("pe: CP ring runs among PEs of one cube (intra-cube). "
|
|
"cube: each CP rank owns cubes (inter-cube ring)."),
|
|
)
|
|
|
|
st.markdown("**Sharding modes**")
|
|
# FFN scope labels show the effective divisor (dynamic per CP/TP/DP).
|
|
_ffn_labels = {
|
|
"TP": f"TP only (div={max(1,tp)})",
|
|
"TP+CP": f"TP*CP (div={max(1,tp*cp)})",
|
|
"TP+CP+DP": f"TP*CP*DP (div={max(1,tp*cp*dp)})",
|
|
}
|
|
_ffn_choice = st.radio(
|
|
"FFN shard scope",
|
|
options=list(_ffn_labels.values()),
|
|
index=1, key="ffn_scope_label", horizontal=True,
|
|
help=("Divisor = number of ranks the FFN weights are split across. "
|
|
"Larger divisor = less memory/PE, more AllReduce cost."),
|
|
)
|
|
# Map label -> raw scope key.
|
|
ffn_shard_scope = {v: k for k, v in _ffn_labels.items()}[_ffn_choice]
|
|
|
|
sip_topology = st.radio(
|
|
"SIP interconnect",
|
|
options=["ring", "mesh2d", "torus2d"],
|
|
index=0, key="sip_topo", horizontal=True,
|
|
help=("ring: 1D chain + wrap. mesh2d: 2D grid, no wrap. "
|
|
"torus2d: 2D grid + wrap."),
|
|
)
|
|
cp_ring_variant = st.radio(
|
|
"CP ring: what rotates",
|
|
options=["kv", "qoml"], index=0, key="cp_ring_variant",
|
|
horizontal=True,
|
|
help=("kv: K,V shards rotate each hop (Q,O,m,l stay local). "
|
|
"Bytes/hop scales with S_local*H_kv - good for prefill.\n"
|
|
"qoml: Q + running (O,m,l) rotate (K,V stay local). "
|
|
"Bytes/hop scales with T_q*H_q - much cheaper for decode."),
|
|
)
|
|
kv_shard_mode = "split"
|
|
if tp > model.h_kv:
|
|
_rep_factor = max(1, tp // model.h_kv)
|
|
kv_shard_mode = st.radio(
|
|
f"KV mode (TP={tp} > H_kv={model.h_kv})",
|
|
options=["split", "replicate"], index=0,
|
|
key="kv_mode", horizontal=True,
|
|
help=(f"split: fractional heads (KV/PE shrinks by {tp/max(1,model.h_kv):.1f}x, "
|
|
f"adds Score AllReduce across {_rep_factor} ranks). "
|
|
f"replicate: whole heads per PE, duplicated {_rep_factor}x, "
|
|
f"no Score AllReduce."),
|
|
)
|
|
|
|
|
|
with st.sidebar.expander("Hardware", expanded=False):
|
|
hw_tab_pe, hw_tab_net = st.tabs(["Per-PE", "Interconnect"])
|
|
with hw_tab_pe:
|
|
hp1, hp2 = st.columns(2)
|
|
with hp1:
|
|
pe_hbm_gb = st.selectbox(
|
|
"PE HBM (GB)", [3.0, 6.0, 12.0, 24.0, 48.0, 96.0],
|
|
index=1, key="pe_hbm",
|
|
)
|
|
bw_hbm = st.selectbox(
|
|
"HBM BW GB/s", [128.0, 256.0, 512.0, 1024.0, 2048.0],
|
|
index=1, key="bw_hbm",
|
|
)
|
|
with hp2:
|
|
peak_tflops = st.selectbox(
|
|
"TFLOPs/PE", [2.0, 4.0, 8.0, 16.0, 32.0, 64.0],
|
|
index=2, key="tflops",
|
|
)
|
|
compute_util = st.selectbox(
|
|
"Compute util", [0.3, 0.5, 0.7, 0.8, 0.9, 1.0],
|
|
index=3, key="util",
|
|
)
|
|
with hw_tab_net:
|
|
n1, n2 = st.columns(2)
|
|
with n1:
|
|
bw_intra = st.selectbox(
|
|
"intra-cube GB/s",
|
|
[128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0],
|
|
index=2, key="bw_intra",
|
|
)
|
|
bw_inter = st.selectbox(
|
|
"inter-cube GB/s",
|
|
[32.0, 64.0, 128.0, 256.0, 512.0, 900.0],
|
|
index=2, key="bw_inter",
|
|
)
|
|
bw_intersip = st.selectbox(
|
|
"inter-SIP GB/s",
|
|
[12.5, 25.0, 50.0, 100.0, 200.0, 400.0],
|
|
index=2, key="bw_intersip",
|
|
)
|
|
with n2:
|
|
alpha_intra = st.selectbox(
|
|
"a_intra ns", [5.0, 10.0, 20.0, 50.0, 100.0],
|
|
index=2, key="a_intra",
|
|
)
|
|
alpha_inter = st.selectbox(
|
|
"a_inter ns", [50.0, 100.0, 200.0, 500.0, 1000.0],
|
|
index=1, key="a_inter",
|
|
)
|
|
alpha_intersip = st.selectbox(
|
|
"a_intersip ns",
|
|
[500.0, 1000.0, 2000.0, 5000.0, 10000.0],
|
|
index=1, key="a_intersip",
|
|
)
|
|
|
|
machine = MachineParams(
|
|
pe_hbm_gb=pe_hbm_gb, peak_tflops_f16=peak_tflops,
|
|
bw_hbm_gbs=bw_hbm, bw_intra_gbs=bw_intra,
|
|
bw_inter_gbs=bw_inter, bw_intersip_gbs=bw_intersip,
|
|
alpha_intra_ns=alpha_intra, alpha_inter_ns=alpha_inter,
|
|
alpha_intersip_ns=alpha_intersip,
|
|
compute_util=compute_util,
|
|
)
|
|
|
|
topo = TopologyConfig(cp=cp, tp=tp, pp=pp, dp=dp, ep=ep,
|
|
s_kv=s_kv, mode=mode,
|
|
kv_shard_mode=kv_shard_mode,
|
|
ffn_shard_scope=ffn_shard_scope,
|
|
sip_topology=sip_topology,
|
|
tp_placement=tp_placement,
|
|
cp_placement=cp_placement,
|
|
cp_ring_variant=cp_ring_variant)
|
|
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
|
|
|
|
|
# ── Main pane ─────────────────────────────────────────────────────
|
|
st.title(f"SIP Analytical Model — {model.name}")
|
|
|
|
st.caption(
|
|
f"Family: **{preset.family or '-'}** | Attn: **{preset.attn_type}** | "
|
|
f"H_q={model.h_q}, H_kv={model.h_kv}, hidden={model.hidden}, "
|
|
f"ffn={model.ffn_dim}, layers={model.layers}"
|
|
)
|
|
if preset.note:
|
|
st.info(f"Note: {preset.note}")
|
|
|
|
|
|
# Recompute auto-suggest with the *current* machine (in case user changed HBM etc.)
|
|
suggestion = auto_suggest(model, machine, s_kv, mode)
|
|
|
|
# Compact top summary: auto-suggest + current in one row of tight bullets.
|
|
_sug_status = "OK" if suggestion.fits else "NO FIT"
|
|
_sug_line = (
|
|
f"**Auto-suggest:** CP={suggestion.cp} - TP={suggestion.tp} - "
|
|
f"PP={suggestion.pp} - PEs={suggestion.pes_used} - SIPs={suggestion.sips_used} "
|
|
f"({_sug_status}, slack {suggestion.slack_gb:.2f} GB / "
|
|
f"{machine.pe_hbm_gb:.1f} GB budget)"
|
|
)
|
|
_cur_line = (
|
|
f"**Current:** CP={cp}@{cp_placement} - TP={tp}@{tp_placement} - "
|
|
f"PP={pp} - DP={dp} - EP={ep} - PEs={topo.total_pes} - "
|
|
f"Cubes={topo.cubes_used} ({topo.pes_per_cube_used}/cube live) - "
|
|
f"SIPs={topo.sips_used} - FFN={cfg.topo.ffn_shard_scope} - "
|
|
f"SIPnet={cfg.topo.sip_topology}"
|
|
)
|
|
sum_l, sum_r = st.columns([1, 1])
|
|
with sum_l:
|
|
st.markdown(_sug_line)
|
|
with sum_r:
|
|
st.markdown(_cur_line)
|
|
|
|
# Consolidated warnings (only if any apply)
|
|
_warnings = []
|
|
if not topo.placement_valid:
|
|
_spill_cubes = (topo.intra_cube_dims + topo.pes_per_cube_hw - 1) // topo.pes_per_cube_hw
|
|
_warnings.append(
|
|
f"Intra-cube demand = {topo.intra_cube_dims} PEs > "
|
|
f"{topo.pes_per_cube_hw} PEs/cube -> spills to {_spill_cubes} cubes "
|
|
f"per group (AllReduce/ring drops from intra-cube "
|
|
f"{machine.bw_intra_gbs:.0f} to inter-cube "
|
|
f"{machine.bw_inter_gbs:.0f} GB/s)."
|
|
)
|
|
if tp > model.h_kv:
|
|
# Compute per-PE KV bytes under each mode so the difference is visible.
|
|
from tests.analytical_visualization.memory_layout import per_pe_kv_cache_bytes
|
|
import copy as _copy
|
|
_cfg_split = FullConfig(model=model, topo=TopologyConfig(
|
|
**{**topo.__dict__, "kv_shard_mode": "split"}), machine=machine)
|
|
_cfg_rep = FullConfig(model=model, topo=TopologyConfig(
|
|
**{**topo.__dict__, "kv_shard_mode": "replicate"}), machine=machine)
|
|
_kv_split_gb = per_pe_kv_cache_bytes(_cfg_split) / 1e9
|
|
_kv_rep_gb = per_pe_kv_cache_bytes(_cfg_rep) / 1e9
|
|
_kv_now = _kv_split_gb if kv_shard_mode == "split" else _kv_rep_gb
|
|
_warnings.append(
|
|
f"KV mode: **{kv_shard_mode}** (TP={tp} > H_kv={model.h_kv}) - "
|
|
f"per-PE KV: split={_kv_split_gb:.2f} GB vs "
|
|
f"replicate={_kv_rep_gb:.2f} GB (currently {_kv_now:.2f} GB). "
|
|
f"Split adds one Score AllReduce over "
|
|
f"{max(1, tp // model.h_kv)} ranks per hop."
|
|
)
|
|
if cfg.kv_replication_needed:
|
|
_warnings.append(f"Head-dim split: {cfg.head_dim_split_factor} ranks share one KV head")
|
|
if topo.tp_spans_cubes > 1:
|
|
_warnings.append(f"TP spans {topo.tp_spans_cubes} cubes "
|
|
f"(AllReduce goes cross-cube)")
|
|
if topo.cp_inter_sip_hops > 0:
|
|
_warnings.append(f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
|
|
f"at {machine.bw_intersip_gbs:.0f} GB/s")
|
|
if _warnings:
|
|
st.warning(" - ".join(_warnings))
|
|
|
|
|
|
# ── Tabs ─────────────────────────────────────────────────────────
|
|
tab_layout, tab_memory, tab_stages, tab_compare = st.tabs([
|
|
"Physical layout", "Memory breakdown", "Per-stage latency",
|
|
"Save & compare",
|
|
])
|
|
|
|
|
|
# ── TAB 1: Physical layout ──────────────────────────────────────
|
|
with tab_layout:
|
|
# Row A: pipeline diagram (color-coded) + per-stage table (formulas).
|
|
st.markdown("**Per-layer pipeline (color = dominant bound)**")
|
|
fig_pipe = draw_pipeline(cfg)
|
|
st.pyplot(fig_pipe, width='stretch')
|
|
plt.close(fig_pipe)
|
|
|
|
with st.expander("Per-stage latency (attention + FFN) with formulas",
|
|
expanded=True):
|
|
import html as _html
|
|
_stages_attn = all_stages(cfg)
|
|
_stages_ffn = all_ffn_stages(cfg)
|
|
_bound_color = {
|
|
"compute": "#d6eaf8",
|
|
"memory": "#fdebd0",
|
|
"comm": "#f5b7b1",
|
|
"trivial": "#eaeded",
|
|
}
|
|
|
|
def _vsplit(formula: str) -> str:
|
|
"""Break 'X = Y = Z' into vertical `= X` / `= Y` / `= Z` lines.
|
|
If the string contains a 'PURPOSE:' preamble ending with a line
|
|
of '---', keep that preamble verbatim and only vsplit the tail.
|
|
"""
|
|
if not formula or formula in ("-", "0"):
|
|
return formula
|
|
_preamble = ""
|
|
_body = formula
|
|
if "PURPOSE:" in formula and "\n---\n" in formula:
|
|
_preamble, _body = formula.split("\n---\n", 1)
|
|
_preamble = _preamble.rstrip() + "\n---"
|
|
pieces = [p.strip() for p in _body.split(" = ")]
|
|
_vsplit_body = "\n".join(f"= {p}" for p in pieces)
|
|
return f"{_preamble}\n{_vsplit_body}" if _preamble else _vsplit_body
|
|
|
|
# Prefill: per-hop ring is concurrent with S5/S6/S7 compute -
|
|
# annotate them. In decode the O/m/l all-reduce is already
|
|
# folded into the S8 row, so no separate annotation needed.
|
|
_per_hop_ring_us = 0.0
|
|
_ring_variant_desc = ""
|
|
if topo.mode == "prefill" and topo.cp > 1:
|
|
_c1 = next((_x for _x in _stages_attn
|
|
if _x.name.startswith("C1 ")), None)
|
|
if _c1 is not None and _c1.comm_s > 0:
|
|
_hops = max(1, topo.cp - 1)
|
|
_per_hop_ring_us = (_c1.comm_s * 1e6) / _hops
|
|
_ring_variant_desc = ("K/V" if topo.cp_ring_variant == "kv"
|
|
else "Q+O/m/l")
|
|
|
|
def _fmt_time(sec: float) -> str:
|
|
"""Auto-scale time: ns for < 1 us, us for < 10 ms, else ms."""
|
|
if sec <= 0:
|
|
return "0"
|
|
ns = sec * 1e9
|
|
if ns < 1000:
|
|
return f"{ns:.1f} ns"
|
|
us = sec * 1e6
|
|
if us < 1000:
|
|
return f"{us:.2f} us"
|
|
ms = sec * 1e3
|
|
return f"{ms:.2f} ms"
|
|
|
|
def _combined_formula(s):
|
|
parts = []
|
|
if s.compute_s > 0 or (s.flops and s.flops > 0):
|
|
parts.append(
|
|
f"Compute:\n{_vsplit(s.flops_formula)}\n"
|
|
f"-> {_fmt_time(s.compute_s)}"
|
|
)
|
|
if s.memory_s > 0 or (s.mem_bytes and s.mem_bytes > 0):
|
|
parts.append(
|
|
f"Memory:\n{_vsplit(s.mem_formula)}\n"
|
|
f"-> {_fmt_time(s.memory_s)}"
|
|
)
|
|
if s.comm_s > 0 or (s.comm_bytes and s.comm_bytes > 0):
|
|
# If this is a concurrent-comm stage (e.g. C1 CP ring in
|
|
# prefill) that gets overlapped with compute, visible < raw.
|
|
# Show both so the '0' in the Time column makes sense.
|
|
_tail = f"-> raw comm: {_fmt_time(s.comm_s)}"
|
|
if s.visible_s < s.comm_s * 0.999:
|
|
_tail += (
|
|
f"\n-> visible: {_fmt_time(s.visible_s)} "
|
|
f"(concurrent S5-S7 compute hides "
|
|
f"{_fmt_time(s.comm_s - s.visible_s)} of it)"
|
|
)
|
|
parts.append(
|
|
f"Comm:\n{_vsplit(s.comm_formula)}\n{_tail}"
|
|
)
|
|
# Prefill: S5/S6/S7 run concurrently with the CP ring.
|
|
if (_per_hop_ring_us > 0
|
|
and s.name.startswith(("S5 ", "S6 ", "S7 "))):
|
|
# convert us to seconds for _fmt_time
|
|
_per_hop_s = _per_hop_ring_us / 1e6
|
|
parts.append(
|
|
f"Concurrent CP ring ({_ring_variant_desc}, per hop):\n"
|
|
f"~{_fmt_time(_per_hop_s)} in flight while this compute runs\n"
|
|
f"(aggregated in C1 row; overlap kept only excess vs compute)"
|
|
)
|
|
return "\n\n".join(parts) if parts else "-"
|
|
|
|
# Compact CSS: rows auto-size to text; tight padding; columns fit content.
|
|
_stage_table_css = """
|
|
<style>
|
|
table.stage-tbl { border-collapse: collapse; width: auto;
|
|
font-size: 11.5px; margin-bottom: 6px; }
|
|
table.stage-tbl th { text-align: left; padding: 3px 8px;
|
|
background: #e9ecef; border: 1px solid #ced4da;
|
|
font-weight: 600; }
|
|
table.stage-tbl td { padding: 3px 8px; border: 1px solid #dee2e6;
|
|
vertical-align: top; }
|
|
table.stage-tbl td.stg { white-space: nowrap; font-weight: 600; }
|
|
table.stage-tbl td.bnd { text-align: center; font-weight: 600;
|
|
text-transform: uppercase; font-size: 10px;
|
|
white-space: nowrap; }
|
|
table.stage-tbl td.us { text-align: right; font-family: Consolas, monospace;
|
|
white-space: nowrap; }
|
|
table.stage-tbl pre { margin: 0; font-family: Consolas, monospace;
|
|
font-size: 11px; line-height: 1.15;
|
|
white-space: pre; }
|
|
</style>
|
|
"""
|
|
|
|
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'<tr style="background:{bg};">'
|
|
f'<td class="stg">{_html.escape(s.name)}</td>'
|
|
f'<td><pre>{formula}</pre></td>'
|
|
f'<td class="bnd">{s.bound}</td>'
|
|
f'<td class="us">{_fmt_time(s.visible_s)}</td>'
|
|
f'</tr>'
|
|
)
|
|
html = (
|
|
_stage_table_css
|
|
+ f'<b>{_html.escape(title)}</b>'
|
|
+ '<table class="stage-tbl">'
|
|
+ '<thead><tr><th>Stage</th><th>Formula</th>'
|
|
'<th>Bound</th><th>Time</th></tr></thead>'
|
|
+ '<tbody>' + ''.join(rows_html) + '</tbody></table>'
|
|
)
|
|
# Prefer st.html (reliable) with a markdown fallback.
|
|
if hasattr(st, "html"):
|
|
st.html(html)
|
|
else:
|
|
st.markdown(html, unsafe_allow_html=True)
|
|
_t = sum(s.visible_s for s in stages)
|
|
st.caption(
|
|
f"{title} total: **{_fmt_time(_t)}**. "
|
|
f"Compute time = FLOPs / ({cfg.machine.peak_tflops_f16:.1f} TFLOPs "
|
|
f"x {cfg.machine.compute_util:.0%} util). "
|
|
f"Memory time = bytes / {cfg.machine.bw_hbm_gbs:.0f} GB/s HBM. "
|
|
f"Bound = max(compute, memory, comm) - row color: "
|
|
f"blue=compute, orange=memory, red=comm, grey=trivial."
|
|
)
|
|
|
|
_render_table(_stages_attn, "Attention")
|
|
_render_table(_stages_ffn, "FFN")
|
|
|
|
with st.expander("Symbol glossary (what does T_q, d_h, div... mean?)",
|
|
expanded=False):
|
|
_glossary = [
|
|
("Model dims", [
|
|
("d", f"hidden dim (current: {model.hidden})"),
|
|
("d_h", f"per-head dim (current: {model.d_head})"),
|
|
("H_q", f"total query heads (current: {model.h_q})"),
|
|
("H_kv", f"total KV heads (current: {model.h_kv})"),
|
|
("ffn", f"FFN intermediate dim (current: {model.ffn_dim})"),
|
|
("L", f"transformer layers (current: {model.layers})"),
|
|
("b", f"bytes per element (dtype size) "
|
|
f"(current: {model.bytes_per_elem} -> "
|
|
f"{'fp8' if model.bytes_per_elem == 1 else 'bf16/fp16' if model.bytes_per_elem == 2 else 'fp32'})"),
|
|
]),
|
|
("Workload", [
|
|
("S_kv", f"global KV cache length / context (current: {topo.s_kv:,})"),
|
|
("S_local", f"per-CP-rank KV length = S_kv / CP "
|
|
f"(current: {topo.s_local:,})"),
|
|
("T_q", f"query tokens processed this step "
|
|
f"(decode=1, prefill=S_local; current: {topo.T_q})"),
|
|
("mode", f"decode | prefill (current: {topo.mode})"),
|
|
]),
|
|
("Parallelism (degrees)", [
|
|
("TP", f"tensor parallelism (attn heads / FFN cols) "
|
|
f"(current: {topo.tp}, on {topo.tp_placement})"),
|
|
("CP", f"context parallelism (sequence axis) "
|
|
f"(current: {topo.cp}, on {topo.cp_placement})"),
|
|
("PP", f"pipeline parallelism (layer stages) (current: {topo.pp})"),
|
|
("DP", f"data parallelism (model replicas) (current: {topo.dp})"),
|
|
("EP", f"expert parallelism (MoE) (current: {topo.ep})"),
|
|
("div", f"FFN shard divisor = product of dims in "
|
|
f"'FFN shard scope' (current: {cfg.ffn_shard_divisor})"),
|
|
("split", f"# ranks sharing one KV head via head-dim split "
|
|
f"(TP/H_kv when TP > H_kv; current: "
|
|
f"{cfg.head_dim_split_factor if cfg.kv_replication_needed else 1})"),
|
|
]),
|
|
("Hardware / machine", [
|
|
("BW_HBM", f"per-PE HBM bandwidth "
|
|
f"(current: {machine.bw_hbm_gbs:.0f} GB/s)"),
|
|
("BW_intra", f"intra-cube PE<->PE bandwidth "
|
|
f"(current: {machine.bw_intra_gbs:.0f} GB/s)"),
|
|
("BW_inter", f"inter-cube (D2D) bandwidth "
|
|
f"(current: {machine.bw_inter_gbs:.0f} GB/s)"),
|
|
("BW_intersip", f"inter-SIP (C2C) bandwidth "
|
|
f"(current: {machine.bw_intersip_gbs:.0f} GB/s)"),
|
|
("alpha_intra", f"per-hop latency intra-cube "
|
|
f"(current: {machine.alpha_intra_ns:.0f} ns)"),
|
|
("alpha_inter", f"per-hop latency inter-cube "
|
|
f"(current: {machine.alpha_inter_ns:.0f} ns)"),
|
|
("alpha_intersip", f"per-hop latency inter-SIP "
|
|
f"(current: {machine.alpha_intersip_ns:.0f} ns)"),
|
|
("peak TFLOPs", f"per-PE peak compute "
|
|
f"(current: {machine.peak_tflops_f16:.1f} TFLOPs f16)"),
|
|
("util", f"achievable compute utilization "
|
|
f"(current: {machine.compute_util:.0%})"),
|
|
]),
|
|
("Stage prefixes", [
|
|
("S1 .. S10", "attention forward stages (RMSNorm, W_Q, W_K+W_V, "
|
|
"KV append, Q.K^T, softmax, P.V, merge, "
|
|
"normalize, W_O)"),
|
|
("F1 .. F5", "FFN forward stages (RMSNorm, W_gate, W_up, "
|
|
"SwiGLU, W_down)"),
|
|
("C1", "CP K/V ring (K,V passed hop-by-hop)"),
|
|
("C2", "TP AllReduce on W_O output"),
|
|
("C3", "Score AllReduce when TP > H_kv with head-dim split"),
|
|
("CF1", "FFN AllReduce over the FFN shard scope"),
|
|
]),
|
|
("Notation in formulas", [
|
|
("H_q / TP", "Q heads per PE (or per TP rank)"),
|
|
("H_kv / TP", "KV heads per PE (fractional allowed in split mode)"),
|
|
("ffn / div", "FFN cols/rows per PE after sharding"),
|
|
("2 * A * B * C", "GEMM FLOPs = 2 * M * K * N (multiply + add per fma)"),
|
|
("bytes * b", "byte count (b = bytes/element)"),
|
|
("-> X us", "resulting latency after dividing by BW or peak"),
|
|
]),
|
|
]
|
|
for section, items in _glossary:
|
|
st.markdown(f"**{section}**")
|
|
_rows = "".join(
|
|
f'<tr><td style="padding:2px 8px;font-family:Consolas,monospace;'
|
|
f'white-space:nowrap;color:#0d47a1;font-weight:600;">{sym}</td>'
|
|
f'<td style="padding:2px 8px;">{desc}</td></tr>'
|
|
for sym, desc in items
|
|
)
|
|
st.html(
|
|
f'<table style="font-size:11.5px;border-collapse:collapse;'
|
|
f'margin-bottom:6px;">{_rows}</table>'
|
|
)
|
|
|
|
st.divider()
|
|
|
|
# Row B: topology diagram (full width - it can be tall for multi-SIP)
|
|
with st.expander(
|
|
f"Physical layout - {topo.total_pes} PEs, {topo.cubes_used} cubes, "
|
|
f"{topo.sips_used} SIP(s), inter-SIP: {topo.sip_topology}",
|
|
expanded=True,
|
|
):
|
|
fig_topo = draw_topology(cfg)
|
|
st.pyplot(fig_topo, width='stretch')
|
|
plt.close(fig_topo)
|
|
l1, l2, l3, l4 = st.columns(4)
|
|
l1.metric("Cubes", topo.cubes_used)
|
|
l2.metric("SIPs", topo.sips_used)
|
|
l3.metric("CP intra-SIP hops", topo.cp_intra_sip_hops)
|
|
l4.metric("CP inter-SIP hops", topo.cp_inter_sip_hops)
|
|
|
|
st.divider()
|
|
|
|
# Row C: two-column - LEFT per-PE memory (pie + summary), RIGHT tensor sharding
|
|
mem_layout = compute_memory(cfg)
|
|
col_mem, col_shard = st.columns([1, 1], gap="medium")
|
|
|
|
with col_mem:
|
|
st.markdown("**Per-PE memory footprint**")
|
|
labels = ["Weights", "KV cache", "Transient", "Slack"]
|
|
vals = [mem_layout.weights_bytes, mem_layout.kv_cache_bytes,
|
|
mem_layout.transient_bytes, mem_layout.slack_bytes]
|
|
colors = ["#3a86ff", "#8338ec", "#ffbe0b", "#adb5bd"]
|
|
fig_mem_l, ax_mem_l = plt.subplots(figsize=(4.2, 4.2))
|
|
ax_mem_l.pie(vals,
|
|
labels=[f"{l}\n{v/1e9:.2f} GB"
|
|
for l, v in zip(labels, vals)],
|
|
colors=colors, autopct="%1.1f%%", startangle=90,
|
|
textprops={"fontsize": 8})
|
|
ax_mem_l.set_title(f"Per-PE (budget: "
|
|
f"{mem_layout.budget_bytes/1e9:.1f} GB)",
|
|
fontsize=10)
|
|
st.pyplot(fig_mem_l, width='stretch')
|
|
plt.close(fig_mem_l)
|
|
if mem_layout.over_budget:
|
|
st.error(f"OVER BUDGET by "
|
|
f"{(mem_layout.used_bytes - mem_layout.budget_bytes)/1e9:.2f} GB")
|
|
else:
|
|
st.caption(
|
|
f"Weights {mem_layout.weights_bytes/1e9:.2f} GB - "
|
|
f"KV {mem_layout.kv_cache_bytes/1e9:.2f} GB "
|
|
f"(S_local={topo.s_local:,}) - "
|
|
f"Transient {mem_layout.transient_bytes/1e9:.2f} GB - "
|
|
f"Slack {mem_layout.slack_bytes/1e9:.2f} GB"
|
|
)
|
|
|
|
with col_shard:
|
|
st.markdown("**Tensor sharding (which dim is split)**")
|
|
st.caption(
|
|
"W_Q/W_K/W_V/W_gate/W_up: **output cols**. "
|
|
"W_O/W_down: **input rows**. "
|
|
"KV cache: **rows by CP** x **cols by TP**."
|
|
)
|
|
s_c1, s_c2 = st.columns([1, 2])
|
|
with s_c1:
|
|
my_pe_view = st.number_input("PE index",
|
|
min_value=0,
|
|
max_value=max(0, topo.tp - 1),
|
|
value=0, key="my_pe_view")
|
|
with s_c2:
|
|
show_physical = st.checkbox(
|
|
"Enlarge + show physical PE/cube mapping",
|
|
value=False, key="tsv_physical",
|
|
help=("Larger figure with each shard cell labelled by owning "
|
|
"cube/PE. Especially useful for KV cache when CP is "
|
|
"split across cubes and PEs."),
|
|
)
|
|
fig_shard = draw_tensor_sharding(cfg, my_pe=my_pe_view,
|
|
show_physical=show_physical)
|
|
st.pyplot(fig_shard, width='stretch')
|
|
plt.close(fig_shard)
|
|
|
|
st.divider()
|
|
|
|
# Row D: PE-level view (wide, full-width) inside an expander
|
|
with st.expander(
|
|
f"PE-level view of one CP group (TP={topo.tp}, "
|
|
f"H_q={model.h_q}, H_kv={model.h_kv})",
|
|
expanded=False,
|
|
):
|
|
st.caption(
|
|
f"One CP rank's {topo.tp} PE(s). Attention weights replicated "
|
|
f"across CP groups; KV sharded by CP (S/{topo.cp} tokens) x TP heads; "
|
|
f"FFN by TP x EP."
|
|
)
|
|
fig_pe = draw_pe_layout(cfg)
|
|
st.pyplot(fig_pe, width='stretch')
|
|
plt.close(fig_pe)
|
|
|
|
# Row E: Replication + Optimization panel (side-by-side)
|
|
st.divider()
|
|
col_rep, col_opt = st.columns([1, 1], gap="medium")
|
|
|
|
with col_rep:
|
|
st.markdown("**Replicated / duplicated storage**")
|
|
rep_entries = replication_report(cfg)
|
|
if not rep_entries:
|
|
st.caption("No duplicated storage detected - every tensor is uniquely sharded.")
|
|
else:
|
|
rep_rows = []
|
|
total_waste = 0
|
|
for e in rep_entries:
|
|
rep_rows.append({
|
|
"Tensor": e.tensor,
|
|
"Per-PE (GB)": round(e.per_pe_bytes / 1e9, 3),
|
|
"Replicated": e.replicated_across,
|
|
"Copies": e.copies,
|
|
"Waste/PE (GB)": round(e.wasted_bytes / 1e9, 3),
|
|
})
|
|
total_waste += e.wasted_bytes
|
|
st.dataframe(pd.DataFrame(rep_rows), width='stretch',
|
|
hide_index=True)
|
|
st.caption(
|
|
f"Total per-PE duplicated bytes: "
|
|
f"**{total_waste/1e9:.2f} GB** "
|
|
f"(vs {mem_layout.budget_bytes/1e9:.1f} GB budget). "
|
|
f"Each 'copy' represents one identical block stored on "
|
|
f"another rank."
|
|
)
|
|
|
|
with col_opt:
|
|
st.markdown("**Optimization opportunities**")
|
|
hints = optimization_hints(cfg)
|
|
if not hints:
|
|
st.caption("No obvious improvements at this configuration.")
|
|
else:
|
|
_icons = {"warn": ":warning:", "info": ":bulb:", "good": ":white_check_mark:"}
|
|
_cats = {"space": "SPACE", "comm": "COMM",
|
|
"compute": "COMPUTE", "layout": "LAYOUT"}
|
|
for h in hints:
|
|
icon = _icons.get(h.severity, ":bulb:")
|
|
cat = _cats.get(h.category, h.category.upper())
|
|
st.markdown(f"{icon} **[{cat}]** {h.message}")
|
|
|
|
st.divider()
|
|
|
|
# Row F: Per-PE breakdown table (collapsible)
|
|
with st.expander("Per-PE breakdown table", expanded=False):
|
|
rows = []
|
|
for pe_id in range(topo.tp):
|
|
q_heads = _q_heads_for_pe(pe_id, topo.tp, model.h_q)
|
|
kv_heads, kv_note = _kv_heads_for_pe(
|
|
pe_id, topo.tp, model.h_kv, topo.kv_shard_mode,
|
|
)
|
|
bytes_ = _per_pe_bytes(cfg, pe_id)
|
|
weights_bytes = sum(v for k, v in bytes_.items()
|
|
if k not in ("KV cache", "Transient"))
|
|
rows.append({
|
|
"PE": pe_id,
|
|
"Q heads": (f"{q_heads[0]}-{q_heads[-1]}"
|
|
if len(q_heads) > 1 else str(q_heads[0])
|
|
if q_heads else "-"),
|
|
"KV heads": (str(kv_heads[0]) + (f" ({kv_note})" if kv_note else "")
|
|
if kv_heads else "-"),
|
|
"W_Q (MB)": round(bytes_["W_Q"] / 1e6, 2),
|
|
"W_K (MB)": round(bytes_["W_K"] / 1e6, 2),
|
|
"W_V (MB)": round(bytes_["W_V"] / 1e6, 2),
|
|
"W_O (MB)": round(bytes_["W_O"] / 1e6, 2),
|
|
"FFN (MB)": round((bytes_["W_gate"] + bytes_["W_up"]
|
|
+ bytes_["W_down"]) / 1e6, 2),
|
|
"Weights total (GB)": round(weights_bytes / 1e9, 3),
|
|
"KV cache (GB)": round(bytes_["KV cache"] / 1e9, 3),
|
|
"Transient (MB)": round(bytes_["Transient"] / 1e6, 2),
|
|
})
|
|
st.dataframe(pd.DataFrame(rows), width='stretch', hide_index=True)
|
|
|
|
|
|
# ── TAB 2: Memory breakdown ─────────────────────────────────────
|
|
with tab_memory:
|
|
mem = compute_memory(cfg)
|
|
layers_per_stage = (model.layers + pp - 1) // pp
|
|
|
|
parallelism_str = (
|
|
f"CP={cp}, TP={tp}, PP={pp}, DP={dp}, EP={ep}"
|
|
f" | layers={model.layers}, layers/stage={layers_per_stage}"
|
|
f" | KV mode: {kv_shard_mode}"
|
|
)
|
|
st.caption(f"**Parallelism:** {parallelism_str}")
|
|
|
|
# ─── ATTENTION (per-layer) ─────────────────────────────────
|
|
st.subheader(f"Attention weights (per layer) — {parallelism_str}")
|
|
attn_rows = attention_weight_rows(cfg)
|
|
display_attn = [{k: v for k, v in r.items() if not k.startswith("_")}
|
|
for r in attn_rows]
|
|
st.dataframe(pd.DataFrame(display_attn), width='stretch', hide_index=True)
|
|
|
|
attn_bytes_1L_global = sum(r["_p_global"] for r in attn_rows) * model.bytes_per_elem
|
|
attn_bytes_1L_pe = sum(r["_p_per_pe"] for r in attn_rows) * model.bytes_per_elem
|
|
attn_bytes_all_global = attn_bytes_1L_global * model.layers
|
|
attn_bytes_all_pe = attn_bytes_1L_pe * layers_per_stage
|
|
|
|
ac1, ac2, ac3, ac4 = st.columns(4)
|
|
ac1.metric("Attn / layer (global)", f"{attn_bytes_1L_global/1e6:.2f} MB")
|
|
ac2.metric(f"Attn all {model.layers} layers", f"{attn_bytes_all_global/1e9:.3f} GB")
|
|
ac3.metric("Attn / layer (per PE)", f"{attn_bytes_1L_pe/1e6:.2f} MB")
|
|
ac4.metric(f"Attn per PE ({layers_per_stage} layers/stage)",
|
|
f"{attn_bytes_all_pe/1e9:.3f} GB")
|
|
|
|
# ─── FFN / MoE (per-layer) ─────────────────────────────────
|
|
is_moe = "MoE" in (preset.family + preset.note)
|
|
ffn_label = "FFN (activated experts approx)" if is_moe else "FFN"
|
|
st.subheader(f"{ffn_label} weights (per layer) — {parallelism_str}")
|
|
ffn_rows = ffn_weight_rows(cfg)
|
|
display_ffn = [{k: v for k, v in r.items() if not k.startswith("_")}
|
|
for r in ffn_rows]
|
|
st.dataframe(pd.DataFrame(display_ffn), width='stretch', hide_index=True)
|
|
|
|
ffn_bytes_1L_global = sum(r["_p_global"] for r in ffn_rows) * model.bytes_per_elem
|
|
ffn_bytes_1L_pe = sum(r["_p_per_pe"] for r in ffn_rows) * model.bytes_per_elem
|
|
ffn_bytes_all_global = ffn_bytes_1L_global * model.layers
|
|
ffn_bytes_all_pe = ffn_bytes_1L_pe * layers_per_stage
|
|
|
|
fc1, fc2, fc3, fc4 = st.columns(4)
|
|
fc1.metric("FFN / layer (global)", f"{ffn_bytes_1L_global/1e6:.2f} MB")
|
|
fc2.metric(f"FFN all {model.layers} layers", f"{ffn_bytes_all_global/1e9:.3f} GB")
|
|
fc3.metric("FFN / layer (per PE)", f"{ffn_bytes_1L_pe/1e6:.2f} MB")
|
|
fc4.metric(f"FFN per PE ({layers_per_stage} layers/stage)",
|
|
f"{ffn_bytes_all_pe/1e9:.3f} GB")
|
|
|
|
# ─── KV cache (per-layer) ─────────────────────────────────
|
|
st.subheader(f"KV cache (per layer, S_kv = {s_kv:,}) — {parallelism_str}")
|
|
kv_rows = kv_cache_rows(cfg)
|
|
display_kv = [{k: v for k, v in r.items() if not k.startswith("_")}
|
|
for r in kv_rows]
|
|
st.dataframe(pd.DataFrame(display_kv), width='stretch', hide_index=True)
|
|
|
|
kv_bytes_1L_global = sum(r["_p_global"] for r in kv_rows) * model.bytes_per_elem
|
|
kv_bytes_1L_pe = sum(r["_p_per_pe"] for r in kv_rows) * model.bytes_per_elem
|
|
kv_bytes_all_global = kv_bytes_1L_global * model.layers
|
|
kv_bytes_all_pe = kv_bytes_1L_pe * layers_per_stage
|
|
|
|
kc1, kc2, kc3, kc4 = st.columns(4)
|
|
kc1.metric("KV / layer (global)", f"{kv_bytes_1L_global/1e6:.2f} MB")
|
|
kc2.metric(f"KV all {model.layers} layers", f"{kv_bytes_all_global/1e9:.3f} GB")
|
|
kc3.metric("KV / layer (per PE)", f"{kv_bytes_1L_pe/1e6:.2f} MB")
|
|
kc4.metric(f"KV per PE ({layers_per_stage} layers/stage)",
|
|
f"{kv_bytes_all_pe/1e9:.3f} GB")
|
|
|
|
st.divider()
|
|
|
|
# ─── FULL-MODEL TOTALS ───────────────────────────────────
|
|
st.subheader("Full-model totals (all layers)")
|
|
total_w_global = attn_bytes_all_global + ffn_bytes_all_global
|
|
total_all_global = total_w_global + kv_bytes_all_global
|
|
total_per_pe = attn_bytes_all_pe + ffn_bytes_all_pe + kv_bytes_all_pe
|
|
tc1, tc2, tc3, tc4 = st.columns(4)
|
|
tc1.metric("Weights (all layers, global)", f"{total_w_global/1e9:.2f} GB")
|
|
tc2.metric("KV (all layers, global)", f"{kv_bytes_all_global/1e9:.2f} GB")
|
|
tc3.metric("Weights + KV (global)", f"{total_all_global/1e9:.2f} GB")
|
|
tc4.metric("Per-PE total (W+KV, layers/stage)",
|
|
f"{total_per_pe/1e9:.2f} GB")
|
|
|
|
# ─── Per-PE memory pie ───────────────────────────────────
|
|
st.subheader("Per-PE memory footprint (with transient + slack)")
|
|
pie_col, info_col = st.columns([1, 1])
|
|
with pie_col:
|
|
labels = ["Weights", "KV cache", "Transient", "Slack"]
|
|
vals = [mem.weights_bytes, mem.kv_cache_bytes,
|
|
mem.transient_bytes, mem.slack_bytes]
|
|
colors = ["#3a86ff", "#8338ec", "#ffbe0b", "#adb5bd"]
|
|
|
|
fig_mem, ax_mem = plt.subplots(figsize=(6, 6))
|
|
ax_mem.pie(vals,
|
|
labels=[f"{l}\n{v/1e9:.2f} GB" for l, v in zip(labels, vals)],
|
|
colors=colors, autopct="%1.1f%%", startangle=90)
|
|
ax_mem.set_title(f"Per-PE (budget: {mem.budget_bytes/1e9:.1f} GB)")
|
|
st.pyplot(fig_mem, width='stretch')
|
|
plt.close(fig_mem)
|
|
with info_col:
|
|
if mem.over_budget:
|
|
st.error(f"OVER BUDGET by "
|
|
f"{(mem.used_bytes - mem.budget_bytes)/1e9:.2f} GB")
|
|
else:
|
|
st.success(f"{mem.slack_bytes/1e9:.2f} GB slack")
|
|
|
|
st.markdown(
|
|
f"- Weights: **{mem.weights_bytes/1e9:.2f} GB** "
|
|
f"(layers/PP = {layers_per_stage})\n"
|
|
f"- KV cache: **{mem.kv_cache_bytes/1e9:.2f} GB** "
|
|
f"(S_local = {topo.s_local:,})\n"
|
|
f"- Transient: {mem.transient_bytes/1e9:.2f} GB\n"
|
|
f"- Slack: {mem.slack_bytes/1e9:.2f} GB"
|
|
)
|
|
|
|
|
|
# ── TAB 3: Per-stage latency ────────────────────────────────────
|
|
with tab_stages:
|
|
st.subheader("Per-stage attention latency")
|
|
stages = all_stages(cfg)
|
|
|
|
rows = []
|
|
for s in stages:
|
|
rows.append({
|
|
"Stage": s.name,
|
|
"Bound": s.bound,
|
|
"Compute (us)": round(s.compute_s * 1e6, 3),
|
|
"Memory (us)": round(s.memory_s * 1e6, 3),
|
|
"Comm (us)": round(s.comm_s * 1e6, 3),
|
|
"Visible (us)": round(s.visible_s * 1e6, 3),
|
|
"Formula": s.formula,
|
|
})
|
|
df = pd.DataFrame(rows)
|
|
st.dataframe(df, width='stretch', hide_index=True)
|
|
|
|
total_visible_s = sum(s.visible_s for s in stages)
|
|
total_layers = total_visible_s * model.layers
|
|
st.markdown(
|
|
f"**Per-layer attention (analytical): {total_visible_s*1e6:.2f} us** \n"
|
|
f"**Full-model attention ({model.layers} layers): "
|
|
f"{total_layers*1e6:.2f} us = {total_layers*1e3:.2f} ms**"
|
|
)
|
|
|
|
st.subheader("Stage cost breakdown")
|
|
fig_bar, ax_bar = plt.subplots(figsize=(12, 5))
|
|
names = [s.name for s in stages]
|
|
cmp_vals = [s.compute_s * 1e6 for s in stages]
|
|
mem_vals = [s.memory_s * 1e6 for s in stages]
|
|
comm_vals = [s.comm_s * 1e6 for s in stages]
|
|
|
|
x = range(len(names))
|
|
ax_bar.bar(x, cmp_vals, label="Compute", color="#3a86ff")
|
|
ax_bar.bar(x, mem_vals, bottom=cmp_vals, label="Memory", color="#ffbe0b")
|
|
ax_bar.bar(x, comm_vals,
|
|
bottom=[c + m for c, m in zip(cmp_vals, mem_vals)],
|
|
label="Comm", color="#d90429")
|
|
ax_bar.set_xticks(x)
|
|
ax_bar.set_xticklabels(names, rotation=45, ha="right", fontsize=8)
|
|
ax_bar.set_ylabel("Time (us)")
|
|
ax_bar.set_title(f"Stage breakdown - mode={mode}, T_q={topo.T_q}")
|
|
ax_bar.legend()
|
|
ax_bar.grid(axis="y", alpha=0.3)
|
|
plt.tight_layout()
|
|
st.pyplot(fig_bar, width='stretch')
|
|
plt.close(fig_bar)
|
|
|
|
|
|
# ── TAB 4: Save & compare ───────────────────────────────────────
|
|
def _snapshot_current_config():
|
|
"""Capture a compact snapshot of the current cfg + key metrics."""
|
|
_stages_attn = all_stages(cfg)
|
|
_stages_ffn = all_ffn_stages(cfg)
|
|
attn_us = sum(s.visible_s for s in _stages_attn) * 1e6
|
|
ffn_us = sum(s.visible_s for s in _stages_ffn) * 1e6
|
|
mem = compute_memory(cfg)
|
|
# Per-stage latencies keyed by short prefix (S1, S2, ..., C1, F1, ..., CF1).
|
|
# Value = seconds (raw), so display can format uniformly across snapshots.
|
|
_by_prefix_attn = {s.name.split()[0]: s.visible_s for s in _stages_attn}
|
|
_by_prefix_ffn = {s.name.split()[0]: s.visible_s for s in _stages_ffn}
|
|
# Preserve the full descriptive name too so a hover/second column can
|
|
# show what the abbreviation means for this particular snapshot.
|
|
_fullname_attn = {s.name.split()[0]: s.name for s in _stages_attn}
|
|
_fullname_ffn = {s.name.split()[0]: s.name for s in _stages_ffn}
|
|
return {
|
|
"model": model.name,
|
|
"mode": topo.mode,
|
|
"s_kv": topo.s_kv,
|
|
"cp": topo.cp, "tp": topo.tp,
|
|
"pp": topo.pp, "dp": topo.dp, "ep": topo.ep,
|
|
"cp_placement": topo.cp_placement,
|
|
"tp_placement": topo.tp_placement,
|
|
"ffn_scope": topo.ffn_shard_scope,
|
|
"kv_mode": topo.kv_shard_mode,
|
|
"sip_topology": topo.sip_topology,
|
|
"cp_ring_variant": topo.cp_ring_variant,
|
|
"pes": topo.total_pes,
|
|
"cubes": topo.cubes_used,
|
|
"sips": topo.sips_used,
|
|
"pe_hbm_gb": machine.pe_hbm_gb,
|
|
"peak_tflops": machine.peak_tflops_f16,
|
|
"bw_hbm_gbs": machine.bw_hbm_gbs,
|
|
"bw_inter_gbs": machine.bw_inter_gbs,
|
|
"bw_intersip_gbs": machine.bw_intersip_gbs,
|
|
"weights_gb": round(mem.weights_bytes / 1e9, 3),
|
|
"kv_gb": round(mem.kv_cache_bytes / 1e9, 3),
|
|
"transient_gb": round(mem.transient_bytes / 1e9, 3),
|
|
"used_gb": round(mem.used_bytes / 1e9, 3),
|
|
"slack_gb": round(mem.slack_bytes / 1e9, 3),
|
|
"over_budget": mem.over_budget,
|
|
"attn_us": round(attn_us, 3),
|
|
"ffn_us": round(ffn_us, 3),
|
|
"layer_us": round(attn_us + ffn_us, 3),
|
|
"stages_attn": _by_prefix_attn,
|
|
"stages_ffn": _by_prefix_ffn,
|
|
"stage_names_attn": _fullname_attn,
|
|
"stage_names_ffn": _fullname_ffn,
|
|
}
|
|
|
|
|
|
with tab_compare:
|
|
st.markdown(
|
|
"Save the current configuration under a name, then load it back "
|
|
"later or compare multiple saved configs side-by-side."
|
|
)
|
|
if "saved_cfgs" not in st.session_state:
|
|
st.session_state["saved_cfgs"] = {} # name -> snapshot dict
|
|
|
|
# Default name: next free config1, config2, ...
|
|
def _next_config_name():
|
|
i = 1
|
|
while f"config{i}" in st.session_state.get("saved_cfgs", {}):
|
|
i += 1
|
|
return f"config{i}"
|
|
|
|
sv_c1, sv_c2, sv_c3 = st.columns([3, 1, 1])
|
|
with sv_c1:
|
|
_save_name = st.text_input(
|
|
"Config name",
|
|
value=_next_config_name(),
|
|
key=f"cfg_save_name_{len(st.session_state.get('saved_cfgs', {}))}",
|
|
)
|
|
with sv_c2:
|
|
if st.button("Save current", width='stretch'):
|
|
if _save_name.strip():
|
|
st.session_state["saved_cfgs"][_save_name.strip()] = \
|
|
_snapshot_current_config()
|
|
st.success(f"Saved '{_save_name.strip()}'")
|
|
st.rerun()
|
|
else:
|
|
st.warning("Give it a name first.")
|
|
with sv_c3:
|
|
if st.button("Clear all", width='stretch'):
|
|
st.session_state["saved_cfgs"] = {}
|
|
st.rerun()
|
|
|
|
_saved_all = st.session_state.get("saved_cfgs", {})
|
|
if not _saved_all:
|
|
st.info("No saved configurations yet.")
|
|
else:
|
|
st.markdown(f"**{len(_saved_all)} saved configuration(s)** - "
|
|
f"pick which to compare below")
|
|
|
|
# Pick which configs to include in the comparison (default all)
|
|
_picked = st.multiselect(
|
|
"Compare",
|
|
options=list(_saved_all.keys()),
|
|
default=list(_saved_all.keys()),
|
|
key="compare_picked",
|
|
)
|
|
_saved = {n: _saved_all[n] for n in _picked}
|
|
|
|
# Delete individual
|
|
with st.expander("Manage (delete individual)", expanded=False):
|
|
for _n in list(_saved_all.keys()):
|
|
d1, d2, d3 = st.columns([5, 1, 1])
|
|
d1.text(_n)
|
|
if d2.button("Rename", key=f"rn_btn_{_n}"):
|
|
st.session_state[f"rn_active_{_n}"] = True
|
|
if d3.button("Delete", key=f"del_{_n}"):
|
|
del st.session_state["saved_cfgs"][_n]
|
|
st.rerun()
|
|
if st.session_state.get(f"rn_active_{_n}"):
|
|
new_name = st.text_input(
|
|
f"New name for '{_n}'", value=_n,
|
|
key=f"rn_txt_{_n}",
|
|
)
|
|
if st.button("Confirm rename", key=f"rn_ok_{_n}"):
|
|
if new_name and new_name != _n and new_name not in _saved_all:
|
|
st.session_state["saved_cfgs"][new_name] = \
|
|
st.session_state["saved_cfgs"].pop(_n)
|
|
st.session_state[f"rn_active_{_n}"] = False
|
|
st.rerun()
|
|
|
|
if not _saved:
|
|
st.info("No configurations selected. Tick some above to compare.")
|
|
st.stop()
|
|
|
|
# Comparison table: rows = metrics, columns = config names
|
|
_key_order = [
|
|
("Model / workload", ["model", "mode", "s_kv"]),
|
|
("Parallelism", ["cp", "tp", "pp", "dp", "ep",
|
|
"cp_placement", "tp_placement",
|
|
"ffn_scope", "kv_mode", "cp_ring_variant",
|
|
"sip_topology"]),
|
|
("Physical", ["pes", "cubes", "sips"]),
|
|
("Hardware knobs", ["pe_hbm_gb", "peak_tflops",
|
|
"bw_hbm_gbs", "bw_inter_gbs",
|
|
"bw_intersip_gbs"]),
|
|
("Memory / PE (GB)", ["weights_gb", "kv_gb", "transient_gb",
|
|
"used_gb", "slack_gb", "over_budget"]),
|
|
("Per-layer latency (us)", ["attn_us", "ffn_us", "layer_us"]),
|
|
]
|
|
|
|
# Build one wide DataFrame - metrics x saved names
|
|
_names = list(_saved.keys())
|
|
_rows_out = []
|
|
for _sec, _keys in _key_order:
|
|
_rows_out.append({"metric": f"= {_sec}",
|
|
**{n: "" for n in _names}})
|
|
for k in _keys:
|
|
_rows_out.append({
|
|
"metric": k,
|
|
**{n: _saved[n].get(k, "-") for n in _names},
|
|
})
|
|
_df_cmp = pd.DataFrame(_rows_out)
|
|
|
|
# Highlight best (lowest) for latency + used_gb rows, worst (over_budget True) red.
|
|
def _cmp_style(row):
|
|
styles = [""] * len(row)
|
|
m = row["metric"]
|
|
vals = row[1:]
|
|
try:
|
|
nums = {k: float(v) for k, v in vals.items()
|
|
if isinstance(v, (int, float)) and not isinstance(v, bool)}
|
|
except Exception:
|
|
nums = {}
|
|
if m in ("layer_us", "attn_us", "ffn_us", "used_gb",
|
|
"weights_gb", "kv_gb"):
|
|
if nums:
|
|
best = min(nums, key=nums.get)
|
|
for i, (k, v) in enumerate(vals.items(), start=1):
|
|
if k == best:
|
|
styles[i] = "background-color: #d4edda; font-weight:600;"
|
|
if m == "slack_gb" and nums:
|
|
best = max(nums, key=nums.get)
|
|
for i, (k, v) in enumerate(vals.items(), start=1):
|
|
if k == best:
|
|
styles[i] = "background-color: #d4edda; font-weight:600;"
|
|
if m == "over_budget":
|
|
for i, (k, v) in enumerate(vals.items(), start=1):
|
|
if v is True or v == "True":
|
|
styles[i] = "background-color: #f8d7da; font-weight:600;"
|
|
if m.startswith("= "):
|
|
styles = ["background-color: #e9ecef; font-weight:600;"] * len(row)
|
|
return styles
|
|
|
|
_styled_cmp = _df_cmp.style.apply(_cmp_style, axis=1)
|
|
st.dataframe(_styled_cmp, width='stretch', hide_index=True,
|
|
row_height=28)
|
|
st.caption(
|
|
"Green cell = best value in that row (lowest latency / memory, "
|
|
"highest slack). Red = over-budget. Grey rows = section headers."
|
|
)
|
|
|
|
# ── Per-stage side-by-side comparison ────────────────────
|
|
_fmt = _fmt_time # reuse the ns/us/ms auto-formatter
|
|
_short_label = {
|
|
"S1": "S1 RMSNorm", "S2": "S2 W_Q", "S3": "S3 W_K+W_V",
|
|
"S4": "S4 KV append", "S5": "S5 Q.K^T", "S6": "S6 softmax",
|
|
"S7": "S7 P.V", "S8": "S8 merge (+ AR if decode)",
|
|
"S9": "S9 norm O/l", "S10": "S10 W_O",
|
|
"C1": "C1 CP ring (prefill only)",
|
|
"C2": "C2 TP AR (W_O)",
|
|
"C3": "C3 Score AR (head-split)",
|
|
"F1": "F1 RMSNorm", "F2": "F2 W_gate", "F3": "F3 W_up",
|
|
"F4": "F4 SwiGLU", "F5": "F5 W_down", "CF1": "CF1 FFN AR",
|
|
}
|
|
|
|
def _stage_rows_across(prefix_order, key):
|
|
"""Build one row per stage prefix, with each saved config as a col."""
|
|
rows = []
|
|
for pfx in prefix_order:
|
|
row = {"stage": _short_label.get(pfx, pfx)}
|
|
sec_values = {}
|
|
for n, snap in _saved.items():
|
|
v = snap.get(key, {}).get(pfx)
|
|
if v is None:
|
|
row[n] = "-"
|
|
else:
|
|
row[n] = _fmt(v)
|
|
sec_values[n] = v
|
|
rows.append(row)
|
|
return rows
|
|
|
|
def _stage_style_row(row):
|
|
styles = [""] * len(row)
|
|
# Try to find min seconds — reconstruct from formatted strings
|
|
# is fragile; recompute from saved secs directly.
|
|
pfx = row["stage"].split()[0]
|
|
secs = {}
|
|
for n, snap in _saved.items():
|
|
key = "stages_attn" if pfx in _attn_prefixes else "stages_ffn"
|
|
v = snap.get(key, {}).get(pfx)
|
|
if v is not None and v > 0:
|
|
secs[n] = v
|
|
if secs:
|
|
best = min(secs, key=secs.get)
|
|
for i, name in enumerate(row.index[1:], start=1):
|
|
if name == best:
|
|
styles[i] = "background-color: #d4edda; font-weight:600;"
|
|
return styles
|
|
|
|
_attn_prefixes = ["S1", "S2", "S3", "S4", "S5", "S6", "S7",
|
|
"S8", "S9", "S10", "C1", "C2", "C3"]
|
|
_ffn_prefixes = ["F1", "F2", "F3", "F4", "F5", "CF1"]
|
|
|
|
st.markdown("**Attention per-stage latency (side-by-side)**")
|
|
_df_attn = pd.DataFrame(_stage_rows_across(_attn_prefixes, "stages_attn"))
|
|
st.dataframe(_df_attn.style.apply(_stage_style_row, axis=1),
|
|
width='stretch', hide_index=True, row_height=28)
|
|
|
|
st.markdown("**FFN per-stage latency (side-by-side)**")
|
|
_df_ffn = pd.DataFrame(_stage_rows_across(_ffn_prefixes, "stages_ffn"))
|
|
st.dataframe(_df_ffn.style.apply(_stage_style_row, axis=1),
|
|
width='stretch', hide_index=True, row_height=28)
|
|
|
|
st.caption(
|
|
"Times auto-scale (ns/us/ms). Green cell = fastest for that stage. "
|
|
"'-' means the stage is absent for that config (e.g. C1 in "
|
|
"decode, since the O/m/l all-reduce is folded into S8)."
|
|
)
|