fix(replay): keep timing and capacity on the selected database version

python/aisimulate/src/aisimulate/capacity.py

@@ -53,10 +53,14 @@ def finish_lowering(value):
 
     lowered = dict(raw)
     timing = lowered.get("timing_model")
+    timing_system_roots = None
     if isinstance(timing, dict) and timing.get("type") == "external" and timing.get("provider") == "aic":
         authored = timing.get("config")
         if not isinstance(authored, dict):
             raise ValueError("external AIC timing config must be a mapping")
+        timing_system_roots = authored.get("systems_paths")
+        if not timing_system_roots and authored.get("systems_path") is not None:
+            timing_system_roots = [authored["systems_path"]]
         if "estimation_mode" in authored or "estimator_config" in authored:
             from aisimulate_core.sdk import RustForwardPassPerfModel
 
@@ -161,6 +165,16 @@ def finish_lowering(value):
     if not model:
         raise ValueError("AIC KV cache capacity estimation requires aic_model_path in engine args")
 
+    capacity_systems_path = lowered.get("systems_path")
+    if timing_system_roots:
+        from .sweeper.forward_pass_estimator import resolve_systems_paths
+
+        resolved_roots = resolve_systems_paths(timing_system_roots)
+        if resolved_roots:
+            # Native timing gives systems_paths precedence over systems_path;
+            # capacity preflight must query that same first effective root.
+            capacity_systems_path = resolved_roots[0]
+
     lowered["num_gpu_blocks"] = estimate_num_gpu_blocks(
         backend_name=backend,
         system=lowered.get("aic_system") or _DEFAULT_AIC_SYSTEM,
@@ -178,7 +192,11 @@ def finish_lowering(value):
         gpu_memory_utilization=lowered.get("gpu_memory_utilization"),
         mem_fraction_static=lowered.get("mem_fraction_static"),
         free_gpu_memory_fraction=lowered.get("free_gpu_memory_fraction"),
-        backend_version=lowered.get("aic_backend_version"),
+        backend_version=(
+            lowered.get("aic_backend_version")
+            if lowered.get("aic_backend_version") is not None
+            else lowered.get("backend_version")
+        ),
         pp_size=(lowered.get("aic_pp_size") if lowered.get("aic_pp_size") is not None else 1),
         moe_tp_size=lowered.get("aic_moe_tp_size"),
         moe_ep_size=lowered.get("aic_moe_ep_size"),
@@ -193,7 +211,7 @@ def finish_lowering(value):
             for name in ("moe_backend", "attention_backend", "enable_eplb", "wideep_num_slots")
             if lowered.get(f"aic_{name}") is not None
         },
-        systems_path=lowered.get("systems_path"),
+        systems_path=capacity_systems_path,
         cuda_graph_reserved_bytes=lowered.get("cuda_graph_reserved_bytes", 0),
         **({"diagnostics": memory_diagnostics} if memory_diagnostics is not None else {}),
     )
@@ -247,6 +265,15 @@ def estimate_num_gpu_blocks(
         estimate_num_gpu_blocks as aic_estimate_num_gpu_blocks,
     )
 
+    if backend_version is None:
+        from aisimulate_core.sdk.perf_database import get_latest_database_version
+
+        # Use the maintained current slot (or the latest version in a custom
+        # legacy root), matching the database used for this capacity estimate.
+        backend_version = get_latest_database_version(system, backend_name, systems_paths=systems_path)
+        if backend_version is None:
+            raise ValueError(f"no perf database for system={system!r}, backend={backend_name!r}")
+
     validate_moe_controls(enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots)
 
     if backend_name == "trtllm":
@@ -268,9 +295,7 @@ def estimate_num_gpu_blocks(
             model_path,
             system,
             backend_name,
-            backend_version=(
-                backend_version if backend_version is not None else DEFAULT_BACKEND_VERSIONS[backend_name]
-            ),
+            backend_version=backend_version,
             scheduler_block_size=block_size,
             max_num_tokens=max_num_batched_tokens,
             max_batch_size=max_num_sequences,

python/aisimulate/src/aisimulate/config/epd.py

@@ -40,6 +40,8 @@ def validate_epd_prediction_mapping(value: dict, spec: ReplaySpec) -> None:
             raise ValueError("encoder parameters or resolved database version changed")
         if engine.model != encoder.model or engine.backend != encoder.backend:
             raise ValueError("encoder model/backend identity changed")
+        if spec.backend_deployment.backend_version and engine.backend_version is None:
+            raise ValueError("selected language backend version was dropped")
         source = prediction.traffic.source
         if not isinstance(source, SyntheticSource) or source.images is None:
             raise ValueError("fixed image workload was dropped")
@@ -82,10 +84,10 @@ def validate_epd_prediction_mapping(value: dict, spec: ReplaySpec) -> None:
 
 def _language_execution(spec: ReplaySpec) -> dict:
     """Normalize compiler/Sweeper spellings at the existing execution boundary."""
-    from aisimulate_core.sdk.perf_database import resolve_query_version
+    from aisimulate_core.sdk.perf_database import get_latest_database_version, resolve_query_version
 
-    from ..capacity import DEFAULT_BACKEND_VERSIONS
     from ..runner import _materialize_engine_role
+    from ..sweeper.forward_pass_estimator import resolve_systems_paths
     from .engine import EstimatorPolicyConfig, SchedulerPredictionConfig
 
     deployment = spec.backend_deployment
@@ -104,10 +106,24 @@ def _language_execution(spec: ReplaySpec) -> dict:
         rank.setdefault("prefill_schedule_interval", SchedulerPredictionConfig().prefill_schedule_interval)
         rank.setdefault("prefill_decode_interval", SchedulerPredictionConfig().prefill_decode_interval)
         timing = rank["timing_model"]["config"]
+        roots = timing.get("systems_paths")
+        if not roots and timing.get("systems_path") is not None:
+            roots = [timing["systems_path"]]
+        resolved_roots = list(resolve_systems_paths(roots)) if roots is not None else None
+        version = timing.get("backend_version")
+        if not version:
+            version = get_latest_database_version(
+                timing["system"],
+                timing["backend"],
+                **({"systems_paths": resolved_roots} if resolved_roots is not None else {}),
+            )
+            if version is None:
+                raise ValueError(f"no perf database for system={timing['system']!r}, backend={timing['backend']!r}")
         timing["backend_version"] = resolve_query_version(
             timing["system"],
             timing["backend"],
-            timing.get("backend_version") or DEFAULT_BACKEND_VERSIONS[deployment.backend],
+            version,
+            systems_paths=resolved_roots,
         )
         timing.setdefault("cuda_graph_reserved_bytes", 0)
         # Legacy Sweeper descriptors omit the default policy; the compiler

python/aisimulate/src/aisimulate/runner.py

@@ -1182,6 +1182,42 @@ def _materialize_engine_role(
     num_gpu_blocks_is_explicit = False
     if "rank" not in role_config:
         num_gpu_blocks_is_explicit = role_config.get("num_gpu_blocks") is not None
+        for name in ("aic_backend_version", "backend_version"):
+            if role_config.get(name) is None:
+                role_config.pop(name, None)
+        # Propagate authored versions before capacity preflight. Otherwise a
+        # deployment pin reaches timing only after capacity has used current.
+        timing = role_config.get("timing_model")
+        timing_config = None
+        if (
+            isinstance(timing, dict)
+            and timing.get("type") == "external"
+            and timing.get("provider") == "aic"
+            and isinstance(timing.get("config"), dict)
+        ):
+            timing_config = timing["config"]
+        version = next(
+            (
+                value
+                for value in (
+                    role_config.get("aic_backend_version"),
+                    role_config.get("backend_version"),
+                    timing_config.get("backend_version") if timing_config is not None else None,
+                    deployment_backend_version or None,
+                )
+                if value is not None
+            ),
+            None,
+        )
+        if version is not None:
+            if (
+                not num_gpu_blocks_is_explicit
+                and role_config.get("aic_backend") is not None
+                and not any(name in role_config for name in ("aic_backend_version", "backend_version"))
+            ):
+                role_config["aic_backend_version"] = version
+            if timing_config is not None and timing_config.get("backend_version") is None:
+                role_config["timing_model"] = {**timing, "config": {**timing_config, "backend_version": version}}
         role_config = materialize_aic_num_gpu_blocks(
             role_config,
             **({"memory_diagnostics": role_memory} if role_memory is not None else {}),
@@ -1367,6 +1403,8 @@ def resolved_version(value):
             roots = identity.get("systems_paths")
             if roots is None and identity.get("systems_path") is not None:
                 roots = [identity["systems_path"]]
+            if roots is None and aic_timing_overrides.get("systems_path") is not None:
+                roots = [aic_timing_overrides["systems_path"]]
             return resolve_query_version(
                 identity.get("system", system),
                 backend,

tests/test_aic.py

@@ -15,8 +15,17 @@
 ]
 
 
+@pytest.mark.parametrize(
+    "version_fields",
+    [
+        {"backend_version": "test-version"},
+        {"aic_backend_version": "test-version"},
+        {"aic_backend_version": None, "backend_version": "test-version"},
+    ],
+)
 def test_materializer_sets_rank_local_capacity_without_forwarding_nextn(
     monkeypatch,
+    version_fields,
 ) -> None:
     calls = []
 
@@ -29,6 +38,7 @@ def estimate(**kwargs):
         {
             "engine_type": "vllm",
             "aic_backend": "vllm",
+            **version_fields,
             "aic_system": "h200_sxm",
             "aic_model_path": "test-model",
             "aic_attention_dp_size": 2,
@@ -44,6 +54,7 @@ def estimate(**kwargs):
     assert lowered["num_gpu_blocks"] == 46000
     assert lowered["dp_size"] == 2
     assert calls[0]["attention_dp_size"] == 2
+    assert calls[0]["backend_version"] == "test-version"
     assert calls[0]["pp_size"] == 3
     assert calls[0]["systems_path"] == "/tmp/custom-systems.yaml"
     assert calls[0]["max_num_sequences"] == 37
@@ -60,9 +71,16 @@ def estimate(*args, **kwargs):
         calls.append((args, kwargs))
         return 123
 
-    from aisimulate_core.sdk import memory
+    from aisimulate_core.sdk import memory, perf_database
 
     monkeypatch.setattr(memory, "estimate_num_gpu_blocks", estimate)
+    version_calls = []
+
+    def latest(system, backend, *, systems_paths):
+        version_calls.append((system, backend, systems_paths))
+        return "test-current"
+
+    monkeypatch.setattr(perf_database, "get_latest_database_version", latest)
     blocks = aic.estimate_num_gpu_blocks(
         backend_name="vllm",
         system="h200_sxm",
@@ -81,7 +99,8 @@ def estimate(*args, **kwargs):
     assert blocks == 123
     args, kwargs = calls[0]
     assert args == ("test-model", "h200_sxm", "vllm")
-    assert kwargs["backend_version"] == "0.19.0"
+    assert kwargs["backend_version"] == "test-current"
+    assert version_calls == [("h200_sxm", "vllm", "/tmp/custom-systems.yaml")]
     assert kwargs["memory_fraction_kind"] == "of_total"
     assert kwargs["memory_fraction_value"] == 0.9
     assert kwargs["pp_size"] == 3
@@ -104,6 +123,11 @@ def test_explicit_capacity_is_preserved_without_estimation(monkeypatch) -> None:
         "aic_backend": "vllm",
         "aic_model_path": "test-model",
         "num_gpu_blocks": 17,
+        "timing_model": {
+            "type": "external",
+            "provider": "aic",
+            "config": {"systems_paths": ["/path/that/does/not/exist"]},
+        },
     }
     assert aic.materialize_aic_num_gpu_blocks(raw) == raw
 
@@ -137,3 +161,96 @@ def estimate(**kwargs):
     assert calls[0]["gpu_memory_utilization"] == 0.0
     assert calls[0]["mem_fraction_static"] == 0.0
     assert calls[0]["free_gpu_memory_fraction"] == 0.0
+
+
+@pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"])
+@pytest.mark.parametrize("timing", ["default", "fixed", "polynomial"])
+def test_prediction_omitted_version_uses_current_database_for_capacity(monkeypatch, backend, timing):
+    from aisimulate.compiler import prediction_to_replay_spec
+    from aisimulate.config import CorePredictionConfig
+    from aisimulate_core.sdk import perf_database
+
+    expected = perf_database.get_version_slots("h200_sxm", backend)["current"]
+    get_database = perf_database.get_database
+    versions = []
+
+    def recorded_database(*args, **kwargs):
+        database = get_database(*args, **kwargs)
+        versions.append(database.version)
+        return database
+
+    monkeypatch.setattr(perf_database, "get_database", recorded_database)
+    timing_config = {"type": timing}
+    if timing == "fixed":
+        timing_config.update(prefill_ms=1, decode_ms=1)
+    config = CorePredictionConfig.model_validate(
+        {
+            "engine": {
+                "model": "Qwen/Qwen3-32B-FP8",
+                "hardware": "h200_sxm",
+                "backend": backend,
+                "context_length": 4096,
+                "workers": {"aggregated": {"parallelism": {"tensor": 2}, "timing": timing_config}},
+            }
+        }
+    )
+    deployment = prediction_to_replay_spec(config).backend_deployment
+    args = aic.materialize_aic_num_gpu_blocks(deployment.agg_engine_args)
+    assert args["num_gpu_blocks"] > 0
+    assert versions and set(versions) == {expected}
+    assert config.engine.backend_version is None
+    if timing == "default":
+        assert args["timing_model"]["config"]["backend_version"] == expected
+    else:
+        assert args["timing_model"]["type"] == timing
+
+
+@pytest.mark.parametrize("version", ["current", "next", "0.24.0", "0.19.0"])
+def test_capacity_preserves_explicit_backend_version(monkeypatch, version):
+    from aisimulate_core.sdk import memory, perf_database
+
+    monkeypatch.setattr(
+        perf_database,
+        "get_latest_database_version",
+        lambda *_args, **_kwargs: pytest.fail("explicit versions must not be replaced"),
+    )
+    versions = []
+
+    def estimate(*args, **kwargs):
+        versions.append(kwargs["backend_version"])
+        return 128
+
+    monkeypatch.setattr(memory, "estimate_num_gpu_blocks", estimate)
+    assert (
+        aic.estimate_num_gpu_blocks(
+            backend_name="vllm",
+            backend_version=version,
+            system="h200_sxm",
+            model_path="test-model",
+            tp_size=1,
+            block_size=64,
+            max_num_batched_tokens=4096,
+        )
+        == 128
+    )
+    assert versions == [version]
+
+
+def test_capacity_without_a_database_fails_before_estimation(monkeypatch):
+    from aisimulate_core.sdk import memory, perf_database
+
+    monkeypatch.setattr(perf_database, "get_latest_database_version", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(
+        memory,
+        "estimate_num_gpu_blocks",
+        lambda *_args, **_kwargs: pytest.fail("missing versions must fail before estimation"),
+    )
+    with pytest.raises(ValueError, match="no perf database.*unsupported-test-gpu.*vllm"):
+        aic.estimate_num_gpu_blocks(
+            backend_name="vllm",
+            system="unsupported-test-gpu",
+            model_path="test-model",
+            tp_size=1,
+            block_size=64,
+            max_num_batched_tokens=4096,
+        )