diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 80bcd7ef6..165ddb320 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use crate::FpmQueryCoverage; use crate::common::enums::{DatabaseMode, TransferPolicy}; use crate::common::error::AicError; +use crate::fpm::FpmExecutionBinding; use crate::operators::RuntimeContext; use crate::operators::base::PerformanceResult; use crate::operators::{FpmForwardOp, FpmPhase, Op}; @@ -344,6 +345,9 @@ pub struct Engine { /// `(nextn + 1)` exactly as Python `_run_generation_phase:200` /// (`batch_size = batch_size * (model._nextn + 1)`). 0 disables scaling. nextn: u32, + /// Measured target+draft round identity. This is parsed once from the + /// existing EngineConfig.extra map, not a new serialized Op field. + fpm_execution_profile: Option, } impl std::fmt::Debug for Engine { @@ -352,6 +356,7 @@ impl std::fmt::Debug for Engine { .field("context_ops", &self.context_ops.len()) .field("generation_ops", &self.generation_ops.len()) .field("nextn", &self.nextn) + .field("fpm_execution_profile", &self.fpm_execution_profile) .finish_non_exhaustive() } } @@ -363,6 +368,15 @@ impl Engine { /// caller (`AicEngineBuilder` / `from_spec_bytes`) is responsible for /// having loaded the matching `PerfDatabase` from `spec.engine`'s identity. pub fn build(spec: EngineSpec, db: Arc) -> Result { + let fpm_execution_profile = Self::validate_fpm_spec(&spec)?; + Self::build_validated(spec, db, fpm_execution_profile) + } + + fn build_validated( + spec: EngineSpec, + db: Arc, + fpm_execution_profile: Option, + ) -> Result { if spec.engine.quantization.fpm_fmha_dtype.is_some() && spec.engine.forward_model.as_deref() != Some("fpm") { @@ -384,7 +398,6 @@ impl Engine { .as_ref() .and_then(|s| s.nextn) .unwrap_or(0); - Self::validate_fpm_spec(&spec)?; super::readiness::validate_glm53_kv( &db, spec.context_ops.iter().chain(&spec.generation_ops), @@ -400,6 +413,7 @@ impl Engine { generation_ops: spec.generation_ops, db, nextn, + fpm_execution_profile, }) } @@ -430,9 +444,25 @@ impl Engine { Ok(()) } - fn validate_fpm_spec(spec: &EngineSpec) -> Result<(), AicError> { - Self::reject_sd_round_ops(&spec.context_ops)?; - Self::reject_sd_round_ops(&spec.generation_ops)?; + fn validate_fpm_spec(spec: &EngineSpec) -> Result, AicError> { + let execution: Option = spec + .engine + .extra + .get("fpm_execution_profile") + .map(|value| { + serde_json::from_str(value).map_err(|error| { + AicError::InvalidEngineConfig(format!( + "invalid full SD-round FPM execution binding: {error}" + )) + }) + }) + .transpose()?; + if let Some(binding) = &execution { + binding.validate()?; + } else { + Self::reject_sd_round_ops(&spec.context_ops)?; + Self::reject_sd_round_ops(&spec.generation_ops)?; + } let nextn = spec .engine .speculative @@ -461,6 +491,7 @@ impl Engine { } let any_fpm = contains_fpm(&spec.context_ops) || contains_fpm(&spec.generation_ops); if any_fpm + || execution.is_some() || spec.engine.fpm_parquet_path.is_some() || spec.engine.forward_model.as_deref() == Some("fpm") { @@ -494,6 +525,26 @@ impl Engine { (Some(Op::FpmForward(p)), Some(Op::FpmForward(d))) => (p, d), _ => unreachable!("validated leading FPM operations"), }; + // An execution selector is never enough to enable round semantics: + // low-level specs must bind the same typed identity explicitly. + for op in [prefill, decode] { + let selected = op + .match_identity + .get(16) + .map(|selector| FpmExecutionBinding::parse_selector(selector)) + .transpose()? + .flatten(); + if selected != execution { + return Err(AicError::InvalidEngineConfig( + "full SD-round FPM selector and execution binding must match on both phases" + .into(), + )); + } + } + if let Some(binding) = &execution { + Self::validate_full_round_spec(spec, binding, prefill, decode, nextn)?; + return Ok(execution); + } let expected_width = nextn.checked_add(1).ok_or_else(|| { AicError::InvalidEngineConfig("nextn+1 exceeds the u32 verify width".into()) })?; @@ -516,6 +567,80 @@ impl Engine { )); } } + Ok(execution) + } + + fn validate_full_round_spec( + spec: &EngineSpec, + binding: &FpmExecutionBinding, + prefill: &FpmForwardOp, + decode: &FpmForwardOp, + nextn: u32, + ) -> Result<(), AicError> { + let config = &spec.engine; + let parallel = &config.parallel; + let tp = parallel.tp_size; + let moe_tp = parallel.moe_tp_size.unwrap_or(1); + let moe_ep = parallel.moe_ep_size.unwrap_or(1); + if config.backend != crate::BackendKind::Vllm + || config.forward_model.as_deref() != Some("fpm") + || config.database_mode != DatabaseMode::Silicon + || config.decoder_replay + || config.moe_kernel_source.is_some() + || tp == 0 + || parallel.pp_size != 1 + || parallel.attention_dp_size.unwrap_or(1) != 1 + || parallel.cp_size.unwrap_or(1) != 1 + || parallel.dcp_size.unwrap_or(1) != 1 + || moe_ep != 1 + || (moe_tp != 1 && moe_tp != tp) + || nextn != binding.execution_profile.num_draft_tokens + { + return Err(AicError::InvalidEngineConfig( + "full SD-round FPM requires matching nextn, SILICON vLLM TP with DP1/PP1/CP1 and no EP/decoder replay" + .into(), + )); + } + let selector = binding.selector()?; + let topology = [tp, 1, 1, moe_tp, moe_ep, 1].map(|value| value.to_string()); + for (op, count) in [ + (prefill, spec.context_ops.len()), + (decode, spec.generation_ops.len()), + ] { + if count != 1 + || op.verify_width != 1 + || op.interpolation != crate::operators::fpm_forward::FpmInterpolation::Direct + || !op.sol_ops.is_empty() + || op.dcp_size.is_some_and(|value| value != 1) + || op.model_path != config.model_name + || op.match_identity.len() != 19 + || op.match_identity[5..11] != topology + || op.match_identity[14] != "False" + || op.match_identity[13] != "False" + || op.match_identity[16] != selector + || op.match_identity[17] != "none" + || op.match_identity[18] != "text" + { + return Err(AicError::InvalidEngineConfig( + "full SD-round FPM requires exact bound direct phase selectors, verify_width=1, and no draft or SOL operations" + .into(), + )); + } + } + if prefill.match_identity != decode.match_identity { + return Err(AicError::InvalidEngineConfig( + "full SD-round FPM prefill and decode identities must agree".into(), + )); + } + Ok(()) + } + + fn reject_full_round_surface(&self, surface: &str) -> Result<(), AicError> { + if self.fpm_execution_profile.is_some() { + return Err(AicError::InvalidEngineConfig(format!( + "full SD-round FPM does not support {surface}; use pure-prefill or pure-decode ForwardPassMetrics with exact iteration totals" + ))); + } Ok(()) } @@ -552,7 +677,7 @@ impl Engine { systems_root: &std::path::Path, ) -> Result { let spec = EngineSpec::from_bincode(bytes)?; - Self::validate_fpm_spec(&spec)?; + let fpm_execution_profile = Self::validate_fpm_spec(&spec)?; Self::validate_engine_database_mode(spec.engine.database_mode)?; let version = spec.engine.backend_version.as_deref().ok_or_else(|| { AicError::InvalidEngineConfig( @@ -598,7 +723,7 @@ impl Engine { spec.engine.fpm_parquet_path.as_deref(), )? .with_mode(spec.engine.database_mode, transfer_policy); - Engine::build(spec, Arc::new(db)) + Self::build_validated(spec, Arc::new(db), fpm_execution_profile) } fn validate_engine_database_mode(database_mode: DatabaseMode) -> Result<(), AicError> { @@ -678,6 +803,12 @@ impl Engine { mode: StaticMode, stride: u32, ) -> Result { + if runtime.beam_width != 1 { + self.reject_full_round_surface("beam-search static phases")?; + } + if !matches!(mode, StaticMode::Context) { + self.reject_full_round_surface("static generation trajectories")?; + } let context_ms = match mode { StaticMode::Context | StaticMode::Both => self.run_context_phase(runtime)?, StaticMode::Generation => 0.0, @@ -721,7 +852,10 @@ impl Engine { "query coverage requires a direct FPM engine".into(), )); }; - if !ctx_tail.is_empty() || !gen_tail.is_empty() || self.nextn > 0 { + if !ctx_tail.is_empty() + || !gen_tail.is_empty() + || (self.nextn > 0 && self.fpm_execution_profile.is_none()) + { return Err(AicError::InvalidEngineConfig( "query coverage does not support speculative FPM".into(), )); @@ -788,6 +922,7 @@ impl Engine { stride: u32, mut on_op: impl FnMut(usize, &Op, PerformanceResult), ) -> Result { + self.reject_full_round_surface("static generation trajectories")?; let bs = runtime .batch_size .saturating_mul(self.nextn.saturating_add(1)); @@ -1062,6 +1197,7 @@ impl Engine { gen_seq_imbalance_correction_scale: f64, mut on_op: impl FnMut(MixedPass, &Op, PerformanceResult), ) -> Result<[f64; 4], AicError> { + self.reject_full_round_surface("static mixed-step composition")?; if ctx_tokens == 0 && gen_tokens == 0 { return Ok([0.0; 4]); } @@ -1385,6 +1521,7 @@ impl Engine { osl: u32, gen_seq_imbalance_correction_scale: f64, ) -> Result { + self.reject_full_round_surface("mean-context static decode steps")?; if gen_tokens == 0 { return Ok(0.0); } @@ -1561,6 +1698,7 @@ impl Engine { prefix: u32, prefill: bool, ) -> Result, AicError> { + self.reject_full_round_surface("operation diagnostics")?; use super::diagnostics::{ ExecutedFallback, OperationDetails, SolDiagnostics, StaticOperationDiagnostics, }; @@ -1752,6 +1890,7 @@ impl Engine { mode: StaticMode, stride: u32, ) -> Result<(Vec, Vec), AicError> { + self.reject_full_round_surface("per-operation static phases")?; let mut context = PerOpFold::for_plan("context", &self.context_plan); if matches!(mode, StaticMode::Context | StaticMode::Both) { if runtime.prefix >= runtime.isl { @@ -1847,6 +1986,7 @@ impl Engine { seq_imbalance_correction_scale: f64, gen_seq_imbalance_correction_scale: f64, ) -> Result { + self.reject_full_round_surface("per-operation mixed steps")?; // Preserve the FPM component contract while reporting the actual target // and draft operations independently. The same execution path supplies // scalar costs, energy, provenance, and fallback metadata. @@ -1948,6 +2088,7 @@ impl Engine { osl: u32, gen_seq_imbalance_correction_scale: f64, ) -> Result, AicError> { + self.reject_full_round_surface("per-operation decode steps")?; let mut out = PerOpFold::new("generation"); if gen_tokens == 0 { return Ok(out.into_values()); @@ -1981,6 +2122,7 @@ impl Engine { seq_imbalance_correction_scale: f64, x_override: Option, ) -> Result, AicError> { + self.reject_full_round_surface("per-operation context queries")?; let mut out = PerOpFold::new("context"); for &i in indices { let op = self.context_ops.get(i).ok_or_else(|| { @@ -2015,6 +2157,7 @@ impl Engine { prefix: u32, x_override: Option, ) -> Result, AicError> { + self.reject_full_round_surface("per-operation generation queries")?; let mut out = PerOpFold::new("generation"); for &i in indices { let op = self.generation_ops.get(i).ok_or_else(|| { @@ -2053,6 +2196,7 @@ impl Engine { imbalance_correction_scale: f64, x_override: Option, ) -> Result, AicError> { + self.reject_full_round_surface("ad-hoc operation queries")?; let ops: Vec = serde_json::from_str(ops_json).map_err(|e| { AicError::InvalidEngineConfig(format!("evaluate_ops_json: invalid op list JSON: {e}")) })?; @@ -2101,6 +2245,7 @@ impl Engine { imbalance_correction_scale: f64, visual_block_upper_triangle: bool, ) -> Result, AicError> { + self.reject_full_round_surface("ad-hoc attention kernel queries")?; if visual_block_upper_triangle && prefix != 0 { return Err(AicError::InvalidEngineConfig( "visual-block attention kernel evaluation requires prefix=0".into(), @@ -2155,6 +2300,7 @@ impl Engine { imbalance_correction_scale: f64, x_override: Option, ) -> Result, AicError> { + self.reject_full_round_surface("SOL operation queries")?; let ops: Vec = serde_json::from_str(ops_json).map_err(|e| { AicError::InvalidEngineConfig(format!( "evaluate_ops_sol_json: invalid op list JSON: {e}" @@ -2233,6 +2379,11 @@ impl Engine { "at least one attention-DP rank metric required".to_string(), )); } + if self.fpm_execution_profile.is_some() && metrics_by_rank.len() != 1 { + return Err(AicError::InvalidForwardPassMetrics( + "full SD-round FPM requires exactly one attention-DP rank".into(), + )); + } for metrics in metrics_by_rank { validate_forward_pass_metrics(metrics)?; } @@ -2326,7 +2477,15 @@ impl Engine { // `_get_fpm_mix_step_latency` at the telemetry counts (already packed, // so no `(nextn + 1)` anywhere; speculative FPM is rejected below). if let Some((prefill_op, decode_op, ctx_tail, gen_tail)) = self.fpm_split() { - if !ctx_tail.is_empty() || !gen_tail.is_empty() || self.nextn > 0 { + if self.fpm_execution_profile.is_some() && has_prefill && has_decode { + return Err(AicError::InvalidForwardPassMetrics( + "full SD-round FPM has no qualified mixed-iteration composition".into(), + )); + } + if !ctx_tail.is_empty() + || !gen_tail.is_empty() + || (self.nextn > 0 && self.fpm_execution_profile.is_none()) + { // The telemetry counts are packed AR semantics; the hybrid // speculative shape (draft tails / widened verify) has no // defined mapping here yet. @@ -3893,6 +4052,262 @@ mod tests { assert!(error.to_string().contains("SD-round"), "{error}"); } + fn full_round_binding() -> FpmExecutionBinding { + serde_json::from_value(serde_json::json!({ + "model_revision": "1".repeat(40), + "execution_profile": { + "kind": "speculative_round", + "method": "eagle3", + "num_draft_tokens": 3, + "draft_model": "org/draft-model", + "draft_revision": "2".repeat(40), + "timing_scope": "model_step_cuda" + } + })) + .unwrap() + } + + fn full_round_spec(tmp: &std::path::Path) -> EngineSpec { + use crate::operators::fpm_forward::FpmInterpolation; + use crate::perf_database::fpm_forward::tests::{NativeCudaSampleSpec, RowSpec, write_pair}; + let binding = full_round_binding(); + let selector = binding.selector().unwrap(); + // The existing RowSpec fixture stores static strings. The few small + // test selectors live for this test process, not in production state. + let selector: &'static str = Box::leak(selector.into_boxed_str()); + let rows: Vec<_> = [ + ("prefill", 1, 256, 0, 11.0), + ("prefill", 1, 512, 0, 17.0), + ("prefill", 1, 256, 128, 13.0), + ("prefill", 1, 512, 128, 19.0), + ("decode", 1, 0, 4096, 7.0), + ("decode", 1, 0, 8192, 9.0), + ("decode", 4, 0, 4096, 31.0), + ("decode", 4, 0, 8192, 33.0), + ] + .into_iter() + .map(|(kind, batch, prefill, kv, latency)| RowSpec { + workload_kind: kind, + batch_size: batch, + total_prefill_tokens: prefill, + total_kv_read_tokens: kv, + latency_ms: latency, + execution: Some(["", selector, "none", "text"]), + kv_seed_regime: (kv > 0).then_some("real_kv"), + native_cuda: Some(NativeCudaSampleSpec { + scope: "model_step_cuda", + reduction: "single_sample", + source_sample_count: 1, + }), + ..RowSpec::default() + }) + .collect(); + let mut spec = build_fpm_spec(tmp, Some(3)); + write_pair(tmp, &rows); + spec.engine.model_name = "org/model-a".into(); + spec.engine.parallel.tp_size = 4; + spec.engine.parallel.moe_tp_size = Some(4); + spec.engine.parallel.moe_ep_size = Some(1); + spec.engine.extra.insert( + "fpm_execution_profile".into(), + serde_json::to_string(&binding).unwrap(), + ); + for op in spec.context_ops.iter_mut().chain(&mut spec.generation_ops) { + let Op::FpmForward(op) = op else { + unreachable!() + }; + op.interpolation = FpmInterpolation::Direct; + op.match_identity[16] = selector.into(); + } + spec + } + + fn full_round_engine(tmp: &std::path::Path) -> Engine { + let spec = full_round_spec(tmp); + Engine::from_spec_bytes(&spec.to_bincode().unwrap(), tmp).unwrap() + } + + #[test] + fn full_round_raw_counts_price_one_measured_round_without_draft_recharge() { + let tmp = tempfile::tempdir().unwrap(); + let engine = full_round_engine(tmp.path()); + assert_eq!( + engine.nextn, 3, + "scheduler identity must retain its real depth" + ); + // Hand-authored measured rows: B1/KV4096 is a COMPLETE 7ms K3 round; + // the B4 row is deliberately 31ms. Widening or charging draft again + // cannot satisfy these independent table oracles. + assert_eq!(engine.predict_decode_latency_total(1, 4096).unwrap(), 7.0); + assert_eq!(engine.predict_decode_latency_total(4, 4096).unwrap(), 31.0); + let mut coverage = FpmQueryCoverage::default(); + assert_eq!( + engine + .predict_decode_latency_total_with_coverage(1, 4096, Some(&mut coverage)) + .unwrap(), + 7.0 + ); + // Measured SD-prefill rows include their own draft-prefill work. + assert_eq!(engine.predict_prefill_latency(1, 256, 0).unwrap(), 11.0); + assert_eq!( + engine + .predict_prefill_detailed(1, 384, 128, Some(&mut coverage)) + .unwrap() + .latency_ms, + Some(13.0) + ); + let prefill = ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_prefill_requests: 1, + sum_prefill_tokens: 256, + sum_prefill_kv_tokens: 128, + ..Default::default() + }, + ..Default::default() + }; + assert_eq!(engine.forward_pass_time_ms(&[prefill]).unwrap(), 13.0); + let context = engine + .run_static(&runtime(1, 256, 1), StaticMode::Context, 32) + .unwrap(); + assert_eq!(context.total_ms, 11.0); + assert_eq!( + engine + .forward_pass_time_ms(&[ForwardPassMetrics::default()]) + .unwrap(), + 0.0 + ); + } + + #[test] + fn full_round_unsupported_surfaces_fail_before_pricing_work() { + let tmp = tempfile::tempdir().unwrap(); + let engine = full_round_engine(tmp.path()); + for osl in [2, 32] { + let err = engine.predict_decode_latency(1, 4096, osl).unwrap_err(); + assert!(err.to_string().contains("static generation"), "{err}"); + } + assert!( + engine + .run_static(&runtime(1, 256, 2), StaticMode::Both, 32) + .is_err() + ); + assert!( + engine + .run_static_per_op(&runtime(1, 256, 1), StaticMode::Context, 32) + .is_err() + ); + assert!(engine.decode_step_latency(1, 4096, 2, 1.0).is_err()); + assert!(engine.decode_step_per_op(1, 4096, 2, 1.0).is_err()); + assert!( + engine + .mixed_step_latency(256, 1, 256, 2, 0, 1.0, 1.0) + .is_err() + ); + assert!( + engine + .mixed_step_breakdown_per_op(256, 1, 256, 2, 0, 1.0, 1.0) + .is_err() + ); + assert!( + engine + .evaluate_generation_ops(&[0], 4, 4096, 1.0, 0, None) + .is_err() + ); + assert!( + engine + .evaluate_ops_json("[]", true, 1, 256, 0, 1.0, None) + .is_err() + ); + let mixed = ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_prefill_requests: 1, + sum_prefill_tokens: 256, + num_decode_requests: 1, + sum_decode_kv_tokens: 4096, + ..Default::default() + }, + ..Default::default() + }; + let error = engine.forward_pass_time_ms(&[mixed]).unwrap_err(); + assert!(error.to_string().contains("mixed-iteration"), "{error}"); + assert!( + engine + .forward_pass_time_ms(&[ + ForwardPassMetrics::default(), + ForwardPassMetrics::default() + ]) + .is_err() + ); + } + + #[test] + fn full_round_forged_specs_are_rejected_before_database_loading() { + use crate::operators::fpm_forward::FpmInterpolation; + for case in 0..15 { + let tmp = tempfile::tempdir().unwrap(); + let mut spec = full_round_spec(tmp.path()); + match case { + 0 => { + spec.engine.extra.remove("fpm_execution_profile"); + } + 1 => { + spec.engine.speculative.as_mut().unwrap().nextn = Some(2); + } + 2 => { + spec.engine.parallel.attention_dp_size = Some(2); + } + 3 => { + spec.engine.parallel.moe_ep_size = Some(2); + } + 4 => { + spec.engine.database_mode = DatabaseMode::Hybrid; + } + 5 => { + spec.generation_ops.push(context_ops().remove(0)); + } + 6 => { + spec.engine.backend = BackendKind::Sglang; + } + 7 => { + spec.engine.parallel.dcp_size = Some(2); + } + 8 => { + spec.engine.parallel.cp_size = Some(2); + } + _ => { + let Op::FpmForward(op) = &mut spec.generation_ops[0] else { + unreachable!() + }; + match case { + 9 => op.verify_width = 4, + 10 => op.interpolation = FpmInterpolation::Sol, + 11 => op.match_identity.truncate(15), + 12 => op.sol_ops.push(context_ops().remove(0)), + 13 => op.match_identity[16] = "full".into(), + 14 => { + let mut other = full_round_binding(); + other.model_revision = "3".repeat(40); + op.match_identity[16] = other.selector().unwrap(); + } + _ => unreachable!(), + } + } + } + // A nonexistent path proves validation runs before database IO. + spec.engine.fpm_parquet_path = Some("/missing/full-round-test.parquet".into()); + let error = + Engine::from_spec_bytes(&spec.to_bincode().unwrap(), tmp.path()).unwrap_err(); + assert!( + matches!(error, AicError::InvalidEngineConfig(_)), + "case {case}: {error}" + ); + assert!( + error.to_string().contains("full SD-round"), + "case {case}: {error}" + ); + } + } + #[test] fn identity_only_sd_selector_in_nested_stage_is_rejected_before_loading() { let tmp = tempfile::tempdir().unwrap(); diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index cd2f4540d..3e8f7593d 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -297,7 +297,9 @@ impl ForwardPassPerfModelConfig { "fpm_profile requires a literal backend_version", )); } - if self.nextn != 0 || self.speculation.is_some() { + if (self.nextn != 0 && self.validate_fpm_execution()?.is_none()) + || self.speculation.is_some() + { return Err(invalid_config( "FPM profiles support plain autoregressive decoder-only execution", )); @@ -466,7 +468,11 @@ impl ForwardPassPerfModelConfig { )); } } - if self.estimation_mode == EstimationMode::FpmInterpolation && self.nextn != 0 { + let execution = self.validate_fpm_execution()?; + if self.estimation_mode == EstimationMode::FpmInterpolation + && self.nextn != 0 + && execution.is_none() + { return Err(invalid_config( "estimation_mode='fpm_interpolation' does not support MTP speculative decoding", )); @@ -497,6 +503,42 @@ impl ForwardPassPerfModelConfig { } self.estimator_config.validate() } + + pub(crate) fn validate_fpm_execution( + &self, + ) -> Result, AicError> { + let Some(binding) = self.fpm_execution_binding()? else { + return Ok(None); + }; + if self.estimation_mode != EstimationMode::FpmInterpolation + || self.estimator_config.fpm_interpolation.method + != super::FpmInterpolationMethod::Direct + || self.fallback_policy != ForwardPassFallbackPolicy::Deny + || self.database_mode != DatabaseMode::Silicon + { + return Err(invalid_config( + "speculative-round FPM requires explicit fpm_interpolation, direct method, SILICON data and fallback_policy='deny'", + )); + } + if self.nextn != binding.execution_profile.num_draft_tokens || self.speculation.is_some() { + return Err(invalid_config( + "speculative-round FPM nextn must match num_draft_tokens and cannot use generic speculation", + )); + } + if self.backend != BackendKind::Vllm + || self.pp != 1 + || self.attention_dp != 1 + || self.cp_size.is_some_and(|cp| cp != 1) + || self.dcp.is_some() + || self.moe_ep_size.is_some_and(|ep| ep != 1) + || self.moe_tp_size.is_some_and(|tp| tp != 1 && tp != self.tp) + { + return Err(invalid_config( + "speculative-round FPM supports vLLM pure TP with PP1/DP1/CP1 and no recorded DCP", + )); + } + Ok(Some(binding)) + } } pub(super) fn invalid_config(message: impl Into) -> AicError { diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index 83fda719f..54fbc61f1 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -1111,7 +1111,10 @@ fn build_native_candidate( "FPM interpolation does not support EPLB, slots, moe_backend or moe_kernel_source overrides".into(), )); } - if config.estimation_mode == EstimationMode::FpmInterpolation && config.nextn != 0 { + if config.estimation_mode == EstimationMode::FpmInterpolation + && config.nextn != 0 + && config.validate_fpm_execution()?.is_none() + { return Err(AicError::UnsupportedModel( "FPM interpolation does not support MTP".into(), )); diff --git a/crates/core/src/perfmodel/fpm/resources.rs b/crates/core/src/perfmodel/fpm/resources.rs index e9ef31e12..dcda1dc32 100644 --- a/crates/core/src/perfmodel/fpm/resources.rs +++ b/crates/core/src/perfmodel/fpm/resources.rs @@ -5,7 +5,9 @@ use serde::{Deserialize, Serialize}; -use super::{ForwardPassPerfModel, ForwardPassPerfModelConfig, ForwardPassWorkerType}; +use super::{ + ForwardPassPerfModel, ForwardPassPerfModelConfig, ForwardPassWorkerType, FpmExecutionBinding, +}; use crate::{AicError, BackendKind, MemoryBreakdown}; const MAX_EXACT_BYTES: u64 = 1 << 53; @@ -338,15 +340,86 @@ impl ForwardPassPerfModelConfig { if profile.get("model").and_then(serde_json::Value::as_str) != Some(self.model.as_str()) { return Err(invalid("FPM profile model identity mismatch")); } + let execution = self.validate_fpm_execution()?; if self.backend != BackendKind::Vllm || self.pp != 1 - || self.nextn != 0 + || (self.nextn != 0 && execution.is_none()) || self.speculation.is_some() { return Err(invalid( "FPM resources require vLLM PP1 plain autoregressive execution", )); } + let deployment = self + .fpm_deployment()? + .ok_or_else(|| invalid("FPM deployment requires a profile"))?; + let resources: FpmResourceConfig = serde_json::from_value( + deployment + .get("resources") + .cloned() + .ok_or_else(|| invalid("FPM deployment requires resources"))?, + ) + .map_err(|error| invalid(format!("FPM resource profile: {error}")))?; + resources.validate()?; + if execution.is_some() + && (resources.cache_layout != FpmCacheLayout::Linear + || resources.runtime_memory.is_none()) + { + return Err(invalid( + "speculative-round FPM resources require observed runtime memory and a linear cache", + )); + } + Ok(Some(resources)) + } + + pub(crate) fn fpm_execution_binding(&self) -> Result, AicError> { + let Some(profile) = &self.fpm_profile else { + return Ok(None); + }; + // Do not change validation order for historical profiles without the + // optional execution identity. + let has_execution = profile + .get("deployments") + .and_then(serde_json::Value::as_array) + .is_some_and(|deployments| { + deployments.iter().any(|deployment| { + deployment + .get("execution_profile") + .is_some_and(|value| !value.is_null()) + }) + }); + if !has_execution { + return Ok(None); + } + let deployment = self + .fpm_deployment()? + .ok_or_else(|| invalid("FPM execution identity requires a deployment"))?; + let Some(execution) = deployment + .get("execution_profile") + .filter(|value| !value.is_null()) + else { + return Ok(None); + }; + let binding = FpmExecutionBinding { + model_revision: profile + .get("model_revision") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| invalid("FPM execution identity requires model_revision"))? + .to_owned(), + execution_profile: serde_json::from_value(execution.clone()) + .map_err(|error| invalid(format!("FPM execution identity: {error}")))?, + }; + binding.validate()?; + Ok(Some(binding)) + } + + fn fpm_deployment(&self) -> Result, AicError> { + let Some(profile) = &self.fpm_profile else { + return Ok(None); + }; + if profile.get("model").and_then(serde_json::Value::as_str) != Some(self.model.as_str()) { + return Err(invalid("FPM profile model identity mismatch")); + } let deployments = profile .get("deployments") .and_then(serde_json::Value::as_array) @@ -406,23 +479,7 @@ impl ForwardPassPerfModelConfig { if matching.next().is_some() { return Err(invalid("duplicate FPM deployment resources")); } - if deployment - .get("execution_profile") - .is_some_and(|execution| !execution.is_null()) - { - return Err(invalid( - "speculative-round FPM profiles are identity-only; execution is not enabled", - )); - } - let resources: FpmResourceConfig = serde_json::from_value( - deployment - .get("resources") - .cloned() - .ok_or_else(|| invalid("FPM deployment requires resources"))?, - ) - .map_err(|error| invalid(format!("FPM resource profile: {error}")))?; - resources.validate()?; - Ok(Some(resources)) + Ok(Some(deployment)) } } @@ -673,6 +730,87 @@ mod tests { config } + fn round_config() -> ForwardPassPerfModelConfig { + let mut config = role_config(); + config.nextn = 3; + config.estimation_mode = super::super::EstimationMode::FpmInterpolation; + config.estimator_config.fpm_interpolation.method = + super::super::FpmInterpolationMethod::Direct; + let profile = config.fpm_profile.as_mut().unwrap(); + profile.insert("model_revision".into(), "1".repeat(40).into()); + for deployment in profile["deployments"].as_array_mut().unwrap() { + deployment["execution_profile"] = serde_json::json!({ + "kind": "speculative_round", "method": "eagle3", "num_draft_tokens": 3, + "draft_model": "test/draft", "draft_revision": "2".repeat(40), + "timing_scope": "model_step_cuda" + }); + let resources = deployment["resources"].as_object_mut().unwrap(); + resources.insert("cache_layout".into(), "linear".into()); + resources.insert("kv_bytes_per_token".into(), 10.into()); + resources.remove("cache_groups"); + } + config + } + + #[test] + fn speculative_round_keeps_observed_capacity_and_identity_after_reload() { + let config = round_config(); + config.validate().unwrap(); + let reloaded: ForwardPassPerfModelConfig = + serde_json::from_str(&serde_json::to_string(&config).unwrap()).unwrap(); + assert_eq!(reloaded.nextn, 3); + assert_eq!( + config.fpm_execution_binding().unwrap(), + reloaded.fpm_execution_binding().unwrap() + ); + let resources = reloaded.fpm_resources().unwrap().unwrap(); + let mut query = request(); + query.cuda_graph_reserved_bytes = 0; + query.tolerance_fraction = None; + let budget = resources.estimate_budget(&query, 4096).unwrap(); + // The recorded usable pool is 600 bytes at 10 bytes/token. SD has + // already affected that observation; neither draft depth nor TP is a + // second capacity multiplier or memory deduction. + assert_eq!(budget.total_kv_size_bytes, 600); + assert_eq!(budget.total_kv_size_tokens, Some(60)); + assert_eq!(budget.memory_breakdown, None); + query.cuda_graph_reserved_bytes = 20; + assert!(resources.estimate_budget(&query, 4096).is_err()); + } + + #[test] + fn speculative_round_requires_matching_depth_direct_selection_and_observed_memory() { + let config = round_config(); + for depth in [0, 2, 4] { + let mut invalid = config.clone(); + invalid.nextn = depth; + assert!(invalid.validate().is_err()); + assert!(invalid.fpm_resources().is_err()); + } + let mut invalid = config.clone(); + invalid.estimation_mode = super::super::EstimationMode::Auto; + assert!(invalid.validate().is_err()); + invalid = config.clone(); + invalid.estimator_config.fpm_interpolation.method = + super::super::FpmInterpolationMethod::Sol; + assert!(invalid.validate().is_err()); + invalid = config.clone(); + invalid.fallback_policy = super::super::ForwardPassFallbackPolicy::Allow; + assert!(invalid.validate().is_err()); + for replacement in [resources(), runtime_resources()] { + let mut invalid = config.clone(); + invalid.fpm_profile.as_mut().unwrap()["deployments"][0]["resources"] = + serde_json::to_value(replacement).unwrap(); + assert!(invalid.fpm_resources().is_err()); + } + let mut invalid = config; + let resources = invalid.fpm_profile.as_mut().unwrap()["deployments"][0]["resources"] + .as_object_mut() + .unwrap(); + resources.remove("runtime_memory"); + assert!(invalid.fpm_resources().is_err()); + } + #[test] fn profile_resources_select_independent_roles_after_round_trip() { let config = role_config(); diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 274b475b8..52510bdd9 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -1038,8 +1038,12 @@ fn aic_capacity_kwargs<'py>( "cuda_graph_reserved_bytes", config.cuda_graph_reserved_bytes, )?; - // Capacity intentionally omits NextN until AIC's Eagle memory model no - // longer returns negative KV capacity. Timing compilation still uses it. + // Legacy analytical capacity still omits NextN. Profile resources have + // already been validated by materialize_aic_capacity; nonzero depth here + // identifies measured SD memory, whose recorded pool must not be re-priced. + if config.fpm_profile.is_some() && config.nextn > 0 { + kwargs.set_item("nextn", config.nextn)?; + } kwargs.set_item("systems_path", config.systems_path.as_deref())?; Ok(kwargs) } @@ -4099,6 +4103,40 @@ mod tests { } } + #[test] + fn measured_profile_capacity_forwards_depth_without_changing_analytical_defaults() { + pyo3::prepare_freethreaded_python(); + let role = aggregated_role(&ReplayEngineConfig::default()); + let mut config = aic_config(); + Python::with_gil(|py| { + for depth in [0, 3] { + config.nextn = depth; + let kwargs = aic_capacity_kwargs(py, &config, &role).unwrap(); + assert!(!kwargs.contains("nextn").unwrap()); + } + // Selection/resource validity is enforced by materialize_aic_capacity. + // This unit test isolates only the previously missing kwargs hop. + config.fpm_profile = Some(serde_json::Map::new()); + let kwargs = aic_capacity_kwargs(py, &config, &role).unwrap(); + assert_eq!( + kwargs + .get_item("nextn") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + 3 + ); + config.nextn = 0; + assert!( + !aic_capacity_kwargs(py, &config, &role) + .unwrap() + .contains("nextn") + .unwrap() + ); + }); + } + #[test] fn grouped_load_and_timing_domain_preserve_byte_capacity() { let mut role = aggregated_role(&ReplayEngineConfig::default()); diff --git a/docs/perf-model/methods/whole-forward.md b/docs/perf-model/methods/whole-forward.md index 8c7dcf010..9028d1599 100644 --- a/docs/perf-model/methods/whole-forward.md +++ b/docs/perf-model/methods/whole-forward.md @@ -34,9 +34,9 @@ execution_profile: The initial identity contract permits vLLM pure TP with DP1/PP1/CP1 and 1–5 draft tokens. Acceptance length belongs to replay progress, not this timing selector. Omitting `execution_profile` preserves existing AR profile JSON and -cell keys. Identity publication alone does not enable execution: canonical -construction rejects these profiles with either `nextn: 0` or nonzero `nextn` -until the corresponding consumer is supported. +cell keys. Consumption requires explicit `estimation_mode: fpm_interpolation`, +`method: direct`, `database_mode: SILICON`, `fallback_policy: deny`, and `nextn` +equal to `num_draft_tokens`. Ordinary AR profiles cannot supply SD-round timing. Rust produces the exact table `execution_profile` string by prefixing `speculative_round_v1:` to its canonical serialization of @@ -71,6 +71,39 @@ not a valid conversion. Cached prefill and decode must retain the existing `real_kv` provenance requirement. Missing or fabricated state cannot be healed into these measured-only cells. +### Consuming a measured round + +Use the same canonical `ForwardPassPerfModel.best_available(config)` path as +ordinary FPM profiles. Both prefill and decode cells must carry the identical +execution binding. The decode lookup uses request count and actual total past +KV tokens, without multiplying the request count by `nextn + 1`. The measured +latency already includes target verification, sampling and draft work; no draft +operation tail is added. Prefill must likewise use the matching measured SD +prefill cell, not a relabelled AR prefill cell. + +Resources must use a linear cache with observed `runtime_memory`. The recorded +usable KV pool already reflects draft and graph allocations; capacity sizing +does not subtract them again. Pending or analytically declared memory alone is +insufficient. Grouped caches, generic `speculation` configurations, recorded +DCP and non-pure-TP topologies remain unsupported on this route. Use `nextn` +with the exact recorded `execution_profile`; generic speculation cost graphs +must not be composed with measurements that already include the draft work. + +Supported queries are pure-prefill and pure-decode `ForwardPassMetrics`, prefill +static/detailed queries, and replay's decode query with actual total KV tokens. +Raw mixed FPM, static multi-round generation, per-operation and ad-hoc operation +queries fail explicitly. Ordinary replay retains its existing separate P/D +phase composition; this is an approximation for mixed iterations, not a measured +mixed-round model. Partial-width decode rounds still use the configured depth's +timing. The existing replay prefill callback uses integer per-request mean +lengths, so ragged batches can lose remainder tokens; raw FPM totals do not have +that approximation. Enable query coverage to identify holes rather than +extrapolating. + +The SD identity does not expand the guided `onboard` collector or its replay +validation scope. Prepare and validate the measured table separately, then use +ordinary `predict` with the [replay SD controls](../../replay/engine/speculation.md#measured-sd-round-fpm). + ## Direct and SOL-assisted methods - `method: direct` needs a verified resource profile and `database_mode: SILICON`. diff --git a/docs/replay/engine/speculation.md b/docs/replay/engine/speculation.md index 87a794cc4..5ecf83728 100644 --- a/docs/replay/engine/speculation.md +++ b/docs/replay/engine/speculation.md @@ -33,7 +33,7 @@ verification round, including draft operations for schemes that model them. | `engine.speculation.num_speculative_tokens` | Required for flat ngram input | fixed | Integer from 1 to 5. | | `engine.speculation.acceptance_rates` | Required for flat ngram input | fixed | Exactly one conditional acceptance probability per draft token, each in `[0, 1]`. | | `engine.speculation.seed` | `42` | fixed | Unsigned 64-bit sampling seed. | -| `engine.nextn` | `0` (off) | fixed | MTP/EAGLE draft tokens, 0 to 5. | +| `engine.nextn` | `0` (off) | fixed | MTP/EAGLE draft tokens, 0 to 5. For a measured SD-round FPM profile, must equal its recorded draft depth. | | `engine.nextn_accepted` | Unset | fixed | Required when `nextn > 0`. Expected accepted **draft** tokens per round, excluding the base token, in `[0, nextn]`. | `recommend` retains the flat ngram and legacy `nextn` inputs, fixed for every @@ -197,7 +197,7 @@ draft readiness. TensorRT-LLM retains its existing scheduling approximation. These are scheduling/progress semantics, not a new timing estimator: the current timing callback still prices the configured verification round for a partial-width decode. It does not interpolate dynamically by actual draft -width or enable measured full-SD-round FPM ingestion. These limits matter +width. Measured full-SD-round profiles use the separate timing contract below. These limits matter when benchmarking tiny token budgets. HBM-only agentic MTP/EAGLE replay is supported by the offline engine stack, @@ -207,3 +207,28 @@ Rejected with MTP: grouped caches, G2/G3 offload, and state-cache prefix matchin with `prefix_match_unit`. Combinations accepted by the timing model are not necessarily supported by Replay; see [feature support](../features.md). + +### Measured SD-round FPM + +Ordinary `predict` can consume an explicitly tagged EAGLE3 profile through +`engine.fpm_profile`, including HBM-only AgentX traffic. Each deployment records +the [SD execution identity](../../perf-model/methods/whole-forward.md#recorded-speculative-round-identity) +and observed linear-cache `runtime_memory`. Select `fpm_interpolation`, direct +interpolation and deny fallback in the worker's timing configuration; `nextn` +must match the recorded depth. For three drafts and mean total acceptance length +2.78, use `nextn: 3` and `nextn_accepted: 1.78`. Leave `engine.speculation` +unset: the registered-scheme configuration above models analytical draft work +and cannot be combined with a measured full-round profile. + +Acceptance determines sampled output progress, not a latency multiplier. Each +decode query consumes one measured complete round at the request count and +actual past-KV total, without another draft charge. Prefill uses its own matching +SD cell. Mixed scheduled iterations retain replay's separate P/D phase-cost +composition; the data does not establish mixed-round fidelity. Small token +budgets can produce partial-width rounds that are still priced at the configured +depth. Query coverage and a measured serving comparison remain separate checks. + +The recorded usable pool is already net of draft/graph allocations. Do not add +another reservation or scalar capacity override. This path requires vLLM pure +TP, PP1/DP1/CP1, no recorded DCP, linear cache and HBM-only replay. It does not +enable grouped-cache or offload SD, nor expand guided `onboard` collection. diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 02e05efa7..b6f9d72ec 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -201,6 +201,32 @@ def finish_lowering(value): # canonical lowering has already pinned the selected root above. capacity_systems_path = resolved_roots[0] + speculative_profile_kwargs = {} + if lowered.get("aic_fpm_profile") is not None: + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + deployment = load_fpm_profile(lowered["aic_fpm_profile"]).select( + model=model, + system=lowered.get("aic_system") or _DEFAULT_AIC_SYSTEM, + backend=backend, + backend_version=lowered.get("aic_backend_version", lowered.get("backend_version")), + worker_type=lowered.get("worker_type", "aggregated"), + tp_size=lowered.get("aic_tp_size", 1), + pp_size=lowered.get("aic_pp_size", 1), + attention_dp_size=dp, + moe_tp_size=lowered.get("aic_moe_tp_size"), + moe_ep_size=lowered.get("aic_moe_ep_size"), + cp_size=lowered.get("aic_cp_size", 1), + ) + if deployment.execution_profile is not None: + depth = lowered.get("aic_nextn", 0) + if canonical_result is not None: + depth = lowered["timing_model"]["config"]["nextn"] + for name in ("aic_nextn", "nextn"): + if name in lowered and lowered[name] != depth: + raise ValueError(f"{name} conflicts with canonical timing configuration") + speculative_profile_kwargs["nextn"] = depth + lowered["num_gpu_blocks"] = estimate_num_gpu_blocks( backend_name=backend, system=lowered.get("aic_system") or _DEFAULT_AIC_SYSTEM, @@ -248,6 +274,7 @@ def finish_lowering(value): ), **({"context_length": lowered.get("max_model_len")} if lowered.get("aic_fpm_profile") is not None else {}), **({"diagnostics": memory_diagnostics} if memory_diagnostics is not None else {}), + **speculative_profile_kwargs, ) return finish_lowering(lowered) @@ -356,12 +383,13 @@ def estimate_num_gpu_blocks( fpm_profile: dict[str, Any] | None = None, worker_type: str = "aggregated", context_length: int | None = None, + nextn: int | None = None, ) -> int: """Estimate per-rank KV blocks using the replay-wide AIC contract. - NextN is intentionally absent. AIC currently can return negative KV - capacity for Eagle when speculative-decoding state is included. Timing - compilation still receives NextN; only capacity estimation omits it. + Legacy analytical estimates omit NextN. Measured speculative profiles + instead carry their depth to validate the observed pool, without estimating + additional draft-model memory. """ if backend_name not in DEFAULT_BACKEND_VERSIONS: @@ -384,6 +412,8 @@ def estimate_num_gpu_blocks( raise ValueError(f"no perf database for system={system!r}, backend={backend_name!r}") validate_moe_controls(enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots) + if nextn is not None and fpm_profile is None: + raise ValueError("nextn capacity forwarding requires an FPM profile") if backend_name == "trtllm": memory_fraction_kind = "of_free" @@ -438,6 +468,7 @@ def estimate_num_gpu_blocks( **({"fpm_profile": fpm_profile, "worker_type": worker_type} if fpm_profile is not None else {}), **({"context_length": context_length} if context_length is not None else {}), **({"diagnostics": diagnostics} if diagnostics is not None else {}), + **({"nextn": nextn} if nextn is not None else {}), ) ) diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index a7dac0eeb..14f3e22ce 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -491,7 +491,8 @@ def compile_engine( Rust constructor supplies an already resolved ``sol`` or ``direct`` method. Profile precision is authoritative on either route, with conflicting caller overrides rejected. ``cp_size`` completes the profile - identity; profiles currently support only CP1/PP1 and no speculation. + identity; profiles support CP1/PP1, with explicitly identified speculative + rounds confined to direct interpolation. """ if not isinstance(decoder_replay, bool): raise InvalidEngineConfigurationError("decoder_replay must be a boolean") @@ -544,7 +545,7 @@ def compile_engine( raise InvalidEngineConfigurationError("FPM profiles do not support moe_kernel_source") if dcp_size is not None: raise InvalidEngineConfigurationError("fpm_profile does not describe recorded DCP identity") - if nextn or speculation is not None: + if speculation is not None: raise InvalidEngineConfigurationError("FPM profiles support plain autoregressive decoder-only execution") literal_version = _literal_backend_version(system, backend, backend_version, systems_path, None) deployment = profile.select( @@ -560,6 +561,13 @@ def compile_engine( moe_ep_size=moe_ep_size, cp_size=cp_size, ) + if deployment.execution_profile is None: + if nextn: + raise InvalidEngineConfigurationError( + "FPM profiles support plain autoregressive decoder-only execution" + ) + elif forward_model != "fpm" or interpolation != "direct": + raise InvalidEngineConfigurationError("speculative-round FPM profiles require direct FPM compilation") deployment.validate_overrides( gemm_quant_mode=gemm_quant_mode, moe_quant_mode=moe_quant_mode, @@ -595,6 +603,7 @@ def compile_engine( fpm_fmha_quant_mode=fpm_fmha_quant_mode, worker_type=worker_type, fpm_options=fpm_config.options, + nextn=nextn, ) return bytes(aisimulate_core.engine_spec_bincode_from_json(spec_json)) gemm_quant_mode = deployment.gemm_quant_mode @@ -733,11 +742,15 @@ def _direct_fpm_spec_json( fpm_fmha_quant_mode: str | None = None, worker_type: str = "aggregated", fpm_options: dict | None = None, + nextn: int = 0, ) -> str: """Build whole-forward timing operations from metadata, without a graph.""" if _database_mode_name(None, database_mode) != "SILICON": raise ValueError("direct FPM interpolation requires database_mode='SILICON'; it has no analytical SOL model") identity = deployment.match_identity() + execution = deployment.execution_profile + if execution is not None: + identity += list(deployment.execution_identity(profile.model_revision)) if fpm_fmha_quant_mode is not None: if fpm_fmha_quant_mode not in {"bfloat16", "fp8", "fp8_block"}: raise ValueError("fpm_fmha_quant_mode must be bfloat16, fp8, or fp8_block") @@ -776,6 +789,12 @@ def _direct_fpm_spec_json( "estimator_config": json.dumps({"fpm_interpolation": controls}), }, } + if execution is not None: + engine["nextn"] = nextn + engine["extra"]["fpm_execution_profile"] = json.dumps( + {"model_revision": profile.model_revision, "execution_profile": execution.model_dump(mode="json")}, + sort_keys=True, + ) def phase_op(phase: str) -> dict: return { diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py index 1ff25cd31..c1c2051f9 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py @@ -15,6 +15,7 @@ from aisimulate_core.fpm_profile import FpmModelProfile as FpmModelProfile from aisimulate_core.fpm_profile import FpmResourceProfile as FpmResourceProfile from aisimulate_core.fpm_profile import FpmRuntimeMemoryProfile as FpmRuntimeMemoryProfile +from aisimulate_core.fpm_profile import FpmSpeculativeExecutionProfile as FpmSpeculativeExecutionProfile from aisimulate_core.fpm_profile import load_fpm_profile as load_fpm_profile diff --git a/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py b/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py index 4bd811200..1dec84739 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py @@ -1164,8 +1164,6 @@ def estimate_kv_cache( if request_occupancy_tokens is not None and fpm_profile is None: raise ValueError("request_occupancy_tokens requires an FPM profile") if fpm_profile is not None: - if nextn: - raise ValueError("FPM profile resources support plain autoregressive execution; nextn must be 0") profile = load_fpm_profile(fpm_profile) if backend_version is None: raise ValueError("FPM profile memory estimation requires the profile's literal backend_version") @@ -1182,6 +1180,8 @@ def estimate_kv_cache( moe_ep_size=moe_ep_size, cp_size=cp_size, ) + if nextn and deployment.execution_profile is None: + raise ValueError("FPM profile resources support plain autoregressive execution; nextn must be 0") deployment.validate_overrides( gemm_quant_mode=gemm_quant_mode, moe_quant_mode=moe_quant_mode, @@ -1208,6 +1208,7 @@ def estimate_kv_cache( resources.cache_layout == "grouped" or resources.runtime_memory is not None or request_occupancy_tokens is not None + or deployment.execution_profile is not None ): from aisimulate_core.sdk.rust_engine_step import RustForwardPassPerfModel @@ -1223,7 +1224,21 @@ def estimate_kv_cache( "attention_dp": attention_dp_size, "moe_tp_size": deployment.moe_tp, "moe_ep_size": deployment.moe_ep, + "cp_size": cp_size, + # This API's default DCP1 means no sharding, not a recorded + # DCP selector. Preserve nontrivial DCP for native validation. + **({"dcp": dcp_size} if dcp_size != 1 else {}), "fpm_profile": profile.model_dump(mode="json"), + **( + { + "nextn": nextn, + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "estimator_config": {"fpm_interpolation": {"method": "direct"}}, + } + if deployment.execution_profile is not None + else {} + ), }, { "total_gpu_capacity_bytes": capacity, diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py index ee7eed011..3667e2a63 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py @@ -273,6 +273,8 @@ def test_unknown_profile_auto_compiles_without_operations(profile_dict, direct_c assert json.loads(spec["engine"]["extra"]["fpm_profile"]) == profile.model_dump(mode="json") assert json.loads(spec["engine"]["extra"]["estimator_config"])["fpm_interpolation"]["method"] == "direct" assert spec["engine"]["activation_dtype"] == "fp8" + assert "nextn" not in spec["engine"] + assert "fpm_execution_profile" not in spec["engine"]["extra"] for key, phase in (("context_ops", "prefill"), ("generation_ops", "decode")): assert len(spec[key]) == 1 op = spec[key][0]["FpmForward"] @@ -812,6 +814,139 @@ def runtime_profile(profile_dict): return profile_dict +@pytest.fixture +def speculative_profile(runtime_profile, monkeypatch): + runtime_profile["model_revision"] = "1" * 40 + runtime_profile["deployments"][0]["execution_profile"] = { + "kind": "speculative_round", + "method": "eagle3", + "num_draft_tokens": 3, + "draft_model": "test/draft", + "draft_revision": "2" * 40, + "timing_scope": "model_step_cuda", + } + + def download(model, filename, *, revision): + assert (model, filename, revision) == (runtime_profile["model"], "config.json", "1" * 40) + return {"architectures": [runtime_profile["architecture"]]} + + monkeypatch.setattr(utils, "_download_hf_json", download) + return runtime_profile + + +def test_speculative_round_direct_spec_preserves_identity_and_raw_batch_width(speculative_profile, direct_compile): + spec = direct_compile(speculative_profile, fpm_interpolation="direct", nextn=3) + profile = load_fpm_profile(speculative_profile) + deployment = profile.deployments[0] + binding = { + "model_revision": profile.model_revision, + "execution_profile": deployment.execution_profile.model_dump(mode="json"), + } + assert spec["engine"]["nextn"] == 3 + assert json.loads(spec["engine"]["extra"]["fpm_execution_profile"]) == binding + saved = json.loads(spec["engine"]["extra"]["fpm_profile"]) + assert _normalize(_request(saved, "direct", nextn=3))["fpm_profile"] == saved + expected_identity = deployment.match_identity() + list(deployment.execution_identity(profile.model_revision)) + assert len(expected_identity) == 19 + for phase in ("context_ops", "generation_ops"): + assert len(spec[phase]) == 1 + op = spec[phase][0]["FpmForward"] + assert op["match_identity"] == expected_identity + assert op["verify_width"] == 1 # Full-round curves already use request count, not verify-token count. + assert op["interpolation"] == "direct" + assert op["sol_ops"] == [] + + +@pytest.mark.parametrize( + "overrides", + [ + {"nextn": 0}, + {"nextn": 2}, + {"estimation_mode": "auto"}, + {"fallback_policy": "allow"}, + {"nextn": 0, "speculation": {"kind": "mtp", "params": {"depth": 3}}}, + {"nextn": 0, "speculation": {"kind": "ngram", "params": {"num_speculative_tokens": 3}}}, + ], +) +def test_speculative_round_profile_rejects_implicit_or_mismatched_execution(speculative_profile, overrides): + config = _request(speculative_profile, "direct", nextn=3) + config.update(overrides) + with pytest.raises(ValueError): + _normalize(config) + + +def test_speculative_round_memory_forwards_depth_and_uses_observed_pool( + speculative_profile, profile_memory, monkeypatch +): + calls = [] + estimate = RustForwardPassPerfModel.estimate_cache_budget + + def observe(config, request): + calls.append(copy.deepcopy(config)) + return estimate(config, request) + + monkeypatch.setattr(RustForwardPassPerfModel, "estimate_cache_budget", observe) + monkeypatch.setattr(memory.KVCacheEstimator, "_estimate_from_breakdown", _fail_graph) + result = profile_memory(nextn=3, context_length=1024, gpu_memory_capacity_bytes_override=2000) + assert calls[0]["nextn"] == 3 + assert calls[0]["estimation_mode"] == "fpm_interpolation" + assert calls[0]["fallback_policy"] == "deny" + assert calls[0]["estimator_config"]["fpm_interpolation"]["method"] == "direct" + assert ( + calls[0]["fpm_profile"]["deployments"][0]["execution_profile"] + == (speculative_profile["deployments"][0]["execution_profile"]) + ) + assert result["total_kv_size_bytes"] == 600 # Recorded pool, independent of the 2000-byte device bound. + assert result["total_kv_size_tokens"] == 60 # The fixture records ten bytes per token. + assert result["memory_breakdown"] is None + with pytest.raises(ValueError): + profile_memory(nextn=0, context_length=1024, allow_naive_fallback=True) + + +def test_speculative_round_memory_rejects_declared_budget_without_analytical_fallback( + speculative_profile, profile_memory, monkeypatch +): + resources = speculative_profile["deployments"][0]["resources"] + resources.pop("runtime_memory") + resources.update(weights_bytes=100, activations_bytes=20, runtime_overhead_bytes=30, comm_overhead_bytes=50) + monkeypatch.setattr(memory.KVCacheEstimator, "_estimate_from_breakdown", _fail_graph) + monkeypatch.setattr(memory.NaiveKVCacheEstimator, "from_model_path", _fail_graph) + with pytest.raises(ValueError, match="runtime|observed"): + profile_memory(nextn=3, context_length=1024, allow_naive_fallback=True) + + +@pytest.mark.parametrize("entry", [memory.estimate_kv_cache, memory.estimate_num_gpu_blocks]) +@pytest.mark.parametrize("dcp", [None, 1, 2]) +def test_speculative_round_standalone_memory_preserves_dcp(speculative_profile, profile_memory, entry, dcp): + kwargs = dict( + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.8, + tp_size=2, + moe_tp_size=2, + nextn=3, + fpm_profile=speculative_profile, + context_length=1024, + allow_naive_fallback=True, + ) + if dcp is not None: + kwargs["dcp_size"] = dcp + if entry is memory.estimate_num_gpu_blocks: + kwargs["scheduler_block_size"] = 16 + + def estimate(): + return entry("test/unknown-decoder", "test_gpu", "vllm", "0.25.1", **kwargs) + + if dcp == 2: + with pytest.raises(ValueError, match="DCP"): + estimate() + elif entry is memory.estimate_kv_cache: + assert estimate()["total_kv_size_tokens"] == 60 # Observed 600-byte pool / 10 bytes per token. + else: + assert estimate() == 3 # floor(60 tokens / 16 tokens per block), not another draft deduction. + + def test_runtime_memory_selects_exact_role_capacity_and_scheduler(runtime_profile, profile_memory): prefill = runtime_profile["deployments"][0] prefill["worker_type"] = "prefill" diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py b/python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py index f5fccacb8..374d62273 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py @@ -188,14 +188,17 @@ def test_unrelated_speculative_deployment_does_not_block_selected_ar(monkeypatch assert after["fpm_profile"]["deployments"][1]["execution_profile"] == execution() -@pytest.mark.parametrize("nextn", [0, 3]) +@pytest.mark.parametrize( + "nextn,error", + [(0, "nextn must match"), (2, "nextn must match"), (3, "observed runtime memory")], +) @pytest.mark.parametrize("entry", ["normalize_config", "best_available"]) -def test_identity_publication_does_not_enable_unimplemented_execution(nextn, entry, monkeypatch): +def test_speculative_execution_rejects_depth_mismatch_and_pending_memory(nextn, error, entry, monkeypatch): monkeypatch.setattr( fpm_profile, "_load_pinned_checkpoint_config", lambda *_: {"architectures": ["TestDecoderForCausalLM"]} ) config = canonical_request(profile(), nextn=nextn) - with pytest.raises(ValueError, match="identity-only|does not support MTP"): + with pytest.raises(ValueError, match=error): getattr(RustForwardPassPerfModel, entry)(config) diff --git a/tests/test_fpm_profile_workflow.py b/tests/test_fpm_profile_workflow.py index 90ccee1f9..7a1aa75c4 100644 --- a/tests/test_fpm_profile_workflow.py +++ b/tests/test_fpm_profile_workflow.py @@ -206,6 +206,120 @@ def test_prediction_materializes_rank_local_capacity_without_a_model(profile, ti assert reserved["num_gpu_blocks"] < lowered["num_gpu_blocks"] +def test_speculative_round_prediction_preserves_progress_identity_and_observed_capacity(profile, timing_systems): + import pyarrow as pa + import pyarrow.parquet as pq + + from aisimulate import _runtime + from aisimulate.runner import EngineReplayRunnerFactory + from aisimulate.sweeper import ReplayOutputRequirements + from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + profile["model_revision"] = "1" * 40 + profile["deployments"] = profile["deployments"][:1] + _observed_profile(profile, context=4096) + execution = { + "kind": "speculative_round", + "method": "eagle3", + "num_draft_tokens": 3, + "draft_model": "test/draft", + "draft_revision": "2" * 40, + "timing_scope": "model_step_cuda", + } + profile["deployments"][0]["execution_profile"] = execution + deployment = load_fpm_profile(profile).deployments[0] + execution_identity = dict( + zip(EXECUTION_COLUMNS, deployment.execution_identity(profile["model_revision"]), strict=True) + ) + path = Path(timing_systems) / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + rows = [row for row in pq.read_table(path).to_pylist() if row["dp"] == 1] + for row in rows: + row.update( + **execution_identity, + measurement_policy="vllm_native_fpm_model_step_cuda_v1", + timing_scope="model_step_cuda", + measurement_reduction="single_sample", + source_sample_count=1, + # Deliberately synthetic per-round times, not an analytical estimate. + latency_ms=7.0 if row["workload_kind"] == "prefill" else 11.0, + ) + pq.write_table(pa.Table.from_pylist(rows), path) + metadata = json.loads(path.with_suffix(".metadata.json").read_text()) + metadata.update( + schema_version=7, + measurement_policy="vllm_native_fpm_model_step_cuda_v1", + timing_scope="model_step_cuda", + row_count=len(rows), + parquet_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + ) + path.with_suffix(".metadata.json").write_text(json.dumps(metadata)) + public = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(profile), + "systems_paths": [timing_systems], + "context_length": 4096, + "nextn": 3, + "nextn_accepted": 1.78, + }, + "traffic": { + "source": {"type": "synthetic", "input_tokens": 1, "output_tokens": 7}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 2}, + }, + } + ) + restored = CorePredictionConfig.model_validate_json(public.model_dump_json()) + spec = prediction_to_replay_spec(restored) + diagnostics = {} + materialized = materialize_aic_num_gpu_blocks( + spec.backend_deployment.agg_engine_args, memory_diagnostics=diagnostics + ) + # Observed 1 GiB / 512 bytes per token / 64 tokens per block; no draft or graph charge is added. + assert materialized["num_gpu_blocks"] == 32768 + assert diagnostics["total_kv_size_bytes"] == 1024**3 + assert diagnostics["memory_breakdown"] is None + canonical_only = deepcopy(spec.backend_deployment.agg_engine_args) + canonical_only.pop("aic_nextn") + assert materialize_aic_num_gpu_blocks(canonical_only)["num_gpu_blocks"] == 32768 + for legacy_name in ("aic_nextn", "nextn"): + conflicting_depth = deepcopy(canonical_only) + conflicting_depth[legacy_name] = 2 + with pytest.raises(ValueError, match="conflicts with canonical"): + materialize_aic_num_gpu_blocks(conflicting_depth) + seen = [] + + class NativeCapture: + def run_replay_json(self, serialized): + seen.append(json.loads(serialized)) + return _runtime.run_replay_json(serialized) + + report = ( + EngineReplayRunnerFactory(runtime=NativeCapture()) + .create(0) + .run(spec, output_requirements=ReplayOutputRequirements(include_raw_report=True, capture_per_request=True)) + ) + rank = seen[0]["spec"]["engine"]["rank"] + assert rank["aic_nextn"] == 3 + # 1.78 accepted drafts plus the mandatory base token means 2.78 output tokens per round. + assert [float(rate) for rate in rank["aic_nextn_accept_rates"].split(",")] == [1.0, 0.78, 0.0] + assert rank["num_gpu_blocks"] == 32768 + assert rank["timing_model"]["config"]["fpm_profile"]["deployments"][0]["execution_profile"] == execution + assert report.metrics["completed_requests"] == 2 + assert report.metrics["mean_ttft_ms"] == pytest.approx(7.0) + records = report.metadata["native_report"]["per_request"] + assert len(records) == 2 + for record in records: + assert record["output_length"] == 7 + assert record["ttft_ms"] == pytest.approx(7.0) + # One prefill token leaves six outputs: 2-3 tokens per round require + # two or three measured 11 ms rounds, with no extra draft charge. + rounds = (record["e2e_latency_ms"] - 7.0) / 11.0 + assert rounds == pytest.approx(2.0) or rounds == pytest.approx(3.0) + assert record["itl_ms"] == pytest.approx(11.0 * rounds / 6.0) + + @pytest.mark.parametrize("memory_state", ["pending", "runtime"]) def test_profile_memory_cannot_be_bypassed_by_explicit_blocks(profile, timing_systems, memory_state): config = CorePredictionConfig.model_validate({"engine": {**_engine(profile), "systems_paths": [timing_systems]}})