diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 86e379a..0e54dea 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -252,7 +252,19 @@ with st.sidebar.expander("Parallelism", expanded=True): st.session_state["tp"] = _snap(_sb_sug.tp, _TP_OPTS) st.session_state["pp"] = _snap(_sb_sug.pp, _PP_OPTS) st.session_state["dp"] = 1 + st.session_state["ep"] = 1 st.session_state["cp_placement"] = _sb_sug.cp_placement + # auto_suggest only explores (cp, tp, pp, cp_placement); reset + # every other TopologyConfig field to its dataclass default so + # the applied topology matches what memory-min actually scored + # (else stale tp_placement="cube" etc. explodes cubes_used). + st.session_state["tp_placement"] = "pe" + st.session_state["kv_mode"] = "split" + st.session_state["cp_ring_variant"] = "kv" + # ffn_scope_label is a dynamic string (embeds current tp/cp) — + # dropping the key lets the radio fall back to index=1 which is + # "TP+CP", the TopologyConfig default. + st.session_state.pop("ffn_scope_label", None) st.session_state["_pl_active_scope"] = _sb_scope st.rerun() diff --git a/tests/analytical_visualization/test_auto_explore.py b/tests/analytical_visualization/test_auto_explore.py index 06cf1cf..dd164c9 100644 --- a/tests/analytical_visualization/test_auto_explore.py +++ b/tests/analytical_visualization/test_auto_explore.py @@ -19,8 +19,9 @@ from tests.analytical_visualization.auto_explore import ( run_auto_explore, score_config, ) +from tests.analytical_visualization.autosuggest import auto_suggest from tests.analytical_visualization.model_config import ( - FullConfig, MachineParams, + FullConfig, MachineParams, TopologyConfig, ) from tests.analytical_visualization.model_presets import PRESETS @@ -208,6 +209,39 @@ def test_ffn_only_latency_less_than_full_and_different_from_attn(): ) +def test_auto_suggest_cubes_match_default_topology(): + """A TopologyConfig built from auto_suggest's returned (cp, tp, pp, + cp_placement) with all other fields left at their dataclass defaults + must produce the same cubes_used the Suggestion reports. Guards the + sidebar apply: any session_state key auto_suggest doesn't explore + (tp_placement, kv_shard_mode, cp_ring_variant, ep, ffn_shard_scope) + must be reset to the default before rerun, else the applied topology + won't match what the memory-min search actually scored. + """ + machine = MachineParams() + matrix = [ + ("Llama 3 70B", 8192), + ("Llama 3 70B", 32768), + ("Llama 3 8B", 131072), + ("Qwen 3 32B", 32768), + ] + for name, skv in matrix: + m = PRESETS[name].model + for ia, ff in [(True, False), (False, True), (True, True)]: + sug = auto_suggest(m, machine, s_kv=skv, mode="decode", + include_attention=ia, include_ffn=ff) + topo = TopologyConfig( + cp=sug.cp, tp=sug.tp, pp=sug.pp, + s_kv=skv, mode="decode", b=1, + cp_placement=sug.cp_placement, + ) + assert topo.cubes_used == sug.cubes_used, ( + f"{name} skv={skv} scope=(ia={ia},ff={ff}): " + f"topo.cubes_used={topo.cubes_used} != " + f"sug.cubes_used={sug.cubes_used}" + ) + + def test_parallelism_sensitivity_includes_baseline_value(): """Each row's values contain the baseline_value.""" from tests.analytical_visualization.auto_explore import (