diff --git a/README.md b/README.md index f820cc822f..2e43290be5 100644 --- a/README.md +++ b/README.md @@ -265,7 +265,7 @@ APIs: - [Estimator/FPE Python and Rust SDK](docs/core-api.md) - [AIC-compatible modeled-power contract (semantics only)](docs/power-model.md) -- [FPM collection-to-prediction workflow](python/aisimulate/docs/fpm/end-to-end-workflow.md) +- [Self-benchmarking and FPM onboarding guide](python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md) - [Replay SDK and artifact contract](crates/core/src/replay/README.md) - [Sweeper SDK](docs/sweeper/overview.md) - [Legacy CLI reference](docs/cli/legacy-aic-user-guide.md) diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 05f0156784..43029ffd16 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -49,7 +49,7 @@ pub use perfmodel::{ }; pub use perfmodel::{ CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, OpLevelConfig, - RegressionFitConfig, RegressionFitKind, + RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, }; #[cfg(feature = "python")] diff --git a/crates/core/src/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index 9976d27c5d..c6bc96317f 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -114,7 +114,12 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // - 22 (GLM-5.2 VR200 pilot): exact observed-MoE selection, prefill graph // identity and two appended composite operators extend the schema-21 layout. // The pilot and AIC-1781 concurrently claimed 21; reject both older layouts. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 22; +// - 23 (DCP identity): ParallelMapping gained optional recorded dcp_size. +// DCP and the pilot concurrently claimed 22; reject both older layouts. +// JSON defaults preserve unrecorded DCP. +// - 24 (typed FPM DCP): FpmForwardOp carries dcp_size separately from the +// base matching identity, so control flow never parses the string tuple. +pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 24; /// Static engine identity and setup information carried by an /// [`crate::perfmodel::engine::spec::EngineSpec`]. @@ -261,6 +266,8 @@ pub struct ParallelMapping { /// re-derived from this field. #[serde(default)] pub cp_size: Option, + #[serde(default)] + pub dcp_size: Option, } /// Precision/quantization dtypes. Flattened into [`EngineConfig`]. Field diff --git a/crates/core/src/perfmodel/engine/readiness.rs b/crates/core/src/perfmodel/engine/readiness.rs index 0c64701f08..c555109e90 100644 --- a/crates/core/src/perfmodel/engine/readiness.rs +++ b/crates/core/src/perfmodel/engine/readiness.rs @@ -163,9 +163,11 @@ impl Availability<'_> { } TokenScale(scale) => return self.op(&scale.op), FpmForward(fpm) => { - self.db - .fpm_forward - .select_cell(&fpm.match_identity, &fpm.model_path)?; + self.db.fpm_forward.select_cell( + &fpm.match_identity, + &fpm.model_path, + fpm.dcp_size, + )?; return Ok(()); } // These families have no analytic/empirical implementation. diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index e09f2a18ce..4206a65a63 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -2485,6 +2485,7 @@ mod tests { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { + dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(1), @@ -2555,6 +2556,7 @@ mod tests { config.backend = BackendKind::Sglang; config.backend_version = Some("0.5.18+nvinternal.rubin.0.8full.66997102".into()); config.parallel = ParallelMapping { + dcp_size: None, tp_size: 4, pp_size: 1, attention_dp_size: Some(1), @@ -3530,6 +3532,7 @@ mod tests { config.systems_path = Some(tmp.to_path_buf()); let fpm_op = |phase: FpmPhase| { Op::FpmForward(FpmForwardOp { + dcp_size: None, name: format!("fpm_forward_{}", phase.as_str()), phase, model_path: "org/model-a".to_string(), @@ -3601,6 +3604,7 @@ mod tests { use crate::perf_database::fpm_forward::tests::default_identity; let db = PerfDatabase::load(&systems_root(), "b200_sxm", "vllm", "0.24.0").unwrap(); let fpm_op = Op::FpmForward(FpmForwardOp { + dcp_size: None, name: "fpm_forward_prefill".into(), phase: FpmPhase::Prefill, model_path: "org/model-a".into(), @@ -3659,6 +3663,7 @@ mod tests { )); let fpm_op = |phase: FpmPhase| { Op::FpmForward(FpmForwardOp { + dcp_size: None, name: format!("fpm_forward_{}", phase.as_str()), phase, model_path: "org/model-a".to_string(), @@ -4204,6 +4209,7 @@ mod tests { let hidden = Op::Overlap(crate::operators::OverlapOp::new( "hidden", vec![Op::FpmForward(FpmForwardOp { + dcp_size: None, name: "fpm_forward_prefill".into(), phase: FpmPhase::Prefill, model_path: "org/model-a".into(), @@ -4522,6 +4528,7 @@ mod tests { "0.25.1", )); let mut op = FpmForwardOp { + dcp_size: None, name: "fpm_forward_decode".into(), phase: FpmPhase::Decode, model_path: "org/model-a".into(), @@ -4568,6 +4575,7 @@ mod tests { )); let fpm_op = |phase: FpmPhase, width: u32| { Op::FpmForward(FpmForwardOp { + dcp_size: None, name: format!("fpm_forward_{}", phase.as_str()), phase, model_path: "org/model-a".to_string(), diff --git a/crates/core/src/perfmodel/engine/spec.rs b/crates/core/src/perfmodel/engine/spec.rs index b3236794cd..4ef523675e 100644 --- a/crates/core/src/perfmodel/engine/spec.rs +++ b/crates/core/src/perfmodel/engine/spec.rs @@ -625,6 +625,7 @@ mod tests { // Recursive like Overlap/Fallback: sol_ops carries the model's // original granular list, so the round-trip must preserve nesting. crate::operators::FpmForwardOp { + dcp_size: None, name: "fpm_forward_prefill".into(), phase: crate::operators::FpmPhase::Prefill, model_path: "org/model-a".into(), @@ -850,6 +851,7 @@ mod tests { moe_kernel_source: None, kv_block_size: Some(64), parallel: ParallelMapping { + dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(8), @@ -1045,7 +1047,7 @@ mod tests { vec![OpSpec::FpmForward(fpm_forward())], vec![OpSpec::Moe(moe())], ); - assert_eq!(spec.schema_version, 22); + assert_eq!(spec.schema_version, 24); let mut bytes = spec.to_bincode().unwrap(); assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); @@ -1350,8 +1352,11 @@ mod tests { fn merged_schema_preserves_fpm_selectors_and_pilot_payloads() { let mut fpm_config = sample_engine_config(); fpm_config.forward_model = Some("fpm".into()); + fpm_config.parallel.dcp_size = Some(2); fpm_config.quantization.fpm_fmha_dtype = Some(DataType::Fp8); - let fpm_spec = EngineSpec::new(fpm_config, vec![OpSpec::FpmForward(fpm_forward())], vec![]); + let mut fpm_op = fpm_forward(); + fpm_op.dcp_size = Some(2); + let fpm_spec = EngineSpec::new(fpm_config, vec![OpSpec::FpmForward(fpm_op)], vec![]); let mut pilot_config = sample_engine_config(); pilot_config.prefill_graph_profile = @@ -1375,18 +1380,18 @@ mod tests { for spec in [fpm_spec, pilot_spec] { let bytes = spec.to_bincode().unwrap(); assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); - for previous_version in [20u32, 21] { + for previous_version in [20u32, 21, 22, 23] { let mut stale = bytes.clone(); stale[..4].copy_from_slice(&previous_version.to_le_bytes()); - // Both branches claimed 21 after distinct schema-20 layouts. - // Reject either version before decoding even a missing payload. + // DCP and the pilot claimed 22 for different layouts. + // Reject all older versions before decoding even a missing payload. for input in [stale.as_slice(), &stale[..4]] { assert!(matches!( EngineSpec::from_bincode(input), Err(AicError::UnsupportedSchemaVersion { kind: "EngineSpec", got, - expected: 22, + expected: 24, }) if got == previous_version )); } @@ -1412,9 +1417,25 @@ mod tests { Err(AicError::UnsupportedSchemaVersion { got: 19, expected: ENGINE_SPEC_SCHEMA_VERSION, - .. }) )); } + #[test] + fn recorded_dcp_and_legacy_absence_round_trip() { + let mut spec = handshake_spec(); + spec.engine.parallel.dcp_size = Some(2); + let decoded = EngineSpec::from_bincode(&spec.to_bincode().unwrap()).unwrap(); + assert_eq!(decoded.engine.parallel.dcp_size, Some(2)); + let mut legacy = serde_json::to_value(&spec.engine).unwrap(); + legacy.as_object_mut().unwrap().remove("dcp_size"); + let legacy: EngineConfig = serde_json::from_value(legacy).unwrap(); + assert_eq!(legacy.parallel.dcp_size, None); + let mut previous = spec.to_bincode().unwrap(); + previous[..4].copy_from_slice(&23u32.to_le_bytes()); + assert!(matches!( + EngineSpec::from_bincode(&previous), + Err(AicError::UnsupportedSchemaVersion { got: 23, .. }) + )); + } } diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index d7e79d2986..457fc9238b 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -94,6 +94,9 @@ pub struct ForwardPassPerfModelConfig { pub pp: u32, #[serde(default = "one", alias = "attention_dp_size")] pub attention_dp: u32, + /// Recorded decode context parallelism within the TP group. + #[serde(default)] + pub dcp: Option, #[serde(default)] pub moe_tp_size: Option, #[serde(default)] @@ -171,6 +174,7 @@ impl ForwardPassPerfModelConfig { tp: 1, pp: 1, attention_dp: 1, + dcp: None, moe_tp_size: None, moe_ep_size: None, gemm_quant_mode: None, @@ -277,6 +281,17 @@ impl ForwardPassPerfModelConfig { if self.tp == 0 || self.pp == 0 || self.attention_dp == 0 { return Err(invalid_config("tp, pp, and attention_dp must be positive")); } + if self.dcp.is_some_and(|dcp| dcp == 0 || self.tp % dcp != 0) { + return Err(invalid_config("dcp must be positive and divide tp")); + } + self.estimator_config + .fpm_interpolation + .validate_quant_modes( + self.fmha_quant_mode + .as_deref() + .or(self.fpm_fmha_quant_mode.as_deref()), + self.comm_quant_mode.as_deref(), + )?; if self.moe_tp_size.is_some() != self.moe_ep_size.is_some() { return Err(invalid_config( "moe_tp_size and moe_ep_size must be configured together", @@ -430,7 +445,7 @@ impl ForwardPassPerfModelConfig { } } -fn invalid_config(message: impl Into) -> AicError { +pub(super) fn invalid_config(message: impl Into) -> AicError { AicError::InvalidEngineConfig(format!( "invalid forward pass perf model config: {}", message.into() @@ -661,6 +676,32 @@ mod tests { .is_err() ); } + #[test] + fn recorded_dcp_is_validated_and_survives_round_trip() { + let cfg = + config(serde_json::json!({"tp": 8, "dcp": 8, "moe_tp_size": 8, "moe_ep_size": 1})); + cfg.validate().unwrap(); + let restored: ForwardPassPerfModelConfig = + serde_json::from_str(&serde_json::to_string(&cfg).unwrap()).unwrap(); + assert_eq!(restored.dcp, Some(8)); + assert_eq!(config(serde_json::json!({})).dcp, None); + for value in [0, 3, 16] { + assert!( + config(serde_json::json!({"tp": 8, "dcp": value})) + .validate() + .is_err() + ); + } + } + + #[test] + fn unrecorded_profile_quantization_rejects_explicit_overrides() { + let mut cfg = config(serde_json::json!({"estimation_mode": "fpm_interpolation", + "estimator_config": {"fpm_interpolation": {"text_only": true, "unrecorded_quant_modes": ["fmha", "comm"]}}})); + cfg.validate().unwrap(); + cfg.fmha_quant_mode = Some("bfloat16".into()); + assert!(cfg.validate().is_err()); + } #[test] fn graph_profile_pins_identity_and_rejects_unqualified_selection() { diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs index 787d324f6c..ed0482056a 100644 --- a/crates/core/src/perfmodel/fpm/estimator.rs +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -37,9 +37,47 @@ pub struct OpLevelConfig { #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct FpmInterpolationConfig { + /// The profile covers text prefill/decode; encoder weights remain resident. + #[serde(skip_serializing_if = "is_false")] + pub text_only: bool, /// External parquet and its same-stem metadata sidecar. #[serde(skip_serializing_if = "Option::is_none")] pub fpm_parquet_path: Option, + /// Match null profile identities only for these unspecified quant modes. + #[serde(skip_serializing_if = "Vec::is_empty")] + pub unrecorded_quant_modes: Vec, +} + +fn is_false(value: &bool) -> bool { + !*value +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum UnrecordedFpmQuantMode { + Fmha, + Comm, +} + +impl FpmInterpolationConfig { + pub(crate) fn validate_quant_modes( + &self, + fmha: Option<&str>, + comm: Option<&str>, + ) -> Result<(), AicError> { + for mode in &self.unrecorded_quant_modes { + let explicit = match mode { + UnrecordedFpmQuantMode::Fmha => fmha, + UnrecordedFpmQuantMode::Comm => comm, + }; + if explicit.is_some() { + return Err(super::config::invalid_config( + "an unrecorded FPM quant mode cannot have an explicit quantization override", + )); + } + } + Ok(()) + } } /// These weights currently affect regression. Native correction keeps its diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index 256d0d1ee6..4e524df1f3 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -328,6 +328,17 @@ impl ForwardPassPerfModel { let mut failures = Vec::new(); let mut last_error = None; for mode in config.candidate_modes() { + if config.dcp.is_some_and(|dcp| dcp > 1) + && (mode != EstimationMode::FpmInterpolation + || config.backend != crate::BackendKind::Vllm) + { + let error = AicError::UnsupportedModel( + "DCP timing requires measured vLLM FPM interpolation".into(), + ); + failures.push(format!("{mode:?}: {error}")); + last_error = Some(error); + continue; + } if config.speculation.is_some() && mode != EstimationMode::OpLevel { let error = AicError::UnsupportedModel("ngram speculation requires op_level timing".into()); diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index 6d8451231b..e03562ec3f 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -33,6 +33,26 @@ fn systems_root() -> PathBuf { const TEST_MODEL: &str = "MiniMaxAI/MiniMax-M2.5"; +#[test] +fn dcp_cannot_construct_explicit_regression() { + let mut config = crate::ForwardPassPerfModelConfig::new( + TEST_MODEL, + "b200_sxm", + BackendKind::Vllm, + ForwardPassWorkerType::Aggregated, + ); + config.tp = 8; + config.dcp = Some(8); + config.estimation_mode = crate::EstimationMode::FpmRegression; + let error = ForwardPassPerfModel::best_available(config).unwrap_err(); + assert!(matches!(error, AicError::UnsupportedModel(_))); + assert!( + error + .to_string() + .contains("measured vLLM FPM interpolation") + ); +} + /// Hand-built context op list against the b200_sxm/vllm/0.24.0 perf tables /// (same fixture pattern as `engine/runtime.rs` and `py.rs` tests). fn context_ops() -> Vec { @@ -114,6 +134,7 @@ fn fixture_engine_config() -> EngineConfig { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { + dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(1), diff --git a/crates/core/src/perfmodel/memory.rs b/crates/core/src/perfmodel/memory.rs index ba7f9bf73f..66d338a85a 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -515,6 +515,7 @@ mod tests { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { + dcp_size: None, tp_size: 1, pp_size: 1, attention_dp_size: Some(1), diff --git a/crates/core/src/perfmodel/mod.rs b/crates/core/src/perfmodel/mod.rs index 30490f2c92..05870ef399 100644 --- a/crates/core/src/perfmodel/mod.rs +++ b/crates/core/src/perfmodel/mod.rs @@ -55,7 +55,7 @@ pub use fpm::{ }; pub use fpm::{ CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, OpLevelConfig, - RegressionFitConfig, RegressionFitKind, + RegressionFitConfig, RegressionFitKind, UnrecordedFpmQuantMode, }; // Forward-pass metrics telemetry types and schema version, plus the // crate-internal validation helper. Re-exported at the crate root so existing diff --git a/crates/core/src/perfmodel/operators/attention.rs b/crates/core/src/perfmodel/operators/attention.rs index fbd1329391..4594edfbc1 100644 --- a/crates/core/src/perfmodel/operators/attention.rs +++ b/crates/core/src/perfmodel/operators/attention.rs @@ -184,7 +184,7 @@ pub struct ContextAttentionOp { /// lanes, density-ranked donor tiers, `"default"`, and the table's own /// leftover lanes — and it is REPLAYED VERBATIM here: no re-deriving, no /// extending, no sorting. Appended at the struct TAIL because bincode - /// payloads are positional (current ENGINE_SPEC_SCHEMA_VERSION 22). + /// payloads are positional (current ENGINE_SPEC_SCHEMA_VERSION 24). #[serde(default = "default_lane_order")] pub lane_order: Vec, /// Whether the fused prefill kernel applies rotary position embeddings. @@ -197,7 +197,7 @@ pub struct ContextAttentionOp { /// Lane precedence for ops built without an explicit order (Rust-side /// constructors and hand-written JSON fixtures predating the `lane_order` -/// field — introduced at schema v8, current ENGINE_SPEC_SCHEMA_VERSION 22). +/// field — introduced at schema v8, current ENGINE_SPEC_SCHEMA_VERSION 24). /// Mirrors the Python fallback in `_attention_lane_order` for an /// unresolvable database: the always-valid `("default",)`. pub(crate) fn default_lane_order() -> Vec { @@ -385,7 +385,7 @@ pub struct GenerationAttentionOp { pub kv_cache_dtype: KvCacheQuantMode, /// Kernel-source lane precedence; see /// [`ContextAttentionOp::lane_order`] (appended at the struct TAIL — - /// bincode payloads are positional, current ENGINE_SPEC_SCHEMA_VERSION 22). + /// bincode payloads are positional, current ENGINE_SPEC_SCHEMA_VERSION 24). #[serde(default = "default_lane_order")] pub lane_order: Vec, /// Per-head RMSNorm on Q and K before decode attention. Appended at the diff --git a/crates/core/src/perfmodel/operators/fpm_forward.rs b/crates/core/src/perfmodel/operators/fpm_forward.rs index 42885a8774..12061213a6 100644 --- a/crates/core/src/perfmodel/operators/fpm_forward.rs +++ b/crates/core/src/perfmodel/operators/fpm_forward.rs @@ -59,11 +59,12 @@ impl FpmPhase { /// One whole-model forward pass for a single phase. /// -/// `match_identity` is the 19-string cell identity (legacy schema-6 requests use 15) in +/// `match_identity` is the 19-string base cell identity (legacy schema-6 requests use 15) in /// [`FPM_CELL_MATCH_COLUMNS`](crate::perf_database::fpm_forward::FPM_CELL_MATCH_COLUMNS) /// order, computed by the PYTHON producer via `_norm_identity` (None -> "", /// Enum -> `.name`) so Rust compares strings verbatim with no re-normalization -/// drift. `sol_ops` is the model's original op-level list for this phase — +/// drift. Recorded DCP is carried separately in `dcp_size`. `sol_ops` is the +/// model's original op-level list for this phase — /// the roofline source, serialized recursively like `Overlap`/`Fallback` /// children. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -95,6 +96,9 @@ pub struct FpmForwardOp { /// before match_identity[2] was selected; it is NOT observed runtime precision. #[serde(default)] pub original_fmha_quant_mode: Option, + /// Recorded DCP selects the measured cell and gates unmodeled SOL transfer. + #[serde(default)] + pub dcp_size: Option, } fn default_verify_width() -> u32 { @@ -106,12 +110,21 @@ fn data_err(msg: String) -> AicError { } impl FpmForwardOp { + fn sol_at(&self, db: &PerfDatabase, coords: &[f64]) -> Result { + if self.dcp_size.is_some_and(|dcp| dcp > 1) { + return Err(AicError::UnsupportedModel( + "DCP FPM supports measured interpolation, not op-level SOL transfer".into(), + )); + } + sol_total(&self.sol_ops, self.phase, db, coords) + } + fn select_cell<'a>(&self, db: &'a PerfDatabase) -> Result<&'a FpmForwardCell, AicError> { // Exact matching remains authoritative, including the recorded FMHA label. // A different precision cell is a miss; the selector never rewrites it. - let cell = db - .fpm_forward - .select_cell(&self.match_identity, &self.model_path)?; + let cell = + db.fpm_forward + .select_cell(&self.match_identity, &self.model_path, self.dcp_size)?; if let Some(original) = &self.original_fmha_quant_mode { if let Some(warning) = cell.fmha_selector_warning(original) { // The perfmodel has no installed logging facade. Emit a visible, @@ -353,8 +366,8 @@ impl FpmForwardOp { let candidate: Vec = std::iter::once(max as f64) .chain(coords[1..].iter().copied()) .collect(); - let true_sol = sol_total(&self.sol_ops, self.phase, db, coords); - let ceiling_sol = sol_total(&self.sol_ops, self.phase, db, &candidate); + let true_sol = self.sol_at(db, coords); + let ceiling_sol = self.sol_at(db, &candidate); if let (Ok(t), Ok(c)) = (true_sol, ceiling_sol) { if t.is_finite() && c.is_finite() && t > 0.0 && c > 0.0 { // True shape is never costlier than the clamped @@ -415,7 +428,7 @@ impl FpmForwardOp { // and the error names the op. let sol_failure: std::cell::RefCell> = std::cell::RefCell::new(None); let sol = |sol_coords: &[f64]| -> f64 { - match sol_total(&self.sol_ops, self.phase, db, sol_coords) { + match self.sol_at(db, sol_coords) { Ok(v) => v, Err(err) => { let mut slot = sol_failure.borrow_mut(); @@ -666,6 +679,7 @@ mod tests { fn op(phase: FpmPhase) -> FpmForwardOp { FpmForwardOp { + dcp_size: None, name: format!("fpm_forward_{}", phase.as_str()), phase, model_path: "org/model-a".to_string(), @@ -678,6 +692,17 @@ mod tests { } } + #[test] + fn recorded_dcp_disables_unmodeled_sol_transfer() { + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &default_rows()); + let db = db_with_pair(tmp.path()); + let mut measured = op(FpmPhase::Prefill); + measured.dcp_size = Some(4); + let error = measured.sol_at(&db, &[1.0, 1024.0, 0.0]).unwrap_err(); + assert!(error.to_string().contains("not op-level SOL transfer")); + } + fn ctx(batch_size: u32, s: u32, prefix: u32) -> RuntimeContext { RuntimeContext { batch_size, @@ -1451,7 +1476,11 @@ mod tests { selected.original_fmha_quant_mode = Some("fp8".into()); let cell = db .fpm_forward - .select_cell(&selected.match_identity, &selected.model_path) + .select_cell( + &selected.match_identity, + &selected.model_path, + selected.dcp_size, + ) .unwrap(); let message = cell.fmha_selector_warning("fp8").unwrap(); assert!(message.contains("original_model_mode=\"fp8\"")); diff --git a/crates/core/src/perfmodel/perf_database/fpm_forward.rs b/crates/core/src/perfmodel/perf_database/fpm_forward.rs index 91feb01c9d..2dbc01ceb3 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -55,9 +55,9 @@ pub const FPM_FORWARD_SCHEMA_NAME: &str = "aic_fpm_forward_perf"; pub const FPM_FORWARD_SCHEMA_VERSION: u64 = 7; pub const FPM_FORWARD_COORDINATE_SYSTEM: &str = "iteration_totals_balanced_v1"; pub const FPM_FORWARD_PARTITION_POLICY: &str = "balanced_v1"; -/// The only measurement policy the collector publishes; pinned in the -/// sidecar gate (a pair measured under a different regime is structural). +/// Legacy collector policy; self-benchmark repeats use a per-row contract. pub const FPM_FORWARD_MEASUREMENT_POLICY: &str = "dynamo_native_single_sample_v1"; +const FPM_ROW_MEASUREMENT_POLICY: &str = "per_row_single_sample_or_median_of_3"; /// `kv_seed_regime` marker for rows whose KV state the collector could not /// reach through the real kvwarm chain and fabricated instead. Such rows /// are measurements of the wrong regime (observed 2-3.7x inflated on 4-GPU @@ -86,6 +86,7 @@ pub const FPM_FAKE_FALLBACK_RAW_ENV: &str = "AIC_FPM_FAKE_FALLBACK_RAW"; /// Identity columns that select a cell, in row-column order (`model_path` is /// handled separately; `weight_quantization` is deliberately excluded). +/// Recorded DCP is appended separately; legacy identities retain this base arity. /// The last four are the schema-v6 explicit backend identity: "auto" = the /// engine decided; the `enable_*` columns are real parquet booleans, /// normalized to "True"/"False" (Python `str(bool)`) for comparison. @@ -328,14 +329,21 @@ impl FpmForwardTable { Ok(self.loaded()?.kept_fake_fallback) } - /// Cell selection, mirroring Python `FPMForwardOp._select_cell` exactly: - /// strict equality on the 15-column identity and `model_path` (no - /// fallback), then a hard ambiguity guard. + /// Exact base identity, recorded DCP and model-path matching, followed by + /// an ambiguity guard. DCP remains typed until composing the lookup key. pub fn select_cell( &self, match_identity: &[String], model_path: &str, + dcp_size: Option, ) -> Result<&FpmForwardCell, AicError> { + if match_identity.len() != 15 && match_identity.len() != FPM_CELL_MATCH_COLUMNS.len() { + return Err(structural( + "FPM match_identity must contain only base identity fields; supply DCP through dcp_size".into(), + )); + } + let dcp_label = dcp_size.map(|value| value.to_string()); + let requested = || match_identity.iter().chain(dcp_label.iter()); let cells = self.cells()?; // Exact matching on every identity dimension (D1 resolved: the match // identity carries no architecture fingerprint, so borrowing the sole @@ -343,13 +351,15 @@ impl FpmForwardTable { let matches: Vec<&FpmForwardCell> = cells .iter() .filter(|cell| { - let legacy = match_identity.len() == 15 + let legacy = dcp_size.is_none() + && match_identity.len() == 15 && cell.match_identity[..15] == *match_identity && cell.match_identity[15..] .iter() .map(String::as_str) .eq(LEGACY_EXECUTION_IDENTITY); - (cell.match_identity == match_identity || legacy) && cell.model_path == model_path + (cell.match_identity.iter().eq(requested()) || legacy) + && cell.model_path == model_path }) .collect(); if matches.is_empty() { @@ -362,7 +372,8 @@ impl FpmForwardTable { available.truncate(8); let identity: Vec = FPM_CELL_MATCH_COLUMNS .iter() - .zip(match_identity) + .chain(std::iter::once(&"dcp")) + .zip(requested()) .map(|(c, v)| format!("{c}={v:?}")) .collect(); return Err(structural(format!( @@ -404,15 +415,19 @@ fn sha256_file(path: &Path) -> Result { Ok(format!("{:x}", hasher.finalize())) } -/// Sidecar validation, mirroring `_validate_sidecar` check-for-check. Returns -/// the sidecar's `row_count` for the post-read cross-check. +struct SidecarContract { + schema_version: u64, + row_count: Option, + per_row_policy: bool, + selector_dcp: Option>, +} fn validate_sidecar( metadata_path: &Path, parquet_path: &Path, system: &str, backend: &str, version: &str, -) -> Result<(Option, u64), AicError> { +) -> Result { if !metadata_path.exists() { return Err(structural(format!( "FPM database is missing its metadata sidecar: {}. \ @@ -460,11 +475,13 @@ fn validate_sidecar( metadata_path.display() ))); } - if metadata.get("measurement_policy").and_then(|v| v.as_str()) - != Some(FPM_FORWARD_MEASUREMENT_POLICY) - { + let measurement_policy = metadata.get("measurement_policy").and_then(|v| v.as_str()); + if !matches!( + measurement_policy, + Some(FPM_FORWARD_MEASUREMENT_POLICY | FPM_ROW_MEASUREMENT_POLICY) + ) { return Err(structural(format!( - "unsupported FPM measurement_policy={:?} (expected {FPM_FORWARD_MEASUREMENT_POLICY:?}): {}", + "unsupported FPM measurement_policy={:?} (expected {FPM_FORWARD_MEASUREMENT_POLICY:?} or {FPM_ROW_MEASUREMENT_POLICY:?}): {}", metadata.get("measurement_policy"), metadata_path.display() ))); @@ -494,10 +511,32 @@ fn validate_sidecar( parquet_path.parent().unwrap_or(parquet_path).display() ))); } - Ok(( - json_uint(metadata.get("row_count")), - schema_version.unwrap(), - )) + let selector_dcp = metadata + .get("configuration_selector") + .map(|selector| { + let selector = selector + .as_object() + .ok_or_else(|| structural("FPM configuration_selector must be an object".into()))?; + selector + .get("dcp") + .map(|value| { + value + .as_u64() + .and_then(|v| u32::try_from(v).ok()) + .filter(|v| *v > 0) + .ok_or_else(|| { + structural("FPM selector dcp must be a positive integer".into()) + }) + }) + .transpose() + }) + .transpose()?; + Ok(SidecarContract { + schema_version: schema_version.unwrap(), + row_count: json_uint(metadata.get("row_count")), + per_row_policy: measurement_policy == Some(FPM_ROW_MEASUREMENT_POLICY), + selector_dcp, + }) } /// Python `metadata.get(k) != n` compares by VALUE: a JSON `5.0` equals the @@ -555,8 +594,7 @@ fn load_pair( return Ok(None); } let metadata_path = parquet_path.with_extension("metadata.json"); - let (sidecar_row_count, schema_version) = - validate_sidecar(&metadata_path, parquet_path, system, backend, version)?; + let sidecar = validate_sidecar(&metadata_path, parquet_path, system, backend, version)?; let reader = PerfReader::open(parquet_path)?; // Physical row-key columns (collector contract), in order. @@ -581,7 +619,7 @@ fn load_pair( for name in str_cols { str_idx.insert(name, reader.col(name)?); } - if schema_version >= 7 { + if sidecar.schema_version >= 7 { for name in &FPM_CELL_MATCH_COLUMNS[15..] { str_idx.insert(*name, reader.col(name)?); } @@ -610,13 +648,19 @@ fn load_pair( // KV-seed provenance column (additive; absent in pairs that predate it). // Resolved once here; per-row reads treat null the same as absence. let kv_seed_col = reader.col_optional("kv_seed_regime"); + let dcp_col = reader.col_optional("dcp"); + let row_policy_col = reader.col_optional("measurement_policy"); + let repeats_col = reader.col_optional("measurement_repeats"); + if sidecar.per_row_policy && (row_policy_col.is_none() || repeats_col.is_none()) { + return Err(structural("FPM per-row measurement policy requires measurement_policy and measurement_repeats columns".into())); + } // Python checks the sidecar row_count against the FULL row list before any // per-row validation (`load_fpm_forward_data`: read_table -> row_count -> // empty -> per-row loop); mirror that error precedence with a cheap count // pass so a wrong-count pair reports the count, not the first bad row. let actual_row_count = reader.rows()?.count() as u64; - match sidecar_row_count { + match sidecar.row_count { Some(expected) if expected == actual_row_count => {} other => { return Err(structural(format!( @@ -735,7 +779,7 @@ fn load_pair( ))); } - let match_identity: Vec = FPM_CELL_MATCH_COLUMNS + let mut match_identity: Vec = FPM_CELL_MATCH_COLUMNS .iter() .map(|name| { if str_idx.contains_key(name) { @@ -752,6 +796,34 @@ fn load_pair( } }) .collect::>()?; + let dcp = dcp_col.map(|col| row.u32(col)).transpose()?; + if let Some(dcp) = dcp { + let tp = get_int("tp")?; + if dcp == 0 || tp == 0 || tp % dcp != 0 { + return Err(structural(format!( + "FPM row {index} dcp must be positive and divide tp" + ))); + } + match_identity.push(dcp.to_string()); + } + if sidecar.selector_dcp.is_some_and(|declared| declared != dcp) { + return Err(structural(format!( + "FPM row {index} dcp disagrees with the sidecar selector" + ))); + } + if sidecar.per_row_policy { + let policy = row.str_optional(row_policy_col)?; + let repeats = row.u32_optional(repeats_col)?; + if !matches!( + (policy, repeats), + (Some(FPM_FORWARD_MEASUREMENT_POLICY), Some(1)) + | (Some("kvwarm_median_of_3"), Some(3)) + ) { + return Err(structural(format!( + "FPM row {index} measurement policy/repeats are inconsistent" + ))); + } + } // The string backend knobs must be present: "auto" or a pinned name. for (offset, name) in [(11usize, "moe_backend"), (12usize, "attention_backend")] { if match_identity[offset].is_empty() { @@ -762,7 +834,7 @@ fn load_pair( } // Full physical row key (collector contract) for duplicate detection, // in Python's _ROW_KEY_COLUMNS order. - let row_key: Vec = vec![ + let mut row_key: Vec = vec![ get_str("cell_id")?, get_str("model_path")?, get_str("system")?, @@ -795,12 +867,14 @@ fn load_pair( partition_policy, ]; + row_key.push(dcp.map(|value| value.to_string()).unwrap_or_default()); + // Seed provenance: fake-fallback rows stay in the row set (their // coordinates, the grid, and the domain gate are untouched); only // their latency value is healed by the replacement pass below, after // the duplicate/collision checks. let kv_seed_regime = row.str_optional(kv_seed_col)?.unwrap_or(""); - if !match_identity[15..] + if !match_identity[15..FPM_CELL_MATCH_COLUMNS.len()] .iter() .map(String::as_str) .eq(LEGACY_EXECUTION_IDENTITY) @@ -1212,6 +1286,8 @@ pub(crate) mod tests { /// fixture sets it, omits the column entirely — the pre-column /// legacy layout). pub kv_seed_regime: Option<&'static str>, + pub dcp: Option, + pub measurement: Option<(&'static str, u32)>, pub execution: Option<[&'static str; 4]>, } @@ -1235,6 +1311,8 @@ pub(crate) mod tests { system: "b200_sxm", backend: "vllm", kv_seed_regime: None, + dcp: None, + measurement: None, execution: None, } } @@ -1300,6 +1378,8 @@ pub(crate) mod tests { // The provenance column is written only when a fixture row sets it, // so default fixtures exercise the pre-column legacy layout. let has_kv_seed = rows.iter().any(|r| r.kv_seed_regime.is_some()); + let has_dcp = rows.iter().any(|r| r.dcp.is_some()); + let has_measurement = rows.iter().any(|r| r.measurement.is_some()); let has_execution = rows.iter().any(|r| r.execution.is_some()); let schema = "message schema { REQUIRED BINARY cell_id (UTF8); @@ -1338,6 +1418,16 @@ pub(crate) mod tests { } else { schema.to_string() }; + let mut schema = schema; + if has_dcp { + schema.insert_str(schema.rfind('}').unwrap(), "REQUIRED INT64 dcp;\n"); + } + if has_measurement { + schema.insert_str( + schema.rfind('}').unwrap(), + "REQUIRED BINARY measurement_policy (UTF8);\nREQUIRED INT64 measurement_repeats;\n", + ); + } let schema = if has_execution { schema.replace( "REQUIRED BINARY workload_kind (UTF8);", @@ -1497,6 +1587,31 @@ pub(crate) mod tests { .expect("write"); col.close().expect("close"); } + if has_dcp { + let values: Vec = rows.iter().map(|r| i64::from(r.dcp.unwrap_or(0))).collect(); + let mut col = rg.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&values, None, None) + .unwrap(); + col.close().unwrap(); + } + if has_measurement { + let values = str_col(&|r| r.measurement.map(|m| m.0).unwrap_or("").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.measurement.map(|m| m.1).unwrap_or(0))) + .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"); @@ -1563,7 +1678,7 @@ pub(crate) mod tests { assert_eq!(table.parquet_path(), external); assert!( table - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .is_ok() ); } @@ -1591,14 +1706,14 @@ pub(crate) mod tests { let table = loaded_table(tmp.path()); let mut identity = default_identity(4); identity[15..].clone_from_slice(&execution.map(str::to_string)); - assert!(table.select_cell(&identity, "org/model-a").is_ok()); + assert!(table.select_cell(&identity, "org/model-a", None).is_ok()); assert!( table - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .is_err() ); identity[16] = "decoder_bounded".to_string(); - assert!(table.select_cell(&identity, "org/model-a").is_err()); + assert!(table.select_cell(&identity, "org/model-a", None).is_err()); let prefill_rows: Vec = default_rows() .into_iter() .filter(|r| r.workload_kind == "prefill") @@ -1670,7 +1785,7 @@ pub(crate) mod tests { write_pair(tmp.path(), &default_rows()); assert!( loaded_table(tmp.path()) - .select_cell(&default_identity(4)[..15], "org/model-a") + .select_cell(&default_identity(4)[..15], "org/model-a", None) .is_ok() ); write_pair_with(tmp.path(), &default_rows(), |m| { @@ -1679,6 +1794,106 @@ pub(crate) mod tests { assert!(loaded_table(tmp.path()).cells().is_err()); } + #[test] + fn dcp_cells_do_not_match_plain_tp_or_each_other() { + let tmp = tempfile::tempdir().unwrap(); + // Same coordinates and cell ID; recorded DCP must distinguish the cells. + let rows = [ + RowSpec { + dcp: Some(1), + latency_ms: 7.0, + ..RowSpec::default() + }, + RowSpec { + dcp: Some(4), + latency_ms: 3.0, + ..RowSpec::default() + }, + ]; + write_pair(tmp.path(), &rows); + let table = loaded_table(tmp.path()); + assert_eq!(table.cells().unwrap().len(), 2); + assert!( + table + .select_cell(&default_identity(4), "org/model-a", None) + .is_err() + ); + for dcp in [1, 4] { + assert!( + table + .select_cell(&default_identity(4), "org/model-a", Some(dcp)) + .is_ok() + ); + // A string tail cannot hide DCP from the typed SOL/admission guards. + let mut identity = default_identity(4); + identity.push(dcp.to_string()); + assert!( + table + .select_cell(&identity, "org/model-a", None) + .unwrap_err() + .to_string() + .contains("supply DCP through dcp_size") + ); + } + write_pair_with(tmp.path(), &rows[..1], |meta| { + meta.insert( + "configuration_selector".into(), + serde_json::json!({"dcp": 4}), + ); + }); + assert!( + loaded_table(tmp.path()) + .cells() + .unwrap_err() + .to_string() + .contains("sidecar selector") + ); + } + + #[test] + fn mixed_measurement_policy_validates_each_record() { + let tmp = tempfile::tempdir().unwrap(); + let rows = [ + RowSpec { + measurement: Some((FPM_FORWARD_MEASUREMENT_POLICY, 1)), + ..RowSpec::default() + }, + RowSpec { + total_kv_read_tokens: 8192, + measurement: Some(("kvwarm_median_of_3", 3)), + ..RowSpec::default() + }, + ]; + let mixed_policy = |meta: &mut serde_json::Map| { + meta.insert( + "measurement_policy".into(), + FPM_ROW_MEASUREMENT_POLICY.into(), + ); + }; + write_pair_with(tmp.path(), &rows, mixed_policy); + assert_eq!(loaded_table(tmp.path()).cells().unwrap().len(), 1); + let invalid = [RowSpec { + measurement: Some(("kvwarm_median_of_3", 1)), + ..RowSpec::default() + }]; + write_pair_with(tmp.path(), &invalid, mixed_policy); + assert!( + loaded_table(tmp.path()) + .cells() + .unwrap_err() + .to_string() + .contains("policy/repeats") + ); + write_pair_with(tmp.path(), &default_rows(), mixed_policy); + assert!( + loaded_table(tmp.path()) + .cells() + .unwrap_err() + .to_string() + .contains("requires measurement_policy") + ); + } + #[test] fn absent_parquet_is_soft_not_collected() { let tmp = tempfile::tempdir().expect("tmpdir"); @@ -2216,14 +2431,16 @@ pub(crate) mod tests { let identity = default_identity(4); // Exact path selects its own cell. - let cell = table.select_cell(&identity, "org/model-a").expect("select"); + let cell = table + .select_cell(&identity, "org/model-a", None) + .expect("select"); assert_eq!(cell.model_path, "org/model-a"); // Unknown path: never borrows a collected cell (D1 exact-only). - let err = table.select_cell(&identity, "org/other").unwrap_err(); + let err = table.select_cell(&identity, "org/other", None).unwrap_err(); assert!(err.to_string().contains("never substitutes"), "{err}"); // Unknown identity: the no-match error listing what was collected. let err = table - .select_cell(&default_identity(8), "org/model-a") + .select_cell(&default_identity(8), "org/model-a", None) .unwrap_err(); assert!(err.to_string().contains("No FPM cell matches"), "{err}"); } @@ -2235,7 +2452,7 @@ pub(crate) mod tests { let tmp = tempfile::tempdir().expect("tmpdir"); write_pair(tmp.path(), &default_rows()); let err = loaded_table(tmp.path()) - .select_cell(&default_identity(4), "some/other-model") + .select_cell(&default_identity(4), "some/other-model", None) .unwrap_err(); assert!(err.to_string().contains("never substitutes"), "{err}"); } @@ -2287,7 +2504,7 @@ pub(crate) mod tests { let tmp = tempfile::tempdir().expect("tmpdir"); write_pair(tmp.path(), &flat_prefill_rows()); let cell_max = loaded_table(tmp.path()) - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .expect("select") .prefill_batch_clamp_max; assert_eq!(cell_max, Some(4)); @@ -2295,7 +2512,7 @@ pub(crate) mod tests { let tmp = tempfile::tempdir().expect("tmpdir"); write_pair(tmp.path(), &default_rows()); let cell_max = loaded_table(tmp.path()) - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .expect("select") .prefill_batch_clamp_max; assert_eq!(cell_max, None); @@ -2307,7 +2524,7 @@ pub(crate) mod tests { let tmp = tempfile::tempdir().expect("tmpdir"); write_pair(tmp.path(), &bumpy); let cell_max = loaded_table(tmp.path()) - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .expect("select") .prefill_batch_clamp_max; assert_eq!(cell_max, None); @@ -2319,7 +2536,7 @@ pub(crate) mod tests { write_pair(tmp.path(), &cliff_decode_rows()); let table = loaded_table(tmp.path()); let cell = table - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .expect("select"); // Rungs = batches whose (x, x+1) pair was collected. assert_eq!(cell.decode_rungs, vec![496, 512]); @@ -2335,7 +2552,7 @@ pub(crate) mod tests { write_pair(tmp.path(), &default_rows()); let table = loaded_table(tmp.path()); let cell = table - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .expect("select"); assert!(cell.decode_rungs.is_empty()); assert_eq!(cell.decode_batches, vec![8, 16]); @@ -2355,7 +2572,7 @@ pub(crate) mod tests { .collect(); write_pair(tmp.path(), &rows); let err = loaded_table(tmp.path()) - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .unwrap_err(); assert!(err.to_string().contains("No FPM cell matches"), "{err}"); } @@ -2378,11 +2595,11 @@ pub(crate) mod tests { let mut wideep_identity = default_identity(4); wideep_identity[13] = "True".to_string(); let cell = table - .select_cell(&wideep_identity, "org/model-a") + .select_cell(&wideep_identity, "org/model-a", None) .expect("select"); assert_eq!(cell.match_identity[13], "True"); let err = table - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .unwrap_err(); assert!(err.to_string().contains("No FPM cell matches"), "{err}"); } diff --git a/crates/core/src/perfmodel/perf_database/mod.rs b/crates/core/src/perfmodel/perf_database/mod.rs index 622839d200..90ab1c0108 100644 --- a/crates/core/src/perfmodel/perf_database/mod.rs +++ b/crates/core/src/perfmodel/perf_database/mod.rs @@ -1144,7 +1144,7 @@ mod tests { assert!( db.fpm_forward - .select_cell(&default_identity(4), "org/model-a") + .select_cell(&default_identity(4), "org/model-a", None) .is_ok() ); } diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 2adc2489c9..3a560998cd 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1100,6 +1100,8 @@ struct EngineBuildRequest { tp_size: u32, pp_size: u32, attention_dp_size: u32, + dcp_size: Option, + fpm_options: crate::FpmInterpolationConfig, moe_tp_size: Option, moe_ep_size: Option, gemm_quant_mode: Option, @@ -1154,6 +1156,8 @@ impl AicEngineBuilder { tp_size: 1, pp_size: 1, attention_dp_size: 1, + dcp_size: None, + fpm_options: crate::FpmInterpolationConfig::default(), moe_tp_size: None, moe_ep_size: None, gemm_quant_mode: None, @@ -1530,6 +1534,15 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result, + comm_quant_mode: Option<&str>, + ) -> PyResult { + let options: crate::FpmInterpolationConfig = + serde_json::from_str(options_json).map_err(|e| { + PyValueError::new_err(format!("invalid FPM interpolation options: {e}")) + })?; + options + .validate_quant_modes(fmha_quant_mode, comm_quant_mode) + .map_err(aic_to_py)?; + // The compilation adapter needs resolved values, not the compact public + // serialization that omits defaults. Python never supplies these defaults. + Ok(serde_json::json!({ + "text_only": options.text_only, + "fpm_parquet_path": options.fpm_parquet_path, + "unrecorded_quant_modes": options.unrecorded_quant_modes, + }) + .to_string()) + } + /// Migration adapter for previously saved flat tuning options. #[staticmethod] fn legacy_estimator_config(options_json: &str) -> PyResult { @@ -1965,6 +2007,7 @@ impl PyForwardPassPerfModel { tp: request.tp_size, pp: request.pp_size, attention_dp: request.attention_dp_size, + dcp: request.dcp_size, moe_tp_size: request.moe_tp_size, moe_ep_size: request.moe_ep_size, gemm_quant_mode: request.gemm_quant_mode, @@ -1995,6 +2038,7 @@ impl PyForwardPassPerfModel { enable_shared_layer: request.shared_layer, strict_provenance: legacy.strict_provenance, }; + config.validate().map_err(aic_to_py)?; serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) } @@ -2264,6 +2308,7 @@ mod tests { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { + dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(1), diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index edf5613d3d..f5da5d0bbc 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -172,6 +172,8 @@ struct AicTimingConfig { #[serde(default = "one")] attention_dp: u32, #[serde(default)] + dcp: Option, + #[serde(default)] moe_tp_size: Option, #[serde(default)] moe_ep_size: Option, @@ -316,6 +318,7 @@ impl AicTimingConfig { tp: self.tp, pp: self.pp, attention_dp: self.attention_dp, + dcp: self.dcp, moe_tp_size: self.moe_tp_size, moe_ep_size: self.moe_ep_size, gemm_quant_mode: self.gemm_dtype.clone(), @@ -900,6 +903,10 @@ fn materialize_aic_capacity( if capacity_is_explicit || role.rank.state_cache.is_some() { return Ok(()); } + ensure!( + config.dcp.is_none_or(|dcp| dcp == 1), + "DCP FPM replay requires explicit KV block capacity; automatic DCP/hybrid sizing is unsupported" + ); let blocks = estimate(config, role)?; ensure!(blocks > 0, "AIC estimated zero KV-cache blocks"); role.rank.num_gpu_blocks = blocks; @@ -2710,6 +2717,7 @@ mod tests { // test engine is synthetic; the production constructor is unchanged. let missing_op = |phase| { Op::FpmForward(FpmForwardOp { + dcp_size: None, name: "missing".into(), phase, model_path: "missing-test-model".into(), @@ -2822,6 +2830,7 @@ mod tests { backend_version: None, pp: 1, attention_dp: 1, + dcp: None, moe_tp_size: None, moe_ep_size: None, gemm_dtype: None, @@ -2913,6 +2922,33 @@ mod tests { assert_eq!(role.rank.num_gpu_blocks, 17); } + #[test] + fn dcp_capacity_requires_explicit_blocks() { + let mut config = aic_config(); + config.tp = 4; + config.dcp = Some(4); + let mut role = aggregated_role(&ReplayEngineConfig::default()); + role.tensor_parallel_size = 4; + role.rank.num_gpu_blocks = 17; + + let error = materialize_aic_capacity(&config, &mut role, false, |_, _| { + panic!("DCP must be rejected before estimating capacity") + }) + .unwrap_err(); + assert!(error.to_string().contains("explicit KV block capacity")); + assert_eq!(role.rank.num_gpu_blocks, 17); + + materialize_aic_capacity(&config, &mut role, true, |_, _| { + panic!("explicit DCP capacity must not invoke the estimator") + }) + .unwrap(); + assert_eq!(role.rank.num_gpu_blocks, 17); + + config.dcp = Some(1); + materialize_aic_capacity(&config, &mut role, false, |_, _| Ok(321)).unwrap(); + assert_eq!(role.rank.num_gpu_blocks, 321); + } + #[test] fn replay_provider_rejects_graph_profile_before_native_construction() { let mut config = aic_config(); diff --git a/crates/core/tests/perfmodel/memory_round_trip.rs b/crates/core/tests/perfmodel/memory_round_trip.rs index 9776a286cd..b4ef917e11 100644 --- a/crates/core/tests/perfmodel/memory_round_trip.rs +++ b/crates/core/tests/perfmodel/memory_round_trip.rs @@ -93,6 +93,7 @@ fn request(tolerance_fraction: Option) -> KvCacheEstimateRequest { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { + dcp_size: None, tp_size: 1, pp_size: 1, attention_dp_size: Some(1), diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index c535ef03c1..5157021eea 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -144,7 +144,9 @@ mod tests { // v20: FpmForwardOp gained original_fmha_quant_mode for selector diagnostics. // v21: EngineConfig and MoeOp gained exact moe_kernel_source identity. // v22: observed MoE selection and exact prefill graph composites. - assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 22); + // v23: ParallelMapping gained optional recorded DCP identity. + // v24: FpmForwardOp carries typed DCP separately from matching strings. + assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 24); assert_eq!(FPM_VERSION, 1); assert_eq!(ForwardPassMetrics::default().version, FPM_VERSION); } diff --git a/docs/core-api.md b/docs/core-api.md index 78738d59d2..05fce8a944 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -318,7 +318,14 @@ nested paths. The supported namespaces are: (5), `factor_bounds` (min 0.5, max 2.0), and the existing `max_num_tokens` (8192), `max_batch_size` (512), and `max_kv_tokens` (2000000) ranges. - `op_level`: optional `decode_workload_distribution` selects a measured decode-MoE distribution, and `prefill_graph_profile` selects a qualified direct-prefill graph composition. Saved configurations retain the resolved immutable `prefill_graph_profile_id`, which is validated on reload. Unknown fields are rejected. -- `fpm_interpolation`: optional `fpm_parquet_path` selects an external FPM parquet with its same-stem metadata sidecar. Unknown fields are rejected. +- `fpm_interpolation`: `text_only` (false) permits text prefill/decode profiles + for multimodal architectures while retaining encoder weights. It does not + supply encoder timing. `unrecorded_quant_modes` (empty) may contain `fmha` + and/or `comm` to match an explicitly unrecorded precision field in a profile. + The corresponding top-level quant mode must remain unset. This selects null + profile values exactly; it does not make precision matching a wildcard. + `fpm_parquet_path` selects an external FPM parquet with its same-stem metadata + sidecar. Unknown fields are rejected. Engine replay rank arguments accept `decode_workload_distribution` (alias `aic_decode_workload_distribution`) only with AIC timing. An active selector paired with a non-AIC timing model, including fixed or polynomial timing, is rejected. The AFD companion's fixed timing and legacy estimator paths also reject active selectors because they cannot apply the profile. `None` preserves ordinary timing in these paths. @@ -340,6 +347,75 @@ return `None`. `tune_with_fpms()` preserves the established FPM observation contract. Native construction still uses Python model compilation; estimator selection, regression, correction, and latency computation are owned by Rust. +### DCP self-benchmark profiles + +`dcp` is an optional recorded decode-context-parallel dimension within TP. +It must be positive and divide `tp`; it does not multiply GPU or MoE group +counts. Missing DCP and explicit DCP1 remain distinct FPM identities. An +explicit DCP8 request cannot consume ordinary TP8 or unrecorded-DCP data. +DCP greater than one currently supports measured vLLM `fpm_interpolation` +timing; op-level DCP and SOL-dependent transfer paths report unsupported. + +FPM v6 accepts optional `dcp` in the Parquet identity and sidecar selector. +The loader also accepts `per_row_single_sample_or_median_of_3` sidecars when +every row declares a consistent `measurement_policy`/`measurement_repeats` +pair: `dynamo_native_single_sample_v1`/1 or `kvwarm_median_of_3`/3. Values are +already aggregated by the producer; the loader preserves their latency. + +For a Kimi K3 text profile, an explicit configuration can be: + +```python +config = ForwardPassPerfModelConfig( + model="moonshotai/Kimi-K3", system="gb300", backend="vllm", + backend_version="0.29.0", worker_type="aggregated", + tp=8, pp=1, attention_dp=1, moe_tp_size=8, moe_ep_size=1, dcp=8, + gemm_quant_mode="bfloat16", moe_quant_mode="w4a16_mxfp4", + kvcache_quant_mode="fp8", attention_backend="FLASHINFER_MLA", + estimation_mode="fpm_interpolation", fallback_policy="deny", + systems_paths=("/absolute/profile/systems",), + estimator_config={ + "fpm_interpolation": { + "text_only": True, + "unrecorded_quant_modes": ["fmha", "comm"], + }, + "correction": {"enabled": False}, + }, +) +model = RustForwardPassPerfModel.best_available(config) +``` + +Only use `unrecorded_quant_modes` for fields that the selected profile actually +leaves unrecorded. The runtime `FLASHINFER_MLA` label is retained for exact FPM +matching while the compiler uses its internal FlashInfer backend description. +The model architecture remains registered once; adding a parallel configuration +does not require another model class. + +Place the reviewed primary pair at +`systems/data/gb300/vllm/0.29.0/fpm_forward_perf.{parquet,metadata.json}` and copy +the matching hardware YAML into the systems root. Keep source hashes and the +pinned dataset revision with the profile. Do not combine synthetic-attention +boundary points or nonuniform layouts with a balanced primary profile. + +Replay YAML passes recorded DCP through +`engine.workers..parallelism.decode_context`. Per-worker `timing` accepts +the canonical quant-mode fields and `attention_backend`, alongside estimator +selection and controls. DCP FPM replay requires explicit fixed KV block capacity; +automatic DCP/hybrid capacity sizing is not implemented. Host offload or P/D +transfer also requires explicit KV bytes per token with DCP. +`decode_context` is currently supported by AISimulate's `--stack engine` only. +Dynamo Replay and Planner do not yet support this field; their configuration +propagation, cache identity, and dependency version need a downstream update. +These timing precision/backend overrides are prediction-only; recommendation +rejects them until its feasibility preflight supports the same identity. +Supplying a fixed pool does not add KDA checkpoint, eviction, or chunk-alignment +fidelity to the generic Replay cache/scheduler. This API change enables timing +consumption, not full hybrid-cache simulation or multimodal prediction from +text-only measurements. + +The positional engine-spec wire version remains unchanged: the engine identity +is JSON-encoded and the FPM match identity is already variable-length. Legacy +configuration and profiles without DCP remain accepted as unrecorded DCP. + ### Migrating saved configuration Use `ForwardPassPerfModelConfig.from_legacy_engine_config(old_config, diff --git a/docs/deepseek-v41-storage.md b/docs/deepseek-v41-storage.md index b2a1647812..cb160f11ff 100644 --- a/docs/deepseek-v41-storage.md +++ b/docs/deepseek-v41-storage.md @@ -66,5 +66,5 @@ Other backends retain the explicitly unqualified `logical_fp4` inventory and traffic, avoiding a guessed physical cache precision. The operator serializes `kv_cache_layout` independently of `fmha_quant_mode`. Legacy JSON without this field defaults to the theoretical layout; fresh SGLang graphs always name the -physical layout. Engine wire schema 22 rejects earlier bincode schemas before decoding +physical layout. Engine wire schema 24 rejects earlier bincode schemas before decoding the new positional field. Model observations and calibration data are unchanged. diff --git a/docs/deepseek-v41.md b/docs/deepseek-v41.md index 90528c4237..782a55a052 100644 --- a/docs/deepseek-v41.md +++ b/docs/deepseek-v41.md @@ -94,7 +94,7 @@ The `logical_fp4` layout (288-byte compressed main, 68-byte index, FP8 window, 890-byte global slope) remains the explicit theoretical estimate for vLLM and TRT-LLM; their runtime storage is unqualified. Both inventories exclude allocator page padding and spare pages. See [source proof and limits](deepseek-v41-storage.md) -for full-context scoring, cache read/write accounting, and schema 22 compatibility. +for full-context scoring, cache read/write accounting, and schema 24 compatibility. Engram's two GPU-resident hash tables include FP8 block scales and TP row sharding. Their full resident size is independent of tokens accessed. Lookup @@ -131,5 +131,7 @@ identities. FPM selector diagnostics append a positional field to `FpmForwardOp` in schema 20. Exact MoE kernel-source identity extends the layout in schema 21, and the Rubin pilot adds observed-MoE selection and prefill graph composites in schema 22. -Schema-21 and earlier bincode inputs are rejected before decoding; legacy JSON +DCP identity extends the combined layout in schema 23. +Typed FPM DCP extends the operator payload in schema 24. +Schema-23 and earlier bincode inputs are rejected before decoding; legacy JSON may still use the documented field defaults. diff --git a/python/aisimulate/docs/add_a_new_model.md b/python/aisimulate/docs/add_a_new_model.md index 2e6c3e72ba..b2e13314cd 100644 --- a/python/aisimulate/docs/add_a_new_model.md +++ b/python/aisimulate/docs/add_a_new_model.md @@ -103,7 +103,7 @@ python/aisimulate/src/aisimulate_core/systems/data////< Do not relabel another backend version's data as newly measured. Update query version/support metadata only when the new cell meets its collection and consumer contracts. For whole-forward profiles, use the separate -[FPM collection-to-prediction workflow](fpm/end-to-end-workflow.md). +[Self-benchmarking and FPM onboarding guide](fpm/self-benchmarking-and-onboarding.md). ## 4. Validate the intended public path diff --git a/python/aisimulate/docs/fpm/README.md b/python/aisimulate/docs/fpm/README.md index dfefb5198d..fbeb0abea0 100644 --- a/python/aisimulate/docs/fpm/README.md +++ b/python/aisimulate/docs/fpm/README.md @@ -5,16 +5,23 @@ SPDX-License-Identifier: Apache-2.0 # Forward-pass models -Start with the [end-to-end FPM workflow](end-to-end-workflow.md) to collect -whole-forward measurements, publish a performance-data pair, load it through -the SDK, and run an AISimulate prediction. The guide includes prerequisites, -commands, expected artifacts, acceptance checks, and recovery steps. +Start with the [self-benchmark and FPM onboarding guide](self-benchmarking-and-onboarding.md). +Its seven steps cover support checks, measurement planning, collection, profile +validation, canonical model construction, Replay integration, and accuracy +validation. It explains which engine-iteration differences self-collection can +capture and which behaviors, including PP, need separate simulator support. + +Self-benchmark collection currently supports vLLM configurations that pass +model/runtime validation; new architectures can require benchmark adaptation. +SGLang and TensorRT-LLM support is coming soon. The guide includes the collected +[Kimi K3 TP8+DCP8 profile](self-benchmarking-and-onboarding.md#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile) +and a new MiniMax collection campaign as worked examples of the general procedure. There are two distinct workflows: | Workflow | Input | Consumer | | --- | --- | --- | -| Offline whole-forward FPM | Collector-produced `fpm_forward_perf.parquet` and its metadata sidecar | `forward_model="fpm"`: lookup, interpolation, and supported SOL transfer | +| Offline whole-forward FPM | Validated `fpm_forward_perf.parquet` and its metadata sidecar | `best_available` with `estimation_mode="fpm_interpolation"`: lookup, interpolation, and supported SOL transfer | | Online regression | Observed per-iteration, per-rank telemetry | A role-bound model updated with `tune_with_fpms` | Offline FPM does not require an additional regression-training step. Predicting diff --git a/python/aisimulate/docs/fpm/end-to-end-workflow.md b/python/aisimulate/docs/fpm/end-to-end-workflow.md deleted file mode 100644 index c78cd04618..0000000000 --- a/python/aisimulate/docs/fpm/end-to-end-workflow.md +++ /dev/null @@ -1,585 +0,0 @@ - - -# Collect FPM data and use it in AISimulate - -This guide follows one offline whole-forward FPM campaign from planning to -request-level prediction. It uses MiniMax-M2.7 on four H200 GPUs with pure tensor -parallelism, then consumes the newly collected data from an external directory. -It assumes an existing, working GPU collection environment. Fill in the -collection inputs from that environment before executing GPU steps. - -```text -Freeze plan -> GPU smoke -> GPU collection -> validate and publish pair - | - v -Independent measurements <- compare <- AIS predict <- SDK FPM query -``` - -The collector measures complete forward passes in a Dynamo/vLLM runtime. It -publishes `fpm_forward_perf.parquet` together with -`fpm_forward_perf.metadata.json`. The Rust performance model loads this pair, -selects a matching cell, and uses exact lookup, interpolation, and supported SOL -transfer. There is no separate `train` or `fit` command in this workflow. -[Online telemetry regression](aic-fpm-regression-design.md) is a separate model. - -## Scope and expected outcomes - -| Stage | Runs on | Expected result | -| --- | --- | --- | -| Plan | Machine running Collector | Frozen model, topology, sampling policy, and phase cells; no Kubernetes workload | -| Smoke | GPU cluster | Minimal prefill and decode execution passes; no formal database publication | -| Collect and publish | Machine running Collector + GPU cluster | All planned cells pass; a matching Parquet/metadata pair is published; owned resources are removed | -| SDK query | Machine running AISimulate | The new pair loads and returns a finite, positive forward latency in milliseconds | -| AIS predict | Machine running AISimulate | `prediction.json` and, when requested, `requests.jsonl` for the configured traffic | -| Accuracy validation | Machine running AISimulate + independent GPU run | Matched prediction/measurement errors and coverage, assessed against thresholds chosen for the target use case | - -The collector currently supports vLLM with `PP=CP=1`. This example uses -`tp=4, pp=1, dp=1, moe_tp=4, moe_ep=1, cp=1`, default checkpoint quantization, -and one worker replica. Other supported topologies need their own collected -cells. The examples specify expected artifact structure and success conditions; -they do not promise a fixed latency, row count, collection duration, or accuracy. - -## 1. Check the AISimulate commands - -Activate an existing environment containing the current AISimulate package, -then check the commands from the repository root. If AISimulate is not installed, -follow the [developer setup](../../../../DEVELOPMENT.md) first. Installing the -development extras is not a required step for each collection campaign. - -```bash -export AIS_REPO="$PWD" -aisimulate predict --help -cd "$AIS_REPO/python/aisimulate" -python3 -m collector.fpm_forward --help -``` - -Use the same activated Bash or zsh session for the remaining commands. The -machine running Collector needs model configuration metadata and -`kubectl` access to the target namespace. It does not need a local GPU. - -For this example, use four H200 GPUs on one node. Obtain the namespace, -checkpoint PVC/path, collection image reference, and Generator target release -from the existing collection configuration. Record the AIS commit, image -digest, model revision, GPU type, and resolved configuration with the results. - -Before launching, review [recovery and cleanup](#8-recovery-and-cleanup), which -uses the campaign's generated manifests. Preserve logs before manual -teardown and verify that no resources from this campaign remain. - -## 2. Choose the example's output paths - -This guide uses an external directory to keep the example's results separate -from bundled databases and previous campaigns. The directory name and layout -are choices for this walkthrough, not mandatory FPM setup steps. The later -commands use these paths consistently; replace the absolute path below with a -new persistent directory. - -```bash -export FPM_RUN=/absolute/new/path/m27-h200-tp4 -mkdir -p "$(dirname "$FPM_RUN")" -mkdir "$FPM_RUN" -mkdir -p "$FPM_RUN/systems/data" -cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/h200_sxm.yaml" \ - "$FPM_RUN/systems/" -cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/query_versions.yaml" \ - "$FPM_RUN/systems/" -``` - -`h200_sxm.yaml` provides hardware specifications for the later SDK/prediction -steps when they load this external systems root. `query_versions.yaml` preserves -the version-slot policy chosen for this example. Neither file is measured FPM -data, and copying them does not collect data or train a model. The collector's -`--fpm-database-root` override is optional; the later examples select it -explicitly so they publish and consume data from the same directory. - -Save the following as `$FPM_RUN/collector-k8s.yaml`, replacing every -`REPLACE_...` value with the corresponding existing collection setting: - -```yaml -K8sConfig: - k8s_namespace: REPLACE_WITH_EXISTING_NAMESPACE - k8s_image: REPLACE_WITH_EXISTING_COLLECTION_IMAGE_DIGEST - k8s_pvc_name: REPLACE_WITH_MODEL_PVC - k8s_pvc_mount_path: /models - k8s_model_path_in_pvc: REPLACE_WITH_RELATIVE_CHECKPOINT_DIRECTORY - k8s_hf_home: /tmp/fpm-hf-cache -``` - -The dedicated collector accepts deployment-only `K8sConfig` fields. In -particular, do not copy `Workers.agg.extra_cli_args` or `ServiceConfig` from a -standalone Generator example into this file. The accepted fields are defined in -[the collector input resolver](../../collector/fpm_forward/entry.py). - -## 3. Freeze and inspect the plan - -Use the Generator target release from the existing collection configuration -and define the campaign arguments once: - -```bash -export AIS_DYNAMO_RELEASE=REPLACE_WITH_MATCHING_DYNAMO_RELEASE -cd "$AIS_REPO/python/aisimulate" -fpm_args=( - --backend vllm - --model-path MiniMaxAI/MiniMax-M2.7 - --gpu h200_sxm - --fpm-max-gpus 4 - --fpm-gpu-counts 4 - --fpm-parallel-presets pure_tp - --fpm-kv-cache-dtypes auto - --fpm-max-prefill-isl 8192 - --fpm-max-prefill-cudagraph-size 2048 - --dynamo-version "$AIS_DYNAMO_RELEASE" - --generator-config "$FPM_RUN/collector-k8s.yaml" - --checkpoint-dir "$FPM_RUN/checkpoints" - --fpm-artifact-root "$FPM_RUN/artifacts" - --fpm-database-root "$FPM_RUN/systems/data" -) -python3 -m collector.fpm_forward "${fpm_args[@]}" \ - --plan-only > "$FPM_RUN/plan.json" -``` - -If the checkpoint exists only inside the target Pod and its metadata cannot be -resolved on the machine running Collector, add `--fpm-model-config /absolute/config.json` -to `fpm_args`. Keep the canonical model ID in `--model-path`. - -**Expected:** exit code 0 and JSON describing the frozen plan, without creating -Kubernetes workloads. Inspect the model/config hash, system, topology, -quantization, phase cells, sampling policy, and admission decisions before -continuing. For this single-topology, single-KV-setting example, expect one -prefill cell and one decode cell if both are admitted. The plan is authoritative; -do not proceed with an empty plan or unexplained omissions. -`counts.cells` should be 2 for that case, while `counts.points` is -`"runtime-determined"`: runtime initialization determines the actual point set. - -The sampling settings have distinct meanings: - -- `--fpm-max-prefill-isl 8192` bounds the scheduled prefill **new-token** axis. - It does not set runtime context length. -- `--fpm-max-prefill-cudagraph-size 2048` is the prefill capture-size policy. - Align it with the deployment being modeled before freezing a formal campaign. -- The collector currently sets runtime `max_model_len=-1` for vLLM auto-fit and - has no CLI override for that setting. Record the resolved value with the - collection results. -- The default is five global warm-up iterations and one measurement per point. - Published `latency_ms` uses the maximum rank wall time, converted to ms; - it is not a mean or median over repeated measurements. - -The complete option list is available from `python3 -m collector.fpm_forward ---help` and [its argument definitions](../../collector/fpm_forward/config.py). - -### Include multiple GPU counts and parallel configurations - -Both GPU counts and parallel presets accept multiple values. To expand this -MiniMax campaign, replace these three entries in `fpm_args` before generating -the plan: - -```text ---fpm-max-gpus 8 ---fpm-gpu-counts 2 4 8 ---fpm-parallel-presets pure_tp tep dep -``` - -For each selected GPU count `N`, the MoE presets request these configurations -(`PP=CP=1`): - -| Preset | Attention TP | Attention DP | MoE TP | MoE EP | -| --- | --- | --- | --- | --- | -| `pure_tp` | N | 1 | N | 1 | -| `tep` | N | 1 | 1 | N | -| `dep` | 1 | N | 1 | N | - -The planner combines the selected counts and presets, deduplicates identical -topologies, and applies model/backend and memory admission. It then creates -prefill and decode cells for admitted topology, KV dtype, and backend-policy -combinations. Read `plan.json` for the actual configurations and cell count; -do not assume every requested combination is admitted. Dense models use the -`tp` preset; MoE `pure_tp` requires support in the model's capability profile. -See [topology enumeration](../../collector/fpm_forward/topology.py). - -`--fpm-max-gpus 8` is the maximum GPU count for one cell. The runner executes -cells sequentially; multiple Pods belonging to one cell execute together. -Selecting multiple counts/presets does not launch all configurations -simultaneously. - -Regenerate and inspect `plan.json` after changing these options. If the original -campaign has already run, first prepare a fresh campaign directory as in step 2 -and redefine `fpm_args` with the new paths. The later TP4 SDK/predict examples -still select only their matching cell from the larger database; they do not -automatically compare all collected configurations. - -## 4. Run smoke, then formal collection - -First read the number of cells from the inspected plan and execute the minimal -profile for every cell. This works for both the two-cell TP4 example and an -expanded campaign: - -```bash -FPM_CELL_COUNT=$(python3 -c \ - 'import json, sys; n = len(json.load(open(sys.argv[1]))["cells"]); assert n > 0; print(n)' \ - "$FPM_RUN/plan.json") -python3 -m collector.fpm_forward "${fpm_args[@]}" --smoke --limit "$FPM_CELL_COUNT" -``` - -**Expected:** exit code 0, all selected smoke cells pass, raw runtime evidence is saved, -and no formal Parquet pair is published. `--smoke` without `--limit` runs only -the first cell. `--limit 2` is sufficient for the original two-cell example, -but in an expanded plan the first two cells can both be prefill: the planner -lists all prefill cells before decode cells. Use the full plan count to check -both phases of every admitted configuration. -In `checkpoints/fpm_forward_smoke.json`, check `smoke.status == "passed"`, -`smoke.cell_count == FPM_CELL_COUNT`, and `smoke.formal_database_written == false`. - -Only after the smoke command succeeds and its evidence is checked, run: - -```bash -python3 -m collector.fpm_forward "${fpm_args[@]}" -``` - -Formal collection must omit both `--smoke` and `--limit`. Smoke and formal -checkpoints/artifacts are isolated automatically, so the roots above can be -shared; smoke samples are not reused as formal measurements. - -**Expected:** all frozen cells pass and the formal pair is published. Check the -formal checkpoint's database status, that `missing_cells` is empty, and that -`published_cells` covers the planned cells. In this fresh database, -`skipped_first_publisher_wins` should be empty. Exit code 0 alone is insufficient -to prove a new attempt replaced existing data: publication preserves a cell -already owned by another attempt. `--fpm-publish-partial` is an explicit opt-in -for incomplete campaigns and still returns a nonzero result. - -For this fresh, two-cell campaign, the formal checkpoint should contain: - -```json -{ - "database": { - "status": "passed", - "published_cells": 2, - "plan_cells": 2, - "missing_cells": [], - "skipped_first_publisher_wins": [] - } -} -``` - -This is an excerpt of the expected fields, not a complete checkpoint file. -For an expanded campaign, use its actual plan count in place of 2; a complete -fresh publication has `published_cells == plan_cells == FPM_CELL_COUNT`. - -The resulting directory has this structure; `` is the first -16 characters of the plan hash: - -```text -$FPM_RUN/ - plan.json - checkpoints/ - fpm_forward_smoke.json - fpm_forward.json - artifacts// - collection-plan.json - run-manifest.json - cells//... - smoke/... - systems/ - h200_sxm.yaml - query_versions.yaml - data/h200_sxm/vllm// - fpm_forward_perf.parquet - fpm_forward_perf.metadata.json -``` - -Retain generated manifests/scripts, raw benchmark files, resolved configs, -provenance, logs, checkpoints, and the published pair. Inspect -`kv_seed_regime`; do not assume every row contains real-KV evidence. See the -[collector contract](../../collector/README.md#whole-forward-fpm-campaign) for -accepted regimes and publication behavior. - -## 5. Inspect the published pair - -The version directory is the actual Pod's `importlib.metadata.version("vllm")`, -also recorded in the published rows/metadata. It is not the AIS version or the -`--dynamo-version` used to select templates. Read it from the artifacts before -setting this variable: - -```bash -export FPM_VLLM_VERSION=REPLACE_WITH_ACTUAL_COLLECTED_VLLM_VERSION -export FPM_PARQUET="$FPM_RUN/systems/data/h200_sxm/vllm/$FPM_VLLM_VERSION/fpm_forward_perf.parquet" -python3 - "$FPM_PARQUET" > "$FPM_RUN/pair-inspection.json" <<'PY' -import hashlib -import json -import sys -from pathlib import Path - -import pyarrow.parquet as pq - -parquet = Path(sys.argv[1]) -metadata = json.loads(parquet.with_suffix(".metadata.json").read_text()) -with parquet.open("rb") as stream: - digest = hashlib.file_digest(stream, "sha256").hexdigest() -assert metadata["schema_name"] == "aic_fpm_forward_perf" -assert metadata["schema_version"] == 6 -assert metadata["parquet_sha256"] == digest -table = pq.read_table(parquet) -assert table.num_rows == metadata["row_count"] and table.num_rows > 0 -identity = [ - "model_path", "system", "backend", "backend_version", "workload_kind", - "tp", "pp", "dp", "moe_tp", "moe_ep", "cp", "gemm_quant_mode", - "moe_quant_mode", "fmha_quant_mode", "comm_quant_mode", "kv_cache_dtype", - "moe_backend", "attention_backend", "enable_wideep", "enable_eplb", -] -frame = table.to_pandas() -report = { - "hash_match": True, - "row_count": table.num_rows, - "phase_cells": frame[identity].drop_duplicates().to_dict(orient="records"), - "kv_seed_regimes": frame["kv_seed_regime"].value_counts().to_dict(), -} -print(json.dumps(report, indent=2)) -PY -``` - -**Expected:** exit code 0, `hash_match: true`, a positive row count, and phase -identities matching the campaign. This structural check does not replace the -native loader's validation or prove interpolation coverage. Inspect the -`batch_size`, `total_prefill_tokens`, and `total_kv_read_tokens` coordinates -before selecting queries; axis minima/maxima alone do not prove that all -interior points can be answered. - -Keep these paths distinct: - -| Setting | Value in this guide | -| --- | --- | -| Collector `--fpm-database-root` | `$FPM_RUN/systems/data` | -| SDK `systems_path` / Python systems root | `$FPM_RUN/systems` | -| Pair directory | `$FPM_RUN/systems/data/h200_sxm/vllm/$FPM_VLLM_VERSION` | - -Whole-forward FPM has no extra model-family directory between the system and -backend. Always transfer the metadata sidecar together with the Parquet file. - -## 6. Load the new data and query one forward pass - -This query explicitly selects the external root. The example asks for a -four-request prefill with 1,024 new tokens per request and no cached prefix; -confirm this coordinate is covered by the new collection. - -```bash -AIC_ALLOW_UNLISTED_VERSIONS=1 \ -python3 - "$FPM_RUN/systems" "$FPM_VLLM_VERSION" <<'PY' -import json -import math -import sys -from pathlib import Path - -from aisimulate_core.sdk import EngineHandle - -engine = EngineHandle.compile( - "MiniMaxAI/MiniMax-M2.7", "h200_sxm", "vllm", - backend_version=sys.argv[2], - systems_path=str(Path(sys.argv[1]).resolve(strict=True)), - forward_model="fpm", - tp_size=4, pp_size=1, attention_dp_size=1, - moe_tp_size=4, moe_ep_size=1, -) -latency_ms = engine.predict_prefill_latency(bs=4, isl=1024, prefix=0) -assert math.isfinite(latency_ms) and latency_ms > 0 -print(json.dumps({"prefill_ms": latency_ms})) -PY -``` - -**Expected:** exit code 0 and a finite, positive `prefill_ms`. No serving GPU is -used for this query. A successful lookup confirms this query can consume the -pair, not that all replay steps will be in-domain or that the prediction is -accurate. - -The model, hardware, backend version, parallel shape, quantization, and backend -knobs must match the collected identity. This example uses checkpoint defaults. -For explicitly collected nondefault precision, use the SDK's -`gemm_quant_mode`, `moe_quant_mode`, `fmha_quant_mode`, `kvcache_quant_mode`, and -`comm_quant_mode` arguments with the corresponding enum names as strings. See -[the compile API](../../src/aisimulate_core/sdk/engine.py). - -This guide copies the version policy file and uses literal backend versions. -If the collected version is already queryable, omit -`AIC_ALLOW_UNLISTED_VERSIONS=1`. The override is a transitional measure, not a -support guarantee. Do not rename version directories or edit identities to -bypass a mismatch. External trees without `query_versions.yaml` use a different -policy that disables the version-slot gate; omitting the file is not needed for -this workflow. - -## 7. Run AIS predict using the same external data - -Save this as `$FPM_RUN/predict.yaml`. Replace the backend version with the -collected version. Confirm block size, context length, scheduler limits, and -prefix-cache policy against the resolved runtime configuration and the target -workload. These settings are explicit simulation inputs, not reconstructed -automatically from the FPM table. - -```yaml -engine: - mode: aggregated - model: MiniMaxAI/MiniMax-M2.7 - hardware: h200_sxm - backend: vllm - backend_version: "REPLACE_WITH_ACTUAL_COLLECTED_VLLM_VERSION" - context_length: 8192 - workers: - aggregated: - parallelism: - replicas: 1 - tensor: 4 - pipeline: 1 - attention_data: 1 - moe_tensor: 4 - moe_expert: 1 - scheduler: - max_batched_tokens: 8192 - max_sequences: 256 - kv_cache: - block_size: 64 - prefix_caching: false - capacity: {type: default} - timing: {type: default, forward_model: fpm} -traffic: - source: {type: synthetic, input_tokens: 1024, output_tokens: 32} - load: {type: concurrency, concurrency: 4} - stop: {requests: 8} -``` - -The unified CLI currently has no `--systems-paths` option. The following -inline Python example sets existing Python and Rust root-selection APIs to the -same directory, then invokes `predict` in that process. It needs no additional -helper file and does not add a new CLI flag: - -```bash -AIC_ALLOW_UNLISTED_VERSIONS=1 \ -python3 - "$FPM_RUN/systems" "$FPM_RUN/predict.yaml" \ - "$FPM_RUN/prediction-fpm" <<'PY' -import os -import runpy -import sys -from pathlib import Path - -root = str(Path(sys.argv[1]).resolve(strict=True)) -config = str(Path(sys.argv[2]).resolve(strict=True)) -output = str(Path(sys.argv[3]).resolve()) -os.environ["AICONFIGURATOR_SYSTEMS_PATH"] = root - -from aisimulate_core.sdk.perf_database import set_systems_paths - -set_systems_paths([root]) -print(f"FPM systems_root={root}", file=sys.stderr) -sys.argv = [ - "aisimulate", "predict", "--stack", "engine", - "--config", config, "--output-dir", output, "--capture-per-request", -] -runpy.run_module("aisimulate", run_name="__main__") -PY -``` - -**Expected:** exit code 0 and a new `prediction-fpm/` directory containing: - -- `prediction.json`: the durable prediction report; check the completion count - against the eight requested completions and inspect the reported metrics. -- `requests.jsonl`: per-request records enabled by `--capture-per-request`. - -Keep the input YAML, printed external root, package/source revision, and pair -hash with the report so its data source can be reproduced. Use a new output -directory for another prediction; the CLI rejects nonempty output directories -unless overwrite is explicitly requested. - -Every simulated step needs a matching identity and supported coordinates. A -missing cell fails rather than silently switching to `op_level`. With default -capacity estimation, FPM also caps KV capacity at the collected decode-KV -ceiling. Inspect that cap when explaining concurrency and throughput. - -This example covers `predict --stack engine` in one process. It does not cover -propagating Python root overrides to multiprocessing `recommend` workers. The -unified YAML also has no explicit per-operation quantization overrides; use the -SDK or compatibility CLI for those cells. See the -[migration limitations](../../../../docs/cli/migrate-from-aiconfigurator.md). - -## 8. Recovery and cleanup - -Use the same frozen arguments and roots to resume: - -```bash -python3 -m collector.fpm_forward "${fpm_args[@]}" --resume -``` - -Before retrying failed cells, copy their existing raw artifacts and logs to a -separate archive. The runner removes old raw/log directories for cells it -re-executes. After preserving that evidence: - -```bash -python3 -m collector.fpm_forward "${fpm_args[@]}" --resume --resume-retry-failed -``` - -A completed, validated schema-v6 publication is terminal on resume even if raw -artifacts were pruned. Changing the model, image, sampling policy, or source can -change the frozen plan. Use fresh campaign, checkpoint, artifact, and database -roots for independent recollection or A/B comparisons. Do not run concurrent -campaigns with colliding workload names. - -The collector normally cleans up each cell's workload in its finalization path; -cleanup failure is an error. After an abnormal interruption, preserve available -logs and inspect the manifests recorded under this campaign's cell artifacts. -For each owned manifest that still has resources, run: - -```bash -kubectl delete -f /absolute/path/to/this/cell/k8s_deploy.yaml \ - --ignore-not-found --cascade=foreground --wait=true --timeout=180s -kubectl get -f /absolute/path/to/this/cell/k8s_deploy.yaml -``` - -**Expected:** the resources are absent. Also check child Pods using the actual -labels/owner references in that manifest. -If deletion times out, cleanup is not complete; inspect and remove the remaining -owned resources before ending the campaign. -Do not delete unrelated workloads or the namespace. Default `/results` storage -is Pod-local `emptyDir`, so deleting Pods before salvaging results loses those -files. - -## 9. Validate accuracy separately - -Passing the preceding stages establishes that the collection and prediction -path works. Accuracy requires independent measurements under matched conditions. - -1. Freeze a validation workload with held-out shapes or traces. Record model and - runtime versions, GPU/system, parallelism, quantization, CUDA Graph policy, - context/KV limits, prefix policy, input/output lengths, and arrival or - concurrency settings. Preserve the raw measurements. -2. For forward-level validation, compare predicted and measured latency for the - same iteration coordinates, phase, rank aggregation, and timing boundary. - Report prefill and decode separately. A collected-point lookup is a data-path - check, not an independent accuracy test. -3. For request-level validation, compare AIS predictions against a real serving - benchmark with the same traffic. Define TTFT, ITL/TPOT, throughput units, and - aggregation identically. Include ragged batches, mixed/chunked execution, - and higher concurrency when they are part of the intended deployment. -4. Report query coverage and failed/out-of-domain cases alongside errors. For - positive measured latencies `m_i` and predictions `p_i`, per-point absolute - percentage error is `100 * abs(p_i - m_i) / m_i`; MAPE averages those errors. - WAPE is `100 * sum(abs(p_i - m_i)) / sum(m_i)` and weights larger latencies - more heavily. Select acceptance thresholds before inspecting the results. - -**Expected:** a reproducible comparison with explicit conditions, coverage, -errors, and pass/fail criteria. FPM forward error alone does not establish -TTFT/TPOT or throughput accuracy. The -[E2E Accuracy Overview](https://ai-dynamo.org/aisimulate/e2e-accuracy/) reports existing -accuracy results; it is not a substitute for validating a new cell. See the -[snapshot and regeneration details](../../../../pages/e2e-accuracy/README.md) for how that evidence -is produced. - -## FPM data and prediction troubleshooting - -| Symptom | Check and next action | -| --- | --- | -| Plan fails to resolve a model | Supply a real local config with `--fpm-model-config`; preserve the canonical model identity | -| Smoke succeeds but no Parquet appears | Expected for smoke; inspect both phase results before formal collection | -| Partial publication or skipped existing cells | Inspect `missing_cells` and `skipped_first_publisher_wins`; use a new database root for independent recollection | -| Pair hash/schema validation fails | Preserve both files, recover a matching published pair, and do not edit the sidecar to conceal a mismatch | -| FPM lookup fails | Check root layout, actual backend version, full identity, and supported coordinates; some queries also require supported SOL operations | -| SDK query works but replay fails | Inspect the later simulated shapes and decode-KV ceiling; one successful prefill query does not establish replay coverage | -| Prediction output directory is not empty | Select a new directory to retain the previous report | diff --git a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md new file mode 100644 index 0000000000..4e27719a55 --- /dev/null +++ b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md @@ -0,0 +1,1098 @@ + + +# Guide: self-benchmark and onboard a forward-pass performance model + +This guide explains how to collect engine forward-pass measurements, turn them +into an FPM data profile, and use that profile in AISimulate. **Self-collection** +means running the engine's self-benchmark on your target configuration. +**Onboarding** means validating and loading its measured profile through the +canonical performance-model API, then using it for prediction. + +The procedure has seven steps. The model, hardware, and engine settings are +inputs to that procedure. [Kimi K3 TP8+DCP8](#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile) +and [MiniMax-M2.7 TP4](#example-b-collect-a-minimax-m27-tp4-profile) are worked +examples at the end; their values are not defaults for other deployments. + +## When self-collection is useful + +Use self-collection when you need iteration latency for a specific engine +configuration that the existing performance data does not represent well: + +- A custom engine build, newer kernel, quantization mode, or CUDA Graph policy + changes forward-pass execution time. +- Your GPU/topology or workload shapes lack measured coverage. +- You need the combined effect of operations, launch overhead, and communication + within the measured engine iteration, rather than composing their costs from + separate operation tables. +- You want to calibrate a supported configuration against the runtime you will + deploy, while keeping the rest of AISimulate's serving simulation unchanged. + +For example, collecting Kimi TP8+DCP8 measures that configuration's prefill and +decode costs, including DCP communication inside the timing boundary. A TP8 +profile without DCP is a different measurement identity and cannot substitute +for those observations. If suitable measured data already exists, reuse it +starting at step 4 after checking its provenance and coverage. + +## What self-collection covers + +FPM onboarding supplies **customized engine step times** for the measured +configuration and covered workload shapes. AISimulate Replay uses these timings +with its existing scheduling, routing, prefill–decode transfer, and KV-cache +models. Special or customized mechanisms that differ from supported Replay +behavior, such as layer-wise KV transfer or a custom scheduling/overlap policy, +require corresponding Replay-side changes; collecting new step times alone does +not implement those mechanisms. + +The current collection workflow requires **PP=1**. Validate TTFT, ITL, and +end-to-end throughput separately against the target serving configuration; +matching engine step times alone does not establish their accuracy. + +## Current support and architecture readiness + +| Backend | Self-benchmark collection status in this workflow | +| --- | --- | +| vLLM | Available for supported configurations; validate the exact model, engine build, and benchmark path before a campaign | +| SGLang | Coming soon; not supported by this self-benchmark collection workflow today | +| TensorRT-LLM | Coming soon; not supported by this self-benchmark collection workflow today | + +This table describes **self-benchmark collection**, not AISimulate's broader +operation-level modeling support. The dedicated Collector currently uses +`PP=CP=1`; its topology presets do not enumerate DCP. The Kimi DCP8 example +imports a profile produced by an adapted vLLM/Dynamo self-benchmark path. + +**Self-benchmarking does not automatically support every model or architecture.** +An engine being able to serve a model does not prove the benchmark can seed its +state, schedule the requested work, and time it correctly. New architectures or +engine versions can require changes in: + +- **Benchmark input and state preparation:** attention KV warmup, recurrent or + convolution state, per-request state allocation, and distributed cache layout. +- **Scheduling and measurement hooks:** requested versus actually scheduled + token counts, graph paths, completion checks, and per-rank timing aggregation. +- **AISimulate model construction:** architecture/configuration parsing, model + registration, resident weights, and memory/cache descriptions used by the + consumer. A collection plan's bootstrap template is not proof of consumer support. + +Use a small smoke run to qualify those paths before collecting the full surface. +If it fails or silently schedules different work, fix and revalidate the +benchmark adapter before using the measurements. For an architecture absent +from AISimulate, follow [How to add a new model](../add_a_new_model.md) as well. +An already registered architecture needs no new class or parallelization mapping +merely to add another supported measured configuration. + +## Workflow at a glance + +| Step | Action | Required result | +| --- | --- | --- | +| [1](#1-check-support-and-prepare-the-environment) | Check backend/model support and prepare the environment | Compatible engine, benchmark, and AISimulate builds; known adaptation gaps | +| [2](#2-freeze-the-engine-configuration-and-measurement-plan) | Freeze configuration, timing semantics, and sampling coverage | Reproducible configuration and a plan covering the intended workloads | +| [3](#3-smoke-test-and-collect) | Smoke test, adapt if needed, then collect | Validated prefill/decode measurements plus raw evidence and provenance | +| [4](#4-validate-and-install-the-fpm-profile) | Validate the data contract and install the pair | A schema-v6 Parquet/metadata pair in a local systems root | +| [5](#5-construct-and-check-the-performance-model) | Construct the canonical model and check queries | Matching identity, ready estimator, and verified measured-point lookups | +| [6](#6-connect-the-model-to-replay) | Supply the same profile plus serving configuration | A reproducible prediction configuration and successful Replay smoke | +| [7](#7-validate-coverage-and-accuracy) | Check coverage and compare held-out measurements | Errors and limitations reported separately from calibration checks | + +With a validated profile from someone else, review steps 1–2 and start execution +at step 4. New GPU collection follows all seven steps. A failure in runtime or +simulator support is an adaptation task; additional sampling alone will not +resolve it. + + + +## 1. Check support and prepare the environment + +Use Python 3.11–3.13 and install AISimulate from the latest `main`. +For an existing checkout, run `git switch main` and start from the pull command: + +```bash +git clone --branch main https://github.com/ai-dynamo/aisimulate.git +cd aisimulate +git pull --ff-only origin main +uv sync --project python/aisimulate --extra dev +source python/aisimulate/.venv/bin/activate +``` + +Source installation requires `uv`, Cargo/Rust, and a C/C++ compiler/linker. +See the [installation guide](../../../../docs/installation.md#use-current-source) +for platform requirements. No serving GPU or model weights are needed to query +an existing profile or run the engine-stack simulation. + +From the repository root, record the source and package in the activated session: + +```bash +export AIS_REPO="$PWD" +git rev-parse HEAD +python -c 'from importlib.metadata import version; print(version("aisimulate"))' +aisimulate predict --help +``` + +Use the same activated Bash or Zsh session for the remaining commands. +For existing data, review its recorded identity and conditions in step 2, then +continue at step 4. For new measurements, continue through steps 2 and 3. + +For **new GPU collection**, additionally check: + +```bash +cd "$AIS_REPO/python/aisimulate" +python -m collector.fpm_forward --help +``` + +The collection host needs model configuration metadata and `kubectl` access. +The target runtime needs the model weights, allocated GPUs, and matching engine +and benchmark code. The Kubernetes Collector additionally needs an existing +namespace, checkpoint PVC/path, collection image, and Generator target release. +Record the image digest, model revision, GPU type, and resolved configuration. +Before launching, read [recovery and cleanup](#recovery-and-cleanup). + +## 2. Freeze the engine configuration and measurement plan + +Define the target deployment before measuring it. Keep a record of: + +- **Identity:** model/config/checkpoint revision, GPU and interconnect topology, + engine version and image/source revision, TP/PP/attention-DP/MoE TP/EP, and DCP + when used. Record actual GPU count separately from logical parallel dimensions. +- **Execution policy:** weight and KV precision, attention/MoE backend settings, + CUDA Graph capture policy, speculation, model context limit, token budget, + maximum active requests, and prefix-cache policy. +- **Measurement contract:** which engine work is timed, units, rank aggregation, + warmup and repeat policy, and the source/initialization of every cache or + recurrent state consumed by a measurement. +- **Coverage:** prefill request/new-token/prefix coordinates and decode + batch/context coordinates needed by the intended Replay workload. + +The balanced profile uses iteration totals: prefill coordinates are +`(batch_size, total_prefill_tokens, total_kv_read_tokens)`; decode coordinates +are `(batch_size, total_kv_read_tokens)`. A rectangular min/max range does not +prove that all interior queries can be interpolated. Include graph boundaries +and the intended batch/context range in the plan. + +For an existing profile, compare these requirements with its manifest and +resolved engine configuration. Keep unknown fields explicit rather than +inventing values or changing the recorded identity to obtain a match. + +**Example — Kimi:** TP8+DCP8 uses eight GPUs, `max_num_seqs=32`, and a custom +vLLM build. Its data covers that measured engine configuration. It neither +provides a TP8-only profile nor qualifies PP>1 or arbitrary mixed/ragged shapes. +**Example — MiniMax:** [B1–B2](#b1-prepare-the-campaign-step-2) show configuration +and plan generation for a new four-H200 campaign. + +**Result:** a frozen plan/configuration with known state-seeding and coverage +assumptions. Use a separate data identity or campaign when changing those conditions. + +## 3. Smoke test and collect + +First execute a small prefill and decode measurement for each intended +configuration. Check that the engine initializes, the requested work is actually +scheduled, timings are finite and use the intended boundary, and all required +ranks participate. Inspect whether attention KV and any recurrent state are +warmed, synthesized, or skipped. A successful launch or prefill alone does not +qualify decode. + +If the architecture needs benchmark changes, implement them and repeat the +smoke before running the sweep. In particular, a token-based attention capacity +limit may not account for per-request linear/recurrent cache state. State +initialization can change performance; document and validate synthetic-state +measurements instead of treating them as equivalent to real request history. + +After smoke passes, run the full planned surface. Save raw FPMs, actual scheduled +shapes, effective engine configuration, versions/hashes, repeat counts, and +failure/skip reasons. Keep incomplete or alternate-state results distinguishable +from the primary profile. Validate that the intended phases and cells were +published; exit code 0 by itself is not evidence that old cells were replaced. + +The [MiniMax collection example](#b3-run-smoke-then-formal-collection-step-3) +shows the Collector commands and publication checks. Use +[recovery and cleanup](#recovery-and-cleanup) for retries and owned resources. + +**Result:** measurements and provenance that satisfy the frozen plan, with +missing or unsupported regions reported explicitly. + + + +## 4. Validate and install the FPM profile + +The runtime consumes a pair of files, not raw self-benchmark JSON: + +```text +/ + .yaml + query_versions.yaml + attention_lane_defaults.yaml + data//// + fpm_forward_perf.parquet + fpm_forward_perf.metadata.json +``` + +Use the [whole-forward publication contract](../../collector/README.md#whole-forward-fpm-campaign) +to produce schema `aic_fpm_forward_perf`, version 6. Collector performs this +conversion for its campaigns. Importing raw output from another benchmark path +requires an explicit conversion preserving the recorded workload and identity. +The `systems_paths` configuration points to ``, not its `data/` +subdirectory. Whole-forward profiles use the layout above; operation tables may +have an additional family directory. + +Check the file hash, schema, row count, phase coverage, identity fields, units, +and measurement policy. A recorded DCP value must be positive and divide TP; +missing DCP and explicit DCP1 remain different identities. For per-row repeated +measurements, the sidecar policy and each row's policy/repeat count must agree. +Keep the producer's aggregated latency; importing the pair does not train or +recompute it. + +Keep the hardware/query-policy files and source manifest with the pair. Preserve +both files together when copying or renaming them to the canonical filenames. +A dataset may contain comparator profiles or nonuniform workloads: choose the +appropriate primary profile instead of merging every artifact into one table. + +[Example A1](#a1-import-the-measured-pair-step-4) downloads our published DCP8 +profile at a pinned HF revision and checks its hashes. [Example B4](#b4-inspect-the-published-pair-step-4) +checks a newly collected pair. Publishing a pair does not automatically bundle +it into an AISimulate release: use an external systems root or submit it through +the repository's data-publication process. + +**Result:** a validated local profile and its provenance. No separate registry +service or regression-training command is needed for measured FPM interpolation. + + + +## 5. Construct and check the performance model + +Use `RustForwardPassPerfModel.best_available(ForwardPassPerfModelConfig(...))`. +Supply the exact model, system, backend/version, worker role, topology, precision, +and backend identity of the profile, together with its `systems_paths` root. +Use `estimation_mode="fpm_interpolation"` and `fallback_policy="deny"` to +select measured FPM data explicitly. The general defaults are `auto` + `deny`; +auto tries op-level, FPM interpolation, then regression at construction. + +Check `diagnostics()` for readiness, the selected estimator, resolved identity, +and systems root. Query a known measured prefill point and a known measured +decode point with `estimate_forward_pass_time_ms`. FPM inputs describe scheduled +iteration totals; they are not request TTFT/ITL. With correction disabled, verify +that measured-point queries reproduce their stored latency before testing +interpolation. Save the resolved provenance with your results. + +For profiles that require them: + +- `dcp` records decode context parallelism within the TP group; it does not + multiply physical GPU count. Current DCP timing uses measured vLLM FPM; + SOL-dependent transfer paths are unsupported. +- `estimator_config.fpm_interpolation.text_only` permits language timing for + a multimodal architecture while retaining encoder weights. It supplies no + vision execution timing. +- `unrecorded_quant_modes` can select null FMHA/communication identity fields. + It is an exact null match, not a wildcard; leave the corresponding explicit + quantization overrides unset. + +A missing cell or unsupported query is a coverage/support error. Do not change +its topology, precision, or version labels to make it load. The selected model +does not switch estimators on a query miss, and untrained regression cannot run +offline prediction. + +[Example A2](#a2-query-the-canonical-model-step-5) checks both phases of the +collected Kimi profile. [Example B5](#b5-query-the-canonical-model-step-5) shows +the same API for MiniMax. See the [Core API](../../../../docs/core-api.md#choosing-a-forward-pass-api) +for the complete contract. + +**Result:** a ready estimator with the intended data identity and checked queries. + + + +## 6. Connect the model to Replay + +Configure `aisimulate predict` with the same data root, estimator selection, +identity, and engine build assumptions. Then supply the serving behavior that +an FPM table does not define: worker layout, scheduler token/request limits, +context limit, block geometry and capacity, prefix-cache policy, traffic, and +any transfer/offload settings. + +Relevant YAML fields include: + +| Purpose | Configuration | +| --- | --- | +| Local profile and selection | `engine.systems_paths`, `engine.estimation_mode`, `engine.fallback_policy` | +| Estimator-specific controls | `engine.estimator_config`, or per-role `timing.estimator_config` | +| Recorded DCP | `engine.workers..parallelism.decode_context` | +| Precision and attention implementation | Per-role `timing.gemm_quant_mode`, `moe_quant_mode`, `fmha_quant_mode`, `kvcache_quant_mode`, `comm_quant_mode`, `attention_backend` | +| Scheduling and KV state | Per-role `scheduler` and `kv_cache` | + +Start with a small in-domain workload and inspect completed requests, generated +tokens, estimator provenance, and GPU counts. Expand it only after validating +query coverage. New FPM measurements do not remove unsupported scheduler or +memory behavior. DCP currently requires explicit fixed KV capacity and explicit +bytes per token if using offload/P-D transfer. The generic cache does not model +KDA checkpoint/eviction behavior or native hybrid prefill chunk alignment. + +For `decode_context`, use **`--stack engine`**. Dynamo Replay and Planner do not +yet support this field; their DCP integration requires a separate downstream update. + +Prediction accepts the timing identity overrides above for regular language +workers with default timing. Recommendation currently rejects those overrides, +and DCP is not a recommendation search dimension. [Example A3](#a3-run-and-check-replay-steps-67) +and [Example B6](#b6-run-replay-step-6) provide complete prediction inputs. + +**Result:** a reproducible Replay configuration and a successful smoke run; +accuracy and unsupported regions are checked in step 7. + + + +## 7. Validate coverage and accuracy + +First verify exact measured-point queries, then test the interpolation and +Replay shapes required by your workload. Report unsupported queries as coverage +gaps. Successfully loading a table and replaying its calibration points does not +establish prediction accuracy; use independent measurements under matched conditions. + +1. Freeze a validation workload with held-out shapes or traces. Record model and + runtime versions, GPU/system, parallelism, quantization, CUDA Graph policy, + context/KV limits, prefix policy, input/output lengths, and arrival or + concurrency settings. Preserve the raw measurements. +2. For forward-level validation, compare predicted and measured latency for the + same iteration coordinates, phase, rank aggregation, and timing boundary. + Report prefill and decode separately. A collected-point lookup is a data-path + check, not an independent accuracy test. +3. For request-level validation, compare AIS predictions against a real serving + benchmark with the same traffic. Define TTFT, ITL/TPOT, throughput units, and + aggregation identically. Include ragged batches, mixed/chunked execution, + and higher concurrency when they are part of the intended deployment. +4. Report query coverage and failed/out-of-domain cases alongside errors. For + positive measured latencies `m_i` and predictions `p_i`, per-point absolute + percentage error is `100 * abs(p_i - m_i) / m_i`; MAPE averages those errors. + WAPE is `100 * sum(abs(p_i - m_i)) / sum(m_i)` and weights larger latencies + more heavily. Select acceptance thresholds before inspecting the results. + +**Result:** a reproducible comparison with explicit conditions, coverage, +errors, and pass/fail criteria. FPM forward error alone does not establish +TTFT/TPOT or throughput accuracy. The +[E2E Accuracy Overview](https://ai-dynamo.org/aisimulate/e2e-accuracy/) reports existing +accuracy results; it is not a substitute for validating a new cell. See the +[snapshot and regeneration details](../../../../pages/e2e-accuracy/README.md) for how that evidence +is produced. + +Retain known limitations with the profile so subsequent users can decide +whether it fits their workload. + + + +## Example A: onboard the collected Kimi K3 TP8+DCP8 profile + +This worked example illustrates steps 4–7 using a profile already collected +through Dynamo self-benchmarking. Complete step 1 first and review the retained +collection configuration as described in step 2. It runs on CPU and downloads +the published +[Kimi K3 profile](https://huggingface.co/datasets/nvidia/aisimulate-fpm-dataset/tree/6fad3f9a0df5a24603108dcea0d201259254b904/data/moonshotai--Kimi-K3/gb300/vllm/0.29.0/tp8-dcp8), +checks one measured point from each phase, and completes a small Replay run. +It models **eight GB300 GPUs**: DCP8 reuses the TP8 group. + +### A1. Import the measured pair (step 4) + +Choose a new output directory. The public download uses Python's standard +library; no Hugging Face login or additional download package is required. + +```bash +export FPM_RUN="$PWD/kimi-k3-fpm" +mkdir "$FPM_RUN" +python - <<'PY' +import hashlib +import json +import os +import shutil +from importlib.resources import files +from pathlib import Path +from urllib.request import urlopen + +run = Path(os.environ["FPM_RUN"]).resolve() +revision = "6fad3f9a0df5a24603108dcea0d201259254b904" +base = f"https://huggingface.co/datasets/nvidia/aisimulate-fpm-dataset/resolve/{revision}/" +config_path = "data/moonshotai--Kimi-K3/gb300/vllm/0.29.0/tp8-dcp8" + +def download(path): + with urlopen(base + path, timeout=60) as response: + return response.read() + +manifest_bytes = download(f"{config_path}/manifest.json") +manifest = json.loads(manifest_bytes) +primary = next(item for item in manifest["fpm"] if item["role"] == "primary") +parquet = download(primary["path"]) +sidecar = download(primary["metadata_path"]) +metadata = json.loads(sidecar) +assert hashlib.sha256(parquet).hexdigest() == primary["sha256"] == metadata["parquet_sha256"] +assert metadata["schema_name"] == "aic_fpm_forward_perf" and metadata["schema_version"] == 6 +assert primary["row_count"] == metadata["row_count"] == 669 +assert manifest["dcp"] == metadata["configuration_selector"]["dcp"] == 8 + +systems = run / "systems" +target = systems / "data/gb300/vllm/0.29.0" +target.mkdir(parents=True) +(target / "fpm_forward_perf.parquet").write_bytes(parquet) +(target / "fpm_forward_perf.metadata.json").write_bytes(sidecar) +packaged = files("aisimulate_core") / "systems" +for name in ("gb300.yaml", "query_versions.yaml", "attention_lane_defaults.yaml"): + shutil.copyfile(str(packaged / name), systems / name) +(run / "manifest.json").write_bytes(manifest_bytes) +(run / "engine-config.json").write_bytes(download(f"{config_path}/fpm/provenance/engine-config.json")) +(run / "dataset-revision.txt").write_text(revision + "\n") +print(f"Installed {primary['row_count']} primary rows under {systems}") +PY +``` + +**Expected:** 669 primary rows: 338 balanced prefill points and 331 decode +points. Only the filenames change to the loader's canonical names; the Parquet +and sidecar contents remain unchanged. Retain the manifest, dataset revision, +and engine configuration beside the imported data. The separate comparator +with synthetic attention KV and the nonuniform prefill records are not part of +this primary profile. + +The decode measurements use warmed MLA KV and **random KDA state**. They are +calibration data, not held-out accuracy evidence. See the retained +`engine-config.json` for measurement conditions and known unknowns. + +This profile uses a custom vLLM `0.29.0` build. Allow its literal version to be +queried outside the copied version-slot policy for this session: + +```bash +export AIC_ALLOW_UNLISTED_VERSIONS=1 +``` + +This permits version selection; model/topology/precision matching stays strict. +For another profile whose backend version is already queryable, omit the override. + +### A2. Query the canonical model (step 5) + +Run this in the same session. The script compares its predictions with the +stored measurements and records the model's resolved configuration. + +```bash +python - <<'PY' +import json +import math +import os +from pathlib import Path + +import pyarrow.parquet as pq +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + +run = Path(os.environ["FPM_RUN"]).resolve() +config = ForwardPassPerfModelConfig( + model="moonshotai/Kimi-K3", system="gb300", backend="vllm", + backend_version="0.29.0", worker_type="aggregated", + tp=8, pp=1, attention_dp=1, moe_tp_size=8, moe_ep_size=1, dcp=8, + gemm_quant_mode="bfloat16", moe_quant_mode="w4a16_mxfp4", + kvcache_quant_mode="fp8", attention_backend="FLASHINFER_MLA", + estimation_mode="fpm_interpolation", fallback_policy="deny", + systems_paths=(str(run / "systems"),), + estimator_config={ + "fpm_interpolation": {"text_only": True, "unrecorded_quant_modes": ["fmha", "comm"]}, + "correction": {"enabled": False}, + }, +) +rows = pq.read_table(run / "systems/data/gb300/vllm/0.29.0/fpm_forward_perf.parquet").to_pylist() +assert len(rows) == 669 +model = RustForwardPassPerfModel.best_available(config) +try: + for phase, scheduled in ( + ("prefill", {"num_prefill_requests": 1, "sum_prefill_tokens": 128, "sum_prefill_kv_tokens": 0}), + ("decode", {"num_decode_requests": 1, "sum_decode_kv_tokens": 128}), + ): + row = next(row for row in rows if row["workload_kind"] == phase and row["batch_size"] == 1 + and row["total_prefill_tokens"] == (128 if phase == "prefill" else 0) + and row["total_kv_read_tokens"] == (0 if phase == "prefill" else 128)) + latency = model.estimate_forward_pass_time_ms({"version": 1, "scheduled_requests": scheduled}) + assert latency is not None and math.isclose(latency, row["latency_ms"], rel_tol=1e-9) + print(json.dumps({"phase": phase, "predicted_ms": latency, "measured_ms": row["latency_ms"]})) + diagnostics = model.diagnostics() + assert diagnostics["readiness"] == "ready" + (run / "estimator-provenance.json").write_text(json.dumps(diagnostics, indent=2) + "\n") +finally: + model.close() +PY +``` + +**Expected:** approximately 66.1715 ms for the selected prefill point and +10.5622 ms for the selected decode point. These are forward-pass latencies; +request TTFT and end-to-end latency also depend on scheduling and traffic. + +`text_only` enables the language profile for Kimi's multimodal architecture +while retaining resident encoder weights. `unrecorded_quant_modes` matches +only null FMHA/communication identities in this dataset. Leave the corresponding +top-level quant fields unset. Missing DCP, DCP1, and DCP8 are distinct identities; +an explicit DCP8 request never substitutes another configuration. + +### A3. Run and check Replay (steps 6–7) + +Create the complete prediction YAML. The unquoted heredoc below expands +`FPM_RUN` to the absolute systems path; YAML itself does not expand shell variables. + +```bash +cat > "$FPM_RUN/predict.yaml" < + +### B1. Prepare the campaign (step 2) + +This example uses an external directory to keep its results separate +from bundled databases and previous campaigns. The directory name and layout +are choices for this walkthrough, not mandatory FPM setup steps. The later +commands use these paths consistently; replace the absolute path below with a +new persistent directory. + +```bash +export FPM_RUN=/absolute/new/path/m27-h200-tp4 +mkdir -p "$(dirname "$FPM_RUN")" +mkdir "$FPM_RUN" +mkdir -p "$FPM_RUN/systems/data" +cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/h200_sxm.yaml" \ + "$FPM_RUN/systems/" +cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/query_versions.yaml" \ + "$FPM_RUN/systems/" +cp "$AIS_REPO/python/aisimulate/src/aisimulate_core/systems/attention_lane_defaults.yaml" \ + "$FPM_RUN/systems/" +``` + +`h200_sxm.yaml` provides hardware specifications for the later SDK/prediction +steps when they load this external systems root. `query_versions.yaml` preserves +the version-slot policy, and `attention_lane_defaults.yaml` preserves backend +lane defaults. These files describe hardware and query policy; they contain no +measured FPM rows. The collector's +`--fpm-database-root` override is optional; the later examples select it +explicitly so they publish and consume data from the same directory. + +Save the following as `$FPM_RUN/collector-k8s.yaml`, replacing every +`REPLACE_...` value with the corresponding existing collection setting: + +```yaml +K8sConfig: + k8s_namespace: REPLACE_WITH_EXISTING_NAMESPACE + k8s_image: REPLACE_WITH_EXISTING_COLLECTION_IMAGE_DIGEST + k8s_pvc_name: REPLACE_WITH_MODEL_PVC + k8s_pvc_mount_path: /models + k8s_model_path_in_pvc: REPLACE_WITH_RELATIVE_CHECKPOINT_DIRECTORY + k8s_hf_home: /tmp/fpm-hf-cache +``` + +The dedicated collector accepts deployment-only `K8sConfig` fields. In +particular, do not copy `Workers.agg.extra_cli_args` or `ServiceConfig` from a +standalone Generator example into this file. The accepted fields are defined in +[the collector input resolver](../../collector/fpm_forward/entry.py). + + + +### B2. Freeze and inspect the plan (step 2) + +Use the Generator target release from the existing collection configuration +and define the campaign arguments once: + +```bash +export AIS_DYNAMO_RELEASE=REPLACE_WITH_MATCHING_DYNAMO_RELEASE +cd "$AIS_REPO/python/aisimulate" +fpm_args=( + --backend vllm + --model-path MiniMaxAI/MiniMax-M2.7 + --gpu h200_sxm + --fpm-max-gpus 4 + --fpm-gpu-counts 4 + --fpm-parallel-presets pure_tp + --fpm-kv-cache-dtypes auto + --fpm-max-prefill-isl 8192 + --fpm-max-prefill-cudagraph-size 2048 + --dynamo-version "$AIS_DYNAMO_RELEASE" + --generator-config "$FPM_RUN/collector-k8s.yaml" + --checkpoint-dir "$FPM_RUN/checkpoints" + --fpm-artifact-root "$FPM_RUN/artifacts" + --fpm-database-root "$FPM_RUN/systems/data" +) +python3 -m collector.fpm_forward "${fpm_args[@]}" \ + --plan-only > "$FPM_RUN/plan.json" +``` + +If the checkpoint exists only inside the target Pod and its metadata cannot be +resolved on the machine running Collector, add `--fpm-model-config /absolute/config.json` +to `fpm_args`. Keep the canonical model ID in `--model-path`. + +**Expected:** exit code 0 and JSON describing the frozen plan, without creating +Kubernetes workloads. Inspect the model/config hash, system, topology, +quantization, phase cells, sampling policy, and admission decisions before +continuing. For this single-topology, single-KV-setting example, expect one +prefill cell and one decode cell if both are admitted. The plan is authoritative; +do not proceed with an empty plan or unexplained omissions. +`counts.cells` should be 2 for that case, while `counts.points` is +`"runtime-determined"`: runtime initialization determines the actual point set. + +The sampling settings have distinct meanings: + +- `--fpm-max-prefill-isl 8192` bounds the scheduled prefill **new-token** axis. + It does not set runtime context length. +- `--fpm-max-prefill-cudagraph-size 2048` is the prefill capture-size policy. + Align it with the deployment being modeled before freezing a formal campaign. +- The collector currently sets runtime `max_model_len=-1` for vLLM auto-fit and + has no CLI override for that setting. Record the resolved value with the + collection results. +- The default is five global warm-up iterations and one measurement per point. + Published `latency_ms` uses the maximum rank wall time, converted to ms; + it is not a mean or median over repeated measurements. + +The complete option list is available from `python3 -m collector.fpm_forward +--help` and [its argument definitions](../../collector/fpm_forward/config.py). + +#### Include multiple GPU counts and parallel configurations + +Both GPU counts and parallel presets accept multiple values. To expand this +MiniMax campaign, replace these three entries in `fpm_args` before generating +the plan: + +```text +--fpm-max-gpus 8 +--fpm-gpu-counts 2 4 8 +--fpm-parallel-presets pure_tp tep dep +``` + +For each selected GPU count `N`, the MoE presets request these configurations +(`PP=CP=1`): + +| Preset | Attention TP | Attention DP | MoE TP | MoE EP | +| --- | --- | --- | --- | --- | +| `pure_tp` | N | 1 | N | 1 | +| `tep` | N | 1 | 1 | N | +| `dep` | 1 | N | 1 | N | + +The planner combines the selected counts and presets, deduplicates identical +topologies, and applies model/backend and memory admission. It then creates +prefill and decode cells for admitted topology, KV dtype, and backend-policy +combinations. Read `plan.json` for the actual configurations and cell count; +do not assume every requested combination is admitted. Dense models use the +`tp` preset; MoE `pure_tp` requires support in the model's capability profile. +See [topology enumeration](../../collector/fpm_forward/topology.py). + +`--fpm-max-gpus 8` is the maximum GPU count for one cell. The runner executes +cells sequentially; multiple Pods belonging to one cell execute together. +Selecting multiple counts/presets does not launch all configurations +simultaneously. + +Regenerate and inspect `plan.json` after changing these options. If the original +campaign has already run, first prepare a fresh campaign directory as in B1 +and redefine `fpm_args` with the new paths. The B5–B6 TP4 SDK/predict examples +still select only their matching cell from the larger database; they do not +automatically compare all collected configurations. + + + +### B3. Run smoke, then formal collection (step 3) + +First read the number of cells from the inspected plan and execute the minimal +profile for every cell. This works for both the two-cell TP4 example and an +expanded campaign: + +```bash +FPM_CELL_COUNT=$(python3 -c \ + 'import json, sys; n = len(json.load(open(sys.argv[1]))["cells"]); assert n > 0; print(n)' \ + "$FPM_RUN/plan.json") +python3 -m collector.fpm_forward "${fpm_args[@]}" --smoke --limit "$FPM_CELL_COUNT" +``` + +**Expected:** exit code 0, all selected smoke cells pass, raw runtime evidence is saved, +and no formal Parquet pair is published. `--smoke` without `--limit` runs only +the first cell. `--limit 2` is sufficient for the original two-cell example, +but in an expanded plan the first two cells can both be prefill: the planner +lists all prefill cells before decode cells. Use the full plan count to check +both phases of every admitted configuration. +In `checkpoints/fpm_forward_smoke.json`, check `smoke.status == "passed"`, +`smoke.cell_count == FPM_CELL_COUNT`, and `smoke.formal_database_written == false`. + +Only after the smoke command succeeds and its evidence is checked, run: + +```bash +python3 -m collector.fpm_forward "${fpm_args[@]}" +``` + +Formal collection must omit both `--smoke` and `--limit`. Smoke and formal +checkpoints/artifacts are isolated automatically, so the roots above can be +shared; smoke samples are not reused as formal measurements. + +**Expected:** all frozen cells pass and the formal pair is published. Check the +formal checkpoint's database status, that `missing_cells` is empty, and that +`published_cells` covers the planned cells. In this fresh database, +`skipped_first_publisher_wins` should be empty. Exit code 0 alone is insufficient +to prove a new attempt replaced existing data: publication preserves a cell +already owned by another attempt. `--fpm-publish-partial` is an explicit opt-in +for incomplete campaigns and still returns a nonzero result. + +For this fresh, two-cell campaign, the formal checkpoint should contain: + +```json +{ + "database": { + "status": "passed", + "published_cells": 2, + "plan_cells": 2, + "missing_cells": [], + "skipped_first_publisher_wins": [] + } +} +``` + +This is an excerpt of the expected fields, not a complete checkpoint file. +For an expanded campaign, use its actual plan count in place of 2; a complete +fresh publication has `published_cells == plan_cells == FPM_CELL_COUNT`. + +The resulting directory has this structure; `` is the first +16 characters of the plan hash: + +```text +$FPM_RUN/ + plan.json + checkpoints/ + fpm_forward_smoke.json + fpm_forward.json + artifacts// + collection-plan.json + run-manifest.json + cells//... + smoke/... + systems/ + h200_sxm.yaml + query_versions.yaml + attention_lane_defaults.yaml + data/h200_sxm/vllm// + fpm_forward_perf.parquet + fpm_forward_perf.metadata.json +``` + +Retain generated manifests/scripts, raw benchmark files, resolved configs, +provenance, logs, checkpoints, and the published pair. Inspect +`kv_seed_regime`; do not assume every row contains real-KV evidence. See the +[collector contract](../../collector/README.md#whole-forward-fpm-campaign) for +accepted regimes and publication behavior. + +### B4. Inspect the published pair (step 4) + +The version directory is the actual Pod's `importlib.metadata.version("vllm")`, +also recorded in the published rows/metadata. It is not the AIS version or the +`--dynamo-version` used to select templates. Read it from the artifacts before +setting this variable: + +```bash +export FPM_VLLM_VERSION=REPLACE_WITH_ACTUAL_COLLECTED_VLLM_VERSION +export FPM_PARQUET="$FPM_RUN/systems/data/h200_sxm/vllm/$FPM_VLLM_VERSION/fpm_forward_perf.parquet" +python3 - "$FPM_PARQUET" > "$FPM_RUN/pair-inspection.json" <<'PY' +import hashlib +import json +import sys +from pathlib import Path + +import pyarrow.parquet as pq + +parquet = Path(sys.argv[1]) +metadata = json.loads(parquet.with_suffix(".metadata.json").read_text()) +with parquet.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() +assert metadata["schema_name"] == "aic_fpm_forward_perf" +assert metadata["schema_version"] == 6 +assert metadata["parquet_sha256"] == digest +table = pq.read_table(parquet) +assert table.num_rows == metadata["row_count"] and table.num_rows > 0 +identity = [ + "model_path", "system", "backend", "backend_version", "workload_kind", + "tp", "pp", "dp", "moe_tp", "moe_ep", "cp", "gemm_quant_mode", + "moe_quant_mode", "fmha_quant_mode", "comm_quant_mode", "kv_cache_dtype", + "moe_backend", "attention_backend", "enable_wideep", "enable_eplb", +] +identity += ["dcp"] if "dcp" in table.column_names else [] +frame = table.to_pandas() +report = { + "hash_match": True, + "row_count": table.num_rows, + "phase_cells": frame[identity].drop_duplicates().to_dict(orient="records"), + "kv_seed_regimes": frame["kv_seed_regime"].value_counts().to_dict() if "kv_seed_regime" in frame else {}, +} +print(json.dumps(report, indent=2)) +PY +``` + +**Expected:** exit code 0, `hash_match: true`, a positive row count, and phase +identities matching the campaign. This structural check does not replace the +native loader's validation or prove interpolation coverage. Inspect the +`batch_size`, `total_prefill_tokens`, and `total_kv_read_tokens` coordinates +before selecting queries; axis minima/maxima alone do not prove that all +interior points can be answered. + +Keep these paths distinct: + +| Setting | Value in this guide | +| --- | --- | +| Collector `--fpm-database-root` | `$FPM_RUN/systems/data` | +| Canonical SDK/Replay `systems_paths` entry | `$FPM_RUN/systems` | +| Pair directory | `$FPM_RUN/systems/data/h200_sxm/vllm/$FPM_VLLM_VERSION` | + +Whole-forward FPM has no extra model-family directory between the system and +backend. Always transfer the metadata sidecar together with the Parquet file. + +### B5. Query the canonical model (step 5) + +This query explicitly selects the external root. The example asks for a +four-request prefill with 1,024 new tokens per request and no cached prefix; +confirm this coordinate is covered by the new collection. + +```bash +AIC_ALLOW_UNLISTED_VERSIONS=1 \ +python3 - "$FPM_RUN/systems" "$FPM_VLLM_VERSION" <<'PY' +import json +import math +import sys +from pathlib import Path + +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + +config = ForwardPassPerfModelConfig( + model="MiniMaxAI/MiniMax-M2.7", system="h200_sxm", backend="vllm", + worker_type="aggregated", backend_version=sys.argv[2], + systems_paths=(str(Path(sys.argv[1]).resolve(strict=True)),), + estimation_mode="fpm_interpolation", fallback_policy="deny", + tp=4, pp=1, attention_dp=1, moe_tp_size=4, moe_ep_size=1, + estimator_config={"correction": {"enabled": False}}, +) +model = RustForwardPassPerfModel.best_available(config) +try: + latency_ms = model.estimate_forward_pass_time_ms({ + "version": 1, + "scheduled_requests": { + "num_prefill_requests": 4, + "sum_prefill_tokens": 4096, + "sum_prefill_kv_tokens": 0, + }, + }) + assert latency_ms is not None and math.isfinite(latency_ms) and latency_ms > 0 + print(json.dumps({"prefill_ms": latency_ms, "provenance": model.diagnostics()["provenance"]})) +finally: + model.close() +PY +``` + +**Expected:** exit code 0 and a finite, positive `prefill_ms`. No serving GPU is +used for this query. A successful lookup confirms this query can consume the +pair, not that all replay steps will be in-domain or that the prediction is +accurate. + +The model, hardware, backend version, parallel shape, quantization, and backend +knobs must match the collected identity. This example uses checkpoint defaults. +For explicitly collected nondefault precision, use the SDK's +`gemm_quant_mode`, `moe_quant_mode`, `fmha_quant_mode`, `kvcache_quant_mode`, and +`comm_quant_mode` arguments with the corresponding enum names as strings. See +[the canonical API](../../../../docs/core-api.md#choosing-a-forward-pass-api). + +This guide copies the version policy file and uses literal backend versions. +If the collected version is already queryable, omit +`AIC_ALLOW_UNLISTED_VERSIONS=1`. The override is a transitional measure, not a +support guarantee. Do not rename version directories or edit identities to +bypass a mismatch. External trees without `query_versions.yaml` use a different +policy that disables the version-slot gate; omitting the file is not needed for +this workflow. + +### B6. Run Replay (step 6) + +Save this as `$FPM_RUN/predict.yaml`. Replace the backend version with the +collected version and `systems_paths` with the absolute `$FPM_RUN/systems` +directory. Confirm block size, context length, scheduler limits, and +prefix-cache policy against the resolved runtime configuration and the target +workload. These settings are explicit simulation inputs, not reconstructed +automatically from the FPM table. + +```yaml +engine: + mode: aggregated + model: MiniMaxAI/MiniMax-M2.7 + hardware: h200_sxm + backend: vllm + backend_version: "REPLACE_WITH_ACTUAL_COLLECTED_VLLM_VERSION" + context_length: 8192 + systems_paths: ["REPLACE_WITH_ABSOLUTE_SYSTEMS_DIRECTORY"] + estimation_mode: fpm_interpolation + fallback_policy: deny + estimator_config: + correction: {enabled: false} + workers: + aggregated: + parallelism: + replicas: 1 + tensor: 4 + pipeline: 1 + attention_data: 1 + moe_tensor: 4 + moe_expert: 1 + scheduler: + max_batched_tokens: 8192 + max_sequences: 256 + kv_cache: + block_size: 64 + prefix_caching: false + capacity: {type: default} + timing: {type: default} +traffic: + source: {type: synthetic, input_tokens: 1024, output_tokens: 32} + load: {type: concurrency, concurrency: 4} + stop: {requests: 8} +``` + +`engine.systems_paths` selects the data root for this prediction. No global +Python root overrides are needed. Run the CLI directly: + +```bash +AIC_ALLOW_UNLISTED_VERSIONS=1 \ +aisimulate predict --stack engine --config "$FPM_RUN/predict.yaml" \ + --output-dir "$FPM_RUN/prediction-fpm" --capture-per-request +``` + +**Expected:** exit code 0 and a new `prediction-fpm/` directory containing: + +- `prediction.json`: the durable prediction report; check the completion count + against the eight requested completions and inspect the reported metrics. +- `requests.jsonl`: per-request records enabled by `--capture-per-request`. + +Keep the input YAML, resolved estimator provenance, package/source revision, +and pair hash with the report so its data source can be reproduced. Use a new output +directory for another prediction; the CLI rejects nonempty output directories +unless overwrite is explicitly requested. + +Every simulated step needs a matching identity and supported coordinates. A +missing cell fails rather than silently switching to `op_level`. With default +capacity estimation, FPM also caps KV capacity at the collected decode-KV +ceiling. Inspect that cap when explaining concurrency and throughput. + + + +## Recovery and cleanup + +For a Collector campaign such as Example B, use the same frozen `fpm_args` +and output roots to resume: + +```bash +python3 -m collector.fpm_forward "${fpm_args[@]}" --resume +``` + +Before retrying failed cells, copy their existing raw artifacts and logs to a +separate archive. The runner removes old raw/log directories for cells it +re-executes. After preserving that evidence: + +```bash +python3 -m collector.fpm_forward "${fpm_args[@]}" --resume --resume-retry-failed +``` + +A completed, validated schema-v6 publication is terminal on resume even if raw +artifacts were pruned. Changing the model, image, sampling policy, or source can +change the frozen plan. Use fresh campaign, checkpoint, artifact, and database +roots for independent recollection or A/B comparisons. Do not run concurrent +campaigns with colliding workload names. + +The collector normally cleans up each cell's workload in its finalization path; +cleanup failure is an error. After an abnormal interruption, preserve available +logs and inspect the manifests recorded under this campaign's cell artifacts. +For each owned manifest that still has resources, run: + +```bash +kubectl delete -f /absolute/path/to/this/cell/k8s_deploy.yaml \ + --ignore-not-found --cascade=foreground --wait=true --timeout=180s +kubectl get -f /absolute/path/to/this/cell/k8s_deploy.yaml +``` + +**Expected:** the resources are absent. Also check child Pods using the actual +labels/owner references in that manifest. +If deletion times out, cleanup is not complete; inspect and remove the remaining +owned resources before ending the campaign. +Do not delete unrelated workloads or the namespace. Default `/results` storage +is Pod-local `emptyDir`, so deleting Pods before salvaging results loses those +files. + +## FPM data and prediction troubleshooting + +| Symptom | Check and next action | +| --- | --- | +| Plan fails to resolve a model | Supply a real local config with `--fpm-model-config`; preserve the canonical model identity | +| Smoke succeeds but no Parquet appears | Expected for smoke; inspect both phase results before formal collection | +| Partial publication or skipped existing cells | Inspect `missing_cells` and `skipped_first_publisher_wins`; use a new database root for independent recollection | +| Pair hash/schema validation fails | Preserve both files, recover a matching published pair, and do not edit the sidecar to conceal a mismatch | +| FPM lookup fails | Check root layout, actual backend version, full identity, and supported coordinates; some queries also require supported SOL operations | +| SDK query works but replay fails | Inspect the later simulated shapes and decode-KV ceiling; one successful prefill query does not establish replay coverage | +| Prediction output directory is not empty | Select a new directory to retain the previous report | diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 5fb67d474f..21b85cdf73 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -156,6 +156,10 @@ def finish_lowering(value): if lowered.get("num_gpu_blocks") is not None: return finish_lowering(lowered) + if canonical_result is not None and (resolved.get("dcp") or 1) > 1: + raise ValueError( + "DCP FPM replay requires explicit KV block capacity; automatic DCP/hybrid sizing is unsupported" + ) backend = lowered.get("aic_backend") if backend is None: return finish_lowering(lowered) diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index abe0c8d6c1..fbdbcdaaab 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -374,6 +374,7 @@ def _parallel_mapping(worker: WorkerPredictionConfig, *, prefix: str) -> dict[st f"{prefix}attention_dp": parallel.attention_data, f"{prefix}moe_tp": parallel.moe_tensor, f"{prefix}moe_ep": parallel.moe_expert, + **({f"{prefix}dcp": parallel.decode_context} if parallel.decode_context is not None else {}), } @@ -405,6 +406,19 @@ def _worker_performance_model_metadata( if getattr(engine, field) is not None }, } + if parallel.decode_context is not None: + config["dcp"] = parallel.decode_context + for field in ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kvcache_quant_mode", + "comm_quant_mode", + "attention_backend", + ): + value = getattr(worker.timing, field) + if value is not None: + config[field] = value config["database_mode"] = worker.timing.database_mode or engine.database_mode if engine.speculation is not None: config["speculation"] = engine.speculation.cost_config() @@ -428,6 +442,27 @@ def _worker_engine_args( cache = worker.kv_cache capacity = cache.capacity memory_fraction = capacity.memory_fraction + identity_fields = ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kvcache_quant_mode", + "comm_quant_mode", + "attention_backend", + ) + if (engine.mode == "afd" or engine.workers.encoder is not None) and any( + getattr(worker.timing, field) is not None for field in identity_fields + ): + raise ValueError("explicit timing identity requires the canonical forward-pass provider") + if parallel.decode_context is not None: + if parallel.tensor % parallel.decode_context: + raise ValueError("decode_context must divide tensor parallelism") + if worker.timing.type != "default" or engine.mode == "afd" or engine.workers.encoder is not None: + raise ValueError("DCP requires the canonical forward-pass timing provider") + if parallel.decode_context > 1 and capacity.type != "fixed": + raise ValueError( + "DCP FPM replay requires explicit KV block capacity; automatic DCP/hybrid sizing is unsupported" + ) if capacity.type == "default" and memory_fraction is None: memory_fraction = 0.88 if backend == "sglang" else 0.9 block_size = cache.block_size @@ -534,11 +569,15 @@ def _worker_engine_args( tp=parallel.tensor, pp=parallel.pipeline, attention_dp=parallel.attention_data, + dcp=parallel.decode_context, moe_tp_size=parallel.moe_tensor if sharded_moe else None, moe_ep_size=parallel.moe_expert if sharded_moe else None, kv_block_size=block_size, nextn=engine.nextn, - **{name: getattr(engine, name) for name in ENGINE_MODEL_CONTROL_FIELDS}, + **{ + name: getattr(timing, name, None) if getattr(timing, name, None) is not None else getattr(engine, name) + for name in ENGINE_MODEL_CONTROL_FIELDS + }, speculation=engine.speculation.cost_config() if engine.speculation is not None else None, estimation_mode=timing.estimation_mode or engine.estimation_mode, fallback_policy=timing.fallback_policy or engine.fallback_policy, @@ -603,13 +642,15 @@ def _resolve_kv_bytes_per_token( if configured != "auto": return configured parallel = worker.parallelism + if (parallel.decode_context or 1) > 1: + raise ValueError("DCP KV transfer sizing requires explicit bytes_per_token") return estimate_kv_bytes_per_token( engine.model, tp_size=parallel.tensor, pp_size=parallel.pipeline, moe_tp_size=parallel.moe_tensor, moe_ep_size=parallel.moe_expert, - **({"kvcache_quant_mode": engine.kvcache_quant_mode} if engine.kvcache_quant_mode else {}), + kvcache_quant_mode=worker.timing.kvcache_quant_mode or engine.kvcache_quant_mode, ) diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 4f0708baba..0ad4815441 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -72,7 +72,7 @@ class AFDSearchRecommendationConfig(StrictModel): max_candidates: PositiveInt = 10_000 -class ParallelismPredictionConfig(StrictModel): +class ParallelismPresetConfig(StrictModel): replicas: PositiveInt = 1 tensor: PositiveInt = 1 pipeline: PositiveInt = 1 @@ -81,6 +81,10 @@ class ParallelismPredictionConfig(StrictModel): moe_expert: PositiveInt = 1 +class ParallelismPredictionConfig(ParallelismPresetConfig): + decode_context: Annotated[int, Field(strict=True, gt=0)] | None = None + + class SchedulerPredictionConfig(StrictModel): max_batched_tokens: PositiveInt = 8192 max_sequences: PositiveInt = 256 @@ -248,6 +252,12 @@ def cost_config(self) -> dict[str, Any]: class TimingConfig(StrictModel): + gemm_quant_mode: str | None = None + moe_quant_mode: str | None = None + fmha_quant_mode: str | None = None + kvcache_quant_mode: str | None = None + comm_quant_mode: str | None = None + attention_backend: str | None = None type: Literal["default", "fixed", "polynomial"] = "default" forward_model: Literal["op_level", "fpm"] = Field(default="op_level", exclude=True) fpm_parquet_path: str | None = None @@ -279,6 +289,18 @@ def _preserve_legacy_estimator_selection(cls, value): @model_validator(mode="after") def _validate_timing(self) -> TimingConfig: + if self.type != "default" and any( + getattr(self, field) is not None + for field in ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kvcache_quant_mode", + "comm_quant_mode", + "attention_backend", + ) + ): + raise ValueError("quantization and backend identity require default timing") if self.estimation_mode == "fpm_interpolation": self.forward_model = "fpm" elif self.estimation_mode == "op_level": @@ -528,7 +550,7 @@ def _validate_roles(self) -> EnginePredictionConfig: class ParallelismRecommendationConfig(StrictModel): - preset: Literal["default", False] | list[ParallelismPredictionConfig] | dict[str, Any] = "default" + preset: Literal["default", False] | list[ParallelismPresetConfig] | dict[str, Any] = "default" replicas: ParallelDomain | None = None tensor: ParallelDomain | None = None pipeline: ParallelDomain | None = None @@ -638,6 +660,25 @@ class WorkerRecommendationConfig(StrictModel): timing: TimingConfig = Field(default_factory=TimingConfig) startup_seconds: float = Field(default=0.0, ge=0.0) + @model_validator(mode="after") + def _reject_timing_identity_overrides(self): + if any( + getattr(self.timing, field) is not None + for field in ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kvcache_quant_mode", + "comm_quant_mode", + "attention_backend", + ) + ): + raise ValueError( + "explicit timing quantization/backend identity is prediction-only; " + "recommendation preflight does not support these overrides" + ) + return self + class EncoderRecommendationConfig(StrictModel): """Finite encoder domains; scalar values pin a single choice.""" diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index 23d454a66e..cccb8e4549 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -881,6 +881,20 @@ def _candidate_prediction( timing["database_mode"] = resolved["database_mode"] policy = resolved["transfer_policy"] timing["transfer_policy"] = list(policy) if policy is not None else None + timing.update( + { + field: resolved[field] + for field in ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kvcache_quant_mode", + "comm_quant_mode", + "attention_backend", + ) + if resolved.get(field) is not None + } + ) kv_cache = { "block_size": block_size, "prefix_caching": sample[f"{role}_enable_prefix_caching"], @@ -906,6 +920,11 @@ def _candidate_prediction( "attention_data": sample[f"{prefix}attention_dp"], "moe_tensor": sample[f"{prefix}moe_tp"], "moe_expert": sample[f"{prefix}moe_ep"], + **( + {"decode_context": estimator.config["dcp"]} + if estimator is not None and estimator.config.get("dcp") is not None + else {} + ), }, "scheduler": { "max_batched_tokens": sample[f"{role}_max_num_batched_tokens"], diff --git a/python/aisimulate/src/aisimulate/sdk/fpm_config.py b/python/aisimulate/src/aisimulate/sdk/fpm_config.py new file mode 100644 index 0000000000..cb6b379b38 --- /dev/null +++ b/python/aisimulate/src/aisimulate/sdk/fpm_config.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility alias for aisimulate_core.sdk.fpm_config.""" + +from aisimulate.sdk._compat import alias_module as _alias_module + +_alias_module(__name__, "aisimulate_core.sdk.fpm_config") diff --git a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py index b2865c1de3..a9bc17039c 100644 --- a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py +++ b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py @@ -85,6 +85,7 @@ def _request(self, sample: dict[str, Any], role: str) -> ForwardPassPerfModelCon tp=int(sample[f"{prefix}tp"]), pp=int(sample[f"{prefix}pp"]), attention_dp=int(sample[f"{prefix}attention_dp"]), + dcp=sample.get(f"{prefix}dcp"), moe_tp_size=moe_tp if moe_tp * moe_ep > 1 else None, moe_ep_size=moe_ep if moe_tp * moe_ep > 1 else None, nextn=int(nextn or 0), @@ -159,6 +160,7 @@ def _resolve(self, request: ForwardPassPerfModelConfig, role: str) -> ForwardPas "tp", "pp", "attention_dp", + "dcp", "moe_tp_size", "moe_ep_size", "nextn", diff --git a/python/aisimulate/src/aisimulate_core/_native.pyi b/python/aisimulate/src/aisimulate_core/_native.pyi index 58aff624a0..e7b013df8c 100644 --- a/python/aisimulate/src/aisimulate_core/_native.pyi +++ b/python/aisimulate/src/aisimulate_core/_native.pyi @@ -173,6 +173,10 @@ class RustForwardPassPerfModel: @staticmethod def normalize_config(config_json: str) -> str: ... @staticmethod + def _normalize_fpm_options( + options_json: str, fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None + ) -> str: ... + @staticmethod def legacy_estimator_config(options_json: str) -> str: ... @staticmethod def migrate_legacy_config( diff --git a/python/aisimulate/src/aisimulate_core/sdk/config.py b/python/aisimulate/src/aisimulate_core/sdk/config.py index 26e5659e4e..56f48bbea1 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config.py @@ -8,6 +8,7 @@ from typing import TypeVar, Union from aisimulate_core.sdk import common +from aisimulate_core.sdk.fpm_config import FpmCompileConfig from aisimulate_core.sdk.speculation.base import SpeculationConfig KernelBackendT = TypeVar("KernelBackendT", bound=StrEnum) @@ -69,6 +70,8 @@ class ModelConfig: Model configuration. """ + dcp_size: int | None = field(default=None, kw_only=True) + fpm_config: FpmCompileConfig | None = field(default=None, kw_only=True) tp_size: int = 1 pp_size: int = 1 gemm_quant_mode: common.GEMMQuantMode | None = None @@ -156,6 +159,10 @@ class ModelConfig: fpm_fmha_quant_mode: common.FMHAQuantMode | None = field(default=None, kw_only=True) def __post_init__(self) -> None: + if self.dcp_size is not None and ( + type(self.dcp_size) is not int or self.dcp_size <= 0 or self.tp_size % self.dcp_size + ): + raise ValueError("dcp_size must be positive and divide tp_size") self.moe_backend = normalize_kernel_backend(self.moe_backend, common.MoEBackend, "moe_backend") self.attention_backend = normalize_kernel_backend( self.attention_backend, diff --git a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py index 1723860363..6e4452e91a 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py @@ -47,6 +47,7 @@ def build_model_config( *, moe_kernel_source: str | None = None, fpm_fmha_quant_mode: str | None = None, + dcp_size: int | None = None, decode_workload_distribution: str | None = None, prefill_graph_profile: str | None = None, ) -> ModelConfig: @@ -59,6 +60,7 @@ def build_model_config( ) return ModelConfig( tp_size=tp_size, + dcp_size=dcp_size, pp_size=pp_size, attention_dp_size=attention_dp_size, moe_tp_size=moe_tp_size, diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index 5677fab48e..a898b38248 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -141,6 +141,8 @@ # - 22 (GLM-5.2 VR200 pilot): exact observed-MoE selection, prefill graph # identity and two appended composite operators extend the schema-21 layout. # The pilot and AIC-1781 concurrently claimed 21; reject both older layouts. +# - 23 (DCP identity): optional recorded dcp_size in ParallelMapping. +# - 24 (typed FPM DCP): dcp_size moves out of the op matching tuple. # Single owner: the Rust crate constant. Python re-exports it for # diagnostics/tests instead of declaring a twin to keep in sync. ENGINE_SPEC_SCHEMA_VERSION = aisimulate_core.engine_spec_schema_version() @@ -172,8 +174,9 @@ def _fpm_spec_dict(op: FPMForwardOp) -> dict: "phase": op._phase, "model_path": op._model_path, "match_identity": list(op._match_identity), + "dcp_size": op._dcp_size, "original_fmha_quant_mode": op._original_fmha_quant_mode, - "weight_bytes": op._weight_bytes, + "weight_bytes": op._resident_weight_bytes, # Speculative verify width for the equivalent-AR decode mapping # (1 = plain AR). Set by the fpm hybrid rewrite in models when a # draft scheme is materialized. @@ -343,6 +346,7 @@ def _engine_config_dict( "moe_tp_size": _opt_int(getattr(cfg, "moe_tp_size", None)), "moe_ep_size": _opt_int(getattr(cfg, "moe_ep_size", None)), "cp_size": _opt_int(getattr(cfg, "cp_size", None)), + "dcp_size": _opt_int(getattr(cfg, "dcp_size", None)), # QuantizationConfig (flattened) "weight_dtype": _rust_quant_to_dtype(getattr(cfg, "gemm_quant_mode", None)), "moe_dtype": _rust_moe_quant_to_dtype(getattr(cfg, "moe_quant_mode", None)), @@ -441,6 +445,8 @@ def compile_engine( tp_size: int = 1, pp_size: int = 1, attention_dp_size: int = 1, + dcp_size: int | None = None, + fpm_options: dict | None = None, moe_tp_size: int | None = None, moe_ep_size: int | None = None, gemm_quant_mode: str | None = None, @@ -483,7 +489,7 @@ def compile_engine( logical batches are exploratory and queries outside 1..32 fail. Missing or mismatched approved profile data raises ``DecodeMoeProfileError``. - Engine schema 22 persists the exact-profile policy on MoE operators. Default + Engine schema 24 persists the exact-profile policy on MoE operators. Default OpSpec JSON includes that new false field, and older binary specs require recompilation; default representations are therefore not byte-identical. """ @@ -507,11 +513,28 @@ def compile_engine( if shared_layer: raise PrefillGraphProfileError("prefill_graph_profile does not allow shared-source inheritance") shared_layer = False + from aisimulate_core.sdk.fpm_config import FpmCompileConfig from aisimulate_core.sdk.speculation import SpeculationConfig + recorded_attention_backend = attention_backend + if backend == "vllm" and attention_backend == "FLASHINFER_MLA": + attention_backend = "flashinfer" resolved_moe_tp = moe_tp_size if moe_tp_size is not None else 1 resolved_moe_ep = moe_ep_size if moe_ep_size is not None else 1 try: + fpm_config = FpmCompileConfig( + fpm_options, + attention_backend=recorded_attention_backend, + fmha_quant_mode=fmha_quant_mode if fmha_quant_mode is not None else fpm_fmha_quant_mode, + comm_quant_mode=comm_quant_mode, + ) + options_path = fpm_config.options["fpm_parquet_path"] + if options_path is not None: + if fpm_parquet_path is not None and fpm_parquet_path != options_path: + raise ValueError("conflicting fpm_parquet_path and fpm_options.fpm_parquet_path") + if forward_model != "fpm": + raise ValueError("fpm_parquet_path requires forward_model='fpm'") + fpm_parquet_path = options_path validate_moe_controls( model_path=model_path, enable_eplb=enable_eplb, @@ -522,6 +545,7 @@ def compile_engine( resolved_speculation = SpeculationConfig(**speculation) if speculation is not None else None model_config = build_model_config( tp_size=tp_size, + dcp_size=dcp_size, pp_size=pp_size, attention_dp_size=attention_dp_size, moe_tp_size=resolved_moe_tp, @@ -542,6 +566,7 @@ def compile_engine( wideep_num_slots=wideep_num_slots, speculation=resolved_speculation, ) + model_config.fpm_config = fpm_config # Apply MTP BEFORE get_model so the walked op lists carry the # (L+nextn)/L compute scale; accepted-token progress is applied above core. apply_nextn(model_config, nextn) diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_config.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_config.py new file mode 100644 index 0000000000..748e04a397 --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_config.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal FPM compilation context; Rust owns the interpolation schema.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass + + +@dataclass(frozen=True, init=False) +class FpmCompileConfig: + """Carry a validated immutable options payload and its recorded backend label.""" + + _options_json: str + attention_backend: str | None + + def __init__( + self, + options: dict | None = None, + *, + attention_backend: str | None = None, + fmha_quant_mode: str | None = None, + comm_quant_mode: str | None = None, + ) -> None: + import aisimulate_core + + normalized = aisimulate_core.RustForwardPassPerfModel._normalize_fpm_options( + json.dumps({} if options is None else options), fmha_quant_mode, comm_quant_mode + ) + object.__setattr__(self, "_options_json", normalized) + object.__setattr__(self, "attention_backend", attention_backend) + + @property + def options(self) -> dict: + """Return a copy of the complete Rust-resolved options, including defaults.""" + return json.loads(self._options_json) + + def cache_identity(self) -> dict: + return {"options": self.options, "attention_backend": self.attention_backend} + + +def resolve_fpm_config(model_config) -> FpmCompileConfig: + """Resolve legacy direct model construction without Python-owned defaults.""" + if model_config.fpm_config is None: + model_config.fpm_config = FpmCompileConfig() + return model_config.fpm_config diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py index 55cb2c9f14..4ab809be23 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py @@ -81,6 +81,7 @@ def _apply_forward_model_fpm(model: BaseModel, backend_name: str = "vllm") -> Ba """Centralized fpm rewrite: each phase list becomes exactly one whole-model op. No model class rewrites its own lists; metadata, parallelism, and the public model type are unchanged.""" + from aisimulate_core.sdk.fpm_config import resolve_fpm_config from aisimulate_core.sdk.operations.fpm_forward import _CELL_MATCH_COLUMNS, FPMForwardOp if getattr(model.config, "decoder_replay", False) and "execution_profile" not in _CELL_MATCH_COLUMNS: @@ -89,7 +90,7 @@ def _apply_forward_model_fpm(model: BaseModel, backend_name: str = "vllm") -> Ba # bounded decoder tail from a full forward at the same coordinates. raise NotImplementedError("decoder_replay requires FPM tables with execution_profile identity") - if model.encoder_ops: + if model.encoder_ops and not resolve_fpm_config(model.config).options["text_only"]: raise NotImplementedError( f"forward_model='fpm' does not support encoder/multimodal models " f"(model_family={model.model_family!r} has encoder ops). Use forward_model='op_level'." @@ -124,27 +125,32 @@ def _apply_forward_model_fpm(model: BaseModel, backend_name: str = "vllm") -> Ba draft_context_ops = [op for op in model.context_ops if op._name.startswith("draft_")] draft_generation_ops = [op for op in model.generation_ops if op._name.startswith("draft_")] weight_bytes = model.get_resident_weights_bytes() - prefill_op = FPMForwardOp("prefill", model.config, model.model_path, sol_ops=context_ops, weight_bytes=weight_bytes) - decode_op = FPMForwardOp( - "decode", model.config, model.model_path, sol_ops=generation_ops, weight_bytes=weight_bytes - ) - if has_draft_scheme: - decode_op._verify_width = int(model.verify_width) - model.context_ops = [prefill_op, *draft_context_ops] - model.generation_ops = [decode_op, *draft_generation_ops] from aisimulate_core.sdk.fpm_identity import execution_identity identity = execution_identity( getattr(model, "raw_config", {}), decoder_replay=getattr(model.config, "decoder_replay", False), backend=backend_name, - # The SDK supports this prediction contract; the producer separately + # These are the SDK prediction assumptions; the producer separately # verifies actual runtime residency and token-only requests. engram_cpu_offload=False, input_modality="text", ) - for op in (prefill_op, decode_op): - op._match_identity = (*op._match_identity[:15], *identity) + prefill_op = FPMForwardOp( + "prefill", model.config, model.model_path, sol_ops=context_ops, weight_bytes=weight_bytes, execution=identity + ) + decode_op = FPMForwardOp( + "decode", model.config, model.model_path, sol_ops=generation_ops, weight_bytes=weight_bytes, execution=identity + ) + # Compiled text specs omit encoder ops; retain their resident weights there. + # Python memory accounting still sums encoder_ops separately. + resident_weights = weight_bytes + float(sum(op.get_weights() for op in model.encoder_ops)) + prefill_op._resident_weight_bytes = resident_weights + decode_op._resident_weight_bytes = resident_weights + if has_draft_scheme: + decode_op._verify_width = int(model.verify_width) + model.context_ops = [prefill_op, *draft_context_ops] + model.generation_ops = [decode_op, *draft_generation_ops] model.forward_model = "fpm" return model @@ -167,6 +173,8 @@ def get_model( rewrites each phase list to a single whole-model ``FPMForwardOp``. """ forward_model = getattr(model_config, "forward_model", "op_level") or "op_level" + if (model_config.dcp_size or 1) > 1 and (forward_model != "fpm" or backend_name != "vllm"): + raise NotImplementedError("DCP timing requires measured vLLM FPM interpolation") if forward_model not in _FORWARD_MODELS: raise InvalidEngineConfigurationError( f"Unknown forward_model: {forward_model!r}. Valid values: {', '.join(_FORWARD_MODELS)}" diff --git a/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py b/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py index 232634f3a3..312e6584fa 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py +++ b/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py @@ -37,6 +37,7 @@ from enum import Enum +from aisimulate_core.sdk.fpm_config import resolve_fpm_config from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY from aisimulate_core.sdk.operations.base import PythonOperation @@ -108,6 +109,8 @@ def __init__( sol_fn=None, weight_bytes: float = 0.0, sol_ops: list | None = None, + *, + execution: tuple[str, ...] = LEGACY_EXECUTION_IDENTITY, ) -> None: """``sol_ops`` — the model's ORIGINAL op-level list for this phase — rides the compiled spec so the Rust FPM SOL roofline derives from the @@ -135,28 +138,37 @@ def __init__( self._phase = phase self._model_path = str(model_path) self._weight_bytes = float(weight_bytes) + self._resident_weight_bytes = self._weight_bytes + fpm_config = resolve_fpm_config(model_config) + unrecorded = fpm_config.options["unrecorded_quant_modes"] fmha_selector = getattr(model_config, "fpm_fmha_quant_mode", None) self._original_fmha_quant_mode = None if fmha_selector is None else _norm_identity(model_config.fmha_quant_mode) - self._match_identity = ( - _norm_identity(model_config.gemm_quant_mode), - _norm_identity(model_config.moe_quant_mode), - _norm_identity(model_config.fmha_quant_mode if fmha_selector is None else fmha_selector), - _norm_identity(model_config.comm_quant_mode), - _norm_identity(model_config.kvcache_quant_mode), - _norm_identity(model_config.tp_size), - _norm_identity(model_config.pp_size), - _norm_identity(model_config.attention_dp_size), - _norm_identity(model_config.moe_tp_size if model_config.moe_tp_size is not None else 1), - _norm_identity(model_config.moe_ep_size if model_config.moe_ep_size is not None else 1), - _norm_identity(model_config.cp_size), - _norm_backend_request(getattr(model_config, "moe_backend", None)), + identity = { + "gemm_quant_mode": model_config.gemm_quant_mode, + "moe_quant_mode": model_config.moe_quant_mode, + "fmha_quant_mode": None + if "fmha" in unrecorded + else (model_config.fmha_quant_mode if fmha_selector is None else fmha_selector), + "comm_quant_mode": None if "comm" in unrecorded else model_config.comm_quant_mode, + "kv_cache_dtype": model_config.kvcache_quant_mode, + "tp": model_config.tp_size, + "pp": model_config.pp_size, + "dp": model_config.attention_dp_size, + "moe_tp": model_config.moe_tp_size if model_config.moe_tp_size is not None else 1, + "moe_ep": model_config.moe_ep_size if model_config.moe_ep_size is not None else 1, + "cp": model_config.cp_size, + "moe_backend": _norm_backend_request(getattr(model_config, "moe_backend", None)), # ModelConfig spells the engine default out ("flashinfer"); the # collector records engine-decided knobs as "auto". - _norm_backend_request(getattr(model_config, "attention_backend", None), engine_default="flashinfer"), - _norm_identity(bool(getattr(model_config, "enable_wideep", False))), - _norm_identity(bool(getattr(model_config, "enable_eplb", False))), - ) - self._match_identity += LEGACY_EXECUTION_IDENTITY + "attention_backend": _norm_backend_request( + fpm_config.attention_backend or model_config.attention_backend, engine_default="flashinfer" + ), + "enable_wideep": bool(getattr(model_config, "enable_wideep", False)), + "enable_eplb": bool(getattr(model_config, "enable_eplb", False)), + **dict(zip(EXECUTION_COLUMNS, execution, strict=True)), + } + self._match_identity = tuple(_norm_identity(identity[name]) for name in _CELL_MATCH_COLUMNS) + self._dcp_size = model_config.dcp_size self._sol_ops = list(sol_ops) # Speculative verify width for the equivalent-AR decode mapping # (1 = plain AR). Set post-construction by the fpm model rewrite for diff --git a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py index f3a525f88e..612be02bdc 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -111,6 +111,7 @@ class ForwardPassPerfModelConfig: tp: int = 1 pp: int = 1 attention_dp: int = 1 + dcp: int | None = dataclass_field(default=None, kw_only=True) moe_tp_size: int | None = None moe_ep_size: int | None = None gemm_quant_mode: str | None = None @@ -1174,6 +1175,7 @@ def _speculation_identity(model_config: Any) -> str | None: def _engine_config_json(model: Any, database: Any) -> str: model_config = model.config + fpm_config = getattr(model_config, "fpm_config", None) # Forward only the MTP draft length. The aic-core layer models iteration compute cost; # accepted-token progress belongs to the upper prediction layer. nextn = getattr(model, "_nextn", None) @@ -1191,6 +1193,7 @@ def _engine_config_json(model: Any, database: Any) -> str: "attention_dp_size": _optional_int(getattr(model_config, "attention_dp_size", None)), # Part of the engine identity so cp variants get distinct cached handles. "cp_size": _optional_int(getattr(model_config, "cp_size", None)), + "dcp_size": _optional_int(getattr(model_config, "dcp_size", None)), "weight_dtype": _quant_to_dtype(getattr(model_config, "gemm_quant_mode", None)), "moe_dtype": _moe_quant_to_dtype(getattr(model_config, "moe_quant_mode", None)), "activation_dtype": _quant_to_dtype(getattr(model_config, "fmha_quant_mode", None)), @@ -1239,6 +1242,7 @@ def _engine_config_json(model: Any, database: Any) -> str: "comm": _raw_quant_name(getattr(model_config, "comm_quant_mode", None)), }, "model_config": { + "fpm_config": fpm_config.cache_identity() if fpm_config is not None else None, "decoder_replay": bool(getattr(model_config, "decoder_replay", False)), "cp_style": getattr(model_config, "cp_style", None), "workload_distribution": getattr(model_config, "workload_distribution", None), diff --git a/python/aisimulate/tests/cross_package/test_core_public_api.py b/python/aisimulate/tests/cross_package/test_core_public_api.py index bc4426a87a..b7b0c8f66a 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -93,6 +93,7 @@ def test_stable_function_signatures() -> None: assert str(inspect.signature(compile_engine)) == ( "(model_path: 'str', system: 'str', backend: 'str', backend_version: 'str | None' = None, *, " "tp_size: 'int' = 1, pp_size: 'int' = 1, attention_dp_size: 'int' = 1, " + "dcp_size: 'int | None' = None, fpm_options: 'dict | None' = None, " "moe_tp_size: 'int | None' = None, moe_ep_size: 'int | None' = None, " "gemm_quant_mode: 'str | None' = None, moe_quant_mode: 'str | None' = None, " "kvcache_quant_mode: 'str | None' = None, fmha_quant_mode: 'str | None' = None, " diff --git a/python/aisimulate/tests/cross_package/test_import_contract.py b/python/aisimulate/tests/cross_package/test_import_contract.py index 593d62fcc7..ba4fe99f54 100644 --- a/python/aisimulate/tests/cross_package/test_import_contract.py +++ b/python/aisimulate/tests/cross_package/test_import_contract.py @@ -31,6 +31,7 @@ "engine", "engine_table_view", "errors", + "fpm_config", "fpm_dataset", "fpm_identity", "inference_summary", @@ -220,8 +221,10 @@ def test_fpm_forward_op_keeps_legacy_constructor_layout() -> None: from aisimulate.sdk.operations import FPMForwardOp - params = list(inspect.signature(FPMForwardOp.__init__).parameters) - assert params == ["self", "phase", "model_config", "model_path", "sol_fn", "weight_bytes", "sol_ops"] + signature = inspect.signature(FPMForwardOp.__init__) + positional = [name for name, param in signature.parameters.items() if param.kind != inspect.Parameter.KEYWORD_ONLY] + assert positional == ["self", "phase", "model_config", "model_path", "sol_fn", "weight_bytes", "sol_ops"] + assert signature.parameters["execution"].kind == inspect.Parameter.KEYWORD_ONLY def test_representative_from_imports_return_canonical_objects() -> None: diff --git a/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py b/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py index 6fa3a1eac2..6efdb90dce 100644 --- a/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py +++ b/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py @@ -243,11 +243,11 @@ def _fetch(database, attribute, *, shared_layer=None): monkeypatch.setattr(_etv, "fetch_attention_lane_density", _fetch) -def test_engine_spec_schema_version_is_twenty_two(): - """Exact MoE source and pilot layouts require a combined wire version.""" +def test_engine_spec_schema_version_is_twenty_four(): + """Typed FPM DCP changes the positional operator payload.""" from aisimulate.sdk import engine - assert engine.ENGINE_SPEC_SCHEMA_VERSION == 22 + assert engine.ENGINE_SPEC_SCHEMA_VERSION == 24 def test_lanes_outside_the_known_vocabulary_stay_reachable(): diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 72440fd69c..48854ce34a 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -42,6 +42,7 @@ from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel from aisimulate_core.sdk.engine import EngineHandle, compile_engine from aisimulate_core.sdk.errors import DecodeMoeProfileError +from aisimulate_core.sdk.fpm_config import FpmCompileConfig from aisimulate_core.sdk.operations.fpm_forward import _CELL_MATCH_COLUMNS pytestmark = pytest.mark.unit @@ -56,6 +57,50 @@ _CORE_SYSTEMS = os.path.join(os.path.dirname(aisimulate_core.__file__), "systems") +def test_fpm_compile_context_uses_rust_defaults_without_python_fallbacks(monkeypatch): + resolved = {"text_only": True, "unrecorded_quant_modes": ["fmha"], "fpm_parquet_path": "/profile.parquet"} + calls = [] + + class RustOptions: + @staticmethod + def _normalize_fpm_options(payload, fmha, comm): + calls.append((json.loads(payload), fmha, comm)) + return json.dumps(resolved) + + monkeypatch.setattr(aisimulate_core, "RustForwardPassPerfModel", RustOptions) + context = FpmCompileConfig(attention_backend="FLASHINFER_MLA") + assert calls == [({}, None, None)] + assert context.options == resolved + assert context.cache_identity() == {"options": resolved, "attention_backend": "FLASHINFER_MLA"} + + +def test_fpm_compile_context_survives_model_copy_and_cannot_be_mutated(): + import copy + import pickle + from dataclasses import FrozenInstanceError, replace + + context = FpmCompileConfig({"text_only": True, "unrecorded_quant_modes": ["fmha", "comm"]}) + model_config = _model_config(forward_model="fpm", fpm_config=context) + expected = context.cache_identity() + context.options["unrecorded_quant_modes"].clear() + with pytest.raises(FrozenInstanceError): + context.attention_backend = "other" + for copied in (copy.deepcopy(model_config), pickle.loads(pickle.dumps(model_config)), replace(model_config)): + assert copied.fpm_config.cache_identity() == expected + + +def test_fpm_execution_identity_and_dcp_are_independent_fields(): + from aisimulate_core.sdk.engine import _fpm_spec_dict + from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS + + execution = ("checkpoint-digest", "decoder_bounded", "hbm_tp_sharded", "text") + op = FPMForwardOp("prefill", _model_config(tp_size=4, dcp_size=4), MODEL_PATH, sol_ops=[], execution=execution) + identity = dict(zip(_CELL_MATCH_COLUMNS, op._match_identity, strict=True)) + assert tuple(identity[field] for field in EXECUTION_COLUMNS) == execution + assert "dcp" not in identity + assert _fpm_spec_dict(op)["FpmForward"]["dcp_size"] == 4 + + def _row( workload_kind: str, batch_size: int, @@ -695,6 +740,242 @@ def test_genonly_step_works_with_single_op(self, fpm_session): assert total == pytest.approx(7.0) +def test_text_only_kimi_profile_preserves_encoder_weights_and_dcp_identity(): + cfg = _model_config( + forward_model="fpm", + tp_size=8, + moe_tp_size=8, + moe_ep_size=1, + dcp_size=8, + fpm_config=FpmCompileConfig( + {"text_only": True, "unrecorded_quant_modes": ["fmha", "comm"]}, attention_backend="FLASHINFER_MLA" + ), + ) + model = models.get_model("moonshotai/Kimi-K3", cfg, "vllm") + assert model.encoder_ops + assert sum(op.get_weights() for op in model.encoder_ops) > 0 + assert all(isinstance(op, FPMForwardOp) for op in (*model.context_ops, *model.generation_ops)) + identity = model.context_ops[0]._match_identity + assert identity[2:4] == ("", "") + assert identity[12] == "FLASHINFER_MLA" + assert len(identity) == len(_CELL_MATCH_COLUMNS) + assert model.context_ops[0]._dcp_size == 8 + assert cfg.total_gpus_per_worker == 8 + from aisimulate_core.sdk.engine import _fpm_spec_dict + + prefill = model.context_ops[0] + assert _fpm_spec_dict(prefill)["FpmForward"]["dcp_size"] == 8 + resident = _fpm_spec_dict(prefill)["FpmForward"]["weight_bytes"] + assert resident == prefill.get_weights() + sum(op.get_weights() for op in model.encoder_ops) + + +@pytest.fixture +def kimi_fpm_profile(tmp_path, request): + from aisimulate_core.sdk import ForwardPassPerfModelConfig + + recorded_dcp = getattr(request, "param", 8) + + systems = tmp_path / "systems" + systems.mkdir() + shutil.copyfile(os.path.join(_CORE_SYSTEMS, "h200_sxm.yaml"), systems / "h200_sxm.yaml") + data = systems / "data" / SYSTEM / BACKEND / VERSION + identity = { + "tp": "8", + "moe_tp": "8", + "moe_ep": "1", + "fmha_quant_mode": "", + "comm_quant_mode": "", + "moe_quant_mode": "w4a16_mxfp4", + "kv_cache_dtype": "fp8", + "attention_backend": "FLASHINFER_MLA", + } + rows = [ + _row("prefill", 1, 128, 0, 7.0, model_path="moonshotai/Kimi-K3", identity=identity), + _row("decode", 1, 0, 128, 3.0, model_path="moonshotai/Kimi-K3", identity=identity), + ] + for row in rows: + row.update(measurement_policy="kvwarm_median_of_3", measurement_repeats=3) + if recorded_dcp is not None: + row["dcp"] = recorded_dcp + # Independent oracle: the synthetic measured records above define 7 ms + # prefill and 3 ms decode, irrespective of the model's operator estimates. + _write_pair( + str(data), + rows, + sidecar_overrides={ + "measurement_policy": "per_row_single_sample_or_median_of_3", + "configuration_selector": {"dcp": recorded_dcp} if recorded_dcp is not None else {}, + }, + ) + return ForwardPassPerfModelConfig( + model="moonshotai/Kimi-K3", + system=SYSTEM, + backend=BACKEND, + backend_version=VERSION, + worker_type="aggregated", + tp=8, + moe_tp_size=8, + moe_ep_size=1, + dcp=recorded_dcp, + kvcache_quant_mode="fp8", + attention_backend="FLASHINFER_MLA", + estimation_mode="fpm_interpolation", + systems_paths=(str(systems),), + estimator_config={ + "fpm_interpolation": {"text_only": True, "unrecorded_quant_modes": ["fmha", "comm"]}, + "correction": {"enabled": False}, + }, + ) + + +@pytest.mark.parametrize("kimi_fpm_profile", [None, 1, 8], indirect=True) +def test_canonical_kimi_dcp_profile_queries_measured_points(kimi_fpm_profile): + from aisimulate_core.sdk import RustForwardPassPerfModel + + config = kimi_fpm_profile + model = RustForwardPassPerfModel.best_available(config) + try: + assert model.estimate_forward_pass_time_ms( + { + "version": 1, + "scheduled_requests": { + "num_prefill_requests": 1, + "sum_prefill_tokens": 128, + }, + } + ) == pytest.approx(7.0) + assert model.estimate_forward_pass_time_ms( + { + "version": 1, + "scheduled_requests": { + "num_decode_requests": 1, + "sum_decode_kv_tokens": 128, + }, + } + ) == pytest.approx(3.0) + assert model.diagnostics()["provenance"]["config"]["dcp"] == config.dcp + finally: + model.close() + from dataclasses import replace + + from aisimulate_core.sdk.errors import PerfDataNotAvailableError + + for dcp in (None, 1, 8): + if dcp == config.dcp: + continue + with pytest.raises(PerfDataNotAvailableError, match="No FPM cell matches"): + RustForwardPassPerfModel.best_available(replace(config, dcp=dcp)) + automatic = RustForwardPassPerfModel.best_available(replace(config, estimation_mode="auto")) + try: + provenance = automatic.diagnostics()["provenance"] + assert provenance["selected_estimation_mode"] == "fpm_interpolation" + assert provenance["config"]["dcp"] == config.dcp + finally: + automatic.close() + + +def test_session_cache_preserves_fpm_identity(kimi_fpm_profile): + from dataclasses import replace + + from aisimulate.sdk.inference_session import InferenceSession + from aisimulate_core.sdk.config import RuntimeConfig + from aisimulate_core.sdk.errors import PerfDataNotAvailableError + from aisimulate_core.sdk.rust_engine_step import _engine_handle_cache_clear + + config = _model_config( + tp_size=8, + moe_tp_size=8, + moe_ep_size=1, + dcp_size=8, + forward_model="fpm", + fpm_config=FpmCompileConfig( + {"text_only": True, "unrecorded_quant_modes": ["fmha", "comm"]}, attention_backend="FLASHINFER_MLA" + ), + kvcache_quant_mode=common.KVCacheQuantMode.fp8, + moe_quant_mode=common.MoEQuantMode.w4a16_mxfp4, + ) + database = PerfDatabase(SYSTEM, BACKEND, VERSION, systems_root=kimi_fpm_profile.systems_paths[0]) + runtime = RuntimeConfig(batch_size=1, isl=128, osl=2) + + def prefill(model_config): + model = models.get_model(kimi_fpm_profile.model, model_config, BACKEND) + return InferenceSession(model, database, get_backend(BACKEND)).run_static_latency_only(runtime, "static_ctx") + + _engine_handle_cache_clear() + try: + assert prefill(config) == pytest.approx(7.0) # synthetic measured prefill row + for different in ( + replace(config, dcp_size=1), + replace(config, fpm_config=FpmCompileConfig({"text_only": True}, attention_backend="FLASHINFER_MLA")), + replace(config, fpm_config=FpmCompileConfig(config.fpm_config.options, attention_backend="auto")), + ): + with pytest.raises(PerfDataNotAvailableError, match="No FPM cell matches"): + prefill(different) + with pytest.raises(NotImplementedError, match="encoder/multimodal"): + prefill(replace(config, fpm_config=FpmCompileConfig())) + assert prefill(config) == pytest.approx(7.0) + finally: + _engine_handle_cache_clear() + + +def test_session_cache_separates_text_only_mode(fpm_session): + from dataclasses import replace + + from aisimulate_core.sdk.rust_engine_step import _cached_engine_handle, _engine_handle_cache_clear + + plain, database, _, isl, _ = fpm_session + text_only = models.get_model( + plain.model_path, replace(plain.config, fpm_config=FpmCompileConfig({"text_only": True})), BACKEND + ) + _engine_handle_cache_clear() + try: + plain_handle = _cached_engine_handle(plain, database) + text_handle = _cached_engine_handle(text_only, database) + assert text_handle is not plain_handle + # Both valid modes query the fixture's measured 40 ms prefill row. + for handle in (plain_handle, text_handle): + assert handle.predict_prefill_latency(bs=2, isl=isl, prefix=0) == pytest.approx(40.0) + assert _cached_engine_handle(plain, database) is plain_handle + finally: + _engine_handle_cache_clear() + + +def test_dcp_cannot_select_regression_without_a_profile(kimi_fpm_profile): + from dataclasses import replace + from pathlib import Path + + from aisimulate_core.sdk import RustForwardPassPerfModel + + config = kimi_fpm_profile + data = Path(config.systems_paths[0]) / "data" / SYSTEM / BACKEND / VERSION + for name in ("fpm_forward_perf.parquet", "fpm_forward_perf.metadata.json"): + (data / name).unlink() + for mode, fallback in (("fpm_regression", "deny"), ("auto", "deny"), ("fpm_interpolation", "allow")): + with pytest.raises(ValueError, match="DCP timing requires measured vLLM FPM interpolation"): + RustForwardPassPerfModel.best_available(replace(config, estimation_mode=mode, fallback_policy=fallback)) + + +@pytest.mark.parametrize( + "options,overrides,error", + [ + ({"unrecorded_quant_modes": ["fmha"]}, {"fmha_quant_mode": "bfloat16"}, "explicit quantization override"), + ({"unrecorded_quant_modes": ["fmha"]}, {"fpm_fmha_quant_mode": "fp8"}, "explicit quantization override"), + ({"unrecorded_quant_modes": ["comm"]}, {"comm_quant_mode": "half"}, "explicit quantization override"), + ({"unrecorded_quant_modes": ["unknown_mode"]}, {}, "unknown variant"), + (False, {}, "invalid FPM interpolation options"), + ], +) +def test_direct_compile_validates_fpm_options(options, overrides, error, monkeypatch): + from aisimulate_core.sdk import compile_engine + + def unexpected_model_build(*args, **kwargs): + pytest.fail("invalid options must be rejected before model construction") + + monkeypatch.setattr("aisimulate_core.sdk.engine.get_model", unexpected_model_build) + with pytest.raises(ValueError, match=error): + compile_engine(MODEL_PATH, SYSTEM, BACKEND, forward_model="fpm", fpm_options=options, **overrides) + + def test_explicit_selector_emits_matched_cell_warning_once(fpm_session, capfd): from aisimulate_core.sdk.rust_engine_step import _cached_engine_handle diff --git a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py index 886804c81f..973975149e 100644 --- a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py +++ b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py @@ -1010,11 +1010,11 @@ def test_forward_pass_config_preserves_existing_positional_arguments(tmp_path: P "wideep_num_slots": 64, } config = ForwardPassPerfModelConfig(*legacy_fields.values()) - assert vars(config) == {**legacy_fields, "fpm_fmha_quant_mode": None, "moe_kernel_source": None} + assert vars(config) == {**legacy_fields, "dcp": None, "fpm_fmha_quant_mode": None, "moe_kernel_source": None} source = " source_with_spaces " pinned = ForwardPassPerfModelConfig(*legacy_fields.values(), moe_kernel_source=source) - assert vars(pinned) == {**legacy_fields, "fpm_fmha_quant_mode": None, "moe_kernel_source": source} + assert vars(pinned) == {**legacy_fields, "dcp": None, "fpm_fmha_quant_mode": None, "moe_kernel_source": source} assert json.loads(json.dumps(pinned.to_dict()))["moe_kernel_source"] == source diff --git a/tests/sweeper/test_forward_pass_estimator.py b/tests/sweeper/test_forward_pass_estimator.py index a8d96f4d1c..ae856f2af6 100644 --- a/tests/sweeper/test_forward_pass_estimator.py +++ b/tests/sweeper/test_forward_pass_estimator.py @@ -599,6 +599,25 @@ def test_mixed_timing_still_enforces_cold_regression_on_default_role(): ForwardPassEstimatorResolver(space).resolve_candidate(sample) +def test_legacy_migration_validates_recorded_dcp(): + legacy = { + "schema_version": 1, + "model_name": "Qwen/Qwen3-32B", + "system_name": "h200_sxm", + "backend": "vllm", + "tp_size": 8, + "pp_size": 1, + "forward_model": "fpm", + } + for dcp in (None, 1, 8): + payload = legacy if dcp is None else {**legacy, "dcp_size": dcp} + config = ForwardPassPerfModelConfig.from_legacy_engine_config(payload, "decode") + assert config.tp == 8 and config.dcp == dcp + for dcp in (0, 3): + with pytest.raises(ValueError, match="dcp must be positive and divide tp"): + ForwardPassPerfModelConfig.from_legacy_engine_config({**legacy, "dcp_size": dcp}, "decode") + + def test_canonical_operation_diagnostics_include_native_sol_and_provenance(): model = RustForwardPassPerfModel.best_available( ForwardPassPerfModelConfig( diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index e30a4c835e..4696b15731 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -362,6 +362,7 @@ def test_custom_parallel_preset_lowers_as_flat_atomic_choices() -> None: } ) + config = CoreRecommendationConfig.model_validate(config.model_dump(mode="json")) lowered = recommendation_to_sweeper(config) assert lowered.search_space.flat_parallel_modes == ["agg"] @@ -1566,6 +1567,125 @@ def test_role_systems_roots_reach_prediction_and_search_preflight(tmp_path, mode assert branch.deployment_mode == ("agg" if mode == "aggregated" else "disagg") +def test_prediction_dcp_identity_roundtrip_and_explicit_capacity(): + from aisimulate.compiler import prediction_to_replay_spec + + raw = { + "engine": { + "model": "moonshotai/Kimi-K3", + "hardware": "gb300", + "backend": "vllm", + "backend_version": "0.29.0", + "context_length": 1048576, + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "workers": { + "aggregated": { + "parallelism": {"tensor": 8, "decode_context": 8, "moe_tensor": 8}, + "kv_cache": { + "block_size": 12288, + "capacity": {"type": "fixed", "blocks": 2175}, + }, + "scheduler": {"max_sequences": 32}, + } + }, + } + } + config = CorePredictionConfig.model_validate(raw) + restored = CorePredictionConfig.model_validate(config.model_dump(mode="json")) + deployment = prediction_to_replay_spec(restored).backend_deployment + timing = deployment.agg_engine_args["timing_model"]["config"] + assert timing["dcp"] == 8 and timing["tp"] == 8 + assert timing["estimation_mode"] == "fpm_interpolation" + assert deployment.agg_engine_args["num_gpu_blocks"] == 2175 + assert deployment.parallel_config["dcp"] == 8 + raw["engine"]["workers"]["aggregated"]["kv_cache"]["capacity"] = {"type": "default"} + with pytest.raises(ValueError, match="explicit KV block capacity"): + prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)) + cache = raw["engine"]["workers"]["aggregated"]["kv_cache"] + cache.update(capacity={"type": "fixed", "blocks": 2175}, host_offload={"num_host_blocks": 32}) + with pytest.raises(ValueError, match="explicit bytes_per_token"): + prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)) + cache["bytes_per_token"] = 1728 + deployment = prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)).backend_deployment + assert deployment.agg_engine_args["kv_cache_bytes_per_token"] == 1728 + + +@pytest.mark.parametrize( + "timing", + [ + { + "type": "fixed", + "prefill_ms": 1, + "decode_ms": 1, + "attention_backend": "FLASHINFER_MLA", + }, + {"type": "polynomial", "kvcache_quant_mode": "fp8"}, + ], +) +def test_prediction_rejects_identity_that_custom_timing_would_ignore(timing): + raw = _engine() + raw["workers"]["aggregated"] = {"timing": timing} + with pytest.raises(ValidationError, match="identity require default timing"): + CorePredictionConfig.model_validate({"engine": raw}) + + +@pytest.mark.parametrize("role", ["aggregated", "prefill", "decode"]) +def test_recommendation_rejects_unsupported_timing_identity(role): + from aisimulate.config.engine import WorkersRecommendationConfig + + with pytest.raises(ValidationError, match="identity is prediction-only"): + WorkersRecommendationConfig.model_validate( + { + role: { + "timing": { + "gemm_quant_mode": "fp8", + "kvcache_quant_mode": "fp8", + "attention_backend": "flashinfer", + } + } + } + ) + + +@pytest.mark.parametrize("mode", ["aggregated", "disaggregated"]) +def test_prediction_fp8_kv_transfer_uses_cache_precision(mode): + from aisimulate.compiler import prediction_to_replay_spec + + roles = ["aggregated"] if mode == "aggregated" else ["prefill", "decode"] + engine = { + "mode": mode, + "model": "Qwen/Qwen3-32B", + "hardware": "h200_sxm", + "backend": "vllm", + "backend_version": "0.24.0", + "context_length": 4096, + "workers": { + role: { + "parallelism": {"tensor": 2}, + "timing": {"kvcache_quant_mode": "fp8"}, + "kv_cache": { + "capacity": {"type": "fixed", "blocks": 256}, + **({"host_offload": {"num_host_blocks": 256}} if mode == "aggregated" else {}), + }, + } + for role in roles + }, + } + if mode == "disaggregated": + engine["kv_transfer"] = { + "bytes_per_token": "auto", + "bandwidth_gb_per_second": 50, + } + deployment = prediction_to_replay_spec(CorePredictionConfig.model_validate({"engine": engine})).backend_deployment + for role in roles: + args = getattr(deployment, f"{'agg' if role == 'aggregated' else role}_engine_args") + assert args["timing_model"]["config"]["kvcache_quant_mode"] == "fp8" + key = "kv_cache_bytes_per_token" if mode == "aggregated" else "kv_transfer_bytes_per_token" + # Qwen3-32B: 64 layers, K+V, 4 KV heads/rank at TP2, 128 dimensions, 1 byte FP8. + assert args[key] == 65536 + + @pytest.mark.parametrize( "option", [ diff --git a/tests/test_host_offload_cli.py b/tests/test_host_offload_cli.py index 6384557024..f9e4f244b2 100644 --- a/tests/test_host_offload_cli.py +++ b/tests/test_host_offload_cli.py @@ -185,6 +185,7 @@ def estimate(_model: str, **shape) -> int: "pp_size": 1, "moe_tp_size": 1, "moe_ep_size": 1, + "kvcache_quant_mode": None, } ]