fix(forge-gemm): MoE runtime key, E2E faults, warm-replay nogit, model path

src/hyperloom/agents/kernel/tools/forge_gemm_tuning.py

@@ -63,6 +63,7 @@ def _build_cmd(args: dict[str, Any]) -> list[str]:
     _add_opt(cmd, args, "global_timeout", "--global-timeout")
     _add_opt(cmd, args, "tuner", "--tuner")
     _add_opt(cmd, args, "untuned_csv", "--untuned-csv")
+    _add_opt(cmd, args, "moe_untuned_csv", "--moe-untuned-csv")
     _add_opt(cmd, args, "shapes_json", "--shapes-json")
     _add_opt(cmd, args, "tunableop_input", "--tunableop-input")
     _add_opt(cmd, args, "kernel_signature_log", "--kernel-signature-log")

src/hyperloom/common/model_paths.py

@@ -19,7 +19,9 @@
 
 from __future__ import annotations
 
+import os
 from pathlib import Path
+from typing import Any
 
 
 def _identity_leaf(seg: str) -> str:
@@ -128,3 +130,60 @@
     if isinstance(hit, str) and Path(hit).is_file():
         return Path(hit).parent
     return None
+
+
+def resolve_serving_model_path(raw: str) -> str:
+    """Resolve a session model identity to a path suitable for launching servers.
+
+    Precedence mirrors ``run_hyperloom.sbatch``: an existing directory wins,
+    then ``HL_MODEL_BASE/<repo-tail>``, then the HuggingFace hub cache via
+    :func:`resolve_local_model_dir`. When nothing resolves, the original
+    string is returned unchanged.
+    """
+    text = str(raw or "").strip()
+    if not text:
+        return ""
+    try:
+        direct = Path(text).expanduser()
+        if direct.is_dir():
+            return str(direct)
+    except OSError:
+        pass
+    base = os.environ.get("HL_MODEL_BASE", "").strip()
+    if base:
+        leaf = text.rstrip("/").split("/")[-1]
+        if leaf:
+            candidate = Path(base) / leaf
+            try:
+                if candidate.is_dir():
+                    return str(candidate)
+            except OSError:
+                pass
+    resolved = resolve_local_model_dir(text)
+    if resolved is not None:
+        return str(resolved)
+    return text
+
+
+def resolve_session_model_path(
+    *,
+    params: dict[str, Any] | None = None,
+    state_model_path: str = "",
+    for_serving: bool = False,
+) -> str:
+    """Unified session model-path precedence for executors and handlers.
+
+    Order: ``params['model_path']`` → ``$MODEL_PATH`` → ``state.model_path``.
+    When ``for_serving`` is true, :func:`resolve_serving_model_path` is applied
+    to the chosen raw value.
+    """
+    raw = (
+        str((params or {}).get("model_path") or "").strip()
+        or os.environ.get("MODEL_PATH", "").strip()
+        or str(state_model_path or "").strip()
+    )
+    if not raw:
+        return ""
+    if for_serving:
+        return resolve_serving_model_path(raw)
+    return raw

src/hyperloom/inference_optimizer/breakdown/collectors/kernels.py

@@ -1297,10 +1297,15 @@ def collect_gemm_tuning(state: dict[str, Any]) -> dict[str, Any]:
         if not isinstance(raw, dict):
             continue
         engine = _resolve_gemm_engine(raw)
+        e2e_gain_pct = _to_float(raw.get("e2e_gain_pct"))
         speedup = _to_float(raw.get("best_speedup"))
         gain_pct: float | None = None
         tuned_tput: float | None = None
-        if speedup is not None:
+        if e2e_gain_pct is not None:
+            gain_pct = e2e_gain_pct
+            if baseline_tput is not None:
+                tuned_tput = baseline_tput * (1.0 + e2e_gain_pct / 100.0)
+        elif speedup is not None:
             gain_pct = (speedup - 1.0) * 100.0
             if baseline_tput is not None:
                 tuned_tput = baseline_tput * speedup
@@ -1331,6 +1336,7 @@ def collect_gemm_tuning(state: dict[str, Any]) -> dict[str, Any]:
             "gpu_type": str(raw.get("gpu_type") or gpu_type),
             "baseline_tput": baseline_tput,
             "best_speedup": speedup,
+            "e2e_gain_pct": e2e_gain_pct,
             "gain_pct": gain_pct,
             "tuned_tput": tuned_tput,
             "tuned_file": tuned_file,

src/hyperloom/inference_optimizer/cli/__init__.py

@@ -1799,7 +1799,9 @@ async def _run_optimize(args: argparse.Namespace) -> int:
             )
             sys.exit(2)
         # Re-export so subprocess executors inject the resolved model into the Magpie YAML, not its hardcoded model.
-        os.environ["MODEL_PATH"] = str(args.model)
+        from hyperloom.common.model_paths import resolve_serving_model_path
+
+        os.environ["MODEL_PATH"] = resolve_serving_model_path(str(args.model)) or str(args.model)
 
         # Quantization prelude (one-shot, before any session/baseline work):
         # if --quantize was passed, quantize the source model now and rewrite

src/hyperloom/inference_optimizer/tests/test_coordinator_gemm_promote_units.py

@@ -2012,6 +2012,112 @@ async def test_does_not_e2e_validate_missing_aiter_candidate(
             "candidate_artifact_missing"
         )
 
+    @pytest.mark.asyncio
+    async def test_integrate_bench_fault_not_recorded_as_zero_gain_revert(
+        self, tmp_path, monkeypatch
+    ):
+        """A server that never booted is an integrate fault, not a 0% REVERT."""
+        coord = _coord(tmp_path, baseline_tput=100.0, framework="sglang")
+        phase = KernelPhase(coord)
+        fmoe_candidate = tmp_path / "fmoe.csv"
+        dense_candidate = tmp_path / "dense.csv"
+        fmoe_candidate.write_text("token,model_dim\n1,2\n", encoding="utf-8")
+        dense_candidate.write_text("M,N,K\n1,2,3\n", encoding="utf-8")
+        calls: list[dict] = []
+
+        async def _fake_integrate(payload, *, session_dir):
+            calls.append(payload)
+            if payload["kernel_id"] == "gemm_tune_fmoe_ck":
+                return {
+                    "status": "failed",
+                    "error_class": "bench_exception",
+                    "decision": "REVERT",
+                    "error": "re-baseline did not succeed",
+                }
+            return {"status": "ok", "decision": "KEEP", "new_tput": 120.0, "gain_pct": 9.09}
+
+        monkeypatch.setattr(krh_mod, "integrate_handler", _fake_integrate)
+        monkeypatch.setattr(explore_mod, "_compute_explore_variant_timeout", lambda **_k: 61)
+        monkeypatch.setattr(
+            phase,
+            "_merge_gemm_candidate_with_runtime",
+            lambda _env_var, env_value: env_value,
+        )
+
+        result = {
+            "backend": "forge",
+            "tuners_run": [
+                {
+                    "status": "ok",
+                    "tuner": "fmoe_ck",
+                    "improved_shapes": 2,
+                    "env_var": "AITER_CONFIG_FMOE",
+                    "env_value": str(fmoe_candidate),
+                },
+                {
+                    "status": "ok",
+                    "tuner": "dense_bf16",
+                    "improved_shapes": 1,
+                    "env_var": "AITER_CONFIG_DENSE",
+                    "env_value": str(dense_candidate),
+                },
+            ],
+        }
+
+        await phase._validate_gemm_tuning_e2e(result)
+
+        assert len(calls) == 3
+        assert result["e2e_results"]["faults"][0]["reason"] == "integrate_fault:bench_exception"
+        assert result["e2e_results"]["faults"][0]["fault_attempts"] == 2
+        assert result["e2e_results"]["reverted"] == []
+        assert result["e2e_results"]["kept"][0]["tuner"] == "dense_bf16"
+        assert result["decision"] == "KEEP"
+
+    @pytest.mark.asyncio
+    async def test_integrate_fault_retries_once_before_verdict(self, tmp_path, monkeypatch):
+        coord = _coord(tmp_path, baseline_tput=100.0, framework="sglang")
+        phase = KernelPhase(coord)
+        dense_candidate = tmp_path / "dense.csv"
+        dense_candidate.write_text("M,N,K\n1,2,3\n", encoding="utf-8")
+        calls: list[dict] = []
+
+        async def _fake_integrate(payload, *, session_dir):
+            calls.append(payload)
+            if len(calls) == 1:
+                return {
+                    "status": "failed",
+                    "error_class": "bench_exception",
+                    "decision": "REVERT",
+                    "error": "re-baseline did not succeed",
+                }
+            return {"status": "ok", "decision": "KEEP", "new_tput": 110.0, "gain_pct": 10.0}
+
+        monkeypatch.setattr(krh_mod, "integrate_handler", _fake_integrate)
+        monkeypatch.setattr(
+            phase,
+            "_merge_gemm_candidate_with_runtime",
+            lambda _env_var, env_value: env_value,
+        )
+        result = {
+            "backend": "forge",
+            "tuners_run": [
+                {
+                    "status": "ok",
+                    "tuner": "dense_bf16",
+                    "improved_shapes": 1,
+                    "env_var": "AITER_CONFIG_DENSE",
+                    "env_value": str(dense_candidate),
+                },
+            ],
+        }
+
+        await phase._validate_gemm_tuning_e2e(result)
+
+        assert len(calls) == 2
+        assert result["e2e_results"]["faults"] == []
+        assert result["e2e_results"]["kept"][0]["tuner"] == "dense_bf16"
+        assert result["decision"] == "KEEP"
+
     @pytest.mark.asyncio
     async def test_a_stopped_run_leaves_its_tuners_unjudged(self, tmp_path, monkeypatch):
         """A clock that ran out is not a verdict on the tuners it interrupted."""
@@ -2182,7 +2288,7 @@ async def test_handles_no_candidates_without_rewriting_raw_result(self, tmp_path
         assert coord.shared_state.optimization_stack == []
 
     @pytest.mark.asyncio
-    async def test_records_integrate_exception_as_revert(self, tmp_path, monkeypatch):
+    async def test_records_integrate_exception_as_fault(self, tmp_path, monkeypatch):
         coord = _coord(tmp_path, baseline_tput=100.0, framework="sglang")
         phase = KernelPhase(coord)
         dense_candidate = tmp_path / "dense.csv"
@@ -2191,7 +2297,13 @@ async def test_records_integrate_exception_as_revert(self, tmp_path, monkeypatch
         async def _raise_integrate(*_args, **_kwargs):
             raise RuntimeError("integrate failed")
 
-        monkeypatch.setattr(krh_mod, "integrate_handler", _raise_integrate)
+        calls: list[str] = []
+
+        async def _counting_raise(*_args, **_kwargs):
+            calls.append("boom")
+            raise RuntimeError("integrate failed")
+
+        monkeypatch.setattr(krh_mod, "integrate_handler", _counting_raise)
         monkeypatch.setattr(
             phase,
             "_merge_gemm_candidate_with_runtime",
@@ -2213,9 +2325,15 @@ async def _raise_integrate(*_args, **_kwargs):
 
         await phase._validate_gemm_tuning_e2e(result)
 
-        assert result["decision"] == "REVERT"
-        assert result["micro_decision"] == "candidate_no_e2e_gain"
-        assert "integrate failed" in result["e2e_results"]["reverted"][0]["reason"]
+        assert result["status"] == "failed"
+        assert result["micro_decision"] == "integrate_fault"
+        assert result["e2e_gain_pct"] is None
+        fault = result["e2e_results"]["faults"][0]
+        assert fault["reason"] == "integrate_fault:handler_exception"
+        assert fault["fault"] is True
+        assert fault["fault_attempts"] == 2
+        assert len(calls) == 2
+        assert result["e2e_results"]["reverted"] == []
 
 
 class TestBf16DenseFallback:
@@ -3103,7 +3221,7 @@ async def test_all_revert_resets_and_marks_no_gain(self, tmp_path, monkeypatch):
         assert result["requires_e2e_validation"] is False
 
     @pytest.mark.asyncio
-    async def test_integrate_exception_reverts_tuner(self, tmp_path, monkeypatch):
+    async def test_integrate_exception_records_fault_not_revert(self, tmp_path, monkeypatch):
         coord = _coord(tmp_path, baseline_tput=100.0, framework="sglang")
 
         async def _boom(payload, *, session_dir):
@@ -3128,10 +3246,13 @@ async def _boom(payload, *, session_dir):
         }
         await coord._validate_gemm_tuning_e2e(result)
 
-        assert result["decision"] == "REVERT"
-        reverted = result["e2e_results"]["reverted"]
-        assert len(reverted) == 1
-        assert reverted[0]["reason"].startswith("RuntimeError")
+        assert result["status"] == "failed"
+        assert result["micro_decision"] == "integrate_fault"
+        assert result["e2e_gain_pct"] is None
+        faults = result["e2e_results"]["faults"]
+        assert len(faults) == 1
+        assert faults[0]["reason"] == "integrate_fault:handler_exception"
+        assert result["e2e_results"]["reverted"] == []
 
     @pytest.mark.asyncio
     async def test_timeout_fallback_when_explore_helper_raises(self, tmp_path, monkeypatch):

src/hyperloom/inference_optimizer/tests/test_forge_gemm_durable_persist.py

@@ -99,3 +99,45 @@
     assert dst.is_file()  # copy committed despite the snapshot failure
     assert out["AITER_CONFIG_GEMM_A8W8_BLOCKSCALE"] == str(dst)  # repoint SURVIVES
     assert snap == ""  # snapshot dir empty (it failed), but durability is kept
+
+
+def test_persist_fmoe_csv_uses_tuned_fmoe_stem(tmp_path, monkeypatch):
+    aiter_pkg = _fake_aiter(monkeypatch, tmp_path)
+    ws = tmp_path / "ws"
+    ws.mkdir()
+    src = ws / "tuned_fmoe.csv"
+    src.write_text("cu_num,token,model_dim,inter_dim,quantType\n304,16,4096,512,14\n", encoding="utf-8")
+
+    extra = {"AITER_CONFIG_FMOE": str(src)}
+    out, snap = rh._persist_forge_gemm_csv_durably(
+        extra, model_path="/models/DeepSeek-V4-Flash", session_dir=ws
+    )
+
+    dst = aiter_pkg / "configs" / "model_configs" / "tuned_fmoe_deepseek-v4-flash.csv"
+    assert dst.is_file()
+    assert out["AITER_CONFIG_FMOE"] == str(dst)
+    assert snap and Path(snap).is_dir()
+
+
+def test_persist_copies_dense_and_fmoe_together(tmp_path, monkeypatch):
+    aiter_pkg = _fake_aiter(monkeypatch, tmp_path)
+    ws = tmp_path / "ws"
+    ws.mkdir()
+    dense = ws / "dense.csv"
+    dense.write_text("gfx,M,N,K,splitK\ngfx950,64,5120,5120,2\n", encoding="utf-8")
+    fmoe = ws / "fmoe.csv"
+    fmoe.write_text("cu_num,token\n304,16\n", encoding="utf-8")
+
+    extra = {
+        "AITER_CONFIG_GEMM_A8W8_BLOCKSCALE": str(dense),
+        "AITER_CONFIG_FMOE": str(fmoe),
+    }
+    out, snap = rh._persist_forge_gemm_csv_durably(
+        extra, model_path="/models/Qwen3-14B-FP8", session_dir=ws
+    )
+
+    assert out["AITER_CONFIG_GEMM_A8W8_BLOCKSCALE"].endswith(
+        "a8w8_blockscale_tuned_gemm_qwen3-14b-fp8.csv"
+    )
+    assert out["AITER_CONFIG_FMOE"].endswith("tuned_fmoe_qwen3-14b-fp8.csv")
+    assert snap and (Path(snap) / "manifest.json").is_file()

src/hyperloom/inference_optimizer/tests/test_geak_breakdown_unit.py

@@ -632,3 +632,29 @@ async def test_sweep_via_geak_requires_existing_bench_script(tmp_path: Path) ->
 
     assert result["status"] == "failed"
     assert result["error_class"] == "missing_bench_script"
+
+
+def test_collect_gemm_tuning_prefers_e2e_gain_over_micro_speedup() -> None:
+    from hyperloom.inference_optimizer.breakdown.collectors.kernels import collect_gemm_tuning
+
+    out = collect_gemm_tuning(
+        {
+            "baseline_tput": 1000.0,
+            "gemm_tuning_attempts": [
+                {
+                    "engine": "forge",
+                    "status": "complete",
+                    "decision": "KEEP",
+                    "best_speedup": 1.5,
+                    "e2e_gain_pct": 9.26,
+                    "e2e_validated": True,
+                    "tuned_file": "/tmp/tuned.csv",
+                }
+            ],
+        }
+    )
+
+    run = out["runs"][0]
+    assert run["gain_pct"] == pytest.approx(9.26)
+    assert run["tuned_tput"] == pytest.approx(1092.6)
+    assert run["best_speedup"] == pytest.approx(1.5)