From b8e1a3322f0cc513a088467b607336ab7889f36e Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 13:58:29 -0700 Subject: [PATCH] analytical-viz: attention-only vs attn+FFN via two sweep buttons MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- tests/analytical_visualization/app.py | 232 +++++++++++++----- .../analytical_visualization/auto_explore.py | 35 ++- .../analytical_visualization/auto_hardware.py | 29 ++- .../test_auto_explore.py | 27 ++ .../test_auto_hardware.py | 16 ++ 5 files changed, 264 insertions(+), 75 deletions(-) diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index dfb3c80..57e4a71 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -419,9 +419,14 @@ if _warnings: # ── Tabs ───────────────────────────────────────────────────────── -tab_layout, tab_memory, tab_stages, tab_compare, tab_auto, tab_hw = st.tabs([ +# One tab per auto feature; each has TWO sweep buttons (Attention only / +# Attn + FFN/MoE) so the user can toggle scope without switching tabs. +# Auto Suggest is renamed "Parallelism" since it only varies parallelism +# knobs (hardware is held fixed at the sidebar values). +tab_auto, tab_layout, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([ + "Auto Suggest Parallelism", "Physical layout", "Memory breakdown", "Per-stage latency", - "Save & compare", "Auto Explore", "Auto Hardware", + "Save & compare", "Auto Hardware", ]) @@ -1284,8 +1289,17 @@ with tab_compare: ) -# ── TAB 5: Auto Explore ───────────────────────────────────────── -with tab_auto: +# ── 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, @@ -1298,34 +1312,80 @@ with tab_auto: "(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." + ) - ax_l, ax_r = st.columns([3, 1]) - with ax_l: - st.markdown( - f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}** | " - f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**" + 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 = st.columns([1, 1, 2]) + with _b1: + _run_attn = st.button("Run sweep — Attention", + type="primary", width='stretch', + key="_auto_run_attn") + with _b2: + _run_full = st.button("Run sweep — Attn + FFN/MoE", + type="primary", width='stretch', + key="_auto_run_full") + + if _run_attn: + with st.spinner("Sweeping ~28k configs (Attention only)..."): + _r = run_auto_explore( + model, machine, s_kv=s_kv, mode=mode, include_ffn=False, + ) + st.session_state["_auto_explore_result_attn"] = _r + st.session_state["_auto_explore_ctx_attn"] = ( + model.name, s_kv, mode, machine.pe_hbm_gb, ) - with ax_r: - _run_sweep = st.button("Run sweep", type="primary", width='stretch') - - if _run_sweep or st.session_state.get("_auto_explore_result") is not None: - if _run_sweep: - with st.spinner("Sweeping ~28k configs..."): - _res = run_auto_explore(model, machine, s_kv=s_kv, mode=mode) - st.session_state["_auto_explore_result"] = _res - st.session_state["_auto_explore_ctx"] = ( - model.name, s_kv, mode, machine.pe_hbm_gb, + st.session_state["_auto_explore_active"] = "attn" + if _run_full: + with st.spinner("Sweeping ~28k configs (Attn + FFN/MoE)..."): + _r = run_auto_explore( + model, machine, s_kv=s_kv, mode=mode, include_ffn=True, ) - _res = st.session_state["_auto_explore_result"] + st.session_state["_auto_explore_result_full"] = _r + st.session_state["_auto_explore_ctx_full"] = ( + model.name, s_kv, mode, machine.pe_hbm_gb, + ) + st.session_state["_auto_explore_active"] = "full" - # Warn if the cached result is stale relative to the current config. - _cached_ctx = st.session_state.get("_auto_explore_ctx") - _cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb) - if _cached_ctx != _cur_ctx: - st.warning( - "Sweep result is from a different model/workload. " - "Click **Run sweep** to refresh." - ) + _active = st.session_state.get("_auto_explore_active") + if _active is None: + st.info( + "Click one of the run buttons above to enumerate valid " + "parallelism configurations and rank them on the Pareto " + "frontier. Takes ~5–10 s. Both button results are cached so " + "you can flip between the two scopes without re-running." + ) + return + + # Downstream code expects _suffix, _label, include_ffn, _res, _ctx_key. + include_ffn = (_active == "full") + _suffix = "_full" if include_ffn else "_attn" + _label = "Attn + FFN/MoE" if include_ffn else "Attention only" + _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:,}") @@ -1431,12 +1491,13 @@ with tab_auto: _sens_row = st.number_input( "Baseline row #", min_value=0, max_value=len(_pareto_by_lat) - 1, - value=0, step=1, key="_auto_sens_row", + 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_ffn=include_ffn, ) import matplotlib.pyplot as _plt @@ -1487,10 +1548,11 @@ with tab_auto: with _lc1: _pick = st.number_input( "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1, - value=0, step=1, key="_auto_pick_row", + value=0, step=1, key=f"_auto_pick_row{_suffix}", ) with _lc2: - if st.button("Load into sidebar", type="secondary"): + if st.button("Load into sidebar", type="secondary", + key=f"_auto_load_btn{_suffix}"): _s = _pareto_by_lat[int(_pick)] # Set the sidebar's session_state keys directly. The # sidebar reads these on the next rerun. @@ -1527,8 +1589,17 @@ with tab_auto: ) -# ── TAB 6: Auto Hardware ───────────────────────────────────────── -with tab_hw: +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, ) @@ -1538,11 +1609,16 @@ with tab_hw: "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?" + "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, hx_r = st.columns([2, 2, 1]) + hx_l, hx_m = st.columns([1, 1]) with hx_l: st.markdown( f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}**" @@ -1557,28 +1633,64 @@ with tab_hw: "balanced (default): 64 HW × ~2k parallelism, ~10-20s.\n" "coarse: 729 HW × ~2k parallelism, ~2-5 min."), ) - with hx_r: - _run_hw = st.button("Run joint sweep", type="primary", - width='stretch', key="_hw_run_btn") - if _run_hw or st.session_state.get("_hw_result") is not None: - if _run_hw: - with st.spinner(f"Sweeping HW × parallelism ({_depth})..."): - _hw_res = joint_explore(model, s_kv=s_kv, - mode=mode, depth=_depth) - st.session_state["_hw_result"] = _hw_res - st.session_state["_hw_ctx"] = ( - model.name, s_kv, mode, _depth, - ) - _hw_res = st.session_state["_hw_result"] + _bh1, _bh2, _bh3 = st.columns([1, 1, 2]) + with _bh1: + _run_hw_attn = st.button("Run joint sweep — Attention", + type="primary", width='stretch', + key="_hw_run_attn") + with _bh2: + _run_hw_full = st.button("Run joint sweep — Attn + FFN/MoE", + type="primary", width='stretch', + key="_hw_run_full") - _cached_ctx = st.session_state.get("_hw_ctx") - _cur_ctx = (model.name, s_kv, mode, _depth) - if _cached_ctx != _cur_ctx: - st.warning( - "Sweep result is from a different model/workload/depth. " - "Click **Run joint sweep** to refresh." - ) + if _run_hw_attn: + with st.spinner(f"Sweeping HW × parallelism ({_depth}, Attention only)..."): + _r = joint_explore(model, s_kv=s_kv, mode=mode, + depth=_depth, include_ffn=False) + st.session_state["_hw_result_attn"] = _r + st.session_state["_hw_ctx_attn"] = ( + model.name, s_kv, mode, _depth, + ) + st.session_state["_hw_active"] = "attn" + if _run_hw_full: + with st.spinner(f"Sweeping HW × parallelism ({_depth}, Attn + FFN/MoE)..."): + _r = joint_explore(model, s_kv=s_kv, mode=mode, + depth=_depth, include_ffn=True) + st.session_state["_hw_result_full"] = _r + st.session_state["_hw_ctx_full"] = ( + model.name, s_kv, mode, _depth, + ) + st.session_state["_hw_active"] = "full" + + _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. Both button results are " + "cached so you can flip between the two scopes without re-running." + ) + return + + include_ffn = (_active == "full") + _suffix = "_full" if include_ffn else "_attn" + _label = "Attn + FFN/MoE" if include_ffn else "Attention only" + _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) @@ -1691,11 +1803,11 @@ with tab_hw: _pick_hw = st.number_input( "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1, - value=0, step=1, key="_hw_pick_row", + value=0, step=1, key=f"_hw_pick_row{_suffix}", ) with _lch2: if st.button("Load into sidebar", type="secondary", - key="_hw_load_btn"): + key=f"_hw_load_btn{_suffix}"): _s = _pareto_by_lat[int(_pick_hw)] _h = _s.hardware _p = _s.parallelism @@ -1748,3 +1860,7 @@ with tab_hw: "Click **Run joint sweep** to explore hardware and parallelism " "co-design space for this model + workload." ) + + +with tab_hw: + _render_auto_hardware_tab() diff --git a/tests/analytical_visualization/auto_explore.py b/tests/analytical_visualization/auto_explore.py index 0099b91..4b60f4f 100644 --- a/tests/analytical_visualization/auto_explore.py +++ b/tests/analytical_visualization/auto_explore.py @@ -164,7 +164,7 @@ def enumerate_configs( # ── Scoring ────────────────────────────────────────────────────────── -def _sum_visible_latency(cfg: FullConfig) -> float: +def _sum_visible_latency(cfg: FullConfig, include_ffn: bool = True) -> float: """Total single-request latency (seconds) across all model layers. A single request traverses every layer sequentially, whether the layers @@ -176,23 +176,30 @@ def _sum_visible_latency(cfg: FullConfig) -> float: 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. + + ``include_ffn=False`` restricts to attention stages only — useful for + isolating attention-kernel tuning from FFN cost. """ attn = sum(s.visible_s for s in all_stages(cfg)) - ffn = sum(s.visible_s for s in all_ffn_stages(cfg)) + 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) -> float: +def _efficiency(cfg: FullConfig, latency_s: float, + 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) + + ``include_ffn=False`` mirrors the same restriction as + :func:`_sum_visible_latency` — attention stages only. """ if latency_s <= 0: return 0.0 attn = all_stages(cfg) - ffn = all_ffn_stages(cfg) + 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) @@ -209,19 +216,24 @@ def _efficiency(cfg: FullConfig, latency_s: float) -> float: return math.sqrt(compute_util * bw_util) -def score_config(cfg: FullConfig) -> ConfigScore: +def score_config(cfg: FullConfig, 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. + + ``include_ffn=False`` restricts latency + efficiency to attention stages + only. Memory feasibility is unchanged — still checks weights+KV+transient + fit in per-PE HBM since the model still exists physically. """ mem = compute_memory(cfg) placement_ok = cfg.topo.placement_valid - latency_s = _sum_visible_latency(cfg) + latency_s = _sum_visible_latency(cfg, include_ffn=include_ffn) throughput = 1.0 / latency_s if latency_s > 0 else 0.0 - efficiency = _efficiency(cfg, latency_s) if latency_s > 0 else 0.0 + efficiency = (_efficiency(cfg, latency_s, include_ffn=include_ffn) + if latency_s > 0 else 0.0) fits = not mem.over_budget reason = "" @@ -333,6 +345,7 @@ def compute_parallelism_sensitivity( machine: MachineParams, s_kv: int, mode: str, + 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 @@ -364,7 +377,7 @@ def compute_parallelism_sensitivity( # 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) + score = score_config(cfg, include_ffn=include_ffn) latencies.append(score.total_latency_ns) fits.append(score.fits_memory and score.placement_valid) rows.append(ParallelismSensitivityRow( @@ -386,19 +399,23 @@ def run_auto_explore( machine: MachineParams, s_kv: int, mode: str = "decode", + 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. + + ``include_ffn=False`` restricts to attention stages only — useful for + isolating attention-kernel tuning independent of FFN cost. """ 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)) + all_scores.append(score_config(cfg, include_ffn=include_ffn)) all_scores.sort(key=lambda s: s.total_latency_ns) pareto = pareto_frontier(all_scores) diff --git a/tests/analytical_visualization/auto_hardware.py b/tests/analytical_visualization/auto_hardware.py index f51d730..d57aa87 100644 --- a/tests/analytical_visualization/auto_hardware.py +++ b/tests/analytical_visualization/auto_hardware.py @@ -235,16 +235,18 @@ def _iter_reduced_parallelism( def _best_parallelism_for_hw( model: ModelConfig, machine: MachineParams, - s_kv: int, mode: str, + s_kv: int, mode: str, 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. + ``include_ffn`` is forwarded to score_config so the "latency" ranked on + matches the caller's attention-only vs full-model 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) + s = score_config(cfg, 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: @@ -254,7 +256,7 @@ def _best_parallelism_for_hw( def _best_parallelism_two_stage( model: ModelConfig, machine: MachineParams, - s_kv: int, mode: str, + s_kv: int, mode: str, 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) @@ -270,7 +272,7 @@ def _best_parallelism_two_stage( 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) + s = score_config(cfg, include_ffn=include_ffn) return s if s.fits_memory and s.placement_valid else None @@ -308,6 +310,7 @@ def compute_sensitivity( baseline_hw: HardwareCandidate, parallelism: ConfigScore, model: ModelConfig, s_kv: int, mode: str, + 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 @@ -318,7 +321,7 @@ def compute_sensitivity( baseline_cfg = FullConfig( model=model, topo=baseline_topo, machine=baseline_machine, ) - baseline_score = score_config(baseline_cfg) + baseline_score = score_config(baseline_cfg, include_ffn=include_ffn) baseline_latency = baseline_score.total_latency_ns for knob in _SENSITIVITY_KNOBS: @@ -327,7 +330,7 @@ def compute_sensitivity( doubled_cfg = FullConfig( model=model, topo=baseline_topo, machine=doubled_hw.as_machine(), ) - doubled_score = score_config(doubled_cfg) + doubled_score = score_config(doubled_cfg, include_ffn=include_ffn) rows.append(SensitivityRow( knob=knob, baseline_value=getattr(baseline_hw, knob), @@ -348,19 +351,28 @@ def joint_explore( s_kv: int, mode: str, depth: str = "balanced", + 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. + + ``include_ffn=False`` restricts scoring to attention stages only — + 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) + par = _best_parallelism_two_stage( + model, machine, s_kv, mode, include_ffn=include_ffn, + ) else: - par = _best_parallelism_for_hw(model, machine, s_kv, mode) + par = _best_parallelism_for_hw( + model, machine, s_kv, mode, include_ffn=include_ffn, + ) if par is None: continue all_scores.append(JointScore( @@ -379,6 +391,7 @@ def joint_explore( best = all_scores[0] sensitivity = compute_sensitivity( best.hardware, best.parallelism, model, s_kv, mode, + include_ffn=include_ffn, ) return JointExploreResult( diff --git a/tests/analytical_visualization/test_auto_explore.py b/tests/analytical_visualization/test_auto_explore.py index 01ab681..28c70d6 100644 --- a/tests/analytical_visualization/test_auto_explore.py +++ b/tests/analytical_visualization/test_auto_explore.py @@ -161,6 +161,33 @@ def test_parallelism_sensitivity_has_all_five_knobs(): 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_ffn=False) + assert res.total_feasible > 0 + assert len(res.pareto_scores) > 0 + + def test_parallelism_sensitivity_includes_baseline_value(): """Each row's values contain the baseline_value.""" from tests.analytical_visualization.auto_explore import ( diff --git a/tests/analytical_visualization/test_auto_hardware.py b/tests/analytical_visualization/test_auto_hardware.py index 02b128d..399c35f 100644 --- a/tests/analytical_visualization/test_auto_hardware.py +++ b/tests/analytical_visualization/test_auto_hardware.py @@ -105,6 +105,22 @@ def test_sensitivity_all_knobs_monotone_non_worsening(): ) +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."""