From 06960c59f14ceb946a8c515573ce468481514d4b Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 16:43:57 -0700 Subject: [PATCH 01/13] feat: define speculative-round FPM identity and native CUDA data contract Signed-off-by: hongkuanz --- crates/core/src/lib.rs | 4 + crates/core/src/perfmodel/fpm/config.rs | 19 ++ crates/core/src/perfmodel/fpm/execution.rs | 186 ++++++++++++++ crates/core/src/perfmodel/fpm/mod.rs | 5 + crates/core/src/perfmodel/mod.rs | 4 + .../perfmodel/perf_database/fpm_forward.rs | 239 +++++++++++++++++- crates/core/src/perfmodel/py.rs | 18 ++ docs/perf-model/methods/whole-forward.md | 57 +++++ .../src/aisimulate_core/fpm_profile.py | 26 ++ .../src/aisimulate_core/sdk/fpm_identity.py | 10 + ...test_fpm_speculative_execution_identity.py | 172 +++++++++++++ 11 files changed, 737 insertions(+), 3 deletions(-) create mode 100644 crates/core/src/perfmodel/fpm/execution.rs create mode 100644 python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index c73203efc..0d942c1a5 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -54,6 +54,10 @@ pub use perfmodel::{ RegressionSamplingConfig, RegressionUpdatePolicy, SplineFitConfig, SplineSearchConfig, UnrecordedFpmQuantMode, }; +pub use perfmodel::{ + FPM_SPECULATIVE_ROUND_PREFIX, FpmExecutionBinding, FpmExecutionKind, FpmExecutionProfile, + FpmSpeculativeMethod, FpmTimingScope, +}; #[cfg(feature = "python")] pub use perfmodel::{ diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 9510bc5c7..bf0fd6f19 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -248,6 +248,25 @@ impl ForwardPassPerfModelConfig { pub(crate) fn resolve_with_registration(mut self) -> Result<(Self, Option), AicError> { self.estimator_config.resolve_defaults(); self.validate()?; + // Identity publication precedes the consumer implementation. Do not let + // either nextn=0 or nextn>0 consume a full-round profile as AR timing. + if self + .fpm_profile + .as_ref() + .and_then(|profile| 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(|execution| !execution.is_null()) + }) + }) + { + return Err(invalid_config( + "speculative-round FPM profiles are identity-only; execution is not enabled", + )); + } let registered = if self.fpm_profile.is_some() { self.estimator_config .fpm_interpolation diff --git a/crates/core/src/perfmodel/fpm/execution.rs b/crates/core/src/perfmodel/fpm/execution.rs new file mode 100644 index 000000000..45638721b --- /dev/null +++ b/crates/core/src/perfmodel/fpm/execution.rs @@ -0,0 +1,186 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Exact recorded execution identity; no timing math or alternative model factory. + +use serde::{Deserialize, Serialize}; + +use crate::AicError; + +pub const FPM_SPECULATIVE_ROUND_PREFIX: &str = "speculative_round_v1:"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmExecutionKind { + SpeculativeRound, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmSpeculativeMethod { + Eagle3, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FpmTimingScope { + ModelStepCuda, +} + +/// A complete target-plus-draft iteration, measured with the native CUDA timer. +/// Acceptance affects replay progress and is deliberately not a timing selector. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FpmExecutionProfile { + pub kind: FpmExecutionKind, + pub method: FpmSpeculativeMethod, + pub num_draft_tokens: u32, + pub draft_model: String, + pub draft_revision: String, + pub timing_scope: FpmTimingScope, +} + +fn invalid(message: impl Into) -> AicError { + AicError::InvalidEngineConfig(format!( + "invalid FPM execution identity: {}", + message.into() + )) +} + +fn pinned_revision(name: &str, value: &str) -> Result<(), AicError> { + if value.len() != 40 + || !value + .bytes() + .all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase()) + { + return Err(invalid(format!( + "{name} must be a lowercase 40-hex immutable checkpoint revision" + ))); + } + Ok(()) +} + +impl FpmExecutionProfile { + pub fn validate(&self) -> Result<(), AicError> { + if !(1..=5).contains(&self.num_draft_tokens) { + return Err(invalid("num_draft_tokens must be in 1..=5")); + } + if self.draft_model.trim().is_empty() + || self.draft_model.trim() != self.draft_model + || self.draft_model.chars().any(char::is_control) + { + return Err(invalid( + "draft_model must be nonblank, trimmed, and contain no control characters", + )); + } + pinned_revision("draft_revision", &self.draft_revision) + } +} + +/// Parent checkpoint revision and the deployment's full-round identity. +/// Field order defines the version-1 selector serialization; consumers must use +/// `selector`, not invent a JSON ordering or use only a draft/checkpoint hash. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FpmExecutionBinding { + pub model_revision: String, + pub execution_profile: FpmExecutionProfile, +} + +impl FpmExecutionBinding { + pub fn validate(&self) -> Result<(), AicError> { + pinned_revision("model_revision", &self.model_revision)?; + self.execution_profile.validate() + } + + pub fn selector(&self) -> Result { + self.validate()?; + let json = serde_json::to_string(self).map_err(|error| invalid(error.to_string()))?; + Ok(format!("{FPM_SPECULATIVE_ROUND_PREFIX}{json}")) + } + + /// Ordinary AR and other established execution labels keep their behavior. + /// The speculative namespace is reserved: malformed/unknown versions and + /// noncanonical encodings fail instead of becoming silently separate cells. + pub fn parse_selector(selector: &str) -> Result, AicError> { + if !selector.starts_with("speculative_round") { + return Ok(None); + } + let payload = selector + .strip_prefix(FPM_SPECULATIVE_ROUND_PREFIX) + .ok_or_else(|| invalid("unsupported speculative_round selector version"))?; + let binding: Self = + serde_json::from_str(payload).map_err(|error| invalid(error.to_string()))?; + if binding.selector()? != selector { + return Err(invalid( + "speculative_round selector must use canonical serialization", + )); + } + Ok(Some(binding)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn binding() -> FpmExecutionBinding { + FpmExecutionBinding { + model_revision: "1".repeat(40), + execution_profile: FpmExecutionProfile { + kind: FpmExecutionKind::SpeculativeRound, + method: FpmSpeculativeMethod::Eagle3, + num_draft_tokens: 3, + draft_model: "test/draft".into(), + draft_revision: "2".repeat(40), + timing_scope: FpmTimingScope::ModelStepCuda, + }, + } + } + + #[test] + fn selector_pins_every_identity_without_changing_legacy_labels() { + let original = binding(); + let selector = original.selector().unwrap(); + assert_eq!( + FpmExecutionBinding::parse_selector(&selector).unwrap(), + Some(original.clone()) + ); + for label in ["full", "decoder_bounded", "prefix_swa_replay", ""] { + assert_eq!(FpmExecutionBinding::parse_selector(label).unwrap(), None); + } + for (field, value) in [ + ("model_revision", "3".repeat(40)), + ("draft_revision", "4".repeat(40)), + ("draft_model", "test/other-draft".into()), + ("num_draft_tokens", "4".into()), + ] { + let mut changed = original.clone(); + match field { + "model_revision" => changed.model_revision = value, + "draft_revision" => changed.execution_profile.draft_revision = value, + "draft_model" => changed.execution_profile.draft_model = value, + _ => changed.execution_profile.num_draft_tokens = 4, + } + assert_ne!(changed.selector().unwrap(), selector); + } + } + + #[test] + fn malformed_or_noncanonical_speculative_identity_is_rejected() { + let selector = binding().selector().unwrap(); + for bad in [ + selector.replacen(':', ": ", 1), + selector.replace("_v1:", "_v2:"), + selector.replace("\"eagle3\"", "\"ngram\""), + selector.replace("model_step_cuda", "scheduler_wall_time"), + selector.replace("\"num_draft_tokens\":3", "\"num_draft_tokens\":true"), + selector.replace("\"num_draft_tokens\":3", "\"num_draft_tokens\":0"), + selector.replace("\"num_draft_tokens\":3", "\"num_draft_tokens\":6"), + selector.replace("\"kind\":", "\"unknown\":1,\"kind\":"), + selector.replace(&"2".repeat(40), "main"), + ] { + assert!(FpmExecutionBinding::parse_selector(&bad).is_err(), "{bad}"); + } + } +} diff --git a/crates/core/src/perfmodel/fpm/mod.rs b/crates/core/src/perfmodel/fpm/mod.rs index a6f71beea..945e4ea0e 100644 --- a/crates/core/src/perfmodel/fpm/mod.rs +++ b/crates/core/src/perfmodel/fpm/mod.rs @@ -28,6 +28,7 @@ mod correction; pub(crate) mod coverage; mod estimator; mod evidence; +mod execution; mod learned; mod metrics; mod model; @@ -47,6 +48,10 @@ pub use coverage::{FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQue pub use estimator::*; pub(crate) use evidence::accumulate_fpm_estimates; pub use evidence::*; +pub use execution::{ + FPM_SPECULATIVE_ROUND_PREFIX, FpmExecutionBinding, FpmExecutionKind, FpmExecutionProfile, + FpmSpeculativeMethod, FpmTimingScope, +}; pub use learned::LEARNED_SCHEMA_VERSION; pub(crate) use metrics::validate_forward_pass_metrics; pub use metrics::{FPM_VERSION, ForwardPassMetrics, QueuedRequestMetrics, ScheduledRequestMetrics}; diff --git a/crates/core/src/perfmodel/mod.rs b/crates/core/src/perfmodel/mod.rs index 05ef85cb9..162390a74 100644 --- a/crates/core/src/perfmodel/mod.rs +++ b/crates/core/src/perfmodel/mod.rs @@ -69,6 +69,10 @@ pub use fpm::{ // `crate::ForwardPassMetrics` / `crate::FPM_VERSION` references (in `py.rs`, // `engine/runtime.rs`) keep resolving after the types moved into `fpm`. pub(crate) use fpm::validate_forward_pass_metrics; +pub use fpm::{ + FPM_SPECULATIVE_ROUND_PREFIX, FpmExecutionBinding, FpmExecutionKind, FpmExecutionProfile, + FpmSpeculativeMethod, FpmTimingScope, +}; pub use fpm::{FPM_VERSION, ForwardPassMetrics, QueuedRequestMetrics, ScheduledRequestMetrics}; // KV-cache memory API. Top-level surface, not a method on // `AicEngine`: estimation runs once at startup, separate from the latency path. diff --git a/crates/core/src/perfmodel/perf_database/fpm_forward.rs b/crates/core/src/perfmodel/perf_database/fpm_forward.rs index 751630be4..ebb3fc2ce 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -62,6 +62,7 @@ pub const FPM_FORWARD_PARTITION_POLICY: &str = "balanced_v1"; /// contract and the backend-bound native-FPM median contracts below; median /// latency is already reduced. pub const FPM_FORWARD_MEASUREMENT_POLICY: &str = "dynamo_native_single_sample_v1"; +pub const FPM_NATIVE_CUDA_MEASUREMENT_POLICY: &str = "vllm_native_fpm_model_step_cuda_v1"; const FPM_ROW_MEASUREMENT_POLICY: &str = "per_row_single_sample_or_median_of_3"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -74,6 +75,9 @@ enum MeasurementPolicy { /// Default serving configuration (radix cache on): native SGLang /// `--enable-forward-pass-metrics` rank-0 wall time, real radix seeding. SglangNativeFpmMedian, + /// A genuine native CUDA observation, optionally selected as an actual + /// upper-median source sample with that sample's own workload coordinates. + NativeCudaActualSample, } impl MeasurementPolicy { @@ -83,12 +87,13 @@ impl MeasurementPolicy { Self::PerRow => FPM_ROW_MEASUREMENT_POLICY, Self::VllmNativeFpmMedian => "vllm_native_fpm_prefix_seed_median_v1", Self::SglangNativeFpmMedian => "sglang_native_fpm_radix_seed_median_v1", + Self::NativeCudaActualSample => FPM_NATIVE_CUDA_MEASUREMENT_POLICY, } } fn state_protocol(self) -> Option<&'static str> { match self { - Self::SingleSample | Self::PerRow => None, + Self::SingleSample | Self::PerRow | Self::NativeCudaActualSample => None, Self::VllmNativeFpmMedian => Some("glm53flash_prefix_cache_real_seed_v1"), Self::SglangNativeFpmMedian => Some("glm53flash_sglang_radix_real_seed_v1"), } @@ -96,7 +101,7 @@ impl MeasurementPolicy { fn timing_boundary(self) -> Option<&'static str> { match self { - Self::SingleSample | Self::PerRow => None, + Self::SingleSample | Self::PerRow | Self::NativeCudaActualSample => None, Self::VllmNativeFpmMedian => Some("dynamo_vllm_instrumented_scheduler_fpm_wall_time"), Self::SglangNativeFpmMedian => Some("sglang_native_fpm_rank0_device_timer_wall_time"), } @@ -539,6 +544,14 @@ fn validate_sidecar( (Some("sglang_native_fpm_radix_seed_median_v1"), Some(7), "sglang") => { MeasurementPolicy::SglangNativeFpmMedian } + (Some(FPM_NATIVE_CUDA_MEASUREMENT_POLICY), Some(7), "vllm") => { + if metadata.get("timing_scope").and_then(|v| v.as_str()) != Some("model_step_cuda") { + return Err(structural( + "native CUDA FPM sidecar requires timing_scope=model_step_cuda".into(), + )); + } + MeasurementPolicy::NativeCudaActualSample + } _ => { return Err(structural(format!( "unsupported FPM measurement_policy={:?} for schema {schema_version:?} backend {backend:?}: {}", @@ -741,6 +754,15 @@ fn load_pair( } else { None }; + let native_cuda_columns = if measurement_policy == MeasurementPolicy::NativeCudaActualSample { + Some([ + reader.col("timing_scope")?, + reader.col("measurement_reduction")?, + reader.col("source_sample_count")?, + ]) + } else { + None + }; // Python checks the sidecar row_count against the FULL row list before any // per-row validation (`load_fpm_forward_data`: read_table -> row_count -> @@ -767,7 +789,7 @@ fn load_pair( let mut rows: Vec = Vec::new(); for (index, row) in reader.rows()?.enumerate() { let row = row?; - if median_columns.is_some() + if (median_columns.is_some() || native_cuda_columns.is_some()) && row.str_optional(row_policy_col)? != Some(measurement_policy.name()) { return Err(structural(format!( @@ -775,6 +797,20 @@ fn load_pair( measurement_policy.name() ))); } + if let Some([scope, reduction, count]) = native_cuda_columns { + let count = row.u32(count)?; + let reduction = row.str(reduction)?; + if row.str(scope)? != "model_step_cuda" + || !matches!( + (reduction, count), + ("single_sample", 1) | ("actual_upper_median", 2..) + ) + { + return Err(structural(format!( + "FPM row {index}: native CUDA timing/reduction/count contract mismatch" + ))); + } + } if let Some([global, warmup, measured, protocol, boundary]) = median_columns { for (name, col, expected) in [ ("global_warmup_iterations", global, 0), @@ -929,7 +965,38 @@ fn load_pair( } }) .collect::>()?; + let speculative = crate::FpmExecutionBinding::parse_selector(&match_identity[16]) + .map_err(|error| structural(format!("FPM row {index}: {error}")))?; + if speculative.is_some() + != (measurement_policy == MeasurementPolicy::NativeCudaActualSample) + { + return Err(structural(format!( + "FPM row {index}: speculative-round identity requires the native CUDA measurement policy, and vice versa" + ))); + } + if speculative.is_some() + && (backend != "vllm" + || get_int("tp")? == 0 + || get_int("pp")? != 1 + || get_int("dp")? != 1 + || get_int("cp")? != 1 + || get_int("moe_ep")? != 1 + || get_int("moe_tp")? != get_int("tp")? + || match_identity[13] != "False" + || match_identity[14] != "False" + || match_identity[17] != "none" + || match_identity[18] != "text") + { + return Err(structural(format!( + "FPM row {index}: speculative-round data requires vLLM pure TP, DP1/PP1/CP1 and text execution" + ))); + } let dcp = dcp_col.map(|col| row.u32(col)).transpose()?; + if speculative.is_some() && dcp.is_some_and(|value| value != 1) { + return Err(structural(format!( + "FPM row {index}: speculative-round data requires DCP1" + ))); + } if let Some(dcp) = dcp { let tp = get_int("tp")?; if dcp == 0 || tp == 0 || tp % dcp != 0 { @@ -1451,6 +1518,13 @@ pub(crate) mod tests { pub measurement: Option<(&'static str, u32)>, pub execution: Option<[&'static str; 4]>, pub median: Option, + pub native_cuda: Option, + } + + pub(crate) struct NativeCudaSampleSpec { + pub scope: &'static str, + pub reduction: &'static str, + pub source_sample_count: u32, } pub(crate) struct MeasurementSpec { @@ -1486,6 +1560,7 @@ pub(crate) mod tests { measurement: None, execution: None, median: None, + native_cuda: None, } } } @@ -1554,6 +1629,7 @@ pub(crate) mod tests { let has_measurement = rows.iter().any(|r| r.measurement.is_some()); let has_execution = rows.iter().any(|r| r.execution.is_some()); let has_median = rows.iter().any(|r| r.median.is_some()); + let has_native_cuda = rows.iter().any(|r| r.native_cuda.is_some()); let schema = "message schema { REQUIRED BINARY cell_id (UTF8); REQUIRED BINARY model_path (UTF8); @@ -1628,6 +1704,11 @@ pub(crate) mod tests { } else { schema }; + let schema = if has_native_cuda { + schema.replace('}', "REQUIRED BINARY measurement_policy (UTF8);\nREQUIRED BINARY timing_scope (UTF8);\nREQUIRED BINARY measurement_reduction (UTF8);\nREQUIRED INT64 source_sample_count;\n}") + } else { + schema + }; let schema = Arc::new(parse_message_type(&schema).expect("schema must parse")); let file = std::fs::File::create(&parquet_path).expect("create parquet"); let mut writer = @@ -1833,6 +1914,33 @@ pub(crate) mod tests { col.close().unwrap(); } } + if has_native_cuda { + for field in 0..3 { + let values = str_col(&|r| { + let sample = r.native_cuda.as_ref().unwrap(); + [ + FPM_NATIVE_CUDA_MEASUREMENT_POLICY, + sample.scope, + sample.reduction, + ][field] + .to_owned() + }); + let mut col = rg.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&values, None, None) + .unwrap(); + col.close().unwrap(); + } + let values: Vec = rows + .iter() + .map(|r| i64::from(r.native_cuda.as_ref().unwrap().source_sample_count)) + .collect(); + let mut col = rg.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&values, None, None) + .unwrap(); + col.close().unwrap(); + } rg.close().expect("close row group"); writer.close().expect("close writer"); @@ -1865,6 +1973,13 @@ pub(crate) mod tests { sidecar.insert("system".into(), "b200_sxm".into()); sidecar.insert("backend".into(), "vllm".into()); sidecar.insert("backend_version".into(), "0.25.1".into()); + if has_native_cuda { + sidecar.insert( + "measurement_policy".into(), + FPM_NATIVE_CUDA_MEASUREMENT_POLICY.into(), + ); + sidecar.insert("timing_scope".into(), "model_step_cuda".into()); + } mutate_sidecar(&mut sidecar); std::fs::write( parquet_path.with_extension("metadata.json"), @@ -1880,6 +1995,124 @@ pub(crate) mod tests { FpmForwardTable::new_with_replacement(dir.to_path_buf(), "b200_sxm", "vllm", "0.25.1", true) } + fn speculative_rows() -> Vec { + let binding: crate::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": "test/draft", "draft_revision": "2".repeat(40), + "timing_scope": "model_step_cuda" + } + })) + .unwrap(); + let selector: &'static str = Box::leak(binding.selector().unwrap().into_boxed_str()); + default_rows() + .into_iter() + .map(|mut row| { + row.execution = Some(["", selector, "none", "text"]); + row.kv_seed_regime = Some("real_kv"); + row.native_cuda = Some(NativeCudaSampleSpec { + scope: "model_step_cuda", + reduction: "actual_upper_median", + source_sample_count: 20, + }); + row + }) + .collect() + } + + #[test] + fn speculative_round_cells_select_exactly_and_do_not_scale_actual_samples() { + let tmp = tempfile::tempdir().unwrap(); + let rows = speculative_rows(); + write_pair(tmp.path(), &rows); + let table = loaded_table(tmp.path()); + let mut identity = default_identity(4); + identity[16] = rows[0].execution.unwrap()[1].to_owned(); + let cell = table.select_cell(&identity, "org/model-a", None).unwrap(); + // Fixture records 7 ms at B8/KV4096. Twenty contributing source + // observations select one actual sample; never divide its time by 20. + assert_eq!(cell.direct_decode[&8][&4096], 7.0); + assert!( + table + .select_cell(&default_identity(4), "org/model-a", None) + .is_err() + ); + let binding = crate::FpmExecutionBinding::parse_selector(&identity[16]) + .unwrap() + .unwrap(); + for changed in [ + identity[16].replace("\"num_draft_tokens\":3", "\"num_draft_tokens\":4"), + identity[16].replace("test/draft", "test/other"), + identity[16].replace(&binding.model_revision, &"3".repeat(40)), + identity[16].replace(&binding.execution_profile.draft_revision, &"4".repeat(40)), + ] { + identity[16] = changed; + assert!(table.select_cell(&identity, "org/model-a", None).is_err()); + } + } + + #[test] + fn speculative_native_policy_accepts_unreduced_real_samples() { + let tmp = tempfile::tempdir().unwrap(); + let mut rows = speculative_rows(); + for row in &mut rows { + row.native_cuda = Some(NativeCudaSampleSpec { + scope: "model_step_cuda", + reduction: "single_sample", + source_sample_count: 1, + }); + } + write_pair(tmp.path(), &rows); + assert_eq!(loaded_table(tmp.path()).cells().unwrap().len(), 1); + } + + #[test] + fn speculative_native_policy_rejects_clock_relabeling_and_bad_reductions() { + for mutation in [ + "legacy_policy", + "ar_identity", + "row_clock", + "sidecar_clock", + "fake_kv", + "zero_count", + "single_count", + "unknown_reduction", + "malformed_selector", + ] { + let tmp = tempfile::tempdir().unwrap(); + let mut rows = speculative_rows(); + for row in &mut rows { + match mutation { + "ar_identity" => row.execution = Some(LEGACY_EXECUTION_IDENTITY), + "row_clock" => row.native_cuda.as_mut().unwrap().scope = "scheduler_wall_time", + "fake_kv" => row.kv_seed_regime = Some("fake_fallback"), + "zero_count" => row.native_cuda.as_mut().unwrap().source_sample_count = 0, + "single_count" => row.native_cuda.as_mut().unwrap().reduction = "single_sample", + "unknown_reduction" => row.native_cuda.as_mut().unwrap().reduction = "mean", + "malformed_selector" => { + row.execution = Some(["", "speculative_round_v1:{}", "none", "text"]) + } + _ => {} + } + } + write_pair_with(tmp.path(), &rows, |metadata| { + if mutation == "legacy_policy" { + metadata.insert( + "measurement_policy".into(), + FPM_FORWARD_MEASUREMENT_POLICY.into(), + ); + } + if mutation == "sidecar_clock" { + metadata.insert("timing_scope".into(), "scheduler_wall_time".into()); + } + }); + assert!( + loaded_table(tmp.path()).cells().is_err(), + "accepted {mutation}" + ); + } + } + #[test] fn external_parquet_path_supports_an_arbitrary_stem() { let tmp = tempfile::tempdir().expect("tmpdir"); diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 566339ce0..9e3df26c3 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1961,6 +1961,24 @@ impl PyForwardPassPerfModel { serde_json::to_string(&estimate).map_err(|e| PyValueError::new_err(e.to_string())) } + /// Normalize recorded execution identity without constructing a timing model. + #[staticmethod] + fn _normalize_fpm_execution_profile( + model_revision: &str, + execution_profile_json: &str, + ) -> PyResult { + let execution_profile: crate::FpmExecutionProfile = + serde_json::from_str(execution_profile_json).map_err(|error| { + PyValueError::new_err(format!("invalid FPM execution profile: {error}")) + })?; + crate::FpmExecutionBinding { + model_revision: model_revision.to_owned(), + execution_profile, + } + .selector() + .map_err(aic_to_py) + } + /// Share typed FPM option validation with the compilation adapter. #[staticmethod] #[pyo3(signature = (options_json, fmha_quant_mode=None, comm_quant_mode=None, has_profile=false))] diff --git a/docs/perf-model/methods/whole-forward.md b/docs/perf-model/methods/whole-forward.md index 48896735d..e22898b8b 100644 --- a/docs/perf-model/methods/whole-forward.md +++ b/docs/perf-model/methods/whole-forward.md @@ -14,6 +14,63 @@ Use [FPM self-service](../fpm-self-service/README.md) to prepare, review, collec or import a profile. The existing [implementation guide](../fpm-self-service/implementation.md) owns artifact validation, publication, checkpoints, and recovery. +## Recorded speculative-round identity + +A profile deployment may record the following optional identity. The parent +profile's `model_revision` and the draft revision must be lowercase 40-hex +immutable checkpoint commits. This schema describes a **complete target-plus-draft +round**, not an AR target-verification approximation: + +```yaml +execution_profile: + kind: speculative_round + method: eagle3 + num_draft_tokens: 3 + draft_model: example/draft-model + draft_revision: "2222222222222222222222222222222222222222" + timing_scope: model_step_cuda +``` + +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. + +Rust produces the exact table `execution_profile` string by prefixing +`speculative_round_v1:` to its canonical serialization of +`{"model_revision": ..., "execution_profile": ...}`. The descriptor's field +order is `kind`, `method`, `num_draft_tokens`, `draft_model`, `draft_revision`, +`timing_scope`. Use the profile's `execution_identity(model_revision)` helper; +do not hand-build a selector or substitute another checkpoint, draft, depth, +or timing boundary. The resulting four-column suffix is +`["", selector, "none", "text"]`; the existing Op layout and binary schema do +not change. Malformed, unknown-version, or noncanonical speculative selectors +are invalid data, not a fallback opportunity. + +### Native-CUDA measured rows + +Speculative-round rows require schema 7, backend `vllm`, sidecar +`measurement_policy: vllm_native_fpm_model_step_cuda_v1`, and sidecar +`timing_scope: model_step_cuda`. Every row carries that same measurement policy +and scope, plus: + +| Field | Meaning | +|---|---| +| `measurement_reduction` | `single_sample` or `actual_upper_median` | +| `source_sample_count` | Integer 1 for a raw sample; at least 2 for an actual upper median | +| `latency_ms` | The selected real observation's complete native CUDA elapsed time | + +An actual upper median selects a real sample and retains **that sample's own +workload coordinates**. It does not average latencies, use an unmeasured nominal +coordinate, or divide the recorded time by the sample count. Preserve raw +observations and selection provenance alongside the table. The policy does not +admit AR-labelled rows or scheduler/host wall-time rows; renaming a CPU clock is +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. + ## Direct and SOL-assisted methods - `method: direct` needs a verified resource profile and `database_mode: SILICON`. diff --git a/python/aisimulate/src/aisimulate_core/fpm_profile.py b/python/aisimulate/src/aisimulate_core/fpm_profile.py index 774eea838..b7e941bb0 100644 --- a/python/aisimulate/src/aisimulate_core/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/fpm_profile.py @@ -161,6 +161,17 @@ def validate_envelope(self, *, max_num_tokens: int, max_batch_size: int) -> None ) +class FpmSpeculativeExecutionProfile(_ProfileModel): + """Explicit full-round timing identity; Rust validates its canonical selector.""" + + kind: Literal["speculative_round"] + method: Literal["eagle3"] + num_draft_tokens: Annotated[int, Field(ge=1, le=5)] + draft_model: _Nonempty + draft_revision: Annotated[str, Field(pattern=r"^[0-9a-f]{40}$")] + timing_scope: Literal["model_step_cuda"] + + class FpmDeploymentProfile(_ProfileModel): """One hardware, runtime, topology, precision and optional serving-role identity.""" @@ -186,6 +197,7 @@ class FpmDeploymentProfile(_ProfileModel): enable_eplb: Literal[False] = False # Absent in historical shared profiles; never infer role-specific evidence. worker_type: Literal["prefill", "decode", "aggregated"] | None = None + execution_profile: FpmSpeculativeExecutionProfile | None = None resources: FpmResourceProfile @model_serializer(mode="wrap") @@ -193,6 +205,8 @@ def _preserve_legacy_identity(self, handler): result = handler(self) if self.worker_type is None: result.pop("worker_type", None) + if self.execution_profile is None: + result.pop("execution_profile", None) return result @field_validator("pp", "cp", mode="before") @@ -223,6 +237,10 @@ def _literal_version(cls, value: str | None) -> str | None: @model_validator(mode="after") def _deployment_identity(self) -> FpmDeploymentProfile: + if self.execution_profile is not None and not ( + self.dp == 1 and self.pp == 1 and self.cp == 1 and self.moe_ep == 1 and self.moe_tp == self.tp + ): + raise ValueError("speculative-round FPM identity currently requires vLLM pure TP with DP1/PP1/CP1") modes = { "gemm_quant_mode": quantization.GEMMQuantMode, "moe_quant_mode": quantization.MoEQuantMode, @@ -254,6 +272,14 @@ def match_identity(self) -> list[str]: str(self.enable_eplb), ] + def execution_identity(self, model_revision: str) -> tuple[str, ...]: + """The existing four-column suffix; absent execution preserves AR identity.""" + from .sdk.fpm_identity import LEGACY_EXECUTION_IDENTITY, speculative_execution_identity + + if self.execution_profile is None: + return LEGACY_EXECUTION_IDENTITY + return speculative_execution_identity(model_revision, self.execution_profile.model_dump(mode="json")) + def quantization_kwargs(self) -> dict[str, str]: return { "gemm_quant_mode": self.gemm_quant_mode, diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py index 1b9dac8ad..3835e44d6 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py @@ -13,6 +13,16 @@ LEGACY_EXECUTION_IDENTITY = ("", "full", "none", "text") +def speculative_execution_identity(model_revision: str, execution_profile: dict) -> tuple[str, ...]: + """Carry the Rust-validated full-round selector; do not infer timing facts.""" + import aisimulate_core + + selector = aisimulate_core.RustForwardPassPerfModel._normalize_fpm_execution_profile( + model_revision, json.dumps(execution_profile) + ) + return ("", selector, "none", "text") + + def execution_identity( raw_config: dict, *, 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 new file mode 100644 index 000000000..9b32e1c2e --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Recorded SD identity is exact, native-normalized, and cannot alias AR data.""" + +import copy +import json + +import pytest +from pydantic import ValidationError + +import aisimulate_core +from aisimulate_core.fpm_profile import FpmModelProfile, FpmSpeculativeExecutionProfile +from aisimulate_core.sdk import RustForwardPassPerfModel +from aisimulate_core.sdk.fpm_identity import LEGACY_EXECUTION_IDENTITY, speculative_execution_identity + +pytestmark = pytest.mark.unit + + +def execution(): + return { + "kind": "speculative_round", + "method": "eagle3", + "num_draft_tokens": 3, + "draft_model": "test/draft", + "draft_revision": "2" * 40, + "timing_scope": "model_step_cuda", + } + + +def profile(*, speculative=True): + deployment = { + "system": "test_gpu", + "backend": "vllm", + "backend_version": "0.28.0", + "tp": 4, + "dp": 1, + "moe_tp": 4, + "moe_ep": 1, + "gemm_quant_mode": "nvfp4", + "moe_quant_mode": "nvfp4", + "fmha_quant_mode": "fp8", + "comm_quant_mode": "half", + "kv_cache_dtype": "fp8", + "resources": { + "cache_layout": "linear", + "kv_bytes_per_token": 128, + "max_num_tokens": 16384, + "max_batch_size": 1024, + "provenance": "Synthetic pending-resource metadata, not silicon measurements.", + }, + } + if speculative: + deployment["execution_profile"] = execution() + return { + "schema_version": 1, + "model": "test/target", + "model_revision": "1" * 40, + "architecture": "TestDecoderForCausalLM", + "context_length": 1048576, + "num_experts": 128, + "provenance": "Synthetic identity fixture.", + "deployments": [deployment], + } + + +def test_native_selector_is_self_describing_and_normalizes_input_key_order(): + deployment = FpmModelProfile.model_validate(profile()).deployments[0] + suffix = deployment.execution_identity("1" * 40) + assert suffix[0] == "" and suffix[2:] == ("none", "text") + assert suffix[1].startswith("speculative_round_v1:") + assert json.loads(suffix[1].split(":", 1)[1]) == { + "model_revision": "1" * 40, + "execution_profile": execution(), + } + reversed_keys = dict(reversed(list(execution().items()))) + assert speculative_execution_identity("1" * 40, reversed_keys) == suffix + assert len(deployment.match_identity()) == 15 + + +def test_legacy_profile_json_and_identity_do_not_gain_default_execution_fields(): + value = FpmModelProfile.model_validate(profile(speculative=False)) + before = value.model_dump_json() + assert "execution_profile" not in before + assert value.deployments[0].execution_identity("legacy-tag-remains-valid") == LEGACY_EXECUTION_IDENTITY + explicit_none = copy.deepcopy(profile(speculative=False)) + explicit_none["deployments"][0]["execution_profile"] = None + assert FpmModelProfile.model_validate(explicit_none).model_dump_json() == before + assert len(value.deployments[0].match_identity()) == 15 + + +@pytest.mark.parametrize( + "field,value", + [ + ("num_draft_tokens", 0), + ("num_draft_tokens", 6), + ("num_draft_tokens", True), + ("num_draft_tokens", 3.0), + ("method", "ngram"), + ("kind", "target_verify"), + ("timing_scope", "scheduler_wall_time"), + ("draft_revision", "main"), + ("extra_field", "ignored"), + ], +) +def test_python_and_rust_reject_invalid_descriptor(field, value): + data = execution() | {field: value} + with pytest.raises((ValueError, ValidationError)): + FpmSpeculativeExecutionProfile.model_validate(data) + with pytest.raises(ValueError): + speculative_execution_identity("1" * 40, data) + + +@pytest.mark.parametrize( + "field,value", + [ + ("num_draft_tokens", 4), + ("draft_model", "test/other"), + ("draft_revision", "3" * 40), + ], +) +def test_draft_identity_changes_never_alias(field, value): + original = speculative_execution_identity("1" * 40, execution()) + assert speculative_execution_identity("1" * 40, execution() | {field: value}) != original + assert speculative_execution_identity("4" * 40, execution()) != original + + +@pytest.mark.parametrize("model_revision", ["main", "fixture-v1", "F" * 40, "1" * 39]) +def test_recorded_target_revision_is_immutable(model_revision): + with pytest.raises(ValueError, match="model_revision"): + speculative_execution_identity(model_revision, execution()) + + +def test_profile_rejects_attention_dp_or_expert_parallel_for_initial_contract(): + data = profile() + data["deployments"][0].update(tp=1, dp=4, moe_tp=1, moe_ep=4) + with pytest.raises(ValidationError, match="pure TP"): + FpmModelProfile.model_validate(data) + + +@pytest.mark.parametrize("nextn", [0, 3]) +@pytest.mark.parametrize("entry", ["normalize_config", "best_available"]) +def test_identity_publication_does_not_enable_unimplemented_execution(nextn, entry): + config = { + "model": "test/target", + "system": "b200_sxm", + "backend": "vllm", + "backend_version": "0.28.0", + "worker_type": "aggregated", + "tp": 4, + "nextn": nextn, + "fpm_profile": profile(), + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "estimator_config": {"fpm_interpolation": {"method": "direct"}}, + } + with pytest.raises(ValueError, match="identity-only|does not support MTP"): + getattr(RustForwardPassPerfModel, entry)(config) + + +def test_identity_helper_delegates_without_python_defaults(monkeypatch): + calls = [] + + class Native: + @staticmethod + def _normalize_fpm_execution_profile(revision, payload): + calls.append((revision, json.loads(payload))) + return "native-selector" + + monkeypatch.setattr(aisimulate_core, "RustForwardPassPerfModel", Native) + assert speculative_execution_identity("pinned", {"supplied": True}) == ("", "native-selector", "none", "text") + assert calls == [("pinned", {"supplied": True})] From 4b527d2375f477bdd4a1e7f7e205a67ffb8ed7fc Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 16:44:09 -0700 Subject: [PATCH 02/13] feat: consume measured SD rounds through canonical replay profiles Signed-off-by: hongkuanz --- crates/core/src/perfmodel/engine/runtime.rs | 421 +++++++++++++++++- crates/core/src/perfmodel/fpm/config.rs | 46 +- crates/core/src/perfmodel/fpm/model.rs | 5 +- crates/core/src/perfmodel/fpm/resources.rs | 168 ++++++- python/aisimulate/src/aisimulate/capacity.py | 31 +- .../src/aisimulate_core/sdk/engine.py | 23 +- .../src/aisimulate_core/sdk/fpm_profile.py | 1 + .../aisimulate_core/sdk/memory/kv_cache.py | 15 +- .../tests/unit/sdk/test_fpm_profile.py | 96 ++++ tests/test_fpm_profile_workflow.py | 90 ++++ 10 files changed, 869 insertions(+), 27 deletions(-) diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 06d985fe2..791975e9f 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}; @@ -310,6 +311,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 { @@ -318,6 +322,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() } } @@ -329,6 +334,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") { @@ -350,7 +364,6 @@ impl Engine { .as_ref() .and_then(|s| s.nextn) .unwrap_or(0); - Self::validate_fpm_spec(&spec)?; super::readiness::validate_glm53_sources( &db, spec.context_ops.iter().chain(&spec.generation_ops), @@ -360,10 +373,26 @@ impl Engine { generation_ops: spec.generation_ops, db, nextn, + fpm_execution_profile, }) } - fn validate_fpm_spec(spec: &EngineSpec) -> Result<(), AicError> { + 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()?; + } let nextn = spec .engine .speculative @@ -391,6 +420,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") { @@ -424,6 +454,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()) })?; @@ -446,6 +496,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(()) } @@ -482,7 +606,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( @@ -528,7 +652,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> { @@ -608,6 +732,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, @@ -651,7 +781,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(), )); @@ -715,6 +848,7 @@ impl Engine { stride: u32, mut on_op: impl FnMut(&Op, PerformanceResult), ) -> Result { + self.reject_full_round_surface("static generation trajectories")?; let bs = runtime .batch_size .saturating_mul(self.nextn.saturating_add(1)); @@ -982,6 +1116,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]); } @@ -1305,6 +1440,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); } @@ -1481,6 +1617,7 @@ impl Engine { prefix: u32, prefill: bool, ) -> Result, AicError> { + self.reject_full_round_surface("operation diagnostics")?; use super::diagnostics::{ ExecutedFallback, OperationDetails, SolDiagnostics, StaticOperationDiagnostics, }; @@ -1658,6 +1795,7 @@ impl Engine { mode: StaticMode, stride: u32, ) -> Result<(Vec, Vec), AicError> { + self.reject_full_round_surface("per-operation static phases")?; let mut context = PerOpFold::new("context"); if matches!(mode, StaticMode::Context | StaticMode::Both) { if runtime.prefix >= runtime.isl { @@ -1750,6 +1888,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. @@ -1851,6 +1990,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()); @@ -1884,6 +2024,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(|| { @@ -1918,6 +2059,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(|| { @@ -1956,6 +2098,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}")) })?; @@ -2003,6 +2146,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(), @@ -2057,6 +2201,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}" @@ -2134,6 +2279,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)?; } @@ -2227,7 +2377,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. @@ -3534,6 +3692,257 @@ mod tests { assert!(err.to_string().contains("exactly one FpmForward"), "{err}"); } + 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::{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"), + ..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}" + ); + } + } + /// The marginal-decode mixed composition, exact arithmetic over the /// fixture rows: the prefill component prices the step's SCHEDULED TOTAL /// (ctx + gen tokens) on the prefill curve; decode is the in-curve lerp diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 9510bc5c7..9a5e58c13 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -262,7 +262,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", )); @@ -429,7 +431,11 @@ impl ForwardPassPerfModelConfig { return Err(invalid_config("ngram speculation requires op_level timing")); } } - 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", )); @@ -460,6 +466,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 ngram 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 7cededbc9..5bddba5c8 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -1099,7 +1099,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 e8c732041..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,15 +479,7 @@ impl ForwardPassPerfModelConfig { if matching.next().is_some() { return Err(invalid("duplicate FPM deployment resources")); } - 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)) } } @@ -665,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/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 02e05efa7..299dd0015 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -201,6 +201,26 @@ 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: + speculative_profile_kwargs["nextn"] = lowered.get("aic_nextn", 0) + lowered["num_gpu_blocks"] = estimate_num_gpu_blocks( backend_name=backend, system=lowered.get("aic_system") or _DEFAULT_AIC_SYSTEM, @@ -248,6 +268,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 +377,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 +406,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 +462,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 cbb6cbe6d..d00957ed1 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py @@ -1160,8 +1160,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") @@ -1178,6 +1176,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, @@ -1204,6 +1204,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 @@ -1220,6 +1221,16 @@ def estimate_kv_cache( "moe_tp_size": deployment.moe_tp, "moe_ep_size": deployment.moe_ep, "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 a5a138f69..36f65e050 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,100 @@ 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"}], +) +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) + + 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/tests/test_fpm_profile_workflow.py b/tests/test_fpm_profile_workflow.py index 90ccee1f9..b69499625 100644 --- a/tests/test_fpm_profile_workflow.py +++ b/tests/test_fpm_profile_workflow.py @@ -206,6 +206,96 @@ 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_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 + 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) + 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 + + @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]}}) From 364b6afe3c7cc53b93df5f83bb6276f21b6bb9b8 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 16:49:47 -0700 Subject: [PATCH 03/13] test: qualify measured SD round timing and document replay limits Signed-off-by: hongkuanz --- crates/core/src/perfmodel/engine/runtime.rs | 7 ++++- docs/perf-model/methods/whole-forward.md | 34 +++++++++++++++++++-- docs/replay/engine/speculation.md | 27 ++++++++++++++-- 3 files changed, 62 insertions(+), 6 deletions(-) diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 791975e9f..a72ee8a4d 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -3709,7 +3709,7 @@ mod tests { fn full_round_spec(tmp: &std::path::Path) -> EngineSpec { use crate::operators::fpm_forward::FpmInterpolation; - use crate::perf_database::fpm_forward::tests::{RowSpec, write_pair}; + 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 @@ -3734,6 +3734,11 @@ mod tests { 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(); diff --git a/docs/perf-model/methods/whole-forward.md b/docs/perf-model/methods/whole-forward.md index e22898b8b..56fda108a 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,34 @@ 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, ngram, recorded DCP and non-pure-TP topologies +remain unsupported on this route. + +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. 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 25aa406f8..df7aaff41 100644 --- a/docs/replay/engine/speculation.md +++ b/docs/replay/engine/speculation.md @@ -26,7 +26,7 @@ draft tokens. | `engine.speculation.num_speculative_tokens` | Required | fixed | Integer from 1 to 5. | | `engine.speculation.acceptance_rates` | Required | fixed | Exactly one value 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` keeps these values fixed for every candidate; it does not search @@ -115,7 +115,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, @@ -125,3 +125,26 @@ 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`. + +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. From 057af0f93db945877b19dd42704c895434b483cf Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 16:50:56 -0700 Subject: [PATCH 04/13] fix: scope SD identity guard to the selected deployment Signed-off-by: hongkuanz --- crates/core/src/perfmodel/fpm/config.rs | 19 ------ crates/core/src/perfmodel/fpm/resources.rs | 8 +++ .../perfmodel/perf_database/fpm_forward.rs | 29 ++++++++- .../src/aisimulate_core/fpm_profile.py | 2 +- ...test_fpm_speculative_execution_identity.py | 59 +++++++++++++++---- 5 files changed, 83 insertions(+), 34 deletions(-) diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index bf0fd6f19..9510bc5c7 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -248,25 +248,6 @@ impl ForwardPassPerfModelConfig { pub(crate) fn resolve_with_registration(mut self) -> Result<(Self, Option), AicError> { self.estimator_config.resolve_defaults(); self.validate()?; - // Identity publication precedes the consumer implementation. Do not let - // either nextn=0 or nextn>0 consume a full-round profile as AR timing. - if self - .fpm_profile - .as_ref() - .and_then(|profile| 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(|execution| !execution.is_null()) - }) - }) - { - return Err(invalid_config( - "speculative-round FPM profiles are identity-only; execution is not enabled", - )); - } let registered = if self.fpm_profile.is_some() { self.estimator_config .fpm_interpolation diff --git a/crates/core/src/perfmodel/fpm/resources.rs b/crates/core/src/perfmodel/fpm/resources.rs index e8c732041..e9ef31e12 100644 --- a/crates/core/src/perfmodel/fpm/resources.rs +++ b/crates/core/src/perfmodel/fpm/resources.rs @@ -406,6 +406,14 @@ 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") diff --git a/crates/core/src/perfmodel/perf_database/fpm_forward.rs b/crates/core/src/perfmodel/perf_database/fpm_forward.rs index ebb3fc2ce..8706ee822 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -981,7 +981,7 @@ fn load_pair( || get_int("dp")? != 1 || get_int("cp")? != 1 || get_int("moe_ep")? != 1 - || get_int("moe_tp")? != get_int("tp")? + || (get_int("moe_tp")? != 1 && get_int("moe_tp")? != get_int("tp")?) || match_identity[13] != "False" || match_identity[14] != "False" || match_identity[17] != "none" @@ -1506,6 +1506,7 @@ pub(crate) mod tests { pub backend_version: &'static str, pub partition_policy: &'static str, pub tp: u32, + pub moe_tp: Option, /// Overrides the coordinate-derived cell_id (collision tests). pub cell_id: Option<&'static str>, pub system: &'static str, @@ -1552,6 +1553,7 @@ pub(crate) mod tests { backend_version: "0.25.1", partition_policy: FPM_FORWARD_PARTITION_POLICY, tp: 4, + moe_tp: None, cell_id: None, system: "b200_sxm", backend: "vllm", @@ -1716,7 +1718,13 @@ pub(crate) mod tests { .expect("writer"); let mut rg = writer.next_row_group().expect("row group"); - let identity = |r: &RowSpec| default_identity(r.tp); + let identity = |r: &RowSpec| { + let mut values = default_identity(r.tp); + if let Some(moe_tp) = r.moe_tp { + values[8] = moe_tp.to_string(); + } + values + }; let str_col = |f: &dyn Fn(&RowSpec) -> String| -> Vec { rows.iter() .map(|r| ByteArray::from(f(r).as_str())) @@ -2066,6 +2074,23 @@ pub(crate) mod tests { assert_eq!(loaded_table(tmp.path()).cells().unwrap().len(), 1); } + #[test] + fn speculative_dense_tp_identity_allows_unused_moe_tp_one() { + let tmp = tempfile::tempdir().unwrap(); + let mut rows = speculative_rows(); + for row in &mut rows { + row.moe_tp = Some(1); + } + write_pair(tmp.path(), &rows); + let table = loaded_table(tmp.path()); + let mut identity = default_identity(4); + identity[8] = "1".into(); + identity[16] = rows[0].execution.unwrap()[1].to_owned(); + assert!(table.select_cell(&identity, "org/model-a", None).is_ok()); + identity[8] = "4".into(); + assert!(table.select_cell(&identity, "org/model-a", None).is_err()); + } + #[test] fn speculative_native_policy_rejects_clock_relabeling_and_bad_reductions() { for mutation in [ diff --git a/python/aisimulate/src/aisimulate_core/fpm_profile.py b/python/aisimulate/src/aisimulate_core/fpm_profile.py index b7e941bb0..53e7647eb 100644 --- a/python/aisimulate/src/aisimulate_core/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/fpm_profile.py @@ -238,7 +238,7 @@ def _literal_version(cls, value: str | None) -> str | None: @model_validator(mode="after") def _deployment_identity(self) -> FpmDeploymentProfile: if self.execution_profile is not None and not ( - self.dp == 1 and self.pp == 1 and self.cp == 1 and self.moe_ep == 1 and self.moe_tp == self.tp + self.dp == 1 and self.pp == 1 and self.cp == 1 and self.moe_ep == 1 and self.moe_tp in (1, self.tp) ): raise ValueError("speculative-round FPM identity currently requires vLLM pure TP with DP1/PP1/CP1") modes = { 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 9b32e1c2e..f9a5d4027 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 @@ -11,7 +11,7 @@ import aisimulate_core from aisimulate_core.fpm_profile import FpmModelProfile, FpmSpeculativeExecutionProfile -from aisimulate_core.sdk import RustForwardPassPerfModel +from aisimulate_core.sdk import RustForwardPassPerfModel, fpm_profile from aisimulate_core.sdk.fpm_identity import LEGACY_EXECUTION_IDENTITY, speculative_execution_identity pytestmark = pytest.mark.unit @@ -30,7 +30,7 @@ def execution(): def profile(*, speculative=True): deployment = { - "system": "test_gpu", + "system": "b200_sxm", "backend": "vllm", "backend_version": "0.28.0", "tp": 4, @@ -138,22 +138,57 @@ def test_profile_rejects_attention_dp_or_expert_parallel_for_initial_contract(): FpmModelProfile.model_validate(data) -@pytest.mark.parametrize("nextn", [0, 3]) -@pytest.mark.parametrize("entry", ["normalize_config", "best_available"]) -def test_identity_publication_does_not_enable_unimplemented_execution(nextn, entry): - config = { - "model": "test/target", - "system": "b200_sxm", - "backend": "vllm", - "backend_version": "0.28.0", +def test_dense_tp_profile_accepts_unused_moe_tp_one(): + data = profile() + data["num_experts"] = 0 + data["deployments"][0]["moe_tp"] = 1 + deployment = FpmModelProfile.model_validate(data).deployments[0] + assert deployment.match_identity()[5:11] == ["4", "1", "1", "1", "1", "1"] + assert deployment.execution_identity(data["model_revision"])[1].startswith("speculative_round_v1:") + + +def canonical_request(data, *, nextn=0): + selected = data["deployments"][0] + return { + "model": data["model"], + "system": selected["system"], + "backend": selected["backend"], + "backend_version": selected["backend_version"], "worker_type": "aggregated", - "tp": 4, + "tp": selected["tp"], + "moe_tp_size": selected["moe_tp"], + "moe_ep_size": selected["moe_ep"], "nextn": nextn, - "fpm_profile": profile(), + "fpm_profile": data, "estimation_mode": "fpm_interpolation", "fallback_policy": "deny", "estimator_config": {"fpm_interpolation": {"method": "direct"}}, } + + +def test_unrelated_speculative_deployment_does_not_block_selected_ar(monkeypatch): + monkeypatch.setattr( + fpm_profile, "_load_pinned_checkpoint_config", lambda *_: {"architectures": ["TestDecoderForCausalLM"]} + ) + ar = profile(speculative=False) + before = RustForwardPassPerfModel.normalize_config(canonical_request(ar)) + both = copy.deepcopy(ar) + other = profile()["deployments"][0] + other.update(tp=2, moe_tp=2) + both["deployments"].append(other) + after = RustForwardPassPerfModel.normalize_config(canonical_request(both)) + assert after["nextn"] == before["nextn"] == 0 + assert after["fpm_profile"]["deployments"][0] == before["fpm_profile"]["deployments"][0] + assert after["fpm_profile"]["deployments"][1]["execution_profile"] == execution() + + +@pytest.mark.parametrize("nextn", [0, 3]) +@pytest.mark.parametrize("entry", ["normalize_config", "best_available"]) +def test_identity_publication_does_not_enable_unimplemented_execution(nextn, 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"): getattr(RustForwardPassPerfModel, entry)(config) From 3762a2d8f0f33f118f2af6bf542b812a982aebe9 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 16:51:41 -0700 Subject: [PATCH 05/13] docs: distinguish exact FPM totals from replay prefill approximations Signed-off-by: hongkuanz --- docs/perf-model/methods/whole-forward.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/perf-model/methods/whole-forward.md b/docs/perf-model/methods/whole-forward.md index 56fda108a..9b05dbbfd 100644 --- a/docs/perf-model/methods/whole-forward.md +++ b/docs/perf-model/methods/whole-forward.md @@ -93,7 +93,10 @@ 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. Enable query coverage to identify holes rather than extrapolating. +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 From ba6eb2cdc18f16ca70736f3f1cc3c35140a79797 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 16:53:56 -0700 Subject: [PATCH 06/13] fix: forward measured SD depth through native replay capacity Signed-off-by: hongkuanz --- crates/core/src/python.rs | 42 ++++++++++++++++++- ...test_fpm_speculative_execution_identity.py | 9 ++-- 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 230757a58..d5eff3292 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -1020,8 +1020,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) } @@ -3899,6 +3903,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/python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py b/python/aisimulate/tests/unit/sdk/test_fpm_speculative_execution_identity.py index f9a5d4027..a90464765 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 @@ -182,14 +182,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) From 592c39a0e0544b75ac880bfa53d8ae0524ea262d Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 16:55:46 -0700 Subject: [PATCH 07/13] fix: preserve canonical SD depth through prediction capacity lowering Signed-off-by: hongkuanz --- python/aisimulate/src/aisimulate/capacity.py | 8 +++++- tests/test_fpm_profile_workflow.py | 26 +++++++++++++++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 299dd0015..b6f9d72ec 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -219,7 +219,13 @@ def finish_lowering(value): cp_size=lowered.get("aic_cp_size", 1), ) if deployment.execution_profile is not None: - speculative_profile_kwargs["nextn"] = lowered.get("aic_nextn", 0) + 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, diff --git a/tests/test_fpm_profile_workflow.py b/tests/test_fpm_profile_workflow.py index b69499625..7a1aa75c4 100644 --- a/tests/test_fpm_profile_workflow.py +++ b/tests/test_fpm_profile_workflow.py @@ -212,6 +212,7 @@ def test_speculative_round_prediction_preserves_progress_identity_and_observed_c 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 @@ -279,6 +280,14 @@ def test_speculative_round_prediction_preserves_progress_identity_and_observed_c 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: @@ -286,7 +295,11 @@ 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) + 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. @@ -294,6 +307,17 @@ def run_replay_json(self, serialized): 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"]) From e3e30ffb60fe2ddf4a19c7b24ebb5dc4f662aaf5 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 17:09:16 -0700 Subject: [PATCH 08/13] fix: reject SD-round data in identity-only EngineSpecs Signed-off-by: hongkuanz --- crates/core/src/perfmodel/engine/runtime.rs | 81 +++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 06d985fe2..956e29295 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -364,6 +364,29 @@ impl Engine { } fn validate_fpm_spec(spec: &EngineSpec) -> Result<(), AicError> { + fn contains_sd_round(ops: &[Op]) -> bool { + ops.iter().any(|op| match op { + Op::FpmForward(op) => { + op.match_identity + .get(16) + .is_some_and(|selector| selector.starts_with("speculative_round")) + || contains_sd_round(&op.sol_ops) + } + Op::TokenScale(op) => contains_sd_round(std::slice::from_ref(&op.op)), + Op::Overlap(op) => contains_sd_round(&op.group_a) || contains_sd_round(&op.group_b), + Op::Fallback(op) => { + contains_sd_round(std::slice::from_ref(&op.primary)) + || contains_sd_round(&op.fallback) + } + _ => false, + }) + } + if contains_sd_round(&spec.context_ops) || contains_sd_round(&spec.generation_ops) { + return Err(AicError::InvalidEngineConfig( + "full SD-round FPM tables are identity-only; EngineSpec execution is not enabled" + .into(), + )); + } let nextn = spec .engine .speculative @@ -3534,6 +3557,64 @@ mod tests { assert!(err.to_string().contains("exactly one FpmForward"), "{err}"); } + #[test] + fn identity_only_sd_tables_cannot_be_consumed_by_a_public_ar_spec() { + use crate::operators::fpm_forward::FpmInterpolation; + use crate::perf_database::fpm_forward::tests::{ + NativeCudaSampleSpec, default_rows, write_pair, + }; + let tmp = tempfile::tempdir().unwrap(); + let binding: crate::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": "test/draft", "draft_revision": "2".repeat(40), + "timing_scope": "model_step_cuda" + } + })) + .unwrap(); + let selector: &'static str = Box::leak(binding.selector().unwrap().into_boxed_str()); + let mut spec = build_fpm_spec(tmp.path(), Some(0)); + 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); + let rows: Vec<_> = default_rows() + .into_iter() + .map(|mut row| { + row.execution = Some(["", selector, "none", "text"]); + row.kv_seed_regime = Some("real_kv"); + row.native_cuda = Some(NativeCudaSampleSpec { + scope: "model_step_cuda", + reduction: "single_sample", + source_sample_count: 1, + }); + row + }) + .collect(); + write_pair(tmp.path(), &rows); + for op in spec.context_ops.iter_mut().chain(&mut spec.generation_ops) { + let Op::FpmForward(op) = op else { + unreachable!() + }; + op.match_identity[16] = selector.into(); + op.interpolation = FpmInterpolation::Direct; + } + let result = Engine::from_spec_bytes(&spec.to_bincode().unwrap(), tmp.path()); + if let Ok(engine) = result { + // Independent fixture oracle: this 7ms row is a complete K3 round, + // but the forged public spec declares AR (nextn=0, verify_width=1). + assert_eq!(engine.predict_decode_latency_total(8, 4096).unwrap(), 7.0); + panic!("identity-only PR consumed a measured K3 round through an AR EngineSpec"); + } + let error = result.unwrap_err(); + assert!(matches!(error, AicError::InvalidEngineConfig(_)), "{error}"); + assert!(error.to_string().contains("SD-round"), "{error}"); + spec.engine.fpm_parquet_path = Some("/missing/identity-only-sd.parquet".into()); + let error = Engine::from_spec_bytes(&spec.to_bincode().unwrap(), tmp.path()).unwrap_err(); + assert!(matches!(error, AicError::InvalidEngineConfig(_)), "{error}"); + assert!(error.to_string().contains("SD-round"), "{error}"); + } + /// The marginal-decode mixed composition, exact arithmetic over the /// fixture rows: the prefill component prices the step's SCHEDULED TOTAL /// (ctx + gen tokens) on the prefill curve; decode is the in-curve lerp From 49d6965162d6d745b2dcf1fc9c7677c30bc2dcd6 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 09:33:55 -0700 Subject: [PATCH 09/13] fix: close speculative identity validation gaps Signed-off-by: hongkuanz --- crates/core/src/perfmodel/engine/runtime.rs | 27 +++++++++++++++++++ .../src/aisimulate_core/fpm_profile.py | 7 +++++ ...test_fpm_speculative_execution_identity.py | 6 +++++ 3 files changed, 40 insertions(+) diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index f4a0b7b6d..3033fd9cb 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -409,6 +409,7 @@ impl Engine { || contains_sd_round(&op.sol_ops) } Op::TokenScale(op) => contains_sd_round(std::slice::from_ref(&op.op)), + Op::Dsv41Stage(op) => contains_sd_round(&op.children), Op::Overlap(op) => contains_sd_round(&op.group_a) || contains_sd_round(&op.group_b), Op::Fallback(op) => { contains_sd_round(std::slice::from_ref(&op.primary)) @@ -441,6 +442,7 @@ impl Engine { ops.iter().any(|op| match op { Op::FpmForward(_) => true, Op::TokenScale(o) => contains_fpm(std::slice::from_ref(&o.op)), + Op::Dsv41Stage(o) => contains_fpm(&o.children), Op::Overlap(o) => contains_fpm(&o.group_a) || contains_fpm(&o.group_b), Op::Fallback(o) => { contains_fpm(std::slice::from_ref(&o.primary)) || contains_fpm(&o.fallback) @@ -3880,6 +3882,31 @@ mod tests { assert!(error.to_string().contains("SD-round"), "{error}"); } + #[test] + fn identity_only_sd_selector_in_nested_stage_is_rejected_before_loading() { + let tmp = tempfile::tempdir().unwrap(); + let mut spec = build_fpm_spec(tmp.path(), None); + spec.engine.forward_model = None; + spec.engine.fpm_parquet_path = None; + for ops in [&mut spec.context_ops, &mut spec.generation_ops] { + let Op::FpmForward(op) = &mut ops[0] else { + unreachable!() + }; + op.match_identity[16] = "speculative_round_v1:invalid".into(); + *ops = vec![Op::Dsv41Stage(crate::operators::Dsv41StageOp { + name: "nested_fpm".into(), + is_context: op.phase == FpmPhase::Prefill, + decoder_replay: false, + bounded: false, + window_size: 128, + children: std::mem::take(ops), + })]; + } + let error = Engine::from_spec_bytes(&spec.to_bincode().unwrap(), tmp.path()).unwrap_err(); + assert!(matches!(error, AicError::InvalidEngineConfig(_)), "{error}"); + assert!(error.to_string().contains("SD-round"), "{error}"); + } + /// The marginal-decode mixed composition, exact arithmetic over the /// fixture rows: the prefill component prices the step's SCHEDULED TOTAL /// (ctx + gen tokens) on the prefill curve; decode is the in-curve lerp diff --git a/python/aisimulate/src/aisimulate_core/fpm_profile.py b/python/aisimulate/src/aisimulate_core/fpm_profile.py index 53e7647eb..f35642de5 100644 --- a/python/aisimulate/src/aisimulate_core/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/fpm_profile.py @@ -171,6 +171,13 @@ class FpmSpeculativeExecutionProfile(_ProfileModel): draft_revision: Annotated[str, Field(pattern=r"^[0-9a-f]{40}$")] timing_scope: Literal["model_step_cuda"] + @field_validator("draft_model") + @classmethod + def _draft_model(cls, value: str) -> str: + if value != value.strip() or any(ord(char) < 32 or 127 <= ord(char) <= 159 for char in value): + raise ValueError("draft_model must be nonblank, trimmed, and contain no control characters") + return value + class FpmDeploymentProfile(_ProfileModel): """One hardware, runtime, topology, precision and optional serving-role identity.""" 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 f9a5d4027..f5fccacb8 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 @@ -100,6 +100,12 @@ def test_legacy_profile_json_and_identity_do_not_gain_default_execution_fields() ("kind", "target_verify"), ("timing_scope", "scheduler_wall_time"), ("draft_revision", "main"), + ("draft_model", " test/draft"), + ("draft_model", "test/draft "), + ("draft_model", "test/\x00draft"), + ("draft_model", "test/\x7fdraft"), + ("draft_model", "test/\x85draft"), + ("draft_model", "test/\x9fdraft"), ("extra_field", "ignored"), ], ) From fbcd4271fed8144763ea5285155b56870b7e2b73 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 09:36:15 -0700 Subject: [PATCH 10/13] test: pin generic speculation boundary for measured rounds Signed-off-by: hongkuanz --- crates/core/src/perfmodel/fpm/config.rs | 2 +- docs/perf-model/methods/whole-forward.md | 6 ++++-- docs/replay/engine/speculation.md | 4 +++- python/aisimulate/tests/unit/sdk/test_fpm_profile.py | 9 ++++++++- 4 files changed, 16 insertions(+), 5 deletions(-) diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 281d625be..3e8f7593d 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -522,7 +522,7 @@ impl ForwardPassPerfModelConfig { } 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 ngram speculation", + "speculative-round FPM nextn must match num_draft_tokens and cannot use generic speculation", )); } if self.backend != BackendKind::Vllm diff --git a/docs/perf-model/methods/whole-forward.md b/docs/perf-model/methods/whole-forward.md index 4b303ad75..9028d1599 100644 --- a/docs/perf-model/methods/whole-forward.md +++ b/docs/perf-model/methods/whole-forward.md @@ -84,8 +84,10 @@ 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, ngram, recorded DCP and non-pure-TP topologies -remain unsupported on this route. +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. diff --git a/docs/replay/engine/speculation.md b/docs/replay/engine/speculation.md index 905a5ed28..5ecf83728 100644 --- a/docs/replay/engine/speculation.md +++ b/docs/replay/engine/speculation.md @@ -216,7 +216,9 @@ the [SD execution identity](../../perf-model/methods/whole-forward.md#recorded-s 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`. +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 diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py index e73b7ba17..0e6a21fc2 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py @@ -859,7 +859,14 @@ def test_speculative_round_direct_spec_preserves_identity_and_raw_batch_width(sp @pytest.mark.parametrize( "overrides", - [{"nextn": 0}, {"nextn": 2}, {"estimation_mode": "auto"}, {"fallback_policy": "allow"}], + [ + {"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) From 3dacb8ddc8d43a761acf548d7012d5b0002745fc Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 17:01:33 -0700 Subject: [PATCH 11/13] fix: preserve DCP in measured profile memory queries Signed-off-by: hongkuanz --- .../aisimulate_core/sdk/memory/kv_cache.py | 4 +++ .../tests/unit/sdk/test_fpm_profile.py | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+) 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 0d19e6594..1dec84739 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py @@ -1224,6 +1224,10 @@ 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"), **( { diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py index 0e6a21fc2..3667e2a63 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py @@ -915,6 +915,38 @@ def test_speculative_round_memory_rejects_declared_budget_without_analytical_fal 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" From 6c560ca46e1a8402ba84fa253068bbced5ee3956 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 17:02:23 -0700 Subject: [PATCH 12/13] fix: guard ad-hoc FPM execution identities Signed-off-by: hongkuanz --- crates/core/src/perfmodel/engine/runtime.rs | 15 ++- .../tests/unit/sdk/test_fpm_forward.py | 97 +++++++++++++++++++ 2 files changed, 108 insertions(+), 4 deletions(-) diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 3033fd9cb..4b2e8a95c 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -399,7 +399,7 @@ impl Engine { }) } - fn validate_fpm_spec(spec: &EngineSpec) -> Result<(), AicError> { + fn reject_sd_round_ops(ops: &[Op]) -> Result<(), AicError> { fn contains_sd_round(ops: &[Op]) -> bool { ops.iter().any(|op| match op { Op::FpmForward(op) => { @@ -418,12 +418,17 @@ impl Engine { _ => false, }) } - if contains_sd_round(&spec.context_ops) || contains_sd_round(&spec.generation_ops) { + if contains_sd_round(ops) { return Err(AicError::InvalidEngineConfig( - "full SD-round FPM tables are identity-only; EngineSpec execution is not enabled" - .into(), + "full SD-round FPM requires an explicitly bound whole-round execution path".into(), )); } + 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)?; let nextn = spec .engine .speculative @@ -2047,6 +2052,7 @@ impl Engine { let ops: Vec = serde_json::from_str(ops_json).map_err(|e| { AicError::InvalidEngineConfig(format!("evaluate_ops_json: invalid op list JSON: {e}")) })?; + Self::reject_sd_round_ops(&ops)?; let mut out = PerOpFold::new(if is_context { "context" } else { "generation" }); for op in &ops { let r = if is_context { @@ -2150,6 +2156,7 @@ impl Engine { "evaluate_ops_sol_json: invalid op list JSON: {e}" )) })?; + Self::reject_sd_round_ops(&ops)?; let sol_db = self.db.sol_full_view(); let mut out = PerOpSolFold::default(); for op in &ops { diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 9bd2859a7..4ce798b6a 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -199,6 +199,103 @@ def _model_config(**overrides) -> sdk_config.ModelConfig: return sdk_config.ModelConfig(**defaults) +@pytest.fixture +def fpm_database_probe(tmp_path, request): + """Invented K3 timings with real loader identity, not measured GPU data.""" + from aisimulate_core.sdk.engine import _fpm_spec_dict + from aisimulate_core.sdk.fpm_identity import ( + EXECUTION_COLUMNS, + LEGACY_EXECUTION_IDENTITY, + speculative_execution_identity, + ) + + speculative = getattr(request, "param", True) + execution = speculative_execution_identity( + "1" * 40, + { + "kind": "speculative_round", + "method": "eagle3", + "num_draft_tokens": 3, + "draft_model": "test/draft", + "draft_revision": "2" * 40, + "timing_scope": "model_step_cuda", + }, + ) + if not speculative: + execution = LEGACY_EXECUTION_IDENTITY + shutil.copy(Path(_CORE_SYSTEMS) / f"{SYSTEM}.yaml", tmp_path / f"{SYSTEM}.yaml") + policy = "vllm_native_fpm_model_step_cuda_v1" if speculative else "dynamo_native_single_sample_v1" + rows, ops = [], {} + for phase, tokens, kv, latency in [("prefill", 32, 0, 11.0), ("decode", 0, 128, 7.0)]: + op = FPMForwardOp(phase, _model_config(), MODEL_PATH, sol_ops=[], execution=execution) + identity = dict(zip(_CELL_MATCH_COLUMNS, op._match_identity, strict=True)) + row = _row(phase, 1, tokens, kv, latency if speculative else 2.0, identity=identity) + row.update(dict(zip(EXECUTION_COLUMNS, execution, strict=True))) + if speculative: + row.update( + kv_seed_regime="real_kv", + measurement_policy=policy, + timing_scope="model_step_cuda", + measurement_reduction="single_sample", + source_sample_count=1, + ) + rows.append(row) + ops[phase] = _fpm_spec_dict(op) + ops[phase]["FpmForward"]["interpolation"] = "direct" + _write_pair( + str(tmp_path / "data" / SYSTEM / BACKEND / VERSION), + rows, + sidecar_overrides={ + "schema_version": 7, + "measurement_policy": policy, + **({"timing_scope": "model_step_cuda"} if speculative else {}), + }, + ) + database = PerfDatabase(SYSTEM, BACKEND, VERSION, systems_root=str(tmp_path)) + return EngineHandle.for_database(database, systems_path=str(tmp_path)), ops + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +@pytest.mark.parametrize("entry", ["evaluate_ops_json", "evaluate_ops_sol_json"]) +@pytest.mark.parametrize("shape", ["direct", "nested", "sol_ops"]) +def test_database_probe_rejects_ad_hoc_sd_rounds(fpm_database_probe, phase, entry, shape): + handle, ops = fpm_database_probe + op = ops[phase] + if shape == "nested": + op = {"TokenScale": {"op": op, "numerator": 1, "denominator": 1}} + op = {"Overlap": {"name": "nested_overlap", "group_a": [], "group_b": [op]}} + op = {"Fallback": {"name": "nested_fallback", "primary": op, "fallback": []}} + op = { + "Dsv41Stage": { + "name": "nested_stage", + "is_context": phase == "prefill", + "decoder_replay": False, + "bounded": False, + "window_size": 128, + "children": [op], + } + } + elif shape == "sol_ops": + outer = dict(op["FpmForward"], sol_ops=[op]) + outer["match_identity"] = list(outer["match_identity"]) + outer["match_identity"][16] = "full" + op = {"FpmForward": outer} + with pytest.raises(ValueError, match="SD-round"): + getattr(handle, entry)( + json.dumps([op]), is_context=phase == "prefill", batch_size=1, s=32 if phase == "prefill" else 128 + ) + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +@pytest.mark.parametrize("fpm_database_probe", [False], indirect=True) +def test_database_probe_keeps_ad_hoc_ar_queries(fpm_database_probe, phase): + handle, ops = fpm_database_probe + result = handle.evaluate_ops_json( + json.dumps([ops[phase]]), is_context=phase == "prefill", batch_size=1, s=32 if phase == "prefill" else 128 + ) + assert result[0][1] == 2.0 # Synthetic AR row, not the 7/11 ms full K3 round. + + _NATIVE_FPM_CONTRACTS = { "vllm": ( "vllm_native_fpm_prefix_seed_median_v1", From b7abbd15e9a4a65d6c7d6a2a79ba374f17f67263 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 17:01:24 -0700 Subject: [PATCH 13/13] fix: reject unsupported speculative collection profiles Signed-off-by: hongkuanz --- .../collector/fpm_forward/planner.py | 7 ++ .../collector/test_fpm_profile_collection.py | 75 +++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/python/aisimulate/collector/fpm_forward/planner.py b/python/aisimulate/collector/fpm_forward/planner.py index b348757ea..7c9ac9a72 100644 --- a/python/aisimulate/collector/fpm_forward/planner.py +++ b/python/aisimulate/collector/fpm_forward/planner.py @@ -1043,6 +1043,13 @@ def _validate_profile_identities( cp_size=topology.cp, worker_type=worker_type or "aggregated", ) + if deployment.execution_profile is not None: + raise ValueError( + "FPM collection does not support speculative execution_profile for the selected deployment " + f"{system}/{backend}/{deployment.backend_version}, topology={deployment.parallel_tuple}, " + f"worker_type={worker_type or 'aggregated'}; refusing to collect an autoregressive workload " + "for a speculative profile" + ) deployments.append(deployment) for kv_dtype in capability.dtype.kv_cache_dtypes: for policy in policies: 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 b708a7aff..a173298f0 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_profile_collection.py @@ -573,6 +573,81 @@ def test_profile_cannot_override_resolved_precision_or_backend(field, value, no_ _plan(profile) +def _pure_tp_collection_profile(): + profile = _profile() + profile.update(model="MiniMaxAI/MiniMax-M2.7", architecture="MiniMaxM2ForCausalLM", model_revision="1" * 40) + deployment = profile["deployments"][0] + deployment.update( + system="h200_sxm", + tp=4, + dp=1, + moe_tp=4, + moe_ep=1, + gemm_quant_mode="fp8_block", + moe_quant_mode="fp8_block", + fmha_quant_mode="fp8", + kv_cache_dtype="fp8", + ) + profile["deployments"] = [deployment] + options = replace( + FPMCollectionOptions.from_args(cli._parser().parse_args(_argv(profile))), + max_gpus=4, + gpu_counts=(4,), + parallel_presets=("pure_tp",), + ) + execution = { + "kind": "speculative_round", + "method": "eagle3", + "num_draft_tokens": 3, + "draft_model": "test/draft", + "draft_revision": "2" * 40, + "timing_scope": "model_step_cuda", + } + return profile, options, execution + + +@pytest.mark.parametrize("worker_type", [None, "prefill", "decode", "aggregated"]) +def test_selected_speculative_profile_is_rejected_before_collection_admission( + monkeypatch, no_models_or_timing_data, worker_type +): + profile, options, execution = _pure_tp_collection_profile() + profile["deployments"][0].update(execution_profile=execution, worker_type=worker_type) + options = replace(options, worker_type=worker_type) + monkeypatch.setattr( + planner, + "filter_memory_infeasible_topologies", + lambda **_kwargs: pytest.fail("unsupported speculative collection reached admission"), + ) + with pytest.raises(ValueError, match="collection.*execution_profile"): + _plan( + profile, + options=options, + system="h200_sxm", + selected_ops={"attention_context", "attention_generation"}, + ) + + +@pytest.mark.parametrize("unselected_by", ["system", "worker_type"]) +def test_unselected_speculative_deployment_does_not_block_ar_collection(no_models_or_timing_data, unselected_by): + profile, options, execution = _pure_tp_collection_profile() + selected = profile["deployments"][0] + selected["worker_type"] = "prefill" + options = replace(options, worker_type="prefill") + unselected = copy.deepcopy(selected) + unselected["execution_profile"] = execution + unselected[unselected_by] = "b200_sxm" if unselected_by == "system" else "decode" + profile["deployments"].append(unselected) + plan = _plan( + profile, + options=options, + system="h200_sxm", + selected_ops={"attention_context", "attention_generation"}, + ) + assert plan.cells + assert all(plan.deployment_profile(cell).execution_profile is None for cell in plan.cells) + assert all(cell.execution_identity == ("", "full", "none", "text") for cell in plan.cells) + + @pytest.mark.parametrize("kv_dtype", ["bfloat16", "fp8"]) def test_minimax_selected_kv_precision_reaches_plan_and_worker(tmp_path, no_models_or_timing_data, kv_dtype): profile = _profile()