From acb5b4624656e6b8caaaf8dc2a8062dc1d735e88 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 09:39:42 -0700 Subject: [PATCH 01/12] feat(perfmodel): load DCP self-benchmark profiles Signed-off-by: hongkuanz --- crates/core/src/lib.rs | 2 +- crates/core/src/perfmodel/config.rs | 2 + crates/core/src/perfmodel/engine/runtime.rs | 1 + crates/core/src/perfmodel/engine/spec.rs | 14 +- crates/core/src/perfmodel/fpm/config.rs | 48 ++++ crates/core/src/perfmodel/fpm/estimator.rs | 22 +- crates/core/src/perfmodel/fpm/model.rs | 8 + crates/core/src/perfmodel/fpm/tests.rs | 1 + crates/core/src/perfmodel/memory.rs | 1 + crates/core/src/perfmodel/mod.rs | 2 +- .../src/perfmodel/operators/fpm_forward.rs | 31 ++- .../perfmodel/perf_database/fpm_forward.rs | 225 ++++++++++++++++-- crates/core/src/perfmodel/py.rs | 19 ++ crates/core/src/python.rs | 8 + .../core/tests/perfmodel/memory_round_trip.rs | 1 + docs/core-api.md | 71 +++++- .../src/aiconfigurator_core/sdk/config.py | 8 + .../sdk/config_builders.py | 2 + .../src/aiconfigurator_core/sdk/engine.py | 13 +- .../sdk/models/__init__.py | 9 +- .../sdk/operations/fpm_forward.py | 11 +- .../sdk/rust_engine_step.py | 1 + python/aisimulate/src/aisimulate/aic.py | 4 + python/aisimulate/src/aisimulate/compiler.py | 37 +++ .../src/aisimulate/config/engine.py | 19 ++ python/aisimulate/src/aisimulate/recommend.py | 19 ++ .../sweeper/forward_pass_estimator.py | 2 + .../cross_package/test_core_public_api.py | 1 + .../tests/unit/sdk/test_fpm_forward.py | 118 +++++++++ tests/test_cli_config.py | 48 ++++ 30 files changed, 718 insertions(+), 30 deletions(-) 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 723b1ca756..fc16c8a571 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -208,6 +208,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/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index b6cd090b18..94562584d3 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -1933,6 +1933,7 @@ mod tests { forward_model: 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/engine/spec.rs b/crates/core/src/perfmodel/engine/spec.rs index 442525ef45..38709beb91 100644 --- a/crates/core/src/perfmodel/engine/spec.rs +++ b/crates/core/src/perfmodel/engine/spec.rs @@ -769,6 +769,7 @@ mod tests { forward_model: None, kv_block_size: Some(64), parallel: ParallelMapping { + dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(8), @@ -1166,9 +1167,20 @@ mod tests { EngineSpec::from_bincode(&bytes), Err(AicError::UnsupportedSchemaVersion { got: 17, - expected: 18, + 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); + } } diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 37822037eb..f791e040ff 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)] @@ -155,6 +158,7 @@ impl ForwardPassPerfModelConfig { tp: 1, pp: 1, attention_dp: 1, + dcp: None, moe_tp_size: None, moe_ep_size: None, gemm_quant_mode: None, @@ -217,6 +221,24 @@ 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")); + } + for mode in &self + .estimator_config + .fpm_interpolation + .unrecorded_quant_modes + { + let explicit = match mode { + super::UnrecordedFpmQuantMode::Fmha => &self.fmha_quant_mode, + super::UnrecordedFpmQuantMode::Comm => &self.comm_quant_mode, + }; + if explicit.is_some() { + return Err(invalid_config( + "an unrecorded FPM quant mode cannot have an explicit quantization override", + )); + } + } 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", @@ -410,4 +432,30 @@ 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()); + } } diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs index a856f9751d..04fb90a7d4 100644 --- a/crates/core/src/perfmodel/fpm/estimator.rs +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -23,8 +23,26 @@ pub struct EstimatorConfig { pub struct OpLevelConfig {} #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct FpmInterpolationConfig {} +#[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, + /// 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, +} /// These weights currently affect regression. Native correction keeps its /// established workload coordinates until a replacement is accuracy-qualified. diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index b99d4525b3..fb214141f0 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -645,6 +645,14 @@ impl ForwardPassPerfModel { fn build_native_candidate( config: &ForwardPassPerfModelConfig, ) -> Result<(Engine, PathBuf), AicError> { + if config.dcp.is_some_and(|dcp| dcp > 1) + && (config.estimation_mode != EstimationMode::FpmInterpolation + || config.backend != crate::BackendKind::Vllm) + { + return Err(AicError::UnsupportedModel( + "DCP timing requires measured vLLM FPM interpolation".into(), + )); + } if config.estimation_mode == EstimationMode::FpmInterpolation && config.nextn != 0 { return Err(AicError::UnsupportedModel( "FPM interpolation does not support MTP".into(), diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index 78c22682c3..2883e5b1c6 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -108,6 +108,7 @@ fn fixture_engine_config() -> EngineConfig { forward_model: 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 491dde5dd6..df76215c2d 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -509,6 +509,7 @@ mod tests { forward_model: 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 a563e2e3db..d433434949 100644 --- a/crates/core/src/perfmodel/mod.rs +++ b/crates/core/src/perfmodel/mod.rs @@ -54,7 +54,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/fpm_forward.rs b/crates/core/src/perfmodel/operators/fpm_forward.rs index 407a904142..a047ea343c 100644 --- a/crates/core/src/perfmodel/operators/fpm_forward.rs +++ b/crates/core/src/perfmodel/operators/fpm_forward.rs @@ -102,6 +102,20 @@ fn data_err(msg: String) -> AicError { } impl FpmForwardOp { + fn sol_at(&self, db: &PerfDatabase, coords: &[f64]) -> Result { + if self + .match_identity + .get(crate::perf_database::fpm_forward::FPM_CELL_MATCH_COLUMNS.len()) + .and_then(|v| v.parse::().ok()) + .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) + } + /// Mirror of Python `FPMForwardOp.query`: validate kwargs, map to /// iteration-total coordinates, resolve against the selected cell. pub fn query( @@ -341,8 +355,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 @@ -403,7 +417,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(); @@ -665,6 +679,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.match_identity.push("4".into()); + 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, diff --git a/crates/core/src/perfmodel/perf_database/fpm_forward.rs b/crates/core/src/perfmodel/perf_database/fpm_forward.rs index 17d8735bfb..f83d734a3f 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -54,9 +54,9 @@ pub const FPM_FORWARD_SCHEMA_NAME: &str = "aic_fpm_forward_perf"; pub const FPM_FORWARD_SCHEMA_VERSION: u64 = 6; 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 @@ -85,6 +85,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. @@ -284,6 +285,7 @@ impl FpmForwardTable { available.truncate(8); let identity: Vec = FPM_CELL_MATCH_COLUMNS .iter() + .chain(std::iter::once(&"dcp")) .zip(match_identity) .map(|(c, v)| format!("{c}={v:?}")) .collect(); @@ -326,15 +328,18 @@ 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 { + 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, AicError> { +) -> Result { if !metadata_path.exists() { return Err(structural(format!( "FPM database is missing its metadata sidecar: {}. \ @@ -381,11 +386,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() ))); @@ -415,7 +422,31 @@ fn validate_sidecar( parquet_path.parent().unwrap_or(parquet_path).display() ))); } - Ok(json_uint(metadata.get("row_count"))) + 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 { + 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 @@ -473,8 +504,7 @@ fn load_pair( return Ok(None); } let metadata_path = parquet_path.with_extension("metadata.json"); - let sidecar_row_count = - 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. @@ -523,13 +553,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!( @@ -648,7 +684,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) { @@ -660,6 +696,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() { @@ -670,7 +734,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")?, @@ -699,6 +763,8 @@ 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 @@ -1104,6 +1170,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)>, } impl Default for RowSpec { @@ -1126,6 +1194,8 @@ pub(crate) mod tests { system: "b200_sxm", backend: "vllm", kv_seed_regime: None, + dcp: None, + measurement: None, } } } @@ -1186,6 +1256,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 schema = "message schema { REQUIRED BINARY cell_id (UTF8); REQUIRED BINARY model_path (UTF8); @@ -1223,6 +1295,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 = Arc::new(parse_message_type(&schema).expect("schema must parse")); let file = std::fs::File::create(&parquet_path).expect("create parquet"); let mut writer = @@ -1358,6 +1440,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"); @@ -1411,6 +1518,94 @@ pub(crate) mod tests { ) } + #[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") + .is_err() + ); + for dcp in [1, 4] { + let mut identity = default_identity(4); + identity.push(dcp.to_string()); + assert!(table.select_cell(&identity, "org/model-a").is_ok()); + } + 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"); diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index b3cb648cff..34fd6091aa 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1039,6 +1039,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, @@ -1084,6 +1086,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, @@ -1372,6 +1376,15 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result) -> En tp_size: config.parallel.tp_size, pp_size: config.parallel.pp_size, attention_dp_size: config.parallel.attention_dp_size.unwrap_or(1), + dcp_size: config.parallel.dcp_size, + fpm_options: crate::FpmInterpolationConfig::default(), moe_tp_size: config.parallel.moe_tp_size, moe_ep_size: config.parallel.moe_ep_size, gemm_quant_mode: gemm_quant_name(config.quantization.weight_dtype.as_ref()) @@ -1700,6 +1717,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, @@ -1932,6 +1950,7 @@ mod tests { forward_model: 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 c1fd2b6a7f..55a965760f 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -160,6 +160,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, @@ -261,6 +263,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(), @@ -672,6 +675,10 @@ fn materialize_aic_capacity( if capacity_is_explicit { 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; @@ -2042,6 +2049,7 @@ mod tests { backend_version: None, pp: 1, attention_dp: 1, + dcp: None, moe_tp_size: None, moe_ep_size: None, gemm_dtype: None, diff --git a/crates/core/tests/perfmodel/memory_round_trip.rs b/crates/core/tests/perfmodel/memory_round_trip.rs index 81ecd3a282..7dc29d8c17 100644 --- a/crates/core/tests/perfmodel/memory_round_trip.rs +++ b/crates/core/tests/perfmodel/memory_round_trip.rs @@ -88,6 +88,7 @@ fn request(tolerance_fraction: Option) -> KvCacheEstimateRequest { forward_model: None, kv_block_size: None, parallel: ParallelMapping { + dcp_size: None, tp_size: 1, pp_size: 1, attention_dp_size: Some(1), diff --git a/docs/core-api.md b/docs/core-api.md index 26c304e202..6227812af1 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -204,8 +204,13 @@ nested paths. The supported namespaces are: - `correction`: `enabled` (true), independent `sampling`, `min_observations` (5), `factor_bounds` (min 0.5, max 2.0), and the existing `max_num_tokens` (8192), `max_batch_size` (512), and `max_kv_tokens` (2000000) ranges. -- `op_level` and `fpm_interpolation`: reserved typed namespaces with no - additional knobs yet; unknown fields are rejected. +- `op_level`: reserved typed namespace with no additional knobs yet. +- `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. Sampling defaults to `bins_per_axis: [4, 4]` and `max_observations: 64` per logical store. Rectangular grids are supported. Regression uses dynamic @@ -225,6 +230,68 @@ 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. 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/python/aisimulate/src/aiconfigurator_core/sdk/config.py b/python/aisimulate/src/aiconfigurator_core/sdk/config.py index 997bf0e3cd..a630e575f6 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/config.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/config.py @@ -60,6 +60,10 @@ class ModelConfig: Model configuration. """ + dcp_size: int | None = field(default=None, kw_only=True) + fpm_text_only: bool = field(default=False, kw_only=True) + fpm_unrecorded_quant_modes: tuple[str, ...] = field(default=(), kw_only=True) + fpm_attention_backend: str | None = field(default=None, kw_only=True) tp_size: int = 1 pp_size: int = 1 gemm_quant_mode: common.GEMMQuantMode | None = None @@ -132,6 +136,10 @@ class ModelConfig: system: str | None = None 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/aiconfigurator_core/sdk/config_builders.py b/python/aisimulate/src/aiconfigurator_core/sdk/config_builders.py index a6b0c3a4ed..d7dcce2dc7 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/config_builders.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/config_builders.py @@ -40,10 +40,12 @@ def build_model_config( enable_encoder_dp: bool = True, attention_backend: str | None = None, speculation=None, + dcp_size: int | None = None, ) -> ModelConfig: """Build a ModelConfig with optional quant mode overrides.""" 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/aiconfigurator_core/sdk/engine.py b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py index 8ad136b877..49ba4155bb 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/engine.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py @@ -159,7 +159,7 @@ def _fpm_spec_dict(op: FPMForwardOp) -> dict: "phase": op._phase, "model_path": op._model_path, "match_identity": list(op._match_identity), - "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. @@ -324,6 +324,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)), @@ -407,6 +408,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, @@ -437,12 +440,17 @@ def compile_engine( # does not take a model_path (quant inference is done inside `get_model`). from aiconfigurator_core.sdk.speculation import SpeculationConfig + fpm_options = fpm_options or {} + recorded_attention_backend = attention_backend + if forward_model == "fpm" and 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: 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, @@ -456,6 +464,9 @@ def compile_engine( attention_backend=attention_backend, speculation=resolved_speculation, ) + model_config.fpm_text_only = fpm_options.get("text_only", False) + model_config.fpm_unrecorded_quant_modes = tuple(fpm_options.get("unrecorded_quant_modes", ())) + model_config.fpm_attention_backend = recorded_attention_backend # 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/aiconfigurator_core/sdk/models/__init__.py b/python/aisimulate/src/aiconfigurator_core/sdk/models/__init__.py index 42279e12bb..59bbc55cca 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/models/__init__.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/models/__init__.py @@ -69,7 +69,7 @@ def _apply_forward_model_fpm(model: BaseModel) -> BaseModel: public model type are unchanged.""" from aiconfigurator_core.sdk.operations.fpm_forward import FPMForwardOp - if model.encoder_ops: + if model.encoder_ops and not model.config.fpm_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'." @@ -108,6 +108,11 @@ def _apply_forward_model_fpm(model: BaseModel) -> BaseModel: decode_op = FPMForwardOp( "decode", model.config, model.model_path, sol_ops=generation_ops, weight_bytes=weight_bytes ) + # 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] @@ -134,6 +139,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/aiconfigurator_core/sdk/operations/fpm_forward.py b/python/aisimulate/src/aiconfigurator_core/sdk/operations/fpm_forward.py index 8874a609b2..ff0968fec5 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/operations/fpm_forward.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/operations/fpm_forward.py @@ -125,11 +125,12 @@ def __init__( self._phase = phase self._model_path = str(model_path) self._weight_bytes = float(weight_bytes) + self._resident_weight_bytes = self._weight_bytes self._match_identity = ( _norm_identity(model_config.gemm_quant_mode), _norm_identity(model_config.moe_quant_mode), - _norm_identity(model_config.fmha_quant_mode), - _norm_identity(model_config.comm_quant_mode), + "" if "fmha" in model_config.fpm_unrecorded_quant_modes else _norm_identity(model_config.fmha_quant_mode), + "" if "comm" in model_config.fpm_unrecorded_quant_modes else _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), @@ -140,10 +141,14 @@ def __init__( _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_backend_request( + model_config.fpm_attention_backend or model_config.attention_backend, engine_default="flashinfer" + ), _norm_identity(bool(getattr(model_config, "enable_wideep", False))), _norm_identity(bool(getattr(model_config, "enable_eplb", False))), ) + if model_config.dcp_size is not None: + self._match_identity += (str(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/aiconfigurator_core/sdk/rust_engine_step.py b/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py index 7682727942..9ae5406d4c 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aiconfigurator_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 diff --git a/python/aisimulate/src/aisimulate/aic.py b/python/aisimulate/src/aisimulate/aic.py index 8f1cd255b0..4d8f80767b 100644 --- a/python/aisimulate/src/aisimulate/aic.py +++ b/python/aisimulate/src/aisimulate/aic.py @@ -146,6 +146,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 39cd59208f..fddc76a31c 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -353,6 +353,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 {}), } @@ -370,11 +371,24 @@ def _worker_performance_model_metadata( "model_path": engine.model, "tp_size": parallel.tensor, "attention_dp_size": parallel.attention_data, + **({"dcp": parallel.decode_context} if parallel.decode_context is not None else {}), "moe_tp_size": parallel.moe_tensor if sharded_moe else None, "moe_ep_size": parallel.moe_expert if sharded_moe else None, "nextn": None, **({"speculation": engine.speculation.cost_config()} if engine.speculation is not None else {}), "forward_model": worker.timing.forward_model, + **{ + field: getattr(worker.timing, field) + for field in ( + "gemm_quant_mode", + "moe_quant_mode", + "fmha_quant_mode", + "kvcache_quant_mode", + "comm_quant_mode", + "attention_backend", + ) + if getattr(worker.timing, field) is not None + }, }, } @@ -391,6 +405,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 @@ -478,9 +513,11 @@ 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, + **{field: getattr(timing, field) for field in identity_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, diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 4407fe8c60..5221a6606f 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -64,6 +64,7 @@ class AFDSearchRecommendationConfig(StrictModel): class ParallelismPredictionConfig(StrictModel): + decode_context: Annotated[int, Field(strict=True, gt=0)] | None = None replicas: PositiveInt = 1 tensor: PositiveInt = 1 pipeline: PositiveInt = 1 @@ -157,6 +158,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) estimation_mode: Literal["auto", "op_level", "fpm_interpolation", "fpm_regression"] | None = None @@ -187,6 +194,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": diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index 96a2926abd..c221371768 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -853,6 +853,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"], @@ -878,6 +892,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/sweeper/forward_pass_estimator.py b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py index c7582502a6..631fabdfa8 100644 --- a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py +++ b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py @@ -76,6 +76,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), @@ -141,6 +142,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/tests/cross_package/test_core_public_api.py b/python/aisimulate/tests/cross_package/test_core_public_api.py index 16e153f45b..df75b5620a 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -91,6 +91,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/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 573c5b7d71..09cea2516f 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -603,3 +603,121 @@ def test_genonly_step_works_with_single_op(self, fpm_session): ) assert per_op["fpm_forward_decode"] == pytest.approx(7.0) 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_text_only=True, + fpm_unrecorded_quant_modes=("fmha", "comm"), + fpm_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 identity[-1] == "8" + assert cfg.total_gpus_per_worker == 8 + from aiconfigurator_core.sdk.engine import _fpm_spec_dict + + prefill = model.context_ops[0] + resident = _fpm_spec_dict(prefill)["FpmForward"]["weight_bytes"] + assert resident == prefill.get_weights() + sum(op.get_weights() for op in model.encoder_ops) + + +def test_canonical_kimi_dcp_profile_queries_measured_points(tmp_path): + from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + 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(dcp=8, measurement_policy="kvwarm_median_of_3", measurement_repeats=3) + # 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": 8}, + }, + ) + config = 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=8, + 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}, + }, + ) + 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"] == 8 + finally: + model.close() + from dataclasses import replace + + from aiconfigurator_core.sdk.errors import PerfDataNotAvailableError + + for dcp in (None, 1): + 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"] == 8 + finally: + automatic.close() diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index 66e3b37be9..f8ef5d5137 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -1304,3 +1304,51 @@ def test_role_systems_roots_reach_prediction_and_search_preflight(tmp_path, mode (branch,) = enumerate_branches(smart, max_seq_len=4096) assert branch.parallel_configs 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)) + + +@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}) From f5c2c5e3678b5f107513eedd3db8ad57108ec271 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 10:43:11 -0700 Subject: [PATCH 02/12] fix(perfmodel): preserve DCP and timing configuration Signed-off-by: hongkuanz --- docs/core-api.md | 12 ++- .../src/aiconfigurator_core/sdk/engine.py | 2 +- .../sdk/rust_engine_step.py | 4 + python/aisimulate/src/aisimulate/aic.py | 7 ++ python/aisimulate/src/aisimulate/compiler.py | 3 + .../src/aisimulate/config/engine.py | 28 ++++++- .../tests/unit/sdk/test_fpm_forward.py | 74 ++++++++++++++++--- tests/test_cli_config.py | 61 +++++++++++++++ 8 files changed, 174 insertions(+), 17 deletions(-) diff --git a/docs/core-api.md b/docs/core-api.md index 6227812af1..49d74365b7 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -283,10 +283,14 @@ 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. 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. +automatic DCP/hybrid capacity sizing is not implemented. Host offload or P/D +transfer also requires explicit KV bytes per token with DCP. +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 diff --git a/python/aisimulate/src/aiconfigurator_core/sdk/engine.py b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py index 49ba4155bb..d26c17065b 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/engine.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py @@ -442,7 +442,7 @@ def compile_engine( fpm_options = fpm_options or {} recorded_attention_backend = attention_backend - if forward_model == "fpm" and backend == "vllm" and attention_backend == "FLASHINFER_MLA": + 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 diff --git a/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py b/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py index 9ae5406d4c..fe03ecf7bd 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/rust_engine_step.py @@ -1157,6 +1157,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)), @@ -1203,6 +1204,9 @@ def _engine_config_json(model: Any, database: Any) -> str: "comm": _raw_quant_name(getattr(model_config, "comm_quant_mode", None)), }, "model_config": { + "fpm_text_only": getattr(model_config, "fpm_text_only", False), + "fpm_unrecorded_quant_modes": getattr(model_config, "fpm_unrecorded_quant_modes", ()), + "fpm_attention_backend": getattr(model_config, "fpm_attention_backend", None), "cp_style": getattr(model_config, "cp_style", None), "workload_distribution": getattr(model_config, "workload_distribution", None), "overwrite_num_layers": getattr(model_config, "overwrite_num_layers", None), diff --git a/python/aisimulate/src/aisimulate/aic.py b/python/aisimulate/src/aisimulate/aic.py index 4d8f80767b..361da5d965 100644 --- a/python/aisimulate/src/aisimulate/aic.py +++ b/python/aisimulate/src/aisimulate/aic.py @@ -289,6 +289,7 @@ def estimate_kv_bytes_per_token( pp_size: int, moe_tp_size: int = 1, moe_ep_size: int = 1, + kvcache_quant_mode: str | None = None, ) -> int: """Derive per-rank KV bytes/token from the resolved Hugging Face config.""" @@ -302,6 +303,12 @@ def estimate_kv_bytes_per_token( moe_ep_size=moe_ep_size, allow_hf_config_download=True, ) + quant_mode = _quant_mode_name("kvcache", kvcache_quant_mode) + if quant_mode is not None: + from aiconfigurator_core.sdk.common import KVCacheQuantMode + + # This instance estimates KV geometry only, never model weights. + estimator.dtype_bytes = KVCacheQuantMode[quant_mode].value.memory value = estimator.kv_bytes_per_token() if value is None or value <= 0: raise ValueError(f"could not derive KV bytes per token for model {model_name!r}") diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index fddc76a31c..732948b142 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -572,12 +572,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=worker.timing.kvcache_quant_mode, ) diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 5221a6606f..e546bcc5c3 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -63,8 +63,7 @@ class AFDSearchRecommendationConfig(StrictModel): max_candidates: PositiveInt = 10_000 -class ParallelismPredictionConfig(StrictModel): - decode_context: Annotated[int, Field(strict=True, gt=0)] | None = None +class ParallelismPresetConfig(StrictModel): replicas: PositiveInt = 1 tensor: PositiveInt = 1 pipeline: PositiveInt = 1 @@ -73,6 +72,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 @@ -398,7 +401,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 @@ -501,6 +504,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/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 09cea2516f..c6c1ddd20a 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -632,8 +632,11 @@ def test_text_only_kimi_profile_preserves_encoder_weights_and_dcp_identity(): assert resident == prefill.get_weights() + sum(op.get_weights() for op in model.encoder_ops) -def test_canonical_kimi_dcp_profile_queries_measured_points(tmp_path): - from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel +@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() @@ -654,7 +657,9 @@ def test_canonical_kimi_dcp_profile_queries_measured_points(tmp_path): _row("decode", 1, 0, 128, 3.0, model_path="moonshotai/Kimi-K3", identity=identity), ] for row in rows: - row.update(dcp=8, measurement_policy="kvwarm_median_of_3", measurement_repeats=3) + 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( @@ -662,10 +667,10 @@ def test_canonical_kimi_dcp_profile_queries_measured_points(tmp_path): rows, sidecar_overrides={ "measurement_policy": "per_row_single_sample_or_median_of_3", - "configuration_selector": {"dcp": 8}, + "configuration_selector": {"dcp": recorded_dcp} if recorded_dcp is not None else {}, }, ) - config = ForwardPassPerfModelConfig( + return ForwardPassPerfModelConfig( model="moonshotai/Kimi-K3", system=SYSTEM, backend=BACKEND, @@ -674,7 +679,7 @@ def test_canonical_kimi_dcp_profile_queries_measured_points(tmp_path): tp=8, moe_tp_size=8, moe_ep_size=1, - dcp=8, + dcp=recorded_dcp, kvcache_quant_mode="fp8", attention_backend="FLASHINFER_MLA", estimation_mode="fpm_interpolation", @@ -684,6 +689,13 @@ def test_canonical_kimi_dcp_profile_queries_measured_points(tmp_path): "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( @@ -704,20 +716,64 @@ def test_canonical_kimi_dcp_profile_queries_measured_points(tmp_path): }, } ) == pytest.approx(3.0) - assert model.diagnostics()["provenance"]["config"]["dcp"] == 8 + assert model.diagnostics()["provenance"]["config"]["dcp"] == config.dcp finally: model.close() from dataclasses import replace from aiconfigurator_core.sdk.errors import PerfDataNotAvailableError - for dcp in (None, 1): + 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"] == 8 + assert provenance["config"]["dcp"] == config.dcp finally: automatic.close() + + +def test_session_cache_preserves_fpm_identity(kimi_fpm_profile): + from dataclasses import replace + + from aiconfigurator.sdk.config import RuntimeConfig + from aiconfigurator.sdk.inference_session import InferenceSession + from aiconfigurator_core.sdk.errors import PerfDataNotAvailableError + from aiconfigurator_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_text_only=True, + fpm_unrecorded_quant_modes=("fmha", "comm"), + fpm_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_unrecorded_quant_modes=()), + replace(config, fpm_attention_backend="auto"), + ): + with pytest.raises(PerfDataNotAvailableError, match="No FPM cell matches"): + prefill(different) + assert prefill(config) == pytest.approx(7.0) + finally: + _engine_handle_cache_clear() diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index f8ef5d5137..c80ef60b37 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -307,6 +307,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"] @@ -1338,6 +1339,13 @@ def test_prediction_dcp_identity_roundtrip_and_explicit_capacity(): 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( @@ -1352,3 +1360,56 @@ def test_prediction_rejects_identity_that_custom_timing_would_ignore(timing): 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 From 5ab66c4f32e6d0e1bf3bf74fc43a7a9c44780adf Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 11:16:15 -0700 Subject: [PATCH 03/12] docs(fpm): update onboarding for measured DCP profiles Signed-off-by: hongkuanz --- python/aisimulate/docs/fpm/README.md | 12 +- .../docs/fpm/end-to-end-workflow.md | 444 ++++++++++++++---- 2 files changed, 365 insertions(+), 91 deletions(-) diff --git a/python/aisimulate/docs/fpm/README.md b/python/aisimulate/docs/fpm/README.md index 0afd89ada8..5fcfd35a3c 100644 --- a/python/aisimulate/docs/fpm/README.md +++ b/python/aisimulate/docs/fpm/README.md @@ -5,16 +5,18 @@ 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 [end-to-end FPM workflow](end-to-end-workflow.md) to import +published self-benchmark data or collect new whole-forward measurements, load +the performance-data pair through the canonical SDK, and run an AISimulate +prediction. The [Kimi K3 TP8+DCP8 quickstart](end-to-end-workflow.md#use-an-existing-profile-kimi-k3-tp8dcp8) +uses existing data and runs entirely on CPU. The guide also includes collection +prerequisites, expected artifacts, acceptance checks, and recovery steps. 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 index 211f4d70c1..c277df36cc 100644 --- a/python/aisimulate/docs/fpm/end-to-end-workflow.md +++ b/python/aisimulate/docs/fpm/end-to-end-workflow.md @@ -5,25 +5,42 @@ SPDX-License-Identifier: Apache-2.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. +Use measured forward-pass data to construct a performance model and predict +request-level serving behavior. Choose the path that matches your starting point: + +| Starting point | Follow | +| --- | --- | +| Published self-benchmark data; no collection environment needed | [Check the package](#1-check-the-aisimulate-commands), then [Kimi K3 TP8+DCP8 quickstart](#use-an-existing-profile-kimi-k3-tp8dcp8) | +| Your own schema-v6 Parquet/metadata pair | [Inspect the pair](#5-inspect-the-published-pair), [construct the model](#6-load-the-new-data-and-query-one-forward-pass), then [predict](#7-run-ais-predict-using-the-same-external-data); substitute your recorded identity | +| New measurements to collect | [Prepare the environment](#1-check-the-aisimulate-commands), then follow steps 2–9 for MiniMax-M2.7 on four H200 GPUs | ```text -Freeze plan -> GPU smoke -> GPU collection -> validate and publish pair - | - v -Independent measurements <- compare <- AIS predict <- SDK FPM query +Existing measured pair ------------------------+ + | +GPU plan -> smoke -> self-benchmark -> publish -+-> SDK query -> AIS predict + | + v + Independent accuracy check ``` -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. +The data contract is `fpm_forward_perf.parquet` plus its matching +`fpm_forward_perf.metadata.json`. The Rust performance model selects a matching +cell and uses measured lookup/interpolation and, where supported, SOL transfer. +An existing model architecture needs no new model class or parallelization +mapping to consume another measured configuration. Place the pair in a systems +root and pass its exact identity to +`RustForwardPassPerfModel.best_available(ForwardPassPerfModelConfig(...))`. +There is no separate FPM registration service or `train` command. + +Raw benchmark JSON is not a loadable FPM pair. Convert it to the +[publication schema](../../collector/README.md#whole-forward-fpm-campaign) +with its measurement provenance before step 5. Collector performs this +conversion for the campaigns it runs in steps 2–4. + +A new architecture still requires a model definition; see +[How to add a new model](../add_a_new_model.md). +[Online telemetry regression](aic-fpm-regression-design.md), which learns from +`tune_with_fpms`, is a separate workflow. ## Scope and expected outcomes @@ -36,38 +53,283 @@ transfer. There is no separate `train` or `fit` command in this workflow. | 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. +The collection walkthrough uses vLLM with +`tp=4, pp=1, dp=1, moe_tp=4, moe_ep=1, cp=1`, checkpoint-default quantization, +and one worker replica. The dedicated Collector currently supports `PP=CP=1`; +its topology presets do not collect DCP configurations. The DCP quickstart +consumes an existing Dynamo self-benchmark profile. ## 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. +Use Python 3.11–3.13 and an AISimulate build containing +[PR #284](https://github.com/ai-dynamo/aisimulate/pull/284). A package version +alone does not establish that an unreleased feature is included. If you already +have that build, activate its environment and continue with the checks below. +For a source installation while the PR is open: + +```bash +git clone https://github.com/ai-dynamo/aisimulate.git +cd aisimulate +git fetch origin pull/284/head +git switch --detach FETCH_HEAD +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"))' +python -c 'from aisimulate_core.sdk import ForwardPassPerfModelConfig; assert "dcp" in ForwardPassPerfModelConfig.__dataclass_fields__' aisimulate predict --help +``` + +Use the same activated Bash or Zsh session for the remaining commands. +For existing data, continue with the quickstart below or step 5. + +For **new GPU collection**, additionally check: + +```bash cd "$AIS_REPO/python/aisimulate" -python3 -m collector.fpm_forward --help +python -m collector.fpm_forward --help +``` + +The collection host needs model configuration metadata and `kubectl` access. +The MiniMax example needs four H200 GPUs on one node and an existing namespace, +checkpoint PVC/path, collection image, and matching Generator target release. +Record the image digest, model revision, GPU type, and resolved configuration. +Before launching, read [recovery and cleanup](#8-recovery-and-cleanup). + +## Use an existing profile: Kimi K3 TP8+DCP8 + +This CPU-only quickstart 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. + +### Download and install the measured pair + +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("aiconfigurator_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 ``` -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. +**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. -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. +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: -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. +```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. + +### Query prefill and decode through the canonical API + +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. + +### Run a small Kimi Replay + +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" </ fpm_forward_perf.parquet fpm_forward_perf.metadata.json @@ -329,12 +595,13 @@ identity = [ "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(), + "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 @@ -352,7 +619,7 @@ 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` | +| 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 @@ -372,19 +639,30 @@ import math import sys from pathlib import Path -from aisimulate_core.sdk import EngineHandle +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel -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, +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}}, ) -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})) +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 ``` @@ -398,7 +676,7 @@ 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/aiconfigurator_core/sdk/engine.py). +[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 @@ -411,7 +689,8 @@ 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 +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. @@ -424,6 +703,11 @@ engine: 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: @@ -440,42 +724,20 @@ engine: block_size: 64 prefix_caching: false capacity: {type: default} - timing: {type: default, forward_model: fpm} + timing: {type: default} 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: +`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 \ -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 aiconfigurator_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 +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: @@ -484,8 +746,8 @@ PY 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 +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. @@ -494,11 +756,21 @@ 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). +New configurations use `estimation_mode` and `fallback_policy`. Without +explicit selection, the defaults are `auto` + `deny`: construction tries +op-level, measured FPM interpolation, then regression. For a measured-profile +workflow, explicitly select `fpm_interpolation` + `deny` as above. An untrained +regression cannot run offline prediction, and a selected estimator never switches +models on a query miss. + +For nondefault precision, put `gemm_quant_mode`, `moe_quant_mode`, +`fmha_quant_mode`, `kvcache_quant_mode`, `comm_quant_mode`, and +`attention_backend` under `engine.workers..timing`, using the exact +recorded identity. These overrides require default timing and regular language +workers. Recommendation rejects them until its feasibility preflight supports +the same identity. See the [Core API](../../../../docs/core-api.md) for estimator +controls and the [DCP quickstart](#use-an-existing-profile-kimi-k3-tp8dcp8) for +a complete configuration with text-only and unrecorded-precision settings. ## 8. Recovery and cleanup From 030b520327faa75b8097809c14209ecb27bd76b7 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 11:36:45 -0700 Subject: [PATCH 04/12] docs(fpm): organize onboarding around collection scope and workflow Signed-off-by: hongkuanz --- python/aisimulate/docs/fpm/README.md | 17 +- .../docs/fpm/end-to-end-workflow.md | 507 +++++++++++++----- 2 files changed, 394 insertions(+), 130 deletions(-) diff --git a/python/aisimulate/docs/fpm/README.md b/python/aisimulate/docs/fpm/README.md index 5fcfd35a3c..43ac9c1378 100644 --- a/python/aisimulate/docs/fpm/README.md +++ b/python/aisimulate/docs/fpm/README.md @@ -5,12 +5,17 @@ SPDX-License-Identifier: Apache-2.0 # Forward-pass models -Start with the [end-to-end FPM workflow](end-to-end-workflow.md) to import -published self-benchmark data or collect new whole-forward measurements, load -the performance-data pair through the canonical SDK, and run an AISimulate -prediction. The [Kimi K3 TP8+DCP8 quickstart](end-to-end-workflow.md#use-an-existing-profile-kimi-k3-tp8dcp8) -uses existing data and runs entirely on CPU. The guide also includes collection -prerequisites, expected artifacts, acceptance checks, and recovery steps. +Start with the [self-benchmark and FPM onboarding guide](end-to-end-workflow.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](end-to-end-workflow.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: diff --git a/python/aisimulate/docs/fpm/end-to-end-workflow.md b/python/aisimulate/docs/fpm/end-to-end-workflow.md index c277df36cc..2842a3e7a5 100644 --- a/python/aisimulate/docs/fpm/end-to-end-workflow.md +++ b/python/aisimulate/docs/fpm/end-to-end-workflow.md @@ -3,63 +3,122 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# Collect FPM data and use it in AISimulate - -Use measured forward-pass data to construct a performance model and predict -request-level serving behavior. Choose the path that matches your starting point: - -| Starting point | Follow | +# 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 + +An FPM predicts **engine iteration / forward-pass latency** for scheduled work. +It supplies the cost of an iteration to the simulator; the simulator still +decides which work runs, when it runs, and how requests progress. + +| Difference you want to model | What self-collection can supply | What still needs separate support | +| --- | --- | --- | +| Kernel implementation, fusion, quantization, CUDA Graph/eager behavior, engine-version changes | Measured iteration costs for the recorded configuration and shapes | A working benchmark path and matching profile identity | +| TP/EP/DCP execution within a supported engine iteration | The observed compute and collective costs inside the measurement boundary | Correct topology, rank aggregation, cache state, and runtime/consumer support | +| Batch size, new-token count, and context length | A measured latency surface and supported interpolation within its coverage | Coverage of the actual query shapes; mixed/ragged workload composition remains a modeling assumption | +| **Pipeline parallelism (PP)** | A scalar forward latency is insufficient to establish pipeline behavior | Stage execution, microbatch flow, overlap, bubbles, and inter-stage communication/scheduling must be modeled separately. The current collection workflow requires PP=1. | +| Request scheduling, routing, arrival patterns, or prefill/decode disaggregation | An iteration timing input to the serving simulation | Scheduler/router behavior, queues, placement, and P/D transfer models | +| KV capacity, allocation, prefix reuse, eviction, or hybrid recurrent state | Timing measured under the recorded cache-state conditions | The corresponding memory and cache-lifecycle model; collecting timing does not implement those semantics | + +**Self-collection cannot add PP support by itself.** A PP=1 profile must not be +scaled or relabeled as PP>1. Even measurements from a pipeline-enabled engine +would require a supported stage/scheduling contract before Replay could use +them to simulate pipeline throughput and latency. + +Likewise, TTFT, ITL, and end-to-end throughput are outputs of the serving +simulation and workload, not direct guarantees from a forward-pass table. +The measured timing boundary matters: record it, including rank aggregation; +do not substitute client latency or pure kernel time for engine wall time. + +## Current support and architecture readiness + +| Backend | Self-benchmark collection status in this workflow | | --- | --- | -| Published self-benchmark data; no collection environment needed | [Check the package](#1-check-the-aisimulate-commands), then [Kimi K3 TP8+DCP8 quickstart](#use-an-existing-profile-kimi-k3-tp8dcp8) | -| Your own schema-v6 Parquet/metadata pair | [Inspect the pair](#5-inspect-the-published-pair), [construct the model](#6-load-the-new-data-and-query-one-forward-pass), then [predict](#7-run-ais-predict-using-the-same-external-data); substitute your recorded identity | -| New measurements to collect | [Prepare the environment](#1-check-the-aisimulate-commands), then follow steps 2–9 for MiniMax-M2.7 on four H200 GPUs | - -```text -Existing measured pair ------------------------+ - | -GPU plan -> smoke -> self-benchmark -> publish -+-> SDK query -> AIS predict - | - v - Independent accuracy check -``` - -The data contract is `fpm_forward_perf.parquet` plus its matching -`fpm_forward_perf.metadata.json`. The Rust performance model selects a matching -cell and uses measured lookup/interpolation and, where supported, SOL transfer. -An existing model architecture needs no new model class or parallelization -mapping to consume another measured configuration. Place the pair in a systems -root and pass its exact identity to -`RustForwardPassPerfModel.best_available(ForwardPassPerfModelConfig(...))`. -There is no separate FPM registration service or `train` command. - -Raw benchmark JSON is not a loadable FPM pair. Convert it to the -[publication schema](../../collector/README.md#whole-forward-fpm-campaign) -with its measurement provenance before step 5. Collector performs this -conversion for the campaigns it runs in steps 2–4. - -A new architecture still requires a model definition; see -[How to add a new model](../add_a_new_model.md). -[Online telemetry regression](aic-fpm-regression-design.md), which learns from -`tune_with_fpms`, is a separate workflow. - -## Scope and expected outcomes - -| Stage | Runs on | Expected result | +| 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. + +Kimi K3 is one such example: its MLA+KDA state and DCP layout needed benchmark +adaptation. The published decode profile uses warmed MLA KV and random KDA +state. That records a specific timing experiment; it does not establish full +KDA-history or cache-lifecycle fidelity in Replay. + +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 | | --- | --- | --- | -| 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 | +| [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. -The collection walkthrough uses vLLM with -`tp=4, pp=1, dp=1, moe_tp=4, moe_ep=1, cp=1`, checkpoint-default quantization, -and one worker replica. The dedicated Collector currently supports `PP=CP=1`; -its topology presets do not collect DCP configurations. The DCP quickstart -consumes an existing Dynamo self-benchmark profile. + -## 1. Check the AISimulate commands +## 1. Check support and prepare the environment Use Python 3.11–3.13 and an AISimulate build containing [PR #284](https://github.com/ai-dynamo/aisimulate/pull/284). A package version @@ -92,7 +151,8 @@ aisimulate predict --help ``` Use the same activated Bash or Zsh session for the remaining commands. -For existing data, continue with the quickstart below or step 5. +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: @@ -102,19 +162,248 @@ python -m collector.fpm_forward --help ``` The collection host needs model configuration metadata and `kubectl` access. -The MiniMax example needs four H200 GPUs on one node and an existing namespace, -checkpoint PVC/path, collection image, and matching Generator target release. +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](#8-recovery-and-cleanup). +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. -## Use an existing profile: Kimi K3 TP8+DCP8 + -This CPU-only quickstart downloads the published +## 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. + +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. -### Download and install the measured pair +### 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. @@ -187,7 +476,7 @@ 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. -### Query prefill and decode through the canonical API +### 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. @@ -248,7 +537,7 @@ 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. -### Run a small Kimi Replay +### 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. @@ -312,7 +601,7 @@ engine block size of 64 is not this logical block size. The recorded engine's `max_num_seqs` is 32. Keep these values tied to this profile, not to Kimi models in general. -This quickstart checks timing consumption. DCP Replay requires explicit capacity; +This example checks timing consumption. DCP Replay requires explicit capacity; if enabling host offload or P/D transfer, also provide explicit bytes per token. The generic Replay cache does not model KDA checkpoint/eviction behavior or native hybrid prefill chunk alignment. Prefix caching is disabled here to keep @@ -325,15 +614,25 @@ Only measured points and supported interpolation are available for DCP; SOL-dependent transfer paths fail explicitly. Arbitrary cached-prefix/chunk shapes may be unsupported even inside the axis minima/maxima. Expand traffic only after checking query coverage and perform an -[independent accuracy check](#9-validate-accuracy-separately) before using the +[independent accuracy check](#7-validate-coverage-and-accuracy) before using the results for deployment decisions. -The existing-data quickstart is complete. Continue below only to collect a new -MiniMax profile; use a separate campaign directory. +These checks verify that the profile loads and the example workload runs. Use +step 7 to validate broader workloads. Example B shows a new collection campaign; +use a separate directory if running both examples. + +## Example B: collect a MiniMax-M2.7 TP4 profile + +This example illustrates steps 2–6 with a new vLLM collection on four H200 GPUs. +Complete step 1 first. It assumes a working GPU collection environment; replace +the namespace, checkpoint, image, and release placeholders with its values. +This example does not collect DCP. Use step 7 for accuracy validation afterward. + + -## 2. Choose the example's output paths +### B1. Prepare the campaign (step 2) -This guide uses an external directory to keep the example's results separate +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 @@ -378,7 +677,9 @@ 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 + + +### B2. Freeze and inspect the plan (step 2) Use the Generator target release from the existing collection configuration and define the campaign arguments once: @@ -435,7 +736,7 @@ The sampling settings have distinct meanings: 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 +#### 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 @@ -470,12 +771,14 @@ 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 +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. -## 4. Run smoke, then formal collection + + +### 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 @@ -562,7 +865,7 @@ provenance, logs, checkpoints, and the published pair. Inspect [collector contract](../../collector/README.md#whole-forward-fpm-campaign) for accepted regimes and publication behavior. -## 5. Inspect the published pair +### 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 @@ -625,7 +928,7 @@ Keep these paths distinct: 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 +### 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; @@ -686,7 +989,7 @@ 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 +### 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` @@ -756,25 +1059,12 @@ 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. -New configurations use `estimation_mode` and `fallback_policy`. Without -explicit selection, the defaults are `auto` + `deny`: construction tries -op-level, measured FPM interpolation, then regression. For a measured-profile -workflow, explicitly select `fpm_interpolation` + `deny` as above. An untrained -regression cannot run offline prediction, and a selected estimator never switches -models on a query miss. - -For nondefault precision, put `gemm_quant_mode`, `moe_quant_mode`, -`fmha_quant_mode`, `kvcache_quant_mode`, `comm_quant_mode`, and -`attention_backend` under `engine.workers..timing`, using the exact -recorded identity. These overrides require default timing and regular language -workers. Recommendation rejects them until its feasibility preflight supports -the same identity. See the [Core API](../../../../docs/core-api.md) for estimator -controls and the [DCP quickstart](#use-an-existing-profile-kimi-k3-tp8dcp8) for -a complete configuration with text-only and unrecorded-precision settings. + -## 8. Recovery and cleanup +## Recovery and cleanup -Use the same frozen arguments and roots to resume: +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 @@ -813,37 +1103,6 @@ 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 | From 51e0f5e8c12ef8300013f96e3c7df25f31862caa Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 11:38:59 -0700 Subject: [PATCH 05/12] docs(fpm): name the self-benchmarking onboarding guide explicitly Signed-off-by: hongkuanz --- README.md | 2 +- python/aisimulate/docs/add_a_new_model.md | 2 +- python/aisimulate/docs/fpm/README.md | 4 ++-- ...to-end-workflow.md => self-benchmarking-and-onboarding.md} | 0 4 files changed, 4 insertions(+), 4 deletions(-) rename python/aisimulate/docs/fpm/{end-to-end-workflow.md => self-benchmarking-and-onboarding.md} (100%) diff --git a/README.md b/README.md index 199f50ec26..889cf5de5b 100644 --- a/README.md +++ b/README.md @@ -255,7 +255,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) - [AIConfigurator compatibility Python API](python/aisimulate/README.md#python-api) diff --git a/python/aisimulate/docs/add_a_new_model.md b/python/aisimulate/docs/add_a_new_model.md index 34757afc23..40d32b0481 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/aiconfigurator_core/systems/data/// Date: Fri, 18 Sep 2026 11:40:11 -0700 Subject: [PATCH 06/12] docs(fpm): shorten self-collection scope explanation Signed-off-by: hongkuanz --- .../fpm/self-benchmarking-and-onboarding.md | 28 ++++--------------- 1 file changed, 6 insertions(+), 22 deletions(-) diff --git a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md index 2842a3e7a5..a44f4fa8e2 100644 --- a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md +++ b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md @@ -38,28 +38,12 @@ starting at step 4 after checking its provenance and coverage. ## What self-collection covers -An FPM predicts **engine iteration / forward-pass latency** for scheduled work. -It supplies the cost of an iteration to the simulator; the simulator still -decides which work runs, when it runs, and how requests progress. - -| Difference you want to model | What self-collection can supply | What still needs separate support | -| --- | --- | --- | -| Kernel implementation, fusion, quantization, CUDA Graph/eager behavior, engine-version changes | Measured iteration costs for the recorded configuration and shapes | A working benchmark path and matching profile identity | -| TP/EP/DCP execution within a supported engine iteration | The observed compute and collective costs inside the measurement boundary | Correct topology, rank aggregation, cache state, and runtime/consumer support | -| Batch size, new-token count, and context length | A measured latency surface and supported interpolation within its coverage | Coverage of the actual query shapes; mixed/ragged workload composition remains a modeling assumption | -| **Pipeline parallelism (PP)** | A scalar forward latency is insufficient to establish pipeline behavior | Stage execution, microbatch flow, overlap, bubbles, and inter-stage communication/scheduling must be modeled separately. The current collection workflow requires PP=1. | -| Request scheduling, routing, arrival patterns, or prefill/decode disaggregation | An iteration timing input to the serving simulation | Scheduler/router behavior, queues, placement, and P/D transfer models | -| KV capacity, allocation, prefix reuse, eviction, or hybrid recurrent state | Timing measured under the recorded cache-state conditions | The corresponding memory and cache-lifecycle model; collecting timing does not implement those semantics | - -**Self-collection cannot add PP support by itself.** A PP=1 profile must not be -scaled or relabeled as PP>1. Even measurements from a pipeline-enabled engine -would require a supported stage/scheduling contract before Replay could use -them to simulate pipeline throughput and latency. - -Likewise, TTFT, ITL, and end-to-end throughput are outputs of the serving -simulation and workload, not direct guarantees from a forward-pass table. -The measured timing boundary matters: record it, including rank aggregation; -do not substitute client latency or pure kernel time for engine wall time. +Self-collection supplies **engine-iteration latency** for the measured configuration +and covered workload shapes. It cannot provide PP pipeline behavior, request +scheduling/routing, prefill–decode transfers, or KV-cache capacity/lifecycle +modeling; these require separate simulator support. The current collection +workflow requires **PP=1**, and iteration timings alone do not establish TTFT, +ITL, or end-to-end throughput accuracy. ## Current support and architecture readiness From 11d5428fa78977f1b8606ecc1fc5c76e9d5d20cf Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 11:43:28 -0700 Subject: [PATCH 07/12] docs(fpm): remove Kimi adaptation example from support overview Signed-off-by: hongkuanz --- .../aisimulate/docs/fpm/self-benchmarking-and-onboarding.md | 5 ----- 1 file changed, 5 deletions(-) diff --git a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md index a44f4fa8e2..b3f3c9b6b6 100644 --- a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md +++ b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md @@ -71,11 +71,6 @@ engine versions can require changes in: registration, resident weights, and memory/cache descriptions used by the consumer. A collection plan's bootstrap template is not proof of consumer support. -Kimi K3 is one such example: its MLA+KDA state and DCP layout needed benchmark -adaptation. The published decode profile uses warmed MLA KV and random KDA -state. That records a specific timing experiment; it does not establish full -KDA-history or cache-lifecycle fidelity in Replay. - 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 9c5ab033c69fcfdbb616723fdc6da4843498921e Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 11:44:51 -0700 Subject: [PATCH 08/12] docs(fpm): prepare the environment from latest main Signed-off-by: hongkuanz --- .../docs/fpm/self-benchmarking-and-onboarding.md | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md index b3f3c9b6b6..1158279f70 100644 --- a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md +++ b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md @@ -99,17 +99,13 @@ resolve it. ## 1. Check support and prepare the environment -Use Python 3.11–3.13 and an AISimulate build containing -[PR #284](https://github.com/ai-dynamo/aisimulate/pull/284). A package version -alone does not establish that an unreleased feature is included. If you already -have that build, activate its environment and continue with the checks below. -For a source installation while the PR is open: +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 https://github.com/ai-dynamo/aisimulate.git +git clone --branch main https://github.com/ai-dynamo/aisimulate.git cd aisimulate -git fetch origin pull/284/head -git switch --detach FETCH_HEAD +git pull --ff-only origin main uv sync --project python/aisimulate --extra dev source python/aisimulate/.venv/bin/activate ``` @@ -125,7 +121,6 @@ From the repository root, record the source and package in the activated session export AIS_REPO="$PWD" git rev-parse HEAD python -c 'from importlib.metadata import version; print(version("aisimulate"))' -python -c 'from aisimulate_core.sdk import ForwardPassPerfModelConfig; assert "dcp" in ForwardPassPerfModelConfig.__dataclass_fields__' aisimulate predict --help ``` From f7934d4218e497181ee678631b47123510608155 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 12:03:03 -0700 Subject: [PATCH 09/12] test(fpm): cover DCP capacity and text-only cache isolation Signed-off-by: hongkuanz --- crates/core/src/python.rs | 27 +++++++++++++++++++ docs/core-api.md | 3 +++ .../fpm/self-benchmarking-and-onboarding.md | 3 +++ .../tests/unit/sdk/test_fpm_forward.py | 22 +++++++++++++++ 4 files changed, 55 insertions(+) diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 55a965760f..2a20723165 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -2133,6 +2133,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 cuda_graph_reservation_reaches_capacity_rematerialization() { let mut config = aic_config(); diff --git a/docs/core-api.md b/docs/core-api.md index 49d74365b7..21d0ada3ee 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -285,6 +285,9 @@ 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 diff --git a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md index 1158279f70..c4b1be8ab6 100644 --- a/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md +++ b/python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md @@ -319,6 +319,9 @@ 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) diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index c6c1ddd20a..43721a6dd0 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -774,6 +774,28 @@ def prefill(model_config): ): with pytest.raises(PerfDataNotAvailableError, match="No FPM cell matches"): prefill(different) + with pytest.raises(NotImplementedError, match="encoder/multimodal"): + prefill(replace(config, fpm_text_only=False)) 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 aiconfigurator_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_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() From 8b13d63711f4e95017c719f8fbe6051659832e9d Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Fri, 18 Sep 2026 16:24:42 -0700 Subject: [PATCH 10/12] fix(perfmodel): enforce validation across construction paths Signed-off-by: hongkuanz --- crates/core/src/perfmodel/fpm/config.rs | 21 ++++------- crates/core/src/perfmodel/fpm/estimator.rs | 21 +++++++++++ crates/core/src/perfmodel/fpm/model.rs | 19 +++++----- crates/core/src/perfmodel/fpm/tests.rs | 20 +++++++++++ crates/core/src/perfmodel/py.rs | 19 ++++++++++ .../_aiconfigurator_core.pyi | 4 +++ .../src/aiconfigurator_core/sdk/engine.py | 6 +++- .../tests/unit/sdk/test_fpm_forward.py | 35 +++++++++++++++++++ tests/sweeper/test_forward_pass_estimator.py | 19 ++++++++++ 9 files changed, 140 insertions(+), 24 deletions(-) diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index f791e040ff..bcf6294e51 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -224,21 +224,12 @@ impl ForwardPassPerfModelConfig { if self.dcp.is_some_and(|dcp| dcp == 0 || self.tp % dcp != 0) { return Err(invalid_config("dcp must be positive and divide tp")); } - for mode in &self - .estimator_config + self.estimator_config .fpm_interpolation - .unrecorded_quant_modes - { - let explicit = match mode { - super::UnrecordedFpmQuantMode::Fmha => &self.fmha_quant_mode, - super::UnrecordedFpmQuantMode::Comm => &self.comm_quant_mode, - }; - if explicit.is_some() { - return Err(invalid_config( - "an unrecorded FPM quant mode cannot have an explicit quantization override", - )); - } - } + .validate_quant_modes( + self.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", @@ -303,7 +294,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() diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs index 04fb90a7d4..6c45805e7b 100644 --- a/crates/core/src/perfmodel/fpm/estimator.rs +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -44,6 +44,27 @@ pub enum UnrecordedFpmQuantMode { 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 /// established workload coordinates until a replacement is accuracy-qualified. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index fb214141f0..ed2c840b90 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -317,6 +317,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()); @@ -645,14 +656,6 @@ impl ForwardPassPerfModel { fn build_native_candidate( config: &ForwardPassPerfModelConfig, ) -> Result<(Engine, PathBuf), AicError> { - if config.dcp.is_some_and(|dcp| dcp > 1) - && (config.estimation_mode != EstimationMode::FpmInterpolation - || config.backend != crate::BackendKind::Vllm) - { - return Err(AicError::UnsupportedModel( - "DCP timing requires measured vLLM FPM interpolation".into(), - )); - } if config.estimation_mode == EstimationMode::FpmInterpolation && config.nextn != 0 { return Err(AicError::UnsupportedModel( "FPM interpolation does not support MTP".into(), diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index 2883e5b1c6..4849fc8807 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 { diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 34fd6091aa..aa6a2e637f 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1667,6 +1667,24 @@ impl PyForwardPassPerfModel { serde_json::to_string(&config).map_err(|e| PyValueError::new_err(e.to_string())) } + /// Share typed FPM option validation with the compilation adapter. + #[staticmethod] + #[pyo3(signature = (options_json, fmha_quant_mode=None, comm_quant_mode=None))] + fn _normalize_fpm_options( + options_json: &str, + fmha_quant_mode: Option<&str>, + 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)?; + serde_json::to_string(&options).map_err(|e| PyValueError::new_err(e.to_string())) + } + /// Migration adapter for previously saved flat tuning options. #[staticmethod] fn legacy_estimator_config(options_json: &str) -> PyResult { @@ -1742,6 +1760,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())) } diff --git a/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi b/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi index 7e4a285980..719446e286 100644 --- a/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi +++ b/python/aisimulate/src/aiconfigurator_core/_aiconfigurator_core.pyi @@ -169,6 +169,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/aiconfigurator_core/sdk/engine.py b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py index d26c17065b..9c99719bcc 100644 --- a/python/aisimulate/src/aiconfigurator_core/sdk/engine.py +++ b/python/aisimulate/src/aiconfigurator_core/sdk/engine.py @@ -440,13 +440,17 @@ def compile_engine( # does not take a model_path (quant inference is done inside `get_model`). from aiconfigurator_core.sdk.speculation import SpeculationConfig - fpm_options = fpm_options or {} 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_options = json.loads( + aiconfigurator_core.RustForwardPassPerfModel._normalize_fpm_options( + json.dumps({} if fpm_options is None else fpm_options), fmha_quant_mode, comm_quant_mode + ) + ) resolved_speculation = SpeculationConfig(**speculation) if speculation is not None else None model_config = build_model_config( tp_size=tp_size, diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 43721a6dd0..d4c12ee040 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -799,3 +799,38 @@ def test_session_cache_separates_text_only_mode(fpm_session): 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": ["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("aiconfigurator_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) diff --git a/tests/sweeper/test_forward_pass_estimator.py b/tests/sweeper/test_forward_pass_estimator.py index 0de5bb384d..83b9d0b827 100644 --- a/tests/sweeper/test_forward_pass_estimator.py +++ b/tests/sweeper/test_forward_pass_estimator.py @@ -597,3 +597,22 @@ def test_mixed_timing_still_enforces_cold_regression_on_default_role(): } with pytest.raises(ForwardPassEstimatorResolutionError, match="prefill is not ready"): 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") From 2dcfbec6d7e270bd9282ec96ca962064a81df339 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Tue, 29 Sep 2026 08:58:26 -0700 Subject: [PATCH 11/12] test: align host-offload default precision contract Signed-off-by: hongkuanz --- tests/test_host_offload_cli.py | 1 + 1 file changed, 1 insertion(+) 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, } ] From 24b9af3cbe5b6e0d4da41ce85ab0c7212d28ae84 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Tue, 29 Sep 2026 09:51:17 -0700 Subject: [PATCH 12/12] refactor: carry resolved FPM config and typed DCP Signed-off-by: hongkuanz --- crates/core/src/perfmodel/config.rs | 4 +- crates/core/src/perfmodel/engine/readiness.rs | 8 +- crates/core/src/perfmodel/engine/runtime.rs | 6 ++ crates/core/src/perfmodel/engine/spec.rs | 15 ++-- .../core/src/perfmodel/operators/attention.rs | 6 +- .../src/perfmodel/operators/fpm_forward.rs | 30 ++++---- .../perfmodel/perf_database/fpm_forward.rs | 73 ++++++++++++------- .../core/src/perfmodel/perf_database/mod.rs | 2 +- crates/core/src/perfmodel/py.rs | 9 ++- crates/core/src/python.rs | 1 + crates/tests/public-api/src/lib.rs | 3 +- docs/deepseek-v41-storage.md | 2 +- docs/deepseek-v41.md | 5 +- .../src/aisimulate/sdk/fpm_config.py | 8 ++ .../src/aisimulate_core/sdk/config.py | 5 +- .../src/aisimulate_core/sdk/engine.py | 23 +++--- .../src/aisimulate_core/sdk/fpm_config.py | 48 ++++++++++++ .../aisimulate_core/sdk/models/__init__.py | 34 ++++----- .../sdk/operations/fpm_forward.py | 51 +++++++------ .../aisimulate_core/sdk/rust_engine_step.py | 5 +- .../cross_package/test_import_contract.py | 7 +- .../unit/sdk/database/test_attention_lanes.py | 6 +- .../tests/unit/sdk/test_fpm_forward.py | 71 +++++++++++++++--- 23 files changed, 292 insertions(+), 130 deletions(-) create mode 100644 python/aisimulate/src/aisimulate/sdk/fpm_config.py create mode 100644 python/aisimulate/src/aisimulate_core/sdk/fpm_config.py diff --git a/crates/core/src/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index 432017bdf2..c6bc96317f 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -117,7 +117,9 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // - 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. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 23; +// - 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`]. 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 3af5f4618f..4206a65a63 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -3532,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(), @@ -3603,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(), @@ -3661,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(), @@ -4206,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(), @@ -4524,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(), @@ -4570,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 dbf37bd754..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(), @@ -1046,7 +1047,7 @@ mod tests { vec![OpSpec::FpmForward(fpm_forward())], vec![OpSpec::Moe(moe())], ); - assert_eq!(spec.schema_version, 23); + assert_eq!(spec.schema_version, 24); let mut bytes = spec.to_bincode().unwrap(); assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); @@ -1353,7 +1354,9 @@ mod tests { 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 = @@ -1377,7 +1380,7 @@ mod tests { for spec in [fpm_spec, pilot_spec] { let bytes = spec.to_bincode().unwrap(); assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); - for previous_version in [20u32, 21, 22] { + for previous_version in [20u32, 21, 22, 23] { let mut stale = bytes.clone(); stale[..4].copy_from_slice(&previous_version.to_le_bytes()); // DCP and the pilot claimed 22 for different layouts. @@ -1388,7 +1391,7 @@ mod tests { Err(AicError::UnsupportedSchemaVersion { kind: "EngineSpec", got, - expected: 23, + expected: 24, }) if got == previous_version )); } @@ -1429,10 +1432,10 @@ mod tests { 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(&22u32.to_le_bytes()); + previous[..4].copy_from_slice(&23u32.to_le_bytes()); assert!(matches!( EngineSpec::from_bincode(&previous), - Err(AicError::UnsupportedSchemaVersion { got: 22, .. }) + Err(AicError::UnsupportedSchemaVersion { got: 23, .. }) )); } } diff --git a/crates/core/src/perfmodel/operators/attention.rs b/crates/core/src/perfmodel/operators/attention.rs index 31e2229575..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 23). + /// 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 23). +/// 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 23). + /// 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 dd37894d06..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 { @@ -107,12 +111,7 @@ fn data_err(msg: String) -> AicError { impl FpmForwardOp { fn sol_at(&self, db: &PerfDatabase, coords: &[f64]) -> Result { - if self - .match_identity - .get(crate::perf_database::fpm_forward::FPM_CELL_MATCH_COLUMNS.len()) - .and_then(|v| v.parse::().ok()) - .is_some_and(|dcp| dcp > 1) - { + 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(), )); @@ -123,9 +122,9 @@ impl FpmForwardOp { 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, @@ -680,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(), @@ -698,7 +698,7 @@ mod tests { write_pair(tmp.path(), &default_rows()); let db = db_with_pair(tmp.path()); let mut measured = op(FpmPhase::Prefill); - measured.match_identity.push("4".into()); + 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")); } @@ -1476,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 43488f4530..2dbc01ceb3 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -329,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 @@ -344,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() { @@ -364,7 +373,7 @@ impl FpmForwardTable { let identity: Vec = FPM_CELL_MATCH_COLUMNS .iter() .chain(std::iter::once(&"dcp")) - .zip(match_identity) + .zip(requested()) .map(|(c, v)| format!("{c}={v:?}")) .collect(); return Err(structural(format!( @@ -1669,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() ); } @@ -1697,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") @@ -1776,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| { @@ -1806,13 +1815,25 @@ pub(crate) mod tests { assert_eq!(table.cells().unwrap().len(), 2); assert!( table - .select_cell(&default_identity(4), "org/model-a") + .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").is_ok()); + 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( @@ -2410,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}"); } @@ -2429,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}"); } @@ -2481,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)); @@ -2489,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); @@ -2501,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); @@ -2513,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]); @@ -2529,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]); @@ -2549,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}"); } @@ -2572,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 c4566b334e..3a560998cd 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1922,7 +1922,14 @@ impl PyForwardPassPerfModel { options .validate_quant_modes(fmha_quant_mode, comm_quant_mode) .map_err(aic_to_py)?; - serde_json::to_string(&options).map_err(|e| PyValueError::new_err(e.to_string())) + // 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. diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index f7c1cc32f7..f5da5d0bbc 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -2717,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(), diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index 9ff3cd3644..5157021eea 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -145,7 +145,8 @@ mod tests { // v21: EngineConfig and MoeOp gained exact moe_kernel_source identity. // v22: observed MoE selection and exact prefill graph composites. // v23: ParallelMapping gained optional recorded DCP identity. - assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 23); + // 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/deepseek-v41-storage.md b/docs/deepseek-v41-storage.md index 5c604f4c48..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 23 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 5e6b6e8b9a..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 23 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 @@ -132,5 +132,6 @@ FPM selector diagnostics append a positional field to `FpmForwardOp` in schema 2 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. DCP identity extends the combined layout in schema 23. -Schema-22 and earlier bincode inputs are rejected before decoding; legacy JSON +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/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_core/sdk/config.py b/python/aisimulate/src/aisimulate_core/sdk/config.py index e5dc624461..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) @@ -70,9 +71,7 @@ class ModelConfig: """ dcp_size: int | None = field(default=None, kw_only=True) - fpm_text_only: bool = field(default=False, kw_only=True) - fpm_unrecorded_quant_modes: tuple[str, ...] = field(default=(), kw_only=True) - fpm_attention_backend: str | 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 diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index c755f8343b..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,6 +174,7 @@ 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._resident_weight_bytes, # Speculative verify width for the equivalent-AR decode mapping @@ -486,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 23 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. """ @@ -510,6 +513,7 @@ 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 @@ -518,14 +522,13 @@ def compile_engine( 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_options = json.loads( - aisimulate_core.RustForwardPassPerfModel._normalize_fpm_options( - json.dumps({} if fpm_options is None else fpm_options), - fmha_quant_mode if fmha_quant_mode is not None else fpm_fmha_quant_mode, - comm_quant_mode, - ) + 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_options.get("fpm_parquet_path") + 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") @@ -563,9 +566,7 @@ def compile_engine( wideep_num_slots=wideep_num_slots, speculation=resolved_speculation, ) - model_config.fpm_text_only = fpm_options.get("text_only", False) - model_config.fpm_unrecorded_quant_modes = tuple(fpm_options.get("unrecorded_quant_modes", ())) - model_config.fpm_attention_backend = recorded_attention_backend + 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 623619979c..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 and not model.config.fpm_text_only: + 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,9 +125,22 @@ 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) + 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, + # These are the SDK prediction assumptions; the producer separately + # verifies actual runtime residency and token-only requests. + engram_cpu_offload=False, + input_modality="text", + ) + 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 + "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. @@ -137,20 +151,6 @@ def _apply_forward_model_fpm(model: BaseModel, backend_name: str = "vllm") -> Ba 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_COLUMNS, 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 - # verifies actual runtime residency and token-only requests. - engram_cpu_offload=False, - input_modality="text", - ) - for op in (prefill_op, decode_op): - # Replace execution fields without dropping the optional DCP selector. - op._match_identity = (*op._match_identity[:15], *identity, *op._match_identity[15 + len(EXECUTION_COLUMNS) :]) model.forward_model = "fpm" return model 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 1972c840fa..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 @@ -136,34 +139,36 @@ def __init__( 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), - "" - if "fmha" in model_config.fpm_unrecorded_quant_modes - else _norm_identity(model_config.fmha_quant_mode if fmha_selector is None else fmha_selector), - "" if "comm" in model_config.fpm_unrecorded_quant_modes else _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( - model_config.fpm_attention_backend or model_config.attention_backend, engine_default="flashinfer" + "attention_backend": _norm_backend_request( + fpm_config.attention_backend or model_config.attention_backend, 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 - if model_config.dcp_size is not None: - self._match_identity += (str(model_config.dcp_size),) + "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 8e8d9dff18..612be02bdc 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -1175,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) @@ -1241,9 +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_text_only": getattr(model_config, "fpm_text_only", False), - "fpm_unrecorded_quant_modes": getattr(model_config, "fpm_unrecorded_quant_modes", ()), - "fpm_attention_backend": getattr(model_config, "fpm_attention_backend", None), + "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_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 222fdec82a..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_three(): - """DCP and the 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 == 23 + 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 50b200c3b6..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, @@ -702,9 +747,9 @@ def test_text_only_kimi_profile_preserves_encoder_weights_and_dcp_identity(): moe_tp_size=8, moe_ep_size=1, dcp_size=8, - fpm_text_only=True, - fpm_unrecorded_quant_modes=("fmha", "comm"), - fpm_attention_backend="FLASHINFER_MLA", + 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 @@ -713,11 +758,13 @@ def test_text_only_kimi_profile_preserves_encoder_weights_and_dcp_identity(): identity = model.context_ops[0]._match_identity assert identity[2:4] == ("", "") assert identity[12] == "FLASHINFER_MLA" - assert identity[-1] == "8" + 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) @@ -841,9 +888,9 @@ def test_session_cache_preserves_fpm_identity(kimi_fpm_profile): moe_ep_size=1, dcp_size=8, forward_model="fpm", - fpm_text_only=True, - fpm_unrecorded_quant_modes=("fmha", "comm"), - fpm_attention_backend="FLASHINFER_MLA", + 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, ) @@ -859,13 +906,13 @@ def prefill(model_config): assert prefill(config) == pytest.approx(7.0) # synthetic measured prefill row for different in ( replace(config, dcp_size=1), - replace(config, fpm_unrecorded_quant_modes=()), - replace(config, fpm_attention_backend="auto"), + 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_text_only=False)) + prefill(replace(config, fpm_config=FpmCompileConfig())) assert prefill(config) == pytest.approx(7.0) finally: _engine_handle_cache_clear() @@ -877,7 +924,9 @@ def test_session_cache_separates_text_only_mode(fpm_session): 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_text_only=True), BACKEND) + 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)