diff --git a/.github/release-gates.json b/.github/release-gates.json new file mode 100644 index 000000000..f3c2afdd8 --- /dev/null +++ b/.github/release-gates.json @@ -0,0 +1,8 @@ +{ + "pending_migrations": [ + { + "pull_request": "https://github.com/ai-dynamo/dynamo/pull/14065", + "requirement": "Refresh and merge Dynamo's Planner and wheel-smoke migration to best_available(config), with required worker_type, nested estimator_config, and the released-wheel compatibility bridge. Validate against the canonical wheel before clearing this gate and releasing the API." + } + ] +} diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index feeaf4c72..19d82cbed 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -130,6 +130,20 @@ jobs: } } core.setFailed('commit_sha is not on main or any release/* branch'); + # Both scheduled and approved manual publication require completed migrations. + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + ref: ${{ steps.target.outputs.sha }} + persist-credentials: false + # Keep the checker/current policy trusted while also checking the exact + # publication target's migration declarations. Neither can bypass the other. + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + ref: ${{ github.sha }} + path: release-policy + persist-credentials: false + - name: Require completed downstream API migrations + run: python3 release-policy/scripts/check_release_migrations.py --target-gates .github/release-gates.json # Computed once here so both arch legs stamp the same date even if the # run crosses a UTC midnight. - name: Compute unique nightly version diff --git a/AGENTS.md b/AGENTS.md index f44ea52a8..e3536c165 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,6 +2,13 @@ This file adds explicit repository-wide development guards. +## Performance Model Changes + +Before changing a performance model, its configuration, or a caller in Rust, +Python, CLI, Sweeper, Replay, or Planner, MUST read and follow +[`perfmodel-api.md`](python/aisimulate/.claude/rules/perfmodel-api.md). +This includes new features and configuration migrations. + ## Required First Step Before making any change under: diff --git a/Cargo.lock b/Cargo.lock index df2d9ffd3..b3f9b1447 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -54,6 +54,7 @@ dependencies = [ "rustc-hash", "serde", "serde_json", + "serde_path_to_error", "serde_yaml", "sha2", "slotmap", @@ -1247,6 +1248,17 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + [[package]] name = "serde_yaml" version = "0.9.34+deprecated" diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index b5bf95f64..9884c861a 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -879,4 +879,5 @@ Copyright 2018- The Hugging Face team. All rights reserved. - License: Apache-2.0; NVIDIA maintainer confirmed permission to migrate and publish this code under Apache-2.0. - Modified: development-only two-predictor evaluation, public overview export, - GitHub Pages presentation, and local import paths. No Plotly assets included. + GitHub Pages presentation, local import paths, and canonical estimator API + adaptation with older-wheel compatibility. No Plotly assets included. diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index f04a99d34..aac04c6c1 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -35,6 +35,7 @@ rand_chacha = "0.3.1" rustc-hash = "2" serde = { workspace = true } serde_json = { workspace = true } +serde_path_to_error = "0.1" serde_yaml = "0.9" sha2 = "0.10" slotmap = "1" diff --git a/crates/core/parity_tests/perfmodel/test_engine_step_parity.py b/crates/core/parity_tests/perfmodel/test_engine_step_parity.py index 099d01175..cdebbfe0c 100644 --- a/crates/core/parity_tests/perfmodel/test_engine_step_parity.py +++ b/crates/core/parity_tests/perfmodel/test_engine_step_parity.py @@ -673,7 +673,10 @@ def __call__(self): def _quiet_call(func, *args, **kwargs): """Keep interpolation loader chatter out of parity test output.""" - with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + with ( + contextlib.redirect_stdout(io.StringIO()), + contextlib.redirect_stderr(io.StringIO()), + ): return func(*args, **kwargs) @@ -1135,7 +1138,9 @@ def _comparison_metrics( return {name: (python_metrics[name], rust_metrics[name]) for name in rust_metrics} -def _static_comparison_metrics(case: EngineStepParityCase) -> dict[str, tuple[float, float]]: +def _static_comparison_metrics( + case: EngineStepParityCase, +) -> dict[str, tuple[float, float]]: return _comparison_metrics(case, "static") @@ -1145,11 +1150,15 @@ def _mixed_step_comparison_metrics( return _comparison_metrics(case, "mixed") -def _agg_comparison_metrics(case: EngineStepParityCase) -> dict[str, tuple[float, float]]: +def _agg_comparison_metrics( + case: EngineStepParityCase, +) -> dict[str, tuple[float, float]]: return _comparison_metrics(case, "agg") -def _disagg_comparison_metrics(case: EngineStepParityCase) -> dict[str, tuple[float, float]]: +def _disagg_comparison_metrics( + case: EngineStepParityCase, +) -> dict[str, tuple[float, float]]: return _comparison_metrics(case, "disagg") @@ -1432,7 +1441,14 @@ class TestRustEngineHandleDatabasePolicyIdentity: """ def _static_ctx_ms(self, model, view) -> float: - rc = config.RuntimeConfig(batch_size=1, beam_width=1, isl=1024, osl=8, prefix=0, engine_step_backend="rust") + rc = config.RuntimeConfig( + batch_size=1, + beam_width=1, + isl=1024, + osl=8, + prefix=0, + engine_step_backend="rust", + ) ctx_latency, _gen, *_ = rust_engine_step.estimate_static_latency_breakdown_with_rust( model, view, rc, "static_ctx", 1, 1.0 ) @@ -2332,7 +2348,12 @@ def test_hybrid_ngram_mixed_pin(self, fpm_systems_root): model, backend, database = self._build_ngram() rc = config.RuntimeConfig(batch_size=4, beam_width=1, isl=1024, osl=2) - est = backend.run_mixed(model, database, rc, MixedStepInput(context_tokens=1024, num_decode_requests=1)) + est = backend.run_mixed( + model, + database, + rc, + MixedStepInput(context_tokens=1024, num_decode_requests=1), + ) value = est.latency_ms if _FPM_HYBRID_MIXED_FROZEN is None: print(f"\nPIN mixed: {value!r}") @@ -2342,6 +2363,77 @@ def test_hybrid_ngram_mixed_pin(self, fpm_systems_root): class TestRustEngineStepFpmParity: + @pytest.mark.parametrize("database_mode", ["SILICON", "HYBRID"]) + def test_auto_skips_absent_op_tables(self, fpm_systems_root, database_mode): + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + cfg = ForwardPassPerfModelConfig( + model=_FPM_MODEL, + system="b200_sxm", + backend="vllm", + backend_version=_FPM_VERSION, + worker_type="decode", + tp=2, + moe_tp_size=1, + moe_ep_size=2, + gemm_quant_mode="fp8_block", + moe_quant_mode="fp8_block", + fmha_quant_mode="bfloat16", + kvcache_quant_mode="fp8", + comm_quant_mode="half", + systems_paths=(str(fpm_systems_root),), + database_mode=database_mode, + ) + model = RustForwardPassPerfModel.best_available(cfg) + try: + provenance = model.diagnostics()["provenance"] + assert provenance["selected_estimation_mode"] == "fpm_interpolation" + assert "gemm_perf.parquet" in provenance["selection_failures"][0] + # Exact measured fixture row: four decode requests, 4,100 total KV. + assert model.estimate_forward_pass_time_ms( + { + "version": 1, + "scheduled_requests": {"num_decode_requests": 4, "sum_decode_kv_tokens": 4100}, + } + ) == pytest.approx(4.5) + finally: + model.close() + from dataclasses import replace + + from aiconfigurator_core.sdk.errors import PerfDataNotAvailableError + + with pytest.raises(PerfDataNotAvailableError, match="required op-level data unavailable"): + RustForwardPassPerfModel.best_available(replace(cfg, estimation_mode="op_level")) + + def test_sol_does_not_require_op_tables(self, fpm_systems_root): + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + model = RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig( + model="Qwen/Qwen3-32B", + system="b200_sxm", + backend="vllm", + backend_version=_FPM_VERSION, + worker_type="decode", + tp=2, + systems_paths=(str(fpm_systems_root),), + database_mode="SOL", + ) + ) + try: + assert model.diagnostics()["provenance"]["selected_estimation_mode"] == "op_level" + assert ( + model.estimate_forward_pass_time_ms( + { + "version": 1, + "scheduled_requests": {"num_decode_requests": 4, "sum_decode_kv_tokens": 4100}, + } + ) + > 0 + ) + finally: + model.close() + """forward_model='fpm' regression vs the frozen Python reference. The Python FPM walk is gone (Phase 2 PR-3); the live side below is the @@ -2417,7 +2509,11 @@ def test_fpm_arena_selects_the_fpm_engine(self, fpm_systems_root, monkeypatch): "nextn": None, "forward_model": "fpm", } - model = RustForwardPassPerfModel.from_native(config) + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig + + model = RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig.from_legacy_engine_config(config, "aggregated") + ) decode_only = [ { "version": 1, @@ -2471,7 +2567,13 @@ def test_fpm_spec_tags(self, fpm_systems_root, monkeypatch): ("static_gen", 4, 1024, 2, 0), # exact decode hit at B=4 ("static_gen", 2, 1024, 2, 0), # uncollected batch -> transfer (SOL) ("static_gen", 4, 9_000_000, 2, 0), # out of domain -> both error - ("static_ctx", 16, 256, 1, 0), # above the batch ceiling -> pure clamp (kv/T = 0) + ( + "static_ctx", + 16, + 256, + 1, + 0, + ), # above the batch ceiling -> pure clamp (kv/T = 0) ("static_ctx", 16, 320, 1, 256), # high KV pressure -> SOL-rescaled clamp ], ) @@ -2490,7 +2592,12 @@ def test_fpm_static_parity(self, fpm_systems_root, monkeypatch, mode, batch, isl (0, 4, 1024, 2), # gen-only keeps full decode (0, 600, 100, 2), # gen-only across the decode regime boundary (eager side) (1024, 0, 1024, 2), # prefill-only chunk - (4096, 0, 256, 1), # 16 whole prefills: certified batch clamp to the ceiling + ( + 4096, + 0, + 256, + 1, + ), # 16 whole prefills: certified batch clamp to the ceiling ], ) def test_fpm_mixed_step_parity(self, fpm_systems_root, monkeypatch, ctx_tokens, gen_tokens, isl, osl): diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 470e4cab4..05f015678 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -35,18 +35,32 @@ pub use replay::{ReplayReport, ReplaySpec, Replayer}; // the explicit `ReplayEngineConfig` alias above. pub use perfmodel::EngineConfig; pub use perfmodel::{ - AicError, BackendKind, DataType, DatabaseMode, ENGINE_CONFIG_SCHEMA_VERSION, - ENGINE_SPEC_SCHEMA_VERSION, EstimateSource, FPM_VERSION, ForwardPassMetrics, - ForwardPassPerfDiagnostics, ForwardPassPerfModel, ForwardPassPerfOptions, - ForwardPassPerfReadiness, ForwardPassPerfSource, ForwardPassRegressionStoreDiagnostics, - ForwardPassRegressionWorkloadKind, ForwardPassWorkerType, KvCacheEstimate, + AicError, BackendKind, CorrectionConfig, DataType, DatabaseMode, ENGINE_CONFIG_SCHEMA_VERSION, + ENGINE_SPEC_SCHEMA_VERSION, EstimateSource, EstimationMode, EstimatorConfig, FPM_VERSION, + ForwardPassFallbackPolicy, ForwardPassMetrics, ForwardPassPerfDiagnostics, + ForwardPassPerfModel, ForwardPassPerfModelConfig, ForwardPassPerfOptions, + ForwardPassPerfProvenance, ForwardPassPerfReadiness, ForwardPassPerfSource, + ForwardPassRegressionStoreDiagnostics, ForwardPassRegressionWorkloadKind, + ForwardPassSpeculationConfig, ForwardPassWorkerType, FpmRegressionConfig, KvCacheEstimate, KvCacheEstimateAdjusted, KvCacheEstimateError, KvCacheEstimateOptions, KvCacheEstimateRequest, KvCacheMemoryFraction, MemoryBreakdown, ParallelMapping, QuantizationConfig, - QueuedRequestMetrics, ScheduledRequestMetrics, SpeculativeConfig, + QueuedRequestMetrics, RegressionFeatureWeights, SamplingConfig, ScheduledRequestMetrics, + SpeculativeConfig, +}; +pub use perfmodel::{ + CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, OpLevelConfig, + RegressionFitConfig, RegressionFitKind, }; #[cfg(feature = "python")] -pub use perfmodel::{AicEngine, AicEngineBuilder, estimate_kv_cache}; +pub use perfmodel::{ + AicEngine, + // Low-level Rust embedder API for compiled step-latency handles. This is + // deliberately not registered on the Python module and is not an + // alternative ForwardPassPerfModel construction boundary. + AicEngineBuilder, + estimate_kv_cache, +}; // The imported perf-model sources historically used additional crate-root // module paths. Keep these module aliases crate-private so the mirror subtree diff --git a/crates/core/src/perfmodel/engine/mod.rs b/crates/core/src/perfmodel/engine/mod.rs index b66f3e219..daa52cd61 100644 --- a/crates/core/src/perfmodel/engine/mod.rs +++ b/crates/core/src/perfmodel/engine/mod.rs @@ -9,6 +9,7 @@ //! static-inference composition. PyO3 bindings over `Engine` live in //! [`crate::py`]. +mod readiness; pub mod runtime; pub mod spec; diff --git a/crates/core/src/perfmodel/engine/readiness.rs b/crates/core/src/perfmodel/engine/readiness.rs new file mode 100644 index 000000000..fafd848e6 --- /dev/null +++ b/crates/core/src/perfmodel/engine/readiness.rs @@ -0,0 +1,515 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Construction-time table availability. Shape, slice and interpolation-domain +//! coverage remains a query-time check; no synthetic workload is queried here. + +use std::collections::HashMap; + +use crate::AicError; +use crate::common::enums::{DatabaseMode, GemmQuantMode, TransferKind}; +use crate::config::PerfSource; +use crate::operators::Op; +use crate::perf_database::{PerfDatabase, kernel_source_ok, parquet_loader::PerfReader}; + +pub(super) fn validate<'a>( + db: &PerfDatabase, + ops: impl Iterator, +) -> Result<(), AicError> { + let mut check = Availability { + db, + tables: HashMap::new(), + }; + for op in ops { + check.op(op)?; + } + Ok(()) +} + +struct Availability<'a> { + db: &'a PerfDatabase, + tables: HashMap, +} + +impl Availability<'_> { + fn has(&mut self, name: &str) -> Result { + if let Some(found) = self.tables.get(name) { + return Ok(*found); + } + let sources = match name { + "nccl_perf.parquet" => self + .db + .communication + .nccl_root() + .map(|root| vec![PerfSource(root.join(name), None)]) + .unwrap_or_default(), + "oneccl_perf.parquet" => self + .db + .communication + .oneccl_root() + .map(|root| vec![PerfSource(root.join(name), None)]) + .unwrap_or_default(), + "dsv4_megamoe_module_perf.parquet" => { + vec![PerfSource( + self.db.dsv4_megamoe.primary_path().to_owned(), + None, + )] + } + _ => self + .db + .source_resolver + .sources_for(name, &self.db.data_root)?, + }; + let mut found = false; + for source in sources { + if !source.path().is_file() { + continue; + } + let reader = PerfReader::open(source.path())?; + let kernel = reader.col_optional("kernel_source"); + for row in reader.rows()? { + if kernel_source_ok(source.kernel_sources(), kernel, &row?)? { + found = true; + break; + } + } + if found { + break; + } + } + self.tables.insert(name.to_owned(), found); + Ok(found) + } + + fn any(&mut self, names: &[&str]) -> Result<(), AicError> { + for name in names { + if self.has(name)? { + return Ok(()); + } + } + Err(AicError::PerfDatabase(format!( + "required op-level data unavailable for {:?}: no rows in {} at {}", + self.db.database_mode, + names.join(" or "), + self.db.data_root.display(), + ))) + } + + fn op(&mut self, op: &Op) -> Result<(), AicError> { + use Op::*; + // CSA context parallelism consumes both sparse correction tables even + // when its base module uses the analytic SOL. HCA never consumes them. + if let Dsv4Context(module) = op + && module.cp_size > 1 + && module.attn_kind == crate::perf_database::dsv4::AttnKind::Csa + { + self.any(&["dsv4_paged_mqa_logits_module_perf.parquet"])?; + self.any(&["dsv4_csa_topk_calib_perf.parquet"])?; + } + let sol = matches!( + self.db.database_mode, + DatabaseMode::Sol | DatabaseMode::SolFull + ); + match op { + Overlap(group) => { + for child in group.group_a.iter().chain(&group.group_b) { + self.op(child)?; + } + return Ok(()); + } + Fallback(group) => { + let primary = if self.db.database_mode == DatabaseMode::Hybrid { + validate(&self.db.silicon_view(), [&*group.primary].into_iter()) + } else { + self.op(&group.primary) + }; + match primary { + Ok(()) => return Ok(()), + Err(error) + if error.is_missing_perf_data() || matches!(error, AicError::Io { .. }) => + { + for child in &group.fallback { + self.op(child)?; + } + return Ok(()); + } + Err(error) => return Err(error), + } + } + TokenScale(scale) => return self.op(&scale.op), + FpmForward(fpm) => { + self.db + .fpm_forward + .select_cell(&fpm.match_identity, &fpm.model_path)?; + return Ok(()); + } + // These families have no analytic/empirical implementation. + MoeAllToAll(_) | MoeExpertCompute(_) | Dsv4MegaMoe(_) + if !matches!( + self.db.database_mode, + DatabaseMode::Silicon | DatabaseMode::Hybrid + ) => + { + return Err(AicError::UnsupportedModel(format!( + "{} requires silicon data; {:?} is unsupported", + op.name(), + self.db.database_mode, + ))); + } + _ if sol => return Ok(()), + _ => {} + } + match op { + Gemm(gemm) => { + self.any(&["gemm_perf.parquet"])?; + if gemm.quant_mode == GemmQuantMode::Fp8Static { + self.any(&["computescale_perf.parquet"])?; + if gemm.low_precision_input { + self.any(&["scale_matrix_perf.parquet"])?; + } + } + Ok(()) + } + ContextAttention(_) => self.any(&["context_attention_perf.parquet"]), + GenerationAttention(_) => self.any(&["generation_attention_perf.parquet"]), + EncoderAttention(_) => self.any(&["encoder_attention_perf.parquet"]), + ContextMla(_) => self.any(&["context_mla_perf.parquet"]), + GenerationMla(_) => self.any(&["generation_mla_perf.parquet"]), + MlaModuleContext(_) => self.any(&["mla_context_module_perf.parquet"]), + MlaModuleGeneration(_) => self.any(&["mla_generation_module_perf.parquet"]), + MlaBmm(_) => self.any(&["mla_bmm_perf.parquet"]), + Moe(_) => self.any(&["moe_perf.parquet"]), + // State-space kernels explicitly fall back to their analytic SOL + // on missing tables in every database mode. + Mamba2(_) | Gdn(_) | Kda(_) => Ok(()), + Mhc(_) => self.any(&["mhc_module_perf.parquet"]), + WideEpContextMla(_) => self.any(&["wideep_context_mla_perf.parquet"]), + WideEpGenerationMla(_) => self.any(&["wideep_generation_mla_perf.parquet"]), + DsaContext(module) => { + self.any(&["dsa_context_module_perf.parquet"])?; + if module.cp_size > 1 { + let prefix = + crate::perf_database::dsa::dsa_sparse_file_prefix(&module.architecture); + self.any(&[&format!("{prefix}_mqa_logits_module_perf.parquet")])?; + self.any(&[&format!("{prefix}_topk_module_perf.parquet")])?; + } + Ok(()) + } + DsaGeneration(_) => self.any(&["dsa_generation_module_perf.parquet"]), + MsaContext(_) | MsaGeneration(_) => { + let phase = if matches!(op, MsaContext(_)) { + "context" + } else { + "generation" + }; + let own = format!("msa_{phase}_module_perf.parquet"); + if self.db.database_mode != DatabaseMode::Empirical && self.has(&own)? { + return Ok(()); + } + if self.db.database_mode != DatabaseMode::Silicon + && self.db.transfer_policy.contains(TransferKind::XOp) + { + return self.any(&[&format!("dsa_{phase}_module_perf.parquet")]); + } + self.any(&[&own]) + } + Dsv4Context(module) | Dsv4Generation(module) => { + use crate::perf_database::dsv4::AttnKind; + let kind = match module.attn_kind { + AttnKind::Csa => "csa", + AttnKind::Hca => "hca", + }; + let phase = if matches!(op, Dsv4Context(_)) { + "context" + } else { + "generation" + }; + self.any(&[&format!("dsv4_{kind}_{phase}_module_perf.parquet")]) + } + Dsv4MegaMoe(_) => self.any(&["dsv4_megamoe_module_perf.parquet"]), + MoeAllToAll(_) => self.any(&[ + "moe_a2a_perf.parquet", + "wideep_deepep_ll_perf.parquet", + "wideep_deepep_normal_perf.parquet", + "trtllm_alltoall_perf.parquet", + ]), + MoeExpertCompute(_) => self.any(&[ + "moe_expert_compute_perf.parquet", + "wideep_context_moe_perf.parquet", + "wideep_generation_moe_perf.parquet", + "wideep_moe_perf.parquet", + "moe_perf.parquet", + ]), + CustomAllReduce(comm) if comm.tp_size > 1 => { + if self.db.system_spec.node.num_gpus_per_node == 72 + && comm.tp_size > 4 + && matches!( + self.db.database_mode, + DatabaseMode::Silicon | DatabaseMode::Hybrid + ) + { + self.any(&["nccl_perf.parquet", "oneccl_perf.parquet"]) + } else { + self.any(&["custom_allreduce_perf.parquet"]) + } + } + Nccl(comm) if comm.num_gpus > 1 => { + self.any(&["nccl_perf.parquet", "oneccl_perf.parquet"]) + } + // Dispatch chooses its comm legs from topology at query time. A + // single-rank graph needs no communication table. + MoeDispatch(dispatch) if dispatch.moe_tp_size * dispatch.moe_ep_size > 1 => { + self.any(&[ + "custom_allreduce_perf.parquet", + "nccl_perf.parquet", + "oneccl_perf.parquet", + "trtllm_alltoall_perf.parquet", + ]) + } + Vision(_) => { + self.any(&["gemm_perf.parquet"])?; + self.any(&["encoder_attention_perf.parquet"]) + } + Embedding(_) | Elementwise(_) | P2P(_) | CustomAllReduce(_) | Nccl(_) + | MoeDispatch(_) => Ok(()), + Overlap(_) | Fallback(_) | TokenScale(_) | FpmForward(_) => Ok(()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::enums::TransferPolicy; + use crate::config::PerfDbSources; + use crate::operators::dsv4::{Dsv4MegaMoeOp, Dsv4ModuleOp}; + use crate::operators::op::{FallbackOp, OverlapOp}; + use crate::operators::{EmbeddingOp, GemmOp}; + use crate::perf_database::energy_test_fixtures::{Col, write_parquet}; + use crate::perfmodel::engine::{Engine, spec::EngineSpec}; + use std::sync::Arc; + + fn systems() -> tempfile::TempDir { + let root = tempfile::tempdir().unwrap(); + let spec = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../python/aisimulate/src/aiconfigurator_core/systems/b200_sxm.yaml"); + std::fs::copy(spec, root.path().join("b200_sxm.yaml")).unwrap(); + std::fs::create_dir_all(root.path().join("data/b200_sxm/vllm/0.24.0")).unwrap(); + root + } + + fn engine_readiness(db: PerfDatabase, op: Op) -> Result<(), AicError> { + let config = serde_json::from_value(serde_json::json!({ + "schema_version": crate::ENGINE_CONFIG_SCHEMA_VERSION, + "model_name": "readiness-fixture", + "system_name": "b200_sxm", "backend": "vllm", "backend_version": "0.24.0", + "tp_size": 1, "pp_size": 1, "database_mode": db.database_mode, + })) + .unwrap(); + let (context, generation) = if matches!(op, Op::Dsv4Generation(_)) { + (vec![], vec![op]) + } else { + (vec![op], vec![]) + }; + Engine::build(EngineSpec::new(config, context, generation), Arc::new(db))? + .validate_forward_pass_readiness() + } + + fn dsv4_op(kind: &str, context: bool, cp: u32) -> Op { + let module: Dsv4ModuleOp = serde_json::from_value(serde_json::json!({ + "name": "attention", "scale_factor": 1.0, "attn_kind": kind, + "num_heads": 64, "native_heads": 64, "tp_size": 1, "cp_size": cp, + "kv_cache_dtype": "bfloat16", "fmha_quant_mode": "bfloat16", + "gemm_quant_mode": "bfloat16", "architecture": "DeepseekV4ForCausalLM", + })) + .unwrap(); + if context { + Op::Dsv4Context(module) + } else { + Op::Dsv4Generation(module) + } + } + + fn table(path: &std::path::Path) { + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + write_parquet(path, &[Col::Str("kernel_source", vec!["fixture"])]); + } + + #[test] + fn dsv4_requires_its_phase_module_not_an_auxiliary_table() { + for kind in ["Csa", "Hca"] { + for context in [true, false] { + let root = systems(); + let data = root.path().join("data/b200_sxm/vllm/0.24.0"); + table(&data.join("dsv4_paged_mqa_logits_module_perf.parquet")); + let load = + || PerfDatabase::load(root.path(), "b200_sxm", "vllm", "0.24.0").unwrap(); + let phase = if context { "context" } else { "generation" }; + let required = format!("dsv4_{}_{phase}_module_perf.parquet", kind.to_lowercase()); + let error = engine_readiness(load(), dsv4_op(kind, context, 1)).unwrap_err(); + assert!(error.to_string().contains(&required), "{error}"); + table(&data.join(required)); + engine_readiness(load(), dsv4_op(kind, context, 1)).unwrap(); + } + } + } + + #[test] + fn csa_context_parallel_requires_both_sparse_tables() { + let mqa = "dsv4_paged_mqa_logits_module_perf.parquet"; + let topk = "dsv4_csa_topk_calib_perf.parquet"; + for mode in [ + DatabaseMode::Silicon, + DatabaseMode::Hybrid, + DatabaseMode::Empirical, + DatabaseMode::Sol, + ] { + for present in [vec![], vec![mqa], vec![topk], vec![mqa, topk]] { + let root = systems(); + let data = root.path().join("data/b200_sxm/vllm/0.24.0"); + if mode != DatabaseMode::Sol { + table(&data.join("dsv4_csa_context_module_perf.parquet")); + table(&data.join("dsv4_hca_context_module_perf.parquet")); + } + for name in &present { + table(&data.join(name)); + } + let load = || { + PerfDatabase::load(root.path(), "b200_sxm", "vllm", "0.24.0") + .unwrap() + .with_mode(mode, TransferPolicy::ALL) + }; + assert_eq!( + engine_readiness(load(), dsv4_op("Csa", true, 2)).is_ok(), + present.len() == 2, + "{mode:?}, {present:?}" + ); + engine_readiness(load(), dsv4_op("Hca", true, 2)).unwrap(); + } + } + } + + #[test] + fn megamoe_readiness_uses_only_the_resolved_primary() { + let root = systems(); + let name = "dsv4_megamoe_module_perf.parquet"; + let family = root.path().join("data/b200_sxm/moe/vllm/0.24.0").join(name); + table(&family); + let op = || { + Op::Dsv4MegaMoe(Dsv4MegaMoeOp { + name: "megamoe".into(), + scale_factor: 1.0, + hidden_size: 4096, + inter_size: 2048, + topk: 8, + num_experts: 256, + moe_tp_size: 1, + moe_ep_size: 8, + quant_mode: crate::common::enums::MoeQuantMode::Fp8, + workload_distribution: "balanced".into(), + is_context: true, + source_policy: "primary".into(), + pre_dispatch: "none".into(), + num_fused_shared_experts: 0, + kernel_source: "fixture".into(), + kernel_dtype: "fp8".into(), + }) + }; + let db = PerfDatabase::load(root.path(), "b200_sxm", "vllm", "0.24.0").unwrap(); + assert_eq!(db.dsv4_megamoe.primary_path(), family); + engine_readiness(db, op()).unwrap(); + + let custom = root.path().join("custom.parquet"); + table(&custom); + for primary in [&custom, &root.path().join("missing.parquet")] { + // MegaMoE ignores row filters and additional sources at query time. + let sources = PerfDbSources::from([( + name.into(), + vec![ + PerfSource(primary.clone(), Some(vec!["excluded".into()])), + PerfSource(family.clone(), None), + ], + )]); + let db = PerfDatabase::load_with_sources( + root.path(), + "b200_sxm", + "vllm", + "0.24.0", + &sources, + ) + .unwrap(); + assert_eq!(db.dsv4_megamoe.primary_path(), primary); + assert_eq!(engine_readiness(db, op()).is_ok(), primary == &custom); + } + } + + #[test] + fn absent_tables_respect_modes_and_composite_fallbacks() { + let root = systems(); + let gemm = Op::Gemm(GemmOp::new("gemm", 16, 16, GemmQuantMode::Bfloat16)); + for mode in [ + DatabaseMode::Silicon, + DatabaseMode::Hybrid, + DatabaseMode::Empirical, + DatabaseMode::Sol, + ] { + let db = PerfDatabase::load(root.path(), "b200_sxm", "vllm", "0.24.0") + .unwrap() + .with_mode(mode, TransferPolicy::ALL); + assert_eq!( + validate(&db, [&gemm].into_iter()).is_ok(), + mode == DatabaseMode::Sol + ); + let fallback = Op::Fallback(FallbackOp::new( + "fallback", + gemm.clone(), + vec![Op::Embedding(EmbeddingOp::new( + "embedding", + 16, + 16, + GemmQuantMode::Bfloat16, + ))], + )); + validate(&db, [&fallback].into_iter()).unwrap(); + let overlap = Op::Overlap(OverlapOp::new( + "overlap", + vec![fallback], + vec![gemm.clone()], + )); + assert_eq!( + validate(&db, [&overlap].into_iter()).is_ok(), + mode == DatabaseMode::Sol + ); + } + } + + #[test] + fn table_presence_honors_source_filters_and_empty_tables() { + let root = systems(); + let data = root + .path() + .join("data/b200_sxm/vllm/0.24.0/gemm_perf.parquet"); + let gemm = Op::Gemm(GemmOp::new("gemm", 16, 16, GemmQuantMode::Bfloat16)); + for rows in [vec![], vec!["allowed"]] { + write_parquet(&data, &[Col::Str("kernel_source", rows.clone())]); + for filter in ["allowed", "excluded"] { + let sources = PerfDbSources::from([( + "gemm_perf.parquet".into(), + vec![PerfSource(data.clone(), Some(vec![filter.into()]))], + )]); + let db = PerfDatabase::load_with_sources( + root.path(), + "b200_sxm", + "vllm", + "0.24.0", + &sources, + ) + .unwrap(); + assert_eq!( + validate(&db, [&gemm].into_iter()).is_ok(), + !rows.is_empty() && filter == "allowed" + ); + } + } + } +} diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 8b0fdd9a4..91a1423f5 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -508,6 +508,22 @@ impl Engine { Ok(()) } + pub(crate) fn validate_forward_pass_readiness(&self) -> Result<(), AicError> { + let Some((prefill, decode)) = self.fpm_ops() else { + return super::readiness::validate( + &self.db, + self.context_ops.iter().chain(&self.generation_ops), + ); + }; + self.db + .fpm_forward + .select_cell(&prefill.match_identity, &prefill.model_path)?; + self.db + .fpm_forward + .select_cell(&decode.match_identity, &decode.model_path)?; + Ok(()) + } + /// Shared perf database handle. pub fn database(&self) -> &Arc { &self.db diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs new file mode 100644 index 000000000..37822037e --- /dev/null +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -0,0 +1,413 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Canonical construction contract for [`super::ForwardPassPerfModel`]. + +use std::path::PathBuf; + +use serde::{Deserialize, Serialize}; + +use super::{EstimatorConfig, ForwardPassWorkerType}; +use crate::common::enums::{DatabaseMode, TransferPolicy}; +use crate::{AicError, BackendKind}; + +const fn one() -> u32 { + 1 +} + +/// Estimator selection. Auto searches the fixed priority list independently +/// of the fallback policy used for an explicitly requested estimator. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum EstimationMode { + #[default] + Auto, + OpLevel, + #[serde(alias = "fpm")] + FpmInterpolation, + FpmRegression, +} + +impl EstimationMode { + pub(crate) fn native_name(self) -> Option<&'static str> { + match self { + Self::OpLevel => Some("op_level"), + Self::FpmInterpolation => Some("fpm"), + Self::Auto | Self::FpmRegression => None, + } + } +} + +/// Additional fallback after an explicitly selected estimator is unavailable. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ForwardPassFallbackPolicy { + #[default] + #[serde(alias = "error")] + Deny, + Allow, + /// Legacy saved requests permitted only a direct regression fallback. + /// Retaining this value preserves their behavior when replayed. + #[serde(rename = "regression")] + LegacyRegression, +} + +/// Target-verification cost; acceptance and scheduler progress belong to replay. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde( + tag = "kind", + content = "params", + rename_all = "snake_case", + deny_unknown_fields +)] +pub enum ForwardPassSpeculationConfig { + Ngram { num_speculative_tokens: u32 }, +} + +impl ForwardPassSpeculationConfig { + pub fn num_speculative_tokens(&self) -> u32 { + match self { + Self::Ngram { + num_speculative_tokens, + } => *num_speculative_tokens, + } + } +} + +/// Immutable model identity and selection policy for a forward-pass estimator. +/// +/// This is the one public construction schema shared by Rust, Python, Replay, +/// Sweeper, and Planner. Estimator-specific tuning controls live in `estimator_config`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ForwardPassPerfModelConfig { + pub model: String, + pub system: String, + pub backend: BackendKind, + pub worker_type: ForwardPassWorkerType, + #[serde(default)] + pub backend_version: Option, + + #[serde(default = "one", alias = "tp_size")] + pub tp: u32, + #[serde(default = "one", alias = "pp_size")] + pub pp: u32, + #[serde(default = "one", alias = "attention_dp_size")] + pub attention_dp: u32, + #[serde(default)] + pub moe_tp_size: Option, + #[serde(default)] + pub moe_ep_size: Option, + + #[serde(default, alias = "gemm_dtype")] + pub gemm_quant_mode: Option, + #[serde(default, alias = "moe_dtype")] + pub moe_quant_mode: Option, + #[serde(default, alias = "fmha_dtype")] + pub fmha_quant_mode: Option, + #[serde(default, alias = "kv_cache_dtype")] + pub kvcache_quant_mode: Option, + #[serde(default, alias = "comm_dtype")] + pub comm_quant_mode: Option, + + #[serde(default)] + pub nextn: u32, + #[serde(default)] + pub speculation: Option, + #[serde(default)] + pub kv_block_size: Option, + #[serde(default)] + #[serde(alias = "forward_model")] + pub estimation_mode: EstimationMode, + #[serde(default)] + pub database_mode: DatabaseMode, + /// Explicit transfer-kind tokens. `None` means the core default (all). + #[serde(default)] + pub transfer_policy: Option>, + /// Ordered request-scoped systems roots. Empty uses normal package/env discovery. + #[serde(default)] + pub systems_paths: Vec, + #[serde(default)] + pub fallback_policy: ForwardPassFallbackPolicy, + #[serde(default)] + pub estimator_config: EstimatorConfig, + #[serde(default)] + pub attention_backend: Option, + #[serde(default)] + pub enable_shared_layer: Option, + #[serde(default)] + pub strict_provenance: bool, +} + +impl ForwardPassPerfModelConfig { + pub fn new( + model: impl Into, + system: impl Into, + backend: BackendKind, + worker_type: ForwardPassWorkerType, + ) -> Self { + Self { + model: model.into(), + system: system.into(), + backend, + worker_type, + backend_version: None, + tp: 1, + pp: 1, + attention_dp: 1, + moe_tp_size: None, + moe_ep_size: None, + gemm_quant_mode: None, + moe_quant_mode: None, + fmha_quant_mode: None, + kvcache_quant_mode: None, + comm_quant_mode: None, + nextn: 0, + speculation: None, + kv_block_size: None, + estimation_mode: EstimationMode::Auto, + database_mode: DatabaseMode::default(), + transfer_policy: None, + systems_paths: Vec::new(), + fallback_policy: ForwardPassFallbackPolicy::Deny, + estimator_config: EstimatorConfig::default(), + attention_backend: None, + enable_shared_layer: None, + strict_provenance: false, + } + } + + pub fn candidate_modes(&self) -> Vec { + use EstimationMode::*; + let priority = [OpLevel, FpmInterpolation, FpmRegression]; + if self.estimation_mode == Auto { + return priority.to_vec(); + } + let mut candidates = vec![self.estimation_mode]; + match self.fallback_policy { + ForwardPassFallbackPolicy::Allow => candidates.extend( + priority + .into_iter() + .filter(|mode| *mode != self.estimation_mode), + ), + ForwardPassFallbackPolicy::LegacyRegression + if self.estimation_mode != FpmRegression => + { + candidates.push(FpmRegression) + } + _ => {} + } + candidates + } + + pub(crate) fn validate(&self) -> Result<(), AicError> { + if self.model.trim().is_empty() { + return Err(invalid_config("model cannot be empty")); + } + if self.system.trim().is_empty() { + return Err(invalid_config("system cannot be empty")); + } + if self + .backend_version + .as_ref() + .is_some_and(|value| value.trim().is_empty()) + { + return Err(invalid_config("backend_version cannot be empty")); + } + if self.tp == 0 || self.pp == 0 || self.attention_dp == 0 { + return Err(invalid_config("tp, pp, and attention_dp must be positive")); + } + if self.moe_tp_size.is_some() != self.moe_ep_size.is_some() { + return Err(invalid_config( + "moe_tp_size and moe_ep_size must be configured together", + )); + } + if let (Some(moe_tp), Some(moe_ep)) = (self.moe_tp_size, self.moe_ep_size) { + if moe_tp == 0 || moe_ep == 0 { + return Err(invalid_config( + "moe_tp_size and moe_ep_size must be positive", + )); + } + if u64::from(self.tp) * u64::from(self.attention_dp) + != u64::from(moe_tp) * u64::from(moe_ep) + { + return Err(invalid_config( + "topology requires tp * attention_dp == moe_tp_size * moe_ep_size", + )); + } + } + if self.nextn > 5 { + return Err(invalid_config("nextn must be in 0..=5")); + } + if let Some(speculation) = &self.speculation { + if self.nextn != 0 { + return Err(invalid_config( + "ngram speculation cannot be combined with nextn", + )); + } + if self.backend != BackendKind::Vllm { + return Err(invalid_config("ngram speculation requires backend=vllm")); + } + if !(1..=5).contains(&speculation.num_speculative_tokens()) { + return Err(invalid_config( + "ngram num_speculative_tokens must be in 1..=5", + )); + } + if !matches!( + self.estimation_mode, + EstimationMode::Auto | EstimationMode::OpLevel + ) { + return Err(invalid_config("ngram speculation requires op_level timing")); + } + } + if self.estimation_mode == EstimationMode::FpmInterpolation && self.nextn != 0 { + return Err(invalid_config( + "estimation_mode='fpm_interpolation' does not support MTP speculative decoding", + )); + } + TransferPolicy::from_wire(self.transfer_policy.as_deref()).map_err(invalid_config)?; + for root in &self.systems_paths { + if root.to_str().is_none() { + return Err(invalid_config("systems_paths entries must be valid UTF-8")); + } + if !root.is_dir() { + return Err(invalid_config(format!( + "systems_paths entry is not an existing directory: {}", + root.display() + ))); + } + } + self.estimator_config.validate() + } +} + +fn invalid_config(message: impl Into) -> AicError { + AicError::InvalidEngineConfig(format!( + "invalid forward pass perf model config: {}", + message.into() + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config(value: serde_json::Value) -> ForwardPassPerfModelConfig { + let mut base = serde_json::json!({ + "model": "Qwen/Qwen3-32B", "system": "h200_sxm", + "backend": "vllm", "worker_type": "aggregated" + }); + base.as_object_mut() + .unwrap() + .extend(value.as_object().unwrap().clone()); + serde_json::from_value(base).unwrap() + } + + #[test] + fn auto_selects_all_modes_even_when_fallback_is_denied() { + let cfg = config(serde_json::json!({})); + assert_eq!(cfg.fallback_policy, ForwardPassFallbackPolicy::Deny); + assert_eq!( + cfg.candidate_modes(), + vec![ + EstimationMode::OpLevel, + EstimationMode::FpmInterpolation, + EstimationMode::FpmRegression + ] + ); + cfg.validate().unwrap(); + } + + #[test] + fn prompt_lookup_is_distinct_from_mtp_and_validated_before_selection() { + let mut cfg = config(serde_json::json!({ + "speculation": {"kind": "ngram", "params": {"num_speculative_tokens": 2}} + })); + cfg.validate().unwrap(); + assert_eq!(cfg.nextn, 0); + assert_eq!( + serde_json::from_str::( + &serde_json::to_string(&cfg).unwrap() + ) + .unwrap(), + cfg + ); + for mode in [ + EstimationMode::FpmInterpolation, + EstimationMode::FpmRegression, + ] { + cfg.estimation_mode = mode; + assert!(cfg.validate().is_err()); + } + cfg.estimation_mode = EstimationMode::Auto; + cfg.nextn = 2; + assert!(cfg.validate().is_err()); + cfg.nextn = 0; + cfg.backend = BackendKind::Sglang; + assert!(cfg.validate().is_err()); + cfg.backend = BackendKind::Vllm; + for depth in [0, 6] { + cfg.speculation = Some(ForwardPassSpeculationConfig::Ngram { + num_speculative_tokens: depth, + }); + assert!(cfg.validate().is_err()); + } + } + + #[test] + fn explicit_selection_and_legacy_fallback_keep_distinct_semantics() { + let cfg = config(serde_json::json!({"estimation_mode": "fpm_interpolation"})); + assert_eq!( + cfg.candidate_modes(), + vec![EstimationMode::FpmInterpolation] + ); + let cfg = config( + serde_json::json!({"estimation_mode": "fpm_interpolation", "fallback_policy": "allow"}), + ); + assert_eq!( + cfg.candidate_modes(), + vec![ + EstimationMode::FpmInterpolation, + EstimationMode::OpLevel, + EstimationMode::FpmRegression + ] + ); + let cfg = config( + serde_json::json!({"forward_model": "op_level", "fallback_policy": "regression"}), + ); + assert_eq!( + cfg.candidate_modes(), + vec![EstimationMode::OpLevel, EstimationMode::FpmRegression] + ); + } + + #[test] + fn role_and_typed_estimator_controls_are_validated() { + assert!( + serde_json::from_value::( + serde_json::json!({"model":"m","system":"s","backend":"vllm"}) + ) + .is_err() + ); + let cfg = config(serde_json::json!({"moe_tp_size": 1, "moe_ep_size": 2})); + assert!(cfg.validate().is_err()); + let cfg = config( + serde_json::json!({"estimator_config": {"fpm_regression": {"sampling": {"bins_per_axis": [4,16], "max_observations": 128}}, "correction": {"sampling": {"max_observations": 32}}}}), + ); + cfg.validate().unwrap(); + assert_eq!( + cfg.estimator_config.regression_options().max_observations, + 128 + ); + assert_eq!( + cfg.estimator_config.correction_options().max_observations, + 32 + ); + assert!( + serde_json::from_value::( + serde_json::json!({"fpm_regression": {"unknown": 1}}) + ) + .is_err() + ); + } +} diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs new file mode 100644 index 000000000..a856f9751 --- /dev/null +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -0,0 +1,245 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Typed estimator-specific controls. Defaults preserve the existing algorithms. + +use serde::{Deserialize, Serialize}; + +use super::options::{ForwardPassPerfOptions, validate_options}; +use crate::AicError; + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct EstimatorConfig { + pub features: RegressionFeatureWeights, + pub op_level: OpLevelConfig, + pub fpm_interpolation: FpmInterpolationConfig, + pub fpm_regression: FpmRegressionConfig, + pub correction: CorrectionConfig, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OpLevelConfig {} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FpmInterpolationConfig {} + +/// These weights currently affect regression. Native correction keeps its +/// established workload coordinates until a replacement is accuracy-qualified. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct RegressionFeatureWeights { + #[serde(with = "super::options::regression_weight_serde")] + pub attention_kv_weight: f64, + #[serde(with = "super::options::regression_weight_serde")] + pub prefill_attention_pair_weight: f64, + #[serde(with = "super::options::regression_weight_serde")] + pub ffn_token_weight: f64, +} + +impl Default for RegressionFeatureWeights { + fn default() -> Self { + Self { + attention_kv_weight: 1.0, + prefill_attention_pair_weight: 1.0, + ffn_token_weight: 1.0, + } + } +} + +/// Per-store sample retention. For the legacy one-dimensional correction grid, +/// the product of the two axis counts supplies its number of bins. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct SamplingConfig { + pub bins_per_axis: [usize; 2], + pub max_observations: usize, +} + +impl Default for SamplingConfig { + fn default() -> Self { + Self { + bins_per_axis: [4, 4], + max_observations: 64, + } + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegressionFitKind { + #[default] + StandardizedNnls, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct RegressionFitConfig { + pub kind: RegressionFitKind, + pub singular_ridge_scale: f64, +} + +impl Default for RegressionFitConfig { + fn default() -> Self { + Self { + kind: RegressionFitKind::StandardizedNnls, + singular_ridge_scale: 1e-9, + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct FpmRegressionConfig { + pub sampling: SamplingConfig, + pub min_observations: usize, + pub fit: RegressionFitConfig, +} + +impl Default for FpmRegressionConfig { + fn default() -> Self { + Self { + sampling: SamplingConfig::default(), + min_observations: 5, + fit: RegressionFitConfig::default(), + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct CorrectionFactorBounds { + pub min: Option, + pub max: Option, +} + +impl Default for CorrectionFactorBounds { + fn default() -> Self { + Self { + min: Some(0.5), + max: Some(2.0), + } + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CorrectionFeatureSpace { + #[default] + LegacyWorkload, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct CorrectionConfig { + pub enabled: bool, + pub feature_space: CorrectionFeatureSpace, + pub sampling: SamplingConfig, + pub min_observations: usize, + pub factor_bounds: CorrectionFactorBounds, + pub max_num_tokens: u32, + pub max_batch_size: u32, + pub max_kv_tokens: u32, +} + +impl Default for CorrectionConfig { + fn default() -> Self { + Self { + enabled: true, + feature_space: CorrectionFeatureSpace::LegacyWorkload, + sampling: SamplingConfig::default(), + min_observations: 5, + factor_bounds: CorrectionFactorBounds::default(), + max_num_tokens: 8192, + max_batch_size: 512, + max_kv_tokens: 2_000_000, + } + } +} + +impl EstimatorConfig { + pub(crate) fn correction_options(&self) -> ForwardPassPerfOptions { + let config = &self.correction; + ForwardPassPerfOptions { + bucket_shape: Some(config.sampling.bins_per_axis), + max_observations: config.sampling.max_observations, + min_observations: config.min_observations, + min_faster_correction_factor: config.factor_bounds.min, + max_slower_correction_factor: config.factor_bounds.max, + max_num_tokens: config.max_num_tokens, + max_batch_size: config.max_batch_size, + max_kv_tokens: config.max_kv_tokens, + ..ForwardPassPerfOptions::default() + } + } + + pub(crate) fn regression_options(&self) -> ForwardPassPerfOptions { + let config = &self.fpm_regression; + ForwardPassPerfOptions { + bucket_shape: Some(config.sampling.bins_per_axis), + max_observations: config.sampling.max_observations, + min_observations: config.min_observations, + regression_attention_kv_weight: self.features.attention_kv_weight, + regression_prefill_attention_pair_weight: self.features.prefill_attention_pair_weight, + regression_ffn_token_weight: self.features.ffn_token_weight, + regression_ridge_scale: config.fit.singular_ridge_scale, + ..ForwardPassPerfOptions::default() + } + } + + pub(crate) fn validate(&self) -> Result<(), AicError> { + let ridge = self.fpm_regression.fit.singular_ridge_scale; + if !ridge.is_finite() || ridge < 0.0 { + return Err(AicError::InvalidEngineConfig("estimator_config.fpm_regression.fit.singular_ridge_scale must be finite and nonnegative".into())); + } + validate_options(&self.correction_options()).map_err(|error| { + AicError::InvalidEngineConfig(format!("estimator_config.correction: {error}")) + })?; + validate_options(&self.regression_options()).map_err(|error| { + AicError::InvalidEngineConfig(format!("estimator_config.fpm_regression: {error}")) + }) + } + + /// Convert legacy flat tuning options without changing their behavior. + pub fn from_legacy(options: ForwardPassPerfOptions) -> Result { + validate_options(&options)?; + let bins = options.bucket_shape.unwrap_or_else(|| { + let n = super::samples::integer_sqrt(options.bucket_count); + [n, n] + }); + let sampling = SamplingConfig { + bins_per_axis: bins, + max_observations: options.max_observations, + }; + Ok(Self { + features: RegressionFeatureWeights { + attention_kv_weight: options.regression_attention_kv_weight, + prefill_attention_pair_weight: options.regression_prefill_attention_pair_weight, + ffn_token_weight: options.regression_ffn_token_weight, + }, + fpm_regression: FpmRegressionConfig { + sampling: sampling.clone(), + min_observations: options.min_observations, + fit: RegressionFitConfig { + singular_ridge_scale: options.regression_ridge_scale, + ..RegressionFitConfig::default() + }, + }, + correction: CorrectionConfig { + sampling, + min_observations: options.min_observations, + factor_bounds: CorrectionFactorBounds { + min: options.min_faster_correction_factor, + max: options.max_slower_correction_factor, + }, + max_num_tokens: options.max_num_tokens, + max_batch_size: options.max_batch_size, + max_kv_tokens: options.max_kv_tokens, + ..CorrectionConfig::default() + }, + ..Self::default() + }) + } +} diff --git a/crates/core/src/perfmodel/fpm/mod.rs b/crates/core/src/perfmodel/fpm/mod.rs index 77567e672..65144f25b 100644 --- a/crates/core/src/perfmodel/fpm/mod.rs +++ b/crates/core/src/perfmodel/fpm/mod.rs @@ -8,9 +8,9 @@ //! [`crate::perfmodel::engine::Engine::forward_pass_time_ms`]. The online correction / //! regression / diagnostics / readiness logic is engine-agnostic. //! -//! Constructing a native model crosses into Python exactly once to compile the -//! model into an [`crate::perfmodel::engine::spec::EngineSpec`] (mirroring -//! [`crate::perfmodel::AicEngineBuilder`]); after that the hot path +//! Native candidates currently use Python to compile the model into an +//! [`crate::perfmodel::engine::spec::EngineSpec`]. Selection may try several +//! candidates or data roots; after construction the hot path //! (`estimate_forward_pass_time_ms` / `tune_with_fpms`) is pure Rust over the //! `Engine` with no Python re-entry. //! @@ -22,7 +22,9 @@ //! - [`samples`]: shared bucketed-sample infrastructure. //! - [`options`]: tuning controls. +mod config; mod correction; +mod estimator; mod metrics; mod model; mod options; @@ -32,11 +34,16 @@ mod samples; #[cfg(test)] mod tests; +pub use config::{ + EstimationMode, ForwardPassFallbackPolicy, ForwardPassPerfModelConfig, + ForwardPassSpeculationConfig, +}; +pub use estimator::*; pub(crate) use metrics::validate_forward_pass_metrics; pub use metrics::{FPM_VERSION, ForwardPassMetrics, QueuedRequestMetrics, ScheduledRequestMetrics}; pub use model::{ - ForwardPassPerfDiagnostics, ForwardPassPerfModel, ForwardPassPerfReadiness, - ForwardPassPerfSource, ForwardPassRegressionStoreDiagnostics, + ForwardPassPerfDiagnostics, ForwardPassPerfModel, ForwardPassPerfProvenance, + ForwardPassPerfReadiness, ForwardPassPerfSource, ForwardPassRegressionStoreDiagnostics, ForwardPassRegressionWorkloadKind, ForwardPassWorkerType, }; pub use options::ForwardPassPerfOptions; diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index d18db7cb6..b99d4525b 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -8,21 +8,17 @@ //! through [`crate::perfmodel::engine::Engine::forward_pass_time_ms`]. The online //! correction / regression / diagnostics / readiness logic is engine-agnostic. -#[cfg(feature = "python")] -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::sync::Arc; use serde::{Deserialize, Serialize}; -#[cfg(feature = "python")] -use crate::perfmodel::EngineConfig; use crate::perfmodel::engine::Engine; use crate::{AicError, ForwardPassMetrics}; +use super::config::{EstimationMode, ForwardPassFallbackPolicy, ForwardPassPerfModelConfig}; use super::correction::CorrectionBuckets; use super::metrics::validate_forward_pass_metrics; -#[cfg(feature = "python")] -use super::options::validate_options; use super::options::{ForwardPassPerfOptions, validate_regression_options}; use super::regression::BucketedRegression; use super::samples::{AxisRange, StoreStats, WithOptions}; @@ -47,6 +43,20 @@ pub struct ForwardPassPerfDiagnostics { pub correction_ready_buckets: usize, /// Fallback reason when `best_available` had to use regression instead of native AIC. pub last_warning: Option, + /// Exact immutable configuration and selected systems root used at construction. + pub provenance: Option, +} + +/// Resolved construction provenance pinned by Replay, Sweeper, and Planner. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +pub struct ForwardPassPerfProvenance { + /// Canonical config with the exact backend version and explicit transfer policy. + pub config: ForwardPassPerfModelConfig, + /// Root that supplied the native engine, or `None` for regression fallback. + pub selected_systems_root: Option, + pub requested_estimation_mode: EstimationMode, + pub selected_estimation_mode: EstimationMode, + pub selection_failures: Vec, } /// Prediction backend currently used by `ForwardPassPerfModel`. @@ -232,6 +242,8 @@ pub struct ForwardPassPerfModel { mode: ForwardPassPerfMode, options: ForwardPassPerfOptions, last_warning: Option, + provenance: Option, + is_correction_enabled: bool, } #[derive(Clone, Debug)] @@ -250,48 +262,10 @@ enum ForwardPassPerfMode { } impl ForwardPassPerfModel { - /// API: - /// `ForwardPassPerfModel::from_native(config, options) -> Result` - /// - /// Description: create a strict native AIC forward-pass model. - /// - /// Compiles `config` into an [`Engine`] by crossing into Python once - /// (mirroring [`crate::AicEngineBuilder`]): `compile_engine` walks the model - /// and returns bincoded spec bytes, then [`Engine::from_spec_bytes`] loads - /// the matching perf database. This constructor fails if `config` cannot be - /// compiled. Use `best_available` when unsupported native configs should - /// fall back to the learned regression model. - #[cfg(feature = "python")] - pub fn from_native( - config: EngineConfig, - options: ForwardPassPerfOptions, - ) -> Result { - validate_options(&options)?; - let engine = build_engine_via_python(&config, None)?; - Ok(Self::from_engine(Arc::new(engine), options)) - } - - /// API: - /// `ForwardPassPerfModel::from_native_with_roots(config, options, systems_root) -> Result` - /// - /// Description: create a strict native AIC forward-pass model with an - /// explicit `systems/` data root (forwarded to `compile_engine` and used to - /// load the perf database). Same tuning and failure behavior as - /// `from_native`. - #[cfg(feature = "python")] - pub fn from_native_with_roots( - config: EngineConfig, - options: ForwardPassPerfOptions, - systems_root: impl AsRef, - ) -> Result { - validate_options(&options)?; - let engine = build_engine_via_python(&config, Some(systems_root.as_ref()))?; - Ok(Self::from_engine(Arc::new(engine), options)) - } - /// Internal: build a native model directly from an already-compiled - /// [`Engine`]. Holds the actual native-mode logic; the public `from_native` - /// constructors compile the `Engine` (crossing into Python) and call this. + /// [`Engine`]. Holds the actual native-mode logic; the public + /// [`Self::best_available`] constructor compiles the `Engine` (crossing + /// into Python) and calls this. /// Used by the `#[cfg(test)]` suite to construct a native model from a /// hand-built fixture `Engine` without Python. pub(crate) fn from_engine(engine: Arc, options: ForwardPassPerfOptions) -> Self { @@ -302,6 +276,8 @@ impl ForwardPassPerfModel { }, options, last_warning: None, + provenance: None, + is_correction_enabled: true, } } @@ -315,7 +291,7 @@ impl ForwardPassPerfModel { /// selected logical store has a usable fit from at least /// `options.min_observations` tuning samples. /// Correction factor getters always return `None` in this mode. - pub fn from_regression( + pub(crate) fn from_regression( worker_type: ForwardPassWorkerType, options: ForwardPassPerfOptions, ) -> Result { @@ -327,66 +303,77 @@ impl ForwardPassPerfModel { }, options, last_warning: None, + provenance: None, + is_correction_enabled: true, }) } - /// API: - /// `ForwardPassPerfModel::best_available(config, worker_type, options) -> Result` - /// - /// Description: create a native model when possible, otherwise fall back to - /// regression. - /// - /// Fallback reason is preserved in `diagnostics().last_warning`. The - /// A successful native construction preserves native workload inference - /// and ignores `worker_type` and regression-only weights. A fallback is - /// bound to `worker_type` and validates those weights when it is created. - #[cfg(feature = "python")] - pub fn best_available( - config: EngineConfig, - worker_type: ForwardPassWorkerType, - options: ForwardPassPerfOptions, - ) -> Result { - match Self::from_native(config, options.clone()) { - Ok(model) => Ok(model), - Err(err) if can_fallback_to_regression(&err) => { - Self::regression_with_warning(worker_type, options, err) + /// Construct and pin one estimator. Auto searches all modes even when + /// fallback is denied; explicit modes use the requested fallback policy. + /// A regression model may be constructed before it has enough observations. + pub fn best_available(config: ForwardPassPerfModelConfig) -> Result { + config.validate()?; + let requested_estimation_mode = config.estimation_mode; + let mut failures = Vec::new(); + let mut last_error = None; + for mode in config.candidate_modes() { + if config.speculation.is_some() && mode != EstimationMode::OpLevel { + let error = + AicError::UnsupportedModel("ngram speculation requires op_level timing".into()); + failures.push(format!("{mode:?}: {error}")); + last_error = Some(error); + continue; } - Err(err) => Err(err), - } - } - - /// API: - /// `ForwardPassPerfModel::best_available_with_roots(config, worker_type, options, systems_root) -> Result` - /// - /// Description: create a `best_available` model with an explicit `systems/` - /// data root. - #[cfg(feature = "python")] - pub fn best_available_with_roots( - config: EngineConfig, - worker_type: ForwardPassWorkerType, - options: ForwardPassPerfOptions, - systems_root: impl AsRef, - ) -> Result { - match Self::from_native_with_roots(config, options.clone(), systems_root) { - Ok(model) => Ok(model), - Err(err) if can_fallback_to_regression(&err) => { - Self::regression_with_warning(worker_type, options, err) + if mode == EstimationMode::FpmRegression { + let options = config.estimator_config.regression_options(); + let mut model = Self::from_regression(config.worker_type, options)?; + let mut resolved = config.clone(); + resolved.estimation_mode = mode; + resolved.fallback_policy = ForwardPassFallbackPolicy::Deny; + model.last_warning = (!failures.is_empty()).then(|| failures.join("; ")); + model.provenance = Some(ForwardPassPerfProvenance { + config: resolved, + selected_systems_root: None, + requested_estimation_mode, + selected_estimation_mode: mode, + selection_failures: failures, + }); + return Ok(model); + } + let mut candidate = config.clone(); + candidate.estimation_mode = mode; + match build_native_candidate(&candidate) { + Ok((engine, root)) => { + candidate.backend_version = Some(engine.database().version.clone()); + candidate.database_mode = engine.database().database_mode; + candidate.transfer_policy = + Some(transfer_policy_tokens(engine.database().transfer_policy)); + candidate.systems_paths = vec![root.clone()]; + candidate.fallback_policy = ForwardPassFallbackPolicy::Deny; + let mut model = Self::from_engine( + Arc::new(engine), + candidate.estimator_config.correction_options(), + ); + model.is_correction_enabled = candidate.estimator_config.correction.enabled; + model.last_warning = (!failures.is_empty()).then(|| failures.join("; ")); + model.provenance = Some(ForwardPassPerfProvenance { + config: candidate, + selected_systems_root: Some(root), + requested_estimation_mode, + selected_estimation_mode: mode, + selection_failures: failures, + }); + return Ok(model); + } + Err(err) if can_fallback_to_regression(&err) => { + failures.push(format!("{mode:?}: {err}")); + last_error = Some(err); + } + Err(err) => return Err(err), } - Err(err) => Err(err), } - } - - #[cfg(feature = "python")] - fn regression_with_warning( - worker_type: ForwardPassWorkerType, - options: ForwardPassPerfOptions, - err: AicError, - ) -> Result { - let mut model = Self::from_regression(worker_type, options)?; - model.last_warning = Some(format!( - "native forward-pass estimator unavailable; using fallback regression: {err}" - )); - Ok(model) + Err(last_error + .unwrap_or_else(|| AicError::UnsupportedModel("no estimator candidates".into()))) } /// API: @@ -424,9 +411,13 @@ impl ForwardPassPerfModel { }; let native = engine.forward_pass_time_ms(metrics_by_rank)?; let corrected = native - * corrections - .store(feature.workload_kind) - .correction_factor_for(&feature.x); + * if self.is_correction_enabled { + corrections + .store(feature.workload_kind) + .correction_factor_for(&feature.x) + } else { + 1.0 + }; Ok(Some(corrected)) } ForwardPassPerfMode::Regression { @@ -476,13 +467,22 @@ impl ForwardPassPerfModel { &mut self, iterations: &[Vec], ) -> Result<(), AicError> { - let Self { mode, options, .. } = self; + let Self { + mode, + options, + is_correction_enabled, + .. + } = self; for metrics_by_rank in iterations { match mode { ForwardPassPerfMode::Native { engine, corrections, } => { + if !*is_correction_enabled { + let _ = IterationFeatures::from_metrics(metrics_by_rank)?; + continue; + } let Some(observation) = IterationObservation::from_metrics(metrics_by_rank)? else { continue; @@ -532,6 +532,7 @@ impl ForwardPassPerfModel { retained_observations: corrections.observation_count(), correction_ready_buckets: ready_buckets, last_warning: self.last_warning.clone(), + provenance: self.provenance.clone(), } } ForwardPassPerfMode::Regression { regression, .. } => { @@ -548,6 +549,7 @@ impl ForwardPassPerfModel { retained_observations: regression.observation_count(), correction_ready_buckets: 0, last_warning: self.last_warning.clone(), + provenance: self.provenance.clone(), } } } @@ -619,6 +621,18 @@ impl ForwardPassPerfModel { &self.options } + pub(crate) fn native_engine(&self) -> Option> { + match &self.mode { + ForwardPassPerfMode::Native { engine, .. } => Some(Arc::clone(engine)), + ForwardPassPerfMode::Regression { .. } => None, + } + } + + /// Exact immutable construction identity and selected systems root. + pub fn provenance(&self) -> Option<&ForwardPassPerfProvenance> { + self.provenance.as_ref() + } + fn correction_factors(&self) -> Vec { match &self.mode { ForwardPassPerfMode::Native { corrections, .. } => corrections.correction_factors(), @@ -627,37 +641,72 @@ impl ForwardPassPerfModel { } } -/// Build a compiled [`Engine`] from an [`EngineConfig`] by crossing into Python -/// once to run `aiconfigurator.sdk.engine.compile_engine`, then loading the -/// matching perf database via [`Engine::from_spec_bytes`]. This is the internal -/// `EngineConfig` counterpart to [`crate::AicEngineBuilder`] and maps its -/// modular fields onto the flat `compile_engine` kwargs. -/// -/// `systems_root` overrides the bundled `systems/` dir for BOTH the -/// `compile_engine` call (`systems_path` kwarg) and the Rust-side perf-DB load. +#[cfg(feature = "python")] +fn build_native_candidate( + config: &ForwardPassPerfModelConfig, +) -> Result<(Engine, PathBuf), AicError> { + if config.estimation_mode == EstimationMode::FpmInterpolation && config.nextn != 0 { + return Err(AicError::UnsupportedModel( + "FPM interpolation does not support MTP".into(), + )); + } + let mut last_error = None; + for root in resolve_systems_roots(config)? { + match build_engine_via_python(config, &root).and_then(|engine| { + engine.validate_forward_pass_readiness()?; + Ok(engine) + }) { + Ok(engine) => return Ok((engine, root)), + Err(err) if can_fallback_to_regression(&err) => last_error = Some(err), + Err(err) => return Err(err), + } + } + Err(last_error.unwrap_or_else(|| AicError::DataRoot("no systems roots".into()))) +} + +#[cfg(not(feature = "python"))] +fn build_native_candidate( + _config: &ForwardPassPerfModelConfig, +) -> Result<(Engine, PathBuf), AicError> { + Err(AicError::UnsupportedModel( + "native engine construction requires the python feature".into(), + )) +} + #[cfg(feature = "python")] fn build_engine_via_python( - config: &EngineConfig, - systems_root: Option<&Path>, + config: &ForwardPassPerfModelConfig, + systems_root: &std::path::Path, ) -> Result { - // `compile_engine`'s `systems_path` kwarg: explicit override -> config's - // own `systems_path` -> None (Python resolves it). - let systems_path: Option = systems_root - .map(PathBuf::from) - .or_else(|| config.systems_path.clone()); - // A non-UTF-8 override path cannot be passed through the Python kwarg; fail - // loudly rather than silently dropping the override. - let systems_path_str = match systems_path.as_ref() { - Some(p) => Some(p.to_str().ok_or_else(|| { - AicError::InvalidEngineConfig(format!( - "systems_path is not valid UTF-8: {}", - p.display() - )) - })?), - None => None, - }; - - crate::py::compile_engine_to_engine(config, systems_path_str) + let systems_path = systems_root.to_str().ok_or_else(|| { + AicError::InvalidEngineConfig(format!( + "systems_path is not valid UTF-8: {}", + systems_root.display() + )) + })?; + crate::py::compile_forward_pass_model_to_engine(config, systems_path) +} + +#[cfg(feature = "python")] +fn resolve_systems_roots(config: &ForwardPassPerfModelConfig) -> Result, AicError> { + if !config.systems_paths.is_empty() { + return Ok(config.systems_paths.clone()); + } + crate::py::resolve_forward_pass_systems_roots() +} + +fn transfer_policy_tokens(policy: crate::common::enums::TransferPolicy) -> Vec { + use crate::common::enums::TransferKind; + [ + (TransferKind::XShape, "xshape"), + (TransferKind::XQuant, "xquant"), + (TransferKind::XProfile, "xprofile"), + (TransferKind::XOp, "xop"), + ] + .into_iter() + .filter(|(kind, _)| policy.contains(*kind)) + .map(|(_, token)| token.to_string()) + .collect() } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -981,7 +1030,6 @@ impl WorkloadStores { /// used for hard caller/config errors (e.g. a non-UTF-8 `systems_path`, invalid /// FPM options, a malformed spec). Those must surface rather than silently /// degrade `best_available` to regression mode. -#[cfg(feature = "python")] fn can_fallback_to_regression(err: &AicError) -> bool { matches!( err, diff --git a/crates/core/src/perfmodel/fpm/options.rs b/crates/core/src/perfmodel/fpm/options.rs index d4a813b13..e80f98be2 100644 --- a/crates/core/src/perfmodel/fpm/options.rs +++ b/crates/core/src/perfmodel/fpm/options.rs @@ -29,6 +29,7 @@ pub(crate) const DEFAULT_REGRESSION_FFN_TOKEN_WEIGHT: f64 = 1.0; /// per inferred workload kind; regression retains one set per logical store /// (one for dedicated roles, four for Aggregated). #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] pub struct ForwardPassPerfOptions { /// Maximum retained observations across the feature-space buckets in each /// logical store. The cap applies per inferred workload kind for Native @@ -62,6 +63,10 @@ pub struct ForwardPassPerfOptions { /// Target bucket count for workload-specific sample retirement and correction lookup. #[serde(default = "default_bucket_count")] pub bucket_count: usize, + #[serde(default)] + pub bucket_shape: Option<[usize; 2]>, + #[serde(default = "default_regression_ridge_scale")] + pub regression_ridge_scale: f64, /// Upper bound for the `sum_prefill_tokens` correction axis. /// /// Used by prefill and mixed/agg workload kinds. The lower bound is always `0`. @@ -111,7 +116,7 @@ pub struct ForwardPassPerfOptions { /// through the options JSON as these exact strings so Native construction can /// continue to ignore regression-only weights while Regression construction /// still reports its existing field-specific validation error. -mod regression_weight_serde { +pub(crate) mod regression_weight_serde { use std::fmt; use serde::{ @@ -123,7 +128,7 @@ mod regression_weight_serde { const POSITIVE_INFINITY: &str = "Infinity"; const NEGATIVE_INFINITY: &str = "-Infinity"; - pub(super) fn serialize(value: &f64, serializer: S) -> Result + pub(crate) fn serialize(value: &f64, serializer: S) -> Result where S: Serializer, { @@ -138,7 +143,7 @@ mod regression_weight_serde { } } - pub(super) fn deserialize<'de, D>(deserializer: D) -> Result + pub(crate) fn deserialize<'de, D>(deserializer: D) -> Result where D: Deserializer<'de>, { @@ -209,6 +214,8 @@ impl Default for ForwardPassPerfOptions { min_faster_correction_factor: default_min_faster_correction_factor(), max_slower_correction_factor: default_max_slower_correction_factor(), bucket_count: DEFAULT_BUCKET_COUNT, + bucket_shape: None, + regression_ridge_scale: 1e-9, max_num_tokens: DEFAULT_MAX_NUM_TOKENS, max_batch_size: DEFAULT_MAX_BATCH_SIZE, max_kv_tokens: DEFAULT_MAX_KV_TOKENS, @@ -261,6 +268,14 @@ pub(crate) fn validate_options(options: &ForwardPassPerfOptions) -> Result<(), A "min_observations must be <= max_observations", )); } + if let Some(shape) = options.bucket_shape { + if shape.contains(&0) || shape[0].checked_mul(shape[1]).is_none() { + return Err(invalid_perf_options( + "bins_per_axis must be positive and have a representable product", + )); + } + return Ok(()); + } let sqrt = integer_sqrt(options.bucket_count); if sqrt * sqrt != options.bucket_count { return Err(invalid_perf_options( @@ -274,6 +289,11 @@ pub(crate) fn validate_regression_options( options: &ForwardPassPerfOptions, ) -> Result<(), AicError> { validate_options(options)?; + if !options.regression_ridge_scale.is_finite() || options.regression_ridge_scale < 0.0 { + return Err(invalid_perf_options( + "singular_ridge_scale must be finite and nonnegative", + )); + } for (name, value) in [ ( "regression_attention_kv_weight", @@ -467,3 +487,7 @@ mod tests { } } } + +fn default_regression_ridge_scale() -> f64 { + 1e-9 +} diff --git a/crates/core/src/perfmodel/fpm/regression.rs b/crates/core/src/perfmodel/fpm/regression.rs index 86227263d..751418903 100644 --- a/crates/core/src/perfmodel/fpm/regression.rs +++ b/crates/core/src/perfmodel/fpm/regression.rs @@ -25,6 +25,7 @@ pub(crate) struct BucketedRegression { samples: BucketedSamples, min_observations: usize, fit: Option, + ridge_scale: f64, } impl StoreStats for BucketedRegression { @@ -43,6 +44,7 @@ impl BucketedRegression { samples: BucketedSamples::new_dynamic(options, FEATURE_DIMENSION), min_observations: options.min_observations, fit: None, + ridge_scale: options.regression_ridge_scale, } } @@ -71,7 +73,7 @@ impl BucketedRegression { .into_iter() .map(|(_, observation)| observation) .collect::>(); - self.fit = fit_regression(&retained, self.min_observations); + self.fit = fit_regression_with_ridge(&retained, self.min_observations, self.ridge_scale); true } @@ -186,9 +188,18 @@ struct StandardizedObservation { observed_ms: f64, } +#[cfg(test)] fn fit_regression( observations: &[RegressionObservation], min_observations: usize, +) -> Option { + fit_regression_with_ridge(observations, min_observations, 1e-9) +} + +fn fit_regression_with_ridge( + observations: &[RegressionObservation], + min_observations: usize, + ridge_scale: f64, ) -> Option { if observations.len() < min_observations { return None; @@ -224,7 +235,9 @@ fn fit_regression( (active_mask & (1usize << mask_axis) != 0).then_some(*feature_axis) }) .collect::>(); - let Some(fit) = fit_linear_active_set(&standardized, standardization, &fitted_axes) else { + let Some(fit) = + fit_linear_active_set(&standardized, standardization, &fitted_axes, ridge_scale) + else { continue; }; let squared_error = standardized @@ -263,6 +276,7 @@ fn fit_linear_active_set( observations: &[StandardizedObservation], standardization: Standardization, fitted_axes: &[usize], + ridge_scale: f64, ) -> Option { let size = fitted_axes.len() + 1; let mut lhs = vec![vec![0.0_f64; size]; size]; @@ -283,7 +297,7 @@ fn fit_linear_active_set( } let solution = solve_linear_system(lhs.clone(), rhs.clone()) - .or_else(|| solve_regularized_linear_system(lhs, rhs))?; + .or_else(|| solve_regularized_linear_system(lhs, rhs, ridge_scale))?; if !solution.iter().all(|value| value.is_finite()) || solution[1..].iter().any(|coefficient| *coefficient < 0.0) { @@ -302,14 +316,18 @@ fn fit_linear_active_set( } /// Retry a singular normal equation with ridge regularization on slopes only. -fn solve_regularized_linear_system(mut lhs: Vec>, rhs: Vec) -> Option> { +fn solve_regularized_linear_system( + mut lhs: Vec>, + rhs: Vec, + ridge_scale: f64, +) -> Option> { let scale = lhs .iter() .enumerate() .map(|(axis, row)| row[axis].abs()) .sum::() .max(1.0); - let ridge = scale * 1e-9; + let ridge = scale * ridge_scale; for (axis, row) in lhs.iter_mut().enumerate().skip(1) { row[axis] += ridge; } @@ -574,6 +592,9 @@ mod tests { #[test] fn collinear_active_axes_use_slope_regularization_with_free_intercept() { + let mut config = crate::EstimatorConfig::default(); + config.fpm_regression.fit.singular_ridge_scale = 0.25; + let options = config.regression_options(); let observations = (1..=6) .map(|attention| RegressionObservation { raw_x: [attention as f64, 2.0 * attention as f64], @@ -593,8 +614,15 @@ mod tests { // The two standardized columns are identical, so the unregularized // normal equation is singular. The fallback regularizes only the two // slopes, leaving the free intercept at the population target mean. - let regularized_fit = - fit_linear_active_set(&standardized, standardization, &[0, 1]).unwrap(); + // trace(X'X)=18, lambda=18/4=4.5, so the combined slope is + // shrunk by 12/(12+4.5)=8/11. At x=7: 17.5 + 10.5*8/11 = 553/22. + let regularized_fit = fit_linear_active_set( + &standardized, + standardization, + &[0, 1], + options.regression_ridge_scale, + ) + .unwrap(); assert!( regularized_fit .coefficients @@ -602,13 +630,20 @@ mod tests { .all(|slope| *slope > 0.0) ); assert_close(regularized_fit.intercept, 17.5, 1e-12); - assert_close(regularized_fit.predict(&[7.0, 14.0]).unwrap(), 28.0, 1e-7); + assert_close( + regularized_fit.predict(&[7.0, 14.0]).unwrap(), + 553.0 / 22.0, + 1e-12, + ); - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new(&options); + assert_eq!(regression.ridge_scale, 0.25); for observation in observations { assert!(regression.add_observation(observation.raw_x, observation.observed_ms)); } assert!(regression.is_ready()); + // The full NNLS search can select a nonsingular single-axis face, + // which fits these exactly collinear observations without shrinkage. assert_close(regression.predict(&[7.0, 14.0]).unwrap(), 28.0, 1e-7); } diff --git a/crates/core/src/perfmodel/fpm/samples.rs b/crates/core/src/perfmodel/fpm/samples.rs index f6b13af13..8a6665db2 100644 --- a/crates/core/src/perfmodel/fpm/samples.rs +++ b/crates/core/src/perfmodel/fpm/samples.rs @@ -31,7 +31,7 @@ pub(crate) struct BucketedSamples { axis_min: Vec, axis_max: Vec, fixed_bounds: bool, - buckets_per_axis: usize, + buckets_per_axis: Vec, max_observations: usize, } @@ -52,10 +52,16 @@ impl AxisRange { impl BucketedSamples { pub(crate) fn new_dynamic(options: &ForwardPassPerfOptions, ndim: usize) -> Self { - let buckets_per_axis = if ndim == 1 { - options.bucket_count + let buckets_per_axis = if let Some(shape) = options.bucket_shape { + if ndim == 1 { + vec![shape[0] * shape[1]] + } else { + shape.to_vec() + } + } else if ndim == 1 { + vec![options.bucket_count] } else { - integer_sqrt(options.bucket_count) + vec![integer_sqrt(options.bucket_count); ndim] }; Self { buckets: HashMap::new(), @@ -63,7 +69,7 @@ impl BucketedSamples { axis_min: vec![f64::INFINITY; ndim], axis_max: vec![f64::NEG_INFINITY; ndim], fixed_bounds: false, - buckets_per_axis: buckets_per_axis.max(1), + buckets_per_axis, max_observations: options.max_observations, } } @@ -118,8 +124,9 @@ impl BucketedSamples { if hi <= lo { 0 } else { - let idx = ((*value - lo) / (hi - lo) * self.buckets_per_axis as f64) as isize; - idx.clamp(0, self.buckets_per_axis as isize - 1) as usize + let idx = + ((*value - lo) / (hi - lo) * self.buckets_per_axis[i] as f64) as isize; + idx.clamp(0, self.buckets_per_axis[i] as isize - 1) as usize } }) .collect() @@ -150,8 +157,8 @@ impl BucketedSamples { return None; } - let idx = ((*value - lo) / (hi - lo) * self.buckets_per_axis as f64) as isize; - key.push(idx.clamp(0, self.buckets_per_axis as isize - 1) as usize); + let idx = ((*value - lo) / (hi - lo) * self.buckets_per_axis[i] as f64) as isize; + key.push(idx.clamp(0, self.buckets_per_axis[i] as isize - 1) as usize); } Some(key) } diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index 749a4c017..1e51ebb24 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -572,6 +572,13 @@ fn options_default_directional_correction_factors() { assert_eq!(no_ceiling.max_slower_correction_factor, None); } +#[test] +fn options_reject_unknown_fields() { + let err = + serde_json::from_str::(r#"{"min_observation": 5}"#).unwrap_err(); + assert!(err.to_string().contains("unknown field"), "{err}"); +} + #[test] fn options_validate_directional_correction_factors() { let model = regression_model( diff --git a/crates/core/src/perfmodel/mod.rs b/crates/core/src/perfmodel/mod.rs index 12f7f8fe7..a563e2e3d 100644 --- a/crates/core/src/perfmodel/mod.rs +++ b/crates/core/src/perfmodel/mod.rs @@ -7,8 +7,11 @@ //! `compile_engine` walks the model once and emits an [`engine::spec::EngineSpec`] //! (op lists + [`EngineConfig`] identity); the Rust [`engine::Engine`] executes //! it without re-entering Python. With the `python` feature enabled, -//! [`AicEngineBuilder`] is the preferred Rust → Python → Rust embedded build -//! entry point and [`AicEngine`] is the PyO3 hot-path pyclass. +//! [`AicEngineBuilder`] is the low-level Rust → Python → Rust build path for +//! embedders that need an [`AicEngine`] step-latency handle (for example Dynamo +//! Mocker). It is not a forward-pass estimator constructor: Planner, Replay, +//! and Sweeper use [`ForwardPassPerfModel::best_available`] and the canonical +//! [`ForwardPassPerfModelConfig`] contract. //! //! This directory remains a stable mirror of the former AIConfigurator Rust //! crate. Keeping the imported implementation behind one namespace makes @@ -42,9 +45,16 @@ pub use common::{AicError, enums::DatabaseMode}; // planner / Mocker) can use it natively; also exposed to Python via the // `RustForwardPassPerfModel` pyclass in `py.rs`. pub use fpm::{ - ForwardPassPerfDiagnostics, ForwardPassPerfModel, ForwardPassPerfOptions, - ForwardPassPerfReadiness, ForwardPassPerfSource, ForwardPassRegressionStoreDiagnostics, - ForwardPassRegressionWorkloadKind, ForwardPassWorkerType, + CorrectionConfig, EstimationMode, EstimatorConfig, ForwardPassFallbackPolicy, + ForwardPassPerfDiagnostics, ForwardPassPerfModel, ForwardPassPerfModelConfig, + ForwardPassPerfOptions, ForwardPassPerfProvenance, ForwardPassPerfReadiness, + ForwardPassPerfSource, ForwardPassRegressionStoreDiagnostics, + ForwardPassRegressionWorkloadKind, ForwardPassSpeculationConfig, ForwardPassWorkerType, + FpmRegressionConfig, RegressionFeatureWeights, SamplingConfig, +}; +pub use fpm::{ + CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, OpLevelConfig, + RegressionFitConfig, RegressionFitKind, }; // Forward-pass metrics telemetry types and schema version, plus the // crate-internal validation helper. Re-exported at the crate root so existing @@ -61,7 +71,8 @@ pub use memory::{ KvCacheEstimateOptions, KvCacheEstimateRequest, KvCacheMemoryFraction, MemoryBreakdown, }; // PyO3 bindings. `AicEngine` is the Python -> Rust hot-path pyclass; -// `AicEngineBuilder` is the Rust -> Python -> Rust entry point. They must be +// `AicEngineBuilder` is the low-level Rust -> Python -> Rust compiled-engine +// entry point. It does not construct a forward-pass estimator. They must be // `pub`-re-exported here because the `py` module itself is private. #[cfg(feature = "python")] pub use py::{AicEngine, AicEngineBuilder}; diff --git a/crates/core/src/perfmodel/perf_database/dsv4_megamoe.rs b/crates/core/src/perfmodel/perf_database/dsv4_megamoe.rs index e4308f2a8..1fa9e1a9f 100644 --- a/crates/core/src/perfmodel/perf_database/dsv4_megamoe.rs +++ b/crates/core/src/perfmodel/perf_database/dsv4_megamoe.rs @@ -41,7 +41,7 @@ //! `PerformanceResult(latency, energy=energy)`). use std::collections::BTreeMap; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::OnceLock; use super::axis_curve::LeafAxisCurve; @@ -106,6 +106,11 @@ impl Dsv4MegaMoeTable { } } + /// The single source selected for both readiness and lazy table loading. + pub(crate) fn primary_path(&self) -> &Path { + &self.primary_path + } + /// Query the measured MegaMoE routed-module value (latency ms + /// power/energy) at `num_tokens` (rank-LOCAL token count — callers must /// NOT pre-multiply by attention_dp_size). Mirrors the table body of diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 0297bc242..b3cb648cf 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -148,6 +148,23 @@ fn parse_mode(mode: &str) -> PyResult { } } +/// Discover the ordered roots using the same SDK/environment policy as the Python facade. +pub(crate) fn resolve_forward_pass_systems_roots() -> Result, AicError> { + Python::with_gil(|py| { + py.import("aiconfigurator_core.sdk.rust_engine_step")? + .getattr("_resolve_forward_pass_systems_paths")? + .call1((Vec::::new(),))? + .extract::>() + }) + .map_err(|error: PyErr| { + if Python::with_gil(|py| error.is_instance_of::(py)) { + AicError::InvalidEngineConfig(format!("resolve systems paths: {error}")) + } else { + AicError::DataRoot(format!("resolve systems paths: {error}")) + } + }) +} + /// Resolve the bundled `systems/` directory for [`AicEngine::from_spec`]. /// /// Mirrors the systems-root half of `DataRoots::discover` but does NOT require @@ -156,7 +173,7 @@ fn parse_mode(mode: &str) -> PyResult { /// Precedence: explicit `systems_path` arg → `AICONFIGURATOR_SYSTEMS_PATH` env /// → the installed core wheel's SDK resource path → repo-relative /// `python/aisimulate/src/aiconfigurator_core/systems`. -fn resolve_systems_root(systems_path: Option<&str>) -> PyResult { +pub(crate) fn resolve_systems_root(systems_path: Option<&str>) -> PyResult { if let Some(p) = systems_path { return Ok(PathBuf::from(p)); } @@ -195,6 +212,10 @@ pub struct AicEngine { } impl AicEngine { + pub(crate) fn from_shared_engine(inner: Arc) -> Self { + Self { inner } + } + /// Internal constructor shared by [`AicEngine::from_spec`] and /// [`AicEngineBuilder::build`]. fn new(engine: Engine) -> Self { @@ -1027,6 +1048,7 @@ struct EngineBuildRequest { comm_quant_mode: Option, attention_backend: Option, nextn: u32, + speculation: Option, kv_block_size: Option, systems_path: Option, forward_model: Option, @@ -1071,6 +1093,7 @@ impl AicEngineBuilder { comm_quant_mode: None, attention_backend: None, nextn: 0, + speculation: None, kv_block_size: None, systems_path: None, forward_model: None, @@ -1363,6 +1386,12 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result Result Result { + let forward_model = config.estimation_mode.native_name().ok_or_else(|| { + AicError::InvalidEngineConfig( + "a native estimator mode must be selected before compilation".into(), + ) + })?; + compile_engine_from_request(EngineBuildRequest { + model_path: config.model.clone(), + system: config.system.clone(), + backend: config.backend.as_str().to_owned(), + backend_version: config.backend_version.clone(), + tp_size: config.tp, + pp_size: config.pp, + attention_dp_size: config.attention_dp, + moe_tp_size: config.moe_tp_size, + moe_ep_size: config.moe_ep_size, + gemm_quant_mode: config.gemm_quant_mode.clone(), + moe_quant_mode: config.moe_quant_mode.clone(), + kvcache_quant_mode: config.kvcache_quant_mode.clone(), + fmha_quant_mode: config.fmha_quant_mode.clone(), + comm_quant_mode: config.comm_quant_mode.clone(), + attention_backend: config.attention_backend.clone(), + nextn: config.nextn, + speculation: config.speculation.clone(), + kv_block_size: config.kv_block_size, + systems_path: Some(systems_path.to_owned()), + forward_model: Some(forward_model.to_owned()), + database_mode: Some(config.database_mode.as_str().to_owned()), + shared_layer: config.enable_shared_layer, + transfer_policy: config.transfer_policy.clone(), + strict_provenance: Some(config.strict_provenance), + }) +} + /// Build a compiled [`Engine`] from a modular [`EngineConfig`] (the /// `fpm::ForwardPassPerfModel::from_native` entry point). Maps the /// config's nested fields onto the shared [`EngineBuildRequest`]. @@ -1433,6 +1524,7 @@ fn engine_build_request(config: &EngineConfig, systems_path: Option<&str>) -> En // Attention backend is not carried on EngineConfig; let Python resolve it. attention_backend: None, nextn, + speculation: None, kv_block_size: config.kv_block_size, systems_path: systems_path.map(str::to_owned), forward_model: config.forward_model.clone(), @@ -1516,31 +1608,13 @@ pub struct PyForwardPassPerfModel { inner: crate::ForwardPassPerfModel, } -/// Parse the optional options JSON into [`ForwardPassPerfOptions`], defaulting -/// when `None`/empty. Serde fills missing fields from the per-field defaults. -/// Raw JSON callers represent nonfinite regression weights with the exact -/// quoted sentinels `"NaN"`, `"Infinity"`, and `"-Infinity"`. -fn parse_fpm_options(options_json: Option<&str>) -> PyResult { - match options_json { - None => Ok(crate::ForwardPassPerfOptions::default()), - Some(s) if s.trim().is_empty() => Ok(crate::ForwardPassPerfOptions::default()), - Some(s) => serde_json::from_str(s) - .map_err(|e| PyValueError::new_err(format!("invalid options JSON: {e}"))), - } -} - -/// Parse the exact worker-type spellings shared with the Dynamo planner. -/// Aliases are intentionally rejected so model identity cannot depend on which -/// API surface constructed it. -fn parse_fpm_worker_type(worker_type: &str) -> PyResult { - match worker_type { - "prefill" => Ok(crate::ForwardPassWorkerType::Prefill), - "decode" => Ok(crate::ForwardPassWorkerType::Decode), - "aggregated" => Ok(crate::ForwardPassWorkerType::Aggregated), - _ => Err(PyValueError::new_err(format!( - "invalid worker_type {worker_type:?}: expected 'prefill', 'decode', or 'aggregated'" - ))), - } +fn parse_forward_pass_config(config_json: &str) -> PyResult { + let mut de = serde_json::Deserializer::from_str(config_json); + let config = serde_path_to_error::deserialize(&mut de) + .map_err(|e| PyValueError::new_err(format!("invalid forward-pass config at {e}")))?; + de.end() + .map_err(|e| PyValueError::new_err(format!("invalid forward-pass config: {e}")))?; + Ok(config) } /// Parse an FPM payload JSON into one iteration's per-rank list. Accepts either @@ -1560,50 +1634,97 @@ fn parse_fpm_iteration(fpm_json: &str) -> PyResult) -> PyResult { - let config: EngineConfig = serde_json::from_str(config_json) - .map_err(|e| PyValueError::new_err(format!("invalid engine config JSON: {e}")))?; - let options = parse_fpm_options(options_json)?; - let inner = crate::ForwardPassPerfModel::from_native(config, options).map_err(aic_to_py)?; + fn best_available(config_json: &str) -> PyResult { + let config = parse_forward_pass_config(config_json)?; + let inner = crate::ForwardPassPerfModel::best_available(config).map_err(aic_to_py)?; Ok(Self { inner }) } - /// `RustForwardPassPerfModel.best_available(config_json, worker_type, - /// options_json=None)`: - /// native when possible, else regression fallback (reason in - /// `diagnostics()["last_warning"]`). + /// Expand and validate the canonical schema without constructing an engine. #[staticmethod] - #[pyo3(signature = (config_json, worker_type, options_json=None))] - fn best_available( - config_json: &str, - worker_type: &str, - options_json: Option<&str>, - ) -> PyResult { - let config: EngineConfig = serde_json::from_str(config_json) - .map_err(|e| PyValueError::new_err(format!("invalid engine config JSON: {e}")))?; - let worker_type = parse_fpm_worker_type(worker_type)?; - let options = parse_fpm_options(options_json)?; - let inner = crate::ForwardPassPerfModel::best_available(config, worker_type, options) - .map_err(aic_to_py)?; - Ok(Self { inner }) + fn normalize_config(config_json: &str) -> PyResult { + let config = parse_forward_pass_config(config_json)?; + config.validate().map_err(aic_to_py)?; + serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) } - /// `RustForwardPassPerfModel.from_regression(worker_type, - /// options_json=None)`: - /// regression-only model (no native engine, no Python compile). + /// Migration adapter for previously saved flat tuning options. #[staticmethod] - #[pyo3(signature = (worker_type, options_json=None))] - fn from_regression(worker_type: &str, options_json: Option<&str>) -> PyResult { - let worker_type = parse_fpm_worker_type(worker_type)?; - let options = parse_fpm_options(options_json)?; - let inner = crate::ForwardPassPerfModel::from_regression(worker_type, options) - .map_err(aic_to_py)?; - Ok(Self { inner }) + fn legacy_estimator_config(options_json: &str) -> PyResult { + let options: crate::ForwardPassPerfOptions = serde_json::from_str(options_json) + .map_err(|e| PyValueError::new_err(format!("invalid legacy options: {e}")))?; + let config = crate::EstimatorConfig::from_legacy(options).map_err(aic_to_py)?; + config.validate().map_err(aic_to_py)?; + serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) + } + + /// Translate a saved EngineConfig and flat options into the canonical schema. + #[staticmethod] + #[pyo3(signature = (config_json, worker_type, options_json=None, allow_regression=false))] + fn migrate_legacy_config( + config_json: &str, + worker_type: &str, + options_json: Option<&str>, + allow_regression: bool, + ) -> PyResult { + let legacy: EngineConfig = + serde_json::from_str(config_json).map_err(|e| PyValueError::new_err(e.to_string()))?; + let options = options_json + .map(serde_json::from_str::) + .transpose() + .map_err(|e| PyValueError::new_err(e.to_string()))? + .unwrap_or_default(); + let request = engine_build_request( + &legacy, + legacy.systems_path.as_ref().and_then(|path| path.to_str()), + ); + let worker_type = serde_json::from_value(serde_json::Value::String(worker_type.to_owned())) + .map_err(|e| PyValueError::new_err(format!("invalid worker_type: {e}")))?; + let estimation_mode = match request.forward_model.as_deref().unwrap_or("op_level") { + "op_level" => crate::EstimationMode::OpLevel, + "fpm" => crate::EstimationMode::FpmInterpolation, + value => { + return Err(PyValueError::new_err(format!( + "unknown legacy forward_model {value:?}" + ))); + } + }; + let config = crate::ForwardPassPerfModelConfig { + model: request.model_path, + system: request.system, + backend: legacy.backend, + worker_type, + backend_version: request.backend_version, + tp: request.tp_size, + pp: request.pp_size, + attention_dp: request.attention_dp_size, + moe_tp_size: request.moe_tp_size, + moe_ep_size: request.moe_ep_size, + gemm_quant_mode: request.gemm_quant_mode, + moe_quant_mode: request.moe_quant_mode, + fmha_quant_mode: request.fmha_quant_mode, + kvcache_quant_mode: request.kvcache_quant_mode, + comm_quant_mode: request.comm_quant_mode, + nextn: request.nextn, + speculation: request.speculation, + kv_block_size: request.kv_block_size, + estimation_mode, + database_mode: legacy.database_mode, + transfer_policy: request.transfer_policy, + systems_paths: legacy.systems_path.into_iter().collect(), + fallback_policy: if allow_regression { + crate::ForwardPassFallbackPolicy::LegacyRegression + } else { + crate::ForwardPassFallbackPolicy::Deny + }, + estimator_config: crate::EstimatorConfig::from_legacy(options).map_err(aic_to_py)?, + attention_backend: request.attention_backend, + enable_shared_layer: request.shared_layer, + strict_provenance: legacy.strict_provenance, + }; + serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) } /// Estimate one forward-pass iteration in ms. `fpm_json` is one iteration as diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index cdffe9c49..c1fd2b6a7 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -21,11 +21,15 @@ use crate::replay::{ load_agentic_mooncake, load_weka_agentic_graph_with_options, }, }; +use crate::{ + EstimationMode, EstimatorConfig, ForwardPassFallbackPolicy, ForwardPassPerfModel, + ForwardPassPerfModelConfig, ForwardPassWorkerType, +}; use anyhow::{Context, Result, anyhow, ensure}; use pyo3::exceptions::{PyMemoryError, PyRuntimeError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict, PyModule}; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; #[derive(Debug, Deserialize)] #[serde(untagged)] @@ -147,7 +151,7 @@ struct AicTimingConfig { model: String, backend: String, system: String, - #[serde(alias = "tp_size")] + #[serde(default = "one", alias = "tp_size")] tp: u32, #[serde(default)] backend_version: Option, @@ -172,7 +176,7 @@ struct AicTimingConfig { #[serde(default)] nextn: u32, #[serde(default)] - speculation: Option, + speculation: Option, #[serde(default)] kv_block_size: Option, #[serde(default)] @@ -187,25 +191,26 @@ struct AicTimingConfig { systems_path: Option, #[serde(default)] forward_model: Option, -} - -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(rename_all = "snake_case")] -enum PromptLookupKind { - Ngram, -} - -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(deny_unknown_fields)] -struct NgramTimingConfig { - kind: PromptLookupKind, - params: NgramTimingParams, -} - -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(deny_unknown_fields)] -struct NgramTimingParams { - num_speculative_tokens: u32, + #[serde(default)] + worker_type: Option, + #[serde(default)] + estimation_mode: Option, + #[serde(default)] + fallback_policy: ForwardPassFallbackPolicy, + #[serde(default)] + estimator_config: EstimatorConfig, + #[serde(default)] + database_mode: crate::DatabaseMode, + #[serde(default)] + transfer_policy: Option>, + #[serde(default)] + systems_paths: Vec, + #[serde(default)] + attention_backend: Option, + #[serde(default)] + enable_shared_layer: Option, + #[serde(default)] + strict_provenance: bool, } const fn one() -> u32 { @@ -213,6 +218,71 @@ const fn one() -> u32 { } impl AicTimingConfig { + fn estimator_request( + &self, + worker_type: ForwardPassWorkerType, + ) -> Result { + ensure!( + self.worker_type + .is_none_or(|configured| configured == worker_type), + "estimator worker_type does not match replay role" + ); + let legacy_mode = match self.forward_model.as_deref() { + Some("fpm") => Some(EstimationMode::FpmInterpolation), + Some("op_level") => Some(EstimationMode::OpLevel), + Some(other) => return Err(anyhow!("unknown legacy forward_model {other:?}")), + None => None, + }; + if let (Some(explicit), Some(legacy)) = (self.estimation_mode, legacy_mode) { + ensure!( + explicit == legacy, + "forward_model conflicts with estimation_mode" + ); + } + let mode = self + .estimation_mode + .or(legacy_mode) + .unwrap_or(if self.worker_type.is_some() { + EstimationMode::Auto + } else { + EstimationMode::OpLevel + }); + let roots = if self.systems_paths.is_empty() { + self.systems_path.iter().map(PathBuf::from).collect() + } else { + self.systems_paths.clone() + }; + Ok(ForwardPassPerfModelConfig { + model: self.model.clone(), + system: self.system.clone(), + backend: serde_json::from_value(serde_json::Value::String(self.backend.clone()))?, + backend_version: self.backend_version.clone(), + worker_type, + tp: self.tp, + pp: self.pp, + attention_dp: self.attention_dp, + moe_tp_size: self.moe_tp_size, + moe_ep_size: self.moe_ep_size, + gemm_quant_mode: self.gemm_dtype.clone(), + moe_quant_mode: self.moe_dtype.clone(), + fmha_quant_mode: self.fmha_dtype.clone(), + kvcache_quant_mode: self.kv_cache_dtype.clone(), + comm_quant_mode: self.comm_dtype.clone(), + nextn: self.nextn, + speculation: self.speculation.clone(), + kv_block_size: self.kv_block_size, + estimation_mode: mode, + fallback_policy: self.fallback_policy, + estimator_config: self.estimator_config.clone(), + database_mode: self.database_mode, + transfer_policy: self.transfer_policy.clone(), + systems_paths: roots, + attention_backend: self.attention_backend.clone(), + enable_shared_layer: self.enable_shared_layer, + strict_provenance: self.strict_provenance, + }) + } + fn speculative_depth(&self) -> Result { let Some(speculation) = &self.speculation else { return Ok(self.nextn); @@ -229,7 +299,7 @@ impl AicTimingConfig { self.forward_model.as_deref().unwrap_or("op_level") == "op_level", "ngram speculation requires op_level timing" ); - let depth = speculation.params.num_speculative_tokens; + let depth = speculation.num_speculative_tokens(); ensure!( (1..=5).contains(&depth), "ngram num_speculative_tokens must be in 1..=5" @@ -310,74 +380,43 @@ struct AicTimingModel { } impl AicTimingModel { - fn build(config: AicTimingConfig) -> Result { - ensure!( - !config.model.trim().is_empty(), - "AIC timing config field \"model\" cannot be empty" - ); - ensure!( - !config.system.trim().is_empty(), - "AIC timing config field \"system\" cannot be empty" - ); + fn build(config: &mut AicTimingConfig, worker_type: ForwardPassWorkerType) -> Result { config.validate_parallel_shape()?; - ensure!( - matches!(config.backend.as_str(), "vllm" | "sglang" | "trtllm"), - "unsupported AIC backend {:?}; expected vllm, sglang, or trtllm", - config.backend - ); - ensure!( - !config.resolved_backend_version().is_empty(), - "AIC backend version cannot be empty" - ); config.resolved_memory_fraction()?; - - let use_fpm_decode_totals = config.forward_model.as_deref() == Some("fpm"); + let model = ForwardPassPerfModel::best_available(config.estimator_request(worker_type)?) + .context("AIC timing provider could not construct the requested estimator")?; + let provenance = model + .provenance() + .context("canonical estimator omitted construction provenance")?; + config.backend_version = provenance.config.backend_version.clone(); + config.systems_path = provenance + .selected_systems_root + .as_ref() + .map(|path| path.to_string_lossy().into_owned()); + config.systems_paths = provenance.config.systems_paths.clone(); + config.database_mode = provenance.config.database_mode; + config.transfer_policy = provenance.config.transfer_policy.clone(); + config.estimation_mode = Some(provenance.selected_estimation_mode); + config.worker_type = Some(worker_type); + config.forward_model = None; + config.fallback_policy = ForwardPassFallbackPolicy::Deny; + config.estimator_config = provenance.config.estimator_config.clone(); + let use_fpm_decode_totals = + provenance.selected_estimation_mode == EstimationMode::FpmInterpolation; + let native = model.native_engine().context( + "AIC regression estimator is not ready: offline replay requires trained observations or a native estimator" + )?; let (engine, fpm_decode_kv_ceiling) = Python::with_gil(|py| -> PyResult<_> { - let sdk = PyModule::import(py, "aiconfigurator_core.sdk.engine")?; - let kwargs = PyDict::new(py); - kwargs.set_item("backend_version", config.resolved_backend_version())?; - kwargs.set_item("tp_size", config.tp)?; - kwargs.set_item("pp_size", config.pp)?; - kwargs.set_item("attention_dp_size", config.attention_dp)?; - kwargs.set_item("moe_tp_size", config.moe_tp_size)?; - kwargs.set_item("moe_ep_size", config.moe_ep_size)?; - kwargs.set_item("gemm_quant_mode", config.gemm_dtype.as_deref())?; - kwargs.set_item("moe_quant_mode", config.moe_dtype.as_deref())?; - kwargs.set_item("fmha_quant_mode", config.fmha_dtype.as_deref())?; - kwargs.set_item("kvcache_quant_mode", config.kv_cache_dtype.as_deref())?; - kwargs.set_item("comm_quant_mode", config.comm_dtype.as_deref())?; - kwargs.set_item("nextn", config.nextn)?; - if let Some(speculation) = &config.speculation { - let value = PyModule::import(py, "json")?.call_method1( - "loads", - (serde_json::to_string(speculation).expect("serializable ngram config"),), - )?; - kwargs.set_item("speculation", value)?; - } - kwargs.set_item("kv_block_size", config.kv_block_size)?; - kwargs.set_item("systems_path", config.systems_path.as_deref())?; - kwargs.set_item("forward_model", config.forward_model.as_deref())?; - let spec = sdk.getattr("compile_engine")?.call( - ( - config.model.as_str(), - config.system.as_str(), - config.backend.as_str(), - ), - Some(&kwargs), - )?; - let aic = PyModule::import(py, "aiconfigurator_core")? - .getattr("AicEngine")? - .call_method1("from_spec", (spec, config.systems_path.as_deref()))?; - let fpm_decode_kv_ceiling = if use_fpm_decode_totals { - aic.call_method0("fpm_decode_kv_ceiling")? + let engine = Py::new(py, crate::AicEngine::from_shared_engine(native))?.into_any(); + let ceiling = if use_fpm_decode_totals { + engine + .bind(py) + .call_method0("fpm_decode_kv_ceiling")? .extract::>()? } else { None }; - Ok((aic.unbind(), fpm_decode_kv_ceiling)) - }) - .map_err(|error| { - anyhow!("AIC timing provider could not compile the requested engine: {error}") + Ok((engine, ceiling)) })?; Ok(Self { engine, @@ -682,6 +721,7 @@ fn role_capacity_is_explicit(engine_value: &serde_json::Value, role: Option<&str fn resolve_role_timing( role: &mut ReplayRoleConfig, capacity_is_explicit: bool, + worker_type: ForwardPassWorkerType, ) -> Result>> { let TimingModelConfig::External { provider, config } = role.rank.timing_model.clone() else { return Ok(None); @@ -691,20 +731,40 @@ fn resolve_role_timing( "native timing provider {provider:?} is not installed; only \"aic\" is \ available in the AISimulate runtime" ); - let config: AicTimingConfig = + let mut config: AicTimingConfig = serde_json::from_value(config).context("invalid AIC timing provider configuration")?; + let timing = AicTimingModel::build(&mut config, worker_type)?; materialize_aic_capacity( &config, role, capacity_is_explicit, estimate_aic_num_gpu_blocks, )?; - let timing = AicTimingModel::build(config)?; cap_role_capacity_to_fpm_decode_domain( role, timing.fpm_decode_kv_ceiling, capacity_is_explicit, )?; + let mut resolved = serde_json::to_value(config.estimator_request(worker_type)?)?; + if let Some(object) = resolved.as_object_mut() { + for (name, value) in [ + ("gpu_memory_utilization", config.gpu_memory_utilization), + ("mem_fraction_static", config.mem_fraction_static), + ("free_gpu_memory_fraction", config.free_gpu_memory_fraction), + ] { + if let Some(value) = value { + object.insert(name.to_owned(), serde_json::json!(value)); + } + } + object.insert( + "cuda_graph_reserved_bytes".to_owned(), + serde_json::json!(config.cuda_graph_reserved_bytes), + ); + } + role.rank.timing_model = TimingModelConfig::External { + provider, + config: resolved, + }; Ok(Some(Arc::new(timing))) } @@ -1334,6 +1394,37 @@ fn scale_power_phase( Ok(phase) } +fn infer_aic_timing_topology(engine: &mut serde_json::Value) -> Result<()> { + let timing = engine.get("rank").and_then(|rank| rank.get("timing_model")); + if let Some(timing) = timing + && timing.get("type").and_then(serde_json::Value::as_str) == Some("external") + && timing.get("provider").and_then(serde_json::Value::as_str) == Some("aic") + { + let config = timing.get("config").cloned().unwrap_or_default(); + for (target, fields) in [ + ("tensor_parallel_size", ["tp", "tp_size"]), + ("dp_size", ["attention_dp", "attention_dp_size"]), + ] { + if let Some(value) = config.get(fields[0]).or_else(|| config.get(fields[1])) { + if let Some(explicit) = engine.get(target) { + ensure!( + explicit == value, + "{target} conflicts with canonical AIC timing topology" + ); + } else if let Some(object) = engine.as_object_mut() { + object.insert(target.to_owned(), value.clone()); + } + } + } + } + for role in ["prefill", "decode"] { + if let Some(child) = engine.get_mut(role) { + infer_aic_timing_topology(child)?; + } + } + Ok(()) +} + fn execute_json(payload: &str, capture_artifacts: bool) -> Result { let (mut spec, mut traffic) = match serde_json::from_str(payload).context("invalid AISimulate execution ReplaySpec")? { @@ -1360,6 +1451,7 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { ); } let serialized_engine = spec.engine.clone(); + infer_aic_timing_topology(&mut spec.engine)?; let mut engine_config: ReplayEngineConfig = if spec.engine.is_null() { ReplayEngineConfig::default() } else { @@ -1378,7 +1470,11 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { let capacity_is_explicit = role .num_gpu_blocks_is_explicit .unwrap_or_else(|| role_capacity_is_explicit(&serialized_engine, None)); - let timing = resolve_role_timing(&mut role, capacity_is_explicit)?; + let timing = resolve_role_timing( + &mut role, + capacity_is_explicit, + ForwardPassWorkerType::Aggregated, + )?; if let Some(timing) = timing.as_ref() { power_sources.push(TimingPowerSource::new(Arc::clone(timing), &role.rank)); } @@ -1435,8 +1531,16 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { let decode_capacity_is_explicit = decode .num_gpu_blocks_is_explicit .unwrap_or_else(|| role_capacity_is_explicit(&serialized_engine, Some("decode"))); - let prefill_timing = resolve_role_timing(&mut prefill, prefill_capacity_is_explicit)?; - let decode_timing = resolve_role_timing(&mut decode, decode_capacity_is_explicit)?; + let prefill_timing = resolve_role_timing( + &mut prefill, + prefill_capacity_is_explicit, + ForwardPassWorkerType::Prefill, + )?; + let decode_timing = resolve_role_timing( + &mut decode, + decode_capacity_is_explicit, + ForwardPassWorkerType::Decode, + )?; if let Some(timing) = prefill_timing.as_ref() { power_sources.push(TimingPowerSource::new(Arc::clone(timing), &prefill.rank)); } @@ -1893,6 +1997,42 @@ mod tests { } } + #[test] + fn canonical_timing_infers_topology_and_rejects_explicit_conflicts() { + let mut engine = serde_json::json!({"rank": {"timing_model": {"type":"external", "provider":"aic", "config":{"tp":2, "attention_dp":4}}}}); + infer_aic_timing_topology(&mut engine).unwrap(); + assert_eq!(engine["tensor_parallel_size"], 2); + assert_eq!(engine["dp_size"], 4); + engine["tensor_parallel_size"] = serde_json::json!(1); + assert!(infer_aic_timing_topology(&mut engine).is_err()); + } + + #[test] + fn canonical_and_legacy_default_selection_are_distinct() { + let mut config = aic_config(); + assert_eq!( + config + .estimator_request(ForwardPassWorkerType::Aggregated) + .unwrap() + .estimation_mode, + EstimationMode::OpLevel + ); + config.worker_type = Some(ForwardPassWorkerType::Aggregated); + assert_eq!( + config + .estimator_request(ForwardPassWorkerType::Aggregated) + .unwrap() + .estimation_mode, + EstimationMode::Auto + ); + config.forward_model = Some("typo".into()); + assert!( + config + .estimator_request(ForwardPassWorkerType::Aggregated) + .is_err() + ); + } + fn aic_config() -> AicTimingConfig { AicTimingConfig { model: "test-model".into(), @@ -1916,6 +2056,16 @@ mod tests { mem_fraction_static: None, free_gpu_memory_fraction: None, cuda_graph_reserved_bytes: 0, + worker_type: None, + estimation_mode: None, + fallback_policy: ForwardPassFallbackPolicy::Deny, + estimator_config: EstimatorConfig::default(), + database_mode: crate::DatabaseMode::default(), + transfer_policy: None, + systems_paths: Vec::new(), + attention_backend: None, + enable_shared_layer: None, + strict_provenance: false, systems_path: None, forward_model: None, } @@ -1924,11 +2074,8 @@ mod tests { #[test] fn ngram_timing_requires_matching_scheduler_depth_and_no_mtp() { let mut config = aic_config(); - config.speculation = Some(NgramTimingConfig { - kind: PromptLookupKind::Ngram, - params: NgramTimingParams { - num_speculative_tokens: 2, - }, + config.speculation = Some(crate::ForwardPassSpeculationConfig::Ngram { + num_speculative_tokens: 2, }); let mut role = aggregated_role(&ReplayEngineConfig::default()); assert!(materialize_aic_capacity(&config, &mut role, true, |_, _| unreachable!()).is_err()); @@ -1937,12 +2084,9 @@ mod tests { config.nextn = 2; assert!(config.validate_parallel_shape().is_err()); config.nextn = 0; - config - .speculation - .as_mut() - .unwrap() - .params - .num_speculative_tokens = 6; + config.speculation = Some(crate::ForwardPassSpeculationConfig::Ngram { + num_speculative_tokens: 6, + }); assert!(config.validate_parallel_shape().is_err()); } @@ -1952,7 +2096,9 @@ mod tests { serde_json::json!({"kind": "mtp", "params": {"num_speculative_tokens": 2}}), serde_json::json!({"kind": "ngram", "params": {"num_speculative_tokens": 2, "trigger_rate": 0.5}}), ] { - assert!(serde_json::from_value::(payload).is_err()); + assert!( + serde_json::from_value::(payload).is_err() + ); } } diff --git a/crates/tests/public-api/Cargo.lock b/crates/tests/public-api/Cargo.lock index ae54df26d..9e14ecfb9 100644 --- a/crates/tests/public-api/Cargo.lock +++ b/crates/tests/public-api/Cargo.lock @@ -44,6 +44,7 @@ dependencies = [ "rustc-hash", "serde", "serde_json", + "serde_path_to_error", "serde_yaml", "sha2", "slotmap", @@ -1084,6 +1085,17 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + [[package]] name = "serde_yaml" version = "0.9.34+deprecated" diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index e9b917536..d6fc71bbb 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -8,8 +8,8 @@ use std::path::Path; use aiconfigurator_core::{ - AicEngine, AicEngineBuilder, AicError, BackendKind, DatabaseMode, EngineConfig, - ForwardPassPerfModel, ForwardPassPerfOptions, ForwardPassRegressionStoreDiagnostics, + AicEngine, AicEngineBuilder, AicError, BackendKind, DatabaseMode, EstimationMode, ForwardPassPerfModelConfig, EstimatorConfig, + ForwardPassPerfModel, ForwardPassRegressionStoreDiagnostics, ForwardPassWorkerType, KvCacheEstimateRequest, }; @@ -44,7 +44,10 @@ pub fn build_engine(builder: AicEngineBuilder) -> Result { /// Compile the forward-pass model's public constructor and telemetry type. pub fn regression_model() -> Result { - ForwardPassPerfModel::from_regression(ForwardPassWorkerType::Aggregated, regression_options()) + let mut config = ForwardPassPerfModelConfig::new("test/model", "test-system", BackendKind::Vllm, ForwardPassWorkerType::Aggregated); + config.estimation_mode = EstimationMode::FpmRegression; + config.estimator_config = regression_options(); + ForwardPassPerfModel::best_available(config) } /// Per-store diagnostics remain accessible without changing the summary type. @@ -55,36 +58,21 @@ pub fn regression_stores( } /// Construct and expose every public regression-weight option from an external crate. -pub fn regression_options() -> ForwardPassPerfOptions { - ForwardPassPerfOptions { - regression_attention_kv_weight: 2.0, - regression_prefill_attention_pair_weight: 3.0, - regression_ffn_token_weight: 4.0, - ..ForwardPassPerfOptions::default() - } +pub fn regression_options() -> EstimatorConfig { + let mut config = EstimatorConfig::default(); + config.features.attention_kv_weight = 2.0; + config.features.prefill_attention_pair_weight = 3.0; + config.features.ffn_token_weight = 4.0; + config } -/// Compile the fallback-capable constructor with its required worker type. -/// This is not called because native construction embeds Python. -pub fn best_available_model(config: EngineConfig) -> Result { - ForwardPassPerfModel::best_available( - config, - ForwardPassWorkerType::Decode, - ForwardPassPerfOptions::default(), - ) +pub fn best_available_model(config: ForwardPassPerfModelConfig) -> Result { + ForwardPassPerfModel::best_available(config) } -/// Compile the explicit-systems-root constructor from an external crate. -pub fn best_available_model_with_roots( - config: EngineConfig, - systems_root: impl AsRef, -) -> Result { - ForwardPassPerfModel::best_available_with_roots( - config, - ForwardPassWorkerType::Prefill, - ForwardPassPerfOptions::default(), - systems_root, - ) +pub fn best_available_model_with_roots(mut config: ForwardPassPerfModelConfig, systems_root: impl AsRef) -> Result { + config.systems_paths = vec![systems_root.as_ref().to_path_buf()]; + ForwardPassPerfModel::best_available(config) } /// Keep the KV request type in the external-consumer contract without @@ -225,8 +213,8 @@ mod tests { #[test] fn regression_weight_fields_are_public() { let options = regression_options(); - assert_eq!(options.regression_attention_kv_weight, 2.0); - assert_eq!(options.regression_prefill_attention_pair_weight, 3.0); - assert_eq!(options.regression_ffn_token_weight, 4.0); + assert_eq!(options.features.attention_kv_weight, 2.0); + assert_eq!(options.features.prefill_attention_pair_weight, 3.0); + assert_eq!(options.features.ffn_token_weight, 4.0); } } diff --git a/docs/cli/migrate-from-aiconfigurator.md b/docs/cli/migrate-from-aiconfigurator.md index a8fee20c7..0fb962629 100644 --- a/docs/cli/migrate-from-aiconfigurator.md +++ b/docs/cli/migrate-from-aiconfigurator.md @@ -236,7 +236,7 @@ required result, keep the AIC command above. - [4.3 Traces and multi-turn sessions](#43-replay-traces-and-multi-turn-sessions) - [4.4 Cache capacity and host offload](#44-model-cache-capacity-and-host-offload) - [4.5 Dynamo routing and planning](#45-include-dynamo-routing-and-planning) -- [4.6 Op-level and FPM timing](#46-select-op-level-or-whole-forward-fpm-timing) +- [4.6 Select and configure performance estimators](#46-select-and-configure-performance-estimators) - [4.7 Analytical EPD](#47-predict-and-search-analytical-epd) - [4.8 Heterogeneous P/D hardware](#48-migrate-heterogeneous-pd-hardware) - [4.9 AFD](#49-afd-translation) @@ -435,7 +435,11 @@ scaling limits and the search's candidate GPU budget are separate controls. -### 4.6 Select op-level or whole-forward FPM timing + + +### 4.6 Select and configure performance estimators + +#### 4.6.1 Select an estimator **Before — AIC estimates a batch using collected whole-forward profiles:** @@ -453,20 +457,79 @@ AIC_ALLOW_UNLISTED_VERSIONS=1 aiconfigurator cli estimate \ ```bash AIC_ALLOW_UNLISTED_VERSIONS=1 aisimulate predict \ --config tests/e2e/configs/unified_cli/predict/fpm/01-minimax-m27-h200-tp4-fpm.yaml \ - --set engine.workers.aggregated.timing.forward_model=fpm \ + --set engine.workers.aggregated.timing.estimation_mode=fpm_interpolation \ + --set engine.workers.aggregated.timing.fallback_policy=deny \ --output-dir ./minimax-fpm ``` **Result to inspect:** `minimax-fpm/prediction.json` contains the serving prediction using whole-forward profiles. -**What changed:** AIC's `--forward-model` becomes a per-role -`engine.workers..timing.forward_model` setting. `op_level` remains the default. The AISimulate -fixture uses four in-flight requests rather than fixing every scheduler batch to four. The AIC -command pins FMHA precision to match the collected profile. Both commands pin vLLM 0.25.1, which -requires the shown unlisted-version override. FPM requires -`timing.type: default` and matching profile coverage; a missing profile fails explicitly. See the -[FPM guide](../../python/aisimulate/docs/fpm/README.md). +**What changed:** AIC's `--forward-model fpm` maps to the per-role +`engine.workers..timing.estimation_mode: fpm_interpolation` setting. The defaults are +`estimation_mode: auto` and `fallback_policy: deny`. Auto searches +`op_level -> fpm_interpolation -> fpm_regression` during construction, including with deny. +For an explicit mode, deny prevents switching estimators; allow tries that mode first, then the +remaining estimators in the same global priority order. Invalid configuration does not trigger +fallback. An untrained regression is not ready for offline simulation, and queries do not +silently switch estimators after construction. Saved `timing.forward_model` inputs are migrated +with their explicit mode and strict selection preserved; use `estimation_mode` in new configs. + +The AISimulate fixture uses four in-flight requests rather than fixing every scheduler batch to +four. The AIC command pins FMHA precision to match the collected profile. Both commands pin +vLLM 0.25.1, which requires the shown unlisted-version override. FPM requires +`timing.type: default` and matching profile coverage; the explicit FPM/deny example fails when +that profile is missing. See the [FPM guide](../../python/aisimulate/docs/fpm/README.md). + +#### 4.6.2 Configure data policies and estimator tuning + +**Choose performance-data and transfer policies.** This uses `HYBRID`, conservative transfer, and +the bundled system definitions: + +```bash +aiconfigurator cli estimate \ + --model-path meta-llama/Meta-Llama-3.1-8B \ + --system h200_sxm --backend vllm --backend-version 0.24.0 \ + --estimate-mode agg --batch-size 64 --tp-size 2 \ + --isl 1024 --osl 128 \ + --database-mode HYBRID --transfer-policy conservative --systems-paths default \ + --detail source +``` + +**Result to inspect:** the estimate and per-operation source breakdown show which data supplied +the prediction. A policy choice does not guarantee coverage. `SOL` selects theoretical estimates; +custom system directories can be added to `--systems-paths`. See +[database modes](legacy-aic-user-guide.md#database-mode) and +[system paths](legacy-aic-user-guide.md#systems-paths). + +**After — carry the policy into serving prediction and recommendation:** + +```bash +aisimulate predict --config prediction.yaml \ + --set engine.database_mode=HYBRID \ + --set engine.transfer_policy=conservative \ + --set 'engine.systems_paths=[default]' \ + --set engine.estimation_mode=auto --set engine.fallback_policy=deny \ + --output-dir ./policy-prediction + +aisimulate recommend --config budget-search.yaml \ + --set engine.database_mode=HYBRID \ + --set engine.transfer_policy=conservative \ + --set 'engine.systems_paths=[default]' \ + --output-dir ./policy-search +``` + +Use an ordered list of existing directories plus `default` to search custom data before the +bundled root. These controls require regular aggregated/disaggregated language workers with +default timing in every role. AFD, analytical encoder pools, and fixed/polynomial providers keep +their existing paths. Recommendation YAML pins each selected role's effective data root, version, +policy, estimator mode, and full estimator configuration for a subsequent prediction. + +`engine.estimator_config` carries supported regression/correction controls; see the +[canonical API](../core-api.md#estimator-controls). Shared role-based correction is deferred: +current native correction stores and the latest role-bound regression routing remain unchanged. +The unified CLI does not expose AIC's per-operation source breakdown shown above; keep the +compatibility CLI or SDK for that diagnostic. @@ -1236,27 +1299,11 @@ alone does not establish support for an entire CLI workflow. ### 5.6 Estimator controls and speculative decoding -Backend version and op-level/FPM selection have unified mappings. The following controls still -require AIC or the estimator SDK. - -**Choose performance-data and transfer policies.** This uses `HYBRID`, conservative transfer, and -the bundled system definitions: - -```bash -aiconfigurator cli estimate \ - --model-path meta-llama/Meta-Llama-3.1-8B \ - --system h200_sxm --backend vllm --backend-version 0.24.0 \ - --estimate-mode agg --batch-size 64 --tp-size 2 \ - --isl 1024 --osl 128 \ - --database-mode HYBRID --transfer-policy conservative --systems-paths default \ - --detail source -``` - -**Result to inspect:** the estimate and per-operation source breakdown show which data supplied -the prediction. A policy choice does not guarantee coverage. `SOL` selects theoretical estimates; -custom system directories can be added to `--systems-paths`. See -[database modes](legacy-aic-user-guide.md#database-mode) and -[system paths](legacy-aic-user-guide.md#systems-paths). +For supported estimator selection, fallback, database/transfer policies, system roots, and +regression/correction tuning, see [section 4.6](#46-select-and-configure-performance-estimators). +The controls below continue to use AIC or the estimator SDK: explicit quantization/kernel +selectors, exact cached-prefix assumptions, and non-ngram speculative decoding. Per-operation source +diagnostics also remain on the compatibility CLI or SDK. **Pin quantization and select an attention implementation.** For a dense-model decode estimate with explicit BF16 compute/cache settings and the framework's default attention implementation: diff --git a/docs/core-api.md b/docs/core-api.md index afca148cd..9ff4938c7 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -124,114 +124,127 @@ minor API update. ## Choosing a forward-pass API -For adaptive forward-pass modeling, use -`RustForwardPassPerfModel.best_available(config, worker_type, options=None)` -from Python or -`ForwardPassPerfModel::best_available(config, worker_type, options)` from Rust. -`worker_type` is an immutable property of the engine model, not an inferred -property of one FPM iteration. Python accepts exactly `"prefill"`, `"decode"`, -or `"aggregated"`; Rust uses -`ForwardPassWorkerType::{Prefill, Decode, Aggregated}`. - -This path uses the native AIC estimate when the native estimator can be built, -learns online correction factors from FPM observations, and falls back to the -regression associated with `worker_type` for eligible native build or -data-availability failures. These include unsupported models and missing or -unreadable model, system, or performance data, plus malformed system YAML. A -successful native build keeps the existing workload-kind inference and -correction behavior; `worker_type` and regression-only weights do not alter -native estimates. Check -`diagnostics()` to determine whether the active source is `aic`, -`aic_with_correction`, or `fallback_regression`, and to inspect any fallback -warning. - -Native online corrections default to an absolute factor range of `[0.5, 2.0]`. -Pass `None` explicitly as `min_faster_correction_factor` or -`max_slower_correction_factor` in the options dictionary to remove the bound -in that direction. Regression fallback ignores both options. - -Use `from_native(...)` instead when native AIC support is required and an -unsupported configuration or native data failure should surface rather than -fall back. This strict-native constructor does not take `worker_type`. - -`EngineConfig.database_mode` selects `SILICON`, `HYBRID`, `EMPIRICAL`, or -`SOL` for native forward-pass construction. The Python dictionary form uses -those uppercase strings; Rust uses `DatabaseMode`. `EMPIRICAL` always uses the -SOL/util path and is intended for research estimates when matching silicon -data is unavailable. Set `enable_shared_layer=true` when empirical estimation -may reuse sibling or cross-version calibration data; `transfer_policy` can -limit the allowed transfer kinds. `strict_provenance=true` makes missing, -malformed, or incomplete collection and reuse metadata fail the database load. -`SOL_FULL` remains a per-call diagnostic and is rejected as an engine default. -Version-slot validation is unchanged; raw versions outside the maintained slots -still require the existing explicit SDK or environment escape hatch. - -`AicEngineBuilder` exposes `.database_mode(...)`, `.shared_layer(...)`, -`.transfer_policy(...)`, and `.strict_provenance(...)`. Python `compile_engine` -accepts the corresponding keyword names; `shared_layer` is serialized as -`EngineConfig.enable_shared_layer`. - -Use `RustForwardPassPerfModel.from_regression(worker_type, options=None)` or -`ForwardPassPerfModel::from_regression(worker_type, options)` for a -regression-only model. The caller creates one instance per worker and passes -its fixed regression role. Prefill and Decode each own one retained sample -bucket and fit; Aggregated owns four buckets selected from the workload across -all active ranks. The default limit of 64 observations and minimum of five -apply independently to each bucket, for at most 256 retained observations in -an Aggregated predictor. Its feature axes are consistently ordered as -`[critical attention, global FFN/MoE]`; bucket retention uses `log1p` of those -raw features, while fitting uses standardized raw features. The optional -regression weights below default to `1.0` and must be finite and strictly -positive: - -- `regression_attention_kv_weight` (`alpha`); -- `regression_prefill_attention_pair_weight` (`beta`); -- `regression_ffn_token_weight` (`gamma`). - -The ergonomic Python facade accepts ordinary Python floats. For these three -fields only, it marshals nonfinite values on a shallow copy of the options -dictionary as the exact valid-JSON string sentinels `"NaN"`, `"Infinity"`, and -`"-Infinity"`; finite values remain JSON numbers. Callers of the JSON-oriented -raw PyO3 API may use the same sentinels directly. The sentinels preserve values -for backend-dependent validation rather than making them valid regression -weights: `from_native` and a successful native `best_available` ignore all -three fields, while `from_regression` and a fallback `best_available` reject a -decoded nonfinite value with the corresponding field-specific error. Other -strings and value types remain invalid. - -Use `regression_store_diagnostics()` to inspect each bucket's label, readiness, -and retained count. Summary `diagnostics()` reports the total count and whether -any bucket has a fit; a query for a different, cold bucket can still return -`None`. Native models return an empty list from the new method. - -The ownership model, routing examples, capacity semantics, formulas, and -compatibility rules are explained in the -[FPM regression design](../python/aisimulate/docs/fpm/aic-fpm-regression-design.md). - -`AicEngineBuilder` serves a different purpose: it constructs the strict native -Rust engine for direct public prefill and decode latency calls. It does not -provide regression fallback or online correction, so it is not a replacement -for `best_available(...)`. +Use `RustForwardPassPerfModel.best_available(config)` from Python or +`ForwardPassPerfModel::best_available(config)` from Rust. The canonical +`ForwardPassPerfModelConfig` owns model, hardware, backend, topology, data +policies, a required immutable `worker_type`, and the complete nested +`estimator_config`. Worker roles are `prefill`, `decode`, and `aggregated`. ```python -from aisimulate_core.sdk import RustForwardPassPerfModel - -# Engine-config and per-rank FPM dictionary setup is omitted here. -model = RustForwardPassPerfModel.best_available(config, "decode") -diagnostics = model.diagnostics() -print(diagnostics["source"]) -if diagnostics["last_warning"] is not None: - print(diagnostics["last_warning"]) - -estimate_ms = model.estimate_forward_pass_time_ms(metrics_by_rank) -if estimate_ms is None: - # Regression fallback starts without observations for this worker type. - # Supply observed FPM iterations with positive wall_time until the configured - # min_observations threshold is reached, then retry the estimate. - model.tune_with_fpms(observed_iterations) # Observed-iteration setup omitted. - estimate_ms = model.estimate_forward_pass_time_ms(metrics_by_rank) +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + +config = ForwardPassPerfModelConfig( + model="Qwen/Qwen3-32B", + system="h200_sxm", + backend="vllm", + worker_type="aggregated", + tp=2, + estimation_mode="auto", + fallback_policy="deny", + estimator_config={ + "features": {"attention_kv_weight": 1.0}, + "fpm_regression": { + "sampling": {"bins_per_axis": [4, 16], "max_observations": 128}, + "min_observations": 5, + }, + "correction": {"enabled": True}, + }, +) +model = RustForwardPassPerfModel.best_available(config) +print(model.diagnostics()["provenance"]) ``` +### Selection and fallback + +`estimation_mode` defaults to `auto`; `fallback_policy` defaults to `deny`. +Auto always searches `op_level`, `fpm_interpolation`, then `fpm_regression`, +including when fallback is denied. For an explicit mode, deny permits only +that estimator; allow tries the requested estimator first, then the remaining +estimators in the same global priority order. Each native mode tries the +ordered `systems_paths` before moving to another estimator. Omitted roots preserve +configured SDK discovery (or the systems-path environment override when the SDK +uses its packaged default); an explicit `default` entry selects the packaged root. +The resolved paths are shared by preflight, construction and capacity estimation. +Selection occurs +at construction; queries do not silently switch estimators on a data-domain error. + +Invalid caller configuration does not trigger fallback. A constructed regression +may be unready: nonempty queries return `None` until their selected workload +store has a usable fit. Offline prediction/recommendation reject an untrained +regression instead of fabricating a latency. The current aggregated regression +retains its four workload stores; dedicated roles retain one each. + +The returned provenance records the requested and selected modes, failed +selection attempts, effective backend version, data policy, selected root, +and complete estimator configuration. Its resolved config pins the selected +mode with deny so saved replay input repeats that selection. + +Prompt-lookup verification uses the same constructor: set `speculation` to +`{"kind": "ngram", "params": {"num_speculative_tokens": 2}}` in Python/JSON, or +`ForwardPassSpeculationConfig::Ngram { num_speculative_tokens: 2 }` in Rust. +It supports vLLM op-level timing with 1–5 draft tokens and `nextn: 0`; auto can +select op-level but cannot fall back to an unsupported speculative estimator. +The cost configuration is retained in provenance and saved recommendations. +Acceptance rates and the scheduler seed stay in the CLI/Replay speculation +configuration; they do not change the model's target-verification graph. + +### Estimator controls + +`estimator_config` is passed intact through the Python facade, CLI, Sweeper, +and Replay, then parsed and validated in Rust. Unknown fields report their +nested paths. The supported namespaces are: + +- `features`: `attention_kv_weight`, `prefill_attention_pair_weight`, and + `ffn_token_weight`, each defaulting to 1.0. These currently affect regression + only; positive finite values are required when regression is constructed. +- `fpm_regression`: independent `sampling`, `min_observations` (5), and `fit`. + The fit kind is `standardized_nnls`, with a free intercept and nonnegative + slopes. `singular_ridge_scale` defaults to `1e-9` and applies only when retrying + a singular equation. +- `correction`: `enabled` (true), independent `sampling`, `min_observations` + (5), `factor_bounds` (min 0.5, max 2.0), and the existing `max_num_tokens` + (8192), `max_batch_size` (512), and `max_kv_tokens` (2000000) ranges. +- `op_level` and `fpm_interpolation`: reserved typed namespaces with no + additional knobs yet; unknown fields are rejected. + +Sampling defaults to `bins_per_axis: [4, 4]` and `max_observations: 64` per +logical store. Rectangular grids are supported. Regression uses dynamic +`log1p` retention coordinates and fits standardized raw features. Correction +uses fixed raw workload coordinates; its one-dimensional prefill grid uses +the product of the two axis counts. Retention evicts the oldest sample from +the most populated cell when the store exceeds its budget. + +Correction explicitly reports `feature_space: legacy_workload`. Its existing +prefill/decode/mixed stores and median-ratio calculation remain unchanged at +default settings. A shared role-based correction space requires separate +accuracy validation and is not accepted as a configuration value in this release. + +Use `regression_store_diagnostics()` for per-store counts/readiness. Summary +readiness means at least one store is ready; another cold store can still +return `None`. `tune_with_fpms()` preserves the established FPM observation +contract. Native construction still uses Python model compilation; estimator +selection, regression, correction, and latency computation are owned by Rust. + +### Migrating saved configuration + +Use `ForwardPassPerfModelConfig.from_legacy_engine_config(old_config, +worker_type, old_options, allow_regression=False)` to convert a saved flat +EngineConfig and tuning options. It pins the old explicit native mode instead +of changing it to auto. Set `allow_regression=True` only for an old caller that +allowed direct regression fallback; the migration preserves that two-mode +order rather than adding interpolation. Legacy `forward_model: fpm` maps to +`fpm_interpolation`, and `fallback_policy: error` maps to deny. The deprecated +`regression` policy remains readable for these saved direct-fallback requests. + +Previously saved CLI timing with `forward_model` retains explicit selection +and deny. Newly authored requests without a selection use auto. `ForwardPassPerfOptions` +is retained as a legacy migration value type; new construction has one complete +config and no separate options argument. The raw PyO3 class also exposes +`normalize_config` and migration helpers for JSON-oriented consumers. + +The [FPM regression design](../python/aisimulate/docs/fpm/aic-fpm-regression-design.md) +explains the retained workload routing and feature mathematics. + ## Stable Rust facade New embedded consumers should construct engines with @@ -242,9 +255,10 @@ Python once to compile an engine specification. Calls on the returned crate-root types remain re-exported during the migration window, but the `perfmodel` namespace is canonical for new code. -Standalone binaries must enable the crate's `embed-python` feature; applications -hosted by an initialized Python interpreter do not. In either case, the matching -`aisimulate` wheel must be importable. See the +Native engine construction in standalone binaries requires the crate's `embed-python` +feature; applications hosted by an initialized Python interpreter do not need +auto-initialization. The matching `aisimulate` wheel must be importable for native +construction. Explicit `fpm_regression` construction works without the Python feature. See the [crate README](../crates/core/README.md) for setup and usage examples. The flat `build_aic_engine` adapter was removed from `main`; consumers must use @@ -254,7 +268,7 @@ The supported `aisimulate_core::perfmodel` Rust surface is grouped as follows: - compiled engine: `AicEngineBuilder`, `AicEngine`, `AicError`; - forward-pass estimation: `ForwardPassPerfModel`, - `ForwardPassWorkerType`, `ForwardPassPerfOptions`, + `ForwardPassWorkerType`, `ForwardPassPerfModelConfig`, `EstimatorConfig`, diagnostics/readiness/source types, and the `ForwardPassMetrics` telemetry types; - KV-cache estimation: `estimate_kv_cache`, `KvCacheEstimateRequest`, @@ -310,19 +324,36 @@ past-KV coordinate rather than the op-level mean-context coordinate. - A supported facade name is not removed or given a new required parameter without a documented deprecation path. The package is pre-1.0, so an unavoidable incompatible API change also requires a minor-version bump. -- Requiring `worker_type` in `from_regression` and `best_available` is a - documented incompatible change for the next minor release. Adding the three - public regression-weight fields to `ForwardPassPerfOptions` is also a Rust - source break for downstream exhaustive struct literals. Rust callers should - prefer update syntax such as - `ForwardPassPerfOptions { bucket_count: 16, ..Default::default() }` so future - option fields do not require source changes. This feature does not itself - change package versions; release coordination must keep the crate and wheel - versions aligned. +- The canonical `best_available(config)` API replaces the old positional worker + role/options signatures and separate estimator constructors. This is a source + migration: use `ForwardPassPerfModelConfig::new(...)` in Rust or the SDK config + class in Python, and use the explicit migration helper for saved EngineConfig + values. Downstream Dynamo callers must migrate before this API is released; + keep the crate and wheel versions aligned at the coordinated minor release. +- Publication is blocked by [the release gate](../.github/release-gates.json) + until [Dynamo #14065](https://github.com/ai-dynamo/dynamo/pull/14065) is refreshed, + merged, and its Planner/wheel smoke validated against this API. Both scheduled + nightly CI and approved manual dispatch run `scripts/check_release_migrations.py` + before staging and the downstream publish trigger. Manual dispatch may select + a main/release commit, but both the workflow revision's policy and the selected + commit's migration declarations must pass before publication. The checker runs + from the workflow revision; missing or malformed target gates fail closed. + Clear the pending entry in a reviewed change only after the migration evidence + is available. - The raw PyO3 class and ergonomic SDK wrapper intentionally share the name `RustForwardPassPerfModel`; callers should import from `aisimulate_core.sdk` unless they specifically need the JSON-oriented native binding. +The standalone AIC 0.12 compatibility distributions form a separate package +lineage: `aiconfigurator` pins `aiconfigurator-core==0.12.0`, and that core builds +its own native binding. They do not depend on the AISimulate wheel or crate. +Their retained binding declarations are not downstream callers of this API. +The [replacement installation](../README.md#upgrade-from-standalone-aiconfigurator) +removes both old distributions and installs AISimulate's complete migrated +application and core. Namespace compatibility does not preserve the removed +constructor signatures; direct SDK callers must perform the source migration +above. Release gates track consumers of the new artifacts, such as Dynamo. + ## CI contract Every change is checked from three consumer viewpoints: diff --git a/docs/sweeper/architecture.md b/docs/sweeper/architecture.md index 505f3b3de..45bfd73c4 100644 --- a/docs/sweeper/architecture.md +++ b/docs/sweeper/architecture.md @@ -63,7 +63,8 @@ flowchart TD C --> D["Resolve configured providers"] D --> E["Generate namespaced search dimensions"] E --> F["Ask sampler for suggestions"] - F --> G["Materialize backend and adapter config"] + F --> R["Resolve exact per-role estimator identities through Core"] + R --> G["Materialize backend and adapter config"] G --> H["Build ReplaySpec"] H --> I["Worker-local Runner executes replay"] I --> J["Score and tell sampler"] diff --git a/docs/sweeper/configuration.md b/docs/sweeper/configuration.md index d661217a4..c2b3c7c02 100644 --- a/docs/sweeper/configuration.md +++ b/docs/sweeper/configuration.md @@ -26,6 +26,17 @@ search_space: gpu_budget: 32 deployment_mode: [disagg, agg] backend: [vllm, sglang] + backend_version: + vllm: current + sglang: current + database_mode: HYBRID + transfer_policy: balanced + estimation_mode: auto + fallback_policy: deny + estimator_config: + fpm_regression: + min_observations: 5 + systems_paths: [default] adapters: example.policy: @@ -67,6 +78,13 @@ configuration for each candidate. | `decode_hardware_sku` | `None` | optional disaggregated-decode system override; inherits `hardware_sku` | | `deployment_mode` | `[disagg, agg]` | deployment branches to search | | `backend` | `[vllm]` | engine backends to search | +| `backend_version` | `None` | exact version for one backend, or a per-backend version mapping; omitted backends resolve once to latest | +| `database_mode` | `SILICON` | forward-pass estimator data-source policy; see below | +| `transfer_policy` | `None` (all) | Core-owned empirical-transfer preset or tier list; used only by `HYBRID` and `EMPIRICAL` | +| `estimation_mode` | `auto` | search op-level, FPM interpolation, then regression; or select an explicit mode | +| `fallback_policy` | `deny` | constrain explicit estimator selection; auto always searches the full priority order | +| `estimator_config` | `{}` (Core defaults) | runtime tuning controls such as observation limits, regression buckets, correction bounds, and workload-axis capacity | +| `systems_paths` | omitted | preserve configured SDK/environment discovery; an explicit list sets ordered request-scoped roots, with `default` selecting the packaged Core root | | `gpu_budget` | `32` | maximum GPUs per candidate | | `min_gpu_budget` | `None` | optional lower bound during enumeration | | `context_length` | `None` | optional KV-feasibility sequence length | @@ -137,6 +155,38 @@ current performance-model adapter also requires concrete positive `isl` and `osl synthetic request rate or absolute concurrency rather than a trace-only workload. The complete topology and capability contract is documented in [AFD Topology Contract](afd-topology.md). +Estimator controls are parsed once, then each concrete regular language-worker candidate +resolves through Core after TP, PP, attention-DP, MoE parallelism, and block size are known. +`estimation_mode` defaults to `auto` and searches op-level, FPM interpolation, then regression. +`fallback_policy` defaults to `deny`: it constrains explicit modes and does not disable auto +selection. `estimator_config` carries the complete nested configuration described in the +[Core API](../core-api.md#estimator-controls). Regression must have a ready fit before an +offline trial can execute; unavailable or cold estimators fail candidate materialization. + +Core pins the effective backend version, estimator mode, transfer policy, selected system root, +and all estimator controls. Replay and candidate artifacts preserve that resolved identity. +Generated prediction YAML records role-specific pins and can be passed directly to `predict`. +AFD, analytical encoder, and custom fixed/polynomial timing paths retain their existing providers; +nondefault public estimator policies require regular language workers with default timing. + +Database modes choose the source of each operation estimate: + +| Mode | Resolution | +|---|---| +| `SILICON` | collected performance data and supported interpolation only | +| `HYBRID` | collected data first; calibrated empirical estimation for uncovered operations | +| `EMPIRICAL` | calibrated estimation for every operation (`latency = SOL / utilization`) | +| `SOL` | uncalibrated analytic speed-of-light estimate | + +`SOL_FULL` is a low-level per-call diagnostic, not an estimator default; public +prediction and Sweeper configuration reject it. + +`transfer_policy` is not a search dimension. It selects which fixed Core transfer kinds the +empirical forward-pass estimator may use when its own calibration slice is missing. It accepts `off`, +`conservative`, `balanced`, or `aggressive`, or an explicit list containing `xshape`, `xquant`, +`xprofile`, and `xop`. Core validates the request and the Sweeper records the normalized explicit +policy in every candidate. The field is ignored by `SILICON` and `SOL`. + ## Pinned Parallel Configurations Pinning `parallel_configs` requires exactly one deployment mode. An aggregated entry is one shape: diff --git a/python/aisimulate/.claude/rules/perfmodel-api.md b/python/aisimulate/.claude/rules/perfmodel-api.md new file mode 100644 index 000000000..0fc9105a7 --- /dev/null +++ b/python/aisimulate/.claude/rules/perfmodel-api.md @@ -0,0 +1,59 @@ +--- +description: Canonical performance-model interface and extension requirements. +paths: + - "crates/core/src/perfmodel/**" + - "crates/core/src/python.rs" + - "crates/tests/public-api/**" + - "python/aisimulate/src/aiconfigurator_core/**" + - "python/aisimulate/src/aiconfigurator/**" + - "python/aisimulate/src/aisimulate/**" + - "python/aisimulate/tests/**" + - "tests/**" + - "scripts/fpm_accuracy/**" +--- + + + +# One Performance Model Interface + +`ForwardPassPerfModel::best_available(ForwardPassPerfModelConfig)` in Rust and +`RustForwardPassPerfModel.best_available(config)` in Python are the **only +production performance-model construction interface**. All new performance-model +features MUST be exposed through this interface and its returned model. CLI, +Sweeper, Replay, Planner, and other consumers must use the same contract. + +- Extend the Rust-owned typed configuration and model API first. Put estimator + controls in `estimator_config`; keep model, hardware, backend, topology, and + immutable `worker_type` in the canonical identity. Rust owns defaults, + validation, selection, and estimation behavior. Python bindings and callers + pass the complete configuration through without a second set of defaults or + independent estimator selection. +- Do not add alternative constructors, caller-specific factory functions, or + separate configuration schemas for new features. Existing low-level engine + builders, compilation helpers, and compatibility adapters are implementation + or migration machinery; they must not become a parallel public path for new + performance-model features. Existing custom timing, AFD, and encoder providers + do not authorize bypassing this interface for new estimator functionality. +- Preserve the complete resolved identity through candidate construction, + caching, metadata, replay, and saved configurations. Hardware/version/capacity + preflight must honor the same role-specific data roots. Resolve every default + timing role even when another role uses a custom timing provider; unsupported + combinations must fail explicitly rather than discard settings. +- Defaults are `estimation_mode: auto` and `fallback_policy: deny`. Auto always + searches `op_level -> fpm_interpolation -> fpm_regression`, including with deny. + Deny restricts an explicit mode; allow tries it first and then the remaining + global priority. Invalid configuration must not trigger fallback. An untrained + regression reports not-ready; offline simulation must reject it. Construction + pins the selected estimator; queries do not silently switch models. +- Migrate legacy configuration at the input boundary. New configuration must + preserve its selection and all controls after serialization and reload. Never + infer an explicit user choice from a serialized legacy default. +- Cover the public Rust/Python contract and affected caller paths with behavioral + regression tests, including saved-config round trips. Estimation math and + data-selection changes must also follow [Rust core parity rules](rust-core/parity.md). + +See the [Core API](../../../../docs/core-api.md#choosing-a-forward-pass-api) for +the schema and supported controls. diff --git a/python/aisimulate/.claude/rules/repo-guide.md b/python/aisimulate/.claude/rules/repo-guide.md index 0e4f1f42b..631629645 100644 --- a/python/aisimulate/.claude/rules/repo-guide.md +++ b/python/aisimulate/.claude/rules/repo-guide.md @@ -23,6 +23,8 @@ unless the task actually targets `src/aiconfigurator/generator/`. ## Governed areas +- Performance models and their Rust/Python/CLI/Sweeper/Replay/Planner consumers + → read [perfmodel-api.md](perfmodel-api.md) before extending or migrating an API. - Editing `src/aiconfigurator/generator/**` → generator rules auto-load; entry point `.claude/rules/generator-development.md`. - Editing `collector/**` → collector rules auto-load; read diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index b5bf95f64..9884c861a 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -879,4 +879,5 @@ Copyright 2018- The Hugging Face team. All rights reserved. - License: Apache-2.0; NVIDIA maintainer confirmed permission to migrate and publish this code under Apache-2.0. - Modified: development-only two-predictor evaluation, public overview export, - GitHub Pages presentation, and local import paths. No Plotly assets included. + GitHub Pages presentation, local import paths, and canonical estimator API + adaptation with older-wheel compatibility. No Plotly assets included. diff --git a/python/aisimulate/docs/fpm/aic-fpm-regression-design.md b/python/aisimulate/docs/fpm/aic-fpm-regression-design.md index 859cdbeed..0e5cacdb0 100644 --- a/python/aisimulate/docs/fpm/aic-fpm-regression-design.md +++ b/python/aisimulate/docs/fpm/aic-fpm-regression-design.md @@ -213,7 +213,7 @@ Aggregated regression predictor: |---|---:|---| | `max_observations` | 64 | Maximum number of retained observations | | `min_observations` | 5 | Minimum retained observations before a fit can be ready | -| `bucket_count` | 16 | Internal retention grid: four cells on each of the two feature axes | +| `sampling.bins_per_axis` | `[4, 4]` | Internal retention grid, independently configurable per axis | The resulting capacity is: @@ -261,12 +261,16 @@ The public Rust method returns a vector of as JSON, and the Python facade returns a list of dictionaries: ```python -from aisimulate_core.sdk import RustForwardPassPerfModel +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel -model = RustForwardPassPerfModel.from_regression( - "aggregated", - {"max_observations": 64, "min_observations": 5, "bucket_count": 16}, -) +model = RustForwardPassPerfModel.best_available(ForwardPassPerfModelConfig( + model="Qwen/Qwen3-32B", system="h200_sxm", backend="vllm", + worker_type="aggregated", estimation_mode="fpm_regression", + estimator_config={"fpm_regression": { + "sampling": {"max_observations": 64, "bins_per_axis": [4, 4]}, + "min_observations": 5, + }}, +)) summary = model.diagnostics() buckets = model.regression_store_diagnostics() @@ -335,19 +339,21 @@ The caller supplies complete, unique, iteration-aligned rank sets and correct worker identities. This PR does not add expected-rank-count or worker-lifecycle enforcement inside AISim. Gym retains its existing parser exclusions. -The following signatures already exist and remain available: +Construction uses one canonical config in both languages: ```text -Rust: from_regression(worker_type, options) - best_available(config, worker_type, options) - best_available_with_roots(config, worker_type, options, systems_root) -Python: from_regression(worker_type, options=None) - best_available(config, worker_type, options=None) +Rust: ForwardPassPerfModel::best_available(config) +Python: RustForwardPassPerfModel.best_available(config) ``` -`from_native` and `from_native_with_roots` remain role-free. Each regression -instance keeps its role for its lifetime. A caller that needs another role -creates a new predictor. +`config.worker_type` is required and fixed for the model's lifetime. Select +`estimation_mode="fpm_regression"` for regression-only construction. Auto searches +op-level, whole-forward interpolation, then regression even when the default +`fallback_policy="deny"` is in effect; deny constrains an explicit estimator +selection. Sampling and fitting controls live under `estimator_config.fpm_regression`. +The workload routing and fits described here remain unchanged. Native correction +continues to use its legacy workload stores pending accuracy validation of a +role-based replacement. See the [canonical API](../../../../docs/core-api.md#choosing-a-forward-pass-api). ## 9. Validation and limits @@ -509,11 +515,11 @@ cross-rank summation, and derived features must be finite and nonnegative. The options are construction-time knobs and all default to `1.0`: -| Formula | `ForwardPassPerfOptions` field | +| Formula | `estimator_config.features` field | |---|---| -| $\alpha$ | `regression_attention_kv_weight` | -| $\beta$ | `regression_prefill_attention_pair_weight` | -| $\gamma$ | `regression_ffn_token_weight` | +| $\alpha$ | `attention_kv_weight` | +| $\beta$ | `prefill_attention_pair_weight` | +| $\gamma$ | `ffn_token_weight` | - $\alpha$ scales KV-token-related attention work for every role: $\widetilde H$ for Prefill, $K$ for Decode, and $\widetilde H+K$ for @@ -529,16 +535,10 @@ constructing a new model, because it changes the feature and retention-cell coordinates. To preserve learned history, replay the original per-rank FPM observations into the new model. -The ergonomic Python facade accepts ordinary Python floats and marshals the -three weight fields on a shallow copy of the caller's options dictionary. -Finite values remain JSON numbers; nonfinite values use the exact valid-JSON -string sentinels `"NaN"`, `"Infinity"`, and `"-Infinity"`. JSON-oriented raw -PyO3 callers may send those same sentinels. Rust deserialization maps only -those exact strings back to their corresponding `f64` values; unknown strings -and other value types are invalid. This transport does not relax validation: -`from_native` and a successful native `best_available` ignore all three -weights, whereas `from_regression` and a fallback `best_available` reject a -decoded nonfinite value with the corresponding field-specific error. +New configuration uses the nested feature-weight names above. The legacy +EngineConfig/options migration adapter retains the old flat names and exact +nonfinite JSON sentinels. These weights are ignored by a selected native model; +regression construction rejects nonpositive or nonfinite weights. ## Reference: retention and fit pipeline @@ -552,7 +552,7 @@ $$ The existing dynamic two-dimensional grid consumes these continuous `f64` coordinates unchanged. Its bounds expand and trigger rebucketing, never shrink, and its fattest-cell eviction policy still enforces the per-bucket sample -cap. With the default `bucket_count=16`, the grid is $4\times4$. Buckets +cap. With the default `sampling.bins_per_axis=[4, 4]`, the grid is $4\times4$. Buckets choose which observations survive; they are not local predictors and are not queried during estimation. diff --git a/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi b/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi index 49bdd37d8..7e4a28598 100644 --- a/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi +++ b/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi @@ -165,13 +165,15 @@ class AicEngine: class RustForwardPassPerfModel: @staticmethod - def from_native(config_json: str, options_json: str | None = None) -> RustForwardPassPerfModel: ... + def best_available(config_json: str) -> RustForwardPassPerfModel: ... @staticmethod - def best_available( - config_json: str, worker_type: str, options_json: str | None = None - ) -> RustForwardPassPerfModel: ... + def normalize_config(config_json: str) -> str: ... @staticmethod - def from_regression(worker_type: str, options_json: str | None = None) -> RustForwardPassPerfModel: ... + def legacy_estimator_config(options_json: str) -> str: ... + @staticmethod + def migrate_legacy_config( + config_json: str, worker_type: str, options_json: str | None = None, allow_regression: bool = False + ) -> str: ... def estimate_forward_pass_time_ms(self, fpm_json: str) -> float | None: ... def tune_with_fpms(self, iterations_json: str) -> None: ... def diagnostics(self) -> str: ... diff --git a/python/aisimulate/src/aiconfigurator_core/sdk/__init__.py b/python/aisimulate/src/aiconfigurator_core/sdk/__init__.py index a9f129355..0367386f5 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/__init__.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/__init__.py @@ -21,6 +21,8 @@ __all__ = [ "AttentionBackend", "EngineHandle", + "ForwardPassPerfModelConfig", + "ForwardPassPerfOptions", "MoEBackend", "ModelConfig", "RuntimeConfig", @@ -33,6 +35,14 @@ _PUBLIC_EXPORTS = { "AttentionBackend": ("aiconfigurator_core.sdk.common", "AttentionBackend"), "EngineHandle": ("aiconfigurator_core.sdk.engine", "EngineHandle"), + "ForwardPassPerfModelConfig": ( + "aiconfigurator_core.sdk.rust_engine_step", + "ForwardPassPerfModelConfig", + ), + "ForwardPassPerfOptions": ( + "aiconfigurator_core.sdk.rust_engine_step", + "ForwardPassPerfOptions", + ), "ModelConfig": ("aiconfigurator_core.sdk.config", "ModelConfig"), "MoEBackend": ("aiconfigurator_core.sdk.common", "MoEBackend"), "RuntimeConfig": ("aiconfigurator_core.sdk.config", "RuntimeConfig"), @@ -69,4 +79,8 @@ def __dir__() -> list[str]: from aiconfigurator_core.sdk.config import ModelConfig, RuntimeConfig from aiconfigurator_core.sdk.engine import EngineHandle, compile_engine from aiconfigurator_core.sdk.memory import estimate_kv_cache, estimate_num_gpu_blocks - from aiconfigurator_core.sdk.rust_engine_step import RustForwardPassPerfModel + from aiconfigurator_core.sdk.rust_engine_step import ( + ForwardPassPerfModelConfig, + ForwardPassPerfOptions, + RustForwardPassPerfModel, + ) diff --git a/python/aisimulate/src/aiconfigurator_core/sdk/engine.py b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py index d1d4fbae3..8ad136b87 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/engine.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py @@ -45,6 +45,7 @@ import aiconfigurator_core from aiconfigurator_core.sdk.config_builders import apply_nextn, build_model_config +from aiconfigurator_core.sdk.errors import InvalidEngineConfigurationError as InvalidEngineConfigurationError from aiconfigurator_core.sdk.models import get_model from aiconfigurator_core.sdk.operations import FPMForwardOp from aiconfigurator_core.sdk.operations.base import Operation @@ -436,27 +437,30 @@ def compile_engine( # does not take a model_path (quant inference is done inside `get_model`). from aiconfigurator_core.sdk.speculation import SpeculationConfig - resolved_speculation = SpeculationConfig(**speculation) if speculation is not None else None resolved_moe_tp = moe_tp_size if moe_tp_size is not None else 1 resolved_moe_ep = moe_ep_size if moe_ep_size is not None else 1 - model_config = build_model_config( - tp_size=tp_size, - pp_size=pp_size, - attention_dp_size=attention_dp_size, - moe_tp_size=resolved_moe_tp, - moe_ep_size=resolved_moe_ep, - gemm_quant_mode=gemm_quant_mode, - kvcache_quant_mode=kvcache_quant_mode, - fmha_quant_mode=fmha_quant_mode, - moe_quant_mode=moe_quant_mode, - comm_quant_mode=comm_quant_mode, - forward_model=forward_model, - attention_backend=attention_backend, - speculation=resolved_speculation, - ) - # Apply MTP BEFORE get_model so the walked op lists carry the - # (L+nextn)/L compute scale; accepted-token progress is applied above core. - apply_nextn(model_config, nextn) + try: + resolved_speculation = SpeculationConfig(**speculation) if speculation is not None else None + model_config = build_model_config( + tp_size=tp_size, + pp_size=pp_size, + attention_dp_size=attention_dp_size, + moe_tp_size=resolved_moe_tp, + moe_ep_size=resolved_moe_ep, + gemm_quant_mode=gemm_quant_mode, + kvcache_quant_mode=kvcache_quant_mode, + fmha_quant_mode=fmha_quant_mode, + moe_quant_mode=moe_quant_mode, + comm_quant_mode=comm_quant_mode, + forward_model=forward_model, + attention_backend=attention_backend, + speculation=resolved_speculation, + ) + # Apply MTP BEFORE get_model so the walked op lists carry the + # (L+nextn)/L compute scale; accepted-token progress is applied above core. + apply_nextn(model_config, nextn) + except (ValueError, TypeError, KeyError) as exc: + raise InvalidEngineConfigurationError(str(exc)) from exc model = get_model(model_path, model_config, backend) # Slot policy FIRST, tolerance second: resolve the requested version to a diff --git a/python/aisimulate/src/aiconfigurator_core/sdk/errors.py b/python/aisimulate/src/aiconfigurator_core/sdk/errors.py index b6051a7e3..9216ffbcf 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/errors.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/errors.py @@ -8,6 +8,10 @@ import dataclasses +class InvalidEngineConfigurationError(ValueError): + """Invalid caller controls; estimator fallback must not hide these errors.""" + + class NoResultsError(RuntimeError): """Base class for *expected* "the sweep produced no results" outcomes. diff --git a/python/aisimulate/src/aiconfigurator_core/sdk/models/__init__.py b/python/aisimulate/src/aiconfigurator_core/sdk/models/__init__.py index eb6a7d986..42279e12b 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/models/__init__.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/models/__init__.py @@ -31,6 +31,7 @@ import pkgutil from aiconfigurator_core.sdk import config +from aiconfigurator_core.sdk.errors import InvalidEngineConfigurationError from aiconfigurator_core.sdk.models.base import _MODEL_REGISTRY, BaseModel from aiconfigurator_core.sdk.models.helpers import ( _apply_model_quant_defaults, @@ -134,7 +135,9 @@ def get_model( """ forward_model = getattr(model_config, "forward_model", "op_level") or "op_level" if forward_model not in _FORWARD_MODELS: - raise ValueError(f"Unknown forward_model: {forward_model!r}. Valid values: {', '.join(_FORWARD_MODELS)}") + raise InvalidEngineConfigurationError( + f"Unknown forward_model: {forward_model!r}. Valid values: {', '.join(_FORWARD_MODELS)}" + ) # Shallow-copy so mutations below don't poison the @cache'd original. model_info = dict(_get_model_info(model_path)) @@ -151,7 +154,10 @@ def get_model( ) _apply_model_quant_defaults(model_config, raw_config, architecture, backend_name) if check_is_moe(model_path, model_info=model_info): - model_config.resolve_moe_parallelism() + try: + model_config.resolve_moe_parallelism() + except (ValueError, TypeError, KeyError) as exc: + raise InvalidEngineConfigurationError(str(exc)) from exc if model_config.overwrite_num_layers > 0: model_info["layers"] = model_config.overwrite_num_layers @@ -179,7 +185,7 @@ def get_model( ) # sglang CP requires the attention side to be pure CP (no concurrent attn TP/DP). if backend_name == "sglang" and (model_config.tp_size != 1 or model_config.attention_dp_size != 1): - raise ValueError( + raise InvalidEngineConfigurationError( f"sglang CP requires tp_size=1 and attention_dp_size=1 when cp_size>1 " f"(CP and attention TP/DP are mutually exclusive on sglang). Got " f"tp_size={model_config.tp_size}, attention_dp_size={model_config.attention_dp_size}, " diff --git a/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py b/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py index 57769a84b..768272794 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py @@ -21,6 +21,9 @@ import os import threading from collections import OrderedDict +from collections.abc import Mapping +from dataclasses import asdict, dataclass +from dataclasses import field as dataclass_field from importlib import resources as pkg_resources from pathlib import Path from typing import Any @@ -96,6 +99,87 @@ class RustEngineUnsupportedError(RuntimeError): error-symmetric between the engines.""" +@dataclass(frozen=True) +class ForwardPassPerfModelConfig: + """Canonical immutable identity and selection policy for the estimator.""" + + model: str + system: str + backend: str + worker_type: str + backend_version: str | None = None + tp: int = 1 + pp: int = 1 + attention_dp: int = 1 + moe_tp_size: int | None = None + moe_ep_size: int | None = None + gemm_quant_mode: str | None = None + moe_quant_mode: str | None = None + fmha_quant_mode: str | None = None + kvcache_quant_mode: str | None = None + comm_quant_mode: str | None = None + nextn: int = 0 + speculation: dict[str, Any] | None = None + kv_block_size: int | None = None + estimation_mode: str = "auto" + database_mode: str = "SILICON" + transfer_policy: str | tuple[str, ...] | None = None + systems_paths: tuple[str, ...] = () + fallback_policy: str = "deny" + estimator_config: dict[str, Any] = dataclass_field(default_factory=dict) + attention_backend: str | None = None + enable_shared_layer: bool | None = None + strict_provenance: bool = False + + @classmethod + def from_legacy_engine_config( + cls, + config: Mapping[str, Any], + worker_type: str, + options: ForwardPassPerfOptions | Mapping[str, Any] | None = None, + *, + allow_regression: bool = False, + ) -> ForwardPassPerfModelConfig: + """Migrate saved configuration while preserving strict/native-only selection.""" + import aiconfigurator_core + + payload = aiconfigurator_core.RustForwardPassPerfModel.migrate_legacy_config( + _json_dumps(dict(config)), + worker_type, + _optional_json_dumps(options.to_dict() if isinstance(options, ForwardPassPerfOptions) else options), + allow_regression, + ) + return cls(**json.loads(payload)) + + def to_dict(self) -> dict[str, Any]: + payload = asdict(self) + payload["transfer_policy"] = _resolve_forward_pass_transfer_policy(self.transfer_policy) + payload["systems_paths"] = _resolve_forward_pass_systems_paths(self.systems_paths) + return payload + + +@dataclass(frozen=True) +class ForwardPassPerfOptions: + """Runtime observation, regression, correction, and capacity controls.""" + + max_observations: int = 64 + min_observations: int = 5 + min_faster_correction_factor: float | None = 0.5 + max_slower_correction_factor: float | None = 2.0 + bucket_count: int = 16 + max_num_tokens: int = 8192 + max_batch_size: int = 512 + max_kv_tokens: int = 2_000_000 + regression_attention_kv_weight: float = 1.0 + regression_prefill_attention_pair_weight: float = 1.0 + regression_ffn_token_weight: float = 1.0 + bucket_shape: tuple[int, int] | None = None + regression_ridge_scale: float = 1e-9 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + class RustForwardPassPerfModel: """Facade over the compiled Rust forward-pass perf model (PR #1152). @@ -140,7 +224,8 @@ class RustForwardPassPerfModel: iter1_rank1]]``. Each backend derives its own cross-rank features, and the maximum positive ``wall_time`` across ranks is the latency target. - Correction grids use fixed constructor-time ranges from ``options``: + Correction grids use fixed constructor-time ranges from + ``config.estimator_config["correction"]``: ``max_num_tokens`` bounds ``sum_prefill_tokens`` and defaults to ``8192``, ``max_batch_size`` bounds ``num_decode_requests`` and defaults to ``512``, and ``max_kv_tokens`` bounds ``sum_decode_kv_tokens`` and defaults to @@ -166,79 +251,29 @@ def __init__(self, inner: Any) -> None: self._inner = inner @classmethod - def from_native( - cls, - config: dict[str, Any], - options: dict[str, Any] | None = None, - ) -> RustForwardPassPerfModel: - """API: ``RustForwardPassPerfModel.from_native(config, options=None)``. - - Description: create a strict native AIC forward-pass model. + def best_available(cls, config: ForwardPassPerfModelConfig | Mapping[str, Any]) -> RustForwardPassPerfModel: + """Construct one estimator from the canonical role-bound configuration. - Crosses into the Rust core, which compiles ``config`` via - ``aiconfigurator_core.sdk.engine.compile_engine``. Raises if the config is - unsupported by the native estimator. Use ``best_available()`` when - unsupported configs should fall back to the learned regression model. + Auto searches op_level, fpm_interpolation, then fpm_regression even + with fallback_policy=deny. Explicit modes default to strict selection. + Native correction retains its existing workload feature space. """ - _configure_default_data_roots() import aiconfigurator_core - inner = aiconfigurator_core.RustForwardPassPerfModel.from_native( - _json_dumps(config), - _optional_json_dumps(options), - ) - return cls(inner) - - @classmethod - def best_available( - cls, - config: dict[str, Any], - worker_type: str, - options: dict[str, Any] | None = None, - ) -> RustForwardPassPerfModel: - """API: ``best_available(config, worker_type, options=None)``. - - Description: create a native model when possible, otherwise fall back to - regression bound to ``worker_type``. Fallback reason is available from - ``diagnostics()["last_warning"]``. ``worker_type`` must be exactly - ``"prefill"``, ``"decode"``, or ``"aggregated"``; a successful native - construction does not otherwise use it. - """ - worker_type = _validate_worker_type(worker_type) - _configure_default_data_roots() - import aiconfigurator_core - - inner = aiconfigurator_core.RustForwardPassPerfModel.best_available( - _json_dumps(config), - worker_type, - _optional_json_dumps(options), - ) - return cls(inner) - - @classmethod - def from_regression( - cls, - worker_type: str, - options: dict[str, Any] | None = None, - ) -> RustForwardPassPerfModel: - """API: ``from_regression(worker_type, options=None)``. - - Description: create a regression-only forward-pass model bound to one - engine-level ``worker_type``. It returns ``None`` for non-empty estimates - until the selected workload store has a ready fit from samples provided through - ``tune_with_fpms()``. Correction factor getters return ``None`` in this - mode. ``worker_type`` must be exactly ``"prefill"``, ``"decode"``, or - ``"aggregated"``. - """ - worker_type = _validate_worker_type(worker_type) - _configure_default_data_roots() - import aiconfigurator_core - - inner = aiconfigurator_core.RustForwardPassPerfModel.from_regression( - worker_type, - _optional_json_dumps(options), - ) - return cls(inner) + payload = config.to_dict() if isinstance(config, ForwardPassPerfModelConfig) else dict(config) + if payload.get("estimation_mode") != "fpm_regression" or payload.get("systems_paths"): + payload["systems_paths"] = _resolve_forward_pass_systems_paths(tuple(payload.get("systems_paths") or ())) + if "transfer_policy" in payload: + payload["transfer_policy"] = _resolve_forward_pass_transfer_policy(payload["transfer_policy"]) + estimator_config = payload.get("estimator_config") + if isinstance(estimator_config, Mapping) and isinstance(estimator_config.get("features"), Mapping): + features = dict(estimator_config["features"]) + for name in ("attention_kv_weight", "prefill_attention_pair_weight", "ffn_token_weight"): + weight = features.get(name) + if isinstance(weight, float) and not math.isfinite(weight): + features[name] = "NaN" if math.isnan(weight) else "Infinity" if weight > 0 else "-Infinity" + payload["estimator_config"] = {**estimator_config, "features": features} + return cls(aiconfigurator_core.RustForwardPassPerfModel.best_available(_json_dumps(payload))) def estimate_forward_pass_time_ms(self, metrics: dict[str, Any] | list[dict[str, Any]]) -> float | None: """API: ``model.estimate_forward_pass_time_ms(metrics) -> float | None``. @@ -322,10 +357,10 @@ def _json_dumps(value: Any) -> str: return json.dumps(value, separators=(",", ":"), sort_keys=True) -def _optional_json_dumps(value: dict[str, Any] | None) -> str | None: +def _optional_json_dumps(value: Mapping[str, Any] | None) -> str | None: if value is None: return None - wire_value = value.copy() + wire_value = dict(value) for field in _REGRESSION_WEIGHT_FIELDS: weight = wire_value.get(field) if isinstance(weight, float) and not math.isfinite(weight): @@ -344,6 +379,36 @@ def _validate_worker_type(worker_type: str) -> str: return worker_type +def _resolve_forward_pass_systems_paths(entries: tuple[str, ...]) -> list[str]: + packaged = os.fspath(pkg_resources.files("aiconfigurator_core") / "systems") + if not entries: + from aiconfigurator_core.sdk.perf_database import get_systems_paths + + configured = get_systems_paths() + env_root = os.environ.get("AICONFIGURATOR_SYSTEMS_PATH") + entries = tuple([env_root] if configured == [packaged] and env_root is not None else configured) + resolved: list[str] = [] + for entry in entries: + if not isinstance(entry, str) or not entry.strip(): + raise ValueError("systems_paths entries must be nonempty strings") + path = packaged if entry.lower() == "default" else os.path.abspath(os.path.expanduser(entry)) + if not os.path.isdir(path): + raise ValueError(f"forward-pass systems path is not a directory: {path}") + if path not in resolved: + resolved.append(path) + return resolved + + +def _resolve_forward_pass_transfer_policy(value: Any) -> list[str] | None: + if value is None: + return None + if isinstance(value, str): + from aiconfigurator_core.sdk.common import resolve_transfer_policy + + return sorted(kind.value for kind in resolve_transfer_policy(value)) + return [getattr(token, "value", str(token)) for token in value] + + def _normalize_tuning_iterations(iterations: dict[str, Any] | list[Any]) -> list[Any]: if isinstance(iterations, dict): return [[iterations]] diff --git a/python/aisimulate/src/aisimulate/aic.py b/python/aisimulate/src/aisimulate/aic.py index 53b6fa6cd..8f1cd255b 100644 --- a/python/aisimulate/src/aisimulate/aic.py +++ b/python/aisimulate/src/aisimulate/aic.py @@ -32,7 +32,106 @@ def materialize_aic_num_gpu_blocks( ) -> dict[str, Any]: """Return engine arguments with rank-local AIC KV capacity materialized.""" + canonical_result = None + + def finish_lowering(value): + if canonical_result is None: + return value + result = dict(canonical_result) + for name in ("timing_model", "num_gpu_blocks", "tensor_parallel_size", "dp_size"): + if name in value: + result[name] = value[name] + for name in ( + "gpu_memory_utilization", + "mem_fraction_static", + "free_gpu_memory_fraction", + "cuda_graph_reserved_bytes", + "systems_path", + ): + result.pop(name, None) + return result + lowered = dict(raw) + timing = lowered.get("timing_model") + 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") + if "estimation_mode" in authored or "estimator_config" in authored: + from aiconfigurator_core.sdk import RustForwardPassPerfModel + + memory_fields = { + key: value + for key, value in authored.items() + if key + in { + "gpu_memory_utilization", + "mem_fraction_static", + "free_gpu_memory_fraction", + "cuda_graph_reserved_bytes", + } + } + for name in ( + "gpu_memory_utilization", + "mem_fraction_static", + "free_gpu_memory_fraction", + "cuda_graph_reserved_bytes", + ): + if name in lowered: + if name in memory_fields and memory_fields[name] != lowered[name]: + raise ValueError(f"{name} conflicts with canonical timing configuration") + memory_fields[name] = lowered[name] + request = {key: value for key, value in authored.items() if key not in memory_fields} + model = RustForwardPassPerfModel.best_available(request) + try: + diagnostics = model.diagnostics() + finally: + model.close() + if diagnostics.get("readiness") != "ready": + raise ValueError("regression estimator is not ready; replay requires training observations") + canonical_result = dict(raw) + resolved = diagnostics["provenance"]["config"] + lowered["timing_model"] = {**timing, "config": {**resolved, **memory_fields}} + for names, expected in ( + (("tensor_parallel_size", "aic_tp_size"), resolved["tp"]), + (("dp_size", "aic_attention_dp_size"), resolved["attention_dp"]), + ): + for name in names: + if name in lowered and lowered[name] != expected: + raise ValueError(f"{name} conflicts with canonical timing topology") + lowered[names[0]] = expected + lowered[names[1]] = expected + lowered.update( + { + key: value + for key, value in { + "aic_model_path": resolved["model"], + "aic_system": resolved["system"], + "aic_backend": resolved["backend"], + "aic_backend_version": resolved["backend_version"], + "aic_pp_size": resolved["pp"], + "aic_moe_tp_size": resolved["moe_tp_size"], + "aic_moe_ep_size": resolved["moe_ep_size"], + }.items() + if value is not None + } + ) + for target, source in ( + ("aic_gemm_dtype", "gemm_quant_mode"), + ("aic_moe_dtype", "moe_quant_mode"), + ("aic_fmha_dtype", "fmha_quant_mode"), + ("aic_kv_cache_dtype", "kvcache_quant_mode"), + ("aic_comm_dtype", "comm_quant_mode"), + ): + if resolved[source] is not None: + lowered[target] = resolved[source] + if resolved["systems_paths"]: + lowered["systems_path"] = resolved["systems_paths"][0] + for name, value in memory_fields.items(): + if name in lowered and lowered[name] != value: + raise ValueError(f"{name} conflicts with canonical timing configuration") + lowered.setdefault(name, value) + attention_dp = lowered.get("aic_attention_dp_size") dp = attention_dp or 1 configured_dp = lowered.get("dp_size") or 1 @@ -46,10 +145,10 @@ def materialize_aic_num_gpu_blocks( lowered["dp_size"] = dp if lowered.get("num_gpu_blocks") is not None: - return lowered + return finish_lowering(lowered) backend = lowered.get("aic_backend") if backend is None: - return lowered + return finish_lowering(lowered) if not isinstance(backend, str) or backend not in DEFAULT_BACKEND_VERSIONS: supported = ", ".join(sorted(DEFAULT_BACKEND_VERSIONS)) raise ValueError( @@ -90,7 +189,7 @@ def materialize_aic_num_gpu_blocks( cuda_graph_reserved_bytes=lowered.get("cuda_graph_reserved_bytes", 0), **({"diagnostics": memory_diagnostics} if memory_diagnostics is not None else {}), ) - return lowered + return finish_lowering(lowered) def estimate_num_gpu_blocks( diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index ea5a90416..39cd59208 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -46,6 +46,7 @@ def prediction_to_replay_spec( workload=workload, afd_performance_model=afd_performance_model, ) + deployment = _pin_estimator_version_aliases(deployment) if config.engine.mode == "disaggregated": assert config.engine.workers.prefill is not None for adapter in (adapter_specs or {}).values(): @@ -70,6 +71,54 @@ def prediction_to_replay_spec( ) +def _pin_estimator_version_aliases(deployment: BackendDeploymentSpec) -> BackendDeploymentSpec: + if deployment.backend_version not in {"current", "previous", "next"}: + return deployment + from aiconfigurator_core.sdk import RustForwardPassPerfModel + + updates = {} + metadata = dict(deployment.performance_model_metadata) + versions = set() + for role, field in ( + ("aggregated", "agg_engine_args"), + ("prefill", "prefill_engine_args"), + ("decode", "decode_engine_args"), + ): + args = getattr(deployment, field) + timing = (args or {}).get("timing_model", {}) + if timing.get("provider") != "aic" or "estimation_mode" not in timing.get("config", {}): + continue + config = dict(timing["config"]) + memory = { + name: config.pop(name) + for name in ( + "gpu_memory_utilization", + "mem_fraction_static", + "free_gpu_memory_fraction", + "cuda_graph_reserved_bytes", + ) + if name in config + } + model = RustForwardPassPerfModel.best_available(config) + try: + diagnostics = model.diagnostics() + finally: + model.close() + if diagnostics["readiness"] != "ready": + raise ValueError("regression estimator is not ready; replay requires training observations") + resolved = diagnostics["provenance"]["config"] + versions.add(resolved["backend_version"]) + updates[field] = {**args, "timing_model": {**timing, "config": {**resolved, **memory}}} + metadata[role] = {"provider": "aic", "config": resolved, "selection": diagnostics} + if len(versions) > 1: + raise ValueError( + f"estimator version alias resolves to different backend versions across roles: {sorted(versions)}" + ) + if versions: + updates.update(backend_version=versions.pop(), performance_model_metadata=metadata) + return replace(deployment, **updates) + + def _prediction_encoder(config: CorePredictionConfig, workload): from .sweeper.config import EncoderSearch, Workload from .sweeper.epd import resolve_encoder_pools @@ -118,18 +167,28 @@ def _deployment( if mode == "disagg": from aiconfigurator_core.sdk.perf_database import load_system_spec + from .sweeper.forward_pass_estimator import resolve_systems_paths + workers = (engine.workers.prefill, engine.workers.decode) + root_kwargs = {} for role, worker in zip(("prefill", "decode"), workers, strict=True): - if worker is not None and worker.hardware is not None and not load_system_spec(worker.hardware): + paths = engine.systems_paths + if worker is not None and worker.timing.systems_paths is not None: + paths = worker.timing.systems_paths + root_kwargs[role] = {"systems_paths": list(resolve_systems_paths(paths))} + if ( + worker is not None + and worker.hardware is not None + and not load_system_spec(worker.hardware, **root_kwargs[role]) + ): raise ValueError(f"unknown workers.{role}.hardware {worker.hardware!r}: no system configuration found") if engine.backend_version is None and any( - worker is not None and worker.hardware is not None for worker in workers + worker is not None and (worker.hardware is not None or worker.timing.systems_paths is not None) + for worker in workers ): versions = { - worker.hardware or engine.hardware: resolve_backend_version( - worker.hardware or engine.hardware, engine.backend - ) - for worker in workers + role: resolve_backend_version(worker.hardware or engine.hardware, engine.backend, **root_kwargs[role]) + for role, worker in zip(("prefill", "decode"), workers, strict=True) if worker is not None } if len(set(versions.values())) != 1: @@ -403,6 +462,50 @@ def _worker_engine_args( "aic_moe_ep_size", ): payload.pop(name, None) + if worker.timing.type == "default" and engine.mode != "afd" and engine.workers.encoder is None: + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig + + from .sweeper.forward_pass_estimator import resolve_systems_paths + + timing = worker.timing + sharded_moe = parallel.moe_tensor * parallel.moe_expert > 1 + canonical = ForwardPassPerfModelConfig( + model=engine.model, + system=worker.hardware or engine.hardware, + backend=backend, + backend_version=engine.backend_version, + worker_type=role, + tp=parallel.tensor, + pp=parallel.pipeline, + attention_dp=parallel.attention_data, + moe_tp_size=parallel.moe_tensor if sharded_moe else None, + moe_ep_size=parallel.moe_expert if sharded_moe else None, + kv_block_size=block_size, + speculation=engine.speculation.cost_config() if engine.speculation is not None else None, + estimation_mode=timing.estimation_mode or engine.estimation_mode, + fallback_policy=timing.fallback_policy or engine.fallback_policy, + estimator_config=timing.estimator_config + if timing.estimator_config is not None + else engine.estimator_config, + database_mode=timing.database_mode or engine.database_mode, + transfer_policy=timing.transfer_policy if timing.transfer_policy is not None else engine.transfer_policy, + systems_paths=resolve_systems_paths(timing.systems_paths or engine.systems_paths), + ) + timing_config = canonical.to_dict() + for key in ( + "gpu_memory_utilization", + "mem_fraction_static", + "free_gpu_memory_fraction", + "cuda_graph_reserved_bytes", + ): + if key in payload: + timing_config[key] = payload[key] + payload["timing_model"] = {"type": "external", "provider": "aic", "config": timing_config} + payload["tensor_parallel_size"] = parallel.tensor + payload["dp_size"] = parallel.attention_data + for key in tuple(payload): + if key.startswith("aic_") and key != "aic_nextn": + payload.pop(key) host_offload = cache.host_offload if host_offload is not None: payload["kv_cache_bytes_per_token"] = _resolve_kv_bytes_per_token( diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 6d8e2dfdb..4407fe8c6 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -158,12 +158,40 @@ def cost_config(self) -> dict[str, Any]: class TimingConfig(StrictModel): type: Literal["default", "fixed", "polynomial"] = "default" - forward_model: Literal["op_level", "fpm"] = "op_level" + forward_model: Literal["op_level", "fpm"] = Field(default="op_level", exclude=True) + estimation_mode: Literal["auto", "op_level", "fpm_interpolation", "fpm_regression"] | None = None + fallback_policy: Literal["deny", "allow"] | None = None + estimator_config: dict[str, Any] | None = None + systems_paths: list[str] | None = Field(default=None, min_length=1) + database_mode: Literal["SILICON", "HYBRID", "EMPIRICAL", "SOL"] | None = None + transfer_policy: str | list[str] | None = None + + @field_validator("database_mode", mode="before") + @classmethod + def _normalize_database_mode(cls, value): + return value.upper() if isinstance(value, str) else value + + @model_validator(mode="before") + @classmethod + def _preserve_legacy_estimator_selection(cls, value): + if isinstance(value, dict) and value.get("type", "default") == "default" and "forward_model" in value: + value = dict(value) + legacy = {"op_level": "op_level", "fpm": "fpm_interpolation"}.get(value["forward_model"]) + if legacy is not None and "estimation_mode" not in value: + value["estimation_mode"] = legacy + value.setdefault("fallback_policy", "deny") + return value + prefill_ms: float | None = Field(default=None, ge=0.0) decode_ms: float | None = Field(default=None, ge=0.0) @model_validator(mode="after") def _validate_timing(self) -> TimingConfig: + if self.estimation_mode == "fpm_interpolation": + self.forward_model = "fpm" + elif self.estimation_mode == "op_level": + self.forward_model = "op_level" + if self.type == "fixed": if self.prefill_ms is None or self.decode_ms is None: raise ValueError("fixed timing requires prefill_ms and decode_ms") @@ -225,7 +253,83 @@ class KvTransferConfig(StrictModel): timing_mode: Literal["full_prompt", "destination_missing"] = "destination_missing" -class EnginePredictionConfig(StrictModel): +class EstimatorPolicyConfig(StrictModel): + database_mode: Literal["SILICON", "HYBRID", "EMPIRICAL", "SOL"] = "SILICON" + transfer_policy: str | list[str] | None = None + systems_paths: list[str] | None = None + estimation_mode: Literal["auto", "op_level", "fpm_interpolation", "fpm_regression"] = "auto" + fallback_policy: Literal["deny", "allow"] = "deny" + estimator_config: dict[str, Any] = Field(default_factory=dict) + + @field_validator("database_mode", mode="before") + @classmethod + def _normalize_database_mode(cls, value): + return value.upper() if isinstance(value, str) else value + + @field_validator("systems_paths") + @classmethod + def _nonempty_system_roots(cls, value): + if value is None: + return value + if not value or any(not path.strip() for path in value): + raise ValueError("systems_paths must contain at least one nonempty root") + return value + + @model_validator(mode="after") + def _supported_estimator_policies(self): + mode = getattr(self, "mode", "aggregated") + modes = mode.choices if hasattr(mode, "choices") else [mode] + workers = getattr(self, "workers", None) + custom_policy = ( + self.database_mode != "SILICON" + or self.transfer_policy is not None + or self.systems_paths not in (None, ["default"]) + or self.estimation_mode != "auto" + or self.fallback_policy != "deny" + or bool(self.estimator_config) + ) + roles = [getattr(workers, role, None) for role in ("aggregated", "prefill", "decode")] + unsupported_provider = "afd" in modes or getattr(workers, "encoder", None) is not None + if unsupported_provider and any( + worker is not None + and ( + bool(worker.timing.estimator_config) + or worker.timing.systems_paths is not None + or worker.timing.database_mode is not None + or worker.timing.transfer_policy is not None + or worker.timing.fallback_policy == "allow" + or worker.timing.estimation_mode not in {None, "op_level", "fpm_interpolation"} + ) + for worker in roles + ): + raise ValueError("estimator policies require regular language workers with default timing") + if custom_policy and ( + "afd" in modes + or getattr(workers, "encoder", None) is not None + or any(worker is not None and worker.timing.type != "default" for worker in roles) + ): + raise ValueError("estimator policies require regular language workers with default timing in every role") + for worker in roles: + if ( + worker is not None + and worker.timing.type != "default" + and any( + getattr(worker.timing, name) is not None + for name in ( + "estimation_mode", + "fallback_policy", + "estimator_config", + "systems_paths", + "database_mode", + "transfer_policy", + ) + ) + ): + raise ValueError("estimator settings require default timing") + return self + + +class EnginePredictionConfig(EstimatorPolicyConfig): mode: EngineMode = "aggregated" model: str hardware: str @@ -400,14 +504,14 @@ class WorkersRecommendationConfig(StrictModel): decode: WorkerRecommendationConfig | None = None -class EngineRecommendationConfig(StrictModel): +class EngineRecommendationConfig(EstimatorPolicyConfig): mode: EngineMode | Choices[EngineMode] = Field( default_factory=lambda: Choices[EngineMode](choices=["aggregated", "disaggregated"]) ) model: str hardware: str backend: Backend | Choices[Backend] = Field(default_factory=lambda: Choices[Backend](choices=["vllm", "sglang"])) - backend_version: str | None = None + backend_version: str | dict[str, str] | None = None context_length: PositiveInt | Literal["max"] = "max" speculation: NgramSpeculationConfig | None = None workers: WorkersRecommendationConfig = Field(default_factory=WorkersRecommendationConfig) diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index d193638ff..96a2926ab 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -152,6 +152,34 @@ def recommendation_to_sweeper( "min_gpu_budget": optimization.constraints.min_candidate_gpus, "context_length": (resolve_model_context_length(model) if context == "max" else context), } + for name in ( + "database_mode", + "transfer_policy", + "systems_paths", + "estimation_mode", + "fallback_policy", + "estimator_config", + ): + if name in engine: + search_space[name] = deepcopy(engine[name]) + search_space["role_estimator_controls"] = { + ("agg" if role == "aggregated" else role): { + name: deepcopy(raw.get("timing", {})[name]) + for name in ( + "estimation_mode", + "fallback_policy", + "estimator_config", + "systems_paths", + "database_mode", + "transfer_policy", + ) + if name in raw.get("timing", {}) + } + for role, raw in workers.items() + if role in {"aggregated", "prefill", "decode"} + and not set(modes) & {"afd", "afd+pd"} + and workers.get("encoder") is None + } if engine.get("speculation") is not None: search_space["speculation"] = deepcopy(engine["speculation"]) for role in ("prefill", "decode"): @@ -402,7 +430,10 @@ def _role_search_space( result[f"{legacy_role}_timing_model"] = {"type": "polynomial"} else: result[f"{legacy_role}_timing_model"] = None - result[f"{legacy_role}_forward_model"] = timing.get("forward_model", "op_level") + if timing.get("estimation_mode") is not None: + result[f"{legacy_role}_forward_model"] = ( + "fpm" if timing["estimation_mode"] == "fpm_interpolation" else "op_level" + ) result[f"{legacy_role}_startup_time"] = raw.get("startup_seconds", 0) # Remove empty internal maps so legacy serialization remains concise. if not result["engine_float_ranges"]: @@ -745,6 +776,17 @@ def _candidate_prediction( "workers": {}, } raw_engine = source.engine.model_dump(mode="python", exclude_none=True) + if deployment.forward_pass_estimators: + for name in ( + "database_mode", + "transfer_policy", + "systems_paths", + "estimation_mode", + "fallback_policy", + "estimator_config", + ): + if name in raw_engine: + engine[name] = deepcopy(raw_engine[name]) if sample.get("speculation") is not None: engine["speculation"] = deepcopy(sample["speculation"]) if deployment.encoder is not None: @@ -797,6 +839,20 @@ def _candidate_prediction( timing = deepcopy(timing_model) else: timing = {"type": "default", "forward_model": sample.get(f"{role}_forward_model") or "op_level"} + estimator = deployment.forward_pass_estimators.get(role) + if estimator is not None: + resolved = estimator.config + timing = { + "type": "default", + "estimation_mode": resolved["estimation_mode"], + "fallback_policy": "deny", + "estimator_config": deepcopy(resolved["estimator_config"]), + } + # Per-role roots can differ; preserve them next to the role's timing. + timing["systems_paths"] = list(resolved["systems_paths"]) + timing["database_mode"] = resolved["database_mode"] + policy = resolved["transfer_policy"] + timing["transfer_policy"] = list(policy) if policy is not None else None kv_cache = { "block_size": block_size, "prefix_caching": sample[f"{role}_enable_prefix_caching"], diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index 9f45378f5..3076d2002 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -985,7 +985,14 @@ def _execution_target_model( if isinstance(metadata, Mapping): config = metadata.get("config") if isinstance(config, Mapping): - model = config.get("model_path") + model = config.get("model", config.get("model_path")) + if isinstance(model, str) and model.strip(): + return model.strip() + timing = raw_engine_args.get("timing_model") + if isinstance(timing, Mapping) and timing.get("provider") == "aic": + config = timing.get("config") + if isinstance(config, Mapping): + model = config.get("model") if isinstance(model, str) and model.strip(): return model.strip() model = raw_engine_args.get("aic_model_path") @@ -1168,6 +1175,23 @@ def _materialize_engine_role( # the AIC timing config because the native runtime rematerializes inferred # capacity before execution. cuda_graph_reserved_bytes = role_config.pop("cuda_graph_reserved_bytes", None) + timing_rank = role_config.get("rank", role_config) + timing = timing_rank.get("timing_model") if isinstance(timing_rank, dict) else None + if isinstance(timing, dict) and timing.get("type") == "external" and timing.get("provider") == "aic": + canonical = timing.get("config") + if not isinstance(canonical, dict): + raise ValueError("external AIC timing config must be a mapping") + for aliases, fields in ( + (("tensor_parallel_size", "aic_tp_size"), ("tp", "tp_size")), + (("dp_size", "aic_attention_dp_size"), ("attention_dp", "attention_dp_size")), + ): + expected = canonical.get(fields[0], canonical.get(fields[1])) + if expected is not None: + for alias in aliases: + if alias in role_config and role_config[alias] != expected: + raise ValueError(f"{alias} conflicts with canonical AIC timing topology") + if not any(alias in role_config for alias in aliases): + role_config[aliases[0]] = expected raw_dp_size = _pop_matching_aliases(role_config, "attention DP", ("dp_size", "aic_attention_dp_size"), 1) raw_tp_size = _pop_matching_aliases(role_config, "tensor parallel", ("tensor_parallel_size", "aic_tp_size"), 1) dp_size = _positive_int( @@ -1297,7 +1321,30 @@ def _materialize_engine_role( "rank.num_gpu_blocks explicitly or use AIC timing" ) if deployment_backend_version: + + def resolved_version(value): + if value not in {"current", "previous", "next"}: + return value + from aiconfigurator_core.sdk.perf_database import resolve_query_version + + from .sweeper.forward_pass_estimator import resolve_systems_paths + + identity = timing_model.get("config", {}) if isinstance(timing_model, dict) else {} + roots = identity.get("systems_paths") + if roots is None and identity.get("systems_path") is not None: + roots = [identity["systems_path"]] + return resolve_query_version( + identity.get("system", system), + backend, + value, + systems_paths=list(resolve_systems_paths(roots)), + ) + + if uses_aic_timing: + deployment_backend_version = resolved_version(deployment_backend_version) configured_version = aic_timing_overrides.get("backend_version") + if uses_aic_timing: + configured_version = resolved_version(configured_version) if configured_version is not None and configured_version != deployment_backend_version: raise ValueError( f"engine provider {role} backend version {configured_version!r} " @@ -1307,6 +1354,10 @@ def _materialize_engine_role( if uses_aic_timing: timing_config = timing_model.get("config") if isinstance(timing_model, dict) else None timing_backend_version = timing_config.get("backend_version") if isinstance(timing_config, dict) else None + timing_backend_version = resolved_version(timing_backend_version) + if timing_backend_version is not None: + timing_model = {**timing_model, "config": {**timing_config, "backend_version": timing_backend_version}} + rank["timing_model"] = timing_model if timing_backend_version is not None and timing_backend_version != deployment_backend_version: raise ValueError( f"engine provider {role} timing_model.config.backend_version=" diff --git a/python/aisimulate/src/aisimulate/sweeper/__init__.py b/python/aisimulate/src/aisimulate/sweeper/__init__.py index 0051c06d5..b5297e3e6 100644 --- a/python/aisimulate/src/aisimulate/sweeper/__init__.py +++ b/python/aisimulate/src/aisimulate/sweeper/__init__.py @@ -97,6 +97,8 @@ "enumerate_afd_topologies": (".afd_parallel", "enumerate_afd_topologies"), "evaluate_afd_phase": (".afd_engine", "evaluate_afd_phase"), "build_backend_deployment": (".deploy", "build_backend_deployment"), + "ForwardPassEstimatorResolver": (".forward_pass_estimator", "ForwardPassEstimatorResolver"), + "ForwardPassEstimatorResolutionError": (".forward_pass_estimator", "ForwardPassEstimatorResolutionError"), "NoPerfDatabase": (".kv_estimate", "NoPerfDatabase"), "estimate_kv_tokens": (".kv_estimate", "estimate_kv_tokens"), "feasible_shape_tokens": (".kv_estimate", "feasible_shape_tokens"), @@ -181,6 +183,8 @@ def __getattr__(name: str) -> Any: "DisaggParallelConfig", "EncoderPoolSpec", "EncoderSearch", + "ForwardPassEstimatorResolutionError", + "ForwardPassEstimatorResolver", "HookCapability", "ImageWorkload", "InfeasibleCandidate", diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index 0a6c18eb7..8a616ab25 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -25,7 +25,7 @@ from typing import Any, Literal import yaml -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_serializer, model_validator from ..config.engine import NgramSpeculationConfig @@ -474,7 +474,7 @@ class SearchSpace(BaseModel): # deployment: branch + backend + legal parallel shapes deployment_mode: list[str] = ["disagg", "agg"] # branches to explore; pin with one backend: list[str] = ["vllm"] # vllm | sglang | trtllm - backend_version: str | None = None + backend_version: str | dict[str, str] | None = None parallel_configs: list[dict[str, Any]] = Field(default_factory=list) # generated when empty parallel_configs_by_mode: dict[str, list[dict[str, Any]]] = Field(default_factory=dict) flat_parallel_modes: list[str] = Field(default_factory=list) @@ -484,6 +484,13 @@ class SearchSpace(BaseModel): # pinned model_name: str # HF id or private model name hardware_sku: str # e.g. "h200_sxm" + database_mode: Literal["SILICON", "HYBRID", "EMPIRICAL", "SOL"] = "SILICON" + transfer_policy: str | list[str] | None = None + systems_paths: list[str] | None = Field(default=None, min_length=1) + estimation_mode: Literal["auto", "op_level", "fpm_interpolation", "fpm_regression"] = "auto" + fallback_policy: Literal["deny", "allow"] = "deny" + estimator_config: dict[str, Any] = Field(default_factory=dict) + role_estimator_controls: dict[str, dict[str, Any]] = Field(default_factory=dict) prefill_hardware_sku: str | None = Field(default=None, min_length=1) decode_hardware_sku: str | None = Field(default=None, min_length=1) gpu_budget: int = 32 # max GPUs per candidate @@ -520,7 +527,7 @@ class SearchSpace(BaseModel): prefill_native_host_offload: dict[str, Any] | None = None prefill_num_gpu_blocks: int | None = None prefill_timing_model: dict[str, Any] | None = None - prefill_forward_model: str = "op_level" # AIC forward-pass model: op_level | fpm + prefill_forward_model: str = "op_level" prefill_startup_time: float | None = None # decode engine (disagg branch): scheduler batching capacity @@ -534,7 +541,7 @@ class SearchSpace(BaseModel): decode_native_host_offload: dict[str, Any] | None = None decode_num_gpu_blocks: int | None = None decode_timing_model: dict[str, Any] | None = None - decode_forward_model: str = "op_level" # AIC forward-pass model: op_level | fpm + decode_forward_model: str = "op_level" decode_startup_time: float | None = None # agg engine (agg branch): scheduler batching capacity @@ -548,7 +555,7 @@ class SearchSpace(BaseModel): agg_native_host_offload: dict[str, Any] | None = None agg_num_gpu_blocks: int | None = None agg_timing_model: dict[str, Any] | None = None - agg_forward_model: str = "op_level" # AIC forward-pass model: op_level | fpm + agg_forward_model: str = "op_level" agg_startup_time: float | None = None kv_transfer_bytes_per_token: int | str | None = None kv_transfer_bandwidth: float | None = None @@ -558,6 +565,34 @@ class SearchSpace(BaseModel): engine_log_discrete: list[str] = Field(default_factory=list) engine_integer_log_ranges: dict[str, list[int]] = Field(default_factory=dict) + @model_validator(mode="before") + @classmethod + def _migrate_legacy_estimator_selection(cls, value): + if not isinstance(value, dict): + return value + modes = value.get("deployment_mode") + if (isinstance(modes, list) and any(mode in ("afd", "afd+pd") for mode in modes)) or value.get( + "encoder" + ) is not None: + return value + value = dict(value) + raw_controls = value.get("role_estimator_controls", {}) + if not isinstance(raw_controls, dict) or any( + not isinstance(settings, dict) for settings in raw_controls.values() + ): + return value + controls = {role: dict(settings) for role, settings in raw_controls.items()} + for role in ("agg", "prefill", "decode"): + legacy = value.get(f"{role}_forward_model") + if legacy == "op_level" and value.get(f"{role}_timing_model") is not None: + continue + if isinstance(legacy, str) and legacy in {"op_level", "fpm"}: + controls.setdefault(role, {}).setdefault( + "estimation_mode", "fpm_interpolation" if legacy == "fpm" else "op_level" + ) + value["role_estimator_controls"] = controls + return value + @model_validator(mode="after") def _validate_search_choices(self) -> SearchSpace: """Every backend dimension is a non-empty subset of its allowed choices.""" @@ -602,8 +637,26 @@ def _validate_search_choices(self) -> SearchSpace: value = getattr(self, field_name) if value not in FORWARD_MODEL_CHOICES: raise ValueError(f"{field_name} has invalid choice {value!r}; allowed: {list(FORWARD_MODEL_CHOICES)}") + if not self._uses_legacy_estimator_provider(): + for role in ("agg", "prefill", "decode"): + mode = self.role_estimator_controls.get(role, {}).get("estimation_mode", self.estimation_mode) + setattr(self, f"{role}_forward_model", "fpm" if mode == "fpm_interpolation" else "op_level") return self + def _uses_legacy_estimator_provider(self) -> bool: + return bool(set(self.deployment_mode) & {"afd", "afd+pd"}) or self.encoder is not None + + @model_serializer(mode="wrap") + def _serialize_estimator_selection(self, handler): + result = handler(self) + # Ordinary workers serialize only the canonical controls. Existing AFD + # and encoder providers still consume the legacy selector, including + # when a saved search is reloaded. + if not self._uses_legacy_estimator_provider(): + for role in ("agg", "prefill", "decode"): + result.pop(f"{role}_forward_model", None) + return result + @model_validator(mode="after") def _validate_role_hardware(self) -> SearchSpace: """Role-specific hardware is an override for ordinary P/D only.""" @@ -623,6 +676,9 @@ def hardware_sku_for(self, role: str) -> str: return self.decode_hardware_sku or self.hardware_sku raise ValueError(f"unknown engine role {role!r}") + def systems_paths_for(self, role: str) -> list[str] | None: + return self.role_estimator_controls.get(role, {}).get("systems_paths", self.systems_paths) + @field_validator( "afd_tp_a_candidates", "afd_batch_size_candidates", @@ -702,6 +758,109 @@ def _validate_afd_contract(self) -> SearchSpace: raise ValueError(f"afd_companion_parallel_configs[{index}].replicas must be positive") return self + @field_validator("database_mode", mode="before") + @classmethod + def _normalize_estimator_database_mode(cls, value): + return value.upper() if isinstance(value, str) else value + + @field_validator("systems_paths") + @classmethod + def _validate_estimator_roots(cls, value): + if value is None: + return value + if any(not path.strip() for path in value): + raise ValueError("systems_paths entries must be nonempty") + return value + + @model_validator(mode="after") + def _validate_estimator_controls(self): + allowed = { + "estimation_mode", + "fallback_policy", + "estimator_config", + "systems_paths", + "database_mode", + "transfer_policy", + } + for role, controls in self.role_estimator_controls.items(): + if role not in {"agg", "prefill", "decode"}: + raise ValueError(f"unknown estimator role {role!r}; expected agg, prefill, or decode") + if set(controls) - allowed: + raise ValueError(f"unknown estimator override for role {role!r}: {sorted(set(controls) - allowed)}") + if controls and self._uses_legacy_estimator_provider(): + raise ValueError( + "role_estimator_controls require regular language workers; " + "AFD and encoder providers are unsupported" + ) + if controls and getattr(self, f"{role}_timing_model") is not None: + raise ValueError(f"{role} estimator settings require default timing") + mode = controls.get("database_mode") + if isinstance(mode, str) and mode.upper() == "SOL_FULL": + raise ValueError( + f"role_estimator_controls.{role}.database_mode cannot be SOL_FULL; it is a per-call diagnostic" + ) + if "systems_paths" in controls: + paths = controls["systems_paths"] + if ( + not isinstance(paths, list) + or not paths + or any(not isinstance(p, str) or not p.strip() for p in paths) + ): + raise ValueError(f"{role} systems_paths must contain nonempty strings") + nondefault = ( + self.database_mode != "SILICON" + or self.transfer_policy is not None + or self.systems_paths not in (None, ["default"]) + or self.estimation_mode != "auto" + or self.fallback_policy != "deny" + or bool(self.estimator_config) + ) + roles = ({"agg"} if "agg" in self.deployment_mode else set()) | ( + {"prefill", "decode"} if "disagg" in self.deployment_mode else set() + ) + if nondefault and ( + set(self.deployment_mode) & {"afd", "afd+pd"} + or self.encoder is not None + or any(getattr(self, f"{role}_timing_model") is not None for role in roles) + ): + raise ValueError("estimator policies require regular language workers with default timing in every role") + return self + + @model_validator(mode="after") + def _validate_backend_versions(self) -> SearchSpace: + configured = list(dict.fromkeys(self.backend)) + if isinstance(self.backend_version, str): + if not self.backend_version.strip(): + raise ValueError("backend_version must be a non-empty version") + if len(configured) != 1: + raise ValueError( + "a string backend_version requires exactly one configured backend; " + "use a {backend: version} mapping for a multi-backend search" + ) + self.backend_version = self.backend_version.strip() + elif isinstance(self.backend_version, dict): + unknown = sorted(set(self.backend_version) - set(configured)) + if unknown: + raise ValueError(f"backend_version contains unconfigured backend(s): {unknown}") + invalid = [ + backend + for backend, version in self.backend_version.items() + if not isinstance(version, str) or not version.strip() + ] + if invalid: + raise ValueError(f"backend_version needs a non-empty version for {sorted(invalid)}") + self.backend_version = {backend: version.strip() for backend, version in self.backend_version.items()} + return self + + def requested_backend_version(self, backend: str) -> str | None: + """Return the version pin for ``backend``; ``None`` means resolve latest.""" + + if isinstance(self.backend_version, str): + return self.backend_version + if isinstance(self.backend_version, dict): + return self.backend_version.get(backend) + return None + @model_validator(mode="after") def _validate_gpu_budget(self) -> SearchSpace: """A minimum GPU budget must be positive and within the maximum budget.""" diff --git a/python/aisimulate/src/aisimulate/sweeper/deploy.py b/python/aisimulate/src/aisimulate/sweeper/deploy.py index 8cb090ce4..5771054af 100644 --- a/python/aisimulate/src/aisimulate/sweeper/deploy.py +++ b/python/aisimulate/src/aisimulate/sweeper/deploy.py @@ -5,11 +5,13 @@ from __future__ import annotations +from collections.abc import Mapping +from copy import deepcopy from typing import Any from ..aic import estimate_kv_bytes_per_token, materialize_aic_num_gpu_blocks from ..config.engine import NgramSpeculationConfig -from .replay import BackendDeploymentSpec, EncoderPoolSpec +from .replay import BackendDeploymentSpec, EncoderPoolSpec, ForwardPassEstimatorSpec def _role_prefix(role: str) -> str: @@ -50,7 +52,13 @@ def _performance_model_metadata(sample: dict[str, Any], role: str, *, backend_ve return {"provider": "aic", "config": config} -def _engine_args_payload(sample: dict[str, Any], role: str, *, backend_version: str) -> dict[str, Any]: +def _engine_args_payload( + sample: dict[str, Any], + role: str, + *, + backend_version: str, + forward_pass_estimator: ForwardPassEstimatorSpec | None = None, +) -> dict[str, Any]: """Build the runner-neutral engine argument payload for one role.""" prefix = _role_prefix(role) tp = int(sample[f"{prefix}tp"]) @@ -118,6 +126,17 @@ def _engine_args_payload(sample: dict[str, Any], role: str, *, backend_version: "aic_forward_model", ): payload.pop(name, None) + if forward_pass_estimator is not None and sample.get(f"{role}_timing_model") is None: + payload["timing_model"] = {"type": "external", "provider": "aic", "config": dict(forward_pass_estimator.config)} + if memory_fraction_field in payload: + payload["timing_model"]["config"][memory_fraction_field] = payload[memory_fraction_field] + payload["tensor_parallel_size"] = tp + payload["dp_size"] = attention_dp + if forward_pass_estimator.performance_data_root: + payload["systems_path"] = forward_pass_estimator.performance_data_root + for key in tuple(payload): + if key.startswith("aic_") and key != "aic_nextn": + payload.pop(key) host_offload = sample.get(f"{role}_native_host_offload") if host_offload is not None: configured_bytes = sample[f"{role}_kv_bytes_per_token"] @@ -160,8 +179,10 @@ def build_backend_deployment( *, backend_version: str, encoder: EncoderPoolSpec | None = None, + forward_pass_estimators: Mapping[str, ForwardPassEstimatorSpec] | None = None, ) -> BackendDeploymentSpec: """Build the Dynamo-independent backend part of a :class:`ReplaySpec`.""" + forward_pass_estimators = dict(forward_pass_estimators or {}) mode = sample["deployment_mode"] if encoder is not None and mode not in {"agg", "disagg"}: raise ValueError("analytical EPD supports only agg/disagg language deployments; AFD is unsupported") @@ -211,6 +232,7 @@ def build_backend_deployment( ) common = { + "forward_pass_estimators": forward_pass_estimators, "encoder": encoder, "deployment_mode": mode, "backend": sample["backend"], @@ -252,14 +274,32 @@ def build_backend_deployment( for role in (("agg",) if mode == "agg" else ("prefill", "decode")) }, } + for role, estimator in forward_pass_estimators.items(): + common["performance_model_metadata"]["aggregated" if role == "agg" else role] = { + "provider": "aic", + "config": deepcopy(estimator.config), + "selection": deepcopy(estimator.diagnostics), + } if mode == "agg": return BackendDeploymentSpec( - agg_engine_args=_engine_args_payload(sample, "agg", backend_version=backend_version), + agg_engine_args=_engine_args_payload( + sample, + "agg", + backend_version=backend_version, + forward_pass_estimator=forward_pass_estimators.get("agg"), + ), num_workers=int(sample["replicas"]), **common, ) - prefill_args = _engine_args_payload(sample, "prefill", backend_version=backend_version) - decode_args = _engine_args_payload(sample, "decode", backend_version=backend_version) + prefill_args = _engine_args_payload( + sample, + "prefill", + backend_version=backend_version, + forward_pass_estimator=forward_pass_estimators.get("prefill"), + ) + decode_args = _engine_args_payload( + sample, "decode", backend_version=backend_version, forward_pass_estimator=forward_pass_estimators.get("decode") + ) return BackendDeploymentSpec( prefill_engine_args=prefill_args, decode_engine_args=decode_args, diff --git a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py new file mode 100644 index 000000000..c7582502a --- /dev/null +++ b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolve exact per-role Sweeper forward-pass estimator contracts.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from typing import Any + +from aiconfigurator_core.sdk import ( + ForwardPassPerfModelConfig, + RustForwardPassPerfModel, +) + +from .config import SearchSpace +from .deploy import _role_hardware_sku +from .replay import ForwardPassEstimatorSpec + + +class ForwardPassEstimatorResolutionError(ValueError): + """A configured forward-pass estimator identity cannot be resolved exactly.""" + + +def resolve_systems_paths(configured: list[str] | None) -> tuple[str, ...]: + """Expand and validate request-scoped system roots without setting globals.""" + + from aiconfigurator_core.sdk.rust_engine_step import _resolve_forward_pass_systems_paths + + return tuple(_resolve_forward_pass_systems_paths(tuple(configured or ()))) + + +_DEFAULT_BLOCK_SIZE = {"vllm": 64, "sglang": 1, "trtllm": 32} + + +def _role_prefix(role: str) -> str: + return "" if role == "agg" else f"{role}_" + + +class ForwardPassEstimatorResolver: + """Resolve and cache Core-owned estimator identities for concrete roles. + + The search-space request alone is not an estimator identity: whole-forward + FPM cells include topology and backend block size. ``resolve_candidate`` is + called after a suggestion has been unrolled but before its ``ReplaySpec`` + reaches a runner. Every exact role request crosses the canonical Core + constructor, and the returned config/provenance is carried unchanged. + """ + + def __init__(self, search_space: SearchSpace) -> None: + self._search_space = search_space + self._resolved: dict[str, ForwardPassEstimatorSpec] = {} + + def _request(self, sample: dict[str, Any], role: str) -> ForwardPassPerfModelConfig: + backend = str(sample["backend"]) + prefix = _role_prefix(role) + moe_tp = int(sample[f"{prefix}moe_tp"]) + moe_ep = int(sample[f"{prefix}moe_ep"]) + block_size = sample[f"{role}_block_size"] + if block_size is None: + block_size = _DEFAULT_BLOCK_SIZE[backend] + transfer_policy: Any = self._search_space.transfer_policy + if isinstance(transfer_policy, list): + transfer_policy = tuple(transfer_policy) + controls = self._search_space.role_estimator_controls.get(role, {}) + nextn = sample.get("aic_nextn") + if nextn is None: + nextn = self._search_space.aic_nextn + return ForwardPassPerfModelConfig( + model=self._search_space.model_name, + system=_role_hardware_sku(sample, role), + backend=backend, + worker_type="aggregated" if role == "agg" else role, + backend_version=self._search_space.requested_backend_version(backend), + tp=int(sample[f"{prefix}tp"]), + pp=int(sample[f"{prefix}pp"]), + attention_dp=int(sample[f"{prefix}attention_dp"]), + moe_tp_size=moe_tp if moe_tp * moe_ep > 1 else None, + moe_ep_size=moe_ep if moe_tp * moe_ep > 1 else None, + nextn=int(nextn or 0), + speculation=self._search_space.speculation.cost_config() + if self._search_space.speculation is not None + else None, + kv_block_size=int(block_size), + estimation_mode=controls.get("estimation_mode", self._search_space.estimation_mode), + database_mode=controls.get("database_mode", self._search_space.database_mode), + transfer_policy=controls.get("transfer_policy", transfer_policy), + systems_paths=resolve_systems_paths(self._search_space.systems_paths_for(role)), + fallback_policy=controls.get("fallback_policy", self._search_space.fallback_policy), + estimator_config=controls.get("estimator_config", self._search_space.estimator_config), + ) + + def _resolve(self, request: ForwardPassPerfModelConfig, role: str) -> ForwardPassEstimatorSpec: + request_payload = vars(request) + cache_key = json.dumps(request.to_dict(), sort_keys=True) + cached = self._resolved.get(cache_key) + if cached is not None: + return deepcopy(cached) + + model: RustForwardPassPerfModel | None = None + try: + model = RustForwardPassPerfModel.best_available(request) + diagnostics = model.diagnostics() + except Exception as exc: + raise ForwardPassEstimatorResolutionError( + "Core cannot construct the exact forward-pass estimator for " + f"role={role}, model={request.model}, system={request.system}, " + f"backend={request.backend}, tp={request.tp}, pp={request.pp}, " + f"attention_dp={request.attention_dp}, moe_tp={request.moe_tp_size}, " + f"moe_ep={request.moe_ep_size}, kv_block_size={request.kv_block_size}: {exc}" + ) from exc + finally: + if model is not None: + model.close() + + provenance = diagnostics.get("provenance") + if not isinstance(provenance, dict) or not isinstance(provenance.get("config"), dict): + raise ForwardPassEstimatorResolutionError( + f"Core returned no resolved provenance for {request.system}/{request.backend}/{role}" + ) + resolved_config = dict(provenance["config"]) + if diagnostics.get("readiness") != "ready": + raise ForwardPassEstimatorResolutionError( + f"estimator for {role} is not ready; regression requires training observations" + ) + if not resolved_config.get("backend_version"): + raise ForwardPassEstimatorResolutionError( + f"Core did not resolve an exact backend version for {request.system}/{request.backend}/{role}" + ) + selected_root = provenance.get("selected_systems_root") + if selected_root and resolved_config.get("systems_paths") != [str(selected_root)]: + raise ForwardPassEstimatorResolutionError( + "Core provenance must pin its selected systems root in the resolved config: " + f"config={resolved_config.get('systems_paths')!r}, selected={selected_root!r}" + ) + for field in ( + "model", + "system", + "backend", + "tp", + "pp", + "attention_dp", + "moe_tp_size", + "moe_ep_size", + "nextn", + "speculation", + "kv_block_size", + "worker_type", + ): + if resolved_config.get(field) != request_payload.get(field): + raise ForwardPassEstimatorResolutionError( + f"Core changed exact candidate field {field!r}: " + f"requested={request_payload.get(field)!r}, resolved={resolved_config.get(field)!r}" + ) + spec = ForwardPassEstimatorSpec( + config=resolved_config, + options=None, + diagnostics=diagnostics, + ) + self._resolved[cache_key] = deepcopy(spec) + return spec + + def resolve_candidate(self, sample: dict[str, Any]) -> dict[str, ForwardPassEstimatorSpec]: + """Resolve every concrete engine role before replay trial execution.""" + + if sample["deployment_mode"] not in {"agg", "disagg"} or self._search_space.encoder is not None: + return {} + roles = ("agg",) if sample["deployment_mode"] == "agg" else ("prefill", "decode") + resolved = { + role: self._resolve(self._request(sample, role), role) + for role in roles + if sample.get(f"{role}_timing_model") is None + } + versions = {spec.backend_version for spec in resolved.values()} + if len(versions) > 1: + raise ForwardPassEstimatorResolutionError( + f"Core resolved inconsistent backend versions across candidate roles: {sorted(versions)}" + ) + return resolved diff --git a/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py b/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py index 878d6a407..a573a7e37 100644 --- a/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py +++ b/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py @@ -29,6 +29,7 @@ from aiconfigurator_core.sdk.memory import estimate_kv_cache from aiconfigurator_core.sdk.perf_database import get_latest_database_version +from .forward_pass_estimator import resolve_systems_paths from .parallel_enum import ParallelShape # Runtime knobs the estimate needs; defaults mirror AIC's memory-estimation tests. @@ -46,9 +47,31 @@ def memory_fraction_kind(backend: str) -> str: return "of_free" if backend == "trtllm" else "of_total" -def resolve_backend_version(hardware_sku: str, backend: str) -> str: - """Latest perf-DB version for the SKU/backend (required by the native estimate).""" - version = get_latest_database_version(hardware_sku, backend) +def resolve_backend_version( + hardware_sku: str, + backend: str, + *, + requested_version: str | None = None, + systems_paths: list[str] | None = None, +) -> str: + """Pinned or latest perf-DB version for the native memory estimate.""" + + resolved_paths = list(resolve_systems_paths(systems_paths)) if systems_paths is not None else None + if requested_version is not None: + from aiconfigurator_core.sdk.perf_database import get_supported_databases + + available = get_supported_databases(**({"systems_paths": resolved_paths} if resolved_paths is not None else {})) + versions = available.get(hardware_sku, {}).get(backend, []) + if requested_version not in versions: + raise NoPerfDatabase( + f"backend_version={requested_version!r} is unavailable for " + f"hardware_sku={hardware_sku!r}, backend={backend!r}; " + f"available versions: {versions}" + ) + return requested_version + version = get_latest_database_version( + hardware_sku, backend, **({"systems_paths": resolved_paths} if resolved_paths is not None else {}) + ) if version is None: raise NoPerfDatabase( f"no perf database for hardware_sku={hardware_sku!r}, backend={backend!r}; " @@ -64,6 +87,7 @@ def estimate_kv_tokens( hardware_sku: str, backend: str, backend_version: str, + systems_paths: list[str] | None = None, max_num_tokens: int = DEFAULT_MAX_NUM_TOKENS, max_batch_size: int = DEFAULT_MAX_BATCH_SIZE, memory_fraction: float = DEFAULT_MEMORY_FRACTION, @@ -90,6 +114,7 @@ def estimate_kv_tokens( moe_tp_size=shape.moe_tp, moe_ep_size=shape.moe_ep, nextn=nextn, + systems_path=(list(resolve_systems_paths(systems_paths)) if systems_paths is not None else None), allow_naive_fallback=False, ) except ValueError as exc: @@ -114,6 +139,7 @@ def feasible_shape_tokens( backend: str, max_seq_len: int, backend_version: str | None = None, + systems_paths: list[str] | None = None, max_num_tokens: int = DEFAULT_MAX_NUM_TOKENS, max_batch_size: int = DEFAULT_MAX_BATCH_SIZE, memory_fraction: float = DEFAULT_MEMORY_FRACTION, @@ -125,7 +151,7 @@ def feasible_shape_tokens( per distinct shape, so repeated shapes across replica counts are free. """ if backend_version is None: - backend_version = resolve_backend_version(hardware_sku, backend) + backend_version = resolve_backend_version(hardware_sku, backend, systems_paths=systems_paths) feasible: dict[ParallelShape, int] = {} for shape in dict.fromkeys(shapes): # dedup, preserve first-seen order tokens = estimate_kv_tokens( @@ -134,6 +160,7 @@ def feasible_shape_tokens( hardware_sku=hardware_sku, backend=backend, backend_version=backend_version, + systems_paths=systems_paths, max_num_tokens=max_num_tokens, max_batch_size=max_batch_size, memory_fraction=memory_fraction, diff --git a/python/aisimulate/src/aisimulate/sweeper/kv_load.py b/python/aisimulate/src/aisimulate/sweeper/kv_load.py index 6a6a8ddce..0cc5c83d9 100644 --- a/python/aisimulate/src/aisimulate/sweeper/kv_load.py +++ b/python/aisimulate/src/aisimulate/sweeper/kv_load.py @@ -10,6 +10,7 @@ from typing import Any from .config import Workload +from .forward_pass_estimator import resolve_systems_paths from .kv_estimate import estimate_kv_tokens from .parallel_enum import DisaggParallelConfig, ParallelShape, ReplicaParallelConfig @@ -36,6 +37,7 @@ def _per_rank_capacity_tokens( hardware_sku: str, backend: str, backend_version: str, + systems_paths: tuple[str, ...], max_num_tokens: int, max_batch_size: int, memory_fraction: float, @@ -47,6 +49,7 @@ def _per_rank_capacity_tokens( hardware_sku=hardware_sku, backend=backend, backend_version=backend_version, + systems_paths=list(systems_paths), max_num_tokens=max_num_tokens, max_batch_size=max_batch_size, memory_fraction=memory_fraction, @@ -75,12 +78,25 @@ def _role_capacity_tokens( if fixed_blocks is not None: per_rank_tokens = int(fixed_blocks) * block_size else: + resolved = sample.get("forward_pass_estimators", {}).get(role, {}).get("config") + if resolved is None: + timing = sample.get(f"{role}_timing_model") + if isinstance(timing, dict) and timing.get("type") == "external" and timing.get("provider") == "aic": + resolved = timing.get("config") + if not isinstance(resolved, dict): + raise ValueError(f"{role} external AIC timing config must be a mapping") + else: + resolved = {} + roots = resolved.get("systems_paths") + if not roots and resolved.get("systems_path"): + roots = [resolved["systems_path"]] per_rank_tokens = _per_rank_capacity_tokens( config.shape, - model_name=str(sample["model_name"]), - hardware_sku=str(sample.get(f"{role}_hardware_sku") or sample["hardware_sku"]), - backend=str(sample["backend"]), - backend_version=backend_version, + model_name=str(resolved.get("model", resolved.get("model_path", sample["model_name"]))), + hardware_sku=str(resolved.get("system", sample.get(f"{role}_hardware_sku") or sample["hardware_sku"])), + backend=str(resolved.get("backend", sample["backend"])), + backend_version=resolved.get("backend_version", backend_version), + systems_paths=resolve_systems_paths(roots), max_num_tokens=int(sample[f"{role}_max_num_batched_tokens"]), max_batch_size=int(sample[f"{role}_max_num_seqs"]), memory_fraction=float(sample[f"{role}_gpu_memory_utilization"]), diff --git a/python/aisimulate/src/aisimulate/sweeper/model_hw.py b/python/aisimulate/src/aisimulate/sweeper/model_hw.py index 72672267b..b1486f231 100644 --- a/python/aisimulate/src/aisimulate/sweeper/model_hw.py +++ b/python/aisimulate/src/aisimulate/sweeper/model_hw.py @@ -65,7 +65,13 @@ class ModelHardware: num_experts: int = 0 -def resolve_model_hardware(model_name: str, hardware_sku: str, *, backend: str) -> ModelHardware: +def resolve_model_hardware( + model_name: str, + hardware_sku: str, + *, + backend: str, + systems_paths: list[str] | None = None, +) -> ModelHardware: """Read the model weights + SKU spec (via AIC) to derive is_moe / mla / wideep and the model's max context length.""" model_config = get_model_config_from_model_path(model_name) @@ -76,7 +82,9 @@ def resolve_model_hardware(model_name: str, hardware_sku: str, *, backend: str) max_context = model_config.get("context") num_experts = int(model_config.get("num_experts") or model_config.get("n_routed_experts") or 0) - system_spec = perf_database.load_system_spec(hardware_sku) + system_spec = perf_database.load_system_spec( + hardware_sku, **({"systems_paths": systems_paths} if systems_paths is not None else {}) + ) if not system_spec: raise ValueError( f"unknown hardware_sku {hardware_sku!r}: no system config found on AIConfigurator Core's systems path" @@ -117,6 +125,7 @@ def parallel_configs_for( max_batch_size: int = DEFAULT_MAX_BATCH_SIZE, memory_fraction: float = DEFAULT_MEMORY_FRACTION, role_runtime: dict[str, tuple[int, int, float] | tuple[int, int, float, int | None]] | None = None, + systems_paths: list[str] | None = None, ) -> list[ReplicaParallelConfig] | list[DisaggParallelConfig]: """Resolve the model/hardware, then enumerate the parallel configs that fit the GPU budget and can hold a ``max_seq_len``-token sequence. @@ -136,7 +145,12 @@ def parallel_configs_for( :class:`NoViableParallelConfig` when no shape can hold the sequence within the budget. """ - mh = resolve_model_hardware(model_name, hardware_sku, backend=backend) + mh = resolve_model_hardware( + model_name, + hardware_sku, + backend=backend, + systems_paths=systems_paths, + ) seq_len = max_seq_len if max_seq_len is not None else mh.max_context if seq_len is None: raise ValueError(f"max_seq_len is required: {model_name} config exposes no max context length") @@ -180,6 +194,7 @@ def feasible_for(role: str, shapes): hardware_sku=hardware_sku, backend=backend, backend_version=backend_version, + systems_paths=systems_paths, max_seq_len=seq_len, max_num_tokens=role_tokens, max_batch_size=role_batch, diff --git a/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index e55bb3458..03aabb528 100644 --- a/python/aisimulate/src/aisimulate/sweeper/replay.py +++ b/python/aisimulate/src/aisimulate/sweeper/replay.py @@ -21,6 +21,59 @@ REPLAY_SPEC_API_VERSION = 1 +@dataclass(frozen=True) +class ForwardPassEstimatorSpec: + """Resolved output of Core's canonical forward-pass constructor. + + The config is the sole estimator identity carried by Sweeper and Replay. + Convenience properties below are projections, never independently authored + values. Diagnostics preserve Core's selection result for artifacts. + """ + + config: dict[str, JSONValue] + options: dict[str, JSONValue] | None = None + diagnostics: dict[str, JSONValue] = field(default_factory=dict) + + @property + def model_path(self) -> str: + return str(self.config["model"]) + + @property + def system(self) -> str: + return str(self.config["system"]) + + @property + def backend(self) -> str: + return str(self.config["backend"]) + + @property + def backend_version(self) -> str: + return str(self.config["backend_version"]) + + @property + def database_mode(self) -> str: + return str(self.config["database_mode"]) + + @property + def transfer_policy(self) -> tuple[str, ...]: + return tuple(str(value) for value in self.config.get("transfer_policy") or ()) + + @property + def forward_model(self) -> str: + return str(self.config["estimation_mode"]) + + @property + def systems_paths(self) -> tuple[str, ...]: + return tuple(str(value) for value in self.config.get("systems_paths") or ()) + + @property + def performance_data_root(self) -> str: + provenance = self.diagnostics.get("provenance") + if isinstance(provenance, dict): + return str(provenance.get("selected_systems_root") or "") + return "" + + @dataclass(frozen=True) class EncoderPoolSpec: """Resolved analytical EPD pool. Missing power is unavailable, never zero watts.""" @@ -95,6 +148,7 @@ class BackendDeploymentSpec: num_prefill_workers: int = 0 num_decode_workers: int = 0 performance_model_metadata: dict[str, JSONValue] = field(default_factory=dict) + forward_pass_estimators: dict[str, ForwardPassEstimatorSpec] = field(default_factory=dict) encoder: EncoderPoolSpec | None = None diff --git a/python/aisimulate/src/aisimulate/sweeper/search.py b/python/aisimulate/src/aisimulate/sweeper/search.py index f649abb7e..4629b6bf9 100644 --- a/python/aisimulate/src/aisimulate/sweeper/search.py +++ b/python/aisimulate/src/aisimulate/sweeper/search.py @@ -48,6 +48,7 @@ from .deploy import build_backend_deployment from .discovery import resolve_providers from .epd import add_encoder_choices, resolve_encoder_catalog +from .forward_pass_estimator import ForwardPassEstimatorResolver from .kv_estimate import resolve_backend_version from .kv_load import InfeasibleKVCapacity, resolve_kv_load from .provider import ( @@ -575,6 +576,7 @@ def _materialize_one( afd_performance_model: AFDPerformanceModel | None = None, prediction_config_factory: Callable[[dict[str, Any], ReplaySpec], dict[str, Any]] | None = None, encoder_catalog: Mapping[str, Any] | None = None, + estimator_resolver: ForwardPassEstimatorResolver | None = None, ) -> tuple[_PreparedCandidate | None, _EvalResult | None]: """Build a complete replay specification on the main process.""" epd_snapshot = None @@ -603,7 +605,13 @@ def _materialize_one( runner_metadata={}, config_snapshot=epd_snapshot, ) - backend_version = config.search_space.backend_version + estimator_resolver = estimator_resolver or ForwardPassEstimatorResolver(config.search_space) + forward_pass_estimators = estimator_resolver.resolve_candidate(sample) + backend_version = ( + next(iter(forward_pass_estimators.values())).backend_version + if forward_pass_estimators + else config.search_space.requested_backend_version(selection["backend"]) + ) if backend_version is None: if sample["deployment_mode"] == "disagg": role_hardware = {role: sample[f"{role}_hardware_sku"] for role in ("prefill", "decode")} @@ -627,6 +635,7 @@ def _materialize_one( # on the candidate so downstream artifact generation cannot independently # select a different backend version. sample["backend_version"] = backend_version + sample["forward_pass_estimators"] = {role: asdict(value) for role, value in forward_pass_estimators.items()} concurrency = config.workload.concurrency workload_payload = config.workload.model_dump(mode="json") if "traffic_load" in selection: @@ -668,7 +677,9 @@ def _materialize_one( sample["concurrency"] = concurrency if encoder is not None: epd_snapshot = deepcopy(sample) - backend_deployment = build_backend_deployment(sample, backend_version=backend_version, encoder=encoder) + backend_deployment = build_backend_deployment( + sample, backend_version=backend_version, encoder=encoder, forward_pass_estimators=forward_pass_estimators + ) if sample["deployment_mode"] in {"afd", "afd+pd"}: backend_deployment = attach_afd_measurements( backend_deployment, @@ -1090,6 +1101,7 @@ def run( prediction_config_factory = self._prediction_config_factory afd_performance_model = self._afd_performance_model + estimator_resolver = ForwardPassEstimatorResolver(config.search_space) goal = config.goal capabilities = runner_factory.capabilities() capabilities.require_replay_spec_version(REPLAY_SPEC_API_VERSION) @@ -1534,6 +1546,7 @@ def _run_branch_round(state: _BranchSearchState) -> None: prepared, build_result = _materialize_one( suggestion.selection, suggestion.parallel_config, + estimator_resolver=estimator_resolver, config=config, goal=goal, providers=resolved_providers, diff --git a/python/aisimulate/src/aisimulate/sweeper/search_space.py b/python/aisimulate/src/aisimulate/sweeper/search_space.py index 46f26a4b5..7f92428a4 100644 --- a/python/aisimulate/src/aisimulate/sweeper/search_space.py +++ b/python/aisimulate/src/aisimulate/sweeper/search_space.py @@ -209,6 +209,13 @@ def _runner_supports_parallel_config( ) +def _estimator_root_kwargs(search_space, role): + paths = search_space.systems_paths_for(role) + from .forward_pass_estimator import resolve_systems_paths + + return {"systems_paths": list(resolve_systems_paths(paths))} + + def _role_runtime(search_space, backend: str, role: str) -> tuple[int, int, float, int | None]: token_name = f"{role}_max_num_batched_tokens" sequence_name = f"{role}_max_num_seqs" @@ -244,9 +251,12 @@ def _heterogeneous_disagg_configs( ) -> list[DisaggParallelConfig]: """Enumerate each P/D role against its own hardware, then apply the shared budget.""" role_hardware = {role: search_space.hardware_sku_for(role) for role in ("prefill", "decode")} - backend_version = search_space.backend_version + backend_version = search_space.requested_backend_version(backend) if backend_version is None: - role_versions = {role: resolve_backend_version(hardware, backend) for role, hardware in role_hardware.items()} + role_versions = { + role: resolve_backend_version(hardware, backend, **_estimator_root_kwargs(search_space, role)) + for role, hardware in role_hardware.items() + } if len(set(role_versions.values())) != 1: raise NoPerfDatabase( "heterogeneous P/D hardware requires one common backend_version; " @@ -268,6 +278,7 @@ def _heterogeneous_disagg_configs( min_gpu_budget=None, max_seq_len=max_seq_len, role_runtime={"agg": _role_runtime(search_space, backend, role)}, + **_estimator_root_kwargs(search_space, role), ) except (NoPerfDatabase, NoViableParallelConfig) as exc: raise type(exc)(f"{role} hardware_sku={hardware!r}: {exc}") from exc @@ -384,7 +395,7 @@ def _afd_branch( gpu_budget=ss.gpu_budget, deployment_mode="agg", backend=backend, - backend_version=ss.backend_version, + backend_version=ss.requested_backend_version(backend), min_gpu_budget=None, max_seq_len=max_seq_len, role_runtime={"agg": _role_runtime(ss, backend, companion_role)}, @@ -574,7 +585,10 @@ def matches_custom(config: _ParallelConfig) -> bool: ): continue try: - if deployment_mode == "disagg" and (ss.hardware_sku_for("prefill") != ss.hardware_sku_for("decode")): + if deployment_mode == "disagg" and ( + ss.hardware_sku_for("prefill") != ss.hardware_sku_for("decode") + or ss.systems_paths_for("prefill") != ss.systems_paths_for("decode") + ): legal = _heterogeneous_disagg_configs( ss, backend=backend, @@ -587,10 +601,11 @@ def matches_custom(config: _ParallelConfig) -> bool: gpu_budget=ss.gpu_budget, deployment_mode=deployment_mode, backend=backend, - backend_version=ss.backend_version, + backend_version=ss.requested_backend_version(backend), min_gpu_budget=ss.min_gpu_budget, max_seq_len=max_seq_len, role_runtime=_runtime_by_role(ss, backend, deployment_mode), + **_estimator_root_kwargs(ss, "agg" if deployment_mode == "agg" else "prefill"), ) except (NoPerfDatabase, NoViableParallelConfig): continue # backend unusable for this mode -> drop it from the search diff --git a/python/aisimulate/tests/cross_package/test_core_public_api.py b/python/aisimulate/tests/cross_package/test_core_public_api.py index c932a00a7..16e153f45 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -26,6 +26,8 @@ EXPECTED_FACADE = { "AttentionBackend", "EngineHandle", + "ForwardPassPerfModelConfig", + "ForwardPassPerfOptions", "ModelConfig", "MoEBackend", "RuntimeConfig", @@ -36,6 +38,20 @@ } +def _raw_regression_model(worker_type, options_json=None, *, cls=None): + cls = cls or aiconfigurator_core.RustForwardPassPerfModel + config = { + "model": "test/model", + "system": "test", + "backend": "vllm", + "worker_type": worker_type, + "estimation_mode": "fpm_regression", + } + if options_json is not None: + config["estimator_config"] = json.loads(cls.legacy_estimator_config(options_json)) + return cls.best_available(json.dumps(config)) + + def test_sdk_facade_import_is_lazy_in_a_fresh_interpreter() -> None: script = """ import sys @@ -90,49 +106,22 @@ def test_stable_function_signatures() -> None: ) assert "scheduler_block_size" in inspect.signature(estimate_num_gpu_blocks).parameters assert "memory_fraction_kind" in inspect.signature(estimate_kv_cache).parameters - assert str(inspect.signature(RustForwardPassPerfModel.from_regression)) == ( - "(worker_type: 'str', options: 'dict[str, Any] | None' = None) -> 'RustForwardPassPerfModel'" - ) - assert str(inspect.signature(RustForwardPassPerfModel.best_available)) == ( - "(config: 'dict[str, Any]', worker_type: 'str', " - "options: 'dict[str, Any] | None' = None) -> 'RustForwardPassPerfModel'" - ) + assert list(inspect.signature(RustForwardPassPerfModel.best_available).parameters) == ["config"] + assert not hasattr(RustForwardPassPerfModel, "from_regression") + assert not hasattr(RustForwardPassPerfModel, "from_native") @pytest.mark.parametrize("worker_type", ["agg", "Prefill", "AGGREGATED", ""]) def test_raw_fpm_binding_rejects_worker_type_aliases(worker_type: str) -> None: - with pytest.raises(ValueError, match="invalid worker_type"): - aiconfigurator_core.RustForwardPassPerfModel.from_regression(worker_type) - config = json.dumps( - { - "schema_version": 1, - "model_name": "this/model-is-not-compiled", - "system_name": "b200_sxm", - "backend": "vllm", - "backend_version": "0.19.0", - "tp_size": 1, - "pp_size": 1, - "attention_dp_size": 1, - "moe_tp_size": None, - "moe_ep_size": None, - "weight_dtype": None, - "activation_dtype": None, - "moe_dtype": None, - "kv_cache_dtype": None, - "kv_block_size": None, - "nextn": None, - "extra": {}, - } - ) - with pytest.raises(ValueError, match="invalid worker_type"): - aiconfigurator_core.RustForwardPassPerfModel.best_available(config, worker_type) + with pytest.raises(ValueError, match="worker_type"): + _raw_regression_model(worker_type) def test_raw_fpm_binding_requires_worker_type() -> None: - with pytest.raises(TypeError): - aiconfigurator_core.RustForwardPassPerfModel.from_regression() - with pytest.raises(TypeError): - aiconfigurator_core.RustForwardPassPerfModel.best_available("{}") + with pytest.raises(ValueError, match="worker_type"): + aiconfigurator_core.RustForwardPassPerfModel.best_available( + json.dumps({"model": "m", "system": "s", "backend": "vllm"}) + ) @pytest.mark.parametrize("worker_type", ["prefill", "decode", "aggregated"]) @@ -144,9 +133,9 @@ def test_raw_fpm_binding_constructs_every_worker_type_with_weights(worker_type: "regression_ffn_token_weight": 4.0, } ) - model = aiconfigurator_core.RustForwardPassPerfModel.from_regression(worker_type, options) + model = _raw_regression_model(worker_type, options) diagnostics = json.loads(model.diagnostics()) - assert diagnostics == { + assert {key: value for key, value in diagnostics.items() if key != "provenance"} == { "source": "fallback_regression", "readiness": "insufficient_data", "retained_observations": 0, @@ -213,7 +202,7 @@ def _raw_regression_iteration(worker_type: str, index: int) -> dict[str, object] @pytest.mark.parametrize("worker_type", ["prefill", "decode", "aggregated"]) def test_raw_fpm_binding_regression_round_trip(worker_type: str) -> None: - model = aiconfigurator_core.RustForwardPassPerfModel.from_regression( + model = _raw_regression_model( worker_type, '{"min_observations":5}', ) @@ -223,7 +212,7 @@ def test_raw_fpm_binding_regression_round_trip(worker_type: str) -> None: model.tune_with_fpms(json.dumps(iterations)) diagnostics = json.loads(model.diagnostics()) - assert diagnostics == { + assert {key: value for key, value in diagnostics.items() if key != "provenance"} == { "source": "fallback_regression", "readiness": "ready", "retained_observations": 6, @@ -236,7 +225,7 @@ def test_raw_fpm_binding_regression_round_trip(worker_type: str) -> None: def test_raw_fpm_binding_validates_regression_weights() -> None: with pytest.raises(ValueError, match="regression_attention_kv_weight"): - aiconfigurator_core.RustForwardPassPerfModel.from_regression( + _raw_regression_model( "prefill", '{"regression_attention_kv_weight": 0.0}', ) @@ -257,15 +246,15 @@ def test_raw_fpm_binding_decodes_exact_nonfinite_weight_sentinels( ) -> None: """Raw ``options_json`` callers use quoted sentinels rather than non-JSON numbers.""" with pytest.raises(ValueError, match=field): - aiconfigurator_core.RustForwardPassPerfModel.from_regression( + _raw_regression_model( "aggregated", json.dumps({field: sentinel}), ) def test_raw_fpm_binding_rejects_unknown_nonfinite_weight_sentinel() -> None: - with pytest.raises(ValueError, match="invalid options JSON"): - aiconfigurator_core.RustForwardPassPerfModel.from_regression( + with pytest.raises(ValueError, match="invalid legacy options"): + _raw_regression_model( "decode", '{"regression_ffn_token_weight":"Inf"}', ) @@ -313,7 +302,7 @@ def test_regression_bucket_diagnostics_stub_matches_native_contract(namespace: s assert [argument.arg for argument in method.args.args] == ["self"] assert ast.unparse(method.returns) == "str" - native_model = importlib.import_module(namespace).RustForwardPassPerfModel.from_regression("aggregated") + native_model = _raw_regression_model("aggregated", cls=importlib.import_module(namespace).RustForwardPassPerfModel) diagnostics = getattr(native_model, method_name)() assert isinstance(diagnostics, str) assert json.loads(diagnostics) == [ diff --git a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py index 8487fca3e..440c5a57e 100644 --- a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py +++ b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py @@ -24,6 +24,39 @@ _GENERATION_FALLBACK = MoECommFallback(*_GENERATION_FALLBACK_PAYLOAD) +def _regression_model(cls, worker_type, options=None): + import aiconfigurator_core + + config = { + "model": "test/model", + "system": "test", + "backend": "vllm", + "worker_type": worker_type, + "estimation_mode": "fpm_regression", + } + if options is not None: + config["estimator_config"] = json.loads( + aiconfigurator_core.RustForwardPassPerfModel.legacy_estimator_config( + rust_engine_step._optional_json_dumps(options) + ) + ) + return cls.best_available(config) + + +def _native_model(cls, config, options=None): + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig + + return cls.best_available(ForwardPassPerfModelConfig.from_legacy_engine_config(config, "aggregated", options)) + + +def _best_model(cls, config, worker_type, options=None): + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig + + return cls.best_available( + ForwardPassPerfModelConfig.from_legacy_engine_config(config, worker_type, options, allow_regression=True) + ) + + def test_should_use_rust_engine_step_supports_runtime_config_and_env(monkeypatch) -> None: monkeypatch.setenv("AICONFIGURATOR_ENGINE_STEP_BACKEND", "rust") @@ -819,7 +852,7 @@ def test_forward_pass_perf_model_regression_stores_end_to_end() -> None: """The public compiled API keeps each full-rank workload fit independent.""" from aisimulate_core.sdk import RustForwardPassPerfModel - model = RustForwardPassPerfModel.from_regression("aggregated") + model = _regression_model(RustForwardPassPerfModel, "aggregated") kinds = ["pure_decode", "contains_locally_mixed", "cross_rank_aggregated", "pure_prefill"] assert model.regression_store_diagnostics() == [ {"workload_kind": kind, "ready": False, "retained_observations": 0} for kind in kinds @@ -862,133 +895,46 @@ def iteration(kind: str, scale: int) -> list[dict[str, object]]: assert model.estimate_forward_pass_time_ms(iteration(cold_kind, 3)) is None -def test_forward_pass_perf_model_constructors_forward_exact_worker_type(monkeypatch) -> None: - """Regression-capable constructors preserve the engine-level role and - reject aliases before crossing the extension boundary.""" +def test_forward_pass_constructor_passes_complete_typed_request(monkeypatch) -> None: import aiconfigurator_core - calls: list[tuple[str, object, object, object | None]] = [] - - class _FakeRawModel: - @staticmethod - def from_regression(worker_type, options_json=None): - calls.append(("from_regression", worker_type, options_json, None)) - return object() - - @staticmethod - def best_available(config_json, worker_type, options_json=None): - calls.append(("best_available", json.loads(config_json), worker_type, options_json)) - return object() - - monkeypatch.setattr(aiconfigurator_core, "RustForwardPassPerfModel", _FakeRawModel) - monkeypatch.setattr(rust_engine_step, "_configure_default_data_roots", lambda: None) - - for worker_type in ("prefill", "decode", "aggregated"): - rust_engine_step.RustForwardPassPerfModel.from_regression( - worker_type, - {"min_observations": 2}, - ) - rust_engine_step.RustForwardPassPerfModel.best_available( - {"schema_version": 1}, - "aggregated", - {"min_observations": 3}, - ) - - assert calls[:3] == [ - ("from_regression", "prefill", '{"min_observations":2}', None), - ("from_regression", "decode", '{"min_observations":2}', None), - ("from_regression", "aggregated", '{"min_observations":2}', None), - ] - assert calls[3] == ( - "best_available", - {"schema_version": 1}, - "aggregated", - '{"min_observations":3}', - ) - - for invalid in ("agg", "Prefill", ""): - with pytest.raises(ValueError, match="invalid worker_type"): - rust_engine_step.RustForwardPassPerfModel.from_regression(invalid) - with pytest.raises(ValueError, match="invalid worker_type"): - rust_engine_step.RustForwardPassPerfModel.best_available({}, invalid) - - -def test_forward_pass_perf_model_constructors_forward_all_regression_weights(monkeypatch) -> None: - """The facade preserves every public regression knob by its wire name.""" - import aiconfigurator_core - - calls: list[tuple[str, dict[str, object]]] = [] - - class _FakeRawModel: - @staticmethod - def from_regression(worker_type, options_json=None): - calls.append((worker_type, json.loads(options_json))) - return object() + calls = [] + class RawModel: @staticmethod - def best_available(config_json, worker_type, options_json=None): - calls.append((worker_type, json.loads(options_json))) + def best_available(config_json): + calls.append(json.loads(config_json)) return object() - monkeypatch.setattr(aiconfigurator_core, "RustForwardPassPerfModel", _FakeRawModel) - monkeypatch.setattr(rust_engine_step, "_configure_default_data_roots", lambda: None) - options = { - "regression_attention_kv_weight": 2.0, - "regression_prefill_attention_pair_weight": 3.0, - "regression_ffn_token_weight": 4.0, + monkeypatch.setattr(aiconfigurator_core, "RustForwardPassPerfModel", RawModel) + payload = { + "model": "test", + "system": "test", + "backend": "vllm", + "worker_type": "aggregated", + "estimation_mode": "fpm_regression", + "fallback_policy": "deny", + "estimator_config": { + "features": {"attention_kv_weight": 2.0, "prefill_attention_pair_weight": 3.0, "ffn_token_weight": 4.0}, + "fpm_regression": {"sampling": {"bins_per_axis": [4, 16], "max_observations": 128}}, + "correction": {"enabled": False}, + }, } + rust_engine_step.RustForwardPassPerfModel.best_available(payload) + assert calls == [payload] + assert not hasattr(rust_engine_step.RustForwardPassPerfModel, "from_native") + assert not hasattr(rust_engine_step.RustForwardPassPerfModel, "from_regression") - rust_engine_step.RustForwardPassPerfModel.from_regression("aggregated", options) - rust_engine_step.RustForwardPassPerfModel.best_available({}, "decode", options) - - assert calls == [("aggregated", options), ("decode", options)] - - -def test_forward_pass_perf_model_marshals_nonfinite_weights_as_valid_json_without_mutation( - monkeypatch, -) -> None: - """The ergonomic facade owns the float-to-sentinel transport, not its caller.""" - import aiconfigurator_core - - calls: list[str] = [] - - class _FakeRawModel: - @staticmethod - def from_native(config_json, options_json=None): - calls.append(options_json) - return object() - - @staticmethod - def from_regression(worker_type, options_json=None): - calls.append(options_json) - return object() - - @staticmethod - def best_available(config_json, worker_type, options_json=None): - calls.append(options_json) - return object() - - monkeypatch.setattr(aiconfigurator_core, "RustForwardPassPerfModel", _FakeRawModel) - monkeypatch.setattr(rust_engine_step, "_configure_default_data_roots", lambda: None) - options = { - "regression_attention_kv_weight": float("nan"), - "regression_prefill_attention_pair_weight": float("inf"), - "regression_ffn_token_weight": float("-inf"), - } - rust_engine_step.RustForwardPassPerfModel.from_native({}, options) - rust_engine_step.RustForwardPassPerfModel.from_regression("aggregated", options) - rust_engine_step.RustForwardPassPerfModel.best_available({}, "decode", options) +def test_forward_pass_config_requires_role_and_defaults_to_auto_deny() -> None: + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig - expected = ( - '{"regression_attention_kv_weight":"NaN",' - '"regression_ffn_token_weight":"-Infinity",' - '"regression_prefill_attention_pair_weight":"Infinity"}' - ) - assert calls == [expected, expected, expected] - assert math.isnan(options["regression_attention_kv_weight"]) - assert options["regression_prefill_attention_pair_weight"] == math.inf - assert options["regression_ffn_token_weight"] == -math.inf + with pytest.raises(TypeError, match="worker_type"): + ForwardPassPerfModelConfig(model="m", system="s", backend="vllm") + config = ForwardPassPerfModelConfig(model="m", system="s", backend="vllm", worker_type="decode") + assert config.estimation_mode == "auto" + assert config.fallback_policy == "deny" + assert config.estimator_config == {} def _supported_fpm_config() -> dict[str, object]: @@ -1038,7 +984,8 @@ def test_nemotron_super_fp8_native_estimation_uses_packaged_moe_data() -> None: pytest.importorskip("aiconfigurator_core") from aiconfigurator_core.sdk.rust_engine_step import RustForwardPassPerfModel - model = RustForwardPassPerfModel.from_native( + model = _native_model( + RustForwardPassPerfModel, { "schema_version": 1, "model_name": "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8", @@ -1113,7 +1060,8 @@ def test_forward_pass_perf_model_native_default_directional_bounds_end_to_end() from aiconfigurator.sdk.rust_engine_step import RustForwardPassPerfModel config = _supported_fpm_config() - model = RustForwardPassPerfModel.from_native( + model = _native_model( + RustForwardPassPerfModel, config, { "min_observations": 2, @@ -1196,7 +1144,8 @@ def test_forward_pass_perf_model_best_available_falls_back_on_bad_config() -> No from aiconfigurator.sdk.rust_engine_step import RustForwardPassPerfModel config = _unsupported_fpm_config() - model = RustForwardPassPerfModel.best_available( + model = _best_model( + RustForwardPassPerfModel, config, "aggregated", {"min_observations": 2}, @@ -1214,7 +1163,8 @@ def test_forward_pass_perf_model_best_available_falls_back_on_bad_config() -> No } assert model.estimate_forward_pass_time_ms(decode_fpm) is None - prefill_model = RustForwardPassPerfModel.best_available( + prefill_model = _best_model( + RustForwardPassPerfModel, config, "prefill", {"min_observations": 2}, @@ -1235,7 +1185,7 @@ def test_best_available_falls_back_on_malformed_system_yaml(tmp_path: Path) -> N config = _supported_fpm_config() config.update(system_name="broken", systems_path=str(systems_root)) - model = RustForwardPassPerfModel.best_available(config, "aggregated") + model = _best_model(RustForwardPassPerfModel, config, "aggregated") diagnostics = model.diagnostics() assert diagnostics["source"] == "fallback_regression" @@ -1266,9 +1216,9 @@ def test_malformed_performance_yaml_remains_a_perf_database_failure(tmp_path: Pa config.update(system_name="synthetic", systems_path=str(systems_root)) with pytest.raises(PerfDataNotAvailableError, match="perf database error"): - RustForwardPassPerfModel.from_native(config) + _native_model(RustForwardPassPerfModel, config) - model = RustForwardPassPerfModel.best_available(config, "aggregated") + model = _best_model(RustForwardPassPerfModel, config, "aggregated") diagnostics = model.diagnostics() assert diagnostics["source"] == "fallback_regression" assert "perf database error" in diagnostics["last_warning"] @@ -1289,7 +1239,8 @@ def test_best_available_validates_regression_weights_only_after_fallback() -> No }, } - native = RustForwardPassPerfModel.best_available( + native = _best_model( + RustForwardPassPerfModel, _supported_fpm_validation_config(), "prefill", invalid_regression_options, @@ -1299,7 +1250,8 @@ def test_best_available_validates_regression_weights_only_after_fallback() -> No assert native_prediction is not None and native_prediction > 0.0 with pytest.raises(ValueError, match="regression_attention_kv_weight"): - RustForwardPassPerfModel.best_available( + _best_model( + RustForwardPassPerfModel, _unsupported_fpm_config(), "prefill", invalid_regression_options, @@ -1319,8 +1271,8 @@ def test_native_constructors_ignore_all_nonfinite_regression_weights() -> None: } config = _supported_fpm_validation_config() - strict_native = RustForwardPassPerfModel.from_native(config, options) - best_native = RustForwardPassPerfModel.best_available(config, "aggregated", options) + strict_native = _native_model(RustForwardPassPerfModel, config, options) + best_native = _best_model(RustForwardPassPerfModel, config, "aggregated", options) assert strict_native.diagnostics()["source"] == "aic" assert best_native.diagnostics()["source"] == "aic" @@ -1345,9 +1297,10 @@ def test_regression_constructors_decode_and_reject_nonfinite_weight_sentinels( options = {field: value} with pytest.raises(ValueError, match=field): - RustForwardPassPerfModel.from_regression("aggregated", options) + _regression_model(RustForwardPassPerfModel, "aggregated", options) with pytest.raises(ValueError, match=field): - RustForwardPassPerfModel.best_available( + _best_model( + RustForwardPassPerfModel, _unsupported_fpm_config(), "aggregated", options, @@ -1653,7 +1606,6 @@ def test_large_ep_op_graph_compiles_natively(caplog): the scalar engine-step keys and match the Python step on the same config.""" import logging - import math from aiconfigurator.sdk.backends.factory import get_backend from aiconfigurator.sdk.engine import build_engine_spec_json @@ -2149,3 +2101,36 @@ def _decode_step_per_op_with_metadata(self, *args, **kwargs): with pytest.raises(expected, match=str(error)) as caught: getattr(rust_engine_step, entrypoint)(None, None, gen_tokens=1, isl=8, osl=4) assert (caught.value.__cause__ if perf_miss else caught.value) is error + + +@pytest.mark.integration +def test_canonical_native_selection_retries_roots_and_pins_effective_configuration(tmp_path): + from importlib import resources + + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + packaged = str(resources.files("aiconfigurator_core") / "systems") + config = ForwardPassPerfModelConfig( + model="Qwen/Qwen3-32B", + system="h200_sxm", + backend="vllm", + worker_type="aggregated", + tp=2, + estimation_mode="op_level", + fallback_policy="deny", + systems_paths=(str(tmp_path), packaged), + database_mode="HYBRID", + transfer_policy="conservative", + estimator_config={"correction": {"enabled": False}}, + ) + model = RustForwardPassPerfModel.best_available(config) + provenance = model.diagnostics()["provenance"] + assert provenance["selected_systems_root"] == packaged + resolved = provenance["config"] + assert resolved["systems_paths"] == [packaged] + assert resolved["backend_version"] + assert resolved["estimation_mode"] == "op_level" + assert resolved["fallback_policy"] == "deny" + assert resolved["transfer_policy"] == ["xshape"] + assert resolved["estimator_config"]["correction"]["enabled"] is False + assert config.systems_paths == (str(tmp_path), packaged) diff --git a/scripts/check_release_migrations.py b/scripts/check_release_migrations.py new file mode 100644 index 000000000..ec395b9d4 --- /dev/null +++ b/scripts/check_release_migrations.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Block release staging while source-controlled downstream migrations remain.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +GATES = Path(__file__).resolve().parents[1] / ".github" / "release-gates.json" + + +def require_completed_migrations(path: Path = GATES) -> None: + document = json.loads(path.read_text()) + if not isinstance(document, dict) or set(document) != {"pending_migrations"}: + raise ValueError("release gates must declare pending_migrations") + pending = document["pending_migrations"] + if not isinstance(pending, list): + raise ValueError("pending_migrations must be a list") + for item in pending: + if ( + not isinstance(item, dict) + or set(item) != {"pull_request", "requirement"} + or any(not isinstance(value, str) or not value.strip() for value in item.values()) + ): + raise ValueError("each release gate needs a pull_request and requirement") + if pending: + details = "\n".join(f"- {item['pull_request']}: {item['requirement']}" for item in pending) + raise RuntimeError(f"Release publication is blocked by pending downstream migrations:\n{details}") + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--target-gates", type=Path, help="Also require the publication target's migrations to be complete" + ) + args = parser.parse_args(argv) + try: + require_completed_migrations(GATES) + if args.target_gates is not None: + require_completed_migrations(args.target_gates) + except (OSError, ValueError, RuntimeError) as error: + print(error, file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/fpm_accuracy/README.md b/scripts/fpm_accuracy/README.md index 7e0b703e6..ce9b3d05b 100644 --- a/scripts/fpm_accuracy/README.md +++ b/scripts/fpm_accuracy/README.md @@ -22,6 +22,10 @@ that all of these observations were deliberately filtered out. Native AISim imports remain deferred in the adapted adapters so parser and fake predictor tests work without an installed native extension. The real campaign checks that the native SDK is installed before evaluating any case. +When the evaluated wheel provides `ForwardPassPerfModelConfig`, both FPM and +regression use `best_available(config)`, with explicit mode, worker identity, +and migrated tuning options. Legacy constructors are used only to evaluate +older branch wheels that do not provide the canonical configuration type. Hub cache loading supports repository-local blobs and the marked cache-wide shared blob store used by huggingface-hub 1.32. Manifest hashes still bind the diff --git a/scripts/fpm_accuracy/models/aic_predictors.py b/scripts/fpm_accuracy/models/aic_predictors.py index e6e0a0016..861e2a5e7 100644 --- a/scripts/fpm_accuracy/models/aic_predictors.py +++ b/scripts/fpm_accuracy/models/aic_predictors.py @@ -8,6 +8,7 @@ import math from collections.abc import Mapping, Sequence +from dataclasses import replace from typing import Any, ClassVar, Literal from fpm_accuracy.exceptions import ConfigurationError, DependencyError @@ -46,7 +47,15 @@ def create(cls, context: PredictorContext) -> _AicPredictor: if cls.mode == "regression": if not callable(getattr(perf_model, "regression_store_diagnostics", None)): raise DependencyError("This AISim revision does not support worker-scoped regression.") - model = perf_model.from_regression(context.worker_role, options) + config_type = _canonical_config_type() + if config_type is not None: + request = config_type.from_legacy_engine_config( + cls._native_engine_config(context), context.worker_role, options + ) + model = perf_model.best_available(replace(request, estimation_mode="fpm_regression")) + else: + # The evaluator also runs against older released branch wheels. + model = perf_model.from_regression(context.worker_role, options) try: store_diagnostics = getattr(model, "regression_store_diagnostics", None) if not callable(store_diagnostics): @@ -71,7 +80,12 @@ def create(cls, context: PredictorContext) -> _AicPredictor: engine_config["systems_path"] = str(prepared_fpm.systems_root) engine_config["forward_model"] = "fpm" try: - model = perf_model.from_native(engine_config, options) + config_type = _canonical_config_type() + if config_type is not None: + request = config_type.from_legacy_engine_config(engine_config, context.worker_role, options) + model = perf_model.best_available(request) + else: + model = perf_model.from_native(engine_config, options) except Exception: if prepared_fpm is not None: prepared_fpm.close() @@ -156,4 +170,10 @@ def _import_aisim_forward_pass_perf_model() -> Any: return RustForwardPassPerfModel +def _canonical_config_type() -> Any: + from aisimulate_core import sdk + + return getattr(sdk, "ForwardPassPerfModelConfig", None) + + __all__ = ["AicFpmPredictor", "AicRegressionPredictor"] diff --git a/tests/fpm_accuracy/test_evaluation.py b/tests/fpm_accuracy/test_evaluation.py index c9ed02241..fb1c1ea20 100644 --- a/tests/fpm_accuracy/test_evaluation.py +++ b/tests/fpm_accuracy/test_evaluation.py @@ -3,6 +3,7 @@ import json from dataclasses import replace +from types import SimpleNamespace import pytest from fpm_accuracy.evaluate import Metric, choose_variant, evaluate_case @@ -220,6 +221,91 @@ def factory(method, context): assert result["warmup"]["metrics"]["all"]["predicted_count"] == 12 +def test_real_regression_uses_canonical_identity_and_options(case): + pytest.importorskip("aisimulate_core.sdk") + from fpm_accuracy.models.fpt_predictor import PredictorContext + + context = PredictorContext( + worker=case.configuration.worker_config_record, + worker_role="decode", + options={"max_observations": 32, "min_observations": 6}, + ) + predictor = aic_predictors.AicRegressionPredictor.create(context) + try: + config = predictor.diagnostics()["provenance"]["config"] + assert config["worker_type"] == "decode" + assert config["estimation_mode"] == "fpm_regression" + assert config["fallback_policy"] == "deny" + assert config["estimator_config"]["fpm_regression"]["sampling"]["max_observations"] == 32 + assert config["estimator_config"]["fpm_regression"]["min_observations"] == 6 + finally: + predictor.close() + + +@pytest.mark.parametrize("mode", ["fpm", "regression"]) +@pytest.mark.parametrize("canonical", [True, False]) +def test_predictor_uses_canonical_api_or_older_wheel_adapter(case, tmp_path, monkeypatch, mode, canonical): + from fpm_accuracy.models.fpt_predictor import PredictorContext + + if canonical: + pytest.importorskip("aisimulate_core.sdk") + else: + monkeypatch.setattr(aic_predictors, "_canonical_config_type", lambda: None) + calls = [] + + class Model: + @classmethod + def best_available(cls, config): + assert canonical + calls.append(config.to_dict()) + return cls() + + @classmethod + def from_native(cls, config, options): + assert not canonical + calls.append((config, options)) + return cls() + + @classmethod + def from_regression(cls, worker_type, options): + assert not canonical + calls.append((worker_type, options)) + return cls() + + def regression_store_diagnostics(self): + return [] + + def close(self): + pass + + monkeypatch.setattr(aic_predictors, "_import_aisim_forward_pass_perf_model", lambda: Model) + monkeypatch.setattr( + aic_predictors, + "prepare_aic_fpm_database", + lambda *args: SimpleNamespace( + systems_root=tmp_path, + close=lambda: None, + ), + ) + context = PredictorContext( + worker=case.configuration.worker_config_record, + worker_role="decode", + options={"max_observations": 32}, + fpm_artifact=object(), + ) + cls = aic_predictors.AicFpmPredictor if mode == "fpm" else aic_predictors.AicRegressionPredictor + predictor = cls.create(context) + predictor.close() + assert len(calls) == 1 + if canonical: + assert calls[0]["worker_type"] == "decode" + assert calls[0]["estimation_mode"] == ("fpm_interpolation" if mode == "fpm" else "fpm_regression") + assert calls[0]["fallback_policy"] == "deny" + assert calls[0]["estimator_config"]["fpm_regression"]["sampling"]["max_observations"] == 32 + if mode == "fpm": + assert calls[0]["systems_paths"] == [str(tmp_path)] + + def test_micro_mape_and_variant_order(): metric = Metric() metric.add(10, 20, False, False) diff --git a/tests/sweeper/test_forward_pass_estimator.py b/tests/sweeper/test_forward_pass_estimator.py new file mode 100644 index 000000000..0de5bb384 --- /dev/null +++ b/tests/sweeper/test_forward_pass_estimator.py @@ -0,0 +1,599 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canonical construction through the real Python/Rust boundary.""" + +import json +from pathlib import Path + +import pytest + +from aiconfigurator_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + +pytestmark = [pytest.mark.unit, pytest.mark.pre_merge, pytest.mark.gpu_0] + + +@pytest.fixture +def custom_systems(tmp_path): + from importlib.resources import files + + packaged = Path(str(files("aiconfigurator_core") / "systems")) + for entry in packaged.iterdir(): + (tmp_path / entry.name).symlink_to(entry, target_is_directory=entry.is_dir()) + (tmp_path / "review_h200.yaml").write_text((packaged / "h200_sxm.yaml").read_text()) + return tmp_path + + +@pytest.mark.parametrize("discovery", ["sdk", "ordered_sdk", "environment"]) +def test_omitted_roots_preserve_configured_discovery(custom_systems, monkeypatch, discovery): + import aiconfigurator_core + from aiconfigurator_core.sdk import perf_database + from aisimulate.runner import EngineReplayRunnerFactory + from aisimulate.sweeper.config import SmartSearchConfig + from aisimulate.sweeper.search import Sweeper + + monkeypatch.setattr(perf_database, "_SYSTEMS_PATHS", perf_database.get_systems_paths()) + monkeypatch.delenv("AICONFIGURATOR_SYSTEMS_PATH", raising=False) + roots = [str(custom_systems)] + if discovery == "ordered_sdk": + empty = custom_systems / "empty" + empty.mkdir() + roots.insert(0, str(empty)) + if discovery in {"sdk", "ordered_sdk"}: + perf_database.set_systems_paths(roots) + else: + perf_database.set_systems_paths("default") + monkeypatch.setenv("AICONFIGURATOR_SYSTEMS_PATH", str(custom_systems)) + cfg = SmartSearchConfig( + search_space={ + "model_name": "Qwen/Qwen3-32B", + "hardware_sku": "review_h200", + "backend": ["vllm"], + "backend_version": "0.24.0", + "deployment_mode": ["agg"], + "gpu_budget": 2, + "context_length": 4096, + "parallel_configs": [{"tp": 2, "replicas": 1}], + "agg_max_num_batched_tokens": [8192], + "agg_max_num_seqs": [256], + "agg_num_gpu_blocks": 256, + "agg_block_size": 64, + }, + workload={"isl": 128, "osl": 2, "request_count": 1, "concurrency": 1}, + sweep={"max_rounds": 1, "candidates_per_round": 1, "parallel_evals": 1}, + ) + assert cfg.search_space.systems_paths is None + cfg = SmartSearchConfig.model_validate_json(cfg.model_dump_json()) + result = Sweeper(runner_factory=EngineReplayRunnerFactory(), show_progress=False).run(cfg, top_n=None) + assert result.model_dump(mode="json")["counts"]["feasible"] == 1 + raw = aiconfigurator_core.RustForwardPassPerfModel.best_available( + json.dumps( + { + "model": "Qwen/Qwen3-32B", + "system": "review_h200", + "backend": "vllm", + "backend_version": "0.24.0", + "worker_type": "aggregated", + "tp": 2, + } + ) + ) + assert json.loads(raw.diagnostics())["provenance"]["config"]["systems_paths"] == [str(custom_systems)] + from aisimulate.sweeper.forward_pass_estimator import resolve_systems_paths + + assert resolve_systems_paths(None) == tuple(roots) + assert resolve_systems_paths(["default"]) != (str(custom_systems),) + + +def test_prediction_pins_populated_version_slots(custom_systems): + from dataclasses import replace + + from aiconfigurator_core.sdk.perf_database import resolve_query_version + from aisimulate.compiler import prediction_to_replay_spec + from aisimulate.config.cli import CorePredictionConfig + from aisimulate.runner import EngineReplayRunnerFactory + + next_version = resolve_query_version("review_h200", "vllm", "next", systems_paths=[str(custom_systems)]) + for requested, literal in [("current", "0.24.0"), ("next", next_version), ("0.24.0", "0.24.0")]: + cfg = CorePredictionConfig.model_validate( + { + "engine": { + "model": "Qwen/Qwen3-32B", + "hardware": "review_h200", + "backend": "vllm", + "backend_version": requested, + "context_length": 4096, + "systems_paths": [str(custom_systems)], + "workers": { + "aggregated": { + "parallelism": {"tensor": 2}, + "kv_cache": {"capacity": {"type": "fixed", "blocks": 128}}, + } + }, + }, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 8, "output_tokens": 2}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 1}, + }, + } + ) + replay = prediction_to_replay_spec(cfg) + assert replay.backend_deployment.backend_version == literal + runner = EngineReplayRunnerFactory().create(0) + try: + assert runner.run(replay).metrics["completed_requests"] == 1 + if requested == "current": + bad = replace(replay, backend_deployment=replace(replay.backend_deployment, backend_version="0.19.0")) + with pytest.raises(ValueError, match="conflicts with"): + runner.run(bad) + finally: + runner.close() + + +@pytest.mark.parametrize("mode", ["agg", "disagg"]) +@pytest.mark.parametrize("ratio", [0.5, 1.0]) +def test_kv_relative_load_uses_resolved_roots(custom_systems, mode, ratio): + from dataclasses import asdict + + from aisimulate.sweeper.config import SearchSpace, Workload + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + from aisimulate.sweeper.kv_load import resolve_kv_load + from aisimulate.sweeper.parallel_enum import DisaggParallelConfig, ParallelShape, ReplicaParallelConfig + from aisimulate.sweeper.sample import unroll_sample + + role = ReplicaParallelConfig(ParallelShape(tp=2, dp=1, moe_tp=1, moe_ep=1), replicas=1) + parallel = role if mode == "agg" else DisaggParallelConfig(prefill=role, decode=role) + space = SearchSpace( + model_name="Qwen/Qwen3-32B", + hardware_sku="review_h200", + backend=["vllm"], + deployment_mode=[mode], + systems_paths=[str(custom_systems)], + gpu_budget=4, + ) + roles = ("agg",) if mode == "agg" else ("prefill", "decode") + selection = {"deployment_mode": mode, "backend": "vllm"} + for name in roles: + selection.update({f"{name}_max_num_batched_tokens": 8192, f"{name}_max_num_seqs": 256}) + sample = unroll_sample(search_space=space, selection=selection, parallel_config=parallel) + estimators = ForwardPassEstimatorResolver(space).resolve_candidate(sample) + sample["forward_pass_estimators"] = {name: asdict(estimator) for name, estimator in estimators.items()} + result = resolve_kv_load( + sample, + workload=Workload(isl=128, osl=2, kv_load_ratio=ratio, request_count=1), + parallel_config=parallel, + ratio=ratio, + backend_version="0.24.0", + ) + assert result.concurrency == max(1, int(ratio * result.concurrency_capacity)) + assert set(result.role_capacity_tokens) == set(roles) + assert all(tokens > 0 for tokens in result.role_capacity_tokens.values()) + + +def request(**changes): + return ForwardPassPerfModelConfig( + model="Qwen/Qwen3-32B", + system="h200_sxm", + backend="vllm", + worker_type="aggregated", + **changes, + ) + + +def test_legacy_options_preserve_shape_ridge_and_general_mapping(): + from collections.abc import Mapping + + from aiconfigurator_core.sdk import ForwardPassPerfOptions + + class OptionsMapping(Mapping): + def __init__(self, values): + self.values = values + + def __getitem__(self, key): + return self.values[key] + + def __iter__(self): + return iter(self.values) + + def __len__(self): + return len(self.values) + + legacy = { + "schema_version": 1, + "model_name": "m", + "system_name": "h200_sxm", + "backend": "vllm", + "tp_size": 1, + "pp_size": 1, + } + options = ForwardPassPerfOptions(bucket_shape=(2, 8), regression_ridge_scale=0.25) + payload = options.to_dict() + for value in [options, OptionsMapping(payload)]: + migrated = ForwardPassPerfModelConfig.from_legacy_engine_config(legacy, "decode", value) + regression = migrated.estimator_config["fpm_regression"] + assert regression["sampling"]["bins_per_axis"] == [2, 8] + assert regression["fit"]["singular_ridge_scale"] == 0.25 + assert migrated.estimator_config["correction"]["sampling"]["bins_per_axis"] == [2, 8] + assert payload == options.to_dict() + + +@pytest.mark.parametrize("invalid_root", ["", "missing-directory"]) +def test_invalid_discovered_roots_are_configuration_errors(monkeypatch, tmp_path, invalid_root): + from importlib.resources import files + + import aiconfigurator_core + from aiconfigurator_core.sdk import perf_database + + monkeypatch.setattr(perf_database, "_SYSTEMS_PATHS", [str(files("aiconfigurator_core") / "systems")]) + root = str(tmp_path / invalid_root) if invalid_root else "" + monkeypatch.setenv("AICONFIGURATOR_SYSTEMS_PATH", root) + for policy in ["deny", "allow"]: + with pytest.raises(ValueError, match="resolve systems paths"): + aiconfigurator_core.RustForwardPassPerfModel.best_available( + json.dumps( + { + "model": "Qwen/Qwen3-32B", + "system": "h200_sxm", + "backend": "vllm", + "worker_type": "aggregated", + "estimation_mode": "auto", + "fallback_policy": policy, + } + ) + ) + + +def test_compile_engine_classifies_controls_without_wrapping_model_failures(monkeypatch): + from aiconfigurator_core.sdk import engine, models + from aiconfigurator_core.sdk.errors import InvalidEngineConfigurationError + + with pytest.raises(InvalidEngineConfigurationError, match="Unknown forward_model"): + engine.compile_engine("missing-model", "h200_sxm", "vllm", forward_model="typo") + + failure = ValueError("fixture model lookup failed") + + def unavailable(*args, **kwargs): + raise failure + + monkeypatch.setattr(models, "_get_model_info", unavailable) + with pytest.raises(ValueError) as caught: + engine.compile_engine("missing-model", "h200_sxm", "vllm") + assert caught.value is failure + + +def unavailable_compiler(monkeypatch): + from aiconfigurator_core.sdk import engine + + calls = [] + + def unavailable(*args, **kwargs): + calls.append(kwargs["forward_model"]) + raise ValueError("fixture performance data is unavailable") + + monkeypatch.setattr(engine, "compile_engine", unavailable) + return calls + + +def test_auto_searches_priority_with_default_deny(monkeypatch): + calls = unavailable_compiler(monkeypatch) + model = RustForwardPassPerfModel.best_available(request()) + diagnostics = model.diagnostics() + assert calls == ["op_level", "fpm"] + assert diagnostics["source"] == "fallback_regression" + assert diagnostics["readiness"] != "ready" + provenance = diagnostics["provenance"] + assert provenance["requested_estimation_mode"] == "auto" + assert provenance["selected_estimation_mode"] == "fpm_regression" + assert provenance["config"]["fallback_policy"] == "deny" + assert len(provenance["selection_failures"]) == 2 + + +def test_explicit_deny_is_strict_and_allow_uses_remaining_priority(monkeypatch): + calls = unavailable_compiler(monkeypatch) + with pytest.raises(ValueError, match="fixture performance data"): + RustForwardPassPerfModel.best_available(request(estimation_mode="op_level")) + assert calls == ["op_level"] + calls.clear() + model = RustForwardPassPerfModel.best_available( + request(estimation_mode="fpm_interpolation", fallback_policy="allow") + ) + assert calls == ["fpm", "op_level"] + assert model.diagnostics()["provenance"]["selected_estimation_mode"] == "fpm_regression" + + +def test_explicit_regression_has_independent_sampling_and_does_not_compile(monkeypatch): + calls = unavailable_compiler(monkeypatch) + model = RustForwardPassPerfModel.best_available( + request( + estimation_mode="fpm_regression", + estimator_config={ + "fpm_regression": { + "sampling": {"bins_per_axis": [4, 16], "max_observations": 2}, + "min_observations": 2, + }, + "correction": { + "sampling": {"max_observations": 32}, + "min_observations": 8, + }, + }, + ) + ) + for count in (1, 2, 4, 8): + metrics = { + "version": 1, + "worker_id": "test", + "dp_rank": 0, + "counter_id": count, + "wall_time": 0.001 * count, + "scheduled_requests": { + "num_decode_requests": count, + "sum_decode_kv_tokens": count * 64, + }, + } + model.tune_with_fpms(metrics) + diagnostics = model.diagnostics() + assert calls == [] + assert diagnostics["retained_observations"] == 2 + resolved = diagnostics["provenance"]["config"]["estimator_config"] + assert resolved["fpm_regression"]["sampling"]["bins_per_axis"] == [4, 16] + assert resolved["correction"]["sampling"]["max_observations"] == 32 + assert resolved["correction"]["feature_space"] == "legacy_workload" + + +def test_invalid_config_never_falls_back_and_nested_errors_have_paths(monkeypatch): + calls = unavailable_compiler(monkeypatch) + with pytest.raises(ValueError, match=r"estimator_config\.fpm_regression\.fit"): + RustForwardPassPerfModel.best_available(request(estimator_config={"fpm_regression": {"fit": {"bogus": 1}}})) + with pytest.raises(ValueError, match="bins_per_axis"): + RustForwardPassPerfModel.best_available( + request(estimator_config={"fpm_regression": {"sampling": {"bins_per_axis": [0, 4]}}}) + ) + assert calls == [] + + +def test_legacy_selection_does_not_become_auto(): + import aiconfigurator_core + + legacy = { + "schema_version": 1, + "model_name": "Qwen/Qwen3-32B", + "system_name": "h200_sxm", + "backend": "vllm", + "tp_size": 2, + "pp_size": 1, + "forward_model": "fpm", + } + config = ForwardPassPerfModelConfig.from_legacy_engine_config(legacy, "decode") + assert config.estimation_mode == "fpm_interpolation" + assert config.fallback_policy == "deny" + assert config.worker_type == "decode" + normalized = json.loads(aiconfigurator_core.RustForwardPassPerfModel.normalize_config(json.dumps(config.to_dict()))) + assert normalized["tp"] == 2 + assert normalized["estimation_mode"] == "fpm_interpolation" + + +def test_invalid_quantization_does_not_degrade_to_another_estimator(monkeypatch): + from aiconfigurator_core.sdk import engine + + original = engine.compile_engine + calls = [] + + def recording(*args, **kwargs): + calls.append(kwargs["forward_model"]) + return original(*args, **kwargs) + + monkeypatch.setattr(engine, "compile_engine", recording) + with pytest.raises(ValueError, match="definitely_invalid_quant"): + RustForwardPassPerfModel.best_available(request(gemm_quant_mode="definitely_invalid_quant")) + assert calls == ["op_level"] + + +def test_resolver_uses_exact_topology_version_pins_and_isolates_cached_configs(monkeypatch): + import aiconfigurator_core + from aisimulate.sweeper.config import SearchSpace + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + + seen = [] + + class Model: + def __init__(self, config): + seen.append(config) + self.config = json.loads( + aiconfigurator_core.RustForwardPassPerfModel.normalize_config(json.dumps(config.to_dict())) + ) + self.config["backend_version"] = "0.24.0" if config.backend == "vllm" else "0.5.10" + self.config["estimation_mode"] = "op_level" + + def diagnostics(self): + return { + "readiness": "ready", + "provenance": {"config": self.config, "selected_systems_root": self.config["systems_paths"][0]}, + } + + def close(self): + pass + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", Model) + space = SearchSpace( + model_name="Qwen/Qwen3-32B", + hardware_sku="h200_sxm", + backend=["vllm", "sglang"], + backend_version={"vllm": "current"}, + ) + resolver = ForwardPassEstimatorResolver(space) + sample = { + "deployment_mode": "agg", + "model_name": space.model_name, + "hardware_sku": space.hardware_sku, + "backend": "vllm", + "tp": 2, + "pp": 1, + "attention_dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "agg_block_size": 64, + } + first = resolver.resolve_candidate(sample)["agg"] + first.config["estimator_config"]["correction"]["enabled"] = False + cached = resolver.resolve_candidate(sample)["agg"] + assert cached.config["estimator_config"]["correction"]["enabled"] is True + assert len(seen) == 1 + assert seen[0].worker_type == "aggregated" + assert (seen[0].tp, seen[0].kv_block_size, seen[0].backend_version) == (2, 64, "current") + resolver.resolve_candidate({**sample, "tp": 4, "backend": "sglang", "agg_block_size": 1}) + assert (seen[-1].tp, seen[-1].kv_block_size, seen[-1].backend_version) == (4, 1, None) + + +def test_search_rejects_unknown_controls_and_policies_on_custom_timing(): + from aisimulate.sweeper.config import SearchSpace + + common = {"model_name": "m", "hardware_sku": "h200_sxm", "deployment_mode": ["agg"]} + with pytest.raises(ValueError, match="estimation_mode"): + SearchSpace(**common, estimation_mode="typo") + with pytest.raises(ValueError, match="default timing"): + SearchSpace( + **common, database_mode="HYBRID", agg_timing_model={"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0} + ) + with pytest.raises(ValueError, match="unknown estimator override"): + SearchSpace(**common, role_estimator_controls={"agg": {"estimation_mod": "op_level"}}) + with pytest.raises(ValueError, match="unknown estimator role 'aggregated'; expected agg, prefill, or decode"): + SearchSpace(**common, role_estimator_controls={"aggregated": {"estimation_mode": "op_level"}}) + with pytest.raises(ValueError, match="estimator settings require default timing"): + SearchSpace( + **common, + agg_timing_model={"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}, + role_estimator_controls={"agg": {"estimation_mode": "fpm_regression"}}, + ) + + +def test_raw_config_rejects_trailing_json(): + import aiconfigurator_core + + payload = json.dumps(request(estimation_mode="fpm_regression").to_dict()) + " {}" + with pytest.raises(ValueError, match="trailing characters"): + aiconfigurator_core.RustForwardPassPerfModel.normalize_config(payload) + + +@pytest.mark.parametrize("deployment_mode", ["afd", "afd+pd"]) +@pytest.mark.parametrize( + "controls", + [ + {"estimation_mode": "fpm_regression"}, + {"estimation_mode": "fpm_interpolation"}, + {"fallback_policy": "allow"}, + {"estimator_config": {"correction": {"enabled": False}}}, + {"systems_paths": ["default"]}, + {"database_mode": "HYBRID"}, + {"transfer_policy": "conservative"}, + ], +) +def test_afd_rejects_role_estimator_controls(deployment_mode, controls): + from aisimulate.sweeper.config import SearchSpace + + with pytest.raises(ValueError, match="role_estimator_controls require regular language workers"): + SearchSpace( + model_name="m", + hardware_sku="h200_sxm", + deployment_mode=[deployment_mode], + afd_batch_size_candidates=[1], + role_estimator_controls={"decode": controls}, + ) + + +@pytest.mark.parametrize("deployment_mode", ["afd", "afd+pd"]) +@pytest.mark.parametrize("forward_model", ["op_level", "fpm"]) +def test_afd_saved_search_preserves_legacy_selection(deployment_mode, forward_model): + from aisimulate.sweeper.config import SearchSpace + + space = SearchSpace( + model_name="m", + hardware_sku="h200_sxm", + deployment_mode=[deployment_mode], + afd_batch_size_candidates=[1], + prefill_forward_model=forward_model, + ) + for _ in range(2): + assert space.prefill_forward_model == forward_model + assert space.role_estimator_controls == {} + space = SearchSpace.model_validate_json(space.model_dump_json()) + + +@pytest.mark.parametrize( + "selection", + [{}, {"agg_forward_model": "fpm"}, {"estimation_mode": "fpm_regression"}], +) +def test_saved_search_preserves_estimator_selection(selection): + from aisimulate.sweeper.config import SearchSpace + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + + space = SearchSpace(model_name="Qwen/Qwen3-32B", hardware_sku="h200_sxm", **selection) + sample = { + "backend": "vllm", + "hardware_sku": "h200_sxm", + "tp": 2, + "pp": 1, + "attention_dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "agg_block_size": 64, + } + expected = "fpm_interpolation" if "agg_forward_model" in selection else selection.get("estimation_mode", "auto") + for _ in range(2): + resolved = ForwardPassEstimatorResolver(space)._request(sample, "agg") + assert resolved.estimation_mode == expected + assert resolved.fallback_policy == "deny" + serialized = space.model_dump_json() + assert "agg_forward_model" not in json.loads(serialized) + space = SearchSpace.model_validate_json(serialized) + assert space.agg_forward_model == ("fpm" if expected == "fpm_interpolation" else "op_level") + + +def test_mixed_timing_still_enforces_cold_regression_on_default_role(): + from aisimulate.config.cli import CoreRecommendationConfig + from aisimulate.recommend import recommendation_to_sweeper + from aisimulate.sweeper.forward_pass_estimator import ( + ForwardPassEstimatorResolutionError, + ForwardPassEstimatorResolver, + ) + + public = CoreRecommendationConfig.model_validate( + { + "engine": { + "model": "Qwen/Qwen3-32B", + "hardware": "h200_sxm", + "backend": "vllm", + "mode": "disaggregated", + "context_length": 4096, + "workers": { + "prefill": { + "timing": { + "estimation_mode": "fpm_regression", + "fallback_policy": "deny", + } + }, + "decode": {"timing": {"type": "fixed", "prefill_ms": 1, "decode_ms": 1}}, + }, + }, + "optimization": { + "target": "throughput", + "constraints": {"max_candidate_gpus": 4}, + }, + } + ) + space = recommendation_to_sweeper(public).search_space + sample = { + "deployment_mode": "disagg", + "backend": "vllm", + "hardware_sku": "h200_sxm", + "prefill_tp": 2, + "prefill_pp": 1, + "prefill_attention_dp": 1, + "prefill_moe_tp": 1, + "prefill_moe_ep": 1, + "prefill_block_size": 64, + "prefill_timing_model": space.prefill_timing_model, + "decode_timing_model": space.decode_timing_model, + } + with pytest.raises(ForwardPassEstimatorResolutionError, match="prefill is not ready"): + ForwardPassEstimatorResolver(space).resolve_candidate(sample) diff --git a/tests/sweeper/test_kv_load.py b/tests/sweeper/test_kv_load.py index 9a52ad854..a1ec18015 100644 --- a/tests/sweeper/test_kv_load.py +++ b/tests/sweeper/test_kv_load.py @@ -14,6 +14,83 @@ ) +@pytest.mark.parametrize("source", ["resolved", "custom", "legacy_custom"]) +def test_capacity_cache_includes_resolved_root(tmp_path, monkeypatch, source): + from aisimulate.sweeper import kv_load + + roots = [tmp_path / "first", tmp_path / "second"] + for root in roots: + root.mkdir() + calls = [] + + def capacity(*args, systems_paths, **kwargs): + calls.append(systems_paths) + return 6400 if systems_paths == [str(roots[0])] else 12800 + + monkeypatch.setattr(kv_load, "estimate_kv_tokens", capacity) + kv_load._per_rank_capacity_tokens.cache_clear() + sample = _sample("agg") + parallel = ReplicaParallelConfig(ParallelShape(tp=1, dp=1, moe_tp=1, moe_ep=1), replicas=1) + capacities = [] + try: + for root in [roots[0], roots[1], roots[0]]: + if source == "resolved": + sample["forward_pass_estimators"] = {"agg": {"config": {"systems_paths": [str(root)]}}} + sample["agg_timing_model"] = { + "type": "external", + "provider": "aic", + "config": {"systems_paths": ["must-not-be-used"]}, + } + else: + config = {"systems_paths": [str(root)]} if source == "custom" else {"systems_path": str(root)} + sample["agg_timing_model"] = {"type": "external", "provider": "aic", "config": config} + result = resolve_kv_load( + sample, + workload=Workload(isl=100, osl=100, kv_load_ratio=1.0, request_count=1), + parallel_config=parallel, + ratio=1.0, + backend_version="v", + ) + capacities.append(result.role_capacity_tokens["agg"]) + assert capacities == [6400, 12800, 6400] + assert calls == [[str(roots[0])], [str(roots[1])]] + finally: + kv_load._per_rank_capacity_tokens.cache_clear() + + +@pytest.mark.parametrize("custom_role", ["prefill", "decode"]) +def test_mixed_timing_capacity_uses_each_roles_own_root(tmp_path, monkeypatch, custom_role): + from aisimulate.sweeper import kv_load + + sample = _sample("disagg") + roots = {role: tmp_path / role for role in ("prefill", "decode")} + for root in roots.values(): + root.mkdir() + canonical_role = "prefill" if custom_role == "decode" else "decode" + sample["forward_pass_estimators"] = {canonical_role: {"config": {"systems_paths": [str(roots[canonical_role])]}}} + sample[f"{custom_role}_timing_model"] = { + "type": "external", + "provider": "aic", + "config": {"systems_paths": [str(roots[custom_role])]}, + } + seen = [] + + def capacity(*args, systems_paths, **kwargs): + seen.append(systems_paths) + return 6400 + + monkeypatch.setattr(kv_load, "_per_rank_capacity_tokens", capacity) + parallel = ReplicaParallelConfig(ParallelShape(tp=1, dp=1, moe_tp=1, moe_ep=1), replicas=1) + resolve_kv_load( + sample, + workload=Workload(isl=100, osl=100, kv_load_ratio=1.0, request_count=1), + parallel_config=DisaggParallelConfig(prefill=parallel, decode=parallel), + ratio=1.0, + backend_version="v", + ) + assert seen == [(str(roots["prefill"]),), (str(roots["decode"]),)] + + def _sample(mode: str) -> dict: sample = { "deployment_mode": mode, diff --git a/tests/sweeper/test_result.py b/tests/sweeper/test_result.py index e39d6dcae..8ac3c1468 100644 --- a/tests/sweeper/test_result.py +++ b/tests/sweeper/test_result.py @@ -756,6 +756,13 @@ def test_optimizer_guided_result_separates_unsupported_and_runtime_failure(monke assert result.selected_candidates == [] +@pytest.fixture(autouse=True) +def _isolate_estimator_data_for_result_orchestration(monkeypatch): + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + + monkeypatch.setattr(ForwardPassEstimatorResolver, "resolve_candidate", lambda self, sample: {}) + + def test_resource_limited_records_roundtrip_without_claiming_evaluation(): payload = _complete_result().model_dump(mode="json") payload["candidates"].append( diff --git a/tests/sweeper/test_search.py b/tests/sweeper/test_search.py index 5713bf458..ed5acb72b 100644 --- a/tests/sweeper/test_search.py +++ b/tests/sweeper/test_search.py @@ -10,7 +10,11 @@ import aisimulate.sweeper.search as search_mod from aisimulate.sweeper.config import OptimizationGoal, SLATarget, SmartSearchConfig from aisimulate.sweeper.kv_load import KVLoadResolution -from aisimulate.sweeper.parallel_enum import DisaggParallelConfig, ParallelShape, ReplicaParallelConfig +from aisimulate.sweeper.parallel_enum import ( + DisaggParallelConfig, + ParallelShape, + ReplicaParallelConfig, +) from aisimulate.sweeper.replay import ( BackendDeploymentSpec, ReplayReport, @@ -1154,3 +1158,12 @@ def factory(branch, study_id, objectives=None): ("throughput_per_user", True), ] assert all(set(metrics) == {"throughput_per_gpu", "throughput_per_user"} for metrics in seen["sampler"].observed) + + +@pytest.fixture(autouse=True) +def _isolate_estimator_data_for_orchestration(monkeypatch): + # These tests use synthetic models/runners. Native construction is exercised + # by the estimator contract tests and CLI round trips. + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + + monkeypatch.setattr(ForwardPassEstimatorResolver, "resolve_candidate", lambda self, sample: {}) diff --git a/tests/sweeper/test_search_providers.py b/tests/sweeper/test_search_providers.py index e06f19d26..c520da847 100644 --- a/tests/sweeper/test_search_providers.py +++ b/tests/sweeper/test_search_providers.py @@ -9,7 +9,11 @@ import aisimulate.sweeper.search as search_module from aisimulate.sweeper.config import SmartSearchConfig -from aisimulate.sweeper.parallel_enum import DisaggParallelConfig, ParallelShape, ReplicaParallelConfig +from aisimulate.sweeper.parallel_enum import ( + DisaggParallelConfig, + ParallelShape, + ReplicaParallelConfig, +) from aisimulate.sweeper.provider import ( AdapterReplaySpec, AdapterSearchPlan, @@ -753,3 +757,12 @@ def test_adapter_parameter_separator_collisions_are_rejected() -> None: injected={"ambiguous::adapter": _Adapter()}, show_progress=False, ) + + +@pytest.fixture(autouse=True) +def _isolate_estimator_data_for_orchestration(monkeypatch): + # These tests use synthetic models/runners. Native construction is exercised + # by the estimator contract tests and CLI round trips. + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + + monkeypatch.setattr(ForwardPassEstimatorResolver, "resolve_candidate", lambda self, sample: {}) diff --git a/tests/sweeper/test_search_space.py b/tests/sweeper/test_search_space.py index 6667c5b9a..a9310e618 100644 --- a/tests/sweeper/test_search_space.py +++ b/tests/sweeper/test_search_space.py @@ -236,6 +236,7 @@ def fake_parallel_configs( min_gpu_budget=None, max_seq_len=None, role_runtime=None, + systems_paths=None, ): calls.append((deployment_mode, backend)) return [_DISAGG_DP1_CFG] @@ -271,12 +272,13 @@ def fake_parallel_configs( min_gpu_budget=None, max_seq_len=None, role_runtime=None, + systems_paths=None, ): calls.append((hardware, deployment_mode, backend_version, role_runtime)) return [_AGG_CFG if hardware == "h200_sxm" else _DP8_CFG] monkeypatch.setattr("aisimulate.sweeper.search_space.parallel_configs_for", fake_parallel_configs) - monkeypatch.setattr("aisimulate.sweeper.search_space.resolve_backend_version", lambda *args: "1.0") + monkeypatch.setattr("aisimulate.sweeper.search_space.resolve_backend_version", lambda *args, **kwargs: "1.0") config = _config( deployment_mode=["disagg"], hardware_sku="gb200", @@ -297,7 +299,7 @@ def fake_parallel_configs( def test_heterogeneous_disagg_requires_one_common_implicit_backend_version(monkeypatch): - def role_version(hardware, backend): + def role_version(hardware, backend, **kwargs): del backend return {"h200_sxm": "prefill-version", "gb200": "decode-version"}[hardware] @@ -698,6 +700,7 @@ def fake_parallel_configs( min_gpu_budget=None, max_seq_len=None, role_runtime=None, + systems_paths=None, ): if deployment_mode == "disagg": raise NoViableParallelConfig("disagg does not fit") @@ -745,6 +748,7 @@ def fake_parallel_configs( min_gpu_budget=None, max_seq_len=None, role_runtime=None, + systems_paths=None, ): if backend == "vllm": raise NoPerfDatabase("no vLLM perf database") diff --git a/tests/sweeper/test_unified_optimizer.py b/tests/sweeper/test_unified_optimizer.py index 086e94094..a61339539 100644 --- a/tests/sweeper/test_unified_optimizer.py +++ b/tests/sweeper/test_unified_optimizer.py @@ -6,6 +6,8 @@ import time from typing import ClassVar +import pytest + import aisimulate.sweeper.search as search_module from aisimulate.sweeper.config import SmartSearchConfig from aisimulate.sweeper.parallel_enum import ( @@ -384,3 +386,12 @@ def test_candidate_timeout_applies_with_parallelism_one(monkeypatch) -> None: assert result.selected_candidates == [] assert _CountingSampler.suggested == 1 + + +@pytest.fixture(autouse=True) +def _isolate_estimator_data_for_orchestration(monkeypatch): + # These tests use synthetic models/runners. Native construction is exercised + # by the estimator contract tests and CLI round trips. + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + + monkeypatch.setattr(ForwardPassEstimatorResolver, "resolve_candidate", lambda self, sample: {}) diff --git a/tests/test_afd_cli.py b/tests/test_afd_cli.py index 952cb2637..4a98a6094 100644 --- a/tests/test_afd_cli.py +++ b/tests/test_afd_cli.py @@ -165,6 +165,35 @@ def estimator(model, hardware, **kwargs): assert report.metadata["afd_replay"]["companion"]["forward_model"] == expected_model +@pytest.mark.parametrize("mode", ["auto", "fpm_regression"]) +def test_afd_legacy_fpm_does_not_bypass_estimator_policy_validation(mode): + raw = _pure_prediction() + raw["engine"]["afd"].update(phase="decode", combined_with_pd=True) + raw["engine"]["workers"] = { + "prefill": {"parallelism": {"tensor": 2}, "timing": {"forward_model": "fpm", "estimation_mode": mode}} + } + with pytest.raises(ValueError, match="estimator policies"): + CorePredictionConfig.model_validate(raw) + + +@pytest.mark.parametrize("forward_model", ["op_level", "fpm"]) +def test_afd_recommendation_preserves_legacy_companion_selection(forward_model): + from aisimulate.recommend import recommendation_to_sweeper + from aisimulate.sweeper.config import SmartSearchConfig + + raw = _pure_prediction() + raw["engine"]["afd"].update(phase="decode", combined_with_pd=True) + raw["engine"]["afd"].pop("n_a_nodes") + raw["engine"]["afd"].pop("n_f_nodes") + raw["engine"]["workers"] = {"prefill": {"timing": {"forward_model": forward_model}}} + config = CoreRecommendationConfig.model_validate({**raw, "optimization": {}}) + search = recommendation_to_sweeper(config, stack="engine") + for _ in range(2): + assert search.search_space.prefill_forward_model == forward_model + assert search.search_space.role_estimator_controls == {} + search = SmartSearchConfig.model_validate_json(search.model_dump_json()) + + @pytest.mark.parametrize( ("afd", "match"), [ diff --git a/tests/test_ci_workflow_contracts.py b/tests/test_ci_workflow_contracts.py index 58940da94..3a3009342 100644 --- a/tests/test_ci_workflow_contracts.py +++ b/tests/test_ci_workflow_contracts.py @@ -39,6 +39,68 @@ ACTION_ROOT = REPOSITORY_ROOT / ".github" / "actions" +def test_release_migration_gate_blocks_publication_until_reviewed_clearance(tmp_path): + from scripts.check_release_migrations import GATES, require_completed_migrations + + with pytest.raises(RuntimeError, match="dynamo/pull/14065"): + require_completed_migrations(GATES) + path = tmp_path / "gates.json" + path.write_text(json.dumps({"pending_migrations": []})) + require_completed_migrations(path) + for invalid in ({}, {"pending_migrations": None}, {"pending_migrations": [{}]}): + path.write_text(json.dumps(invalid)) + with pytest.raises(ValueError): + require_completed_migrations(path) + path.unlink() + with pytest.raises(FileNotFoundError): + require_completed_migrations(path) + + jobs = _workflow("nightly-ci.yml")["jobs"] + guard = jobs["changes-guard"] + index = next(i for i, step in enumerate(guard["steps"]) if "check_release_migrations.py" in step.get("run", "")) + assert next(i for i, step in enumerate(guard["steps"]) if step.get("id") == "target") < index + assert index < next(i for i, step in enumerate(guard["steps"]) if step.get("id") == "decide") + assert "if" not in guard["steps"][index] + checkouts = [step for step in guard["steps"][:index] if step.get("uses", "").startswith("actions/checkout@")] + assert len(checkouts) == 2 + assert checkouts[0]["with"]["ref"] == "${{ steps.target.outputs.sha }}" + assert "path" not in checkouts[0]["with"] + assert checkouts[1]["with"]["ref"] == "${{ github.sha }}" + assert checkouts[1]["with"]["path"] == "release-policy" + assert all("if" not in step and step["with"]["persist-credentials"] == "false" for step in checkouts) + assert guard["steps"][index]["run"] == ( + "python3 release-policy/scripts/check_release_migrations.py --target-gates .github/release-gates.json" + ) + assert "build-artifacts" in jobs["trigger-gitlab-security"]["needs"] + assert "changes-guard" in jobs["build-artifacts"]["needs"] + assert "needs.changes-guard.outputs.should-build == 'true'" in jobs["build-artifacts"]["if"] + + +@pytest.mark.parametrize( + "current,target,expected", + [ + ("clear", "clear", 0), + ("clear", "pending", 1), + ("pending", "clear", 1), + ("clear", "missing", 1), + ("clear", "malformed", 1), + ], +) +def test_publication_checks_current_policy_and_selected_target(tmp_path, monkeypatch, current, target, expected): + from scripts import check_release_migrations as checker + + paths = {} + for name, state in (("current", current), ("target", target)): + paths[name] = tmp_path / f"{name}.json" + if state == "missing": + continue + pending = [] if state == "clear" else [{"pull_request": "migration/pr/1", "requirement": "Migrate consumer"}] + document = {} if state == "malformed" else {"pending_migrations": pending} + paths[name].write_text(json.dumps(document)) + monkeypatch.setattr(checker, "GATES", paths["current"]) + assert checker.main(["--target-gates", str(paths["target"])]) == expected + + def test_forward_perf_selects_before_allocating_the_benchmark_runner(): workflow = _workflow("performance.yml") assert workflow["on"]["push"] == {"branches": ["pull-request/*"]} diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index 6a24bc85a..66e3b37be 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -130,7 +130,10 @@ def test_prediction_rejects_wrong_backend_scheduler_interval(backend, field, val engine["workers"] = {"prefill": {}, "decode": {}} engine["workers"][role] = {"scheduler": {field: value}} - with pytest.raises(ValidationError, match=rf"workers.{role}.scheduler.{field}.*backend={required_backend}"): + with pytest.raises( + ValidationError, + match=rf"workers.{role}.scheduler.{field}.*backend={required_backend}", + ): CorePredictionConfig.model_validate({"engine": engine}) @@ -407,6 +410,7 @@ def test_engine_scheduler_domains_replace_defaults_and_preserve_log_scale() -> N **_engine(), "mode": "aggregated", "backend_version": "0.19.0", + "backend": "vllm", "context_length": 4096, "workers": { "aggregated": { @@ -820,7 +824,9 @@ def test_prediction_timing_accepts_fpm_forward_model_with_default_timing() -> No {"type": "polynomial", "forward_model": "fpm"}, ], ) -def test_prediction_timing_rejects_fpm_forward_model_without_default_timing(timing: dict) -> None: +def test_prediction_timing_rejects_fpm_forward_model_without_default_timing( + timing: dict, +) -> None: engine = _engine() engine["workers"]["aggregated"] = {"timing": timing} @@ -870,7 +876,12 @@ def _fpm_recommendation() -> CoreRecommendationConfig: "context_length": 2048, "workers": { "aggregated": { - "parallelism": {"preset": False, "tensor": 1, "moe_tensor": 1, "moe_expert": 1}, + "parallelism": { + "preset": False, + "tensor": 1, + "moe_tensor": 1, + "moe_expert": 1, + }, "kv_cache": {"capacity": {"type": "fixed", "blocks": 256}}, "timing": {"type": "default", "forward_model": "fpm"}, } @@ -932,10 +943,67 @@ def test_recommendation_candidate_yaml_round_trips_forward_model() -> None: adapter_sections={}, ) - assert prediction["engine"]["workers"]["aggregated"]["timing"] == {"type": "default", "forward_model": "fpm"} + assert prediction["engine"]["workers"]["aggregated"]["timing"]["estimation_mode"] == "fpm_interpolation" + assert prediction["engine"]["workers"]["aggregated"]["timing"]["fallback_policy"] == "deny" CorePredictionConfig.model_validate(prediction) +@pytest.mark.parametrize("policy", [None, []]) +def test_candidate_preserves_default_vs_disabled_transfers_and_pinned_capacity(policy): + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig + from aisimulate.compiler import prediction_to_replay_spec + from aisimulate.sweeper.replay import ForwardPassEstimatorSpec + + config = _fpm_recommendation() + smart = recommendation_to_sweeper(config) + sample = unroll_sample( + search_space=smart.search_space, + selection={ + "deployment_mode": "agg", + "backend": "vllm", + "agg_max_num_batched_tokens": 8192, + "agg_max_num_seqs": 256, + }, + parallel_config=ReplicaParallelConfig(ParallelShape(tp=1, dp=1, moe_tp=1, moe_ep=1), replicas=1), + ) + resolved = ForwardPassPerfModelConfig( + model="example/model", + system="h200_sxm", + backend="vllm", + worker_type="aggregated", + backend_version="test", + estimation_mode="op_level", + transfer_policy=policy, + ).to_dict() + deployment = build_backend_deployment( + sample, backend_version="test", forward_pass_estimators={"agg": ForwardPassEstimatorSpec(config=resolved)} + ) + assert "gpu_memory_utilization" not in deployment.agg_engine_args["timing_model"]["config"] + prediction = _candidate_prediction( + config, sample, ReplaySpec(backend_deployment=deployment, workload={}, goal={}), adapter_sections={} + ) + replay = prediction_to_replay_spec(CorePredictionConfig.model_validate(prediction)) + timing = replay.backend_deployment.agg_engine_args["timing_model"]["config"] + assert timing["transfer_policy"] == policy + assert "gpu_memory_utilization" not in timing + + +def test_public_estimator_config_rejects_sol_full(): + from aisimulate.config.engine import TimingConfig + from aisimulate.sweeper.config import SearchSpace + + for make in ( + lambda: CorePredictionConfig.model_validate({"engine": {**_engine(), "database_mode": "SOL_FULL"}}), + lambda: SearchSpace(model_name="m", hardware_sku="h200_sxm", database_mode="SOL_FULL"), + lambda: SearchSpace( + model_name="m", hardware_sku="h200_sxm", role_estimator_controls={"agg": {"database_mode": "SOL_FULL"}} + ), + lambda: TimingConfig(database_mode="sol_full"), + ): + with pytest.raises(ValueError, match="database_mode"): + make() + + def test_recommendation_candidate_yaml_spells_out_op_level_like_other_defaults() -> None: config = _fpm_recommendation() raw = config.model_dump(mode="python", exclude_none=True) @@ -962,7 +1030,8 @@ def test_recommendation_candidate_yaml_spells_out_op_level_like_other_defaults() ) # Normalization materializes every schema default into the candidate; forward_model is no exception. - assert prediction["engine"]["workers"]["aggregated"]["timing"] == {"type": "default", "forward_model": "op_level"} + assert prediction["engine"]["workers"]["aggregated"]["timing"]["estimation_mode"] == "op_level" + assert prediction["engine"]["workers"]["aggregated"]["timing"]["fallback_policy"] == "deny" def _pd_hardware_config(*, recommend=False, **overrides): @@ -986,7 +1055,16 @@ def _pd_hardware_config(*, recommend=False, **overrides): @pytest.mark.parametrize("role", ["prefill", "decode"]) @pytest.mark.parametrize( "value", - ["", " ", "auto", " auto ", " gb200 ", "gb200 ", "\th200_sxm", {"choices": ["h200_sxm", "gb200"]}], + [ + "", + " ", + "auto", + " auto ", + " gb200 ", + "gb200 ", + "\th200_sxm", + {"choices": ["h200_sxm", "gb200"]}, + ], ) def test_worker_hardware_requires_concrete_sku(recommend, role, value): schema = CoreRecommendationConfig if recommend else CorePredictionConfig @@ -1038,14 +1116,17 @@ def test_pd_hardware_survives_search_candidate_yaml_and_predict(overrides, expec sample["backend_version"] = "0.24.0" deployment = build_backend_deployment(sample, backend_version="0.24.0") mapping = _candidate_prediction( - source, sample, ReplaySpec(backend_deployment=deployment, workload={}, goal={}), adapter_sections={} + source, + sample, + ReplaySpec(backend_deployment=deployment, workload={}, goal={}), + adapter_sections={}, ) concrete = CorePredictionConfig.model_validate(yaml.safe_load(yaml.safe_dump(mapping))) compiled = prediction_to_replay_spec(concrete).backend_deployment assert concrete.engine.hardware == "h200_sxm" for role, hardware in zip(("prefill", "decode"), expected, strict=True): assert ("hardware" in mapping["engine"]["workers"][role]) == (role in overrides) - assert getattr(compiled, f"{role}_engine_args")["aic_system"] == hardware + assert getattr(compiled, f"{role}_engine_args")["timing_model"]["config"]["system"] == hardware assert compiled.performance_model_metadata[role]["config"]["system"] == hardware assert getattr(deployment, f"{role}_engine_args")["aic_system"] == hardware @@ -1068,7 +1149,7 @@ def test_pd_predict_requires_shared_implicit_backend_version(monkeypatch, same_v calls = [] - def resolve(hardware, backend): + def resolve(hardware, backend, **kwargs): calls.append((hardware, backend)) return "0.24.0" if same_version or hardware == "h200_sxm" else "0.23.0" @@ -1080,7 +1161,7 @@ def resolve(hardware, backend): deployment = prediction_to_replay_spec(config).backend_deployment assert deployment.backend_version == "0.24.0" for role in ("prefill", "decode"): - assert getattr(deployment, f"{role}_engine_args")["aic_backend_version"] == "0.24.0" + assert getattr(deployment, f"{role}_engine_args")["timing_model"]["config"]["backend_version"] == "0.24.0" else: with pytest.raises(ValueError, match="Set engine.backend_version"): prediction_to_replay_spec(config) @@ -1113,14 +1194,14 @@ def test_pd_predict_accepts_worker_hardware_from_configured_system_paths(monkeyp raw = _pd_hardware_config(**{role: "custom_worker_sku"}) config = CorePredictionConfig.model_validate(raw) deployment = prediction_to_replay_spec(config).backend_deployment - assert getattr(deployment, f"{role}_engine_args")["aic_system"] == "custom_worker_sku" + assert getattr(deployment, f"{role}_engine_args")["timing_model"]["config"]["system"] == "custom_worker_sku" @pytest.mark.parametrize("override", [False, True]) def test_pd_predict_keeps_legacy_version_defaults_without_hardware_override(monkeypatch, override): from aisimulate.compiler import prediction_to_replay_spec - def resolve(hardware, backend): + def resolve(hardware, backend, **kwargs): assert hardware == "gb200" return "0.24.0" @@ -1155,3 +1236,71 @@ def test_pd_predict_checks_effective_prefill_hardware_in_router_hook(router_hard else: with pytest.raises(ValueError, match="does not match effective prefill_hardware_sku"): prediction_to_replay_spec(config, adapter_specs={"dynamo.router": spec}) + + +def test_new_default_selection_survives_serialization_without_becoming_legacy_op_level(): + config = CorePredictionConfig.model_validate({"engine": _engine()}) + serialized = config.model_dump(mode="json", exclude_none=True) + timing = serialized["engine"]["workers"]["aggregated"]["timing"] + assert "forward_model" not in timing + reloaded = CorePredictionConfig.model_validate(serialized) + assert reloaded.engine.estimation_mode == "auto" + assert reloaded.engine.workers.aggregated.timing.estimation_mode is None + + +@pytest.mark.parametrize("mode", ["aggregated", "disaggregated"]) +def test_role_systems_roots_reach_prediction_and_search_preflight(tmp_path, mode): + from copy import deepcopy + from importlib.resources import files + from pathlib import Path + + from aisimulate.compiler import prediction_to_replay_spec + from aisimulate.sweeper.search_space import enumerate_branches + + packaged = Path(str(files("aiconfigurator_core") / "systems")) + roles = ("aggregated",) if mode == "aggregated" else ("prefill", "decode") + workers = {} + for role in roles: + root = tmp_path / role + root.mkdir() + for entry in packaged.iterdir(): + (root / entry.name).symlink_to(entry, target_is_directory=entry.is_dir()) + (root / "role_only_gpu.yaml").write_text((packaged / "h200_sxm.yaml").read_text()) + workers[role] = { + "parallelism": {"tensor": 2}, + "kv_cache": {"capacity": {"type": "fixed", "blocks": 256}}, + "timing": {"systems_paths": [str(root)]}, + } + if mode == "disaggregated": + workers[role]["hardware"] = "role_only_gpu" + engine = { + "model": "Qwen/Qwen3-32B", + "hardware": "role_only_gpu" if mode == "aggregated" else "h200_sxm", + "backend": "vllm", + "mode": mode, + "context_length": 4096, + "workers": workers, + } + prediction = CorePredictionConfig.model_validate({"engine": engine}) + deployment = prediction_to_replay_spec(prediction).backend_deployment + for role in roles: + args = getattr(deployment, f"{'agg' if role == 'aggregated' else role}_engine_args") + assert args["timing_model"]["config"]["systems_paths"] == [str(tmp_path / role)] + assert args["timing_model"]["config"]["system"] == "role_only_gpu" + recommendation_engine = deepcopy(engine) + for worker in recommendation_engine["workers"].values(): + worker.pop("parallelism") + worker["kv_cache"] = {"capacity": {"memory_fraction": 0.9}} + recommendation = CoreRecommendationConfig.model_validate( + { + "engine": recommendation_engine, + "optimization": { + "target": "throughput", + "constraints": {"max_candidate_gpus": 4}, + }, + } + ) + smart = recommendation_to_sweeper(recommendation) + (branch,) = enumerate_branches(smart, max_seq_len=4096) + assert branch.parallel_configs + assert branch.deployment_mode == ("agg" if mode == "aggregated" else "disagg") diff --git a/tests/test_epd_cli.py b/tests/test_epd_cli.py index 9b13666a9..3f6fda9aa 100644 --- a/tests/test_epd_cli.py +++ b/tests/test_epd_cli.py @@ -170,6 +170,13 @@ def test_epd_public_schema_rejects_unsupported(kind, recommend): (CoreRecommendationConfig if recommend else CorePredictionConfig).model_validate(raw) +def test_encoder_legacy_fpm_does_not_bypass_estimator_policy_validation(): + raw = _prediction() + raw["engine"]["workers"]["aggregated"]["timing"] = {"forward_model": "fpm", "estimation_mode": "fpm_regression"} + with pytest.raises(ValueError, match="estimator policies"): + CorePredictionConfig.model_validate(raw) + + @pytest.mark.parametrize( "field,value", [ @@ -388,7 +395,7 @@ def test_epd_native_search_preserves_available_backends(monkeypatch, caplog, bac def database(system, backend, version, **kwargs): return None if backend == "vllm" else original_database(system, backend, version, **kwargs) - def latest_version(system, backend): + def latest_version(system, backend, **kwargs): return None if backend == "vllm" else original_version(system, backend) raw = _recommendation() diff --git a/tests/test_ngram_cli.py b/tests/test_ngram_cli.py index a161b83f5..773a082e8 100644 --- a/tests/test_ngram_cli.py +++ b/tests/test_ngram_cli.py @@ -155,6 +155,47 @@ def capture(*args, **kwargs): assert calls and all(call["speculation"] == _COST and call["nextn"] == 0 for call in calls) +def test_canonical_estimator_round_trip_preserves_prompt_lookup_cost(): + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + config = ForwardPassPerfModelConfig( + model="meta-llama/Meta-Llama-3.1-8B", + system="h200_sxm", + backend="vllm", + backend_version="0.24.0", + worker_type="aggregated", + speculation=deepcopy(_COST), + ) + for _ in range(2): + model = RustForwardPassPerfModel.best_available(config) + try: + provenance = model.diagnostics()["provenance"] + assert provenance["selected_estimation_mode"] == "op_level" + assert provenance["config"]["speculation"] == _COST + assert provenance["config"]["nextn"] == 0 + config = ForwardPassPerfModelConfig(**provenance["config"]) + finally: + model.close() + + +@pytest.mark.parametrize("mode", ["fpm_interpolation", "fpm_regression"]) +def test_canonical_prompt_lookup_rejects_unsupported_estimators(mode): + from aiconfigurator_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + with pytest.raises(ValueError, match="ngram speculation requires op_level timing"): + RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig( + model="m", + system="s", + backend="vllm", + worker_type="aggregated", + speculation=deepcopy(_COST), + estimation_mode=mode, + fallback_policy="allow", + ) + ) + + def test_recommend_predict_round_trip_preserves_speculation(tmp_path): raw = yaml.safe_load( (_ROOT / "tests/e2e/configs/unified_cli/recommend/engine/02-custom-preset-throughput-per-gpu.yaml").read_text() diff --git a/tests/test_runner.py b/tests/test_runner.py index c6f4793be..88a48be86 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -261,6 +261,51 @@ def test_runner_preserves_weka_lane_input_without_defaulting_source_block_size() assert "trace_block_size" not in traffic +@pytest.mark.parametrize("model_key", ["model", "model_path"]) +def test_agentic_runner_reads_canonical_model_identity(model_key): + runtime = RecordingRuntime() + engine_args = _engine_args() + engine_args.pop("aic_model_path") + deployment = BackendDeploymentSpec( + deployment_mode="agg", + backend="vllm", + backend_version="test", + agg_engine_args=engine_args, + num_workers=1, + performance_model_metadata={"aggregated": {"provider": "aic", "config": {model_key: "test-model"}}}, + ) + EngineReplayRunnerFactory(runtime=runtime).create(0).run( + _spec( + deployment=deployment, + workload={ + "source_type": "trace", + "load_type": "trace_timestamps", + "trace_path": "agentic.jsonl", + "trace_format": "agentic_mooncake", + }, + ) + ) + assert runtime.execution_spec["traffic"]["execution_model"] == "test-model" + + +@pytest.mark.parametrize("nested_rank", [False, True]) +@pytest.mark.parametrize("payload", [{}, {"config": None}, {"config": []}, {"config": "invalid"}]) +def test_runner_rejects_nonmapping_aic_timing_before_materialization(nested_rank, payload, monkeypatch): + monkeypatch.setattr(aic, "estimate_num_gpu_blocks", lambda **kwargs: pytest.fail("invalid timing reached capacity")) + runtime = RecordingRuntime() + args = _engine_args(timing={"type": "external", "provider": "aic", **payload}) + if nested_rank: + args = {"rank": args} + spec = _spec( + deployment=BackendDeploymentSpec( + deployment_mode="agg", backend="vllm", backend_version="test", agg_engine_args=args, num_workers=1 + ) + ) + with pytest.raises(ValueError, match="external AIC timing config must be a mapping"): + EngineReplayRunnerFactory(runtime=runtime).create(0).run(spec) + assert runtime.execution_spec is None + + def test_runner_rejects_agentic_execution_without_a_target_model(): engine_args = _engine_args() engine_args.pop("aic_model_path") @@ -531,7 +576,8 @@ def test_public_host_offload_config_reaches_native_execution_rank(): } -def test_public_cuda_graph_reservation_reaches_native_capacity(tmp_path, monkeypatch): +@pytest.mark.parametrize("rank_only_controls", [False, True]) +def test_public_cuda_graph_reservation_reaches_native_capacity(tmp_path, monkeypatch, rank_only_controls): reserved_bytes = 14_559_947_612 path = tmp_path / "prediction.yaml" path.write_text( @@ -560,8 +606,27 @@ def estimate(**kwargs): return 321 monkeypatch.setattr(aic, "estimate_num_gpu_blocks", estimate) + import aiconfigurator_core + from aiconfigurator_core.sdk import RustForwardPassPerfModel + + class Estimator: + def __init__(self, config): + self.config = json.loads(aiconfigurator_core.RustForwardPassPerfModel.normalize_config(json.dumps(config))) + self.config.update(backend_version="0.24.0", estimation_mode="op_level", fallback_policy="deny") + + def diagnostics(self): + return {"readiness": "ready", "provenance": {"config": self.config}} + + def close(self): + pass + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", Estimator) public = CorePredictionConfig.from_yaml(path) spec = prediction_to_replay_spec(public) + if rank_only_controls: + timing = spec.backend_deployment.agg_engine_args["timing_model"]["config"] + timing.pop("cuda_graph_reserved_bytes") + timing.pop("gpu_memory_utilization") runtime = RecordingRuntime() assert public.engine.workers.aggregated is not None @@ -1003,6 +1068,32 @@ def test_runner_accepts_matching_backend_version_in_explicit_aic_timing(): assert timing_config["backend_version"] == "0.11.1" +def test_direct_replay_normalizes_both_version_aliases_without_mutating_input(): + runtime = RecordingRuntime() + timing = { + "type": "external", + "provider": "aic", + "config": { + "model": "Qwen/Qwen3-32B", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "current", + "tp": 2, + "attention_dp": 1, + }, + } + args = {"tensor_parallel_size": 2, "rank": {"block_size": 4, "num_gpu_blocks": 128, "timing_model": timing}} + deployment = BackendDeploymentSpec( + deployment_mode="agg", backend="vllm", backend_version="current", agg_engine_args=args, num_workers=1 + ) + EngineReplayRunnerFactory(runtime=runtime).create(0).run(_spec(deployment=deployment)) + actual = runtime.execution_spec["engine"]["rank"]["timing_model"]["config"]["backend_version"] + from aiconfigurator_core.sdk.perf_database import resolve_query_version + + assert actual == resolve_query_version("h200_sxm", "vllm", "current") + assert timing["config"]["backend_version"] == "current" + + def test_runner_rejects_conflicting_backend_version_in_explicit_aic_timing(): timing = { "type": "external",