mukesh
77abc95d78
Two more table sections on the Capacity planning tab: **Section 4 — Practical rules of thumb (by context regime)** - Short (S_kv < L*): pack HBM, B ~ 2·B*, high util, cheap tier - Long (S_kv > L*): fewer users/replica, low B, util drops as 1/(1 + S_kv/L*), pricier - Extreme (S_kv >> L*): dedicated pool, heavy CP, disaggregated prefill, very low util without sparse attention Columns: Regime, Batch strategy, Utilization, Cost/token, Deployment. **Section 5 — Sample deployment templates** Same base model, three different sharding recipes routed to by the API gateway based on request context length: - Config_small : CP=1, TP=8, PP=1 → 8 GPUs, up to 32k, B=64 - Config_medium: CP=4, TP=8, PP=1 → 32 GPUs, up to 128k, B=32 - Config_large : CP=32, TP=8, PP=1 → 256 GPUs, up to 1M, B=4 Columns: Tier, CP, TP, PP, Total GPUs/replica, Max context, Typical B, Best for. Playbook section renumbered to #6. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2918 lines
130 KiB
Python
2918 lines
130 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 dataclasses
|
||
import importlib
|
||
import sys
|
||
|
||
import matplotlib.pyplot as plt
|
||
import pandas as pd
|
||
import streamlit as st
|
||
|
||
# Streamlit's hot-reload only re-runs THIS script; imported sub-modules stay
|
||
# cached in sys.modules across reruns. When we edit auto_explore.py or
|
||
# auto_hardware.py while the Streamlit process is alive, the app would
|
||
# keep using the pre-edit versions until a full Ctrl+C + restart. Force
|
||
# a reload of our own modules on every rerun so signature changes land
|
||
# without a full restart.
|
||
# Order matters: model_config first because everything else imports from it
|
||
# (e.g. TopologyConfig). If we reload downstream modules while model_config
|
||
# is stale, they end up holding the OLD dataclass definition and every
|
||
# `cfg.topo.new_field` lookup blows up with AttributeError.
|
||
for _mod_name in (
|
||
"tests.analytical_visualization.model_config",
|
||
"tests.analytical_visualization.model_presets",
|
||
"tests.analytical_visualization.autosuggest",
|
||
"tests.analytical_visualization.memory_layout",
|
||
"tests.analytical_visualization.stage_latencies",
|
||
"tests.analytical_visualization.stage_shapes",
|
||
"tests.analytical_visualization.chip_roofline",
|
||
"tests.analytical_visualization.auto_explore",
|
||
"tests.analytical_visualization.auto_hardware",
|
||
"tests.analytical_visualization.pe_weight_layout",
|
||
"tests.analytical_visualization.tensor_sharding",
|
||
"tests.analytical_visualization.topology_map",
|
||
"tests.analytical_visualization.pipeline_diagram",
|
||
"tests.analytical_visualization.optimization_report",
|
||
):
|
||
_m = sys.modules.get(_mod_name)
|
||
if _m is not None:
|
||
importlib.reload(_m)
|
||
|
||
from tests.analytical_visualization.model_config import (
|
||
FullConfig, ModelConfig, TopologyConfig, MachineParams,
|
||
)
|
||
from tests.analytical_visualization.model_presets import PRESETS
|
||
from tests.analytical_visualization.stage_latencies import all_stages, all_ffn_stages
|
||
from tests.analytical_visualization.stage_shapes import (
|
||
attn_stage_shape_rows,
|
||
ffn_stage_shape_rows,
|
||
)
|
||
from tests.analytical_visualization.chip_roofline import (
|
||
ai_sensitivity_curve,
|
||
arithmetic_intensity,
|
||
balance_context,
|
||
bound_regime,
|
||
critical_batch,
|
||
good_batch,
|
||
good_context,
|
||
knee_batch,
|
||
max_batch_within_slo,
|
||
memory_budget_curve_vs_batch,
|
||
memory_budget_curve_vs_skv,
|
||
per_token_latency_curve,
|
||
size_deployment,
|
||
step_latency_curve,
|
||
t_com,
|
||
t_mem_long,
|
||
t_mem_short,
|
||
total_active_params,
|
||
utilization_at,
|
||
)
|
||
from tests.analytical_visualization.memory_layout import (
|
||
compute_memory,
|
||
attention_weight_rows,
|
||
ffn_weight_rows,
|
||
kv_cache_rows,
|
||
sum_bytes_all_layers,
|
||
total_weight_bytes_full_model,
|
||
total_kv_bytes_full_model,
|
||
)
|
||
from tests.analytical_visualization.topology_map import draw_topology
|
||
from tests.analytical_visualization.pe_weight_layout import (
|
||
draw_pe_layout, _per_pe_bytes, _q_heads_for_pe, _kv_heads_for_pe,
|
||
)
|
||
from tests.analytical_visualization.tensor_sharding import draw_tensor_sharding
|
||
from tests.analytical_visualization.pipeline_diagram import draw_pipeline
|
||
from tests.analytical_visualization.optimization_report import (
|
||
replication_report, optimization_hints,
|
||
)
|
||
from tests.analytical_visualization.autosuggest import auto_suggest
|
||
|
||
|
||
st.set_page_config(page_title="SIP Attention Analytical Model",
|
||
layout="wide")
|
||
|
||
# Tighter spacing so more info fits per page.
|
||
st.markdown(
|
||
"""
|
||
<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,
|
||
tab_roofline, tab_planning) = st.tabs([
|
||
"Physical layout",
|
||
"Auto Suggest Parallelism",
|
||
"Memory breakdown", "Per-stage latency",
|
||
"Save & compare", "Auto Hardware",
|
||
"Chip roofline & B*",
|
||
"Capacity planning",
|
||
])
|
||
|
||
|
||
# ── 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)
|
||
|
||
# Row G: All tensor shapes (per PE) — one collapsed expander with
|
||
# weight, KV, and per-stage shapes for anyone who wants them.
|
||
with st.expander("All tensor shapes (per PE)", expanded=False):
|
||
st.caption(
|
||
"Every tensor this deployment holds, per PE. Weights + KV "
|
||
"show global vs per-PE shape (one row per tensor, per layer). "
|
||
"Per-stage shapes show input / weight / output tensors each "
|
||
"stage of the forward pass operates on."
|
||
)
|
||
_lp = (model.layers + topo.pp - 1) // topo.pp
|
||
|
||
st.markdown(f"**Attention weights** (per layer, layers/stage = {_lp})")
|
||
_attn_rows_layout = attention_weight_rows(cfg)
|
||
_display_attn_layout = [
|
||
{k: v for k, v in r.items() if not k.startswith("_")}
|
||
for r in _attn_rows_layout
|
||
]
|
||
st.dataframe(pd.DataFrame(_display_attn_layout),
|
||
width='stretch', hide_index=True)
|
||
|
||
st.markdown(f"**FFN weights** (per layer, layers/stage = {_lp})")
|
||
_ffn_rows_layout = ffn_weight_rows(cfg)
|
||
_display_ffn_layout = [
|
||
{k: v for k, v in r.items() if not k.startswith("_")}
|
||
for r in _ffn_rows_layout
|
||
]
|
||
st.dataframe(pd.DataFrame(_display_ffn_layout),
|
||
width='stretch', hide_index=True)
|
||
|
||
st.markdown(f"**KV cache** (per layer, S_kv = {s_kv:,})")
|
||
_kv_rows_layout = kv_cache_rows(cfg)
|
||
_display_kv_layout = [
|
||
{k: v for k, v in r.items() if not k.startswith("_")}
|
||
for r in _kv_rows_layout
|
||
]
|
||
st.dataframe(pd.DataFrame(_display_kv_layout),
|
||
width='stretch', hide_index=True)
|
||
|
||
st.markdown("**Per-stage attention shapes**")
|
||
st.dataframe(pd.DataFrame(attn_stage_shape_rows(cfg)),
|
||
width='stretch', hide_index=True)
|
||
|
||
st.markdown("**Per-stage FFN shapes**")
|
||
st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)),
|
||
width='stretch', hide_index=True)
|
||
|
||
|
||
# ── TAB 2: Memory breakdown ─────────────────────────────────────
|
||
with tab_memory:
|
||
mem = compute_memory(cfg)
|
||
layers_per_stage = (model.layers + pp - 1) // pp
|
||
|
||
parallelism_str = (
|
||
f"CP={cp}, TP={tp}, PP={pp}, DP={dp}, EP={ep}"
|
||
f" | layers={model.layers}, layers/stage={layers_per_stage}"
|
||
f" | KV mode: {kv_shard_mode}"
|
||
)
|
||
st.caption(f"**Parallelism:** {parallelism_str}")
|
||
|
||
# ─── ATTENTION (per-layer) ─────────────────────────────────
|
||
st.subheader(f"Attention weights (per layer) — {parallelism_str}")
|
||
attn_rows = attention_weight_rows(cfg)
|
||
display_attn = [{k: v for k, v in r.items() if not k.startswith("_")}
|
||
for r in attn_rows]
|
||
st.dataframe(pd.DataFrame(display_attn), width='stretch', hide_index=True)
|
||
|
||
attn_bytes_1L_global = sum(r["_p_global"] for r in attn_rows) * model.bytes_per_elem
|
||
attn_bytes_1L_pe = sum(r["_p_per_pe"] for r in attn_rows) * model.bytes_per_elem
|
||
attn_bytes_all_global = attn_bytes_1L_global * model.layers
|
||
attn_bytes_all_pe = attn_bytes_1L_pe * layers_per_stage
|
||
|
||
ac1, ac2, ac3, ac4 = st.columns(4)
|
||
ac1.metric("Attn / layer (global)", f"{attn_bytes_1L_global/1e6:.2f} MB")
|
||
ac2.metric(f"Attn all {model.layers} layers", f"{attn_bytes_all_global/1e9:.3f} GB")
|
||
ac3.metric("Attn / layer (per PE)", f"{attn_bytes_1L_pe/1e6:.2f} MB")
|
||
ac4.metric(f"Attn per PE ({layers_per_stage} layers/stage)",
|
||
f"{attn_bytes_all_pe/1e9:.3f} GB")
|
||
|
||
# ─── FFN / MoE (per-layer) ─────────────────────────────────
|
||
is_moe = "MoE" in (preset.family + preset.note)
|
||
ffn_label = "FFN (activated experts approx)" if is_moe else "FFN"
|
||
st.subheader(f"{ffn_label} weights (per layer) — {parallelism_str}")
|
||
ffn_rows = ffn_weight_rows(cfg)
|
||
display_ffn = [{k: v for k, v in r.items() if not k.startswith("_")}
|
||
for r in ffn_rows]
|
||
st.dataframe(pd.DataFrame(display_ffn), width='stretch', hide_index=True)
|
||
|
||
ffn_bytes_1L_global = sum(r["_p_global"] for r in ffn_rows) * model.bytes_per_elem
|
||
ffn_bytes_1L_pe = sum(r["_p_per_pe"] for r in ffn_rows) * model.bytes_per_elem
|
||
ffn_bytes_all_global = ffn_bytes_1L_global * model.layers
|
||
ffn_bytes_all_pe = ffn_bytes_1L_pe * layers_per_stage
|
||
|
||
fc1, fc2, fc3, fc4 = st.columns(4)
|
||
fc1.metric("FFN / layer (global)", f"{ffn_bytes_1L_global/1e6:.2f} MB")
|
||
fc2.metric(f"FFN all {model.layers} layers", f"{ffn_bytes_all_global/1e9:.3f} GB")
|
||
fc3.metric("FFN / layer (per PE)", f"{ffn_bytes_1L_pe/1e6:.2f} MB")
|
||
fc4.metric(f"FFN per PE ({layers_per_stage} layers/stage)",
|
||
f"{ffn_bytes_all_pe/1e9:.3f} GB")
|
||
|
||
# ─── KV cache (per-layer) ─────────────────────────────────
|
||
st.subheader(f"KV cache (per layer, S_kv = {s_kv:,}) — {parallelism_str}")
|
||
kv_rows = kv_cache_rows(cfg)
|
||
display_kv = [{k: v for k, v in r.items() if not k.startswith("_")}
|
||
for r in kv_rows]
|
||
st.dataframe(pd.DataFrame(display_kv), width='stretch', hide_index=True)
|
||
|
||
kv_bytes_1L_global = sum(r["_p_global"] for r in kv_rows) * model.bytes_per_elem
|
||
kv_bytes_1L_pe = sum(r["_p_per_pe"] for r in kv_rows) * model.bytes_per_elem
|
||
kv_bytes_all_global = kv_bytes_1L_global * model.layers
|
||
kv_bytes_all_pe = kv_bytes_1L_pe * layers_per_stage
|
||
|
||
kc1, kc2, kc3, kc4 = st.columns(4)
|
||
kc1.metric("KV / layer (global)", f"{kv_bytes_1L_global/1e6:.2f} MB")
|
||
kc2.metric(f"KV all {model.layers} layers", f"{kv_bytes_all_global/1e9:.3f} GB")
|
||
kc3.metric("KV / layer (per PE)", f"{kv_bytes_1L_pe/1e6:.2f} MB")
|
||
kc4.metric(f"KV per PE ({layers_per_stage} layers/stage)",
|
||
f"{kv_bytes_all_pe/1e9:.3f} GB")
|
||
|
||
st.divider()
|
||
|
||
# ─── FULL-MODEL TOTALS ───────────────────────────────────
|
||
st.subheader("Full-model totals (all layers)")
|
||
total_w_global = attn_bytes_all_global + ffn_bytes_all_global
|
||
total_all_global = total_w_global + kv_bytes_all_global
|
||
total_per_pe = attn_bytes_all_pe + ffn_bytes_all_pe + kv_bytes_all_pe
|
||
tc1, tc2, tc3, tc4 = st.columns(4)
|
||
tc1.metric("Weights (all layers, global)", f"{total_w_global/1e9:.2f} GB")
|
||
tc2.metric("KV (all layers, global)", f"{kv_bytes_all_global/1e9:.2f} GB")
|
||
tc3.metric("Weights + KV (global)", f"{total_all_global/1e9:.2f} GB")
|
||
tc4.metric("Per-PE total (W+KV, layers/stage)",
|
||
f"{total_per_pe/1e9:.2f} GB")
|
||
|
||
# ─── Per-PE memory pie ───────────────────────────────────
|
||
st.subheader("Per-PE memory footprint (with transient + slack)")
|
||
pie_col, info_col = st.columns([1, 1])
|
||
with pie_col:
|
||
labels = ["Weights", "KV cache", "Transient", "Slack"]
|
||
vals = [mem.weights_bytes, mem.kv_cache_bytes,
|
||
mem.transient_bytes, mem.slack_bytes]
|
||
colors = ["#3a86ff", "#8338ec", "#ffbe0b", "#adb5bd"]
|
||
|
||
fig_mem, ax_mem = plt.subplots(figsize=(6, 6))
|
||
ax_mem.pie(vals,
|
||
labels=[f"{l}\n{v/1e9:.2f} GB" for l, v in zip(labels, vals)],
|
||
colors=colors, autopct="%1.1f%%", startangle=90)
|
||
ax_mem.set_title(f"Per-PE (budget: {mem.budget_bytes/1e9:.1f} GB)")
|
||
st.pyplot(fig_mem, width='stretch')
|
||
plt.close(fig_mem)
|
||
with info_col:
|
||
if mem.over_budget:
|
||
st.error(f"OVER BUDGET by "
|
||
f"{(mem.used_bytes - mem.budget_bytes)/1e9:.2f} GB")
|
||
else:
|
||
st.success(f"{mem.slack_bytes/1e9:.2f} GB slack")
|
||
|
||
st.markdown(
|
||
f"- Weights: **{mem.weights_bytes/1e9:.2f} GB** "
|
||
f"(layers/PP = {layers_per_stage})\n"
|
||
f"- KV cache: **{mem.kv_cache_bytes/1e9:.2f} GB** "
|
||
f"(S_local = {topo.s_local:,})\n"
|
||
f"- Transient: {mem.transient_bytes/1e9:.2f} GB\n"
|
||
f"- Slack: {mem.slack_bytes/1e9:.2f} GB"
|
||
)
|
||
|
||
|
||
# ── TAB 3: Per-stage latency ────────────────────────────────────
|
||
with tab_stages:
|
||
st.subheader("Per-stage attention latency")
|
||
stages = all_stages(cfg)
|
||
|
||
rows = []
|
||
for s in stages:
|
||
rows.append({
|
||
"Stage": s.name,
|
||
"Bound": s.bound,
|
||
"Compute (us)": round(s.compute_s * 1e6, 3),
|
||
"Memory (us)": round(s.memory_s * 1e6, 3),
|
||
"Comm (us)": round(s.comm_s * 1e6, 3),
|
||
"Visible (us)": round(s.visible_s * 1e6, 3),
|
||
"Formula": s.formula,
|
||
})
|
||
df = pd.DataFrame(rows)
|
||
st.dataframe(df, width='stretch', hide_index=True)
|
||
|
||
total_visible_s = sum(s.visible_s for s in stages)
|
||
total_layers = total_visible_s * model.layers
|
||
st.markdown(
|
||
f"**Per-layer attention (analytical): {total_visible_s*1e6:.2f} us** \n"
|
||
f"**Full-model attention ({model.layers} layers): "
|
||
f"{total_layers*1e6:.2f} us = {total_layers*1e3:.2f} ms**"
|
||
)
|
||
|
||
st.subheader("Stage cost breakdown")
|
||
fig_bar, ax_bar = plt.subplots(figsize=(12, 5))
|
||
names = [s.name for s in stages]
|
||
cmp_vals = [s.compute_s * 1e6 for s in stages]
|
||
mem_vals = [s.memory_s * 1e6 for s in stages]
|
||
comm_vals = [s.comm_s * 1e6 for s in stages]
|
||
|
||
x = range(len(names))
|
||
ax_bar.bar(x, cmp_vals, label="Compute", color="#3a86ff")
|
||
ax_bar.bar(x, mem_vals, bottom=cmp_vals, label="Memory", color="#ffbe0b")
|
||
ax_bar.bar(x, comm_vals,
|
||
bottom=[c + m for c, m in zip(cmp_vals, mem_vals)],
|
||
label="Comm", color="#d90429")
|
||
ax_bar.set_xticks(x)
|
||
ax_bar.set_xticklabels(names, rotation=45, ha="right", fontsize=8)
|
||
ax_bar.set_ylabel("Time (us)")
|
||
ax_bar.set_title(f"Stage breakdown - mode={mode}, T_q={topo.T_q}")
|
||
ax_bar.legend()
|
||
ax_bar.grid(axis="y", alpha=0.3)
|
||
plt.tight_layout()
|
||
st.pyplot(fig_bar, width='stretch')
|
||
plt.close(fig_bar)
|
||
|
||
|
||
# ── TAB 4: Save & compare ───────────────────────────────────────
|
||
def _snapshot_current_config():
|
||
"""Capture a compact snapshot of the current cfg + key metrics."""
|
||
_stages_attn = all_stages(cfg)
|
||
_stages_ffn = all_ffn_stages(cfg)
|
||
attn_us = sum(s.visible_s for s in _stages_attn) * 1e6
|
||
ffn_us = sum(s.visible_s for s in _stages_ffn) * 1e6
|
||
mem = compute_memory(cfg)
|
||
# Per-stage latencies keyed by short prefix (S1, S2, ..., C1, F1, ..., CF1).
|
||
# Value = seconds (raw), so display can format uniformly across snapshots.
|
||
_by_prefix_attn = {s.name.split()[0]: s.visible_s for s in _stages_attn}
|
||
_by_prefix_ffn = {s.name.split()[0]: s.visible_s for s in _stages_ffn}
|
||
# Preserve the full descriptive name too so a hover/second column can
|
||
# show what the abbreviation means for this particular snapshot.
|
||
_fullname_attn = {s.name.split()[0]: s.name for s in _stages_attn}
|
||
_fullname_ffn = {s.name.split()[0]: s.name for s in _stages_ffn}
|
||
return {
|
||
"model": model.name,
|
||
"mode": topo.mode,
|
||
"s_kv": topo.s_kv,
|
||
"cp": topo.cp, "tp": topo.tp,
|
||
"pp": topo.pp, "dp": topo.dp, "ep": topo.ep,
|
||
"cp_placement": topo.cp_placement,
|
||
"tp_placement": topo.tp_placement,
|
||
"ffn_scope": topo.ffn_shard_scope,
|
||
"kv_mode": topo.kv_shard_mode,
|
||
"sip_topology": topo.sip_topology,
|
||
"cp_ring_variant": topo.cp_ring_variant,
|
||
"pes": topo.total_pes,
|
||
"cubes": topo.cubes_used,
|
||
"sips": topo.sips_used,
|
||
"pe_hbm_gb": machine.pe_hbm_gb,
|
||
"peak_tflops": machine.peak_tflops_f16,
|
||
"bw_hbm_gbs": machine.bw_hbm_gbs,
|
||
"bw_inter_gbs": machine.bw_inter_gbs,
|
||
"bw_intersip_gbs": machine.bw_intersip_gbs,
|
||
"weights_gb": round(mem.weights_bytes / 1e9, 3),
|
||
"kv_gb": round(mem.kv_cache_bytes / 1e9, 3),
|
||
"transient_gb": round(mem.transient_bytes / 1e9, 3),
|
||
"used_gb": round(mem.used_bytes / 1e9, 3),
|
||
"slack_gb": round(mem.slack_bytes / 1e9, 3),
|
||
"over_budget": mem.over_budget,
|
||
"attn_us": round(attn_us, 3),
|
||
"ffn_us": round(ffn_us, 3),
|
||
"layer_us": round(attn_us + ffn_us, 3),
|
||
"stages_attn": _by_prefix_attn,
|
||
"stages_ffn": _by_prefix_ffn,
|
||
"stage_names_attn": _fullname_attn,
|
||
"stage_names_ffn": _fullname_ffn,
|
||
}
|
||
|
||
|
||
with tab_compare:
|
||
st.markdown(
|
||
"Save the current configuration under a name, then load it back "
|
||
"later or compare multiple saved configs side-by-side."
|
||
)
|
||
if "saved_cfgs" not in st.session_state:
|
||
st.session_state["saved_cfgs"] = {} # name -> snapshot dict
|
||
|
||
# Default name: next free config1, config2, ...
|
||
def _next_config_name():
|
||
i = 1
|
||
while f"config{i}" in st.session_state.get("saved_cfgs", {}):
|
||
i += 1
|
||
return f"config{i}"
|
||
|
||
sv_c1, sv_c2, sv_c3 = st.columns([3, 1, 1])
|
||
with sv_c1:
|
||
_save_name = st.text_input(
|
||
"Config name",
|
||
value=_next_config_name(),
|
||
key=f"cfg_save_name_{len(st.session_state.get('saved_cfgs', {}))}",
|
||
)
|
||
with sv_c2:
|
||
if st.button("Save current", width='stretch'):
|
||
if _save_name.strip():
|
||
st.session_state["saved_cfgs"][_save_name.strip()] = \
|
||
_snapshot_current_config()
|
||
st.success(f"Saved '{_save_name.strip()}'")
|
||
st.rerun()
|
||
else:
|
||
st.warning("Give it a name first.")
|
||
with sv_c3:
|
||
if st.button("Clear all", width='stretch'):
|
||
st.session_state["saved_cfgs"] = {}
|
||
st.rerun()
|
||
|
||
_saved_all = st.session_state.get("saved_cfgs", {})
|
||
if not _saved_all:
|
||
st.info("No saved configurations yet.")
|
||
else:
|
||
st.markdown(f"**{len(_saved_all)} saved configuration(s)** - "
|
||
f"pick which to compare below")
|
||
|
||
# Pick which configs to include in the comparison (default all)
|
||
_picked = st.multiselect(
|
||
"Compare",
|
||
options=list(_saved_all.keys()),
|
||
default=list(_saved_all.keys()),
|
||
key="compare_picked",
|
||
)
|
||
_saved = {n: _saved_all[n] for n in _picked}
|
||
|
||
# Delete individual
|
||
with st.expander("Manage (delete individual)", expanded=False):
|
||
for _n in list(_saved_all.keys()):
|
||
d1, d2, d3 = st.columns([5, 1, 1])
|
||
d1.text(_n)
|
||
if d2.button("Rename", key=f"rn_btn_{_n}"):
|
||
st.session_state[f"rn_active_{_n}"] = True
|
||
if d3.button("Delete", key=f"del_{_n}"):
|
||
del st.session_state["saved_cfgs"][_n]
|
||
st.rerun()
|
||
if st.session_state.get(f"rn_active_{_n}"):
|
||
new_name = st.text_input(
|
||
f"New name for '{_n}'", value=_n,
|
||
key=f"rn_txt_{_n}",
|
||
)
|
||
if st.button("Confirm rename", key=f"rn_ok_{_n}"):
|
||
if new_name and new_name != _n and new_name not in _saved_all:
|
||
st.session_state["saved_cfgs"][new_name] = \
|
||
st.session_state["saved_cfgs"].pop(_n)
|
||
st.session_state[f"rn_active_{_n}"] = False
|
||
st.rerun()
|
||
|
||
if not _saved:
|
||
st.info("No configurations selected. Tick some above to compare.")
|
||
st.stop()
|
||
|
||
# Comparison table: rows = metrics, columns = config names
|
||
_key_order = [
|
||
("Model / workload", ["model", "mode", "s_kv"]),
|
||
("Parallelism", ["cp", "tp", "pp", "dp", "ep",
|
||
"cp_placement", "tp_placement",
|
||
"ffn_scope", "kv_mode", "cp_ring_variant",
|
||
"sip_topology"]),
|
||
("Physical", ["pes", "cubes", "sips"]),
|
||
("Hardware knobs", ["pe_hbm_gb", "peak_tflops",
|
||
"bw_hbm_gbs", "bw_inter_gbs",
|
||
"bw_intersip_gbs"]),
|
||
("Memory / PE (GB)", ["weights_gb", "kv_gb", "transient_gb",
|
||
"used_gb", "slack_gb", "over_budget"]),
|
||
("Per-layer latency (us)", ["attn_us", "ffn_us", "layer_us"]),
|
||
]
|
||
|
||
# Build one wide DataFrame - metrics x saved names
|
||
_names = list(_saved.keys())
|
||
_rows_out = []
|
||
for _sec, _keys in _key_order:
|
||
_rows_out.append({"metric": f"= {_sec}",
|
||
**{n: "" for n in _names}})
|
||
for k in _keys:
|
||
_rows_out.append({
|
||
"metric": k,
|
||
**{n: _saved[n].get(k, "-") for n in _names},
|
||
})
|
||
_df_cmp = pd.DataFrame(_rows_out)
|
||
|
||
# Highlight best (lowest) for latency + used_gb rows, worst (over_budget True) red.
|
||
def _cmp_style(row):
|
||
styles = [""] * len(row)
|
||
m = row["metric"]
|
||
vals = row[1:]
|
||
try:
|
||
nums = {k: float(v) for k, v in vals.items()
|
||
if isinstance(v, (int, float)) and not isinstance(v, bool)}
|
||
except Exception:
|
||
nums = {}
|
||
if m in ("layer_us", "attn_us", "ffn_us", "used_gb",
|
||
"weights_gb", "kv_gb"):
|
||
if nums:
|
||
best = min(nums, key=nums.get)
|
||
for i, (k, v) in enumerate(vals.items(), start=1):
|
||
if k == best:
|
||
styles[i] = "background-color: #d4edda; font-weight:600;"
|
||
if m == "slack_gb" and nums:
|
||
best = max(nums, key=nums.get)
|
||
for i, (k, v) in enumerate(vals.items(), start=1):
|
||
if k == best:
|
||
styles[i] = "background-color: #d4edda; font-weight:600;"
|
||
if m == "over_budget":
|
||
for i, (k, v) in enumerate(vals.items(), start=1):
|
||
if v is True or v == "True":
|
||
styles[i] = "background-color: #f8d7da; font-weight:600;"
|
||
if m.startswith("= "):
|
||
styles = ["background-color: #e9ecef; font-weight:600;"] * len(row)
|
||
return styles
|
||
|
||
_styled_cmp = _df_cmp.style.apply(_cmp_style, axis=1)
|
||
st.dataframe(_styled_cmp, width='stretch', hide_index=True,
|
||
row_height=28)
|
||
st.caption(
|
||
"Green cell = best value in that row (lowest latency / memory, "
|
||
"highest slack). Red = over-budget. Grey rows = section headers."
|
||
)
|
||
|
||
# ── Per-stage side-by-side comparison ────────────────────
|
||
_fmt = _fmt_time # reuse the ns/us/ms auto-formatter
|
||
_short_label = {
|
||
"S1": "S1 RMSNorm", "S2": "S2 W_Q", "S3": "S3 W_K+W_V",
|
||
"S4": "S4 KV append", "S5": "S5 Q.K^T", "S6": "S6 softmax",
|
||
"S7": "S7 P.V", "S8": "S8 merge (+ AR if decode)",
|
||
"S9": "S9 norm O/l", "S10": "S10 W_O",
|
||
"C1": "C1 CP ring (prefill only)",
|
||
"C2": "C2 TP AR (W_O)",
|
||
"C3": "C3 Score AR (head-split)",
|
||
"F1": "F1 RMSNorm", "F2": "F2 W_gate", "F3": "F3 W_up",
|
||
"F4": "F4 SwiGLU", "F5": "F5 W_down", "CF1": "CF1 FFN AR",
|
||
}
|
||
|
||
def _stage_rows_across(prefix_order, key):
|
||
"""Build one row per stage prefix, with each saved config as a col."""
|
||
rows = []
|
||
for pfx in prefix_order:
|
||
row = {"stage": _short_label.get(pfx, pfx)}
|
||
sec_values = {}
|
||
for n, snap in _saved.items():
|
||
v = snap.get(key, {}).get(pfx)
|
||
if v is None:
|
||
row[n] = "-"
|
||
else:
|
||
row[n] = _fmt(v)
|
||
sec_values[n] = v
|
||
rows.append(row)
|
||
return rows
|
||
|
||
def _stage_style_row(row):
|
||
styles = [""] * len(row)
|
||
# Try to find min seconds — reconstruct from formatted strings
|
||
# is fragile; recompute from saved secs directly.
|
||
pfx = row["stage"].split()[0]
|
||
secs = {}
|
||
for n, snap in _saved.items():
|
||
key = "stages_attn" if pfx in _attn_prefixes else "stages_ffn"
|
||
v = snap.get(key, {}).get(pfx)
|
||
if v is not None and v > 0:
|
||
secs[n] = v
|
||
if secs:
|
||
best = min(secs, key=secs.get)
|
||
for i, name in enumerate(row.index[1:], start=1):
|
||
if name == best:
|
||
styles[i] = "background-color: #d4edda; font-weight:600;"
|
||
return styles
|
||
|
||
_attn_prefixes = ["S1", "S2", "S3", "S4", "S5", "S6", "S7",
|
||
"S8", "S9", "S10", "C1", "C2", "C3"]
|
||
_ffn_prefixes = ["F1", "F2", "F3", "F4", "F5", "CF1"]
|
||
|
||
st.markdown("**Attention per-stage latency (side-by-side)**")
|
||
_df_attn = pd.DataFrame(_stage_rows_across(_attn_prefixes, "stages_attn"))
|
||
st.dataframe(_df_attn.style.apply(_stage_style_row, axis=1),
|
||
width='stretch', hide_index=True, row_height=28)
|
||
|
||
st.markdown("**FFN per-stage latency (side-by-side)**")
|
||
_df_ffn = pd.DataFrame(_stage_rows_across(_ffn_prefixes, "stages_ffn"))
|
||
st.dataframe(_df_ffn.style.apply(_stage_style_row, axis=1),
|
||
width='stretch', hide_index=True, row_height=28)
|
||
|
||
st.caption(
|
||
"Times auto-scale (ns/us/ms). Green cell = fastest for that stage. "
|
||
"'-' means the stage is absent for that config (e.g. C1 in "
|
||
"decode, since the O/m/l all-reduce is folded into S8)."
|
||
)
|
||
|
||
|
||
# ── TAB (Auto Suggest Parallelism) — one tab, two sweep buttons ──
|
||
|
||
|
||
def _render_auto_explore_tab():
|
||
"""Render the Auto Suggest Parallelism tab body.
|
||
|
||
Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
|
||
result under its own session_state key. Whichever button was last
|
||
clicked drives the display below; both caches persist so re-clicking
|
||
either without changing model/workload just re-shows the cached view.
|
||
"""
|
||
from tests.analytical_visualization.auto_explore import (
|
||
compute_parallelism_sensitivity,
|
||
run_auto_explore,
|
||
)
|
||
|
||
st.markdown(
|
||
"Sweep every valid parallelism config (CP, TP, PP, DP, kv_shard, "
|
||
"ffn_scope, tp_placement, cp_placement, cp_ring) for the currently-"
|
||
"selected model + workload and extract the 3D Pareto frontier on "
|
||
"(latency ↓, PEs ↓, efficiency ↑). Throughput = 1/latency for a "
|
||
"single request so it collapses with latency and is shown for info only."
|
||
)
|
||
st.caption(
|
||
"**Attention only**: drop FFN / MoE from the summed latency — "
|
||
"useful when isolating attention-kernel tuning (matches what our "
|
||
"kernbench attention sim measures). \n"
|
||
"**Attn + FFN/MoE**: sum every stage — the full per-token cost."
|
||
)
|
||
|
||
st.markdown(
|
||
f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}** | "
|
||
f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**"
|
||
)
|
||
|
||
_b1, _b2, _b3, _b4 = st.columns([1, 1, 1, 1])
|
||
with _b1:
|
||
_run_attn = st.button("Run sweep — Attention",
|
||
type="primary", width='stretch',
|
||
key="_auto_run_attn")
|
||
with _b2:
|
||
_run_ffn = st.button("Run sweep — FFN/MoE",
|
||
type="primary", width='stretch',
|
||
key="_auto_run_ffn")
|
||
with _b3:
|
||
_run_full = st.button("Run sweep — Attn + FFN/MoE",
|
||
type="primary", width='stretch',
|
||
key="_auto_run_full")
|
||
|
||
_scope_meta = {
|
||
"attn": (True, False, "Attention only"),
|
||
"ffn": (False, True, "FFN / MoE only"),
|
||
"full": (True, True, "Attn + FFN/MoE"),
|
||
}
|
||
_button_click = {"attn": _run_attn, "ffn": _run_ffn, "full": _run_full}
|
||
|
||
for _scope, _clicked in _button_click.items():
|
||
if not _clicked:
|
||
continue
|
||
_ia, _ff, _lbl = _scope_meta[_scope]
|
||
with st.spinner(f"Sweeping ~28k configs ({_lbl})..."):
|
||
_r = run_auto_explore(
|
||
model, machine, s_kv=s_kv, mode=mode,
|
||
include_attention=_ia, include_ffn=_ff, b=b_batch,
|
||
)
|
||
st.session_state[f"_auto_explore_result_{_scope}"] = _r
|
||
st.session_state[f"_auto_explore_ctx_{_scope}"] = (
|
||
model.name, s_kv, mode, machine.pe_hbm_gb,
|
||
)
|
||
st.session_state["_auto_explore_active"] = _scope
|
||
|
||
_active = st.session_state.get("_auto_explore_active")
|
||
if _active is None:
|
||
st.info(
|
||
"Click one of the run buttons above to enumerate valid "
|
||
"parallelism configurations and rank them on the Pareto "
|
||
"frontier. Takes ~5–10 s per scope. All button results are "
|
||
"cached so you can flip between scopes without re-running."
|
||
)
|
||
return
|
||
|
||
# Downstream code expects _suffix, _label, include_attention, include_ffn,
|
||
# _res, _ctx_key.
|
||
include_attention, include_ffn, _label = _scope_meta[_active]
|
||
_suffix = f"_{_active}"
|
||
_result_key = f"_auto_explore_result{_suffix}"
|
||
_ctx_key = f"_auto_explore_ctx{_suffix}"
|
||
_res = st.session_state.get(_result_key)
|
||
if _res is None:
|
||
st.info(f"Click **Run sweep — {_label.split(' ')[0]}** to compute this view.")
|
||
return
|
||
|
||
st.markdown(f"### Currently viewing: **{_label}**")
|
||
|
||
# Warn if the cached result is stale relative to the current config.
|
||
_cached_ctx = st.session_state.get(_ctx_key)
|
||
_cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb)
|
||
if _cached_ctx != _cur_ctx:
|
||
st.warning(
|
||
f"Cached {_label} result is from a different model/workload. "
|
||
f"Click **Run sweep — {_label.split(' ')[0]}** to refresh."
|
||
)
|
||
|
||
_m1, _m2, _m3, _m4 = st.columns(4)
|
||
_m1.metric("Enumerated", f"{_res.total_enumerated:,}")
|
||
_m2.metric("Feasible", f"{_res.total_feasible:,}")
|
||
_m3.metric("Pareto configs", len(_res.pareto_scores))
|
||
if _res.pareto_scores:
|
||
_best = min(_res.pareto_scores,
|
||
key=lambda s: (s.total_latency_ns,
|
||
s.pes_used,
|
||
s.hbm_utilization))
|
||
_m4.metric("Best latency", f"{_best.latency_ms:.2f} ms")
|
||
|
||
if not _res.pareto_scores:
|
||
st.error(
|
||
"No feasible configs — every parallelism choice exceeds "
|
||
"per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or "
|
||
"reducing `s_kv`."
|
||
)
|
||
return
|
||
|
||
# ── Pareto scatter: latency vs PEs, coloured by efficiency ─
|
||
import matplotlib.pyplot as _plt
|
||
_fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5))
|
||
|
||
_feas = [s for s in _res.all_scores
|
||
if s.fits_memory and s.placement_valid]
|
||
_axs[0].scatter(
|
||
[s.pes_used for s in _feas],
|
||
[s.latency_ms for s in _feas],
|
||
c="lightgrey", s=8, alpha=0.4, label="feasible",
|
||
)
|
||
_pareto_sorted = sorted(_res.pareto_scores,
|
||
key=lambda s: s.pes_used)
|
||
_sc = _axs[0].scatter(
|
||
[s.pes_used for s in _pareto_sorted],
|
||
[s.latency_ms for s in _pareto_sorted],
|
||
c=[s.efficiency_score for s in _pareto_sorted],
|
||
cmap="viridis", s=60, edgecolors="black", linewidths=0.8,
|
||
label="Pareto",
|
||
)
|
||
_axs[0].plot(
|
||
[s.pes_used for s in _pareto_sorted],
|
||
[s.latency_ms for s in _pareto_sorted],
|
||
"k--", linewidth=0.8, alpha=0.5,
|
||
)
|
||
_axs[0].set_xlabel("PEs used"); _axs[0].set_ylabel("Latency (ms)")
|
||
_axs[0].set_xscale("log", base=2)
|
||
_axs[0].set_yscale("log")
|
||
_axs[0].set_title("Pareto: latency vs PEs")
|
||
_axs[0].grid(alpha=0.3)
|
||
_axs[0].legend(loc="upper right", fontsize=8)
|
||
_fig.colorbar(_sc, ax=_axs[0], label="efficiency")
|
||
|
||
_pareto_by_lat = sorted(_res.pareto_scores,
|
||
key=lambda s: s.total_latency_ns)
|
||
_axs[1].scatter(
|
||
[s.hbm_utilization * 100 for s in _pareto_by_lat],
|
||
[s.latency_ms for s in _pareto_by_lat],
|
||
c=[s.pes_used for s in _pareto_by_lat],
|
||
cmap="plasma", s=60, edgecolors="black", linewidths=0.8,
|
||
)
|
||
_axs[1].set_xlabel("HBM utilization (%)")
|
||
_axs[1].set_ylabel("Latency (ms)")
|
||
_axs[1].set_yscale("log")
|
||
_axs[1].set_title("HBM usage vs latency (Pareto)")
|
||
_axs[1].grid(alpha=0.3)
|
||
_sc2 = _axs[1].collections[0]
|
||
_fig.colorbar(_sc2, ax=_axs[1], label="PEs")
|
||
|
||
_fig.tight_layout()
|
||
st.pyplot(_fig, width='stretch')
|
||
_plt.close(_fig)
|
||
|
||
# ── Pareto table ────────────────────────────────────────────
|
||
st.markdown("**Pareto configurations** (sorted by latency)")
|
||
_rows = []
|
||
for i, s in enumerate(_pareto_by_lat):
|
||
_rows.append({
|
||
"#": i,
|
||
"CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp,
|
||
"kv": s.kv_shard_mode,
|
||
"ffn": s.ffn_shard_scope,
|
||
"tp_place": s.tp_placement,
|
||
"cp_place": s.cp_placement,
|
||
"cp_ring": s.cp_ring_variant,
|
||
"lat (ms)": round(s.latency_ms, 3),
|
||
"eff": round(s.efficiency_score, 4),
|
||
"PEs": s.pes_used,
|
||
"SIPs": s.sips_used,
|
||
"HBM %": round(s.hbm_utilization * 100, 1),
|
||
})
|
||
_df = pd.DataFrame(_rows)
|
||
st.dataframe(_df, width='stretch', hide_index=True)
|
||
|
||
# ── Top-N by memory (smallest → largest) ─────────────────────
|
||
st.markdown(
|
||
"**Top 10 configurations by memory footprint** "
|
||
"(smallest HBM% first). Useful when memory pressure is the "
|
||
"real binding constraint of your deployment."
|
||
)
|
||
_pareto_by_mem = sorted(
|
||
_res.pareto_scores,
|
||
key=lambda s: (s.hbm_utilization, s.pes_used, s.total_latency_ns),
|
||
)[:10]
|
||
_mem_rows = []
|
||
for i, s in enumerate(_pareto_by_mem):
|
||
_mem_rows.append({
|
||
"#": i,
|
||
"HBM %": round(s.hbm_utilization * 100, 1),
|
||
"weights (GB)": round(s.weights_gb, 2),
|
||
"KV (GB)": round(s.kv_gb, 2),
|
||
"PEs": s.pes_used,
|
||
"SIPs": s.sips_used,
|
||
"lat (ms)": round(s.latency_ms, 3),
|
||
"CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp,
|
||
"kv": s.kv_shard_mode,
|
||
"ffn": s.ffn_shard_scope,
|
||
"tp_place": s.tp_placement,
|
||
"cp_place": s.cp_placement,
|
||
})
|
||
_df_mem = pd.DataFrame(_mem_rows)
|
||
st.dataframe(_df_mem, width='stretch', hide_index=True)
|
||
|
||
# ── Parallelism sensitivity subplots ──────────────────────
|
||
st.markdown(
|
||
"**Parallelism sensitivity** — for the *baseline* config picked "
|
||
"below, vary each knob one at a time (holding others fixed). "
|
||
"Solid line = fits memory; red X = infeasible."
|
||
)
|
||
_sens_c1, _sens_c2 = st.columns([1, 4])
|
||
with _sens_c1:
|
||
_sens_row = st.number_input(
|
||
"Baseline row #",
|
||
min_value=0, max_value=len(_pareto_by_lat) - 1,
|
||
value=0, step=1, key=f"_auto_sens_row{_suffix}",
|
||
help="Which Pareto row is the sensitivity computed around?",
|
||
)
|
||
_baseline_for_sens = _pareto_by_lat[int(_sens_row)]
|
||
_sens_rows = compute_parallelism_sensitivity(
|
||
_baseline_for_sens, model, machine, s_kv=s_kv, mode=mode,
|
||
include_attention=include_attention, include_ffn=include_ffn,
|
||
)
|
||
|
||
_fig_sens, _axes_sens = _plt.subplots(1, 5, figsize=(14, 3.2))
|
||
for _ax, _row in zip(_axes_sens, _sens_rows):
|
||
_fit_vals, _fit_lats = [], []
|
||
_bad_vals = []
|
||
for _v, _lat_ns, _fit in zip(_row.values, _row.latencies_ns,
|
||
_row.fits_flags):
|
||
if _lat_ns != _lat_ns:
|
||
continue
|
||
if _fit:
|
||
_fit_vals.append(_v)
|
||
_fit_lats.append(_lat_ns / 1e6)
|
||
else:
|
||
_bad_vals.append((_v, _lat_ns / 1e6))
|
||
if _fit_vals:
|
||
_ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50",
|
||
linewidth=1.4, markersize=4)
|
||
if _bad_vals:
|
||
_ax.scatter(
|
||
[v for v, _ in _bad_vals],
|
||
[l for _, l in _bad_vals],
|
||
marker="x", color="#c0392b", s=40, alpha=0.6,
|
||
label="no fit",
|
||
)
|
||
_ax.axvline(_row.baseline_value, color="#3498db",
|
||
linestyle=":", linewidth=1.0)
|
||
_ax.set_xscale("log")
|
||
_ax.set_yscale("log")
|
||
_ax.set_title(
|
||
f"{_row.knob.upper()} (baseline={_row.baseline_value})",
|
||
fontsize=10,
|
||
)
|
||
_ax.set_xlabel("value")
|
||
_ax.set_ylabel("lat (ms)")
|
||
_ax.grid(alpha=0.3)
|
||
if _bad_vals:
|
||
_ax.legend(loc="best", fontsize=7)
|
||
_fig_sens.tight_layout()
|
||
st.pyplot(_fig_sens, width='stretch')
|
||
_plt.close(_fig_sens)
|
||
|
||
# ── Load-into-sidebar ───────────────────────────────────────
|
||
st.markdown("**Load a Pareto config into the main sidebar sliders**")
|
||
_lc1, _lc2 = st.columns([1, 3])
|
||
with _lc1:
|
||
_pick = st.number_input(
|
||
"Row #", min_value=0, max_value=len(_pareto_by_lat) - 1,
|
||
value=0, step=1, key=f"_auto_pick_row{_suffix}",
|
||
)
|
||
with _lc2:
|
||
if st.button("Load into sidebar", type="secondary",
|
||
key=f"_auto_load_btn{_suffix}"):
|
||
_s = _pareto_by_lat[int(_pick)]
|
||
st.session_state["cp"] = _s.cp
|
||
st.session_state["tp"] = _s.tp
|
||
st.session_state["pp"] = _s.pp
|
||
st.session_state["dp"] = _s.dp
|
||
st.session_state["tp_placement"] = _s.tp_placement
|
||
st.session_state["cp_placement"] = _s.cp_placement
|
||
st.session_state["cp_ring_variant"] = _s.cp_ring_variant
|
||
st.session_state["kv_mode"] = _s.kv_shard_mode
|
||
_tp, _cp, _dp = _s.tp, _s.cp, _s.dp
|
||
_ffn_label_map = {
|
||
"TP": f"TP only (div={max(1, _tp)})",
|
||
"TP+CP": f"TP*CP (div={max(1, _tp * _cp)})",
|
||
"TP+CP+DP": f"TP*CP*DP (div={max(1, _tp * _cp * _dp)})",
|
||
}
|
||
st.session_state["ffn_scope_label"] = \
|
||
_ffn_label_map[_s.ffn_shard_scope]
|
||
st.success(
|
||
f"Loaded row {_pick}: CP={_s.cp} TP={_s.tp} PP={_s.pp} "
|
||
f"DP={_s.dp} → latency {_s.latency_ms:.2f} ms, "
|
||
f"{_s.pes_used} PEs. Flip to another tab to see the "
|
||
f"full breakdown."
|
||
)
|
||
st.rerun()
|
||
|
||
|
||
with tab_auto:
|
||
_render_auto_explore_tab()
|
||
|
||
|
||
# ── TAB (Auto Hardware) — one tab, two sweep buttons ─────────────
|
||
def _render_auto_hardware_tab():
|
||
"""Render the Auto Hardware tab body.
|
||
|
||
Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
|
||
result. Whichever button was last clicked drives the display; both
|
||
caches persist for fast toggling."""
|
||
from tests.analytical_visualization.auto_hardware import (
|
||
_HW_KNOB_DEFAULTS, joint_explore,
|
||
)
|
||
|
||
st.markdown(
|
||
"Sweep hardware knobs (PE HBM, HBM BW, TFLOPs, PE↔PE / D2D / C2C "
|
||
"interconnect BW) **together** with parallelism. For each hardware "
|
||
"candidate, we pick the latency-minimum parallelism that fits, then "
|
||
"Pareto-rank the joint (latency ↓, hardware cost ↓) space. Also "
|
||
"computes per-knob sensitivity: which HW investment gives the "
|
||
"biggest speedup?"
|
||
)
|
||
st.caption(
|
||
"**Attention only**: FFN / MoE stages are excluded from the summed "
|
||
"latency (both in the joint search and sensitivity). \n"
|
||
"**Attn + FFN/MoE**: full per-token cost."
|
||
)
|
||
|
||
hx_l, hx_m = st.columns([1, 1])
|
||
with hx_l:
|
||
st.markdown(
|
||
f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}**"
|
||
)
|
||
with hx_m:
|
||
_depth = st.radio(
|
||
"Sweep depth",
|
||
options=["two_stage", "balanced", "coarse"],
|
||
index=1, key="_hw_depth", horizontal=True,
|
||
help=("two_stage: 1 HW candidate (defaults) + autosuggest "
|
||
"parallelism, ~1s. Great for sensitivity alone.\n"
|
||
"balanced (default): 64 HW × ~2k parallelism, ~10-20s.\n"
|
||
"coarse: 729 HW × ~2k parallelism, ~2-5 min."),
|
||
)
|
||
|
||
_bh1, _bh2, _bh3, _bh4 = st.columns([1, 1, 1, 1])
|
||
with _bh1:
|
||
_run_hw_attn = st.button("Run joint sweep — Attention",
|
||
type="primary", width='stretch',
|
||
key="_hw_run_attn")
|
||
with _bh2:
|
||
_run_hw_ffn = st.button("Run joint sweep — FFN/MoE",
|
||
type="primary", width='stretch',
|
||
key="_hw_run_ffn")
|
||
with _bh3:
|
||
_run_hw_full = st.button("Run joint sweep — Attn + FFN/MoE",
|
||
type="primary", width='stretch',
|
||
key="_hw_run_full")
|
||
|
||
_scope_meta_hw = {
|
||
"attn": (True, False, "Attention only"),
|
||
"ffn": (False, True, "FFN / MoE only"),
|
||
"full": (True, True, "Attn + FFN/MoE"),
|
||
}
|
||
_button_click_hw = {
|
||
"attn": _run_hw_attn, "ffn": _run_hw_ffn, "full": _run_hw_full,
|
||
}
|
||
|
||
for _scope, _clicked in _button_click_hw.items():
|
||
if not _clicked:
|
||
continue
|
||
_ia, _ff, _lbl = _scope_meta_hw[_scope]
|
||
with st.spinner(f"Sweeping HW × parallelism ({_depth}, {_lbl})..."):
|
||
_r = joint_explore(model, s_kv=s_kv, mode=mode, depth=_depth,
|
||
include_attention=_ia, include_ffn=_ff,
|
||
b=b_batch)
|
||
st.session_state[f"_hw_result_{_scope}"] = _r
|
||
st.session_state[f"_hw_ctx_{_scope}"] = (
|
||
model.name, s_kv, mode, _depth,
|
||
)
|
||
st.session_state["_hw_active"] = _scope
|
||
|
||
_active = st.session_state.get("_hw_active")
|
||
if _active is None:
|
||
st.info(
|
||
"Click one of the run buttons to explore hardware × parallelism "
|
||
"co-design for this model + workload. All button results are "
|
||
"cached so you can flip between scopes without re-running."
|
||
)
|
||
return
|
||
|
||
include_attention, include_ffn, _label = _scope_meta_hw[_active]
|
||
_suffix = f"_{_active}"
|
||
_hw_result_key = f"_hw_result{_suffix}"
|
||
_hw_ctx_key = f"_hw_ctx{_suffix}"
|
||
_hw_res = st.session_state.get(_hw_result_key)
|
||
if _hw_res is None:
|
||
st.info(f"Click **Run joint sweep — {_label.split(' ')[0]}** to compute this view.")
|
||
return
|
||
|
||
st.markdown(f"### Currently viewing: **{_label}**")
|
||
|
||
_cached_ctx = st.session_state.get(_hw_ctx_key)
|
||
_cur_ctx = (model.name, s_kv, mode, _depth)
|
||
if _cached_ctx != _cur_ctx:
|
||
st.warning(
|
||
f"Cached {_label} result is from a different model/workload/depth. "
|
||
f"Click **Run joint sweep — {_label.split(' ')[0]}** to refresh."
|
||
)
|
||
|
||
_hm1, _hm2, _hm3, _hm4 = st.columns(4)
|
||
_hm1.metric("HW candidates", _hw_res.total_hw)
|
||
_hm2.metric("Feasible joint", _hw_res.total_joint)
|
||
_hm3.metric("Pareto configs", len(_hw_res.pareto_scores))
|
||
if _hw_res.pareto_scores:
|
||
_best = min(_hw_res.pareto_scores,
|
||
key=lambda s: (s.total_latency_ns,
|
||
s.parallelism.pes_used,
|
||
s.parallelism.hbm_utilization))
|
||
_hm4.metric("Best latency", f"{_best.latency_ms:.2f} ms")
|
||
|
||
if not _hw_res.pareto_scores:
|
||
st.error(
|
||
"No feasible joint configurations — every HW+parallelism "
|
||
"combo exceeds per-PE HBM at this workload."
|
||
)
|
||
return
|
||
if True:
|
||
# ── Panel 1: latency vs cost Pareto ────────────────────────
|
||
import matplotlib.pyplot as _plt
|
||
_fig1, _ax1 = _plt.subplots(figsize=(8, 4))
|
||
_feas = _hw_res.all_scores
|
||
_ax1.scatter(
|
||
[s.cost_score for s in _feas],
|
||
[s.latency_ms for s in _feas],
|
||
c="lightgrey", s=15, alpha=0.5, label="feasible",
|
||
)
|
||
_pareto_by_cost = sorted(_hw_res.pareto_scores,
|
||
key=lambda s: s.cost_score)
|
||
_sc = _ax1.scatter(
|
||
[s.cost_score for s in _pareto_by_cost],
|
||
[s.latency_ms for s in _pareto_by_cost],
|
||
c=[s.parallelism.pes_used for s in _pareto_by_cost],
|
||
cmap="viridis", s=90, edgecolors="black", linewidths=1.0,
|
||
label="Pareto",
|
||
)
|
||
_ax1.plot(
|
||
[s.cost_score for s in _pareto_by_cost],
|
||
[s.latency_ms for s in _pareto_by_cost],
|
||
"k--", linewidth=1.0, alpha=0.6,
|
||
)
|
||
_ax1.set_xlabel("Hardware cost proxy (unitless — 6.0 = all defaults)")
|
||
_ax1.set_ylabel("Latency (ms)")
|
||
_ax1.set_yscale("log")
|
||
_ax1.set_title("Optimal HW+parallelism: latency vs cost")
|
||
_ax1.grid(alpha=0.3)
|
||
_ax1.legend(loc="upper right", fontsize=9)
|
||
_fig1.colorbar(_sc, ax=_ax1, label="PEs")
|
||
_fig1.tight_layout()
|
||
st.pyplot(_fig1, width='stretch')
|
||
_plt.close(_fig1)
|
||
|
||
# ── Panel 2: per-knob sensitivity ──────────────────────────
|
||
st.markdown(
|
||
"**Per-knob sensitivity** (doubling each HW knob from the "
|
||
"latency-optimal baseline). Higher bar = more valuable "
|
||
"investment for this workload."
|
||
)
|
||
_fig2, _ax2 = _plt.subplots(figsize=(8, 3.5))
|
||
_knobs = [r.knob for r in _hw_res.sensitivity]
|
||
_speedups = [r.rel_speedup * 100 for r in _hw_res.sensitivity]
|
||
_colors = ["#2ecc71" if v > 0 else "#95a5a6" for v in _speedups]
|
||
_ax2.barh(_knobs, _speedups, color=_colors,
|
||
edgecolor="black", linewidth=0.5)
|
||
_ax2.invert_yaxis() # biggest speedup on top
|
||
_ax2.set_xlabel("Speedup when knob is doubled (%)")
|
||
_ax2.set_title("Sensitivity — where to invest next")
|
||
_ax2.grid(alpha=0.3, axis="x")
|
||
for i, r in enumerate(_hw_res.sensitivity):
|
||
_ax2.text(
|
||
r.rel_speedup * 100 + 0.5, i,
|
||
f"{r.baseline_value:g} → {r.doubled_value:g}",
|
||
va="center", fontsize=8,
|
||
)
|
||
_fig2.tight_layout()
|
||
st.pyplot(_fig2, width='stretch')
|
||
_plt.close(_fig2)
|
||
|
||
# ── Panel 3: Pareto table (HW spec + parallelism) ──────────
|
||
st.markdown("**Pareto joint configurations** (sorted by latency)")
|
||
_pareto_by_lat = sorted(_hw_res.pareto_scores,
|
||
key=lambda s: s.total_latency_ns)
|
||
_rows = []
|
||
for i, s in enumerate(_pareto_by_lat):
|
||
p = s.parallelism
|
||
h = s.hardware
|
||
_rows.append({
|
||
"#": i,
|
||
"lat (ms)": round(s.latency_ms, 3),
|
||
"cost": round(s.cost_score, 2),
|
||
"PEs": p.pes_used,
|
||
"CP": p.cp, "TP": p.tp, "PP": p.pp, "DP": p.dp,
|
||
"kv": p.kv_shard_mode,
|
||
"PE HBM": f"{h.pe_hbm_gb:g} GB",
|
||
"HBM BW": f"{h.bw_hbm_gbs:g} GB/s",
|
||
"TFLOPs": f"{h.peak_tflops_f16:g}",
|
||
"PE↔PE": f"{h.bw_intra_gbs:g}",
|
||
"D2D": f"{h.bw_inter_gbs:g}",
|
||
"C2C": f"{h.bw_intersip_gbs:g}",
|
||
})
|
||
_df_hw = pd.DataFrame(_rows)
|
||
st.dataframe(_df_hw, width='stretch', hide_index=True)
|
||
|
||
# ── Load hardware into sidebar ─────────────────────────────
|
||
st.markdown(
|
||
"**Load a Pareto HW config into the main sidebar** "
|
||
"(overwrites the Hardware→Per-PE and Hardware→Interconnect "
|
||
"sliders; also loads the paired parallelism)."
|
||
)
|
||
_lch1, _lch2 = st.columns([1, 3])
|
||
with _lch1:
|
||
_pick_hw = st.number_input(
|
||
"Row #", min_value=0,
|
||
max_value=len(_pareto_by_lat) - 1,
|
||
value=0, step=1, key=f"_hw_pick_row{_suffix}",
|
||
)
|
||
with _lch2:
|
||
if st.button("Load into sidebar", type="secondary",
|
||
key=f"_hw_load_btn{_suffix}"):
|
||
_s = _pareto_by_lat[int(_pick_hw)]
|
||
_h = _s.hardware
|
||
_p = _s.parallelism
|
||
# HW knobs (only load if the value is in the selectbox
|
||
# options; otherwise leave as-is).
|
||
_apply_if_in = {
|
||
"pe_hbm": ([3.0, 6.0, 12.0, 24.0, 48.0, 96.0],
|
||
_h.pe_hbm_gb),
|
||
"bw_hbm": ([128.0, 256.0, 512.0, 1024.0, 2048.0],
|
||
_h.bw_hbm_gbs),
|
||
"tflops": ([2.0, 4.0, 8.0, 16.0, 32.0, 64.0],
|
||
_h.peak_tflops_f16),
|
||
"bw_intra": ([128.0, 256.0, 512.0, 1024.0, 2048.0,
|
||
4096.0], _h.bw_intra_gbs),
|
||
"bw_inter": ([32.0, 64.0, 128.0, 256.0, 512.0,
|
||
900.0], _h.bw_inter_gbs),
|
||
"bw_intersip": ([12.5, 25.0, 50.0, 100.0, 200.0,
|
||
400.0], _h.bw_intersip_gbs),
|
||
}
|
||
for _key, (_opts, _val) in _apply_if_in.items():
|
||
if _val in _opts:
|
||
st.session_state[_key] = _val
|
||
# Parallelism knobs
|
||
st.session_state["cp"] = _p.cp
|
||
st.session_state["tp"] = _p.tp
|
||
st.session_state["pp"] = _p.pp
|
||
st.session_state["dp"] = _p.dp
|
||
st.session_state["tp_placement"] = _p.tp_placement
|
||
st.session_state["cp_placement"] = _p.cp_placement
|
||
st.session_state["cp_ring_variant"] = _p.cp_ring_variant
|
||
st.session_state["kv_mode"] = _p.kv_shard_mode
|
||
# Ffn_scope_label reconstruction (same as auto_explore tab)
|
||
_tp2, _cp2, _dp2 = _p.tp, _p.cp, _p.dp
|
||
_ffn_label_map = {
|
||
"TP": f"TP only (div={max(1, _tp2)})",
|
||
"TP+CP": f"TP*CP (div={max(1, _tp2 * _cp2)})",
|
||
"TP+CP+DP": f"TP*CP*DP (div={max(1, _tp2 * _cp2 * _dp2)})",
|
||
}
|
||
st.session_state["ffn_scope_label"] = \
|
||
_ffn_label_map[_p.ffn_shard_scope]
|
||
st.success(
|
||
f"Loaded row {_pick_hw}: latency {_s.latency_ms:.2f} "
|
||
f"ms, cost {_s.cost_score:.2f}, "
|
||
f"{_p.pes_used} PEs. Flip to another tab to see the "
|
||
f"full breakdown."
|
||
)
|
||
st.rerun()
|
||
|
||
|
||
with tab_hw:
|
||
_render_auto_hardware_tab()
|
||
|
||
|
||
# ── TAB 7: Chip roofline & B* ────────────────────────────────────
|
||
with tab_roofline:
|
||
st.subheader("Chip roofline & critical batch size (B*)")
|
||
st.caption(
|
||
"Back-of-envelope answer to 'is my chip well-matched to this "
|
||
"model?' — see how big the batch must be before decode becomes "
|
||
"compute-bound, and how big the context can grow before the KV "
|
||
"bandwidth wall kills your utilization. All numbers per PE, "
|
||
"peak roofline (no compute-util factor), no comm."
|
||
)
|
||
|
||
_n_active = total_active_params(model)
|
||
|
||
# ── Local knobs: S_kv, B, FLOPs, HBM BW ───────────────────────
|
||
# Override sidebar for exploring the plots without changing the
|
||
# rest of the app's config. FLOPs/BW knobs drive ALL roofline
|
||
# math and every plot/formula on this tab (memory-budget stacks
|
||
# excluded — those are HBM-capacity, not BW).
|
||
st.markdown("**Explore: local overrides for the plots below**")
|
||
_kn_a, _kn_b, _kn_c, _kn_d = st.columns(4)
|
||
with _kn_a:
|
||
_skv_min = 128
|
||
_skv_max = 1_000_000
|
||
_rf_skv = st.slider(
|
||
"S_kv (context length)",
|
||
min_value=_skv_min, max_value=_skv_max,
|
||
value=min(int(s_kv), _skv_max),
|
||
step=_skv_min,
|
||
key="_rf_skv_override",
|
||
help="Local to this tab. Drives every plot's S_kv.",
|
||
)
|
||
with _kn_b:
|
||
_b_min = 1
|
||
_b_max = 256
|
||
_rf_b = st.slider(
|
||
"B (batch size)",
|
||
min_value=_b_min, max_value=_b_max,
|
||
value=min(max(1, int(b_batch)), _b_max),
|
||
key="_rf_b_override",
|
||
help="Local to this tab. Marker on plots + PE-memory KV.",
|
||
)
|
||
with _kn_c:
|
||
_rf_flops_tflops = st.slider(
|
||
"FLOPs (TFLOPS)",
|
||
min_value=1.0, max_value=32.0,
|
||
value=float(_default_machine.peak_tflops_f16),
|
||
step=0.5,
|
||
key="_rf_flops_tflops",
|
||
help=("Chip peak TFLOPS for the AI-sensitivity plot. "
|
||
"Higher FLOPs → higher AI and higher B*."),
|
||
)
|
||
with _kn_d:
|
||
_rf_bw_gbs = st.slider(
|
||
"HBM BW (GB/s)",
|
||
min_value=128.0, max_value=1024.0,
|
||
value=float(_default_machine.bw_hbm_gbs),
|
||
step=16.0,
|
||
key="_rf_bw_gbs",
|
||
help=("Chip HBM bandwidth for the AI-sensitivity plot. "
|
||
"Higher BW → lower AI (more bytes to feed per FLOP)."),
|
||
)
|
||
# Build the effective chip from the knobs; every roofline number
|
||
# on this tab is derived from it.
|
||
_scaled_machine = dataclasses.replace(
|
||
_default_machine,
|
||
peak_tflops_f16=_rf_flops_tflops,
|
||
bw_hbm_gbs=_rf_bw_gbs,
|
||
)
|
||
_ai = arithmetic_intensity(_scaled_machine)
|
||
_b_star = critical_batch(_scaled_machine, model)
|
||
_l_star = balance_context(_scaled_machine, model)
|
||
_flops_mult_now = _rf_flops_tflops / _default_machine.peak_tflops_f16
|
||
_bw_mult_now = _rf_bw_gbs / _default_machine.bw_hbm_gbs
|
||
_ai_base = arithmetic_intensity(_default_machine)
|
||
_b_star_base = critical_batch(_default_machine, model)
|
||
st.caption(
|
||
f"**S_kv = {_rf_skv:,}**, **B = {_rf_b}**, "
|
||
f"**FLOPs = {_rf_flops_tflops:.1f} TFLOPS** "
|
||
f"(× {_flops_mult_now:.2f} of base), "
|
||
f"**BW = {_rf_bw_gbs:.0f} GB/s** "
|
||
f"(× {_bw_mult_now:.2f} of base). "
|
||
f"Effective **AI = {_ai:.1f}** (base {_ai_base:.1f}), "
|
||
f"**B\\* = {_b_star:.0f}** (base {_b_star_base:.0f}), "
|
||
f"**L\\* = {_l_star:,.0f}** tokens → S_kv/L\\* = "
|
||
f"**{_rf_skv/_l_star:.2f}**."
|
||
)
|
||
|
||
# ── Headline plots: step latency and cost per token (=÷B) ─────
|
||
_b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256]
|
||
_step_pts = step_latency_curve(_scaled_machine, model, _b_range,
|
||
s_kv=_rf_skv)
|
||
_tok_pts = per_token_latency_curve(_scaled_machine, model, _b_range,
|
||
s_kv=_rf_skv)
|
||
_xs_head = [p.batch for p in _step_pts]
|
||
|
||
_hcol1, _hcol2 = st.columns(2)
|
||
|
||
# Plot A — Latency per decode step (undivided by B)
|
||
with _hcol1:
|
||
st.markdown("**Latency per decode step** (one forward pass)")
|
||
_figA, _axA = plt.subplots(figsize=(6, 4.2))
|
||
_axA.plot(_xs_head, [p.weight_s * 1e3 for p in _step_pts],
|
||
"^-", color="#ffbe0b",
|
||
label="Weight fetch (flat in B)")
|
||
_axA.plot(_xs_head, [p.compute_s * 1e3 for p in _step_pts],
|
||
"o-", color="#3a86ff",
|
||
label="Compute (linear ↑)")
|
||
_axA.plot(_xs_head, [p.kv_s * 1e3 for p in _step_pts],
|
||
"s-", color="#d90429",
|
||
label="KV fetch (linear ↑)")
|
||
_axA.plot(_xs_head, [p.total_s * 1e3 for p in _step_pts],
|
||
"-", color="#212529", linewidth=2.5, label="Total")
|
||
_axA.axvline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6,
|
||
label=f"B = {_rf_b}")
|
||
_axA.set_xscale("log", base=2)
|
||
_axA.set_yscale("log")
|
||
_axA.set_xlabel("Batch size B")
|
||
_axA.set_ylabel("Step time (ms)")
|
||
_axA.set_title("Step latency = weight + compute·B + KV·B")
|
||
_axA.grid(True, which="both", alpha=0.3)
|
||
_axA.legend(fontsize=8, loc="upper left")
|
||
plt.tight_layout()
|
||
st.pyplot(_figA, width='stretch')
|
||
plt.close(_figA)
|
||
|
||
# Plot B — Cost per token = latency ÷ B
|
||
with _hcol2:
|
||
st.markdown("**Cost per token** (= step latency ÷ B)")
|
||
_figB, _axB = plt.subplots(figsize=(6, 4.2))
|
||
_axB.plot(_xs_head, [p.weight_s * 1e3 for p in _tok_pts],
|
||
"^-", color="#ffbe0b",
|
||
label="Weight fetch (shrinks 1/B)")
|
||
_axB.plot(_xs_head, [p.compute_s * 1e3 for p in _tok_pts],
|
||
"o--", color="#3a86ff",
|
||
label="Compute (flat)")
|
||
_axB.plot(_xs_head, [p.kv_s * 1e3 for p in _tok_pts],
|
||
"s--", color="#d90429",
|
||
label="KV fetch (flat)")
|
||
_axB.plot(_xs_head, [p.total_s * 1e3 for p in _tok_pts],
|
||
"-", color="#212529", linewidth=2.5, label="Total")
|
||
_axB.axvline(_b_star, linestyle=":", color="#2e7d32",
|
||
label=f"B* = {_b_star:.0f}")
|
||
_axB.axvline(_rf_b, linestyle="-", color="#7b1fa2", alpha=0.6,
|
||
label=f"B = {_rf_b}")
|
||
_axB.set_xscale("log", base=2)
|
||
_axB.set_yscale("log")
|
||
_axB.set_xlabel("Batch size B")
|
||
_axB.set_ylabel("Per-token time (ms)")
|
||
_axB.set_title("Cost/token = weight/B + compute + KV")
|
||
_axB.grid(True, which="both", alpha=0.3)
|
||
_axB.legend(fontsize=8, loc="upper right")
|
||
plt.tight_layout()
|
||
st.pyplot(_figB, width='stretch')
|
||
plt.close(_figB)
|
||
|
||
st.caption(
|
||
"**Left (step latency)**: how long one forward pass takes. "
|
||
"Grows with B because compute and KV per-sequence grow linearly; "
|
||
"weight fetch is loaded once per step so it's flat. This is the "
|
||
"**SLO / TTFT view** — bigger B → longer step.\n\n"
|
||
"**Right (cost per token)**: step latency divided by B tokens "
|
||
"produced. Weight fetch amortizes (shrinks 1/B); compute and KV "
|
||
"per-sequence stay flat per token. This is the **efficiency / "
|
||
"$-per-token view** — bigger B (up to ~2·B*) → cheaper per token.\n\n"
|
||
"Same underlying decomposition, two divisors — the classic "
|
||
"throughput ↔ latency tradeoff."
|
||
)
|
||
|
||
st.divider()
|
||
|
||
# ── PE memory budget (sharded, per-PE) ────────────────────────
|
||
st.markdown(
|
||
"**PE memory budget — how batch and context fill the HBM**"
|
||
)
|
||
_hbm_gb = _default_machine.pe_hbm_gb
|
||
_skv_sweep = [128, 512, 2048, 8192, 32768, 131072,
|
||
524288, 1_000_000]
|
||
_bud_skv = memory_budget_curve_vs_skv(cfg, _skv_sweep, batch=_rf_b)
|
||
_b_sweep = [1, 2, 4, 8, 16, 32, 64, 128, 256]
|
||
_bud_b = memory_budget_curve_vs_batch(cfg, _b_sweep, s_kv=_rf_skv)
|
||
|
||
_mem_L, _mem_R = st.columns(2)
|
||
|
||
with _mem_L:
|
||
st.markdown(f"vs S_kv (at B = {_rf_b}) — sharded CP={topo.cp}, "
|
||
f"TP={topo.tp}, PP={topo.pp}")
|
||
_figM1, _axM1 = plt.subplots(figsize=(6, 4.2))
|
||
_xM = [p.axis_val for p in _bud_skv]
|
||
_w = [p.weights_gb for p in _bud_skv]
|
||
_k = [p.kv_gb for p in _bud_skv]
|
||
_tr = [p.transient_gb for p in _bud_skv]
|
||
_axM1.stackplot(_xM, _w, _k, _tr,
|
||
labels=["Weights", "KV cache", "Transient"],
|
||
colors=["#ffbe0b", "#d90429", "#3a86ff"],
|
||
alpha=0.8)
|
||
_axM1.axhline(_hbm_gb, color="#212529", linestyle="--",
|
||
linewidth=1.5, label=f"HBM budget = {_hbm_gb} GB")
|
||
_axM1.axvline(_rf_skv, color="#7b1fa2", linewidth=1.5,
|
||
alpha=0.6, label=f"S_kv = {_rf_skv:,}")
|
||
_axM1.set_xscale("log")
|
||
_axM1.set_xlabel("S_kv (tokens)")
|
||
_axM1.set_ylabel("Per-PE memory (GB)")
|
||
_axM1.set_title("Weight is fixed; KV fills the budget as S_kv grows")
|
||
_axM1.grid(True, which="both", alpha=0.3)
|
||
_axM1.legend(fontsize=8, loc="upper left")
|
||
plt.tight_layout()
|
||
st.pyplot(_figM1, width='stretch')
|
||
plt.close(_figM1)
|
||
|
||
with _mem_R:
|
||
st.markdown(f"vs B (at S_kv = {_rf_skv:,}) — same sharding")
|
||
_figM2, _axM2 = plt.subplots(figsize=(6, 4.2))
|
||
_xM2 = [p.axis_val for p in _bud_b]
|
||
_w2 = [p.weights_gb for p in _bud_b]
|
||
_k2 = [p.kv_gb for p in _bud_b]
|
||
_tr2 = [p.transient_gb for p in _bud_b]
|
||
_axM2.stackplot(_xM2, _w2, _k2, _tr2,
|
||
labels=["Weights", "KV cache", "Transient"],
|
||
colors=["#ffbe0b", "#d90429", "#3a86ff"],
|
||
alpha=0.8)
|
||
_axM2.axhline(_hbm_gb, color="#212529", linestyle="--",
|
||
linewidth=1.5, label=f"HBM budget = {_hbm_gb} GB")
|
||
_axM2.axvline(_rf_b, color="#7b1fa2", linewidth=1.5,
|
||
alpha=0.6, label=f"B = {_rf_b}")
|
||
_axM2.set_xscale("log", base=2)
|
||
_axM2.set_xlabel("Batch size B")
|
||
_axM2.set_ylabel("Per-PE memory (GB)")
|
||
_axM2.set_title("Weight fixed; KV grows linearly with B")
|
||
_axM2.grid(True, which="both", alpha=0.3)
|
||
_axM2.legend(fontsize=8, loc="upper left")
|
||
plt.tight_layout()
|
||
st.pyplot(_figM2, width='stretch')
|
||
plt.close(_figM2)
|
||
|
||
_bud_now = memory_budget_curve_vs_skv(cfg, [_rf_skv], batch=_rf_b)[0]
|
||
st.caption(
|
||
f"At **S_kv = {_rf_skv:,}, B = {_rf_b}**: weights = "
|
||
f"**{_bud_now.weights_gb:.2f} GB**, KV = **{_bud_now.kv_gb:.2f} GB**, "
|
||
f"transient = **{_bud_now.transient_gb:.2f} GB**, "
|
||
f"used = **{_bud_now.used_gb:.2f} GB / {_hbm_gb:.1f} GB** "
|
||
f"({'**OVER BUDGET**' if _bud_now.over_budget else f'free = {_bud_now.free_gb:.2f} GB'}). "
|
||
"Weights are fixed by sharding (CP·TP·PP·EP splits them across "
|
||
"PEs); KV per PE = B × S_kv × kv_bpt_per_PE."
|
||
)
|
||
|
||
st.divider()
|
||
|
||
# ── AI sensitivity to FLOPs / BW ──────────────────────────────
|
||
st.markdown("**AI sensitivity — how scaling chip FLOPs or BW moves B\\***")
|
||
_ai_mults = [0.25, 0.5, 1.0, 2.0, 4.0, 8.0]
|
||
_ai_flops_pts = ai_sensitivity_curve(_default_machine, model,
|
||
_ai_mults, axis="flops")
|
||
_ai_bw_pts = ai_sensitivity_curve(_default_machine, model,
|
||
_ai_mults, axis="bw")
|
||
_figAI, (_axAI1, _axAI2) = plt.subplots(1, 2, figsize=(12, 4.2))
|
||
|
||
_axAI1.plot(_ai_mults, [p.ai for p in _ai_flops_pts], "o-",
|
||
color="#3a86ff", linewidth=2, label="Scale FLOPs")
|
||
_axAI1.plot(_ai_mults, [p.ai for p in _ai_bw_pts], "s-",
|
||
color="#d90429", linewidth=2, label="Scale HBM BW")
|
||
_axAI1.axhline(_ai_base, color="#2e7d32", linestyle=":",
|
||
label=f"Base AI = {_ai_base:.1f}")
|
||
_axAI1.axvline(1.0, color="#888", linestyle=":", alpha=0.5)
|
||
_axAI1.set_xscale("log", base=2)
|
||
_axAI1.set_yscale("log")
|
||
_axAI1.set_xlabel("Multiplier on chip parameter")
|
||
_axAI1.set_ylabel("AI (FLOPs/byte)")
|
||
_axAI1.set_title("AI = C / W. FLOPs ↑ raises AI, BW ↑ lowers AI")
|
||
_axAI1.grid(True, which="both", alpha=0.3)
|
||
_axAI1.legend(fontsize=8, loc="best")
|
||
|
||
_axAI2.plot(_ai_mults, [p.b_star for p in _ai_flops_pts], "o-",
|
||
color="#3a86ff", linewidth=2, label="Scale FLOPs")
|
||
_axAI2.plot(_ai_mults, [p.b_star for p in _ai_bw_pts], "s-",
|
||
color="#d90429", linewidth=2, label="Scale HBM BW")
|
||
_axAI2.axhline(_b_star_base, color="#2e7d32", linestyle=":",
|
||
label=f"Base B* = {_b_star_base:.0f}")
|
||
_axAI2.axvline(_flops_mult_now, color="#3a86ff", linewidth=1.2,
|
||
alpha=0.5, label=f"FLOPs × {_flops_mult_now:.2f}")
|
||
_axAI2.axvline(_bw_mult_now, color="#d90429", linewidth=1.2,
|
||
alpha=0.5, label=f"BW × {_bw_mult_now:.2f}",
|
||
linestyle="--")
|
||
_axAI2.set_xscale("log", base=2)
|
||
_axAI2.set_yscale("log")
|
||
_axAI2.set_xlabel("Multiplier on chip parameter")
|
||
_axAI2.set_ylabel("B* (critical batch)")
|
||
_axAI2.set_title("B* = AI · b/2 — same shape as AI (for BF16)")
|
||
_axAI2.grid(True, which="both", alpha=0.3)
|
||
_axAI2.legend(fontsize=8, loc="best")
|
||
|
||
plt.tight_layout()
|
||
st.pyplot(_figAI, width='stretch')
|
||
plt.close(_figAI)
|
||
|
||
st.caption(
|
||
f"**Left**: AI vs chip-parameter multiplier. Scaling **FLOPs** "
|
||
f"raises AI linearly (more compute per byte). Scaling **HBM BW** "
|
||
f"lowers AI (more bytes to feed per FLOP). "
|
||
f"**Right**: B\\* moves in the same direction as AI. "
|
||
f"**Your knob position**: FLOPs × {_flops_mult_now:.2f}, "
|
||
f"BW × {_bw_mult_now:.2f} → **AI = {_ai:.2f}**, "
|
||
f"**B\\* = {_b_star:.0f}** (base AI {_ai_base:.1f}, "
|
||
f"B\\* {_b_star_base:.0f})."
|
||
)
|
||
|
||
st.divider()
|
||
|
||
_regime_now = bound_regime(_scaled_machine, model,
|
||
batch=max(1, _rf_b), s_kv=_rf_skv)
|
||
|
||
# ── KPI cards ─────────────────────────────────────────────────
|
||
_r1, _r2, _r3, _r4 = st.columns(4)
|
||
_r1.metric("AI (FLOPs/byte)", f"{_ai:.1f}",
|
||
help="C/W. Peak FLOPs per byte of HBM bandwidth. "
|
||
"Chip-only — no model-dependent factor.")
|
||
_r2.metric("B* (dense)", f"{_b_star:.0f}",
|
||
help="C*b/(2*W). Batch size where weight-fetch time "
|
||
"equals compute time. Below this you're memory-bound.")
|
||
_r3.metric("L* (tokens)", f"{_l_star:,.0f}",
|
||
help="Context length where KV-read time equals compute "
|
||
"time. Above this, no batch size gets you back to "
|
||
"compute-bound.")
|
||
_r4.metric(f"Regime @ (B={_rf_b}, S_kv={_rf_skv:,})",
|
||
_regime_now.replace("-bound", ""),
|
||
help="Which term dominates cost at the current (B, S_kv) "
|
||
"you've selected via the sliders above.")
|
||
|
||
# ── Recommended B and L (heuristics) ──────────────────────────
|
||
_rec_b = good_batch(_scaled_machine, model, sparsity=1.0)
|
||
_rec_l = good_context(_scaled_machine, model, s_kv=_rf_skv)
|
||
_g1, _g2, _g3 = st.columns(3)
|
||
_g1.metric("Good B (dense)", f"{_rec_b.effective:.0f}",
|
||
help=_rec_b.reason)
|
||
_g2.metric("Good L ceiling", f"{_rec_l.max_efficient:,.0f} tok",
|
||
help="Stay at or below L* to keep decode compute-bound.")
|
||
_g3.metric(f"Utilization @ S_kv={_rf_skv:,}",
|
||
f"{_rec_l.utilization_at*100:.1f}%",
|
||
help="Peak compute utilization at this context: "
|
||
"1 / (1 + S_kv/L*).")
|
||
|
||
# Optional MoE B*: only surface if the preset flags MoE.
|
||
_is_moe = "MoE" in (preset.family + preset.note)
|
||
if _is_moe:
|
||
st.info(
|
||
"This preset is dense-approximated as an MoE. Reiner Pope's "
|
||
"'300 × sparsity' rule: for real sparsity k = N_total/N_active, "
|
||
f"B*_moe ≈ {_b_star:.0f} × k. E.g. DeepSeek 32-of-256 experts "
|
||
f"(k=8) → B* ≈ {_b_star*8:.0f}."
|
||
)
|
||
|
||
st.caption(
|
||
f"**Full-model attention+FFN params:** {_n_active/1e9:.2f} B. "
|
||
f"**Effective chip (from knobs):** "
|
||
f"{_scaled_machine.peak_tflops_f16:.1f} TFLOPs BF16, "
|
||
f"{_scaled_machine.bw_hbm_gbs:.0f} GB/s HBM. "
|
||
f"Base chip from sidebar Hardware: "
|
||
f"{_default_machine.peak_tflops_f16:.0f} TFLOPs, "
|
||
f"{_default_machine.bw_hbm_gbs:.0f} GB/s."
|
||
)
|
||
|
||
st.divider()
|
||
|
||
# ── Regime formulas (short vs long context) ───────────────────
|
||
st.markdown("**Cost-term formulas by regime**")
|
||
_t_com_val = t_com(_scaled_machine, model) * 1e3 # ms
|
||
_t_mem_s_at_bstar = t_mem_short(_scaled_machine, model,
|
||
int(round(_b_star))) * 1e3
|
||
_t_mem_l_at_skv = t_mem_long(_scaled_machine, model, _rf_skv) * 1e3
|
||
_regime_rows = [
|
||
{"Term": "t_com (compute)",
|
||
"Short context (S_kv < L*)": "2 · N / C",
|
||
"Long context (S_kv > L*)": "2 · N / C (same)",
|
||
"Value now (ms)": f"{_t_com_val:.3f}"},
|
||
{"Term": "t_mem — weight fetch",
|
||
"Short context (S_kv < L*)": "N · b / (W · B)",
|
||
"Long context (S_kv > L*)": "N · b / (W · B) — becomes small",
|
||
"Value now (ms)":
|
||
f"{_t_mem_s_at_bstar:.3f} at B = B*"},
|
||
{"Term": "t_mem — KV fetch",
|
||
"Short context (S_kv < L*)": "S_kv · kv_bpt / W — small",
|
||
"Long context (S_kv > L*)": "S_kv · kv_bpt / W — dominates",
|
||
"Value now (ms)": f"{_t_mem_l_at_skv:.3f}"},
|
||
{"Term": "Bottleneck",
|
||
"Short context (S_kv < L*)":
|
||
"Weights (B < B*), else Compute (B ≥ B*)",
|
||
"Long context (S_kv > L*)":
|
||
"KV bandwidth (regardless of B)",
|
||
"Value now (ms)":
|
||
_regime_now.replace("-bound", "")},
|
||
]
|
||
st.dataframe(pd.DataFrame(_regime_rows), width='stretch',
|
||
hide_index=True)
|
||
st.caption(
|
||
"N = active params, b = bytes/elem, W = HBM BW, C = peak FLOPs, "
|
||
"kv_bpt = KV bytes per token = 2 · H_kv · d_head · b · layers. "
|
||
"'Value now' uses the S_kv slider above and B = B\\*."
|
||
)
|
||
|
||
st.divider()
|
||
|
||
# ── Recommended B and L (recap) ───────────────────────────────
|
||
st.markdown("**How to pick a good B and a good S_kv**")
|
||
st.markdown(
|
||
f"- **Good batch**: `B_target = 2 · B* = {_rec_b.effective:.0f}`. "
|
||
"Below B\\*: memory-bound (doubling B halves cost/token). "
|
||
"Above 3·B\\*: diminishing returns; latency keeps rising. "
|
||
"Also cap by your HBM budget: `B_max ≤ (HBM − weights_per_PE) / "
|
||
"(S_kv · kv_bpt_per_PE)` — check the Memory Breakdown tab.\n"
|
||
f"- **Good S_kv**: stay ≤ **L\\* = {_l_star:,.0f} tokens** to "
|
||
"keep decode compute-bound (peak utilization). Past L\\*, "
|
||
f"utilization = 1/(1 + S_kv/L\\*). At S_kv = {_rf_skv:,}, "
|
||
f"you're at ~**{_rec_l.utilization_at*100:.1f}%** utilization."
|
||
)
|
||
|
||
st.divider()
|
||
|
||
# ── Interpretation table ───────────────────────────────────────
|
||
st.markdown("**Formulas & interpretation**")
|
||
_C = _scaled_machine.peak_flops # FLOPs / s (from knobs)
|
||
_W = _scaled_machine.bw_hbm # bytes / s (from knobs)
|
||
_b_elem = model.bytes_per_elem # bytes / weight
|
||
_kv_bpt_full = 2 * model.h_kv * model.d_head * model.bytes_per_elem * model.layers
|
||
|
||
_knee_now = knee_batch(_scaled_machine, model, _rf_skv)
|
||
if _knee_now is None:
|
||
_sub_bknee = (
|
||
f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = "
|
||
f"{_b_star:.0f} / (1 − {_rf_skv/_l_star:.3f}) [≤ 0]"
|
||
)
|
||
_bknee_val = "no knee (S_kv ≥ L*)"
|
||
else:
|
||
_sub_bknee = (
|
||
f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = "
|
||
f"{_b_star:.0f} / {1 - _rf_skv/_l_star:.3f} = "
|
||
f"{_knee_now:.0f}"
|
||
)
|
||
_bknee_val = f"{_knee_now:.0f}"
|
||
|
||
_formula_rows = [
|
||
{"Symbol": "AI",
|
||
"Formula": "C / W",
|
||
"With numbers": f"{_C:.2e} / {_W:.2e}",
|
||
"Meaning": "FLOPs per byte of HBM BW. Chip-only.",
|
||
"Value": f"{_ai:.2f} FLOPs/byte"},
|
||
{"Symbol": "B*",
|
||
"Formula": "C · b / (2 · W)",
|
||
"With numbers": f"{_C:.2e} · {_b_elem} / (2 · {_W:.2e})",
|
||
"Meaning": ("Batch where weight-fetch = compute. Below: "
|
||
"memory-bound. Above: amortized."),
|
||
"Value": f"{_b_star:.0f}"},
|
||
{"Symbol": "B*_moe",
|
||
"Formula": "B* · (N_total / N_active)",
|
||
"With numbers": (f"{_b_star:.0f} · 8 (k=8 example)" if _is_moe
|
||
else "N/A (dense preset)"),
|
||
"Meaning": "MoE fetches all experts; computes on active only.",
|
||
"Value": (f"{_b_star * 8:.0f}" if _is_moe else "1× (dense)")},
|
||
{"Symbol": "L*",
|
||
"Formula": "2 · N_active · W / (C · kv_bpt)",
|
||
"With numbers": (f"2 · {_n_active:.2e} · {_W:.2e} / "
|
||
f"({_C:.2e} · {_kv_bpt_full:,})"),
|
||
"Meaning": ("Context where KV-read = compute. Above: "
|
||
"KV-bound regardless of B."),
|
||
"Value": f"{_l_star:,.0f} tokens"},
|
||
{"Symbol": "B_knee",
|
||
"Formula": "B* / (1 − S_kv/L*)",
|
||
"With numbers": _sub_bknee,
|
||
"Meaning": ("Batch where cost curve bends onto its floor. "
|
||
"Diverges at S_kv = L*."),
|
||
"Value": _bknee_val},
|
||
]
|
||
st.dataframe(pd.DataFrame(_formula_rows),
|
||
width='stretch', hide_index=True)
|
||
|
||
|
||
# ── TAB 8: Capacity planning ─────────────────────────────────────
|
||
with tab_planning:
|
||
st.subheader("Capacity planning & standards")
|
||
st.caption(
|
||
"How hyperscalers decide GPU count from a workload, what SLO "
|
||
"targets they aim for, and what deployment patterns they use. "
|
||
"All numbers below use the sidebar model + the Chip Roofline "
|
||
"tab's effective chip (from its FLOPs/BW knobs, or sidebar "
|
||
"defaults if you haven't opened that tab)."
|
||
)
|
||
|
||
# ── Section 1: three-axis sizing formula + calculator ─────────
|
||
st.markdown("### 1. GPU count = max(A, B, C) × N_replicas")
|
||
|
||
_axes_rows = [
|
||
{"Axis": "A. Capacity floor",
|
||
"Formula": "⌈ N·b / HBM_per_PE ⌉",
|
||
"Meaning": ("Min PEs to hold ONE replica's weights alone. "
|
||
"The bare floor — below this weights don't fit."),
|
||
"Grows with": "Model size (N)"},
|
||
{"Axis": "B. KV headroom",
|
||
"Formula": "⌈ (N·b + users_per_replica · S_kv · kv_bpt) / HBM_per_PE ⌉",
|
||
"Meaning": ("Min PEs to hold weights + all KV of the users "
|
||
"assigned to this replica."),
|
||
"Grows with": "Users × context"},
|
||
{"Axis": "C. Throughput SLO",
|
||
"Formula": ("N_replicas = ⌈ n_users / B_at_SLO ⌉, where "
|
||
"B_at_SLO = largest B s.t. step_latency ≤ TPOT SLO"),
|
||
"Meaning": ("Enough replicas so each carries ≤ B_at_SLO users "
|
||
"and meets per-token latency SLO."),
|
||
"Grows with": "Users, or tighter SLO"},
|
||
]
|
||
st.dataframe(pd.DataFrame(_axes_rows), width='stretch',
|
||
hide_index=True)
|
||
|
||
st.markdown("#### Live calculator")
|
||
|
||
_cp_a, _cp_b, _cp_c = st.columns(3)
|
||
with _cp_a:
|
||
_cp_users = st.number_input(
|
||
"Concurrent users", min_value=1, max_value=100_000,
|
||
value=100, step=1, key="_cp_n_users",
|
||
help="Target simultaneous active decode sequences.",
|
||
)
|
||
with _cp_b:
|
||
_cp_ctx = st.number_input(
|
||
"Avg context / user (tokens)", min_value=1,
|
||
max_value=2_000_000, value=int(s_kv), step=128,
|
||
key="_cp_avg_ctx",
|
||
help="Mean S_kv across the active users.",
|
||
)
|
||
with _cp_c:
|
||
_cp_slo_ms = st.number_input(
|
||
"TPOT SLO (ms/token)", min_value=1, max_value=1000,
|
||
value=30, step=1, key="_cp_tpot_slo_ms",
|
||
help=("Per-token latency target during decode. "
|
||
"Tighter = smaller B per replica, more replicas."),
|
||
)
|
||
|
||
# Use the base sidebar chip for the calculator so it's stable.
|
||
_cp_machine = _default_machine
|
||
_sizing = size_deployment(
|
||
_cp_machine, model,
|
||
n_users=int(_cp_users), avg_ctx=int(_cp_ctx),
|
||
tpot_slo_s=_cp_slo_ms / 1000.0,
|
||
)
|
||
|
||
_s1, _s2, _s3, _s4 = st.columns(4)
|
||
_s1.metric("A. Capacity",
|
||
f"{_sizing.pes_axis_a_capacity} PEs",
|
||
help="Min PEs to hold weights alone.")
|
||
_s2.metric("B. KV headroom",
|
||
f"{_sizing.pes_axis_b_kv} PEs",
|
||
help="Min PEs to hold weights + all replica KV.")
|
||
_s3.metric("C. Throughput",
|
||
(f"{_sizing.n_replicas} replicas"
|
||
if _sizing.b_at_slo > 0
|
||
else "SLO INFEASIBLE"),
|
||
help=(f"B_at_SLO = {_sizing.b_at_slo}; ceil("
|
||
f"{_cp_users}/{_sizing.users_per_replica}) "
|
||
f"= {_sizing.n_replicas} replicas."))
|
||
_s4.metric("Total GPUs",
|
||
f"{_sizing.total_pes}",
|
||
delta=f"binds: {_sizing.binding_axis}",
|
||
help="max(A, B) × N_replicas.")
|
||
|
||
if _sizing.b_at_slo == 0:
|
||
st.error(
|
||
f"SLO infeasible: even B=1 step latency exceeds "
|
||
f"{_cp_slo_ms} ms at S_kv = {_cp_ctx:,}. Loosen the SLO, "
|
||
f"shorten context, or use a faster chip."
|
||
)
|
||
else:
|
||
_binding_hint = {
|
||
"capacity": ("Weights dominate. Adding replicas or bigger-HBM "
|
||
"chips buys the most headroom."),
|
||
"kv": ("KV cache dominates. Shorter context / smaller B / "
|
||
"more aggressive CP sharding help most."),
|
||
"throughput": (f"Latency SLO forces more replicas — one "
|
||
f"replica serves ~{_sizing.b_at_slo} users. "
|
||
"Loosening TPOT or faster chip cuts count."),
|
||
}[_sizing.binding_axis]
|
||
st.caption(
|
||
f"**Binding axis: {_sizing.binding_axis}.** {_binding_hint} "
|
||
f"B_at_SLO = **{_sizing.b_at_slo}** / replica × "
|
||
f"**{_sizing.n_replicas}** replicas = "
|
||
f"**{_sizing.total_pes} GPUs**."
|
||
)
|
||
|
||
st.divider()
|
||
|
||
# ── Section 2: SLO targets by workload ─────────────────────────
|
||
st.markdown("### 2. SLO standards — TTFT & TPOT by workload")
|
||
st.caption(
|
||
"TTFT = Time To First Token (dominated by prefill of the prompt). "
|
||
"TPOT = Time Per Output Token (one decode step per token). "
|
||
"Real workload targets vary widely; the table is a rough guide."
|
||
)
|
||
_slo_rows = [
|
||
{"Use case": "Voice / realtime",
|
||
"TTFT target": "200–300 ms (whole pipeline)",
|
||
"TPOT / ITL target": "10–25 ms",
|
||
"What binds": "TPOT — smooth speech needs tight per-token cadence"},
|
||
{"Use case": "Code completion",
|
||
"TTFT target": "< 100 ms",
|
||
"TPOT / ITL target": "mostly n/a (short response)",
|
||
"What binds": "TTFT — user cursor is waiting"},
|
||
{"Use case": "Interactive chat",
|
||
"TTFT target": "300 ms – 1 s",
|
||
"TPOT / ITL target": "20–50 ms",
|
||
"What binds": "Balanced; both matter to perceived latency"},
|
||
{"Use case": "Agentic / multi-step",
|
||
"TTFT target": "as low as possible — compounds across steps",
|
||
"TPOT / ITL target": "maximize throughput",
|
||
"What binds": "Throughput — many sequential LLM calls"},
|
||
{"Use case": "Batch / offline",
|
||
"TTFT target": "seconds to minutes",
|
||
"TPOT / ITL target": "irrelevant",
|
||
"What binds": "Cost/token — pack B to the ceiling"},
|
||
]
|
||
st.dataframe(pd.DataFrame(_slo_rows), width='stretch',
|
||
hide_index=True)
|
||
|
||
st.divider()
|
||
|
||
# ── Section 3: Hybrid deployment layers (table form) ──────────
|
||
st.markdown("### 3. Same setup or different for short vs long context?")
|
||
st.caption(
|
||
"**Hybrid: one elastic pool + length-tier routing on top + "
|
||
"prefill/decode disaggregation at the frontier.**"
|
||
)
|
||
_layers_rows = [
|
||
{"Layer": "1. Elastic pool",
|
||
"Serves": "Mixed traffic (~90% of requests, chat/RAG/code)",
|
||
"Key mechanism": "PagedAttention (KV in fixed-size pages) + "
|
||
"continuous batching",
|
||
"Admission rule": "Scheduler checks remaining KV page pool per "
|
||
"iteration; admits if user's expected "
|
||
"max_tokens fits",
|
||
"Typical config": "Standard replica (e.g. TP=8, CP=1..4)"},
|
||
{"Layer": "2. Length-tier routing",
|
||
"Serves": "The long-context tail (128k–1M+)",
|
||
"Key mechanism": "API gateway inspects max context, dispatches "
|
||
"to a dedicated pool",
|
||
"Admission rule": "Standard tier for < 32-128k; long-context "
|
||
"tier otherwise",
|
||
"Typical config": ("Long-context replicas have more chips per "
|
||
"instance + heavier CP sharding, lower B "
|
||
"(e.g. TP=8, CP=32)")},
|
||
{"Layer": "3. Disaggregated prefill/decode",
|
||
"Serves": "Frontier deployments (DistServe, Splitwise pattern)",
|
||
"Key mechanism": "Prefill on FLOPs-heavy pool; decode on HBM-BW-"
|
||
"heavy pool; KV moves over NVLink/RDMA",
|
||
"Admission rule": "Prefill goes to a prefill GPU; KV then "
|
||
"handed off to decode pool",
|
||
"Typical config": ("Two GPU pools with different chip mixes "
|
||
"optimized for prefill vs decode "
|
||
"rooflines")},
|
||
]
|
||
st.dataframe(pd.DataFrame(_layers_rows), width='stretch',
|
||
hide_index=True)
|
||
|
||
st.divider()
|
||
|
||
# ── Section 4: Practical rules of thumb by context regime ─────
|
||
st.markdown("### 4. Practical rules of thumb (by context regime)")
|
||
_rules_rows = [
|
||
{"Regime": "Short context (S_kv < L*)",
|
||
"Batch strategy": "Pack as many users as HBM allows; aim for "
|
||
"B ≈ 2·B*",
|
||
"Utilization": "High (near compute floor)",
|
||
"Cost / token": "Low — cheap tier",
|
||
"Deployment": "Standard elastic pool"},
|
||
{"Regime": "Long context (S_kv > L*)",
|
||
"Batch strategy": "Fewer users per replica; more HBM per user; "
|
||
"lower B",
|
||
"Utilization": "Drops as 1 / (1 + S_kv/L*)",
|
||
"Cost / token": "Higher — accept lower utilization",
|
||
"Deployment": "Length-tier replica with more CP sharding"},
|
||
{"Regime": "Extreme long context (S_kv ≫ L*)",
|
||
"Batch strategy": "Dedicated pool, heavy CP; often "
|
||
"disaggregated prefill",
|
||
"Utilization": "Very low without sparse attention / MLA",
|
||
"Cost / token": "Much higher — priced accordingly",
|
||
"Deployment": "Long-context tier + disaggregated prefill/decode"},
|
||
]
|
||
st.dataframe(pd.DataFrame(_rules_rows), width='stretch',
|
||
hide_index=True)
|
||
|
||
st.divider()
|
||
|
||
# ── Section 5: Sample deployment templates ────────────────────
|
||
st.markdown("### 5. Sample deployment templates (same model, different sharding)")
|
||
st.caption(
|
||
"One base model runs on all three configs. What changes is the "
|
||
"sharding (CP/TP/PP) that shapes the effective 'small' vs "
|
||
"'large' setup. The API gateway routes each request to the "
|
||
"right tier based on its context length."
|
||
)
|
||
_tmpl_rows = [
|
||
{"Tier": "Config_small",
|
||
"CP": 1, "TP": 8, "PP": 1,
|
||
"Total GPUs / replica": 8,
|
||
"Max context": "up to 32k",
|
||
"Typical B": 64,
|
||
"Best for": "Chat, RAG, short prompts — high utilization"},
|
||
{"Tier": "Config_medium",
|
||
"CP": 4, "TP": 8, "PP": 1,
|
||
"Total GPUs / replica": 32,
|
||
"Max context": "up to 128k",
|
||
"Typical B": 32,
|
||
"Best for": "Long-doc analysis, code review — standard tier"},
|
||
{"Tier": "Config_large",
|
||
"CP": 32, "TP": 8, "PP": 1,
|
||
"Total GPUs / replica": 256,
|
||
"Max context": "up to 1M",
|
||
"Typical B": 4,
|
||
"Best for": "Whole-repo / long-video / massive-doc — premium tier"},
|
||
]
|
||
st.dataframe(pd.DataFrame(_tmpl_rows), width='stretch',
|
||
hide_index=True)
|
||
|
||
st.divider()
|
||
|
||
# ── Section 6: Binding-axis playbook ──────────────────────────
|
||
st.markdown("### 6. What to do when each axis binds")
|
||
_playbook_rows = [
|
||
{"Binding axis": "A. Capacity (weights)",
|
||
"Symptom": "Weights alone are close to per-PE HBM budget",
|
||
"First lever": "Increase TP or PP to shard weights across more PEs",
|
||
"Second lever": "Move to a bigger-HBM chip; use lower-precision "
|
||
"weights (FP8/INT4)",
|
||
"Cost impact": "Sub-linear — you pay for more chips, but "
|
||
"utilization stays high"},
|
||
{"Binding axis": "B. KV headroom",
|
||
"Symptom": "Weights fit fine but users × context blows HBM",
|
||
"First lever": "Increase CP (shard sequence dim), or reduce "
|
||
"batch (accept fewer users)",
|
||
"Second lever": "GQA / MQA / MLA to shrink kv_bpt; sparse "
|
||
"attention; KV compression (INT4 KV)",
|
||
"Cost impact": "Direct — long-context requests fundamentally "
|
||
"cost more"},
|
||
{"Binding axis": "C. Throughput SLO",
|
||
"Symptom": "Latency budget forces small B per replica",
|
||
"First lever": "Add more replicas (linear scale)",
|
||
"Second lever": "Loosen SLO if product allows; faster chip; "
|
||
"reduce model size; speculative decoding",
|
||
"Cost impact": "Linear in replica count"},
|
||
]
|
||
st.dataframe(pd.DataFrame(_playbook_rows), width='stretch',
|
||
hide_index=True)
|