diff --git a/AGENTS.md b/AGENTS.md index 492da2324..f6e84ca9a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,8 +8,9 @@ When asked to onboard a model for FPM simulation on designated hardware, follow [Onboard with an agent](docs/fpm-self-service.md#onboard-with-an-agent) in the FPM self-service guide. Use the checkout's `aisimulate onboard` CLI and current help. Follow its six stages: inspect the model and target; choose the worker and -workload; derive, review and save the profile; plan collection; collect and verify -data; run prediction/recommendation and report. Start by asking only for a missing +review runtime/collection limits; derive, review and save the profile; plan +collection; collect and verify data; validate replay and run ordinary +prediction/recommendation. Start by asking only for a missing Hugging Face model ID (`organization/model-name`) and target GPU platform. Accept an already supplied local config, profile or checkpoint path instead of requiring a Hub ID; do not ask for both an ID and a config upfront. For a Hub ID, retrieve @@ -17,7 +18,7 @@ the config using available Hub access as described in stage 1. Reuse supplied facts and inspect the config/profile before asking for derivable metadata; defer other questions to their stage and help the user choose the worker topology. At stage 2, inspect `aisimulate onboard init --model-config PATH --suggest-parallel` -with the actual checkpoint/revision, runtime, GPU, interconnect and workload +with the actual checkpoint/revision, runtime, GPU, interconnect and collection options. This read-only JSON preview reports model/hardware-aware TP choices and MoE TP/DEP/TEP alternatives, exact flags, resource sources and missing inputs. Ask for unresolved shared precision/layout facts rather than guessing. Only a @@ -34,6 +35,29 @@ interconnect characteristics; verify actual collection resources before executio Generated predict/recommend configs validate one worker. Deployment replicas and optimization GPU budgets belong to ordinary predict/recommend configurations. +Review context, scheduler and prefill capture limits independently of validation +traffic. Fresh config/profile defaults use the smaller of the declared context +and 256,000 tokens, profile scheduler bounds or 8,192 tokens/256 sequences, and a +2,048-token prefill CUDA graph capture limit. Explain and allow edits to these +initial policies; they establish neither capacity nor timing coverage. Do not +require fixed input/output lengths, concurrency, TTFT or TPOT during intake. +Those flags customize optional synthetic examples only. Maximum sequences does +not reserve maximum context for every sequence. Dynamo self-benchmark owns the +grid generated from capture sizes and runtime bounds; do not invent a separate +AgentX collection grid. + +After verifying the formal data pair, use `aisimulate onboard validate-fpm` +with a local Weka trace and a separate validation output directory. Follow the +guide's pinned AgentX reference and current cold aggregated, one-lane, HBM-only, +non-speculative scope. Preserve target model projection and strict direct FPM +with denied fallback. Read `validation.json` together with native query coverage +and request completion evidence. Coverage counts native lookup resolutions, +including interpolation and failures; cached timing reuse is not another query. +Missing timing stops replay and preserves partial evidence, which cannot certify +the remainder of the trace. Report coverage and silicon accuracy separately. +Reuse verified v2 collection plans for validation-only changes; legacy v1 plans +need a new directory as described in the guide. + Do not reject a checkpoint solely because it is multimodal. Read its unambiguous `text_config`, or its flat text-decoder fields, and explain that FPM models only the text decoder. Multimodal encoders, projectors, preprocessing and other diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 99025ae6d..b52fc98c1 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -49,7 +49,8 @@ pub use perfmodel::{ }; pub use perfmodel::{ CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, FpmInterpolationMethod, - OpLevelConfig, RegressionFitConfig, RegressionFitKind, + FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQueryPurpose, OpLevelConfig, + RegressionFitConfig, RegressionFitKind, }; #[cfg(feature = "python")] diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index fb0662a40..b92d4b740 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -20,8 +20,10 @@ use std::sync::Arc; +use crate::FpmQueryCoverage; use crate::common::enums::{DatabaseMode, TransferPolicy}; use crate::common::error::AicError; +use crate::operators::RuntimeContext; use crate::operators::base::PerformanceResult; use crate::operators::{FpmForwardOp, FpmPhase, Op}; use crate::perf_database::PerfDatabase; @@ -585,6 +587,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_coverage(runtime, None) + } + + fn run_context_phase_with_coverage( + &self, + runtime: &RuntimeConfig, + coverage: Option<&mut FpmQueryCoverage>, + ) -> Result { // Python raises `ValueError` when `effective_isl <= 0`; mirror that. if runtime.prefix >= runtime.isl { return Err(AicError::InvalidEngineConfig(format!( @@ -593,6 +603,33 @@ impl Engine { ))); } let effective_isl = runtime.isl - runtime.prefix; + if let Some(coverage) = coverage { + let Some((prefill, _, ctx_tail, gen_tail)) = self.fpm_split() else { + return Err(AicError::InvalidEngineConfig( + "query coverage requires a direct FPM engine".into(), + )); + }; + if !ctx_tail.is_empty() || !gen_tail.is_empty() || self.nextn > 0 { + return Err(AicError::InvalidEngineConfig( + "query coverage does not support speculative FPM".into(), + )); + } + return prefill + .query_with_coverage( + &self.db, + &RuntimeContext { + batch_size: runtime.batch_size, + beam_width: 1, + s: effective_isl, + prefix: runtime.prefix, + num_tokens: runtime.batch_size * effective_isl, + seq_imbalance_correction_scale: runtime.seq_imbalance_correction_scale, + ..Default::default() + }, + Some(coverage), + ) + .map(|result| result.latency_ms); + } run_context_ops( &self.context_ops, &self.db, @@ -688,6 +725,16 @@ impl Engine { /// Thin shim over [`Self::run_static`] with `mode=Context` (osl is /// irrelevant for the context phase, so it is fixed at 1). pub fn predict_prefill_latency(&self, bs: u32, isl: u32, prefix: u32) -> Result { + self.predict_prefill_latency_with_coverage(bs, isl, prefix, None) + } + + pub(crate) fn predict_prefill_latency_with_coverage( + &self, + bs: u32, + isl: u32, + prefix: u32, + coverage: Option<&mut FpmQueryCoverage>, + ) -> Result { let rt = RuntimeConfig { batch_size: bs, isl, @@ -695,9 +742,7 @@ impl Engine { prefix, ..Default::default() }; - Ok(self - .run_static(&rt, StaticMode::Context, DEFAULT_STATIC_STRIDE)? - .total_ms) + self.run_context_phase_with_coverage(&rt, coverage) } /// Mocker H2: decode-step latency in ms. Pure-Rust inherent method (no @@ -725,14 +770,26 @@ impl Engine { batch_size: u32, total_past_kv_tokens: u32, ) -> Result { - self.forward_pass_time_ms(&[ForwardPassMetrics { - scheduled_requests: crate::ScheduledRequestMetrics { - num_decode_requests: batch_size, - sum_decode_kv_tokens: total_past_kv_tokens, + self.predict_decode_latency_total_with_coverage(batch_size, total_past_kv_tokens, None) + } + + pub(crate) fn predict_decode_latency_total_with_coverage( + &self, + batch_size: u32, + total_past_kv_tokens: u32, + coverage: Option<&mut FpmQueryCoverage>, + ) -> Result { + self.forward_pass_time_ms_with_coverage( + &[ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_decode_requests: batch_size, + sum_decode_kv_tokens: total_past_kv_tokens, + ..Default::default() + }, ..Default::default() - }, - ..Default::default() - }]) + }], + coverage, + ) } /// Highest decode KV-read total covered by a compiled FPM engine. @@ -1971,6 +2028,14 @@ impl Engine { pub fn forward_pass_time_ms( &self, metrics_by_rank: &[ForwardPassMetrics], + ) -> Result { + self.forward_pass_time_ms_with_coverage(metrics_by_rank, None) + } + + pub(crate) fn forward_pass_time_ms_with_coverage( + &self, + metrics_by_rank: &[ForwardPassMetrics], + mut coverage: Option<&mut FpmQueryCoverage>, ) -> Result { if metrics_by_rank.is_empty() { return Err(AicError::InvalidForwardPassMetrics( @@ -1982,7 +2047,7 @@ impl Engine { } let mut max_latency = 0.0_f64; for metrics in metrics_by_rank { - let rank_latency = self.rank_latency_ms(metrics)?; + let rank_latency = self.rank_latency_ms(metrics, coverage.as_deref_mut())?; if rank_latency > max_latency { max_latency = rank_latency; } @@ -1996,7 +2061,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 coverage: Option<&mut FpmQueryCoverage>, + ) -> 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 @@ -2033,24 +2102,26 @@ impl Engine { let mut total = 0.0_f64; if has_prefill { total += prefill_op - .query_totals( + .query_totals_with_coverage( &self.db, &[ sched.num_prefill_requests as f64, sched.sum_prefill_tokens as f64, sched.sum_prefill_kv_tokens as f64, ], + coverage.as_deref_mut(), )? .latency_ms; } if has_decode { let decode_ms = decode_op - .query_totals( + .query_totals_with_coverage( &self.db, &[ sched.num_decode_requests as f64, sched.sum_decode_kv_tokens as f64, ], + coverage.as_deref_mut(), )? .latency_ms; if has_prefill { @@ -2058,10 +2129,11 @@ impl Engine { // `_get_fpm_mix_step_latency` (counts already packed, no // `(nextn + 1)` — speculative FPM was rejected above). let baseline_ms = decode_op - .query_pass_baseline( + .query_pass_baseline_with_coverage( &self.db, sched.num_decode_requests, sched.sum_decode_kv_tokens as f64, + coverage.as_deref_mut(), )? .latency_ms; total += (decode_ms - baseline_ms).max(0.0); diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index a71fabce5..920031ead 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -274,6 +274,16 @@ impl ForwardPassPerfModelConfig { } pub(crate) fn validate(&self) -> Result<(), AicError> { + if self.estimator_config.fpm_interpolation.collect_coverage + && (self.estimation_mode != EstimationMode::FpmInterpolation + || self.fallback_policy != ForwardPassFallbackPolicy::Deny + || self.estimator_config.fpm_interpolation.method + != super::FpmInterpolationMethod::Direct) + { + return Err(invalid_config( + "collect_coverage requires estimation_mode='fpm_interpolation', method='direct', and fallback_policy='deny'", + )); + } if self.model.trim().is_empty() { return Err(invalid_config("model cannot be empty")); } @@ -436,6 +446,66 @@ mod tests { ); } + #[test] + fn coverage_controls_preserve_default_serialization_and_require_strict_direct() { + let default = serde_json::to_value(EstimatorConfig::default()).unwrap(); + assert_eq!( + default["fpm_interpolation"], + serde_json::json!({"method": "auto"}) + ); + let cfg = config(serde_json::json!({ + "estimation_mode": "fpm_interpolation", + "estimator_config": {"fpm_interpolation": {"method": "direct", "collect_coverage": true}} + })); + cfg.validate().unwrap(); + let reloaded: ForwardPassPerfModelConfig = + serde_json::from_value(serde_json::to_value(&cfg).unwrap()).unwrap(); + assert_eq!(cfg, reloaded); + for mode in [ + EstimationMode::Auto, + EstimationMode::OpLevel, + EstimationMode::FpmRegression, + ] { + let mut invalid = cfg.clone(); + invalid.estimation_mode = mode; + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("collect_coverage") + ); + } + for method in [ + super::super::FpmInterpolationMethod::Auto, + super::super::FpmInterpolationMethod::Sol, + ] { + let mut invalid = cfg.clone(); + invalid.estimator_config.fpm_interpolation.method = method; + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("collect_coverage") + ); + } + for fallback in [ + ForwardPassFallbackPolicy::Allow, + ForwardPassFallbackPolicy::LegacyRegression, + ] { + let mut invalid = cfg.clone(); + invalid.fallback_policy = fallback; + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("collect_coverage") + ); + } + } + #[test] fn engine_controls_round_trip_with_identity_and_policy() { let cfg = config(serde_json::json!({ diff --git a/crates/core/src/perfmodel/fpm/coverage.rs b/crates/core/src/perfmodel/fpm/coverage.rs new file mode 100644 index 000000000..2eecba262 --- /dev/null +++ b/crates/core/src/perfmodel/fpm/coverage.rs @@ -0,0 +1,175 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Bounded evidence from actual direct-FPM lookups. This is not a replay +//! completion or prediction-accuracy assessment. + +use std::collections::BTreeMap; +use std::sync::{Mutex, MutexGuard}; + +use serde::{Deserialize, Serialize}; + +use crate::AicError; +use crate::operators::fpm_forward::{FpmForwardOp, FpmPhase}; +use crate::perf_database::fpm_forward::{ + FPM_CELL_MATCH_COLUMNS, FPM_DECODE_AXES, FPM_PREFILL_AXES, +}; + +const GAP_LIMIT: usize = 128; + +/// Counts of native table lookup resolutions, not requests or iterations. +/// Reusing a cached timing does not perform another native lookup. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct FpmQueryCoverageCounts { + pub measured: u64, + pub interpolated: u64, + pub unsupported: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmQueryPurpose { + ForwardPass, + /// A mixed pass subtracts the measured decode curve floor, or an + /// interpolation of floors, using the same covered rows as decode. + MixedDecodeBaseline, +} + +/// One unsupported query, retained after the estimation call returns its error. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmQueryGap { + pub phase: String, + pub purpose: FpmQueryPurpose, + pub model_path: String, + pub cell_identity: BTreeMap, + pub cell_ids: Vec, + pub coordinates: BTreeMap, + pub reason: String, + pub occurrences: u64, +} + +/// Cumulative evidence for one model instance. A snapshot cannot certify that +/// replay completed: the replay consumer must report completion separately. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct FpmQueryCoverage { + /// Explicit unit prevents callers from interpreting these as replay counts. + pub counting_unit: String, + pub queries: FpmQueryCoverageCounts, + pub prefill: FpmQueryCoverageCounts, + pub decode: FpmQueryCoverageCounts, + pub mixed_decode_baseline: FpmQueryCoverageCounts, + pub gaps: Vec, + pub gap_limit: usize, + /// Failed lookup occurrences omitted after the distinct-gap limit. + pub omitted_gap_queries: u64, +} + +impl Default for FpmQueryCoverage { + fn default() -> Self { + Self { + counting_unit: "native_lookup_resolutions".into(), + queries: Default::default(), + prefill: Default::default(), + decode: Default::default(), + mixed_decode_baseline: Default::default(), + gaps: Vec::new(), + gap_limit: GAP_LIMIT, + omitted_gap_queries: 0, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum DirectFpmResolution { + Measured, + Interpolated, +} + +impl FpmQueryCoverageCounts { + fn record(&mut self, result: Result) { + match result { + Ok(DirectFpmResolution::Measured) => self.measured += 1, + Ok(DirectFpmResolution::Interpolated) => self.interpolated += 1, + Err(_) => self.unsupported += 1, + } + } +} + +impl FpmQueryCoverage { + pub(crate) fn record( + &mut self, + op: &FpmForwardOp, + cell_ids: &[String], + coords: &[f64], + purpose: FpmQueryPurpose, + result: Result, + ) { + self.queries.record(result); + match purpose { + FpmQueryPurpose::MixedDecodeBaseline => &mut self.mixed_decode_baseline, + FpmQueryPurpose::ForwardPass => match op.phase { + FpmPhase::Prefill => &mut self.prefill, + FpmPhase::Decode => &mut self.decode, + }, + } + .record(result); + let Err(error) = result else { return }; + let axes: &[&str] = match op.phase { + FpmPhase::Prefill => &FPM_PREFILL_AXES, + FpmPhase::Decode => &FPM_DECODE_AXES, + }; + let mut gap = FpmQueryGap { + phase: op.phase.as_str().into(), + purpose, + model_path: op.model_path.clone(), + cell_identity: FPM_CELL_MATCH_COLUMNS + .iter() + .zip(&op.match_identity) + .map(|(name, value)| ((*name).into(), value.clone())) + .collect(), + cell_ids: cell_ids.to_vec(), + coordinates: axes + .iter() + .zip(coords) + .map(|(k, v)| ((*k).into(), *v)) + .collect(), + reason: error.to_string(), + occurrences: 0, + }; + if let Some(existing) = self.gaps.iter_mut().find(|entry| { + entry.phase == gap.phase + && entry.purpose == gap.purpose + && entry.model_path == gap.model_path + && entry.cell_identity == gap.cell_identity + && entry.cell_ids == gap.cell_ids + && entry.coordinates == gap.coordinates + && entry.reason == gap.reason + }) { + existing.occurrences += 1; + } else if self.gaps.len() < GAP_LIMIT { + gap.occurrences = 1; + self.gaps.push(gap); + } else { + self.omitted_gap_queries += 1; + } + } +} + +#[derive(Debug, Default)] +pub(crate) struct FpmCoverageState(Mutex); + +impl Clone for FpmCoverageState { + fn clone(&self) -> Self { + // Model clones may be used for another replay/candidate. The loaded + // engine is shared, but diagnostics must start empty for each clone. + Self::default() + } +} + +impl FpmCoverageState { + pub(crate) fn lock(&self) -> Result, AicError> { + self.0.lock().map_err(|_| { + AicError::InvalidEngineConfig("FPM query coverage accumulator was poisoned".into()) + }) + } +} diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs index 7332ab7d9..b26940824 100644 --- a/crates/core/src/perfmodel/fpm/estimator.rs +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -26,6 +26,14 @@ pub struct OpLevelConfig {} #[serde(default, deny_unknown_fields)] pub struct FpmInterpolationConfig { pub method: FpmInterpolationMethod, + /// Record bounded query coverage on the returned canonical model. + /// Requires explicit direct FPM with fallback denied. + #[serde(skip_serializing_if = "is_false")] + pub collect_coverage: bool, +} + +fn is_false(value: &bool) -> bool { + !value } /// Construction-time interpolation selection. Native operators receive only diff --git a/crates/core/src/perfmodel/fpm/mod.rs b/crates/core/src/perfmodel/fpm/mod.rs index 65144f25b..9a151b378 100644 --- a/crates/core/src/perfmodel/fpm/mod.rs +++ b/crates/core/src/perfmodel/fpm/mod.rs @@ -24,6 +24,7 @@ mod config; mod correction; +pub(crate) mod coverage; mod estimator; mod metrics; mod model; @@ -38,6 +39,7 @@ pub use config::{ EstimationMode, ForwardPassFallbackPolicy, ForwardPassPerfModelConfig, ForwardPassSpeculationConfig, }; +pub use coverage::{FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQueryPurpose}; pub use estimator::*; pub(crate) use metrics::validate_forward_pass_metrics; pub use metrics::{FPM_VERSION, ForwardPassMetrics, QueuedRequestMetrics, ScheduledRequestMetrics}; diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index ef8a41ddd..c4eb27d7a 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -18,6 +18,7 @@ use crate::{AicError, ForwardPassMetrics}; use super::config::{EstimationMode, ForwardPassFallbackPolicy, ForwardPassPerfModelConfig}; use super::correction::CorrectionBuckets; +use super::coverage::{FpmCoverageState, FpmQueryCoverage}; use super::metrics::validate_forward_pass_metrics; use super::options::{ForwardPassPerfOptions, validate_regression_options}; use super::regression::BucketedRegression; @@ -244,6 +245,7 @@ pub struct ForwardPassPerfModel { last_warning: Option, provenance: Option, is_correction_enabled: bool, + pub(super) query_coverage: Option, } #[derive(Clone, Debug)] @@ -278,6 +280,7 @@ impl ForwardPassPerfModel { last_warning: None, provenance: None, is_correction_enabled: true, + query_coverage: None, } } @@ -305,6 +308,7 @@ impl ForwardPassPerfModel { last_warning: None, provenance: None, is_correction_enabled: true, + query_coverage: None, }) } @@ -364,6 +368,11 @@ impl ForwardPassPerfModel { candidate.estimator_config.correction_options(), ); model.is_correction_enabled = candidate.estimator_config.correction.enabled; + model.query_coverage = candidate + .estimator_config + .fpm_interpolation + .collect_coverage + .then(FpmCoverageState::default); model.last_warning = (!failures.is_empty()).then(|| failures.join("; ")); model.provenance = Some(ForwardPassPerfProvenance { config: candidate, @@ -418,7 +427,13 @@ impl ForwardPassPerfModel { let Some(feature) = IterationFeatures::from_metrics(metrics_by_rank)? else { return Ok(Some(0.0)); }; - let native = engine.forward_pass_time_ms(metrics_by_rank)?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + let native = engine + .forward_pass_time_ms_with_coverage(metrics_by_rank, coverage.as_deref_mut())?; let corrected = native * if self.is_correction_enabled { corrections @@ -480,6 +495,7 @@ impl ForwardPassPerfModel { mode, options, is_correction_enabled, + query_coverage, .. } = self; for metrics_by_rank in iterations { @@ -496,7 +512,14 @@ impl ForwardPassPerfModel { else { continue; }; - let native = engine.forward_pass_time_ms(metrics_by_rank)?; + let mut coverage = query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + let native = engine.forward_pass_time_ms_with_coverage( + metrics_by_rank, + coverage.as_deref_mut(), + )?; corrections .store_mut(observation.feature.workload_kind) .add_observation(observation.feature.x, observation.wall_time_ms, native); @@ -630,6 +653,75 @@ impl ForwardPassPerfModel { &self.options } + /// Uncorrected static prefill latency, preserving the compiled engine's + /// existing `(batch, full input length, cached prefix)` contract. Unlike + /// telemetry estimation this does not apply learned online correction. + /// Enabled coverage records the actual native lookup on this model. + pub fn predict_prefill_latency( + &self, + batch_size: u32, + isl: u32, + prefix: u32, + ) -> Result { + let engine = self.require_native_engine()?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + engine.predict_prefill_latency_with_coverage( + batch_size, + isl, + prefix, + coverage.as_deref_mut(), + ) + } + + /// Uncorrected decode latency using exact past-KV totals, excluding the + /// current input token per request. Shares the existing engine contract. + pub fn predict_decode_latency_total( + &self, + batch_size: u32, + total_past_kv_tokens: u32, + ) -> Result { + let engine = self.require_native_engine()?; + let mut coverage = self + .query_coverage + .as_ref() + .map(FpmCoverageState::lock) + .transpose()?; + engine.predict_decode_latency_total_with_coverage( + batch_size, + total_past_kv_tokens, + coverage.as_deref_mut(), + ) + } + + /// Largest collected decode KV total for the selected FPM cell. Reading + /// this bound is not a timing lookup and does not add coverage evidence. + pub fn fpm_decode_kv_ceiling(&self) -> Result, AicError> { + self.require_native_engine()?.fpm_decode_kv_ceiling() + } + + /// Snapshot of actual native direct-FPM lookup resolutions, including + /// failed lookups. `None` means collection was disabled. Counts do not + /// include timings reused by an external cache, or calls on a separate + /// low-level engine handle. Replay completion must be checked separately. + pub fn fpm_query_coverage(&self) -> Result, AicError> { + self.query_coverage + .as_ref() + .map(|state| state.lock().map(|value| value.clone())) + .transpose() + } + + fn require_native_engine(&self) -> Result, AicError> { + self.native_engine().ok_or_else(|| { + AicError::InvalidEngineConfig( + "static timing requires a native forward-pass estimator".into(), + ) + }) + } + /// 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. diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index bcc04a8c5..90f3bac9e 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -140,6 +140,237 @@ fn native_model(options: ForwardPassPerfOptions) -> ForwardPassPerfModel { ForwardPassPerfModel::from_engine(fixture_engine(), options) } +/// Synthetic measured timings with no SOL operations. The fixture's expected +/// answers are its literal measurements and raw linear midpoints. +fn direct_coverage_model(dir: &std::path::Path) -> ForwardPassPerfModel { + use crate::operators::fpm_forward::{FpmForwardOp, FpmInterpolation, FpmPhase}; + use crate::perf_database::fpm_forward::tests::{RowSpec, default_identity, write_pair}; + let mut rows = Vec::new(); + for (tokens, latency) in [(100, 10.0), (200, 20.0)] { + for (kv, bump) in [(0, 0.0), (200, 4.0)] { + rows.push(RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: tokens, + total_kv_read_tokens: kv, + latency_ms: latency + bump, + ..Default::default() + }); + } + } + for (batch, kv, latency) in [ + (8, 64, 8.0), + (8, 128, 16.0), + (9, 64, 9.0), + (9, 128, 17.0), + (16, 64, 16.0), + (16, 128, 24.0), + (17, 64, 170.0), + (17, 128, 178.0), + ] { + rows.push(RowSpec { + batch_size: batch, + total_kv_read_tokens: kv, + latency_ms: latency, + kv_seed_regime: Some("real_kv"), + ..Default::default() + }); + } + rows.push(RowSpec { + batch_size: 8, + total_kv_read_tokens: 256, + latency_ms: 99.0, + kv_seed_regime: Some("fake_fallback"), + ..Default::default() + }); + write_pair(dir, &rows); + let mut db = PerfDatabase::load(&systems_root(), "b200_sxm", "vllm", "0.24.0").unwrap(); + db.set_fpm_forward_for_test(crate::perf_database::FpmForwardTable::new( + dir.to_path_buf(), + "b200_sxm", + "vllm", + "0.25.1", + )); + let op = |phase: FpmPhase| { + Op::FpmForward(FpmForwardOp { + 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![], + }) + }; + let spec = EngineSpec::new( + fixture_engine_config(), + vec![op(FpmPhase::Prefill)], + vec![op(FpmPhase::Decode)], + ); + let engine = Engine::build(spec, Arc::new(db)).unwrap(); + let mut model = ForwardPassPerfModel::from_engine(Arc::new(engine), Default::default()); + model.query_coverage = Some(Default::default()); + model +} + +#[test] +fn direct_coverage_records_real_resolutions_and_mixed_baselines() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + assert_eq!(model.predict_prefill_latency(1, 100, 0).unwrap(), 10.0); + // At 150 new / 100 cached tokens, each axis is halfway between rows. + assert_eq!(model.predict_prefill_latency(1, 250, 100).unwrap(), 17.0); + assert_eq!(model.predict_decode_latency_total(8, 64).unwrap(), 8.0); + assert_eq!(model.predict_decode_latency_total(8, 96).unwrap(), 12.0); + let mixed = ForwardPassMetrics { + scheduled_requests: ScheduledRequestMetrics { + num_prefill_requests: 1, + sum_prefill_tokens: 100, + num_decode_requests: 12, + sum_decode_kv_tokens: 96, + ..Default::default() + }, + ..Default::default() + }; + // Same graph bracket rows 9/16 give decode=16, floor=12. Prefill=10. + assert_eq!( + model.estimate_forward_pass_time_ms(&[mixed]).unwrap(), + Some(14.0) + ); + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.counting_unit, "native_lookup_resolutions"); + assert_eq!(report.queries.measured, 3); + assert_eq!(report.queries.interpolated, 4); + assert_eq!(report.prefill.measured, 2); + assert_eq!(report.prefill.interpolated, 1); + assert_eq!(report.decode.measured, 1); + assert_eq!(report.decode.interpolated, 2); + assert_eq!(report.mixed_decode_baseline.interpolated, 1); + assert_eq!(report.queries.unsupported, 0); + assert!(report.gaps.is_empty()); +} + +#[test] +fn direct_coverage_keeps_partial_failures_and_excludes_fake_rows() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + let mixed = ForwardPassMetrics { + scheduled_requests: ScheduledRequestMetrics { + num_prefill_requests: 1, + sum_prefill_tokens: 100, + num_decode_requests: 8, + sum_decode_kv_tokens: 256, + ..Default::default() + }, + ..Default::default() + }; + let error = model.estimate_forward_pass_time_ms(&[mixed]).unwrap_err(); + assert!(error.to_string().contains("genuine FPM points")); + assert!(model.predict_decode_latency_total(8, 256).is_err()); + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.prefill.measured, 1); + assert_eq!(report.decode.unsupported, 2); + assert_eq!(report.mixed_decode_baseline, Default::default()); + assert_eq!(report.gaps.len(), 1); + let gap = &report.gaps[0]; + assert_eq!(gap.occurrences, 2); + assert_eq!(gap.phase, "decode"); + assert_eq!(gap.model_path, "org/model-a"); + assert_eq!(gap.cell_identity["tp"], "4"); + assert!(!gap.cell_ids.is_empty()); + assert_eq!(gap.coordinates["batch_size"], 8.0); + assert_eq!(gap.coordinates["total_kv_read_tokens"], 256.0); + assert_eq!(gap.reason, error.to_string()); + assert_eq!(model.fpm_decode_kv_ceiling().unwrap(), Some(128)); + assert_eq!(model.fpm_query_coverage().unwrap().unwrap(), report); +} + +#[test] +fn direct_coverage_is_bounded_and_owned_by_each_model() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + model.predict_prefill_latency(1, 100, 0).unwrap(); + let clone = model.clone(); + assert_eq!( + clone.fpm_query_coverage().unwrap().unwrap().queries, + Default::default() + ); + assert_eq!( + model + .fpm_query_coverage() + .unwrap() + .unwrap() + .queries + .measured, + 1 + ); + assert!(clone.predict_prefill_latency(3, 100, 0).is_err()); + assert_eq!( + model + .fpm_query_coverage() + .unwrap() + .unwrap() + .queries + .unsupported, + 0 + ); + for kv in 1000..1140 { + assert!(model.predict_decode_latency_total(8, kv).is_err()); + } + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.queries.unsupported, 140); + assert_eq!(report.gaps.len(), report.gap_limit); + assert_eq!(report.omitted_gap_queries, 12); + // A retained gap continues counting duplicates even after the cap. + assert!(model.predict_decode_latency_total(8, 1000).is_err()); + let report = model.fpm_query_coverage().unwrap().unwrap(); + assert_eq!(report.gaps[0].occurrences, 2); + assert_eq!(report.omitted_gap_queries, 12); +} + +#[test] +fn direct_coverage_static_wrappers_preserve_engine_answers_and_errors() { + let tmp = tempfile::tempdir().unwrap(); + let model = direct_coverage_model(tmp.path()); + let engine = model.native_engine().unwrap(); + for (batch, isl, prefix) in [ + (1, 100, 0), + (1, 250, 100), + (1, 400, 200), + (0, 100, 0), + (1, 100, 100), + (3, 100, 0), + ] { + assert_eq!( + model + .predict_prefill_latency(batch, isl, prefix) + .map_err(|err| err.to_string()), + engine + .predict_prefill_latency(batch, isl, prefix) + .map_err(|err| err.to_string()), + ); + } + for (batch, kv) in [(0, 0), (8, 64), (8, 96), (12, 96), (8, 256)] { + assert_eq!( + model + .predict_decode_latency_total(batch, kv) + .map_err(|err| err.to_string()), + engine + .predict_decode_latency_total(batch, kv) + .map_err(|err| err.to_string()), + ); + } + assert_eq!( + model.fpm_decode_kv_ceiling().unwrap(), + engine.fpm_decode_kv_ceiling().unwrap() + ); + let mut disabled = model.clone(); + disabled.query_coverage = None; + assert_eq!(disabled.predict_prefill_latency(1, 250, 100).unwrap(), 17.0); + assert_eq!(disabled.fpm_query_coverage().unwrap(), None); +} + fn regression_model( worker_type: ForwardPassWorkerType, options: ForwardPassPerfOptions, diff --git a/crates/core/src/perfmodel/mod.rs b/crates/core/src/perfmodel/mod.rs index 86b98291f..cb65a2f45 100644 --- a/crates/core/src/perfmodel/mod.rs +++ b/crates/core/src/perfmodel/mod.rs @@ -54,7 +54,8 @@ pub use fpm::{ }; pub use fpm::{ CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, FpmInterpolationMethod, - OpLevelConfig, RegressionFitConfig, RegressionFitKind, + FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQueryPurpose, OpLevelConfig, + RegressionFitConfig, RegressionFitKind, }; // Forward-pass metrics telemetry types and schema version, plus the // crate-internal validation helper. Re-exported at the crate root so existing diff --git a/crates/core/src/perfmodel/operators/fpm_forward.rs b/crates/core/src/perfmodel/operators/fpm_forward.rs index 0c8a9f93a..b9f6ab426 100644 --- a/crates/core/src/perfmodel/operators/fpm_forward.rs +++ b/crates/core/src/perfmodel/operators/fpm_forward.rs @@ -38,6 +38,7 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use crate::common::error::AicError; +use crate::fpm::coverage::{DirectFpmResolution, FpmQueryCoverage, FpmQueryPurpose}; use crate::operators::op::{Op, RuntimeContext}; use crate::operators::{PerformanceResult, Source}; use crate::perf_database::PerfDatabase; @@ -125,6 +126,15 @@ impl FpmForwardOp { &self, db: &PerfDatabase, ctx: &RuntimeContext, + ) -> Result { + self.query_with_coverage(db, ctx, None) + } + + pub(crate) fn query_with_coverage( + &self, + db: &PerfDatabase, + ctx: &RuntimeContext, + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { if self.verify_width == 0 || (self.phase == FpmPhase::Prefill && self.verify_width != 1) { return Err(data_err(format!( @@ -153,9 +163,6 @@ impl FpmForwardOp { ctx.beam_width ))); } - let cell = db - .fpm_forward - .select_cell(&self.match_identity, &self.model_path)?; let b = batch_size as f64; let coords: Vec = match self.phase { FpmPhase::Prefill => { @@ -174,7 +181,7 @@ impl FpmForwardOp { vec![b, b / w * s as f64] } }; - self.resolve(db, cell, &coords) + self.query_totals_with_coverage(db, &coords, coverage) } /// Resolve at RAW per-rank iteration totals — the table's native @@ -187,6 +194,15 @@ impl FpmForwardOp { &self, db: &PerfDatabase, coords: &[f64], + ) -> Result { + self.query_totals_with_coverage(db, coords, None) + } + + pub(crate) fn query_totals_with_coverage( + &self, + db: &PerfDatabase, + coords: &[f64], + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { let expected = match self.phase { FpmPhase::Prefill => 3, @@ -199,6 +215,14 @@ impl FpmForwardOp { coords ))); } + if let Some(coverage) = coverage { + return self.query_direct_with_coverage( + db, + coords, + FpmQueryPurpose::ForwardPass, + coverage, + ); + } let cell = db .fpm_forward .select_cell(&self.match_identity, &self.model_path)?; @@ -242,6 +266,16 @@ impl FpmForwardOp { db: &PerfDatabase, batch_size: u32, total_kv: f64, + ) -> Result { + self.query_pass_baseline_with_coverage(db, batch_size, total_kv, None) + } + + pub(crate) fn query_pass_baseline_with_coverage( + &self, + db: &PerfDatabase, + batch_size: u32, + total_kv: f64, + coverage: Option<&mut FpmQueryCoverage>, ) -> Result { if self.phase != FpmPhase::Decode { return Err(data_err(format!( @@ -254,11 +288,21 @@ impl FpmForwardOp { "invalid FPM baseline query: batch_size={batch_size}" ))); } + if let Some(coverage) = coverage { + return self.query_direct_with_coverage( + db, + &[batch_size as f64, total_kv], + FpmQueryPurpose::MixedDecodeBaseline, + coverage, + ); + } let cell = db .fpm_forward .select_cell(&self.match_identity, &self.model_path)?; if self.interpolation == FpmInterpolation::Direct { - return self.resolve_direct_decode(cell, &[batch_size as f64, total_kv], true); + return self + .resolve_direct_decode(cell, &[batch_size as f64, total_kv], true) + .map(|(result, _)| result); } let Some(domain) = cell.decode_domain else { return Err(data_err(format!( @@ -343,7 +387,8 @@ impl FpmForwardOp { return match self.phase { FpmPhase::Prefill => self.resolve_direct_prefill(cell, coords), FpmPhase::Decode => self.resolve_direct_decode(cell, coords, false), - }; + } + .map(|(result, _)| result); } // Data-certified prefill batch clamp (mirrors Python _resolve): the // regime coordinate is the token TOTAL, which stays untouched — the @@ -476,11 +521,48 @@ impl FpmForwardOp { Ok(PerformanceResult::new(latency, Source::Silicon)) } + fn query_direct_with_coverage( + &self, + db: &PerfDatabase, + coords: &[f64], + purpose: FpmQueryPurpose, + coverage: &mut FpmQueryCoverage, + ) -> Result { + let mut cell_ids = &[][..]; + let result = (|| { + if self.interpolation != FpmInterpolation::Direct { + return Err(AicError::InvalidEngineConfig( + "query coverage requires direct FPM interpolation".into(), + )); + } + let cell = db + .fpm_forward + .select_cell(&self.match_identity, &self.model_path)?; + cell_ids = cell.cell_ids.as_slice(); + match self.phase { + FpmPhase::Prefill => self.resolve_direct_prefill(cell, coords), + FpmPhase::Decode => self.resolve_direct_decode( + cell, + coords, + purpose == FpmQueryPurpose::MixedDecodeBaseline, + ), + } + })(); + coverage.record( + self, + cell_ids, + coords, + purpose, + result.as_ref().map(|(_, kind)| *kind), + ); + result.map(|(result, _)| result) + } + fn resolve_direct_prefill( &self, cell: &FpmForwardCell, coords: &[f64], - ) -> Result { + ) -> Result<(PerformanceResult, DirectFpmResolution), AicError> { 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( @@ -489,13 +571,21 @@ impl FpmForwardOp { "batch_size must be a positive integer", )); } - if let Some(value) = cell + if let Some((curve, 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)) + .and_then(|curve| direct_curve_value(curve, tokens).map(|value| (curve, value))) { - return self.direct_result(cell, coords, value); + let resolution = + if tokens == (tokens as u32) as f64 && curve.contains_key(&(tokens as u32)) { + DirectFpmResolution::Measured + } else { + DirectFpmResolution::Interpolated + }; + return self + .direct_result(cell, coords, value) + .map(|result| (result, resolution)); } let covered: Vec<(u32, f64)> = cell .direct_prefill @@ -532,6 +622,7 @@ impl FpmForwardOp { let weight = (kv - lo as f64) / (hi as f64 - lo as f64); let latency = lo_ms + (hi_ms - lo_ms) * weight; self.direct_result(cell, coords, latency) + .map(|result| (result, DirectFpmResolution::Interpolated)) } fn direct_result( @@ -574,7 +665,7 @@ impl FpmForwardOp { cell: &FpmForwardCell, coords: &[f64], baseline: bool, - ) -> Result { + ) -> Result<(PerformanceResult, DirectFpmResolution), AicError> { let (lo_row, hi_row) = self.direct_decode_rows(cell, coords)?; let row_value = |row: u32| { let curve = &cell.direct_decode[&row]; @@ -591,7 +682,19 @@ impl FpmForwardOp { 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 }; + // A baseline uses each selected row's own measured KV floor. Its + // requested KV controls row coverage, not the floor coordinate. + let measured = lo_row == hi_row + && (baseline + || (coords[1] == (coords[1] as u32) as f64 + && cell.direct_decode[&lo_row].contains_key(&(coords[1] as u32)))); + let resolution = if measured { + DirectFpmResolution::Measured + } else { + DirectFpmResolution::Interpolated + }; self.direct_result(cell, coords, latency) + .map(|result| (result, resolution)) } /// One shared coverage decision for decode queries and their mixed-pass diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 3050694b2..30b4ae130 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1935,6 +1935,42 @@ impl PyForwardPassPerfModel { .map_err(|e| PyValueError::new_err(format!("diagnostics serialize: {e}"))) } + /// Bounded direct-FPM native lookup evidence, or JSON null when disabled. + fn fpm_query_coverage(&self) -> PyResult { + let coverage = self.inner.fpm_query_coverage().map_err(aic_to_py)?; + serde_json::to_string(&coverage) + .map_err(|error| PyValueError::new_err(format!("FPM coverage serialize: {error}"))) + } + + fn predict_prefill_latency( + &self, + py: Python<'_>, + batch_size: u32, + isl: u32, + prefix: u32, + ) -> PyResult { + py.allow_threads(|| self.inner.predict_prefill_latency(batch_size, isl, prefix)) + .map_err(aic_to_py) + } + + fn predict_decode_latency_total( + &self, + py: Python<'_>, + batch_size: u32, + total_past_kv_tokens: u32, + ) -> PyResult { + py.allow_threads(|| { + self.inner + .predict_decode_latency_total(batch_size, total_past_kv_tokens) + }) + .map_err(aic_to_py) + } + + fn fpm_decode_kv_ceiling(&self, py: Python<'_>) -> PyResult> { + py.allow_threads(|| self.inner.fpm_decode_kv_ceiling()) + .map_err(aic_to_py) + } + /// Regression store labels, readiness and retained counts as JSON. /// Includes cold stores; native AIC models return an empty list. fn regression_store_diagnostics(&self) -> PyResult { diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index f5a021fe8..1e95789a9 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -405,7 +405,8 @@ type PhaseEvidenceKey = (u32, u32, u32, u32, bool); struct AicTimingModel { engine: Py, - diagnostic_model: Option, + diagnostic_model: Option>, + fpm_model: Option>, decoder_replay: bool, use_fpm_decode_totals: bool, fpm_decode_kv_ceiling: Option, @@ -417,8 +418,10 @@ impl AicTimingModel { fn build(config: &mut AicTimingConfig, worker_type: ForwardPassWorkerType) -> Result { config.validate_parallel_shape()?; config.resolved_memory_fraction()?; - let model = ForwardPassPerfModel::best_available(config.estimator_request(worker_type)?) - .context("AIC timing provider could not construct the requested estimator")?; + let model = Arc::new( + ForwardPassPerfModel::best_available(config.estimator_request(worker_type)?) + .context("AIC timing provider could not construct the requested estimator")?, + ); let provenance = model .provenance() .context("canonical estimator omitted construction provenance")?; @@ -448,20 +451,17 @@ impl AicTimingModel { let native = model.native_engine().context( "AIC regression estimator is not ready: offline replay requires trained observations or a native estimator" )?; - let (engine, fpm_decode_kv_ceiling) = Python::with_gil(|py| -> PyResult<_> { - let engine = Py::new(py, crate::AicEngine::from_shared_engine(native))?.into_any(); - let ceiling = if use_fpm_decode_totals { - engine - .bind(py) - .call_method0("fpm_decode_kv_ceiling")? - .extract::>()? - } else { - None - }; - Ok((engine, ceiling)) + let engine = Python::with_gil(|py| -> PyResult<_> { + Ok(Py::new(py, crate::AicEngine::from_shared_engine(native))?.into_any()) })?; + let fpm_decode_kv_ceiling = if use_fpm_decode_totals { + model.fpm_decode_kv_ceiling()? + } else { + None + }; Ok(Self { engine, + fpm_model: use_fpm_decode_totals.then(|| Arc::clone(&model)), diagnostic_model: Some(model), decoder_replay: config.decoder_replay, use_fpm_decode_totals, @@ -607,6 +607,11 @@ impl TimingModel for AicTimingModel { self.record_evidence(evidence, true)?; return Ok(latency_ms); } + if let Some(model) = &self.fpm_model { + return model + .predict_prefill_latency(batch_size, mean_isl, mean_prefix) + .context("AIC prefill prediction failed"); + } Python::with_gil(|py| { self.engine .bind(py) @@ -633,6 +638,11 @@ 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.fpm_model { + return model + .predict_decode_latency_total(batch_size, total_past_kv_tokens) + .context("AIC decode prediction failed"); + } return Python::with_gil(|py| { self.engine .bind(py) @@ -821,6 +831,7 @@ fn resolve_role_timing( capacity_is_explicit: bool, worker_type: ForwardPassWorkerType, capture_performance_diagnostics: bool, + coverage_sources: &mut Vec<(ForwardPassWorkerType, Arc)>, ) -> Result>> { let TimingModelConfig::External { provider, config } = role.rank.timing_model.clone() else { return Ok(None); @@ -833,6 +844,11 @@ fn resolve_role_timing( let mut config: AicTimingConfig = serde_json::from_value(config).context("invalid AIC timing provider configuration")?; let mut timing = AicTimingModel::build(&mut config, worker_type)?; + if let Some(model) = &timing.fpm_model + && model.fpm_query_coverage()?.is_some() + { + coverage_sources.push((worker_type, Arc::clone(model))); + } if !capture_performance_diagnostics { timing.diagnostic_model = None; } @@ -1577,6 +1593,76 @@ fn infer_aic_timing_topology(engine: &mut serde_json::Value) -> Result<()> { Ok(()) } +#[derive(Debug)] +struct ReplayCoverageFailure { + error: anyhow::Error, + coverage: serde_json::Value, +} + +impl std::fmt::Display for ReplayCoverageFailure { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "{}", self.error) + } +} + +impl std::error::Error for ReplayCoverageFailure { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.error.as_ref()) + } +} + +fn replay_fpm_coverage( + sources: &[(ForwardPassWorkerType, Arc)], + report: Option<&crate::replay::ReplayReport>, +) -> Result> { + if sources.is_empty() { + return Ok(None); + } + let mut measured = 0u64; + let mut interpolated = 0u64; + let mut unsupported = 0u64; + let mut roles = Vec::with_capacity(sources.len()); + for (role, model) in sources { + let coverage = model + .fpm_query_coverage()? + .context("requested FPM query coverage is unavailable")?; + measured += coverage.queries.measured; + interpolated += coverage.queries.interpolated; + unsupported += coverage.queries.unsupported; + roles.push(serde_json::json!({"worker_type": role, "coverage": coverage})); + } + let replay_completed = report.is_some_and(|report| { + let counts = &report.request_counts; + counts.num_requests > 0 + && counts.completed_requests == counts.num_requests + && report.agentic_graph.as_ref().is_none_or(|graph| { + graph.node_count == counts.completed_requests + && report + .agentic_play_outcomes + .as_ref() + .is_some_and(|outcomes| { + outcomes.len() == graph.play_count + && outcomes.iter().all(|outcome| { + matches!( + outcome.status, + crate::replay::loadgen::AgenticPlayStatus::Completed + ) + }) + }) + }) + }); + let covered = replay_completed && measured + interpolated > 0 && unsupported == 0; + Ok(Some(serde_json::json!({ + "schema_version": 1, + "status": if covered { "covered" } else { "incomplete" }, + "replay_completed": replay_completed, + "counting_unit": "native_lookup_resolutions", + "queries": {"measured": measured, "interpolated": interpolated, "unsupported": unsupported}, + "roles": roles, + "accuracy": "not_assessed", + }))) +} + fn execute_json(payload: &str, capture_artifacts: bool) -> Result { let (mut spec, mut traffic, capture_performance_diagnostics) = match serde_json::from_str(payload).context("invalid AISimulate execution ReplaySpec")? { @@ -1619,141 +1705,161 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { ReplayTopology::Disaggregated { .. } => 2, }; let mut power_sources = Vec::with_capacity(expected_power_sources); - - let (mut report, artifacts, resolved_weka_timestamp_basis) = match spec.topology.clone() { - ReplayTopology::Aggregated { .. } => { - let mut role = aggregated_role(&engine_config); - let capacity_is_explicit = role - .num_gpu_blocks_is_explicit - .unwrap_or_else(|| role_capacity_is_explicit(&serialized_engine, None)); - let timing = resolve_role_timing( - &mut role, - capacity_is_explicit, - ForwardPassWorkerType::Aggregated, - capture_performance_diagnostics, - )?; - if let Some(timing) = timing.as_ref() { - power_sources.push(TimingPowerSource::new(Arc::clone(timing), &role.rank)); - } - engine_config.dp_size = role.dp_size; - engine_config.tensor_parallel_size = role.tensor_parallel_size; - engine_config.num_gpu_blocks_is_explicit = role.num_gpu_blocks_is_explicit; - engine_config.rank = role.rank; - if let Some(traffic) = traffic.as_mut() - && let ReplayTopology::Aggregated { workers } = &spec.topology - && let Some(concurrency) = resolve_kv_capacity_concurrency( - traffic, - &ReplayRoleConfig { - dp_size: engine_config.dp_size, - tensor_parallel_size: engine_config.tensor_parallel_size, - num_gpu_blocks_is_explicit: engine_config.num_gpu_blocks_is_explicit, - rank: engine_config.rank.clone(), - }, - workers.initial_workers, - )? - { - spec.max_in_flight = Some(concurrency); - } - spec.engine = serde_json::to_value(&engine_config) - .context("serializing materialized native engine descriptor")?; - let built_input = traffic - .map(|traffic| build_runtime_input(traffic, engine_config.rank.block_size, true)) - .transpose()?; - if let Some(built) = &built_input { - validate_public_agentic_engine(&built.input, &engine_config.rank)?; - } - let resolved_basis = built_input - .as_ref() - .and_then(|built| built.weka_nested_timestamp_basis); - let input = built_input.map(|built| built.input); - let factory = timing.map_or_else( - ReplayEngineFactory::new, - ReplayEngineFactory::with_timing_model, - ); - run_with_input(spec, factory, input, capture_artifacts) - .map(|(report, artifacts)| (report, artifacts, resolved_basis)) - } - ReplayTopology::Disaggregated { .. } => { - let mut prefill = engine_config - .prefill - .clone() - .unwrap_or_else(|| aggregated_role(&engine_config)); - let mut decode = engine_config - .decode - .clone() - .unwrap_or_else(|| aggregated_role(&engine_config)); - let prefill_capacity_is_explicit = prefill - .num_gpu_blocks_is_explicit - .unwrap_or_else(|| role_capacity_is_explicit(&serialized_engine, Some("prefill"))); - let decode_capacity_is_explicit = decode - .num_gpu_blocks_is_explicit - .unwrap_or_else(|| role_capacity_is_explicit(&serialized_engine, Some("decode"))); - let prefill_timing = resolve_role_timing( - &mut prefill, - prefill_capacity_is_explicit, - ForwardPassWorkerType::Prefill, - capture_performance_diagnostics, - )?; - let decode_timing = resolve_role_timing( - &mut decode, - decode_capacity_is_explicit, - ForwardPassWorkerType::Decode, - capture_performance_diagnostics, - )?; - if let Some(timing) = prefill_timing.as_ref() { - power_sources.push(TimingPowerSource::new(Arc::clone(timing), &prefill.rank)); - } - if let Some(timing) = decode_timing.as_ref() { - power_sources.push(TimingPowerSource::new(Arc::clone(timing), &decode.rank)); - } - engine_config.prefill = Some(prefill); - engine_config.decode = Some(decode); - if let Some(traffic) = traffic.as_mut() - && let ReplayTopology::Disaggregated { decode, .. } = &spec.topology - && let Some(concurrency) = resolve_kv_capacity_concurrency( - traffic, - engine_config - .decode - .as_ref() - .expect("decode role was materialized"), - decode.initial_workers, - )? - { - spec.max_in_flight = Some(concurrency); + let mut coverage_sources = Vec::new(); + + let replay_result = (|| -> Result<_> { + match spec.topology.clone() { + ReplayTopology::Aggregated { .. } => { + let mut role = aggregated_role(&engine_config); + let capacity_is_explicit = role + .num_gpu_blocks_is_explicit + .unwrap_or_else(|| role_capacity_is_explicit(&serialized_engine, None)); + let timing = resolve_role_timing( + &mut role, + capacity_is_explicit, + ForwardPassWorkerType::Aggregated, + capture_performance_diagnostics, + &mut coverage_sources, + )?; + if let Some(timing) = timing.as_ref() { + power_sources.push(TimingPowerSource::new(Arc::clone(timing), &role.rank)); + } + engine_config.dp_size = role.dp_size; + engine_config.tensor_parallel_size = role.tensor_parallel_size; + engine_config.num_gpu_blocks_is_explicit = role.num_gpu_blocks_is_explicit; + engine_config.rank = role.rank; + if let Some(traffic) = traffic.as_mut() + && let ReplayTopology::Aggregated { workers } = &spec.topology + && let Some(concurrency) = resolve_kv_capacity_concurrency( + traffic, + &ReplayRoleConfig { + dp_size: engine_config.dp_size, + tensor_parallel_size: engine_config.tensor_parallel_size, + num_gpu_blocks_is_explicit: engine_config.num_gpu_blocks_is_explicit, + rank: engine_config.rank.clone(), + }, + workers.initial_workers, + )? + { + spec.max_in_flight = Some(concurrency); + } + spec.engine = serde_json::to_value(&engine_config) + .context("serializing materialized native engine descriptor")?; + let built_input = traffic + .map(|traffic| { + build_runtime_input(traffic, engine_config.rank.block_size, true) + }) + .transpose()?; + if let Some(built) = &built_input { + validate_public_agentic_engine(&built.input, &engine_config.rank)?; + } + let resolved_basis = built_input + .as_ref() + .and_then(|built| built.weka_nested_timestamp_basis); + let input = built_input.map(|built| built.input); + let factory = timing.map_or_else( + ReplayEngineFactory::new, + ReplayEngineFactory::with_timing_model, + ); + run_with_input(spec, factory, input, capture_artifacts) + .map(|(report, artifacts)| (report, artifacts, resolved_basis)) } - spec.engine = serde_json::to_value(&engine_config) - .context("serializing materialized native engine descriptor")?; - let built_input = traffic - .map(|traffic| { - build_runtime_input( + ReplayTopology::Disaggregated { .. } => { + let mut prefill = engine_config + .prefill + .clone() + .unwrap_or_else(|| aggregated_role(&engine_config)); + let mut decode = engine_config + .decode + .clone() + .unwrap_or_else(|| aggregated_role(&engine_config)); + let prefill_capacity_is_explicit = + prefill.num_gpu_blocks_is_explicit.unwrap_or_else(|| { + role_capacity_is_explicit(&serialized_engine, Some("prefill")) + }); + let decode_capacity_is_explicit = + decode.num_gpu_blocks_is_explicit.unwrap_or_else(|| { + role_capacity_is_explicit(&serialized_engine, Some("decode")) + }); + let prefill_timing = resolve_role_timing( + &mut prefill, + prefill_capacity_is_explicit, + ForwardPassWorkerType::Prefill, + capture_performance_diagnostics, + &mut coverage_sources, + )?; + let decode_timing = resolve_role_timing( + &mut decode, + decode_capacity_is_explicit, + ForwardPassWorkerType::Decode, + capture_performance_diagnostics, + &mut coverage_sources, + )?; + if let Some(timing) = prefill_timing.as_ref() { + power_sources.push(TimingPowerSource::new(Arc::clone(timing), &prefill.rank)); + } + if let Some(timing) = decode_timing.as_ref() { + power_sources.push(TimingPowerSource::new(Arc::clone(timing), &decode.rank)); + } + engine_config.prefill = Some(prefill); + engine_config.decode = Some(decode); + if let Some(traffic) = traffic.as_mut() + && let ReplayTopology::Disaggregated { decode, .. } = &spec.topology + && let Some(concurrency) = resolve_kv_capacity_concurrency( traffic, engine_config - .prefill + .decode .as_ref() - .expect("prefill role was materialized") - .rank - .block_size, - false, - ) - }) - .transpose()?; - let resolved_basis = built_input - .as_ref() - .and_then(|built| built.weka_nested_timestamp_basis); - let input = built_input.map(|built| built.input); - run_with_input( - spec, - ReplayEngineFactory::with_optional_role_timing_models( - prefill_timing, - decode_timing, - ), - input, - capture_artifacts, - ) - .map(|(report, artifacts)| (report, artifacts, resolved_basis)) + .expect("decode role was materialized"), + decode.initial_workers, + )? + { + spec.max_in_flight = Some(concurrency); + } + spec.engine = serde_json::to_value(&engine_config) + .context("serializing materialized native engine descriptor")?; + let built_input = traffic + .map(|traffic| { + build_runtime_input( + traffic, + engine_config + .prefill + .as_ref() + .expect("prefill role was materialized") + .rank + .block_size, + false, + ) + }) + .transpose()?; + let resolved_basis = built_input + .as_ref() + .and_then(|built| built.weka_nested_timestamp_basis); + let input = built_input.map(|built| built.input); + run_with_input( + spec, + ReplayEngineFactory::with_optional_role_timing_models( + prefill_timing, + decode_timing, + ), + input, + capture_artifacts, + ) + .map(|(report, artifacts)| (report, artifacts, resolved_basis)) + } } - } - .context("AISimulate replay failed")?; + .map_err(Into::into) + })() + .context("AISimulate replay failed"); + let (mut report, artifacts, resolved_weka_timestamp_basis) = match replay_result { + Ok(result) => result, + Err(error) => { + if let Some(coverage) = replay_fpm_coverage(&coverage_sources, None)? { + return Err(ReplayCoverageFailure { error, coverage }.into()); + } + return Err(error); + } + }; let (timing_evidence, unavailable_reason) = if power_sources.len() == expected_power_sources { ( replay_timing_evidence(&power_sources)?, @@ -1777,6 +1883,9 @@ 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 let Some(coverage) = replay_fpm_coverage(&coverage_sources, Some(&report))? { + report_json["fpm_query_coverage"] = coverage; + } if capture_performance_diagnostics { report_json["performance_diagnostics"] = replay_performance_diagnostics(timing_evidence.as_ref()); @@ -1848,7 +1957,7 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { } fn replay_python_error(error: anyhow::Error) -> PyErr { - if error.chain().any(|cause| { + let exception = if error.chain().any(|cause| { matches!( cause.downcast_ref::(), Some(crate::replay::ReplayError::ResourceLimited(_)) @@ -1857,7 +1966,15 @@ fn replay_python_error(error: anyhow::Error) -> PyErr { PyMemoryError::new_err(format!("{error:#}")) } else { PyRuntimeError::new_err(format!("{error:#}")) + }; + if let Some(failure) = error.downcast_ref::() { + Python::with_gil(|py| { + let _ = exception + .value(py) + .setattr("fpm_query_coverage", failure.coverage.to_string()); + }); } + exception } /// Execute one canonical serialized ReplaySpec and return serialized report JSON. @@ -2154,6 +2271,7 @@ mod tests { AicTimingModel { engine, diagnostic_model: None, + fpm_model: None, decoder_replay: false, use_fpm_decode_totals, fpm_decode_kv_ceiling: None, @@ -2966,10 +3084,10 @@ mod tests { .into_any() }); let mut timing = timing_model(engine, false); - timing.diagnostic_model = Some(ForwardPassPerfModel::from_engine( + timing.diagnostic_model = Some(Arc::new(ForwardPassPerfModel::from_engine( Arc::new(native), ForwardPassPerfOptions::default(), - )); + ))); for mode in ["static_ctx", "static_gen"] { let error = timing .predict_phase_evidence(1, 128, 2, 0, mode) diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index 22f55d95c..b6edb1572 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -11,7 +11,7 @@ use aisimulate_core::{ AicEngine, AicEngineBuilder, AicError, BackendKind, DatabaseMode, EstimationMode, EstimatorConfig, ForwardPassPerfModel, ForwardPassPerfModelConfig, ForwardPassRegressionStoreDiagnostics, ForwardPassWorkerType, FpmInterpolationMethod, - KvCacheEstimateRequest, + FpmQueryCoverage, KvCacheEstimateRequest, }; /// Compile the ergonomic engine builder without starting embedded Python. @@ -80,6 +80,19 @@ pub fn profile_model( ForwardPassPerfModel::best_available(config) } +/// Coverage is owned by the model returned by the one construction API. +pub fn fpm_coverage(model: &ForwardPassPerfModel) -> Result, AicError> { + model.fpm_query_coverage() +} + +pub fn covered_prefill(model: &ForwardPassPerfModel) -> Result { + model.predict_prefill_latency(1, 128, 0) +} + +pub fn covered_decode(model: &ForwardPassPerfModel) -> Result { + model.predict_decode_latency_total(1, 128) +} + pub fn best_available_model( config: ForwardPassPerfModelConfig, ) -> Result { @@ -243,6 +256,11 @@ mod tests { } /// Detailed phase evidence is reachable through the canonical model. -pub fn operation_diagnostics(model: &ForwardPassPerfModel) -> Result, AicError> { +pub fn operation_diagnostics( + model: &ForwardPassPerfModel, +) -> Result< + Vec, + AicError, +> { model.static_phase_diagnostics(1, 128, 0, true) } diff --git a/docs/core-api.md b/docs/core-api.md index aba641df4..7fda510c2 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -263,6 +263,10 @@ nested paths. The supported namespaces are: SOL for a registered architecture, or direct interpolation for an unknown architecture with a valid profile. Without a profile, auto retains SOL. Explicit SOL requires a registered analytical model; direct requires a profile. +- `fpm_interpolation.collect_coverage`: `false` (default). Opt in to bounded + evidence from actual direct-FPM lookups, as described below. It requires + explicit `estimation_mode: fpm_interpolation`, resolved `method: direct`, and + `fallback_policy: deny`. It is omitted from normalized serialization when false. - `op_level`: a reserved typed namespace; unknown fields are rejected. The top-level `fpm_profile` contains the complete profile dictionary: pinned @@ -306,6 +310,70 @@ return `None`. `tune_with_fpms()` preserves the established FPM observation contract. Native construction still uses Python model compilation; estimator selection, regression, correction, and latency computation are owned by Rust. +### Direct-FPM query coverage + +Enable coverage through the same canonical `best_available(config)` construction: + +```yaml +estimation_mode: fpm_interpolation +fallback_policy: deny +estimator_config: + fpm_interpolation: + method: direct + collect_coverage: true +``` + +The rest of the configuration must identify the exact deployment, including its +complete `fpm_profile`, literal backend version and data roots. CLI prediction +places these timing controls under `engine.workers..timing`. Coverage does +not change point selection, interpolation, correction or latency, and it does +not permit fallback or extrapolation after a missing query. + +The returned model exposes `fpm_query_coverage()` as +`Result, AicError>` in Rust and `dict[str, Any] | None` +in the Python SDK. The raw PyO3 method returns a JSON string, using `null` when +disabled. Evidence belongs to that model instance; cloning a Rust model starts +a fresh accumulator so candidate evaluations do not share counts. The public +Rust facade exports `FpmQueryCoverage`, `FpmQueryCoverageCounts`, `FpmQueryGap` +and `FpmQueryPurpose`. + +A snapshot has cumulative `queries`, `prefill`, `decode` and +`mixed_decode_baseline` counts, each with `measured`, `interpolated` and +`unsupported`. Classification comes from the native lookup that supplied or +rejected the timing, including mixed-pass decode-floor lookups. Counts describe +`native_lookup_resolutions`; they are not request or replay-iteration counts, +and an external timing-cache hit adds no lookup. The snapshot retains up to 128 +distinct gaps with phase, purpose, model/cell identity, query coordinates, reason +and occurrence count. `omitted_gap_queries` counts failed occurrences beyond +that storage limit; unsupported totals are not truncated. Snapshot reads do not +reset evidence, and lookup errors remain errors with partial evidence available. +The snapshot alone does not assess replay completion or prediction accuracy. + +For existing static phase consumers, the same returned model provides these +uncorrected native timings, in milliseconds: + +| Method | Coordinate contract | +| --- | --- | +| `predict_prefill_latency(batch_size, isl, prefix)` | Full input length and cached-prefix length per request; direct-FPM totals are `(batch_size, batch_size * (isl - prefix), batch_size * prefix)`. | +| `predict_decode_latency_total(batch_size, total_past_kv_tokens)` | Exact total past-KV across the decode batch, preserving the native FPM coordinate without a mean-context conversion. | +| `fpm_decode_kv_ceiling()` | Largest collected decode KV total, or `None`; reading this bound does not record a timing lookup or prove shape-specific coverage. | + +These methods preserve the existing engine's phase behavior and record actual +direct lookup evidence when enabled. Empty work does not invent a measured +query. Whole-iteration `estimate_forward_pass_time_ms()` remains the API for +scheduled per-rank FPM telemetry, including mixed work and its existing online +correction behavior. + +Ordinary replay uses this returned model for FPM timing and emits +`fpm_query_coverage` in its report when collection is enabled. The CLI also saves +`fpm-coverage.json`, including partial evidence when a native timing error stops +replay; the Python exception retains a `fpm_query_coverage` attribute for that +failure path. A passing replay coverage status requires completed requests, +nonempty resolved queries and no unsupported lookup. It is distinct from +operation/energy evidence, which whole-model FPM does not provide. See +[FPM replay validation](fpm-self-service.md#validate-fpm-query-coverage-with-agentx-replay) +for the stricter whole-corpus completion checks and saved onboarding artifacts. + ### Migrating saved configuration Use `ForwardPassPerfModelConfig.from_legacy_engine_config(old_config, diff --git a/docs/fpm-self-service.md b/docs/fpm-self-service.md index 58615d684..1bb480e18 100644 --- a/docs/fpm-self-service.md +++ b/docs/fpm-self-service.md @@ -5,7 +5,9 @@ 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 interconnect, plans one TP, DEP, or TEP worker, and derives the minimum GPUs required to collect that worker's timings. It produces ordinary `predict` and `recommend` configurations for validating that worker against your collected FPM data. It builds on AISimulate's 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 interconnect, plans one TP, DEP, or TEP worker, and derives the minimum GPUs required to collect that worker's timings. You review the resource and collection limits, collect whole-forward timings through Dynamo self-benchmark, then validate their query coverage using ordinary trace replay. It also produces ordinary `predict` and `recommend` configurations for the selected worker. + +Collection limits and validation traffic are separate inputs. Dynamo self-benchmark generates its sampling grid from CUDA graph sizes and runtime bounds. AgentX traces exercise the resulting FPM library through replay; their variable request lengths do not require a fixed input/output length or latency target during onboarding. 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. @@ -18,11 +20,11 @@ For a request such as "Help me onboard my model for FPM simulation on my target | Stage | Required result | | --- | --- | | 1. Inspect the model and target | Accessible config/profile and packaged hardware specification identified; supported metadata read and gaps recorded. | -| 2. Choose the worker and workload | Checkpoint identity/revision, runtime, interconnect, exact worker topology and initial workload selected; minimum collection GPUs derived. Runtime compatibility remains unchecked. | +| 2. Choose the worker and collection limits | Checkpoint identity/revision, runtime, interconnect, exact worker topology, context and scheduler/capture limits selected; minimum collection GPUs derived. Runtime compatibility remains unchecked. | | 3. Derive, review and save the profile | Exact resource/precision values and assumptions reviewed and accepted; final request contains the complete profile and provenance. | | 4. Plan collection | Validated saved plan, generated configurations and collector preview; sampling scope and remaining execution prerequisites explained. | | 5. Collect and verify data | Matching formal Parquet/metadata pair verified, with provenance and available phase cells recorded. | -| 6. Run predict/recommend and report | Single-worker prediction and recommendation exercised; results, coverage and failures reported, with accuracy stated separately. | +| 6. Validate replay and run predict/recommend | Cold aggregated trace replay checks direct-FPM query coverage; completed and incomplete results are distinguished. Ordinary prediction/recommendation examples remain available; accuracy is assessed separately. | At each transition or blocker, give a short update: **Stage N/6 — name; result or blocker; next action.** Ask only for missing information needed for the current stage, using facts and decisions already supplied. Continue independent authorized work while an answer is pending. Move on when the required evidence is available; stage transitions do not require another approval. Honor prior authorization, including authorization to execute collection, while preserving the profile acceptance step below. @@ -40,11 +42,13 @@ A natively multimodal checkpoint can be onboarded for its text decoder. Explain The agent handles checkout and environment checks: record the branch/commit, follow [development setup](../DEVELOPMENT.md#initial-setup), activate the environment, and inspect `aisimulate onboard --help` and `aisimulate onboard init --help`. Check later subcommands before using them. Report missing commands/options as a version mismatch, rather than asking the user to supply unsupported inputs. -### 2. Choose the worker and workload +### 2. Choose the worker and collection limits Resolve only the remaining checkpoint identifier and immutable revision, literal vLLM version, target GPU platform and interconnect. Use available repository/deployment metadata before asking the user. A model label, config hash or example revision does not establish a checkpoint pin. Do not ask for total available GPUs, node allocation, GPUs per node or replica budgets during onboarding; those are choices for actual prediction or recommendation runs. -Establish input/output lengths, concurrency, context and latency targets here. If there is no workload preference, show the [planning defaults](#create-the-request) and cap the context at the model's known limit; adjust token lengths if needed to fit that limit. With a local model config, inspect the [read-only topology preview](#preview-and-choose-parallelism) using the actual identity, runtime, target and workload. Resolve shared precision/layout facts explicitly and rerun the preview. Present the default and alternatives with their resource assumptions and unresolved fields. A candidate marked `estimated_fit` has a complete declared/estimated byte budget; it is not a performance ranking, runtime qualification or a measurement. If no default is available, explain why and select an exact candidate before asking for its per-rank bounds. Do not silently assume topology or precision. +Review the [runtime and collection defaults](#runtime-and-collection-limits): per-request context, scheduled token budget, maximum sequences and prefill CUDA graph capture limit. Fresh config/profile-based setup caps context at the smaller of the declared limit and the 256,000-token AgentX reference bound. Explain this initial policy and allow edits; it does not establish memory fit or timing coverage. Do not ask for fixed input/output lengths, concurrency, TTFT or TPOT as required collection inputs. Those flags customize optional synthetic validation examples. + +With a local model config, inspect the [read-only topology preview](#preview-and-choose-parallelism) using the actual identity, runtime, target and collection limits. Resolve shared precision/layout facts explicitly and rerun the preview. Present the default and alternatives with their resource assumptions and unresolved fields. A candidate marked `estimated_fit` has a complete declared/estimated byte budget; it is not a performance ranking, runtime qualification or a measurement. If no default is available, explain why and select an exact candidate before asking for its per-rank bounds. Do not silently assume topology or precision. Help choose one initial TP configuration, or the relevant TP/DEP/TEP configuration for MoE, using the preview's exact flags or the [supported topology flags](#create-the-request). Honor an explicit topology even when it lies outside the automatic shortlist. Explain the selected tuple; do not require the user to know every parallelism field upfront. Derive its minimum collection GPUs as attention TP times attention DP: TP4 requires four GPUs, while DEP8 requires eight. This requirement does not declare available capacity or establish runtime placement or compatibility. Each plan collects one exact topology; use separate requests and output directories to investigate alternatives. @@ -55,7 +59,7 @@ Use [a local model config](#start-from-a-local-model-config), or [a supplied pro | Agent environment | Review and save behavior | | --- | --- | | Terminal or agent tool with a PTY | Run `aisimulate onboard init --model-config /path/to/config.json --interactive --output support-request.yaml`. Relay unresolved prompts and the final profile to the user. Apply requested `edit` actions and return the revised profile to the user for review. Enter CLI `accept` only after the user explicitly accepts those exact values; honor any existing explicit acceptance of those same values. `cancel`, Ctrl-C or EOF creates no new request and preserves any existing output, even with `--overwrite`. This final review is specific to `--model-config --interactive`; supplying `--fpm-profile` does not add it. | -| Headless or noninteractive agent | Supply identity/workload flags and `--resource-overrides` as needed, without `--interactive`. Missing required inputs exit 2 without saving; use the diagnostics to ask for the missing facts. A successful command writes immediately. Initially write to a separate path such as `draft-request.yaml` and show the embedded profile, sources and scope for user review. Apply edits in the inputs/overrides and repeat draft review until accepted; then rerun the unchanged reviewed inputs to a new final request path and verify that it matches the accepted draft before planning. The draft name is only a file convention; it has no special CLI status. Review a supplied `--fpm-profile` in the same way. | +| Headless or noninteractive agent | Supply identity/collection flags and `--resource-overrides` as needed, without `--interactive`. Missing required inputs exit 2 without saving; use the diagnostics to ask for the missing facts. A successful command writes immediately. Initially write to a separate path such as `draft-request.yaml` and show the embedded profile, sources and scope for user review. Apply edits in the inputs/overrides and repeat draft review until accepted; then rerun the unchanged reviewed inputs to a new final request path and verify that it matches the accepted draft before planning. The draft name is only a file convention; it has no special CLI status. Review a supplied `--fpm-profile` in the same way. | Do not pipe answers into `--interactive`: it requires a terminal. Scripted setup has no built-in acceptance prompt. Keep draft files separate from the final request and preserve prior outputs when revising a deployment. @@ -65,7 +69,7 @@ Model identity, runtime and topology are not profile-review edit fields. If they Follow [Plan, preview, and explicitly execute](#plan-preview-and-explicitly-execute) using a new output directory. Inspect `support-plan.json`, the embedded/saved profile, generated prediction/recommendation configs and `commands.json`. Run `onboard collect-fpm` without `--execute` to print the collector command; this does not run the collector's own plan or check the target runtime. Inspect the read-only collector plan when its input environment is available and identify any missing prerequisites. -Explain the minimum collection GPUs, rank-local scheduler/resource envelope and sampling scope. One onboarding plan selects one parallel tuple and generates single-worker validation configs. Use separate requests/output directories for additional tuples. Actual collection resources and placement must be checked in the collector environment before execution. The synthetic request count does not bound timing samples or collection duration. +Explain the minimum collection GPUs, rank-local scheduler/resource envelope and sampling scope. Dynamo self-benchmark owns the point grid; the plan supplies runtime and capture limits, without deriving a second grid from trace requests. One onboarding plan selects one parallel tuple and generates single-worker validation configs. Use separate requests/output directories for additional tuples. Actual collection resources and placement must be checked in the collector environment before execution. The synthetic request count does not bound timing samples or collection duration. ### 5. Collect and verify data @@ -73,17 +77,19 @@ First inspect any existing timing data for a matching deployment. Reuse a verifi For both reused and new data, [inspect the published pair](../python/aisimulate/docs/fpm/end-to-end-workflow.md#5-inspect-the-published-pair): verify hashes, schema and identities, including actual runtime, topology, precision and available prefill/decode cells. Keep it at the generated configs' local systems path. Record any historical checkpoint-revision uncertainty in provenance; a declared revision does not prove that old measurements used it. A successful preview or smoke run is not formal data, and a matching pair does not prove all simulated queries are covered. -### 6. Run predict/recommend and report +### 6. Validate replay and run predict/recommend + +[Validate a local AgentX trace](#validate-fpm-query-coverage-with-agentx-replay) through `onboard validate-fpm`. It writes an inspectable ordinary prediction config, runs cold aggregated replay, and saves coverage evidence separately from the collection plan. Use the selected complete local corpus or one complete play for a first check, and state that selection. Keep `engine.systems_paths`, the target model projection and the direct-FPM profile intact. Report the command's exit status and exact missing coordinates. Missing timing stops replay; a partial report does not audit the rest of the trace. -[Run the generated ordinary configurations](#run-the-generated-ordinary-configurations) for single-worker prediction and recommendation validation. Keep `engine.systems_paths` and the direct-FPM profile intact. Report each command's exit status and results; missing cells and out-of-domain queries need their exact coordinates. For deployment prediction or optimization, set the desired replicas and GPU budget in the ordinary runtime configs. Do not change precision labels or silently switch timing methods to obtain a result. Return to stage 5 to address missing data or to stage 2 if the user changes the selected worker/workload scope. Report the simulation stage as incomplete while required queries fail. +[Run the generated ordinary configurations](#run-the-generated-ordinary-configurations) when a synthetic prediction or recommendation check is useful. For deployment prediction or optimization, set the desired replicas and GPU budget in the ordinary runtime configs. Do not change precision labels or silently switch timing methods to obtain a result. Return to stage 5 to address missing data, or stage 2 if the selected deployment or collection limits change. A different validation trace alone does not invalidate collected timings. Report the simulation stage as incomplete while required queries fail. At handoff, include the checkout revision, final request/profile, plan directory, data pair/provenance, exact commands/exit statuses and all result paths. Distinguish estimated memory fit and CPU planning from actual target-runtime checks, formal data/coverage, and completed simulations. For accuracy, report an independent matched silicon comparison if performed, or explicitly **not assessed**. Successful simulation is not evidence of accuracy; an accuracy study is not a mandatory additional collection campaign for onboarding. ### Resume from existing work -Inspect the saved request/profile, plan, data pair and results before deciding where to resume. Validate that they still match the checkout's CLI, selected deployment and workload; use the existing plan checks described below. An accepted final request can start at stage 4, a valid plan at stage 5, and a verified matching data pair at stage 6. A draft or a saved file without evidence of acceptance still needs stage 3 review. Do not repeat accepted decisions or rerun completed work without a reason. +Inspect the saved request/profile, plan, data pair and results before deciding where to resume. Validate that they still match the checkout's CLI, selected deployment and collection settings; use the existing plan checks described below. An accepted final request can start at stage 4, a valid plan at stage 5, and a verified matching data pair at stage 6. A draft or a saved file without evidence of acceptance still needs stage 3 review. Do not repeat accepted decisions or rerun completed work without a reason. -Preserve completed artifacts when blocked and report the current stage, specific missing input and next action. Changed deployment or workload inputs invalidate dependent estimates, plans and results; return to the affected stage and use a new output directory. Collection's existing `--resume` is for a matching collector checkpoint as described below, not a general onboarding-stage resume command. +Preserve completed artifacts when blocked and report the current stage, specific missing input and next action. Changed deployment, resource profile or collection bounds require review and a new collection directory. Validation-only changes can reuse the verified collection plan as described below; use a separate results directory for each replay. Collection's existing `--resume` is for a matching collector checkpoint as described below, not a general onboarding-stage resume command. ## Create the request @@ -93,7 +99,7 @@ Use an environment installed from this checkout; see [development setup](../DEVE aisimulate onboard init --interactive --output support-request.yaml ``` -Enter the actual model identifier or checkpoint path, pinned model revision, dense/MoE kind, pinned vLLM version, GPU system and interconnect. Then choose a TP size and validation workload. Setup derives and displays the GPUs required for the selected worker. Supplied options skip their prompts. Enter accepts displayed defaults; invalid values can be corrected; Ctrl-C or end-of-input cancels without saving. Existing files require `--overwrite`. +Enter the actual model identifier or checkpoint path, pinned model revision, dense/MoE kind, pinned vLLM version, GPU system and interconnect. Choose a worker and review its context, scheduler and capture limits. Setup derives and displays the GPUs required for that worker. Supplied options skip their prompts. Enter accepts displayed defaults; invalid values can be corrected; Ctrl-C or end-of-input cancels without saving. Existing files require `--overwrite`. Both guided and scripted setup use onboarding. `--profile onboarding` is an optional spelling of the same behavior. For automation, supply the identity flags directly: @@ -110,7 +116,9 @@ aisimulate onboard init \ Replace the model, runtime and hardware inputs with your deployment's values. Framework support currently selects vLLM. Optional tokenizer, chat-template, and AISimulate revisions are recorded only when supplied. -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. With `--model-config`, the default pilot context is capped at the config's known context limit, and omitted parallelism flags trigger model/hardware-aware suggestions. Headless setup selects only a fully assessed default; otherwise it exits without saving. Identifier-only and supplied-profile setup retain TP1 when parallelism is omitted. These are planning defaults, not measured model capacity or latency. Change them with the corresponding flags shown by `aisimulate onboard init --help`. +With `--model-config`, omitted parallelism flags trigger model/hardware-aware suggestions. Headless setup selects only a fully assessed default; otherwise it exits without saving. Identifier-only and supplied-profile setup retain TP1 when parallelism is omitted. Review the default limits below and change them through `aisimulate onboard init --help` options or config-based profile review. + +The optional `predict/pilot.yaml` and `recommend/pilot.yaml` examples use 1,024 input tokens, 128 output tokens, concurrency 1, four requests, TTFT target 1,000 ms and TPOT target 100 ms. Their workload/SLA flags do not size memory, choose the collector grid or change the collection identity. The example token lengths must still fit the selected runtime context. Guided intake does not prompt for these synthetic workload or SLA fields. Onboarding takes no GPU-pool or node-allocation inputs. The selected topology establishes a minimum collection requirement, not a reservation or a check that those GPUs are available. Advanced flags include `--request-count`, `--objective`, `--seed`, and `--sm`. Replica counts and optimization budgets remain configurable in ordinary `predict` and `recommend` inputs. @@ -124,6 +132,21 @@ Onboarding takes no GPU-pool or node-allocation inputs. The selected topology es The first profile implementation supports vLLM text decoders, including the text portion of multimodal checkpoints, with PP1, CP1 and linear KV storage. AFD, encoder pools, speculative decoding, nonlinear recurrent state, wide EP and EPLB need additional metadata/semantics and are rejected by this route. +### Runtime and collection limits + +These settings define the runtime envelope used for profile sizing, collection and generated configurations. They describe different constraints; none independently proves that a trace is covered. + +| Setting | Scope and initial value | +| --- | --- | +| `--context-length` | Maximum input plus output tokens for one request. Fresh config/profile-based setup uses `min(declared context, 256000)`; identifier-only setup uses 256,000 until reviewed. The saved field remains `search.context_length`. | +| `--max-num-tokens` | Scheduled token budget per attention-DP rank. Use the supplied profile's bound, or an initial policy of 8,192. Saved as `collection.max_num_tokens` when explicitly set. | +| `--max-batch-size` | Scheduler sequence bound per attention-DP rank. Use the supplied profile's bound, or an initial policy of 256. It is independent of validation concurrency. Saved as `collection.max_batch_size` when explicitly set. | +| `--max-prefill-cudagraph-size` | Prefill CUDA graph capture limit passed to the collector, initially 2,048. Match the target serving configuration. Saved as `collection.max_prefill_cudagraph_size` when explicitly set. | + +Profile resource bounds must cover the selected scheduler envelope. Conflicting declarations fail rather than silently clipping the request. With config-based setup, editing scheduler bounds recomputes dependent estimates; review those estimates again. An explicit byte override remains your declared bound until you edit it. + +The 256,000-token context policy is motivated by the [AgentX reference subset](#validate-fpm-query-coverage-with-agentx-replay), whose request limit is 256,000, not 262,144. It is a starting point for review, not an assertion that every maximum-length request fits or that the FPM table covers it. The scheduler sequence limit does not reserve maximum context for all sequences simultaneously. Automatic KV capacity comes from the remaining rank-local GPU memory after non-KV resources and any explicit CUDA graph reservation; replay then applies actual request/cache occupancy. + ## Start from a local model config Use a local Hugging Face-style JSON configuration to fill supported metadata and create the existing FPM profile: @@ -134,7 +157,7 @@ aisimulate onboard init \ --interactive --output support-request.yaml ``` -Setup reads that file without downloading a checkpoint, importing model code, constructing an analytical model, or launching GPU work. It displays source information and derived inputs, then asks for unresolved values. Missing model metadata is collected before workload options so the model's context limit can bound the validation run. Config identity hints skip their ordinary prompts; explicit CLI identity options take precedence and conflicts can be corrected. A pinned checkpoint revision, literal runtime version, GPU system and interconnect still need your input when absent. The config's SHA-256 records the local source; it is not a checkpoint revision. +Setup reads that file without downloading a checkpoint, importing model code, constructing an analytical model, or launching GPU work. It displays source information and derived inputs, then asks for unresolved values. Missing model metadata is collected before runtime/collection options so the model's context limit can bound the selected envelope. Config identity hints skip their ordinary prompts; explicit CLI identity options take precedence and conflicts can be corrected. A pinned checkpoint revision, literal runtime version, GPU system and interconnect still need your input when absent. The config's SHA-256 records the local source; it is not a checkpoint revision. With no parallelism flags, guided setup asks for missing shared precision/layout facts, shows a small topology shortlist with resource reasons, and lets you choose one worker. Enter selects a displayed default only when all required profile inputs are resolved and its estimated bytes fit the budget. Otherwise a candidate number is required before setup asks for its remaining per-rank bounds. Any explicit parallelism flag bypasses suggestions and retains the existing topology validation. The chosen topology then follows the same profile review, edit, accept and cancellation flow below. @@ -158,7 +181,7 @@ provenance: User-declared bound; replace with the actual source of your estimate Review action (accept/edit/cancel): accept ``` -Changing an input recomputes dependent estimates: for example, changing `kv_cache_dtype` updates inferred `kv_bytes_per_token`, and changing `max_num_tokens` updates inferred activation bytes. Explicit values remain in place until you edit those fields themselves. If an edit makes a required estimate unavailable, setup asks for that value before returning to review. Invalid individual answers can be corrected; an edit that conflicts with the config or complete request is rejected with the reason, and the previous profile is retained. Model identity and topology remain the declared deployment. +Changing an input recomputes dependent estimates: for example, changing `kv_cache_dtype` updates inferred `kv_bytes_per_token`, and changing `max_num_tokens` updates inferred activation bytes. The review also exposes `runtime_context_length` and `max_prefill_cudagraph_size` as editable collection settings; `context_length` edits the model profile's declared maximum. Explicit values remain in place until you edit those fields themselves. If an edit makes a required estimate unavailable, setup asks for that value before returning to review. Invalid individual answers can be corrected; an edit that conflicts with the config or complete request is rejected with the reason, and the previous profile is retained. Model identity and topology remain the declared deployment. Only an explicit `accept` saves the reviewed request; Enter alone does not accept it. You can make repeated edits before accepting. Choose `cancel` at topology selection or review, or press Ctrl-C or send end-of-input at any prompt, to exit 130 without creating the output directory or replacing an existing request, including with `--overwrite`. This review step applies to `--model-config --interactive`; scripted setup never prompts. @@ -178,7 +201,7 @@ runtime_overhead_bytes: 1073741824 comm_overhead_bytes: 268435456 kv_bytes_per_token: 262144 max_num_tokens: 8192 -max_batch_size: 1 +max_batch_size: 256 provenance: Illustrative file shape only; replace values and this explanation with their actual source. ``` @@ -193,9 +216,9 @@ aisimulate onboard init \ --tensor-parallel 4 --output support-request.yaml ``` -Replace the model, revision, runtime and hardware inputs with the worker you will actually run. `--model-config` and `--fpm-profile` are mutually exclusive; `--resource-overrides` requires `--model-config` and also works with guided setup. Flat per-rank byte overrides require explicit topology flags because the same byte bound cannot be transferred or rescaled across candidate tuples. Shared precision/layout overrides can be used for automatic suggestions. Scripted setup never reads terminal input. It exits 2 without writing a request when required inputs remain unresolved or no fully assessed automatic default exists. Validation first lists all missing or invalid target identity and workload options. Once that stage is valid, it lists unresolved profile fields, or candidate-specific gaps and exact topology flags, and asks for the necessary inputs or guided setup. +Replace the model, revision, runtime and hardware inputs with the worker you will actually run. `--model-config` and `--fpm-profile` are mutually exclusive; `--resource-overrides` requires `--model-config` and also works with guided setup. Flat per-rank byte overrides require explicit topology flags because the same byte bound cannot be transferred or rescaled across candidate tuples. Shared precision/layout overrides can be used for automatic suggestions. Scripted setup never reads terminal input. It exits 2 without writing a request when required inputs remain unresolved or no fully assessed automatic default exists. Validation first lists all missing or invalid target identity and collection options. Once that stage is valid, it lists unresolved profile fields, or candidate-specific gaps and exact topology flags, and asks for the necessary inputs or guided setup. -The flat override fields are `architecture`, `context_length`, `num_experts`, every precision and resource field shown above, `moe_backend`, `attention_backend`, and optional `provenance`. Unknown fields, duplicate fields, invalid types, unsupported values, and incompatible cache semantics are rejected. Overrides are recorded with per-field provenance rather than discarded. Profile `context_length` is the model's declared maximum; the CLI `--context-length` selects the smaller pilot limit and cannot exceed it. A scheduler envelope is per attention-DP rank and defaults to 8,192 tokens and the pilot concurrency; `max_num_tokens` and `max_batch_size` can declare different bounds. +The flat override fields are `architecture`, `context_length`, `num_experts`, every precision and resource field shown above, `moe_backend`, `attention_backend`, and optional `provenance`. Unknown fields, duplicate fields, invalid types, unsupported values, and incompatible cache semantics are rejected. Overrides are recorded with per-field provenance rather than discarded. Profile `context_length` is the model's declared maximum; the CLI `--context-length` selects a runtime limit that cannot exceed it. Scheduler limits are per attention-DP rank, as described in [Runtime and collection limits](#runtime-and-collection-limits). Derivation is deliberately limited. Unambiguous architecture, context, expert count and model kind can be read independently of memory support. The initial supported estimates are: @@ -213,7 +236,7 @@ The saved request embeds the complete profile and its provenance. The original m ### Preview and choose parallelism -`--suggest-parallel` emits one JSON report to stdout without creating a request, plan, timing data or output directory. It requires `--model-config` and the actual target identity; config hints can supply the model label and kind, but checkpoint revision and runtime version are never invented. Use your intended workload, context and shared precision declarations for meaningful resource accounting: +`--suggest-parallel` emits one JSON report to stdout without creating a request, plan, timing data or output directory. It requires `--model-config` and the actual target identity; config hints can supply the model label and kind, but checkpoint revision and runtime version are never invented. Use your intended runtime/collection limits and shared precision declarations for meaningful resource accounting: ```bash aisimulate onboard init \ @@ -222,7 +245,6 @@ aisimulate onboard init \ --model-revision YOUR_IMMUTABLE_REVISION \ --framework-version YOUR_PINNED_VLLM_VERSION \ --gpu h200_sxm --interconnect nvswitch \ - --input-tokens 1024 --output-tokens 128 --concurrency 1 \ --suggest-parallel ``` @@ -230,7 +252,7 @@ The report includes hardware/config provenance, exact `cli_flags`, required GPUs Dense decoders consider TP; MoE decoders consider pure TP, DEP and TEP. The shortlist contains at most two choices per family and six distinct tuples overall, preferring the smallest estimated fit and the next one. Widths are powers of two within the packaged node domain. A declared NVLink/NVSwitch fabric can use a packaged fast rack domain when its inter-node bandwidth matches or exceeds its intra-node bandwidth; a declared `none` interconnect permits only one GPU automatically. This is hardware architecture, not an available GPU allocation. Explicit topology choices may exceed the shortlist. -The memory precheck compares complete per-rank resource estimates against 90% of packaged GPU memory. It reserves the chosen full context for each concurrently resident sequence on every rank, using the smaller of workload concurrency and the profile's per-rank batch limit, plus one cached token per rank for planner headroom. It does not assume that attention-DP routing balances requests. Missing precision, weights, cache, activations or reservations prevent a default; a known lower bound can reject an oversized candidate but cannot prove fit. DEP/TEP communication storage, many quantized or custom decoder layouts, and large TP activation envelopes need explicit bounds after selection. CUDA graph reservations and non-text components remain outside these estimates. Only known geometry constraints are checked; generated-plan admission and actual serving-runtime checks remain authoritative. The shortlist does not establish timing coverage, optimal performance or measured accuracy. +The memory precheck compares complete per-rank resource estimates against 90% of packaged GPU memory. It reserves one full runtime context plus one cached token per rank for planner headroom. It does not multiply context by the maximum scheduler sequence count or assume balanced attention-DP routing. Actual total KV capacity is computed from remaining memory; concurrent requests must share that capacity. Missing precision, weights, cache, activations or reservations prevent a default; a known lower bound can reject an oversized candidate but cannot prove fit. DEP/TEP communication storage, many quantized or custom decoder layouts, and large TP activation envelopes need explicit bounds after selection. CUDA graph reservations and non-text components remain outside these estimates. Only known geometry constraints are checked; generated-plan admission and actual serving-runtime checks remain authoritative. The shortlist does not establish timing coverage, optimal performance or measured accuracy. ## Provide identity and resource metadata @@ -267,16 +289,29 @@ 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. 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. On an existing plan, `--overwrite` can repair missing generated files for the identical request; it 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 and effective collection limits with their sources. 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 deployment or collection envelope. -If an AISimulate update changes generated guidance, recreate the plan in a new output directory: repair compares generated files byte for byte and does not migrate existing plans. Guidance changes do not change the saved-request identity checks used by collection. +The v2 plan's `request_id` identifies collection inputs: deployment identity, exact profile and collection settings. Its separate `validation_id` also includes synthetic workload/SLA, recommendation objective and seed. To change only those validation inputs, copy the request to a new file, edit that copy, then run `onboard plan --config UPDATED_REQUEST --output-dir EXISTING_PLAN --overwrite`. The command first validates the original saved request and every existing generated file against the saved plan, then refreshes validation examples and their hashes. It can also repair missing generated files for the same collection. It preserves collected data, checkpoints and prior result directories. Changed collection/profile inputs or modified existing generated files are rejected; do not edit files inside the plan to bypass those checks. + +If an AISimulate update changes generated guidance, recreate the plan in a new output directory: validation compares generated files byte for byte and does not migrate existing plans. A validation trace is supplied separately to `validate-fpm`, so selecting another trace does not require regenerating the collection plan. The plan reports `fpm.collection_gpus_required`, derived as attention TP times attention DP. Collection uses only that worker width and the matching `tp`, `pure_tp`, `dep`, or `tep` preset. The generated prediction and recommendation configs both use one worker. Recommendation has one exact preset and one trial; its `optimization.constraints.max_candidate_gpus` equals the selected worker's requirement so even workers wider than the runtime default are representable. That derived cap describes this validation example and does not declare the user's total available GPUs. There is no replica expansion or deployment optimization during onboarding. -Earlier draft `aisimulate-support-request/v1` files included `identity.gpu_count`, `identity.node_count`, `identity.gpus_per_node` and `search.max_candidates`. The draft request schema keeps its version name but now rejects those obsolete fields explicitly. Copy an old request to a new file, remove those fields, review the retained topology/profile and regenerate the plan in a new output directory. Do not edit or overwrite the old plan, timings or checkpoints to reuse its identity. Regenerated plan summaries report `fpm.collection_gpus_required` and `search.candidates[].required_gpus` in place of `fpm.worker_gpus` and `search.candidates[].total_gpus`. The corresponding `onboard init` flags (`--gpu-count`, `--node-count`, `--gpus-per-node`, `--max-candidates`) are no longer accepted. Existing ordinary `predict` and `recommend` configurations remain valid; their replica and GPU-budget APIs are unchanged. +Saved `aisimulate-support-request/v1` requests are read as v2 while preserving their declared context and profile resource bounds; an omitted legacy context retains 16,384. Legacy validation concurrency is not reused as a collection sequence limit. Review these retained and newly independent settings before use. Old `aisimulate-support-plan/v1` directories must be regenerated in a **new directory** because their collection commands depended on the synthetic validation workload. Do not overwrite old timings or checkpoints to reuse the old plan identity. + +Some earlier v1 drafts also included `identity.gpu_count`, `identity.node_count`, `identity.gpus_per_node` and `search.max_candidates`; those obsolete fields remain rejected. Copy the request, remove them and review the retained topology/profile before regeneration. Plan summaries use `fpm.collection_gpus_required` and `search.candidates[].required_gpus`. The old `onboard init` allocation flags (`--gpu-count`, `--node-count`, `--gpus-per-node`, `--max-candidates`) are not accepted. Ordinary `predict` and `recommend` replica and GPU-budget APIs are unchanged. Before execution, prepare the real checkpoint and the pinned runtime using the existing [FPM collection guide](../python/aisimulate/docs/fpm/end-to-end-workflow.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. +For a source checkout installed with `uv sync`, activate its environment and expose the collector source package before running generated collector commands. From the repository root: + +```bash +source python/aisimulate/.venv/bin/activate +export PYTHONPATH="$PWD/python/aisimulate${PYTHONPATH:+:$PYTHONPATH}" +``` + +Alternatively, run the collector from the `python/aisimulate` source directory with that environment active. The installed release wheel includes the collector and needs no source-path adjustment. These steps make the command importable; they do not provide the target deployment or GPUs. + After those prerequisites are ready, explicitly launch collection from that environment: ```bash @@ -297,11 +332,55 @@ 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. 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. +Dynamo self-benchmark generates the actual point grid after runtime initialization, using CUDA graph sizes and available runtime bounds. The onboarding command supplies `--fpm-max-model-len`, `--fpm-max-num-batched-tokens` and `--fpm-max-num-seqs` to both prefill and decode workers. Prefill sampling uses the reviewed scheduled-token and sequence limits, plus the capture limit; `--fpm-max-prefill-isl` is a total new-token budget, not a per-request context limit. DEP does not multiply these rank-local settings by GPU count. The runtime chooses the actual sampling coordinates; onboarding does not build a separate trace-derived grid or promise a sample count or collection duration. + +Direct collector callers can set the same shared runtime flags. With a supplied FPM profile, omitted limits use that profile's bounds; without one, an omitted context retains the existing vLLM `-1` auto-fit behavior. Explicit limits must be positive and cannot exceed the supplied profile or contradict the prefill sampling bounds. 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. Validate the worker's actual FPM query shapes afterward, including summed queries for DEP. + +## Validate FPM query coverage with AgentX replay + +After verifying a matching formal FPM pair, run the existing cold aggregated replay path against a local Weka JSON or JSONL file: + +```bash +aisimulate onboard validate-fpm \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support \ + --trace /path/to/traces.jsonl \ + --validation-output-dir ./aisimulate-validation/agentx +``` + +The output directory must be separate from, and neither inside nor above, the collection directory. Use a fresh directory or `--overwrite` to replace prior validation outputs. The command verifies the saved collection plan, writes `predict.yaml`, and invokes ordinary `aisimulate predict` with per-request evidence. It does not download traces or run GPU collection. The reviewed target model, topology, runtime, profile, scheduler limits and systems paths remain fixed; source model labels in a trace do not replace the configured target model. + +The current validation scope is **cold aggregated replay, one client lane, HBM-only cache and no speculative decoding**. The cache starts cold and normal prefix reuse can accumulate during replay. Nested timestamps are interpreted relative to the root play (`nested_timestamp_basis: absolute`). The command replays the complete supplied file; selecting one complete play is useful for a first check, but does not establish coverage of a larger corpus. Preserve the full play and its dependencies when preparing a subset. Ordinary prediction's host-memory checks still apply; a larger corpus can require a larger host, and a preflight rejection remains incomplete validation. The current path builds on existing vLLM/SGLang replay support; onboarding collection remains vLLM. Seeded cache snapshots, explicit warmup and expanded replay modes are follow-up work after the relevant AgentX changes, including [#207](https://github.com/ai-dynamo/aisimulate/pull/207) and [#235](https://github.com/ai-dynamo/aisimulate/pull/235); they are not prerequisites for this workflow. + +The reference corpus is [semianalysisai/cc-traces-weka-062126-256k](https://huggingface.co/datasets/semianalysisai/cc-traces-weka-062126-256k/tree/8fecd2fc56694469f758f0afbbb6335ad3043740). To download that pinned revision separately, with Hugging Face access available: + +```bash +hf download semianalysisai/cc-traces-weka-062126-256k traces.jsonl \ + --type dataset \ + --revision 8fecd2fc56694469f758f0afbbb6335ad3043740 \ + --local-dir ./agentx-reference +``` + +The expected `traces.jsonl` SHA-256 is `e39cd2ff3eba21d4a3664be51da743ac3d2149a1933898cafc7bfeac8147eeef`. Record the file hash and selection when using a derived subset. No trace corpus is bundled or fetched automatically by AISimulate. + +### Read the coverage result + +| Artifact | Meaning | +| --- | --- | +| `predict.yaml` | Inspectable ordinary prediction input with strict direct FPM and native coverage collection enabled. | +| `validation.json` | Overall `covered` or `incomplete` result, issues, collection identity, trace/config/FPM file hashes and replay scope. Accuracy remains `not_assessed`. | +| `prediction/fpm-coverage.json` | Native lookup counts and bounded missing-query details, also saved when a timing lookup fails after replay starts. | +| `prediction/prediction.json` | Ordinary prediction report when produced, including per-request completion, AgentX graph and target-model projection. | + +`measured` means the native lookup resolved from an exact measured point; `interpolated` means it used supported direct interpolation; `unsupported` means the native query failed. Counts are **native lookup resolutions**, not requests, replay iterations or tokens. Reusing a cached timing does not make another lookup. Each role records prefill, decode and any mixed-pass decode-baseline counts. Missing entries identify phase, purpose, cell identity, coordinates, reason and occurrence count. Each model retains at most 128 distinct gaps and reports omitted gap occurrences; its total unsupported count remains complete. + +An exit status of 0 and overall `validation.json` status `covered` require a nonempty completed replay, all selected requests and plays completed with their requested output, and at least one supported native query with no unsupported queries. An early failure, failed or aborted request, empty query set or interruption cannot pass. Inspect `validation.json` first: an individual native snapshot does not certify replay completion. A missing query stops ordinary prediction and returns nonzero while retaining available evidence. That partial report describes only the attempted portion; it does not identify every missing point in the remainder of the corpus. + +For a missing decode query, inspect its exact batch size and total past-KV coordinate together. A high maximum KV value elsewhere in a table does not prove coverage for batch size 1 at that context, or for another batch/capture region. Prefill likewise needs the requested batch and prompt/KV support. Use these gaps to review the collection envelope and matching data, preserving strict direct interpolation and denied fallback. Successful coverage establishes that the selected replay could obtain timings; it does not establish predictive accuracy. Assess accuracy separately with matched silicon measurements. ## Run the generated ordinary configurations -After the selected model execution route is available and formal data collection is complete: +The generated synthetic configurations are optional small prediction/recommendation examples after the selected model route and formal data are available: ```bash aisimulate predict \ diff --git a/python/aisimulate/collector/fpm_forward/config.py b/python/aisimulate/collector/fpm_forward/config.py index 4a25f867b..49f5ef5ef 100644 --- a/python/aisimulate/collector/fpm_forward/config.py +++ b/python/aisimulate/collector/fpm_forward/config.py @@ -199,15 +199,68 @@ class FPMCollectionOptions: # vLLM 0.24 auto-fits this limit after model/CUDA-graph profiling, and # Dynamo records the resolved value in ``limits.max_model_len``. vllm_max_model_len: int = VLLM_AUTO_FIT_MAX_MODEL_LEN - max_prefill_isl: int = FPM_MAX_PREFILL_ISL + max_num_batched_tokens: int | None = None + max_num_seqs: int | None = None + max_prefill_isl: int | None = None max_prefill_batch_size: int | None = None max_prefill_cudagraph_size: int = FPM_MAX_PREFILL_CUDAGRAPH_SIZE + def __post_init__(self) -> None: + for name, value in ( + ( + "max-model-len", + None if self.vllm_max_model_len == VLLM_AUTO_FIT_MAX_MODEL_LEN else self.vllm_max_model_len, + ), + ("max-num-batched-tokens", self.max_num_batched_tokens), + ("max-num-seqs", self.max_num_seqs), + ): + if value is not None and (type(value) is not int or value < 1): + raise ValueError(f"--fpm-{name} must be a positive integer") + if ( + self.max_prefill_isl is not None + and self.max_num_batched_tokens is not None + and self.max_prefill_isl > self.max_num_batched_tokens + ): + raise ValueError("--fpm-max-prefill-isl exceeds --fpm-max-num-batched-tokens") + if ( + self.max_prefill_batch_size is not None + and self.max_num_seqs is not None + and self.max_prefill_batch_size > self.max_num_seqs + ): + raise ValueError("--fpm-max-prefill-batch-size exceeds --fpm-max-num-seqs") + self.validate_scheduler_limits() + + def validate_scheduler_limits( + self, *, profile_max_num_tokens: int | None = None, profile_max_batch_size: int | None = None + ) -> None: + """Check known effective phase limits without inventing runtime defaults.""" + shared_tokens = self.max_num_batched_tokens or profile_max_num_tokens + shared_sequences = self.max_num_seqs or profile_max_batch_size + for phase, tokens, sequences in ( + ("decode", shared_tokens, shared_sequences), + ( + "prefill", + self.max_prefill_isl or shared_tokens or FPM_MAX_PREFILL_ISL, + self.max_prefill_batch_size or shared_sequences, + ), + ): + # vLLM v0.25.1 vllm/config/scheduler.py:286 rejects token budgets below + # max_num_seqs: every scheduled sequence needs at least one token. + if tokens is not None and sequences is not None and tokens < sequences: + raise ValueError( + f"FPM {phase} max_num_batched_tokens ({tokens}) must be at least max_num_seqs ({sequences}); " + "edit the shared/prefill limits or profile bounds" + ) + @property def prefill_sampling(self) -> PrefillSamplingProfile: return PrefillSamplingProfile.build( - max_isl=self.max_prefill_isl, - max_batch_size=self.max_prefill_batch_size, + max_isl=self.max_prefill_isl + if self.max_prefill_isl is not None + else (self.max_num_batched_tokens or FPM_MAX_PREFILL_ISL), + max_batch_size=self.max_prefill_batch_size + if self.max_prefill_batch_size is not None + else self.max_num_seqs, max_cudagraph_capture_size=self.max_prefill_cudagraph_size, ) @@ -271,7 +324,14 @@ def from_args(cls, args: argparse.Namespace) -> FPMCollectionOptions: if getattr(args, "fpm_warmup_iterations", None) is None else args.fpm_warmup_iterations ), - max_prefill_isl=getattr(args, "fpm_max_prefill_isl", None) or FPM_MAX_PREFILL_ISL, + vllm_max_model_len=( + VLLM_AUTO_FIT_MAX_MODEL_LEN + if getattr(args, "fpm_max_model_len", None) is None + else args.fpm_max_model_len + ), + max_num_batched_tokens=getattr(args, "fpm_max_num_batched_tokens", None), + max_num_seqs=getattr(args, "fpm_max_num_seqs", None), + max_prefill_isl=getattr(args, "fpm_max_prefill_isl", None), max_prefill_batch_size=getattr(args, "fpm_max_prefill_batch_size", None), max_prefill_cudagraph_size=( getattr(args, "fpm_max_prefill_cudagraph_size", None) or FPM_MAX_PREFILL_CUDAGRAPH_SIZE @@ -279,7 +339,7 @@ def from_args(cls, args: argparse.Namespace) -> FPMCollectionOptions: ) def to_dict(self) -> dict[str, object]: - return { + result = { "max_gpus": self.max_gpus, "gpu_counts": list(self.gpu_counts), "parallel_presets": list(self.parallel_presets), @@ -303,6 +363,13 @@ def to_dict(self) -> dict[str, object]: "point_source": "dynamo_native_self_benchmark", "prefill_sampling": self.prefill_sampling.to_dict(), } + # Preserve the existing frozen-plan representation when the new shared + # runtime limits are absent. + if self.max_num_batched_tokens is not None: + result["max_num_batched_tokens"] = self.max_num_batched_tokens + if self.max_num_seqs is not None: + result["max_num_seqs"] = self.max_num_seqs + return result def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: @@ -398,6 +465,24 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: f"(default: {FPM_WARMUP_ITERATIONS})." ), ) + group.add_argument( + "--fpm-max-model-len", + type=_positive_int, + default=None, + help="Maximum context length for both phases; defaults to the profile context or vLLM auto-fit (-1).", + ) + group.add_argument( + "--fpm-max-num-batched-tokens", + type=_positive_int, + default=None, + help="Shared vLLM scheduled-token limit for prefill and decode; may be narrowed for prefill.", + ) + group.add_argument( + "--fpm-max-num-seqs", + type=_positive_int, + default=None, + help="Shared vLLM sequence limit for prefill and decode; may be narrowed for prefill.", + ) group.add_argument( "--fpm-max-prefill-isl", dest="fpm_max_prefill_isl", @@ -406,7 +491,8 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: help=( "Maximum total scheduled prefill new-token axis; the batch=1 points also " "cover this per-request new-token length " - f"(default: {FPM_MAX_PREFILL_ISL})." + "(defaults to the shared token limit, otherwise " + f"{FPM_MAX_PREFILL_ISL} within supplied profile bounds)." ), ) group.add_argument( @@ -414,7 +500,7 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: dest="fpm_max_prefill_batch_size", type=_positive_int, default=None, - help="Optional prefill max_num_seqs override; defaults to the Dynamo/vLLM runtime value.", + help="Optional narrower prefill max_num_seqs; defaults to shared/profile bounds or the Dynamo/vLLM value.", ) group.add_argument( "--fpm-max-prefill-cudagraph-size", @@ -535,6 +621,9 @@ def reject_fpm_arguments_without_fpm(args: argparse.Namespace) -> None: "fpm_enable_wideep", "fpm_enable_eplb", "fpm_warmup_iterations", + "fpm_max_model_len", + "fpm_max_num_batched_tokens", + "fpm_max_num_seqs", "fpm_max_prefill_isl", "fpm_max_prefill_batch_size", "fpm_max_prefill_cudagraph_size", diff --git a/python/aisimulate/collector/fpm_forward/planner.py b/python/aisimulate/collector/fpm_forward/planner.py index 97d9b403d..4ac19f1fb 100644 --- a/python/aisimulate/collector/fpm_forward/planner.py +++ b/python/aisimulate/collector/fpm_forward/planner.py @@ -15,14 +15,14 @@ import logging import os import subprocess -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path, PurePosixPath from typing import Any from aisimulate_core.sdk.fpm_profile import FpmDeploymentProfile, FpmModelProfile, load_fpm_profile from .capabilities import ModelCapabilityProfile, ResolvedDTypeProfile, resolve_model_capability -from .config import PARALLEL_AXES, FPMCollectionOptions +from .config import FPM_MAX_PREFILL_ISL, PARALLEL_AXES, VLLM_AUTO_FIT_MAX_MODEL_LEN, FPMCollectionOptions from .memory_admission import TopologyMemoryDecision, filter_memory_infeasible_topologies from .topology import enumerate_fpm_topologies, topology_strategy from .types import ParallelTopology @@ -649,7 +649,41 @@ def build_collection_plan( ) 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) + deployments = _validate_profile_identities( + profile, capability, candidate_topologies, policies, model_path, system, backend + ) + if options.vllm_max_model_len > profile.context_length: + raise ValueError( + f"--fpm-max-model-len={options.vllm_max_model_len} exceeds " + f"FPM profile context_length={profile.context_length}" + ) + max_prefill_isl = options.max_prefill_isl + if max_prefill_isl is None: + max_prefill_isl = options.max_num_batched_tokens + if max_prefill_isl is None: + max_prefill_isl = min( + [FPM_MAX_PREFILL_ISL, *(deployment.resources.max_num_tokens for deployment in deployments)] + ) + options = replace( + options, + vllm_max_model_len=( + profile.context_length + if options.vllm_max_model_len == VLLM_AUTO_FIT_MAX_MODEL_LEN + else options.vllm_max_model_len + ), + max_prefill_isl=max_prefill_isl, + ) + # Validate the shared decode envelope as well as narrower prefill + # controls before memory admission can queue or drop any deployment. + for deployment in deployments: + resources = deployment.resources + options.validate_scheduler_limits( + profile_max_num_tokens=resources.max_num_tokens, profile_max_batch_size=resources.max_batch_size + ) + resources.validate_envelope( + max_num_tokens=options.max_num_batched_tokens or max_prefill_isl, + max_batch_size=options.max_num_seqs or options.max_prefill_batch_size or resources.max_batch_size, + ) topologies, topology_memory_admission = filter_memory_infeasible_topologies( backend=backend, model_path=model_path, @@ -658,7 +692,7 @@ def build_collection_plan( topologies=candidate_topologies, max_new_tokens=options.prefill_sampling.max_total_prefill_tokens, fpm_profile=profile, - max_batch_size=options.max_prefill_batch_size, + max_batch_size=options.prefill_sampling.max_batch_size, ) weight_quantization = capability.dtype.gemm_quant_mode runnable_dtype_pairs = { @@ -763,7 +797,7 @@ def _validate_profile_identities( model_path: str, system: str, backend: str, -) -> None: +) -> tuple[FpmDeploymentProfile, ...]: """Require a resource declaration for every requested cell, before admission.""" if profile.architecture != capability.architecture: raise ValueError( @@ -782,6 +816,7 @@ def _validate_profile_identities( "enable_wideep", "enable_eplb", ) + deployments = [] for topology in topologies: deployment = profile.select( model=model_path, @@ -795,6 +830,7 @@ def _validate_profile_identities( moe_ep_size=topology.moe_ep, cp_size=topology.cp, ) + deployments.append(deployment) for kv_dtype in capability.dtype.kv_cache_dtypes: for policy in policies: resolved = [ @@ -818,3 +854,4 @@ def _validate_profile_identities( f"FPM collection profile identity mismatch: {conflicts}; supply a profile matching " "the collection configuration. The profile does not override serving dispatch." ) + return tuple(deployments) diff --git a/python/aisimulate/collector/fpm_forward/runner.py b/python/aisimulate/collector/fpm_forward/runner.py index c813fae4d..77014fdf7 100644 --- a/python/aisimulate/collector/fpm_forward/runner.py +++ b/python/aisimulate/collector/fpm_forward/runner.py @@ -935,16 +935,27 @@ def _cell_generator_overrides( "--max-model-len", str(plan.options.vllm_max_model_len), ] - if cell.workload_kind == "prefill" and not smoke: + max_num_tokens = getattr(plan.options, "max_num_batched_tokens", None) + max_num_seqs = getattr(plan.options, "max_num_seqs", None) + 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 + max_num_tokens = resources.max_num_tokens if max_num_tokens is None else max_num_tokens + max_num_seqs = resources.max_batch_size if max_num_seqs is None else max_num_seqs + if cell.workload_kind == "prefill": profile = plan.options.prefill_sampling + if not smoke or getattr(plan.options, "max_prefill_isl", None) is not None: + max_num_tokens = profile.max_total_prefill_tokens + if profile.max_batch_size is not None: + max_num_seqs = profile.max_batch_size + if cell.workload_kind == "prefill" and not smoke: compilation_config = { "cudagraph_capture_sizes": list(profile.cudagraph_capture_sizes), "max_cudagraph_capture_size": profile.max_cudagraph_capture_size, } scheduler_args.extend( [ - "--max-num-batched-tokens", - str(profile.max_total_prefill_tokens), "--compilation-config", json.dumps(compilation_config, sort_keys=True, separators=(",", ":")), "--prefill-max-new-token-samples", @@ -953,8 +964,6 @@ def _cell_generator_overrides( str(profile.max_kv_read_token_samples), ] ) - if profile.max_batch_size is not None: - scheduler_args.extend(["--max-num-seqs", str(profile.max_batch_size)]) elif smoke: if cell.workload_kind == "prefill": scheduler_args.extend( @@ -976,16 +985,12 @@ 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)]) + # Native point generation sees the runtime bounds before any cases are + # queued. Each flag is emitted once, after resolving the phase's bounds. + if max_num_tokens is not None: + scheduler_args.extend(["--max-num-batched-tokens", str(max_num_tokens)]) + if max_num_seqs is not None: + scheduler_args.extend(["--max-num-seqs", str(max_num_seqs)]) model_args = [] architecture = getattr(getattr(plan, "capability", None), "architecture", None) if architecture == "GlmMoeDsaForCausalLM": diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py index 0bbad6213..6fea69076 100644 --- a/python/aisimulate/src/aisimulate/main.py +++ b/python/aisimulate/src/aisimulate/main.py @@ -33,6 +33,7 @@ format_prediction_stdout, format_recommendation_stdout, prepare_output_directory, + write_fpm_coverage, write_prediction_report, write_recommendation_csv, write_recommendation_result, @@ -131,9 +132,14 @@ def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: capture_performance_diagnostics=bool({"time", "source"}.intersection(args.detail)), ), ) - except (KeyboardInterrupt, ResourceLimitError): - raise - except Exception as exc: + except BaseException as exc: + coverage = getattr(exc, "fpm_query_coverage", None) + if isinstance(coverage, str): + coverage = json.loads(coverage) + if isinstance(coverage, dict): + write_fpm_coverage(root, {**coverage, "status": "incomplete", "error": str(exc)}) + if isinstance(exc, (KeyboardInterrupt, ResourceLimitError)) or not isinstance(exc, Exception): + raise raise _CliExecutionError(f"{type(exc).__name__}: {exc}") from exc finally: mark_shutdown() @@ -149,6 +155,8 @@ def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: # JSON stdout, like prediction.json, must identify the approximation. native["summary"]["metric_semantics"] = report.metadata["metric_semantics"] native["summary"]["total_gpus"] = report.metadata["total_gpus"] + if isinstance(native.get("fpm_query_coverage"), dict): + write_fpm_coverage(root, native["fpm_query_coverage"]) summary = prediction_summary(native) summary.update(normalize_power_summary(report.metrics)) if "summary" in native: diff --git a/python/aisimulate/src/aisimulate/output.py b/python/aisimulate/src/aisimulate/output.py index acacd8c6a..f2503a41e 100644 --- a/python/aisimulate/src/aisimulate/output.py +++ b/python/aisimulate/src/aisimulate/output.py @@ -34,6 +34,7 @@ def prepare_output_directory(path: str | Path, *, overwrite: bool) -> Path: if overwrite: for name in ( "prediction.json", + "fpm-coverage.json", "recommendation.json", "recommendation.csv", "requests.jsonl", @@ -60,6 +61,14 @@ def write_prediction_report(root: Path, report: dict[str, Any]) -> Path: return path +def write_fpm_coverage(root: Path, coverage: dict[str, Any]) -> Path: + """Persist native lookup evidence even when the replay fails.""" + + path = root / "fpm-coverage.json" + path.write_text(json.dumps(coverage, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return path + + def write_requests(root: Path, records: list[dict[str, Any]]) -> Path: path = root / "requests.jsonl" with path.open("w", encoding="utf-8") as output: diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index e9edc50f5..fc4f6c139 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -443,7 +443,10 @@ def run( from .resources import ResourceLimitError except ImportError: raise error from None - raise ResourceLimitError(str(error)) from error + failure = ResourceLimitError(str(error)) + if hasattr(error, "fpm_query_coverage"): + failure.fpm_query_coverage = error.fpm_query_coverage + raise failure from error if not isinstance(report_json, str): raise InvalidRunnerError("AISimulate engine replay runtime report must be a JSON string") try: diff --git a/python/aisimulate/src/aisimulate/support/cli.py b/python/aisimulate/src/aisimulate/support/cli.py index 683825e1c..c91f77bf1 100644 --- a/python/aisimulate/src/aisimulate/support/cli.py +++ b/python/aisimulate/src/aisimulate/support/cli.py @@ -34,7 +34,16 @@ ) from .fpm import run_fpm from .plan import create_plan, request_id -from .schema import FPMDeployment, SearchProfile, SloSpec, SupportIdentity, SupportRequest, WorkloadSpec +from .schema import ( + AGENTX_REFERENCE_CONTEXT, + CollectionSpec, + FPMDeployment, + SearchProfile, + SloSpec, + SupportIdentity, + SupportRequest, + WorkloadSpec, +) from .topology import TopologySuggestions, suggest_topologies @@ -50,13 +59,14 @@ def add_support_parser(subparsers: Any) -> None: actions = onboard.add_subparsers(dest="support_action", required=True) init = actions.add_parser( "init", - help="Declare a model, target hardware, worker topology, and FPM workload.", + help="Declare a model, target hardware, worker topology, and FPM collection limits.", description=( "Use --interactive for terminal prompts. --model-config derives supported local model metadata; " "supply unresolved profile fields with --resource-overrides. Otherwise scripted setup requires " "--model, --model-revision, --model-kind, --framework-version, --gpu, and --interconnect. " "Collection GPUs are derived from the selected topology. Model-config setup suggests model/hardware-aware " - "topologies; other routes default to TP1. A small synthetic workload is the default. " + "topologies; other routes default to TP1. Dynamo self-benchmark generates the collection grid. " + "Synthetic workload and SLA options customize optional validation examples only. " "Set deployment replicas and GPU budgets in ordinary predict/recommend configs." ), ) @@ -98,13 +108,28 @@ def add_support_parser(subparsers: Any) -> None: action="store_true", help="Preview model-config topology suggestions as JSON; no prompts or writes, including to --output.", ) - 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).") - init.add_argument("--request-count", type=int, help="Synthetic request count (default: 4).") - init.add_argument("--context-length", type=int, help="Validation context limit in tokens (default: 16384).") - init.add_argument("--ttft-ms", type=float, help="Target time to first token in ms (default: 1000).") - init.add_argument("--tpot-ms", type=float, help="Target time per output token in ms (default: 100).") + init.add_argument("--input-tokens", type=int, help="Synthetic validation input tokens (default: 1024).") + init.add_argument("--output-tokens", type=int, help="Synthetic validation output tokens (default: 128).") + init.add_argument("--concurrency", type=int, help="Synthetic validation concurrency (default: 1).") + init.add_argument("--request-count", type=int, help="Synthetic validation request count (default: 4).") + init.add_argument( + "--context-length", + type=int, + help="Runtime per-request context limit; defaults to min(model/profile context, AgentX reference 256000).", + ) + init.add_argument( + "--max-num-tokens", type=int, help="Rank-local scheduled token budget; profile bound or initial policy 8192." + ) + init.add_argument( + "--max-batch-size", type=int, help="Rank-local scheduler sequence bound; profile bound or initial policy 256." + ) + init.add_argument( + "--max-prefill-cudagraph-size", + type=int, + help="Collector prefill CUDA graph capture limit; default 2048. Match the target runtime.", + ) + init.add_argument("--ttft-ms", type=float, help="Synthetic validation TTFT target in ms (default: 1000).") + init.add_argument("--tpot-ms", type=float, help="Synthetic validation TPOT target in ms (default: 100).") init.add_argument( "--objective", choices=get_args(SearchProfile.model_fields["objective"].annotation), @@ -145,6 +170,9 @@ def add_support_parser(subparsers: Any) -> None: deployment.add_argument("--model-cache", metavar="NAME[:MOUNT[:SUBPATH]]", help="Model-cache PVC and mount.") deployment.add_argument("--transport", choices=("nvlink", "ib", "efa")) deployment.add_argument("--image-pull-secret", help="Kubernetes secret for pulling the collector image.") + from .validation import add_validation_parser + + add_validation_parser(actions) def _values(args: argparse.Namespace, model: Any) -> dict[str, Any]: @@ -157,6 +185,7 @@ def _request_from_args(args: argparse.Namespace) -> SupportRequest: "identity": _values(args, SupportIdentity), "workload": {**_values(args, WorkloadSpec), "slo": _values(args, SloSpec)}, "search": _values(args, SearchProfile), + "collection": _values(args, CollectionSpec), **({"fpm_profile": load_yaml(args.fpm_profile)} if getattr(args, "fpm_profile", None) else {}), } ) @@ -177,6 +206,7 @@ def _request_target(path: str | Path, *, overwrite: bool) -> Path: def _write_request(request: SupportRequest, path: str | Path, *, overwrite: bool) -> Path: + request.scheduler_limits() target = _request_target(path, overwrite=overwrite) payload = yaml.safe_dump(request.model_dump(mode="json", exclude_none=True), sort_keys=False) target.parent.mkdir(parents=True, exist_ok=True) @@ -212,15 +242,18 @@ def _print(value: dict[str, Any], output_format: str = "table") -> None: "gpu": ("Target GPU system name (for example h200_sxm)", str), "interconnect": ("GPU interconnect (for example nvswitch, pcie, or none)", str), "tensor_parallel": ("Attention tensor-parallel size", int), +} +_CORRECTION_PROMPTS = { + **_PROMPTS, "input_tokens": ("Input tokens per request", int), "output_tokens": ("Output tokens per request", int), "concurrency": ("Concurrent requests", int), - "context_length": ("Validation context limit in tokens", int), + "context_length": ("Runtime per-request context limit in tokens", int), "ttft_ms": ("Target time to first token (ms)", float), "tpot_ms": ("Target time per output token (ms)", float), -} -_CORRECTION_PROMPTS = { - **_PROMPTS, + "max_num_tokens": ("Rank-local scheduled token budget", int), + "max_batch_size": ("Rank-local scheduler sequence bound", int), + "max_prefill_cudagraph_size": ("Prefill CUDA graph capture limit", int), "framework": ("Runtime (vllm)", str), "tokenizer_revision": ("Pinned tokenizer revision", str), "chat_template_revision": ("Pinned chat-template revision", str), @@ -239,7 +272,7 @@ def _prompt(args: argparse.Namespace, name: str) -> None: label, convert = _CORRECTION_PROMPTS[name] default = getattr(args, name) if default is None: - for model in (SupportIdentity, WorkloadSpec, SloSpec, SearchProfile): + for model in (SupportIdentity, WorkloadSpec, SloSpec, SearchProfile, CollectionSpec): field = model.model_fields.get(name) if field is not None and not field.is_required(): default = field.default @@ -267,7 +300,7 @@ def _require_terminal() -> None: def _guided_request(args: argparse.Namespace, *, skip_topology: bool = False) -> SupportRequest: _require_terminal() print("Onboard a model for FPM simulation on a target hardware platform.") - print("Select one TP, DEP, or TEP worker and a small synthetic workload on vLLM.") + print("Select one TP, DEP, or TEP worker and review collection limits for vLLM.") print("Enter accepts a displayed default. Ctrl-C cancels without saving. Supplied options skip their prompts.") for name in _PROMPTS: if skip_topology and name == "tensor_parallel": @@ -396,7 +429,7 @@ def _review_config_profile( draft: ProfileDraft, ) -> SupportRequest: while True: - identity, workload = request.identity, request.workload + identity = request.identity print("Review FPM profile before saving:") print(f" Modeling scope: {config.notes['modeling_scope']}") print( @@ -406,10 +439,11 @@ def _review_config_profile( ) print(f" Parallelism: {request.parallelism()}") print(f" Collection GPUs required: {request.worker_gpus} for one selected worker; availability is unchecked.") - print( - f" Workload: {workload.input_tokens}/{workload.output_tokens} tokens, concurrency {workload.concurrency}, " - f"{workload.request_count} requests; context limit {request.search.context_length} tokens" - ) + settings = request.collection_settings() + print(f" Runtime context limit: {request.search.context_length} tokens per request.") + print(f" Collection: {json.dumps(settings, sort_keys=True)}") + print(" Dynamo generates the grid; exact points and total KV capacity resolve at runtime.") + print(" Synthetic validation traffic and SLAs do not determine these collection limits.") print(" Resource *_bytes values are bytes per rank; kv_bytes_per_token is bytes per cached token per rank.") print(" max_num_tokens/max_batch_size are per-rank scheduler limits. Estimates require runtime verification.") for name, value in draft.resolved.items(): @@ -429,23 +463,43 @@ def _review_config_profile( return request if action == "cancel": raise KeyboardInterrupt - print("Editable fields: " + ", ".join(sorted(OVERRIDE_FIELDS))) + collection_fields = {"runtime_context_length", "max_prefill_cudagraph_size"} + print("Editable fields: " + ", ".join(sorted(OVERRIDE_FIELDS | collection_fields))) while True: name = input("Field to edit: ").strip() - if name in OVERRIDE_FIELDS: + if name in OVERRIDE_FIELDS | collection_fields: break print("Choose an editable field by its displayed name.") # Derive from explicit inputs again so inferred dependents can change. # Only publish the staged values after the entire request validates. - staged = {**overrides, **_read_profile_value(name)} + staged = dict(overrides) + payload = request.model_dump(exclude={"fpm_profile"}) + if name in collection_fields: + while True: + try: + value = int(input(f"{name} (positive integer): ").strip()) + if value <= 0: + raise ValueError + break + except ValueError: + print("Enter a positive integer.") + if name == "runtime_context_length": + payload["search"]["context_length"] = value + else: + payload["collection"][name] = value + else: + staged.update(_read_profile_value(name)) + if name in {"max_num_tokens", "max_batch_size"}: + payload["collection"][name] = staged[name] try: - updated = derive_profile(config, request, staged) + edited_request = SupportRequest.model_validate(payload) + updated = derive_profile(config, edited_request, staged) while updated.missing: field = min(updated.missing, key=_memory_field) - staged, updated = _prompt_profile_value(config, request, staged, updated, field) + staged, updated = _prompt_profile_value(config, edited_request, staged, updated, field) if updated.profile is None: raise ValueError("model config did not produce a complete FPM profile") - validated = SupportRequest.model_validate({**request.model_dump(), "fpm_profile": updated.profile}) + validated = SupportRequest.model_validate({**edited_request.model_dump(), "fpm_profile": updated.profile}) except ValueError as exc: print(f"Edit rejected: {exc}. Previous profile values retained.") continue @@ -483,9 +537,7 @@ def _config_inputs(args: argparse.Namespace) -> tuple[ModelConfig, dict[str, Any ("moe" if experts > 0 else "dense") if experts is not None else config.suggestions.get("model_kind") ) if args.context_length is None and "context_length" in preview.resolved: - args.context_length = min( - preview.resolved["context_length"], SearchProfile.model_fields["context_length"].default - ) + args.context_length = min(preview.resolved["context_length"], AGENTX_REFERENCE_CONTEXT) if args.interactive: request = _guided_request(args, skip_topology=not _explicit_topology(args)) else: @@ -498,7 +550,7 @@ def _config_inputs(args: argparse.Namespace) -> tuple[ModelConfig, dict[str, Any option = "--" + name.replace("_", "-") if name in _CORRECTION_PROMPTS else "request" details.append(f" {option}: {error['msg']}") raise ValueError( - "Complete target identity and workload options before dependent profile derivation:\n" + "Complete target identity and collection options before dependent profile derivation:\n" + "\n".join(details) + "\nThen supply unresolved profile fields with --resource-overrides PATH (flat JSON/YAML), " "or use --interactive." @@ -653,14 +705,13 @@ def _init(args: argparse.Namespace) -> int: except (EOFError, KeyboardInterrupt): print("Setup cancelled.", file=sys.stderr) return 130 - workload = request.workload print( f"Scope: FPM simulation for {request.identity.model} ({request.identity.model_kind}), " f"vLLM {request.identity.framework_version}, " f"{request.identity.gpu}; {request.identity.interconnect}; " - f"{request.parallel_preset} {request.parallelism()}, {workload.input_tokens}/{workload.output_tokens} tokens, " - f"concurrency {workload.concurrency}, {workload.request_count} requests." + f"{request.parallel_preset} {request.parallelism()}; runtime context {request.search.context_length} tokens." ) + print("Collection settings: " + json.dumps(request.collection_settings(), sort_keys=True)) print(f"Collection GPUs required: {request.worker_gpus} for one selected worker; availability is unchecked.") if args.model_config: print( @@ -715,6 +766,10 @@ def _plan(args: argparse.Namespace) -> int: def run_support_command(args: argparse.Namespace) -> int: + if args.support_action == "validate-fpm": + from .validation import run_validation + + return run_validation(args) if args.support_action == "init": return _init(args) if args.support_action == "plan": diff --git a/python/aisimulate/src/aisimulate/support/config_profile.py b/python/aisimulate/src/aisimulate/support/config_profile.py index c36e9e07f..e38c49fa3 100644 --- a/python/aisimulate/src/aisimulate/support/config_profile.py +++ b/python/aisimulate/src/aisimulate/support/config_profile.py @@ -689,14 +689,16 @@ def derive_profile( values[key], sources[key] = value, source for key in ("moe_backend", "attention_backend"): values[key], sources[key] = "auto", "existing FPM deployment schema default: automatic backend selection" - values["max_num_tokens"] = 8192 - sources["max_num_tokens"] = "onboarding scheduler policy: rank-local max_batched_tokens=8192" - values["max_batch_size"] = request.workload.concurrency if request else 256 - sources["max_batch_size"] = ( - "onboarding rank-local workload concurrency" - if request - else "initial onboarding max_sequences=256; resolved to workload concurrency after identity input" - ) + for key, default in (("max_num_tokens", 8192), ("max_batch_size", 256)): + declared = getattr(request.collection, key) if request is not None else None + if declared is not None and key in supplied and declared != supplied[key]: + raise ValueError(f"collection.{key} conflicts with the resource override; provide one consistent bound") + values[key] = declared if declared is not None else default + sources[key] = ( + "user collection override; per-rank runtime bound" + if declared is not None + else f"initial vLLM collection policy: rank-local {key}={default}; independent of validation traffic" + ) for key, value in supplied.items(): if key in ("architecture", "num_experts") and key in values and values[key] != value: raise ValueError(f"{key} override conflicts with the source config: {value!r} != {values[key]!r}") diff --git a/python/aisimulate/src/aisimulate/support/fpm.py b/python/aisimulate/src/aisimulate/support/fpm.py index e8ca547b9..ad60ac06c 100644 --- a/python/aisimulate/src/aisimulate/support/fpm.py +++ b/python/aisimulate/src/aisimulate/support/fpm.py @@ -31,16 +31,11 @@ def fpm_cli_args( 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") + scheduler = request.scheduler_limits() + max_prefill_tokens = scheduler["max_batched_tokens"] + max_prefill_batch = scheduler["max_sequences"] + 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", @@ -57,6 +52,12 @@ def fpm_cli_args( str(request.worker_gpus), "--fpm-parallel-presets", request.parallel_preset, + "--fpm-max-model-len", + str(request.search.context_length), + "--fpm-max-num-batched-tokens", + str(max_prefill_tokens), + "--fpm-max-num-seqs", + str(max_prefill_batch), "--fpm-max-prefill-isl", str(max_prefill_tokens), "--fpm-max-prefill-batch-size", @@ -68,6 +69,8 @@ def fpm_cli_args( "--fpm-database-root", str(root / "systems/data"), ] + if request.collection.max_prefill_cudagraph_size is not None: + command.extend(("--fpm-max-prefill-cudagraph-size", str(request.collection.max_prefill_cudagraph_size))) if profile is not None: command.extend( ( diff --git a/python/aisimulate/src/aisimulate/support/plan.py b/python/aisimulate/src/aisimulate/support/plan.py index 94aa13fca..f8d829211 100644 --- a/python/aisimulate/src/aisimulate/support/plan.py +++ b/python/aisimulate/src/aisimulate/support/plan.py @@ -10,6 +10,7 @@ import json import os import stat +import tempfile from collections.abc import Iterator from contextlib import contextmanager from pathlib import Path @@ -26,12 +27,22 @@ def request_id(request: SupportRequest) -> str: - """Bind the entire declared request, including workload and search settings.""" + """Identify reusable collection inputs, independent of evaluation traffic.""" - payload = json.dumps(request.model_dump(mode="json", exclude_none=True), sort_keys=True, separators=(",", ":")) + collection = request.model_dump(mode="json", exclude_none=True) + collection.pop("workload") + for key in ("objective", "seed"): + collection["search"].pop(key) + payload = json.dumps(collection, sort_keys=True, separators=(",", ":")) return "onboarding-" + hashlib.sha256(payload.encode()).hexdigest() +def validation_id(request: SupportRequest) -> str: + """Bind generated validation examples to the complete saved request.""" + payload = json.dumps(request.model_dump(mode="json", exclude_none=True), sort_keys=True, separators=(",", ":")) + return "validation-" + hashlib.sha256(payload.encode()).hexdigest() + + def _configs( request: SupportRequest, root: Path ) -> tuple[dict[str, Any], dict[str, dict[str, Any]], list[dict[str, int]]]: @@ -137,8 +148,9 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] prediction, recommendations, presets = _configs(request, root) commands = _commands(request, root, list(recommendations)) plan = { - "schema_version": "aisimulate-support-plan/v1", + "schema_version": "aisimulate-support-plan/v2", "request_id": request_id(request), + "validation_id": validation_id(request), "search": { "candidate_count": len(presets), "candidates": [{"parallelism": preset, "required_gpus": request.worker_gpus} for preset in presets], @@ -157,10 +169,12 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] "Actual available GPUs, placement and collector runtime resources remain unchecked." ), "plan_command": commands["fpm_plan_local"], + "runtime_limits": request.collection_settings(), "sampling": ( - "Prefill is bounded by max(2, input_tokens * concurrency) and concurrency. Decode sampling and " - "runtime context auto-fitting use the collector's existing profile. The synthetic request count " - "does not bound timing samples. Smoke and limited runs do not publish formal FPM data." + "Dynamo self-benchmark generates the prefill/decode grid from CUDA graph sizes and the reviewed " + "runtime limits. The new-token budget, per-request context, scheduler sequence bound and total KV " + "capacity are separate quantities. Exact points resolve at runtime; validation traffic and SLAs " + "do not change these collection settings. Smoke and limited runs do not publish formal FPM data." ), }, "prerequisites": [ @@ -226,7 +240,7 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] 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") + raise ValueError("declared FPM resources leave insufficient rank-local KV capacity for the runtime context") plan["resources"] = estimate plan["search"]["baseline_rule"] = ( "Selected profile deployment; declared topology and rank-local resource bounds pass CPU admission. " @@ -246,13 +260,6 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] "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." @@ -270,6 +277,9 @@ def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any] ).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() + plan["generated_files_sha256"] = { + path.as_posix(): hashlib.sha256(content).hexdigest() for path, content in documents.items() + } documents[Path("support-plan.json")] = (json.dumps(plan, indent=2, sort_keys=True) + "\n").encode() return plan, documents @@ -318,7 +328,7 @@ def _check_paths(root: Path) -> None: pending.extend(path.iterdir()) -def check_plan(request: SupportRequest, root: Path) -> None: +def check_plan(request: SupportRequest, root: Path, *, allow_missing: bool = False) -> None: """Verify saved identity before collecting or reusing any local output.""" _check_paths(root) @@ -328,27 +338,40 @@ def check_plan(request: SupportRequest, root: Path) -> None: raise ValueError( f"a readable onboarding plan is required in {root}; run aisimulate onboard plan first" ) from exc - if not isinstance(prior, dict) or prior.get("request_id") != request_id(request): + if not isinstance(prior, dict) or prior.get("schema_version") != "aisimulate-support-plan/v2": + raise ValueError( + f"legacy or invalid onboarding plan in {root}; regenerate in a new output directory. " + "Legacy plans used validation traffic to choose collection bounds and cannot be reused implicitly." + ) + if prior.get("request_id") != request_id(request): raise ValueError(f"refusing to mix a different request identity in {root}; choose a new output directory") try: saved = SupportRequest.from_yaml(root / "request.yaml") except ValueError as exc: 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") + if request_id(saved) != prior["request_id"] or validation_id(saved) != prior.get("validation_id"): + raise ValueError(f"saved request identity in {root} was modified or differs from the requested plan") + expected, documents = _plan_documents(saved, root) + if prior != expected: + raise ValueError( + f"generated plan input {root / 'support-plan.json'} was modified; choose a new output directory" + ) + # Recompute the manifest rather than trusting hashes in an edited plan. An + # evaluation-only refresh must never bless modified collector inputs. + for relative, content in documents.items(): + destination = root / relative + if destination.is_symlink(): + raise ValueError(f"refusing symlinked plan output {destination}") + if not destination.exists() and allow_missing: + continue + if not destination.is_file() or destination.read_bytes() != content: + raise ValueError( + f"generated plan input {destination} was modified or missing; choose a new output directory" + ) def create_plan(request: SupportRequest, output_dir: str | Path, *, overwrite: bool = False) -> dict[str, Any]: - """Write a new plan, or repair missing files for an identical saved request.""" + """Write a plan or safely refresh validation inputs for the same collection.""" root = Path(output_dir).expanduser().resolve() plan, documents = _plan_documents(request, root) @@ -358,23 +381,33 @@ def create_plan(request: SupportRequest, output_dir: str | Path, *, overwrite: b nonempty = any(path.name != _LOCK_NAME for path in root.iterdir()) if nonempty: if not overwrite: - raise ValueError(f"output directory {root} is nonempty; use overwrite only for the same request") - check_plan(request, root) + raise ValueError(f"output directory {root} is nonempty; use overwrite only for the same collection") + check_plan(request, root, allow_missing=True) # Inspect every generated file before writing any file. Existing data, - # checkpoints, results, and edited inputs are never replaced or deleted. + # checkpoints and results are never replaced or deleted. Only verified + # generated inputs may change after an evaluation-only request edit. pending: dict[Path, bytes] = {} for relative, content in documents.items(): destination = root / relative if destination.is_symlink(): raise ValueError(f"refusing symlinked plan output {destination}") - if not destination.exists(): + if not destination.exists() or destination.read_bytes() != content: pending[relative] = content - elif destination.read_bytes() != content: - raise ValueError(f"generated plan input {destination} was modified; choose a new output directory") (root / "systems/data").mkdir(parents=True, exist_ok=True) for relative, content in pending.items(): destination = root / relative destination.parent.mkdir(parents=True, exist_ok=True) - with destination.open("xb") as handle: - handle.write(content) + if not destination.exists(): + with destination.open("xb") as handle: + handle.write(content) + else: + temporary = None + try: + with tempfile.NamedTemporaryFile(dir=destination.parent, delete=False) as handle: + temporary = Path(handle.name) + handle.write(content) + os.replace(temporary, destination) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) return plan diff --git a/python/aisimulate/src/aisimulate/support/schema.py b/python/aisimulate/src/aisimulate/support/schema.py index 8c5c4b160..a5c14cb96 100644 --- a/python/aisimulate/src/aisimulate/support/schema.py +++ b/python/aisimulate/src/aisimulate/support/schema.py @@ -16,6 +16,8 @@ from aisimulate.config.common import PositiveFiniteFloat, PositiveStrictInt, StrictModel, load_yaml from aisimulate.fpm_profile import FpmModelProfile +AGENTX_REFERENCE_CONTEXT = 256000 + class SupportIdentity(StrictModel): model: str @@ -92,7 +94,9 @@ class SearchProfile(StrictModel): attention_data_parallel: PositiveStrictInt | None = None moe_tensor_parallel: PositiveStrictInt | None = None moe_expert_parallel: PositiveStrictInt | None = None - context_length: PositiveStrictInt = 16384 + # Historical field location retained for saved requests. This is a runtime + # limit, independent of the optional synthetic validation workload. + context_length: PositiveStrictInt = AGENTX_REFERENCE_CONTEXT objective: Literal[ "throughput", "throughput_per_gpu", @@ -106,6 +110,21 @@ class SearchProfile(StrictModel): seed: int = Field(default=42, strict=True, ge=0) +class CollectionSpec(StrictModel): + """Rank-local runtime bounds; Dynamo generates the measurement grid.""" + + max_num_tokens: PositiveStrictInt | None = None + max_batch_size: PositiveStrictInt | None = None + max_prefill_cudagraph_size: PositiveStrictInt | None = None + + @field_validator("max_num_tokens") + @classmethod + def _collector_token_minimum(cls, value: int | None) -> int | None: + if value is not None and value < 2: + raise ValueError("FPM collection requires max_num_tokens >= 2") + return value + + _DNS_LABEL = r"[a-z0-9](?:[-a-z0-9]{0,61}[a-z0-9])?" _DNS_SUBDOMAIN = rf"{_DNS_LABEL}(?:\.{_DNS_LABEL})*" @@ -143,10 +162,11 @@ def _model_cache(cls, value: str | None) -> str | None: class SupportRequest(StrictModel): - schema_version: Literal["aisimulate-support-request/v1"] = "aisimulate-support-request/v1" + schema_version: Literal["aisimulate-support-request/v2"] = "aisimulate-support-request/v2" identity: SupportIdentity workload: WorkloadSpec = Field(default_factory=WorkloadSpec) search: SearchProfile = Field(default_factory=SearchProfile) + collection: CollectionSpec = Field(default_factory=CollectionSpec) fpm_profile: FpmModelProfile | None = None @model_validator(mode="before") @@ -167,6 +187,25 @@ def _legacy_allocation(cls, values: Any) -> Any: "Copy the request, remove these fields, and regenerate the plan in a new output directory. " "Set deployment replicas and GPU budgets in ordinary predict/recommend configs." ) + if values.get("schema_version") == "aisimulate-support-request/v1": + search = values.get("search", {}) + values = {**values, "schema_version": "aisimulate-support-request/v2"} + if isinstance(search, Mapping): + search = dict(search) + search.setdefault("context_length", 16384) + values["search"] = search + # Preserve explicit legacy limits and profile resource bounds. Do + # not infer collection limits from a saved validation workload. + profile = values.get("fpm_profile") + search = values.get("search", {}) + if profile is not None and isinstance(search, Mapping) and "context_length" not in search: + context = ( + profile.get("context_length") + if isinstance(profile, Mapping) + else getattr(profile, "context_length", None) + ) + if type(context) is int and context > 0: + values = {**values, "search": {**search, "context_length": min(context, AGENTX_REFERENCE_CONTEXT)}} return values def parallelism(self) -> dict[str, int]: @@ -215,11 +254,46 @@ def profile_deployment(self): 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} + limits = { + "max_batched_tokens": self.collection.max_num_tokens + or (deployment.resources.max_num_tokens if deployment is not None else 8192), + "max_sequences": self.collection.max_batch_size + or (deployment.resources.max_batch_size if deployment is not None else 256), + } + if limits["max_batched_tokens"] < limits["max_sequences"]: + raise ValueError( + f"resolved collection max_num_tokens ({limits['max_batched_tokens']}) must be at least " + f"max_batch_size ({limits['max_sequences']}); edit the collection or profile bounds" + ) + return limits + + def collection_settings(self) -> dict[str, Any]: + """Resolved runtime inputs and the source of each proposed bound.""" + scheduler = self.scheduler_limits() + profile = self.fpm_profile is not None return { - "max_batched_tokens": min(8192, deployment.resources.max_num_tokens), - "max_sequences": min(self.workload.concurrency, deployment.resources.max_batch_size), + "context_length": self.search.context_length, + **scheduler, + "max_prefill_cudagraph_size": self.collection.max_prefill_cudagraph_size or 2048, + "sources": { + "context_length": ( + "reviewed runtime limit; model/profile context capped at the 256000-token AgentX reference " + "unless explicitly overridden" + ), + "max_batched_tokens": "user collection override" + if self.collection.max_num_tokens is not None + else "reviewed profile resource bound" + if profile + else "initial vLLM collection policy (8192); review for the target runtime", + "max_sequences": "user collection override" + if self.collection.max_batch_size is not None + else "reviewed profile resource bound" + if profile + else "initial vLLM collection policy (256); independent of validation concurrency", + "max_prefill_cudagraph_size": "user collection override" + if self.collection.max_prefill_cudagraph_size is not None + else "collector default (2048); review against the target runtime CUDA graph configuration", + }, } @model_validator(mode="after") @@ -237,7 +311,11 @@ def _shape(self) -> SupportRequest: 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() + deployment = self.profile_deployment() + scheduler = self.scheduler_limits() + deployment.resources.validate_envelope( + max_num_tokens=scheduler["max_batched_tokens"], max_batch_size=scheduler["max_sequences"] + ) 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"): diff --git a/python/aisimulate/src/aisimulate/support/topology.py b/python/aisimulate/src/aisimulate/support/topology.py index afa74017a..7fd8f201b 100644 --- a/python/aisimulate/src/aisimulate/support/topology.py +++ b/python/aisimulate/src/aisimulate/support/topology.py @@ -106,7 +106,7 @@ def to_dict(self) -> dict[str, Any]: class TopologySuggestions: config_sha256: str identity: dict[str, Any] - workload: dict[str, Any] + collection: dict[str, Any] context_length: int overrides: dict[str, Any] hardware: HardwareEnvelope @@ -119,7 +119,7 @@ def to_dict(self) -> dict[str, Any]: return { "config_sha256": self.config_sha256, "identity": dict(self.identity), - "workload": dict(self.workload), + "collection": dict(self.collection), "context_length": self.context_length, "overrides": dict(self.overrides), "hardware": asdict(self.hardware), @@ -239,10 +239,13 @@ def _assess( family, sizes, "rejected", (str(error),), None, None, None, hardware.memory_budget_bytes, None ) resolved = draft.resolved - resident = min(request.workload.concurrency, resolved["max_batch_size"]) + # Admission establishes room for one declared full-context request. The + # scheduler's sequence limit is not an allocation of full contexts; runtime + # total KV capacity is the remaining memory divided by bytes per token. + resident = 1 known = sum(resolved[field] for field in _NON_KV_BYTES if field in resolved) if "kv_bytes_per_token" in resolved: - known += resolved["kv_bytes_per_token"] * (request.search.context_length * resident + 1) + known += resolved["kv_bytes_per_token"] * (request.search.context_length + 1) complete_resources = resolved.keys() >= _RANK_BYTES estimated = known if complete_resources else None if known > hardware.memory_budget_bytes: @@ -260,7 +263,7 @@ def _assess( else: status = "estimated_fit" reasons = ( - f"Complete declared/estimated per-rank resources including full-context cache need {known} bytes, " + f"Complete declared/estimated per-rank resources including one full-context cache need {known} bytes, " f"within the 90% per-GPU budget of {hardware.memory_budget_bytes} bytes; " "this is a conservative admission precheck, not runtime qualification.", ) @@ -274,7 +277,7 @@ def suggest_topologies( ) -> TopologySuggestions: """Shortlist one worker's starting points, without user allocation or timing data. - The request supplies actual deployment identity, workload and context. Its + The request supplies actual deployment identity, collection bounds and context. Its topology is replaced during enumeration. Exact per-rank bounds and attached deployment profiles must instead use the existing explicit-topology route. """ @@ -308,10 +311,18 @@ def suggest_topologies( candidates.extend((passing or pending)[:2]) passing = [candidate for candidate in candidates if candidate.status == "estimated_fit"] default = min(passing, key=lambda candidate: candidate.required_gpus, default=None) + payload = request.model_dump() + payload["collection"].update( + {field: supplied[field] for field in ("max_num_tokens", "max_batch_size") if field in supplied} + ) + collection = SupportRequest.model_validate(payload).collection_settings() + for field, scheduler_field in (("max_num_tokens", "max_batched_tokens"), ("max_batch_size", "max_sequences")): + if field in supplied: + collection["sources"][scheduler_field] = "user resource override; shared rank-local scheduler bound" return TopologySuggestions( config_sha256=config.sha256, identity=request.identity.model_dump(mode="json"), - workload=request.workload.model_dump(mode="json"), + collection=collection, context_length=request.search.context_length, overrides=supplied, hardware=hardware, @@ -321,10 +332,10 @@ def suggest_topologies( assumptions=( hardware.assumption, "Enumerate powers-of-two widths within the fast domain; explicit topology choices may exceed it.", - "Cache reserves search.context_length tokens per concurrently resident sequence on each rank; " - "rank-local concurrency is min(workload.concurrency, max_batch_size), " - "without assuming balanced attention-DP routing. Reserve one additional cached token per rank " - "to satisfy the planner's strict context-capacity headroom requirement.", + "Check room for one full search.context_length request plus one cached token of strict capacity " + "headroom per rank. max_batch_size is a scheduler bound, not that many full-context cache allocations. " + "Total KV capacity is estimated from the remaining per-rank memory; this check does not guarantee " + "a particular concurrent workload fits or assume balanced attention-DP routing.", "The 90% byte budget reuses onboarding planning policy. CUDA graph reservations are excluded; " "generated plan and serving-runtime admission remain authoritative.", "Only known config geometry constraints are checked. Unknown architecture constraints, " diff --git a/python/aisimulate/src/aisimulate/support/validation.py b/python/aisimulate/src/aisimulate/support/validation.py new file mode 100644 index 000000000..6911cce2c --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/validation.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate collected FPM timings through ordinary cold aggregated replay.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + +import yaml + +from aisimulate.config import CorePredictionConfig +from aisimulate.config.common import load_yaml +from aisimulate.output import prepare_output_directory + +from .plan import check_plan, plan_lock, request_id +from .schema import SupportRequest + + +def add_validation_parser(actions: Any) -> None: + parser = actions.add_parser( + "validate-fpm", + help="Check collected direct-FPM query coverage with a local Weka/AgentX trace.", + description=( + "Run ordinary predict with cold aggregated replay, one client lane, HBM-only cache and no speculation. " + "The trace may contain one play or a selected corpus. All selected requests must finish and every " + "native timing lookup must resolve for coverage to pass. Accuracy is not assessed. " + "This command does not download traces or collect GPU measurements." + ), + ) + parser.add_argument("-c", "--config", required=True, help="Reviewed onboarding request.") + parser.add_argument("--output-dir", required=True, help="Existing matching collection plan directory.") + parser.add_argument("--trace", required=True, help="Local Weka JSON/JSONL file; replay its full contents.") + parser.add_argument( + "--validation-output-dir", required=True, help="Separate directory for config, replay and coverage evidence." + ) + parser.add_argument("--overwrite", action="store_true", help="Replace previous validation outputs.") + + +def _file_identity(path: Path) -> dict[str, Any]: + with path.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + return {"path": str(path), "sha256": digest, "size_bytes": path.stat().st_size} + + +def _fpm_artifacts(plan: Path) -> list[dict[str, Any]]: + return [ + _file_identity(path) for path in sorted((plan / "systems/data").rglob("fpm_forward_perf.*")) if path.is_file() + ] + + +def _prediction_config(request: SupportRequest, plan: Path, trace: Path) -> dict[str, Any]: + """Retain the reviewed deployment and replace only evaluation inputs.""" + + prediction = load_yaml(plan / "predict/pilot.yaml") + prediction["traffic"] = { + "source": { + "type": "trace", + "paths": [str(trace)], + "format": "weka", + "nested_timestamp_basis": "absolute", + }, + "load": {"type": "trace_timestamps", "agentic_lanes": 1}, + } + prediction["evaluation"] = {} + engine = prediction["engine"] + if engine["mode"] != "aggregated" or engine.get("speculation") is not None: + raise ValueError("FPM replay validation requires aggregated serving with speculation disabled") + worker = engine["workers"]["aggregated"] + timing = worker["timing"] + timing["estimation_mode"] = "fpm_interpolation" + timing["fallback_policy"] = "deny" + interpolation = timing.setdefault("estimator_config", {}).setdefault("fpm_interpolation", {}) + interpolation.update(method="direct", collect_coverage=True) + if engine["model"] != request.identity.model: + raise ValueError("saved prediction model differs from the reviewed collection identity") + CorePredictionConfig.model_validate(prediction) + return prediction + + +def _completion_issues(report: dict[str, Any]) -> list[str]: + """Check terminal evidence, including requests that a failed play never dispatched.""" + + issues = [] + graph = report.get("agentic_graph", {}) + count = graph.get("node_count", 0) + if not isinstance(count, int) or count <= 0: + return ["replay did not report a nonempty AgentX request graph"] + if report.get("num_requests") != count or report.get("completed_requests") != count: + issues.append("not every selected request completed") + records = report.get("per_request", []) + if len(records) != count or any( + row.get("terminal_status") != "completed" or row.get("output_length") != row.get("requested_output_length") + for row in records + ): + issues.append("request evidence is missing, failed, or has incomplete output") + outcomes = report.get("agentic_play_outcomes", []) + if ( + not outcomes + or len(outcomes) != graph.get("play_count") + or any(row.get("status") != "completed" or row.get("settled_at_ms") is None for row in outcomes) + ): + issues.append("not every selected play reached completed settlement") + return issues + + +def _write_report(path: Path, report: dict[str, Any]) -> None: + path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def run_validation(args: argparse.Namespace) -> int: + request = SupportRequest.from_yaml(args.config) + plan = Path(args.output_dir).expanduser().resolve() + trace = Path(args.trace).expanduser().resolve() + output = Path(args.validation_output_dir).expanduser().resolve() + if not trace.is_file() or trace.suffix.lower() not in {".json", ".jsonl"}: + raise ValueError("--trace must be a local Weka JSON or JSONL file") + if output == plan or plan in output.parents or output in plan.parents: + raise ValueError("validation output must be separate from the collection plan directory") + for source in (Path(args.config).expanduser().resolve(), trace): + if source == output or output in source.parents: + raise ValueError(f"validation output must not contain the input file {source}") + with plan_lock(plan): + check_plan(request, plan) + prediction = _prediction_config(request, plan, trace) + collection = json.loads((plan / "support-plan.json").read_text(encoding="utf-8")) + for relative in ("predict.yaml", "validation.json", "prediction"): + if (output / relative).is_symlink(): + raise ValueError(f"refusing symlinked validation output {output / relative}") + prepare_output_directory(output, overwrite=args.overwrite) + prepare_output_directory(output / "prediction", overwrite=args.overwrite) + prediction_path = output / "predict.yaml" + prediction_path.write_text(yaml.safe_dump(prediction, sort_keys=False), encoding="utf-8") + report_path = output / "validation.json" + report: dict[str, Any] = { + "schema_version": 1, + "status": "incomplete", + "accuracy": "not_assessed", + "collection_id": request_id(request), + "collection_plan": _file_identity(plan / "support-plan.json"), + "trace": _file_identity(trace), + "fpm_artifacts": _fpm_artifacts(plan), + "scope": { + "replay": "cold_aggregated", + "agentic_lanes": 1, + "nested_timestamp_basis": "absolute", + "cache": "hbm_only", + "speculation": "disabled", + "model_projection": request.identity.model, + "selected_corpus": "complete_local_file", + }, + "prediction_config": _file_identity(prediction_path), + "prediction_report": str(output / "prediction/prediction.json"), + "coverage_report": str(output / "prediction/fpm-coverage.json"), + "collection_schema_version": collection.get("schema_version"), + "issues": ["replay has not completed"], + } + _write_report(report_path, report) + # Use the existing supervised CLI: the same compiler, capacity checks, + # timing provider, source importer and host limits apply to ordinary predict. + from aisimulate.supervision import main as predict + + command = [ + "predict", + "--config", + str(prediction_path), + "--output-dir", + str(output / "prediction"), + "--capture-per-request", + "--format", + "json", + ] + if args.overwrite: + command.append("--overwrite") + try: + exit_code = predict(command) + except SystemExit as exc: + exit_code = int(exc.code) if isinstance(exc.code, int) else 1 + except KeyboardInterrupt: + report["issues"] = ["replay interrupted; coverage remains incomplete"] + _write_report(report_path, report) + raise + report["prediction_exit_code"] = exit_code + coverage_path = Path(report["coverage_report"]) + coverage = json.loads(coverage_path.read_text(encoding="utf-8")) if coverage_path.is_file() else None + report["fpm_query_coverage"] = coverage + issues = [] + if exit_code != 0: + issues.append(f"predict exited with status {exit_code}; replay coverage is incomplete") + prediction_report = Path(report["prediction_report"]) + if prediction_report.is_file(): + replay = json.loads(prediction_report.read_text(encoding="utf-8")) + issues.extend(_completion_issues(replay)) + report["agentic_graph"] = replay.get("agentic_graph") + report["agentic_model_projection"] = replay.get("agentic_model_projection") + else: + issues.append("no completed prediction report was produced") + if not coverage or coverage.get("status") != "covered": + issues.append( + "native direct-FPM query coverage did not pass; inspect the coverage report for missing coordinates" + ) + if _file_identity(trace) != report["trace"]: + issues.append("trace file changed during validation") + if _fpm_artifacts(plan) != report["fpm_artifacts"]: + issues.append("FPM library files changed during validation") + with plan_lock(plan): + check_plan(request, plan) + report["issues"] = issues + report["status"] = "covered" if not issues else "incomplete" + _write_report(report_path, report) + print(f"FPM replay coverage: {report['status']}; accuracy: not assessed. Saved {report_path}") + return 0 if not issues else (exit_code or 1) diff --git a/python/aisimulate/src/aisimulate_core/_native.pyi b/python/aisimulate/src/aisimulate_core/_native.pyi index 9d0e0098f..d9fede017 100644 --- a/python/aisimulate/src/aisimulate_core/_native.pyi +++ b/python/aisimulate/src/aisimulate_core/_native.pyi @@ -178,6 +178,10 @@ class RustForwardPassPerfModel: def tune_with_fpms(self, iterations_json: str) -> None: ... def static_phase_diagnostics(self, batch_size: int, context_length: int, prefix: int, prefill: bool) -> str: ... def diagnostics(self) -> str: ... + def fpm_query_coverage(self) -> str: ... + def predict_prefill_latency(self, batch_size: int, isl: int, prefix: int) -> float: ... + def predict_decode_latency_total(self, batch_size: int, total_past_kv_tokens: int) -> float: ... + def fpm_decode_kv_ceiling(self) -> int | None: ... def regression_store_diagnostics(self) -> str: ... def min_correction_factor(self) -> float | None: ... def max_correction_factor(self) -> float | None: ... 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 bae258a5c..f446c7869 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -325,6 +325,27 @@ def diagnostics(self) -> dict[str, Any]: """ return json.loads(self._inner.diagnostics()) + def fpm_query_coverage(self) -> dict[str, Any] | None: + """Bounded direct-FPM lookup evidence, or None when collection is disabled. + + Counts describe native lookups, including errors and mixed-pass + baselines. External timing-cache hits are not new lookups. A snapshot + does not establish replay completion or prediction accuracy. + """ + return json.loads(self._inner.fpm_query_coverage()) + + def predict_prefill_latency(self, batch_size: int, isl: int, prefix: int) -> float: + """Uncorrected native prefill timing at full input and cached-prefix lengths.""" + return self._inner.predict_prefill_latency(batch_size, isl, prefix) + + def predict_decode_latency_total(self, batch_size: int, total_past_kv_tokens: int) -> float: + """Uncorrected native decode timing at the exact past-KV batch total.""" + return self._inner.predict_decode_latency_total(batch_size, total_past_kv_tokens) + + def fpm_decode_kv_ceiling(self) -> int | None: + """Largest collected decode KV total; reading it does not record a lookup.""" + return self._inner.fpm_decode_kv_ceiling() + def regression_store_diagnostics(self) -> list[dict[str, Any]]: """Return each regression store's label, readiness, and retained count. diff --git a/python/aisimulate/tests/unit/collector/test_fpm_forward.py b/python/aisimulate/tests/unit/collector/test_fpm_forward.py index bc0ac9eb3..9b37b5de4 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_forward.py @@ -17,7 +17,12 @@ import pytest from collector.fpm_forward.capabilities import resolve_model_capability -from collector.fpm_forward.config import FPMCollectionOptions, PrefillSamplingProfile, add_fpm_arguments +from collector.fpm_forward.config import ( + FPMCollectionOptions, + PrefillSamplingProfile, + add_fpm_arguments, + reject_fpm_arguments_without_fpm, +) from collector.fpm_forward.database import ( aggregate_cell, validate_formal_database_commit, @@ -152,7 +157,7 @@ def test_options_leave_point_generation_to_dynamo(): assert options.parallel_presets == ("auto",) assert options.to_dict()["point_source"] == "dynamo_native_self_benchmark" assert options.to_dict()["measurement_repeats"] == 1 - assert options.max_prefill_isl == 8192 + assert options.max_prefill_isl is None assert options.max_prefill_batch_size is None assert options.vllm_max_model_len == -1 assert options.prefill_sampling.max_total_prefill_tokens == 8192 @@ -184,6 +189,38 @@ def test_prefill_limits_expose_no_cli_aliases(): assert "--fpm-max-prefill-bs" not in help_text +@pytest.mark.parametrize("option", ["max-model-len", "max-num-batched-tokens", "max-num-seqs"]) +@pytest.mark.parametrize("value", ["0", "-1", "1.5"]) +def test_runtime_limit_cli_requires_positive_integers(option, value): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + with pytest.raises(SystemExit) as error: + parser.parse_args([f"--fpm-{option}", value]) + assert error.value.code == 2 + + +@pytest.mark.parametrize("option", ["max-model-len", "max-num-batched-tokens", "max-num-seqs"]) +def test_runtime_limits_require_fpm_collection(option): + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + args = parser.parse_args([f"--fpm-{option}", "64"]) + args.ops = ["gemm"] + with pytest.raises(ValueError, match="FPM-only arguments require --ops fpm_forward"): + reject_fpm_arguments_without_fpm(args) + + +@pytest.mark.parametrize( + "arguments, message", + [ + ({"fpm_max_num_batched_tokens": 1024, "fpm_max_prefill_isl": 2048}, "prefill-isl exceeds"), + ({"fpm_max_num_seqs": 8, "fpm_max_prefill_batch_size": 16}, "prefill-batch-size exceeds"), + ], +) +def test_prefill_limits_cannot_exceed_shared_limits(arguments, message): + with pytest.raises(ValueError, match=message): + FPMCollectionOptions.from_args(_args(**arguments)) + + def test_prefill_sampling_profile_keeps_vllm_strides_and_exact_endpoint(): short = PrefillSamplingProfile.build(max_isl=1000, max_batch_size=16) diff --git a/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py b/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py index 2facdd7a0..c0dcf183b 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py @@ -481,7 +481,9 @@ def test_resource_envelope_errors_do_not_fail_open(field, value, no_models_or_ti profile = _profile() profile["deployments"][0]["resources"][field] = value options = replace( - FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), max_prefill_batch_size=4 + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), + max_prefill_isl=8192, + max_prefill_batch_size=4, ) with pytest.raises(ValueError, match="resource envelope exceeded"): _plan(profile, options=options) @@ -506,6 +508,9 @@ def test_profile_bounds_native_scheduling_before_cases_are_queued(phase, smoke, 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" + # A profile's declared context is now authoritative; legacy profile + # callers previously used -1 and could auto-fit beyond that contract. + assert arguments[arguments.index("--max-model-len") + 1] == "8192" def test_profile_preserves_explicit_prefill_envelope(no_models_or_timing_data): @@ -523,6 +528,159 @@ def test_profile_preserves_explicit_prefill_envelope(no_models_or_timing_data): assert arguments[arguments.index("--max-num-seqs") + 1] == expected_batch +@pytest.mark.parametrize("with_profile", [False, True]) +@pytest.mark.parametrize("narrow_prefill", [False, True]) +@pytest.mark.parametrize("smoke", [False, True]) +@pytest.mark.parametrize( + "max_tokens,max_sequences,prefill_tokens,prefill_sequences", [(4096, 64, 1024, 4), (256, 256, 128, 128)] +) +def test_cli_runtime_limits_reach_both_rendered_workers( + tmp_path, + monkeypatch, + with_profile, + narrow_prefill, + smoke, + max_tokens, + max_sequences, + prefill_tokens, + prefill_sequences, +): + """Exercise parsing, plan resolution and real Generator output without launching GPUs.""" + profile = _profile() + argv = [ + *_argv(profile), + "--fpm-max-model-len", + "4096", + "--fpm-max-num-batched-tokens", + str(max_tokens), + "--fpm-max-num-seqs", + str(max_sequences), + ] + if with_profile: + path = tmp_path / "profile.json" + path.write_text(json.dumps(profile)) + argv.extend(["--fpm-model-profile", str(path)]) + if narrow_prefill: + argv.extend( + ["--fpm-max-prefill-isl", str(prefill_tokens), "--fpm-max-prefill-batch-size", str(prefill_sequences)] + ) + if smoke: + argv.append("--smoke") + + def render_without_launch(_args, resolved): + plan, overrides = resolved + for phase in ("prefill", "decode"): + cell = next(cell for cell in plan.cells if cell.workload_kind == phase) + target = tmp_path / phase + target.mkdir() + runner._render_cell(plan, cell, target, overrides, smoke=_args.smoke) + return [] + + monkeypatch.setattr(cli, "run_resolved", render_without_launch) + assert cli.main(argv) == 0 + for phase in ("prefill", "decode"): + script = (tmp_path / phase / "run.sh").read_text() + tokens, sequences = ( + (prefill_tokens, prefill_sequences) + if narrow_prefill and phase == "prefill" + else (max_tokens, max_sequences) + ) + for option, value in ( + ("--max-model-len", 4096), + ("--max-num-batched-tokens", tokens), + ("--max-num-seqs", sequences), + ): + assert script.count(option) == 1 + assert f"{option} {value}" in script + + +@pytest.mark.parametrize("smoke", [False, True]) +def test_omitted_prefill_defaults_stay_within_selected_profile_bounds(no_models_or_timing_data, smoke): + profile = _profile() + for deployment in profile["deployments"]: + deployment["resources"].update(max_num_tokens=4096, max_batch_size=32) + plan = _plan(profile) + assert plan.options.prefill_sampling.max_total_prefill_tokens == 4096 + for cell in plan.cells: + arguments = runner._cell_generator_overrides(plan, cell, {}, smoke=smoke)["params"]["agg"]["extra_cli_args"] + for option, value in ( + ("--max-model-len", "8192"), + ("--max-num-batched-tokens", "4096"), + ("--max-num-seqs", "32"), + ): + assert arguments.count(option) == 1 + assert arguments[arguments.index(option) + 1] == value + + +@pytest.mark.parametrize( + "with_profile,resource_limits,limits", + [ + (False, {}, ["--fpm-max-num-batched-tokens", "128", "--fpm-max-num-seqs", "256"]), + (True, {}, ["--fpm-max-num-batched-tokens", "128", "--fpm-max-num-seqs", "256"]), + ( + False, + {}, + ["--fpm-max-num-batched-tokens", "4096", "--fpm-max-num-seqs", "256", "--fpm-max-prefill-isl", "128"], + ), + (False, {}, ["--fpm-max-num-seqs", "16384"]), + (True, {}, ["--fpm-max-num-batched-tokens", "128"]), + (True, {}, ["--fpm-max-prefill-isl", "128"]), + (True, {"max_num_tokens": 128, "max_batch_size": 256}, []), + ( + True, + {"max_num_tokens": 128, "max_batch_size": 256}, + ["--fpm-max-prefill-batch-size", "8"], + ), + ], +) +@pytest.mark.parametrize("smoke", [False, True]) +def test_cli_rejects_scheduler_tokens_below_sequences_before_execution( + tmp_path, monkeypatch, capsys, with_profile, resource_limits, limits, smoke +): + profile = _profile() + for deployment in profile["deployments"]: + deployment["resources"].update(resource_limits) + argv = [*_argv(profile), *limits] + if with_profile: + path = tmp_path / "profile.json" + path.write_text(json.dumps(profile)) + argv.extend(["--fpm-model-profile", str(path)]) + if smoke: + argv.append("--smoke") + monkeypatch.setattr(cli, "run_resolved", lambda *_args: pytest.fail("collection execution started")) + + with pytest.raises(SystemExit) as error: + cli.main(argv) + + assert error.value.code == 2 + message = capsys.readouterr().err + assert "max_num_batched_tokens" in message + assert "max_num_seqs" in message + + +@pytest.mark.parametrize( + "limits, message", + [ + (["--fpm-max-model-len", "8193"], "context"), + (["--fpm-max-num-batched-tokens", "8192", "--fpm-max-prefill-isl", "1024"], "resource envelope exceeded"), + (["--fpm-max-num-seqs", "128", "--fpm-max-prefill-batch-size", "4"], "resource envelope exceeded"), + (["--fpm-max-prefill-isl", "8192"], "resource envelope exceeded"), + (["--fpm-max-prefill-batch-size", "128"], "resource envelope exceeded"), + ], +) +def test_cli_rejects_profile_limit_overshoots_before_execution(tmp_path, monkeypatch, capsys, limits, message): + profile = _profile() + for deployment in profile["deployments"]: + deployment["resources"].update(max_num_tokens=4096, max_batch_size=64) + path = tmp_path / "profile.json" + path.write_text(json.dumps(profile)) + monkeypatch.setattr(cli, "run_resolved", lambda *_args: pytest.fail("collection execution started")) + with pytest.raises(SystemExit) as error: + cli.main([*_argv(profile), "--fpm-model-profile", str(path), *limits]) + assert error.value.code == 2 + assert message in capsys.readouterr().err + + @pytest.mark.parametrize("actual_version", ["0.25.1", "0.24.0"]) def test_profile_runtime_version_is_observed_before_execution(tmp_path, monkeypatch, actual_version): resource = object.__new__(runner.KubernetesCellRunner) diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py index 7e269c73f..6b838ff0b 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py @@ -769,6 +769,122 @@ def make_root(name, kind): return make_root +def test_direct_coverage_is_native_persists_errors_and_round_trips(profile_dict, measured_profile_roots): + root = measured_profile_roots("coverage", "mixed") + config = _request(profile_dict, "direct", systems_paths=[root]) + plain = RustForwardPassPerfModel.best_available(config) + assert plain.fpm_query_coverage() is None + assert ( + "collect_coverage" not in plain.diagnostics()["provenance"]["config"]["estimator_config"]["fpm_interpolation"] + ) + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = True + model = RustForwardPassPerfModel.best_available(config) + assert model.predict_prefill_latency(1, 1, 0) == plain.predict_prefill_latency(1, 1, 0) == 2.0 + assert model.predict_decode_latency_total(1, 1) == plain.predict_decode_latency_total(1, 1) == 3.0 + assert model.fpm_decode_kv_ceiling() == 1 # The fake KV=64 row is not coverage. + for _ in range(2): + with pytest.raises(PerfDataNotAvailableError, match="genuine FPM points"): + model.predict_decode_latency_total(1, 64) + report = model.fpm_query_coverage() + assert report["counting_unit"] == "native_lookup_resolutions" + assert report["queries"] == {"measured": 2, "interpolated": 0, "unsupported": 2} + assert len(report["gaps"]) == 1 + gap = report["gaps"][0] + assert gap["occurrences"] == 2 + assert gap["coordinates"] == {"batch_size": 1.0, "total_kv_read_tokens": 64.0} + assert gap["cell_identity"]["attention_backend"] == "future_attention" + assert gap["model_path"] == profile_dict["model"] + saved = json.loads(json.dumps(model.diagnostics()["provenance"]["config"])) + assert saved["estimator_config"]["fpm_interpolation"]["collect_coverage"] is True + reloaded = RustForwardPassPerfModel.best_available(saved) + assert reloaded.fpm_query_coverage()["queries"] == {"measured": 0, "interpolated": 0, "unsupported": 0} + assert reloaded.predict_decode_latency_total(1, 1) == 3.0 + assert model.fpm_query_coverage() == report + + +@pytest.mark.parametrize( + ("field", "value"), + [("estimation_mode", "auto"), ("estimation_mode", "op_level"), ("fallback_policy", "allow")], +) +def test_direct_coverage_rejects_estimator_fallback(profile_dict, field, value): + config = _request(profile_dict, "direct", **{field: value}) + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = True + with pytest.raises(ValueError, match="collect_coverage"): + _normalize(config) + + +@pytest.mark.parametrize( + ("model_id", "system", "tp", "dp", "moe_tp", "moe_ep", "gemm", "fmha"), + [ + ("MiniMaxAI/MiniMax-M2.7", "h200_sxm", 4, 1, 4, 1, "fp8_block", "bfloat16"), + ("nvidia/GLM-5.2-NVFP4", "b200_sxm", 1, 8, 1, 8, "nvfp4", "fp8"), + ("nvidia/GLM-5.2-NVFP4", "b200_sxm", 8, 1, 1, 8, "nvfp4", "fp8"), + ], + ids=["tp4", "dep8", "tep8"], +) +def test_coverage_wrappers_preserve_existing_engine_queries_on_real_cells( + profile_dict, monkeypatch, model_id, system, tp, dp, moe_tp, moe_ep, gemm, fmha +): + import pyarrow.parquet as pq + + import aisimulate_core + + # Retained real timing identity; synthetic resource bounds isolate this + # regression from GPU memory qualification and analytical model classes. + monkeypatch.setenv("AIC_ALLOW_UNLISTED_VERSIONS", "1") + monkeypatch.setattr(engine, "get_model", _fail_graph) + monkeypatch.setattr(engine, "build_model_config", _fail_graph) + profile_dict["model"] = model_id + 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, + ) + config = _request(profile_dict, "direct") + config["estimator_config"]["fpm_interpolation"]["collect_coverage"] = True + covered = RustForwardPassPerfModel.best_available(config) + original = aisimulate_core.AicEngine.from_spec( + engine.compile_engine( + model_id, + system, + "vllm", + "0.25.1", + tp_size=tp, + attention_dp_size=dp, + moe_tp_size=moe_tp, + moe_ep_size=moe_ep, + forward_model="fpm", + fpm_profile=profile_dict, + fpm_interpolation="direct", + ) + ) + data = Path(aisimulate_core.__file__).parent / "systems/data" / system / "vllm/0.25.1/fpm_forward_perf.parquet" + rows = pq.read_table( + data, filters=[("model_path", "=", model_id), ("tp", "=", tp), ("dp", "=", dp), ("batch_size", "=", 1)] + ).to_pylist() + for cached in (False, True): + row = next( + row for row in rows if row["workload_kind"] == "prefill" and bool(row["total_kv_read_tokens"]) == cached + ) + prefix = row["total_kv_read_tokens"] + isl = row["total_prefill_tokens"] + prefix + assert covered.predict_prefill_latency(1, isl, prefix) == original.predict_prefill_latency(1, isl, prefix) + assert covered.predict_prefill_latency(1, isl, prefix) == pytest.approx(row["latency_ms"]) + row = next(row for row in rows if row["workload_kind"] == "decode" and row["kv_seed_regime"] != "fake_fallback") + assert covered.predict_decode_latency_total( + 1, row["total_kv_read_tokens"] + ) == original.predict_decode_latency_total(1, row["total_kv_read_tokens"]) + assert covered.predict_decode_latency_total(0, 0) == original.predict_decode_latency_total(0, 0) == 0.0 + assert covered.fpm_decode_kv_ceiling() == original.fpm_decode_kv_ceiling() + report = covered.fpm_query_coverage() + assert report["queries"] == {"measured": 5, "interpolated": 0, "unsupported": 0} + + @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") diff --git a/python/aisimulate/tests/unit/test_onboard_model_config.py b/python/aisimulate/tests/unit/test_onboard_model_config.py index 2200ef303..0c8bccdfd 100644 --- a/python/aisimulate/tests/unit/test_onboard_model_config.py +++ b/python/aisimulate/tests/unit/test_onboard_model_config.py @@ -17,6 +17,8 @@ import aisimulate.main as cli from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.support.fpm import fpm_cli_args +from aisimulate.support.plan import request_id from aisimulate.support.schema import SupportRequest pytestmark = pytest.mark.unit @@ -68,6 +70,186 @@ } +def test_validation_changes_preserve_profile_and_collection_command(tmp_path, capsys): + source, resources = _files(tmp_path) + requests = [] + for name, validation in ( + ("small", {}), + ( + "large", + {"input_tokens": 4096, "output_tokens": 256, "concurrency": 16, "request_count": 32, "ttft_ms": 5000}, + ), + ): + output = tmp_path / f"{name}.yaml" + assert cli.main(_args(output, source, resources, **validation)) == 0 + requests.append(SupportRequest.from_yaml(output)) + first, second = requests + assert first.workload != second.workload + assert first.fpm_profile == second.fpm_profile + assert request_id(first) == request_id(second) + assert fpm_cli_args(first, output_dir=tmp_path / "plan", plan_only=True) == fpm_cli_args( + second, output_dir=tmp_path / "plan", plan_only=True + ) + assert first.scheduler_limits() == {"max_batched_tokens": 4096, "max_sequences": 8} + + +@pytest.mark.parametrize("model_context,expected", [(32768, 32768), (262144, 256000)]) +def test_model_context_and_reference_cap_define_initial_runtime_context(tmp_path, model_context, expected): + source, resources = _files(tmp_path, config={**_CONFIG, "max_position_embeddings": model_context}) + output = tmp_path / "request.yaml" + assert cli.main(_args(output, source, resources)) == 0 + request = SupportRequest.from_yaml(output) + assert request.search.context_length == expected + assert request.fpm_profile.context_length == model_context + command = fpm_cli_args(request, output_dir=tmp_path / "plan", plan_only=True) + assert command[command.index("--fpm-max-model-len") + 1] == str(expected) + assert command[command.index("--fpm-max-prefill-isl") + 1] == "4096" + + +def test_guided_collection_edit_reestimates_resources_before_acceptance(tmp_path, monkeypatch): + source, resources = _files( + tmp_path, + config={**_CONFIG, "hidden_size": 4096, "num_attention_heads": 32, "intermediate_size": 8192}, + overrides={"fmha_quant_mode": "bfloat16", "comm_quant_mode": "half", "kv_cache_dtype": "bfloat16"}, + ) + output = tmp_path / "request.yaml" + prompts = _terminal( + monkeypatch, + [ + "edit", + "max_num_tokens", + "4096", + "edit", + "runtime_context_length", + "8192", + "edit", + "max_prefill_cudagraph_size", + "512", + "accept", + ], + ) + assert cli.main(_args(output, source, resources, max_num_tokens=16384) + ["--interactive"]) == 0 + request = SupportRequest.from_yaml(output) + profile = request.profile_deployment() + assert request.search.context_length == 8192 + assert request.collection.max_num_tokens == profile.resources.max_num_tokens == 4096 + assert request.collection.max_prefill_cudagraph_size == 512 + assert profile.resources.activations_bytes == 2 * 4096 * 4096 * 11 + assert request.scheduler_limits() == {"max_batched_tokens": 4096, "max_sequences": 256} + assert not any("Input tokens" in prompt or "Concurrent requests" in prompt for prompt in prompts) + command = fpm_cli_args(request, output_dir=tmp_path / "plan", plan_only=True) + for option, expected in ( + ("--fpm-max-model-len", "8192"), + ("--fpm-max-num-batched-tokens", "4096"), + ("--fpm-max-num-seqs", "256"), + ("--fpm-max-prefill-cudagraph-size", "512"), + ): + assert command[command.index(option) + 1] == expected + + +def test_conflicting_resource_and_collection_bounds_fail_before_saving(tmp_path): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + with pytest.raises(SystemExit): + cli.main(_args(output, source, resources, max_num_tokens=8192)) + assert not output.exists() + + +@pytest.mark.parametrize( + "route,limits", + [ + (route, limits) + for route in ("identifier", "profile", "config") + for limits in ({"max_num_tokens": 128, "max_batch_size": 256}, {"max_num_tokens": 128}) + ] + + [("profile", {}), ("config", {})], +) +def test_invalid_resolved_scheduler_cannot_be_saved_or_planned(tmp_path, capsys, route, limits): + source, resources = _files(tmp_path, overrides={**_OVERRIDES, "max_num_tokens": 8192, "max_batch_size": 256}) + baseline = tmp_path / "baseline.yaml" + assert cli.main(_args(baseline, source, resources)) == 0 + payload = yaml.safe_load(baseline.read_text()) + profile = payload["fpm_profile"] + output = tmp_path / "invalid.yaml" + if route == "config": + resources.write_text(yaml.safe_dump({**_OVERRIDES, "max_num_tokens": 128, "max_batch_size": 256})) + command = _args(output, source, resources, **limits) + else: + command = _args(output, source, None, model=_CONFIG["_name_or_path"], model_kind="dense", **limits) + command.remove(str(source)) + command.remove("--model-config") + if route == "profile": + if not limits: + profile["deployments"][0]["resources"]["max_num_tokens"] = 128 + profile_path = tmp_path / "profile.json" + profile_path.write_text(json.dumps(profile)) + command.extend(["--fpm-profile", str(profile_path)]) + + with pytest.raises(SystemExit) as error: + cli.main(command) + assert error.value.code == 2 + message = capsys.readouterr().err + assert "max_num_tokens" in message and "max_batch_size" in message + assert not output.exists() + + payload["collection"] = limits + if route == "identifier": + payload.pop("fpm_profile") + elif not limits: + payload["fpm_profile"]["deployments"][0]["resources"]["max_num_tokens"] = 128 + request_path = tmp_path / "invalid-saved.yaml" + request_path.write_text(yaml.safe_dump(payload)) + original = request_path.read_bytes() + plan = tmp_path / "invalid-plan" + with pytest.raises(SystemExit) as error: + cli.main(["onboard", "plan", "--config", str(request_path), "--output-dir", str(plan)]) + assert error.value.code == 2 + message = capsys.readouterr().err + assert "max_num_tokens" in message and "max_batch_size" in message + assert request_path.read_bytes() == original + assert not plan.exists() + + +@pytest.mark.parametrize("suggest_parallel", [False, True]) +def test_config_small_token_budget_uses_reviewed_sequence_bound(tmp_path, capsys, suggest_parallel): + overrides = {**_OVERRIDES, "max_num_tokens": 128} + if suggest_parallel: + # Automatic topology selection requires shared inputs rather than rank-local byte overrides. + overrides = {key: value for key, value in overrides.items() if not key.endswith("_bytes")} + overrides.pop("kv_bytes_per_token") + source, resources = _files(tmp_path, overrides=overrides) + output = tmp_path / "request.yaml" + command = _args(output, source, resources, max_num_tokens=128, tensor_parallel=None if suggest_parallel else 1) + if suggest_parallel: + command.append("--suggest-parallel") + + assert cli.main(command) == 0 + + if suggest_parallel: + assert json.loads(capsys.readouterr().out)["default"] is not None + assert not output.exists() + else: + request = SupportRequest.from_yaml(output) + assert request.scheduler_limits() == {"max_batched_tokens": 128, "max_sequences": 8} + root = tmp_path / "plan" + assert cli.main(["onboard", "plan", "--config", str(output), "--output-dir", str(root)]) == 0 + command = fpm_cli_args(request, output_dir=root, plan_only=True) + assert command[command.index("--fpm-max-num-batched-tokens") + 1] == "128" + assert command[command.index("--fpm-max-num-seqs") + 1] == "8" + + +def test_guided_config_intake_needs_no_validation_lengths_concurrency_or_slas(tmp_path, monkeypatch): + source, resources = _files(tmp_path) + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["accept"]) + omitted = dict.fromkeys(("input_tokens", "output_tokens", "concurrency", "ttft_ms", "tpot_ms")) + assert cli.main(_args(output, source, resources, **omitted) + ["--interactive"]) == 0 + assert prompts == ["Review action (accept/edit/cancel): "] + saved = SupportRequest.from_yaml(output) + assert saved.workload.concurrency == 1 + assert saved.scheduler_limits()["max_sequences"] == 8 + + def _files(tmp_path, *, config=None, overrides=None, suffix="yaml"): source = tmp_path / "config.json" source.write_text(json.dumps(_CONFIG if config is None else config)) @@ -146,7 +328,7 @@ def test_scripted_config_embeds_complete_profile_and_sources(tmp_path, monkeypat assert request.identity.model_revision == _IDENTITY["model_revision"] assert profile.architecture == "LlamaForCausalLM" assert profile.context_length == 32768 - assert request.search.context_length == 16384 + assert request.search.context_length == 32768 assert profile.num_experts == 0 assert profile.deployments[0].resources.weights_bytes == _OVERRIDES["weights_bytes"] provenance = profile.provenance + profile.deployments[0].resources.provenance @@ -729,7 +911,7 @@ def respond(prompt): resources = SupportRequest.from_yaml(output).profile_deployment().resources assert resources.kv_bytes_per_token == 512 assert resources.max_num_tokens == 8192 - assert resources.max_batch_size == 1 + assert resources.max_batch_size == 256 def test_guided_memory_accepts_exact_units_and_reprompts_invalid_quantities(tmp_path, monkeypatch, capsys): diff --git a/python/aisimulate/tests/unit/test_onboard_topology.py b/python/aisimulate/tests/unit/test_onboard_topology.py index 4a57f7f6c..d440316a7 100644 --- a/python/aisimulate/tests/unit/test_onboard_topology.py +++ b/python/aisimulate/tests/unit/test_onboard_topology.py @@ -162,8 +162,8 @@ def test_real_cli_preview_is_json_and_never_validates_or_writes_output_target(tm assert report["identity"]["model"] == "deployment/actual-model" assert report["identity"]["model_revision"] == _OPTIONS["model_revision"] assert report["identity"]["framework_version"] == _OPTIONS["framework_version"] - assert report["workload"]["request_count"] == 3 - assert report["workload"]["concurrency"] == 2 + assert "workload" not in report + assert report["collection"]["max_sequences"] == 256 assert report["context_length"] == 4096 assert report["config_sha256"] == hashlib.sha256(source.read_bytes()).hexdigest() assert {path.relative_to(tmp_path): path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} == before diff --git a/python/aisimulate/tests/unit/test_support_cli.py b/python/aisimulate/tests/unit/test_support_cli.py index a16eef09b..582f589a3 100644 --- a/python/aisimulate/tests/unit/test_support_cli.py +++ b/python/aisimulate/tests/unit/test_support_cli.py @@ -189,6 +189,24 @@ def unexpected_dispatch(*args, **kwargs): assert not list(tmp_path.iterdir()) +@pytest.mark.parametrize("search", [123, True, [], None, [["context_length", 4096]]]) +def test_v1_request_rejects_malformed_search_without_writing_plan(tmp_path, capsys, search): + source = tmp_path / "request.yaml" + source.write_text( + yaml.safe_dump({"schema_version": "aisimulate-support-request/v1", "identity": _REQUIRED, "search": search}) + ) + original = source.read_bytes() + output = tmp_path / "plan" + + with pytest.raises(SystemExit) as result: + cli.main(["onboard", "plan", "--config", str(source), "--output-dir", str(output)]) + + assert result.value.code == 2 + assert "search" in capsys.readouterr().err + assert source.read_bytes() == original + assert not output.exists() + + def test_guided_and_scripted_setup_produce_the_same_request(tmp_path, monkeypatch, capsys) -> None: guided = tmp_path / "guided request.yaml" scripted = tmp_path / "scripted.yaml" @@ -201,8 +219,10 @@ def test_guided_and_scripted_setup_produce_the_same_request(tmp_path, monkeypatc assert cli.main(_init_args(scripted, tensor_parallel=2)) == 0 assert SupportRequest.from_yaml(guided) == SupportRequest.from_yaml(scripted) - assert any("Input tokens per request [1024]" in prompt for prompt in prompts) - assert any("Target time to first token (ms) [1000.0]" in prompt for prompt in prompts) + assert not any("tokens" in prompt.lower() or "concurrent" in prompt.lower() for prompt in prompts) + assert not any( + "time to first token" in prompt.lower() or "time per output token" in prompt.lower() for prompt in prompts + ) request = SupportRequest.from_yaml(guided) assert request.workload.request_count == 4 assert request.worker_gpus == 2 diff --git a/python/aisimulate/tests/unit/test_support_config_profile.py b/python/aisimulate/tests/unit/test_support_config_profile.py index 243bd0626..3c44039f4 100644 --- a/python/aisimulate/tests/unit/test_support_config_profile.py +++ b/python/aisimulate/tests/unit/test_support_config_profile.py @@ -266,7 +266,7 @@ def test_dense_tensor_counts_include_replicated_norms_and_kv_heads(tmp_path, tp, assert resources.weights_bytes == weights assert resources.kv_bytes_per_token == kv assert resources.max_num_tokens == 8192 - assert resources.max_batch_size == 2 + assert resources.max_batch_size == 256 assert "estimate" in draft.sources["weights_bytes"] assert "exact" in draft.sources["kv_bytes_per_token"] assert "estimate" in draft.sources["runtime_overhead_bytes"] diff --git a/python/aisimulate/tests/unit/test_support_plan.py b/python/aisimulate/tests/unit/test_support_plan.py index 0b991d83e..5db657cb0 100644 --- a/python/aisimulate/tests/unit/test_support_plan.py +++ b/python/aisimulate/tests/unit/test_support_plan.py @@ -23,7 +23,7 @@ from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig from aisimulate.support.fpm import run_fpm -from aisimulate.support.plan import create_plan, plan_lock, request_id +from aisimulate.support.plan import check_plan, create_plan, plan_lock, request_id, validation_id from aisimulate.support.schema import SupportRequest pytestmark = pytest.mark.unit @@ -274,8 +274,8 @@ def test_plan_collects_one_chosen_worker_and_generates_validation_configs(tmp_pa assert command[command.index("--fpm-gpu-counts") + 1] == "4" assert command[command.index("--fpm-max-gpus") + 1] == "4" assert command[command.index("--fpm-parallel-presets") + 1] == preset - assert command[command.index("--fpm-max-prefill-isl") + 1] == "1024" - assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "1" + assert command[command.index("--fpm-max-prefill-isl") + 1] == "8192" + assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "256" assert "--plan-only" in command assert "--execute" in commands["fpm_run_local"] @@ -480,22 +480,24 @@ def create(self, worker_id): assert timing.systems_paths == prediction.engine.systems_paths -def test_prefill_bounds_follow_concurrent_workload_and_collector_minimum(tmp_path): +def test_collection_bounds_are_independent_of_synthetic_validation(tmp_path): request = _request(workload={"input_tokens": 10, "concurrency": 3, "request_count": 4}) command = create_plan(request, tmp_path)["fpm"]["plan_command"] - assert command[command.index("--fpm-max-prefill-isl") + 1] == "30" - assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "3" + assert command[command.index("--fpm-max-prefill-isl") + 1] == "8192" + assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "256" + assert command[command.index("--fpm-max-model-len") + 1] == "256000" small = _request(workload={"input_tokens": 1}) command = create_plan(small, tmp_path / "small")["fpm"]["plan_command"] - assert command[command.index("--fpm-max-prefill-isl") + 1] == "2" + assert command[command.index("--fpm-max-prefill-isl") + 1] == "8192" @pytest.mark.parametrize( "updates", [ {"search": {"tensor_parallel": 2}}, - {"search": {"seed": 7}}, - {"workload": {"input_tokens": 64}}, + {"search": {"context_length": 65536}}, + {"collection": {"max_num_tokens": 16384}}, + {"collection": {"max_prefill_cudagraph_size": 512}}, {"identity": {"model_revision": "other-checkpoint"}}, {"identity": {"framework_version": "0.26.0"}}, {"identity": {"tokenizer_revision": "tokenizer-v2"}}, @@ -530,6 +532,81 @@ def test_same_request_overwrite_preserves_timings_and_rejects_modified_inputs(tm assert data.read_bytes() == b"timings" +@pytest.mark.parametrize( + "updates", + [ + {"workload": {"input_tokens": 4096, "output_tokens": 256, "concurrency": 8, "request_count": 16}}, + {"workload": {"slo": {"ttft_ms": 2500.0, "tpot_ms": 75.0}}}, + {"search": {"seed": 7, "objective": "goodput"}}, + ], +) +def test_validation_refresh_keeps_collection_artifacts_and_updates_examples(tmp_path, updates): + request = _request() + original = create_plan(request, tmp_path) + data = tmp_path / "systems/data/collected.parquet" + data.write_bytes(b"collected timings") + checkpoint = tmp_path / "fpm-checkpoint" + checkpoint.mkdir() + (checkpoint / "campaign.json").write_bytes(b"immutable checkpoint") + changed = _request(**updates) + assert request_id(changed) == request_id(request) + assert validation_id(changed) != validation_id(request) + + updated = create_plan(changed, tmp_path, overwrite=True) + + assert updated["request_id"] == original["request_id"] + assert updated["fpm"] == original["fpm"] + assert updated["validation_id"] != original["validation_id"] + assert data.read_bytes() == b"collected timings" + assert (checkpoint / "campaign.json").read_bytes() == b"immutable checkpoint" + assert SupportRequest.from_yaml(tmp_path / "request.yaml") == changed + prediction = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(tmp_path / "recommend/pilot.yaml") + assert prediction.traffic.source.input_tokens == changed.workload.input_tokens + assert prediction.traffic.load.concurrency == changed.workload.concurrency + assert recommendation.optimizer.seed == changed.search.seed + check_plan(changed, tmp_path) + + +@pytest.mark.parametrize("relative", ["request.yaml", "predict/pilot.yaml", "commands.json", "support-plan.json"]) +def test_validation_refresh_does_not_bless_tampered_generated_files(tmp_path, relative): + create_plan(_request(), tmp_path) + target = tmp_path / relative + target.write_bytes(target.read_bytes() + b"\n") + before = _file_contents(tmp_path) + with pytest.raises(ValueError, match="modified|identity"): + create_plan(_request(workload={"input_tokens": 4096}), tmp_path, overwrite=True) + assert _file_contents(tmp_path) == before + + +@pytest.mark.parametrize( + "search,expected_context", + [(None, 16384), ({}, 16384), ({"tensor_parallel": 4, "context_length": 8192}, 8192)], +) +def test_v1_request_migration_preserves_declared_context_but_requires_a_new_plan(tmp_path, search, expected_context): + import yaml + + legacy = _request().model_dump(mode="json", exclude_none=True) + legacy["schema_version"] = "aisimulate-support-request/v1" + legacy.pop("collection") + if search is None: + legacy.pop("search") + else: + legacy["search"] = search + source = tmp_path / "old-request.yaml" + source.write_text(yaml.safe_dump(legacy)) + migrated = SupportRequest.from_yaml(source) + assert migrated.schema_version == "aisimulate-support-request/v2" + assert migrated.search.context_length == expected_context + assert migrated.scheduler_limits() == {"max_batched_tokens": 8192, "max_sequences": 256} + root = tmp_path / "old-plan" + root.mkdir() + (root / "support-plan.json").write_text(json.dumps({"schema_version": "aisimulate-support-plan/v1"})) + with pytest.raises(ValueError, match="new output directory"): + create_plan(migrated, root, overwrite=True) + assert create_plan(migrated, tmp_path / "new-plan")["fpm"]["runtime_limits"]["context_length"] == expected_context + + @pytest.mark.parametrize("modification", ["retired_command", "missing_execute", "formatting"]) def test_overwrite_rejects_modified_commands_before_repairing_missing_files(tmp_path, modification): request = _request() diff --git a/python/aisimulate/tests/unit/test_support_topology.py b/python/aisimulate/tests/unit/test_support_topology.py index 156ca30f7..f64cd9401 100644 --- a/python/aisimulate/tests/unit/test_support_topology.py +++ b/python/aisimulate/tests/unit/test_support_topology.py @@ -262,7 +262,7 @@ def test_known_weight_lower_bound_can_reject_without_precision_cache_or_overhead } <= rejected.missing.keys() -def test_cache_precheck_uses_full_context_and_conservative_rank_local_concurrency(tmp_path): +def test_cache_precheck_requires_one_full_context_independent_of_scheduler_batch_limit(tmp_path): config = _config(tmp_path, architectures=["MixtralForCausalLM"], model_type="mixtral", num_local_experts=4) request = _request("moe", concurrency=8) report = suggest_topologies(config, request, _runtime(max_batch_size=3, max_num_tokens=64)) @@ -272,22 +272,23 @@ def test_cache_precheck_uses_full_context_and_conservative_rank_local_concurrenc for field in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes") if field in dep4.draft.resolved ) - assert dep4.resident_sequences_per_rank == 3 - assert dep4.known_required_bytes - known_non_kv == 64 * (4096 * 3 + 1) + assert dep4.resident_sequences_per_rank == 1 + assert dep4.known_required_bytes - known_non_kv == 64 * (4096 + 1) assert dep4.draft.resolved["max_num_tokens"] == 64 assert report.context_length == 4096 - assert report.workload["concurrency"] == 8 - assert any("without assuming balanced attention-DP routing" in item for item in report.assumptions) + assert report.collection["max_sequences"] == 3 + assert report.collection["max_batched_tokens"] == 64 + assert "workload" not in report.to_dict() + assert any("not guarantee" in item for item in report.assumptions) -def test_context_and_concurrency_can_rule_out_previously_small_workers(tmp_path, monkeypatch): +def test_validation_concurrency_does_not_change_topology_fit(tmp_path, monkeypatch): _hardware(tmp_path, monkeypatch, gpu={"mem_capacity": 85_000_000}) config = _config(tmp_path) light = suggest_topologies(config, _request(concurrency=1), _runtime()) assert light.default.required_gpus == 1 heavy = suggest_topologies(config, _request(concurrency=16), _runtime()) - assert heavy.default.required_gpus == 2 - assert heavy.rejected_candidates[0].required_gpus == 1 + assert heavy.to_dict() == light.to_dict() @pytest.mark.parametrize("extra_token", [0, 1]) diff --git a/python/aisimulate/tests/unit/test_support_validation.py b/python/aisimulate/tests/unit/test_support_validation.py new file mode 100644 index 000000000..ee210c480 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_validation.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Onboarding validates actual replay lookups, including failed and partial runs.""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +import yaml + +from aisimulate import main as cli +from aisimulate import supervision +from aisimulate.support.plan import create_plan +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def validation_case(tmp_path, monkeypatch, request): + import pyarrow as pa + import pyarrow.parquet as pq + + from aisimulate_core.sdk import engine + + def reject_graph(*_args, **_kwargs): + raise AssertionError("FPM replay must not construct an analytical model") + + monkeypatch.setattr(engine, "get_model", reject_graph) + monkeypatch.setattr(engine, "build_model_config", reject_graph) + profile = { + "schema_version": 1, + "model": "test/onboarding-coverage", + "model_revision": "synthetic-v1", + "architecture": "UnregisteredDecoderForCausalLM", + "context_length": 512, + "num_experts": 0, + "provenance": "Synthetic test metadata; not silicon qualification.", + "deployments": [ + { + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "tp": 2, + "dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "gemm_quant_mode": "fp8", + "moe_quant_mode": "fp8", + "fmha_quant_mode": "bfloat16", + "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": 128, + "max_batch_size": 2, + "provenance": "Declared synthetic rank-local bounds.", + }, + } + ], + } + topology = getattr(request, "param", "tp") + parallel = {"tensor_parallel": 2, "context_length": 512} + if topology in {"dep", "tep"}: + profile["num_experts"] = 4 + profile["deployments"][0].update(tp=1 if topology == "dep" else 2, dp=2 if topology == "dep" else 1, moe_ep=2) + parallel.update( + tensor_parallel=profile["deployments"][0]["tp"], + attention_data_parallel=profile["deployments"][0]["dp"], + moe_tensor_parallel=1, + moe_expert_parallel=2, + ) + request = SupportRequest.model_validate( + { + "identity": { + "model": profile["model"], + "model_revision": profile["model_revision"], + "model_kind": "dense" if topology == "tp" else "moe", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "interconnect": "nvswitch", + }, + "search": parallel, + "collection": {"max_num_tokens": 128, "max_batch_size": 2}, + "workload": {"input_tokens": 64, "output_tokens": 3, "request_count": 1}, + "fpm_profile": profile, + } + ) + plan = tmp_path / "collection" + create_plan(request, plan) + deployment = request.profile_deployment().model_dump(mode="json", exclude={"resources"}) + rows = [] + for batch in (1, 2): + for phase, tokens, kv_values in ( + ("prefill", (1, 64, 128, 256), (0, 64, 128, 512)), + ("decode", (0,), (1, 2, 64, 128, 1024)), + ): + for tokens_in_batch in tokens: + for kv in kv_values: + rows.append( + { + **deployment, + "model_path": profile["model"], + "cell_id": f"synthetic-{phase}", + "weight_quantization": "synthetic", + "workload_kind": phase, + "partition_policy": "balanced_v1", + "batch_size": batch, + "total_prefill_tokens": tokens_in_batch, + "total_kv_read_tokens": kv, + "latency_ms": 1 + batch + tokens_in_batch / 128 + kv / 1024, + "kv_seed_regime": "real_kv", + } + ) + table = plan / "systems/data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + table.parent.mkdir(parents=True) + pq.write_table(pa.Table.from_pylist(rows), table) + metadata = { + "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(table.read_bytes()).hexdigest(), + } + table.with_suffix(".metadata.json").write_text(json.dumps(metadata)) + trace = tmp_path / "trace.jsonl" + play = { + "id": "synthetic-play", + "models": ["trace/source-model"], + "block_size": 64, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "s", "model": "trace/source-model", "in": 64, "out": 3, "hash_ids": [1]}, + { + "t": 0.1, + "type": "s", + "model": "trace/source-model", + "in": 128, + "out": 3, + "hash_ids": [1, 2], + }, + ], + } + trace.write_text(json.dumps(play) + "\n") + output = tmp_path / "validation" + args = [ + "onboard", + "validate-fpm", + "--config", + str(plan / "request.yaml"), + "--output-dir", + str(plan), + "--trace", + str(trace), + "--validation-output-dir", + str(output), + ] + # Exercise the public CLI compiler/runner/native runtime in-process so the + # no-analytical-graph assertion above also covers estimator construction. + monkeypatch.setattr(supervision, "main", cli.main) + return args, request, plan, trace, output + + +@pytest.mark.parametrize("validation_case", ["tp", "dep", "tep"], indirect=True) +def test_onboard_validation_replays_profile_without_graph_and_preserves_timing(validation_case, tmp_path): + args, request, plan, _trace, output = validation_case + saved_inputs = {path: path.read_bytes() for path in plan.rglob("*") if path.is_file()} + assert cli.main(args) == 0 + validation = json.loads((output / "validation.json").read_text()) + covered = json.loads((output / "prediction/prediction.json").read_text()) + assert validation["status"] == "covered" + assert validation["accuracy"] == "not_assessed" + assert validation["fpm_query_coverage"]["queries"]["measured"] > 0 + assert validation["fpm_query_coverage"]["queries"]["interpolated"] > 0 + assert validation["fpm_query_coverage"]["queries"]["unsupported"] == 0 + assert covered["completed_requests"] == 2 + assert covered["agentic_model_projection"]["target_model"] == request.identity.model + assert covered["agentic_graph"]["source_models"] == ["trace/source-model"] + prediction = yaml.safe_load((output / "predict.yaml").read_text()) + engine = prediction["engine"] + assert engine["fpm_profile"] == request.fpm_profile.model_dump(mode="json") + assert engine["workers"]["aggregated"]["scheduler"] == request.scheduler_limits() + assert saved_inputs == {path: path.read_bytes() for path in saved_inputs} + + # Coverage is diagnostic only: cold/warm prefix behavior and every native + # request output remain identical when collection is disabled. + ordinary = copy.deepcopy(prediction) + ordinary["engine"]["workers"]["aggregated"]["timing"]["estimator_config"]["fpm_interpolation"].pop( + "collect_coverage" + ) + ordinary_path = tmp_path / "ordinary.yaml" + ordinary_path.write_text(yaml.safe_dump(ordinary)) + ordinary_output = tmp_path / "ordinary" + assert ( + cli.main( + ["predict", "--config", str(ordinary_path), "--output-dir", str(ordinary_output), "--capture-per-request"] + ) + == 0 + ) + uncollected = json.loads((ordinary_output / "prediction.json").read_text()) + assert "fpm_query_coverage" not in uncollected + for key in ("per_request", "agentic_graph", "agentic_play_outcomes", "prefix_cache_reused_ratio"): + assert covered[key] == uncollected[key] + + +def test_onboard_validation_keeps_missing_query_evidence_on_native_failure(validation_case): + import pyarrow as pa + import pyarrow.parquet as pq + + args, _request, plan, trace, output = validation_case + play = json.loads(trace.read_text()) + play["requests"][0]["in"] = 384 + play["requests"][0]["hash_ids"] = list(range(6)) + trace.write_text(json.dumps(play) + "\n") + table = plan / "systems/data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + rows = [ + row + for row in pq.read_table(table).to_pylist() + if row["workload_kind"] != "prefill" or row["total_kv_read_tokens"] <= 128 + ] + pq.write_table(pa.Table.from_pylist(rows), table) + metadata_path = table.with_suffix(".metadata.json") + metadata = json.loads(metadata_path.read_text()) + metadata.update(row_count=len(rows), parquet_sha256=hashlib.sha256(table.read_bytes()).hexdigest()) + metadata_path.write_text(json.dumps(metadata)) + # The third prefill chunk needs a longer KV bracket than this table has. + assert cli.main(args) != 0 + validation = json.loads((output / "validation.json").read_text()) + assert validation["status"] == "incomplete" + assert validation["issues"] + assert validation["prediction_exit_code"] != 0 + coverage = validation["fpm_query_coverage"] + assert coverage["queries"]["unsupported"] > 0 + assert coverage["queries"]["measured"] > 0 + assert coverage["roles"][0]["coverage"]["gaps"][0]["coordinates"]["total_kv_read_tokens"] == 256 + + +def test_onboard_validation_rejects_changed_collection_identity_before_replay(validation_case, monkeypatch): + args, _request, plan, _trace, output = validation_case + config = yaml.safe_load((plan / "request.yaml").read_text()) + config["identity"]["model_revision"] = "different-revision" + supplied = plan.parent / "different.yaml" + supplied.write_text(yaml.safe_dump(config)) + args[args.index("--config") + 1] = str(supplied) + monkeypatch.setattr(supervision, "main", lambda _args: pytest.fail("replay must not start")) + with pytest.raises(SystemExit) as error: + cli.main(args) + assert error.value.code == 2 + assert not output.exists() + + +@pytest.mark.parametrize("flag", ["--trace", "--config"]) +def test_validation_overwrite_preserves_inputs_in_output_directory(validation_case, flag): + args, _request, _plan, _trace, output = validation_case + source = Path(args[args.index(flag) + 1]) + before = source.read_bytes() + output.mkdir() + input_path = output / "prediction.json" + input_path.write_bytes(before) + args[args.index(flag) + 1] = str(input_path) + with pytest.raises(SystemExit) as error: + cli.main([*args, "--overwrite"]) + assert error.value.code == 2 + assert input_path.read_bytes() == before + + +@pytest.mark.parametrize("failure", ["failed_request", "incomplete_play", "truncated_output", "zero_queries"]) +def test_validation_never_accepts_partial_replay(validation_case, monkeypatch, failure): + args, _request, _plan, _trace, output = validation_case + + def partial_prediction(command): + assert command[0] == "predict" + assert json.loads((output / "validation.json").read_text())["status"] == "incomplete" + report = { + "num_requests": 1, + "completed_requests": 1, + "agentic_graph": {"node_count": 1, "play_count": 1}, + "per_request": [{"terminal_status": "completed", "output_length": 2, "requested_output_length": 2}], + "agentic_play_outcomes": [{"status": "completed", "settled_at_ms": 5}], + } + coverage = {"status": "covered", "queries": {"measured": 1, "interpolated": 0, "unsupported": 0}} + if failure == "failed_request": + report["per_request"][0]["terminal_status"] = "failed" + elif failure == "incomplete_play": + report["agentic_play_outcomes"][0]["status"] = "incomplete" + elif failure == "truncated_output": + report["per_request"][0]["output_length"] = 1 + else: + coverage.update(status="incomplete", queries={"measured": 0, "interpolated": 0, "unsupported": 0}) + destination = Path(command[command.index("--output-dir") + 1]) + (destination / "prediction.json").write_text(json.dumps(report)) + (destination / "fpm-coverage.json").write_text(json.dumps(coverage)) + return 0 + + monkeypatch.setattr(supervision, "main", partial_prediction) + assert cli.main(args) == 1 + report = json.loads((output / "validation.json").read_text()) + assert report["status"] == "incomplete" + assert report["issues"] + + +def test_validation_persists_incomplete_status_before_interrupted_prediction(validation_case, monkeypatch): + args, _request, _plan, _trace, output = validation_case + + def interrupt(_command): + assert json.loads((output / "validation.json").read_text())["status"] == "incomplete" + raise KeyboardInterrupt + + monkeypatch.setattr(supervision, "main", interrupt) + assert cli.main(args) == 130 + report = json.loads((output / "validation.json").read_text()) + assert report["status"] == "incomplete" + assert "interrupted" in report["issues"][0] diff --git a/tests/test_fpm_profile_workflow.py b/tests/test_fpm_profile_workflow.py index d8f2d520e..c65e00e3c 100644 --- a/tests/test_fpm_profile_workflow.py +++ b/tests/test_fpm_profile_workflow.py @@ -640,8 +640,10 @@ def no_timing(*args, **kwargs): 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"): + before = {p.relative_to(tmp_path): p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} + with pytest.raises(ValueError, match="fpm-model-profile[.]json was modified or missing"): check_plan(request, tmp_path) + assert {p.relative_to(tmp_path): p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} == before @pytest.mark.parametrize( @@ -710,10 +712,14 @@ def test_generated_glm_collector_plan_preserves_backend_and_rank_local_bounds( assert collected["counts"]["cells"] == 2 assert {cell["parallel_strategy"] for cell in collected["cells"]} == {preset} assert collected["options"]["moe_backend"] == moe_backend + assert collected["options"]["vllm_max_model_len"] == 8192 + assert collected["options"]["max_num_batched_tokens"] == max_tokens + assert collected["options"]["max_num_seqs"] == max_sequences sampling = collected["options"]["prefill_sampling"] - expected_tokens = min(max_tokens, 1024 * max_sequences) - assert sampling["max_total_prefill_tokens"] == expected_tokens + assert sampling["max_total_prefill_tokens"] == max_tokens assert sampling["max_batch_size"] == max_sequences + assert collected["point_generation"]["method"] == "native_self_benchmark" + assert collected["point_generation"]["planned_point_count"] is None 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 @@ -721,11 +727,28 @@ def test_generated_glm_collector_plan_preserves_backend_and_rank_local_bounds( 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"] + + updated_input = request.model_dump(mode="json") + updated_input["workload"].update(input_tokens=64, output_tokens=16, concurrency=2, request_count=4) + updated_request = SupportRequest.model_validate(updated_input) + updated_plan = create_plan(updated_request, tmp_path, overwrite=True) + assert updated_plan["request_id"] == plan["request_id"] + assert updated_plan["validation_id"] != plan["validation_id"] + assert updated_plan["fpm"] == plan["fpm"] + assert updated_plan["resources"] == plan["resources"] + updated_command = json.loads((tmp_path / "commands.json").read_text())["fpm_plan_local"] + assert updated_command == command + assert collector_cli.main(updated_command[3:]) == 0 + assert json.loads(capsys.readouterr().out) == collected + updated_prediction = CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml") + assert updated_prediction.traffic.source.input_tokens == 64 + assert updated_prediction.traffic.source.output_tokens == 16 + assert updated_prediction.traffic.load.concurrency == 2 + assert updated_prediction.traffic.stop.requests == 4 def test_onboard_rejects_profile_below_collector_minimum_token_axis(profile, tmp_path): - profile["deployments"][0]["resources"]["max_num_tokens"] = 1 + profile["deployments"][0]["resources"].update(max_num_tokens=1, max_batch_size=1) request = SupportRequest.model_validate( { "identity": { @@ -740,8 +763,10 @@ def test_onboard_rejects_profile_below_collector_minimum_token_axis(profile, tmp "fpm_profile": profile, } ) + output = tmp_path / "plan" with pytest.raises(ValueError, match="FPM collection requires.*at least 2"): - create_plan(request, tmp_path) + create_plan(request, output) + assert not output.exists() @pytest.mark.parametrize(