From 8fa758797e7089e6a2b8fc50f4d4020c19969af2 Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Wed, 24 Jun 2026 15:21:22 +0800
Subject: [PATCH] Guard generic topology search from introducing EP

---
 src/aituner/harness.py  |   5 ++
 tests/test_core_flow.py | 103 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+)

diff --git a/src/aituner/harness.py b/src/aituner/harness.py
index 5392a9d..a8f50bf 100644
--- a/src/aituner/harness.py
+++ b/src/aituner/harness.py
@@ -1482,6 +1482,11 @@ def _legal_topology_points(
         and current_ep > 1
     ):
         ep_values = [current_ep]
+    elif not current_enable_ep and current_ep <= 1:
+        # Generic topology search should not introduce EP. EP is a separate
+        # intervention family and needs direct EP-specific evidence; otherwise
+        # TTFT/TPOT topology probes can waste trials on launch-risky EP variants.
+        ep_values = [1]
     elif constraints is not None and constraints.allowed_expert_parallel_sizes:
         ep_values = sorted(set(constraints.allowed_expert_parallel_sizes))
     elif "expert-parallel-size" in tunable:
diff --git a/tests/test_core_flow.py b/tests/test_core_flow.py
index 1c22a3d..55b4c2b 100644
--- a/tests/test_core_flow.py
+++ b/tests/test_core_flow.py
@@ -1979,6 +1979,109 @@ class CoreFlowTests(unittest.TestCase):
             self.assertIn("ttft_prefill", context["bottleneck_hypotheses"][0]["name"])
             self.assertFalse(context["harness_stop"]["should_stop"])
 
+    def test_profile_driven_topology_does_not_introduce_ep_for_ttft(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            study_path = _write_study_assets(
+                tmp_path,
+                engine_overrides={
+                    "base_flags": {"host": "127.0.0.1", "port": 8000},
+                    "tunable_flags": [
+                        "tensor-parallel-size",
+                        "data-parallel-size",
+                        "expert-parallel-size",
+                        "enable-expert-parallel",
+                    ],
+                    "topology_constraints": {
+                        "allowed_tensor_parallel_sizes": [1, 2, 4],
+                        "allowed_data_parallel_sizes": [1],
+                        "allowed_expert_parallel_sizes": [1, 2],
+                        "allowed_tp_dp_products": [1, 2, 4],
+                        "require_ep_size_leq_tp_dp_product": True,
+                        "require_ep_size_divides_tp_dp_product": True,
+                        "require_enable_expert_parallel_when_ep_gt_one": True,
+                    },
+                },
+            )
+            result_paths: list[Path] = []
+            for idx in range(1, 4):
+                result_path = tmp_path / f"trial-000{idx}.json"
+                result_path.write_text(
+                    json.dumps(
+                        {
+                            "status": "completed",
+                            "best_sampling_u": 0.25,
+                            "best_request_rate": 2.0,
+                            "best_pass_rate": 1.0,
+                            "probes": [
+                                {
+                                    "threshold": 0.5,
+                                    "feasible": False,
+                                    "payload": {
+                                        "request_count": 100,
+                                        "pass_rate": 0.6,
+                                        "request_rate": 4.0,
+                                        "early_stop_reason": "slo_pass_rate_unrecoverable",
+                                        "latency_summary": {
+                                            "failed_reason_counts": {"ttft_ms>2000": 40}
+                                        },
+                                    },
+                                }
+                            ],
+                        }
+                    ),
+                    encoding="utf-8",
+                )
+                result_paths.append(result_path)
+            study = load_study_spec(study_path)
+            context = build_harness_context(
+                study=study,
+                window_summary={"prompt_tokens_p95": 8192},
+                state=StudyState(
+                    study_id=study.study_id,
+                    best_trial_id="trial-0002",
+                    best_request_rate=4.0,
+                    best_request_rate_per_gpu=2.0,
+                    trials=[
+                        TrialSummary(
+                            trial_id="trial-0001",
+                            status="completed",
+                            best_request_rate=2.0,
+                            best_request_rate_per_gpu=2.0,
+                            result_path=str(result_paths[0]),
+                            config_patch={"env_patch": {}, "flag_patch": {}},
+                        ),
+                        TrialSummary(
+                            trial_id="trial-0002",
+                            status="completed",
+                            best_request_rate=4.0,
+                            best_request_rate_per_gpu=2.0,
+                            result_path=str(result_paths[1]),
+                            config_patch={
+                                "env_patch": {},
+                                "flag_patch": {"tensor-parallel-size": 2},
+                            },
+                        ),
+                        TrialSummary(
+                            trial_id="trial-0003",
+                            status="completed",
+                            best_request_rate=4.0,
+                            best_request_rate_per_gpu=1.0,
+                            result_path=str(result_paths[2]),
+                            config_patch={
+                                "env_patch": {},
+                                "flag_patch": {"tensor-parallel-size": 4},
+                            },
+                        ),
+                    ],
+                ),
+            )
+            candidate_actions = context["experiment_plan"]["candidate_actions"]
+            for action in candidate_actions:
+                patch = action["config_patch"]["flag_patch"]
+                self.assertNotIn("enable-expert-parallel", patch)
+                self.assertNotIn("expert-parallel-size", patch)
+
     def test_profile_driven_planner_prefers_decode_concurrency_relief(self) -> None:
         with tempfile.TemporaryDirectory() as tmp:
             tmp_path = Path(tmp)