mukesh
adc14c84af
New stage_shapes module reports per-PE input / weight / output tensor
shapes for every attention (S1..S10, C1..C3) and FFN (F1..F5, CF1)
stage. Streamlit's Stages tab picks these up below the existing
latency table + bar chart as two separate dataframes ('Attention
stages' and 'FFN stages'), so anyone who wants to see the exact
sharded tensor dims per stage can, without cluttering the latency
view.
Rows mirror the conditional structure of all_stages / all_ffn_stages:
S8 merge only when CP>1; C1 CP ring only in prefill with CP>1; C3
Score AR only when kv_shard_mode='split' and TP>H_kv; C2 TP AR only
when TP>1; CF1 FFN AR only when the FFN divisor > 1. Batch B flows
into every activation dim; weights stay B-invariant.
test_stage_shapes covers column presence, conditional row insertion
(decode omits C1; CP=1 omits S8; prefill+CP>1 inserts C1), batch
scaling of the leading dim, and that S5 output references both T_q
and S_local for the current cfg.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2115 lines
93 KiB
Python
2115 lines
93 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 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.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.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")
|
||
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 = st.tabs([
|
||
"Physical layout",
|
||
"Auto Suggest Parallelism",
|
||
"Memory breakdown", "Per-stage latency",
|
||
"Save & compare", "Auto Hardware",
|
||
])
|
||
|
||
|
||
# ── 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 = """
|
||
<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."
|
||
)
|
||
|
||
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'<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)
|
||
|
||
# ── Per-stage tensor shapes (per PE) ───────────────────────────
|
||
st.divider()
|
||
st.subheader("Per-stage tensor shapes (per PE)")
|
||
st.caption(
|
||
"Input / weight / output tensor dimensions each stage operates on, "
|
||
"*per PE* (already sharded by CP, TP, PP, EP as configured). Rows "
|
||
"match the latency table above; comm-only stages show payload in "
|
||
"the Input column."
|
||
)
|
||
st.markdown("**Attention stages**")
|
||
st.dataframe(pd.DataFrame(attn_stage_shape_rows(cfg)),
|
||
width='stretch', hide_index=True)
|
||
st.markdown("**FFN stages**")
|
||
st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)),
|
||
width='stretch', hide_index=True)
|
||
|
||
|
||
# ── 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()
|