12 Commits

Author SHA1 Message Date
mukesh c2b42999d8 analytical-viz: reorder tabs — Physical layout back to first
New tab order:
  Physical layout / Auto Suggest Parallelism / Memory breakdown /
  Per-stage latency / Save & compare / Auto Hardware

Physical layout is what a user typically wants to see first when
loading the app, so it takes the leftmost slot again. Auto Suggest
Parallelism moves to second — still prominent, still usable as a
first step, but doesn't push the layout view behind it.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 14:26:07 -07:00
mukesh 9afeb6bc16 analytical-viz: 3 sidebar apply-buttons (Attn / FFN / Attn+FFN)
Replaces the single memory-only "Apply auto-suggest" button in the
sidebar Parallelism expander with three latency-optimal buttons:
"Attention", "FFN/MoE", "Attn+FFN".

Each clicked button runs run_auto_explore for its scope, picks the
latency-minimum Pareto config, snaps the parallelism knobs to the
sidebar's selectbox option sets (CP/TP/PP/DP), and loads the other
knobs directly (tp_placement, cp_placement, cp_ring_variant, kv_mode,
ffn_scope_label). Also sets _pl_active_scope so the Physical Layout
tab's stage-table filter follows the scope automatically.

The caption above the buttons still shows the memory-only autosuggest
values as a reference — separately labeled "(memory-min)" to avoid
confusion with the latency-optimal buttons below.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 14:24:06 -07:00
mukesh bccc90db82 analytical-viz: Physical Layout scope filter for stage tables
Clicking one of the three Physical Layout tab buttons now persists the
chosen scope in session_state["_pl_active_scope"] and filters the
per-stage latency tables accordingly:

  - attn scope → show only "Attention" stage table
  - ffn scope  → show only "FFN" stage table
  - full scope → show both (default; also matches a fresh sidebar-driven
                  config with no button click yet)

Added a "Layout scope: {label}" header so the user can tell at a glance
which scope's config is loaded.

The pipeline diagram + topology map + weight/tensor sharding + PE
layout continue to show the full model layout — those diagrams give
context that stays useful even when the user is focusing on attn or
FFN. Deeper filtering (e.g., attention-only pipeline stripe) can come
later if needed.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 14:21:27 -07:00
mukesh 1a6d52c58a analytical-viz: force-reload our modules on Streamlit rerun
Streamlit's hot-reload re-executes app.py on every interaction, but
sub-modules imported from app.py stay cached in sys.modules across
reruns. When we edit auto_explore.py or auto_hardware.py while
Streamlit is alive, the app keeps using the pre-edit version until a
full Ctrl+C + restart — leading to confusing "unexpected keyword
argument" errors after a signature change.

Force importlib.reload() on our own modules at the top of app.py so
future signature changes land without a full restart. Only reloads if
the module is already in sys.modules (first run just imports normally).

Applied to:
  - tests.analytical_visualization.auto_explore
  - tests.analytical_visualization.auto_hardware
  - tests.analytical_visualization.autosuggest
  - tests.analytical_visualization.stage_latencies
  - tests.analytical_visualization.memory_layout

Third-party modules (streamlit, matplotlib, pandas, numpy) NOT reloaded
— unnecessary and slow.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 14:17:13 -07:00
mukesh 09721040ca analytical-viz: 3 sweep buttons at top of Physical Layout tab
Shortcut: click any of {Attention, FFN/MoE, Attn+FFN/MoE} → runs
run_auto_explore for that scope, picks the latency-minimum Pareto config,
loads it into the sidebar's session_state (cp, tp, pp, dp, tp_placement,
cp_placement, cp_ring_variant, kv_mode, ffn_scope_label), then st.rerun().

Because the Physical Layout tab reads model + machine + cfg from the
sidebar, the sidebar update automatically redraws the pipeline diagram,
per-stage table, and everything below. No preview / parallel-display
state — WYSIWYG.

Ffn_scope_label reconstruction handles the dynamic "(div=…)" suffix the
sidebar shows (same pattern as auto_explore + auto_hardware tabs' "Load
into sidebar" widgets).

Verified: app.py parses; 24 pytest tests still pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 14:11:30 -07:00
mukesh bf5b659a3d analytical-viz: add FFN-only scope + 3rd sweep button
Both auto tabs now offer three sweep scopes via three buttons instead
of two:
  - Run sweep — Attention          → include_attention=T, include_ffn=F
  - Run sweep — FFN/MoE            → include_attention=F, include_ffn=T
  - Run sweep — Attn + FFN/MoE     → include_attention=T, include_ffn=T

Each button caches its result under its own session_state key; the most
recently clicked button drives the display. All three caches persist so
users can flip between scopes without re-running.

Core changes:
  auto_explore.py + auto_hardware.py:
    - New include_attention: bool = True param alongside include_ffn
    - _sum_visible_latency, _efficiency, score_config, run_auto_explore,
      compute_parallelism_sensitivity, joint_explore, compute_sensitivity,
      _best_parallelism_for_hw, _best_parallelism_two_stage all wired.
    - Attention and FFN are additive with no overlap: measured 7.35 ms
      (attn) + 5.50 ms (ffn) = 12.85 ms (full) for Llama 70B decode 128K.

Bug fix (drive-by): the display block in both tab render functions was
incorrectly nested inside the "cached ctx is stale" warning branch, so
metrics/scatter/table/sensitivity/load only rendered when the cache
was stale. Un-indented and removed the dead trailing else-info block.

Verified:
  - 24 pytest tests pass (added 1 new for FFN-only scope invariants)
  - Smoke: Llama 70B decode 128K all three scopes produce sensible Pareto:
      * Attention only: 7.35 ms (14 pareto configs)
      * FFN / MoE only: 5.50 ms (34 pareto configs)
      * Attn + FFN/MoE: 12.85 ms (6 pareto configs)
    Sums add up exactly, confirming no overlap in stage summation.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 14:09:24 -07:00
mukesh b8e1a3322f analytical-viz: attention-only vs attn+FFN via two sweep buttons
Both auto tabs now accept a scope choice at run time:
  - "Run sweep — Attention"        → include_ffn=False
  - "Run sweep — Attn + FFN/MoE"   → include_ffn=True

Each button runs an independent sweep and caches its result under its
own session_state key. The most recently clicked button determines the
displayed view; both caches persist so users can flip between the two
scopes without re-running.

Tabs:
  - Renamed "Auto Explore" → "Auto Suggest Parallelism"
    (accurately reflects that it only varies parallelism knobs; HW is
    held at the sidebar values).
  - "Auto Hardware" tab unchanged.
  - Still 6 top-level tabs; no additional tabs added.

Core changes:
  auto_explore.py:
    - New include_ffn: bool = True parameter on _sum_visible_latency,
      _efficiency, score_config, run_auto_explore, compute_parallelism_
      sensitivity. False drops all FFN stages from the summed latency.

  auto_hardware.py:
    - New include_ffn: bool = True parameter on joint_explore,
      compute_sensitivity, _best_parallelism_for_hw, _best_parallelism_
      two_stage. Forwards to score_config.

Both defaults keep existing tests byte-identical.

Verified:
  - 23 pytest tests pass (added 3 new: attn-only latency lower, attn-only
    Pareto non-empty, joint HW attn-only faster than full).
  - Smoke: Llama 70B decode 128K:
      * Attn+FFN best latency: 12.85 ms (unchanged)
      * Attention-only best:    7.35 ms (~57% of full)
      * Both sensitivities top-rank bw_hbm_gbs (physics preserved).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 13:58:29 -07:00
mukesh 91b63eb9f6 analytical-viz: parallelism sensitivity chart in Auto Explore tab
Adds compute_parallelism_sensitivity() + ParallelismSensitivityRow to
auto_explore.py. For a baseline ConfigScore (picked from the Pareto set),
sweeps each parallelism knob (CP, TP, PP, DP, EP) individually while
holding others fixed. Reports latency + memory-fit per value.

Sweep values start at 1 and step by 2 (multiples of 2 rather than only
powers of 2), giving finer granularity for the visual than the enumerator's
sparser set. CP goes up to 256 (per real-deployment scale), TP to 64,
PP to 32.

New UI panel in Auto Explore tab: 5 subplots (log/log), one per knob:
  - Solid line = latency where the config fits memory
  - Red X markers = infeasible (out of budget)
  - Dotted vertical line = baseline value
User picks which Pareto row is the "baseline" via a number input; the
sensitivity chart re-computes around it.

Behaviour caveat noted during verification: for large PP the FFN AR can
cross a SIP boundary (uses stage_latencies.py:730 sips_used tier
selection), producing a step-up in latency. This is the existing model's
choice, honestly reflected in the chart. Whether the FFN AR should span
only one PP stage's ranks (thus not cross SIPs) is a separate discussion
about stage_latencies.py.

Verified:
- 11 pytest tests pass (added 2 new for parallelism sensitivity)
- Smoke: at Llama 70B decode 128K with baseline CP=8/TP=16/PP=1/DP=1:
  * CP sweep: 4 fits, 8 = baseline optimum, 16+ overshoots memory
  * TP sweep: 8/16 fit, 16 = baseline optimum, 32 slower (spans SIPs)
  * PP sweep: 1 = baseline, 2+ slower due to sips_used tier drop for FFN AR
  * DP sweep: same shape as PP
  * EP sweep: monotone decreasing (bigger EP = smaller per-PE FFN)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 13:39:51 -07:00
mukesh 55c20bd1a6 analytical-viz: add Auto Hardware Streamlit tab
Consumes auto_hardware.joint_explore(). UI:

  - Sweep-depth radio: two_stage / balanced (default) / coarse
  - Run joint sweep button + spinner
  - 4 metric cards: HW candidates, feasible joint, Pareto count,
    best latency
  - Panel 1 (Pareto scatter): latency vs hardware cost proxy, feasible
    in grey, Pareto colored by PE count, dashed line connects Pareto in
    cost order — this is the "optimal HW config for the model" plot
  - Panel 2 (sensitivity bar chart): per-knob relative speedup when
    doubled from the baseline. Green bars, sorted biggest first;
    annotated with the baseline → doubled values. Answers "where to
    invest next?"
  - Panel 3 (table): sortable Pareto joint configs with parallelism +
    HW spec columns (PE HBM, HBM BW, TFLOPs, PE↔PE, D2D, C2C)
  - Load into sidebar: picks a Pareto row and syncs both the HW
    selectboxes (Per-PE + Interconnect sub-tabs) AND the parallelism
    sliders. Values only get loaded when the target is one of the
    selectbox options; otherwise the sidebar keeps its current value.

Session-state cache under _hw_result keyed by (model, s_kv, mode, depth);
if any of these drift, a warning suggests refreshing.

Verified: app.py parses; auto_hardware smoke run on Llama 70B decode
128K in ~20s produces sensible HW co-design signal (HBM BW 33% speedup,
everything else <1.5%).

Next: verification across Qwen 3 8B, Mixtral 8x7B (Commit 6).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 13:19:13 -07:00
mukesh 6ef09dd6f5 analytical-viz: auto_hardware.py — joint HW×parallelism explore
New module extends auto_explore into the hardware co-design space. For a
fixed model + workload, sweeps hardware knobs (pe_hbm_gb, bw_hbm_gbs,
peak_tflops_f16, bw_intra_gbs, bw_inter_gbs, bw_intersip_gbs) and, for
each hardware candidate, searches parallelism for the latency-minimum
that fits memory. Returns:

  - all_scores: every (hw, parallelism) pair that fits, sorted by latency
  - pareto_scores: 2D Pareto frontier on (latency ↓, cost_score ↓)
  - sensitivity: per-knob rel_speedup when doubled from the best-fast HW
    baseline. Ranks which HW knob gives the biggest speedup — a co-design
    signal.

Three sweep depths trade coverage for time:
  - two_stage: 1 HW candidate (defaults) × autosuggest's memory-min
    parallelism. Fast (~1s), useful for the sensitivity ranking alone.
  - balanced: 64 HW × ~2k reduced-parallelism configs = ~130k joint evals,
    ~10-20s. Default UI setting.
  - coarse:   729 HW × ~2k configs = ~1.4M joint evals, ~2-5 min.

Reduced parallelism sweep for the inner loop: CP × TP × PP × DP ×
kv_shard_mode (1,920 configs), other 4 knobs held at latency-friendly
defaults (ffn_shard_scope='TP+CP', tp_placement='cube', cp_placement='pe',
cp_ring_variant='qoml' for decode, 'kv' for prefill). Full 28,800-config
auto_explore per HW would take 6+ minutes — too slow.

Cost proxy: sum of (knob / knob_default). 6.0 at defaults. Not dollars —
a rough capability score where higher = "more spec'd hardware".

Verified:
- 9 pytest tests pass:
    * enumeration counts match expected (1, 64, 729)
    * default cost_score = 6.0
    * Pareto non-dominated + subset of all_scores
    * every knob is monotone-non-worsening when doubled
    * for Llama 70B decode, bw_hbm_gbs tops the sensitivity ranking
      (physically correct: memory-bound workload)
- Smoke: Llama 70B decode 128K balanced sweep in ~20s produces
  5 Pareto configs; best 7.57 ms with 1024 GB/s HBM BW. Doubling
  HBM BW gives 33% additional speedup; every other knob < 1.5%.

Next: Streamlit UI tab consuming this in Commit 5.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 13:16:40 -07:00
mukesh 2834383700 analytical-viz: add Auto Explore Streamlit tab
New tab consumes auto_explore.run_auto_explore(). UI:
  - Header showing current model + workload + per-PE HBM budget
  - Run sweep button (spinner while ~28k configs run in ~5-10s)
  - 4 metric cards: enumerated, feasible, pareto count, best latency
  - 2-panel scatter:
    * latency vs PEs (feasible in grey, Pareto coloured by efficiency,
      dashed line connects Pareto in PE order to show the trade-off curve)
    * HBM utilization vs latency for Pareto, coloured by PE count
  - Sortable Pareto table
  - "Load into sidebar" widget: pick a row, sets the sidebar
    session_state keys (cp, tp, pp, dp, tp_placement, cp_placement,
    cp_ring_variant, kv_mode, ffn_scope_label) and st.rerun()s so the
    user can flip to another tab and see the full breakdown.

Session-state caches the last sweep result under _auto_explore_result;
if the model/workload/HBM change without re-running, a warning suggests
refreshing.

Ffn_scope_label mapping is dynamic (contains substituted divisors), so
the Load button reconstructs the exact label using the target
cp/tp/dp values before assigning to session_state["ffn_scope_label"].

Verified:
- app.py parses cleanly
- All 9 pytest tests in test_auto_explore.py still pass
- Smoke: matplotlib + pandas + auto_explore imports round-trip

Next: verification pass across Qwen 3 8B and Mixtral 8x7B; polish.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 13:02:08 -07:00
mukesh 85f6716fc1 analytical-viz: auto_explore.py — 9-knob Pareto search
Extends the memory-only autosuggest to search the full 9-dimensional
parallelism space (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
tp_placement, cp_placement, cp_ring_variant) and rank feasible configs
on a 3D Pareto frontier: (latency ↓, pes_used ↓, efficiency ↑).

Throughput is stored on ConfigScore for display but is deliberately NOT
a Pareto axis because for a single-request analysis it collapses to
1 / latency, which would collapse the frontier.

Reuses existing physics (stage_latencies.all_stages + all_ffn_stages,
memory_layout.compute_memory) — no new formulas.

Single-request latency formula fix: PP does NOT reduce single-request
decode/prefill latency because the request has to traverse every layer
sequentially regardless of pipeline depth. The initial version had
latency ~ per_layer × layers_per_stage, which incorrectly rewarded high
PP. Corrected to latency ~ per_layer × model.layers.

Enumerator prunes:
  - PP > model.layers
  - TP > 4 × h_q
  - ffn_shard_scope contains 'DP' when dp=1 (redundant)
  - cp_ring_variant='qoml' when cp=1 (no-op)

Full sweep on Llama 3.1 70B: ~28,800 configs enumerated in ~7s, ~7k-10k
feasible (varies with S_kv), 2-7 unique Pareto configs. Faster context
lengths produce richer frontiers; at 1M, memory forces a single
dominant config (128 PEs, HBM 85%).

Verified:
- 9 pytest tests pass (enumerate, score, Pareto, subset invariants)
- Manual: Llama 70B decode at 8K/64K/128K/1M produces physically
  sensible Pareto (CP=8/TP=16 wins latency; smaller-PE options
  appear at longer context up to memory limits)

Next: Streamlit tab UI in app.py + verification against more presets.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-28 12:58:10 -07:00
5 changed files with 1985 additions and 11 deletions
+752 -11
View File
@@ -6,10 +6,30 @@ Run:
""" """
from __future__ import annotations from __future__ import annotations
import importlib
import sys
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import pandas as pd import pandas as pd
import streamlit as st 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.
for _mod_name in (
"tests.analytical_visualization.auto_explore",
"tests.analytical_visualization.auto_hardware",
"tests.analytical_visualization.autosuggest",
"tests.analytical_visualization.stage_latencies",
"tests.analytical_visualization.memory_layout",
):
_m = sys.modules.get(_mod_name)
if _m is not None:
importlib.reload(_m)
from tests.analytical_visualization.model_config import ( from tests.analytical_visualization.model_config import (
FullConfig, ModelConfig, TopologyConfig, MachineParams, FullConfig, ModelConfig, TopologyConfig, MachineParams,
) )
@@ -156,14 +176,68 @@ if st.session_state.get("_last_preset") != preset_name:
with st.sidebar.expander("Parallelism", expanded=True): with st.sidebar.expander("Parallelism", expanded=True):
st.caption( st.caption(
f"Auto-suggest: CP={_default_suggestion.cp}, " f"Auto-suggest (memory-min): CP={_default_suggestion.cp}, "
f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp}" f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp}"
) )
if st.button("Apply auto-suggest", width='stretch'): # Three latency-optimal apply buttons — pick a scope, load the
st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS) # latency-min Pareto config into the sidebar sliders.
st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS) from tests.analytical_visualization.auto_explore import (
st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS) run_auto_explore as _run_auto_explore_sb,
st.rerun() )
_sb_scope_meta = {
"attn": (True, False, "Attention"),
"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_{_sb_scope}",
help=f"Apply latency-optimal parallelism for "
f"{_sb_scope_meta[_sb_scope][2]} scope.",
)
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"Finding latency-optimal parallelism ({_sb_lbl})..."):
_sb_res = _run_auto_explore_sb(
model, _default_machine, s_kv=s_kv, mode=mode,
include_attention=_sb_ia, include_ffn=_sb_ff,
)
if not _sb_res.pareto_scores:
st.error(
f"No feasible parallelism for {_sb_lbl}. Try raising "
"`pe_hbm_gb` or reducing `s_kv`."
)
else:
_sb_best = min(_sb_res.pareto_scores,
key=lambda s: s.total_latency_ns)
st.session_state["cp"] = _snap(_sb_best.cp, _CP_OPTS)
st.session_state["tp"] = _snap(_sb_best.tp, _TP_OPTS)
st.session_state["pp"] = _snap(_sb_best.pp, _PP_OPTS)
st.session_state["dp"] = _snap(_sb_best.dp, _DP_OPTS)
st.session_state["tp_placement"] = _sb_best.tp_placement
st.session_state["cp_placement"] = _sb_best.cp_placement
st.session_state["cp_ring_variant"] = _sb_best.cp_ring_variant
st.session_state["kv_mode"] = _sb_best.kv_shard_mode
_sb_tp2 = _snap(_sb_best.tp, _TP_OPTS)
_sb_cp2 = _snap(_sb_best.cp, _CP_OPTS)
_sb_dp2 = _snap(_sb_best.dp, _DP_OPTS)
_sb_ffn_map = {
"TP": f"TP only (div={max(1, _sb_tp2)})",
"TP+CP": f"TP*CP (div={max(1, _sb_tp2 * _sb_cp2)})",
"TP+CP+DP": f"TP*CP*DP (div={max(1, _sb_tp2 * _sb_cp2 * _sb_dp2)})",
}
st.session_state["ffn_scope_label"] = \
_sb_ffn_map[_sb_best.ffn_shard_scope]
# Track for the Physical Layout scope filter, so the tables
# below also filter to this scope.
st.session_state["_pl_active_scope"] = _sb_scope
st.rerun()
# Sharding dims — arrange in 2x3 grid to save vertical space # Sharding dims — arrange in 2x3 grid to save vertical space
p_row1 = st.columns(3) p_row1 = st.columns(3)
@@ -419,14 +493,110 @@ if _warnings:
# ── Tabs ───────────────────────────────────────────────────────── # ── Tabs ─────────────────────────────────────────────────────────
tab_layout, tab_memory, tab_stages, tab_compare = st.tabs([ # One tab per auto feature; each has TWO sweep buttons (Attention only /
"Physical layout", "Memory breakdown", "Per-stage latency", # Attn + FFN/MoE) so the user can toggle scope without switching tabs.
"Save & compare", # Auto Suggest is renamed "Parallelism" since it only varies parallelism
# knobs (hardware is held fixed at the sidebar values).
tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
"Physical layout",
"Auto Suggest Parallelism",
"Memory breakdown", "Per-stage latency",
"Save & compare", "Auto Hardware",
]) ])
# ── TAB 1: Physical layout ────────────────────────────────────── # ── TAB 1: Physical layout ──────────────────────────────────────
with tab_layout: with tab_layout:
# Quick shortcut: three buttons that pick the latency-optimal 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 latency-optimal "
"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 latency-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,
)
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:
_pl_best = min(_pl_res.pareto_scores,
key=lambda s: 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 latency-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). # Row A: pipeline diagram (color-coded) + per-stage table (formulas).
st.markdown("**Per-layer pipeline (color = dominant bound)**") st.markdown("**Per-layer pipeline (color = dominant bound)**")
fig_pipe = draw_pipeline(cfg) fig_pipe = draw_pipeline(cfg)
@@ -584,8 +754,10 @@ with tab_layout:
f"blue=compute, orange=memory, red=comm, grey=trivial." f"blue=compute, orange=memory, red=comm, grey=trivial."
) )
_render_table(_stages_attn, "Attention") if _pl_show_attn:
_render_table(_stages_ffn, "FFN") _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?)", with st.expander("Symbol glossary (what does T_q, d_h, div... mean?)",
expanded=False): expanded=False):
@@ -1282,3 +1454,572 @@ with tab_compare:
"'-' means the stage is absent for that config (e.g. C1 in " "'-' means the stage is absent for that config (e.g. C1 in "
"decode, since the O/m/l all-reduce is folded into S8)." "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,
)
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 ~510 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)
_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)
# ── 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)
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)
_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()
@@ -0,0 +1,456 @@
"""Auto-explore parallelism configuration space and rank by Pareto frontier.
Extends the existing ``autosuggest`` (memory-only, single winner) to the
full 9-knob search (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
tp_placement, cp_placement, cp_ring_variant) with four objectives:
- min total_latency_ns (single-request decode/prefill step)
- max throughput_tok_s (naive tokens/sec = 1 / latency for single-req)
- max efficiency_score (geo-mean of compute + BW utilization)
- min pes_used
Reuses ``stage_latencies`` and ``memory_layout`` as the physics; adds
enumeration + 4D Pareto sort on top.
Analytical model runs in ~microseconds per config, so the full sweep
(~5-10k feasible configs after pruning) completes in a few seconds.
"""
from __future__ import annotations
import math
from collections.abc import Iterator
from dataclasses import dataclass, field, replace
from .memory_layout import compute_memory
from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
from .stage_latencies import all_ffn_stages, all_stages
# ── Parallelism sensitivity sweep values (see compute_parallelism_sensitivity)
# Multiples of 2 (finer grid than powers of 2 alone), capped at values that
# are physically meaningful for the modelled topology.
_PARALLELISM_SWEEP_VALUES = {
"cp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64, 96, 128, 192, 256),
"tp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64),
"pp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 32),
"dp": (1, 2, 4, 6, 8, 12, 16),
"ep": (1, 2, 4, 6, 8, 12, 16, 32),
}
# ── Search space ─────────────────────────────────────────────────────
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
_PP_OPTIONS = (1, 2, 4, 8, 16)
_DP_OPTIONS = (1, 2, 4)
_KV_SHARD_MODES = ("split", "replicate")
_FFN_SHARD_SCOPES = ("TP", "TP+CP", "TP+CP+DP")
_TP_PLACEMENTS = ("pe", "cube")
_CP_PLACEMENTS = ("cube", "pe")
_CP_RING_VARIANTS = ("kv", "qoml")
# ── Result types ─────────────────────────────────────────────────────
@dataclass
class ConfigScore:
"""A single (config, computed-metrics) tuple. All floats in SI units
(seconds, tokens/sec, dimensionless) unless suffixed otherwise."""
# ── Config (the 9 knobs) ─────
cp: int
tp: int
pp: int
dp: int
kv_shard_mode: str
ffn_shard_scope: str
tp_placement: str
cp_placement: str
cp_ring_variant: str
# ── Objectives ─────
total_latency_ns: float # ↓ minimize
throughput_tok_s: float # ↑ maximize
efficiency_score: float # ↑ maximize (0..1)
pes_used: int # ↓ minimize
# ── Info-only (not ranked on) ─────
hbm_utilization: float # bytes_used / budget (0..1+; may exceed 1 if over-budget)
weights_gb: float
kv_gb: float
transient_gb: float
sips_used: int
fits_memory: bool
placement_valid: bool
reason: str = ""
@property
def latency_us(self) -> float:
return self.total_latency_ns / 1e3
@property
def latency_ms(self) -> float:
return self.total_latency_ns / 1e6
def as_topology(self, s_kv: int, mode: str) -> TopologyConfig:
"""Reconstruct the TopologyConfig this score was computed for."""
return TopologyConfig(
cp=self.cp, tp=self.tp, pp=self.pp, dp=self.dp,
s_kv=s_kv, mode=mode,
kv_shard_mode=self.kv_shard_mode,
ffn_shard_scope=self.ffn_shard_scope,
tp_placement=self.tp_placement,
cp_placement=self.cp_placement,
cp_ring_variant=self.cp_ring_variant,
)
@dataclass
class AutoExploreResult:
model_name: str
s_kv: int
mode: str
total_enumerated: int # after basic domain pruning
total_feasible: int # after memory + placement checks
all_scores: list[ConfigScore] = field(default_factory=list) # every feasible config
pareto_scores: list[ConfigScore] = field(default_factory=list) # non-dominated set
# ── Enumeration + pruning ────────────────────────────────────────────
def enumerate_configs(
model: ModelConfig,
s_kv: int,
mode: str,
) -> Iterator[TopologyConfig]:
"""Yield every domain-valid TopologyConfig.
Domain rules (fast pruning; no memory check yet):
- PP ≤ model.layers (can't have more stages than layers)
- TP ≤ 4 × model.h_q (unrealistic head-dim splits above this)
- Skip ffn_shard_scope containing 'DP' when DP=1 (redundant with plain 'TP+CP')
- Skip cp_ring_variant='qoml' when CP=1 (no ring, variant is a no-op)
"""
for cp in _CP_OPTIONS:
for tp in _TP_OPTIONS:
if tp > 4 * model.h_q:
continue
for pp in _PP_OPTIONS:
if pp > model.layers:
continue
for dp in _DP_OPTIONS:
for kv_mode in _KV_SHARD_MODES:
for ffn_scope in _FFN_SHARD_SCOPES:
if "DP" in ffn_scope and dp == 1:
continue
for tp_place in _TP_PLACEMENTS:
for cp_place in _CP_PLACEMENTS:
for cp_ring in _CP_RING_VARIANTS:
if cp == 1 and cp_ring == "qoml":
continue
yield TopologyConfig(
cp=cp, tp=tp, pp=pp, dp=dp,
s_kv=s_kv, mode=mode,
kv_shard_mode=kv_mode,
ffn_shard_scope=ffn_scope,
tp_placement=tp_place,
cp_placement=cp_place,
cp_ring_variant=cp_ring,
)
# ── Scoring ──────────────────────────────────────────────────────────
def _sum_visible_latency(
cfg: FullConfig,
include_attention: bool = True,
include_ffn: bool = True,
) -> float:
"""Total single-request latency (seconds) across all model layers.
A single request traverses every layer sequentially, whether the layers
sit on one PP stage or are spread across many:
- PP=1 : one rank holds all L layers, latency = L × per_layer
- PP=K : K ranks each hold L/K layers, but request crosses all K
stages sequentially → still L × per_layer
PP therefore does NOT reduce single-request latency; it only improves
throughput under batching. This function is the single-request cost, so
we multiply by full model.layers regardless of PP.
Scope selection (both default True — full per-token cost):
- ``include_attention=True, include_ffn=True`` → full transformer
- ``include_attention=True, include_ffn=False`` → attention only
- ``include_attention=False, include_ffn=True`` → FFN / MoE only
- ``include_attention=False, include_ffn=False`` → zero (rejected caller-side)
"""
attn = sum(s.visible_s for s in all_stages(cfg)) if include_attention else 0.0
ffn = sum(s.visible_s for s in all_ffn_stages(cfg)) if include_ffn else 0.0
per_layer = attn + ffn
return per_layer * cfg.model.layers
def _efficiency(cfg: FullConfig, latency_s: float,
include_attention: bool = True,
include_ffn: bool = True) -> float:
"""Geo-mean of compute-util and BW-util. Range ~ (0, 1].
- compute_util = achieved_flops / (peak_flops × pes × latency)
- bw_util = achieved_bytes / (peak_bw × pes × latency)
Scope flags mirror :func:`_sum_visible_latency`.
"""
if latency_s <= 0:
return 0.0
attn = all_stages(cfg) if include_attention else []
ffn = all_ffn_stages(cfg) if include_ffn else []
layers = math.ceil(cfg.model.layers / cfg.topo.pp)
total_flops = layers * sum(s.flops for s in attn + ffn)
total_bytes = layers * sum(s.mem_bytes for s in attn + ffn)
pes = cfg.topo.total_pes
peak_flops = cfg.machine.peak_flops * pes
peak_bw = cfg.machine.bw_hbm * pes
compute_util = total_flops / (peak_flops * latency_s) if peak_flops > 0 else 0.0
bw_util = total_bytes / (peak_bw * latency_s) if peak_bw > 0 else 0.0
compute_util = min(1.0, max(0.0, compute_util))
bw_util = min(1.0, max(0.0, bw_util))
# Geo-mean; if either is 0 the score is 0 (avoids overrewarding lopsided configs).
return math.sqrt(compute_util * bw_util)
def score_config(cfg: FullConfig,
include_attention: bool = True,
include_ffn: bool = True) -> ConfigScore:
"""Compute all 4 objectives + info fields for one config.
Feasibility (memory + placement) is stored but does NOT gate scoring —
infeasible configs get returned with fits_memory=False so callers can
filter or display them.
Scope flags (default: full transformer) restrict *latency + efficiency*
to attention only, FFN only, or both. Memory feasibility is unchanged —
still checks weights+KV+transient fit in per-PE HBM since the model
still exists physically regardless of what the caller is scoring.
"""
mem = compute_memory(cfg)
placement_ok = cfg.topo.placement_valid
latency_s = _sum_visible_latency(
cfg, include_attention=include_attention, include_ffn=include_ffn,
)
throughput = 1.0 / latency_s if latency_s > 0 else 0.0
efficiency = (
_efficiency(cfg, latency_s,
include_attention=include_attention, include_ffn=include_ffn)
if latency_s > 0 else 0.0
)
fits = not mem.over_budget
reason = ""
if not fits:
reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
f"exceeds per-PE budget ({mem.budget_bytes/1e9:.2f} GB)")
elif not placement_ok:
reason = (f"intra-cube demand ({cfg.topo.intra_cube_dims}) "
f"exceeds PEs/cube ({cfg.topo.pes_per_cube_hw})")
return ConfigScore(
cp=cfg.topo.cp, tp=cfg.topo.tp, pp=cfg.topo.pp, dp=cfg.topo.dp,
kv_shard_mode=cfg.topo.kv_shard_mode,
ffn_shard_scope=cfg.topo.ffn_shard_scope,
tp_placement=cfg.topo.tp_placement,
cp_placement=cfg.topo.cp_placement,
cp_ring_variant=cfg.topo.cp_ring_variant,
total_latency_ns=latency_s * 1e9,
throughput_tok_s=throughput,
efficiency_score=efficiency,
pes_used=cfg.topo.total_pes,
hbm_utilization=mem.used_bytes / mem.budget_bytes if mem.budget_bytes > 0 else 0.0,
weights_gb=mem.weights_bytes / 1e9,
kv_gb=mem.kv_cache_bytes / 1e9,
transient_gb=mem.transient_bytes / 1e9,
sips_used=cfg.topo.sips_used,
fits_memory=fits,
placement_valid=placement_ok,
reason=reason,
)
# ── Pareto sort ──────────────────────────────────────────────────────
def _dominates(a: ConfigScore, b: ConfigScore) -> bool:
"""True iff a is no worse than b on every axis AND strictly better on ≥ 1.
Axes (with direction) — 3D Pareto:
total_latency_ns ↓ (a ≤ b)
pes_used ↓ (a ≤ b) — proxy for cost / deployment size
efficiency_score ↑ (a ≥ b) — geo-mean of compute+BW utilization
Note: ``throughput_tok_s`` is deliberately NOT a Pareto axis. For a
single-request analysis, throughput = 1 / latency, so it's collinear
with latency and would collapse the frontier. It stays on ConfigScore
for display but doesn't participate in domination.
"""
no_worse = (
a.total_latency_ns <= b.total_latency_ns
and a.pes_used <= b.pes_used
and a.efficiency_score >= b.efficiency_score
)
if not no_worse:
return False
strictly_better = (
a.total_latency_ns < b.total_latency_ns
or a.pes_used < b.pes_used
or a.efficiency_score > b.efficiency_score
)
return strictly_better
def pareto_frontier(scores: list[ConfigScore]) -> list[ConfigScore]:
"""Extract non-dominated set on (latency↓, pe↓, throughput↑, efficiency↑).
Only feasible configs (fits_memory AND placement_valid) participate; the
non-feasible are excluded from the frontier (they can still be returned in
``all_scores`` for informational display).
O(N²) — fine for N up to ~50k with early exits.
"""
feasible = [s for s in scores if s.fits_memory and s.placement_valid]
frontier: list[ConfigScore] = []
for i, a in enumerate(feasible):
dominated = False
for j, b in enumerate(feasible):
if i == j:
continue
if _dominates(b, a):
dominated = True
break
if not dominated:
frontier.append(a)
return frontier
# ── Parallelism sensitivity ─────────────────────────────────────────
@dataclass
class ParallelismSensitivityRow:
"""One knob's sweep result: latency + fit at every tested value.
All *other* parallelism knobs (including HW) are held at their baseline
values, so this isolates the effect of just this one knob.
"""
knob: str # "cp" / "tp" / "pp" / "dp" / "ep"
values: list[int] # tested values (from _PARALLELISM_SWEEP_VALUES)
latencies_ns: list[float] # one per value; NaN when infeasible
fits_flags: list[bool] # per-value memory+placement feasibility
baseline_value: int # what the baseline config uses for this knob
baseline_latency_ns: float
def compute_parallelism_sensitivity(
baseline: ConfigScore,
model: ModelConfig,
machine: MachineParams,
s_kv: int,
mode: str,
include_attention: bool = True,
include_ffn: bool = True,
) -> list[ParallelismSensitivityRow]:
"""For each parallelism knob (CP, TP, PP, DP, EP), sweep its values
holding the OTHER knobs fixed at the baseline. Reports latency + memory
fit per value.
Answers: 'if I only change this one knob, how does latency and memory
fit change?' — the complement to auto_hardware.compute_sensitivity
which does the same for hardware knobs.
"""
baseline_topo = baseline.as_topology(s_kv, mode)
rows: list[ParallelismSensitivityRow] = []
for knob, values in _PARALLELISM_SWEEP_VALUES.items():
baseline_val = getattr(baseline_topo, knob, 1)
# EP isn't stored on ConfigScore's as_topology output today (dp
# attribute exists but ep does not always). Fall back to 1.
latencies: list[float] = []
fits: list[bool] = []
for v in values:
# Skip nonsensical values that would violate hard bounds.
if knob == "pp" and v > model.layers:
latencies.append(float("nan"))
fits.append(False)
continue
if knob == "tp" and v > 4 * model.h_q:
latencies.append(float("nan"))
fits.append(False)
continue
# Build a variant TopologyConfig with just this one knob changed.
trial = replace(baseline_topo, **{knob: v})
cfg = FullConfig(model=model, topo=trial, machine=machine)
score = score_config(
cfg, include_attention=include_attention, include_ffn=include_ffn,
)
latencies.append(score.total_latency_ns)
fits.append(score.fits_memory and score.placement_valid)
rows.append(ParallelismSensitivityRow(
knob=knob,
values=list(values),
latencies_ns=latencies,
fits_flags=fits,
baseline_value=int(baseline_val),
baseline_latency_ns=baseline.total_latency_ns,
))
return rows
# ── Top-level driver ─────────────────────────────────────────────────
def run_auto_explore(
model: ModelConfig,
machine: MachineParams,
s_kv: int,
mode: str = "decode",
include_attention: bool = True,
include_ffn: bool = True,
) -> AutoExploreResult:
"""Enumerate all configs, score each, extract Pareto frontier.
Returns both the full ``all_scores`` list (for the table view) and the
``pareto_scores`` subset (for the scatter/highlight view). Both are sorted
by total_latency_ns ascending.
Scope: at least one of ``include_attention`` / ``include_ffn`` must be
True. Setting both False would give zero latency for every config —
the caller should not do that.
"""
all_scores: list[ConfigScore] = []
total_enumerated = 0
for topo in enumerate_configs(model, s_kv, mode):
total_enumerated += 1
cfg = FullConfig(model=model, topo=topo, machine=machine)
all_scores.append(score_config(
cfg, include_attention=include_attention, include_ffn=include_ffn,
))
all_scores.sort(key=lambda s: s.total_latency_ns)
pareto = pareto_frontier(all_scores)
pareto.sort(key=lambda s: s.total_latency_ns)
feasible_count = sum(1 for s in all_scores if s.fits_memory and s.placement_valid)
return AutoExploreResult(
model_name=model.name,
s_kv=s_kv,
mode=mode,
total_enumerated=total_enumerated,
total_feasible=feasible_count,
all_scores=all_scores,
pareto_scores=pareto,
)
@@ -0,0 +1,418 @@
"""Joint hardware × parallelism exploration.
For a fixed model + workload, sweep both:
- hardware parameters (MachineParams: pe_hbm_gb, bw_hbm_gbs, peak_tflops_f16,
bw_intra_gbs, bw_inter_gbs, bw_intersip_gbs), and
- parallelism knobs (a reduced subset of auto_explore's 9 knobs)
Then rank on:
- latency ↓
- hardware cost proxy ↓ (normalized sum of knob-over-default)
Also computes per-knob sensitivity — how much latency drops when each
hardware knob is doubled from its baseline value. Useful for HW co-design
("which knob to invest in first?").
The default 'balanced' depth: 64 HW candidates × ~2k parallelism configs =
~130k joint evaluations, ~1-5 s at analytical speeds. Coarse and
two-stage variants trade coverage for speed.
"""
from __future__ import annotations
import math
from collections.abc import Iterator
from dataclasses import dataclass, field, replace
from .auto_explore import ConfigScore, score_config
from .autosuggest import auto_suggest
from .memory_layout import compute_memory
from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
# ── Hardware search space ────────────────────────────────────────────
#
# For each depth level, values per knob. Defaults are always included so the
# baseline configuration always appears in the sweep.
_HW_KNOB_DEFAULTS = {
"pe_hbm_gb": 6.0,
"bw_hbm_gbs": 256.0,
"peak_tflops_f16": 8.0,
"bw_intra_gbs": 512.0,
"bw_inter_gbs": 128.0,
"bw_intersip_gbs": 50.0,
}
_HW_VALUES_BALANCED = {
"pe_hbm_gb": [6.0, 12.0],
"bw_hbm_gbs": [256.0, 512.0],
"peak_tflops_f16": [8.0, 16.0],
"bw_intra_gbs": [512.0, 1024.0],
"bw_inter_gbs": [128.0, 256.0],
"bw_intersip_gbs": [50.0, 100.0],
}
_HW_VALUES_COARSE = {
"pe_hbm_gb": [6.0, 12.0, 24.0],
"bw_hbm_gbs": [256.0, 512.0, 1024.0],
"peak_tflops_f16": [8.0, 16.0, 32.0],
"bw_intra_gbs": [512.0, 1024.0, 2048.0],
"bw_inter_gbs": [128.0, 256.0, 512.0],
"bw_intersip_gbs": [50.0, 100.0, 200.0],
}
# For sensitivity: "double each knob independently from baseline."
_SENSITIVITY_KNOBS = list(_HW_KNOB_DEFAULTS.keys())
# ── Reduced parallelism search space (for per-HW inner loop) ─────────
#
# The full auto_explore has 28,800 configs; using it inside a HW loop is
# too slow. Reduce to CP × TP × PP × DP × kv_shard_mode with common
# defaults for the remaining 4 knobs (~2k configs).
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
_PP_OPTIONS = (1, 2, 4, 8, 16)
_DP_OPTIONS = (1, 2, 4)
_KV_SHARD_MODES = ("split", "replicate")
# ── Result types ─────────────────────────────────────────────────────
@dataclass
class HardwareCandidate:
pe_hbm_gb: float
bw_hbm_gbs: float
peak_tflops_f16: float
bw_intra_gbs: float
bw_inter_gbs: float
bw_intersip_gbs: float
def as_machine(self) -> MachineParams:
"""Build a MachineParams from this HW candidate, holding alphas +
compute_util at their MachineParams defaults."""
return MachineParams(
pe_hbm_gb=self.pe_hbm_gb,
bw_hbm_gbs=self.bw_hbm_gbs,
peak_tflops_f16=self.peak_tflops_f16,
bw_intra_gbs=self.bw_intra_gbs,
bw_inter_gbs=self.bw_inter_gbs,
bw_intersip_gbs=self.bw_intersip_gbs,
)
@property
def cost_score(self) -> float:
"""Normalized cost proxy: sum of (knob / default). 6.0 at defaults.
This is NOT dollars — it's a rough silicon-area / capability proxy. A
higher cost_score means "more capable hardware" (more HBM, faster BW,
etc.). Under identical latency, lower cost_score wins.
"""
return sum(
getattr(self, k) / _HW_KNOB_DEFAULTS[k]
for k in _HW_KNOB_DEFAULTS
)
@dataclass
class JointScore:
"""One (hardware, parallelism) pair with its computed latency + cost."""
hardware: HardwareCandidate
parallelism: ConfigScore
total_latency_ns: float
cost_score: float
@property
def latency_ms(self) -> float:
return self.total_latency_ns / 1e6
@dataclass
class SensitivityRow:
knob: str
baseline_value: float
doubled_value: float
baseline_latency_ns: float
doubled_latency_ns: float
@property
def rel_speedup(self) -> float:
"""1 - (doubled/baseline). Positive = doubling this knob speeds things up."""
if self.baseline_latency_ns <= 0:
return 0.0
return 1.0 - (self.doubled_latency_ns / self.baseline_latency_ns)
@dataclass
class JointExploreResult:
model_name: str
s_kv: int
mode: str
depth: str
total_hw: int
total_joint: int
all_scores: list[JointScore] = field(default_factory=list)
pareto_scores: list[JointScore] = field(default_factory=list)
sensitivity: list[SensitivityRow] = field(default_factory=list)
# ── Hardware enumeration ─────────────────────────────────────────────
def enumerate_hardware(depth: str = "balanced") -> Iterator[HardwareCandidate]:
"""Yield HardwareCandidates for the given sweep depth.
depth ∈ {"two_stage", "balanced", "coarse"}:
- two_stage: only the default HW (1 candidate; parallelism sweep dominates)
- balanced: 2 values per knob → 2^6 = 64 candidates
- coarse: 3 values per knob → 3^6 = 729 candidates
"""
if depth == "two_stage":
vals = {k: [v] for k, v in _HW_KNOB_DEFAULTS.items()}
elif depth == "coarse":
vals = _HW_VALUES_COARSE
else: # balanced (default)
vals = _HW_VALUES_BALANCED
for pe_hbm in vals["pe_hbm_gb"]:
for bw_hbm in vals["bw_hbm_gbs"]:
for tflops in vals["peak_tflops_f16"]:
for bw_intra in vals["bw_intra_gbs"]:
for bw_inter in vals["bw_inter_gbs"]:
for bw_intersip in vals["bw_intersip_gbs"]:
yield HardwareCandidate(
pe_hbm_gb=pe_hbm,
bw_hbm_gbs=bw_hbm,
peak_tflops_f16=tflops,
bw_intra_gbs=bw_intra,
bw_inter_gbs=bw_inter,
bw_intersip_gbs=bw_intersip,
)
# ── Reduced parallelism search per HW ────────────────────────────────
def _iter_reduced_parallelism(
model: ModelConfig, s_kv: int, mode: str,
) -> Iterator[TopologyConfig]:
"""Reduced sweep: CP × TP × PP × DP × kv_shard_mode.
Other 4 knobs held at defaults chosen to be latency-friendly for the
decode/prefill common case:
- ffn_shard_scope = "TP+CP" (common good choice; trades a small AR
for lower per-PE weight bytes)
- tp_placement = "cube" (packs TP across cubes; UCIe D2D)
- cp_placement = "pe" (packs CP inside cubes; intra NoC)
- cp_ring_variant = "qoml" for decode, "kv" for prefill
Total: 8 × 6 × 5 × 4 × 2 = 1,920 configs per HW candidate.
"""
cp_ring = "qoml" if mode == "decode" else "kv"
for cp in _CP_OPTIONS:
for tp in _TP_OPTIONS:
if tp > 4 * model.h_q:
continue
for pp in _PP_OPTIONS:
if pp > model.layers:
continue
for dp in _DP_OPTIONS:
for kv_mode in _KV_SHARD_MODES:
# cp_ring=qoml with cp=1 is a no-op; fall back to kv.
ring = "kv" if cp == 1 else cp_ring
yield TopologyConfig(
cp=cp, tp=tp, pp=pp, dp=dp,
s_kv=s_kv, mode=mode,
kv_shard_mode=kv_mode,
ffn_shard_scope="TP+CP",
tp_placement="cube",
cp_placement="pe",
cp_ring_variant=ring,
)
def _best_parallelism_for_hw(
model: ModelConfig, machine: MachineParams,
s_kv: int, mode: str,
include_attention: bool = True, include_ffn: bool = True,
) -> ConfigScore | None:
"""Return the latency-minimum feasible parallelism for this HW.
Feasibility: fits memory + placement_valid. Returns None if nothing fits.
Scope flags are forwarded to score_config so the "latency" ranked on
matches the caller's attention/FFN/full choice.
"""
best: ConfigScore | None = None
for topo in _iter_reduced_parallelism(model, s_kv, mode):
cfg = FullConfig(model=model, topo=topo, machine=machine)
s = score_config(cfg, include_attention=include_attention,
include_ffn=include_ffn)
if not (s.fits_memory and s.placement_valid):
continue
if best is None or s.total_latency_ns < best.total_latency_ns:
best = s
return best
def _best_parallelism_two_stage(
model: ModelConfig, machine: MachineParams,
s_kv: int, mode: str,
include_attention: bool = True, include_ffn: bool = True,
) -> ConfigScore | None:
"""Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it."""
sug = auto_suggest(model, machine, s_kv, mode)
if not sug.fits:
return None
topo = TopologyConfig(
cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1,
s_kv=s_kv, mode=mode,
kv_shard_mode="split",
ffn_shard_scope="TP+CP",
tp_placement="cube",
cp_placement="pe",
cp_ring_variant="qoml" if mode == "decode" and sug.cp > 1 else "kv",
)
cfg = FullConfig(model=model, topo=topo, machine=machine)
s = score_config(cfg, include_attention=include_attention,
include_ffn=include_ffn)
return s if s.fits_memory and s.placement_valid else None
# ── Pareto over joint (latency, cost) ────────────────────────────────
def _pareto_2d(scores: list[JointScore]) -> list[JointScore]:
"""Non-dominated set on (total_latency_ns ↓, cost_score ↓). O(N²)."""
frontier: list[JointScore] = []
for i, a in enumerate(scores):
dominated = False
for j, b in enumerate(scores):
if i == j:
continue
no_worse = (
b.total_latency_ns <= a.total_latency_ns
and b.cost_score <= a.cost_score
)
strictly_better = (
b.total_latency_ns < a.total_latency_ns
or b.cost_score < a.cost_score
)
if no_worse and strictly_better:
dominated = True
break
if not dominated:
frontier.append(a)
return frontier
# ── Per-knob sensitivity ─────────────────────────────────────────────
def compute_sensitivity(
baseline_hw: HardwareCandidate,
parallelism: ConfigScore,
model: ModelConfig, s_kv: int, mode: str,
include_attention: bool = True, include_ffn: bool = True,
) -> list[SensitivityRow]:
"""For each HW knob, double it (holding others at baseline) and measure
the latency change. Same parallelism used throughout so we isolate the
HW knob's effect."""
rows: list[SensitivityRow] = []
baseline_machine = baseline_hw.as_machine()
baseline_topo = parallelism.as_topology(s_kv, mode)
baseline_cfg = FullConfig(
model=model, topo=baseline_topo, machine=baseline_machine,
)
baseline_score = score_config(
baseline_cfg, include_attention=include_attention, include_ffn=include_ffn,
)
baseline_latency = baseline_score.total_latency_ns
for knob in _SENSITIVITY_KNOBS:
doubled_val = 2.0 * getattr(baseline_hw, knob)
doubled_hw = replace(baseline_hw, **{knob: doubled_val})
doubled_cfg = FullConfig(
model=model, topo=baseline_topo, machine=doubled_hw.as_machine(),
)
doubled_score = score_config(
doubled_cfg, include_attention=include_attention, include_ffn=include_ffn,
)
rows.append(SensitivityRow(
knob=knob,
baseline_value=getattr(baseline_hw, knob),
doubled_value=doubled_val,
baseline_latency_ns=baseline_latency,
doubled_latency_ns=doubled_score.total_latency_ns,
))
# Sort by biggest speedup first (most sensitive knob at the top).
rows.sort(key=lambda r: -r.rel_speedup)
return rows
# ── Top-level driver ─────────────────────────────────────────────────
def joint_explore(
model: ModelConfig,
s_kv: int,
mode: str,
depth: str = "balanced",
include_attention: bool = True,
include_ffn: bool = True,
) -> JointExploreResult:
"""Sweep HW candidates × parallelism, return joint Pareto + sensitivity.
Sensitivity is computed around the LATENCY-MINIMUM joint point (the
"best fast" config), doubling each HW knob one at a time.
Scope flags select which stages contribute to the summed latency —
both the per-HW parallelism search and the sensitivity ranking use
the same restriction.
"""
all_scores: list[JointScore] = []
for hw in enumerate_hardware(depth):
machine = hw.as_machine()
if depth == "two_stage":
par = _best_parallelism_two_stage(
model, machine, s_kv, mode,
include_attention=include_attention, include_ffn=include_ffn,
)
else:
par = _best_parallelism_for_hw(
model, machine, s_kv, mode,
include_attention=include_attention, include_ffn=include_ffn,
)
if par is None:
continue
all_scores.append(JointScore(
hardware=hw,
parallelism=par,
total_latency_ns=par.total_latency_ns,
cost_score=hw.cost_score,
))
all_scores.sort(key=lambda s: s.total_latency_ns)
pareto = _pareto_2d(all_scores)
pareto.sort(key=lambda s: s.total_latency_ns)
sensitivity: list[SensitivityRow] = []
if all_scores:
best = all_scores[0]
sensitivity = compute_sensitivity(
best.hardware, best.parallelism, model, s_kv, mode,
include_attention=include_attention, include_ffn=include_ffn,
)
return JointExploreResult(
model_name=model.name,
s_kv=s_kv,
mode=mode,
depth=depth,
total_hw=sum(1 for _ in enumerate_hardware(depth)),
total_joint=len(all_scores),
all_scores=all_scores,
pareto_scores=pareto,
sensitivity=sensitivity,
)
@@ -0,0 +1,227 @@
"""Interface + smoke tests for auto_explore.
Covers:
- enumerate_configs yields valid TopologyConfigs with expected pruning
- score_config returns a ConfigScore with all fields populated
- pareto_frontier is a subset of feasible and is non-empty for a
reasonable model+workload
- run_auto_explore returns a coherent AutoExploreResult (all_scores
sorted by latency, pareto ⊆ feasible ⊆ all_scores)
"""
from __future__ import annotations
import pytest
from tests.analytical_visualization.auto_explore import (
ConfigScore,
enumerate_configs,
pareto_frontier,
run_auto_explore,
score_config,
)
from tests.analytical_visualization.model_config import (
FullConfig, MachineParams,
)
from tests.analytical_visualization.model_presets import PRESETS
# ── Enumeration ─────────────────────────────────────────────────────
def test_enumerate_yields_valid_configs():
"""Enumerator produces TopologyConfigs with all 9 knobs set."""
model = PRESETS["Llama 3.1 70B"].model
it = enumerate_configs(model, s_kv=8192, mode="decode")
first = next(it)
for f in ("cp", "tp", "pp", "dp", "kv_shard_mode",
"ffn_shard_scope", "tp_placement", "cp_placement",
"cp_ring_variant"):
assert getattr(first, f) is not None, f"knob {f} not set"
assert first.s_kv == 8192
assert first.mode == "decode"
def test_enumerate_prunes_pp_beyond_layers():
"""PP > model.layers is skipped by the enumerator."""
model = PRESETS["Llama 3.1 70B"].model # layers=80
for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
assert cfg.pp <= model.layers
def test_enumerate_prunes_ffn_dp_when_dp_is_1():
"""ffn_shard_scope containing 'DP' is skipped when dp=1 (redundant)."""
model = PRESETS["Llama 3.1 70B"].model
for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
if cfg.dp == 1:
assert "DP" not in cfg.ffn_shard_scope
# ── Scoring ─────────────────────────────────────────────────────────
def test_score_returns_populated_config_score():
"""score_config returns a fully-populated ConfigScore."""
model = PRESETS["Llama 3.1 70B"].model
topo = next(enumerate_configs(model, s_kv=8192, mode="decode"))
machine = MachineParams()
cfg = FullConfig(model=model, topo=topo, machine=machine)
score = score_config(cfg)
assert isinstance(score, ConfigScore)
assert score.total_latency_ns > 0
assert score.pes_used == topo.total_pes
assert 0.0 <= score.efficiency_score <= 1.0
assert score.hbm_utilization >= 0.0
# ── Pareto ──────────────────────────────────────────────────────────
def test_pareto_subset_of_feasible():
"""Pareto scores are always feasible (memory + placement)."""
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
for p in res.pareto_scores:
assert p.fits_memory
assert p.placement_valid
def test_pareto_non_empty_when_feasible_configs_exist():
"""When at least one config fits, Pareto must have ≥ 1 entry."""
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
assert res.total_feasible > 0
assert len(res.pareto_scores) > 0
def test_pareto_frontier_non_dominated():
"""No Pareto entry is dominated by another."""
from tests.analytical_visualization.auto_explore import _dominates
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
for i, a in enumerate(res.pareto_scores):
for j, b in enumerate(res.pareto_scores):
if i == j:
continue
assert not _dominates(b, a), (
f"Pareto entry {i} is dominated by entry {j}"
)
# ── End-to-end sanity ───────────────────────────────────────────────
def test_run_auto_explore_shape():
"""all_scores sorted asc by latency; pareto ⊆ feasible ⊆ all_scores."""
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
assert res.total_enumerated == len(res.all_scores)
assert res.total_feasible == sum(
1 for s in res.all_scores if s.fits_memory and s.placement_valid
)
assert len(res.pareto_scores) <= res.total_feasible
latencies = [s.total_latency_ns for s in res.all_scores]
assert latencies == sorted(latencies)
def test_pp_does_not_reduce_single_request_latency():
"""A single request traverses all layers regardless of PP.
Latency-optimal Pareto configs should NOT prefer PP>1 for decode."""
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
# Sort Pareto by latency; the fastest should not need PP>1 to win.
fastest = res.pareto_scores[0]
assert fastest.pp == 1, (
f"expected PP=1 for the fastest decode config; got PP={fastest.pp}"
)
# ── Parallelism sensitivity ─────────────────────────────────────────
def test_parallelism_sensitivity_has_all_five_knobs():
"""compute_parallelism_sensitivity returns one row per knob."""
from tests.analytical_visualization.auto_explore import (
compute_parallelism_sensitivity,
)
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
baseline = res.pareto_scores[0]
rows = compute_parallelism_sensitivity(
baseline, model, machine, s_kv=8192, mode="decode",
)
knobs = {r.knob for r in rows}
assert knobs == {"cp", "tp", "pp", "dp", "ep"}
def test_attention_only_latency_is_lower_than_full():
"""include_ffn=False must produce lower latency than include_ffn=True
for the same model+workload (FFN cost is dropped)."""
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
include_ffn=True)
r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
include_ffn=False)
best_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
best_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
assert best_attn.total_latency_ns < best_full.total_latency_ns, (
f"attn-only ({best_attn.latency_us:.2f} us) should be less than "
f"full ({best_full.latency_us:.2f} us)"
)
def test_attention_only_pareto_non_empty():
"""The attention-only sweep must still produce a non-empty Pareto set."""
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode",
include_attention=True, include_ffn=False)
assert res.total_feasible > 0
assert len(res.pareto_scores) > 0
def test_ffn_only_latency_less_than_full_and_different_from_attn():
"""FFN-only latency must be > 0, < full-model latency, and != attn-only."""
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
include_attention=True, include_ffn=True)
r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
include_attention=True, include_ffn=False)
r_ffn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
include_attention=False, include_ffn=True)
b_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
b_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
b_ffn = min(r_ffn.pareto_scores, key=lambda s: s.total_latency_ns)
assert b_ffn.total_latency_ns > 0
assert b_ffn.total_latency_ns < b_full.total_latency_ns
assert abs(b_ffn.total_latency_ns - b_attn.total_latency_ns) > 1.0, (
"FFN-only and attention-only latencies should not be identical"
)
def test_parallelism_sensitivity_includes_baseline_value():
"""Each row's values contain the baseline_value."""
from tests.analytical_visualization.auto_explore import (
compute_parallelism_sensitivity,
)
model = PRESETS["Llama 3.1 70B"].model
machine = MachineParams()
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
baseline = res.pareto_scores[0]
rows = compute_parallelism_sensitivity(
baseline, model, machine, s_kv=8192, mode="decode",
)
for r in rows:
# Baseline value must be sweepable OR mapped to something in values.
assert r.baseline_value in r.values or r.baseline_value == 1, (
f"knob {r.knob}: baseline_value={r.baseline_value} not in {r.values}"
)
@@ -0,0 +1,132 @@
"""Interface + invariant tests for auto_hardware."""
from __future__ import annotations
from tests.analytical_visualization.auto_hardware import (
HardwareCandidate,
JointScore,
_HW_KNOB_DEFAULTS,
enumerate_hardware,
joint_explore,
)
from tests.analytical_visualization.model_presets import PRESETS
# ── Enumeration ─────────────────────────────────────────────────────
def test_enumerate_two_stage_yields_one():
"""Two-stage depth yields exactly 1 HW candidate (defaults)."""
hws = list(enumerate_hardware("two_stage"))
assert len(hws) == 1
for knob, default_val in _HW_KNOB_DEFAULTS.items():
assert getattr(hws[0], knob) == default_val
def test_enumerate_balanced_yields_64():
"""Balanced = 2 values × 6 knobs = 2^6 = 64 candidates."""
hws = list(enumerate_hardware("balanced"))
assert len(hws) == 64
def test_enumerate_coarse_yields_729():
"""Coarse = 3 values × 6 knobs = 3^6 = 729 candidates."""
hws = list(enumerate_hardware("coarse"))
assert len(hws) == 729
def test_default_cost_score_is_6():
"""At every knob = its default, cost_score = 6 (one per knob)."""
hw = HardwareCandidate(
pe_hbm_gb=_HW_KNOB_DEFAULTS["pe_hbm_gb"],
bw_hbm_gbs=_HW_KNOB_DEFAULTS["bw_hbm_gbs"],
peak_tflops_f16=_HW_KNOB_DEFAULTS["peak_tflops_f16"],
bw_intra_gbs=_HW_KNOB_DEFAULTS["bw_intra_gbs"],
bw_inter_gbs=_HW_KNOB_DEFAULTS["bw_inter_gbs"],
bw_intersip_gbs=_HW_KNOB_DEFAULTS["bw_intersip_gbs"],
)
assert hw.cost_score == 6.0
# ── Joint explore + Pareto ──────────────────────────────────────────
def test_joint_explore_returns_pareto_non_empty():
"""Given a reasonable model+workload, at least one HW+parallelism fits."""
model = PRESETS["Llama 3.1 70B"].model
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
assert res.total_joint >= 1
assert len(res.pareto_scores) >= 1
def test_pareto_subset_of_all_scores():
"""Every Pareto entry is present in all_scores."""
model = PRESETS["Llama 3.1 70B"].model
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
pareto_ids = {id(s) for s in res.pareto_scores}
all_ids = {id(s) for s in res.all_scores}
assert pareto_ids.issubset(all_ids)
def test_pareto_non_dominated():
"""No Pareto entry is dominated by another Pareto entry on (lat, cost)."""
model = PRESETS["Llama 3.1 70B"].model
res = joint_explore(model, s_kv=131072, mode="decode", depth="balanced")
for i, a in enumerate(res.pareto_scores):
for j, b in enumerate(res.pareto_scores):
if i == j:
continue
no_worse = (
b.total_latency_ns <= a.total_latency_ns
and b.cost_score <= a.cost_score
)
strictly_better = (
b.total_latency_ns < a.total_latency_ns
or b.cost_score < a.cost_score
)
assert not (no_worse and strictly_better), (
f"Pareto entry {i} is dominated by entry {j}"
)
# ── Sensitivity ─────────────────────────────────────────────────────
def test_sensitivity_all_knobs_monotone_non_worsening():
"""Doubling any HW knob should not slow the sim down (rel_speedup ≥ 0
within floating-point tolerance)."""
model = PRESETS["Llama 3.1 70B"].model
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
assert len(res.sensitivity) == 6
for row in res.sensitivity:
# Allow a tiny slop for floating-point but disallow real regressions.
assert row.rel_speedup >= -1e-9, (
f"knob {row.knob}: doubling slowed latency from "
f"{row.baseline_latency_ns} to {row.doubled_latency_ns}"
)
def test_joint_explore_attention_only_faster_than_full():
"""include_ffn=False produces smaller best-latency than include_ffn=True
for the same HW+model."""
model = PRESETS["Llama 3.1 70B"].model
r_full = joint_explore(model, s_kv=131072, mode="decode",
depth="two_stage", include_ffn=True)
r_attn = joint_explore(model, s_kv=131072, mode="decode",
depth="two_stage", include_ffn=False)
b_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
b_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
assert b_attn.total_latency_ns < b_full.total_latency_ns, (
f"attn-only ({b_attn.latency_ms:.2f}ms) should be less than "
f"full ({b_full.latency_ms:.2f}ms)"
)
def test_sensitivity_hbm_bw_dominant_for_llama_decode():
"""For Llama 70B decode (memory-bound), HBM BW should be the top
sensitivity knob — doubling it gives more speedup than any other knob."""
model = PRESETS["Llama 3.1 70B"].model
res = joint_explore(model, s_kv=131072, mode="decode", depth="balanced")
assert res.sensitivity[0].knob == "bw_hbm_gbs", (
f"expected bw_hbm_gbs to top the sensitivity ranking for Llama 70B "
f"decode; got {res.sensitivity[0].knob}"
)