diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 6c3b33bb0..95489e4ee 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -33,6 +33,8 @@ areas: - python/aisimulate/ - python/aisimulate/src/aisimulate/capacity.py - tests/test_aic.py + - tests/test_fpm_profile_workflow.py + - tests/test_fpm_query_evidence.py - label: sweeper github_team: "@ai-dynamo/aisimulate-sweeper-codeowners" path_globs: @@ -158,6 +160,10 @@ shared: owners: [forward-pass-engine, maintainers] - glob: tests/test_aic.py owners: [forward-pass-engine, maintainers] + - glob: tests/test_fpm_profile_workflow.py + owners: [forward-pass-engine, maintainers] + - glob: tests/test_fpm_query_evidence.py + owners: [forward-pass-engine, replay, maintainers] - glob: python/aisimulate/src/aisimulate/sweeper/ owners: [sweeper, maintainers] - glob: docs/sweeper/ diff --git a/CODEOWNERS b/CODEOWNERS index 274892cd0..277f513a2 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -45,6 +45,8 @@ /crates/tests/public-api/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners /crates/core/src/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners /crates/core/tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners +/tests/test_fpm_query_evidence.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners +/tests/test_fpm_profile_workflow.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners /crates/core/parity_tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners /python/aisimulate/src/aisimulate/capacity.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @@ -179,12 +181,14 @@ /crates/core/src/engine/timing.rs @ai-dynamo/aisimulate-mocker-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/.gitattributes @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/pyproject.toml @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-mocker-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_fpm_query_evidence.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_resource_scheduler.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_state_cache_config.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-mocker-codeowners @ai-dynamo/access-aisimulate-maintain /crates/tests/public-api/Cargo.toml @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/.coderabbit.yaml @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/eslint.config.js @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_e2e_accuracy_nightly.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_fpm_profile_workflow.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_resource_supervision.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/access-aisimulate-maintain /crates/core/parity_tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_ci_workflow_contracts.py @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain diff --git a/crates/core/src/engine/timing.rs b/crates/core/src/engine/timing.rs index f1351b62e..0edb558e7 100644 --- a/crates/core/src/engine/timing.rs +++ b/crates/core/src/engine/timing.rs @@ -208,6 +208,11 @@ impl TimingOperationEvidence { left.fallbacks.push(fallback); } } + crate::perfmodel::fpm::accumulate_fpm_estimates( + &mut left.fpm_estimates, + right.fpm_estimates, + ) + .map_err(anyhow::Error::msg)?; Some(left) } (Some(mut details), None) | (None, Some(mut details)) => { @@ -734,6 +739,12 @@ pub trait TimingModel: Send + Sync { None } + /// Latency-only whole-model FPM evidence is separate from energy evidence. + /// It must not change the provider's existing power publication contract. + fn fpm_evidence_summary(&self) -> Option { + None + } + /// Start a new measurement epoch without changing timing predictions or /// provider caches. Called only after preparation work has settled. /// @@ -912,6 +923,7 @@ mod tests { sol: None, sol_unavailable_reason: Some("unsupported".into()), fallbacks: Vec::new(), + fpm_estimates: Vec::new(), }) } } @@ -1021,6 +1033,7 @@ mod tests { TimingOperationEvidence::new("test", 2.0, None, TimingEvidenceSource::Estimated) .unwrap(); op.details = Some(OperationDetails { + fpm_estimates: Vec::new(), sol: has_sol.then_some(SolDiagnostics { latency_ms: 1.0, math_ms: 0.0, @@ -1079,6 +1092,7 @@ mod tests { ) .unwrap(); operation.details = Some(OperationDetails { + fpm_estimates: Vec::new(), sol: None, sol_unavailable_reason: Some("unsupported operation".into()), fallbacks: vec![valid.clone(), fallback], @@ -1100,6 +1114,7 @@ mod tests { .unwrap(); let mut detailed = plain.clone(); detailed.details = Some(OperationDetails { + fpm_estimates: Vec::new(), sol: Some(SolDiagnostics { latency_ms: 1.0, math_ms: 0.0, @@ -1455,6 +1470,7 @@ mod tests { }), sol_unavailable_reason: None, fallbacks: Vec::new(), + fpm_estimates: Vec::new(), }); } let incoming = ValidatedTimingPhase::from_operations(vec![op]).unwrap(); diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 43029ffd1..d918d4afa 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -48,8 +48,8 @@ pub use perfmodel::{ SpeculativeConfig, }; pub use perfmodel::{ - CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, OpLevelConfig, - RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, + CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, FpmInterpolationMethod, + OpLevelConfig, RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, }; #[cfg(feature = "python")] @@ -75,3 +75,8 @@ pub(crate) use perfmodel::{ #[cfg(feature = "python")] #[allow(unused_imports)] pub(crate) use perfmodel::{py, py_ops}; + +pub use perfmodel::{ + DirectFpmQueryEvidence, ForwardPassEstimate, FpmCoordinates, FpmEstimateEvidence, + FpmMeasurementSupport, FpmQueryResolution, FpmRankEstimate, +}; diff --git a/crates/core/src/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index c6bc96317..2cc4e44b3 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -119,7 +119,10 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // JSON defaults preserve unrecorded DCP. // - 24 (typed FPM DCP): FpmForwardOp carries dcp_size separately from the // base matching identity, so control flow never parses the string tuple. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 24; +// - 25 (FPM decoupling): FpmForwardOp also carries the SOL/direct interpolation +// selector. The decoupling branch independently used 21 for this positional +// field; the combined layout differs from every prior schema, including 24. +pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 25; /// Static engine identity and setup information carried by an /// [`crate::perfmodel::engine::spec::EngineSpec`]. diff --git a/crates/core/src/perfmodel/engine/diagnostics.rs b/crates/core/src/perfmodel/engine/diagnostics.rs index 820037067..afcf51ace 100644 --- a/crates/core/src/perfmodel/engine/diagnostics.rs +++ b/crates/core/src/perfmodel/engine/diagnostics.rs @@ -27,6 +27,9 @@ pub struct OperationDetails { pub sol_unavailable_reason: Option, /// Executed measurement substitutions; empty means no recorded substitution. pub fallbacks: Vec, + /// Whole-forward estimates, counted by identical result and replay scale. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub fpm_estimates: Vec, } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] diff --git a/crates/core/src/perfmodel/engine/readiness.rs b/crates/core/src/perfmodel/engine/readiness.rs index c555109e9..b2aa3e8c2 100644 --- a/crates/core/src/perfmodel/engine/readiness.rs +++ b/crates/core/src/perfmodel/engine/readiness.rs @@ -10,6 +10,7 @@ use crate::AicError; use crate::common::enums::{DatabaseMode, GemmQuantMode, MoeQuantMode, TransferKind}; use crate::config::PerfSource; use crate::operators::Op; +use crate::operators::fpm_forward::{FpmForwardOp, FpmInterpolation, FpmPhase}; use crate::perf_database::moe::MoeKernel; use crate::perf_database::{PerfDatabase, kernel_source_ok, parquet_loader::PerfReader}; @@ -27,6 +28,28 @@ pub(super) fn validate<'a>( Ok(()) } +pub(super) fn validate_fpm(db: &PerfDatabase, fpm: &FpmForwardOp) -> Result<(), AicError> { + let cell = db + .fpm_forward + .select_cell(&fpm.match_identity, &fpm.model_path, fpm.dcp_size)?; + if fpm.interpolation == FpmInterpolation::Direct + && match fpm.phase { + FpmPhase::Prefill => cell.direct_prefill.is_empty(), + FpmPhase::Decode => cell.direct_decode.is_empty(), + } + { + let phase = fpm.phase.as_str(); + return Err(AicError::PerfDatabase(format!( + "direct FPM {phase} interpolation has no genuine measurements for {:?} at {}. \ + Collect genuine {phase} FPM rows for this model and identity, or provide a \ + systems root containing them.", + fpm.model_path, + db.data_root.display(), + ))); + } + Ok(()) +} + struct Availability<'a> { db: &'a PerfDatabase, tables: HashMap, @@ -162,14 +185,7 @@ impl Availability<'_> { } } TokenScale(scale) => return self.op(&scale.op), - FpmForward(fpm) => { - self.db.fpm_forward.select_cell( - &fpm.match_identity, - &fpm.model_path, - fpm.dcp_size, - )?; - return Ok(()); - } + FpmForward(fpm) => return validate_fpm(self.db, fpm), // These families have no analytic/empirical implementation. MoeAllToAll(_) | MoeExpertCompute(_) | Dsv4MegaMoe(_) if !matches!( @@ -446,6 +462,62 @@ mod tests { write_parquet(path, &[Col::Str("kernel_source", vec!["fixture"])]); } + #[test] + fn direct_fpm_readiness_requires_genuine_rows_only_for_requested_phases() { + use crate::perf_database::fpm_forward::tests::{RowSpec, default_identity, write_pair}; + + for missing_phase in [FpmPhase::Prefill, FpmPhase::Decode] { + for fake_only in [false, true] { + let root = systems(); + let data = root.path().join("data/b200_sxm/vllm/0.25.1"); + std::fs::create_dir_all(&data).unwrap(); + let present_phase = match missing_phase { + FpmPhase::Prefill => FpmPhase::Decode, + FpmPhase::Decode => FpmPhase::Prefill, + }; + let row = |phase: FpmPhase, regime| RowSpec { + workload_kind: phase.as_str(), + total_prefill_tokens: if phase == FpmPhase::Prefill { 8 } else { 0 }, + kv_seed_regime: Some(regime), + ..RowSpec::default() + }; + let mut rows = vec![row(present_phase, "real_kv")]; + if fake_only { + rows.push(row(missing_phase, "fake_fallback")); + } + write_pair(&data, &rows); + let db = PerfDatabase::load(root.path(), "b200_sxm", "vllm", "0.25.1").unwrap(); + let op = |phase: FpmPhase| { + Op::FpmForward(FpmForwardOp { + dcp_size: None, + name: format!("fpm_forward_{}", phase.as_str()), + phase, + model_path: "org/model-a".into(), + match_identity: default_identity(4), + weight_bytes: 0.0, + verify_width: 1, + interpolation: FpmInterpolation::Direct, + sol_ops: vec![], + original_fmha_quant_mode: None, + }) + }; + let present = op(present_phase); + let missing = op(missing_phase); + // A single-phase operation requires only its own measurements. + validate(&db, [&present].into_iter()).unwrap(); + for requested in [vec![&missing], vec![&present, &missing]] { + let error = validate(&db, requested.into_iter()) + .unwrap_err() + .to_string(); + assert!(error.contains(missing_phase.as_str()), "{error}"); + assert!(error.contains("no genuine measurements"), "{error}"); + assert!(error.contains("Collect genuine"), "{error}"); + assert!(error.contains(data.to_str().unwrap()), "{error}"); + } + } + } + } + #[test] fn dsv41_stage_requires_measured_module_for_silicon() { use crate::operators::dsv41::{Dsv41LinearOp, Dsv41StageOp}; diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 4206a65a6..a3dd55eb2 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -735,6 +735,14 @@ impl Engine { /// Python `_run_context_phase` (`base_backend.py:144`): `effective_isl = /// isl - prefix`, validate `> 0`, then one full pass over `context_ops`. fn run_context_phase(&self, runtime: &RuntimeConfig) -> Result { + self.run_context_phase_with(runtime, |_, _| {}) + } + + fn run_context_phase_with( + &self, + runtime: &RuntimeConfig, + on_op: impl FnMut(&Op, PerformanceResult), + ) -> Result { // Python raises `ValueError` when `effective_isl <= 0`; mirror that. if runtime.prefix >= runtime.isl { return Err(AicError::InvalidEngineConfig(format!( @@ -743,7 +751,7 @@ impl Engine { ))); } let effective_isl = runtime.isl - runtime.prefix; - run_context_ops( + run_context_ops_with( &self.context_ops, &self.db, runtime.batch_size, @@ -751,6 +759,7 @@ impl Engine { runtime.prefix, runtime.seq_imbalance_correction_scale, ContextOpFilter::All, + on_op, ) } @@ -855,6 +864,42 @@ impl Engine { } } + /// Preserve static prefill geometry, including explicitly equal per-request + /// lengths for bounded decoder replay. Aggregate telemetry cannot express + /// that information and retains its separate validation. + pub(crate) fn predict_prefill_detailed( + &self, + bs: u32, + isl: u32, + prefix: u32, + ) -> Result { + let runtime = RuntimeConfig { + batch_size: bs, + isl, + prefix, + osl: 1, + ..Default::default() + }; + let mut queries = Vec::new(); + let prefill_ms = self.run_context_phase_with(&runtime, |_, result| { + queries.extend(result.fpm_queries.into_iter().flat_map(|queries| *queries)); + })?; + let latency_ms = prefill_ms + 0.0; // same context + generation sum as run_static + Ok(crate::ForwardPassEstimate { + latency_ms: Some(latency_ms), + native_latency_ms: Some(latency_ms), + correction_factor: Some(1.0), + max_rank: (latency_ms > 0.0).then_some(0), + ranks: vec![crate::FpmRankEstimate { + rank: 0, + latency_ms, + prefill_ms: Some(prefill_ms), + queries, + ..Default::default() + }], + }) + } + /// Mocker H2: decode-step latency in ms. Pure-Rust inherent method (no /// PyO3 `py` token). Thin shim over [`Self::run_static`] with /// `mode=Generation`. Mocker passes `osl=2` (one decode step at @@ -1636,6 +1681,7 @@ impl Engine { energy_wms, source: source.into(), details: OperationDetails { + fpm_estimates: Vec::new(), sol, sol_unavailable_reason, fallbacks, @@ -2155,6 +2201,15 @@ impl Engine { &self, metrics_by_rank: &[ForwardPassMetrics], ) -> Result { + self.forward_pass_estimate(metrics_by_rank, false) + .map(|estimate| estimate.native_latency_ms.unwrap()) + } + + pub(crate) fn forward_pass_estimate( + &self, + metrics_by_rank: &[ForwardPassMetrics], + capture: bool, + ) -> Result { self.require_general_engine()?; if metrics_by_rank.is_empty() { return Err(AicError::InvalidForwardPassMetrics( @@ -2165,13 +2220,30 @@ impl Engine { validate_forward_pass_metrics(metrics)?; } let mut max_latency = 0.0_f64; - for metrics in metrics_by_rank { - let rank_latency = self.rank_latency_ms(metrics)?; + let mut max_rank = None; + let mut ranks = Vec::new(); + for (rank, metrics) in metrics_by_rank.iter().enumerate() { + let mut evidence = crate::FpmRankEstimate { + rank, + ..Default::default() + }; + let rank_latency = self.rank_latency_ms(metrics, capture.then_some(&mut evidence))?; if rank_latency > max_latency { max_latency = rank_latency; + max_rank = Some(rank); + } + if capture && self.fpm_ops().is_some() { + evidence.latency_ms = rank_latency; + ranks.push(evidence); } } - Ok(max_latency) + Ok(crate::ForwardPassEstimate { + latency_ms: Some(max_latency), + native_latency_ms: Some(max_latency), + correction_factor: Some(1.0), + max_rank, + ranks, + }) } /// Dispatch one rank's FPM on its scheduled workload. Literal port of @@ -2180,7 +2252,11 @@ impl Engine { /// [`run_context_ops`]; decode-only -> [`run_generation_ops_step`]. The FPM /// counts pass through unscaled (no `nextn` multiplier — see /// [`Self::forward_pass_time_ms`]). - fn rank_latency_ms(&self, metrics: &ForwardPassMetrics) -> Result { + fn rank_latency_ms( + &self, + metrics: &ForwardPassMetrics, + mut evidence: Option<&mut crate::FpmRankEstimate>, + ) -> Result { let sched = &metrics.scheduled_requests; // Token-based dispatch, aligned with `IterationFeatures` (fpm/model.rs): // a fully prefix-cached payload can retain prefill request/KV metadata @@ -2245,39 +2321,58 @@ impl Engine { // division on each axis. let mut total = 0.0_f64; if has_prefill { - total += prefill_op - .query_totals( - &self.db, - &[ - sched.num_prefill_requests as f64, - sched.sum_prefill_tokens as f64, - sched.sum_prefill_kv_tokens as f64, - ], - )? - .latency_ms; + let result = prefill_op.query_totals( + &self.db, + &[ + sched.num_prefill_requests as f64, + sched.sum_prefill_tokens as f64, + sched.sum_prefill_kv_tokens as f64, + ], + )?; + total += result.latency_ms; + if let Some(evidence) = evidence.as_deref_mut() { + evidence.prefill_ms = Some(result.latency_ms); + evidence + .queries + .extend(result.fpm_queries.into_iter().flat_map(|queries| *queries)); + } } if has_decode { - let decode_ms = decode_op - .query_totals( - &self.db, - &[ - sched.num_decode_requests as f64, - sched.sum_decode_kv_tokens as f64, - ], - )? - .latency_ms; + let result = decode_op.query_totals( + &self.db, + &[ + sched.num_decode_requests as f64, + sched.sum_decode_kv_tokens as f64, + ], + )?; + let decode_ms = result.latency_ms; + if let Some(evidence) = evidence.as_deref_mut() { + evidence.decode_ms = Some(decode_ms); + evidence + .queries + .extend(result.fpm_queries.into_iter().flat_map(|queries| *queries)); + } if has_prefill { // Mixed rank: marginal-decode composition, mirroring // `_get_fpm_mix_step_latency` (counts already packed, no // `(nextn + 1)` — speculative FPM was rejected above). - let baseline_ms = decode_op - .query_pass_baseline( - &self.db, - sched.num_decode_requests, - sched.sum_decode_kv_tokens as f64, - )? - .latency_ms; + let baseline = decode_op.query_pass_baseline( + &self.db, + sched.num_decode_requests, + sched.sum_decode_kv_tokens as f64, + )?; + let baseline_ms = baseline.latency_ms; total += (decode_ms - baseline_ms).max(0.0); + if let Some(evidence) = evidence.as_deref_mut() { + evidence.decode_baseline_ms = Some(baseline_ms); + evidence.marginal_decode_ms = Some((decode_ms - baseline_ms).max(0.0)); + evidence.queries.extend( + baseline + .fpm_queries + .into_iter() + .flat_map(|queries| *queries), + ); + } } else { total += decode_ms; } @@ -3532,6 +3627,7 @@ mod tests { config.systems_path = Some(tmp.to_path_buf()); let fpm_op = |phase: FpmPhase| { Op::FpmForward(FpmForwardOp { + interpolation: Default::default(), dcp_size: None, name: format!("fpm_forward_{}", phase.as_str()), phase, @@ -3604,6 +3700,7 @@ mod tests { use crate::perf_database::fpm_forward::tests::default_identity; let db = PerfDatabase::load(&systems_root(), "b200_sxm", "vllm", "0.24.0").unwrap(); let fpm_op = Op::FpmForward(FpmForwardOp { + interpolation: Default::default(), dcp_size: None, name: "fpm_forward_prefill".into(), phase: FpmPhase::Prefill, @@ -3663,6 +3760,7 @@ mod tests { )); let fpm_op = |phase: FpmPhase| { Op::FpmForward(FpmForwardOp { + interpolation: Default::default(), dcp_size: None, name: format!("fpm_forward_{}", phase.as_str()), phase, @@ -3893,6 +3991,27 @@ mod tests { .context_ms, 43.0 ); + let detailed = engine + .predict_prefill_detailed(input.batch_size, input.isl, input.prefix) + .unwrap(); + assert_eq!(detailed.latency_ms, Some(43.0)); + assert_eq!( + engine + .predict_prefill_latency(input.batch_size, input.isl, input.prefix) + .unwrap(), + 43.0 + ); + // Capturing the known static geometry does not weaken the aggregate + // telemetry guard for multi-request bounded decoder replay. + assert!( + engine + .forward_pass_estimate(&[fpm_replay_guard_metrics(0.0)], true) + .is_err() + ); + for (isl, prefix) in [(0, 0), (512, 512)] { + assert!(engine.predict_prefill_latency(2, isl, prefix).is_err()); + assert!(engine.predict_prefill_detailed(2, isl, prefix).is_err()); + } } #[test] @@ -4059,6 +4178,118 @@ mod tests { /// Mixed telemetry may request a synthetic decode baseline below the KV /// floor of its padded bracket rows. Only that baseline holds each row at /// its measured floor; the actual decode query remains in-range and strict. + #[test] + fn direct_fpm_detailed_result_explains_mixed_rank_max_and_correction() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let rows = vec![ + RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: 4, + total_kv_read_tokens: 0, + latency_ms: 10.0, + ..Default::default() + }, + RowSpec { + batch_size: 2, + total_kv_read_tokens: 4, + latency_ms: 10.0, + ..Default::default() + }, + RowSpec { + batch_size: 2, + total_kv_read_tokens: 20, + latency_ms: 26.0, + ..Default::default() + }, + RowSpec { + batch_size: 4, + total_kv_read_tokens: 8, + latency_ms: 20.0, + ..Default::default() + }, + RowSpec { + batch_size: 4, + total_kv_read_tokens: 24, + latency_ms: 52.0, + ..Default::default() + }, + ]; + let tmp = tempfile::tempdir().unwrap(); + let mut engine = build_fpm_engine_with_rows(tmp.path(), &rows).unwrap(); + for op in engine + .context_ops + .iter_mut() + .chain(engine.generation_ops.iter_mut()) + { + if let Op::FpmForward(op) = op { + op.interpolation = + crate::perfmodel::operators::fpm_forward::FpmInterpolation::Direct; + } + } + let mut model = crate::ForwardPassPerfModel::from_engine( + Arc::new(engine), + crate::ForwardPassPerfOptions { + min_observations: 1, + ..Default::default() + }, + ); + let mixed = ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_prefill_requests: 1, + sum_prefill_tokens: 4, + num_decode_requests: 3, + sum_decode_kv_tokens: 12, + ..Default::default() + }, + ..Default::default() + }; + let decode = ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_decode_requests: 4, + sum_decode_kv_tokens: 24, + ..Default::default() + }, + ..Default::default() + }; + let metrics = [mixed, decode]; + let result = model.estimate_forward_pass_detailed(&metrics).unwrap(); + assert_eq!( + result.latency_ms, + model.estimate_forward_pass_time_ms(&metrics).unwrap() + ); + assert_eq!(result.latency_ms, Some(52.0)); + assert_eq!(result.max_rank, Some(1)); + let rank = &result.ranks[0]; + assert_eq!( + ( + rank.prefill_ms, + rank.decode_ms, + rank.decode_baseline_ms, + rank.marginal_decode_ms, + rank.latency_ms + ), + (Some(10.0), Some(23.0), Some(15.0), Some(8.0), 18.0) + ); + assert_eq!(rank.queries.len(), 3); + let mut observed = metrics.clone(); + observed[1].wall_time = 0.104; + model.tune_with_fpms(&[observed.to_vec()]).unwrap(); + let corrected = model.estimate_forward_pass_detailed(&metrics).unwrap(); + assert_eq!(corrected.correction_factor, Some(2.0)); + assert_eq!(corrected.latency_ms, Some(104.0)); + assert_eq!(corrected.ranks, result.ranks); + assert_eq!( + corrected.latency_ms, + model.estimate_forward_pass_time_ms(&metrics).unwrap() + ); + let empty = model + .estimate_forward_pass_detailed(&[ForwardPassMetrics::default()]) + .unwrap(); + assert_eq!(empty.latency_ms, Some(0.0)); + assert!(empty.ranks.iter().all(|rank| rank.queries.is_empty())); + } + #[test] fn fpm_rank_mixed_baseline_holds_bracket_curve_floors() { use crate::fpm::{ForwardPassMetrics, ScheduledRequestMetrics}; @@ -4209,6 +4440,7 @@ mod tests { let hidden = Op::Overlap(crate::operators::OverlapOp::new( "hidden", vec![Op::FpmForward(FpmForwardOp { + interpolation: Default::default(), dcp_size: None, name: "fpm_forward_prefill".into(), phase: FpmPhase::Prefill, @@ -4528,6 +4760,7 @@ mod tests { "0.25.1", )); let mut op = FpmForwardOp { + interpolation: Default::default(), dcp_size: None, name: "fpm_forward_decode".into(), phase: FpmPhase::Decode, @@ -4575,6 +4808,7 @@ mod tests { )); let fpm_op = |phase: FpmPhase, width: u32| { Op::FpmForward(FpmForwardOp { + interpolation: Default::default(), dcp_size: None, name: format!("fpm_forward_{}", phase.as_str()), phase, diff --git a/crates/core/src/perfmodel/engine/spec.rs b/crates/core/src/perfmodel/engine/spec.rs index 4ef523675..7e62af7eb 100644 --- a/crates/core/src/perfmodel/engine/spec.rs +++ b/crates/core/src/perfmodel/engine/spec.rs @@ -625,6 +625,7 @@ mod tests { // Recursive like Overlap/Fallback: sol_ops carries the model's // original granular list, so the round-trip must preserve nesting. crate::operators::FpmForwardOp { + interpolation: Default::default(), dcp_size: None, name: "fpm_forward_prefill".into(), phase: crate::operators::FpmPhase::Prefill, @@ -1047,7 +1048,7 @@ mod tests { vec![OpSpec::FpmForward(fpm_forward())], vec![OpSpec::Moe(moe())], ); - assert_eq!(spec.schema_version, 24); + assert_eq!(spec.schema_version, ENGINE_SPEC_SCHEMA_VERSION); let mut bytes = spec.to_bincode().unwrap(); assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); @@ -1130,6 +1131,49 @@ mod tests { assert_eq!(decoded.schema_version, ENGINE_SPEC_SCHEMA_VERSION); } + #[test] + fn fpm_interpolation_and_dcp_round_trip_and_stale_layout_rejection() { + let mut op = fpm_forward(); + op.interpolation = crate::operators::fpm_forward::FpmInterpolation::Direct; + op.sol_ops.clear(); + op.original_fmha_quant_mode = None; + op.dcp_size = None; + let spec = EngineSpec::new(sample_engine_config(), vec![], vec![OpSpec::FpmForward(op)]); + let bytes = spec.to_bincode().unwrap(); + assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); + // The final op ends in interpolation (4 bytes), empty sol_ops (8), + // original-FMHA None (1), and recorded-DCP None (1). + let selector = bytes.len() - 14; + assert_eq!(&bytes[selector..selector + 4], &1u32.to_le_bytes()); + // Recreate both branches' incompatible payloads, not just their stamps. + for version in [20u32, 21, 22, 23, 24] { + let mut stale = bytes.clone(); + if version != 21 { + stale.drain(selector..selector + 4); + } + if version < 24 { + stale.pop(); + } + stale[..4].copy_from_slice(&version.to_le_bytes()); + assert!(matches!( + EngineSpec::from_bincode(&stale), + Err(AicError::UnsupportedSchemaVersion { + kind: "EngineSpec", + got, + expected: ENGINE_SPEC_SCHEMA_VERSION, + }) if got == version + )); + } + let mut recorded = spec; + if let OpSpec::FpmForward(op) = &mut recorded.generation_ops[0] { + op.dcp_size = Some(8); + } + assert_eq!( + EngineSpec::from_bincode(&recorded.to_bincode().unwrap()).unwrap(), + recorded + ); + } + /// A buffer too short to even hold the 4-byte version prefix must fail at the /// prefix stage, not deep in the (absent) payload. #[test] @@ -1380,7 +1424,7 @@ mod tests { for spec in [fpm_spec, pilot_spec] { let bytes = spec.to_bincode().unwrap(); assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); - for previous_version in [20u32, 21, 22, 23] { + for previous_version in [20u32, 21, 22, 23, 24] { let mut stale = bytes.clone(); stale[..4].copy_from_slice(&previous_version.to_le_bytes()); // DCP and the pilot claimed 22 for different layouts. @@ -1391,7 +1435,7 @@ mod tests { Err(AicError::UnsupportedSchemaVersion { kind: "EngineSpec", got, - expected: 24, + expected: ENGINE_SPEC_SCHEMA_VERSION, }) if got == previous_version )); } diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 457fc9238..eacde9f89 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -87,6 +87,9 @@ pub struct ForwardPassPerfModelConfig { pub worker_type: ForwardPassWorkerType, #[serde(default)] pub backend_version: Option, + /// Full independent model identity, precision and resource profile. + #[serde(default)] + pub fpm_profile: Option>, #[serde(default = "one", alias = "tp_size")] pub tp: u32, @@ -171,6 +174,7 @@ impl ForwardPassPerfModelConfig { backend, worker_type, backend_version: None, + fpm_profile: None, tp: 1, pp: 1, attention_dp: 1, @@ -226,6 +230,72 @@ impl ForwardPassPerfModelConfig { candidates } + /// Validate profile metadata and pin interpolation without constructing a + /// graph or opening performance data. This also serves configuration-only + /// consumers such as memory planning. + pub(crate) fn resolve(self) -> Result { + self.resolve_with_registration().map(|(config, _)| config) + } + + /// Retain registry availability for native candidate selection without + /// adding transient registry facts to the serialized request. + pub(crate) fn resolve_with_registration(mut self) -> Result<(Self, Option), AicError> { + self.resolve_prefill_graph_profile()?; + self.validate()?; + let registered = if self.fpm_profile.is_some() { + self.estimator_config + .fpm_interpolation + .validate_profile_options()?; + if self.dcp.is_some() { + return Err(invalid_config( + "fpm_profile does not describe recorded DCP identity", + )); + } + if self.backend_version.is_none() { + return Err(invalid_config( + "fpm_profile requires a literal backend_version", + )); + } + if self.nextn != 0 || self.speculation.is_some() { + return Err(invalid_config( + "FPM profiles support plain autoregressive decoder-only execution", + )); + } + #[cfg(feature = "python")] + { + let facts = crate::py::validate_forward_pass_profile(&self)?; + self.fpm_profile = Some(facts.profile); + self.gemm_quant_mode = Some(facts.gemm_quant_mode); + self.moe_quant_mode = Some(facts.moe_quant_mode); + self.fmha_quant_mode = Some(facts.fmha_quant_mode); + self.kvcache_quant_mode = Some(facts.kvcache_quant_mode); + self.comm_quant_mode = Some(facts.comm_quant_mode); + self.attention_backend = Some(facts.attention_backend); + Some(facts.registered) + } + #[cfg(not(feature = "python"))] + return Err(invalid_config( + "FPM profile validation requires the python feature", + )); + } else { + None + }; + self.estimator_config.fpm_interpolation.method = self + .estimator_config + .fpm_interpolation + .method + .resolve(registered)?; + if self.estimator_config.fpm_interpolation.method == super::FpmInterpolationMethod::Direct + && self.database_mode != DatabaseMode::Silicon + { + return Err(invalid_config( + "direct FPM interpolation requires database_mode='SILICON'; it has no analytical SOL model", + )); + } + self.validate()?; + Ok((self, registered)) + } + pub(crate) fn validate(&self) -> Result<(), AicError> { self.validate_prefill_graph_profile()?; if let Some(distribution) = &self.estimator_config.op_level.decode_workload_distribution { @@ -303,8 +373,11 @@ impl ForwardPassPerfModelConfig { "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) + // Profile schema also admits dense TP with unused MoE dimensions + // equal to one, and checks the exact requested deployment below. + if self.fpm_profile.is_none() + && 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", @@ -324,7 +397,7 @@ impl ForwardPassPerfModelConfig { )); } } - if self.estimation_mode == EstimationMode::FpmInterpolation + if (self.estimation_mode == EstimationMode::FpmInterpolation || self.fpm_profile.is_some()) && (self.enable_eplb || self.moe_kernel_source.is_some() || self.wideep_num_slots.is_some() @@ -467,6 +540,38 @@ mod tests { serde_json::from_value(base).unwrap() } + #[test] + fn interpolation_controls_are_typed_and_resolved_before_candidates() { + use super::super::FpmInterpolationMethod::{Auto, Direct, Sol}; + for (requested, registered, expected) in [ + (Auto, None, Sol), + (Auto, Some(true), Sol), + (Auto, Some(false), Direct), + (Sol, None, Sol), + (Sol, Some(true), Sol), + (Direct, Some(true), Direct), + (Direct, Some(false), Direct), + ] { + assert_eq!(requested.resolve(registered).unwrap(), expected); + } + assert!(Direct.resolve(None).is_err()); + assert!(Sol.resolve(Some(false)).is_err()); + let cfg = config(serde_json::json!({})).resolve().unwrap(); + assert_eq!(cfg.estimator_config.fpm_interpolation.method, Sol); + assert_eq!(cfg.estimation_mode, EstimationMode::Auto); + let cfg = config(serde_json::json!({ + "estimation_mode": "fpm_regression", "fallback_policy": "allow", + "estimator_config": {"fpm_interpolation": {"method": "direct"}} + })); + assert!(cfg.resolve().is_err()); + assert!( + serde_json::from_value::(serde_json::json!({ + "fpm_interpolation": {"method": "dierct"} + })) + .is_err() + ); + } + #[test] fn fpm_selector_is_validated_and_preserved_in_canonical_identity() { let cfg = config(serde_json::json!({ @@ -511,6 +616,24 @@ mod tests { assert_eq!(decoded.fallback_policy, ForwardPassFallbackPolicy::Deny); } + #[test] + fn profile_auto_selection_does_not_drop_unsupported_engine_controls() { + for controls in [ + serde_json::json!({"enable_eplb": true}), + serde_json::json!({"wideep_num_slots": 256}), + serde_json::json!({"moe_backend": "deepep_moe", "backend": "sglang"}), + ] { + let mut cfg = config(controls); + cfg.fpm_profile = Some(Default::default()); + assert!( + cfg.validate() + .unwrap_err() + .to_string() + .contains("does not support EPLB") + ); + } + } + #[test] fn auto_selects_all_modes_even_when_fallback_is_denied() { let cfg = config(serde_json::json!({})); diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs index 655b6c635..485622bc2 100644 --- a/crates/core/src/perfmodel/fpm/estimator.rs +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -37,6 +37,7 @@ pub struct OpLevelConfig { #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct FpmInterpolationConfig { + pub method: FpmInterpolationMethod, /// The profile covers text prefill/decode; encoder weights remain resident. #[serde(skip_serializing_if = "is_false")] pub text_only: bool, @@ -60,6 +61,15 @@ pub enum UnrecordedFpmQuantMode { } impl FpmInterpolationConfig { + pub(crate) fn validate_profile_options(&self) -> Result<(), AicError> { + if self.text_only || !self.unrecorded_quant_modes.is_empty() { + return Err(super::config::invalid_config( + "fpm_profile does not support text_only or unrecorded_quant_modes overrides", + )); + } + Ok(()) + } + pub(crate) fn validate_quant_modes( &self, fmha: Option<&str>, @@ -80,6 +90,40 @@ impl FpmInterpolationConfig { } } +/// Construction-time interpolation selection. Native operators receive only +/// the resolved SOL or direct method and never change it during queries. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmInterpolationMethod { + #[default] + Auto, + Sol, + Direct, +} + +impl FpmInterpolationMethod { + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::Auto => "auto", + Self::Sol => "sol", + Self::Direct => "direct", + } + } + + pub(crate) fn resolve(self, registered: Option) -> Result { + match (self, registered) { + (Self::Direct, None) => Err(AicError::InvalidEngineConfig( + "direct FPM interpolation requires an fpm_profile with identity and resource metadata".into(), + )), + (Self::Sol, Some(false)) => Err(AicError::InvalidEngineConfig( + "SOL interpolation requires a registered analytical model class; use estimator_config.fpm_interpolation.method='direct' with the supplied profile".into(), + )), + (Self::Auto, Some(false)) | (Self::Direct, Some(_)) => Ok(Self::Direct), + _ => Ok(Self::Sol), + } + } +} + /// These weights currently affect regression. Native correction keeps its /// established workload coordinates until a replacement is accuracy-qualified. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -219,6 +263,128 @@ impl Default for CorrectionConfig { } impl EstimatorConfig { + /// Merge explicit migration inputs before applying canonical defaults. + /// Serialized defaults from parsing one input must not become choices that + /// conflict with a setting supplied by another input. + #[cfg(any(feature = "python", test))] + pub(crate) fn migrate_legacy_inputs( + estimator_json: Option<&str>, + options_json: Option<&str>, + interpolation_method: Option<&str>, + fpm_parquet_path: Option<&str>, + ) -> Result { + let mut explicit = serde_json::json!({}); + if let Some(json) = estimator_json { + let supplied: serde_json::Value = serde_json::from_str(json) + .map_err(|error| invalid_migration("estimator_config", error))?; + let typed: Self = serde_json::from_str(json) + .map_err(|error| invalid_migration("estimator_config", error))?; + // Normalize scalar representations (for example 1 and 1.0) while + // retaining only fields that were actually supplied. + let normalized = serde_json::to_value(typed) + .map_err(|error| invalid_migration("estimator_config", error))?; + explicit = explicit_fields(&supplied, &normalized); + } + if let Some(json) = options_json { + let supplied: serde_json::Value = serde_json::from_str(json) + .map_err(|error| invalid_migration("legacy options", error))?; + let options: ForwardPassPerfOptions = serde_json::from_str(json) + .map_err(|error| invalid_migration("legacy options", error))?; + let normalized = serde_json::to_value(&options) + .map_err(|error| invalid_migration("legacy options", error))?; + let mut migrated = serde_json::json!({}); + for (legacy, paths) in [ + ( + "max_observations", + &[ + "fpm_regression.sampling.max_observations", + "correction.sampling.max_observations", + ][..], + ), + ( + "min_observations", + &[ + "fpm_regression.min_observations", + "correction.min_observations", + ], + ), + ( + "min_faster_correction_factor", + &["correction.factor_bounds.min"], + ), + ( + "max_slower_correction_factor", + &["correction.factor_bounds.max"], + ), + ("max_num_tokens", &["correction.max_num_tokens"]), + ("max_batch_size", &["correction.max_batch_size"]), + ("max_kv_tokens", &["correction.max_kv_tokens"]), + ( + "regression_attention_kv_weight", + &["features.attention_kv_weight"], + ), + ( + "regression_prefill_attention_pair_weight", + &["features.prefill_attention_pair_weight"], + ), + ( + "regression_ffn_token_weight", + &["features.ffn_token_weight"], + ), + ( + "regression_ridge_scale", + &["fpm_regression.fit.singular_ridge_scale"], + ), + ] { + if supplied.get(legacy).is_some() { + for path in paths { + insert_explicit(&mut migrated, path, normalized[legacy].clone()); + } + } + } + if supplied.get("bucket_count").is_some() || supplied.get("bucket_shape").is_some() { + // Bucket shape overrides count in the legacy schema. Validate + // those controls here because count is absent from the new schema. + validate_options(&ForwardPassPerfOptions { + bucket_count: options.bucket_count, + bucket_shape: options.bucket_shape, + ..ForwardPassPerfOptions::default() + })?; + let bins = options.bucket_shape.unwrap_or_else(|| { + let n = super::samples::integer_sqrt(options.bucket_count); + [n, n] + }); + for path in [ + "fpm_regression.sampling.bins_per_axis", + "correction.sampling.bins_per_axis", + ] { + insert_explicit(&mut migrated, path, serde_json::json!(bins)); + } + } + merge_explicit(&mut explicit, migrated, "estimator_config")?; + } + let mut interpolation = serde_json::json!({}); + if let Some(method) = interpolation_method { + insert_explicit( + &mut interpolation, + "fpm_interpolation.method", + serde_json::json!(method), + ); + } + if let Some(path) = fpm_parquet_path { + insert_explicit( + &mut interpolation, + "fpm_interpolation.fpm_parquet_path", + serde_json::json!(path), + ); + } + merge_explicit(&mut explicit, interpolation, "estimator_config")?; + let config: Self = serde_json::from_value(explicit) + .map_err(|error| invalid_migration("estimator_config", error))?; + config.validate()?; + Ok(config) + } + pub(crate) fn correction_options(&self) -> ForwardPassPerfOptions { let config = &self.correction; ForwardPassPerfOptions { @@ -308,6 +474,200 @@ impl EstimatorConfig { } } +#[cfg(any(feature = "python", test))] +fn invalid_migration(field: &str, error: serde_json::Error) -> AicError { + AicError::InvalidEngineConfig(format!("invalid {field}: {error}")) +} + +#[cfg(any(feature = "python", test))] +fn explicit_fields( + supplied: &serde_json::Value, + normalized: &serde_json::Value, +) -> serde_json::Value { + match supplied { + serde_json::Value::Object(fields) => serde_json::Value::Object( + fields + .iter() + .map(|(key, value)| { + // Compact serde output may omit an explicit default; keep it supplied. + ( + key.clone(), + explicit_fields(value, normalized.get(key).unwrap_or(value)), + ) + }) + .collect(), + ), + _ => normalized.clone(), + } +} + +#[cfg(any(feature = "python", test))] +fn insert_explicit(root: &mut serde_json::Value, path: &str, value: serde_json::Value) { + if let Some((field, rest)) = path.split_once('.') { + if root.get(field).is_none() { + root[field] = serde_json::json!({}); + } + insert_explicit(&mut root[field], rest, value); + } else { + root[path] = value; + } +} + +#[cfg(any(feature = "python", test))] +fn merge_explicit( + target: &mut serde_json::Value, + incoming: serde_json::Value, + path: &str, +) -> Result<(), AicError> { + if let (Some(target_fields), Some(incoming_fields)) = + (target.as_object_mut(), incoming.as_object()) + { + for (key, value) in incoming_fields { + if let Some(existing) = target_fields.get_mut(key) { + merge_explicit(existing, value.clone(), &format!("{path}.{key}"))?; + } else { + target_fields.insert(key.clone(), value.clone()); + } + } + } else if *target != incoming { + if path == "estimator_config.fpm_interpolation.fpm_parquet_path" + && target + .as_str() + .zip(incoming.as_str()) + .is_some_and(|(saved, supplied)| { + std::path::Path::new(saved) == std::path::Path::new(supplied) + }) + { + return Ok(()); + } + return Err(AicError::InvalidEngineConfig(format!( + "conflicting explicit values for {path}" + ))); + } + Ok(()) +} + +#[cfg(test)] +mod migration_tests { + use super::{EstimatorConfig, FpmInterpolationMethod}; + + #[test] + fn migration_preserves_explicit_compact_fpm_defaults() { + let saved = r#"{"fpm_interpolation": { + "method": "direct", "text_only": false, + "unrecorded_quant_modes": [], "fpm_parquet_path": null + }}"#; + let config = EstimatorConfig::migrate_legacy_inputs(Some(saved), None, None, None).unwrap(); + assert_eq!( + config.fpm_interpolation.method, + FpmInterpolationMethod::Direct + ); + assert!(!config.fpm_interpolation.text_only); + assert!(config.fpm_interpolation.unrecorded_quant_modes.is_empty()); + assert_eq!(config.fpm_interpolation.fpm_parquet_path, None); + } + + #[test] + fn migration_defaults_are_applied_after_explicit_inputs_merge() { + let saved = r#"{ + "fpm_interpolation": {"method": "direct"}, + "correction": {"enabled": false, "sampling": {"max_observations": 128}}, + "fpm_regression": {"sampling": {"max_observations": 128}} + }"#; + let options = r#"{"min_observations": 80, "regression_ridge_scale": 0}"#; + let config = + EstimatorConfig::migrate_legacy_inputs(Some(saved), Some(options), None, None).unwrap(); + assert_eq!( + config.fpm_interpolation.method, + FpmInterpolationMethod::Direct + ); + assert!(!config.correction.enabled); + assert_eq!(config.correction.min_observations, 80); + assert_eq!(config.fpm_regression.min_observations, 80); + assert_eq!(config.fpm_regression.fit.singular_ridge_scale, 0.0); + assert_eq!(config.features.attention_kv_weight, 1.0); + let reloaded: EstimatorConfig = + serde_json::from_str(&serde_json::to_string(&config).unwrap()).unwrap(); + assert_eq!(reloaded, config); + } + + #[test] + fn migration_accepts_equal_numeric_array_and_null_controls() { + let saved = r#"{ + "features": {"attention_kv_weight": 1}, + "correction": {"factor_bounds": {"min": null}}, + "fpm_regression": {"sampling": {"bins_per_axis": [4, 16]}} + }"#; + let options = r#"{ + "regression_attention_kv_weight": 1.0, + "min_faster_correction_factor": null, + "bucket_shape": [4, 16] + }"#; + let config = + EstimatorConfig::migrate_legacy_inputs(Some(saved), Some(options), None, None).unwrap(); + assert_eq!(config.features.attention_kv_weight, 1.0); + assert_eq!(config.correction.factor_bounds.min, None); + assert_eq!(config.correction.sampling.bins_per_axis, [4, 16]); + assert_eq!(config.fpm_regression.sampling.bins_per_axis, [4, 16]); + } + + #[test] + fn migration_compares_parquet_paths_without_filesystem_resolution() { + let saved = r#"{"fpm_interpolation":{"fpm_parquet_path":"/tmp/fpm.parquet"}}"#; + for legacy in [ + "/tmp/fpm.parquet", + "/tmp/./fpm.parquet", + "/tmp//fpm.parquet", + ] { + let config = + EstimatorConfig::migrate_legacy_inputs(Some(saved), None, None, Some(legacy)) + .unwrap(); + assert_eq!( + config.fpm_interpolation.fpm_parquet_path.unwrap().to_str(), + Some("/tmp/fpm.parquet") + ); + } + for legacy in ["/tmp/different.parquet", "/tmp/subdir/../fpm.parquet"] { + let error = + EstimatorConfig::migrate_legacy_inputs(Some(saved), None, None, Some(legacy)) + .unwrap_err(); + assert!( + error.to_string().contains( + "conflicting explicit values for estimator_config.fpm_interpolation.fpm_parquet_path" + ), + "{error}" + ); + } + } + + #[test] + fn interpolation_conflicts_include_explicit_auto_but_not_omission() { + let config = EstimatorConfig::migrate_legacy_inputs( + Some(r#"{"fpm_interpolation": {}}"#), + None, + Some("direct"), + None, + ) + .unwrap(); + assert_eq!( + config.fpm_interpolation.method, + FpmInterpolationMethod::Direct + ); + for (canonical, legacy) in [("auto", "direct"), ("direct", "auto"), ("sol", "direct")] { + let saved = format!(r#"{{"fpm_interpolation": {{"method": "{canonical}"}}}}"#); + let error = + EstimatorConfig::migrate_legacy_inputs(Some(&saved), None, Some(legacy), None) + .unwrap_err(); + assert!( + error.to_string().contains( + "conflicting explicit values for estimator_config.fpm_interpolation.method" + ), + "{error}" + ); + } + } +} + pub(super) fn validate_rebuild_interval(interval: Option) -> Result<(), AicError> { if interval == Some(0) { return Err(AicError::InvalidEngineConfig( diff --git a/crates/core/src/perfmodel/fpm/evidence.rs b/crates/core/src/perfmodel/fpm/evidence.rs new file mode 100644 index 000000000..6f8cce55b --- /dev/null +++ b/crates/core/src/perfmodel/fpm/evidence.rs @@ -0,0 +1,97 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Executed direct-FPM lookups and their forward-pass composition. +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmCoordinates { + pub batch_size: f64, + pub total_prefill_tokens: Option, + pub total_kv_read_tokens: f64, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmMeasurementSupport { + pub coordinates: FpmCoordinates, + pub latency_ms: f64, + /// Product of the linear interpolation weights at this measured point. + pub weight: f64, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmQueryResolution { + ExactLookup, + WithinCurveInterpolation, + CrossKvInterpolation, + CrossBatchInterpolation, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct DirectFpmQueryEvidence { + pub phase: String, + pub model_path: String, + /// True for the mixed-pass baseline: support uses each selected row's + /// measured KV floor, while `query` retains the paired decode request. + pub decode_baseline: bool, + pub query: FpmCoordinates, + pub resolution: FpmQueryResolution, + /// Raw lookup latency, before composition, correction or replay speedup. + pub latency_ms: f64, + pub support: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct FpmRankEstimate { + pub rank: usize, + pub latency_ms: f64, + pub prefill_ms: Option, + pub decode_ms: Option, + pub decode_baseline_ms: Option, + /// For a mixed pass, max(decode_ms - decode_baseline_ms, 0). + pub marginal_decode_ms: Option, + pub queries: Vec, +} + +/// Additive detailed result from the canonical returned model. Native timing +/// reduces ranks by maximum, then applies `correction_factor`. Rank lookup +/// evidence is supplied for whole-model FPM; other estimators leave it empty. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ForwardPassEstimate { + pub latency_ms: Option, + pub native_latency_ms: Option, + pub correction_factor: Option, + /// First rank attaining the positive native maximum; None for empty work. + pub max_rank: Option, + pub ranks: Vec, +} + +/// Identical complete estimates are retained once, with the number of timing +/// invocations that used them. No records are truncated. `latency_scale` is +/// replay's synthetic speedup multiplier and does not alter measured support. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmEstimateEvidence { + pub estimate: ForwardPassEstimate, + pub count: u64, + pub latency_scale: f64, +} + +pub(crate) fn accumulate_fpm_estimates( + target: &mut Vec, + incoming: Vec, +) -> Result<(), &'static str> { + for record in incoming { + if let Some(existing) = target.iter_mut().find(|existing| { + existing.estimate == record.estimate && existing.latency_scale == record.latency_scale + }) { + existing.count = existing + .count + .checked_add(record.count) + .ok_or("FPM evidence invocation count overflow")?; + } else { + target.push(record); + } + } + Ok(()) +} diff --git a/crates/core/src/perfmodel/fpm/mod.rs b/crates/core/src/perfmodel/fpm/mod.rs index 65144f25b..945862f9c 100644 --- a/crates/core/src/perfmodel/fpm/mod.rs +++ b/crates/core/src/perfmodel/fpm/mod.rs @@ -25,6 +25,7 @@ mod config; mod correction; mod estimator; +mod evidence; mod metrics; mod model; mod options; @@ -39,6 +40,8 @@ pub use config::{ ForwardPassSpeculationConfig, }; pub use estimator::*; +pub(crate) use evidence::accumulate_fpm_estimates; +pub use evidence::*; pub(crate) use metrics::validate_forward_pass_metrics; pub use metrics::{FPM_VERSION, ForwardPassMetrics, QueuedRequestMetrics, ScheduledRequestMetrics}; pub use model::{ diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index 52b52d591..4e14e1b67 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -314,9 +314,8 @@ impl ForwardPassPerfModel { /// 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(mut config: ForwardPassPerfModelConfig) -> Result { - config.resolve_prefill_graph_profile()?; - config.validate()?; + pub fn best_available(config: ForwardPassPerfModelConfig) -> Result { + let (mut config, registered) = config.resolve_with_registration()?; if let Some(path) = config .estimator_config .fpm_interpolation @@ -331,6 +330,15 @@ impl ForwardPassPerfModel { let mut failures = Vec::new(); let mut last_error = None; for mode in config.candidate_modes() { + if mode == EstimationMode::OpLevel && registered == Some(false) { + let error = AicError::UnsupportedModel(format!( + "op-level timing requires a registered architecture for {:?}", + config.model, + )); + failures.push(format!("{mode:?}: {error}")); + last_error = Some(error); + continue; + } if config.dcp.is_some_and(|dcp| dcp > 1) && (mode != EstimationMode::FpmInterpolation || config.backend != crate::BackendKind::Vllm) @@ -440,14 +448,7 @@ impl ForwardPassPerfModel { return Ok(Some(0.0)); }; let native = engine.forward_pass_time_ms(metrics_by_rank)?; - let corrected = native - * if self.is_correction_enabled { - corrections - .store(feature.workload_kind) - .correction_factor_for(&feature.x) - } else { - 1.0 - }; + let corrected = native * self.correction_factor(corrections, &feature); Ok(Some(corrected)) } ForwardPassPerfMode::Regression { @@ -467,6 +468,44 @@ impl ForwardPassPerfModel { } } + /// Estimate the same iteration with executed direct-FPM support. The + /// native rank maximum and online correction are explicit; measurement + /// weights describe raw lookup values only. No query state is retained. + pub fn estimate_forward_pass_detailed( + &self, + metrics_by_rank: &[ForwardPassMetrics], + ) -> Result { + self.require_general_forward_api()?; + if let ForwardPassPerfMode::Native { + engine, + corrections, + } = &self.mode + { + let Some(feature) = IterationFeatures::from_metrics(metrics_by_rank)? else { + return Ok(crate::ForwardPassEstimate { + latency_ms: Some(0.0), + native_latency_ms: Some(0.0), + correction_factor: Some(1.0), + max_rank: None, + ranks: Vec::new(), + }); + }; + let mut estimate = engine.forward_pass_estimate(metrics_by_rank, true)?; + let factor = self.correction_factor(corrections, &feature); + estimate.correction_factor = Some(factor); + estimate.latency_ms = estimate.native_latency_ms.map(|native| native * factor); + Ok(estimate) + } else { + Ok(crate::ForwardPassEstimate { + latency_ms: self.estimate_forward_pass_time_ms(metrics_by_rank)?, + native_latency_ms: None, + correction_factor: None, + max_rank: None, + ranks: Vec::new(), + }) + } + } + /// API: /// `model.tune_with_fpms(iterations) -> Result<(), AicError>` /// @@ -672,6 +711,20 @@ impl ForwardPassPerfModel { } } + pub(crate) fn static_prefill_detailed( + &self, + batch_size: u32, + input_tokens: u32, + prefix: u32, + ) -> Result { + self.require_general_forward_api()?; + self.native_engine() + .ok_or_else(|| { + AicError::InvalidEngineConfig("static prefill requires a native estimator".into()) + })? + .predict_prefill_detailed(batch_size, input_tokens, prefix) + } + /// Native operation evidence for one static prefill or decode step. Values /// precede learned online correction; SOL is a comparison only. Whole-model /// estimators cannot provide an operation decomposition and fail explicitly. @@ -751,6 +804,20 @@ impl ForwardPassPerfModel { self.provenance.as_ref() } + fn correction_factor( + &self, + corrections: &WorkloadStores, + feature: &IterationFeatures, + ) -> f64 { + if self.is_correction_enabled { + corrections + .store(feature.workload_kind) + .correction_factor_for(&feature.x) + } else { + 1.0 + } + } + fn correction_factors(&self) -> Vec { match &self.mode { ForwardPassPerfMode::Native { corrections, .. } => corrections.correction_factors(), diff --git a/crates/core/src/perfmodel/memory.rs b/crates/core/src/perfmodel/memory.rs index 66d338a85..3830f3742 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -120,6 +120,8 @@ pub enum EstimateSource { Native, /// Post-weight reservation heuristic (`naive_kv_reservation`, default 80%). NaiveFallback, + /// Explicit rank-local bounds supplied in an FPM profile. + Profile, } /// Non-KV memory components, in bytes. Maps AIC's `_get_memory_usage` dict: @@ -250,12 +252,18 @@ fn fetch_python_estimate( req: &KvCacheEstimateRequest, ) -> Result { use pyo3::prelude::*; + use pyo3::types::PyDictMethods; let engine = &req.engine; Python::with_gil(|py| -> PyResult { let engine_mod = py.import("aisimulate.sdk.memory")?; let kwargs = estimate_kwargs(py, req)?; + if engine.extra.contains_key("fpm_profile") { + for (name, mode) in crate::py::fpm_profile_quantization(py, engine)? { + kwargs.set_item(name, mode)?; + } + } let out = engine_mod.call_method( "estimate_kv_cache", @@ -311,6 +319,11 @@ fn estimate_kwargs<'py>( kwargs.set_item("attention_dp_size", parallel.attention_dp_size.unwrap_or(1))?; kwargs.set_item("moe_tp_size", parallel.moe_tp_size)?; kwargs.set_item("moe_ep_size", parallel.moe_ep_size)?; + kwargs.set_item("cp_size", parallel.cp_size.unwrap_or(1))?; + kwargs.set_item( + "fpm_profile", + engine.extra.get("fpm_profile").map(String::as_str), + )?; kwargs.set_item( "gemm_quant_mode", quant.weight_dtype.as_ref().map(dtype_str), @@ -364,6 +377,7 @@ fn estimate_from_dict( let source = match out.get_item("source")?.extract::()?.as_str() { "native" => EstimateSource::Native, "naive_fallback" => EstimateSource::NaiveFallback, + "profile" => EstimateSource::Profile, other => { return Err(PyValueError::new_err(format!( "estimate_kv_cache returned unknown source {other:?}" @@ -552,6 +566,63 @@ mod tests { } } + #[cfg(feature = "embed-python")] + #[test] + fn profile_memory_preserves_exact_quantization_and_rejects_wire_conflicts() { + use crate::DataType; + + pyo3::prepare_freethreaded_python(); + for (gemm_mode, weight_dtype, moe_mode, moe_dtype) in [ + ("fp8_block", DataType::Fp8Block, "nvfp4", DataType::Nvfp4), + ("int8_wo", DataType::Int8, "nvfp4", DataType::Nvfp4), + ("int4_wo", DataType::Int4, "nvfp4", DataType::Nvfp4), + ("sq", DataType::Int8, "nvfp4", DataType::Nvfp4), + ("fp8_ootb", DataType::Fp8, "nvfp4", DataType::Nvfp4), + ("nvfp4", DataType::Nvfp4, "int4_wo", DataType::Int4), + ] { + let mut req = python_forwarding_request(); + req.cuda_graph_reserved_bytes = 0; + req.gpu_memory_capacity_bytes_override = Some(1_000); + req.engine.quantization.weight_dtype = Some(weight_dtype); + req.engine.quantization.moe_dtype = Some(moe_dtype); + req.engine.quantization.activation_dtype = Some(DataType::Fp8); + req.engine.quantization.kv_cache_dtype = Some(DataType::Fp8); + req.engine.extra.insert("fpm_profile".into(), serde_json::json!({ + "schema_version": 1, + "model": "test-model", + "model_revision": "native-memory-fixture-v1", + "architecture": "UnregisteredMemoryFixture", + "context_length": 8192, + "num_experts": 0, + "provenance": "Synthetic native transport fixture; no silicon qualification", + "deployments": [{ + "system": "test-system", "backend": "vllm", "backend_version": "test-version", + "tp": 1, "pp": 1, "dp": 1, "moe_tp": 1, "moe_ep": 1, "cp": 1, + "gemm_quant_mode": gemm_mode, "moe_quant_mode": moe_mode, + "fmha_quant_mode": "fp8", "comm_quant_mode": "half", "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 100, "activations_bytes": 20, "runtime_overhead_bytes": 30, + "comm_overhead_bytes": 50, "kv_bytes_per_token": 10, "cache_layout": "linear", + "max_num_tokens": 8192, "max_batch_size": 256, + "provenance": "Declared rank-local arithmetic fixture; excludes CUDA graphs" + } + }] + }).to_string()); + let estimate = + estimate_kv_cache(req.clone()).expect("profile precision must survive the wire"); + assert_eq!(estimate.source, EstimateSource::Profile); + assert_eq!(estimate.total_kv_size_bytes, 600); + assert_eq!(estimate.total_kv_size_tokens, 60); + + req.engine.quantization.weight_dtype = Some(DataType::Bfloat16); + let error = estimate_kv_cache(req).expect_err("a different wire dtype must fail"); + assert!( + error.to_string().contains("FPM profile identity conflict"), + "{error}" + ); + } + } + #[cfg(feature = "python")] #[test] fn py_dict_forwarding_and_parsing_are_deterministic() { @@ -559,9 +630,20 @@ mod tests { use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods}; pyo3::prepare_freethreaded_python(); - let req = python_forwarding_request(); + let mut req = python_forwarding_request(); + req.engine.extra.insert( + "fpm_profile".into(), + serde_json::json!({"model": "test-model"}).to_string(), + ); Python::with_gil(|py| -> PyResult<()> { let kwargs = estimate_kwargs(py, &req)?; + assert_eq!( + kwargs + .get_item("fpm_profile")? + .unwrap() + .extract::()?, + r#"{"model":"test-model"}"# + ); assert_eq!( kwargs .get_item("cuda_graph_reserved_bytes")? @@ -593,6 +675,11 @@ mod tests { .cuda_graph_reserved_bytes, req.cuda_graph_reserved_bytes ); + out.set_item("source", "profile")?; + assert_eq!( + estimate_from_dict(out.as_any())?.source, + EstimateSource::Profile + ); Ok(()) }) .expect("deterministic PyDict round-trip must succeed"); diff --git a/crates/core/src/perfmodel/mod.rs b/crates/core/src/perfmodel/mod.rs index 05870ef39..f89fbb72f 100644 --- a/crates/core/src/perfmodel/mod.rs +++ b/crates/core/src/perfmodel/mod.rs @@ -54,8 +54,8 @@ pub use fpm::{ FpmRegressionConfig, RegressionFeatureWeights, SamplingConfig, }; pub use fpm::{ - CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, OpLevelConfig, - RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, + CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, FpmInterpolationMethod, + OpLevelConfig, RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, }; // Forward-pass metrics telemetry types and schema version, plus the // crate-internal validation helper. Re-exported at the crate root so existing @@ -104,3 +104,8 @@ pub(crate) fn repo_relative(rel: &str) -> Option { pub fn register_python(module: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult<()> { py::register(module) } + +pub use fpm::{ + DirectFpmQueryEvidence, ForwardPassEstimate, FpmCoordinates, FpmEstimateEvidence, + FpmMeasurementSupport, FpmQueryResolution, FpmRankEstimate, +}; diff --git a/crates/core/src/perfmodel/operators/base.rs b/crates/core/src/perfmodel/operators/base.rs index 6e7214376..585a5c3e2 100644 --- a/crates/core/src/perfmodel/operators/base.rs +++ b/crates/core/src/perfmodel/operators/base.rs @@ -183,6 +183,8 @@ pub struct PerformanceResult { pub source: Source, pub sol: Option, pub moe_comm_fallbacks: MoeCommFallbacks, + /// Raw direct-FPM leaves; composition never rewrites their measured support. + pub fpm_queries: Option>>, } impl PerformanceResult { @@ -193,6 +195,7 @@ impl PerformanceResult { source, sol: None, moe_comm_fallbacks: MoeCommFallbacks::default(), + fpm_queries: None, } } @@ -203,6 +206,7 @@ impl PerformanceResult { source, sol: None, moe_comm_fallbacks: MoeCommFallbacks::default(), + fpm_queries: None, } } @@ -215,6 +219,7 @@ impl PerformanceResult { source: Source::Sol, sol: Some(components), moe_comm_fallbacks: MoeCommFallbacks::default(), + fpm_queries: None, } } @@ -261,6 +266,7 @@ impl PerformanceResult { mem_ms: c.mem_ms * factor, }), moe_comm_fallbacks: self.moe_comm_fallbacks, + fpm_queries: self.fpm_queries, } } @@ -269,6 +275,12 @@ impl PerformanceResult { /// AND energy both 0.0) is a source-neutral identity — the other /// side's tag survives, mirroring Python's zero-identity rule. pub fn plus(self, other: PerformanceResult) -> Self { + let mut fpm_queries = self.fpm_queries; + if let Some(other_queries) = other.fpm_queries { + fpm_queries + .get_or_insert_with(Default::default) + .extend(*other_queries); + } let mut moe_comm_fallbacks = self.moe_comm_fallbacks; moe_comm_fallbacks.extend(other.moe_comm_fallbacks); let (source, sol) = if self.latency_ms == 0.0 && self.energy_wms == 0.0 { @@ -294,6 +306,7 @@ impl PerformanceResult { source, sol, moe_comm_fallbacks, + fpm_queries, } } @@ -310,6 +323,7 @@ impl PerformanceResult { mem_ms: c.mem_ms.max(0.0), }), moe_comm_fallbacks: self.moe_comm_fallbacks, + fpm_queries: self.fpm_queries, } } } diff --git a/crates/core/src/perfmodel/operators/fpm_forward.rs b/crates/core/src/perfmodel/operators/fpm_forward.rs index 12061213a..8832efdcd 100644 --- a/crates/core/src/perfmodel/operators/fpm_forward.rs +++ b/crates/core/src/perfmodel/operators/fpm_forward.rs @@ -15,23 +15,31 @@ //! - prefill coords `(B, B*s, B*prefix)` — `s` is NEW prefill tokens per //! request, `prefix` the past-KV per request; //! - decode coords `(B, B*s)` — one new token per request, `s` the per-request -//! KV length at this decode step; -//! - hard per-axis domain gate BEFORE interpolation (FPM never extrapolates); -//! - ScatteredSites resolution with `own_curve_coverage_fallback=true` and -//! `max_site_distance=2.0`; prefill additionally admits sites within 32 raw -//! KV tokens while retaining the normal log2 gate on batch, so small -//! block-aligned coordinates around zero remain connected. +//! KV length at this decode step. //! -//! The SOL roofline anchoring the interpolation is the model's ORIGINAL -//! op-level list (carried as `sol_ops` on the wire) queried in SOL mode with -//! Python's coordinate back-mapping — see [`sol_total`]. +//! With `interpolation="sol"` (the default), the roofline anchoring transfer +//! is the model's ORIGINAL op-level list (`sol_ops` on the wire) queried in +//! SOL mode with Python's coordinate back-mapping — see [`sol_total`]. +//! SOL applies a hard per-axis domain gate before ScatteredSites resolution +//! with `own_curve_coverage_fallback=true` and `max_site_distance=2.0`; +//! prefill additionally admits sites within 32 raw KV tokens while retaining +//! the log2 gate on batch, so small block-aligned coordinates remain connected. //! -//! NOT the crate's `src/fpm/` (`ForwardPassPerfModel`) module: that is the -//! online-tuning model over Dynamo ForwardPassMetrics telemetry, an unrelated -//! concept that shares the "FPM" abbreviation. +//! `interpolation="direct"` instead uses raw linear interpolation between +//! genuine measured curves. Prefill uses the nearest same-batch lower and +//! upper KV sites whose token curves cover the query, without a KV distance +//! limit; decode stays inside its inferred capture regime. It never +//! evaluates `sol_ops`, clamps an unsupported batch, or uses fabricated KV. +//! +//! The canonical `ForwardPassPerfModel` constructor selects this native +//! operator for `estimation_mode="fpm_interpolation"` and pins its method. + +use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; +use crate::{DirectFpmQueryEvidence, FpmCoordinates, FpmMeasurementSupport, FpmQueryResolution}; + use crate::common::error::AicError; use crate::operators::op::{Op, RuntimeContext}; use crate::operators::{PerformanceResult, Source}; @@ -57,6 +65,16 @@ impl FpmPhase { } } +/// Timing interpolation policy. Existing specs retain the analytical +/// roofline; direct interpolation reads measured times without querying it. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmInterpolation { + #[default] + Sol, + Direct, +} + /// One whole-model forward pass for a single phase. /// /// `match_identity` is the 19-string base cell identity (legacy schema-6 requests use 15) in @@ -91,6 +109,8 @@ pub struct FpmForwardOp { /// separate op-level draft ops). #[serde(default = "default_verify_width")] pub verify_width: u32, + #[serde(default)] + pub interpolation: FpmInterpolation, pub sol_ops: Vec, /// Present only for an explicit table selector. Preserves the model's mode /// before match_identity[2] was selected; it is NOT observed runtime precision. @@ -226,6 +246,14 @@ impl FpmForwardOp { ))); } let cell = self.select_cell(db)?; + if self.interpolation == FpmInterpolation::Direct { + return Ok(cell + .direct_decode + .values() + .filter_map(|curve| curve.keys().next_back()) + .max() + .copied()); + } Ok(cell.decode_domain.as_ref().map(|domain| domain[1].1)) } @@ -257,6 +285,9 @@ impl FpmForwardOp { ))); } let cell = self.select_cell(db)?; + if self.interpolation == FpmInterpolation::Direct { + return self.resolve_direct_decode(cell, &[batch_size as f64, total_kv], true); + } let Some(domain) = cell.decode_domain else { return Err(data_err(format!( "FPM cell {:?} has no decode rows (model_path={:?}).", @@ -336,6 +367,12 @@ impl FpmForwardOp { cell: &FpmForwardCell, coords: &[f64], ) -> Result { + if self.interpolation == FpmInterpolation::Direct { + return match self.phase { + FpmPhase::Prefill => self.resolve_direct_prefill(cell, coords), + FpmPhase::Decode => self.resolve_direct_decode(cell, coords, false), + }; + } // Data-certified prefill batch clamp (mirrors Python _resolve): the // regime coordinate is the token TOTAL, which stays untouched — the // clamped query prices the same side of the capture cliff and is a @@ -467,6 +504,251 @@ impl FpmForwardOp { Ok(PerformanceResult::new(latency, Source::Silicon)) } + fn resolve_direct_prefill( + &self, + cell: &FpmForwardCell, + coords: &[f64], + ) -> Result { + let (batch, tokens, kv) = (coords[0], coords[1], coords[2]); + if !batch.is_finite() || batch.fract() != 0.0 || batch < 1.0 || batch > u32::MAX as f64 { + return Err(self.direct_coverage_err( + cell, + coords, + "batch_size must be a positive integer", + )); + } + if let Some(value) = cell + .direct_prefill + .get(&(batch as u32, kv as u32)) + .filter(|_| kv == (kv as u32) as f64) + .and_then(|curve| direct_curve_value(curve, tokens)) + { + let support = direct_curve_support( + &cell.direct_prefill[&(batch as u32, kv as u32)], + tokens, + |token| fpm_coordinates(self.phase, &[batch, token, kv]), + ); + let resolution = if support.len() == 1 { + FpmQueryResolution::ExactLookup + } else { + FpmQueryResolution::WithinCurveInterpolation + }; + return self.direct_result(cell, coords, value, resolution, support, false); + } + let covered: Vec<(u32, f64)> = cell + .direct_prefill + .range((batch as u32, 0)..=(batch as u32, u32::MAX)) + .filter_map(|(&(.., site_kv), curve)| { + direct_curve_value(curve, tokens).map(|value| (site_kv, value)) + }) + .collect(); + // The map orders sites by KV. Direct interpolation uses the nearest + // token-covered site on each side, with no SOL site-distance guard. + let lower = covered + .iter() + .rev() + .find(|&&(site_kv, _)| (site_kv as f64) < kv); + let upper = covered.iter().find(|&&(site_kv, _)| (site_kv as f64) > kv); + let Some((&(lo, lo_ms), &(hi, hi_ms))) = lower.zip(upper) else { + return Err(self.direct_coverage_err( + cell, + coords, + "no same-batch, two-sided KV bracket whose token curves cover this query", + )); + }; + let weight = (kv - lo as f64) / (hi as f64 - lo as f64); + let latency = lo_ms + (hi_ms - lo_ms) * weight; + let mut support = Vec::new(); + for (site_kv, site_weight) in [(lo, 1.0 - weight), (hi, weight)] { + support.extend( + direct_curve_support( + &cell.direct_prefill[&(batch as u32, site_kv)], + tokens, + |token| fpm_coordinates(self.phase, &[batch, token, site_kv as f64]), + ) + .into_iter() + .map(|mut point| { + point.weight *= site_weight; + point + }), + ); + } + self.direct_result( + cell, + coords, + latency, + FpmQueryResolution::CrossKvInterpolation, + support, + false, + ) + } + + fn direct_result( + &self, + cell: &FpmForwardCell, + coords: &[f64], + latency: f64, + resolution: FpmQueryResolution, + support: Vec, + decode_baseline: bool, + ) -> Result { + if !latency.is_finite() || latency <= 0.0 { + return Err(self.direct_coverage_err( + cell, + coords, + &format!("interpolation produced an invalid latency ({latency})"), + )); + } + let mut result = PerformanceResult::new(latency, Source::Silicon); + result.fpm_queries = Some(Box::new(vec![DirectFpmQueryEvidence { + phase: self.phase.as_str().into(), + model_path: cell.model_path.clone(), + decode_baseline, + query: fpm_coordinates(self.phase, coords), + resolution, + latency_ms: latency, + support, + }])); + Ok(result) + } + + fn direct_coverage_err(&self, cell: &FpmForwardCell, coords: &[f64], reason: &str) -> AicError { + let axes: &[&str] = match self.phase { + FpmPhase::Prefill => &FPM_PREFILL_AXES, + FpmPhase::Decode => &FPM_DECODE_AXES, + }; + let query = axes + .iter() + .zip(coords) + .map(|(axis, value)| format!("{axis}={value}")) + .collect::>() + .join(", "); + data_err(format!( + "FPM direct {} query ({query}) is unsupported for model_path={:?}: {reason}. \ + Collect matching genuine FPM points; direct interpolation never uses SOL or extrapolation.", + self.phase.as_str(), + cell.model_path + )) + } + + fn resolve_direct_decode( + &self, + cell: &FpmForwardCell, + coords: &[f64], + baseline: bool, + ) -> Result { + let (lo_row, hi_row) = self.direct_decode_rows(cell, coords)?; + let row_value = |row: u32| { + let curve = &cell.direct_decode[&row]; + if baseline { + *curve.values().next().expect("nonempty measured curve") + } else { + direct_curve_value(curve, coords[1]).expect("row coverage checked") + } + }; + let lo_ms = row_value(lo_row); + let latency = if lo_row == hi_row { + lo_ms + } else { + let weight = (coords[0] - lo_row as f64) / (hi_row as f64 - lo_row as f64); + lo_ms + (row_value(hi_row) - lo_ms) * weight + }; + let rows = if lo_row == hi_row { + vec![(lo_row, 1.0)] + } else { + let weight = (coords[0] - lo_row as f64) / (hi_row as f64 - lo_row as f64); + vec![(lo_row, 1.0 - weight), (hi_row, weight)] + }; + let mut support = Vec::new(); + for (row, row_weight) in rows { + let curve = &cell.direct_decode[&row]; + let kv = if baseline { + *curve.keys().next().expect("nonempty measured curve") as f64 + } else { + coords[1] + }; + support.extend( + direct_curve_support(curve, kv, |kv| { + fpm_coordinates(self.phase, &[row as f64, kv]) + }) + .into_iter() + .map(|mut point| { + point.weight *= row_weight; + point + }), + ); + } + let resolution = if lo_row != hi_row { + FpmQueryResolution::CrossBatchInterpolation + } else if support.len() == 1 { + FpmQueryResolution::ExactLookup + } else { + FpmQueryResolution::WithinCurveInterpolation + }; + self.direct_result(cell, coords, latency, resolution, support, baseline) + } + + /// One shared coverage decision for decode queries and their mixed-pass + /// baseline. Neither can include a curve the other cannot evaluate. + fn direct_decode_rows( + &self, + cell: &FpmForwardCell, + coords: &[f64], + ) -> Result<(u32, u32), AicError> { + let (batch, kv) = (coords[0], coords[1]); + if !batch.is_finite() || batch.fract() != 0.0 || batch < 1.0 || batch > u32::MAX as f64 { + return Err(self.direct_coverage_err( + cell, + coords, + "batch_size must be a positive integer", + )); + } + if cell + .direct_decode + .get(&(batch as u32)) + .and_then(|curve| direct_curve_value(curve, kv)) + .is_some() + { + return Ok((batch as u32, batch as u32)); + } + let regime = |position: f64| { + cell.direct_decode_rungs + .partition_point(|&r| (r as f64) < position) + }; + let query_regime = regime(batch); + let mut covered = Vec::new(); + for (&row, curve) in &cell.direct_decode { + let distance = ((row as f64).log2() - batch.log2()).abs(); + if row as f64 != batch + && distance <= 2.0 + && regime(row as f64) == query_regime + && direct_curve_value(curve, kv).is_some() + { + covered.push((distance, row)); + } + } + covered.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1))); + covered.truncate(4); + let lower = covered + .iter() + .map(|&(_, row)| row) + .filter(|&row| (row as f64) < batch) + .max(); + let upper = covered + .iter() + .map(|&(_, row)| row) + .filter(|&row| (row as f64) > batch) + .min(); + if let Some(rows) = lower.zip(upper) { + return Ok(rows); + } + Err(self.direct_coverage_err( + cell, + coords, + "no two-sided covered batch bracket within one capture regime", + )) + } + /// Resolve an OFF-LATTICE decode batch by its segment bracket (mirrors /// Python `_decode_bracket`): the engine pads decode batches to capture /// rungs — every rung marked by an (x, x+1) row pair — so a query @@ -563,6 +845,67 @@ impl FpmForwardOp { } } +fn fpm_coordinates(phase: FpmPhase, coords: &[f64]) -> FpmCoordinates { + FpmCoordinates { + batch_size: coords[0], + total_prefill_tokens: (phase == FpmPhase::Prefill).then(|| coords[1]), + total_kv_read_tokens: coords[coords.len() - 1], + } +} + +/// Read support only after coverage and latency have been resolved. Keep this +/// separate from the latency arithmetic so existing evaluation order is intact. +fn direct_curve_support( + curve: &BTreeMap, + coordinate: f64, + coordinates: impl Fn(f64) -> FpmCoordinates, +) -> Vec { + let (&lo, &lo_ms) = curve + .range(..=(coordinate as u32)) + .next_back() + .expect("covered curve"); + if coordinate == lo as f64 { + return vec![FpmMeasurementSupport { + coordinates: coordinates(lo as f64), + latency_ms: lo_ms, + weight: 1.0, + }]; + } + let (&hi, &hi_ms) = curve + .range((std::ops::Bound::Excluded(lo), std::ops::Bound::Unbounded)) + .next() + .expect("covered curve"); + let weight = (coordinate - lo as f64) / (hi as f64 - lo as f64); + vec![ + FpmMeasurementSupport { + coordinates: coordinates(lo as f64), + latency_ms: lo_ms, + weight: 1.0 - weight, + }, + FpmMeasurementSupport { + coordinates: coordinates(hi as f64), + latency_ms: hi_ms, + weight, + }, + ] +} + +/// Exact lookup or raw linear interpolation strictly inside one measured +/// curve. There is deliberately no callback, boundary hold, or SOL input. +fn direct_curve_value(curve: &BTreeMap, coordinate: f64) -> Option { + if !coordinate.is_finite() || coordinate < 0.0 || coordinate > u32::MAX as f64 { + return None; + } + let (&lo, &lo_ms) = curve.range(..=(coordinate as u32)).next_back()?; + if coordinate == lo as f64 { + return Some(lo_ms); + } + let (&hi, &hi_ms) = curve + .range((std::ops::Bound::Excluded(lo), std::ops::Bound::Unbounded)) + .next()?; + Some(lo_ms + (hi_ms - lo_ms) * (coordinate - lo as f64) / (hi as f64 - lo as f64)) +} + /// FPM interpolation configs: prefill sites `(batch, kv)` own the new-token /// curve; decode sites `(batch,)` own the KV curve. Prefill adds a raw-KV /// fallback around zero to the shared `ScatteredSites` defaults. @@ -686,6 +1029,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + interpolation: FpmInterpolation::Sol, original_fmha_quant_mode: None, // Empty sol_ops: exact hits and in-curve lerps never call SOL. sol_ops: vec![], @@ -712,6 +1056,439 @@ mod tests { } } + #[test] + fn direct_query_evidence_preserves_support_and_nested_weights() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let mut rows = Vec::new(); + for (tokens, kv, latency_ms) in [(4, 0, 10.0), (12, 0, 18.0), (4, 8, 30.0), (12, 8, 46.0)] { + rows.push(RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: tokens, + total_kv_read_tokens: kv, + latency_ms, + ..Default::default() + }); + } + for (batch_size, kv, latency_ms) in + [(2, 4, 10.0), (2, 20, 26.0), (4, 8, 20.0), (4, 24, 52.0)] + { + rows.push(RowSpec { + batch_size, + total_kv_read_tokens: kv, + latency_ms, + ..Default::default() + }); + } + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &rows); + let db = db_with_pair(tmp.path()); + let mut prefill = op(FpmPhase::Prefill); + prefill.interpolation = FpmInterpolation::Direct; + // Exact leaf, quarter along the token curve, then one-quarter along KV: + // 12 + (34 - 12) / 4 = 17.5 ms. Combined weights are 9/16,3/16,3/16,1/16. + for (query, latency, resolution, weights) in [ + ( + [1.0, 4.0, 0.0], + 10.0, + FpmQueryResolution::ExactLookup, + vec![1.0], + ), + ( + [1.0, 6.0, 0.0], + 12.0, + FpmQueryResolution::WithinCurveInterpolation, + vec![0.75, 0.25], + ), + ( + [1.0, 6.0, 2.0], + 17.5, + FpmQueryResolution::CrossKvInterpolation, + vec![0.5625, 0.1875, 0.1875, 0.0625], + ), + ] { + let result = prefill.query_totals(&db, &query).unwrap(); + assert_eq!(result.latency_ms, latency); + assert_eq!(result.source, Source::Silicon); + let evidence = &result.fpm_queries.as_ref().unwrap()[0]; + assert_eq!(evidence.resolution, resolution); + assert_eq!( + evidence + .support + .iter() + .map(|point| point.weight) + .collect::>(), + weights + ); + assert_eq!(evidence.query, fpm_coordinates(FpmPhase::Prefill, &query)); + } + let mut decode = op(FpmPhase::Decode); + decode.interpolation = FpmInterpolation::Direct; + let result = decode.query_totals(&db, &[3.0, 12.0]).unwrap(); + assert_eq!(result.latency_ms, 23.0); // midpoint of 18 and 28 ms + let evidence = &result.fpm_queries.as_ref().unwrap()[0]; + assert_eq!( + evidence.resolution, + FpmQueryResolution::CrossBatchInterpolation + ); + assert_eq!( + evidence + .support + .iter() + .map(|p| ( + p.coordinates.batch_size, + p.coordinates.total_kv_read_tokens, + p.latency_ms, + p.weight + )) + .collect::>(), + vec![ + (2.0, 4.0, 10.0, 0.25), + (2.0, 20.0, 26.0, 0.25), + (4.0, 8.0, 20.0, 0.375), + (4.0, 24.0, 52.0, 0.125) + ] + ); + let baseline = decode.query_pass_baseline(&db, 3, 12.0).unwrap(); + assert_eq!(baseline.latency_ms, 15.0); + let evidence = &baseline.fpm_queries.as_ref().unwrap()[0]; + assert!(evidence.decode_baseline); + assert_eq!(evidence.query.total_kv_read_tokens, 12.0); + assert_eq!( + evidence + .support + .iter() + .map(|p| (p.coordinates.total_kv_read_tokens, p.weight)) + .collect::>(), + vec![(4.0, 0.5), (8.0, 0.5)] + ); + assert!(decode.query_totals(&db, &[3.0, 30.0]).is_err()); + } + + #[test] + fn direct_prefill_uses_nearest_token_covered_kv_sites() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let rows = [ + (100, 0, 2.0), + (200, 0, 4.0), + (100, 64, 10.0), + (200, 64, 20.0), + (100, 80, 999.0), + (100, 112, 999.0), + (100, 128, 30.0), + (200, 128, 50.0), + (100, 256, 100.0), + (200, 256, 200.0), + ] + .into_iter() + .map(|(tokens, kv, latency)| RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: tokens, + total_kv_read_tokens: kv, + latency_ms: latency, + ..RowSpec::default() + }) + .collect::>(); + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &rows); + let db = db_with_pair(tmp.path()); + let mut direct = op(FpmPhase::Prefill); + direct.interpolation = FpmInterpolation::Direct; + // T=150 gives 15 ms at KV=64 and 40 ms at KV=128. KV=96 is + // their midpoint; closer sites at KV=80/112 do not cover T=150. + assert_eq!( + direct + .query_totals(&db, &[1.0, 150.0, 96.0]) + .unwrap() + .latency_ms, + 27.5 + ); + } + + #[test] + fn direct_prefill_uses_wider_two_sided_kv_brackets_without_sol() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let rows = [ + (100, 0, 10.0), + (200, 0, 20.0), + (100, 2048, 999.0), + (100, 4096, 30.0), + (200, 4096, 40.0), + ] + .into_iter() + .map(|(tokens, kv, latency)| RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: tokens, + total_kv_read_tokens: kv, + latency_ms: latency, + ..RowSpec::default() + }) + .collect::>(); + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &rows); + let db = db_with_pair(tmp.path()); + let mut spec = serde_json::to_value(op(FpmPhase::Prefill)).unwrap(); + spec["interpolation"] = serde_json::json!("direct"); + let mut direct: FpmForwardOp = serde_json::from_value(spec).unwrap(); + direct.sol_ops = vec![Op::MlaBmm(crate::operators::MlaBmmOp { + name: "unsupported_sol_sentinel".into(), + scale_factor: 1.0, + num_heads: 128, + is_pre: true, + quant_mode: crate::common::enums::GemmQuantMode::Bfloat16, + })]; + // KV=0 is outside the SOL two-octave/32-token site-distance guard. + // First evaluate the two token curves at T=150, then blend raw KV. + let result = direct.query_totals(&db, &[1.0, 150.0, 2048.0]).unwrap(); + assert_eq!(result.latency_ms, 25.0); + assert_eq!(result.source, Source::Silicon); + // The own KV curve does not cover T=150, but its exact T=100 leaf + // still takes precedence over interpolation between other KV sites. + assert_eq!( + direct + .query_totals(&db, &[1.0, 100.0, 2048.0]) + .unwrap() + .latency_ms, + 999.0 + ); + } + + #[test] + fn direct_decode_excludes_healed_fake_fallback_queries_and_baselines() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let rows = [ + (64, 10.0, "real_kv"), + (128, 20.0, "real_kv"), + (256, 99.0, "fake_fallback"), + ] + .into_iter() + .map(|(kv, latency, seed)| RowSpec { + batch_size: 8, + total_kv_read_tokens: kv, + latency_ms: latency, + kv_seed_regime: Some(seed), + ..RowSpec::default() + }) + .collect::>(); + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &rows); + let db = db_with_pair(tmp.path()); + let mut direct = op(FpmPhase::Decode); + direct.interpolation = FpmInterpolation::Direct; + // Legacy SOL keeps its in-memory fake-row healing semantics. + assert_eq!( + op(FpmPhase::Decode) + .query_totals(&db, &[8.0, 256.0]) + .unwrap() + .latency_ms, + 40.0 + ); + for kv in [192.0, 256.0] { + let error = direct.query_totals(&db, &[8.0, kv]).unwrap_err(); + assert!(error.to_string().contains("genuine"), "{error}"); + assert!(direct.query_pass_baseline(&db, 8, kv).is_err()); + } + assert_eq!(direct.decode_kv_ceiling(&db).unwrap(), Some(128)); + assert_eq!( + direct.query_totals(&db, &[8.0, 96.0]).unwrap().latency_ms, + 15.0 + ); + assert_eq!( + direct.query_pass_baseline(&db, 8, 96.0).unwrap().latency_ms, + 10.0 + ); + } + + #[test] + fn direct_decode_brackets_within_capture_regime_and_shares_baseline_coverage() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let rows = [ + (8, 8, 1000.0), + (8, 192, 1000.0), + (9, 64, 9.0), + (9, 192, 29.0), + (16, 64, 16.0), + (16, 128, 26.0), + (17, 64, 2000.0), + (17, 192, 2000.0), + ] + .into_iter() + .map(|(batch, kv, latency)| RowSpec { + batch_size: batch, + total_kv_read_tokens: kv, + latency_ms: latency, + kv_seed_regime: Some("real_kv"), + ..RowSpec::default() + }) + .collect::>(); + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &rows); + let db = db_with_pair(tmp.path()); + let mut direct = op(FpmPhase::Decode); + direct.interpolation = FpmInterpolation::Direct; + // The same padded graph owns rows 9 and 16; neighboring rows 8 and + // 17 deliberately have very different measured costs. + assert_eq!( + direct.query_totals(&db, &[12.0, 96.0]).unwrap().latency_ms, + 17.0 + ); + assert_eq!( + direct + .query_pass_baseline(&db, 12, 96.0) + .unwrap() + .latency_ms, + 12.0 + ); + for kv in [32.0, 160.0] { + // At 160 only row 9 covers the query. Direct cannot inherit the + // legacy one-sided bracket hold or fabricate a mixed baseline. + assert!(direct.query_totals(&db, &[12.0, kv]).is_err()); + assert!(direct.query_pass_baseline(&db, 12, kv).is_err()); + } + } + + #[test] + fn direct_prefill_preserves_own_curves_and_rejects_missing_batches_and_boundaries() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let mut rows = Vec::new(); + for batch in [1, 2, 4] { + for (tokens, latency) in [(100, 10.0), (200, 20.0), (300, 30.0)] { + for (kv, bump) in [(0, 0.0), (64, 4.0), (128, 8.0)] { + rows.push(RowSpec { + workload_kind: "prefill", + batch_size: batch, + total_prefill_tokens: tokens, + total_kv_read_tokens: kv, + latency_ms: latency + bump, + ..RowSpec::default() + }); + } + } + } + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &rows); + let db = db_with_pair(tmp.path()); + let mut direct = op(FpmPhase::Prefill); + direct.interpolation = FpmInterpolation::Direct; + assert_eq!( + direct + .query_totals(&db, &[1.0, 100.0, 0.0]) + .unwrap() + .latency_ms, + 10.0 + ); + assert_eq!( + direct + .query_totals(&db, &[1.0, 150.0, 0.0]) + .unwrap() + .latency_ms, + 15.0 + ); + assert_eq!( + direct + .query_totals(&db, &[1.0, 150.0, 96.0]) + .unwrap() + .latency_ms, + 21.0 + ); + // This fixture certifies the legacy high-batch clamp. It remains a + // valid registered/SOL behavior but is not a measured direct point. + assert!( + op(FpmPhase::Prefill) + .query_totals(&db, &[16.0, 200.0, 0.0]) + .is_ok() + ); + for coords in [ + [3.0, 200.0, 0.0], + [16.0, 200.0, 0.0], + [1.0, 400.0, 32.0], + [1.0, 150.0, -1.0], + [1.0, 150.0, 256.0], + [1.5, 150.0, 64.0], + [1.0, f64::NAN, 64.0], + [1.0, 150.0, f64::INFINITY], + ] { + let error = direct.query_totals(&db, &coords).unwrap_err().to_string(); + assert!(error.contains("FPM direct prefill"), "{error}"); + assert!( + error.contains("Collect matching genuine FPM points"), + "{error}" + ); + } + } + + #[test] + fn omitted_interpolation_keeps_existing_sol_policy() { + let mut spec = serde_json::to_value(op(FpmPhase::Prefill)).unwrap(); + spec.as_object_mut().unwrap().remove("interpolation"); + let decoded: FpmForwardOp = serde_json::from_value(spec.clone()).unwrap(); + assert_eq!(decoded.interpolation, FpmInterpolation::Sol); + spec["interpolation"] = serde_json::json!("unknown"); + assert!(serde_json::from_value::(spec).is_err()); + } + + #[test] + fn direct_rejects_a_fake_only_cell_with_or_without_legacy_healing() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let tmp = tempfile::tempdir().unwrap(); + write_pair( + tmp.path(), + &[RowSpec { + kv_seed_regime: Some("fake_fallback"), + ..RowSpec::default() + }], + ); + for replace in [false, true] { + let mut db = db_with_pair(tmp.path()); + db.set_fpm_forward_for_test( + crate::perf_database::FpmForwardTable::new_with_replacement( + tmp.path().to_path_buf(), + "b200_sxm", + "vllm", + "0.25.1", + replace, + ), + ); + let mut direct = op(FpmPhase::Decode); + direct.interpolation = FpmInterpolation::Direct; + assert!( + op(FpmPhase::Decode) + .query_totals(&db, &[8.0, 4096.0]) + .is_ok() + ); + assert!(direct.query_totals(&db, &[8.0, 4096.0]).is_err()); + assert!(direct.query_pass_baseline(&db, 8, 4096.0).is_err()); + assert_eq!(direct.decode_kv_ceiling(&db).unwrap(), None); + } + } + + #[test] + fn direct_reports_nonfinite_interpolation_as_a_data_error() { + use crate::perf_database::fpm_forward::tests::RowSpec; + let rows = [(100, 1e308), (100_000_000, 1.1e308)] + .into_iter() + .map(|(tokens, latency)| RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: tokens, + total_kv_read_tokens: 0, + latency_ms: latency, + ..RowSpec::default() + }) + .collect::>(); + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &rows); + let db = db_with_pair(tmp.path()); + let mut direct = op(FpmPhase::Prefill); + direct.interpolation = FpmInterpolation::Direct; + let error = direct + .query_totals(&db, &[1.0, 1_000_000.0, 0.0]) + .unwrap_err(); + assert!(error.to_string().contains("invalid latency"), "{error}"); + } + #[test] fn certified_prefill_batch_clamp_routes_to_the_ceiling() { use crate::perf_database::fpm_forward::tests::RowSpec; diff --git a/crates/core/src/perfmodel/perf_database/fpm_forward.rs b/crates/core/src/perfmodel/perf_database/fpm_forward.rs index 2dbc01ceb..6d2c445ed 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -25,6 +25,9 @@ //! `[batch][total_prefill][total_kv]`, decode `[batch][total_kv]` — plus the //! per-phase axis-aligned domain bounding box and a prebuilt //! [`SiteIndex`](super::perf_interp::SiteIndex). +//! Separate direct-timing curves omit every `fake_fallback` row, including +//! healed rows, so measured-only interpolation never inherits extrapolated +//! values or KV ceilings from the legacy SOL table. //! //! Contract notes, all mirrored from Python: //! - An ABSENT parquet is the soft "not collected" case: `cells()` errors only @@ -150,6 +153,13 @@ pub struct FpmForwardCell { pub decode_batches: Vec, pub decode_rungs: Vec, pub decode_curve_bounds: BTreeMap, + /// Direct timing curves `(batch, KV) -> (prefill tokens -> latency)`. + /// Fake-fallback rows are excluded even when the legacy table heals them. + pub direct_prefill: BTreeMap<(u32, u32), BTreeMap>, + /// Direct decode curves contain genuine measurements only; their bounds + /// do not inherit the fabricated rows' extrapolated KV ceilings. + pub direct_decode: BTreeMap>, + pub direct_decode_rungs: Vec, /// Diagnostic state only; no change to table identity, values, or schema. warned_fmha_model_modes: Mutex>, } @@ -1033,6 +1043,9 @@ fn load_pair( decode_batches: Vec::new(), decode_rungs: Vec::new(), decode_curve_bounds: BTreeMap::new(), + direct_prefill: BTreeMap::new(), + direct_decode: BTreeMap::new(), + direct_decode_rungs: Vec::new(), warned_fmha_model_modes: Mutex::new(BTreeSet::new()), }, }); @@ -1040,6 +1053,14 @@ fn load_pair( building.cell.cell_ids.push(row.cell_id.clone()); } if row.workload_kind == "prefill" { + if !row.fake_fallback { + building + .cell + .direct_prefill + .entry((row.batch_size, row.total_kv_read_tokens)) + .or_default() + .insert(row.total_prefill_tokens, row.latency_ms); + } building.cell.prefill.insert( &[ row.batch_size, @@ -1049,6 +1070,14 @@ fn load_pair( row.latency_ms, ); } else { + if !row.fake_fallback { + building + .cell + .direct_decode + .entry(row.batch_size) + .or_default() + .insert(row.total_kv_read_tokens, row.latency_ms); + } building .cell .decode @@ -1060,6 +1089,15 @@ fn load_pair( .into_values() .map(|b| { let mut cell = b.cell; + cell.direct_decode_rungs = cell + .direct_decode + .keys() + .copied() + .filter(|b| { + b.checked_add(1) + .is_some_and(|n| cell.direct_decode.contains_key(&n)) + }) + .collect(); cell.prefill_domain = domain::<3>(&cell.prefill); cell.decode_domain = domain::<2>(&cell.decode); if cell.prefill_domain.is_some() { diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 3a560998c..5aea8651e 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1120,6 +1120,9 @@ struct EngineBuildRequest { kv_block_size: Option, systems_path: Option, forward_model: Option, + fpm_profile: Option, + fpm_interpolation: Option, + cp_size: u32, decode_workload_distribution: Option, prefill_graph_profile: Option, fpm_parquet_path: Option, @@ -1176,6 +1179,9 @@ impl AicEngineBuilder { kv_block_size: None, systems_path: None, forward_model: None, + fpm_profile: None, + fpm_interpolation: None, + cp_size: 1, decode_workload_distribution: None, prefill_graph_profile: None, fpm_parquet_path: None, @@ -1432,7 +1438,7 @@ mod builder_tests { "enable_shared_layer": true, "transfer_policy": ["xshape", "xquant"], "strict_provenance": true, - "extra": {} + "extra": {"fpm_profile": "{\"model\":\"model\"}", "fpm_interpolation": "direct"} })) .expect("deserialize engine config"); @@ -1447,6 +1453,9 @@ mod builder_tests { Some(["xshape".to_owned(), "xquant".to_owned()].as_slice()) ); assert_eq!(request.strict_provenance, Some(true)); + assert_eq!(request.fpm_profile.as_deref(), Some(r#"{"model":"model"}"#)); + assert_eq!(request.fpm_interpolation.as_deref(), Some("direct")); + assert_eq!(request.cp_size, 1); assert_eq!( request.fpm_parquet_path.as_deref(), Some("/artifacts/reviewed-fpm.parquet") @@ -1502,6 +1511,12 @@ fn build_engine_from_request(request: EngineBuildRequest) -> Result Result { + if request.cp_size != 1 { + return Err(AicError::InvalidEngineConfig( + "cp_size must be the integer 1; this SDK entry point does not support context parallelism" + .to_string(), + )); + } if request.fpm_fmha_quant_mode.is_some() && request.forward_model.as_deref() != Some("fpm") { return Err(AicError::InvalidEngineConfig( "fpm_fmha_dtype requires forward_model='fpm'".into(), @@ -1560,6 +1575,9 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result Result, + pub registered: bool, + pub gemm_quant_mode: String, + pub moe_quant_mode: String, + pub fmha_quant_mode: String, + pub kvcache_quant_mode: String, + pub comm_quant_mode: String, + pub attention_backend: String, +} + +pub(crate) fn validate_forward_pass_profile( + config: &crate::ForwardPassPerfModelConfig, +) -> Result { + let json = serde_json::to_string(config) + .map_err(|error| AicError::InvalidEngineConfig(error.to_string()))?; + let facts = Python::with_gil(|py| -> PyResult { + py.import("aisimulate_core.sdk.fpm_profile")? + .call_method1("_validate_forward_pass_profile", (json,))? + .extract() + }) + .map_err(|error| AicError::InvalidEngineConfig(format!("fpm_profile: {error}")))?; + serde_json::from_str(&facts) + .map_err(|error| AicError::InvalidEngineConfig(format!("fpm_profile facts: {error}"))) +} + /// Compile the selected canonical native estimator through the shared builder. pub(crate) fn compile_forward_pass_model_to_engine( config: &crate::ForwardPassPerfModelConfig, @@ -1671,6 +1717,19 @@ pub(crate) fn compile_forward_pass_model_to_engine( kv_block_size: config.kv_block_size, systems_path: Some(systems_path.to_owned()), forward_model: Some(forward_model.to_owned()), + fpm_profile: config + .fpm_profile + .as_ref() + .map(|profile| serde_json::to_string(profile).expect("JSON profile")), + fpm_interpolation: Some( + config + .estimator_config + .fpm_interpolation + .method + .as_str() + .to_owned(), + ), + cp_size: 1, decode_workload_distribution: config .estimator_config .op_level @@ -1729,7 +1788,47 @@ pub(crate) fn compile_engine_to_engine( "fpm_fmha_dtype must be bfloat16, fp8, or fp8_block".into(), )); } - compile_engine_from_request(engine_build_request(config, systems_path)?) + let mut request = engine_build_request(config, systems_path)?; + if config.extra.contains_key("estimator_config") || request.fpm_interpolation.is_some() { + let controls = crate::EstimatorConfig::migrate_legacy_inputs( + config.extra.get("estimator_config").map(String::as_str), + None, + request.fpm_interpolation.as_deref(), + request.fpm_parquet_path.as_deref(), + )?; + request.fpm_options = controls.fpm_interpolation.clone(); + request.fpm_interpolation = match controls.fpm_interpolation.method { + crate::FpmInterpolationMethod::Auto => None, + method => Some(method.as_str().to_owned()), + }; + request.fpm_parquet_path = controls + .fpm_interpolation + .fpm_parquet_path + .map(|path| path.to_str().expect("validated UTF-8 path").to_owned()); + } + if config.extra.contains_key("fpm_profile") { + let mut modes = Python::with_gil(|py| fpm_profile_quantization(py, config)) + .map_err(|e| AicError::InvalidEngineConfig(format!("FPM profile quantization: {e}")))?; + request.gemm_quant_mode = modes.remove("gemm_quant_mode"); + request.moe_quant_mode = modes.remove("moe_quant_mode"); + request.kvcache_quant_mode = modes.remove("kvcache_quant_mode"); + request.fmha_quant_mode = modes.remove("fmha_quant_mode"); + request.comm_quant_mode = modes.remove("comm_quant_mode"); + } + compile_engine_from_request(request) +} + +/// Restore profile precision shared by the native timing and memory bridges. +/// Python checks the wire dtypes before recovering exact quantization modes. +pub(crate) fn fpm_profile_quantization( + py: Python<'_>, + config: &EngineConfig, +) -> PyResult> { + let config_json = serde_json::to_string(config) + .map_err(|e| PyValueError::new_err(format!("invalid engine config: {e}")))?; + py.import("aisimulate_core.sdk.fpm_profile")? + .call_method1("_quantization_from_engine_config", (config_json,))? + .extract() } fn engine_build_request( @@ -1775,6 +1874,9 @@ fn engine_build_request( kv_block_size: config.kv_block_size, systems_path: systems_path.map(str::to_owned), forward_model: config.forward_model.clone(), + fpm_profile: config.extra.get("fpm_profile").cloned(), + fpm_interpolation: config.extra.get("fpm_interpolation").cloned(), + cp_size: config.parallel.cp_size.unwrap_or(1), decode_workload_distribution: None, prefill_graph_profile: config.prefill_graph_profile.clone(), fpm_parquet_path: crate::config::validate_fpm_parquet_path( @@ -1901,24 +2003,28 @@ impl PyForwardPassPerfModel { /// Expand and validate the canonical schema without constructing an engine. #[staticmethod] fn normalize_config(config_json: &str) -> PyResult { - let mut config = parse_forward_pass_config(config_json)?; - config.resolve_prefill_graph_profile().map_err(aic_to_py)?; - config.validate().map_err(aic_to_py)?; + let config = parse_forward_pass_config(config_json)? + .resolve() + .map_err(aic_to_py)?; serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) } /// Share typed FPM option validation with the compilation adapter. #[staticmethod] - #[pyo3(signature = (options_json, fmha_quant_mode=None, comm_quant_mode=None))] + #[pyo3(signature = (options_json, fmha_quant_mode=None, comm_quant_mode=None, has_profile=false))] fn _normalize_fpm_options( options_json: &str, fmha_quant_mode: Option<&str>, comm_quant_mode: Option<&str>, + has_profile: bool, ) -> PyResult { let options: crate::FpmInterpolationConfig = serde_json::from_str(options_json).map_err(|e| { PyValueError::new_err(format!("invalid FPM interpolation options: {e}")) })?; + if has_profile { + options.validate_profile_options().map_err(aic_to_py)?; + } options .validate_quant_modes(fmha_quant_mode, comm_quant_mode) .map_err(aic_to_py)?; @@ -1973,20 +2079,31 @@ impl PyForwardPassPerfModel { } let legacy: EngineConfig = serde_json::from_value(value).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( + let mut request = engine_build_request( &legacy, legacy.systems_path.as_ref().and_then(|path| path.to_str()), ) .map_err(aic_to_py)?; - let mut estimator_config = - crate::EstimatorConfig::from_legacy(options).map_err(aic_to_py)?; - estimator_config.fpm_interpolation.fpm_parquet_path = - request.fpm_parquet_path.as_ref().map(PathBuf::from); + if legacy.extra.contains_key("fpm_profile") { + let mut modes = Python::with_gil(|py| fpm_profile_quantization(py, &legacy))?; + request.gemm_quant_mode = modes.remove("gemm_quant_mode"); + request.moe_quant_mode = modes.remove("moe_quant_mode"); + request.fmha_quant_mode = modes.remove("fmha_quant_mode"); + request.kvcache_quant_mode = modes.remove("kvcache_quant_mode"); + request.comm_quant_mode = modes.remove("comm_quant_mode"); + } + if request.cp_size != 1 { + return Err(PyValueError::new_err( + "cp_size must be the integer 1; this SDK entry point does not support context parallelism", + )); + } + let estimator_config = crate::EstimatorConfig::migrate_legacy_inputs( + legacy.extra.get("estimator_config").map(String::as_str), + options_json, + request.fpm_interpolation.as_deref(), + request.fpm_parquet_path.as_deref(), + ) + .map_err(aic_to_py)?; 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") { @@ -2004,6 +2121,12 @@ impl PyForwardPassPerfModel { backend: legacy.backend, worker_type, backend_version: request.backend_version, + fpm_profile: request + .fpm_profile + .as_deref() + .map(serde_json::from_str) + .transpose() + .map_err(|e| PyValueError::new_err(format!("invalid fpm_profile: {e}")))?, tp: request.tp_size, pp: request.pp_size, attention_dp: request.attention_dp_size, @@ -2076,6 +2199,14 @@ impl PyForwardPassPerfModel { .map_err(aic_to_py) } + fn estimate_forward_pass_detailed(&self, py: Python<'_>, fpm_json: &str) -> PyResult { + let metrics = parse_fpm_iteration(fpm_json)?; + let result = py + .allow_threads(|| self.inner.estimate_forward_pass_detailed(&metrics)) + .map_err(aic_to_py)?; + serde_json::to_string(&result).map_err(|error| PyValueError::new_err(error.to_string())) + } + /// Tune from observed FPM iterations. `iterations_json` is the nested list /// `[[iter0_rank0, ...], [iter1_rank0, ...]]` (the wrapper normalizes the /// convenience forms before calling). Pure-Rust compute (GIL freed). @@ -2333,6 +2464,20 @@ mod tests { } } + #[test] + fn legacy_compilation_with_only_correction_controls_keeps_default_method() { + py_init(); + let mut config = fixture_engine_config(); + config.extra.insert( + "estimator_config".into(), + r#"{"correction":{"enabled":false},"fpm_interpolation":{"text_only":false,"unrecorded_quant_modes":[]}}"#.into(), + ); + for forward_model in [None, Some("fpm".to_owned())] { + config.forward_model = forward_model; + compile_engine_to_engine(&config, systems_root().to_str()).unwrap(); + } + } + #[test] fn legacy_engine_rejects_invalid_fpm_paths() { for path in ["", "/missing/fpm.parquet"] { @@ -2345,6 +2490,29 @@ mod tests { } } + #[test] + fn legacy_compilation_rejects_conflicting_interpolation_before_entering_python() { + for (canonical, legacy) in [("auto", "direct"), ("direct", "auto"), ("sol", "direct")] { + let mut config = fixture_engine_config(); + config.extra.insert( + "estimator_config".into(), + format!(r#"{{"fpm_interpolation": {{"method": "{canonical}"}}}}"#), + ); + config + .extra + .insert("fpm_interpolation".into(), legacy.into()); + let error = compile_engine_to_engine(&config, None) + .err() + .expect("conflicting interpolation methods"); + assert!( + error.to_string().contains( + "conflicting explicit values for estimator_config.fpm_interpolation.method" + ), + "{error}" + ); + } + } + #[cfg(unix)] #[test] fn native_model_rejects_non_utf8_fpm_path() { diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index f5da5d0bb..143a1a2c1 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -211,6 +211,10 @@ struct AicTimingConfig { #[serde(default)] decode_workload_distribution: Option, #[serde(default)] + fpm_profile: Option>, + #[serde(default, skip_serializing)] + fpm_interpolation: Option, + #[serde(default)] fpm_parquet_path: Option, #[serde(default)] decoder_replay: bool, @@ -221,7 +225,7 @@ struct AicTimingConfig { #[serde(default)] fallback_policy: ForwardPassFallbackPolicy, #[serde(default)] - estimator_config: EstimatorConfig, + estimator_config: serde_json::Map, #[serde(default)] database_mode: crate::DatabaseMode, #[serde(default)] @@ -283,7 +287,18 @@ impl AicTimingConfig { } else { self.systems_paths.clone() }; - let mut estimator_config = self.estimator_config.clone(); + if let Some(path) = &self.fpm_parquet_path { + crate::config::validate_fpm_parquet_path( + Some(std::path::Path::new(path)), + mode == EstimationMode::FpmInterpolation, + )?; + } + let mut estimator_config = EstimatorConfig::migrate_legacy_inputs( + Some(&serde_json::to_string(&self.estimator_config)?), + None, + self.fpm_interpolation.map(|method| method.as_str()), + self.fpm_parquet_path.as_deref(), + )?; if let Some(selected) = &self.decode_workload_distribution { ensure!( estimator_config @@ -295,20 +310,6 @@ impl AicTimingConfig { ); estimator_config.op_level.decode_workload_distribution = Some(selected.clone()); } - if let Some(path) = &self.fpm_parquet_path { - crate::config::validate_fpm_parquet_path( - Some(std::path::Path::new(path)), - mode == EstimationMode::FpmInterpolation, - )?; - let configured = &mut estimator_config.fpm_interpolation.fpm_parquet_path; - ensure!( - configured - .as_ref() - .is_none_or(|existing| existing == std::path::Path::new(path)), - "conflicting fpm_parquet_path and estimator_config.fpm_interpolation.fpm_parquet_path" - ); - *configured = Some(path.into()); - } Ok(ForwardPassPerfModelConfig { model: self.model.clone(), system: self.system.clone(), @@ -334,6 +335,7 @@ impl AicTimingConfig { estimation_mode: mode, fallback_policy: self.fallback_policy, estimator_config, + fpm_profile: self.fpm_profile.clone(), database_mode: self.database_mode, transfer_policy: self.transfer_policy.clone(), systems_paths: roots, @@ -424,8 +426,9 @@ impl AicTimingConfig { ); if let (Some(moe_tp), Some(moe_ep)) = (self.moe_tp_size, self.moe_ep_size) { ensure!( - u64::from(self.tp) * u64::from(self.attention_dp) - == u64::from(moe_tp) * u64::from(moe_ep), + self.fpm_profile.is_some() + || u64::from(self.tp) * u64::from(self.attention_dp) + == u64::from(moe_tp) * u64::from(moe_ep), "AIC topology requires tp * attention_dp == moe_tp_size * moe_ep_size" ); } @@ -456,17 +459,22 @@ struct AicTimingModel { use_fpm_decode_totals: bool, fpm_decode_kv_ceiling: Option, evidence: Mutex, + // Capture-only history. Keyed by phase and exact iteration totals, so + // repeated invocations increment a count without cloning prior history. + fpm_evidence: + Mutex>, phase_cache: quick_cache::sync::Cache>, } impl AicTimingModel { fn build(config: &mut AicTimingConfig, worker_type: ForwardPassWorkerType) -> Result { - if config + let request = config.estimator_request(worker_type)?; + if request .estimator_config .op_level .prefill_graph_profile .is_some() - || config + || request .estimator_config .op_level .prefill_graph_profile_id @@ -478,7 +486,7 @@ impl AicTimingModel { } config.validate_parallel_shape()?; config.resolved_memory_fraction()?; - let model = ForwardPassPerfModel::best_available(config.estimator_request(worker_type)?) + let model = ForwardPassPerfModel::best_available(request) .context("AIC timing provider could not construct the requested estimator")?; let provenance = model .provenance() @@ -496,7 +504,16 @@ impl AicTimingModel { config.forward_model = None; config.fpm_parquet_path = None; config.fallback_policy = ForwardPassFallbackPolicy::Deny; - config.estimator_config = provenance.config.estimator_config.clone(); + config.estimator_config = + serde_json::from_value(serde_json::to_value(&provenance.config.estimator_config)?)?; + config.fpm_profile = provenance.config.fpm_profile.clone(); + config.fpm_interpolation = None; + config.gemm_dtype = provenance.config.gemm_quant_mode.clone(); + config.moe_dtype = provenance.config.moe_quant_mode.clone(); + config.fmha_dtype = provenance.config.fmha_quant_mode.clone(); + config.kv_cache_dtype = provenance.config.kvcache_quant_mode.clone(); + config.comm_dtype = provenance.config.comm_quant_mode.clone(); + config.attention_backend = provenance.config.attention_backend.clone(); let use_fpm_decode_totals = provenance.selected_estimation_mode == EstimationMode::FpmInterpolation; let native = model.native_engine().context( @@ -523,6 +540,7 @@ impl AicTimingModel { use_fpm_decode_totals, fpm_decode_kv_ceiling, evidence: Mutex::new(TimingEvidenceAccumulator::default()), + fpm_evidence: Mutex::default(), phase_cache: quick_cache::sync::Cache::new(128), }) } @@ -658,6 +676,32 @@ impl AicTimingModel { .map_err(|_| anyhow!("AIC timing evidence accumulator was poisoned"))?; evidence.record(phase, prefill) } + + fn record_fpm_estimate( + &self, + key: (bool, u32, u32, u32), + estimate: crate::ForwardPassEstimate, + ) -> Result { + let latency_ms = estimate + .latency_ms + .context("native FPM returned no estimate")?; + if latency_ms > 0.0 { + let mut evidence = self + .fpm_evidence + .lock() + .map_err(|_| anyhow!("FPM evidence accumulator was poisoned"))?; + let entry = evidence.entry(key).or_insert(crate::FpmEstimateEvidence { + estimate, + count: 0, + latency_scale: 1.0, + }); + entry.count = entry + .count + .checked_add(1) + .context("FPM evidence invocation count overflow")?; + } + Ok(latency_ms) + } } fn phase_evidence_from_native_entries( @@ -735,6 +779,12 @@ impl TimingModel for AicTimingModel { self.record_evidence(&evidence, true)?; return Ok(latency_ms); } + if let Some(model) = &self.diagnostic_model { + return self.record_fpm_estimate( + (true, batch_size, mean_isl, mean_prefix), + model.static_prefill_detailed(batch_size, mean_isl, mean_prefix)?, + ); + } Python::with_gil(|py| { self.engine .bind(py) @@ -761,6 +811,19 @@ impl TimingModel for AicTimingModel { .min(total_kv_tokens); let batch_size = checked_u32(batch_size, "decode batch size")?; let total_past_kv_tokens = checked_u32(total_past_kv_tokens, "total past KV tokens")?; + if let Some(model) = &self.diagnostic_model { + let estimate = + model.estimate_forward_pass_detailed(&[crate::ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_decode_requests: batch_size, + sum_decode_kv_tokens: total_past_kv_tokens, + ..Default::default() + }, + ..Default::default() + }])?; + return self + .record_fpm_estimate((false, batch_size, 0, total_past_kv_tokens), estimate); + } return Python::with_gil(|py| { self.engine .bind(py) @@ -797,7 +860,52 @@ impl TimingModel for AicTimingModel { .map(|evidence| evidence.snapshot()) } + fn fpm_evidence_summary(&self) -> Option { + if !self.use_fpm_decode_totals || self.diagnostic_model.is_none() { + return None; + } + let history = self.fpm_evidence.lock().ok()?; + let phase = |prefill: bool| { + let fpm_estimates: Vec<_> = history + .iter() + .filter(|(key, _)| key.0 == prefill) + .map(|(_, evidence)| evidence.clone()) + .collect(); + if fpm_estimates.is_empty() { + return TimingPhaseEvidence::default(); + } + let latency_ms = fpm_estimates + .iter() + .map(|record| record.estimate.latency_ms.unwrap_or(0.0) * record.count as f64) + .sum(); + let mut operation = TimingOperationEvidence::new( + if prefill { "fpm_prefill" } else { "fpm_decode" }, + latency_ms, + None, + TimingEvidenceSource::Silicon, + ) + .expect("finite native FPM evidence"); + operation.details = Some(crate::perfmodel::engine::diagnostics::OperationDetails { + sol: None, + sol_unavailable_reason: Some( + "whole-model FPM does not export an operation SOL decomposition".into(), + ), + fallbacks: Vec::new(), + fpm_estimates, + }); + TimingPhaseEvidence::from_operations(vec![operation]) + }; + Some(TimingEvidenceSummary { + prefill: phase(true), + decode: phase(false), + }) + } + fn reset_evidence(&self) -> Result<()> { + self.fpm_evidence + .lock() + .map_err(|_| anyhow!("FPM evidence accumulator was poisoned"))? + .clear(); *self .evidence .lock() @@ -836,6 +944,13 @@ fn estimate_aic_num_gpu_blocks(config: &AicTimingConfig, role: &ReplayRoleConfig kwargs.set_item("attention_backend", config.attention_backend.as_deref())?; kwargs.set_item("enable_eplb", config.enable_eplb)?; kwargs.set_item("wideep_num_slots", config.wideep_num_slots)?; + kwargs.set_item( + "fpm_profile", + config + .fpm_profile + .as_ref() + .map(|profile| serde_json::to_string(profile).expect("JSON profile")), + )?; kwargs.set_item( "cuda_graph_reserved_bytes", config.cuda_graph_reserved_bytes, @@ -1429,9 +1544,20 @@ impl TimingPowerSource { } fn replay_timing_evidence(sources: &[TimingPowerSource]) -> Result> { + replay_timing_evidence_with_fpm(sources, false) +} + +fn replay_timing_evidence_with_fpm( + sources: &[TimingPowerSource], + include_fpm: bool, +) -> Result> { let mut combined = TimingEvidenceSummary::default(); for source in sources { - let Some(summary) = source.timing.evidence_summary() else { + let Some(summary) = source.timing.evidence_summary().or_else(|| { + include_fpm + .then(|| source.timing.fpm_evidence_summary()) + .flatten() + }) else { return Ok(None); }; combined.prefill.try_accumulate(scale_power_phase( @@ -1446,6 +1572,42 @@ fn replay_timing_evidence(sources: &[TimingPowerSource]) -> Result Result> { + let mut providers = Vec::new(); + for (provider_index, source) in sources.iter().enumerate() { + let Some(summary) = source.timing.fpm_evidence_summary() else { + continue; + }; + let mut phases = Vec::new(); + for (name, phase, speedup) in [ + ("prefill", summary.prefill, source.prefill_speedup_ratio), + ("decode", summary.decode, source.decode_speedup_ratio), + ] { + let phase = scale_power_phase(phase, speedup)?; + let operations = phase + .operations + .into_iter() + .filter_map(|operation| { + let details = operation.details?; + (!details.fpm_estimates.is_empty()).then(|| { + serde_json::json!({ + "name": operation.name, "fpm_estimates": details.fpm_estimates, + }) + }) + }) + .collect::>(); + phases.push(serde_json::json!({"name": name, "operations": operations})); + } + providers.push(serde_json::json!({"provider_index": provider_index, "phases": phases})); + } + Ok((!providers.is_empty()).then(|| { + serde_json::json!({ + "schema_version": "1.0", "aggregation": "identical_estimates_with_invocation_counts", + "providers": providers, + }) + })) +} + fn replay_power_stats(summary: &TimingEvidenceSummary) -> Result { let mut combined = summary.prefill.clone(); combined.try_accumulate(summary.decode.clone())?; @@ -1491,12 +1653,12 @@ fn replay_performance_diagnostics(summary: Option<&TimingEvidenceSummary>) -> se let scope = "accumulated_active_forward_pass_per_gpu"; let Some(summary) = summary else { return serde_json::json!({"status": "unavailable", "scope": scope, "latency_unit": "ms", "phases": [], - "unavailable_reason": "selected timing provider does not export operation diagnostics (whole-model FPM and latency-only providers are unsupported)"}); + "unavailable_reason": "selected timing provider does not export operation diagnostics"}); }; let phases = [("prefill", &summary.prefill), ("decode", &summary.decode)].into_iter().map(|(name, phase)| { let mut operations = phase.operations.iter().map(|op| { let sol = op.details.as_ref().and_then(|d| d.sol.as_ref()); - serde_json::json!({"name": op.name, "latency_ms": op.latency_ms, + let mut record = serde_json::json!({"name": op.name, "latency_ms": op.latency_ms, "source": op.source.as_str(), "sol": sol, "sol_unavailable_reason": if sol.is_some() { None } else { @@ -1504,7 +1666,11 @@ fn replay_performance_diagnostics(summary: Option<&TimingEvidenceSummary>) -> se }, "latency_to_sol_ratio": sol.filter(|s| s.latency_ms > 0.0).map(|s| op.latency_ms / s.latency_ms).filter(|r| r.is_finite()), "fallbacks": op.details.as_ref().map(|d| &d.fallbacks), - }) + }); + if let Some(details) = &op.details && !details.fpm_estimates.is_empty() { + record["fpm_estimates"] = serde_json::to_value(&details.fpm_estimates).expect("finite native FPM evidence"); + } + record }).collect::>(); operations.sort_by(|a,b| a["name"].as_str().cmp(&b["name"].as_str())); let sol = phase.operations.iter().map(|op| op.details.as_ref()?.sol.as_ref()).collect::>>() @@ -1700,6 +1866,11 @@ fn scale_power_phase( operation.energy_wms = operation.energy_wms.map(|energy| energy * scale); operation.latency_ms *= scale; operation.covered_latency_ms *= scale; + if let Some(details) = &mut operation.details { + for evidence in &mut details.fpm_estimates { + evidence.latency_scale *= scale; + } + } // SOL compares the same scheduled work, before synthetic speedup; it // remains an unscaled physical baseline. } @@ -1984,8 +2155,12 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { let mut report_json = serde_json::to_value(&report).context("serializing AISimulate replay report summary")?; if capture_performance_diagnostics { + let performance_evidence = replay_timing_evidence_with_fpm(&power_sources, true)?; report_json["performance_diagnostics"] = - replay_performance_diagnostics(timing_evidence.as_ref()); + replay_performance_diagnostics(performance_evidence.as_ref()); + } + if let Some(evidence) = replay_fpm_query_evidence(&power_sources)? { + report_json["fpm_query_evidence"] = evidence; } if report.agentic_graph.is_some() && let Some((input_format, agentic_lanes, execution_model)) = agentic_input @@ -2556,6 +2731,7 @@ mod tests { use_fpm_decode_totals, fpm_decode_kv_ceiling: None, evidence: Mutex::new(TimingEvidenceAccumulator::default()), + fpm_evidence: Mutex::default(), phase_cache: quick_cache::sync::Cache::new(128), } } @@ -2718,6 +2894,7 @@ mod tests { let missing_op = |phase| { Op::FpmForward(FpmForwardOp { dcp_size: None, + interpolation: Default::default(), name: "missing".into(), phase, model_path: "missing-test-model".into(), @@ -2806,6 +2983,34 @@ mod tests { ); } + #[test] + fn replay_interpolation_migration_preserves_explicit_method_presence() { + for canonical in [None, Some("direct"), Some("auto"), Some("sol")] { + let mut config = aic_config(); + config.fpm_interpolation = Some(crate::FpmInterpolationMethod::Direct); + config.estimator_config = serde_json::from_value(serde_json::json!({ + "correction": {"enabled": false}, + "fpm_interpolation": canonical + .map(|method| serde_json::json!({"method": method})) + .unwrap_or_else(|| serde_json::json!({})) + })) + .unwrap(); + let request = config.estimator_request(ForwardPassWorkerType::Aggregated); + if canonical.is_none() || canonical == Some("direct") { + let controls = request.unwrap().estimator_config; + assert_eq!( + controls.fpm_interpolation.method, + crate::FpmInterpolationMethod::Direct + ); + assert!(!controls.correction.enabled); + } else { + assert!(request.unwrap_err().to_string().contains( + "conflicting explicit values for estimator_config.fpm_interpolation.method" + )); + } + } + } + #[test] fn canonical_timing_config_preserves_moe_kernel_source() { let mut config = aic_config(); @@ -2849,7 +3054,7 @@ mod tests { worker_type: None, estimation_mode: None, fallback_policy: ForwardPassFallbackPolicy::Deny, - estimator_config: EstimatorConfig::default(), + estimator_config: serde_json::Map::new(), database_mode: crate::DatabaseMode::default(), transfer_policy: None, systems_paths: Vec::new(), @@ -2862,6 +3067,8 @@ mod tests { strict_provenance: false, systems_path: None, forward_model: None, + fpm_profile: None, + fpm_interpolation: None, decode_workload_distribution: None, fpm_parquet_path: None, decoder_replay: false, @@ -2952,8 +3159,10 @@ mod tests { #[test] fn replay_provider_rejects_graph_profile_before_native_construction() { let mut config = aic_config(); - config.estimator_config.op_level.prefill_graph_profile = - Some(crate::perf_database::prefill_graph::PROFILE_NAME.to_owned()); + config.estimator_config.insert( + "op_level".into(), + serde_json::json!({"prefill_graph_profile": crate::perf_database::prefill_graph::PROFILE_NAME}), + ); for role in [ ForwardPassWorkerType::Prefill, ForwardPassWorkerType::Decode, @@ -3368,6 +3577,7 @@ mod tests { TimingOperationEvidence::new("dispatch", 10.0, None, TimingEvidenceSource::Estimated) .unwrap(); op.details = Some(OperationDetails { + fpm_estimates: Vec::new(), sol: Some(SolDiagnostics { latency_ms: 2.0, math_ms: 0.0, @@ -3468,10 +3678,20 @@ mod tests { "system": "test-system", "tp": 1, "forward_model": "fpm", + "fpm_profile": {"model": "test-model"}, + "fpm_interpolation": "direct", "fpm_parquet_path": "/artifacts/reviewed-fpm.parquet" })) .unwrap(); assert_eq!(config.forward_model.as_deref(), Some("fpm")); + assert_eq!( + config.fpm_profile, + Some(serde_json::from_value(serde_json::json!({"model": "test-model"})).unwrap()) + ); + assert_eq!( + config.fpm_interpolation, + Some(crate::FpmInterpolationMethod::Direct) + ); assert_eq!( config.fpm_parquet_path.as_deref(), Some("/artifacts/reviewed-fpm.parquet") @@ -3522,6 +3742,13 @@ mod tests { assert_eq!(latency, 546_046.0); assert_eq!(timing.evidence_summary(), None); + assert_eq!(timing.fpm_evidence_summary(), None); + // Capture is disabled: even a long sequence of distinct total-KV + // queries cannot allocate a growing evidence history. + for kv in 546_100..547_100 { + timing.predict_decode_ms(35, kv, 15_602, 1_000_000).unwrap(); + } + assert!(timing.fpm_evidence.lock().unwrap().is_empty()); } #[test] diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index b9895f18f..0b4306f46 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -10,7 +10,8 @@ use std::path::Path; use aisimulate_core::{ AicEngine, AicEngineBuilder, AicError, BackendKind, DatabaseMode, EstimationMode, EstimatorConfig, ForwardPassPerfModel, ForwardPassPerfModelConfig, - ForwardPassRegressionStoreDiagnostics, ForwardPassWorkerType, KvCacheEstimateRequest, + ForwardPassRegressionStoreDiagnostics, ForwardPassWorkerType, FpmInterpolationMethod, + KvCacheEstimateRequest, }; /// Compile the ergonomic engine builder without starting embedded Python. @@ -71,6 +72,14 @@ pub fn regression_options() -> EstimatorConfig { config } +/// Independent profile and interpolation method use the one construction API. +pub fn profile_model( + mut config: ForwardPassPerfModelConfig, +) -> Result { + config.estimator_config.fpm_interpolation.method = FpmInterpolationMethod::Direct; + ForwardPassPerfModel::best_available(config) +} + pub fn best_available_model( config: ForwardPassPerfModelConfig, ) -> Result { @@ -105,9 +114,9 @@ pub fn accept_kv_request(request: KvCacheEstimateRequest) -> KvCacheEstimateRequ mod tests { use super::*; use aisimulate_core::{ - ForwardPassMetrics, ForwardPassRegressionWorkloadKind, TimingEvidenceSource, - TimingEvidenceSummary, TimingOperationEvidence, TimingPhaseEvidence, - ENGINE_CONFIG_SCHEMA_VERSION, ENGINE_SPEC_SCHEMA_VERSION, FPM_VERSION, + ENGINE_CONFIG_SCHEMA_VERSION, ENGINE_SPEC_SCHEMA_VERSION, FPM_VERSION, ForwardPassMetrics, + ForwardPassRegressionWorkloadKind, TimingEvidenceSource, TimingEvidenceSummary, + TimingOperationEvidence, TimingPhaseEvidence, }; #[test] @@ -146,7 +155,8 @@ mod tests { // v22: observed MoE selection and exact prefill graph composites. // v23: ParallelMapping gained optional recorded DCP identity. // v24: FpmForwardOp carries typed DCP separately from matching strings. - assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 24); + // v25: SOL/direct interpolation selector combined with typed DCP. + assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 25); assert_eq!(FPM_VERSION, 1); assert_eq!(ForwardPassMetrics::default().version, FPM_VERSION); } @@ -263,9 +273,11 @@ mod tests { stores[0].workload_kind, ForwardPassRegressionWorkloadKind::PureDecode ); - assert!(stores - .iter() - .all(|store| !store.ready && store.retained_observations == 0)); + assert!( + stores + .iter() + .all(|store| !store.ready && store.retained_observations == 0) + ); let _roles = [ ForwardPassWorkerType::Prefill, ForwardPassWorkerType::Decode, diff --git a/docs/cli/prediction-details.schema.json b/docs/cli/prediction-details.schema.json index 75f237e68..bf3005a7c 100644 --- a/docs/cli/prediction-details.schema.json +++ b/docs/cli/prediction-details.schema.json @@ -1031,6 +1031,12 @@ } } ] + }, + "fpm_estimates": { + "type": "array", + "items": { + "$ref": "#/$defs/fpmEstimateRecord" + } } }, "required": [ @@ -1094,6 +1100,12 @@ "null" ], "minimum": 0 + }, + "fpm_estimates": { + "type": "array", + "items": { + "$ref": "#/$defs/fpmEstimateRecord" + } } }, "required": [ @@ -1323,6 +1335,231 @@ } } ] + }, + "fpmCoordinates": { + "type": "object", + "properties": { + "batch_size": { + "type": "number", + "minimum": 0 + }, + "total_prefill_tokens": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "total_kv_read_tokens": { + "type": "number", + "minimum": 0 + } + }, + "required": [ + "batch_size", + "total_prefill_tokens", + "total_kv_read_tokens" + ], + "additionalProperties": false + }, + "fpmMeasurementSupport": { + "type": "object", + "properties": { + "coordinates": { + "$ref": "#/$defs/fpmCoordinates" + }, + "latency_ms": { + "type": "number", + "minimum": 0 + }, + "weight": { + "type": "number", + "minimum": 0, + "maximum": 1 + } + }, + "required": [ + "coordinates", + "latency_ms", + "weight" + ], + "additionalProperties": false + }, + "fpmQuery": { + "type": "object", + "properties": { + "phase": { + "enum": [ + "prefill", + "decode" + ] + }, + "model_path": { + "type": "string", + "minLength": 1 + }, + "decode_baseline": { + "type": "boolean" + }, + "query": { + "$ref": "#/$defs/fpmCoordinates" + }, + "resolution": { + "enum": [ + "exact_lookup", + "within_curve_interpolation", + "cross_kv_interpolation", + "cross_batch_interpolation" + ] + }, + "latency_ms": { + "type": "number", + "minimum": 0 + }, + "support": { + "type": "array", + "items": { + "$ref": "#/$defs/fpmMeasurementSupport" + }, + "minItems": 1 + } + }, + "required": [ + "phase", + "model_path", + "decode_baseline", + "query", + "resolution", + "latency_ms", + "support" + ], + "additionalProperties": false + }, + "fpmRankEstimate": { + "type": "object", + "properties": { + "rank": { + "type": "integer", + "minimum": 0 + }, + "latency_ms": { + "type": "number", + "minimum": 0 + }, + "prefill_ms": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "decode_ms": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "decode_baseline_ms": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "marginal_decode_ms": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "queries": { + "type": "array", + "items": { + "$ref": "#/$defs/fpmQuery" + } + } + }, + "required": [ + "rank", + "latency_ms", + "prefill_ms", + "decode_ms", + "decode_baseline_ms", + "marginal_decode_ms", + "queries" + ], + "additionalProperties": false + }, + "fpmEstimate": { + "type": "object", + "properties": { + "latency_ms": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "native_latency_ms": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "correction_factor": { + "type": [ + "number", + "null" + ], + "minimum": 0 + }, + "max_rank": { + "type": [ + "integer", + "null" + ], + "minimum": 0 + }, + "ranks": { + "type": "array", + "items": { + "$ref": "#/$defs/fpmRankEstimate" + } + } + }, + "required": [ + "latency_ms", + "native_latency_ms", + "correction_factor", + "max_rank", + "ranks" + ], + "additionalProperties": false + }, + "fpmEstimateRecord": { + "type": "object", + "properties": { + "estimate": { + "$ref": "#/$defs/fpmEstimate" + }, + "count": { + "type": "integer", + "minimum": 1 + }, + "latency_scale": { + "type": "number", + "minimum": 0 + } + }, + "required": [ + "estimate", + "count", + "latency_scale" + ], + "additionalProperties": false } }, "$comment": "SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0", diff --git a/docs/core-api.md b/docs/core-api.md index 08243ea6f..6ee9a4738 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -321,6 +321,53 @@ nested paths. The supported namespaces are: - `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. +- `fpm_interpolation.method`: `auto` (default), `sol`, or `direct`. Rust selects + SOL for a registered architecture, or direct interpolation for a verified + architecture without a registered class when a profile is supplied. Without + a profile, auto retains SOL. Explicit SOL requires a registered analytical + model; direct requires a profile and also supports registered architectures. + +The top-level `fpm_profile` contains the complete profile dictionary: pinned +model revision, architecture, context length, expert count, deployment precision +and topology, conservative resource bounds, and provenance. A profile requires +an explicit literal `backend_version` that matches its selected deployment; +slot aliases and omitted versions are rejected. This profile schema does not +declare recorded DCP, so combining `fpm_profile` with an explicit `dcp` is +rejected; measured DCP profiles without `fpm_profile` retain their existing route. +Profile/schema and precision conflicts fail before estimator fallback. Omitted precision fields are filled +from the profile and preserved in the resolved canonical configuration. + +Runtime normalization and construction verify the profile architecture against +checkpoint `config.json` from the local model path or pinned remote revision +before interpolation selection. Missing or malformed architecture metadata and +mismatched declarations fail explicitly. This check reads configuration only: +no weights or analytical graph are required. An architecture with verified +metadata remains valid for direct interpolation without a registered analytical +class. Schema-only profile and application configuration parsing remain lightweight. + +For measured-only timing, pass `estimation_mode="fpm_interpolation"`, +`fallback_policy="deny"`, and +`estimator_config={"fpm_interpolation": {"method": "direct"}}` together with +`fpm_profile`. Direct interpolation requires `database_mode="SILICON"`, emits +whole-forward native operations without SOL operations, and never constructs +an analytical graph. Profile resource estimates and memory planning do not +require timing data or a native timing model. + +Direct readiness requires genuine measurements for each operation's phase; +prefill-only operations do not require decode rows, and vice versa. Construction +continues to later systems roots when a required phase is unavailable. Query +coverage still needs an exact point or supported interpolation. Cross-KV prefill +uses the nearest same-batch lower and upper KV curves that both cover the +requested token count, without a KV distance limit. SOL's site-distance guard +does not apply to this direct bracket. See the [self-service coverage rules](fpm-self-service.md#choose-the-model-execution-route). + +The returned provenance pins both the selected estimation mode and interpolation +method, alongside the complete normalized profile. Reusing its `config` keeps +that selection across serialization and replay. Later timing coverage errors +never switch estimator or interpolation method. A registered model's graph +construction failure does not change SOL to direct; top-level fallback still +follows the configured estimator ordering and policy. + - `op_level`: optional `decode_workload_distribution` selects a measured decode-MoE distribution, and `prefill_graph_profile` selects a qualified direct-prefill graph composition. Saved configurations retain the resolved immutable `prefill_graph_profile_id`, which is validated on reload. Unknown fields are rejected. - `fpm_interpolation`: `text_only` (false) permits text prefill/decode profiles for multimodal architectures while retaining encoder weights. It does not @@ -486,9 +533,10 @@ fidelity to the generic Replay cache/scheduler. This API change enables timing consumption, not full hybrid-cache simulation or multimodal prediction from text-only measurements. -The positional engine-spec wire version remains unchanged: the engine identity -is JSON-encoded and the FPM match identity is already variable-length. Legacy -configuration and profiles without DCP remain accepted as unrecorded DCP. +The positional engine-spec wire format is version 25, including the new FPM +interpolation field. Recompile older binary EngineSpecs: schema 24 and other +incompatible versions are rejected before payload decoding. Legacy configuration +and profiles without DCP remain accepted as unrecorded DCP. ### Migrating saved configuration @@ -500,6 +548,17 @@ 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. +Legacy `extra.fpm_profile` and `extra.fpm_interpolation` migrate to the full +canonical profile and nested interpolation method; newly exported configuration +uses only the canonical fields. + +Migration merges saved estimator controls with explicitly supplied legacy +options before applying ordinary defaults. It preserves disjoint settings and +accepts agreeing overlaps; contradictory explicit values report the canonical +setting's path. The same rule applies to legacy and canonical interpolation +methods, including an explicit `auto`. An omitted field does not override a +saved value. `ForwardPassPerfOptions.to_dict()` serializes only arguments +explicitly supplied to that legacy options object, including explicit defaults. The migration adapter rejects any non-null `prefill_graph_profile`, `prefill_graph_profile_id`, or `decode_workload_distribution` field, including an orphan profile ID. These selectors require the canonical `ForwardPassPerfModelConfig.estimator_config.op_level` configuration; pass a saved canonical configuration directly to `RustForwardPassPerfModel.best_available` to preserve its profile identity and supported API restrictions. Profile-free legacy configurations continue to migrate normally. @@ -557,6 +616,31 @@ missing-data sentinel, not evidence of a zero-power operation. See the aggregation rules, and public output boundary. Typed per-op energy alone does not make unified replay power available. +## Direct FPM query evidence + +The same model returned by `RustForwardPassPerfModel.best_available(config)` exposes `estimate_forward_pass_detailed(metrics)` (also available on the Rust `ForwardPassPerfModel`). Its `latency_ms` equals `estimate_forward_pass_time_ms(metrics)`; the scalar API and four-element per-operation tuple APIs are unchanged. + +```python +result = model.estimate_forward_pass_detailed({ + "scheduled_requests": { + "num_prefill_requests": 1, + "sum_prefill_tokens": 512, + "sum_prefill_kv_tokens": 256, + }, +}) +for rank in result["ranks"]: + for query in rank["queries"]: + print(query["resolution"], query["query"], query["support"]) +``` + +Each direct lookup records `exact_lookup`, `within_curve_interpolation`, `cross_kv_interpolation`, or `cross_batch_interpolation`, its phase and model, requested coordinates, raw latency, and supporting measured coordinates, latencies and weights. Coordinates are iteration totals: `batch_size`, `total_prefill_tokens` (null for decode), and `total_kv_read_tokens`. Nested interpolation weights are multiplied for each supporting measurement. This evidence preserves `source="silicon"`, whose meaning includes interpolation; it does not establish accuracy or observed CUDA graph dispatch. + +Weights describe a raw lookup, before later composition. Each rank records prefill, decode and mixed-pass baseline latencies, plus `marginal_decode_ms = max(decode_ms - decode_baseline_ms, 0)`. A baseline lookup has `decode_baseline=true`: its query is the paired decode request, but its support contains each selected row's actual measured minimum-KV point. `max_rank` is the first zero-based input index attaining the positive native maximum. The result separately records `native_latency_ms`, `correction_factor`, and final `latency_ms`. Empty work has no query evidence. Other estimators also have no direct lookup evidence; an unready regression retains `latency_ms=null`. + +Replay and prediction capture this evidence with `ReplayOutputRequirements(capture_performance_diagnostics=True)`; CLI prediction enables it with `--detail source` or `--detail time`. Use `--format json` or inspect the saved `prediction.json` for the query records and measured support; the default source text output shows operation sources and fallbacks. Public `run_recommendation(..., output_requirements=ReplayOutputRequirements(capture_performance_diagnostics=True))` and `Sweeper(..., output_requirements=...)` forward the same request to each candidate's replay. Evidence survives saved results under `ReplayReport.metadata["fpm_query_evidence"]` and each candidate's `provenance.runner_metadata["fpm_query_evidence"]`. The native replay report carries the same top-level key; source details include `fpm_estimates` on whole-model operations. FPM power remains unavailable and SOL comparison has an explicit unavailable reason. + +Captured replay records are grouped by timing provider and phase. Each `fpm_estimates` entry holds an `estimate`, an invocation `count`, and `latency_scale` for any synthetic replay speedup; raw support is never rescaled. Identical query coordinates within an immutable replay provider share one counted record. Counts describe timing invocations in the measurement epoch, not independent silicon observations or served requests. Distinct queries are retained without truncation, so capture memory grows with distinct coordinates; ordinary runs do not retain query histories. Aggregated topology has provider index 0; disaggregated topology orders prefill before decode. + ## Static phase diagnostics The canonical `ForwardPassPerfModel` returned by `best_available` exposes diff --git a/docs/fpm-self-service.md b/docs/fpm-self-service.md index 71d5fc9f0..d3681ea0e 100644 --- a/docs/fpm-self-service.md +++ b/docs/fpm-self-service.md @@ -5,13 +5,13 @@ SPDX-License-Identifier: Apache-2.0 # FPM self-service -`aisimulate onboard` guides onboarding a new model for FPM simulation on your designated hardware platform. It records the model, runtime, target GPU system and allocation, plans one pure tensor-parallel worker, and produces ordinary `predict` and `recommend` configurations that read your collected FPM data. It builds on AISimulate's existing per-worker FPM support and packaged collector. +`aisimulate onboard` guides onboarding a new model for FPM simulation on your designated hardware platform. It records the model, runtime, target GPU system and allocation, plans one TP, DEP, or TEP worker, and produces ordinary `predict` and `recommend` configurations that read your collected FPM data. It builds on AISimulate's per-worker FPM support and packaged collector. -Planning works before the model has an AISimulate model class or measured FPM timings. A valid request records your choices; model metadata and execution readiness, runtime compatibility, and data readiness remain **unchecked**, and accuracy is **not assessed**. This setup does not run target preflight checks, provision GPU resources, or establish measured accuracy. +Planning works before the model has an AISimulate model class or measured FPM timings. With a supplied FPM profile, planning validates the declared deployment identity and estimates memory admission from your resource bounds. Runtime compatibility and data readiness remain **unchecked**, and accuracy is **not assessed**. This setup does not provision GPUs or run target preflight checks. -In this revision, ordinary FPM `predict` and `recommend` still construct a registered analytical model for resource accounting and SOL-based timing transfer. A separate decoupling change will add a route that uses model/resource metadata and direct interpolation of measured timings without an op-level model class. The existing registered-model/SOL route will remain available. See [Choose the model execution route](#choose-the-model-execution-route) before handing off a new architecture. +Ordinary FPM `predict` and `recommend` accept an inline `engine.fpm_profile` with model identity and rank-local resource bounds. Direct interpolation uses measured timings without constructing an op-level model. Automatic interpolation selection retains SOL for registered architectures. See [Choose the model execution route](#choose-the-model-execution-route) for selection and coverage rules. -The Inkling pilot is planned for NVIDIA GB200 with TP, DEP, and TEP configurations through the class-independent route. This guided planner currently generates pure-TP configurations only. The checkpoint, runtime, allocation, and strategy degrees must be pinned before that pilot; this guide does not establish Inkling readiness or GB200 accuracy. +The Inkling pilot targets NVIDIA GB200 with TP, DEP, and TEP through the class-independent route. Its checkpoint, runtime, allocation, profile resource bounds, and strategy degrees must be pinned before collection. This guide does not establish Inkling readiness or GB200 accuracy. ## Create the request @@ -40,7 +40,43 @@ Replace the model and runtime placeholders with the actual inputs. The example d The default pilot uses 1,024 input tokens, 128 output tokens, concurrency 1, four requests, a 16,384-token context limit, TTFT target 1,000 ms, and TPOT target 100 ms. Scripted setup defaults to TP1 unless `--tensor-parallel` is supplied. These are planning defaults, not measured model capacity or latency. Change them with the corresponding flags shown by `aisimulate onboard init --help`. -One node is the default; GPUs per node then equals `--gpu-count`. For multiple nodes, supply both `--node-count` and `--gpus-per-node`; their product must equal the total allocation. Each pure-TP worker must fit on one node. Advanced flags include `--request-count`, `--max-candidates`, `--objective`, `--seed`, and `--sm`. +One node is the default; GPUs per node then equals `--gpu-count`. For multiple nodes, supply both `--node-count` and `--gpus-per-node`; their product must equal the total allocation. Each selected worker must fit on one node. Advanced flags include `--request-count`, `--max-candidates`, `--objective`, `--seed`, and `--sm`. + +`--tensor-parallel` always means attention TP. Use all MoE dimensions explicitly for DEP and TEP; replicas are independent workers: + +| Worker | CLI parallelism flags | `(TP, PP, attention DP, MoE TP, MoE EP, CP)` | +| --- | --- | --- | +| MoE TP4 | `--tensor-parallel 4 --moe-tensor-parallel 4` | `(4, 1, 1, 4, 1, 1)` | +| DEP8 | `--tensor-parallel 1 --attention-data-parallel 8 --moe-tensor-parallel 1 --moe-expert-parallel 8` | `(1, 1, 8, 1, 8, 1)` | +| TEP8 | `--tensor-parallel 8 --attention-data-parallel 1 --moe-tensor-parallel 1 --moe-expert-parallel 8` | `(8, 1, 1, 1, 8, 1)` | + +The first profile implementation supports vLLM decoder-only models with PP1, CP1 and linear KV storage. Recorded DCP is not represented by this profile schema and is rejected; existing registered-model DCP profiles remain a separate supported path. AFD, encoder pools, speculative decoding, nonlinear recurrent state, wide EP and EPLB need additional metadata/semantics and are rejected by this route. + +## Provide identity and resource metadata + +For a model without an analytical class, create a JSON or YAML FPM profile and pass `--fpm-profile /path/to/model-profile.yaml` to `onboard init`. The request embeds the complete profile; ordinary prediction and recommendation use the same object under `engine.fpm_profile`. The profile schema is [FpmModelProfile](../python/aisimulate/src/aisimulate_core/fpm_profile.py). It rejects missing fields, unknown fields, conflicting identities and mutable revision placeholders. + +| Profile fields | Required meaning | +| --- | --- | +| `schema_version` | Integer `1`. | +| `model`, `model_revision`, `architecture` | Exact timing model identity, pinned checkpoint revision, and architecture identifier matching the checkpoint configuration. An analytical model class is optional for direct interpolation. | +| `context_length`, `num_experts`, `provenance` | Declared context limit, routed expert count (`0` for dense), and how the metadata was obtained. | +| `deployments` | One or more exact deployment records, with one precision/resource identity per hardware/runtime/full parallel tuple. | +| Deployment `system`, `backend`, `backend_version` | GPU system, `vllm`, and literal runtime version. | +| Deployment `tp`, `pp`, `dp`, `moe_tp`, `moe_ep`, `cp` | Complete topology. PP and CP default to `1`; other dimensions are explicit. | +| Deployment precision | Exact `gemm_quant_mode`, `moe_quant_mode`, `fmha_quant_mode`, `comm_quant_mode`, `kv_cache_dtype`. They must match the collected cell; FP8 FMHA and BF16 FMHA are separate identities. | +| Deployment runtime hints | `moe_backend`, `attention_backend` (default `auto`), `enable_wideep`, `enable_eplb` (currently both `false`). | +| Deployment `resources` | Conservative bounds described below, specific to this topology. | + +Each `resources` record requires integer `weights_bytes`, `activations_bytes`, `runtime_overhead_bytes`, `comm_overhead_bytes`, positive `kv_bytes_per_token`, `cache_layout: linear`, positive `max_num_tokens` and `max_batch_size`, and a `provenance` explanation. Obtain these from checkpoint tensor/storage metadata and serving-runtime accounting or conservative explicit estimates. A profile is an input declaration; AISimulate does not certify that its values describe your runtime. + +All resource bytes are **per rank**, and must bound every rank. Scheduler envelope limits are also per attention-DP rank; they are not the summed batch and token totals used by the worker's FPM timing query. Do not copy TP resource values into DEP/TEP merely because GPU counts match. Include all non-KV storage, including quantization overheads. Overhead fields exclude the separate `kv_cache.capacity.cuda_graph_reserved_bytes` reservation. Profiles carry explicit quantization identity; `text_only` and `unrecorded_quant_modes` overrides are rejected on this route. Requests above the declared scheduler envelope fail explicitly. Collection validates both prefill and explicit decode batch limits against that envelope before launching workers. + +Automatic KV admission subtracts these non-KV bounds and CUDA graph reservation from the configured fraction of GPU memory, then divides by the declared KV bytes per token. `context_length: max` uses the profile limit; `bytes_per_token: auto` uses the exact deployment's cache geometry. These operations need a hardware specification and profile, but no FPM timing files or model graph. + +The declared checkpoint revision is preserved for review and replay. Existing FPM tables may not pin a checkpoint revision; supplying one in a profile does not prove that the historical measurements came from that exact revision. Record that limitation in `provenance` when validating existing data. + +Ordinary FPM model construction reads `config.json` from the local checkpoint or the declared remote `model_revision` and verifies the profile's architecture against the checkpoint's text-model architecture before selecting interpolation. Missing or malformed architecture metadata and mismatched names fail explicitly. This verification loads no weights and constructs no analytical model graph. A verified architecture without a registered analytical class remains valid for direct interpolation. Schema-only parsing and onboarding request creation do not require this runtime metadata check. ## Plan, preview, and explicitly execute @@ -53,11 +89,11 @@ aisimulate onboard collect-fpm \ --output-dir ./aisimulate-support ``` -The first command saves the request, `support-plan.json`, `commands.json`, `predict/pilot.yaml`, `recommend/pilot.yaml`, and a local `systems/` directory. The second prints the collector command without launching it. Generated command vectors and printed next commands use absolute output paths and preserve spaces or shell punctuation. Use a separate output directory for each request. `--overwrite` can repair missing generated files for the identical saved request, including an interrupted write before `support-plan.json` was created. Repair requires an intact matching `request.yaml`, rejects changed inputs, and preserves existing collected data. +The first command saves the request, `support-plan.json`, `commands.json`, `predict/pilot.yaml`, `recommend/pilot.yaml`, and a local `systems/` directory. When supplied, the profile is also saved as `fpm-model-profile.json`, included in the collector command, and embedded in prediction/recommendation configs. The plan records the CPU resource estimate. The second command prints the collector invocation without launching it. Generated command vectors and printed next commands use absolute output paths and preserve spaces or shell punctuation. Use a separate output directory for each request. `--overwrite` can repair missing generated files for the identical saved request, including an interrupted write before `support-plan.json` was created. Repair requires an intact matching `request.yaml`, rejects changed inputs, and preserves existing collected data. If a software update changes generated guidance, recreate the plan in a new output directory: repair compares generated files byte for byte and does not migrate older plans. Guidance changes do not change the saved-request identity checks used by collection. -The search uses one selected TP size. By default it evaluates a single worker, so recommendation is not a broad deployment search. `--max-candidates 2` additionally considers the largest count of identical workers that fits the allocation, when that differs from one worker. Each choice gets an independent recommendation config pinned to that replica count with a one-trial budget. The single worker keeps `recommend/pilot.yaml`; the second choice uses `recommend/replicas-N.yaml`, where `N` is its replica count. The plan reports the actual candidate count and lists both config and result paths. Dense collection uses the `tp` preset; MoE uses `pure_tp`; the selected TP size remains exact. +The guided plan uses one selected parallel tuple. By default it evaluates a single worker, so recommendation is not a broad deployment search. `--max-candidates 2` additionally considers the largest count of identical workers that fits the allocation, when that differs from one worker. Each choice gets an independent recommendation config pinned to that replica count with a one-trial budget. The single worker keeps `recommend/pilot.yaml`; the second choice uses `recommend/replicas-N.yaml`, where `N` is its replica count. The plan reports the actual candidate count and lists both config and result paths. Collection uses the matching `tp`, `pure_tp`, `dep`, or `tep` preset and exact worker GPU count. Before execution, prepare the real checkpoint and the pinned runtime using the existing [FPM collection guide](../python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md). The packaged collector invokes a Generator-resolved Dynamo/vLLM deployment and needs the corresponding GPU resources, deployment configuration, permissions, and model access. Invoking its command locally does not create that environment. `commands.json` publishes the guarded `aisimulate onboard collect-fpm --execute` command for collection, alongside a read-only collector planning command. @@ -69,7 +105,21 @@ aisimulate onboard collect-fpm \ --output-dir ./aisimulate-support --execute ``` -Execution requires a matching saved plan. The request records model and runtime revisions; this setup does not download a pinned checkpoint or check the installed runtime before execution. Keep the actual checkpoint and runtime consistent with the request before collecting or predicting. Successful formal collection and resume verify that the published data's pod-reported runtime version exactly matches `framework_version`, including suffixes such as `+cu128`. A mismatch exits 1 and preserves collection artifacts. Use the declared runtime or create a new request and plan for the observed version; generated configs are never silently retargeted. Diagnostic smoke runs do not publish or verify formal data. +Execution requires a matching saved plan. Creating the initial plan records model and runtime revisions without downloading a pinned checkpoint or inspecting the running runtime. During profile-based collection execution, the collector records the observed container's vLLM version before checking it against the profile's literal backend version and before benchmarking. Keep the actual checkpoint consistent with the declared model revision. + +Custom and development builds are supported with an exact full-version pin. Local and development suffixes are part of the identity: `0.25.1+custom` must match `0.25.1+custom`, and `0.25.2.dev3` must match `0.25.2.dev3`. Inspect the version **inside the collection container**: + +```bash +python3 -c "import importlib.metadata; print(importlib.metadata.version('vllm'))" +``` + +Use that complete string for both the request's `--framework-version` and the collection profile's deployment `backend_version`. On a mismatch, the error reports the expected and observed strings, and the observed version remains in `collector-provenance.json`. Recover by using the pinned runtime, or by creating a new collection profile pinned to the observed full version and regenerating the request and collection plan in a new output directory. + +For a profile-based campaign, the collector validates its resolved topology and precision against the supplied profile before execution. It does not relabel an FP8 cell as BF16 or change checkpoint quantization to satisfy the profile. A mismatch reports the conflicting field and requires a matching profile/runtime or a supported collector configuration. Resource admission uses the supplied bounds without constructing an analytical model. Profile contents participate in the collector's frozen-plan identity. + +New profile-based collection currently uses checkpoint-native weight, FMHA and KV precision without consulting op-level timing tables. The requested KV dtype must match that checkpoint-native dtype. An older FPM profile can still be valid for prediction while being unsuitable for new collection: the historical MiniMax/H200 BF16 FMHA fallback cell differs from its checkpoint-native FP8 inference. Use the historical identity to query those timings and a matching native profile for new collection; the collector rejects that mismatch. Publish the new FP8 campaign into a clean, separate dataset using a new onboarding output directory. Existing cell IDs omit FMHA precision, and publication retains the first published run for a cell ID. If the historical BF16 dataset already holds that cell ID, publication skips the new FP8 run. A collection profile must select one literal runtime version for the target hardware/backend. + +Successful formal collection and resume verify that the published data's pod-reported runtime version exactly matches `framework_version`, including suffixes such as `+cu128`. A mismatch exits 1 and preserves collection artifacts. Use the declared runtime or create a new request and plan for the observed version; generated configs are never silently retargeted. Diagnostic smoke runs do not publish or verify formal data. Set deployment options directly on `onboard collect-fpm`: `--dynamo-version VERSION`, `--image IMAGE`, `--namespace NAME`, `--model-cache NAME[:MOUNT[:SUBPATH]]`, `--transport nvlink|ib|efa`, and `--image-pull-secret NAME`. The mount, when supplied, is an absolute container path. Prefer an immutable image digest. Supply the same options when previewing, executing, and resuming; deployment settings are part of the collector's frozen-plan identity, so changed settings require a new output directory. Arbitrary collector arguments and engine overrides are not accepted by this command. @@ -77,7 +127,7 @@ For a diagnostic run, add `--execute --smoke`; `--limit N` also requires `--smok Planning and collection reject concurrent onboarding operations. The persistent `.support.lock` file uses an OS advisory lock; ownership is released when the process exits, including after an abrupt termination. Leave the file in place. Request validation, saved onboarding-plan checks, and collector input resolution exit 2. Failures after collector execution starts, including a frozen checkpoint identity mismatch, exit 1 with a concise message. Interruption exits 130. -The collector narrows initial prefill sampling with the pilot's input-token and concurrency bounds. Decode uses the collector's existing profile; a four-request synthetic pilot does not imply four timing samples or a short decode campaign. Inspect the generated command and collector plan before committing GPU time. Successful formal collection publishes the FPM Parquet file and metadata pair into the plan's local systems data directory; diagnostic success alone does not provide that pair. +The collector narrows initial prefill sampling with the pilot's input-token and concurrency bounds. With an FPM profile, both the sample batch size and total new-token axis are capped by the same rank-local scheduler limits used in the generated prediction and recommendation configs. DEP does not multiply these limits by the GPU count; coverage of the worker's summed FPM queries still requires validation. Decode uses the collector's existing sampling profile within the declared resource envelope; a four-request synthetic pilot does not imply four timing samples or a short decode campaign. Inspect the generated command and collector plan before committing GPU time. Successful formal collection publishes the FPM Parquet file and metadata pair into the plan's local systems data directory; diagnostic success alone does not provide that pair. ## Run the generated ordinary configurations @@ -107,7 +157,7 @@ from aisimulate.support.schema import SupportRequest root = Path("./aisimulate-support").resolve() request = SupportRequest.from_yaml(root / "request.yaml") check_plan(request, root) -max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.search.tensor_parallel) +max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.worker_gpus) replicas = list(dict.fromkeys((1, max_replicas)))[:request.search.max_candidates] statuses = [] for count in replicas: @@ -129,17 +179,39 @@ The loop reports each command's exit status and attempts all commands, including Run either the individual recommendation command or the loop against fresh result directories. The loop also handles a one-candidate plan and does not execute entries from `commands.json`. Results remain separate under `recommend-results/pilot` and, when present, `recommend-results/replicas-N`; compare their objective and latency results for the same workload. This plan does not produce a combined ranking or search additional TP sizes or scheduler settings. -The generated configurations select `engine.workers.aggregated.timing.estimation_mode: fpm_interpolation` with `fallback_policy: deny` and set `engine.systems_paths` to a list containing the plan's absolute local systems directory. The same root supplies hardware and collected FPM data. Recommendation preserves it in exported prediction configs. Moving the plan to another machine requires updating absolute paths or regenerating it there. The single-root `engine.systems_path` input is an alias added by onboarding; saved configurations use the canonical `engine.systems_paths` list. Both spellings resolve relative local paths against the working directory when the configuration is loaded, so saving and reloading from another directory preserves the selected roots. The canonical list's `default` entry continues to select packaged data. A worker’s `timing.systems_paths` overrides engine-level roots; performance metadata records those effective roots. Empty or whitespace-only roots are rejected by the shared root type in both prediction and recommendation configurations. +The generated configurations select `engine.workers.aggregated.timing.estimation_mode: fpm_interpolation` with `fallback_policy: deny`. With a profile, they also set `estimator_config.fpm_interpolation.method: direct`. `engine.systems_paths` contains the plan's absolute local systems directory, which supplies hardware and collected FPM data. Recommendation preserves the resolved root, interpolation method, and complete profile in exported prediction configs. Moving the plan to another machine requires updating absolute paths or regenerating it there. Existing configurations using the single-root `engine.systems_path` input are normalized to `engine.systems_paths` when saved. Both spellings resolve relative local paths against the working directory when the configuration is loaded, so saving and reloading from another directory preserves the selected roots. The canonical list's `default` entry continues to select packaged data. A worker’s `timing.systems_paths` overrides engine-level roots; performance metadata records those effective roots. Empty or whitespace-only roots are rejected by the shared root type in both prediction and recommendation configurations. These configurations target AISimulate's standalone `predict` and `recommend` commands. Dynamo's replay Planner AIC session adapter does not currently forward custom systems roots; using that downstream path requires a separate adapter update. -Ordinary `predict` and `recommend` commands retain their existing behavior and defaults. Their generated configs can be loaded and edited through the public configuration schema. In this revision, missing model registration or data may still prevent execution; a successful simulation is not an accuracy result. Compare its output with an independent run of the same model, runtime, topology, and workload to assess accuracy. +Ordinary `predict` and `recommend` retain their existing defaults when no profile is supplied. Generated configs can be edited through the public schema. A recommendation with a profile enumerates only its declared deployment tuples that fit the GPU budget; it preserves the complete inline profile and interpolation choice in exported prediction configs. Default scheduler search ranges may exceed a profile's envelope, so pin or bound those domains explicitly. A successful simulation is not an accuracy result. Compare its output with an independent run of the same model, runtime, topology, and workload to assess accuracy. ## Choose the model execution route Hand off the saved request, pinned model configuration, and plan. Both routes need a canonical checkpoint identity, effective precision and topology, correct weight and KV-cache accounting, and matching whole-forward FPM measurements. Collected timings alone do not establish memory fit. -- **Registered-model/SOL route, available in this revision:** reuse a compatible analytical class or follow [How to Add a New Model](../python/aisimulate/docs/add_a_new_model.md) when choosing to add one. Verify its operation graph, memory and cache accounting, and native FPM SOL execution. A dedicated class is an option for this route, not the intended prerequisite for every FPM onboarding. -- **Class-independent route, planned in the separate decoupling change:** resolve model and resource metadata without constructing an operation graph. Use direct measured-time interpolation, including wider two-sided KV brackets at the same batch size when both neighboring prompt curves cover the query. Missing brackets or unsupported metadata must produce explicit errors. This route is not implemented by the guided foundation; 2D interpolation remains experimental. +- **Registered-model/SOL route:** reuse a compatible analytical class or follow [How to Add a New Model](../python/aisimulate/docs/add_a_new_model.md) when choosing to add one. Verify its operation graph, memory/cache accounting, and native FPM SOL execution. +- **Class-independent direct route:** supply the identity/resource profile and configure the worker as shown below. The guided planner selects this route whenever a profile is supplied. No operation graph is constructed for resources, timing, or recommendation candidates. + +```yaml +timing: + type: default + estimation_mode: fpm_interpolation + fallback_policy: deny + estimator_config: + fpm_interpolation: + method: direct +``` + +Within `estimator_config.fpm_interpolation`, `method: auto` retains SOL for a registered architecture and chooses direct for a verified architecture without a registered class when a profile is supplied. Explicit `sol` requires a registered class; explicit `direct` requires a profile and is also valid for registered architectures. Rust validates and selects the method through `RustForwardPassPerfModel.best_available(ForwardPassPerfModelConfig)`, then pins it for queries and exported configurations. A model-construction error does not trigger a silent change of method. The engine's top-level `estimation_mode: auto` retains the normal priority of operation-level estimation, FPM interpolation, then FPM regression; providing a profile does not change that priority. + +Direct readiness requires genuine measurements for each operation's phase. Prefill rows cannot make an operation with no genuine decode rows ready, and vice versa. Construction tries later configured systems roots when a root lacks a required phase; if none is usable, it reports the missing phase. An operation that needs only prefill does not require decode measurements, and vice versa. Readiness does not guarantee coverage of every subsequent query. + +Direct timing first uses an exact point or interpolation within a measured curve. For cross-KV prefill interpolation, it selects the nearest measured KV curve on each side at the same batch size, considering only curves that cover the requested total new-token count. This direct bracket has no KV distance limit; SOL's site-distance guard is separate. Decode interpolation respects the measured batch/capture domain. Both phases exclude synthetic `fake_fallback` rows, including healed/extrapolated values. Missing two-sided support, unmeasured prefill batches and out-of-domain queries fail explicitly. The direct route does not apply SOL-dependent prefill batch clamping or general extrapolation; 2D interpolation remains experimental. + +Use `predict --detail source --format json` to inspect executed direct lookups and their measured support, or inspect the saved `prediction.json` after requesting source details. The default text output shows operation sources and fallbacks. The canonical model's `estimate_forward_pass_detailed(metrics)` provides the same lookup evidence with rank reduction, mixed-pass baseline and online-correction context. Replay and recommendation APIs accept `ReplayOutputRequirements(capture_performance_diagnostics=True)` to retain counted query records in their returned metadata. Ordinary runs keep no query history; requested capture retains every distinct query without truncation. See [Direct FPM query evidence](core-api.md#direct-fpm-query-evidence) for the result fields and recommendation API. Exact and interpolated values keep the existing `silicon` source tag. Provenance alone does not qualify interpolation accuracy or prove which CUDA graph a measured forward used; accuracy gates and graph-aware interpolation remain deferred until new measurements establish the relevant behavior. + +Saved legacy interpolation controls merge with explicitly supplied correction, regression, and sampling controls. Disjoint settings and agreeing overlapping values are retained; conflicting explicit values fail with the setting's name. Defaults apply after these explicit inputs merge. See [Migrating saved configuration](core-api.md#migrating-saved-configuration) for the public migration API. + +Per-operation silicon profiling described in the model guide is not required by either FPM route. The workflow collects whole-forward timings, then verifies prediction and recommendation for the exact target deployment. Report timing coverage and interpolation error separately; successful simulation alone does not establish measured accuracy. -Per-operation silicon profiling described in the model guide is not required by either FPM route. The intended self-service workflow collects whole-forward timings, then verifies prediction and recommendation for the exact target deployment. Collection bootstrap can already resolve some unregistered model configurations; that does not establish that this revision's ordinary FPM prediction path can construct them. +Accuracy gates and interpolation corrections remain deferred until new measurements are collected with the provenance from [Dynamo #15110](https://github.com/ai-dynamo/dynamo/pull/15110). Those records support investigation; they do not by themselves establish the cause of earlier interpolation errors. If a reproducible CUDA graph boundary causes the discrepancy, interpolation must respect that boundary. No graph-aware partitioning or accuracy policy is introduced here. diff --git a/docs/python-source-migration.md b/docs/python-source-migration.md index c1382ecd7..0726816ae 100644 --- a/docs/python-source-migration.md +++ b/docs/python-source-migration.md @@ -19,6 +19,11 @@ Application SDK modules that expose core types delegate to the canonical core module objects, preserving registry, exception, and cache identity. Core code does not import the application orchestration or legacy CLI layers. +FPM profile schemas and quantization identities live in the lightweight +`aisimulate_core.fpm_profile` and `aisimulate_core.quantization` modules. +The application re-exports those same types. Importing metadata does not load +the native extension; native exports on `aisimulate_core` load on first use. + ## Replacement imports | Removed import | Replacement | diff --git a/python/aisimulate/collector/fpm_forward/capabilities.py b/python/aisimulate/collector/fpm_forward/capabilities.py index 16f58320c..74e97e849 100644 --- a/python/aisimulate/collector/fpm_forward/capabilities.py +++ b/python/aisimulate/collector/fpm_forward/capabilities.py @@ -183,8 +183,13 @@ def resolve_model_capability( requested_kv_cache_dtypes: tuple[str, ...], model_config_path: str | None = None, database_version: str | None = None, + checkpoint_native_dtypes: bool = False, ) -> ModelCapabilityProfile: - """Resolve a conservative matrix profile from AIC code and perf data.""" + """Resolve checkpoint identity with optional legacy op-data dtype evidence. + + Profile collection supplies a literal runtime version and retains native + checkpoint precision without requiring previously collected timing data. + """ # The CLI maps an empty --fpm-kv-cache-dtypes to ("auto",); a programmatic # caller passing an empty tuple would otherwise crash below on @@ -253,6 +258,8 @@ def resolve_model_capability( f"its AIC-inferred GEMM mode {gemm!r}, got {sorted(requested_weights)}" ) + if checkpoint_native_dtypes and database_version is None: + raise ValueError("checkpoint-native FPM collection requires a literal profile backend_version") version = database_version or get_latest_database_version(system=system, backend=backend) if not version: raise ValueError(f"no AIC database version is available for system={system!r}, backend={backend!r}") @@ -287,31 +294,43 @@ def resolve_model_capability( fmha_resolution_by_kv_dtype={native_kv: "checkpoint_native"}, ), ) - database = get_database(system, backend, version) - if database is None: - raise ValueError(f"failed to load AIC database for system={system!r}, backend={backend!r}, version={version!r}") - supported = getattr(database, "supported_quant_mode", {}) or {} context_op, generation_op = _attention_ops(source_name) + requested_kv = tuple(dict.fromkeys(_normalize_kv_dtype(value) for value in requested_kv_cache_dtypes)) + resolved_kv = tuple(dict.fromkeys(native_kv if value == "auto" else value for value in requested_kv)) + database = None + if checkpoint_native_dtypes: + if resolved_kv != (native_kv,): + raise ValueError( + "profile-based FPM collection currently requires the checkpoint-native KV dtype " + f"{native_kv!r}; got {list(resolved_kv)}" + ) + else: + database = get_database(system, backend, version) + if database is None: + raise ValueError( + f"failed to load AIC database for system={system!r}, backend={backend!r}, version={version!r}" + ) + supported = getattr(database, "supported_quant_mode", {}) or {} - _require_mode(supported, "gemm", gemm) - if is_moe: - _require_mode(supported, "moe", moe) + _require_mode(supported, "gemm", gemm) + if is_moe: + _require_mode(supported, "moe", moe) - context_modes = _supported_modes(supported, context_op) - if not context_modes: - raise ValueError(f"AIC database has no context-attention dtype evidence for template op {context_op!r}") + context_modes = _supported_modes(supported, context_op) + if not context_modes: + raise ValueError(f"AIC database has no context-attention dtype evidence for template op {context_op!r}") - generation_modes = _supported_modes(supported, generation_op) - if not generation_modes: - raise ValueError(f"AIC database has no generation-attention dtype evidence for template op {generation_op!r}") - requested_kv = tuple(dict.fromkeys(_normalize_kv_dtype(value) for value in requested_kv_cache_dtypes)) - resolved_kv = tuple(dict.fromkeys(native_kv if value == "auto" else value for value in requested_kv)) - unsupported_kv = [value for value in resolved_kv if value not in generation_modes] - if unsupported_kv: - raise ValueError( - f"AIC database does not support KV-cache dtype(s) {unsupported_kv} for {generation_op}; " - f"supported modes: {sorted(generation_modes)}" - ) + generation_modes = _supported_modes(supported, generation_op) + if not generation_modes: + raise ValueError( + f"AIC database has no generation-attention dtype evidence for template op {generation_op!r}" + ) + unsupported_kv = [value for value in resolved_kv if value not in generation_modes] + if unsupported_kv: + raise ValueError( + f"AIC database does not support KV-cache dtype(s) {unsupported_kv} for {generation_op}; " + f"supported modes: {sorted(generation_modes)}" + ) # fmha_quant_mode names the AIC data slice a template transfers # utilization from, so it is resolved PER KV DTYPE against joint @@ -326,6 +345,18 @@ def resolve_model_capability( fmha_by_kv: dict[str, str] = {} fmha_resolution_by_kv: dict[str, str] = {} for kv_name in resolved_kv: + if checkpoint_native_dtypes: + # Profile collection uses the checkpoint mode, without borrowing + # an op-table slice. This is a whole-forward identity, not a + # claim that every attention intermediate has this dtype: vLLM + # v0.25.1 (752a3a504485790a2e8491cacbb35c137339ad34), + # vllm/v1/attention/backends/mla/flashmla_sparse.py:879-894 + # selects its FP8-cache route while that route can use BF16 + # prefill intermediates. Actual backend dispatch stays in vLLM. + fmha_by_kv[kv_name] = inferred_fmha + fmha_resolution_by_kv[kv_name] = "checkpoint_native" + continue + assert database is not None joint_modes = context_fmha_supported_modes(database, context_op, common.KVCacheQuantMode[kv_name]) if inferred_fmha in joint_modes: fmha_by_kv[kv_name] = inferred_fmha diff --git a/python/aisimulate/collector/fpm_forward/config.py b/python/aisimulate/collector/fpm_forward/config.py index 826ac7b51..8d62210f7 100644 --- a/python/aisimulate/collector/fpm_forward/config.py +++ b/python/aisimulate/collector/fpm_forward/config.py @@ -404,6 +404,12 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: "FPM forward collection", "Whole-model forward-pass planning, execution, and publication.", ) + group.add_argument( + "--fpm-model-profile", + default=None, + metavar="PATH", + help="JSON or YAML FPM model profile with exact deployment identities and rank-local resource bounds.", + ) group.add_argument( "--fpm-enforce-eager", action="store_true", @@ -651,6 +657,7 @@ def reject_fpm_arguments_without_fpm(args: argparse.Namespace) -> None: "fpm_weight_quantizations", "fpm_kv_cache_dtypes", "fpm_model_config", + "fpm_model_profile", "fpm_tp_sizes", "fpm_pp_sizes", "fpm_dp_sizes", diff --git a/python/aisimulate/collector/fpm_forward/entry.py b/python/aisimulate/collector/fpm_forward/entry.py index 96f226bd3..3abbe1b3c 100644 --- a/python/aisimulate/collector/fpm_forward/entry.py +++ b/python/aisimulate/collector/fpm_forward/entry.py @@ -7,10 +7,13 @@ import argparse import copy +from pathlib import Path from typing import Any import yaml +from aisimulate_core.sdk.fpm_profile import load_fpm_profile + from .config import FPMCollectionOptions from .planner import FPMCollectionPlan, build_collection_plan @@ -140,6 +143,13 @@ def resolve_inputs(args: argparse.Namespace, case_plan) -> ResolvedFPMInputs: options = FPMCollectionOptions.from_args(args) generator_overrides = copy.deepcopy(_load_generator_overrides(args)) + profile_path = getattr(args, "fpm_model_profile", None) + fpm_profile = None + if profile_path is not None: + payload = yaml.safe_load(Path(profile_path).expanduser().read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise ValueError("--fpm-model-profile must contain a JSON or YAML mapping") + fpm_profile = load_fpm_profile(payload) plan = build_collection_plan( backend=args.backend, model_path=case_plan.model_path, @@ -149,6 +159,7 @@ def resolve_inputs(args: argparse.Namespace, case_plan) -> ResolvedFPMInputs: model_architecture=case_plan.model_architecture, has_model_cases=bool(case_plan.model_cases_paths), model_config_path=getattr(args, "fpm_model_config", None), + fpm_profile=fpm_profile, collector_config={}, generator_overrides=generator_overrides, ) diff --git a/python/aisimulate/collector/fpm_forward/memory_admission.py b/python/aisimulate/collector/fpm_forward/memory_admission.py index 9b039e249..df741d2c1 100644 --- a/python/aisimulate/collector/fpm_forward/memory_admission.py +++ b/python/aisimulate/collector/fpm_forward/memory_admission.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Conservative AIC memory admission for FPM parallel topologies.""" +"""Conservative memory admission for FPM parallel topologies.""" from __future__ import annotations @@ -11,6 +11,8 @@ from aisimulate.sdk.memory import KVCacheEstimator from aisimulate_core.sdk.errors import PerfDataNotAvailableError +from aisimulate_core.sdk.fpm_profile import FpmModelProfile +from aisimulate_core.sdk.perf_database import load_system_spec from .capabilities import ModelCapabilityProfile from .types import ParallelTopology @@ -25,9 +27,10 @@ class DTypeMemoryEstimate: estimated_non_kv_bytes: int | None gpu_capacity_bytes: int | None reason: str + provenance: str | None = None def to_dict(self) -> dict[str, object]: - return { + payload = { "kv_cache_dtype": self.kv_cache_dtype, "disposition": self.disposition, "estimated_non_kv_bytes": self.estimated_non_kv_bytes, @@ -39,6 +42,9 @@ def to_dict(self) -> dict[str, object]: ), "reason": self.reason, } + if self.provenance is not None: + payload["provenance"] = self.provenance + return payload @dataclass(frozen=True, slots=True) @@ -48,17 +54,29 @@ class TopologyMemoryDecision: max_new_tokens: int estimates: tuple[DTypeMemoryEstimate, ...] reason: str + source: str = "aic_native_configured_max_new_tokens" + max_batch_size: int = 1 + profile_max_num_tokens: int | None = None def to_dict(self) -> dict[str, object]: + envelope = { + "max_new_tokens": self.max_new_tokens, + "aic_max_num_tokens": self.max_new_tokens, + "max_batch_size": self.max_batch_size, + } + if self.profile_max_num_tokens is not None: + envelope = { + "requested_prefill_tokens": self.max_new_tokens, + "max_num_tokens": self.profile_max_num_tokens, + "max_batch_size": self.max_batch_size, + "scope": "rank_local", + "cuda_graph_reservation": "runtime_profiled_separately", + } return { "topology": self.topology.to_dict(), "disposition": self.disposition, - "source": "aic_native_configured_max_new_tokens", - "activation_envelope": { - "max_new_tokens": self.max_new_tokens, - "aic_max_num_tokens": self.max_new_tokens, - "max_batch_size": 1, - }, + "source": self.source, + "activation_envelope": envelope, "estimates": [estimate.to_dict() for estimate in self.estimates], "reason": self.reason, } @@ -73,6 +91,8 @@ def _estimate_dtype( topology: ParallelTopology, kv_cache_dtype: str, max_new_tokens: int, + fpm_profile: FpmModelProfile | None = None, + max_batch_size: int | None = None, ) -> DTypeMemoryEstimate: # Planner-owned capability data fails closed: resolve_model_capability # guarantees an fmha mapping for every resolved KV dtype, so a missing @@ -84,6 +104,41 @@ def _estimate_dtype( raise ValueError( f"capability profile invariant violated: kv dtype {kv_cache_dtype!r} has no fmha mapping" ) from error + if fpm_profile is not None: + deployment = fpm_profile.select( + model=model_path, + system=system, + backend=backend, + backend_version=capability.aic_database_version, + tp_size=topology.tp, + pp_size=topology.pp, + attention_dp_size=topology.dp, + moe_tp_size=topology.moe_tp, + moe_ep_size=topology.moe_ep, + cp_size=topology.cp, + ) + resources = deployment.resources + resources.validate_envelope( + max_num_tokens=max_new_tokens, + max_batch_size=max_batch_size or resources.max_batch_size, + ) + capacity = math.floor(float(load_system_spec(system)["gpu"]["mem_capacity"])) + if capacity <= 0: + raise ValueError(f"invalid GPU memory capacity for {system!r}: {capacity}") + rejected = resources.non_kv_bytes >= capacity + return DTypeMemoryEstimate( + kv_cache_dtype=kv_cache_dtype, + disposition="rejected" if rejected else "admitted", + estimated_non_kv_bytes=resources.non_kv_bytes, + gpu_capacity_bytes=capacity, + reason=( + "declared rank-local non-KV resource bound is not below GPU capacity" + if rejected + else "declared rank-local non-KV resource bound is below GPU capacity; " + "runtime profiling still determines KV capacity and CUDA-graph memory" + ), + provenance=resources.provenance, + ) try: breakdown = KVCacheEstimator.from_request( model_path, @@ -168,13 +223,16 @@ def filter_memory_infeasible_topologies( capability: ModelCapabilityProfile, topologies: tuple[ParallelTopology, ...], max_new_tokens: int, + fpm_profile: FpmModelProfile | None = None, + max_batch_size: int | None = None, ) -> tuple[tuple[ParallelTopology, ...], tuple[TopologyMemoryDecision, ...]]: """Drop topologies that cannot fit the configured max-new-token envelope. This is intentionally a one-sided generation-time filter. A topology is - rejected only when every requested KV dtype has a successful AIC estimate - and all estimates exceed rank-local physical capacity. Unknown estimates - remain runnable so the target runtime stays authoritative. + rejected only when every requested KV dtype has a successful estimate + and all estimates exceed rank-local physical capacity. Without a profile, + unknown AIC estimates remain runnable. Supplied profiles must explicitly + cover the requested deployment and envelope; validation errors propagate. """ if max_new_tokens < 1: @@ -184,6 +242,22 @@ def filter_memory_infeasible_topologies( admitted = [] rejected_capacity = [] for topology in topologies: + resources = ( + fpm_profile.select( + model=model_path, + system=system, + backend=backend, + backend_version=capability.aic_database_version, + tp_size=topology.tp, + pp_size=topology.pp, + attention_dp_size=topology.dp, + moe_tp_size=topology.moe_tp, + moe_ep_size=topology.moe_ep, + cp_size=topology.cp, + ).resources + if fpm_profile is not None + else None + ) estimates = tuple( _estimate_dtype( backend=backend, @@ -193,6 +267,8 @@ def filter_memory_infeasible_topologies( topology=topology, kv_cache_dtype=kv_cache_dtype, max_new_tokens=max_new_tokens, + fpm_profile=fpm_profile, + max_batch_size=max_batch_size, ) for kv_cache_dtype in capability.dtype.kv_cache_dtypes ) @@ -216,6 +292,9 @@ def filter_memory_infeasible_topologies( max_new_tokens=max_new_tokens, estimates=estimates, reason=reason, + source="fpm_profile_declared" if fpm_profile is not None else "aic_native_configured_max_new_tokens", + max_batch_size=resources.max_batch_size if resources is not None else 1, + profile_max_num_tokens=resources.max_num_tokens if resources is not None else None, ) ) @@ -237,17 +316,19 @@ def filter_memory_infeasible_topologies( f"{best.gpu_capacity_bytes / 2**30:.2f} GiB" ) logger.warning( - "fpm_forward: dropped %d/%d topologies (AIC configured max-new-token non-KV memory " + "fpm_forward: dropped %d/%d topologies (%s configured max-new-token non-KV memory " "exceeds GPU capacity, system=%s, max_new_tokens=%d): %s", len(rejected_capacity), len(topologies), + "FPM profile" if fpm_profile is not None else "AIC", system, max_new_tokens, "; ".join(details), ) if not admitted: + source = "FPM profile" if fpm_profile is not None else "AIC" raise ValueError( - "AIC max-new-token memory admission rejected every FPM topology; " + f"{source} max-new-token memory admission rejected every FPM topology; " f"model={model_path!r}, system={system!r}, max_new_tokens={max_new_tokens}" ) return tuple(admitted), tuple(decisions) diff --git a/python/aisimulate/collector/fpm_forward/planner.py b/python/aisimulate/collector/fpm_forward/planner.py index 27c1680a5..3e3deb194 100644 --- a/python/aisimulate/collector/fpm_forward/planner.py +++ b/python/aisimulate/collector/fpm_forward/planner.py @@ -15,14 +15,15 @@ import logging import os import subprocess -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path, PurePosixPath from typing import Any from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY, execution_identity +from aisimulate_core.sdk.fpm_profile import FpmDeploymentProfile, FpmModelProfile, load_fpm_profile from .capabilities import ModelCapabilityProfile, ResolvedDTypeProfile, resolve_model_capability -from .config import FPMCollectionOptions +from .config import PARALLEL_AXES, FPMCollectionOptions from .memory_admission import TopologyMemoryDecision, filter_memory_infeasible_topologies from .topology import enumerate_fpm_topologies, topology_strategy from .types import ParallelTopology @@ -489,12 +490,35 @@ class FPMCollectionPlan: backend_policies: tuple[BackendPolicy, ...] cells: tuple[FPMCell, ...] sha256: str + _fpm_profile_json: str | None = field(default=None, repr=False) + + @property + def fpm_profile(self) -> FpmModelProfile | None: + """Return detached profile metadata from the frozen collection input.""" + return load_fpm_profile(self._fpm_profile_json) if self._fpm_profile_json is not None else None + + def deployment_profile(self, cell: FPMCell) -> FpmDeploymentProfile | None: + profile = self.fpm_profile + if profile is None: + return None + return profile.select( + model=self.model_path, + system=self.system, + backend=self.backend, + backend_version=self.capability.aic_database_version, + tp_size=cell.topology.tp, + pp_size=cell.topology.pp, + attention_dp_size=cell.topology.dp, + moe_tp_size=cell.topology.moe_tp, + moe_ep_size=cell.topology.moe_ep, + cp_size=cell.topology.cp, + ) def to_dict(self) -> dict[str, object]: explicit_points = ( json.loads(self.options.benchmark_points_json) if self.options.benchmark_points_json is not None else None ) - return { + payload = { "schema_name": "aic_fpm_collection_plan", "schema_version": 11, "backend": self.backend, @@ -552,6 +576,9 @@ def to_dict(self) -> dict[str, object]: }, "sha256": self.sha256, } + if self._fpm_profile_json is not None: + payload["fpm_profile"] = json.loads(self._fpm_profile_json) + return payload def _cell_id( @@ -593,12 +620,31 @@ def build_collection_plan( model_architecture: str | None = None, has_model_cases: bool = True, model_config_path: str | None = None, + fpm_profile: FpmModelProfile | dict[str, Any] | None = None, collector_config: dict[str, Any] | None = None, generator_overrides: dict[str, Any] | None = None, ) -> FPMCollectionPlan: if backend != "vllm": raise ValueError("FPM Generator V1 currently supports only backend=vllm") - collector_config = collector_config or {} + collector_config = dict(collector_config or {}) + profile = load_fpm_profile(fpm_profile) if fpm_profile is not None else None + if profile is not None: + if profile.model != model_path: + raise ValueError( + f"FPM profile model identity mismatch: requested {model_path!r}, profile={profile.model!r}" + ) + versions = { + deployment.backend_version + for deployment in profile.deployments + if deployment.system == system and deployment.backend == backend + } + if "aic_database_version" not in collector_config: + if len(versions) != 1: + raise ValueError( + "FPM collection profile must identify one runtime version for the target " + f"{system}/{backend}; found {sorted(versions)}" + ) + collector_config["aic_database_version"] = next(iter(versions)) generator_config_sha256 = _canonical_hash(generator_overrides or {}) capability = resolve_model_capability( backend=backend, @@ -613,6 +659,7 @@ def build_collection_plan( database_version=( str(collector_config["aic_database_version"]) if "aic_database_version" in collector_config else None ), + checkpoint_native_dtypes=profile is not None, ) execution = execution_identity( capability.model_config.payload, @@ -638,6 +685,18 @@ def build_collection_plan( options=options, allow_pure_tp=capability.allow_pure_tp, ) + policies = _backend_policies(options, collector_config, backend=backend) + if profile is not None: + _validate_profile_identities( + profile, + capability, + candidate_topologies, + policies, + model_path, + system, + backend, + max_decode_batch_size=options.max_decode_batch_size, + ) topologies, topology_memory_admission = filter_memory_infeasible_topologies( backend=backend, model_path=model_path, @@ -645,8 +704,9 @@ def build_collection_plan( capability=capability, topologies=candidate_topologies, max_new_tokens=options.prefill_sampling.max_total_prefill_tokens, + fpm_profile=profile, + max_batch_size=options.max_prefill_batch_size, ) - policies = _backend_policies(options, collector_config, backend=backend) weight_quantization = capability.dtype.gemm_quant_mode runnable_dtype_pairs = { (decision.topology, estimate.kv_cache_dtype) @@ -724,6 +784,10 @@ def build_collection_plan( "policies": [policy.to_dict() for policy in policies], "cells": [cell.to_dict() for cell in cells], } + profile_json = None + if profile is not None: + canonical["fpm_profile"] = profile.model_dump(mode="json") + profile_json = json.dumps(canonical["fpm_profile"], sort_keys=True, separators=(",", ":")) return FPMCollectionPlan( backend=backend, model_path=model_path, @@ -738,4 +802,74 @@ def build_collection_plan( backend_policies=policies, cells=cells, sha256=_canonical_hash(canonical), + _fpm_profile_json=profile_json, ) + + +def _validate_profile_identities( + profile: FpmModelProfile, + capability: ModelCapabilityProfile, + topologies: tuple[ParallelTopology, ...], + policies: tuple[BackendPolicy, ...], + model_path: str, + system: str, + backend: str, + *, + max_decode_batch_size: int | None = None, +) -> None: + """Require a resource declaration for every requested cell, before admission.""" + if profile.architecture != capability.architecture: + raise ValueError( + "FPM collection profile architecture mismatch: " + f"checkpoint={capability.architecture!r}, profile={profile.architecture!r}" + ) + identity_fields = ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "comm_quant_mode", + "kv_cache_dtype", + *PARALLEL_AXES, + "moe_backend", + "attention_backend", + "enable_wideep", + "enable_eplb", + ) + for topology in topologies: + deployment = profile.select( + model=model_path, + system=system, + backend=backend, + backend_version=capability.aic_database_version, + tp_size=topology.tp, + pp_size=topology.pp, + attention_dp_size=topology.dp, + moe_tp_size=topology.moe_tp, + moe_ep_size=topology.moe_ep, + cp_size=topology.cp, + ) + if max_decode_batch_size is not None: + deployment.resources.validate_envelope(max_num_tokens=1, max_batch_size=max_decode_batch_size) + for kv_dtype in capability.dtype.kv_cache_dtypes: + for policy in policies: + resolved = [ + capability.dtype.gemm_quant_mode, + capability.dtype.moe_quant_mode, + capability.dtype.fmha_by_kv_dtype[kv_dtype], + capability.dtype.comm_quant_mode, + kv_dtype, + *(str(getattr(topology, axis)) for axis in PARALLEL_AXES), + *(str(value) for value in backend_identity_columns(policy).values()), + ] + if deployment.match_identity() != resolved: + conflicts = "; ".join( + f"{name}: resolved={actual!r}, profile={expected!r}" + for name, actual, expected in zip( + identity_fields, resolved, deployment.match_identity(), strict=True + ) + if actual != expected + ) + raise ValueError( + f"FPM collection profile identity mismatch: {conflicts}; supply a profile matching " + "the collection configuration. The profile does not override serving dispatch." + ) diff --git a/python/aisimulate/collector/fpm_forward/runner.py b/python/aisimulate/collector/fpm_forward/runner.py index 971a5c273..4c3e1f273 100644 --- a/python/aisimulate/collector/fpm_forward/runner.py +++ b/python/aisimulate/collector/fpm_forward/runner.py @@ -446,6 +446,48 @@ def _run_command( return subprocess.CompletedProcess(args, process.returncode, stdout, stderr) +def _attempt_provenance_command( + *, + cell_id: str, + plan_sha256: str, + attempt_id: str, + expected_backend_version: str | None = None, +) -> list[str]: + """Record the observed runtime before checking a profile's pinned version.""" + + payload = json.dumps( + { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + "cell_id": cell_id, + "plan_sha256": plan_sha256, + "attempt_id": attempt_id, + }, + sort_keys=True, + separators=(",", ":"), + ) + script = ( + "import importlib.metadata, json, pathlib, sys; " + "payload = json.loads(sys.argv[1]); " + "payload['runtime'] = {'backend': 'vllm', " + "'backend_version': importlib.metadata.version('vllm')}; " + f"path = pathlib.Path('{FPM_RESULTS_DIR}') / sys.argv[2]; " + "path.write_text(json.dumps(payload, sort_keys=True) + '\\n')" + ) + version_args = [] + if expected_backend_version is not None: + script += ( + "\nactual = payload['runtime']['backend_version']\n" + "if actual != sys.argv[3]:\n" + " raise RuntimeError(" + "f'FPM profile runtime mismatch: actual={actual!r}, expected={sys.argv[3]!r}. '" + "'Use the pinned runtime, or create a new collection profile with '" + "f'backend_version={actual!r} and regenerate the collection plan in a new output directory.')\n" + ) + version_args.append(expected_backend_version) + return ["python3", "-c", script, payload, COLLECTOR_PROVENANCE_FILENAME, *version_args] + + class KubernetesCellRunner: """Own one generated Pod/LWS/PCS and auxiliary resources through deletion.""" @@ -663,27 +705,15 @@ def prepare_attempt( cell_id: str, plan_sha256: str, attempt_id: str, + expected_backend_version: str | None = None, ) -> None: """Clear stale results and bind every Pod to this Collector attempt.""" - payload = json.dumps( - { - "schema_name": "aic_fpm_collector_provenance", - "schema_version": 1, - "cell_id": cell_id, - "plan_sha256": plan_sha256, - "attempt_id": attempt_id, - }, - sort_keys=True, - separators=(",", ":"), - ) - script = ( - "import importlib.metadata, json, pathlib, sys; " - "payload = json.loads(sys.argv[1]); " - "payload['runtime'] = {'backend': 'vllm', " - "'backend_version': importlib.metadata.version('vllm')}; " - f"path = pathlib.Path('{FPM_RESULTS_DIR}') / sys.argv[2]; " - "path.write_text(json.dumps(payload, sort_keys=True) + '\\n')" + command = _attempt_provenance_command( + cell_id=cell_id, + plan_sha256=plan_sha256, + attempt_id=attempt_id, + expected_backend_version=expected_backend_version, ) for pod in pods: self._exec_checked( @@ -704,11 +734,7 @@ def prepare_attempt( ], timeout=60, ) - self._exec_checked( - pod, - ["python3", "-c", script, payload, COLLECTOR_PROVENANCE_FILENAME], - timeout=60, - ) + self._exec_checked(pod, command, timeout=60) def _remote_result_manifest(self, pod: str) -> dict[str, dict[str, int | str]]: script = ( @@ -1222,6 +1248,16 @@ def _cell_generator_overrides( "2", ] ) + if getattr(plan, "fpm_profile", None) is not None: + deployment_profile = plan.deployment_profile(cell) + assert deployment_profile is not None + resources = deployment_profile.resources + # Native point generation sees the declared scheduler bounds before + # any cases are queued. Never discard runtime cases to fit a profile. + if "--max-num-batched-tokens" not in scheduler_args: + scheduler_args.extend(["--max-num-batched-tokens", str(resources.max_num_tokens)]) + if "--max-num-seqs" not in scheduler_args: + scheduler_args.extend(["--max-num-seqs", str(resources.max_batch_size)]) model_args = [] architecture = getattr(getattr(plan, "capability", None), "architecture", None) if architecture == "DeepseekV41ForCausalLM": @@ -2043,11 +2079,17 @@ def _run_collection_impl( ), ], ) + profile_version = ( + {"expected_backend_version": plan.capability.aic_database_version} + if getattr(plan, "fpm_profile", None) is not None + else {} + ) resource.prepare_attempt( pods, cell_id=cell.cell_id, plan_sha256=plan.sha256, attempt_id=attempt_id, + **profile_version, ) _record_points_receipts(resource, pods, plan, cell, attempt_id, phase="before") phase_marks["stage_s"] = round(time.monotonic() - mark, 3) diff --git a/python/aisimulate/collector/fpm_forward/slurm.py b/python/aisimulate/collector/fpm_forward/slurm.py index 0fffe4c6b..aea47d688 100644 --- a/python/aisimulate/collector/fpm_forward/slurm.py +++ b/python/aisimulate/collector/fpm_forward/slurm.py @@ -171,25 +171,25 @@ def _exec(self, unit: str, command: list[str], *, timeout: int): timeout=timeout, ) - def prepare_attempt(self, pods: list[str], *, cell_id: str, plan_sha256: str, attempt_id: str) -> None: - from .native_artifact import COLLECTOR_PROVENANCE_FILENAME + def prepare_attempt( + self, + pods: list[str], + *, + cell_id: str, + plan_sha256: str, + attempt_id: str, + expected_backend_version: str | None = None, + ) -> None: + from .runner import _attempt_provenance_command - payload = json.dumps( - { - "schema_name": "aic_fpm_collector_provenance", - "schema_version": 1, - "cell_id": cell_id, - "plan_sha256": plan_sha256, - "attempt_id": attempt_id, - } - ) - script = ( - "import importlib.metadata,json,pathlib,sys; p=json.loads(sys.argv[1]); " - "p['runtime']={'backend':'vllm','backend_version':importlib.metadata.version('vllm')}; " - "pathlib.Path('/results',sys.argv[2]).write_text(json.dumps(p,sort_keys=True)+'\\n')" + command = _attempt_provenance_command( + cell_id=cell_id, + plan_sha256=plan_sha256, + attempt_id=attempt_id, + expected_backend_version=expected_backend_version, ) for unit in pods: - self._exec(unit, ["python3", "-c", script, payload, COLLECTOR_PROVENANCE_FILENAME], timeout=300) + self._exec(unit, command, timeout=300) def execute(self, pods: list[str], timeout_seconds: int = 14400) -> None: from .runner import CommandScope, _cancel_preserving_interrupt diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index c7d48bd13..08b05d8b5 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -119,6 +119,7 @@ def finish_lowering(value): "aic_pp_size": resolved["pp"], "aic_moe_tp_size": resolved["moe_tp_size"], "aic_moe_ep_size": resolved["moe_ep_size"], + "aic_fpm_profile": resolved.get("fpm_profile"), }.items() if value is not None } @@ -220,6 +221,7 @@ def finish_lowering(value): }, systems_path=capacity_systems_path, cuda_graph_reserved_bytes=lowered.get("cuda_graph_reserved_bytes", 0), + **({"fpm_profile": lowered["aic_fpm_profile"]} if lowered.get("aic_fpm_profile") is not None else {}), **({"diagnostics": memory_diagnostics} if memory_diagnostics is not None else {}), ) return finish_lowering(lowered) @@ -255,6 +257,7 @@ def estimate_num_gpu_blocks( systems_path: str | list[str] | None = None, cuda_graph_reserved_bytes: int = 0, diagnostics: dict[str, Any] | None = None, + fpm_profile: dict[str, Any] | None = None, ) -> int: """Estimate per-rank KV blocks using the replay-wide AIC contract. @@ -332,6 +335,7 @@ def estimate_num_gpu_blocks( }, systems_path=systems_path, cuda_graph_reserved_bytes=cuda_graph_reserved_bytes, + **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), **({"diagnostics": diagnostics} if diagnostics is not None else {}), ) ) @@ -345,9 +349,31 @@ def estimate_kv_bytes_per_token( moe_tp_size: int = 1, moe_ep_size: int = 1, kvcache_quant_mode: str | None = None, + fpm_profile: dict[str, Any] | None = None, + system: str | None = None, + backend: str = "vllm", + backend_version: str | None = None, + attention_dp_size: int = 1, ) -> int: """Derive per-rank KV bytes/token from the resolved Hugging Face config.""" + if fpm_profile is not None: + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + deployment = load_fpm_profile(fpm_profile).select( + model=model_name, + system=system, + backend=backend, + backend_version=backend_version, + tp_size=tp_size, + pp_size=pp_size, + attention_dp_size=attention_dp_size, + moe_tp_size=moe_tp_size, + moe_ep_size=moe_ep_size, + ) + deployment.validate_overrides(kvcache_quant_mode=kvcache_quant_mode) + return deployment.resources.kv_bytes_per_token + from aisimulate_core.sdk.memory import NaiveKVCacheEstimator estimator = NaiveKVCacheEstimator.from_model_path( diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index 7089d0e09..fb590baf3 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -49,6 +49,27 @@ def prediction_to_replay_spec( afd_performance_model=afd_performance_model, ) deployment = _pin_estimator_version_aliases(deployment) + metadata = dict(deployment.performance_model_metadata) + for role, field in ( + ("aggregated", "agg_engine_args"), + ("prefill", "prefill_engine_args"), + ("decode", "decode_engine_args"), + ): + timing = (getattr(deployment, field) or {}).get("timing_model", {}) + if timing.get("provider") == "aic" and "estimation_mode" in timing.get("config", {}): + identity = { + key: value + for key, value in timing["config"].items() + if key + not in { + "gpu_memory_utilization", + "mem_fraction_static", + "free_gpu_memory_fraction", + "cuda_graph_reserved_bytes", + } + } + metadata[role] = {**metadata.get(role, {}), "provider": "aic", "config": identity} + deployment = replace(deployment, performance_model_metadata=metadata) if config.engine.mode == "disaggregated": assert config.engine.workers.prefill is not None for adapter in (adapter_specs or {}).values(): @@ -515,6 +536,8 @@ def _worker_engine_args( payload["max_model_len"] = ( engine.context_length if isinstance(engine.context_length, int) + else engine.fpm_profile.context_length + if engine.fpm_profile is not None else resolve_model_context_length(engine.model) ) if cache.prefix_match_unit is not None: @@ -567,6 +590,7 @@ def _worker_engine_args( sharded_moe = parallel.moe_tensor * parallel.moe_expert > 1 canonical = ForwardPassPerfModelConfig( model=engine.model, + fpm_profile=engine.fpm_profile.model_dump(mode="json") if engine.fpm_profile is not None else None, system=worker.hardware or engine.hardware, backend=backend, backend_version=engine.backend_version, @@ -658,6 +682,17 @@ def _resolve_kv_bytes_per_token( pp_size=parallel.pipeline, moe_tp_size=parallel.moe_tensor, moe_ep_size=parallel.moe_expert, + **( + { + "fpm_profile": engine.fpm_profile.model_dump(mode="json"), + "system": worker.hardware or engine.hardware, + "backend": engine.backend, + "backend_version": engine.backend_version, + "attention_dp_size": parallel.attention_data, + } + if engine.fpm_profile is not None + else {} + ), kvcache_quant_mode=worker.timing.kvcache_quant_mode or engine.kvcache_quant_mode, ) diff --git a/python/aisimulate/src/aisimulate/config/common.py b/python/aisimulate/src/aisimulate/config/common.py index 1d9fa8df1..493f40066 100644 --- a/python/aisimulate/src/aisimulate/config/common.py +++ b/python/aisimulate/src/aisimulate/config/common.py @@ -38,6 +38,11 @@ def _normalize_systems_root(value: str) -> str: SystemsRoot = Annotated[str, Field(strict=True, min_length=1), AfterValidator(_normalize_systems_root)] +def requested_backend_version(versions: str | dict[str, str] | None, backend: str) -> str | None: + """Return the version pin for one backend; ``None`` means resolve latest.""" + return versions.get(backend) if isinstance(versions, dict) else versions + + T = TypeVar("T") PositiveFiniteFloat = Annotated[float, Field(strict=True, gt=0, allow_inf_nan=False)] PositiveStrictInt = Annotated[int, Field(strict=True, gt=0)] diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index df3658305..f4d1b82b5 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -10,6 +10,8 @@ from pydantic import Field, StrictBool, field_validator, model_validator +from aisimulate.fpm_profile import FpmModelProfile + from .common import ( ENGINE_MODEL_CONTROL_FIELDS, Choices, @@ -19,6 +21,7 @@ SystemsPath, SystemsRoot, is_active_engine_model_control, + requested_backend_version, ) PositiveInt = Annotated[int, Field(strict=True, gt=0)] @@ -495,6 +498,7 @@ def _supported_estimator_policies(self): class EnginePredictionConfig(EstimatorPolicyConfig): mode: EngineMode = "aggregated" model: str + fpm_profile: FpmModelProfile | None = None hardware: str backend: Backend = "vllm" backend_version: str | None = None @@ -523,6 +527,7 @@ def _reject_auto_hardware(cls, value: str) -> str: @model_validator(mode="after") def _validate_roles(self) -> EnginePredictionConfig: + _validate_fpm_profile(self, {self.mode}, {self.backend}) if self.decoder_replay: # Keep the pre-supervision configuration path lightweight. Import # the canonical model identity only when this runtime feature is @@ -710,6 +715,7 @@ class EngineRecommendationConfig(EstimatorPolicyConfig): default_factory=lambda: Choices[EngineMode](choices=["aggregated", "disaggregated"]) ) model: str + fpm_profile: FpmModelProfile | None = None hardware: str backend: Backend | Choices[Backend] = Field(default_factory=lambda: Choices[Backend](choices=["vllm", "sglang"])) backend_version: str | dict[str, str] | None = None @@ -729,6 +735,8 @@ def _validate_nonempty(cls, value: str) -> str: @model_validator(mode="after") def _validate_roles(self) -> EngineRecommendationConfig: modes = set(self.mode.choices) if isinstance(self.mode, Choices) else {self.mode} + backends = set(self.backend.choices) if isinstance(self.backend, Choices) else {self.backend} + _validate_fpm_profile(self, modes, backends) _validate_worker_hardware(modes=modes, workers=self.workers) if "afd" in modes: if modes != {"afd"}: @@ -753,6 +761,55 @@ def _validate_roles(self) -> EngineRecommendationConfig: return self +def _validate_fpm_profile(engine, modes: set[str], backends: set[str]) -> None: + profile = engine.fpm_profile + if profile is None: + return + if engine.model != profile.model: + raise ValueError("engine.model must match engine.fpm_profile.model") + if backends != {"vllm"} or "afd" in modes or engine.workers.encoder is not None: + raise ValueError("FPM profiles support vLLM aggregated/disaggregated decoder workers without AFD or encoders") + backend_version = requested_backend_version(engine.backend_version, "vllm") + if ( + backend_version is None + or not backend_version.strip() + or backend_version.strip() in {"current", "previous", "next"} + ): + raise ValueError("engine.fpm_profile requires a literal engine.backend_version for vllm") + backend_version = backend_version.strip() + if isinstance(engine.backend_version, dict): + engine.backend_version = {**engine.backend_version, "vllm": backend_version} + else: + engine.backend_version = backend_version + if not any(deployment.backend_version == backend_version for deployment in profile.deployments): + raise ValueError("engine.backend_version does not match any FPM profile deployment") + if isinstance(engine.context_length, int) and engine.context_length > profile.context_length: + raise ValueError("engine.context_length exceeds the FPM profile context_length") + for role in ("aggregated", "prefill", "decode"): + worker = getattr(engine.workers, role) + if worker is None: + continue + if worker.timing.type != "default": + raise ValueError("engine.fpm_profile requires default timing for every worker") + if isinstance(worker, WorkerPredictionConfig): + parallel = worker.parallelism + deployment = profile.select( + model=engine.model, + system=worker.hardware or engine.hardware, + backend=engine.backend, + backend_version=backend_version, + tp_size=parallel.tensor, + pp_size=parallel.pipeline, + attention_dp_size=parallel.attention_data, + moe_tp_size=parallel.moe_tensor, + moe_ep_size=parallel.moe_expert, + ) + deployment.resources.validate_envelope( + max_num_tokens=worker.scheduler.max_batched_tokens, + max_batch_size=worker.scheduler.max_sequences, + ) + + def _validate_speculation(engine, *, modes: set[str], backends: set[str]) -> None: if engine.speculation is None: return diff --git a/python/aisimulate/src/aisimulate/detail.py b/python/aisimulate/src/aisimulate/detail.py index b5be9d172..3ec572705 100644 --- a/python/aisimulate/src/aisimulate/detail.py +++ b/python/aisimulate/src/aisimulate/detail.py @@ -115,6 +115,65 @@ def number(value: Any, path: str) -> None: def array(value: Any, path: str) -> None: require(isinstance(value, list), path, "a list") + def fpm_coordinates(value: Any, path: str) -> None: + fields(value, path, {"batch_size", "total_prefill_tokens", "total_kv_read_tokens"}) + number(value["batch_size"], f"{path}.batch_size") + if value["total_prefill_tokens"] is not None: + number(value["total_prefill_tokens"], f"{path}.total_prefill_tokens") + number(value["total_kv_read_tokens"], f"{path}.total_kv_read_tokens") + + def fpm_query(value: Any, path: str) -> None: + fields(value, path, {"phase", "model_path", "decode_baseline", "query", "resolution", "latency_ms", "support"}) + require(value["phase"] in ("prefill", "decode"), f"{path}.phase", "prefill or decode") + require( + isinstance(value["model_path"], str) and bool(value["model_path"]), + f"{path}.model_path", + "a nonempty string", + ) + require(type(value["decode_baseline"]) is bool, f"{path}.decode_baseline", "a boolean") + fpm_coordinates(value["query"], f"{path}.query") + require( + value["resolution"] + in ("exact_lookup", "within_curve_interpolation", "cross_kv_interpolation", "cross_batch_interpolation"), + f"{path}.resolution", + "a supported direct resolution", + ) + number(value["latency_ms"], f"{path}.latency_ms") + array(value["support"], f"{path}.support") + require(bool(value["support"]), f"{path}.support", "nonempty") + for i, support in enumerate(value["support"]): + sp = f"{path}.support[{i}]" + fields(support, sp, {"coordinates", "latency_ms", "weight"}) + fpm_coordinates(support["coordinates"], f"{sp}.coordinates") + number(support["latency_ms"], f"{sp}.latency_ms") + number(support["weight"], f"{sp}.weight") + require(support["weight"] <= 1, f"{sp}.weight", "at most 1") + + def fpm_estimate(value: Any, path: str) -> None: + fields(value, path, {"latency_ms", "native_latency_ms", "correction_factor", "max_rank", "ranks"}) + for key in ("latency_ms", "native_latency_ms", "correction_factor"): + if value[key] is not None: + number(value[key], f"{path}.{key}") + if value["max_rank"] is not None: + require( + type(value["max_rank"]) is int and value["max_rank"] >= 0, + f"{path}.max_rank", + "a nonnegative integer", + ) + array(value["ranks"], f"{path}.ranks") + for i, rank in enumerate(value["ranks"]): + rp = f"{path}.ranks[{i}]" + durations = {"prefill_ms", "decode_ms", "decode_baseline_ms", "marginal_decode_ms"} + fields(rank, rp, {"rank", "latency_ms", "queries"} | durations) + require(type(rank["rank"]) is int and rank["rank"] >= 0, f"{rp}.rank", "a nonnegative integer") + number(rank["latency_ms"], f"{rp}.latency_ms") + for key in durations: + if rank[key] is not None: + number(rank[key], f"{rp}.{key}") + array(rank["queries"], f"{rp}.queries") + for j, query in enumerate(rank["queries"]): + fpm_query(query, f"{rp}.queries[{j}]") + def sol(value: dict[str, Any], path: str) -> None: if value["sol"] is None: string(value["sol_unavailable_reason"], f"{path}.sol_unavailable_reason") @@ -150,10 +209,21 @@ def sol(value: dict[str, Any], path: str) -> None: operation, op, {"name", "latency_ms", "source", "fallbacks", "sol", "sol_unavailable_reason", "latency_to_sol_ratio"}, + {"fpm_estimates"}, ) string(operation["name"], f"{op}.name") string(operation["source"], f"{op}.source") number(operation["latency_ms"], f"{op}.latency_ms") + if "fpm_estimates" in operation: + array(operation["fpm_estimates"], f"{op}.fpm_estimates") + for k, evidence in enumerate(operation["fpm_estimates"]): + ep = f"{op}.fpm_estimates[{k}]" + fields(evidence, ep, {"estimate", "count", "latency_scale"}) + require( + type(evidence["count"]) is int and evidence["count"] > 0, f"{ep}.count", "a positive integer" + ) + number(evidence["latency_scale"], f"{ep}.latency_scale") + fpm_estimate(evidence["estimate"], f"{ep}.estimate") sol(operation, op) if operation["latency_to_sol_ratio"] is not None: number(operation["latency_to_sol_ratio"], f"{op}.latency_to_sol_ratio") @@ -244,7 +314,11 @@ def build_prediction_details(native: dict[str, Any], sections: tuple[str, ...]) { "name": phase["name"], "operations": [ - {key: op[key] for key in ("name", "latency_ms", "source", "fallbacks")} + { + key: op[key] + for key in ("name", "latency_ms", "source", "fallbacks", "fpm_estimates") + if key in op + } for op in phase["operations"] ], } diff --git a/python/aisimulate/src/aisimulate/fpm_profile.py b/python/aisimulate/src/aisimulate/fpm_profile.py new file mode 100644 index 000000000..c81b1ebaa --- /dev/null +++ b/python/aisimulate/src/aisimulate/fpm_profile.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility exports for the lightweight core fpm_profile types.""" + +from aisimulate_core.fpm_profile import FpmDeploymentProfile as FpmDeploymentProfile +from aisimulate_core.fpm_profile import FpmModelProfile as FpmModelProfile +from aisimulate_core.fpm_profile import FpmResourceProfile as FpmResourceProfile +from aisimulate_core.fpm_profile import load_fpm_profile as load_fpm_profile diff --git a/python/aisimulate/src/aisimulate/quantization.py b/python/aisimulate/src/aisimulate/quantization.py new file mode 100644 index 000000000..3fea5d06c --- /dev/null +++ b/python/aisimulate/src/aisimulate/quantization.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility exports for the lightweight core quantization types.""" + +from aisimulate_core.quantization import CommQuantMode as CommQuantMode +from aisimulate_core.quantization import FMHAQuantMode as FMHAQuantMode +from aisimulate_core.quantization import GEMMQuantMode as GEMMQuantMode +from aisimulate_core.quantization import KVCacheQuantMode as KVCacheQuantMode +from aisimulate_core.quantization import MoEQuantMode as MoEQuantMode +from aisimulate_core.quantization import QuantMapping as QuantMapping diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index ff902df5c..e2a58a777 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -24,7 +24,7 @@ from .sweeper.afd_perfmodel import AFDPerformanceModel from .sweeper.config import SmartSearchConfig from .sweeper.provider import InfeasibleCandidate, SweepContext -from .sweeper.replay import ReplaySpec, RunnerFactory +from .sweeper.replay import ReplayOutputRequirements, ReplaySpec, RunnerFactory from .sweeper.result import SweepResult @@ -37,6 +37,7 @@ def run_recommendation( providers: Mapping[str, SimulationConfigAdapter] | None = None, afd_performance_model: AFDPerformanceModel | None = None, show_progress: bool = True, + output_requirements: ReplayOutputRequirements | None = None, ) -> SweepResult: """Run a public recommendation through the existing Sweeper core.""" @@ -49,6 +50,7 @@ def run_recommendation( providers=providers, afd_performance_model=afd_performance_model, show_progress=show_progress, + output_requirements=output_requirements, ) if not in_supervised_process(): return supervised_recommendation(config, kwargs) @@ -64,6 +66,7 @@ def _run_recommendation( providers: Mapping[str, SimulationConfigAdapter] | None = None, afd_performance_model: AFDPerformanceModel | None = None, show_progress: bool = True, + output_requirements: ReplayOutputRequirements | None = None, ) -> SweepResult: from .supervision import checkpoint @@ -102,6 +105,7 @@ def _run_recommendation( adapter_sections = {name: provider.section for name, provider in (providers or {}).items()} sweeper = Sweeper( runner_factory=runner_factory, + output_requirements=output_requirements, providers=compiled_providers, show_progress=show_progress, prediction_config_factory=lambda sample, spec: _candidate_prediction( @@ -151,8 +155,16 @@ def recommendation_to_sweeper( "hardware_sku": hardware, "gpu_budget": optimization.constraints.max_candidate_gpus, "min_gpu_budget": optimization.constraints.min_candidate_gpus, - "context_length": (resolve_model_context_length(model) if context == "max" else context), + "context_length": ( + config.engine.fpm_profile.context_length + if context == "max" and config.engine.fpm_profile is not None + else resolve_model_context_length(model) + if context == "max" + else context + ), } + if engine.get("fpm_profile") is not None: + search_space["fpm_profile"] = engine["fpm_profile"] for name in ( "database_mode", "transfer_policy", @@ -798,6 +810,8 @@ def _candidate_prediction( } if sample.get("systems_paths") is not None: engine["systems_paths"] = sample["systems_paths"] + if sample.get("fpm_profile") is not None: + engine["fpm_profile"] = deepcopy(sample["fpm_profile"]) for name in (*ENGINE_MODEL_CONTROL_FIELDS, "enable_chunked_prefill", "nextn_accepted"): if sample.get(name) is not None: engine[name] = sample[name] diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index f67ad033b..efd627aeb 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -99,6 +99,7 @@ "comm_dtype": ("comm_dtype", "aic_comm_dtype"), "systems_path": ("systems_path",), "forward_model": ("forward_model", "aic_forward_model"), + "fpm_profile": ("fpm_profile", "aic_fpm_profile"), "fpm_parquet_path": ("fpm_parquet_path", "aic_fpm_parquet_path"), "moe_backend": ("aic_moe_backend",), "moe_kernel_source": ("moe_kernel_source", "aic_moe_kernel_source"), @@ -1188,6 +1189,10 @@ def _pop_aic_timing_overrides(rank: dict[str, JSONValue], role: str) -> dict[str continue if target in {"pp", "moe_tp_size", "moe_ep_size", "wideep_num_slots"}: value = _positive_int(value, f"engine provider {role} {target}") + elif target == "fpm_profile": + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + value = load_fpm_profile(value).model_dump(mode="json") elif target in {"enable_eplb", "decoder_replay", "enable_shared_layer", "strict_provenance"}: if not isinstance(value, bool): raise ValueError(f"engine provider {role} {target} must be a boolean") @@ -1987,6 +1992,7 @@ def add(name: str, value: object) -> None: "agentic_phases", "agentic_model_projection", "weka_nested_timestamp_basis", + "fpm_query_evidence", ) if key in payload } diff --git a/python/aisimulate/src/aisimulate/sdk/fpm_profile.py b/python/aisimulate/src/aisimulate/sdk/fpm_profile.py new file mode 100644 index 000000000..0bf5daba0 --- /dev/null +++ b/python/aisimulate/src/aisimulate/sdk/fpm_profile.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility alias for aisimulate_core.sdk.fpm_profile.""" + +from aisimulate.sdk._compat import alias_module as _alias_module + +_alias_module(__name__, "aisimulate_core.sdk.fpm_profile") diff --git a/python/aisimulate/src/aisimulate/support/cli.py b/python/aisimulate/src/aisimulate/support/cli.py index e5a06ce1d..08dc3c2c6 100644 --- a/python/aisimulate/src/aisimulate/support/cli.py +++ b/python/aisimulate/src/aisimulate/support/cli.py @@ -17,6 +17,8 @@ import yaml from pydantic import ValidationError +from aisimulate.config.common import load_yaml + from .fpm import run_fpm from .plan import create_plan, request_id from .schema import FPMDeployment, SearchProfile, SloSpec, SupportIdentity, SupportRequest, WorkloadSpec @@ -57,7 +59,11 @@ def add_support_parser(subparsers: Any) -> None: init.add_argument("--tokenizer-revision") init.add_argument("--chat-template-revision") init.add_argument("--aisimulate-revision", help="Optional pinned AISimulate source revision.") - init.add_argument("--tensor-parallel", type=int, help="GPUs per pure-TP worker (default: 1).") + init.add_argument("--tensor-parallel", type=int, help="Attention tensor-parallel size (default: 1).") + init.add_argument("--attention-data-parallel", type=int, help="Attention data-parallel size (default: 1).") + init.add_argument("--moe-tensor-parallel", type=int, help="Expert tensor-parallel size (default: TP for MoE).") + init.add_argument("--moe-expert-parallel", type=int, help="Expert-parallel size (default: 1).") + init.add_argument("--fpm-profile", help="JSON/YAML identity and resource profile for class-independent FPM.") init.add_argument("--input-tokens", type=int, help="Input tokens per request (default: 1024).") init.add_argument("--output-tokens", type=int, help="Output tokens per request (default: 128).") init.add_argument("--concurrency", type=int, help="Concurrent requests (default: 1).") @@ -121,6 +127,7 @@ def _request_from_args(args: argparse.Namespace) -> SupportRequest: "identity": identity, "workload": {**_values(args, WorkloadSpec), "slo": _values(args, SloSpec)}, "search": _values(args, SearchProfile), + **({"fpm_profile": load_yaml(args.fpm_profile)} if getattr(args, "fpm_profile", None) else {}), } ) @@ -196,6 +203,9 @@ def _print(value: dict[str, Any], output_format: str = "table") -> None: "max_candidates": ("Maximum candidates (1 or 2)", int), "objective": ("Recommendation objective", str), "seed": ("Recommendation search seed", int), + "attention_data_parallel": ("Attention data-parallel size", int), + "moe_tensor_parallel": ("Expert tensor-parallel size", int), + "moe_expert_parallel": ("Expert-parallel size", int), } @@ -266,12 +276,12 @@ def _init(args: argparse.Namespace) -> int: f"Scope: FPM simulation for {request.identity.model} ({request.identity.model_kind}), " f"vLLM {request.identity.framework_version}, " f"{request.identity.gpu_count} {request.identity.gpu} GPU(s) on {request.identity.node_count} node(s); " - f"TP{request.search.tensor_parallel}, {workload.input_tokens}/{workload.output_tokens} tokens, " + f"{request.parallel_preset} {request.parallelism()}, {workload.input_tokens}/{workload.output_tokens} tokens, " f"concurrency {workload.concurrency}, {workload.request_count} requests, " f"up to {request.search.max_candidates} candidate(s)." ) print( - "Request saved. Model integration, runtime compatibility and FPM data are unchecked; " + "Request saved. Model metadata/resources, runtime compatibility and FPM data are unchecked; " "accuracy is not assessed. Setup has not launched GPU work." ) _print( @@ -294,7 +304,8 @@ def _plan(args: argparse.Namespace) -> int: "candidate_count": plan["search"]["candidate_count"], "plan": str(root / "support-plan.json"), "commands": plan["outputs"]["commands"], - "prerequisites": "unchecked", + "prerequisites": "runtime_and_fpm_data_unchecked" if request.fpm_profile is not None else "unchecked", + **({"resources": "estimated_from_declared_profile"} if request.fpm_profile is not None else {}), "accuracy": "not assessed", "next": shlex.join( [ diff --git a/python/aisimulate/src/aisimulate/support/fpm.py b/python/aisimulate/src/aisimulate/support/fpm.py index af53c8fda..b5f642a2d 100644 --- a/python/aisimulate/src/aisimulate/support/fpm.py +++ b/python/aisimulate/src/aisimulate/support/fpm.py @@ -30,6 +30,17 @@ def fpm_cli_args( checkpoint = Path(checkpoint_dir).expanduser().resolve() if checkpoint_dir else root / "fpm-checkpoint" if checkpoint != root / "fpm-checkpoint" and root / "fpm-checkpoint" not in checkpoint.parents: raise ValueError("checkpoint_dir must stay within the plan's fpm-checkpoint directory") + profile = request.profile_deployment() + max_prefill_tokens = max(2, request.workload.input_tokens * request.workload.concurrency) + max_prefill_batch = request.workload.concurrency + if profile is not None: + scheduler = request.scheduler_limits() + max_prefill_batch = scheduler["max_sequences"] + max_prefill_tokens = min( + max(2, request.workload.input_tokens * max_prefill_batch), scheduler["max_batched_tokens"] + ) + if max_prefill_tokens < 2: + raise ValueError("FPM collection requires a rank-local token limit of at least 2 in the resource profile") command = [ "python3", "-m", @@ -41,15 +52,15 @@ def fpm_cli_args( "--gpu", request.identity.gpu, "--fpm-max-gpus", - str(request.search.tensor_parallel), + str(request.worker_gpus), "--fpm-gpu-counts", - str(request.search.tensor_parallel), + str(request.worker_gpus), "--fpm-parallel-presets", - "pure_tp" if request.identity.model_kind == "moe" else "tp", + request.parallel_preset, "--fpm-max-prefill-isl", - str(max(2, request.workload.input_tokens * request.workload.concurrency)), + str(max_prefill_tokens), "--fpm-max-prefill-batch-size", - str(request.workload.concurrency), + str(max_prefill_batch), "--checkpoint-dir", str(checkpoint), "--fpm-artifact-root", @@ -57,6 +68,23 @@ def fpm_cli_args( "--fpm-database-root", str(root / "systems/data"), ] + if profile is not None: + command.extend( + ( + "--model-architecture", + request.fpm_profile.architecture, + "--fpm-weight-quantizations", + profile.gemm_quant_mode, + "--fpm-kv-cache-dtypes", + profile.kv_cache_dtype, + "--fpm-attention-backend", + profile.attention_backend, + "--fpm-moe-backend", + profile.moe_backend, + "--fpm-model-profile", + str(root / "fpm-model-profile.json"), + ) + ) if request.identity.sm is not None: command.extend(("--sm", str(request.identity.sm))) if deployment is not None: diff --git a/python/aisimulate/src/aisimulate/support/plan.py b/python/aisimulate/src/aisimulate/support/plan.py index 46b08c406..4a62c73a5 100644 --- a/python/aisimulate/src/aisimulate/support/plan.py +++ b/python/aisimulate/src/aisimulate/support/plan.py @@ -35,20 +35,13 @@ def request_id(request: SupportRequest) -> str: def _parallelism(request: SupportRequest, replicas: int) -> dict[str, int]: - return { - "replicas": replicas, - "tensor": request.search.tensor_parallel, - "pipeline": 1, - "attention_data": 1, - "moe_tensor": request.search.tensor_parallel if request.identity.model_kind == "moe" else 1, - "moe_expert": 1, - } + return request.parallelism(replicas) def _configs( request: SupportRequest, root: Path ) -> tuple[dict[str, Any], dict[str, dict[str, Any]], list[dict[str, int]]]: - max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.search.tensor_parallel) + max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.worker_gpus) replicas = list(dict.fromkeys((1, max_replicas)))[: request.search.max_candidates] presets = [_parallelism(request, count) for count in replicas] workload = request.workload @@ -75,9 +68,12 @@ def _configs( } # Recommendation otherwise expands scheduler defaults into extra domains. worker = { - "scheduler": {"max_batched_tokens": 8192, "max_sequences": 256}, + "scheduler": request.scheduler_limits(), "timing": {"type": "default", "estimation_mode": "fpm_interpolation", "fallback_policy": "deny"}, } + if request.fpm_profile is not None: + engine["fpm_profile"] = request.fpm_profile.model_dump(mode="json") + worker["timing"]["estimator_config"] = {"fpm_interpolation": {"method": "direct"}} prediction = { **common, "engine": {**engine, "workers": {"aggregated": {**worker, "parallelism": presets[0]}}}, @@ -161,7 +157,7 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] "candidate_count": len(presets), "max_candidates": request.search.max_candidates, "candidates": [ - {"parallelism": preset, "total_gpus": preset["replicas"] * preset["tensor"]} for preset in presets + {"parallelism": preset, "total_gpus": preset["replicas"] * request.worker_gpus} for preset in presets ], "detail": ( "A single selected worker is planned; recommendation does not compare alternative configurations." @@ -174,7 +170,7 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] }, "fpm": { "status": "planned", - "worker_gpus": request.search.tensor_parallel, + "worker_gpus": request.worker_gpus, "plan_command": commands["fpm_plan_local"], "sampling": ( "Prefill is bounded by max(2, input_tokens * concurrency) and concurrency. Decode sampling and " @@ -187,10 +183,10 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] "id": "model_integration", "status": "not_checked", "detail": ( - "Verify pinned model metadata, memory/cache accounting, and chosen parallelism. This revision's " - "FPM predict/recommend uses registered analytical classes and SOL transfer; follow " - "python/aisimulate/docs/add_a_new_model.md for that route. Class-independent FPM with direct " - "timing interpolation is planned separately. Per-operation silicon data is not required." + "Verify pinned model metadata, memory/cache accounting, and chosen parallelism. Without an " + "FPM profile, predict/recommend uses a registered analytical class and SOL transfer; follow " + "python/aisimulate/docs/add_a_new_model.md for that route. Supply an FPM identity/resource " + "profile to use direct interpolation without a class. Per-operation silicon data is not required." ), }, { @@ -223,12 +219,70 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] "systems_root": str(root / "systems"), }, } + if request.fpm_profile is not None: + from aisimulate_core.sdk.memory import estimate_kv_cache + + deployment = request.profile_deployment() + scheduler = prediction["engine"]["workers"]["aggregated"]["scheduler"] + estimate = estimate_kv_cache( + request.identity.model, + request.identity.gpu, + request.identity.framework, + backend_version=request.identity.framework_version, + max_num_tokens=scheduler["max_batched_tokens"], + max_batch_size=scheduler["max_sequences"], + memory_fraction_kind="of_total", + memory_fraction_value=0.9, + tp_size=deployment.tp, + pp_size=deployment.pp, + attention_dp_size=deployment.dp, + moe_tp_size=deployment.moe_tp, + moe_ep_size=deployment.moe_ep, + fpm_profile=request.fpm_profile, + ) + if estimate["total_kv_size_tokens"] <= request.search.context_length: + raise ValueError("declared FPM resources leave insufficient rank-local KV capacity for the pilot context") + plan["resources"] = estimate + plan["search"]["baseline_rule"] = ( + "Selected profile deployment; declared topology and rank-local resource bounds pass CPU admission. " + "Serving-runtime compatibility and silicon accuracy remain unchecked." + ) + plan["prerequisites"][0].update( + status="declared_profile_validated", + detail=( + "The pinned identity, declared deployment, scheduler envelope and resource budget validate without " + "an analytical model class. Verify the profile's resource provenance against the serving runtime. " + "Generated configs use direct interpolation; matching timing coverage is still required." + ), + ) + plan["prerequisites"][1]["detail"] = ( + "Prepare and verify the declared model/tokenizer revisions, visible GPUs, model access, and the " + "packaged Dynamo/Kubernetes/Generator collector runtime. This plan does not download a pinned " + "checkpoint or inspect the running runtime. During collection execution, the collector checks " + "the observed Pod vLLM version against the profile's literal backend version before benchmarking." + ) + plan["fpm"]["sampling"] = ( + "Prefill sampling is capped by the pilot input length and the generated rank-local scheduler limits. " + "These bounds do not establish complete timing coverage of worker-wide FPM queries, including DEP. " + "Decode uses the collector's existing sampling profile within the declared resource envelope. " + "The synthetic request count does not bound timing samples. Smoke and limited runs do not publish " + "formal FPM data." + ) + plan["prerequisites"][2]["detail"] = ( + "Collect matching prefill/decode timings into the local systems tree before prediction or recommendation. " + "The profile resource estimate does not establish timing coverage or measured accuracy." + ) + plan["outputs"]["fpm_model_profile"] = str(root / "fpm-model-profile.json") yaml_documents = { Path("request.yaml"): request.model_dump(mode="json", exclude_none=True), Path("predict/pilot.yaml"): prediction, **{Path(f"recommend/{name}.yaml"): config for name, config in recommendations.items()}, } documents = {path: yaml.safe_dump(data, sort_keys=False).encode() for path, data in yaml_documents.items()} + if request.fpm_profile is not None: + documents[Path("fpm-model-profile.json")] = ( + json.dumps(request.fpm_profile.model_dump(mode="json"), indent=2, sort_keys=True) + "\n" + ).encode() documents[Path("systems") / source_spec.name] = source_spec.read_bytes() documents[Path("commands.json")] = (json.dumps(commands, indent=2, sort_keys=True) + "\n").encode() documents[Path("support-plan.json")] = (json.dumps(plan, indent=2, sort_keys=True) + "\n").encode() @@ -257,6 +311,7 @@ def _check_paths(root: Path) -> None: for relative in ( "request.yaml", "support-plan.json", + "fpm-model-profile.json", "commands.json", "predict", "predict/pilot.yaml", @@ -300,6 +355,15 @@ def _check_saved_request(request: SupportRequest, root: Path) -> None: raise ValueError(f"saved request in {root} was modified or cannot be read") from exc if saved != request: raise ValueError(f"saved request identity in {root} differs from the requested plan") + if request.fpm_profile is not None and (root / "fpm-model-profile.json").exists(): + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + try: + profile = load_fpm_profile((root / "fpm-model-profile.json").read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ValueError(f"saved FPM model profile in {root} is missing or invalid") from exc + if profile != request.fpm_profile: + raise ValueError(f"saved FPM model profile in {root} differs from the requested identity") def create_plan(request: SupportRequest, output_dir: str | Path, *, overwrite: bool = False) -> dict[str, Any]: diff --git a/python/aisimulate/src/aisimulate/support/schema.py b/python/aisimulate/src/aisimulate/support/schema.py index 603f6356c..532a6af58 100644 --- a/python/aisimulate/src/aisimulate/support/schema.py +++ b/python/aisimulate/src/aisimulate/support/schema.py @@ -13,6 +13,7 @@ from pydantic import Field, field_validator, model_validator from aisimulate.config.common import PositiveFiniteFloat, PositiveStrictInt, StrictModel, load_yaml +from aisimulate.fpm_profile import FpmModelProfile class SupportIdentity(StrictModel): @@ -96,6 +97,9 @@ def _request_count(self) -> WorkloadSpec: class SearchProfile(StrictModel): tensor_parallel: PositiveStrictInt = 1 + attention_data_parallel: PositiveStrictInt | None = None + moe_tensor_parallel: PositiveStrictInt | None = None + moe_expert_parallel: PositiveStrictInt | None = None context_length: PositiveStrictInt = 16384 max_candidates: int = Field(default=1, strict=True, ge=1, le=2) objective: Literal[ @@ -152,13 +156,84 @@ class SupportRequest(StrictModel): identity: SupportIdentity workload: WorkloadSpec = Field(default_factory=WorkloadSpec) search: SearchProfile = Field(default_factory=SearchProfile) + fpm_profile: FpmModelProfile | None = None + + def parallelism(self, replicas: int = 1) -> dict[str, int]: + search = self.search + return { + "replicas": replicas, + "tensor": search.tensor_parallel, + "pipeline": 1, + "attention_data": search.attention_data_parallel or 1, + "moe_tensor": search.moe_tensor_parallel + or (search.tensor_parallel if self.identity.model_kind == "moe" else 1), + "moe_expert": search.moe_expert_parallel or 1, + } + + @property + def worker_gpus(self) -> int: + parallel = self.parallelism() + return parallel["tensor"] * parallel["attention_data"] + + @property + def parallel_preset(self) -> str: + parallel = self.parallelism() + if self.identity.model_kind == "dense": + return "tp" + if parallel["attention_data"] > 1: + return "dep" + return "tep" if parallel["moe_expert"] > 1 else "pure_tp" + + def profile_deployment(self): + if self.fpm_profile is None: + return None + parallel = self.parallelism() + return self.fpm_profile.select( + model=self.identity.model, + system=self.identity.gpu, + backend=self.identity.framework, + backend_version=self.identity.framework_version, + tp_size=parallel["tensor"], + pp_size=parallel["pipeline"], + attention_dp_size=parallel["attention_data"], + moe_tp_size=parallel["moe_tensor"], + moe_ep_size=parallel["moe_expert"], + ) + + def scheduler_limits(self) -> dict[str, int]: + """Rank-local limits shared by generated configs and collection bounds.""" + deployment = self.profile_deployment() + if deployment is None: + return {"max_batched_tokens": 8192, "max_sequences": 256} + return { + "max_batched_tokens": min(8192, deployment.resources.max_num_tokens), + "max_sequences": min(self.workload.concurrency, deployment.resources.max_batch_size), + } @model_validator(mode="after") def _shape(self) -> SupportRequest: - if self.search.tensor_parallel > self.identity.gpus_per_node: - raise ValueError("search.tensor_parallel must fit within gpus_per_node") + parallel = self.parallelism() + tp, dp, mtp, ep = (parallel[name] for name in ("tensor", "attention_data", "moe_tensor", "moe_expert")) + if self.identity.model_kind == "dense": + valid = dp == mtp == ep == 1 + else: + valid = ( + (dp == 1 and mtp == tp and ep == 1) or (tp == mtp == 1 and ep == dp) or (dp == mtp == 1 and ep == tp) + ) + if not valid: + raise ValueError("onboarding requires a complete TP, DEP, or TEP topology; set the attention and MoE sizes") + if self.worker_gpus > self.identity.gpus_per_node: + raise ValueError("search.tensor_parallel * attention_data_parallel must fit within gpus_per_node") if self.workload.input_tokens + self.workload.output_tokens > self.search.context_length: raise ValueError("search.context_length must cover the input and output tokens") + if self.fpm_profile is not None: + self.profile_deployment() + if self.identity.model_revision != self.fpm_profile.model_revision: + raise ValueError("identity.model_revision must match fpm_profile.model_revision") + if (self.fpm_profile.num_experts > 0) != (self.identity.model_kind == "moe"): + raise ValueError("identity.model_kind must match fpm_profile.num_experts") + if self.search.context_length > self.fpm_profile.context_length: + raise ValueError("search.context_length exceeds fpm_profile.context_length") return self @classmethod diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index 30865d343..8167bbbae 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -29,7 +29,13 @@ from aisimulate.config.traffic import AgenticSnapshotOptions -from ..config.common import ENGINE_MODEL_CONTROL_FIELDS, SystemsPath, SystemsRoot, is_active_engine_model_control +from ..config.common import ( + ENGINE_MODEL_CONTROL_FIELDS, + SystemsPath, + SystemsRoot, + is_active_engine_model_control, + requested_backend_version, +) from ..config.engine import NgramSpeculationConfig @@ -523,6 +529,7 @@ class SearchSpace(BaseModel): # pinned model_name: str # HF id or private model name hardware_sku: str # e.g. "h200_sxm" + fpm_profile: dict[str, Any] | None = None database_mode: Literal["SILICON", "HYBRID", "EMPIRICAL", "SOL"] = "SILICON" transfer_policy: str | list[str] | None = None systems_paths: list[SystemsRoot] | None = Field(default=None, min_length=1) @@ -651,6 +658,19 @@ def _migrate_legacy_estimator_selection(cls, value): @model_validator(mode="after") def _validate_search_choices(self) -> SearchSpace: """Every backend dimension is a non-empty subset of its allowed choices.""" + if self.fpm_profile is not None: + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + profile = load_fpm_profile(self.fpm_profile) + if profile.model != self.model_name: + raise ValueError("FPM profile model identity does not match the requested model") + if set(self.backend) != {"vllm"} or self._uses_legacy_estimator_provider(): + raise ValueError( + "FPM profiles support vLLM aggregated/disaggregated decoder workers without AFD or encoders" + ) + if any(getattr(self, f"{role}_timing_model") is not None for role in ("agg", "prefill", "decode")): + raise ValueError("fpm_profile requires default timing for every worker") + self.fpm_profile = profile.model_dump(mode="json") if self.systems_path is not None: if self.systems_paths is not None and self.systems_paths != [self.systems_path]: raise ValueError("systems_path conflicts with systems_paths") @@ -956,11 +976,7 @@ def _validate_engine_controls(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 + return requested_backend_version(self.backend_version, backend) @model_validator(mode="after") def _validate_gpu_budget(self) -> SearchSpace: diff --git a/python/aisimulate/src/aisimulate/sweeper/deploy.py b/python/aisimulate/src/aisimulate/sweeper/deploy.py index 129376e9e..612d69072 100644 --- a/python/aisimulate/src/aisimulate/sweeper/deploy.py +++ b/python/aisimulate/src/aisimulate/sweeper/deploy.py @@ -172,6 +172,7 @@ def _engine_args_payload( pp_size=int(sample[f"{prefix}pp"]), moe_tp_size=moe_tp, moe_ep_size=moe_ep, + **_profile_kv_args(sample, role, backend_version), **({"kvcache_quant_mode": sample["kvcache_quant_mode"]} if sample.get("kvcache_quant_mode") else {}), ) if configured_bytes == "auto" @@ -186,6 +187,7 @@ def _engine_args_payload( pp_size=int(sample["prefill_pp"]), moe_tp_size=int(sample["prefill_moe_tp"]), moe_ep_size=int(sample["prefill_moe_ep"]), + **_profile_kv_args(sample, "prefill", backend_version), **({"kvcache_quant_mode": sample["kvcache_quant_mode"]} if sample.get("kvcache_quant_mode") else {}), ) if transfer_geometry == "auto" @@ -201,6 +203,19 @@ def _engine_args_payload( return payload +def _profile_kv_args(sample: dict[str, Any], role: str, backend_version: str) -> dict[str, Any]: + if sample.get("fpm_profile") is None: + return {} + prefix = _role_prefix(role) + return { + "fpm_profile": sample["fpm_profile"], + "system": _role_hardware_sku(sample, role), + "backend": sample["backend"], + "backend_version": backend_version, + "attention_dp_size": int(sample[f"{prefix}attention_dp"]), + } + + def build_backend_deployment( sample: dict[str, Any], *, diff --git a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py index a9bc17039..351b33525 100644 --- a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py +++ b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py @@ -78,6 +78,7 @@ def _request(self, sample: dict[str, Any], role: str) -> ForwardPassPerfModelCon interpolation["fpm_parquet_path"] = path return ForwardPassPerfModelConfig( model=self._search_space.model_name, + fpm_profile=self._search_space.fpm_profile, system=_role_hardware_sku(sample, role), backend=backend, worker_type="aggregated" if role == "agg" else role, @@ -111,14 +112,13 @@ def _resolve(self, request: ForwardPassPerfModelConfig, role: str) -> ForwardPas estimator_config = deepcopy(request.estimator_config) estimator_config["fpm_interpolation"]["fpm_parquet_path"] = str(Path(path).absolute()) request = replace(request, estimator_config=estimator_config) - 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: + request_payload = RustForwardPassPerfModel.normalize_config(request) + cache_key = json.dumps(request_payload, sort_keys=True) + cached = self._resolved.get(cache_key) + if cached is not None: + return deepcopy(cached) model = RustForwardPassPerfModel.best_available(request) diagnostics = model.diagnostics() except Exception as exc: @@ -155,6 +155,7 @@ def _resolve(self, request: ForwardPassPerfModelConfig, role: str) -> ForwardPas ) for field in ( "model", + "fpm_profile", "system", "backend", "tp", diff --git a/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py b/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py index 631859ac8..49f65aa6a 100644 --- a/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py +++ b/python/aisimulate/src/aisimulate/sweeper/kv_estimate.py @@ -20,11 +20,16 @@ (:func:`get_latest_database_version` resolves the version). SKUs without a perf DB raise :class:`NoPerfDatabase`; the naive fallback is intentionally disabled because it mis-models MoE expert sharding. + +An explicit FPM profile instead supplies per-deployment resource bounds. That +memory path requires a hardware specification and literal profile version, +but no timing database or analytical model. """ from __future__ import annotations from collections.abc import Iterable +from typing import Any from aisimulate_core.sdk.memory import estimate_kv_cache from aisimulate_core.sdk.perf_database import get_latest_database_version @@ -92,6 +97,7 @@ def estimate_kv_tokens( max_batch_size: int = DEFAULT_MAX_BATCH_SIZE, memory_fraction: float = DEFAULT_MEMORY_FRACTION, nextn: int = 0, + fpm_profile: dict[str, Any] | None = None, model_controls: dict[str, str | int | bool] | None = None, ) -> int | None: """Per-rank KV-cache capacity (in tokens) for ``shape``, or ``None`` when the @@ -118,6 +124,7 @@ def estimate_kv_tokens( **(model_controls or {}), systems_path=(list(resolve_systems_paths(systems_paths)) if systems_paths is not None else None), allow_naive_fallback=False, + **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), ) except ValueError as exc: msg = str(exc) @@ -145,6 +152,7 @@ def feasible_shape_tokens( max_num_tokens: int = DEFAULT_MAX_NUM_TOKENS, max_batch_size: int = DEFAULT_MAX_BATCH_SIZE, memory_fraction: float = DEFAULT_MEMORY_FRACTION, + fpm_profile: dict[str, Any] | None = None, model_controls: dict[str, str | int | bool] | None = None, nextn: int = 0, ) -> dict[ParallelShape, int]: @@ -168,6 +176,7 @@ def feasible_shape_tokens( max_num_tokens=max_num_tokens, max_batch_size=max_batch_size, memory_fraction=memory_fraction, + **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), **({"model_controls": model_controls} if model_controls else {}), **({"nextn": nextn} if nextn else {}), ) diff --git a/python/aisimulate/src/aisimulate/sweeper/kv_load.py b/python/aisimulate/src/aisimulate/sweeper/kv_load.py index f6a9cb259..00c9cb506 100644 --- a/python/aisimulate/src/aisimulate/sweeper/kv_load.py +++ b/python/aisimulate/src/aisimulate/sweeper/kv_load.py @@ -5,6 +5,7 @@ from __future__ import annotations +import json from dataclasses import dataclass from functools import cache from typing import Any @@ -42,6 +43,7 @@ def _per_rank_capacity_tokens( max_batch_size: int, memory_fraction: float, nextn: int, + fpm_profile_json: str | None = None, model_controls: tuple[tuple[str, str | int | bool], ...] = (), ) -> int: tokens = estimate_kv_tokens( @@ -55,6 +57,7 @@ def _per_rank_capacity_tokens( max_batch_size=max_batch_size, memory_fraction=memory_fraction, nextn=nextn, + **({"fpm_profile": json.loads(fpm_profile_json)} if fpm_profile_json is not None else {}), **({"model_controls": dict(model_controls)} if model_controls else {}), ) if tokens is None: @@ -94,6 +97,7 @@ def _role_capacity_tokens( roots = [resolved["systems_path"]] if not roots: roots = sample.get("systems_paths") + profile = resolved.get("fpm_profile", sample.get("fpm_profile")) per_rank_tokens = _per_rank_capacity_tokens( config.shape, model_name=str(resolved.get("model", resolved.get("model_path", sample["model_name"]))), @@ -104,6 +108,7 @@ def _role_capacity_tokens( 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"]), + fpm_profile_json=json.dumps(profile, sort_keys=True) if profile is not None else None, nextn=int(resolved.get("nextn", sample.get("aic_nextn")) or 0), model_controls=tuple( (name, resolved.get(name, sample.get(name))) diff --git a/python/aisimulate/src/aisimulate/sweeper/model_hw.py b/python/aisimulate/src/aisimulate/sweeper/model_hw.py index 0749eca9e..8e9ee40ca 100644 --- a/python/aisimulate/src/aisimulate/sweeper/model_hw.py +++ b/python/aisimulate/src/aisimulate/sweeper/model_hw.py @@ -15,11 +15,16 @@ workload's ``max_seq_len`` (:mod:`aisimulate.sweeper.kv_estimate`). That is per-shape (TEP / DEP / TP differ at the same GPU count) and uses the real quantized weights — it replaces the old BF16 min-GPU weight floor entirely. + +With an explicit FPM profile, enumeration uses only its declared deployment +tuples and rank-local resource bounds. It does not resolve an analytical model +or infer another topology's resources from its GPU count. """ from __future__ import annotations from dataclasses import dataclass +from typing import Any from aisimulate.generator.naive import _estimate_model_weight_bytes from aisimulate_core.sdk import perf_database @@ -34,6 +39,7 @@ ) from .parallel_enum import ( DisaggParallelConfig, + ParallelShape, ReplicaParallelConfig, enumerate_disagg_configs, enumerate_parallel_configs, @@ -126,6 +132,7 @@ def parallel_configs_for( 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, + fpm_profile: dict[str, Any] | None = None, model_controls: dict[str, str | int | bool] | None = None, nextn: int = 0, ) -> list[ReplicaParallelConfig] | list[DisaggParallelConfig]: @@ -147,13 +154,28 @@ 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, - systems_paths=systems_paths, - ) - seq_len = max_seq_len if max_seq_len is not None else mh.max_context + profile = None + if fpm_profile is not None: + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + profile = load_fpm_profile(fpm_profile) + if profile.model != model_name: + raise ValueError("FPM profile model identity does not match the requested model") + seq_len = max_seq_len if max_seq_len is not None else profile.context_length + if seq_len > profile.context_length: + raise ValueError("max_seq_len exceeds FPM profile context_length") + configs = _profile_parallel_configs( + profile, + hardware_sku=hardware_sku, + backend=backend, + backend_version=backend_version, + gpu_budget=gpu_budget, + min_gpu_budget=min_gpu_budget, + deployment_mode=deployment_mode, + ) + else: + 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") @@ -161,20 +183,21 @@ def parallel_configs_for( # MoE tensor-parallel (moe_ep == 1) is enabled for every MoE model, MLA # included: real deployments (e.g. InferenceX GLM-5, reported as EP=1) run it, # so the search must be able to find it rather than have it filtered out here. - common = dict( - is_moe=mh.is_moe, - backend=backend, - gpu_budget=gpu_budget, - min_gpu_budget=min_gpu_budget, - enable_wideep=mh.enable_wideep, - allow_moe_pure_tp=True, - ) - if deployment_mode == "disagg": - configs = enumerate_disagg_configs(**common) - elif deployment_mode == "agg": - configs = enumerate_parallel_configs(**common) - else: - raise ValueError(f"deployment_mode must be 'agg' or 'disagg', got {deployment_mode!r}") + if profile is None: + common = dict( + is_moe=mh.is_moe, + backend=backend, + gpu_budget=gpu_budget, + min_gpu_budget=min_gpu_budget, + enable_wideep=mh.enable_wideep, + allow_moe_pure_tp=True, + ) + if deployment_mode == "disagg": + configs = enumerate_disagg_configs(**common) + elif deployment_mode == "agg": + configs = enumerate_parallel_configs(**common) + else: + raise ValueError(f"deployment_mode must be 'agg' or 'disagg', got {deployment_mode!r}") # KV-cache validity: keep configs whose every role-shape holds a max_seq_len sequence. def feasible_for(role: str, shapes): @@ -188,6 +211,20 @@ def feasible_for(role: str, shapes): raise ValueError( "role_runtime values must be (tokens, batch, memory) or (tokens, batch, memory, fixed_tokens)" ) + if profile is not None: + for shape in dict.fromkeys(shapes): + deployment = profile.select( + model=model_name, + system=hardware_sku, + backend=backend, + backend_version=backend_version, + tp_size=shape.tp, + pp_size=shape.pp, + attention_dp_size=shape.dp, + moe_tp_size=shape.moe_tp, + moe_ep_size=shape.moe_ep, + ) + deployment.resources.validate_envelope(max_num_tokens=role_tokens, max_batch_size=role_batch) if fixed_tokens is not None: return {shape: fixed_tokens for shape in dict.fromkeys(shapes) if fixed_tokens > seq_len} return feasible_shape_tokens( @@ -201,6 +238,7 @@ def feasible_for(role: str, shapes): max_num_tokens=role_tokens, max_batch_size=role_batch, memory_fraction=role_memory, + **({"fpm_profile": fpm_profile} if fpm_profile is not None else {}), **({"model_controls": model_controls} if model_controls else {}), **({"nextn": nextn} if nextn else {}), ) @@ -218,3 +256,50 @@ def feasible_for(role: str, shapes): f"sequence within {gpu_budget} GPUs ({backend} KV-cache estimate)" ) return kept + + +def _profile_parallel_configs( + profile, + *, + hardware_sku: str, + backend: str, + backend_version: str | None, + gpu_budget: int, + min_gpu_budget: int | None, + deployment_mode: str, +): + """Enumerate only declared deployment shapes; replicas do not change a cell.""" + deployments = [ + deployment + for deployment in profile.deployments + if deployment.system == hardware_sku + and deployment.backend == backend + and deployment.backend_version == backend_version + ] + if not deployments: + raise ValueError(f"no FPM deployment profile for {hardware_sku}/{backend}/{backend_version}") + replicas = [] + for deployment in deployments: + shape = ParallelShape( + tp=deployment.tp, + pp=deployment.pp, + dp=deployment.dp, + moe_tp=deployment.moe_tp, + moe_ep=deployment.moe_ep, + ) + replicas.extend( + ReplicaParallelConfig(shape=shape, replicas=count) + for count in range(1, gpu_budget // shape.gpus_per_worker + 1) + ) + if deployment_mode == "agg": + configs = replicas + elif deployment_mode == "disagg": + configs = [ + DisaggParallelConfig(prefill=prefill, decode=decode) + for prefill in replicas + for decode in replicas + if prefill.total_gpus + decode.total_gpus <= gpu_budget + ] + else: + raise ValueError("FPM profiles support only agg or disagg deployment modes") + return [config for config in configs if min_gpu_budget is None or config.total_gpus >= min_gpu_budget] diff --git a/python/aisimulate/src/aisimulate/sweeper/sample.py b/python/aisimulate/src/aisimulate/sweeper/sample.py index a4dd74aa4..34a46ece2 100644 --- a/python/aisimulate/src/aisimulate/sweeper/sample.py +++ b/python/aisimulate/src/aisimulate/sweeper/sample.py @@ -138,6 +138,8 @@ def unroll_sample( sample[key] = getattr(search_space, key) if search_space.systems_paths is not None: sample["systems_paths"] = search_space.systems_paths + if search_space.fpm_profile is not None: + sample["fpm_profile"] = search_space.fpm_profile if mode == "disagg": sample["prefill_hardware_sku"] = search_space.hardware_sku_for("prefill") diff --git a/python/aisimulate/src/aisimulate/sweeper/search.py b/python/aisimulate/src/aisimulate/sweeper/search.py index e726bd306..4b30bb630 100644 --- a/python/aisimulate/src/aisimulate/sweeper/search.py +++ b/python/aisimulate/src/aisimulate/sweeper/search.py @@ -66,6 +66,7 @@ ) from .replay import ( REPLAY_SPEC_API_VERSION, + ReplayOutputRequirements, ReplayReport, ReplaySpec, Runner, @@ -777,12 +778,18 @@ def _materialize_one( ), None -def _run_replay_detailed(spec: ReplaySpec, runner: Runner) -> _ReplayEvaluation: +def _run_replay_detailed( + spec: ReplaySpec, runner: Runner, output_requirements: ReplayOutputRequirements | None = None +) -> _ReplayEvaluation: """Run one replay while retaining validated runner provenance metadata.""" try: try: - report = runner.run(spec) + report = ( + runner.run(spec) + if output_requirements is None + else runner.run(spec, output_requirements=output_requirements) + ) except ResourceLimitError: raise except Exception as exc: @@ -1066,6 +1073,8 @@ def _score_prepared( def _init_worker( runner_factory: RunnerFactory, + *, + output_requirements: ReplayOutputRequirements | None = None, ) -> None: identity = getattr(mp.current_process(), "_identity", ()) worker_id = int(identity[0]) if identity else 0 @@ -1074,11 +1083,11 @@ def _init_worker( # shutdown. Unlike a plain atexit handler, this matches the ProcessPool # worker lifecycle and lets a runtime release worker-local resources. Finalize(None, runner.close, exitpriority=0) - _WORKER_CTX.update(runner=runner) + _WORKER_CTX.update(runner=runner, output_requirements=output_requirements) def _worker_eval(spec: ReplaySpec) -> _ReplayEvaluation: - return _run_replay_detailed(spec, _WORKER_CTX["runner"]) + return _run_replay_detailed(spec, _WORKER_CTX["runner"], _WORKER_CTX.get("output_requirements")) class Sweeper: @@ -1098,8 +1107,10 @@ def __init__( show_progress: bool = True, prediction_config_factory: Callable[[dict[str, Any], ReplaySpec], dict[str, Any]] | None = None, afd_performance_model: AFDPerformanceModel | None = None, + output_requirements: ReplayOutputRequirements | None = None, ) -> None: self._runner_factory = runner_factory + self._output_requirements = output_requirements self._providers = dict(providers or {}) self._sampler_factory = sampler_factory self._show_progress = show_progress @@ -1131,6 +1142,9 @@ def run( raise ValueError(f"top_n must be positive or None, got {top_n}") retention = CandidateRetention(candidate_retention) runner_factory = self._runner_factory + from functools import partial + + worker_initializer = partial(_init_worker, output_requirements=self._output_requirements) providers = self._providers sampler_factory = self._sampler_factory show_progress = self._show_progress @@ -1233,7 +1247,7 @@ def _new_pool() -> ProcessPoolExecutor: return ProcessPoolExecutor( max_workers=worker_count, mp_context=mp.get_context("spawn"), - initializer=_init_worker, + initializer=worker_initializer, initargs=(runner_factory,), ) @@ -1296,7 +1310,7 @@ def _eval_batch(todo: list[tuple[Suggestion, _PreparedCandidate]]): for index, result in evaluate_waves( [prepared.replay_spec for _, prepared in todo], factory=runner_factory, - initializer=_init_worker, + initializer=worker_initializer, evaluate=_worker_eval, workers=worker_count, timeout=max_eval_seconds, @@ -1327,7 +1341,7 @@ def _eval_batch(todo: list[tuple[Suggestion, _PreparedCandidate]]): suggestion, _score_prepared( prepared, - _run_replay_detailed(prepared.replay_spec, sequential_runner), + _run_replay_detailed(prepared.replay_spec, sequential_runner, self._output_requirements), config=config, goal=goal, ), diff --git a/python/aisimulate/src/aisimulate/sweeper/search_space.py b/python/aisimulate/src/aisimulate/sweeper/search_space.py index 090cba135..f1534b5df 100644 --- a/python/aisimulate/src/aisimulate/sweeper/search_space.py +++ b/python/aisimulate/src/aisimulate/sweeper/search_space.py @@ -292,6 +292,7 @@ def _heterogeneous_disagg_configs( role_runtime={"agg": _role_runtime(search_space, backend, role)}, **_engine_memory_kwargs(search_space), **_estimator_root_kwargs(search_space, role), + **({"fpm_profile": search_space.fpm_profile} if search_space.fpm_profile is not None else {}), ) except (NoPerfDatabase, NoViableParallelConfig) as exc: raise type(exc)(f"{role} hardware_sku={hardware!r}: {exc}") from exc @@ -620,6 +621,7 @@ def matches_custom(config: _ParallelConfig) -> bool: role_runtime=_runtime_by_role(ss, backend, deployment_mode), **_engine_memory_kwargs(ss), **_estimator_root_kwargs(ss, "agg" if deployment_mode == "agg" else "prefill"), + **({"fpm_profile": ss.fpm_profile} if ss.fpm_profile is not None else {}), ) except (NoPerfDatabase, NoViableParallelConfig): continue # backend unusable for this mode -> drop it from the search diff --git a/python/aisimulate/src/aisimulate_core/__init__.py b/python/aisimulate/src/aisimulate_core/__init__.py index 0af0641bb..56393f035 100644 --- a/python/aisimulate/src/aisimulate_core/__init__.py +++ b/python/aisimulate/src/aisimulate_core/__init__.py @@ -3,22 +3,11 @@ """AISimulate estimator and native performance-model API.""" -from aisimulate import __version__ +from __future__ import annotations + +from importlib import import_module -from ._native import ( - AicEngine, - RustForwardPassPerfModel, - _build_smoke, - engine_spec_bincode_from_json, - engine_spec_schema_version, - gemm_quant_util_levels, - moe_quant_util_levels, - op_from_spec_json, - ops_json_from_ops, - resolve_op_sources_report_json, - table_view_attributes, - weights_ops_json, -) +from aisimulate import __version__ __all__ = [ "AicEngine", @@ -35,3 +24,16 @@ "table_view_attributes", "weights_ops_json", ] + + +def __getattr__(name: str) -> object: + """Load native APIs on use while keeping metadata imports lightweight.""" + if name not in __all__ or name == "__version__": + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module("aisimulate_core._native"), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted({*globals(), *__all__}) diff --git a/python/aisimulate/src/aisimulate_core/_native.pyi b/python/aisimulate/src/aisimulate_core/_native.pyi index e7b013df8..6f6bf7699 100644 --- a/python/aisimulate/src/aisimulate_core/_native.pyi +++ b/python/aisimulate/src/aisimulate_core/_native.pyi @@ -174,7 +174,10 @@ class RustForwardPassPerfModel: def normalize_config(config_json: str) -> str: ... @staticmethod def _normalize_fpm_options( - options_json: str, fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None + options_json: str, + fmha_quant_mode: str | None = None, + comm_quant_mode: str | None = None, + has_profile: bool = False, ) -> str: ... @staticmethod def legacy_estimator_config(options_json: str) -> str: ... @@ -184,6 +187,7 @@ class RustForwardPassPerfModel: ) -> str: ... def predict_prefill_latency(self, bs: int, isl: int, prefix: int = 0) -> float: ... def estimate_forward_pass_time_ms(self, fpm_json: str) -> float | None: ... + def estimate_forward_pass_detailed(self, fpm_json: str) -> str: ... def tune_with_fpms(self, iterations_json: str) -> None: ... def static_phase_latency(self, batch_size: int, input_tokens: int, output_tokens: int, prefill: bool) -> float: ... def static_phase_diagnostics(self, batch_size: int, context_length: int, prefix: int, prefill: bool) -> str: ... diff --git a/python/aisimulate/src/aisimulate_core/fpm_profile.py b/python/aisimulate/src/aisimulate_core/fpm_profile.py new file mode 100644 index 000000000..39a0f9ac4 --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/fpm_profile.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit identity and resource bounds for decoder-only FPM simulation. + +Profiles describe every rank of a deployment without constructing operations. +Resource values are declarations supplied by the caller, not GPU measurements +or automatically inferred checkpoint metadata. Their provenance must say how +they were obtained; CUDA graph reservations are supplied separately at runtime. +""" + +from __future__ import annotations + +import json +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +from . import quantization + +_PositiveInt = Annotated[int, Field(gt=0)] +_Bytes = Annotated[int, Field(ge=0, le=2**53)] +_Nonempty = Annotated[str, Field(min_length=1, pattern=r"\S")] +_MUTABLE_REFERENCES = {"main", "master", "head", "latest", "current", "unknown", "tbd", "unpinned"} + + +class _ProfileModel(BaseModel): + model_config = ConfigDict(extra="forbid", strict=True) + + +class FpmResourceProfile(_ProfileModel): + """Conservative rank-local bounds at a declared scheduler envelope. + + ``max_num_tokens`` and ``max_batch_size`` apply to each attention-DP rank, + not the summed worker iteration used to query FPM timings. Overheads must + exclude the separately configured CUDA graph reservation. + """ + + weights_bytes: _Bytes + activations_bytes: _Bytes + runtime_overhead_bytes: _Bytes + comm_overhead_bytes: _Bytes + kv_bytes_per_token: Annotated[int, Field(gt=0, le=2**53)] + cache_layout: Literal["linear"] + max_num_tokens: _PositiveInt + max_batch_size: _PositiveInt + provenance: _Nonempty + + @model_validator(mode="after") + def _exact_total(self) -> FpmResourceProfile: + if self.non_kv_bytes > 2**53: + raise ValueError("total non-KV resource bytes must not exceed 2**53") + return self + + @property + def non_kv_bytes(self) -> int: + return self.weights_bytes + self.activations_bytes + self.runtime_overhead_bytes + self.comm_overhead_bytes + + def validate_envelope(self, *, max_num_tokens: int, max_batch_size: int) -> None: + for name, requested in (("max_num_tokens", max_num_tokens), ("max_batch_size", max_batch_size)): + if isinstance(requested, bool) or not isinstance(requested, int) or requested <= 0: + raise ValueError(f"{name} must be a positive integer, got {requested!r}") + limit = getattr(self, name) + if requested > limit: + raise ValueError( + f"FPM resource envelope exceeded: requested rank-local {name}={requested}, " + f"profile bound={limit}; provide resource bounds for the larger scheduler envelope" + ) + + +class FpmDeploymentProfile(_ProfileModel): + """One exact hardware, runtime, topology, precision and resource identity.""" + + system: _Nonempty + backend: Literal["vllm"] + backend_version: _Nonempty + tp: _PositiveInt + pp: Literal[1] = 1 + dp: _PositiveInt + moe_tp: _PositiveInt + moe_ep: _PositiveInt + cp: Literal[1] = 1 + gemm_quant_mode: _Nonempty + moe_quant_mode: _Nonempty + fmha_quant_mode: _Nonempty + comm_quant_mode: _Nonempty + kv_cache_dtype: _Nonempty + moe_backend: _Nonempty = "auto" + attention_backend: _Nonempty = "auto" + enable_wideep: Literal[False] = False + enable_eplb: Literal[False] = False + resources: FpmResourceProfile + + @field_validator("pp", "cp", mode="before") + @classmethod + def _integer_dimensions(cls, value: Any) -> int: + if type(value) is not int: + raise ValueError("parallel dimensions must be integers, not booleans or floats") + return value + + @field_validator("enable_wideep", "enable_eplb", mode="before") + @classmethod + def _boolean_flags(cls, value: Any) -> bool: + if type(value) is not bool: + raise ValueError("backend feature flags must be booleans") + return value + + @field_validator("backend_version") + @classmethod + def _literal_version(cls, value: str) -> str: + if value.lower() in _MUTABLE_REFERENCES | {"previous", "next"} or any(c.isspace() for c in value): + raise ValueError("FPM profile backend_version must be a literal runtime version") + return value + + @model_validator(mode="after") + def _deployment_identity(self) -> FpmDeploymentProfile: + modes = { + "gemm_quant_mode": quantization.GEMMQuantMode, + "moe_quant_mode": quantization.MoEQuantMode, + "fmha_quant_mode": quantization.FMHAQuantMode, + "comm_quant_mode": quantization.CommQuantMode, + "kv_cache_dtype": quantization.KVCacheQuantMode, + } + for name, enum in modes.items(): + if getattr(self, name) not in enum.__members__: + raise ValueError(f"unknown FPM {name}: {getattr(self, name)!r}") + return self + + @property + def parallel_tuple(self) -> tuple[int, int, int, int, int, int]: + return self.tp, self.pp, self.dp, self.moe_tp, self.moe_ep, self.cp + + def match_identity(self) -> list[str]: + """Return the exact 15-column native FPM identity, with no aliasing.""" + return [ + self.gemm_quant_mode, + self.moe_quant_mode, + self.fmha_quant_mode, + self.comm_quant_mode, + self.kv_cache_dtype, + *(str(value) for value in self.parallel_tuple), + self.moe_backend, + self.attention_backend, + str(self.enable_wideep), + str(self.enable_eplb), + ] + + def quantization_kwargs(self) -> dict[str, str]: + return { + "gemm_quant_mode": self.gemm_quant_mode, + "moe_quant_mode": self.moe_quant_mode, + "fmha_quant_mode": self.fmha_quant_mode, + "comm_quant_mode": self.comm_quant_mode, + "kvcache_quant_mode": self.kv_cache_dtype, + "attention_backend": self.attention_backend, + } + + def validate_overrides(self, **overrides: str | int | bool | None) -> None: + if ( + overrides.get("enable_eplb", False) is not False + or overrides.get("wideep_num_slots") is not None + or overrides.get("moe_backend") not in (None, "default") + ): + raise ValueError("FPM profiles do not support EPLB, slots or moe_backend overrides") + for name, expected in self.quantization_kwargs().items(): + supplied = overrides.get(name) + if supplied is not None and supplied != expected: + raise ValueError( + f"FPM profile identity conflict: {name}={supplied!r}, profile={expected!r}; " + "use matching measured precision and backend settings" + ) + + +class FpmModelProfile(_ProfileModel): + """Serializable FPM metadata independent of an analytical model class.""" + + schema_version: Literal[1] + model: _Nonempty + model_revision: _Nonempty + architecture: _Nonempty + context_length: _PositiveInt + num_experts: Annotated[int, Field(ge=0)] + provenance: _Nonempty + deployments: Annotated[list[FpmDeploymentProfile], Field(min_length=1)] + + @field_validator("schema_version", mode="before") + @classmethod + def _integer_schema(cls, value: Any) -> int: + if type(value) is not int: + raise ValueError("schema_version must be an integer") + return value + + @field_validator("model_revision") + @classmethod + def _pinned_revision(cls, value: str) -> str: + if value.strip().lower() in _MUTABLE_REFERENCES: + raise ValueError("provide a pinned model_revision, not a mutable or unknown placeholder") + return value + + @model_validator(mode="after") + def _unique_deployments(self) -> FpmModelProfile: + identities = set() + for deployment in self.deployments: + tp = deployment.dp == 1 and deployment.moe_ep == 1 and deployment.moe_tp == deployment.tp + dep = deployment.tp == 1 and deployment.moe_tp == 1 and deployment.moe_ep == deployment.dp + tep = deployment.dp == 1 and deployment.moe_tp == 1 and deployment.moe_ep == deployment.tp + dense_tp = deployment.dp == 1 and deployment.moe_ep == 1 and deployment.moe_tp in (1, deployment.tp) + if not (dense_tp if self.num_experts == 0 else tp or dep or tep): + raise ValueError( + "FPM profile requires dense TP or an MoE TP, DEP, or TEP topology with matching world sizes" + ) + identity = ( + deployment.system, + deployment.backend, + deployment.backend_version, + deployment.parallel_tuple, + ) + if identity in identities: + raise ValueError( + "duplicate FPM deployment identity; specify one precision/resource profile per topology" + ) + identities.add(identity) + if deployment.moe_ep > 1 and (self.num_experts == 0 or self.num_experts % deployment.moe_ep): + raise ValueError("num_experts must be positive and divisible by moe_ep for expert parallelism") + return self + + def select( + self, + *, + model: str, + system: str, + backend: str, + backend_version: str, + tp_size: int = 1, + pp_size: int = 1, + attention_dp_size: int = 1, + moe_tp_size: int | None = None, + moe_ep_size: int | None = None, + cp_size: int = 1, + ) -> FpmDeploymentProfile: + if model != self.model: + raise ValueError(f"FPM profile model identity mismatch: requested {model!r}, profile={self.model!r}") + shape = ( + tp_size, + pp_size, + attention_dp_size, + 1 if moe_tp_size is None else moe_tp_size, + 1 if moe_ep_size is None else moe_ep_size, + cp_size, + ) + if any(type(value) is not int or value <= 0 for value in shape): + raise ValueError("FPM deployment dimensions must be positive integers, not booleans") + for deployment in self.deployments: + if ( + deployment.system == system + and deployment.backend == backend + and deployment.backend_version == backend_version + and deployment.parallel_tuple == shape + ): + return deployment + raise ValueError( + f"no matching FPM deployment profile for {model!r}, {system}/{backend}/{backend_version}, " + f"(tp,pp,dp,moe_tp,moe_ep,cp)={shape}; provide explicit resources and precision for that deployment" + ) + + +def load_fpm_profile(value: dict[str, Any] | str | FpmModelProfile) -> FpmModelProfile: + if isinstance(value, FpmModelProfile): + value = value.model_dump(mode="json") + if isinstance(value, str): + value = json.loads(value) + return FpmModelProfile.model_validate(value) diff --git a/python/aisimulate/src/aisimulate_core/quantization.py b/python/aisimulate/src/aisimulate_core/quantization.py new file mode 100644 index 000000000..9d0e5f134 --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/quantization.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared quantization identities without performance-model runtime imports.""" + +from collections import namedtuple +from enum import Enum + +# compute_dtype names the tensor-core pipeline the mode's MMA actually executes +# on ("bfloat16" | "int8" | "fp8" | "fp4"), i.e. which _tc_flops entry of +# the system YAML governs its SOL math. Weight-only modes (int8_wo, int4_wo, +# w4a16_*) dequantize to bf16 before the MMA, so they map to "bfloat16". None +# marks memory-only modes (KV cache, comm) that never query FLOPS. +QuantMapping = namedtuple("QuantMapping", ["memory", "compute", "name", "compute_dtype"]) + + +class GEMMQuantMode(Enum): + """ + GEMM quant mode. + """ + + bfloat16 = QuantMapping(2, 1, "bfloat16", "bfloat16") # w16a16 + int8_wo = QuantMapping(1, 1, "int8_wo", "bfloat16") # w8a16 + int4_wo = QuantMapping(0.5, 1, "int4_wo", "bfloat16") # w4a16 + fp8 = QuantMapping(1, 2, "fp8", "fp8") # w8fp8 + fp8_static = QuantMapping( + 1, 2, "fp8_static", "fp8" + ) # fp8 with static quantization (compute_scale/scale_matrix modeled) + sq = QuantMapping(1, 2, "sq", "int8") # w8int8 + fp8_block = QuantMapping(1, 2, "fp8_block", "fp8") # specific for trtllm torch ds fp8 + fp8_ootb = QuantMapping( + 1, 2, "fp8_ootb", "fp8" + ) # in future, should deprecate this mode as it's specific for trtllm trt backend + nvfp4 = QuantMapping(9 / 16, 4, "nvfp4", "fp4") # nvfp4 on blackwell. 1 fp8 scale per 16 nvfp4 weights. + nvfp4_wo = QuantMapping(9 / 16, 1, "nvfp4_wo", "bfloat16") # nvfp4 sw dequant to bf16 (non-Blackwell) + w4a16_nvfp4 = QuantMapping( + 9 / 16, 1, "w4a16_nvfp4", "bfloat16" + ) # NVFP4 weights + 1 fp8 scale per 16 weights, dequantized into the bf16 MMA lane. + + +class MoEQuantMode(Enum): + """ + MoE quant mode. + """ + + bfloat16 = QuantMapping(2, 1, "bfloat16", "bfloat16") # w16a16 + fp8 = QuantMapping(1, 2, "fp8", "fp8") # w8fp8 + int4_wo = QuantMapping(0.5, 1, "int4_wo", "bfloat16") # w4a16 + fp8_block = QuantMapping(1, 2, "fp8_block", "fp8") # specific for trtllm torch ds fp8 + w4afp8 = QuantMapping(0.5, 2, "w4afp8", "fp8") # specific for trtllm torch ds w4a8 + nvfp4 = QuantMapping(9 / 16, 4, "nvfp4", "fp4") # nvfp4 on blackwell. 1 fp8 scale per 16 nvfp4 weights. + nvfp4_wo = QuantMapping(9 / 16, 1, "nvfp4_wo", "bfloat16") # nvfp4 sw dequant to bf16 (non-Blackwell) + w4a16_mxfp4 = QuantMapping(0.5, 1, "w4a16_mxfp4", "bfloat16") # native data format for gpt oss + w4a8_mxfp4_mxfp8 = QuantMapping(0.5, 2, "w4a8_mxfp4_mxfp8", "fp8") + # mxfp4 weights, mxfp8 activations (recommended for Blackwell) + w4a8_mxfp4_mxfp8_trtllm = QuantMapping(0.5, 2, "w4a8_mxfp4_mxfp8_trtllm", "fp8") + # Blackwell trtllm-gen fused MoE: MXFP4 (E2M1, block-32) weights x MXFP8 (E4M3) + # activations -- the kernel DeepSeek-V4-Pro actually runs in prefill on sm100 + # (bmm_MxE4m3_MxE2m1MxE4m3 ... sm100f, flashinfer trtllm_fp4_block_scale_moe). + # Distinct backend from w4a8_mxfp4_mxfp8 above (flashinfer cutedsl). DSV4 MoE + # weights are stored MXFP4 (I8-packed E2M1 + E8M0 scales), so sglang dispatches + # by GPU: sm100 -> this (trtllm-gen); sm90 -> w4a16_mxfp4_cutlass below. + w4a16_mxfp4_cutlass = QuantMapping(0.5, 1, "w4a16_mxfp4_cutlass", "bfloat16") + # Hopper (sm90) DeepSeek-V4-Pro MoE: flashinfer cutlass SM90 mixed GEMM + # (cutlass_fused_moe(use_w4_group_scaling=True)) -- MXFP4 weights x BF16 + # activations (weight-only). Distinct backend from w4a16_mxfp4 above, which is + # GPT-OSS's triton_kernels mxfp4 path. (DSV4 Hopper silicon data pending.) + w4a16_nvfp4 = QuantMapping(9 / 16, 1, "w4a16_nvfp4", "bfloat16") + # Scale-aware NVFP4 weights dequantized into the BF16 MoE compute lane. + w4a16_mxfp4_humming = QuantMapping(0.5, 1, "w4a16_mxfp4_humming", "bfloat16") + # Native Humming E2M1/E8M0 block-32 experts with unquantized BF16 inputs. + + +class FMHAQuantMode(Enum): + """ + FMHA quant mode. + """ + + bfloat16 = QuantMapping(2, 1, "bfloat16", "bfloat16") + float16 = QuantMapping(2, 1, "float16", "bfloat16") # sglang decode attention uses float16 compute + fp8 = QuantMapping(1, 2, "fp8", "fp8") + fp8_block = QuantMapping(1, 2, "fp8_block", "fp8") # FIXME: specific for sglang wideep + + +class KVCacheQuantMode(Enum): + """ + KVCache quant mode. + """ + + bfloat16 = QuantMapping(2, 0, "bfloat16", None) + int8 = QuantMapping(1, 0, "int8", None) + fp8 = QuantMapping(1, 0, "fp8", None) + + +class CommQuantMode(Enum): + """ + Comm quant mode. + """ + + half = QuantMapping(2, 0, "half", None) + int8 = QuantMapping(1, 0, "int8", None) + fp8 = QuantMapping(1, 0, "fp8", None) diff --git a/python/aisimulate/src/aisimulate_core/sdk/common.py b/python/aisimulate/src/aisimulate_core/sdk/common.py index 5741c9455..4a54ea344 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/common.py +++ b/python/aisimulate/src/aisimulate_core/sdk/common.py @@ -4,7 +4,6 @@ import csv import json import math -from collections import namedtuple from dataclasses import dataclass, field from enum import Enum, StrEnum from functools import cache @@ -12,6 +11,12 @@ from packaging.version import InvalidVersion, Version +from aisimulate_core.quantization import CommQuantMode as CommQuantMode +from aisimulate_core.quantization import FMHAQuantMode as FMHAQuantMode +from aisimulate_core.quantization import GEMMQuantMode as GEMMQuantMode +from aisimulate_core.quantization import KVCacheQuantMode as KVCacheQuantMode +from aisimulate_core.quantization import MoEQuantMode as MoEQuantMode +from aisimulate_core.quantization import QuantMapping as QuantMapping from aisimulate_core.sdk.deepseek_v41 import DeepSeekV41Config as DeepSeekV41Config from aisimulate_core.sdk.errors import MissingSystemFlopsError @@ -1384,102 +1389,6 @@ class PerfDataFilename(Enum): fpm_forward = "fpm_forward_perf.parquet" -# compute_dtype names the tensor-core pipeline the mode's MMA actually executes -# on ("bfloat16" | "int8" | "fp8" | "fp4"), i.e. which _tc_flops entry of -# the system YAML governs its SOL math. Weight-only modes (int8_wo, int4_wo, -# w4a16_*) dequantize to bf16 before the MMA, so they map to "bfloat16". None -# marks memory-only modes (KV cache, comm) that never query FLOPS. -QuantMapping = namedtuple("QuantMapping", ["memory", "compute", "name", "compute_dtype"]) - - -class GEMMQuantMode(Enum): - """ - GEMM quant mode. - """ - - bfloat16 = QuantMapping(2, 1, "bfloat16", "bfloat16") # w16a16 - int8_wo = QuantMapping(1, 1, "int8_wo", "bfloat16") # w8a16 - int4_wo = QuantMapping(0.5, 1, "int4_wo", "bfloat16") # w4a16 - fp8 = QuantMapping(1, 2, "fp8", "fp8") # w8fp8 - fp8_static = QuantMapping( - 1, 2, "fp8_static", "fp8" - ) # fp8 with static quantization (compute_scale/scale_matrix modeled) - sq = QuantMapping(1, 2, "sq", "int8") # w8int8 - fp8_block = QuantMapping(1, 2, "fp8_block", "fp8") # specific for trtllm torch ds fp8 - fp8_ootb = QuantMapping( - 1, 2, "fp8_ootb", "fp8" - ) # in future, should deprecate this mode as it's specific for trtllm trt backend - nvfp4 = QuantMapping(9 / 16, 4, "nvfp4", "fp4") # nvfp4 on blackwell. 1 fp8 scale per 16 nvfp4 weights. - nvfp4_wo = QuantMapping(9 / 16, 1, "nvfp4_wo", "bfloat16") # nvfp4 sw dequant to bf16 (non-Blackwell) - w4a16_nvfp4 = QuantMapping( - 9 / 16, 1, "w4a16_nvfp4", "bfloat16" - ) # NVFP4 weights + 1 fp8 scale per 16 weights, dequantized into the bf16 MMA lane. - - -class MoEQuantMode(Enum): - """ - MoE quant mode. - """ - - bfloat16 = QuantMapping(2, 1, "bfloat16", "bfloat16") # w16a16 - fp8 = QuantMapping(1, 2, "fp8", "fp8") # w8fp8 - int4_wo = QuantMapping(0.5, 1, "int4_wo", "bfloat16") # w4a16 - fp8_block = QuantMapping(1, 2, "fp8_block", "fp8") # specific for trtllm torch ds fp8 - w4afp8 = QuantMapping(0.5, 2, "w4afp8", "fp8") # specific for trtllm torch ds w4a8 - nvfp4 = QuantMapping(9 / 16, 4, "nvfp4", "fp4") # nvfp4 on blackwell. 1 fp8 scale per 16 nvfp4 weights. - nvfp4_wo = QuantMapping(9 / 16, 1, "nvfp4_wo", "bfloat16") # nvfp4 sw dequant to bf16 (non-Blackwell) - w4a16_mxfp4 = QuantMapping(0.5, 1, "w4a16_mxfp4", "bfloat16") # native data format for gpt oss - w4a8_mxfp4_mxfp8 = QuantMapping(0.5, 2, "w4a8_mxfp4_mxfp8", "fp8") - # mxfp4 weights, mxfp8 activations (recommended for Blackwell) - w4a8_mxfp4_mxfp8_trtllm = QuantMapping(0.5, 2, "w4a8_mxfp4_mxfp8_trtllm", "fp8") - # Blackwell trtllm-gen fused MoE: MXFP4 (E2M1, block-32) weights x MXFP8 (E4M3) - # activations -- the kernel DeepSeek-V4-Pro actually runs in prefill on sm100 - # (bmm_MxE4m3_MxE2m1MxE4m3 ... sm100f, flashinfer trtllm_fp4_block_scale_moe). - # Distinct backend from w4a8_mxfp4_mxfp8 above (flashinfer cutedsl). DSV4 MoE - # weights are stored MXFP4 (I8-packed E2M1 + E8M0 scales), so sglang dispatches - # by GPU: sm100 -> this (trtllm-gen); sm90 -> w4a16_mxfp4_cutlass below. - w4a16_mxfp4_cutlass = QuantMapping(0.5, 1, "w4a16_mxfp4_cutlass", "bfloat16") - # Hopper (sm90) DeepSeek-V4-Pro MoE: flashinfer cutlass SM90 mixed GEMM - # (cutlass_fused_moe(use_w4_group_scaling=True)) -- MXFP4 weights x BF16 - # activations (weight-only). Distinct backend from w4a16_mxfp4 above, which is - # GPT-OSS's triton_kernels mxfp4 path. (DSV4 Hopper silicon data pending.) - w4a16_nvfp4 = QuantMapping(9 / 16, 1, "w4a16_nvfp4", "bfloat16") - # Scale-aware NVFP4 weights dequantized into the BF16 MoE compute lane. - w4a16_mxfp4_humming = QuantMapping(0.5, 1, "w4a16_mxfp4_humming", "bfloat16") - # Native Humming E2M1/E8M0 block-32 experts with unquantized BF16 inputs. - - -class FMHAQuantMode(Enum): - """ - FMHA quant mode. - """ - - bfloat16 = QuantMapping(2, 1, "bfloat16", "bfloat16") - float16 = QuantMapping(2, 1, "float16", "bfloat16") # sglang decode attention uses float16 compute - fp8 = QuantMapping(1, 2, "fp8", "fp8") - fp8_block = QuantMapping(1, 2, "fp8_block", "fp8") # FIXME: specific for sglang wideep - - -class KVCacheQuantMode(Enum): - """ - KVCache quant mode. - """ - - bfloat16 = QuantMapping(2, 0, "bfloat16", None) - int8 = QuantMapping(1, 0, "int8", None) - fp8 = QuantMapping(1, 0, "fp8", None) - - -class CommQuantMode(Enum): - """ - Comm quant mode. - """ - - half = QuantMapping(2, 0, "half", None) - int8 = QuantMapping(1, 0, "int8", None) - fp8 = QuantMapping(1, 0, "fp8", None) - - # compute_dtype -> system YAML key. A key's absence from the YAML means the # platform has no hardware for that dtype (see get_quant_tc_flops). _COMPUTE_DTYPE_TO_FLOPS_KEY = { diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index 8e95cae61..505320c8e 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -12,6 +12,10 @@ plain-data ``OpSpec`` wire form, and ships the whole thing across the boundary as a bincode-serialised ``EngineSpec``. +With direct FPM interpolation, an explicit ``FpmModelProfile`` supplies identity +and resources instead. Compilation emits whole-forward timing operations with +empty SOL lists and never constructs an analytical model. + Since the pyo3 op unification the ``Operation`` objects the model layer builds ARE engine ops (Rust pyclasses; ``operations/*.py`` keeps thin data-plane shells). Serialization is therefore Rust-to-Rust: @@ -47,6 +51,11 @@ from aisimulate_core.sdk.config_builders import apply_nextn, build_model_config, validate_moe_controls from aisimulate_core.sdk.deepseek_v41 import MODEL_PATH as DEEPSEEK_V41_MODEL_PATH from aisimulate_core.sdk.errors import InvalidEngineConfigurationError as InvalidEngineConfigurationError +from aisimulate_core.sdk.fpm_profile import ( + FpmDeploymentProfile, + FpmModelProfile, + load_fpm_profile, +) from aisimulate_core.sdk.models import get_model from aisimulate_core.sdk.models.deepseek_v32 import _generation_ops_for_engine from aisimulate_core.sdk.models.helpers import resolve_dsv4_moe_arch, resolve_sglang_mla_compute @@ -143,6 +152,7 @@ # The pilot and AIC-1781 concurrently claimed 21; reject both older layouts. # - 23 (DCP identity): optional recorded dcp_size in ParallelMapping. # - 24 (typed FPM DCP): dcp_size moves out of the op matching tuple. +# - 25 (FPM decoupling): merge the SOL/direct selector into the typed DCP layout. # Single owner: the Rust crate constant. Python re-exports it for # diagnostics/tests instead of declaring a twin to keep in sync. ENGINE_SPEC_SCHEMA_VERSION = aisimulate_core.engine_spec_schema_version() @@ -181,6 +191,7 @@ def _fpm_spec_dict(op: FPMForwardOp) -> dict: # (1 = plain AR). Set by the fpm hybrid rewrite in models when a # draft scheme is materialized. "verify_width": int(getattr(op, "_verify_width", 1) or 1), + "interpolation": "sol", "sol_ops": [json.loads(_as_engine_op(c)._spec_json()) for c in op._sol_ops], } } @@ -470,6 +481,9 @@ def compile_engine( systems_path: str | None = None, forward_model: str | None = None, decoder_replay: bool = False, + fpm_profile: dict | str | FpmModelProfile | None = None, + fpm_interpolation: str | None = None, + cp_size: int = 1, database_mode: str | None = None, shared_layer: bool | None = None, transfer_policy: str | list[str] | None = None, @@ -486,6 +500,13 @@ def compile_engine( decomposed), ``context_ops`` and ``generation_ops`` into OpSpecs and returns the bytes produced by the Rust ``engine_spec_bincode_from_json`` pyfunction. + ``fpm_profile`` supplies independent decoder metadata and rank-local + resource bounds. This is internal compilation plumbing: the canonical + Rust constructor supplies an already resolved ``sol`` or ``direct`` method. + Profile precision is authoritative on either route, with + conflicting caller overrides rejected. ``cp_size`` completes the profile + identity; profiles currently support only CP1/PP1 and no speculation. + ``decode_workload_distribution=None`` preserves the default model and latency behavior. The observed GLM-5.2 NVFP4 pilot profile selects only generation MoE on the exact VR200 SGLang runtime, TP4/MoETP4/EP1, BF16 GEMM/FMHA, FP8 KV and @@ -493,8 +514,8 @@ def compile_engine( logical batches are exploratory and queries outside 1..32 fail. Missing or mismatched approved profile data raises ``DecodeMoeProfileError``. - Engine schema 24 persists the exact-profile policy on MoE operators. Default - OpSpec JSON includes that new false field, and older binary specs require + Engine schema 25 combines the FPM interpolation selector, typed DCP, and + exact-profile policy on MoE operators. Older binary specs require recompilation; default representations are therefore not byte-identical. """ if not isinstance(decoder_replay, bool): @@ -503,34 +524,22 @@ def compile_engine( raise InvalidEngineConfigurationError( f"decoder_replay requires model={DEEPSEEK_V41_MODEL_PATH!r} and backend='sglang'" ) - + if type(cp_size) is not int or cp_size != 1: + raise ValueError("cp_size must be the integer 1; this SDK entry point does not support context parallelism") if fpm_parquet_path is not None: if not fpm_parquet_path: raise ValueError("fpm_parquet_path cannot be empty") if forward_model != "fpm": raise ValueError("fpm_parquet_path requires forward_model='fpm'") - # `_build_model_config` resolves MoE parallelism defaults internally and - # does not take a model_path (quant inference is done inside `get_model`). - if prefill_graph_profile is not None: - from aisimulate_core.sdk.errors import PrefillGraphProfileError - - if shared_layer: - raise PrefillGraphProfileError("prefill_graph_profile does not allow shared-source inheritance") - shared_layer = False from aisimulate_core.sdk.fpm_config import FpmCompileConfig - from aisimulate_core.sdk.speculation import SpeculationConfig - recorded_attention_backend = attention_backend - if backend == "vllm" and attention_backend == "FLASHINFER_MLA": - attention_backend = "flashinfer" - 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 try: fpm_config = FpmCompileConfig( fpm_options, - attention_backend=recorded_attention_backend, + attention_backend=attention_backend, fmha_quant_mode=fmha_quant_mode if fmha_quant_mode is not None else fpm_fmha_quant_mode, comm_quant_mode=comm_quant_mode, + has_profile=fpm_profile is not None, ) options_path = fpm_config.options["fpm_parquet_path"] if options_path is not None: @@ -539,6 +548,102 @@ def compile_engine( if forward_model != "fpm": raise ValueError("fpm_parquet_path requires forward_model='fpm'") fpm_parquet_path = options_path + except (ValueError, TypeError, KeyError) as exc: + raise InvalidEngineConfigurationError(str(exc)) from exc + profile = load_fpm_profile(fpm_profile) if fpm_profile is not None else None + if fpm_interpolation not in (None, "sol", "direct") or (profile is not None and fpm_interpolation is None): + raise InvalidEngineConfigurationError("FPM compilation requires a resolved method from best_available(config)") + if fpm_interpolation == "direct" and profile is None: + raise InvalidEngineConfigurationError("direct FPM interpolation requires an fpm_profile") + if ( + fpm_interpolation is not None + and fpm_options is not None + and "method" in fpm_options + and fpm_options["method"] != fpm_interpolation + ): + raise InvalidEngineConfigurationError("conflicting fpm_interpolation and fpm_options.method") + interpolation = fpm_interpolation + deployment = None + if profile is not None: + if any(value is not None for value in (moe_kernel_source, decode_workload_distribution, prefill_graph_profile)): + raise InvalidEngineConfigurationError( + "FPM profiles do not support moe_kernel_source, decode_workload_distribution or prefill_graph_profile" + ) + if dcp_size is not None: + raise InvalidEngineConfigurationError("fpm_profile does not describe recorded DCP identity") + if nextn or speculation is not None: + raise InvalidEngineConfigurationError("FPM profiles support plain autoregressive decoder-only execution") + literal_version = _literal_backend_version(system, backend, backend_version, systems_path, None) + deployment = profile.select( + model=model_path, + system=system, + backend=backend, + backend_version=literal_version, + tp_size=tp_size, + pp_size=pp_size, + attention_dp_size=attention_dp_size, + moe_tp_size=moe_tp_size, + moe_ep_size=moe_ep_size, + cp_size=cp_size, + ) + deployment.validate_overrides( + gemm_quant_mode=gemm_quant_mode, + moe_quant_mode=moe_quant_mode, + kvcache_quant_mode=kvcache_quant_mode, + fmha_quant_mode=fmha_quant_mode, + comm_quant_mode=comm_quant_mode, + attention_backend=attention_backend, + moe_backend=moe_backend, + enable_eplb=enable_eplb, + wideep_num_slots=wideep_num_slots, + ) + try: + fpm_config = FpmCompileConfig( + fpm_options, + attention_backend=deployment.attention_backend, + fmha_quant_mode=deployment.fmha_quant_mode, + comm_quant_mode=deployment.comm_quant_mode, + has_profile=True, + ) + except (ValueError, TypeError, KeyError) as exc: + raise InvalidEngineConfigurationError(str(exc)) from exc + if forward_model == "fpm" and interpolation == "direct": + spec_json = _direct_fpm_spec_json( + profile, + deployment, + kv_block_size=kv_block_size, + systems_path=systems_path, + database_mode=database_mode, + shared_layer=shared_layer, + transfer_policy=transfer_policy, + strict_provenance=strict_provenance, + fpm_parquet_path=fpm_parquet_path, + fpm_fmha_quant_mode=fpm_fmha_quant_mode, + fpm_options=fpm_config.options, + ) + return bytes(aisimulate_core.engine_spec_bincode_from_json(spec_json)) + gemm_quant_mode = deployment.gemm_quant_mode + moe_quant_mode = deployment.moe_quant_mode + kvcache_quant_mode = deployment.kv_cache_dtype + fmha_quant_mode = deployment.fmha_quant_mode + comm_quant_mode = deployment.comm_quant_mode + attention_backend = None if deployment.attention_backend == "auto" else deployment.attention_backend + + # `_build_model_config` resolves MoE parallelism defaults internally and + # does not take a model_path (quant inference is done inside `get_model`). + if prefill_graph_profile is not None: + from aisimulate_core.sdk.errors import PrefillGraphProfileError + + if shared_layer: + raise PrefillGraphProfileError("prefill_graph_profile does not allow shared-source inheritance") + shared_layer = False + from aisimulate_core.sdk.speculation import SpeculationConfig + + if backend == "vllm" and attention_backend == "FLASHINFER_MLA": + attention_backend = "flashinfer" + 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 + try: validate_moe_controls( model_path=model_path, enable_eplb=enable_eplb, @@ -581,6 +686,8 @@ def compile_engine( resolve_dsv4_moe_arch(model_config, model_path, system_name=system, backend_name=backend) except (ValueError, TypeError, KeyError) as exc: raise InvalidEngineConfigurationError(str(exc)) from exc + if deployment is not None: + model_config.moe_backend = None if deployment.moe_backend == "auto" else deployment.moe_backend # Slot policy FIRST, tolerance second: resolve the requested version to a # literal (raising on unlisted versions / unpopulated aliases) before the @@ -592,6 +699,17 @@ def compile_engine( model_config, model_path, backend, literal_version, load_system_spec(system, systems_path) ) model = get_model(model_path, model_config, backend) + if deployment is not None and forward_model == "fpm": + from aisimulate_core.sdk.fpm_identity import LEGACY_EXECUTION_IDENTITY + + expected_identity = deployment.match_identity() + list(LEGACY_EXECUTION_IDENTITY) + if fpm_fmha_quant_mode is not None: + expected_identity[2] = fpm_fmha_quant_mode + for op in (*model.context_ops, *model.generation_ops): + if not isinstance(op, FPMForwardOp) or list(op._match_identity) != expected_identity: + raise InvalidEngineConfigurationError( + "registered SOL construction does not preserve the requested FPM profile identity" + ) database = _maybe_load_database( system, backend, @@ -620,9 +738,98 @@ def compile_engine( fpm_parquet_path=fpm_parquet_path, ) + if profile is not None: + spec = json.loads(spec_json) + spec["engine"]["extra"].update( + fpm_profile=profile.model_dump_json(), + estimator_config=json.dumps({"fpm_interpolation": {**fpm_config.options, "method": interpolation}}), + ) + spec_json = json.dumps(spec) + return bytes(aisimulate_core.engine_spec_bincode_from_json(spec_json)) +def _direct_fpm_spec_json( + profile: FpmModelProfile, + deployment: FpmDeploymentProfile, + *, + kv_block_size: int | None, + systems_path: str | None, + database_mode: str | None, + shared_layer: bool | None, + transfer_policy: str | list[str] | None, + strict_provenance: bool | None, + fpm_parquet_path: str | None = None, + fpm_fmha_quant_mode: str | None = None, + fpm_options: dict | None = None, +) -> str: + """Build whole-forward timing operations from metadata, without a graph.""" + if _database_mode_name(None, database_mode) != "SILICON": + raise ValueError("direct FPM interpolation requires database_mode='SILICON'; it has no analytical SOL model") + identity = deployment.match_identity() + if fpm_fmha_quant_mode is not None: + if fpm_fmha_quant_mode not in {"bfloat16", "fp8", "fp8_block"}: + raise ValueError("fpm_fmha_quant_mode must be bfloat16, fp8, or fp8_block") + identity[2] = fpm_fmha_quant_mode + controls = {**(fpm_options or {}), "method": "direct"} + if fpm_parquet_path is not None: + controls["fpm_parquet_path"] = fpm_parquet_path + engine = { + "schema_version": ENGINE_CONFIG_SCHEMA_VERSION, + "model_name": profile.model, + "system_name": deployment.system, + "systems_path": systems_path, + "backend": deployment.backend, + "backend_version": deployment.backend_version, + "forward_model": "fpm", + "fpm_parquet_path": fpm_parquet_path, + "fpm_fmha_dtype": _rust_quant_to_dtype(fpm_fmha_quant_mode), + "kv_block_size": kv_block_size, + "tp_size": deployment.tp, + "pp_size": deployment.pp, + "attention_dp_size": deployment.dp, + "moe_tp_size": deployment.moe_tp, + "moe_ep_size": deployment.moe_ep, + "cp_size": deployment.cp, + "weight_dtype": _rust_quant_to_dtype(deployment.gemm_quant_mode), + "moe_dtype": _rust_moe_quant_to_dtype(deployment.moe_quant_mode), + "activation_dtype": _rust_quant_to_dtype(deployment.fmha_quant_mode), + "kv_cache_dtype": _rust_quant_to_dtype(deployment.kv_cache_dtype), + "enable_shared_layer": shared_layer, + "strict_provenance": bool(strict_provenance), + "database_mode": "SILICON", + "transfer_policy": _transfer_policy_tokens(None, transfer_policy), + "extra": { + "fpm_profile": profile.model_dump_json(), + "estimator_config": json.dumps({"fpm_interpolation": controls}), + }, + } + + def phase_op(phase: str) -> dict: + return { + "FpmForward": { + "name": f"fpm_forward_{phase}", + "phase": phase, + "model_path": profile.model, + "match_identity": identity, + "original_fmha_quant_mode": deployment.fmha_quant_mode if fpm_fmha_quant_mode is not None else None, + "weight_bytes": deployment.resources.weights_bytes, + "verify_width": 1, + "sol_ops": [], + "interpolation": "direct", + } + } + + return json.dumps( + { + "schema_version": ENGINE_SPEC_SCHEMA_VERSION, + "engine": engine, + "context_ops": [phase_op("prefill")], + "generation_ops": [phase_op("decode")], + } + ) + + def build_ops_json(ops: Any) -> str: """Serialize an op list to the OpSpec JSON array the ad-hoc op-list evaluation FFI (``AicEngine.evaluate_ops_json``) consumes. diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_config.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_config.py index 748e04a39..e0c8e3452 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/fpm_config.py +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_config.py @@ -23,11 +23,12 @@ def __init__( attention_backend: str | None = None, fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None, + has_profile: bool = False, ) -> None: import aisimulate_core normalized = aisimulate_core.RustForwardPassPerfModel._normalize_fpm_options( - json.dumps({} if options is None else options), fmha_quant_mode, comm_quant_mode + json.dumps({} if options is None else options), fmha_quant_mode, comm_quant_mode, has_profile ) object.__setattr__(self, "_options_json", normalized) object.__setattr__(self, "attention_backend", attention_backend) diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py new file mode 100644 index 000000000..cdbc4779a --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FPM profile compatibility exports and runtime-facing validation bridges.""" + +from __future__ import annotations + +import json +from functools import cache +from pathlib import Path + +from aisimulate_core.fpm_profile import _MUTABLE_REFERENCES +from aisimulate_core.fpm_profile import FpmDeploymentProfile as FpmDeploymentProfile +from aisimulate_core.fpm_profile import FpmModelProfile as FpmModelProfile +from aisimulate_core.fpm_profile import FpmResourceProfile as FpmResourceProfile +from aisimulate_core.fpm_profile import load_fpm_profile as load_fpm_profile + + +def _quantization_from_engine_config(config_json: str) -> dict[str, str]: + """Restore exact profile modes after checking the native wire dtypes. + + EngineConfig stores dtypes rather than quantization modes: ``sq`` and + ``int8_wo`` both serialize as ``int8``. Both native timing and memory + bridges use the profile to recover the mode, but only when every supplied + wire dtype agrees with the selected deployment's canonical encoding. + """ + from aisimulate_core.sdk.rust_engine_step import _moe_quant_to_dtype, _quant_to_dtype + + config = json.loads(config_json) + profile = load_fpm_profile(config["extra"]["fpm_profile"]) + deployment = profile.select( + model=config["model_name"], + system=config["system_name"], + backend=config["backend"], + backend_version=config["backend_version"], + tp_size=config["tp_size"], + pp_size=config["pp_size"], + attention_dp_size=1 if config.get("attention_dp_size") is None else config["attention_dp_size"], + moe_tp_size=config.get("moe_tp_size"), + moe_ep_size=config.get("moe_ep_size"), + cp_size=1 if config.get("cp_size") is None else config["cp_size"], + ) + modes = deployment.quantization_kwargs() + for field, mode, encode in ( + ("weight_dtype", "gemm_quant_mode", _quant_to_dtype), + ("moe_dtype", "moe_quant_mode", _moe_quant_to_dtype), + ("activation_dtype", "fmha_quant_mode", _quant_to_dtype), + ("kv_cache_dtype", "kvcache_quant_mode", _quant_to_dtype), + ): + supplied = config.get(field) + expected = encode(modes[mode]) + if supplied is not None and supplied != expected: + raise ValueError( + f"FPM profile identity conflict: {field}={supplied!r}, " + f"profile {mode}={modes[mode]!r} requires wire dtype {expected!r}" + ) + modes.pop("attention_backend") + return modes + + +@cache +def _load_pinned_checkpoint_config(model: str, revision: str) -> dict | None: + """Reuse immutable remote metadata across candidate normalization and replay.""" + from aisimulate_core.sdk.utils import _download_hf_json + + return _download_hf_json(model, "config.json", revision=revision) + + +def _validate_checkpoint_architecture(profile: FpmModelProfile) -> None: + """Verify raw checkpoint metadata without loading weights or remote code.""" + from aisimulate_core.sdk.utils import HuggingFaceDownloadError, _load_local_config + + checkpoint = Path(profile.model).expanduser() + identity = f"{profile.model!r} at model_revision={profile.model_revision!r}" + try: + if checkpoint.exists() or checkpoint.is_absolute() or profile.model.startswith(("./", "../", "~")): + metadata = _load_local_config(str(checkpoint)) + else: + # Bundled configs do not record their revision and cannot establish + # the architecture of this pinned checkpoint. + metadata = _load_pinned_checkpoint_config(profile.model, profile.model_revision) + except (HuggingFaceDownloadError, OSError, ValueError) as error: + raise ValueError( + f"cannot verify FPM profile architecture: config.json metadata for {identity} " + f"could not be resolved; make the pinned checkpoint configuration available: {error}" + ) from error + if not isinstance(metadata, dict): + raise ValueError(f"FPM profile checkpoint config.json must be a mapping for {identity}") + # Match the collector's text-model view of checkpoint metadata. + text_config = metadata.get("text_config") + if isinstance(text_config, dict): + metadata = {**metadata, **text_config} + architectures = metadata.get("architectures") + if ( + not isinstance(architectures, list) + or not architectures + or not isinstance(architectures[0], str) + or not architectures[0].strip() + ): + raise ValueError( + f"FPM profile checkpoint config.json must declare a nonempty architectures list for {identity}" + ) + if profile.architecture != architectures[0]: + raise ValueError( + f"FPM profile architecture mismatch for {identity}: " + f"checkpoint={architectures[0]!r}, profile={profile.architecture!r}" + ) + + +def _validate_forward_pass_profile(config_json: str) -> str: + """Validate identity against checkpoint metadata and return registration facts. + + This path reads no timing data and does not construct an analytical model. + Estimator defaults and interpolation selection belong to Rust. + """ + config = json.loads(config_json) + profile = load_fpm_profile(config["fpm_profile"]) + version = config.get("backend_version") + if not isinstance(version, str) or version.lower() in _MUTABLE_REFERENCES | {"previous", "next"}: + raise ValueError("fpm_profile requires a literal backend_version") + deployment = profile.select( + model=config["model"], + system=config["system"], + backend=config["backend"], + backend_version=version, + tp_size=config["tp"], + pp_size=config["pp"], + attention_dp_size=config["attention_dp"], + moe_tp_size=config.get("moe_tp_size"), + moe_ep_size=config.get("moe_ep_size"), + ) + deployment.validate_overrides(**{key: config.get(key) for key in deployment.quantization_kwargs()}) + _validate_checkpoint_architecture(profile) + + from aisimulate_core.sdk import common + from aisimulate_core.sdk.models.base import _MODEL_REGISTRY + + family = common.ARCHITECTURE_TO_MODEL_FAMILY.get(profile.architecture, profile.architecture) + return json.dumps( + { + "profile": profile.model_dump(mode="json"), + "registered": family in _MODEL_REGISTRY, + **deployment.quantization_kwargs(), + } + ) diff --git a/python/aisimulate/src/aisimulate_core/sdk/memory.py b/python/aisimulate/src/aisimulate_core/sdk/memory.py index d9f7be288..b6e5b31b2 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory.py @@ -12,7 +12,11 @@ math -- fraction + tolerance validation, the native AIC memory model, the naive fallback, and the tolerance margin -- lives in this module. -Two estimators, each built then asked to ``estimate()``: +Resource sources: + +- **FPM profile**: caller-declared rank-local resource bounds bypass model + construction and timing data. The same budget arithmetic below applies, + after checking the scheduler against the declared resource envelope. - **Native** (:class:`KVCacheEstimator`): ``from_request`` reuses AIC's full backend memory model (``BaseBackend._get_memory_usage`` plus the model's @@ -46,6 +50,7 @@ from aisimulate_core.sdk.backends.factory import get_backend from aisimulate_core.sdk.common import DefaultHFModels from aisimulate_core.sdk.config_builders import apply_nextn, build_model_config, validate_moe_controls, validate_nextn +from aisimulate_core.sdk.fpm_profile import FpmModelProfile, load_fpm_profile from aisimulate_core.sdk.models import get_model from aisimulate_core.sdk.models.helpers import resolve_sglang_mla_compute from aisimulate_core.sdk.utils import ( @@ -1040,6 +1045,8 @@ def estimate_kv_cache( naive_kv_reservation: float = _DEFAULT_NAIVE_KV_RESERVATION, allow_naive_fallback: bool = False, allow_hf_config_download: bool = False, + fpm_profile: dict | str | FpmModelProfile | None = None, + cp_size: int = 1, ) -> dict[str, Any]: """Compute the KV-cache memory estimate (raw + optional tolerance margin). @@ -1078,11 +1085,17 @@ def estimate_kv_cache( model build is unsupported (default off -> the error propagates). allow_hf_config_download: allow the naive fallback to download a ``config.json`` from HuggingFace when it is not local / pre-cached. + fpm_profile: explicit FPM identity and conservative rank-local resource + bounds. This route reads only hardware specifications, requires the + requested scheduler envelope to fit the profile, and never builds + a model or falls back to naive estimation. CUDA graph bytes remain + a separate reservation in addition to declared profile overheads. + cp_size: profile context-parallel identity; currently only CP1 is supported. Returns: A flat dict with ``total_gpu_capacity_bytes``, ``total_kv_size_bytes``, ``kv_size_per_token_bytes``, ``total_kv_size_tokens``, ``source`` - (``"native"`` | ``"naive_fallback"``), ``memory_breakdown`` (dict on the + (``"native"`` | ``"naive_fallback"`` | ``"profile"``), ``memory_breakdown`` (dict on the native path, ``None`` on the fallback), and ``tolerance_adjusted`` (a dict when ``tolerance_fraction`` is set, else ``None``). @@ -1091,6 +1104,8 @@ def estimate_kv_cache( tolerance, no KV budget, insufficient model metadata, or (with the fallback off) an unsupported model/backend. """ + if type(cp_size) is not int or cp_size != 1: + raise ValueError("cp_size must be the integer 1; this SDK entry point does not support context parallelism") _validate_memory_fraction(backend, memory_fraction_kind, memory_fraction_value) _validate_tolerance(tolerance_fraction) _validate_naive_reservation(naive_kv_reservation) @@ -1100,6 +1115,63 @@ def estimate_kv_cache( # Validate the compute-side MTP depth before any fallback path. validate_nextn(nextn) + if fpm_profile is not None: + if nextn: + raise ValueError("FPM profile resources support plain autoregressive execution; nextn must be 0") + profile = load_fpm_profile(fpm_profile) + if backend_version is None: + raise ValueError("FPM profile memory estimation requires the profile's literal backend_version") + deployment = profile.select( + model=model_path, + system=system, + backend=backend, + backend_version=backend_version, + tp_size=tp_size, + pp_size=pp_size, + attention_dp_size=attention_dp_size, + moe_tp_size=moe_tp_size, + moe_ep_size=moe_ep_size, + cp_size=cp_size, + ) + deployment.validate_overrides( + gemm_quant_mode=gemm_quant_mode, + moe_quant_mode=moe_quant_mode, + kvcache_quant_mode=kvcache_quant_mode, + fmha_quant_mode=fmha_quant_mode, + comm_quant_mode=comm_quant_mode, + attention_backend=attention_backend, + moe_backend=moe_backend, + enable_eplb=enable_eplb, + wideep_num_slots=wideep_num_slots, + ) + resources = deployment.resources + resources.validate_envelope(max_num_tokens=max_num_tokens, max_batch_size=max_batch_size) + if resources.non_kv_bytes + cuda_graph_reserved_bytes > _MAX_EXACT_BYTE_COUNT: + raise ValueError("total FPM non-KV bytes including CUDA graphs must not exceed 2**53") + capacity = gpu_memory_capacity_bytes_override + if capacity is None: + capacity = perf_database.load_system_spec(system, systems_paths=systems_path)["gpu"]["mem_capacity"] + estimate = KVCacheEstimator._estimate_from_breakdown( + { + "weights_bytes": resources.weights_bytes, + "activations_bytes": resources.activations_bytes, + "runtime_overhead_bytes": resources.runtime_overhead_bytes, + "comm_overhead_bytes": resources.comm_overhead_bytes, + "non_kv_bytes": resources.non_kv_bytes, + "kv_size_per_token_bytes": resources.kv_bytes_per_token, + "gpu_memory_capacity_bytes": capacity, + }, + is_of_free=is_of_free, + fraction=fraction, + gpu_memory_capacity_bytes_override=gpu_memory_capacity_bytes_override, + cuda_graph_reserved_bytes=cuda_graph_reserved_bytes, + ) + estimate["source"] = "profile" + estimate["resource_provenance"] = resources.provenance + _apply_tolerance(estimate, tolerance_fraction) + estimate.pop("tokens_from_kv_bytes", None) + return estimate + validate_moe_controls( model_path=model_path, enable_eplb=enable_eplb, @@ -1107,7 +1179,6 @@ def estimate_kv_cache( moe_backend=moe_backend, moe_kernel_source=moe_kernel_source, ) - try: native = KVCacheEstimator.from_request( model_path, @@ -1213,6 +1284,8 @@ def estimate_num_gpu_blocks( allow_naive_fallback: bool = False, allow_hf_config_download: bool = False, diagnostics: dict[str, Any] | None = None, + fpm_profile: dict | str | FpmModelProfile | None = None, + cp_size: int = 1, ) -> int: """Convert the KV-cache token capacity to a scheduler block count. @@ -1248,13 +1321,13 @@ def estimate_num_gpu_blocks( system, backend, backend_version=backend_version, - max_num_tokens=int(max_num_tokens), - max_batch_size=int(max_batch_size), + max_num_tokens=max_num_tokens if fpm_profile is not None else int(max_num_tokens), + max_batch_size=max_batch_size if fpm_profile is not None else int(max_batch_size), memory_fraction_kind=memory_fraction_kind, memory_fraction_value=float(memory_fraction_value), - tp_size=int(tp_size), - pp_size=int(pp_size), - attention_dp_size=int(attention_dp_size), + tp_size=tp_size if fpm_profile is not None else int(tp_size), + pp_size=pp_size if fpm_profile is not None else int(pp_size), + attention_dp_size=attention_dp_size if fpm_profile is not None else int(attention_dp_size), moe_tp_size=moe_tp_size, moe_ep_size=moe_ep_size, gemm_quant_mode=gemm_quant_mode, @@ -1267,7 +1340,7 @@ def estimate_num_gpu_blocks( attention_backend=attention_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, - nextn=int(nextn), + nextn=nextn if fpm_profile is not None else int(nextn), systems_path=systems_path, gpu_memory_capacity_bytes_override=gpu_memory_capacity_bytes_override, cuda_graph_reserved_bytes=cuda_graph_reserved_bytes, @@ -1275,6 +1348,8 @@ def estimate_num_gpu_blocks( naive_kv_reservation=float(naive_kv_reservation), allow_naive_fallback=allow_naive_fallback, allow_hf_config_download=allow_hf_config_download, + fpm_profile=fpm_profile, + cp_size=cp_size, ) adjusted = estimate.get("tolerance_adjusted") diff --git a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py index 148e3edff..1cdd5fc35 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -22,7 +22,7 @@ import threading from collections import OrderedDict from collections.abc import Mapping -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, fields from dataclasses import field as dataclass_field from importlib import resources as pkg_resources from pathlib import Path @@ -116,6 +116,7 @@ class ForwardPassPerfModelConfig: backend: str worker_type: str backend_version: str | None = None + fpm_profile: dict[str, Any] | None = dataclass_field(default=None, kw_only=True) tp: int = 1 pp: int = 1 attention_dp: int = 1 @@ -176,7 +177,12 @@ def to_dict(self) -> dict[str, Any]: @dataclass(frozen=True) class ForwardPassPerfOptions: - """Runtime observation, regression, correction, and capacity controls.""" + """Legacy controls; serialization retains only explicit constructor arguments. + + Use ``ForwardPassPerfOptions(**(options.to_dict() | changes))`` to retain + omitted fields when modifying options. ``dataclasses.replace`` supplies + every field to the constructor, making all of its values explicit. + """ max_observations: int = 64 min_observations: int = 5 @@ -192,8 +198,25 @@ class ForwardPassPerfOptions: bucket_shape: tuple[int, int] | None = None regression_ridge_scale: float = 1e-9 + def __new__(cls, *args: Any, **kwargs: Any) -> ForwardPassPerfOptions: + instance = super().__new__(cls) + # Capture presence before the dataclass initializer supplies defaults. + # Comparing values against defaults would discard explicit user choices. + positional = tuple(field.name for field in fields(cls))[: len(args)] + object.__setattr__(instance, "_explicit_fields", frozenset((*positional, *kwargs))) + return instance + + def __setstate__(self, state: dict[str, Any]) -> None: + for name, value in state.items(): + object.__setattr__(self, name, value) + if "_explicit_fields" not in state: + # Older pickles have saved values but no constructor-presence metadata. + object.__setattr__( + self, "_explicit_fields", frozenset(field.name for field in fields(self) if field.name in state) + ) + def to_dict(self) -> dict[str, Any]: - return asdict(self) + return {name: value for name, value in asdict(self).items() if name in self._explicit_fields} class RustForwardPassPerfModel: @@ -286,20 +309,14 @@ def best_available(cls, config: ForwardPassPerfModelConfig | Mapping[str, Any]) """ import aisimulate_core - 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(aisimulate_core.RustForwardPassPerfModel.best_available(_json_dumps(payload))) + return cls(aisimulate_core.RustForwardPassPerfModel.best_available(_forward_pass_config_json(config))) + + @staticmethod + def normalize_config(config: ForwardPassPerfModelConfig | Mapping[str, Any]) -> dict[str, Any]: + """Expand and validate the Rust-owned configuration without constructing a model.""" + import aisimulate_core + + return json.loads(aisimulate_core.RustForwardPassPerfModel.normalize_config(_forward_pass_config_json(config))) def predict_prefill_latency(self, bs: int, isl: int, prefix: int = 0) -> float: """Return latency in ms for a qualified homogeneous graph-prefill shape. @@ -326,6 +343,15 @@ def estimate_forward_pass_time_ms(self, metrics: dict[str, Any] | list[dict[str, """ return self._inner.estimate_forward_pass_time_ms(_json_dumps(metrics)) + def estimate_forward_pass_detailed(self, metrics: dict[str, Any] | list[dict[str, Any]]) -> dict[str, Any]: + """Return timing with direct-FPM support, rank composition and correction. + + Measurement weights explain raw lookups. ``ranks`` records mixed-pass + subtraction and ``max_rank`` identifies the native rank maximum, before + ``correction_factor``. Other estimators have no direct lookup evidence. + """ + return json.loads(self._inner.estimate_forward_pass_detailed(_json_dumps(metrics))) + def tune_with_fpms(self, iterations: dict[str, Any] | list[Any]) -> None: """API: ``model.tune_with_fpms(iterations) -> None``. @@ -411,6 +437,23 @@ def _json_dumps(value: Any) -> str: return json.dumps(value, separators=(",", ":"), sort_keys=True) +def _forward_pass_config_json(config: ForwardPassPerfModelConfig | Mapping[str, Any]) -> str: + 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 _json_dumps(payload) + + def _optional_json_dumps(value: Mapping[str, Any] | None) -> str | None: if value is None: return None diff --git a/python/aisimulate/src/aisimulate_core/sdk/utils.py b/python/aisimulate/src/aisimulate_core/sdk/utils.py index 94d6251ce..1a94f5a25 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/utils.py +++ b/python/aisimulate/src/aisimulate_core/sdk/utils.py @@ -17,6 +17,7 @@ import urllib.request from functools import cache from pathlib import Path +from urllib.parse import quote import yaml @@ -431,9 +432,9 @@ def _get_hf_auth_headers() -> dict[str, str]: return headers -def _download_hf_json(hf_id: str, filename: str, *, raise_on_404: bool = True) -> dict | None: +def _download_hf_json(hf_id: str, filename: str, *, raise_on_404: bool = True, revision: str = "main") -> dict | None: """Download and parse a JSON file from a HuggingFace model repo.""" - url = f"https://huggingface.co/{hf_id}/raw/main/{filename}" + url = f"https://huggingface.co/{hf_id}/raw/{quote(revision, safe='')}/{filename}" try: req = urllib.request.Request(url, headers=_get_hf_auth_headers()) with urllib.request.urlopen(req, timeout=30) as response: 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 e7ee92809..bfa3bad75 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -107,6 +107,8 @@ def test_stable_function_signatures() -> None: "systems_path: 'str | None' = None, " "forward_model: 'str | None' = None, " "decoder_replay: 'bool' = False, " + "fpm_profile: 'dict | str | FpmModelProfile | None' = None, " + "fpm_interpolation: 'str | None' = None, cp_size: 'int' = 1, " "database_mode: 'str | None' = None, shared_layer: 'bool | None' = None, " "transfer_policy: 'str | list[str] | None' = None, " "strict_provenance: 'bool | None' = None, " diff --git a/python/aisimulate/tests/cross_package/test_import_contract.py b/python/aisimulate/tests/cross_package/test_import_contract.py index ba4fe99f5..3fb0a8a37 100644 --- a/python/aisimulate/tests/cross_package/test_import_contract.py +++ b/python/aisimulate/tests/cross_package/test_import_contract.py @@ -9,7 +9,10 @@ import importlib import importlib.resources +import json +import pickle import sys +from pathlib import Path from unittest.mock import patch import pytest @@ -31,6 +34,7 @@ "engine", "engine_table_view", "errors", + "fpm_profile", "fpm_config", "fpm_dataset", "fpm_identity", @@ -133,6 +137,134 @@ def test_legacy_leaf_module_is_canonical_module(module_suffix: str) -> None: assert sys.modules[legacy_name] is sys.modules[canonical_name] +@pytest.mark.parametrize("namespace", ["aisimulate.sdk", "aisimulate_core.sdk"]) +def test_fpm_profile_alias_preserves_module_and_type_identity(namespace: str) -> None: + """All public profile imports must share the canonical classes and state.""" + alias_name = f"{namespace}.fpm_profile" + alias = importlib.import_module(alias_name) + canonical = importlib.import_module("aisimulate_core.sdk.fpm_profile") + lightweight = importlib.import_module("aisimulate.fpm_profile") + core_types = importlib.import_module("aisimulate_core.fpm_profile") + + assert alias is canonical + assert sys.modules[alias_name] is canonical + for name in ("FpmModelProfile", "FpmDeploymentProfile", "FpmResourceProfile"): + assert getattr(alias, name) is getattr(canonical, name) + assert getattr(canonical, name) is getattr(lightweight, name) + assert getattr(canonical, name) is getattr(core_types, name) + legacy_global = f"c{namespace}.fpm_profile\n{name}\n.".encode() + assert pickle.loads(legacy_global) is getattr(lightweight, name) + assert pickle.loads(f"caisimulate.fpm_profile\n{name}\n.".encode()) is getattr(core_types, name) + assert alias.load_fpm_profile is lightweight.load_fpm_profile + + +@pytest.mark.parametrize("namespace", ["aisimulate.sdk", "aisimulate_core.sdk"]) +@pytest.mark.parametrize( + "name", ["GEMMQuantMode", "MoEQuantMode", "FMHAQuantMode", "KVCacheQuantMode", "CommQuantMode", "QuantMapping"] +) +def test_quantization_exports_preserve_shared_types_and_pickles(namespace: str, name: str) -> None: + from aisimulate import quantization + from aisimulate_core import quantization as core_quantization + + alias = importlib.import_module(f"{namespace}.common") + shared = getattr(quantization, name) + assert getattr(alias, name) is shared + assert getattr(core_quantization, name) is shared + # Historic pickle GLOBAL references resolve through the SDK re-exports. + assert pickle.loads(f"c{namespace}.common\n{name}\n.".encode()) is shared + assert pickle.loads(f"caisimulate.quantization\n{name}\n.".encode()) is shared + if name != "QuantMapping": + for member in shared: + assert pickle.loads(pickle.dumps(member)) is member + value = pickle.loads(pickle.dumps(member.value)) + assert type(value) is quantization.QuantMapping + assert value == member.value + + +@pytest.mark.parametrize("namespace", ["aisimulate.sdk", "aisimulate_core.sdk"]) +def test_fpm_profile_alias_instances_load_and_compile(namespace: str, tmp_path: Path) -> None: + """Profiles created through either facade must reach native compilation.""" + from aisimulate_core.sdk import engine + from aisimulate_core.sdk.fpm_profile import FpmModelProfile, load_fpm_profile + + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["ImportContractDecoderForCausalLM"]})) + alias = importlib.import_module(f"{namespace}.fpm_profile") + profile = alias.FpmModelProfile.model_validate( + { + "schema_version": 1, + "model": str(checkpoint), + "model_revision": "import-contract-fixture-v1", + "architecture": "ImportContractDecoderForCausalLM", + "context_length": 4096, + "num_experts": 0, + "provenance": "Synthetic metadata for import compatibility; not a GPU measurement.", + "deployments": [ + { + "system": "test_gpu", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 1, + "dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "gemm_quant_mode": "fp8", + "moe_quant_mode": "fp8", + "fmha_quant_mode": "fp8", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 100, + "activations_bytes": 20, + "runtime_overhead_bytes": 30, + "comm_overhead_bytes": 50, + "kv_bytes_per_token": 10, + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 256, + "provenance": "Synthetic rank-local resource bounds; excludes CUDA graphs.", + }, + } + ], + } + ) + + loaded = load_fpm_profile(profile) + assert type(loaded) is FpmModelProfile + assert loaded.model_dump() == profile.model_dump() + + canonical = json.loads( + engine.aisimulate_core.RustForwardPassPerfModel.normalize_config( + json.dumps( + { + "model": profile.model, + "system": "test_gpu", + "backend": "vllm", + "backend_version": "0.25.1", + "worker_type": "aggregated", + "estimation_mode": "fpm_interpolation", + "fpm_profile": profile.model_dump(mode="json"), + } + ) + ) + ) + assert canonical["fpm_profile"] == loaded.model_dump(mode="json") + assert canonical["estimator_config"]["fpm_interpolation"]["method"] == "direct" + compiled = engine.compile_engine( + profile.model, + "test_gpu", + "vllm", + "0.25.1", + systems_path=str(tmp_path), + forward_model="fpm", + fpm_profile=profile, + fpm_interpolation=canonical["estimator_config"]["fpm_interpolation"]["method"], + ) + assert isinstance(compiled, bytes) + assert compiled + + @pytest.mark.parametrize("package_suffix", ["models", "operations", "speculation"]) def test_legacy_package_reexports_canonical_public_surface(package_suffix: str) -> None: """Package facades preserve child wrappers and export canonical objects.""" diff --git a/python/aisimulate/tests/cross_package/test_single_distribution.py b/python/aisimulate/tests/cross_package/test_single_distribution.py index 45c01dc42..9ab856888 100644 --- a/python/aisimulate/tests/cross_package/test_single_distribution.py +++ b/python/aisimulate/tests/cross_package/test_single_distribution.py @@ -5,6 +5,8 @@ import importlib import importlib.metadata +import subprocess +import sys from pathlib import Path from packaging.requirements import Requirement @@ -12,6 +14,42 @@ APPLICATION_ROOT = Path(__file__).resolve().parents[2] +def test_core_metadata_is_cold_and_lazy_native_exports_retain_identity() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import sys +import aisimulate_core as core +from aisimulate_core import fpm_profile, quantization +blocked = ('aisimulate._runtime', 'aisimulate_core._native', 'aisimulate_core.sdk') +assert not any(name in sys.modules for name in blocked) +assert set(core.__all__) <= set(dir(core)) +try: + core.unknown_export +except AttributeError: + pass +else: + raise AssertionError('unknown native export must fail') +assert 'aisimulate._runtime' not in sys.modules +core.AicEngine +from aisimulate import _runtime +namespace = {} +exec('from aisimulate_core import *', namespace) +for name in core.__all__: + assert namespace[name] is getattr(core, name) + if name != '__version__': + assert getattr(core, name) is getattr(_runtime, name) +""", + ], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_canonical_core_exports_the_unified_native_types() -> None: core = importlib.import_module("aisimulate_core") runtime = importlib.import_module("aisimulate._runtime") diff --git a/python/aisimulate/tests/unit/collector/test_dsv41_isolated.py b/python/aisimulate/tests/unit/collector/test_dsv41_isolated.py index 93680c62c..645ccabfc 100644 --- a/python/aisimulate/tests/unit/collector/test_dsv41_isolated.py +++ b/python/aisimulate/tests/unit/collector/test_dsv41_isolated.py @@ -6,7 +6,6 @@ from pathlib import Path import pytest - from collector.sglang.dsv41_contract import build_manifest, canonical_json from collector.sglang.dsv41_isolated_runner import ( FRAMEWORK_COMMIT, diff --git a/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py b/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py new file mode 100644 index 000000000..110a1190e --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py @@ -0,0 +1,370 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import importlib.metadata +import json +import sys +from dataclasses import replace + +import pytest +import yaml +from collector.fpm_forward import capabilities, cli, memory_admission, planner, runner +from collector.fpm_forward.config import FPMCollectionOptions, reject_fpm_arguments_without_fpm +from collector.fpm_forward.model_capability import load_model_config + +pytestmark = pytest.mark.unit + + +def _profile(): + """Real GLM deployment identities with explicitly declared test-only bounds.""" + return { + "schema_version": 1, + "model": "nvidia/GLM-5.2-NVFP4", + "model_revision": "collector-test-snapshot", + "architecture": "GlmMoeDsaForCausalLM", + "context_length": 8192, + "num_experts": 256, + "provenance": "Collector test declaration; not a measured memory qualification.", + "deployments": [ + { + "system": "b200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": tp, + "dp": dp, + "moe_tp": 1, + "moe_ep": 8, + "gemm_quant_mode": "nvfp4", + "moe_quant_mode": "nvfp4", + "fmha_quant_mode": "fp8", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 100_000_000_000, + "activations_bytes": 1_000_000_000, + "runtime_overhead_bytes": 1_000_000_000, + "comm_overhead_bytes": 1_000_000_000, + "kv_bytes_per_token": 60_216, + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 1024, + "provenance": "Independent test envelope: 100 GB weights plus three 1 GB overhead bounds.", + }, + } + for tp, dp in ((1, 8), (8, 1)) + ], + } + + +def _argv(profile): + return [ + "--model-path", + profile["model"], + "--gpu", + "b200_sxm", + "--fpm-max-gpus", + "8", + "--fpm-gpu-counts", + "8", + "--fpm-parallel-presets", + "dep", + "tep", + "--fpm-kv-cache-dtypes", + "fp8", + ] + + +def _plan(profile, **overrides): + args = cli._parser().parse_args(_argv(profile)) + kwargs = { + "backend": "vllm", + "model_path": profile["model"], + "system": "b200_sxm", + "selected_ops": {"dsa_context_module", "dsa_generation_module"}, + "options": FPMCollectionOptions.from_args(args), + "fpm_profile": profile, + } + kwargs.update(overrides) + return planner.build_collection_plan(**kwargs) + + +@pytest.fixture +def no_models_or_timing_data(monkeypatch): + def forbidden(*_args, **_kwargs): + pytest.fail("profile collection constructed a model or consulted timing data") + + from aisimulate_core.sdk import engine, memory, models, perf_database + from aisimulate_core.sdk.models import base + + monkeypatch.setattr(models, "_MODEL_REGISTRY", {}) + monkeypatch.setattr(base, "_MODEL_REGISTRY", {}) + for module in (models, memory, engine): + monkeypatch.setattr(module, "get_model", forbidden) + monkeypatch.setattr(memory_admission.KVCacheEstimator, "from_request", forbidden) + for module in (capabilities, perf_database): + monkeypatch.setattr(module, "get_database", forbidden) + monkeypatch.setattr(module, "get_latest_database_version", forbidden) + monkeypatch.setattr(capabilities, "context_fmha_supported_modes", forbidden) + monkeypatch.setattr(planner, "_git_revision", lambda: "profile-collection-test") + + +@pytest.mark.parametrize("file_format", ["json", "yaml"]) +def test_glm_cli_plans_without_model_classes_or_any_timing_data( + tmp_path, monkeypatch, capsys, no_models_or_timing_data, file_format +): + profile = _profile() + path = tmp_path / f"profile.{file_format}" + path.write_text(json.dumps(profile) if file_format == "json" else yaml.safe_dump(profile)) + monkeypatch.setenv("COLLECTOR_MODEL_PATH", "") + + assert cli.main([*_argv(profile), "--fpm-model-profile", str(path), "--plan-only"]) == 0 + + output = json.loads(capsys.readouterr().out) + assert output["counts"]["cells"] == 4 + assert output["schema_version"] == 11 + # Schema11 adds the explicit legacy execution/cache/input identity and the + # native point source to cell IDs even for non-DeepSeek-V4.1 campaigns. + assert {cell["cell_id"] for cell in output["cells"]} == { + "fpm-340a2e34da72dc11", + "fpm-cd10cc816446953a", + "fpm-be0eb714954d0f58", + "fpm-dd91802723ac733f", + } + assert {cell["parallel_strategy"] for cell in output["cells"]} == {"dep", "tep"} + assert {cell["resolved_dtypes"]["fmha_quant_mode"] for cell in output["cells"]} == {"fp8"} + assert {cell["resolved_dtypes"]["fmha_resolution"] for cell in output["cells"]} == {"checkpoint_native"} + assert output["fpm_profile"]["model_revision"] == "collector-test-snapshot" + assert {decision["source"] for decision in output["topology_memory_admission"]} == {"fpm_profile_declared"} + for decision in output["topology_memory_admission"]: + assert decision["estimates"][0]["estimated_non_kv_bytes"] == 103_000_000_000 + assert decision["estimates"][0]["provenance"] == profile["deployments"][0]["resources"]["provenance"] + assert decision["activation_envelope"]["max_batch_size"] == 1024 + assert decision["activation_envelope"]["scope"] == "rank_local" + + +def test_unknown_architecture_plans_and_renders_from_real_config(tmp_path, no_models_or_timing_data): + profile = _profile() + config = load_model_config(profile["model"]).payload + config["architectures"] = ["UnregisteredMoeForCausalLM"] + profile.update(model="private-org/new-model", architecture=config["architectures"][0]) + path = tmp_path / "config.json" + path.write_text(json.dumps(config)) + + plan = _plan(profile, model_config_path=str(path), has_model_cases=False) + assert plan.capability.support_level == "bootstrap_template" + cell_dir = tmp_path / "rendered" + cell_dir.mkdir() + runner._render_cell(plan, plan.cells[0], cell_dir, {}) + + script = (cell_dir / "run.sh").read_text() + assert "--model private-org/new-model" in script + assert "--max-num-batched-tokens 8192" in script + assert "--max-num-seqs 1024" in script + + +def test_frozen_profile_content_invalidates_resume_without_changing_cell_ids(tmp_path, no_models_or_timing_data): + source = _profile() + first = _plan(source) + original = first.to_dict() + source["model_revision"] = "different-test-snapshot" + second = _plan(source) + detached = first.fpm_profile + detached.deployments[0].resources.weights_bytes += 10 + + assert first.to_dict() == original + assert first.sha256 != second.sha256 + assert [cell.cell_id for cell in first.cells] == [cell.cell_id for cell in second.cells] + assert _plan(original["fpm_profile"]).sha256 == first.sha256 + checkpoint = tmp_path / "checkpoint.json" + checkpoint.write_text(json.dumps({"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": first.sha256, "cells": {}})) + with pytest.raises(ValueError, match="checkpoint does not match the current frozen plan"): + runner._load_checkpoint(checkpoint, second, resume=True) + + +@pytest.mark.parametrize("field,value", [("fmha_quant_mode", "bfloat16"), ("moe_backend", "flashinfer_cutlass")]) +def test_profile_cannot_override_resolved_precision_or_backend(field, value, no_models_or_timing_data): + profile = _profile() + profile["deployments"][0][field] = value + with pytest.raises(ValueError, match="profile identity mismatch"): + _plan(profile) + + +def test_historical_minimax_transfer_precision_cannot_relabel_new_collection(no_models_or_timing_data): + profile = _profile() + profile.update(model="MiniMaxAI/MiniMax-M2.7", architecture="MiniMaxM2ForCausalLM") + deployment = profile["deployments"][0] + deployment.update( + system="h200_sxm", + tp=4, + dp=1, + moe_tp=4, + moe_ep=1, + gemm_quant_mode="fp8_block", + moe_quant_mode="fp8_block", + fmha_quant_mode="bfloat16", + ) + profile["deployments"] = [deployment] + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), + max_gpus=4, + gpu_counts=(4,), + parallel_presets=("pure_tp",), + ) + with pytest.raises(ValueError, match="fmha_quant_mode: resolved='fp8', profile='bfloat16'"): + _plan(profile, options=options, system="h200_sxm", selected_ops={"attention_context", "attention_generation"}) + + +@pytest.mark.parametrize("field,value", [("model", "other/model"), ("architecture", "OtherArchitecture")]) +def test_profile_model_identity_must_match_real_checkpoint(field, value, no_models_or_timing_data): + profile = _profile() + profile[field] = value + with pytest.raises(ValueError, match="profile (model identity|architecture) mismatch"): + _plan(profile, model_path="nvidia/GLM-5.2-NVFP4") + + +def test_missing_requested_topology_raises_before_memory_admission(monkeypatch, no_models_or_timing_data): + profile = _profile() + profile["deployments"] = profile["deployments"][:1] + monkeypatch.setattr(memory_admission, "load_system_spec", lambda *_args: pytest.fail("admission ran")) + with pytest.raises(ValueError, match="no matching FPM deployment profile"): + _plan(profile) + + +def test_profile_collection_rejects_non_native_kv_dtype(no_models_or_timing_data): + profile = _profile() + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), kv_cache_dtypes=("bfloat16",) + ) + with pytest.raises(ValueError, match="checkpoint-native KV dtype"): + _plan(profile, options=options) + + +@pytest.mark.parametrize("field,value", [("max_num_tokens", 4096), ("max_batch_size", 1)]) +def test_resource_envelope_errors_do_not_fail_open(field, value, no_models_or_timing_data): + profile = _profile() + profile["deployments"][0]["resources"][field] = value + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), max_prefill_batch_size=4 + ) + with pytest.raises(ValueError, match="resource envelope exceeded"): + _plan(profile, options=options) + + +def test_declared_size_capacity_filter_counts_rejected_topology(caplog, no_models_or_timing_data): + profile = _profile() + profile["deployments"][0]["resources"]["weights_bytes"] = 200_000_000_000 + plan = _plan(profile) + assert {cell.parallel_strategy for cell in plan.cells} == {"tep"} + assert [decision.disposition for decision in plan.topology_memory_admission] == ["rejected", "admitted"] + assert "dropped 1/2 topologies" in caplog.text + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +@pytest.mark.parametrize("smoke", [False, True]) +def test_profile_bounds_native_scheduling_before_cases_are_queued(phase, smoke, no_models_or_timing_data): + plan = _plan(_profile()) + cell = next(cell for cell in plan.cells if cell.workload_kind == phase) + arguments = runner._cell_generator_overrides(plan, cell, {}, smoke=smoke)["params"]["agg"]["extra_cli_args"] + assert arguments.count("--max-num-batched-tokens") == 1 + assert arguments[arguments.index("--max-num-batched-tokens") + 1] == "8192" + assert arguments.count("--max-num-seqs") == 1 + assert arguments[arguments.index("--max-num-seqs") + 1] == "1024" + + +def test_profile_preserves_explicit_prefill_envelope(no_models_or_timing_data): + profile = _profile() + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), + max_prefill_isl=2048, + max_prefill_batch_size=4, + ) + plan = _plan(profile, options=options) + for cell in plan.cells: + arguments = runner._cell_generator_overrides(plan, cell, {})["params"]["agg"]["extra_cli_args"] + expected_tokens, expected_batch = ("2048", "4") if cell.workload_kind == "prefill" else ("8192", "1024") + assert arguments[arguments.index("--max-num-batched-tokens") + 1] == expected_tokens + assert arguments[arguments.index("--max-num-seqs") + 1] == expected_batch + + +@pytest.mark.parametrize( + ("expected_version", "actual_version"), + [ + ("0.25.1", "0.25.1"), + ("0.25.1+custom", "0.25.1+custom"), + ("0.25.2.dev3", "0.25.2.dev3"), + ("0.25.2.dev3+custom", "0.25.2.dev3+custom"), + ("0.25.1", "0.24.0"), + ("0.25.1", "0.25.1+custom"), + ("0.25.1+custom", "0.25.1"), + ("0.25.1+custom", "0.25.1+other"), + ("0.25.2.dev3", "0.25.2.dev4"), + ("0.25.2", "0.25.2.dev3"), + ], +) +def test_profile_runtime_version_is_observed_before_execution(tmp_path, monkeypatch, expected_version, actual_version): + resource = object.__new__(runner.KubernetesCellRunner) + monkeypatch.setattr(runner, "FPM_RESULTS_DIR", str(tmp_path)) + monkeypatch.setattr(importlib.metadata, "version", lambda _: actual_version) + + def execute(_pod, command, timeout): + if command[:2] == ["python3", "-c"]: + with monkeypatch.context() as patch: + patch.setattr(sys, "argv", ["-c", *command[3:]]) + exec(command[2], {}) + + resource._exec_checked = execute + arguments = dict( + cell_id="cell", plan_sha256="plan", attempt_id="attempt", expected_backend_version=expected_version + ) + if actual_version != expected_version: + with pytest.raises(RuntimeError, match="profile runtime mismatch") as failure: + resource.prepare_attempt(["pod-0"], **arguments) + message = str(failure.value) + assert f"actual={actual_version!r}, expected={expected_version!r}" in message + assert "Use the pinned runtime" in message + assert f"create a new collection profile with backend_version={actual_version!r}" in message + assert "regenerate the collection plan in a new output directory" in message + else: + resource.prepare_attempt(["pod-0"], **arguments) + provenance = json.loads((tmp_path / "collector-provenance.json").read_text()) + assert provenance == { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + "cell_id": "cell", + "plan_sha256": "plan", + "attempt_id": "attempt", + "runtime": {"backend": "vllm", "backend_version": actual_version}, + } + + +def test_profile_input_is_not_valid_for_op_collection(): + from argparse import Namespace + + with pytest.raises(ValueError, match="--fpm-model-profile"): + reject_fpm_arguments_without_fpm(Namespace(ops=["gemm"], fpm_model_profile="profile.json")) + + +def test_multiple_profile_runtime_versions_fail_explicitly(no_models_or_timing_data): + profile = _profile() + other = copy.deepcopy(profile["deployments"][0]) + other["backend_version"] = "0.24.0" + profile["deployments"].append(other) + with pytest.raises(ValueError, match="must identify one runtime version"): + _plan(profile) + + +@pytest.mark.parametrize("decode_batch", [1024, 1025]) +def test_decode_collection_respects_profile_scheduler_envelope(no_models_or_timing_data, decode_batch): + profile = _profile() + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), + max_decode_batch_size=decode_batch, + ) + if decode_batch > profile["deployments"][0]["resources"]["max_batch_size"]: + with pytest.raises(ValueError, match="max_batch_size"): + _plan(profile, options=options) + else: + assert _plan(profile, options=options).cells diff --git a/python/aisimulate/tests/unit/collector/test_fpm_slurm.py b/python/aisimulate/tests/unit/collector/test_fpm_slurm.py index fea09b50e..5eeacdb83 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_slurm.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_slurm.py @@ -548,3 +548,103 @@ def run(): assert not failures assert results[0].returncode == -signal.SIGTERM assert denied[:2] == [signal.SIGTERM, 0] + + +@pytest.mark.parametrize("actual_version", ["0.25.1", "0.24.0"]) +def test_profile_campaign_observes_slurm_version_before_native_collection(tmp_path, monkeypatch, actual_version): + """Run the real profile campaign; only external Slurm/container operations are faked.""" + from dataclasses import replace + from pathlib import Path + + from collector.fpm_forward import cli + from collector.fpm_forward import runner as campaign + from collector.fpm_forward.config import FPMCollectionOptions + + from .test_fpm_profile_collection import _argv, _plan, _profile + from .test_fpm_runner import _native_payload + + profile = _profile() + options = FPMCollectionOptions.from_args( + cli._parser().parse_args( + [ + *_argv(profile), + "--fpm-executor", + "slurm", + "--fpm-slurm-container-image", + "image@sha256:synthetic", + ] + ) + ) + plan = _plan(profile, options=replace(options, parallel_presets=("tep",))) + commands = [] + executed = [] + monkeypatch.setenv("SLURM_JOB_ID", "1234") + monkeypatch.setattr("collector.fpm_forward.slurm.shutil.which", lambda name: f"/fake/{name}") + monkeypatch.setattr(campaign.time, "sleep", lambda _seconds: None) + import importlib.metadata + + real_version = importlib.metadata.version + monkeypatch.setattr( + importlib.metadata, "version", lambda name: actual_version if name == "vllm" else real_version(name) + ) + + def command(args, **_kwargs): + commands.append(args) + stdout = "" + if args[0] == "scontrol": + stdout = "JobId=1234 JobState=RUNNING NodeList=test-node" if "job" in args else "test-node" + elif args[0] == "squeue": + pass + else: + assert args[0] == "srun", args + assert "--jobid=1234" in args and "--gpus-per-node=8" in args + assert "--container-image=image@sha256:synthetic" in args + mounts = next(value.split("=", 1)[1] for value in args if value.startswith("--container-mounts=")) + raw = Path( + next(value.removesuffix(":/results") for value in mounts.split(",") if value.endswith(":/results")) + ) + assert raw.name == "node0000" + program = args[args.index("env") + 3 :] + if program[:2] == ["python3", "-c"]: + # Execute the emitted program with the container's result mount mapped to disk. + with monkeypatch.context() as patch: + patch.setattr(sys, "argv", ["-c", *program[3:]]) + try: + exec(program[2].replace("/results", str(raw)), {}) + except RuntimeError as error: + raise subprocess.CalledProcessError(1, args, output="prepared", stderr=str(error)) from error + else: + assert program == ["bash", "/tmp/fpm-bench/fpm_exec.sh"] + cell = next(item for item in plan.cells if item.cell_id == raw.parent.parent.name) + executed.append(cell.cell_id) + (raw / "benchmark.json").write_text(json.dumps(_native_payload(phase=cell.workload_kind, rank=0, dp=1))) + return subprocess.CompletedProcess(args, 0, stdout=stdout, stderr="") + + monkeypatch.setattr(campaign, "_run_command", command) + errors = campaign.run_collection( + plan, + generator_overrides={}, + checkpoint_dir=str(tmp_path / "checkpoints"), + artifact_root=str(tmp_path / "artifacts"), + database_root=str(tmp_path / "data"), + resume=False, + retry_failed=False, + ) + checkpoint = json.loads((tmp_path / "checkpoints/fpm_forward.json").read_text()) + if actual_version == "0.25.1": + assert errors == [] + assert executed == [cell.cell_id for cell in plan.cells] + assert checkpoint["database"]["status"] == "passed" + assert {entry["status"] for entry in checkpoint["cells"].values()} == {"passed"} + else: + assert len(errors) == 2 and all(error["classification"] == "campaign_cell_failed" for error in errors) + assert executed == [] and "database" not in checkpoint + for entry in checkpoint["cells"].values(): + failures = list(Path(entry["artifact_dir"]).glob("logs/transport-failures/*/stderr.log")) + assert len(failures) == 1 and "FPM profile runtime mismatch" in failures[0].read_text() + for cell in plan.cells: + entry = checkpoint["cells"][cell.cell_id] + provenance = json.loads((Path(entry["artifact_dir"]) / "raw/node0000/collector-provenance.json").read_text()) + assert provenance["runtime"] == {"backend": "vllm", "backend_version": actual_version} + assert provenance["plan_sha256"] == plan.sha256 and provenance["attempt_id"] == entry["attempt_id"] + assert all(args[0] in {"srun", "squeue", "scontrol"} for args in commands) diff --git a/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py b/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py index 6efdb90dc..3f84cd61f 100644 --- a/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py +++ b/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py @@ -243,11 +243,11 @@ def _fetch(database, attribute, *, shared_layer=None): monkeypatch.setattr(_etv, "fetch_attention_lane_density", _fetch) -def test_engine_spec_schema_version_is_twenty_four(): - """Typed FPM DCP changes the positional operator payload.""" +def test_engine_spec_schema_version_is_twenty_five(): + """Direct interpolation with typed FPM DCP changes the positional operator payload.""" from aisimulate.sdk import engine - assert engine.ENGINE_SPEC_SCHEMA_VERSION == 24 + assert engine.ENGINE_SPEC_SCHEMA_VERSION == 25 def test_lanes_outside_the_known_vocabulary_stay_reachable(): diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 48854ce34..b195c80c0 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -63,13 +63,13 @@ def test_fpm_compile_context_uses_rust_defaults_without_python_fallbacks(monkeyp class RustOptions: @staticmethod - def _normalize_fpm_options(payload, fmha, comm): - calls.append((json.loads(payload), fmha, comm)) + def _normalize_fpm_options(payload, fmha, comm, has_profile): + calls.append((json.loads(payload), fmha, comm, has_profile)) return json.dumps(resolved) monkeypatch.setattr(aisimulate_core, "RustForwardPassPerfModel", RustOptions) context = FpmCompileConfig(attention_backend="FLASHINFER_MLA") - assert calls == [({}, None, None)] + assert calls == [({}, None, None, False)] assert context.options == resolved assert context.cache_identity() == {"options": resolved, "attention_backend": "FLASHINFER_MLA"} @@ -1375,9 +1375,13 @@ def test_fpm_detail_distinguishes_memory_budget_from_runtime_capacity(external_f assert memory["estimated_num_gpu_blocks"] > 0 assert "num_gpu_blocks" not in memory assert set(sections) == {"summary", "memory", "time", "energy", "source"} - assert sections["time"]["diagnostics"]["status"] == "unavailable" - assert sections["source"]["status"] == "unavailable" - assert "whole-model FPM" in sections["source"]["unavailable_reason"] + assert sections["time"]["diagnostics"]["status"] == "available" + assert sections["source"]["status"] == "available" + operations = [op for phase in sections["source"]["phases"] for op in phase["operations"]] + assert operations and all(op["fpm_estimates"] for op in operations) + assert sections["energy"]["status"] == "unsupported" + assert sections["energy"]["diagnostics"]["power_w"] is None + assert sections["energy"]["diagnostics"]["power_coverage"] is None assert sections["time"]["serving_metrics"]["mean_ttft_ms"] > 0 diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_legacy_migration.py b/python/aisimulate/tests/unit/sdk/test_fpm_legacy_migration.py new file mode 100644 index 000000000..8cc8f117a --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_fpm_legacy_migration.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit legacy estimator controls survive migration and saved-config reloads.""" + +import copy +import json +import pickle + +import pytest + +from aisimulate_core.sdk import ForwardPassPerfModelConfig, ForwardPassPerfOptions, RustForwardPassPerfModel + +pytestmark = pytest.mark.unit + + +def _legacy_config(estimator_config=None, **extra): + if estimator_config is not None: + extra["estimator_config"] = json.dumps(estimator_config) + return { + "schema_version": 1, + "model_name": "test/model", + "system_name": "test-system", + "backend": "vllm", + "tp_size": 1, + "pp_size": 1, + "extra": extra, + } + + +def _reload(canonical): + return ForwardPassPerfModelConfig(**json.loads(json.dumps(canonical.to_dict()))) + + +def test_migration_merges_saved_and_caller_controls_before_round_trip(): + legacy = _legacy_config( + { + "fpm_interpolation": {"method": "sol"}, + "features": {"ffn_token_weight": 2}, + "correction": {"enabled": False, "factor_bounds": {"max": 3}}, + "fpm_regression": {"fit": {"kind": "standardized_nnls"}}, + } + ) + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config( + legacy, + "prefill", + { + "max_observations": 128, + "min_observations": 8, + "min_faster_correction_factor": 0.9, + "regression_attention_kv_weight": 7, + "bucket_shape": [4, 16], + "regression_ridge_scale": 0, + }, + ) + controls = RustForwardPassPerfModel.normalize_config(_reload(canonical))["estimator_config"] + assert controls["fpm_interpolation"]["method"] == "sol" + assert controls["features"] == { + "attention_kv_weight": 7, + "prefill_attention_pair_weight": 1, + "ffn_token_weight": 2, + } + assert controls["correction"]["enabled"] is False + assert controls["correction"]["factor_bounds"] == {"min": 0.9, "max": 3} + for store in ("correction", "fpm_regression"): + assert controls[store]["sampling"] == {"bins_per_axis": [4, 16], "max_observations": 128} + assert controls[store]["min_observations"] == 8 + assert controls["fpm_regression"]["fit"] == { + "kind": "standardized_nnls", + "singular_ridge_scale": 0, + "rebuild_interval": None, + } + + +@pytest.mark.parametrize("options", [None, {}, ForwardPassPerfOptions()]) +def test_omitted_caller_fields_do_not_override_saved_controls(options): + saved = { + "correction": {"enabled": False, "sampling": {"max_observations": 128}}, + "fpm_regression": {"fit": {"singular_ridge_scale": 0}}, + "features": {"attention_kv_weight": 7}, + } + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(_legacy_config(saved), "prefill", options) + controls = _reload(canonical).estimator_config + assert controls["correction"]["enabled"] is False + assert controls["correction"]["sampling"]["max_observations"] == 128 + assert controls["fpm_regression"]["sampling"]["max_observations"] == 64 + assert controls["fpm_regression"]["fit"]["singular_ridge_scale"] == 0 + assert controls["features"]["attention_kv_weight"] == 7 + + +def test_legacy_options_serialize_presence_including_explicit_defaults(): + options = ForwardPassPerfOptions(64, min_faster_correction_factor=None, regression_ridge_scale=0) + assert options.min_observations == 5 # Existing attribute defaults remain available. + assert options.to_dict() == { + "max_observations": 64, + "min_faster_correction_factor": None, + "regression_ridge_scale": 0, + } + loaded = ForwardPassPerfOptions(**json.loads(json.dumps(options.to_dict()))) + assert loaded.to_dict() == options.to_dict() + saved = {"features": {"attention_kv_weight": 7}} + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(_legacy_config(saved), "decode", loaded) + controls = _reload(canonical).estimator_config + assert controls["features"]["attention_kv_weight"] == 7 + assert controls["correction"]["factor_bounds"]["min"] is None + assert controls["fpm_regression"]["fit"]["singular_ridge_scale"] == 0 + + +@pytest.mark.parametrize("clone", [copy.copy, copy.deepcopy]) +@pytest.mark.parametrize("controls", [{}, {"max_observations": 64, "regression_ridge_scale": 0}]) +def test_legacy_options_copy_preserves_explicit_fields(clone, controls): + assert clone(ForwardPassPerfOptions(**controls)).to_dict() == controls + + +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +@pytest.mark.parametrize("controls", [{}, {"max_observations": 64, "regression_ridge_scale": 0}]) +def test_legacy_options_pickle_preserves_explicit_fields(protocol, controls): + options = ForwardPassPerfOptions(**controls) + assert pickle.loads(pickle.dumps(options, protocol=protocol)).to_dict() == controls + + +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +def test_legacy_options_pickle_without_presence_preserves_saved_controls(protocol): + options = ForwardPassPerfOptions(max_observations=128, min_observations=8, regression_ridge_scale=0) + # Older pickles saved dataclass attributes without constructor-presence metadata. + object.__delattr__(options, "_explicit_fields") + restored = pickle.loads(pickle.dumps(options, protocol=protocol)) + assert restored.to_dict() == vars(options) + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(_legacy_config(), "prefill", restored) + controls = _reload(canonical).estimator_config + for store in ("correction", "fpm_regression"): + assert controls[store]["sampling"]["max_observations"] == 128 + assert controls[store]["min_observations"] == 8 + assert controls["fpm_regression"]["fit"]["singular_ridge_scale"] == 0 + + +@pytest.mark.parametrize( + ("saved", "options"), + [ + ({"features": {"attention_kv_weight": 1}}, {"regression_attention_kv_weight": 1.0}), + ({"correction": {"factor_bounds": {"min": None}}}, {"min_faster_correction_factor": None}), + ({"correction": {"factor_bounds": {"min": 0.5}}}, {"min_faster_correction_factor": 0.5}), + ({"fpm_regression": {"sampling": {"bins_per_axis": [4, 16]}}}, {"bucket_shape": [4, 16]}), + ({"correction": {"sampling": {"bins_per_axis": [8, 8]}}}, {"bucket_count": 64}), + ], +) +def test_agreeing_overlap_is_accepted(saved, options): + saved = saved | {"fpm_interpolation": {"method": "sol"}} + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(_legacy_config(saved), "decode", options) + assert ( + RustForwardPassPerfModel.normalize_config(_reload(canonical))["estimator_config"] == canonical.estimator_config + ) + + +@pytest.mark.parametrize( + ("saved", "options", "field"), + [ + ( + {"fpm_regression": {"sampling": {"max_observations": 128}}}, + ForwardPassPerfOptions(max_observations=64), + "fpm_regression.sampling.max_observations", + ), + ( + {"features": {"attention_kv_weight": 7}}, + {"regression_attention_kv_weight": 1}, + "features.attention_kv_weight", + ), + ( + {"correction": {"factor_bounds": {"min": 0.5}}}, + {"min_faster_correction_factor": None}, + "correction.factor_bounds.min", + ), + ( + {"correction": {"sampling": {"bins_per_axis": [4, 16]}}}, + {"bucket_shape": [4, 4]}, + "correction.sampling.bins_per_axis", + ), + ( + {"fpm_regression": {"fit": {"singular_ridge_scale": 1e-9}}}, + {"regression_ridge_scale": 0}, + "fpm_regression.fit.singular_ridge_scale", + ), + ], +) +def test_conflicting_overlap_identifies_canonical_field(saved, options, field): + with pytest.raises(ValueError, match=rf"conflicting explicit values for estimator_config\.{field}"): + ForwardPassPerfModelConfig.from_legacy_engine_config(_legacy_config(saved), "prefill", options) + + +def test_sampling_validation_uses_merged_values(): + saved = { + "correction": {"sampling": {"max_observations": 128}}, + "fpm_regression": {"sampling": {"max_observations": 128}}, + } + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config( + _legacy_config(saved), "prefill", {"min_observations": 80} + ) + controls = _reload(canonical).estimator_config + assert controls["correction"]["min_observations"] == 80 + assert controls["fpm_regression"]["min_observations"] == 80 + with pytest.raises(ValueError, match="min_observations must be <= max_observations"): + ForwardPassPerfModelConfig.from_legacy_engine_config( + _legacy_config(saved), "prefill", {"min_observations": 129} + ) + + +@pytest.mark.parametrize("canonical_method", [None, "direct"]) +def test_omitted_and_agreeing_interpolation_methods_accept_legacy_selection(canonical_method): + controls = {"fpm_interpolation": {} if canonical_method is None else {"method": canonical_method}} + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config( + _legacy_config(controls, fpm_interpolation="direct"), "prefill" + ) + assert _reload(canonical).estimator_config["fpm_interpolation"]["method"] == "direct" + + +@pytest.mark.parametrize("canonical_method,legacy_method", [("auto", "direct"), ("direct", "auto"), ("sol", "direct")]) +def test_explicit_interpolation_method_conflicts_are_symmetric(canonical_method, legacy_method): + with pytest.raises(ValueError, match=r"estimator_config\.fpm_interpolation\.method"): + ForwardPassPerfModelConfig.from_legacy_engine_config( + _legacy_config({"fpm_interpolation": {"method": canonical_method}}, fpm_interpolation=legacy_method), + "prefill", + ) + + +@pytest.mark.parametrize("legacy_path", ["/tmp/fpm.parquet", "/tmp/./fpm.parquet", "/tmp//fpm.parquet"]) +def test_equivalent_parquet_paths_merge(legacy_path): + legacy = _legacy_config({"fpm_interpolation": {"fpm_parquet_path": "/tmp/fpm.parquet"}}) + legacy.update(forward_model="fpm", fpm_parquet_path=legacy_path) + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(legacy, "prefill") + assert _reload(canonical).estimator_config["fpm_interpolation"]["fpm_parquet_path"] == "/tmp/fpm.parquet" + + +@pytest.mark.parametrize("legacy_path", ["/tmp/different.parquet", "/tmp/subdir/../fpm.parquet"]) +def test_different_parquet_paths_conflict_without_filesystem_resolution(legacy_path): + legacy = _legacy_config({"fpm_interpolation": {"fpm_parquet_path": "/tmp/fpm.parquet"}}) + legacy.update(forward_model="fpm", fpm_parquet_path=legacy_path) + with pytest.raises(ValueError, match=r"estimator_config\.fpm_interpolation\.fpm_parquet_path"): + ForwardPassPerfModelConfig.from_legacy_engine_config(legacy, "prefill") + + +@pytest.mark.parametrize( + ("saved", "options", "reason"), + [ + ({"correction": {"enabld": False}}, {}, "unknown field"), + ({"fpm_regression": {"sampling": {"bins_per_axis": [4]}}}, {}, "length"), + ({"correction": None}, {}, "invalid type"), + ({}, {"max_observatons": 128}, "unknown field"), + ({}, {"min_observations": 0}, "min_observations must be >= 1"), + ({}, {"max_observations": False}, "invalid type"), + ({}, {"bucket_count": 3}, "bucket_count must be a perfect square"), + ({"fpm_regression": {"sampling": {"bins_per_axis": [0, 4]}}}, {}, "bins_per_axis must be positive"), + ], +) +def test_migration_rejects_unknown_and_invalid_controls(saved, options, reason): + with pytest.raises(ValueError, match=reason): + ForwardPassPerfModelConfig.from_legacy_engine_config(_legacy_config(saved), "prefill", options) + + +@pytest.mark.parametrize("options", [{"text_only": False}, {"unrecorded_quant_modes": []}]) +def test_explicit_skipped_defaults_migrate_without_becoming_null(options): + config = ForwardPassPerfModelConfig.from_legacy_engine_config( + _legacy_config({"fpm_interpolation": options}), "prefill" + ) + normalized = RustForwardPassPerfModel.normalize_config(_reload(config)) + assert normalized["estimator_config"]["fpm_interpolation"]["method"] == "sol" diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_phase_readiness.py b/python/aisimulate/tests/unit/sdk/test_fpm_phase_readiness.py new file mode 100644 index 000000000..63cfd2f9a --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_fpm_phase_readiness.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct FPM construction checks each required phase before selecting a root.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +import aisimulate_core +from aisimulate_core.sdk import RustForwardPassPerfModel +from aisimulate_core.sdk.errors import PerfDataNotAvailableError +from aisimulate_core.sdk.fpm_profile import load_fpm_profile + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def phase_roots(tmp_path, monkeypatch): + """Use real metadata, Parquet loading, compilation and native root selection.""" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + architecture = "UnregisteredDecoderForCausalLM" + (checkpoint / "config.json").write_text(json.dumps({"architectures": [architecture]})) + profile = { + "schema_version": 1, + "model": str(checkpoint), + "model_revision": "phase-readiness-fixture-v1", + "architecture": architecture, + "context_length": 4096, + "num_experts": 64, + "provenance": "Synthetic SDK readiness fixture, not silicon qualification.", + "deployments": [ + { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 1, + "dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "gemm_quant_mode": "nvfp4", + "moe_quant_mode": "nvfp4", + "fmha_quant_mode": "fp8", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 100, + "activations_bytes": 20, + "runtime_overhead_bytes": 30, + "comm_overhead_bytes": 50, + "kv_bytes_per_token": 10, + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 256, + "provenance": "Synthetic per-rank resource declarations.", + }, + } + ], + } + config = { + "model": profile["model"], + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "worker_type": "aggregated", + "tp": 1, + "attention_dp": 1, + "moe_tp_size": 1, + "moe_ep_size": 1, + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "fpm_profile": profile, + "estimator_config": {"fpm_interpolation": {"method": "direct"}}, + } + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + + def make_root(name, *, prefill="genuine", decode="genuine"): + root = tmp_path / name + root.mkdir() + packaged = Path(aisimulate_core.__file__).parent / "systems/h200_sxm.yaml" + (root / packaged.name).write_bytes(packaged.read_bytes()) + identity = load_fpm_profile(profile).deployments[0].model_dump(mode="json", exclude={"resources"}) + rows = [ + { + **identity, + "model_path": profile["model"], + "cell_id": f"synthetic-{phase}", + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": 1, + "total_prefill_tokens": 1 if phase == "prefill" else 0, + "total_kv_read_tokens": 0 if phase == "prefill" else 1, + "latency_ms": 2.0 if phase == "prefill" else 3.0, + "kv_seed_regime": "real_kv" if kind == "genuine" else "fake_fallback", + } + for phase, kind in (("prefill", prefill), ("decode", decode)) + if kind != "absent" + ] + path = root / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + path.parent.mkdir(parents=True) + pq.write_table(pa.Table.from_pylist(rows), path) + path.with_suffix(".metadata.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_forward_perf", + "schema_version": 6, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + ) + ) + return str(root) + + return config, make_root + + +@pytest.mark.parametrize("missing_phase", ["prefill", "decode"]) +@pytest.mark.parametrize("rows", ["absent", "fake"]) +def test_direct_constructor_rejects_missing_genuine_phase(phase_roots, missing_phase, rows): + config, make_root = phase_roots + root = make_root("incomplete", **{missing_phase: rows}) + config["systems_paths"] = [root] + with pytest.raises(PerfDataNotAvailableError, match=f"direct FPM {missing_phase}") as exc: + RustForwardPassPerfModel.best_available(config) + message = str(exc.value) + assert "no genuine measurements" in message + assert f"Collect genuine {missing_phase} FPM rows" in message + assert root in message + + +@pytest.mark.parametrize("missing_phase", ["prefill", "decode"]) +@pytest.mark.parametrize("rows", ["absent", "fake"]) +@pytest.mark.parametrize("fallback", ["deny", "allow"]) +def test_direct_constructor_selects_later_root_with_both_required_phases(phase_roots, missing_phase, rows, fallback): + config, make_root = phase_roots + first = make_root("incomplete", **{missing_phase: rows}) + second = make_root("usable") + config.update(systems_paths=[first, second], fallback_policy=fallback) + model = RustForwardPassPerfModel.best_available(config) + before = model.diagnostics() + assert before["readiness"] == "ready" + provenance = before["provenance"] + assert provenance["selected_systems_root"] == second + assert provenance["config"]["systems_paths"] == [second] + assert provenance["selected_estimation_mode"] == "fpm_interpolation" + assert provenance["config"]["estimator_config"]["fpm_interpolation"]["method"] == "direct" + # The selected root's exact, hand-declared synthetic timings remain queryable. + assert ( + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_prefill_requests": 1, "sum_prefill_tokens": 1}} + ) + == 2.0 + ) + assert ( + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 1}} + ) + == 3.0 + ) + with pytest.raises(PerfDataNotAvailableError, match="direct"): + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 64}} + ) + assert model.diagnostics() == before diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py new file mode 100644 index 000000000..e50c457c1 --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py @@ -0,0 +1,1233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FPM profiles keep identity, memory admission and compilation graph-independent.""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel, engine, fpm_profile, memory, utils +from aisimulate_core.sdk.errors import PerfDataNotAvailableError +from aisimulate_core.sdk.fpm_profile import FpmModelProfile, load_fpm_profile + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def checkpoint_configs(monkeypatch): + fpm_profile._load_pinned_checkpoint_config.cache_clear() + configs = { + "test/unknown-decoder": {"architectures": ["UnregisteredDecoderForCausalLM"]}, + "nvidia/GLM-5.2-NVFP4": {"architectures": ["GlmMoeDsaForCausalLM"]}, + "MiniMaxAI/MiniMax-M2.7": {"architectures": ["MiniMaxM2ForCausalLM"]}, + } + + def download(model, filename, *, revision): + assert filename == "config.json" + assert revision == "profile-test-fixture-v1" + return copy.deepcopy(configs[model]) + + monkeypatch.setattr(utils, "_download_hf_json", download) + yield configs + fpm_profile._load_pinned_checkpoint_config.cache_clear() + + +@pytest.fixture +def profile_dict(checkpoint_configs): + return { + "schema_version": 1, + "model": "test/unknown-decoder", + "model_revision": "profile-test-fixture-v1", + "architecture": "UnregisteredDecoderForCausalLM", + "context_length": 4096, + "num_experts": 64, + "provenance": "Synthetic metadata for SDK behavior tests; not silicon qualification.", + "deployments": [ + { + "system": "test_gpu", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 2, + "dp": 1, + "moe_tp": 2, + "moe_ep": 1, + "gemm_quant_mode": "nvfp4", + "moe_quant_mode": "nvfp4", + "fmha_quant_mode": "fp8", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 100, + "activations_bytes": 20, + "runtime_overhead_bytes": 30, + "comm_overhead_bytes": 50, + "kv_bytes_per_token": 10, + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 256, + "provenance": "Declared per-rank arithmetic fixture; excludes CUDA graphs.", + }, + } + ], + } + + +def _fail_graph(*_args, **_kwargs): + raise AssertionError("analytical model construction or timing database was accessed") + + +def _request(profile, method="auto", **overrides): + profile = json.loads(profile) if isinstance(profile, str) else profile + deployment = profile["deployments"][0] if profile else {} + config = { + "model": profile["model"] if profile else "test/unknown-decoder", + "system": deployment.get("system", "test_gpu"), + "backend": "vllm", + "backend_version": deployment.get("backend_version", "0.25.1"), + "worker_type": "aggregated", + "tp": deployment.get("tp", 2), + "attention_dp": deployment.get("dp", 1), + "moe_tp_size": deployment.get("moe_tp", 2), + "moe_ep_size": deployment.get("moe_ep", 1), + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "fpm_profile": profile, + "estimator_config": {"fpm_interpolation": {"method": method}}, + } + config.update(overrides) + return config + + +def _normalize(config): + return RustForwardPassPerfModel.normalize_config(config) + + +@pytest.mark.parametrize( + "entry_point", [RustForwardPassPerfModel.normalize_config, RustForwardPassPerfModel.best_available] +) +@pytest.mark.parametrize("method", ["auto", "direct"]) +@pytest.mark.parametrize("architecture", ["LlamaForCausalLm", "Qwen2ForCausalLM"]) +def test_profile_architecture_mismatch_fails_before_selection( + profile_dict, checkpoint_configs, monkeypatch, entry_point, method, architecture +): + checkpoint_configs[profile_dict["model"]] = {"architectures": ["LlamaForCausalLM"]} + profile_dict["architecture"] = architecture + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + with pytest.raises(ValueError, match="FPM profile architecture mismatch") as error: + entry_point(_request(profile_dict, method, estimation_mode="auto", fallback_policy="allow")) + assert "LlamaForCausalLM" in str(error.value) + assert architecture in str(error.value) + assert profile_dict["model_revision"] in str(error.value) + + +@pytest.mark.parametrize( + "entry_point", [RustForwardPassPerfModel.normalize_config, RustForwardPassPerfModel.best_available] +) +@pytest.mark.parametrize( + "metadata", + [ + None, + [], + {}, + {"model_type": "llama"}, + {"architectures": []}, + {"architectures": "LlamaForCausalLM"}, + {"architectures": [None]}, + {"architectures": [" "]}, + ], +) +def test_profile_missing_architecture_is_not_unregistered( + profile_dict, checkpoint_configs, monkeypatch, entry_point, metadata +): + checkpoint_configs[profile_dict["model"]] = metadata + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + with pytest.raises(ValueError, match="checkpoint config.json must"): + entry_point(_request(profile_dict, estimation_mode="auto", fallback_policy="allow")) + + +@pytest.mark.parametrize( + "entry_point", [RustForwardPassPerfModel.normalize_config, RustForwardPassPerfModel.best_available] +) +def test_profile_unresolvable_checkpoint_is_not_unregistered(profile_dict, monkeypatch, entry_point): + def unavailable(*_args, **_kwargs): + raise utils.HuggingFaceDownloadError("pinned revision is unavailable") + + monkeypatch.setattr(utils, "_download_hf_json", unavailable) + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + with pytest.raises(ValueError, match="cannot verify FPM profile architecture") as error: + entry_point(_request(profile_dict, estimation_mode="auto", fallback_policy="allow")) + assert profile_dict["model_revision"] in str(error.value) + assert "pinned revision is unavailable" in str(error.value) + + +@pytest.mark.parametrize( + "model,architecture,method", + [ + ("test/unknown-decoder", "UnregisteredDecoderForCausalLM", "direct"), + ("nvidia/GLM-5.2-NVFP4", "GlmMoeDsaForCausalLM", "sol"), + ], +) +def test_profile_architecture_metadata_uses_pinned_revision(profile_dict, monkeypatch, model, architecture, method): + calls = [] + profile_dict.update(model=model, architecture=architecture, model_revision="a" * 40) + + def download(model, filename, *, revision): + calls.append((model, filename, revision)) + return {"architectures": [architecture]} + + monkeypatch.setattr(utils, "_download_hf_json", download) + monkeypatch.setattr(utils, "_load_pre_downloaded_hf_config", _fail_graph) + monkeypatch.setattr(utils, "_parse_hf_config_json", _fail_graph) + resolved = RustForwardPassPerfModel.normalize_config(_request(profile_dict)) + assert RustForwardPassPerfModel.normalize_config(_request(profile_dict)) == resolved + assert calls == [(profile_dict["model"], "config.json", "a" * 40)] + assert resolved["estimator_config"]["fpm_interpolation"]["method"] == method + + +def test_profile_metadata_cache_separates_checkpoint_revisions(profile_dict, monkeypatch): + calls = [] + architectures = {"a" * 40: "UnregisteredDecoderForCausalLM", "b" * 40: "LlamaForCausalLM"} + + def download(model, filename, *, revision): + calls.append((model, filename, revision)) + return {"architectures": [architectures[revision]]} + + monkeypatch.setattr(utils, "_download_hf_json", download) + for revision, method in (("a" * 40, "direct"), ("b" * 40, "sol"), ("a" * 40, "direct")): + profile_dict.update(model_revision=revision, architecture=architectures[revision]) + assert _normalize(_request(profile_dict))["estimator_config"]["fpm_interpolation"]["method"] == method + assert calls == [(profile_dict["model"], "config.json", revision) for revision in architectures] + + +def test_profile_architecture_matches_collector_text_config_view(profile_dict, checkpoint_configs): + checkpoint_configs[profile_dict["model"]] = { + "architectures": ["UnknownForConditionalGeneration"], + "text_config": {"architectures": [profile_dict["architecture"]]}, + } + assert _normalize(_request(profile_dict))["estimator_config"]["fpm_interpolation"]["method"] == "direct" + + +@pytest.mark.parametrize("metadata", [None, "not JSON"]) +def test_local_checkpoint_missing_or_invalid_config_fails_without_download( + profile_dict, monkeypatch, tmp_path, metadata +): + profile_dict["model"] = str(tmp_path) + if metadata is not None: + (tmp_path / "config.json").write_text(metadata) + monkeypatch.setattr(utils, "_download_hf_json", _fail_graph) + with pytest.raises(ValueError, match="cannot verify FPM profile architecture"): + RustForwardPassPerfModel.best_available(_request(profile_dict, fallback_policy="allow")) + + +def test_local_checkpoint_architecture_is_revalidated(profile_dict, monkeypatch, tmp_path): + profile_dict["model"] = str(tmp_path) + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"architectures": [profile_dict["architecture"]]})) + monkeypatch.setattr(utils, "_download_hf_json", _fail_graph) + assert _normalize(_request(profile_dict))["estimator_config"]["fpm_interpolation"]["method"] == "direct" + config_path.write_text(json.dumps({"architectures": ["LlamaForCausalLM"]})) + with pytest.raises(ValueError, match="FPM profile architecture mismatch"): + RustForwardPassPerfModel.best_available(_request(profile_dict, fallback_policy="allow")) + + +@pytest.fixture +def direct_compile(monkeypatch): + monkeypatch.setattr(engine, "get_model", _fail_graph) + monkeypatch.setattr(engine, "build_model_config", _fail_graph) + monkeypatch.setattr(engine, "_maybe_load_database", _fail_graph) + monkeypatch.setattr(engine, "_literal_backend_version", lambda _s, _b, version, *_args: version) + monkeypatch.setattr(engine.aisimulate_core, "engine_spec_bincode_from_json", lambda value: value.encode()) + + def compile_profile(profile, **kwargs): + config = _normalize(_request(profile, kwargs.pop("fpm_interpolation", "auto"), **kwargs)) + return json.loads( + engine.compile_engine( + config["model"], + config["system"], + config["backend"], + config["backend_version"], + tp_size=config["tp"], + moe_tp_size=config["moe_tp_size"], + moe_ep_size=config["moe_ep_size"], + forward_model="fpm", + fpm_profile=config["fpm_profile"], + fpm_interpolation=config["estimator_config"]["fpm_interpolation"]["method"], + **kwargs, + ) + ) + + return compile_profile + + +def test_unknown_profile_auto_compiles_without_operations(profile_dict, direct_compile): + spec = direct_compile(profile_dict) + profile = FpmModelProfile.model_validate(profile_dict) + deployment = profile.deployments[0] + assert json.loads(spec["engine"]["extra"]["fpm_profile"]) == profile.model_dump(mode="json") + assert json.loads(spec["engine"]["extra"]["estimator_config"])["fpm_interpolation"]["method"] == "direct" + assert spec["engine"]["activation_dtype"] == "fp8" + for key, phase in (("context_ops", "prefill"), ("generation_ops", "decode")): + assert len(spec[key]) == 1 + op = spec[key][0]["FpmForward"] + assert op["phase"] == phase + assert op["model_path"] == profile.model + assert op["match_identity"] == deployment.match_identity() + assert op["interpolation"] == "direct" + assert op["sol_ops"] == [] + assert op["weight_bytes"] == 100 + + +@pytest.mark.parametrize("entry_point", [_normalize, RustForwardPassPerfModel.best_available]) +@pytest.mark.parametrize("mode", ["fmha", "comm"]) +def test_profile_precision_cannot_be_selected_as_unrecorded(profile_dict, monkeypatch, entry_point, mode): + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + request = _request(profile_dict, estimation_mode="auto", fallback_policy="allow") + request["estimator_config"]["fpm_interpolation"]["unrecorded_quant_modes"] = [mode] + with pytest.raises(ValueError, match="fpm_profile does not support"): + entry_point(request) + + +@pytest.mark.parametrize("dcp", [1, 2]) +@pytest.mark.parametrize("entry_point", [_normalize, RustForwardPassPerfModel.best_available]) +def test_profile_rejects_unrepresented_recorded_dcp(profile_dict, monkeypatch, dcp, entry_point): + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + with pytest.raises(ValueError, match="does not describe recorded DCP identity"): + entry_point(_request(profile_dict, dcp=dcp, estimation_mode="auto", fallback_policy="allow")) + + +def test_direct_compilation_preserves_external_data_controls(profile_dict, direct_compile): + options = {"text_only": False, "fpm_parquet_path": "/external/profile.parquet"} + spec = json.loads( + engine.compile_engine( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + tp_size=2, + moe_tp_size=2, + moe_ep_size=1, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + fpm_options=options, + ) + ) + controls = json.loads(spec["engine"]["extra"]["estimator_config"])["fpm_interpolation"] + assert controls["method"] == "direct" + assert controls.get("text_only", False) is False + assert controls["fpm_parquet_path"] == spec["engine"]["fpm_parquet_path"] == options["fpm_parquet_path"] + assert spec["context_ops"][0]["FpmForward"]["sol_ops"] == [] + + +@pytest.mark.parametrize( + "options,overrides,error", + [ + ({"unrecorded_quant_modes": ["fmha"]}, {}, "fpm_profile does not support"), + ({"unrecorded_quant_modes": ["comm"]}, {}, "fpm_profile does not support"), + ({}, {"dcp_size": 2}, "does not describe recorded DCP identity"), + ({"method": "sol"}, {}, "conflicting fpm_interpolation and fpm_options.method"), + ], +) +def test_direct_compilation_validates_options_before_short_circuit( + profile_dict, direct_compile, options, overrides, error +): + with pytest.raises(ValueError, match=error): + engine.compile_engine( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + tp_size=2, + moe_tp_size=2, + moe_ep_size=1, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + fpm_options=options, + **overrides, + ) + + +@pytest.mark.parametrize("control", ["moe_kernel_source", "decode_workload_distribution", "prefill_graph_profile"]) +def test_profile_compilation_rejects_op_level_controls_before_graph_construction(profile_dict, direct_compile, control): + with pytest.raises(ValueError, match="FPM profiles do not support"): + engine.compile_engine( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + tp_size=2, + moe_tp_size=2, + moe_ep_size=1, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + **{control: "explicit-selection"}, + ) + + +def test_direct_json_transport_preserves_all_identity_fields(profile_dict, direct_compile): + profile_dict["deployments"][0].update(moe_backend="pinned_moe", attention_backend="pinned_attention") + spec = direct_compile(json.dumps(profile_dict), fpm_interpolation="direct") + assert spec["context_ops"][0]["FpmForward"]["match_identity"] == [ + "nvfp4", + "nvfp4", + "fp8", + "half", + "fp8", + "2", + "1", + "1", + "2", + "1", + "1", + "pinned_moe", + "pinned_attention", + "False", + "False", + ] + + +@pytest.mark.parametrize("override", [{"fmha_quant_mode": "bfloat16"}, {"gemm_quant_mode": "fp8"}]) +def test_precision_conflicts_fail_before_build(profile_dict, direct_compile, override): + with pytest.raises(ValueError, match="FPM profile identity conflict"): + direct_compile(profile_dict, **override) + + +@pytest.mark.parametrize("controls", [{"enable_eplb": True}, {"wideep_num_slots": 256}, {"moe_backend": "megamoe"}]) +def test_profile_engine_controls_fail_before_model_lookup(profile_dict, direct_compile, profile_memory, controls): + with pytest.raises(ValueError): + _normalize(_request(profile_dict, "direct", estimation_mode="auto", **controls)) + with pytest.raises(ValueError, match="FPM profiles do not support"): + engine.compile_engine( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + tp_size=2, + moe_tp_size=2, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + **controls, + ) + with pytest.raises(ValueError, match="FPM profiles do not support"): + profile_memory(**controls) + + +@pytest.mark.parametrize("cp_size", [2, 0, True, 1.0, "1", None]) +@pytest.mark.parametrize("with_profile", [False, True]) +def test_sdk_entry_points_reject_unsupported_cp_before_construction(profile_dict, monkeypatch, cp_size, with_profile): + monkeypatch.setattr(engine, "build_model_config", _fail_graph) + monkeypatch.setattr(memory.KVCacheEstimator, "from_request", _fail_graph) + monkeypatch.setattr(memory.NaiveKVCacheEstimator, "from_model_path", _fail_graph) + kwargs = {"cp_size": cp_size, "fpm_profile": profile_dict if with_profile else None} + with pytest.raises(ValueError, match="cp_size must be the integer 1"): + engine.compile_engine("test/unknown-decoder", "test_gpu", "vllm", "0.25.1", forward_model="fpm", **kwargs) + with pytest.raises(ValueError, match="cp_size must be the integer 1"): + memory.estimate_kv_cache( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.8, + allow_naive_fallback=True, + **kwargs, + ) + + +@pytest.mark.parametrize("kwargs", [{"nextn": 1}, {"cp_size": 2}, {"database_mode": "SOL"}]) +def test_direct_rejects_unsupported_execution(profile_dict, direct_compile, kwargs): + with pytest.raises(ValueError): + direct_compile(profile_dict, **kwargs) + + +@pytest.mark.parametrize("dcp", [1, 2]) +def test_profile_rejects_unrepresented_dcp_identity(profile_dict, dcp): + with pytest.raises(ValueError, match="does not describe recorded DCP identity"): + _normalize(_request(profile_dict, "direct", dcp=dcp)) + with pytest.raises(ValueError, match="does not describe recorded DCP identity"): + engine.compile_engine( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + tp_size=2, + moe_tp_size=2, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + dcp_size=dcp, + ) + + +def test_direct_requires_explicit_resources(direct_compile): + with pytest.raises(ValueError, match="requires an fpm_profile"): + direct_compile(None, fpm_interpolation="direct") + + +def test_explicit_sol_does_not_fall_back_on_unknown_architecture(profile_dict, direct_compile): + with pytest.raises(ValueError, match="requires a registered analytical model class"): + direct_compile(profile_dict, fpm_interpolation="sol") + + +def test_registered_auto_uses_registry_presence_without_constructing(profile_dict, monkeypatch): + from aisimulate_core.sdk.models.base import _MODEL_REGISTRY + + profile = load_fpm_profile(profile_dict) + monkeypatch.setitem(_MODEL_REGISTRY, profile.architecture, object()) + assert _normalize(_request(profile_dict))["estimator_config"]["fpm_interpolation"]["method"] == "sol" + assert _normalize(_request(profile_dict, "direct"))["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert _normalize(_request(None))["estimator_config"]["fpm_interpolation"]["method"] == "sol" + + +def test_registered_build_failure_is_not_direct_fallback(profile_dict, monkeypatch): + from aisimulate_core.sdk.models.base import _MODEL_REGISTRY + + monkeypatch.setitem(_MODEL_REGISTRY, profile_dict["architecture"], object()) + monkeypatch.setattr(engine, "_literal_backend_version", lambda _s, _b, version, *_args: version) + observed = [] + + def broken_model(_path, config, _backend): + observed.append(config.fmha_quant_mode.name) + raise RuntimeError("registered model has an unrelated implementation error") + + monkeypatch.setattr(engine, "get_model", broken_model) + monkeypatch.setattr(engine, "_direct_fpm_spec_json", _fail_graph) + with pytest.raises(ValueError, match="unrelated implementation error"): + RustForwardPassPerfModel.best_available(_request(profile_dict)) + assert observed == ["fp8"] + + +@pytest.mark.parametrize( + "mutation", + [ + lambda p: p.update(unrecognized_field=1), + lambda p: p.update(model_revision="main"), + lambda p: p.update(model_revision="unknown"), + lambda p: p.update(context_length=True), + lambda p: p.update(schema_version=True), + lambda p: p.update(schema_version=1.0), + lambda p: p.update(provenance=" "), + lambda p: p.update(deployments=[]), + lambda p: p["deployments"].append(copy.deepcopy(p["deployments"][0])), + lambda p: p["deployments"][0].update(tp="2"), + lambda p: p["deployments"][0].update(pp=2), + lambda p: p["deployments"][0].update(cp=2), + lambda p: p["deployments"][0].update(cp=True), + lambda p: p["deployments"][0].update(pp=1.0), + lambda p: p["deployments"][0].update(dp=2), + lambda p: p["deployments"][0].update(backend_version="current"), + lambda p: p["deployments"][0].update(fmha_quant_mode="bf16"), + lambda p: p["deployments"][0].update(enable_wideep=True), + lambda p: p["deployments"][0].update(enable_eplb=0), + lambda p: p["deployments"][0]["resources"].update(cache_layout="sliding_window"), + lambda p: p["deployments"][0]["resources"].update(kv_bytes_per_token=0), + lambda p: p["deployments"][0]["resources"].update(weights_bytes=-1), + lambda p: p["deployments"][0]["resources"].update(weights_bytes=1.5), + lambda p: p["deployments"][0]["resources"].update(weights_bytes=2**53), + ], +) +def test_profiles_reject_ambiguous_or_unsupported_metadata(profile_dict, mutation): + mutation(profile_dict) + with pytest.raises(ValidationError): + load_fpm_profile(profile_dict) + + +def test_topologies_have_separate_resources_and_expert_admission(profile_dict): + tep = profile_dict["deployments"][0] + tep.update(tp=8, dp=1, moe_tp=1, moe_ep=8) + dep = copy.deepcopy(tep) + dep.update(tp=1, dp=8) + dep["resources"]["weights_bytes"] = 180 + profile_dict["deployments"].append(dep) + profile = load_fpm_profile(profile_dict) + selected = profile.select( + model=profile.model, + system="test_gpu", + backend="vllm", + backend_version="0.25.1", + tp_size=1, + attention_dp_size=8, + moe_tp_size=1, + moe_ep_size=8, + ) + assert selected.resources.weights_bytes == 180 + assert profile.deployments[0].resources.weights_bytes == 100 + with pytest.raises(ValueError, match="no matching FPM deployment"): + profile.select(model=profile.model, system="test_gpu", backend="vllm", backend_version="0.25.1") + with pytest.raises(ValueError, match="model identity mismatch"): + profile.select(model="other/model", system="test_gpu", backend="vllm", backend_version="0.25.1") + profile_dict["num_experts"] = 63 + with pytest.raises(ValidationError, match="divisible by moe_ep"): + load_fpm_profile(profile_dict) + + +def test_dense_tp_does_not_require_expert_tensor_parallelism(profile_dict): + profile_dict["num_experts"] = 0 + profile_dict["deployments"][0]["moe_tp"] = 1 + assert load_fpm_profile(profile_dict).deployments[0].parallel_tuple == (2, 1, 1, 1, 1, 1) + + +@pytest.mark.parametrize("kwargs", [{"moe_tp_size": 0}, {"moe_ep_size": False}, {"cp_size": True}]) +def test_select_rejects_zero_or_boolean_dimensions(profile_dict, kwargs): + profile = load_fpm_profile(profile_dict) + with pytest.raises(ValueError, match="positive integers"): + profile.select(model=profile.model, system="test_gpu", backend="vllm", backend_version="0.25.1", **kwargs) + + +@pytest.fixture +def profile_memory(profile_dict, monkeypatch): + monkeypatch.setattr(memory, "get_model", _fail_graph) + monkeypatch.setattr(memory, "build_model_config", _fail_graph) + monkeypatch.setattr(memory.perf_database, "get_database", _fail_graph) + monkeypatch.setattr(memory.perf_database, "resolve_query_version", _fail_graph) + monkeypatch.setattr(memory.perf_database, "load_system_spec", lambda *_a, **_k: {"gpu": {"mem_capacity": 1000}}) + monkeypatch.setattr(engine, "_literal_backend_version", lambda _s, _b, version, *_args: version) + + def estimate(**overrides): + args = dict( + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.8, + tp_size=2, + moe_tp_size=2, + fpm_profile=profile_dict, + ) + args.update(overrides) + return memory.estimate_kv_cache("test/unknown-decoder", "test_gpu", "vllm", "0.25.1", **args) + + return estimate + + +def test_profile_memory_uses_declared_bounds_without_graph_or_timing_data(profile_memory): + result = profile_memory(cuda_graph_reserved_bytes=20, tolerance_fraction=0.1) + assert result["source"] == "profile" + assert result["total_gpu_capacity_bytes"] == 1000 + assert result["total_kv_size_bytes"] == 580 # 1000 * .8 - 200 - 20 + assert result["kv_size_per_token_bytes"] == 10 + assert result["total_kv_size_tokens"] == 58 + assert result["tolerance_adjusted"]["total_kv_size_bytes"] == 522 + assert result["tolerance_adjusted"]["total_kv_size_tokens"] == 52 + assert result["memory_breakdown"]["cuda_graph_reserved_bytes"] == 20 + assert "Declared per-rank" in result["resource_provenance"] + json.dumps(result) # The result is replayable; it contains no callables. + + +def test_profile_memory_capacity_override_does_not_read_hardware(profile_memory, monkeypatch): + monkeypatch.setattr(memory.perf_database, "load_system_spec", _fail_graph) + result = profile_memory(gpu_memory_capacity_bytes_override=2000) + assert result["total_kv_size_bytes"] == 1400 + + +def test_profile_block_budget_preserves_resource_provenance(profile_dict, profile_memory): + diagnostics = {} + blocks = memory.estimate_num_gpu_blocks( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + scheduler_block_size=4, + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.8, + tp_size=2, + moe_tp_size=2, + fpm_profile=profile_dict, + diagnostics=diagnostics, + ) + assert blocks == 15 # (1000 * .8 - 200) / 10 / 4 + assert diagnostics["source"] == "profile" + assert diagnostics["resource_provenance"] == profile_dict["deployments"][0]["resources"]["provenance"] + with pytest.raises(ValueError, match="positive integer"): + memory.estimate_num_gpu_blocks( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + scheduler_block_size=4, + max_num_tokens=True, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.8, + tp_size=2, + moe_tp_size=2, + fpm_profile=profile_dict, + ) + + +@pytest.mark.parametrize( + ("kwargs", "reason"), + [ + ({"max_num_tokens": 8193}, "resource envelope exceeded"), + ({"max_batch_size": 257}, "resource envelope exceeded"), + ({"max_num_tokens": True}, "positive integer"), + ({"fmha_quant_mode": "bfloat16"}, "identity conflict"), + ({"attention_backend": "different"}, "identity conflict"), + ({"cp_size": 2}, "cp_size must be the integer 1"), + ({"nextn": 1}, "nextn must be 0"), + ({"gpu_memory_capacity_bytes_override": 200}, "no KV budget"), + ({"memory_fraction_kind": "of_free"}, "incompatible memory fraction"), + ], +) +def test_profile_memory_fails_closed(profile_memory, kwargs, reason): + with pytest.raises(ValueError, match=reason): + profile_memory(**kwargs) + + +@pytest.fixture +def native_profile_config(profile_dict, monkeypatch): + # Use a bundled system/version for native engine construction. Resource and + # precision variants below are transport fixtures, not measured timing cells. + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + profile_dict.update(model="MiniMaxAI/MiniMax-M2.7", architecture="MiniMaxM2ForCausalLM") + profile_dict["deployments"][0].update(system="h200_sxm") + monkeypatch.setattr(engine, "get_model", _fail_graph) + specs = [] + encode = engine.aisimulate_core.engine_spec_bincode_from_json + + def record_spec(value): + specs.append(json.loads(value)) + return encode(value) + + monkeypatch.setattr(engine.aisimulate_core, "engine_spec_bincode_from_json", record_spec) + + def compile_config(): + engine.compile_engine( + profile_dict["model"], + "h200_sxm", + "vllm", + "0.25.1", + tp_size=2, + moe_tp_size=2, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + ) + return specs[-1]["engine"] + + return compile_config, specs + + +@pytest.mark.parametrize( + ("gemm_mode", "moe_mode"), + [ + ("fp8_block", "nvfp4"), + ("int8_wo", "nvfp4"), + ("int4_wo", "nvfp4"), + ("sq", "nvfp4"), + ("fp8_ootb", "nvfp4"), + ("nvfp4", "int4_wo"), + ], +) +def test_legacy_migration_preserves_profile_precision(profile_dict, native_profile_config, gemm_mode, moe_mode): + profile_dict["deployments"][0].update(gemm_quant_mode=gemm_mode, moe_quant_mode=moe_mode) + compile_config, _ = native_profile_config + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(compile_config(), "aggregated") + resolved = _normalize(canonical.to_dict()) + assert resolved["gemm_quant_mode"] == gemm_mode + assert resolved["moe_quant_mode"] == moe_mode + assert resolved["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert resolved["fpm_profile"] == load_fpm_profile(profile_dict).model_dump(mode="json") + assert "extra" not in resolved + + +@pytest.mark.parametrize("field", ["weight_dtype", "moe_dtype", "activation_dtype", "kv_cache_dtype"]) +def test_native_profile_rejects_different_wire_dtype(native_profile_config, field): + compile_config, _ = native_profile_config + config = compile_config() + config[field] = "bfloat16" + with pytest.raises(ValueError, match="FPM profile identity conflict"): + ForwardPassPerfModelConfig.from_legacy_engine_config(config, "prefill") + + +@pytest.mark.parametrize("mode", ["auto", "fpm_interpolation", "fpm_regression"]) +@pytest.mark.parametrize( + "case", ["malformed_profile", "invalid_method", "missing_profile", "precision", "missing_version", "version_alias"] +) +def test_invalid_profile_controls_fail_before_any_fallback(profile_dict, monkeypatch, mode, case): + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + config = _request(profile_dict, "direct", estimation_mode=mode, fallback_policy="allow") + if case == "malformed_profile": + config["fpm_profile"] = {} + elif case == "invalid_method": + config["estimator_config"]["fpm_interpolation"]["method"] = "dierct" + elif case == "missing_profile": + config["fpm_profile"] = None + elif case == "precision": + config["gemm_quant_mode"] = "bfloat16" + elif case == "missing_version": + config["backend_version"] = None + else: + config["backend_version"] = "current" + with pytest.raises(ValueError): + RustForwardPassPerfModel.best_available(config) + + +def test_legacy_selector_migrates_to_nested_control(native_profile_config): + compile_config, _ = native_profile_config + config = compile_config() + config["extra"].pop("estimator_config") + config["extra"]["fpm_interpolation"] = "direct" + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(config, "prefill").to_dict() + assert canonical["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert "fpm_interpolation" not in canonical + assert isinstance(canonical["fpm_profile"], dict) + + +def test_saved_profile_compiler_controls_merge_with_caller_options(native_profile_config): + compile_config, _ = native_profile_config + saved = json.loads(json.dumps(compile_config())) + assert json.loads(saved["extra"]["estimator_config"]) == { + "fpm_interpolation": { + "method": "direct", + "text_only": False, + "unrecorded_quant_modes": [], + "fpm_parquet_path": None, + } + } + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config( + saved, + "prefill", + {"max_observations": 128, "min_faster_correction_factor": 0.9, "regression_attention_kv_weight": 7}, + ) + reloaded = ForwardPassPerfModelConfig(**json.loads(json.dumps(canonical.to_dict()))) + controls = RustForwardPassPerfModel.normalize_config(reloaded)["estimator_config"] + assert controls["fpm_interpolation"]["method"] == "direct" + assert controls["correction"]["sampling"]["max_observations"] == 128 + assert controls["fpm_regression"]["sampling"]["max_observations"] == 128 + assert controls["correction"]["factor_bounds"]["min"] == 0.9 + assert controls["features"]["attention_kv_weight"] == 7 + + +def _write_external_profile_pair(profile, directory): + """Generate routing/identity evidence; these 2/3 ms rows are not silicon data.""" + import pyarrow as pa + import pyarrow.parquet as pq + + deployment = load_fpm_profile(profile).deployments[0] + identity = deployment.model_dump(mode="json", exclude={"resources"}) + rows = [ + { + **identity, + "cell_id": f"synthetic-{phase}-{kv}", + "model_path": profile["model"], + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": 1, + "total_prefill_tokens": 1 if phase == "prefill" else 0, + "total_kv_read_tokens": kv, + "latency_ms": latency, + "kv_seed_regime": "real_kv", + } + for phase, kv, latency in [("prefill", 0, 2.0), ("decode", 1, 3.0), ("decode", 64, 3.0)] + ] + path = directory / "synthetic-profile.parquet" + pq.write_table(pa.Table.from_pylist(rows), path) + path.with_suffix(".metadata.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_forward_perf", + "schema_version": 6, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "system": deployment.system, + "backend": deployment.backend, + "backend_version": deployment.backend_version, + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + ) + ) + return path + + +@pytest.mark.parametrize( + "architecture,method", + [("UnregisteredDecoderForCausalLM", "auto"), ("LlamaForCausalLM", "direct")], +) +def test_verified_local_checkpoint_queries_without_analytical_model( + profile_dict, monkeypatch, tmp_path, architecture, method +): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + # A config-only checkpoint cannot supply weights or a model graph. Remote + # code declarations stay inert; direct timing needs just verified identity. + (checkpoint / "config.json").write_text( + json.dumps({"architectures": [architecture], "auto_map": {"AutoConfig": "untrusted.Configuration"}}) + ) + (checkpoint / "untrusted.py").write_text("raise AssertionError('remote model code was executed')\n") + profile_dict.update(model=str(checkpoint), architecture=architecture) + profile_dict["deployments"][0]["system"] = "h200_sxm" + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + monkeypatch.setattr(utils, "_download_hf_json", _fail_graph) + monkeypatch.setattr(utils, "_parse_hf_config_json", _fail_graph) + monkeypatch.setattr(engine, "get_model", _fail_graph) + monkeypatch.setattr(engine, "build_model_config", _fail_graph) + monkeypatch.setattr(engine, "_maybe_load_database", _fail_graph) + path = _write_external_profile_pair(profile_dict, tmp_path) + config = _request( + profile_dict, + estimator_config={"fpm_interpolation": {"method": method, "fpm_parquet_path": str(path)}}, + ) + resolved = RustForwardPassPerfModel.normalize_config(config) + assert resolved["estimator_config"]["fpm_interpolation"]["method"] == "direct" + forward = RustForwardPassPerfModel.best_available(config) + constructed = forward.diagnostics()["provenance"]["config"] + assert constructed["fpm_profile"] == resolved["fpm_profile"] + assert constructed["estimator_config"] == resolved["estimator_config"] + assert forward.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_prefill_requests": 1, "sum_prefill_tokens": 1}} + ) == pytest.approx(2.0) + assert forward.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 1}} + ) == pytest.approx(3.0) + + +@pytest.mark.parametrize( + ("model", "architecture", "system", "tp", "dp", "moe_tp", "moe_ep", "gemm", "fmha"), + [ + ("MiniMaxAI/MiniMax-M2.7", "MiniMaxM2ForCausalLM", "h200_sxm", 4, 1, 4, 1, "fp8_block", "bfloat16"), + ("nvidia/GLM-5.2-NVFP4", "GlmMoeDsaForCausalLM", "b200_sxm", 1, 8, 1, 8, "nvfp4", "fp8"), + ("nvidia/GLM-5.2-NVFP4", "GlmMoeDsaForCausalLM", "b200_sxm", 8, 1, 1, 8, "nvfp4", "fp8"), + ], +) +def test_external_fpm_cells_query_with_model_construction_disabled( + profile_dict, monkeypatch, tmp_path, model, architecture, system, tp, dp, moe_tp, moe_ep, gemm, fmha +): + # Synthetic timing and resource declarations test the real checkpoint + # identities through external storage with all graph constructors disabled. + # Accuracy against retained silicon measurements is separate validation. + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + profile_dict.update(model=model, architecture=architecture) + profile_dict["deployments"][0].update( + system=system, + tp=tp, + dp=dp, + moe_tp=moe_tp, + moe_ep=moe_ep, + gemm_quant_mode=gemm, + moe_quant_mode=gemm, + fmha_quant_mode=fmha, + ) + monkeypatch.setattr(engine, "get_model", _fail_graph) + monkeypatch.setattr(engine, "build_model_config", _fail_graph) + monkeypatch.setattr(engine, "_maybe_load_database", _fail_graph) + path = _write_external_profile_pair(profile_dict, tmp_path) + forward = RustForwardPassPerfModel.best_available( + _request( + profile_dict, + estimator_config={"fpm_interpolation": {"method": "direct", "fpm_parquet_path": str(path)}}, + ) + ) + resolved = forward.diagnostics()["provenance"]["config"] + assert resolved["estimation_mode"] == "fpm_interpolation" + assert resolved["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert resolved["fpm_profile"] == load_fpm_profile(profile_dict).model_dump(mode="json") + assert resolved["gemm_quant_mode"] == gemm + assert resolved["fmha_quant_mode"] == fmha + prefill_metrics = {"scheduled_requests": {"num_prefill_requests": 1, "sum_prefill_tokens": 1}} + decode_metrics = { + "scheduled_requests": { + "num_decode_requests": 1, + "sum_decode_kv_tokens": 1, + } + } + # Exact expectations are the hand-declared synthetic rows in the external pair. + assert forward.estimate_forward_pass_time_ms(prefill_metrics) == pytest.approx(2.0) + assert forward.estimate_forward_pass_time_ms(decode_metrics) == pytest.approx(3.0) + reloaded = RustForwardPassPerfModel.best_available(json.loads(json.dumps(resolved))) + assert reloaded.estimate_forward_pass_time_ms(prefill_metrics) == pytest.approx(2.0) + assert reloaded.diagnostics()["provenance"]["config"] == resolved + with pytest.raises(PerfDataNotAvailableError, match="direct"): + forward.estimate_forward_pass_time_ms( + { + "scheduled_requests": { + "num_prefill_requests": 1, + "sum_prefill_tokens": 1_000_000, + } + } + ) + assert forward.diagnostics()["provenance"]["config"] == resolved + + +def test_registered_glm_auto_retains_sol_with_explicit_fp8_fmha(profile_dict, monkeypatch, tmp_path): + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + profile_dict.update(model="nvidia/GLM-5.2-NVFP4", architecture="GlmMoeDsaForCausalLM") + profile_dict["deployments"][0].update(system="b200_sxm", tp=1, dp=8, moe_tp=1, moe_ep=8) + real_get_model = engine.get_model + built_models = [] + + def record_model(*args): + model = real_get_model(*args) + built_models.append(model) + return model + + monkeypatch.setattr(engine, "get_model", record_model) + monkeypatch.setattr(engine, "_direct_fpm_spec_json", _fail_graph) + path = _write_external_profile_pair(profile_dict, tmp_path) + forward = RustForwardPassPerfModel.best_available( + _request( + profile_dict, + estimator_config={"fpm_interpolation": {"method": "auto", "fpm_parquet_path": str(path)}}, + ) + ) + resolved = forward.diagnostics()["provenance"]["config"] + assert resolved["estimator_config"]["fpm_interpolation"]["method"] == "sol" + assert len(built_models) == 1 + assert built_models[0].config.fmha_quant_mode.name == "fp8" + assert built_models[0].context_ops[0]._sol_ops + assert ( + forward.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_prefill_requests": 1, "sum_prefill_tokens": 1}} + ) + > 0 + ) + + +def test_profile_topology_accepts_dense_tp_without_moe_partition(profile_dict, monkeypatch): + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + profile_dict["num_experts"] = 0 + profile_dict["deployments"][0]["moe_tp"] = 1 + resolved = _normalize(_request(profile_dict)) + assert resolved["tp"] == 2 + assert resolved["moe_tp_size"] == 1 + assert resolved["estimator_config"]["fpm_interpolation"]["method"] == "direct" + + +@pytest.mark.parametrize("registered", [False, True]) +def test_profile_auto_preserves_global_priority_with_deny(profile_dict, monkeypatch, registered): + from aisimulate_core.sdk.models.base import _MODEL_REGISTRY + + if registered: + monkeypatch.setitem(_MODEL_REGISTRY, profile_dict["architecture"], object()) + attempts = [] + + def unavailable(*_args, **kwargs): + attempts.append(kwargs["forward_model"]) + raise RuntimeError("fixture native timing unavailable") + + monkeypatch.setattr(engine, "compile_engine", unavailable) + forward = RustForwardPassPerfModel.best_available(_request(profile_dict, "direct", estimation_mode="auto")) + provenance = forward.diagnostics()["provenance"] + assert attempts == (["op_level", "fpm"] if registered else ["fpm"]) + assert len(provenance["selection_failures"]) == 2 + assert provenance["selected_estimation_mode"] == "fpm_regression" + assert provenance["config"]["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert ( + forward.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 1}} + ) + is None + ) + + +def test_canonical_profile_round_trip_can_size_memory_without_timing(profile_dict, profile_memory, monkeypatch): + monkeypatch.setattr(engine, "compile_engine", _fail_graph) + resolved = json.loads(json.dumps(_normalize(_request(profile_dict)))) + result = memory.estimate_kv_cache( + resolved["model"], + resolved["system"], + resolved["backend"], + resolved["backend_version"], + tp_size=resolved["tp"], + pp_size=resolved["pp"], + attention_dp_size=resolved["attention_dp"], + moe_tp_size=resolved["moe_tp_size"], + moe_ep_size=resolved["moe_ep_size"], + fpm_profile=resolved["fpm_profile"], + gemm_quant_mode=resolved["gemm_quant_mode"], + moe_quant_mode=resolved["moe_quant_mode"], + fmha_quant_mode=resolved["fmha_quant_mode"], + kvcache_quant_mode=resolved["kvcache_quant_mode"], + comm_quant_mode=resolved["comm_quant_mode"], + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.8, + ) + assert result["source"] == "profile" + assert result["total_kv_size_bytes"] == 600 # 1000 * 0.8 - 200 declared non-KV bytes. + assert result["total_kv_size_tokens"] == 60 # 600 bytes / 10 declared bytes per token. + + +@pytest.fixture +def measured_profile_roots(profile_dict, tmp_path, monkeypatch): + """Small synthetic cells isolate availability from phase/domain coverage.""" + import pyarrow as pa + import pyarrow.parquet as pq + + import aisimulate_core + + profile_dict["deployments"][0].update(system="h200_sxm", attention_backend="future_attention") + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + monkeypatch.setattr(engine, "get_model", _fail_graph) + monkeypatch.setattr(engine, "build_model_config", _fail_graph) + + def make_root(name, kind): + root = tmp_path / name + root.mkdir() + packaged = Path(aisimulate_core.__file__).parent / "systems/h200_sxm.yaml" + (root / packaged.name).write_bytes(packaged.read_bytes()) + identity = load_fpm_profile(profile_dict).deployments[0].model_dump(mode="json", exclude={"resources"}) + coordinates = [] + if kind in {"genuine", "mixed"}: + coordinates.extend([("prefill", 0, "real_kv", 2.0), ("decode", 1, "real_kv", 3.0)]) + if kind in {"fake", "mixed"}: + coordinates.append(("decode", 64, "fake_fallback", 99.0)) + rows = [ + { + **identity, + "model_path": profile_dict["model"], + "cell_id": f"synthetic-{phase}-{kv}", + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": 1, + "total_prefill_tokens": 1 if phase == "prefill" else 0, + "total_kv_read_tokens": kv, + "latency_ms": latency, + "kv_seed_regime": regime, + } + for phase, kv, regime, latency in coordinates + ] + path = root / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + path.parent.mkdir(parents=True) + pq.write_table(pa.Table.from_pylist(rows), path) + path.with_suffix(".metadata.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_forward_perf", + "schema_version": 6, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + ) + ) + return str(root) + + return make_root + + +@pytest.mark.parametrize("fallback", ["deny", "allow"]) +def test_direct_fake_only_cell_is_unavailable(profile_dict, measured_profile_roots, fallback): + root = measured_profile_roots("fake", "fake") + config = _request(profile_dict, "direct", systems_paths=[root], fallback_policy=fallback) + if fallback == "deny": + with pytest.raises(PerfDataNotAvailableError, match="no genuine measurements"): + RustForwardPassPerfModel.best_available(config) + return + model = RustForwardPassPerfModel.best_available(config) + diagnostics = model.diagnostics() + assert diagnostics["readiness"] == "unsupported_config" + assert diagnostics["provenance"]["selected_estimation_mode"] == "fpm_regression" + assert diagnostics["provenance"]["selected_systems_root"] is None + assert "no genuine measurements" in diagnostics["provenance"]["selection_failures"][0] + assert ( + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 64}} + ) + is None + ) + + +@pytest.mark.parametrize("fallback", ["deny", "allow"]) +@pytest.mark.parametrize( + "kinds,selected", + [(["genuine"], 0), (["mixed"], 0), (["fake", "genuine"], 1), (["fake", "mixed"], 1), (["genuine", "fake"], 0)], +) +def test_direct_availability_uses_first_genuine_root_and_pins_queries( + profile_dict, measured_profile_roots, fallback, kinds, selected +): + roots = [measured_profile_roots(str(index), kind) for index, kind in enumerate(kinds)] + model = RustForwardPassPerfModel.best_available( + _request(profile_dict, "direct", systems_paths=roots, fallback_policy=fallback) + ) + before = model.diagnostics() + assert before["readiness"] == "ready" + assert before["provenance"]["selected_systems_root"] == roots[selected] + assert before["provenance"]["config"]["systems_paths"] == [roots[selected]] + # This is the exact, hand-declared genuine decode timing above. + assert ( + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 1}} + ) + == 3.0 + ) + with pytest.raises(PerfDataNotAvailableError, match="direct"): + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 64}} + ) + assert model.diagnostics() == before + + +def test_unknown_profile_outer_auto_skips_analytical_backend_parsing(profile_dict, measured_profile_roots): + root = measured_profile_roots("genuine", "genuine") + model = RustForwardPassPerfModel.best_available( + _request(profile_dict, systems_paths=[root], estimation_mode="auto") + ) + provenance = model.diagnostics()["provenance"] + assert provenance["selected_estimation_mode"] == "fpm_interpolation" + assert provenance["config"]["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert provenance["config"]["attention_backend"] == "future_attention" + assert len(provenance["selection_failures"]) == 1 + assert "registered architecture" in provenance["selection_failures"][0] + assert ( + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 1}} + ) + == 3.0 + ) + + +def test_registered_profile_outer_auto_keeps_op_level_priority_with_inner_direct(profile_dict, monkeypatch): + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + profile_dict.update(model="nvidia/GLM-5.2-NVFP4", architecture="GlmMoeDsaForCausalLM") + # Op-level priority uses a currently shipped op database; FPM0.25.1 pairs + # are external inputs and do not establish op-level data availability. + profile_dict["deployments"][0].update(system="b200_sxm", tp=1, dp=8, moe_tp=1, moe_ep=8, backend_version="0.24.0") + monkeypatch.setattr(engine, "_direct_fpm_spec_json", _fail_graph) + model = RustForwardPassPerfModel.best_available(_request(profile_dict, "direct", estimation_mode="auto")) + provenance = model.diagnostics()["provenance"] + assert provenance["selected_estimation_mode"] == "op_level" + assert provenance["selection_failures"] == [] + assert provenance["config"]["estimator_config"]["fpm_interpolation"]["method"] == "direct" + + +@pytest.mark.parametrize("options", [{"text_only": True}, {"unrecorded_quant_modes": ["fmha"]}]) +def test_profile_rejects_unsupported_fpm_options_at_canonical_and_compile_boundaries(profile_dict, options): + config = _request(profile_dict, "direct") + config["estimator_config"] = {"fpm_interpolation": {"method": "direct", **options}} + with pytest.raises(ValueError, match="fpm_profile does not support"): + _normalize(config) + with pytest.raises(ValueError, match="fpm_profile does not support"): + engine.compile_engine( + "test/unknown-decoder", + "test_gpu", + "vllm", + "0.25.1", + tp_size=2, + moe_tp_size=2, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + fpm_options=options, + ) diff --git a/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py b/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py index 15d283b67..da626b0f5 100644 --- a/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py +++ b/python/aisimulate/tests/unit/sdk/test_prefill_graph_canonical.py @@ -218,6 +218,8 @@ def test_profile_rejects_telemetry_tuning_and_energy_routes_including_empty_inpu for payload in [[], {}, [{"scheduled_requests": []}]]: with pytest.raises((PrefillGraphProfileError, ValueError)): graph_model.estimate_forward_pass_time_ms(payload) + with pytest.raises((PrefillGraphProfileError, ValueError)): + graph_model.estimate_forward_pass_detailed(payload) with pytest.raises(PrefillGraphProfileError, match="only direct"): graph_model.tune_with_fpms([]) for prefill in [False, True]: 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 3bdf44073..fe0bf7450 100644 --- a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py +++ b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py @@ -1014,11 +1014,23 @@ def test_forward_pass_config_preserves_existing_positional_arguments(tmp_path: P "wideep_num_slots": 64, } config = ForwardPassPerfModelConfig(*legacy_fields.values()) - assert vars(config) == {**legacy_fields, "dcp": None, "fpm_fmha_quant_mode": None, "moe_kernel_source": None} + assert vars(config) == { + **legacy_fields, + "fpm_profile": None, + "dcp": None, + "fpm_fmha_quant_mode": None, + "moe_kernel_source": None, + } source = " source_with_spaces " pinned = ForwardPassPerfModelConfig(*legacy_fields.values(), moe_kernel_source=source) - assert vars(pinned) == {**legacy_fields, "dcp": None, "fpm_fmha_quant_mode": None, "moe_kernel_source": source} + assert vars(pinned) == { + **legacy_fields, + "fpm_profile": None, + "dcp": None, + "fpm_fmha_quant_mode": None, + "moe_kernel_source": source, + } assert json.loads(json.dumps(pinned.to_dict()))["moe_kernel_source"] == source diff --git a/python/aisimulate/tests/unit/sdk/test_utils.py b/python/aisimulate/tests/unit/sdk/test_utils.py index d764295f8..cf3a757ce 100644 --- a/python/aisimulate/tests/unit/sdk/test_utils.py +++ b/python/aisimulate/tests/unit/sdk/test_utils.py @@ -7,12 +7,14 @@ Tests HuggingFace config parsing and model config retrieval. """ +import io +import json from types import SimpleNamespace from unittest.mock import patch import pytest -from aisimulate.sdk import common, config +from aisimulate.sdk import common, config, utils from aisimulate.sdk.backends.base_backend import BaseBackend from aisimulate.sdk.models import Gemma4MixModel, HybridMoEModel from aisimulate.sdk.utils import ( @@ -25,6 +27,26 @@ pytestmark = pytest.mark.unit +@pytest.mark.parametrize( + "revision,path_revision", + [(None, "main"), ("a" * 40, "a" * 40), ("release/v1", "release%2Fv1")], +) +def test_download_hf_config_json_revision(monkeypatch, revision, path_revision): + requests = [] + metadata = {"architectures": ["UnregisteredDecoderForCausalLM"]} + + def fetch(request, *, timeout): + requests.append(request.full_url) + assert timeout == 30 + return io.BytesIO(json.dumps(metadata).encode()) + + monkeypatch.setattr(utils, "_get_hf_auth_headers", lambda: {}) + monkeypatch.setattr(utils.urllib.request, "urlopen", fetch) + kwargs = {} if revision is None else {"revision": revision} + assert utils._download_hf_json("test/model", "config.json", **kwargs) == metadata + assert requests == [f"https://huggingface.co/test/model/raw/{path_revision}/config.json"] + + class TestParseHFConfig: """Test HuggingFace config parsing.""" diff --git a/python/aisimulate/tests/unit/test_fpm_config_imports.py b/python/aisimulate/tests/unit/test_fpm_config_imports.py new file mode 100644 index 000000000..04a59b74f --- /dev/null +++ b/python/aisimulate/tests/unit/test_fpm_config_imports.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FPM parsing must finish before resource supervision loads the runtime.""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + + +def _config() -> dict: + return { + "engine": { + "mode": "aggregated", + "model": "test/unregistered-decoder", + "hardware": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "workers": {"aggregated": {"scheduler": {"max_batched_tokens": 1024, "max_sequences": 4}}}, + "fpm_profile": { + "schema_version": 1, + "model": "test/unregistered-decoder", + "model_revision": "import-boundary-fixture-v1", + "architecture": "UnregisteredDecoderForCausalLM", + "context_length": 4096, + "num_experts": 0, + "provenance": "Synthetic metadata for pre-budget validation; no timing measurements.", + "deployments": [ + { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 1, + "dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "gemm_quant_mode": "fp8", + "moe_quant_mode": "fp8", + "fmha_quant_mode": "fp8", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 100, + "activations_bytes": 20, + "runtime_overhead_bytes": 30, + "comm_overhead_bytes": 50, + "kv_bytes_per_token": 10, + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 256, + "provenance": "Declared bounds for import validation only.", + }, + } + ], + }, + } + } + + +_IMPORT_GUARD = """ +import sys +from pathlib import Path + +blocked = ('aisimulate._runtime', 'aisimulate.sweeper', 'numpy', 'pandas', 'aiconfigurator', + 'aiconfigurator_core', 'aisimulate_core._native', 'aisimulate_core.sdk', + 'urllib.request', 'huggingface_hub', 'transformers') + +def assert_lightweight(): + loaded = [name for name in sys.modules if any(name == item or name.startswith(item + '.') for item in blocked)] + assert not loaded, loaded + core = {name for name in sys.modules if name == 'aisimulate_core' or name.startswith('aisimulate_core.')} + assert core <= {'aisimulate_core', 'aisimulate_core.fpm_profile', 'aisimulate_core.quantization'}, core + +class RejectRuntimeImports: + def find_spec(self, fullname, path=None, target=None): + if any(fullname == item or fullname.startswith(item + '.') for item in blocked): + raise AssertionError('pre-budget runtime import: ' + fullname) + +assert_lightweight() +sys.meta_path.insert(0, RejectRuntimeImports()) +""" + +_SUPERVISION_PROBE = ( + _IMPORT_GUARD + + """ +from aisimulate import supervision + +calls = [] +def supervised_child(*args, **kwargs): + assert_lightweight() + calls.append(args) + return {'status': 'completed', 'exit_code': 0} + +supervision.run_process = supervised_child +config, output, command, error = sys.argv[1:] +try: + code = supervision.main([command, '--config', config, '--output-dir', output, '--overwrite']) +except SystemExit as exc: + code = exc.code +assert_lightweight() +if error: + assert code == 2, code + assert not calls, calls + assert (Path(output) / 'existing.txt').read_text() == 'keep existing results' +else: + assert code == 0, code + assert len(calls) == 1, calls + from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig + config_type = CorePredictionConfig if command == 'predict' else CoreRecommendationConfig + parsed = config_type.from_yaml(config) + assert config_type.model_validate_json(parsed.model_dump_json()) == parsed + assert_lightweight() +""" +) + + +@pytest.mark.parametrize("command", ["predict", "recommend"]) +@pytest.mark.parametrize("case", ["no-profile", "profile", "unknown-quant", "boolean-topology", "wrong-model"]) +def test_supervised_profile_validation_is_runtime_free(command: str, case: str, tmp_path: Path) -> None: + """Exercise the real parser and envelope validation before the child starts.""" + raw = _config() + if command == "recommend": + raw["optimization"] = {"constraints": {"max_candidate_gpus": 1}} + error = "" + if case == "no-profile": + del raw["engine"]["fpm_profile"] + elif case == "unknown-quant": + raw["engine"]["fpm_profile"]["deployments"][0]["fmha_quant_mode"] = "unsupported" + error = "unknown FPM fmha_quant_mode" + elif case == "boolean-topology": + raw["engine"]["fpm_profile"]["deployments"][0]["tp"] = True + error = "Input should be a valid integer" + elif case == "wrong-model": + raw["engine"]["model"] = "test/different-decoder" + error = "engine.model must match engine.fpm_profile.model" + config = tmp_path / "config.json" + config.write_text(json.dumps(raw)) + output = tmp_path / "output" + output.mkdir() + (output / "existing.txt").write_text("keep existing results") + + result = subprocess.run( + [sys.executable, "-c", _SUPERVISION_PROBE, str(config), str(output), command, error], + capture_output=True, + text=True, + timeout=15, + ) + assert result.returncode == 0, result.stdout + result.stderr + if error: + assert error in result.stderr + + +def test_onboarding_parser_is_runtime_free() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + _IMPORT_GUARD + + """ +from aisimulate.cli_args import build_parser +parser = build_parser() +try: + parser.parse_args(['onboard', 'init', '--help']) +except SystemExit as exc: + assert exc.code == 0 +assert_lightweight() +""", + ], + capture_output=True, + text=True, + timeout=15, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "--model" in result.stdout diff --git a/python/aisimulate/tests/unit/test_support_plan.py b/python/aisimulate/tests/unit/test_support_plan.py index c07453354..8203a2071 100644 --- a/python/aisimulate/tests/unit/test_support_plan.py +++ b/python/aisimulate/tests/unit/test_support_plan.py @@ -208,8 +208,8 @@ def test_onboarding_plan_and_preview_do_not_import_estimator_runtime(tmp_path): import sys from pathlib import Path -blocked = ("aisimulate._runtime", "aisimulate_core", "aiconfigurator", "aiconfigurator_core", - "numpy", "pandas", "torch", "transformers", "huggingface_hub", "collector") +blocked = ("aisimulate._runtime", "aisimulate_core._native", "aisimulate_core.sdk", "aiconfigurator", + "aiconfigurator_core", "numpy", "pandas", "torch", "transformers", "huggingface_hub", "collector") class NoEstimatorRuntime(importlib.abc.MetaPathFinder): def find_spec(self, fullname, path=None, target=None): diff --git a/tests/sweeper/test_forward_pass_estimator.py b/tests/sweeper/test_forward_pass_estimator.py index ae856f2af..cd94226e0 100644 --- a/tests/sweeper/test_forward_pass_estimator.py +++ b/tests/sweeper/test_forward_pass_estimator.py @@ -181,6 +181,79 @@ def request(**changes): ) +@pytest.mark.parametrize("input_kind", ["dataclass", "mapping"]) +@pytest.mark.parametrize("field", ["attention_kv_weight", "prefill_attention_pair_weight", "ffn_token_weight"]) +@pytest.mark.parametrize("sentinel", ["NaN", "Infinity", "-Infinity"]) +def test_canonical_nonfinite_regression_weights_preserve_native_and_regression_contracts(input_kind, field, sentinel): + from dataclasses import replace + + config = request( + tp=2, + backend_version="0.24.0", + estimation_mode="op_level", + estimator_config={"features": {field: float(sentinel)}}, + ) + payload = config if input_kind == "dataclass" else config.to_dict() + original = json.dumps(config.to_dict() if input_kind == "dataclass" else payload, sort_keys=True) + normalized = RustForwardPassPerfModel.normalize_config(payload) + assert normalized["estimator_config"]["features"][field] == sentinel + model = RustForwardPassPerfModel.best_available(payload) + try: + diagnostics = model.diagnostics() + assert diagnostics["readiness"] == "ready" + assert diagnostics["provenance"]["selected_estimation_mode"] == "op_level" + assert diagnostics["provenance"]["config"]["estimator_config"]["features"][field] == sentinel + assert ( + model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 128}} + ) + > 0 + ) + finally: + model.close() + assert json.dumps(config.to_dict() if input_kind == "dataclass" else payload, sort_keys=True) == original + + regression = replace(config, estimation_mode="fpm_regression") + regression_payload = regression if input_kind == "dataclass" else regression.to_dict() + original = json.dumps(regression.to_dict() if input_kind == "dataclass" else regression_payload, sort_keys=True) + normalized = RustForwardPassPerfModel.normalize_config(regression_payload) + assert normalized["estimator_config"]["features"][field] == sentinel + with pytest.raises(ValueError, match=f"regression_{field}"): + RustForwardPassPerfModel.best_available(regression_payload) + assert ( + json.dumps(regression.to_dict() if input_kind == "dataclass" else regression_payload, sort_keys=True) + == original + ) + + +@pytest.mark.parametrize("field", ["attention_kv_weight", "prefill_attention_pair_weight", "ffn_token_weight"]) +@pytest.mark.parametrize("sentinel", ["NaN", "Infinity", "-Infinity"]) +def test_resolver_preserves_nonfinite_regression_weight_contract(field, sentinel): + from dataclasses import replace + + from aisimulate.sweeper.config import SearchSpace + from aisimulate.sweeper.forward_pass_estimator import ( + ForwardPassEstimatorResolutionError, + ForwardPassEstimatorResolver, + ) + + config = request( + tp=2, + backend_version="0.24.0", + estimation_mode="op_level", + estimator_config={"features": {field: float(sentinel)}}, + ) + resolver = ForwardPassEstimatorResolver(SearchSpace(model_name=config.model, hardware_sku=config.system)) + spec = resolver._resolve(config, "agg") + assert spec.diagnostics["readiness"] == "ready" + assert spec.config["estimation_mode"] == "op_level" + assert spec.config["estimator_config"]["features"][field] == sentinel + spec.config["estimator_config"]["features"][field] = 1.0 + assert resolver._resolve(config, "agg").config["estimator_config"]["features"][field] == sentinel + with pytest.raises(ForwardPassEstimatorResolutionError, match=f"regression_{field}"): + resolver._resolve(replace(config, estimation_mode="fpm_regression"), "agg") + + def test_legacy_options_preserve_shape_ridge_and_general_mapping(): from collections.abc import Mapping diff --git a/tests/test_fpm_profile_workflow.py b/tests/test_fpm_profile_workflow.py new file mode 100644 index 000000000..9786fb4f3 --- /dev/null +++ b/tests/test_fpm_profile_workflow.py @@ -0,0 +1,870 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Class-independent public plumbing; real timing validation is separate.""" + +from __future__ import annotations + +import hashlib +import json +from copy import deepcopy +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from aisimulate.capacity import materialize_aic_num_gpu_blocks +from aisimulate.compiler import prediction_to_replay_spec +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.config.common import ENGINE_MODEL_CONTROL_FIELDS +from aisimulate.recommend import _candidate_prediction, recommendation_to_sweeper +from aisimulate.support.fpm import fpm_cli_args +from aisimulate.support.plan import check_plan, create_plan +from aisimulate.support.schema import SupportRequest +from aisimulate.sweeper.config import SearchSpace +from aisimulate.sweeper.deploy import build_backend_deployment +from aisimulate.sweeper.forward_pass_estimator import ( + ForwardPassEstimatorResolutionError, + ForwardPassEstimatorResolver, +) +from aisimulate.sweeper.model_hw import parallel_configs_for +from aisimulate.sweeper.replay import ReplaySpec +from aisimulate.sweeper.sample import unroll_sample + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def profile(tmp_path_factory): + checkpoint = tmp_path_factory.mktemp("workflow-checkpoint") + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["UnregisteredMoeForCausalLM"]})) + # Deliberately fictional metadata, with declared byte bounds, proves the + # workflow cannot recover resources through an existing model registration. + common = { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "gemm_quant_mode": "fp8_block", + "moe_quant_mode": "fp8_block", + "fmha_quant_mode": "bfloat16", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "weights_bytes": 20 * 1024**3, + "activations_bytes": 1024**3, + "runtime_overhead_bytes": 1024**3, + "comm_overhead_bytes": 1024**3, + "kv_bytes_per_token": 512, + "cache_layout": "linear", + "max_num_tokens": 8192, + "max_batch_size": 8, + "provenance": "Synthetic rank-local bounds for workflow tests; excludes CUDA graph reservation.", + }, + } + return { + "schema_version": 1, + "model": str(checkpoint), + "model_revision": "workflow-fixture-v1", + "architecture": "UnregisteredMoeForCausalLM", + "context_length": 4096, + "num_experts": 8, + "provenance": "Fictional model, no timing or silicon accuracy claim.", + "deployments": [ + {**deepcopy(common), "tp": 2, "dp": 1, "moe_tp": 2, "moe_ep": 1}, + { + **deepcopy(common), + "tp": 1, + "dp": 2, + "moe_tp": 1, + "moe_ep": 2, + "resources": {**common["resources"], "kv_bytes_per_token": 1024, "weights_bytes": 30 * 1024**3}, + }, + ], + } + + +@pytest.fixture(autouse=True) +def forbid_registered_model(monkeypatch): + from aisimulate import capacity, compiler, recommend + from aisimulate.sweeper import model_hw + from aisimulate_core.sdk import engine, memory + + def fail(*args, **kwargs): + pytest.fail("class-independent workflow reached model/config construction") + + monkeypatch.setattr(engine, "get_model", fail) + monkeypatch.setattr(memory.KVCacheEstimator, "from_request", fail) + monkeypatch.setattr(memory.NaiveKVCacheEstimator, "from_model_path", fail) + monkeypatch.setattr(capacity, "resolve_model_context_length", fail) + monkeypatch.setattr(compiler, "resolve_model_context_length", fail) + monkeypatch.setattr(recommend, "resolve_model_context_length", fail) + monkeypatch.setattr(model_hw, "get_model_config_from_model_path", fail) + monkeypatch.setattr(model_hw, "_estimate_model_weight_bytes", fail) + + +@pytest.fixture +def timing_systems(profile, tmp_path, monkeypatch): + """Small synthetic timing cells exercise the public constructor and replay transport.""" + import pyarrow as pa + import pyarrow.parquet as pq + + import aisimulate_core + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + root = tmp_path / "systems" + root.mkdir() + packaged = Path(aisimulate_core.__file__).parent / "systems" + (root / "h200_sxm.yaml").write_bytes((packaged / "h200_sxm.yaml").read_bytes()) + rows = [] + for index, deployment in enumerate(load_fpm_profile(profile).deployments): + identity = deployment.model_dump(mode="json", exclude={"resources"}) + for phase, kv_tokens in (("prefill", 0), ("decode", 0), ("decode", 1), ("decode", 64)): + rows.append( + { + **identity, + "cell_id": f"synthetic-{index}-{phase}-{kv_tokens}", + "model_path": profile["model"], + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": 1, + "total_prefill_tokens": 1 if phase == "prefill" else 0, + "total_kv_read_tokens": kv_tokens, + "latency_ms": 1.0, + "kv_seed_regime": "real_kv", + } + ) + path = root / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + path.parent.mkdir(parents=True) + pq.write_table(pa.Table.from_pylist(rows), path) + path.with_suffix(".metadata.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_forward_perf", + "schema_version": 6, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + ) + ) + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + return str(root) + + +def _worker(*, dep=False): + return { + "parallelism": { + "tensor": 1 if dep else 2, + "attention_data": 2 if dep else 1, + "moe_tensor": 1 if dep else 2, + "moe_expert": 2 if dep else 1, + }, + "scheduler": {"max_batched_tokens": 1024, "max_sequences": 4}, + "timing": { + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "estimator_config": {"fpm_interpolation": {"method": "direct"}}, + }, + } + + +def _engine(profile): + return { + "model": profile["model"], + "fpm_profile": profile, + "hardware": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "workers": {"aggregated": _worker()}, + } + + +def test_prediction_materializes_rank_local_capacity_without_a_model(profile, timing_systems): + config = CorePredictionConfig.model_validate({"engine": {**_engine(profile), "systems_paths": [timing_systems]}}) + spec = prediction_to_replay_spec(config) + args = spec.backend_deployment.agg_engine_args + assert args["max_model_len"] == 4096 + canonical = args["timing_model"]["config"] + assert canonical["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert canonical["fpm_profile"] == config.engine.fpm_profile.model_dump(mode="json") + metadata = spec.backend_deployment.performance_model_metadata["aggregated"]["config"] + assert metadata["fpm_profile"] == canonical["fpm_profile"] + assert metadata["estimator_config"] == canonical["estimator_config"] + diagnostics = {} + lowered = materialize_aic_num_gpu_blocks(args, memory_diagnostics=diagnostics) + assert lowered["num_gpu_blocks"] > 0 + assert diagnostics["source"] == "profile" + reserved_args = deepcopy(args) + reserved_args["cuda_graph_reserved_bytes"] = 1024**3 + reserved_args["timing_model"]["config"]["cuda_graph_reserved_bytes"] = 1024**3 + reserved = materialize_aic_num_gpu_blocks(reserved_args) + assert reserved["num_gpu_blocks"] < lowered["num_gpu_blocks"] + + +def test_disaggregated_transfer_uses_profile_cache_geometry(profile): + engine = _engine(profile) + engine.update( + mode="disaggregated", + workers={"prefill": _worker(dep=True), "decode": _worker()}, + kv_transfer={"bytes_per_token": "auto"}, + ) + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate({"engine": engine})) + assert spec.backend_deployment.prefill_engine_args["kv_transfer_bytes_per_token"] == 1024 + assert spec.backend_deployment.decode_engine_args["kv_transfer_bytes_per_token"] == 1024 + + +@pytest.mark.parametrize("mode,budget,expected", [("agg", 2, 2), ("disagg", 4, 4)]) +def test_profile_candidates_use_only_declared_topologies(profile, mode, budget, expected, monkeypatch): + from aisimulate_core.sdk import RustForwardPassPerfModel + + def no_timing(*args, **kwargs): + pytest.fail("candidate enumeration constructed a timing estimator before collection") + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", no_timing) + candidates = parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode=mode, + gpu_budget=budget, + max_num_tokens=1024, + max_batch_size=4, + fpm_profile=profile, + ) + assert len(candidates) == expected + assert all(candidate.total_gpus <= budget for candidate in candidates) + if mode == "agg": + assert {candidate.shape.strategy for candidate in candidates} == {"tp", "dep"} + + +def test_dense_profile_preserves_unsharded_moe_dimensions(profile): + profile["num_experts"] = 0 + profile["deployments"] = [{**profile["deployments"][0], "moe_tp": 1, "moe_ep": 1}] + candidates = parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=2, + max_num_tokens=1024, + max_batch_size=4, + fpm_profile=profile, + ) + assert [(candidate.shape.tp, candidate.shape.moe_tp, candidate.shape.moe_ep) for candidate in candidates] == [ + (2, 1, 1) + ] + + +@pytest.mark.parametrize("source", ["sample", "resolved"]) +def test_profile_kv_load_cache_keeps_resource_identity_without_model_construction(profile, monkeypatch, source): + from aisimulate.sweeper import kv_load + from aisimulate.sweeper.config import Workload + from aisimulate.sweeper.parallel_enum import ParallelShape, ReplicaParallelConfig + + sample = { + "deployment_mode": "agg", + "model_name": profile["model"], + "hardware_sku": "h200_sxm", + "backend": "vllm", + "agg_block_size": 64, + "agg_max_num_batched_tokens": 1024, + "agg_max_num_seqs": 4, + "agg_gpu_memory_utilization": 0.9, + } + parallel = ReplicaParallelConfig(ParallelShape(tp=2, dp=1, moe_tp=2, moe_ep=1), replicas=1) + larger = deepcopy(profile) + larger["deployments"][0]["resources"]["weights_bytes"] += 10 * 1024**3 + original = kv_load.estimate_kv_tokens + seen = [] + + def estimate(*args, fpm_profile, **kwargs): + seen.append(deepcopy(fpm_profile)) + return original(*args, fpm_profile=fpm_profile, **kwargs) + + monkeypatch.setattr(kv_load, "estimate_kv_tokens", estimate) + kv_load._per_rank_capacity_tokens.cache_clear() + capacities = [] + try: + for resources in (profile, larger, profile): + if source == "resolved": + sample["forward_pass_estimators"] = {"agg": {"config": {"fpm_profile": resources}}} + sample["fpm_profile"] = larger + else: + sample["fpm_profile"] = resources + result = kv_load.resolve_kv_load( + sample, + workload=Workload(isl=1024, osl=128, kv_load_ratio=1.0, request_count=1), + parallel_config=parallel, + ratio=1.0, + backend_version="0.25.1", + ) + capacities.append(result.role_capacity_tokens["agg"]) + assert capacities[0] == capacities[2] > capacities[1] + assert seen == [profile, larger] + finally: + kv_load._per_rank_capacity_tokens.cache_clear() + + +@pytest.mark.parametrize("method", ["auto", "direct"]) +@pytest.mark.parametrize("backend_version", ["0.25.1", {"vllm": " 0.25.1 "}]) +def test_recommendation_preserves_profile_in_exported_prediction( + profile, monkeypatch, method, timing_systems, backend_version +): + from aisimulate_core.sdk import RustForwardPassPerfModel + + engine = _engine(profile) + engine["backend_version"] = backend_version + engine["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + engine["workers"]["aggregated"]["timing"]["estimator_config"]["fpm_interpolation"]["method"] = method + engine["systems_paths"] = [timing_systems] + engine["mode"] = "aggregated" + source = CoreRecommendationConfig.model_validate( + {"engine": engine, "optimization": {"constraints": {"max_candidate_gpus": 2}}} + ) + smart = recommendation_to_sweeper(source) + assert smart.search_space.requested_backend_version("vllm") == "0.25.1" + candidates = parallel_configs_for( + profile["model"], + "h200_sxm", + backend="vllm", + backend_version="0.25.1", + deployment_mode="agg", + gpu_budget=2, + max_num_tokens=1024, + max_batch_size=4, + fpm_profile=profile, + ) + sample = unroll_sample( + search_space=smart.search_space, + selection={ + "deployment_mode": "agg", + "backend": "vllm", + "agg_max_num_batched_tokens": 1024, + "agg_max_num_seqs": 4, + }, + parallel_config=candidates[1], + ) + sample["backend_version"] = "0.25.1" + estimators = ForwardPassEstimatorResolver(smart.search_space).resolve_candidate(sample) + resolved = estimators["agg"].config + assert resolved["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert resolved["fpm_profile"] == source.engine.fpm_profile.model_dump(mode="json") + controls = { + "gemm_quant_mode": "fp8_block", + "moe_quant_mode": "fp8_block", + "fmha_quant_mode": "bfloat16", + "kvcache_quant_mode": "fp8", + "comm_quant_mode": "half", + "attention_backend": "auto", + "moe_backend": None, + "moe_kernel_source": None, + "enable_eplb": False, + "wideep_num_slots": None, + } + assert {name: resolved[name] for name in ENGINE_MODEL_CONTROL_FIELDS} == controls + assert resolved["systems_paths"] == [timing_systems] + from aisimulate_core.sdk.models.base import _MODEL_REGISTRY + + # A fresh environment may now have an analytical class for the same model. + # Export must retain the method actually selected for the evaluated sample. + monkeypatch.setitem(_MODEL_REGISTRY, profile["architecture"], object()) + replay = ReplaySpec( + backend_deployment=build_backend_deployment( + sample, backend_version="0.25.1", forward_pass_estimators=estimators + ), + workload={}, + goal={}, + ) + exported = _candidate_prediction(source, sample, replay, adapter_sections={}) + reloaded = CorePredictionConfig.model_validate(exported) + assert reloaded.engine.fpm_profile == source.engine.fpm_profile + assert reloaded.engine.workers.aggregated.timing.estimator_config["fpm_interpolation"]["method"] == "direct" + assert "forward_model" not in exported["engine"]["workers"]["aggregated"]["timing"] + assert "fpm_interpolation" not in exported["engine"]["workers"]["aggregated"]["timing"] + assert reloaded.engine.workers.aggregated.parallelism.attention_data == 2 + deployment = prediction_to_replay_spec(reloaded).backend_deployment + args = deployment.agg_engine_args + model = RustForwardPassPerfModel.best_available(deployment.performance_model_metadata["aggregated"]["config"]) + try: + identity = model.diagnostics()["provenance"]["config"] + finally: + model.close() + assert {name: identity[name] for name in ENGINE_MODEL_CONTROL_FIELDS} == controls + assert identity["worker_type"] == "aggregated" + assert tuple(identity[name] for name in ("tp", "pp", "attention_dp", "moe_tp_size", "moe_ep_size")) == ( + 1, + 1, + 2, + 1, + 2, + ) + assert materialize_aic_num_gpu_blocks(args)["num_gpu_blocks"] > 0 + + +@pytest.mark.parametrize( + "field,value", + [ + ("gemm_quant_mode", "bfloat16"), + ("moe_quant_mode", "bfloat16"), + ("fmha_quant_mode", "fp8"), + ("kvcache_quant_mode", "bfloat16"), + ("comm_quant_mode", "fp8"), + ("attention_backend", "fa3"), + ("moe_backend", "deepep_moe"), + ("enable_eplb", True), + ("wideep_num_slots", 128), + ], +) +def test_profile_resolver_normalizes_cache_identity_and_rejects_conflicting_overrides( + profile, timing_systems, monkeypatch, field, value +): + from dataclasses import replace + + from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + resolver = ForwardPassEstimatorResolver(SearchSpace(model_name=profile["model"], hardware_sku="h200_sxm")) + request = ForwardPassPerfModelConfig( + model=profile["model"], + fpm_profile=profile, + system="h200_sxm", + backend="vllm", + backend_version="0.25.1", + worker_type="aggregated", + tp=2, + moe_tp_size=2, + moe_ep_size=1, + kv_block_size=64, + systems_paths=(timing_systems,), + estimation_mode="fpm_interpolation", + estimator_config={"fpm_interpolation": {"method": "direct"}}, + ) + first = resolver._resolve(request, "agg") + + def no_build(*args, **kwargs): + pytest.fail("equivalent defaults or conflicting profile override reached estimator construction") + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", no_build) + explicit = replace(request, **{name: first.config[name] for name in ENGINE_MODEL_CONTROL_FIELDS}) + assert resolver._resolve(explicit, "agg").config == first.config + with pytest.raises(ForwardPassEstimatorResolutionError, match="identity conflict|override|EPLB"): + resolver._resolve(replace(explicit, **{field: value}), "agg") + + +@pytest.mark.parametrize("field", [*ENGINE_MODEL_CONTROL_FIELDS, "worker_type", "tp"]) +def test_profile_resolver_rejects_changed_resolved_identity(profile, timing_systems, monkeypatch, field): + from aisimulate_core.sdk import RustForwardPassPerfModel + + engine = _engine(profile) + engine.update(mode="aggregated", systems_paths=[timing_systems]) + engine["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + smart = recommendation_to_sweeper(CoreRecommendationConfig.model_validate({"engine": engine, "optimization": {}})) + sample = { + "deployment_mode": "agg", + "hardware_sku": "h200_sxm", + "backend": "vllm", + "tp": 2, + "pp": 1, + "attention_dp": 1, + "moe_tp": 2, + "moe_ep": 1, + "agg_block_size": 64, + } + diagnostics = RustForwardPassPerfModel.diagnostics + + def changed(model): + result = diagnostics(model) + result["provenance"]["config"][field] = "unexpected" + return result + + monkeypatch.setattr(RustForwardPassPerfModel, "diagnostics", changed) + with pytest.raises(ForwardPassEstimatorResolutionError, match=f"changed exact candidate field '{field}'"): + ForwardPassEstimatorResolver(smart.search_space).resolve_candidate(sample) + + +@pytest.mark.parametrize("role", ["agg", "prefill", "decode"]) +@pytest.mark.parametrize("timing", [{"type": "fixed", "prefill_ms": 1, "decode_ms": 1}, {"type": "polynomial"}]) +def test_profile_search_space_rejects_custom_timing(profile, role, timing): + with pytest.raises(ValidationError, match="fpm_profile requires default timing"): + SearchSpace( + model_name=profile["model"], + hardware_sku="h200_sxm", + backend_version="0.25.1", + fpm_profile=profile, + **{f"{role}_timing_model": timing, f"{role}_num_gpu_blocks": 1024}, + ) + + +@pytest.mark.parametrize( + "updates", + [ + {"deployment_mode": ["afd"]}, + {"deployment_mode": ["afd+pd"]}, + {"encoder": {"hardware_sku": "h200_sxm"}}, + {"backend": ["sglang"]}, + ], +) +def test_profile_search_space_rejects_unsupported_providers(profile, updates): + with pytest.raises(ValidationError, match="FPM profiles support vLLM aggregated/disaggregated"): + SearchSpace( + model_name=profile["model"], + hardware_sku="h200_sxm", + backend_version="0.25.1", + fpm_profile=profile, + **updates, + ) + + +@pytest.mark.parametrize( + "backend_version,error", + [ + (None, "literal"), + ({}, "literal"), + ({"sglang": "0.25.1"}, "literal"), + ({"vllm": " "}, "literal"), + ({"vllm": "current"}, "literal"), + ({"vllm": "previous"}, "literal"), + ({"vllm": "next"}, "literal"), + ({"vllm": "0.24.0"}, "does not match"), + ], +) +def test_profile_recommendation_cli_reports_invalid_version_mapping(profile, tmp_path, capsys, backend_version, error): + import aisimulate.main as cli + + engine = _engine(profile) + engine.update(mode="aggregated", backend_version=backend_version) + engine["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + path = tmp_path / "recommend.json" + path.write_text(json.dumps({"engine": engine})) + with pytest.raises(SystemExit) as exc: + cli.main(["recommend", "--config", str(path), "--output-dir", str(tmp_path / "output")]) + assert exc.value.code == 2 + stderr = capsys.readouterr().err + assert error in stderr + assert "TypeError" not in stderr + assert not (tmp_path / "output").exists() + + +def test_profile_recommendation_cli_accepts_version_mapping_and_exports_literal( + profile, timing_systems, tmp_path, monkeypatch +): + import yaml + + import aisimulate.main as cli + from aisimulate.recommend import _run_recommendation + + monkeypatch.setattr("aisimulate.recommend.run_recommendation", _run_recommendation) + engine = _engine(profile) + engine.update(mode="aggregated", backend_version={"vllm": "0.25.1"}, systems_paths=[timing_systems]) + engine["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + raw = { + "engine": engine, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 2}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 1}, + }, + "optimization": {"constraints": {"max_candidate_gpus": 2}}, + "optimizer": {"algorithm": "random", "max_trials": 1, "parallelism": 1}, + } + path = tmp_path / "recommend.json" + path.write_text(json.dumps(raw)) + output = tmp_path / "output" + assert cli.main(["recommend", "--config", str(path), "--output-dir", str(output), "--format", "json"]) == 0 + exported = yaml.safe_load((output / "recommendations/0001.yaml").read_text()) + assert exported["engine"]["backend_version"] == "0.25.1" + assert CorePredictionConfig.model_validate(exported).engine.fpm_profile.model == profile["model"] + + +@pytest.mark.parametrize( + "update,error", + [ + ({"model": "another/model"}, "match"), + ({"backend_version": None}, "literal"), + ({"context_length": 8192}, "exceeds"), + ({"workers": {"aggregated": {**_worker(), "scheduler": {"max_sequences": 32}}}}, "envelope exceeded"), + ({"workers": {"aggregated": {**_worker(), "parallelism": {"tensor": 4, "moe_tensor": 4}}}}, "no matching"), + ], +) +def test_invalid_profile_predictions_fail_before_runtime(profile, update, error): + with pytest.raises(ValidationError, match=error): + CorePredictionConfig.model_validate({"engine": {**_engine(profile), **update}}) + + +def test_onboard_dep_keeps_full_identity_and_checks_resources_before_collection(profile, tmp_path, monkeypatch): + from aisimulate_core.sdk import RustForwardPassPerfModel + + def no_timing(*args, **kwargs): + pytest.fail("planning constructed a timing estimator before collection") + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", no_timing) + request = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "moe", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "gpu_count": 2, + "gpus_per_node": 2, + "interconnect": "NVLink", + }, + "search": { + "tensor_parallel": 1, + "attention_data_parallel": 2, + "moe_tensor_parallel": 1, + "moe_expert_parallel": 2, + "context_length": 4096, + }, + "fpm_profile": profile, + } + ) + plan = create_plan(request, tmp_path) + assert plan["fpm"]["worker_gpus"] == 2 + assert plan["search"]["candidates"][0]["total_gpus"] == 2 + assert plan["resources"]["source"] == "profile" + assert plan["resources"]["total_kv_size_tokens"] > 4096 + loaded = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") + assert loaded.engine.workers.aggregated.parallelism.attention_data == 2 + assert loaded.engine.workers.aggregated.timing.estimator_config["fpm_interpolation"]["method"] == "direct" + assert loaded.engine.workers.aggregated.timing.estimation_mode == "fpm_interpolation" + assert loaded.engine.workers.aggregated.timing.fallback_policy == "deny" + command = fpm_cli_args(request, output_dir=tmp_path, plan_only=True) + assert command[command.index("--fpm-parallel-presets") + 1] == "dep" + assert command[command.index("--fpm-gpu-counts") + 1] == "2" + assert command[command.index("--fpm-kv-cache-dtypes") + 1] == "fp8" + profile_path = tmp_path / "fpm-model-profile.json" + assert command[command.index("--fpm-model-profile") + 1] == str(profile_path) + tampered = json.loads(profile_path.read_text()) + tampered["deployments"][0]["fmha_quant_mode"] = "fp8" + profile_path.write_text(json.dumps(tampered)) + with pytest.raises(ValueError, match="differs from the requested identity"): + check_plan(request, tmp_path) + + +@pytest.mark.parametrize( + "preset,max_tokens,max_sequences,moe_backend", + [ + ("dep", 4096, 8, "auto"), + ("dep", 8192, 4, "auto"), + ("dep", 8192, 8, "flashinfer_cutlass"), + ("tep", 4096, 4, "flashinfer_cutlass"), + ], +) +def test_generated_glm_collector_plan_preserves_backend_and_rank_local_bounds( + profile, tmp_path, monkeypatch, capsys, preset, max_tokens, max_sequences, moe_backend +): + from collector.fpm_forward import cli as collector_cli + + # Real GLM identity, with synthetic resource bounds: this exercises the + # generated command and collector admission without claiming GPU accuracy. + profile.update( + model="nvidia/GLM-5.2-NVFP4", + architecture="GlmMoeDsaForCausalLM", + context_length=8192, + num_experts=256, + provenance="Real GLM identity with synthetic resource bounds for guided collection tests.", + ) + deployment = profile["deployments"][0] + deployment.update( + system="b200_sxm", + tp=1 if preset == "dep" else 8, + dp=8 if preset == "dep" else 1, + moe_tp=1, + moe_ep=8, + gemm_quant_mode="nvfp4", + moe_quant_mode="nvfp4", + fmha_quant_mode="fp8", + moe_backend=moe_backend, + ) + deployment["resources"].update(max_num_tokens=max_tokens, max_batch_size=max_sequences) + profile["deployments"] = [deployment] + request = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "moe", + "framework_version": "0.25.1", + "gpu": "b200_sxm", + "gpu_count": 8, + "gpus_per_node": 8, + "interconnect": "NVLink", + }, + "search": { + "tensor_parallel": deployment["tp"], + "attention_data_parallel": deployment["dp"], + "moe_tensor_parallel": 1, + "moe_expert_parallel": 8, + "context_length": 8192, + }, + "workload": {"input_tokens": 1024, "output_tokens": 128, "concurrency": 8, "request_count": 8}, + "fpm_profile": profile, + } + ) + plan = create_plan(request, tmp_path) + command = json.loads((tmp_path / "commands.json").read_text())["fpm_plan_local"] + monkeypatch.delenv("COLLECTOR_MODEL_PATH", raising=False) + assert collector_cli.main(command[3:]) == 0 + collected = json.loads(capsys.readouterr().out) + assert collected["counts"]["cells"] == 2 + assert {cell["parallel_strategy"] for cell in collected["cells"]} == {preset} + assert collected["options"]["moe_backend"] == moe_backend + sampling = collected["options"]["prefill_sampling"] + expected_tokens = min(max_tokens, 1024 * max_sequences) + assert sampling["max_total_prefill_tokens"] == expected_tokens + assert sampling["max_batch_size"] == max_sequences + assert collected["topology_memory_admission"][0]["activation_envelope"]["scope"] == "rank_local" + prediction = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") + scheduler = prediction.engine.workers.aggregated.scheduler + assert scheduler.max_batched_tokens == max_tokens + assert scheduler.max_sequences == max_sequences + assert prediction.traffic.load.concurrency == 8 + assert prediction.traffic.stop.requests == 8 + assert "timing coverage" in plan["fpm"]["sampling"] + + +def test_onboard_rejects_profile_below_collector_minimum_token_axis(profile, tmp_path): + profile["deployments"][0]["resources"]["max_num_tokens"] = 1 + request = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "moe", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "gpu_count": 2, + "gpus_per_node": 2, + "interconnect": "NVLink", + }, + "search": {"tensor_parallel": 2, "context_length": 4096}, + "fpm_profile": profile, + } + ) + with pytest.raises(ValueError, match="FPM collection requires.*at least 2"): + create_plan(request, tmp_path) + + +@pytest.mark.parametrize( + "tp,dp,moe_tp,moe_ep,preset", [(2, 1, 2, 1, "pure_tp"), (1, 2, 1, 2, "dep"), (2, 1, 1, 2, "tep")] +) +def test_onboard_cli_embeds_profile_and_preserves_topology(profile, tmp_path, tp, dp, moe_tp, moe_ep, preset): + import aisimulate.main as cli + + profile["deployments"].append({**deepcopy(profile["deployments"][0]), "moe_tp": 1, "moe_ep": 2}) + source = tmp_path / "model-profile.json" + source.write_text(json.dumps(profile)) + request_path = tmp_path / "request.yaml" + assert ( + cli.main( + [ + "onboard", + "init", + "--model", + profile["model"], + "--model-revision", + profile["model_revision"], + "--model-kind", + "moe", + "--framework-version", + "0.25.1", + "--gpu", + "h200_sxm", + "--gpu-count", + "2", + "--interconnect", + "NVLink", + "--context-length", + "4096", + "--fpm-profile", + str(source), + "--tensor-parallel", + str(tp), + "--attention-data-parallel", + str(dp), + "--moe-tensor-parallel", + str(moe_tp), + "--moe-expert-parallel", + str(moe_ep), + "--output", + str(request_path), + ] + ) + == 0 + ) + source.unlink() + request = SupportRequest.from_yaml(request_path) + assert request.fpm_profile.model == profile["model"] + assert request.worker_gpus == 2 + assert request.parallel_preset == preset + assert request.profile_deployment().parallel_tuple == (tp, 1, dp, moe_tp, moe_ep, 1) + plan = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(request_path), "--output-dir", str(plan)]) == 0 + prediction = CorePredictionConfig.from_yaml(plan / "predict/pilot.yaml") + assert prediction.engine.fpm_profile == request.fpm_profile + profile_output = plan / "fpm-model-profile.json" + original = profile_output.read_bytes() + profile_output.unlink() + create_plan(request, plan, overwrite=True) + assert profile_output.read_bytes() == original + + +@pytest.mark.parametrize("selector", [None, "fp8"]) +def test_direct_profile_uses_external_pair_and_preserves_fmha_arithmetic( + profile, timing_systems, tmp_path, monkeypatch, selector +): + import shutil + + import pyarrow as pa + import pyarrow.parquet as pq + + from aisimulate_core.sdk import RustForwardPassPerfModel + + # Move the only measured pair outside the systems tree, so silently losing + # the external path cannot succeed through packaged-data fallback. + source = Path(timing_systems) / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + external = tmp_path / "reviewed-fpm.parquet" + rows = pq.read_table(source).to_pylist() + if selector is not None: + for row in rows: + row["fmha_quant_mode"] = selector + pq.write_table(pa.Table.from_pylist(rows), external) + metadata = json.loads(source.with_suffix(".metadata.json").read_text()) + metadata["parquet_sha256"] = hashlib.sha256(external.read_bytes()).hexdigest() + external.with_suffix(".metadata.json").write_text(json.dumps(metadata)) + shutil.rmtree(Path(timing_systems) / "data") + monkeypatch.chdir(tmp_path) + model = RustForwardPassPerfModel.best_available( + { + "model": profile["model"], + "fpm_profile": profile, + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "worker_type": "aggregated", + "tp": 2, + "moe_tp_size": 2, + "moe_ep_size": 1, + "systems_paths": [timing_systems], + "estimation_mode": "fpm_interpolation", + "fpm_fmha_quant_mode": selector, + "estimator_config": {"fpm_interpolation": {"method": "direct", "fpm_parquet_path": external.name}}, + } + ) + resolved = model.diagnostics()["provenance"]["config"] + assert resolved["fmha_quant_mode"] == "bfloat16" + assert resolved["fpm_fmha_quant_mode"] == selector + assert resolved["estimator_config"]["fpm_interpolation"]["fpm_parquet_path"] == str(external) + # Every synthetic decode row is 1 ms; the unseen KV=2 point interpolates + # between the independently declared KV=1 and KV=64 rows without a graph. + assert model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 2}} + ) == pytest.approx(1.0) diff --git a/tests/test_fpm_query_evidence.py b/tests/test_fpm_query_evidence.py new file mode 100644 index 000000000..5977dfc63 --- /dev/null +++ b/tests/test_fpm_query_evidence.py @@ -0,0 +1,454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Executed public FPM evidence over real metadata and synthetic Parquets.""" + +import hashlib +import json +from copy import deepcopy +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +# Reuse the config-only, verified unregistered architecture and its guards. +import test_fpm_profile_workflow as workflow + +from aisimulate.compiler import prediction_to_replay_spec +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.detail import build_prediction_details +from aisimulate.recommend import run_recommendation +from aisimulate.runner import EngineReplayRunnerFactory, _normalize_engine_replay_report +from aisimulate.sweeper.replay import ReplayOutputRequirements +from aisimulate.sweeper.result import SweepResult +from aisimulate_core.sdk import RustForwardPassPerfModel + +pytestmark = pytest.mark.unit + +profile = workflow.profile +timing_systems = workflow.timing_systems +forbid_registered_model = workflow.forbid_registered_model + + +@pytest.fixture +def evidence_systems(timing_systems): + path = Path(timing_systems) / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + templates = [row for row in pq.read_table(path).to_pylist() if row["workload_kind"] == "prefill"] + rows = [] + # Prefill token interpolation at 6 lies 1/4 along [4,12]; KV=2 is + # 1/4 along [0,8]. The four weights are 9/16,3/16,3/16,1/16. + points = [ + ("prefill", 1, 1, 0, 2.0), + ("prefill", 1, 4, 0, 10.0), + ("prefill", 1, 12, 0, 18.0), + ("prefill", 1, 4, 8, 30.0), + ("prefill", 1, 12, 8, 46.0), + ("decode", 1, 0, 0, 3.0), + ("decode", 1, 0, 1, 4.0), + ("decode", 1, 0, 64, 67.0), + ("decode", 2, 0, 4, 10.0), + ("decode", 2, 0, 20, 26.0), + ("decode", 4, 0, 8, 20.0), + ("decode", 4, 0, 24, 52.0), + # A decreasing measured curve exercises marginal-decode clamping. + ("decode", 6, 0, 12, 30.0), + ("decode", 6, 0, 24, 10.0), + ] + for template in templates: + for phase, batch, tokens, kv, latency in points: + rows.append( + { + **template, + "workload_kind": phase, + "batch_size": batch, + "total_prefill_tokens": tokens, + "total_kv_read_tokens": kv, + "latency_ms": latency, + } + ) + pq.write_table(pa.Table.from_pylist(rows), path) + metadata = json.loads(path.with_suffix(".metadata.json").read_text()) + metadata.update(row_count=len(rows), parquet_sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + path.with_suffix(".metadata.json").write_text(json.dumps(metadata)) + return timing_systems + + +def canonical(profile, systems): + return { + "model": profile["model"], + "fpm_profile": profile, + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "worker_type": "aggregated", + "tp": 2, + "moe_tp_size": 2, + "moe_ep_size": 1, + "systems_paths": [systems], + "estimation_mode": "fpm_interpolation", + "estimator_config": {"fpm_interpolation": {"method": "direct"}, "correction": {"min_observations": 1}}, + } + + +def prefill(tokens=4, kv=0): + return { + "scheduled_requests": {"num_prefill_requests": 1, "sum_prefill_tokens": tokens, "sum_prefill_kv_tokens": kv} + } + + +def decode(batch=3, kv=12): + return {"scheduled_requests": {"num_decode_requests": batch, "sum_decode_kv_tokens": kv}} + + +@pytest.mark.parametrize( + "metrics,latency,resolution,weights,measurements", + [ + (prefill(), 10.0, "exact_lookup", [1.0], [(1, 4, 0, 10)]), + (prefill(6), 12.0, "within_curve_interpolation", [0.75, 0.25], [(1, 4, 0, 10), (1, 12, 0, 18)]), + ( + prefill(6, 2), + 17.5, + "cross_kv_interpolation", + [0.5625, 0.1875, 0.1875, 0.0625], + [(1, 4, 0, 10), (1, 12, 0, 18), (1, 4, 8, 30), (1, 12, 8, 46)], + ), + (decode(2, 4), 10.0, "exact_lookup", [1.0], [(2, None, 4, 10)]), + (decode(2, 12), 18.0, "within_curve_interpolation", [0.5, 0.5], [(2, None, 4, 10), (2, None, 20, 26)]), + ( + decode(), + 23.0, + "cross_batch_interpolation", + [0.25, 0.25, 0.375, 0.125], + [(2, None, 4, 10), (2, None, 20, 26), (4, None, 8, 20), (4, None, 24, 52)], + ), + ], +) +def test_canonical_detailed_query_support( + profile, evidence_systems, metrics, latency, resolution, weights, measurements +): + model = RustForwardPassPerfModel.best_available(canonical(profile, evidence_systems)) + detailed = model.estimate_forward_pass_detailed(metrics) + assert detailed["latency_ms"] == model.estimate_forward_pass_time_ms(metrics) == latency + assert detailed["native_latency_ms"] == latency + assert detailed["correction_factor"] == 1.0 + evidence = detailed["ranks"][0]["queries"][0] + assert evidence["model_path"] == profile["model"] + assert evidence["resolution"] == resolution + assert evidence["latency_ms"] == latency + assert evidence["decode_baseline"] is False + assert [point["weight"] for point in evidence["support"]] == weights + assert [ + ( + point["coordinates"]["batch_size"], + point["coordinates"]["total_prefill_tokens"], + point["coordinates"]["total_kv_read_tokens"], + point["latency_ms"], + ) + for point in evidence["support"] + ] == measurements + assert json.loads(json.dumps(detailed)) == detailed + + +def test_mixed_baseline_rank_max_correction_and_clamp(profile, evidence_systems): + model = RustForwardPassPerfModel.best_available(canonical(profile, evidence_systems)) + mixed = {"scheduled_requests": {**prefill()["scheduled_requests"], **decode()["scheduled_requests"]}} + metrics = [mixed, decode(4, 24)] + detail = model.estimate_forward_pass_detailed(metrics) + assert detail["latency_ms"] == model.estimate_forward_pass_time_ms(metrics) == 52.0 + assert detail["max_rank"] == 1 + rank = detail["ranks"][0] + assert [ + rank[key] for key in ("prefill_ms", "decode_ms", "decode_baseline_ms", "marginal_decode_ms", "latency_ms") + ] == [10, 23, 15, 8, 18] + baseline = rank["queries"][2] + assert baseline["decode_baseline"] is True + assert baseline["query"]["total_kv_read_tokens"] == 12 + assert [(p["coordinates"]["total_kv_read_tokens"], p["weight"]) for p in baseline["support"]] == [ + (4, 0.5), + (8, 0.5), + ] + observed = deepcopy(metrics) + observed[1]["wall_time"] = 0.104 # seconds: twice the 52 ms native maximum + model.tune_with_fpms([observed]) + corrected = model.estimate_forward_pass_detailed(metrics) + assert corrected["correction_factor"] == 2.0 + assert corrected["latency_ms"] == model.estimate_forward_pass_time_ms(metrics) == 104.0 + assert corrected["ranks"] == detail["ranks"] + clamped = {"scheduled_requests": {**prefill()["scheduled_requests"], **decode(6, 24)["scheduled_requests"]}} + detail = model.estimate_forward_pass_detailed(clamped) + assert detail["ranks"][0]["decode_ms"] == 10 + assert detail["ranks"][0]["decode_baseline_ms"] == 30 + assert detail["ranks"][0]["marginal_decode_ms"] == 0 + assert detail["ranks"][0]["latency_ms"] == 10 + + +def test_uncovered_empty_and_non_fpm_results(profile, evidence_systems): + config = canonical(profile, evidence_systems) + model = RustForwardPassPerfModel.best_available(config) + for query in (prefill(20), prefill(4, 20), decode(3, 30)): + for method in (model.estimate_forward_pass_time_ms, model.estimate_forward_pass_detailed): + with pytest.raises(Exception, match="unsupported|outside"): + method(query) + empty = model.estimate_forward_pass_detailed({}) + assert empty == { + "latency_ms": 0.0, + "native_latency_ms": 0.0, + "correction_factor": 1.0, + "max_rank": None, + "ranks": [], + } + config["estimation_mode"] = "fpm_regression" + regression = RustForwardPassPerfModel.best_available(config) + assert regression.estimate_forward_pass_detailed(decode()) == { + "latency_ms": None, + "native_latency_ms": None, + "correction_factor": None, + "max_rank": None, + "ranks": [], + } + + +def prediction(profile, systems): + return { + "engine": {**workflow._engine(profile), "systems_paths": [systems]}, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 3}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 3}, + }, + } + + +def records(evidence): + return [ + record + for provider in evidence["providers"] + for phase in provider["phases"] + for operation in phase["operations"] + for record in operation["fpm_estimates"] + ] + + +def test_native_prediction_replay_and_source_detail_roundtrip(profile, evidence_systems): + import jsonschema + + config = CorePredictionConfig.model_validate(prediction(profile, evidence_systems)) + # Round-trip the saved public config before lowering it to the real runner. + config = CorePredictionConfig.model_validate_json(config.model_dump_json()) + spec = prediction_to_replay_spec(config) + runner = EngineReplayRunnerFactory().create(0) + try: + plain = runner.run(spec) + report = runner.run(spec, output_requirements=ReplayOutputRequirements(capture_performance_diagnostics=True)) + finally: + runner.close() + assert "fpm_query_evidence" not in plain.metadata + wall_clock_metrics = {"wall_time_ms", "processed_tokens_per_s", "processed_output_tokens_per_s"} + assert {k: v for k, v in report.metrics.items() if k not in wall_clock_metrics} == { + k: v for k, v in plain.metrics.items() if k not in wall_clock_metrics + } + assert report.metrics["power_w"] is None + assert report.metrics["power_coverage"] is None + evidence = report.metadata["fpm_query_evidence"] + assert evidence["aggregation"] == "identical_estimates_with_invocation_counts" + estimates = records(evidence) + assert len(estimates) == 3 # prefill emits the first token; two subsequent decode steps + assert all(record["count"] == 3 for record in estimates) + assert sum(record["count"] for record in estimates) == 9 + assert all(record["latency_scale"] == 1.0 for record in estimates) + assert {q["resolution"] for r in estimates for rank in r["estimate"]["ranks"] for q in rank["queries"]} == { + "exact_lookup", + "within_curve_interpolation", + } + native = report.metadata["native_report"] + detail = build_prediction_details(native, ["source"]) + schema = json.loads((Path(__file__).resolve().parents[1] / "docs/cli/prediction-details.schema.json").read_text()) + jsonschema.validate(detail, schema) + jsonschema.validate(build_prediction_details(native, ["time"]), schema) + assert detail["sections"]["source"]["status"] == "available" + operations = [op for phase in detail["sections"]["source"]["phases"] for op in phase["operations"]] + assert all(op["source"] == "silicon" and op["fpm_estimates"] for op in operations) + assert all(op["sol"] is None for phase in native["performance_diagnostics"]["phases"] for op in phase["operations"]) + normalized = _normalize_engine_replay_report(json.loads(json.dumps(native)), include_native_report=False) + assert "native_report" not in normalized.metadata + assert normalized.metadata["fpm_query_evidence"] == evidence + + +@pytest.mark.parametrize("supervised", [False, True]) +def test_public_recommendation_retains_requested_evidence(profile, evidence_systems, monkeypatch, supervised): + # Keep orchestration in this process so the analytical-graph guards remain + # active; estimation, candidate scoring and replay use the real native API. + if not supervised: + monkeypatch.setattr("aisimulate.supervision.in_supervised_process", lambda: True) + monkeypatch.setattr("aisimulate.resources.GuardedRunnerFactory.admit_wave", None) + raw = prediction(profile, evidence_systems) + raw["engine"]["mode"] = "aggregated" + raw["engine"]["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + raw["optimization"] = {"constraints": {"max_candidate_gpus": 2}} + raw["optimizer"] = {"algorithm": "random", "max_trials": 1, "parallelism": 1} + config = CoreRecommendationConfig.model_validate(raw) + result = run_recommendation( + config, + stack="engine", + runner_factory=EngineReplayRunnerFactory(), + show_progress=False, + output_requirements=ReplayOutputRequirements(capture_performance_diagnostics=True), + ) + assert result.counts.feasible == 1 + saved = SweepResult.model_validate_json(result.model_dump_json()) + metadata = saved.candidates[0].provenance.runner_metadata + evidence = metadata["fpm_query_evidence"] + assert records(evidence) + assert all(record["count"] == 3 for record in records(evidence)) + assert saved.candidates[0].metrics["power_w"] is None + + +def test_prediction_cli_saves_direct_query_evidence(profile, evidence_systems, tmp_path, capsys): + import aisimulate.main as cli + + config = tmp_path / "prediction.json" + config.write_text(json.dumps(prediction(profile, evidence_systems))) + output = tmp_path / "output" + assert ( + cli.main( + ["predict", "--config", str(config), "--output-dir", str(output), "--detail", "source", "--format", "json"] + ) + == 0 + ) + displayed = json.loads(capsys.readouterr().out) + assert displayed["details"]["sections"]["source"]["status"] == "available" + saved = json.loads((output / "prediction.json").read_text()) + assert records(saved["fpm_query_evidence"]) + assert saved["details"] == displayed["details"] + + +@pytest.fixture +def native_prediction_report(profile, evidence_systems): + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate(prediction(profile, evidence_systems))) + runner = EngineReplayRunnerFactory().create(0) + try: + report = runner.run(spec, output_requirements=ReplayOutputRequirements(capture_performance_diagnostics=True)) + finally: + runner.close() + return report.metadata["native_report"] + + +@pytest.mark.parametrize( + "path,value", + [ + (("estimate",), {}), + (("estimate", "latency_ms"), float("inf")), + (("estimate", "native_latency_ms"), float("nan")), + (("estimate", "correction_factor"), -1), + (("estimate", "max_rank"), True), + (("estimate", "max_rank"), 0.5), + (("estimate", "ranks"), {}), + (("estimate", "ranks", 0), None), + (("estimate", "ranks", 0, "rank"), -1), + (("estimate", "ranks", 0, "rank"), True), + (("estimate", "ranks", 0, "latency_ms"), True), + (("estimate", "ranks", 0, "prefill_ms"), "1"), + (("estimate", "ranks", 0, "decode_ms"), float("inf")), + (("estimate", "ranks", 0, "decode_baseline_ms"), False), + (("estimate", "ranks", 0, "marginal_decode_ms"), -1), + (("estimate", "ranks", 0, "queries"), {}), + (("estimate", "ranks", 0, "queries", 0), []), + (("estimate", "ranks", 0, "queries", 0, "phase"), "generation"), + (("estimate", "ranks", 0, "queries", 0, "model_path"), 1), + (("estimate", "ranks", 0, "queries", 0, "model_path"), ""), + (("estimate", "ranks", 0, "queries", 0, "decode_baseline"), 0), + (("estimate", "ranks", 0, "queries", 0, "query"), {}), + (("estimate", "ranks", 0, "queries", 0, "query", "batch_size"), True), + (("estimate", "ranks", 0, "queries", 0, "query", "total_prefill_tokens"), -1), + (("estimate", "ranks", 0, "queries", 0, "query", "total_kv_read_tokens"), float("inf")), + (("estimate", "ranks", 0, "queries", 0, "resolution"), "nearest"), + (("estimate", "ranks", 0, "queries", 0, "latency_ms"), float("nan")), + (("estimate", "ranks", 0, "queries", 0, "support"), []), + (("estimate", "ranks", 0, "queries", 0, "support"), {}), + (("estimate", "ranks", 0, "queries", 0, "support", 0), None), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "coordinates"), []), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "coordinates", "batch_size"), -1), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "coordinates", "total_prefill_tokens"), "1"), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "coordinates", "total_kv_read_tokens"), float("nan")), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "latency_ms"), "1"), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "weight"), float("nan")), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "weight"), -0.1), + (("estimate", "ranks", 0, "queries", 0, "support", 0, "weight"), 1.1), + ], +) +def test_prediction_details_reject_malformed_fpm_evidence(native_prediction_report, path, value): + record = native_prediction_report["performance_diagnostics"]["phases"][0]["operations"][0]["fpm_estimates"][0] + for key in path[:-1]: + record = record[key] + record[path[-1]] = value + expected_path = "performance_diagnostics.phases[0].operations[0].fpm_estimates[0]" + "".join( + f"[{key}]" if isinstance(key, int) else f".{key}" for key in path + ) + for section in ("source", "time"): + with pytest.raises(ValueError) as error: + build_prediction_details(native_prediction_report, (section,)) + assert str(error.value).startswith(expected_path) + + +def test_prediction_details_reject_missing_and_unknown_fpm_fields(native_prediction_report): + prefix = ("performance_diagnostics", "phases", 0, "operations", 0, "fpm_estimates", 0) + for path in [ + (), + ("estimate",), + ("estimate", "ranks", 0), + ("estimate", "ranks", 0, "queries", 0), + ("estimate", "ranks", 0, "queries", 0, "query"), + ("estimate", "ranks", 0, "queries", 0, "support", 0), + ("estimate", "ranks", 0, "queries", 0, "support", 0, "coordinates"), + ]: + record = native_prediction_report + for key in (*prefix, *path): + record = record[key] + for field in (*record, "unknown"): + invalid = deepcopy(native_prediction_report) + target = invalid + for key in (*prefix, *path): + target = target[key] + if field == "unknown": + target[field] = None + else: + del target[field] + expected_path = "".join( + f"[{key}]" if isinstance(key, int) else f".{key}" for key in (*prefix, *path, field) + ).lstrip(".") + for section in ("source", "time"): + with pytest.raises(ValueError) as error: + build_prediction_details(invalid, (section,)) + assert str(error.value).startswith(f"{expected_path} must be") + + +def test_prediction_details_preserve_valid_nested_fpm_evidence(profile, evidence_systems, native_prediction_report): + import jsonschema + + config = canonical(profile, evidence_systems) + model = RustForwardPassPerfModel.best_available(config) + estimates = [ + model.estimate_forward_pass_detailed(metrics) + for metrics in (prefill(), prefill(6), prefill(6, 2), decode(), {}) + ] + config["estimation_mode"] = "fpm_regression" + estimates.append(RustForwardPassPerfModel.best_available(config).estimate_forward_pass_detailed(decode())) + operation = native_prediction_report["performance_diagnostics"]["phases"][0]["operations"][0] + operation["fpm_estimates"] = [{"estimate": estimate, "count": 1, "latency_scale": 1.0} for estimate in estimates] + original = deepcopy(native_prediction_report) + details = build_prediction_details(native_prediction_report, ("source", "time")) + schema = json.loads((Path(__file__).resolve().parents[1] / "docs/cli/prediction-details.schema.json").read_text()) + jsonschema.validate(details, schema) + for section in (details["sections"]["source"], details["sections"]["time"]["diagnostics"]): + published = section["phases"][0]["operations"][0]["fpm_estimates"] + assert published == operation["fpm_estimates"] + published[0]["estimate"]["ranks"].clear() + assert native_prediction_report == original + + # FPM evidence remains optional for pre-existing diagnostics producers. + del operation["fpm_estimates"] + details = build_prediction_details(native_prediction_report, ("source", "time")) + jsonschema.validate(details, schema) + for section in (details["sections"]["source"], details["sections"]["time"]["diagnostics"]): + assert "fpm_estimates" not in section["phases"][0]["operations"][0] diff --git a/tests/test_resource_supervision.py b/tests/test_resource_supervision.py index 9dea4e6c6..4dfafc714 100644 --- a/tests/test_resource_supervision.py +++ b/tests/test_resource_supervision.py @@ -516,11 +516,17 @@ def test_supervisor_argument_and_output_setup_do_not_import_runtime(): [ sys.executable, "-c", - ( - "import sys; import aisimulate.supervision, aisimulate.cli_args, aisimulate.output; " - "assert not any(name in sys.modules for name in " - "('aisimulate._runtime', 'aisimulate.sweeper', 'numpy', 'pandas', 'aisimulate_core'))" - ), + """ +import sys +import aisimulate.supervision, aisimulate.cli_args, aisimulate.output + +runtime = ('aisimulate._runtime', 'aisimulate.sweeper', 'numpy', 'pandas') +assert not any(name == root or name.startswith(root + '.') for name in sys.modules for root in runtime) +# Profile validation needs only lightweight metadata before resource admission. +allowed_core = {'aisimulate_core', 'aisimulate_core.fpm_profile', 'aisimulate_core.quantization'} +loaded_core = {name for name in sys.modules if name == 'aisimulate_core' or name.startswith('aisimulate_core.')} +assert loaded_core <= allowed_core, loaded_core - allowed_core +""", ], check=True, timeout=15, diff --git a/tests/test_unified_traffic_runtime.py b/tests/test_unified_traffic_runtime.py index 31474febb..bb06d8db1 100644 --- a/tests/test_unified_traffic_runtime.py +++ b/tests/test_unified_traffic_runtime.py @@ -533,22 +533,59 @@ def test_prediction_spec_lowers_fpm_forward_model_into_canonical_timing() -> Non "/artifacts/reviewed-fpm.parquet" ) assert "aic_forward_model" not in deployment.agg_engine_args - assert deployment.performance_model_metadata["aggregated"]["config"]["forward_model"] == "fpm" - assert ( - deployment.performance_model_metadata["aggregated"]["config"]["fpm_parquet_path"] - == "/artifacts/reviewed-fpm.parquet" + config = deployment.performance_model_metadata["aggregated"]["config"] + assert all(value == timing["config"][name] for name, value in config.items()) + assert config["estimation_mode"] == "fpm_interpolation" + assert config["fallback_policy"] == "deny" + assert config["worker_type"] == "aggregated" + assert "forward_model" not in config + assert config["estimator_config"]["fpm_interpolation"]["fpm_parquet_path"] == ("/artifacts/reviewed-fpm.parquet") + assert config["model"] == parsed.engine.model + assert config["system"] == parsed.engine.hardware + assert config["backend"] == parsed.engine.backend + assert config["backend_version"] == parsed.engine.backend_version + assert (config["tp"], config["pp"], config["attention_dp"], config["moe_tp_size"], config["moe_ep_size"]) == ( + 4, + 1, + 1, + 4, + 1, ) -def test_prediction_spec_lowers_default_timing_into_canonical_auto_selection() -> None: +@pytest.mark.parametrize("explicit_mode", [False, True]) +def test_prediction_spec_omits_the_forward_model_rank_field_for_op_level(monkeypatch, explicit_mode: bool) -> None: + from aisimulate_core.sdk import RustForwardPassPerfModel + + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") engine = _fpm_engine() + engine["backend_version"] = "0.24.0" # A shipped op-level database, independent of external FPM data. engine["workers"]["aggregated"]["timing"] = {"type": "default"} + if explicit_mode: + engine["workers"]["aggregated"]["timing"]["estimation_mode"] = "op_level" parsed = CorePredictionConfig.model_validate({"engine": engine}) deployment = prediction_to_replay_spec(parsed).backend_deployment assert "aic_forward_model" not in deployment.agg_engine_args - assert deployment.agg_engine_args["timing_model"]["config"]["estimation_mode"] == "auto" - assert deployment.performance_model_metadata["aggregated"]["config"]["forward_model"] == "op_level" + config = deployment.performance_model_metadata["aggregated"]["config"] + requested_mode = "op_level" if explicit_mode else "auto" + assert "forward_model" not in config + assert config["estimation_mode"] == requested_mode + assert config["model"] == engine["model"] + assert config["system"] == engine["hardware"] + assert config["backend_version"] == engine["backend_version"] + assert config["worker_type"] == "aggregated" + assert config["tp"] == config["moe_tp_size"] == 4 + model = RustForwardPassPerfModel.best_available(config) + try: + provenance = model.diagnostics()["provenance"] + finally: + model.close() + assert provenance["requested_estimation_mode"] == requested_mode + assert provenance["selected_estimation_mode"] == "op_level" + assert provenance["config"]["estimation_mode"] == "op_level" + for field in ("model", "system", "backend", "backend_version", "worker_type", "tp", "moe_tp_size"): + assert provenance["config"][field] == config[field] def test_prediction_spec_lowers_forward_model_per_role_in_disaggregated_mode() -> None: @@ -573,8 +610,24 @@ def test_prediction_spec_lowers_forward_model_per_role_in_disaggregated_mode() - assert prefill["worker_type"] == "prefill" assert decode["estimation_mode"] == "fpm_interpolation" assert decode["worker_type"] == "decode" - assert deployment.performance_model_metadata["prefill"]["config"]["forward_model"] == "op_level" - assert deployment.performance_model_metadata["decode"]["config"]["forward_model"] == "fpm" + for role, identity in (("prefill", prefill), ("decode", decode)): + config = deployment.performance_model_metadata[role]["config"] + assert all(value == identity[name] for name, value in config.items()) + assert config["estimation_mode"] == ("auto" if role == "prefill" else "fpm_interpolation") + assert config["worker_type"] == role + assert "forward_model" not in config + assert config["fallback_policy"] == "deny" + assert config["model"] == engine["model"] + assert config["system"] == engine["hardware"] + assert config["backend"] == engine["backend"] + assert config["backend_version"] == engine["backend_version"] + assert (config["tp"], config["pp"], config["attention_dp"], config["moe_tp_size"], config["moe_ep_size"]) == ( + 4, + 1, + 1, + 4, + 1, + ) _SMALL_TRAFFIC = {