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
+749 -8
View File
@@ -6,10 +6,30 @@ Run:
"""
from __future__ import annotations
import importlib
import sys
import matplotlib.pyplot as plt
import pandas as pd
import streamlit as st
# Streamlit's hot-reload only re-runs THIS script; imported sub-modules stay
# cached in sys.modules across reruns. When we edit auto_explore.py or
# auto_hardware.py while the Streamlit process is alive, the app would
# keep using the pre-edit versions until a full Ctrl+C + restart. Force
# a reload of our own modules on every rerun so signature changes land
# without a full restart.
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 (
FullConfig, ModelConfig, TopologyConfig, MachineParams,
)
@@ -156,13 +176,67 @@ if st.session_state.get("_last_preset") != preset_name:
with st.sidebar.expander("Parallelism", expanded=True):
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}"
)
if st.button("Apply auto-suggest", width='stretch'):
st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
# Three latency-optimal apply buttons — pick a scope, load the
# latency-min Pareto config into the sidebar sliders.
from tests.analytical_visualization.auto_explore import (
run_auto_explore as _run_auto_explore_sb,
)
_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
@@ -419,14 +493,110 @@ if _warnings:
# ── Tabs ─────────────────────────────────────────────────────────
tab_layout, tab_memory, tab_stages, tab_compare = st.tabs([
"Physical layout", "Memory breakdown", "Per-stage latency",
"Save & compare",
# One tab per auto feature; each has TWO sweep buttons (Attention only /
# Attn + FFN/MoE) so the user can toggle scope without switching tabs.
# Auto Suggest is renamed "Parallelism" since it only varies parallelism
# knobs (hardware is held fixed at the sidebar values).
tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
"Physical layout",
"Auto Suggest Parallelism",
"Memory breakdown", "Per-stage latency",
"Save & compare", "Auto Hardware",
])
# ── TAB 1: Physical layout ──────────────────────────────────────
with tab_layout:
# Quick shortcut: three buttons that pick the 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).
st.markdown("**Per-layer pipeline (color = dominant bound)**")
fig_pipe = draw_pipeline(cfg)
@@ -584,7 +754,9 @@ with tab_layout:
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?)",
@@ -1282,3 +1454,572 @@ with tab_compare:
"'-' 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,
)
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}"
)