diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 3f1bffbe8..803b19137 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -67,6 +67,19 @@ This material is licensed under the Apache License 2.0. The upstream license at the identified revision is available at: https://github.com/ai-dynamo/aiconfigurator/blob/915f590680d8a79fe9c39f6f3a9ff13bc267fcce/LICENSE +## Dynamo V4.1 FPM collection adapter + +`collector/fpm_forward/runtime/dsv41/dsv41_scheduler.py` is modified code +adapted from `components/src/dynamo/vllm/instrumented_scheduler.py` in +https://github.com/ai-dynamo/dynamo/tree/54960177085413259859c88bd34ed0734d4c2ea9. +It adds bounded same-request real-KV collection while preserving the native +benchmark and FPM contracts. Copyright (c) 2025-2026 NVIDIA CORPORATION & +AFFILIATES. All rights reserved. Licensed under Apache-2.0; the upstream +license is preserved in the adapter's adjacent `LICENSE`. The adjacent README +records the inspected vLLM API revision and immutable runtime image/source +hashes. vLLM implementation files are not vendored. The text fixture and +lifecycle tests are original work for this change, with no external corpus. + ## NVIDIA AIConfigurator speculative decoding The speculation SDK, compatibility exports, CLI/task integration, attention and whole-forward FPM operation changes, native bindings, and their tests are adapted and modified from AIConfigurator PR #1563, pinned at commit `6290c161a354da5250c391bd43372b2e9c6f4a51`. Original paths are under `aic-core/src/aiconfigurator_core/sdk/`, `src/aiconfigurator/`, `aic-core/rust/aiconfigurator-core/`, `aic-core/rust/tests/public-api/`, and `tests/`. @@ -535,7 +548,7 @@ SOFTWARE. ## Meta Muse Glimmer model configuration -`src/aiconfigurator_core/model_configs/meta-models--Muse-Glimmer-30B_config.json` +`src/aisimulate_core/model_configs/meta-models--Muse-Glimmer-30B_config.json` is an unmodified copy of `config.json` from the Meta Muse Glimmer model repository at immutable revision `f84ecc3a0ea984a4c04542a84269e3d065350a6e`: @@ -557,8 +570,8 @@ named Qwen model repositories at the immutable revisions shown: | Packaged file | Upstream revision | | --- | --- | -| `aiconfigurator_core/model_configs/Qwen--Qwen3.8-2.4T-A95B_config.json` | `Qwen/Qwen3.8-2.4T-A95B@207bd685a7e3696cfaff12ded7c6a7ea0f88c996` | -| `aiconfigurator_core/model_configs/Qwen--Qwen3.8-2.4T-A95B-FP8_config.json` | `Qwen/Qwen3.8-2.4T-A95B-FP8@d2dc35658bcf77e66643428cb52e774cc3b5bd29` | +| `aisimulate_core/model_configs/Qwen--Qwen3.8-2.4T-A95B_config.json` | `Qwen/Qwen3.8-2.4T-A95B@207bd685a7e3696cfaff12ded7c6a7ea0f88c996` | +| `aisimulate_core/model_configs/Qwen--Qwen3.8-2.4T-A95B-FP8_config.json` | `Qwen/Qwen3.8-2.4T-A95B-FP8@d2dc35658bcf77e66643428cb52e774cc3b5bd29` | Upstream repositories: https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B and @@ -592,10 +605,10 @@ For any questions regarding this license, please contact model-business@notice.q The following bundled model configs are modified copies of Meta Llama 4 checkpoint configuration files: -- `src/aiconfigurator_core/model_configs/meta-llama--Llama-4-Scout-17B-16E-Instruct_config.json` +- `src/aisimulate_core/model_configs/meta-llama--Llama-4-Scout-17B-16E-Instruct_config.json` from revision `92f3b1597a195b523d8d9e5700e57e4fbb8f20d3`: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/92f3b1597a195b523d8d9e5700e57e4fbb8f20d3/config.json -- `src/aiconfigurator_core/model_configs/meta-llama--Llama-4-Maverick-17B-128E-Instruct_config.json` +- `src/aisimulate_core/model_configs/meta-llama--Llama-4-Maverick-17B-128E-Instruct_config.json` from revision `73d14711bcc77c16df3470856949c3764056b617`: https://huggingface.co/meta-llama/Llama-4-Maverick-17B-128E-Instruct/blob/73d14711bcc77c16df3470856949c3764056b617/config.json diff --git a/crates/core/parity_tests/perfmodel/test_engine_step_parity.py b/crates/core/parity_tests/perfmodel/test_engine_step_parity.py index cc7271957..7842f1706 100644 --- a/crates/core/parity_tests/perfmodel/test_engine_step_parity.py +++ b/crates/core/parity_tests/perfmodel/test_engine_step_parity.py @@ -2553,7 +2553,9 @@ def test_fpm_spec_tags(self, fpm_systems_root, monkeypatch): ctx_op = spec["context_ops"][0]["FpmForward"] assert ctx_op["phase"] == "prefill" assert spec["generation_ops"][0]["FpmForward"]["phase"] == "decode" - assert len(ctx_op["match_identity"]) == 15 + assert len(ctx_op["match_identity"]) == 19 + assert ctx_op["match_identity"][-4:] == ["", "full", "none", "text"] + assert spec["generation_ops"][0]["FpmForward"]["match_identity"] == ctx_op["match_identity"] assert ctx_op["sol_ops"], "sol_ops must carry the original granular list" @pytest.mark.parametrize( diff --git a/crates/core/src/engine/common/perf_model.rs b/crates/core/src/engine/common/perf_model.rs index 1ab334d36..84533dba4 100644 --- a/crates/core/src/engine/common/perf_model.rs +++ b/crates/core/src/engine/common/perf_model.rs @@ -42,10 +42,12 @@ impl std::fmt::Debug for PerfModel { } impl PerfModel { - pub(crate) fn prefill_batch_validation_can_fail(&self) -> bool { + // External prediction and duration conversion can fail even when geometry + // validation accepts every batch. Keep admission transactional for all providers. + pub(crate) fn prefill_pass_can_fail(&self) -> bool { match self { Self::Polynomial => false, - Self::External { timing } => timing.prefill_batch_validation_can_fail(), + Self::External { .. } => true, } } diff --git a/crates/core/src/engine/scheduler/sglang/core.rs b/crates/core/src/engine/scheduler/sglang/core.rs index b9c1ccc93..2ecce8ccb 100644 --- a/crates/core/src/engine/scheduler/sglang/core.rs +++ b/crates/core/src/engine/scheduler/sglang/core.rs @@ -778,11 +778,11 @@ impl SglangCore { } } } - // Only providers with fallible geometry validation need to preserve the - // admission state. Normal polynomial and unrestricted AIC passes avoid - // copying radix metadata. Lease checkpoints never become independent owners. + // Preserve admission around validation, prediction and duration conversion. + // The built-in polynomial remains infallible; external providers may fail + // after accepting geometry. Lease checkpoints are not independent owners. let admission_checkpoint = (!self.waiting.is_empty() - && self.config.perf_model.prefill_batch_validation_can_fail()) + && self.config.perf_model.prefill_pass_can_fail()) .then(|| { let waiting = self .waiting diff --git a/crates/core/src/engine/scheduler/sglang/tests.rs b/crates/core/src/engine/scheduler/sglang/tests.rs index 22c62a331..7c3d1dbda 100644 --- a/crates/core/src/engine/scheduler/sglang/tests.rs +++ b/crates/core/src/engine/scheduler/sglang/tests.rs @@ -2658,6 +2658,10 @@ mod admission_validation_rollback { } impl crate::engine::TimingModel for FallibleTiming { + fn prefill_batch_validation_can_fail(&self) -> bool { + !self.fail_in_prediction + } + fn validate_prefill_batch(&self, _: &[(usize, usize)]) -> anyhow::Result<()> { anyhow::ensure!( self.fail_in_prediction || !self.fail.load(Ordering::Relaxed), diff --git a/crates/core/src/engine/timing.rs b/crates/core/src/engine/timing.rs index 9363f600d..91c635ed0 100644 --- a/crates/core/src/engine/timing.rs +++ b/crates/core/src/engine/timing.rs @@ -447,9 +447,10 @@ pub enum TimingModelConfig { /// Implementations may call AIC, interpolate profiler data, or use another /// provider without adding that dependency to `aisimulate-core`. pub trait TimingModel: Send + Sync { - /// Whether admission needs a checkpoint around `validate_prefill_batch`. - /// Existing custom providers default to the safe, fallible contract. Providers - /// opting out must accept every batch geometry in that validation hook. + /// Whether `validate_prefill_batch` may reject geometry. A false return + /// guarantees only that validation accepts every batch, not that prediction + /// or duration conversion is infallible. Admission must still protect those + /// later operations for external providers. fn prefill_batch_validation_can_fail(&self) -> bool { true } diff --git a/crates/core/src/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index f15ecf067..d73f824a9 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -105,7 +105,10 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // - 19 (DeepSeek-V4.1 review): Dsv41AttentionOp gained kv_cache_layout, // separating physical backend KV payload from attention arithmetic precision. // Its appended enum changes positional bincode layout; old JSON defaults only. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 19; +// - 20 (DeepSeek-V4.1 FPM): FpmForwardOp gained original_fmha_quant_mode +// for selector diagnostics. This appends a positional field after the schema-19 +// release; serde defaults support legacy JSON, not legacy bincode. +pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 20; /// Static engine identity and setup information carried by an /// [`crate::perfmodel::engine::spec::EngineSpec`]. @@ -254,6 +257,9 @@ pub struct QuantizationConfig { #[serde(default)] pub moe_dtype: Option, pub activation_dtype: Option, + /// FPM cell selector only; does not override model arithmetic or memory. + #[serde(default)] + pub fpm_fmha_dtype: Option, pub kv_cache_dtype: Option, } diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index f098957e3..ea894c537 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -327,6 +327,13 @@ impl Engine { /// caller (`AicEngineBuilder` / `from_spec_bytes`) is responsible for /// having loaded the matching `PerfDatabase` from `spec.engine`'s identity. pub fn build(spec: EngineSpec, db: Arc) -> Result { + if spec.engine.quantization.fpm_fmha_dtype.is_some() + && spec.engine.forward_model.as_deref() != Some("fpm") + { + return Err(AicError::InvalidEngineConfig( + "fpm_fmha_dtype requires forward_model='fpm'".into(), + )); + } Self::validate_engine_database_mode(spec.engine.database_mode)?; Self::validate_engine_database_mode(db.database_mode)?; if spec.engine.database_mode != db.database_mode { @@ -2027,6 +2034,35 @@ impl Engine { let has_prefill = sched.sum_prefill_tokens > 0; let has_decode = sched.num_decode_requests > 0 || sched.sum_decode_kv_tokens > 0; + // The whole-forward rewrite retains the original stage graph in + // sol_ops. Apply the same replay restriction before its table lookup + // can return early. FPM v1 variance describes whole prompt lengths, + // not the current extends: even equal prompts can have different + // cached prefixes or previously completed chunks. Only one prefill + // request has identifiable (query, prefix) geometry in this input. + let replay_ops = self + .fpm_ops() + .map_or(self.context_ops.as_slice(), |(prefill, _)| { + prefill.sol_ops.as_slice() + }); + if has_prefill + && replay_ops.iter().any(|op| { + matches!(op, + Op::Dsv41Stage(stage) if stage.decoder_replay && stage.bounded) + }) + { + if !sched.var_prefill_length.is_finite() || sched.var_prefill_length < 0.0 { + return Err(AicError::InvalidForwardPassMetrics( + "V4.1 Decoder replay requires finite nonnegative prefill variance".into(), + )); + } + if sched.num_prefill_requests > 1 { + return Err(AicError::InvalidForwardPassMetrics( + "V4.1 Decoder replay requires per-request extend lengths; FPM v1 aggregates with multiple prefill requests cannot identify the tails, even when prompt-length variance is zero".into(), + )); + } + } + // FPM engines never enter the three-pass mix composition (its op-name // filters cannot see a whole-model op). Prefill-only and decode-only // dispatch through the same shared free fns as op-level (the FpmForward @@ -2093,26 +2129,10 @@ impl Engine { } if self.has_dsv41_stages() { - if has_prefill - && sched.num_prefill_requests > 1 - && self.context_ops.iter().any(|op| { - matches!(op, - Op::Dsv41Stage(stage) if stage.decoder_replay && stage.bounded) - }) - { - return Err(AicError::InvalidForwardPassMetrics( - "V4.1 Decoder replay requires per-request extend lengths; FPM v1 aggregates with multiple prefill requests cannot identify the tails, even when prompt-length variance is zero".into(), - )); - } - // FPM v1 variance measures complete prompt lengths, not this - // iteration's extends. Equal prompts can have different cached - // prefixes or completed chunks, so even zero variance cannot prove - // homogeneous tails. Bounded replay only accepts one prefill here; - // explicitly grouped static/mixed workloads keep their own paths. // Retain every scheduled token in balanced aggregate telemetry; // integer averages alone discard the remainder. FPM v1 does not - // carry individual extend lengths; this approximation is only used - // for multiple prefills when decoder replay does not bound them. + // carry individual extend lengths; multiple bounded-replay + // prefills were rejected before the whole-forward dispatch above. let mut prefills = Vec::new(); if has_prefill { let n = sched.num_prefill_requests; @@ -2316,6 +2336,7 @@ mod tests { weight_dtype: None, moe_dtype: None, activation_dtype: None, + fpm_fmha_dtype: None, kv_cache_dtype: None, }, speculative: nextn.map(|n| crate::SpeculativeConfig { nextn: Some(n) }), @@ -3244,6 +3265,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, sol_ops: vec![], }) }; @@ -3314,6 +3336,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, sol_ops: vec![], }); let spec = EngineSpec::new( @@ -3371,6 +3394,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, sol_ops: vec![], }) }; @@ -3382,6 +3406,241 @@ mod tests { Engine::build(spec, Arc::new(db)) } + fn fpm_replay_guard_probe(tmp: &std::path::Path, replay: bool) -> Engine { + use crate::perf_database::fpm_forward::tests::RowSpec; + let rows = [ + RowSpec { + workload_kind: "prefill", + batch_size: 2, + total_prefill_tokens: 511, + total_kv_read_tokens: 0, + latency_ms: 23.0, + ..RowSpec::default() + }, + RowSpec { + workload_kind: "prefill", + batch_size: 2, + total_prefill_tokens: 257, + total_kv_read_tokens: 513, + latency_ms: 31.0, + ..RowSpec::default() + }, + RowSpec { + workload_kind: "prefill", + batch_size: 2, + total_prefill_tokens: 1024, + total_kv_read_tokens: 1024, + latency_ms: 43.0, + ..RowSpec::default() + }, + RowSpec { + workload_kind: "prefill", + batch_size: 1, + total_prefill_tokens: 129, + total_kv_read_tokens: 513, + latency_ms: 17.0, + ..RowSpec::default() + }, + RowSpec { + workload_kind: "decode", + batch_size: 8, + total_prefill_tokens: 0, + total_kv_read_tokens: 4096, + latency_ms: 7.0, + ..RowSpec::default() + }, + ]; + let mut engine = build_fpm_engine_with_rows(tmp, &rows).unwrap(); + let Op::FpmForward(prefill) = &mut engine.context_ops[0] else { + unreachable!() + }; + // The real Python FPM rewrite retains the original V4.1 stage graph. + prefill.sol_ops = dsv41_probe_engine(replay).context_ops; + engine + } + + fn fpm_replay_guard_metrics(variance: f64) -> ForwardPassMetrics { + ForwardPassMetrics { + scheduled_requests: crate::ScheduledRequestMetrics { + num_prefill_requests: 2, + sum_prefill_tokens: 511, + // Prompt-length variance does not identify current extends; + // these totals also address the balanced [256, 255] row. + var_prefill_length: variance, + ..Default::default() + }, + ..Default::default() + } + } + + #[test] + fn fpm_replay_guard_rejects_before_balanced_table_lookup() { + let tmp = tempfile::tempdir().unwrap(); + let engine = fpm_replay_guard_probe(tmp.path(), true); + for variance in [0.0, 16512.25, 0.25, -1.0, f64::NAN, f64::INFINITY] { + let result = engine.forward_pass_time_ms(&[fpm_replay_guard_metrics(variance)]); + assert!( + matches!(result, Err(AicError::InvalidForwardPassMetrics(_))), + "bounded replay variance {variance:?} must not query a balanced table: {result:?}" + ); + } + } + + #[test] + fn fpm_replay_guard_rejects_mixed_prefill_before_table_composition() { + let tmp = tempfile::tempdir().unwrap(); + let engine = fpm_replay_guard_probe(tmp.path(), true); + let mut metrics = fpm_replay_guard_metrics(0.0); + metrics.scheduled_requests.num_decode_requests = 8; + metrics.scheduled_requests.sum_decode_kv_tokens = 4096; + assert!(matches!( + engine.forward_pass_time_ms(&[metrics]), + Err(AicError::InvalidForwardPassMetrics(_)) + )); + } + + #[test] + fn fpm_replay_guard_requires_a_bounded_stage() { + let tmp = tempfile::tempdir().unwrap(); + let mut engine = fpm_replay_guard_probe(tmp.path(), true); + let Op::FpmForward(prefill) = &mut engine.context_ops[0] else { + unreachable!() + }; + for op in &mut prefill.sol_ops { + let Op::Dsv41Stage(stage) = op else { + unreachable!() + }; + stage.bounded = false; + } + assert_eq!( + engine + .forward_pass_time_ms(&[fpm_replay_guard_metrics(16512.25)]) + .unwrap(), + 23.0 + ); + } + + #[test] + fn fpm_replay_guard_preserves_off_and_ordinary_models() { + let tmp = tempfile::tempdir().unwrap(); + let mut engine = fpm_replay_guard_probe(tmp.path(), false); + let metrics = fpm_replay_guard_metrics(16512.25); + assert_eq!( + engine.forward_pass_time_ms(&[metrics.clone()]).unwrap(), + 23.0 + ); + let Op::FpmForward(prefill) = &mut engine.context_ops[0] else { + unreachable!() + }; + prefill.sol_ops.clear(); + assert_eq!(engine.forward_pass_time_ms(&[metrics]).unwrap(), 23.0); + } + + #[test] + fn fpm_replay_guard_rejects_zero_variance_remainders() { + let tmp = tempfile::tempdir().unwrap(); + let engine = fpm_replay_guard_probe(tmp.path(), true); + let mut metrics = fpm_replay_guard_metrics(0.0); + metrics.scheduled_requests.sum_prefill_tokens = 257; + metrics.scheduled_requests.sum_prefill_kv_tokens = 513; + assert!(matches!( + engine.forward_pass_time_ms(&[metrics]), + Err(AicError::InvalidForwardPassMetrics(_)) + )); + } + + #[test] + fn fpm_replay_guard_rejects_equal_prompt_heterogeneous_extends() { + // Both prompts have 1024 tokens, but their (new, prefix) pairs are + // (1, 1023) and (1023, 1). The bounded tails total 129 tokens, while + // the same aggregate addresses a balanced table cell with 256 tails. + let requests = [(1024, 1, 1023), (1024, 1023, 1)]; + assert!( + requests + .iter() + .all(|&(prompt, query, prefix)| { prompt == 1024 && query + prefix == prompt }) + ); + let mut metrics = fpm_replay_guard_metrics(0.0); + metrics.scheduled_requests.sum_prefill_tokens = requests.iter().map(|r| r.1).sum(); + metrics.scheduled_requests.sum_prefill_kv_tokens = requests.iter().map(|r| r.2).sum(); + let tmp = tempfile::tempdir().unwrap(); + let engine = fpm_replay_guard_probe(tmp.path(), true); + // Exercise both the whole-forward rewrite and direct stage engine. + for probe in [&engine, &dsv41_probe_engine(true)] { + let error = probe + .forward_pass_time_ms(std::slice::from_ref(&metrics)) + .unwrap_err(); + assert!(matches!(error, AicError::InvalidForwardPassMetrics(_))); + assert!(error.to_string().contains("multiple prefill requests")); + } + // OFF still consumes exactly the table cell that ON must not borrow. + let off = fpm_replay_guard_probe(tmp.path(), false); + assert_eq!(off.forward_pass_time_ms(&[metrics]).unwrap(), 43.0); + } + + #[test] + fn fpm_replay_guard_preserves_single_prefill_and_rejects_invalid_variance() { + let tmp = tempfile::tempdir().unwrap(); + let engine = fpm_replay_guard_probe(tmp.path(), true); + let mut metrics = fpm_replay_guard_metrics(0.0); + metrics.scheduled_requests.num_prefill_requests = 1; + metrics.scheduled_requests.sum_prefill_tokens = 129; + metrics.scheduled_requests.sum_prefill_kv_tokens = 513; + assert_eq!( + engine + .forward_pass_time_ms(std::slice::from_ref(&metrics)) + .unwrap(), + 17.0 + ); + for variance in [-1.0, f64::NAN, f64::INFINITY] { + metrics.scheduled_requests.var_prefill_length = variance; + let error = engine + .forward_pass_time_ms(std::slice::from_ref(&metrics)) + .unwrap_err(); + assert!(matches!(error, AicError::InvalidForwardPassMetrics(_))); + assert!(error.to_string().contains("finite nonnegative")); + } + } + + #[test] + fn fpm_replay_guard_preserves_explicit_static_geometry() { + let tmp = tempfile::tempdir().unwrap(); + let engine = fpm_replay_guard_probe(tmp.path(), true); + // Static input explicitly declares two identical (new, prefix) + // pairs. It is not the ambiguous FPM v1 telemetry entry point. + let mut input = runtime(2, 1024, 1); + input.prefix = 512; + assert_eq!( + engine + .run_static(&input, StaticMode::Context, 32) + .unwrap() + .context_ms, + 43.0 + ); + } + + #[test] + fn fpm_replay_guard_ignores_prefill_metadata_without_compute() { + let tmp = tempfile::tempdir().unwrap(); + let engine = fpm_replay_guard_probe(tmp.path(), true); + let mut metrics = fpm_replay_guard_metrics(16512.25); + metrics.scheduled_requests.sum_prefill_tokens = 0; + metrics.scheduled_requests.sum_prefill_kv_tokens = 513; + assert_eq!( + engine.forward_pass_time_ms(&[metrics.clone()]).unwrap(), + 0.0 + ); + metrics.scheduled_requests.num_decode_requests = 8; + metrics.scheduled_requests.sum_decode_kv_tokens = 4096; + assert_eq!( + engine.forward_pass_time_ms(&[metrics.clone()]).unwrap(), + 7.0 + ); + metrics.scheduled_requests.num_prefill_requests = 0; + metrics.scheduled_requests.sum_prefill_kv_tokens = 0; + assert_eq!(engine.forward_pass_time_ms(&[metrics]).unwrap(), 7.0); + } + fn cliff_rows() -> Vec { use crate::perf_database::fpm_forward::tests::RowSpec; let mk = |kind: &'static str, batch: u32, prefill: u32, kv: u32, lat: f64| RowSpec { @@ -3680,6 +3939,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, sol_ops: vec![], })], vec![], @@ -3997,6 +4257,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 8, + original_fmha_quant_mode: None, sol_ops: vec![], }; // 8 requests x width 8 arrive as batch = 8 widened tokens with ... @@ -4042,6 +4303,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: width, + original_fmha_quant_mode: None, sol_ops: vec![], }) }; diff --git a/crates/core/src/perfmodel/engine/spec.rs b/crates/core/src/perfmodel/engine/spec.rs index 48d20a46f..856a6a44b 100644 --- a/crates/core/src/perfmodel/engine/spec.rs +++ b/crates/core/src/perfmodel/engine/spec.rs @@ -642,6 +642,7 @@ mod tests { weight_bytes: 1.5e10, // Non-default on purpose: the round-trip must preserve the field. verify_width: 8, + original_fmha_quant_mode: Some("fp8".into()), sol_ops: vec![ OpSpec::Gemm(gemm()), OpSpec::ContextAttention(context_attention()), @@ -749,12 +750,6 @@ mod tests { hc_mult: 4, tp_size: 4, }), - OpSpec::Dsv41Linear(Dsv41LinearOp { - name: "v41_linear".into(), - n: 1152, - k: 5120, - quant_mode: GemmQuantMode::Fp8Block, - }), OpSpec::Dsv41Stage(Dsv41StageOp { name: "v41_stage".into(), is_context: true, @@ -763,6 +758,12 @@ mod tests { window_size: 128, children: vec![OpSpec::Gemm(gemm())], }), + OpSpec::Dsv41Linear(Dsv41LinearOp { + name: "v41_linear".into(), + n: 1152, + k: 5120, + quant_mode: GemmQuantMode::Fp8Block, + }), ]; // Exhaustiveness guard: if a variant is added to `Op`, this match @@ -839,6 +840,7 @@ mod tests { weight_dtype: Some(DataType::Fp8), moe_dtype: Some(DataType::Fp8), activation_dtype: Some(DataType::Fp8), + fpm_fmha_dtype: None, kv_cache_dtype: Some(DataType::Fp8), }, speculative: Some(SpeculativeConfig { nextn: Some(1) }), @@ -902,8 +904,7 @@ mod tests { assert_eq!(MOE_EXPERT_COMPUTE_INDEX, MOE_ALL_TO_ALL_INDEX + 1); assert_eq!(TOKEN_SCALE_INDEX, MOE_EXPERT_COMPUTE_INDEX + 1); // Keep the main-branch TokenScale index; V41 variants append after it. - let mut appended: Vec<_> = all_op_variants().iter().skip(36).map(index_of).collect(); - appended.sort(); + let appended: Vec<_> = all_op_variants().iter().skip(36).map(index_of).collect(); assert_eq!( appended, vec![36, 37, 38, 39, 40], @@ -1108,11 +1109,11 @@ mod tests { Err(AicError::UnsupportedSchemaVersion { kind: "EngineSpec", got: 18, - expected: 19 + expected: ENGINE_SPEC_SCHEMA_VERSION }) )); - // A false schema19 stamp cannot silently use the JSON-only default. - bytes[..4].copy_from_slice(&19u32.to_le_bytes()); + // A false current-schema stamp cannot silently use the JSON-only default. + bytes[..4].copy_from_slice(&ENGINE_SPEC_SCHEMA_VERSION.to_le_bytes()); assert!(matches!( EngineSpec::from_bincode(&bytes), Err(AicError::EngineSpec(_)) @@ -1269,4 +1270,27 @@ mod tests { }) )); } + #[test] + fn fpm_selector_json_default_and_schema19_rejection() { + let selected = fpm_forward(); + assert_eq!(selected.original_fmha_quant_mode.as_deref(), Some("fp8")); + let mut json = serde_json::to_value(selected).unwrap(); + json.as_object_mut() + .unwrap() + .remove("original_fmha_quant_mode"); + let legacy: crate::operators::FpmForwardOp = serde_json::from_value(json).unwrap(); + assert_eq!(legacy.original_fmha_quant_mode, None); + let mut bytes = handshake_spec().to_bincode().unwrap(); + bytes[..4].copy_from_slice(&19u32.to_le_bytes()); + bytes.truncate(4); + assert!(matches!( + EngineSpec::from_bincode(&bytes), + Err(AicError::UnsupportedSchemaVersion { + got: 19, + expected: ENGINE_SPEC_SCHEMA_VERSION, + + .. + }) + )); + } } diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 25a1c8da4..3b802c9b8 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -105,6 +105,9 @@ pub struct ForwardPassPerfModelConfig { pub moe_quant_mode: Option, #[serde(default, alias = "fmha_dtype")] pub fmha_quant_mode: Option, + /// Exact recorded FPM table label; does not override model arithmetic. + #[serde(default, alias = "fpm_fmha_dtype")] + pub fpm_fmha_quant_mode: Option, #[serde(default, alias = "kv_cache_dtype")] pub kvcache_quant_mode: Option, #[serde(default, alias = "comm_dtype")] @@ -169,6 +172,7 @@ impl ForwardPassPerfModelConfig { gemm_quant_mode: None, moe_quant_mode: None, fmha_quant_mode: None, + fpm_fmha_quant_mode: None, kvcache_quant_mode: None, comm_quant_mode: None, nextn: 0, @@ -303,6 +307,18 @@ impl ForwardPassPerfModelConfig { "estimation_mode='fpm_interpolation' does not support MTP speculative decoding", )); } + if let Some(selector) = self.fpm_fmha_quant_mode.as_deref() { + if self.estimation_mode != EstimationMode::FpmInterpolation { + return Err(invalid_config( + "fpm_fmha_quant_mode requires estimation_mode='fpm_interpolation'", + )); + } + if !matches!(selector, "bfloat16" | "fp8" | "fp8_block") { + return Err(invalid_config( + "fpm_fmha_quant_mode must be bfloat16, fp8, or fp8_block", + )); + } + } TransferPolicy::from_wire(self.transfer_policy.as_deref()).map_err(invalid_config)?; for root in &self.systems_paths { if root.to_str().is_none() { @@ -341,6 +357,33 @@ mod tests { serde_json::from_value(base).unwrap() } + #[test] + fn fpm_selector_is_validated_and_preserved_in_canonical_identity() { + let cfg = config(serde_json::json!({ + "estimation_mode": "fpm_interpolation", "fpm_fmha_quant_mode": "fp8" + })); + cfg.validate().unwrap(); + let restored: ForwardPassPerfModelConfig = + serde_json::from_str(&serde_json::to_string(&cfg).unwrap()).unwrap(); + assert_eq!(restored, cfg); + assert_eq!(restored.fpm_fmha_quant_mode.as_deref(), Some("fp8")); + for mode in ["auto", "op_level", "fpm_regression"] { + let invalid = + config(serde_json::json!({"estimation_mode": mode, "fpm_fmha_quant_mode": "fp8"})); + assert!( + invalid + .validate() + .unwrap_err() + .to_string() + .contains("requires estimation_mode='fpm_interpolation'") + ); + } + let invalid = config( + serde_json::json!({"estimation_mode": "fpm_interpolation", "fpm_fmha_quant_mode": "bogus"}), + ); + assert!(invalid.validate().is_err()); + } + #[test] fn engine_controls_round_trip_with_identity_and_policy() { let cfg = config(serde_json::json!({ diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index 23b740697..e75d36ba0 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -122,6 +122,7 @@ fn fixture_engine_config() -> EngineConfig { weight_dtype: None, moe_dtype: None, activation_dtype: None, + fpm_fmha_dtype: None, kv_cache_dtype: None, }, speculative: None, diff --git a/crates/core/src/perfmodel/memory.rs b/crates/core/src/perfmodel/memory.rs index c30377599..820973e23 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -522,6 +522,7 @@ mod tests { weight_dtype: None, moe_dtype: None, activation_dtype: None, + fpm_fmha_dtype: None, kv_cache_dtype: None, }, speculative: None, diff --git a/crates/core/src/perfmodel/operators/attention.rs b/crates/core/src/perfmodel/operators/attention.rs index 1f8d77c55..604e77ace 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 19). + /// payloads are positional (current ENGINE_SPEC_SCHEMA_VERSION 20). #[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 19). +/// field — introduced at schema v8, current ENGINE_SPEC_SCHEMA_VERSION 20). /// 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 19). + /// bincode payloads are positional, current ENGINE_SPEC_SCHEMA_VERSION 20). #[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 c86e8eba3..718e121ae 100644 --- a/crates/core/src/perfmodel/operators/fpm_forward.rs +++ b/crates/core/src/perfmodel/operators/fpm_forward.rs @@ -59,7 +59,7 @@ impl FpmPhase { /// One whole-model forward pass for a single phase. /// -/// `match_identity` is the 11-string cell identity in +/// `match_identity` is the 19-string 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 @@ -91,6 +91,10 @@ pub struct FpmForwardOp { #[serde(default = "default_verify_width")] pub verify_width: u32, pub sol_ops: Vec, + /// Present only for an explicit table selector. Preserves the model's mode + /// before match_identity[2] was selected; it is NOT observed runtime precision. + #[serde(default)] + pub original_fmha_quant_mode: Option, } fn default_verify_width() -> u32 { @@ -102,6 +106,22 @@ fn data_err(msg: String) -> AicError { } 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)?; + 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, + // once-per-cell/model-mode warning rather than a silent log record. + eprintln!("WARNING: {warning}"); + } + } + Ok(cell) + } + /// Mirror of Python `FPMForwardOp.query`: validate kwargs, map to /// iteration-total coordinates, resolve against the selected cell. pub fn query( @@ -136,9 +156,7 @@ impl FpmForwardOp { ctx.beam_width ))); } - let cell = db - .fpm_forward - .select_cell(&self.match_identity, &self.model_path)?; + let cell = self.select_cell(db)?; let b = batch_size as f64; let coords: Vec = match self.phase { FpmPhase::Prefill => { @@ -182,9 +200,7 @@ impl FpmForwardOp { coords ))); } - let cell = db - .fpm_forward - .select_cell(&self.match_identity, &self.model_path)?; + let cell = self.select_cell(db)?; self.resolve(db, cell, coords) } @@ -196,9 +212,7 @@ impl FpmForwardOp { self.phase.as_str() ))); } - let cell = db - .fpm_forward - .select_cell(&self.match_identity, &self.model_path)?; + let cell = self.select_cell(db)?; Ok(cell.decode_domain.as_ref().map(|domain| domain[1].1)) } @@ -229,9 +243,7 @@ impl FpmForwardOp { "invalid FPM baseline query: batch_size={batch_size}" ))); } - let cell = db - .fpm_forward - .select_cell(&self.match_identity, &self.model_path)?; + let cell = self.select_cell(db)?; let Some(domain) = cell.decode_domain else { return Err(data_err(format!( "FPM cell {:?} has no decode rows (model_path={:?}).", @@ -660,6 +672,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, // Empty sol_ops: exact hits and in-curve lerps never call SOL. sol_ops: vec![], } @@ -1350,4 +1363,60 @@ mod tests { ); } } + #[test] + fn selector_warning_preserves_values_and_rejects_wrong_recorded_precision() { + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &default_rows()); + let db = db_with_pair(tmp.path()); + let mut selected = op(FpmPhase::Decode); + // Independent original model mode differs from the recorded table label. + selected.original_fmha_quant_mode = Some("fp8".into()); + let cell = db + .fpm_forward + .select_cell(&selected.match_identity, &selected.model_path) + .unwrap(); + let message = cell.fmha_selector_warning("fp8").unwrap(); + assert!(message.contains("original_model_mode=\"fp8\"")); + assert!(message.contains("selector=\"bfloat16\"")); + assert!(message.contains(&format!("matched_cell_ids={:?}", cell.cell_ids))); + assert!(message.contains("does not independently verify runtime attention precision")); + assert!(cell.fmha_selector_warning("fp8").is_none()); + assert!(cell.fmha_selector_warning("float16").is_some()); + let base = op(FpmPhase::Decode); + let coords = [8.0, 4096.0]; + assert_eq!( + selected.query_totals(&db, &coords).unwrap(), + base.query_totals(&db, &coords).unwrap() + ); + selected.match_identity[2] = "fp8".into(); + // All query routes retain exact matching; no substitution for table precision. + assert!( + selected + .query(&db, &ctx(4, 1024, 0)) + .unwrap_err() + .to_string() + .contains("No FPM cell matches") + ); + assert!( + selected + .query_totals(&db, &coords) + .unwrap_err() + .to_string() + .contains("No FPM cell matches") + ); + assert!( + selected + .query_pass_baseline(&db, 4, 4096.0) + .unwrap_err() + .to_string() + .contains("No FPM cell matches") + ); + assert!( + selected + .decode_kv_ceiling(&db) + .unwrap_err() + .to_string() + .contains("No FPM cell matches") + ); + } } diff --git a/crates/core/src/perfmodel/operators/fpm_sol.rs b/crates/core/src/perfmodel/operators/fpm_sol.rs index ddc48ba4d..2c84c2b5a 100644 --- a/crates/core/src/perfmodel/operators/fpm_sol.rs +++ b/crates/core/src/perfmodel/operators/fpm_sol.rs @@ -116,6 +116,24 @@ pub(crate) fn op_sol_latency_ms( Op::Elementwise(o) => Ok(elementwise_sol(o, spec, x)), Op::ContextAttention(o) => Ok(context_attention_sol(o, spec, batch, s, prefix)), Op::GenerationAttention(o) => Ok(generation_attention_sol(o, spec, batch, s)), + Op::Dsv41Attention(o) => Ok(o.sol(spec, batch, s, prefix)?.latency_ms), + Op::Dsv41Linear(o) => Ok(o.sol(spec, x)?.latency_ms), + Op::Dsv41Mhc(o) => Ok(o.sol(spec, x)?.latency_ms), + Op::Dsv41Engram(o) => Ok(o.sol(spec, x)?.latency_ms), + Op::Dsv41Stage(o) => { + let (stage_s, stage_prefix) = o.scope(s, prefix); + let stage_x = if o.is_context { batch * stage_s } else { x }; + let mut total = 0.0; + for inner in &o.children { + let child_x = if inner.is_logits_gemm() { + batch + } else { + stage_x + }; + total += op_sol_latency_ms(inner, db, child_x, batch, stage_s, stage_prefix)?; + } + Ok(total) + } Op::DsaContext(o) => dsa_context_module_sol(o, spec, batch, s, prefix), Op::DsaGeneration(o) => dsa_generation_module_sol(o, spec, batch, s), Op::Moe(o) => Ok(moe_sol(o, spec, x)), @@ -661,6 +679,52 @@ mod tests { ); } + #[test] + fn v41_stage_roofline_retains_fractional_tokens_and_bounds_only_prefill() { + use crate::operators::{Dsv41MhcOp, Dsv41StageOp}; + + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../python/aisimulate/src/aisimulate_core/systems"); + let db = PerfDatabase::load(&root, "gb200", "vllm", "0.19.0").expect("GB200 db"); + let leaf = Dsv41MhcOp { + name: "mhc".into(), + hidden_size: 32, + hc_mult: 4, + sinkhorn_iters: 20, + }; + let mut stage = Dsv41StageOp { + name: "decoder".into(), + is_context: true, + decoder_replay: true, + bounded: true, + window_size: 128, + children: vec![Op::Dsv41Mhc(leaf.clone())], + }; + // One balanced iteration has 385 new tokens over three requests; + // replay executes 128 each, while a short fractional extend stays short. + approx( + op_sol_latency_ms( + &Op::Dsv41Stage(stage.clone()), + &db, + 385.0, + 3.0, + 385.0 / 3.0, + 17.5, + ) + .unwrap(), + leaf.sol(&db.system_spec, 384.0).unwrap().latency_ms, + ); + approx( + op_sol_latency_ms(&Op::Dsv41Stage(stage.clone()), &db, 2.5, 1.0, 2.5, 17.5).unwrap(), + leaf.sol(&db.system_spec, 2.5).unwrap().latency_ms, + ); + stage.is_context = false; + approx( + op_sol_latency_ms(&Op::Dsv41Stage(stage), &db, 3.0, 3.0, 2048.0, 0.0).unwrap(), + leaf.sol(&db.system_spec, 3.0).unwrap().latency_ms, + ); + } + /// Python oracle: /// PYTHONPATH=aic-core/src python3 -c " /// from aisimulate_core.sdk import perf_database, common diff --git a/crates/core/src/perfmodel/perf_database/fpm_forward.rs b/crates/core/src/perfmodel/perf_database/fpm_forward.rs index 449261810..91feb01c9 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -20,7 +20,7 @@ //! `kv_seed_regime == "fake_fallback"` values (fabricated-KV measurements //! replaced in memory by in-station extrapolation; see //! [`FPM_KV_SEED_FAKE_FALLBACK`]), and grouped into cells keyed by -//! `(model_path, 15 identity columns)`. Each +//! `(model_path, 19 identity columns)`. Each //! cell holds one nested table per phase — prefill //! `[batch][total_prefill][total_kv]`, decode `[batch][total_kv]` — plus the //! per-phase axis-aligned domain bounding box and a prebuilt @@ -40,9 +40,9 @@ //! online-tuning model over Dynamo ForwardPassMetrics telemetry — an //! unrelated concept that also abbreviates to "FPM". -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::path::{Path, PathBuf}; -use std::sync::OnceLock; +use std::sync::{Mutex, OnceLock}; use sha2::{Digest, Sha256}; @@ -52,7 +52,7 @@ use crate::common::error::AicError; pub const FPM_FORWARD_BASENAME: &str = "fpm_forward_perf.parquet"; pub const FPM_FORWARD_SCHEMA_NAME: &str = "aic_fpm_forward_perf"; -pub const FPM_FORWARD_SCHEMA_VERSION: u64 = 6; +pub const FPM_FORWARD_SCHEMA_VERSION: u64 = 7; pub const FPM_FORWARD_COORDINATE_SYSTEM: &str = "iteration_totals_balanced_v1"; pub const FPM_FORWARD_PARTITION_POLICY: &str = "balanced_v1"; /// The only measurement policy the collector publishes; pinned in the @@ -89,7 +89,7 @@ pub const FPM_FAKE_FALLBACK_RAW_ENV: &str = "AIC_FPM_FAKE_FALLBACK_RAW"; /// 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. -pub const FPM_CELL_MATCH_COLUMNS: [&str; 15] = [ +pub const FPM_CELL_MATCH_COLUMNS: [&str; 19] = [ "gemm_quant_mode", "moe_quant_mode", "fmha_quant_mode", @@ -105,15 +105,20 @@ pub const FPM_CELL_MATCH_COLUMNS: [&str; 15] = [ "attention_backend", "enable_wideep", "enable_eplb", + "model_config_sha256", + "execution_profile", + "engram_residency", + "input_modality", ]; +const LEGACY_EXECUTION_IDENTITY: [&str; 4] = ["", "full", "none", "text"]; + pub const FPM_PREFILL_AXES: [&str; 3] = ["batch_size", "total_prefill_tokens", "total_kv_read_tokens"]; pub const FPM_DECODE_AXES: [&str; 2] = ["batch_size", "total_kv_read_tokens"]; /// One collected cell: the tables and domains for a single -/// `(model_path, identity)` tuple (the 15-column identity carries the -/// backend knobs since schema v6). +/// `(model_path, identity)` tuple (schema v6 added backend knobs; schema v7 adds execution identity). #[derive(Debug)] pub struct FpmForwardCell { pub model_path: String, @@ -144,6 +149,24 @@ pub struct FpmForwardCell { pub decode_batches: Vec, pub decode_rungs: Vec, pub decode_curve_bounds: BTreeMap, + /// Diagnostic state only; no change to table identity, values, or schema. + warned_fmha_model_modes: Mutex>, +} + +impl FpmForwardCell { + pub(crate) fn fmha_selector_warning(&self, original: &str) -> Option { + let mut warned = self + .warned_fmha_model_modes + .lock() + .unwrap_or_else(|e| e.into_inner()); + if !warned.insert(original.to_owned()) { + return None; + } + Some(format!( + "FPM table FMHA selector: original_model_mode={original:?}, selector={:?}, matched_cell_ids={:?}, model_path={:?}. Exact recorded-label matching does not independently verify runtime attention precision; arithmetic and memory modes are unchanged.", + self.match_identity[2], self.cell_ids, self.model_path + )) + } } /// One loaded parquet/sidecar pair: the grouped cells plus the @@ -319,7 +342,15 @@ impl FpmForwardTable { // collected model_path could silently answer for a different model). let matches: Vec<&FpmForwardCell> = cells .iter() - .filter(|cell| cell.match_identity == match_identity && cell.model_path == model_path) + .filter(|cell| { + let legacy = 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 + }) .collect(); if matches.is_empty() { let mut available: Vec = cells @@ -381,7 +412,7 @@ fn validate_sidecar( system: &str, backend: &str, version: &str, -) -> Result, AicError> { +) -> Result<(Option, u64), AicError> { if !metadata_path.exists() { return Err(structural(format!( "FPM database is missing its metadata sidecar: {}. \ @@ -412,7 +443,8 @@ fn validate_sidecar( metadata_path.display() ))); } - if json_uint(metadata.get("schema_version")) != Some(FPM_FORWARD_SCHEMA_VERSION) { + let schema_version = json_uint(metadata.get("schema_version")); + if !matches!(schema_version, Some(6 | FPM_FORWARD_SCHEMA_VERSION)) { return Err(structural(format!( "unsupported FPM schema_version={:?} (expected {FPM_FORWARD_SCHEMA_VERSION}): {}", metadata.get("schema_version"), @@ -462,7 +494,10 @@ fn validate_sidecar( parquet_path.parent().unwrap_or(parquet_path).display() ))); } - Ok(json_uint(metadata.get("row_count"))) + Ok(( + json_uint(metadata.get("row_count")), + schema_version.unwrap(), + )) } /// Python `metadata.get(k) != n` compares by VALUE: a JSON `5.0` equals the @@ -520,7 +555,7 @@ fn load_pair( return Ok(None); } let metadata_path = parquet_path.with_extension("metadata.json"); - let sidecar_row_count = + let (sidecar_row_count, schema_version) = validate_sidecar(&metadata_path, parquet_path, system, backend, version)?; let reader = PerfReader::open(parquet_path)?; @@ -546,6 +581,11 @@ fn load_pair( for name in str_cols { str_idx.insert(name, reader.col(name)?); } + if schema_version >= 7 { + for name in &FPM_CELL_MATCH_COLUMNS[15..] { + str_idx.insert(*name, reader.col(name)?); + } + } let int_cols = [ "tp", "pp", @@ -702,6 +742,11 @@ fn load_pair( get_str(name) } else if bool_idx.contains_key(name) { get_bool_identity(name) + } else if let Some(offset) = FPM_CELL_MATCH_COLUMNS[15..] + .iter() + .position(|field| field == name) + { + Ok(LEGACY_EXECUTION_IDENTITY[offset].to_string()) } else { get_int_identity(name) } @@ -739,6 +784,10 @@ fn load_pair( match_identity[12].clone(), match_identity[13].clone(), match_identity[14].clone(), + match_identity[15].clone(), + match_identity[16].clone(), + match_identity[17].clone(), + match_identity[18].clone(), workload_kind.clone(), batch_size.to_string(), total_prefill_tokens.to_string(), @@ -751,6 +800,17 @@ fn load_pair( // their latency value is healed by the replacement pass below, after // the duplicate/collision checks. let kv_seed_regime = row.str_optional(kv_seed_col)?.unwrap_or(""); + if !match_identity[15..] + .iter() + .map(String::as_str) + .eq(LEGACY_EXECUTION_IDENTITY) + && (workload_kind == "decode" || total_kv_read_tokens > 0) + && kv_seed_regime != FPM_KV_SEED_REAL_KV + { + return Err(structural(format!( + "FPM row {index}: config-bound cached prefill/decode requires real_kv provenance" + ))); + } let fake_fallback = kv_seed_regime == FPM_KV_SEED_FAKE_FALLBACK; let real_kv_anchor = kv_seed_regime == FPM_KV_SEED_REAL_KV; @@ -899,6 +959,7 @@ fn load_pair( decode_batches: Vec::new(), decode_rungs: Vec::new(), decode_curve_bounds: BTreeMap::new(), + warned_fmha_model_modes: Mutex::new(BTreeSet::new()), }, }); if !building.cell.cell_ids.contains(&row.cell_id) { @@ -1151,6 +1212,7 @@ pub(crate) mod tests { /// fixture sets it, omits the column entirely — the pre-column /// legacy layout). pub kv_seed_regime: Option<&'static str>, + pub execution: Option<[&'static str; 4]>, } impl Default for RowSpec { @@ -1173,6 +1235,7 @@ pub(crate) mod tests { system: "b200_sxm", backend: "vllm", kv_seed_regime: None, + execution: None, } } } @@ -1191,7 +1254,7 @@ pub(crate) mod tests { .collect() } - /// The 11-string identity every default row carries, in + /// The normalized identity every default row carries, in /// `FPM_CELL_MATCH_COLUMNS` order. pub(crate) fn default_identity(tp: u32) -> Vec { vec![ @@ -1210,10 +1273,14 @@ pub(crate) mod tests { "auto".to_string(), // attention_backend "False".to_string(), // enable_wideep (str(bool)) "False".to_string(), // enable_eplb + "".to_string(), + "full".to_string(), + "none".to_string(), + "text".to_string(), ] } - /// Write the v5-schema parquet + sha256'd sidecar pair into `dir`. + /// Write the legacy-v6 or execution-bound-v7 parquet + sha256'd sidecar pair into `dir`. pub(crate) fn write_pair(dir: &Path, rows: &[RowSpec]) -> PathBuf { write_pair_with(dir, rows, |_| {}) } @@ -1233,6 +1300,7 @@ 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_execution = rows.iter().any(|r| r.execution.is_some()); let schema = "message schema { REQUIRED BINARY cell_id (UTF8); REQUIRED BINARY model_path (UTF8); @@ -1270,6 +1338,18 @@ pub(crate) mod tests { } else { schema.to_string() }; + let schema = if has_execution { + schema.replace( + "REQUIRED BINARY workload_kind (UTF8);", + "REQUIRED BINARY model_config_sha256 (UTF8); + REQUIRED BINARY execution_profile (UTF8); + REQUIRED BINARY engram_residency (UTF8); + REQUIRED BINARY input_modality (UTF8); + REQUIRED BINARY workload_kind (UTF8);", + ) + } else { + schema + }; let schema = Arc::new(parse_message_type(&schema).expect("schema must parse")); let file = std::fs::File::create(&parquet_path).expect("create parquet"); let mut writer = @@ -1349,6 +1429,18 @@ pub(crate) mod tests { .expect("write"); col.close().expect("close"); } + if has_execution { + for index in 0..4 { + let values = str_col(&|r| { + r.execution.unwrap_or(LEGACY_EXECUTION_IDENTITY)[index].to_string() + }); + let mut col = rg.next_column().expect("next col").expect("str col"); + col.typed::() + .write_batch(&values, None, None) + .expect("write"); + col.close().expect("close"); + } + } { let values = str_col(&|r| r.workload_kind.to_string()); let mut col = rg.next_column().expect("next col").expect("str col"); @@ -1412,7 +1504,11 @@ pub(crate) mod tests { sidecar.insert("schema_name".into(), FPM_FORWARD_SCHEMA_NAME.into()); sidecar.insert( "schema_version".into(), - serde_json::Value::from(FPM_FORWARD_SCHEMA_VERSION), + serde_json::Value::from(if has_execution { + FPM_FORWARD_SCHEMA_VERSION + } else { + 6 + }), ); sidecar.insert( "coordinate_system".into(), @@ -1482,6 +1578,107 @@ pub(crate) mod tests { ) } + #[test] + fn schema_seven_execution_identity_is_exact_and_requires_real_decode_kv() { + let tmp = tempfile::tempdir().unwrap(); + let execution = ["abc", "full", "hbm_tp_sharded", "text"]; + let mut rows = default_rows(); + for row in &mut rows { + row.execution = Some(execution); + row.kv_seed_regime = Some("real_kv"); + } + write_pair(tmp.path(), &rows); + 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(&default_identity(4), "org/model-a") + .is_err() + ); + identity[16] = "decoder_bounded".to_string(); + assert!(table.select_cell(&identity, "org/model-a").is_err()); + let prefill_rows: Vec = default_rows() + .into_iter() + .filter(|r| r.workload_kind == "prefill") + .map(|mut r| { + r.execution = Some(execution); + r.kv_seed_regime = Some("fake_fallback"); + r + }) + .collect(); + write_pair(tmp.path(), &prefill_rows); + assert!( + loaded_table(tmp.path()) + .cells() + .unwrap_err() + .to_string() + .contains("requires real_kv") + ); + for row in &mut rows { + row.kv_seed_regime = Some("fake_fallback"); + } + write_pair(tmp.path(), &rows); + assert!( + loaded_table(tmp.path()) + .cells() + .unwrap_err() + .to_string() + .contains("requires real_kv") + ); + } + + #[test] + fn every_nonlegacy_execution_tail_requires_real_kv_even_without_fingerprint() { + let tmp = tempfile::tempdir().unwrap(); + for execution in [ + ["", "decoder_bounded", "none", "text"], + ["", "full", "hbm_tp_sharded", "text"], + ["", "full", "none", "image"], + LEGACY_EXECUTION_IDENTITY, + ] { + for phase in ["prefill", "decode"] { + let mut rows: Vec<_> = default_rows() + .into_iter() + .filter(|r| r.workload_kind == phase) + .collect(); + for row in &mut rows { + row.execution = Some(execution); + row.kv_seed_regime = Some("fake_fallback"); + } + write_pair(tmp.path(), &rows); + let table = loaded_table(tmp.path()); + let result = table.cells(); + if execution == LEGACY_EXECUTION_IDENTITY { + assert!(result.is_ok()); + } else { + assert!(result.unwrap_err().to_string().contains("requires real_kv")); + for row in &mut rows { + row.kv_seed_regime = Some("real_kv"); + } + write_pair(tmp.path(), &rows); + assert!(loaded_table(tmp.path()).cells().is_ok()); + } + } + } + } + + #[test] + fn legacy_identity_is_upgraded_but_schema_seven_cannot_omit_execution_fields() { + let tmp = tempfile::tempdir().unwrap(); + write_pair(tmp.path(), &default_rows()); + assert!( + loaded_table(tmp.path()) + .select_cell(&default_identity(4)[..15], "org/model-a") + .is_ok() + ); + write_pair_with(tmp.path(), &default_rows(), |m| { + m.insert("schema_version".into(), 7.into()); + }); + assert!(loaded_table(tmp.path()).cells().is_err()); + } + #[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 933a04e08..6c6039d08 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1045,6 +1045,7 @@ struct EngineBuildRequest { moe_quant_mode: Option, kvcache_quant_mode: Option, fmha_quant_mode: Option, + fpm_fmha_quant_mode: Option, comm_quant_mode: Option, attention_backend: Option, moe_backend: Option, @@ -1095,6 +1096,7 @@ impl AicEngineBuilder { moe_quant_mode: None, kvcache_quant_mode: None, fmha_quant_mode: None, + fpm_fmha_quant_mode: None, comm_quant_mode: None, attention_backend: None, moe_backend: None, @@ -1207,6 +1209,12 @@ impl AicEngineBuilder { self } + /// Select a whole-forward FPM cell without changing the SOL arithmetic. + pub fn fpm_fmha_dtype(mut self, value: impl Into) -> Self { + self.request.fpm_fmha_quant_mode = Some(value.into()); + self + } + /// Override the FMHA quantization mode. pub fn fmha_quant_mode(mut self, value: impl Into) -> Self { self.request.fmha_quant_mode = Some(value.into()); @@ -1334,6 +1342,8 @@ mod builder_tests { "weight_dtype": null, "moe_dtype": null, "activation_dtype": null, + "fpm_fmha_dtype": "fp8", + "forward_model": "fpm", "kv_cache_dtype": null, "database_mode": "EMPIRICAL", "enable_shared_layer": true, @@ -1346,6 +1356,8 @@ mod builder_tests { let request = engine_build_request(&config, None).unwrap(); assert_eq!(request.database_mode.as_deref(), Some("EMPIRICAL")); + assert_eq!(request.fpm_fmha_quant_mode.as_deref(), Some("fp8")); + assert_eq!(request.fmha_quant_mode, None); assert_eq!(request.shared_layer, Some(true)); assert_eq!( request.transfer_policy.as_deref(), @@ -1383,6 +1395,15 @@ mod builder_tests { if message.contains("SOL_FULL") && message.contains("per-call diagnostic") )); } + + #[test] + fn builder_rejects_fpm_selector_without_fpm_before_loading_python() { + let result = AicEngineBuilder::new("model", "system", BackendKind::Vllm) + .fpm_fmha_dtype("fp8") + .build(); + assert!(matches!(result, Err(AicError::InvalidEngineConfig(message)) + if message.contains("requires forward_model='fpm'"))); + } } /// Construct the public handle from the one canonical build request. @@ -1398,10 +1419,17 @@ fn build_engine_from_request(request: EngineBuildRequest) -> Result Result { + if request.fpm_fmha_quant_mode.is_some() && request.forward_model.as_deref() != Some("fpm") { + return Err(AicError::InvalidEngineConfig( + "fpm_fmha_dtype requires forward_model='fpm'".into(), + )); + } + crate::config::validate_fpm_parquet_path( request.fpm_parquet_path.as_deref().map(Path::new), request.forward_model.as_deref() == Some("fpm"), )?; + if request.database_mode.as_deref() == Some(DatabaseMode::SolFull.as_str()) { return Err(AicError::InvalidEngineConfig( "database mode SOL_FULL is a per-call diagnostic and cannot be an engine default; use SOL instead" @@ -1429,6 +1457,10 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result, ) -> Result { + if config.quantization.fpm_fmha_dtype.is_some() + && fmha_quant_name(config.quantization.fpm_fmha_dtype.as_ref()).is_none() + { + return Err(AicError::InvalidEngineConfig( + "fpm_fmha_dtype must be bfloat16, fp8, or fp8_block".into(), + )); + } compile_engine_from_request(engine_build_request(config, systems_path)?) } @@ -1595,6 +1639,8 @@ fn engine_build_request( .map(str::to_owned), fmha_quant_mode: fmha_quant_name(config.quantization.activation_dtype.as_ref()) .map(str::to_owned), + fpm_fmha_quant_mode: fmha_quant_name(config.quantization.fpm_fmha_dtype.as_ref()) + .map(str::to_owned), // Comm quant is not carried on EngineConfig; let Python default it. comm_quant_mode: None, // Attention backend is not carried on EngineConfig; let Python resolve it. @@ -1795,6 +1841,7 @@ impl PyForwardPassPerfModel { gemm_quant_mode: request.gemm_quant_mode, moe_quant_mode: request.moe_quant_mode, fmha_quant_mode: request.fmha_quant_mode, + fpm_fmha_quant_mode: request.fpm_fmha_quant_mode, kvcache_quant_mode: request.kvcache_quant_mode, comm_quant_mode: request.comm_quant_mode, nextn: request.nextn, @@ -2073,6 +2120,7 @@ mod tests { weight_dtype: None, moe_dtype: None, activation_dtype: None, + fpm_fmha_dtype: None, kv_cache_dtype: None, }, speculative: None, diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 2992aae5e..211a1254c 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -177,6 +177,8 @@ struct AicTimingConfig { moe_dtype: Option, #[serde(default, alias = "fmha_quant_mode")] fmha_dtype: Option, + #[serde(default, alias = "fpm_fmha_quant_mode")] + fpm_fmha_dtype: Option, #[serde(default, alias = "kvcache_quant_mode")] kv_cache_dtype: Option, #[serde(default, alias = "comm_quant_mode")] @@ -299,6 +301,7 @@ impl AicTimingConfig { gemm_quant_mode: self.gemm_dtype.clone(), moe_quant_mode: self.moe_dtype.clone(), fmha_quant_mode: self.fmha_dtype.clone(), + fpm_fmha_quant_mode: self.fpm_fmha_dtype.clone(), kvcache_quant_mode: self.kv_cache_dtype.clone(), comm_quant_mode: self.comm_dtype.clone(), nextn: self.nextn, @@ -2262,6 +2265,7 @@ mod tests { gemm_dtype: None, moe_dtype: None, fmha_dtype: None, + fpm_fmha_dtype: None, kv_cache_dtype: None, comm_dtype: None, nextn: 0, @@ -2837,6 +2841,7 @@ mod tests { let config = serde_json::from_value::(serde_json::json!({ "model": "test-model", "backend": "sglang", "system": "test-system", "tp": 1, "decoder_replay": replay, "database_mode": "SILICON", + "forward_model": "fpm", "fpm_fmha_dtype": "fp8", "enable_shared_layer": false, "strict_provenance": true })) .unwrap(); @@ -2844,6 +2849,7 @@ mod tests { .estimator_request(ForwardPassWorkerType::Aggregated) .unwrap(); assert_eq!(request.decoder_replay, replay); + assert_eq!(request.fpm_fmha_quant_mode.as_deref(), Some("fp8")); assert_eq!(request.database_mode, crate::DatabaseMode::Silicon); assert_eq!(request.enable_shared_layer, Some(false)); assert!(request.strict_provenance); diff --git a/crates/core/tests/perfmodel/memory_round_trip.rs b/crates/core/tests/perfmodel/memory_round_trip.rs index 0af0f5cf8..d2236e4ba 100644 --- a/crates/core/tests/perfmodel/memory_round_trip.rs +++ b/crates/core/tests/perfmodel/memory_round_trip.rs @@ -101,6 +101,7 @@ fn request(tolerance_fraction: Option) -> KvCacheEstimateRequest { weight_dtype: None, moe_dtype: None, activation_dtype: None, + fpm_fmha_dtype: None, kv_cache_dtype: None, }, speculative: None, diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index c115f5c3d..43aac5db8 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -120,7 +120,7 @@ mod tests { // v17: ContextAttentionOp gained apply_rope (Muse Glimmer review // follow-up) — a positional bincode op-layout change. // v18: speculative attention width fields and FpmForward verify_width. - assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 19); + assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 20); 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 e3d4d86b9..2a7a22fd2 100644 --- a/docs/deepseek-v41-storage.md +++ b/docs/deepseek-v41-storage.md @@ -53,5 +53,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 19 rejects schema 18 bincode before decoding +physical layout. Engine wire schema 20 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 d3d152482..f953bac6e 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 19 compatibility. +for full-context scoring, cache read/write accounting, and schema 20 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 @@ -127,3 +127,7 @@ The [historical SOL validation bundle](https://github.com/ai-dynamo/aisimulate/t preserves observed inputs, source identities, replay tools, and result coverage at that commit. Its results remain specific to the recorded source and runtime identities. + +FPM selector diagnostics append a positional field to `FpmForwardOp`, so this +branch uses EngineSpec schema 20. Schema-19 and earlier bincode inputs are rejected +before decoding; legacy JSON may still use the documented field defaults. diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 3f1bffbe8..803b19137 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -67,6 +67,19 @@ This material is licensed under the Apache License 2.0. The upstream license at the identified revision is available at: https://github.com/ai-dynamo/aiconfigurator/blob/915f590680d8a79fe9c39f6f3a9ff13bc267fcce/LICENSE +## Dynamo V4.1 FPM collection adapter + +`collector/fpm_forward/runtime/dsv41/dsv41_scheduler.py` is modified code +adapted from `components/src/dynamo/vllm/instrumented_scheduler.py` in +https://github.com/ai-dynamo/dynamo/tree/54960177085413259859c88bd34ed0734d4c2ea9. +It adds bounded same-request real-KV collection while preserving the native +benchmark and FPM contracts. Copyright (c) 2025-2026 NVIDIA CORPORATION & +AFFILIATES. All rights reserved. Licensed under Apache-2.0; the upstream +license is preserved in the adapter's adjacent `LICENSE`. The adjacent README +records the inspected vLLM API revision and immutable runtime image/source +hashes. vLLM implementation files are not vendored. The text fixture and +lifecycle tests are original work for this change, with no external corpus. + ## NVIDIA AIConfigurator speculative decoding The speculation SDK, compatibility exports, CLI/task integration, attention and whole-forward FPM operation changes, native bindings, and their tests are adapted and modified from AIConfigurator PR #1563, pinned at commit `6290c161a354da5250c391bd43372b2e9c6f4a51`. Original paths are under `aic-core/src/aiconfigurator_core/sdk/`, `src/aiconfigurator/`, `aic-core/rust/aiconfigurator-core/`, `aic-core/rust/tests/public-api/`, and `tests/`. @@ -535,7 +548,7 @@ SOFTWARE. ## Meta Muse Glimmer model configuration -`src/aiconfigurator_core/model_configs/meta-models--Muse-Glimmer-30B_config.json` +`src/aisimulate_core/model_configs/meta-models--Muse-Glimmer-30B_config.json` is an unmodified copy of `config.json` from the Meta Muse Glimmer model repository at immutable revision `f84ecc3a0ea984a4c04542a84269e3d065350a6e`: @@ -557,8 +570,8 @@ named Qwen model repositories at the immutable revisions shown: | Packaged file | Upstream revision | | --- | --- | -| `aiconfigurator_core/model_configs/Qwen--Qwen3.8-2.4T-A95B_config.json` | `Qwen/Qwen3.8-2.4T-A95B@207bd685a7e3696cfaff12ded7c6a7ea0f88c996` | -| `aiconfigurator_core/model_configs/Qwen--Qwen3.8-2.4T-A95B-FP8_config.json` | `Qwen/Qwen3.8-2.4T-A95B-FP8@d2dc35658bcf77e66643428cb52e774cc3b5bd29` | +| `aisimulate_core/model_configs/Qwen--Qwen3.8-2.4T-A95B_config.json` | `Qwen/Qwen3.8-2.4T-A95B@207bd685a7e3696cfaff12ded7c6a7ea0f88c996` | +| `aisimulate_core/model_configs/Qwen--Qwen3.8-2.4T-A95B-FP8_config.json` | `Qwen/Qwen3.8-2.4T-A95B-FP8@d2dc35658bcf77e66643428cb52e774cc3b5bd29` | Upstream repositories: https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B and @@ -592,10 +605,10 @@ For any questions regarding this license, please contact model-business@notice.q The following bundled model configs are modified copies of Meta Llama 4 checkpoint configuration files: -- `src/aiconfigurator_core/model_configs/meta-llama--Llama-4-Scout-17B-16E-Instruct_config.json` +- `src/aisimulate_core/model_configs/meta-llama--Llama-4-Scout-17B-16E-Instruct_config.json` from revision `92f3b1597a195b523d8d9e5700e57e4fbb8f20d3`: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/92f3b1597a195b523d8d9e5700e57e4fbb8f20d3/config.json -- `src/aiconfigurator_core/model_configs/meta-llama--Llama-4-Maverick-17B-128E-Instruct_config.json` +- `src/aisimulate_core/model_configs/meta-llama--Llama-4-Maverick-17B-128E-Instruct_config.json` from revision `73d14711bcc77c16df3470856949c3764056b617`: https://huggingface.co/meta-llama/Llama-4-Maverick-17B-128E-Instruct/blob/73d14711bcc77c16df3470856949c3764056b617/config.json diff --git a/python/aisimulate/collector/fpm_forward/capabilities.py b/python/aisimulate/collector/fpm_forward/capabilities.py index 70fd2e611..16f58320c 100644 --- a/python/aisimulate/collector/fpm_forward/capabilities.py +++ b/python/aisimulate/collector/fpm_forward/capabilities.py @@ -46,6 +46,8 @@ def _attention_template( is_moe: bool, ) -> tuple[str, str]: lowered = (architecture or "").lower() + if model_family == "DEEPSEEKV41": + return "dsv41_module", "moe_dsv41" if model_family == "DEEPSEEKV4": return "dsv4_module", "moe_dsv4" if model_family == "MINIMAXM3": @@ -254,6 +256,37 @@ def resolve_model_capability( version = database_version or get_latest_database_version(system=system, backend=backend) if not version: raise ValueError(f"no AIC database version is available for system={system!r}, backend={backend!r}") + if model_family == "DEEPSEEKV41": + # A whole-forward campaign is what establishes V4.1 timing evidence. + # Op-level V4/MLA tables cannot authorize or substitute its fused path. + requested_kv = tuple(dict.fromkeys(_normalize_kv_dtype(value) for value in requested_kv_cache_dtypes)) + if set(requested_kv) - {"auto", native_kv}: + raise ValueError(f"DeepSeek-V4.1 FPM requires its native KV dtype {native_kv!r}") + return ModelCapabilityProfile( + architecture=architecture, + model_family=model_family, + is_moe=is_moe, + attention_source="dsv41_module", + attention_kind="moe_dsv41", + support_level="native_runtime", + template_id="deepseek_v41_text_native", + template_version=TEMPLATE_VERSION, + support_reason="registered V4.1 config; native runtime must qualify every measured cell", + allow_pure_tp=backend == "vllm", + aic_database_version=str(version), + model_config=resolved_config, + dtype=ResolvedDTypeProfile( + gemm_quant_mode=gemm, + moe_quant_mode=moe, + fmha_quant_mode=inferred_fmha, + comm_quant_mode=common.CommQuantMode.half.name, + native_kv_cache_dtype=native_kv, + kv_cache_dtypes=(native_kv,), + fmha_resolution="checkpoint_native", + fmha_by_kv_dtype={native_kv: inferred_fmha}, + fmha_resolution_by_kv_dtype={native_kv: "checkpoint_native"}, + ), + ) database = get_database(system, backend, version) if database is None: raise ValueError(f"failed to load AIC database for system={system!r}, backend={backend!r}, version={version!r}") diff --git a/python/aisimulate/collector/fpm_forward/config.py b/python/aisimulate/collector/fpm_forward/config.py index 0fe88d9bb..826ac7b51 100644 --- a/python/aisimulate/collector/fpm_forward/config.py +++ b/python/aisimulate/collector/fpm_forward/config.py @@ -6,7 +6,10 @@ from __future__ import annotations import argparse +import hashlib +import json from dataclasses import dataclass +from pathlib import Path FPM_FORWARD_OP = "fpm_forward" FPM_WARMUP_ITERATIONS = 5 @@ -49,6 +52,46 @@ def _optional_size_list(values: list[int] | None) -> tuple[int, ...] | None: return tuple(sorted(set(values))) +def _freeze_benchmark_points(path: str) -> tuple[str, str]: + """Freeze validated coordinates; native Dynamo owns execution/row admission.""" + + def unique_object(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate benchmark-points field: {key}") + result[key] = value + return result + + payload = json.loads(Path(path).expanduser().read_text(encoding="utf-8"), object_pairs_hook=unique_object) + if not isinstance(payload, dict) or set(payload) != {"schema_version", "prefill", "decode"}: + raise ValueError("benchmark-points manifest requires schema_version, prefill and decode") + if type(payload["schema_version"]) is not int or payload["schema_version"] not in (1, 2, 3): + raise ValueError("benchmark-points schema_version must be 1, 2 or 3") + for phase in ("prefill", "decode"): + if not isinstance(payload[phase], list) or any(not isinstance(point, dict) for point in payload[phase]): + raise ValueError(f"benchmark-points {phase} must be a list of point objects") + for point in payload[phase]: + required = {"batch_size": 1, "total_kv_read_tokens": 0} + if phase == "prefill": + required["total_prefill_tokens"] = 1 + elif "total_prefill_tokens" in point: + required["total_prefill_tokens"] = 0 + for field, minimum in required.items(): + value = point.get(field) + if type(value) is not int or value < minimum: + raise ValueError(f"benchmark-points {phase} {field} must be an integer >= {minimum}") + token_field = "total_prefill_tokens" if phase == "prefill" else "total_kv_read_tokens" + if point[token_field] < point["batch_size"]: + raise ValueError(f"benchmark-points {phase} {token_field} must be >= batch_size") + if phase == "decode" and point.get("total_prefill_tokens", 0) != 0: + raise ValueError("benchmark-points decode total_prefill_tokens must be zero") + if not payload["prefill"] and not payload["decode"]: + raise ValueError("benchmark-points manifest must contain at least one point") + canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False) + return canonical, hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + def _powers_of_two_up_to(limit: int) -> tuple[int, ...]: values = [] value = 1 @@ -201,7 +244,15 @@ class FPMCollectionOptions: vllm_max_model_len: int = VLLM_AUTO_FIT_MAX_MODEL_LEN max_prefill_isl: int = FPM_MAX_PREFILL_ISL max_prefill_batch_size: int | None = None + max_decode_batch_size: int | None = None max_prefill_cudagraph_size: int = FPM_MAX_PREFILL_CUDAGRAPH_SIZE + decoder_replay: bool = False + enforce_eager: bool = False + benchmark_points_json: str | None = None + benchmark_points_sha256: str | None = None + executor: str = "kubernetes" + slurm_container_image: str = "" + slurm_container_mounts: tuple[str, ...] = () @property def prefill_sampling(self) -> PrefillSamplingProfile: @@ -247,7 +298,28 @@ def from_args(cls, args: argparse.Namespace) -> FPMCollectionOptions: if {"pp", "cp"}.intersection(requested_axes): raise ValueError("FPM typical-matrix V1 does not vary PP or CP") + executor = getattr(args, "fpm_executor", None) or "kubernetes" + image = getattr(args, "fpm_slurm_container_image", None) or "" + mounts = tuple(getattr(args, "fpm_slurm_container_mount", None) or ()) + if executor == "slurm" and not image: + raise ValueError("--fpm-executor slurm requires --fpm-slurm-container-image") + if executor != "slurm" and (image or mounts): + raise ValueError("Slurm container options require --fpm-executor slurm") + + model_len = getattr(args, "fpm_max_model_len", None) + if model_len is not None and model_len != -1 and model_len < 1: + raise ValueError("--fpm-max-model-len must be positive or -1 for auto-fit") + + points_path = getattr(args, "fpm_benchmark_points_file", None) + points_json = points_sha256 = None + if points_path is not None: + if getattr(args, "smoke", False): + raise ValueError("--fpm-benchmark-points-file cannot be combined with --smoke") + points_json, points_sha256 = _freeze_benchmark_points(points_path) + return cls( + benchmark_points_json=points_json, + benchmark_points_sha256=points_sha256, max_gpus=max_gpus, gpu_counts=tuple(counts), parallel_presets=requested_presets, @@ -271,15 +343,26 @@ def from_args(cls, args: argparse.Namespace) -> FPMCollectionOptions: if getattr(args, "fpm_warmup_iterations", None) is None else args.fpm_warmup_iterations ), + vllm_max_model_len=model_len if model_len is not None else VLLM_AUTO_FIT_MAX_MODEL_LEN, + max_decode_batch_size=getattr(args, "fpm_max_decode_batch_size", None), max_prefill_isl=getattr(args, "fpm_max_prefill_isl", None) or FPM_MAX_PREFILL_ISL, max_prefill_batch_size=getattr(args, "fpm_max_prefill_batch_size", None), max_prefill_cudagraph_size=( getattr(args, "fpm_max_prefill_cudagraph_size", None) or FPM_MAX_PREFILL_CUDAGRAPH_SIZE ), + decoder_replay=bool(getattr(args, "fpm_decoder_replay", False)), + enforce_eager=bool(getattr(args, "fpm_enforce_eager", False)), + executor=executor, + slurm_container_image=image, + slurm_container_mounts=mounts, ) def to_dict(self) -> dict[str, object]: - return { + payload = { + "decoder_replay": self.decoder_replay, + "executor": self.executor, + "slurm_container_image": self.slurm_container_image, + "slurm_container_mounts": list(self.slurm_container_mounts), "max_gpus": self.max_gpus, "gpu_counts": list(self.gpu_counts), "parallel_presets": list(self.parallel_presets), @@ -298,11 +381,20 @@ def to_dict(self) -> dict[str, object]: "cp_sizes": list(self.cp_sizes) if self.cp_sizes is not None else None, "global_warmup_iterations": self.warmup_iterations, "vllm_max_model_len": self.vllm_max_model_len, + "max_decode_batch_size": self.max_decode_batch_size, "warmup_repeats": 0, "measurement_repeats": FPM_MEASUREMENT_REPEATS, "point_source": "dynamo_native_self_benchmark", "prefill_sampling": self.prefill_sampling.to_dict(), } + if self.enforce_eager: + payload["enforce_eager"] = True + if self.benchmark_points_json is not None: + payload["benchmark_points"] = { + "payload": json.loads(self.benchmark_points_json), + "sha256": self.benchmark_points_sha256, + } + return payload def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: @@ -312,6 +404,33 @@ def add_fpm_arguments(parser: argparse.ArgumentParser) -> None: "FPM forward collection", "Whole-model forward-pass planning, execution, and publication.", ) + group.add_argument( + "--fpm-enforce-eager", + action="store_true", + default=None, + help="Require native eager execution with CUDA graphs disabled; included in the frozen plan.", + ) + group.add_argument( + "--fpm-benchmark-points-file", + metavar="PATH", + default=None, + help="Freeze a native Dynamo point manifest into the plan; cannot be combined with --smoke.", + ) + group.add_argument( + "--fpm-decoder-replay", + action="store_true", + default=None, + help="Use true bounded decoder replay; requires verified runtime support.", + ) + group.add_argument( + "--fpm-max-model-len", type=int, default=None, help="vLLM context limit; positive or -1 for runtime auto-fit." + ) + group.add_argument( + "--fpm-max-decode-batch-size", + type=_positive_int, + default=None, + help="Optional native decode max-num-seqs bound.", + ) group.add_argument( "--fpm-max-gpus", type=_positive_int, @@ -465,6 +584,23 @@ def add_fpm_generator_arguments(parser: argparse.ArgumentParser) -> None: """ group = parser.add_argument_group("FPM deployment inputs") + group.add_argument( + "--fpm-executor", + choices=["kubernetes", "slurm"], + default=None, + help="Execution transport; Slurm runs inside an existing allocation.", + ) + group.add_argument( + "--fpm-slurm-container-image", + default=None, + help="Immutable Pyxis image reference or staged squashfs for Slurm.", + ) + group.add_argument( + "--fpm-slurm-container-mount", + action="append", + default=None, + help="Pyxis SOURCE:TARGET mount; repeat for checkpoint/runtime/cache paths.", + ) group.add_argument( "--generator-config", default=None, @@ -532,6 +668,14 @@ def reject_fpm_arguments_without_fpm(args: argparse.Namespace) -> None: "fpm_max_prefill_batch_size", "fpm_max_prefill_cudagraph_size", "fpm_artifact_root", + "fpm_max_model_len", + "fpm_max_decode_batch_size", + "fpm_decoder_replay", + "fpm_enforce_eager", + "fpm_benchmark_points_file", + "fpm_executor", + "fpm_slurm_container_image", + "fpm_slurm_container_mount", "fpm_database_root", "fpm_publish_partial", # Deployment-only Generator inputs are registered unconditionally on diff --git a/python/aisimulate/collector/fpm_forward/database.py b/python/aisimulate/collector/fpm_forward/database.py index a2d851399..920d9a3cc 100644 --- a/python/aisimulate/collector/fpm_forward/database.py +++ b/python/aisimulate/collector/fpm_forward/database.py @@ -9,6 +9,7 @@ import hashlib import json import logging +import math import os import re import tempfile @@ -19,6 +20,7 @@ import yaml from aisimulate.fpm_contract import FPM_RESOLVED_CONFIG_GLOB +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY from .native_artifact import validate_native_collection from .planner import FPMCell, FPMCollectionPlan, backend_identity_columns @@ -47,6 +49,7 @@ "attention_backend", "enable_wideep", "enable_eplb", + *EXECUTION_COLUMNS, "workload_kind", "batch_size", "total_prefill_tokens", @@ -168,7 +171,7 @@ def aggregate_cell( approved_skip_reasons = {"moe_tp_balanced_by_construction"} def _kv_seed_regime(point: dict[str, Any], phase: str) -> str: - if phase == "prefill": + if phase == "prefill" and (not cell.execution_identity[0] or int(point["total_kv_read_tokens"]) == 0): return "n/a" if kvwarm_meta is None: return "legacy" @@ -193,6 +196,12 @@ def _kv_seed_regime(point: dict[str, Any], phase: str) -> str: batch = int(point["batch_size"]) total_prefill = int(point["total_prefill_tokens"]) total_kv = int(point["total_kv_read_tokens"]) + if ( + cell.execution_identity[0] + and (phase == "decode" or total_kv > 0) + and _kv_seed_regime(point, phase) != "real_kv" + ): + raise ValueError("DeepSeek-V4.1 cached prefill/decode publication requires real_kv measurements") rows.append( { "cell_id": cell.cell_id, @@ -215,6 +224,7 @@ def _kv_seed_regime(point: dict[str, Any], phase: str) -> str: "moe_ep": cell.topology.moe_ep, "cp": cell.topology.cp, **backend_identity_columns(cell.backend_policy), + **dict(zip(EXECUTION_COLUMNS, cell.execution_identity, strict=True)), "workload_kind": phase, "batch_size": batch, "total_prefill_tokens": total_prefill, @@ -234,6 +244,9 @@ def _kv_seed_regime(point: dict[str, Any], phase: str) -> str: "collector_attempt_id": collection.collector_attempt_id, "runtime_run_id": collection.runtime_run_id, "runtime_grid_digest": collection.runtime_grid_digest, + "input_text_sha256": (collection.input_provenance or {}).get("text_sha256"), + "input_token_ids_sha256": (collection.input_provenance or {}).get("token_ids_sha256"), + "input_tokenizer_revision": (collection.input_provenance or {}).get("tokenizer_revision"), } ) return rows @@ -268,8 +281,18 @@ def validate_formal_database_commit( parquet_path: Path, metadata_path: Path, plan: FPMCollectionPlan, + *, + expected_attempt_ids: dict[str, str] | None = None, + reused_cell_ids: tuple[str, ...] = (), + expected_cell_rows: dict[str, str] | None = None, + required_cells: tuple[FPMCell, ...] | None = None, ) -> dict[str, Any]: - """Validate the committed formal database pair referenced by a checkpoint.""" + """Validate the sealed table and requested cells, returning per-cell row hashes. + + A terminal checkpoint supplies its published row hashes and attempt IDs. + First-publisher reuse may keep another plan/attempt only for explicitly + recorded reused cells; their identity and coordinate coverage still apply. + """ if not parquet_path.is_file() or not metadata_path.is_file(): raise ValueError( @@ -281,7 +304,7 @@ def validate_formal_database_commit( raise TypeError(f"FPM database commit record must be a mapping: {metadata_path}") expected = { "schema_name": "aic_fpm_forward_perf", - "schema_version": 6, + "schema_version": 7, "system": plan.system, "backend": plan.backend, } @@ -303,14 +326,97 @@ def validate_formal_database_commit( required = {*_ROW_KEY, *_RUN_IDENTITY_FIELDS} missing = sorted(required - set(parquet.schema_arrow.names)) if missing: - raise ValueError(f"committed FPM database is missing schema-v6 columns: {missing}") + raise ValueError(f"committed FPM database is missing schema-v7 columns: {missing}") row_count = payload.get("row_count") if not isinstance(row_count, int) or row_count < 1 or parquet.metadata.num_rows != row_count: raise ValueError( "FPM database row count does not match its commit record: " f"parquet={parquet.metadata.num_rows}, metadata={row_count!r}" ) - return payload + rows = parquet.read().to_pylist() + identities = _run_identities_by_cell(rows, source="committed") + grouped: dict[str, list[dict[str, Any]]] = {} + for row in rows: + grouped.setdefault(row["cell_id"], []).append(row) + cells = plan.cells if required_cells is None else required_cells + if expected_cell_rows is not None and not isinstance(expected_cell_rows, dict): + raise TypeError("FPM checkpoint cell row hashes must be a mapping") + required_ids = {cell.cell_id for cell in cells} + if set(reused_cell_ids) - required_ids: + raise ValueError("FPM first-publisher reuse names cells outside the frozen plan") + explicit_json = getattr(plan.options, "benchmark_points_json", None) + explicit = json.loads(explicit_json) if explicit_json is not None else None + cell_rows = {} + for cell in cells: + selected = grouped.get(cell.cell_id, []) + if not selected: + raise ValueError(f"committed FPM database does not cover required cell {cell.cell_id}") + expected_identity = { + "model_path": plan.model_path, + "system": plan.system, + "backend": plan.backend, + "backend_version": payload.get("backend_version"), + "weight_quantization": cell.weight_quantization, + "gemm_quant_mode": cell.gemm_quant_mode, + "moe_quant_mode": cell.moe_quant_mode, + "fmha_quant_mode": cell.fmha_quant_mode, + "comm_quant_mode": cell.comm_quant_mode, + "fmha_resolution": cell.fmha_resolution, + "kv_cache_dtype": cell.kv_cache_dtype, + "parallel_strategy": cell.parallel_strategy, + "workload_kind": cell.workload_kind, + "partition_policy": "balanced_v1", + **{name: getattr(cell.topology, name) for name in ("tp", "pp", "dp", "moe_tp", "moe_ep", "cp")}, + **backend_identity_columns(cell.backend_policy), + **dict(zip(EXECUTION_COLUMNS, cell.execution_identity, strict=True)), + } + if cell.input_text_sha256: + expected_identity["input_text_sha256"] = cell.input_text_sha256 + for row in selected: + if any(row.get(key) != value for key, value in expected_identity.items()): + raise ValueError(f"committed FPM cell identity differs from frozen plan: {cell.cell_id}") + latency = row.get("latency_ms") + if ( + isinstance(latency, bool) + or not isinstance(latency, int | float) + or not math.isfinite(latency) + or latency <= 0 + ): + raise ValueError(f"committed FPM cell has an invalid latency: {cell.cell_id}") + if cell.cell_id not in reused_cell_ids: + if identities[cell.cell_id][0] != plan.sha256: + raise ValueError(f"committed FPM cell has a different source plan: {cell.cell_id}") + if expected_attempt_ids is not None and identities[cell.cell_id][1] != expected_attempt_ids.get( + cell.cell_id + ): + raise ValueError(f"committed FPM cell has a different Collector attempt: {cell.cell_id}") + coordinates = { + (row["batch_size"], row["total_prefill_tokens"], row["total_kv_read_tokens"]) for row in selected + } + if len(coordinates) != len(selected): + raise ValueError(f"committed FPM cell contains duplicate coordinates: {cell.cell_id}") + # New checkpoints bind the actually achieved native coordinates, + # including legal context clamping. Legacy explicit checkpoints and + # first-publisher reuse must prove requested coverage from the table. + if explicit is not None and ( + cell.cell_id in reused_cell_ids or (expected_attempt_ids is not None and expected_cell_rows is None) + ): + requested = { + (point["batch_size"], point.get("total_prefill_tokens", 0), point["total_kv_read_tokens"]) + for point in explicit[cell.workload_kind] + } + if not requested <= coordinates: + raise ValueError(f"committed FPM cell does not cover requested coordinates: {cell.cell_id}") + canonical_rows = sorted( + json.dumps(row, sort_keys=True, separators=(",", ":"), allow_nan=False) for row in selected + ) + digest = hashlib.sha256("\n".join(canonical_rows).encode()).hexdigest() + if expected_cell_rows is not None and expected_cell_rows.get(cell.cell_id) != digest: + raise ValueError(f"committed FPM cell differs from its published checkpoint rows: {cell.cell_id}") + if expected_attempt_ids is not None and expected_cell_rows is None and explicit is None: + raise ValueError("terminal FPM checkpoint has no sealed cell coverage; raw artifacts must be revalidated") + cell_rows[cell.cell_id] = digest + return {**payload, "cell_rows": cell_rows} def _curated_systems_root() -> Path: @@ -482,13 +588,21 @@ def write_formal_database( # list turns schema drift into a bare KeyError instead of this # actionable error, and silently rots when _ROW_KEY grows. required = {*_ROW_KEY, *_RUN_IDENTITY_FIELDS} + if committed.get("schema_version") == 6: + required -= set(EXECUTION_COLUMNS) + elif committed.get("schema_version") != 7: + raise ValueError("existing FPM database has unsupported schema version") missing = sorted(required - set(table.column_names)) if missing: raise ValueError( - "existing FPM database does not satisfy the attempt-bound schema-v6 row-key " + "existing FPM database does not satisfy the attempt-bound schema-v7 row-key " f"contract (missing columns: {missing}); publish to a clean destination: {parquet_path}" ) - merged.extend(table.to_pylist()) + existing_rows = table.to_pylist() + if committed.get("schema_version") == 6: + for row in existing_rows: + row.update(zip(EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY, strict=True)) + merged.extend(existing_rows) existing_versions = {str(row.get("backend_version") or "") for row in merged} if existing_versions and existing_versions != {version}: raise ValueError( @@ -538,6 +652,8 @@ def write_formal_database( # as null instead. for row in merged: row.setdefault("kv_seed_regime", None) + for field in ("input_text_sha256", "input_token_ids_sha256", "input_tokenizer_revision"): + row.setdefault(field, None) temporary = _temporary_path(parquet_path) temporary_metadata = _temporary_path(metadata_path) @@ -545,7 +661,7 @@ def write_formal_database( pq.write_table(pa.Table.from_pylist(merged), temporary, compression="zstd") metadata = { "schema_name": "aic_fpm_forward_perf", - "schema_version": 6, + "schema_version": 7, "coordinate_system": "iteration_totals_balanced_v1", "measurement_policy": "dynamo_native_single_sample_v1", "warmup_repeats": 0, diff --git a/python/aisimulate/collector/fpm_forward/native_artifact.py b/python/aisimulate/collector/fpm_forward/native_artifact.py index 18178bebd..1794d2344 100644 --- a/python/aisimulate/collector/fpm_forward/native_artifact.py +++ b/python/aisimulate/collector/fpm_forward/native_artifact.py @@ -5,8 +5,10 @@ from __future__ import annotations +import hashlib import json import math +import re from dataclasses import dataclass from pathlib import Path from typing import Any @@ -15,6 +17,7 @@ FPM_BENCHMARK_RESULT_GLOB, FPM_NATIVE_BENCHMARK_RESULT_SCHEMA_VERSION, ) +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS from .planner import FPMCell from .types import KVWARM_STRATEGIES @@ -80,6 +83,131 @@ class NativeCollection: # Engine-reported KV warm-up envelope (warm_eligible/skip_reason/...); # None only for artifacts predating the kvwarm-enabled runtime. kvwarm_meta: dict[str, Any] | None = None + input_provenance: dict[str, Any] | None = None + + +def _validate_execution_provenance(cell: FPMCell, payload: dict[str, Any], path: Path) -> dict[str, Any] | None: + """Config-bound curves require engine evidence for execution and real input.""" + if not cell.execution_identity[0]: + return None + expected = dict(zip(EXECUTION_COLUMNS, cell.execution_identity, strict=True)) + if payload.get("execution_identity") != expected: + raise ValueError(f"native execution identity differs from the frozen V4.1 cell: {path}") + if payload.get("execution_mode") != "eager": + raise ValueError(f"V4.1 native data requires verified eager execution: {path}") + evidence = payload.get("input_provenance") + if not isinstance(evidence, dict) or evidence.get("source") != "tokenizer_text": + raise ValueError(f"V4.1 native result requires tokenizer-generated text provenance: {path}") + for field in ("text_sha256", "token_ids_sha256"): + value = evidence.get(field) + if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value): + raise ValueError(f"V4.1 native result has invalid {field}: {path}") + if cell.input_text_sha256 and evidence["text_sha256"] != cell.input_text_sha256: + raise ValueError(f"V4.1 native input text differs from the frozen corpus: {path}") + if not isinstance(evidence.get("tokenizer_revision"), str) or not evidence["tokenizer_revision"]: + raise ValueError(f"V4.1 native result has no tokenizer revision: {path}") + counts = [evidence.get(field) for field in ("token_count", "unique_token_count")] + if any(not isinstance(value, int) or isinstance(value, bool) or value < 2 for value in counts): + raise ValueError(f"V4.1 native input corpus must contain multiple tokenizer-generated tokens: {path}") + if counts[1] > counts[0]: + raise ValueError(f"V4.1 native input corpus token counts are inconsistent: {path}") + return evidence + + +def _validate_token_streams(payload: dict[str, Any], path: Path) -> None: + """Verify archived real request histories and completed-forward witnesses.""" + manifest = payload["input_provenance"].get("token_stream_manifest") + if not isinstance(manifest, dict): + raise ValueError(f"V4.1 native result lacks token-stream manifest: {path}") + name = manifest.get("file") + if not isinstance(name, str) or Path(name).name != name or not name.endswith(".token-streams.jsonl"): + raise ValueError(f"V4.1 token-stream path must be an adjacent JSONL file: {path}") + raw = path.with_name(name).read_bytes() + if hashlib.sha256(raw).hexdigest() != manifest.get("sha256"): + raise ValueError(f"V4.1 token-stream manifest SHA mismatch: {path}") + lines = raw.splitlines() + schema = manifest.get("schema_version", 1) + if type(schema) is not int or schema not in (1, 2): + raise ValueError(f"V4.1 token-stream schema is unsupported: {path}") + warmups = payload.get("warmup_results", []) + if not isinstance(warmups, list) or (schema == 1 and warmups): + raise ValueError(f"V4.1 warmup histories require token-stream schema 2: {path}") + measured = payload["results"] + if any( + not isinstance(row, dict) + or not isinstance(row.get("point"), dict) + or type(row["point"].get("benchmark_id")) is not int + or row["point"].get("point_type") not in ("prefill", "decode") + or not isinstance(row["point"].get("sample_reasons", []), list) + for row in measured + warmups + ): + raise ValueError(f"V4.1 native result has malformed token-stream points: {path}") + warmup_ids = set() + for row in warmups: + point = row.get("point", {}) + benchmark_id = point.get("benchmark_id") + if ( + type(benchmark_id) is not int + or benchmark_id <= len(measured) + or benchmark_id in warmup_ids + or "eager_warmup" not in point.get("sample_reasons", []) + ): + raise ValueError(f"V4.1 warmup point is not a distinct native eager replica: {path}") + warmup_ids.add(benchmark_id) + if schema == 2 and manifest.get("warmup_benchmark_ids") != sorted(warmup_ids): + raise ValueError(f"V4.1 native eager warmup coverage mismatch: {path}") + if any("eager_warmup" in row["point"].get("sample_reasons", []) for row in measured): + raise ValueError(f"V4.1 warmup timing cannot enter measured results: {path}") + all_results = measured + warmups + if manifest.get("records") != len(lines) or len(lines) != len(all_results): + raise ValueError(f"V4.1 token-stream coverage mismatch: {path}") + streams = {} + for line in lines: + stream = json.loads(line) + if not isinstance(stream, dict): + raise ValueError(f"V4.1 token-stream record must be an object: {path}") + benchmark_id = stream.get("benchmark_id") + if type(benchmark_id) is not int or benchmark_id in streams: + raise ValueError(f"V4.1 token-stream benchmark ID is invalid or duplicated: {path}") + expected_role = "warmup" if benchmark_id in warmup_ids else "measurement" + if schema == 2 and stream.get("sampling_role") != expected_role: + raise ValueError(f"V4.1 token-stream sampling role mismatch: {path}") + streams[benchmark_id] = (stream, hashlib.sha256(line).hexdigest()) + if set(streams) != {row["point"]["benchmark_id"] for row in all_results}: + raise ValueError(f"V4.1 token-stream point identity coverage mismatch: {path}") + for row in all_results: + point = row["point"] + witness = row.get("real_kv_witness") + stream, digest = streams.get(point["benchmark_id"], ({}, None)) + batch = _require_int(point, "batch_size") + decode = point["point_type"] == "decode" + expected_seed = _require_int(point, "total_kv_read_tokens") - (batch if decode else 0) + if ( + not isinstance(witness, dict) + or witness.get("same_request") is not True + or witness.get("allocated_fake_tokens") != 0 + or witness.get("completed_seed_tokens") != expected_seed + or witness.get("token_stream_sha256") != digest + ): + raise ValueError(f"V4.1 native result has invalid completed real-KV witness: {path}") + requests = stream.get("requests") + if not isinstance(requests, list) or len(requests) != batch: + raise ValueError(f"V4.1 token-stream request count mismatch: {path}") + prompt_total = 0 + for index, request in enumerate(requests): + if not isinstance(request, dict) or request.get("request_index") != index: + raise ValueError(f"V4.1 token-stream request order mismatch: {path}") + for field in ("prompt_token_ids", "output_token_ids"): + ids = request.get(field) + if not isinstance(ids, list) or not ids or any(type(token) is not int or token < 0 for token in ids): + raise ValueError(f"V4.1 token-stream has invalid real token IDs: {path}") + prompt_length = len(request["prompt_token_ids"]) + prompt_total += prompt_length + if request.get("computed_tokens") != prompt_length + (2 if decode else 0): + raise ValueError(f"V4.1 token-stream computed-token witness mismatch: {path}") + expected_prompt = expected_seed + (0 if decode else _require_int(point, "total_prefill_tokens")) + if prompt_total != expected_prompt: + raise ValueError(f"V4.1 token-stream prompt history differs from the measured point: {path}") def _validate_collector_provenance( @@ -251,6 +379,7 @@ def validate_native_collection( run_identity: tuple[str, str] | None = None kvwarm_meta: dict[str, Any] | None = None kvwarm_seen: object = _KVWARM_UNSEEN + input_provenance: dict[str, Any] | None = None local_fpms: dict[tuple[int, int], dict[str, Any]] = {} rank_timings: list[tuple[int, float, float]] = [] @@ -279,7 +408,20 @@ def validate_native_collection( rows = payload.get("results") groups = payload.get("iteration_groups") if not isinstance(coverage, dict) or not isinstance(rows, list) or not isinstance(groups, list): - raise TypeError(f"native result is missing coverage/results/iteration_groups: {path}") + raise ValueError(f"native result is missing coverage/results/iteration_groups: {path}") + if any(not isinstance(row, dict) or not isinstance(row.get("point"), dict) for row in rows): + raise ValueError(f"native result has malformed result points: {path}") + evidence = _validate_execution_provenance(cell, payload, path) + if evidence is not None: + _validate_token_streams(payload, path) + if input_provenance is None: + input_provenance = evidence + elif {k: v for k, v in evidence.items() if k != "token_stream_manifest"} != { + k: v for k, v in input_provenance.items() if k != "token_stream_manifest" + } or {k: v for k, v in evidence["token_stream_manifest"].items() if k != "file"} != { + k: v for k, v in input_provenance["token_stream_manifest"].items() if k != "file" + }: + raise ValueError(f"native DP ranks disagree on input provenance: {path}") expected = coverage.get("expected_points") if ( not isinstance(expected, int) @@ -439,4 +581,5 @@ def validate_native_collection( runtime_run_id=run_identity[0], runtime_grid_digest=run_identity[1], kvwarm_meta=kvwarm_meta, + input_provenance=input_provenance, ) diff --git a/python/aisimulate/collector/fpm_forward/planner.py b/python/aisimulate/collector/fpm_forward/planner.py index 33e62b502..27c1680a5 100644 --- a/python/aisimulate/collector/fpm_forward/planner.py +++ b/python/aisimulate/collector/fpm_forward/planner.py @@ -19,6 +19,8 @@ from pathlib import Path, PurePosixPath from typing import Any +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY, execution_identity + from .capabilities import ModelCapabilityProfile, ResolvedDTypeProfile, resolve_model_capability from .config import FPMCollectionOptions from .memory_admission import TopologyMemoryDecision, filter_memory_infeasible_topologies @@ -397,6 +399,8 @@ def _backend_policies( extra_cli_args: list[str] = [] expected_markers: dict[str, str] = {} + if options.enforce_eager: + expected_markers["config.engine_args.enforce_eager"] = "True" if moe != "auto": extra_cli_args += ["--kernel-config", json.dumps({"moe_backend": moe})] expected_markers["config.engine_args.kernel_config.moe_backend"] = moe @@ -444,10 +448,14 @@ class FPMCell: fmha_quant_mode: str | None = None comm_quant_mode: str | None = None fmha_resolution: str | None = None + execution_identity: tuple[str, ...] = LEGACY_EXECUTION_IDENTITY + input_text_sha256: str = "" def to_dict(self) -> dict[str, object]: return { "cell_id": self.cell_id, + "execution_identity": dict(zip(EXECUTION_COLUMNS, self.execution_identity, strict=True)), + "input_text_sha256": self.input_text_sha256, "workload_kind": self.workload_kind, "point_source": "dynamo_native_self_benchmark", "topology": self.topology.to_dict(), @@ -483,9 +491,12 @@ class FPMCollectionPlan: sha256: str def to_dict(self) -> dict[str, object]: + explicit_points = ( + json.loads(self.options.benchmark_points_json) if self.options.benchmark_points_json is not None else None + ) return { "schema_name": "aic_fpm_collection_plan", - "schema_version": 10, + "schema_version": 11, "backend": self.backend, "model_path": self.model_path, "system": self.system, @@ -497,6 +508,8 @@ def to_dict(self) -> dict[str, object]: "point_generation": { "owner": "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler", "method": "native_self_benchmark", + "source": "frozen_explicit_manifest" if explicit_points is not None else "native_auto_grid", + "manifest_sha256": self.options.benchmark_points_sha256, "coordinates": [ "batch_size", "total_prefill_tokens", @@ -506,7 +519,11 @@ def to_dict(self) -> dict[str, object]: "point_admission": "dynamo_live_scheduler", "precondition": "vllm_engine_initialized", "prefill_sampling": self.options.prefill_sampling.to_dict(), - "planned_point_count": None, + "planned_point_count": ( + sum(len(explicit_points.get(phase, [])) for phase in ("prefill", "decode")) + if explicit_points is not None + else None + ), }, "topologies": [ { @@ -547,6 +564,8 @@ def _cell_id( weight_quantization: str, kv_cache_dtype: str, policy: BackendPolicy, + execution: tuple[str, ...] = LEGACY_EXECUTION_IDENTITY, + input_text_sha256: str = "", ) -> str: payload = { "backend": backend, @@ -557,6 +576,8 @@ def _cell_id( "weight_quantization": weight_quantization, "kv_cache_dtype": kv_cache_dtype, **backend_identity_columns(policy), + **dict(zip(EXECUTION_COLUMNS, execution, strict=True)), + "input_text_sha256": input_text_sha256, "point_source": "dynamo_native_self_benchmark", } return f"fpm-{_canonical_hash(payload)[:16]}" @@ -593,6 +614,24 @@ def build_collection_plan( str(collector_config["aic_database_version"]) if "aic_database_version" in collector_config else None ), ) + execution = execution_identity( + capability.model_config.payload, + decoder_replay=options.decoder_replay, + backend=backend, + # The rendered V4.1 collection contract requests text-only HBM Engram. + # The producer must independently attest these actual runtime facts. + engram_cpu_offload=False, + input_modality="text", + ) + if execution[0] and not options.enforce_eager: + raise ValueError("V4.1 FPM collection currently requires --fpm-enforce-eager; graph timing is not qualified") + if options.enforce_eager and not execution[0]: + raise ValueError("explicit eager FPM collection is currently qualified only for DeepSeek V4.1") + input_text_sha256 = ( + hashlib.sha256((Path(__file__).parent / "runtime" / "fpm_text.txt").read_bytes()).hexdigest() + if execution[0] + else "" + ) candidate_topologies = enumerate_fpm_topologies( backend=backend, is_moe=capability.is_moe, @@ -646,7 +685,11 @@ def build_collection_plan( weight_quantization=weight_quantization, kv_cache_dtype=kv_cache_dtype, policy=policy, + execution=execution, + input_text_sha256=input_text_sha256, ), + execution_identity=execution, + input_text_sha256=input_text_sha256, workload_kind=phase, topology=topology, weight_quantization=weight_quantization, diff --git a/python/aisimulate/collector/fpm_forward/runner.py b/python/aisimulate/collector/fpm_forward/runner.py index c5ad6acad..971a5c273 100644 --- a/python/aisimulate/collector/fpm_forward/runner.py +++ b/python/aisimulate/collector/fpm_forward/runner.py @@ -17,6 +17,7 @@ import shutil import signal import subprocess +import sys import threading import time import uuid @@ -51,6 +52,9 @@ logger = logging.getLogger(__name__) CHECKPOINT_SCHEMA = "aic-fpm-collector-checkpoint-v3" +RUNTIME_ENV_FILENAME = "collector-runtime-env.sh" +READINESS_TIMEOUT_ENV = "FPM_READINESS_TIMEOUT_SECONDS" +DEFAULT_READINESS_TIMEOUT_SECONDS = 900 # KV warm-up dominates a decode cell's wall clock (~80 min for a tep4 decode # sweep on MiniMax M2.7); one hour would kill the engine mid-warm-up. 10800 # matches the r15 parity protocol's budget. @@ -176,23 +180,180 @@ def _kubectl_command() -> list[str]: raise RuntimeError("neither kubectl nor tsh is available") -# Live kubectl children, registered so an interrupt can terminate them: a +# Live transport children, registered so an interrupt can terminate them: a # worker thread blocked in a subprocess wait cannot be interrupted by the main # thread's signal, so without this a detached campaign stopped with SIGINT or # SIGTERM would sit in ThreadPoolExecutor joins for up to the full exec # timeout before salvage/teardown could start. _ACTIVE_COMMANDS: set[subprocess.Popen[str]] = set() _ACTIVE_COMMANDS_LOCK = threading.Lock() +_COMMAND_TERMINATION_GRACE_SECONDS = 10.0 + + +_COMMAND_SCOPE = threading.local() + + +class _CommandCancelled(RuntimeError): + """An execution closed admission before this transport could run.""" + + +class CommandScope: + """One execute invocation's launch admission and owned transport groups. + + Workers reserve before Popen. Cancellation closes admission under the same + lock used by registration, so a child returned by an in-flight Popen is + stopped by its launching worker even if it missed the cancellation snapshot. + Pool shutdown then joins those in-flight launch owners. Salvage and later + execute invocations have independent scopes. + """ + + def __init__(self) -> None: + self.lock = threading.Lock() + self.cancelled = False + self.inflight = 0 + self.processes: set[subprocess.Popen[str]] = set() + + def run(self, function, *args): + previous = getattr(_COMMAND_SCOPE, "current", None) + _COMMAND_SCOPE.current = self + try: + with self.lock: + if self.cancelled: + raise _CommandCancelled("FPM execution was cancelled before worker launch") + return function(*args) + finally: + _COMMAND_SCOPE.current = previous + + def reserve(self) -> None: + with self.lock: + if self.cancelled: + raise _CommandCancelled("FPM execution was cancelled before command launch") + self.inflight += 1 + + def register(self, process: subprocess.Popen[str]) -> bool: + with self.lock: + self.processes.add(process) + return self.cancelled + + def finish(self, process: subprocess.Popen[str] | None) -> None: + with self.lock: + self.inflight -= 1 + if process is not None: + self.processes.discard(process) + + def cancel(self) -> None: + with self.lock: + self.cancelled = True + processes = list(self.processes) + _stop_commands(processes) + + +def _cancel_preserving_interrupt(scope: CommandScope, error: BaseException) -> None: + try: + scope.cancel() + except Exception as cleanup_error: + error.add_note(f"FPM transport cleanup failed: {cleanup_error!r}") + logger.error("FPM transport cleanup failed during interruption", exc_info=True) + + +def _darwin_group_gone_after_reap(process: subprocess.Popen[str]) -> bool: + # Darwin may report EPERM for a group containing only an unreaped zombie. + # Reap our direct child, then require ESRCH: its exit alone says nothing + # about surviving descendants, and EPERM can also mean genuine denial. + if sys.platform != "darwin" or process.poll() is None: + return False + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + return True + except PermissionError: + return False + return False + + +def _signal_command(process: subprocess.Popen[str], *, force: bool) -> None: + with suppress(ProcessLookupError): + if os.name == "posix": + # _run_command starts each transport in its own session. Include + # pipe-sharing descendants even when the direct child has exited. + try: + os.killpg(process.pid, signal.SIGKILL if force else signal.SIGTERM) + except PermissionError: + if not _darwin_group_gone_after_reap(process): + raise + elif force: + process.kill() + else: + process.terminate() + + +def _command_group_running(process: subprocess.Popen[str]) -> bool: + direct_running = process.poll() is None + if os.name != "posix": + return direct_running + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + return False + except PermissionError: + if _darwin_group_gone_after_reap(process): + return False + raise + return True + + +def _stop_commands(processes: list[subprocess.Popen[str]]) -> None: + errors: list[Exception] = [] + + def send(process, *, force): + try: + _signal_command(process, force=force) + except OSError as error: + # On Darwin the owner may still be reaping a zombie while holding + # Popen's wait lock. Defer TERM EPERM to the group grace probes; + # persistent denial still fails, and KILL failures are never hidden. + if force or sys.platform != "darwin" or not isinstance(error, PermissionError): + errors.append(error) + + for process in processes: + send(process, force=False) + deadline = time.monotonic() + _COMMAND_TERMINATION_GRACE_SECONDS + pending = list(processes) + while pending: + running = [] + for process in pending: + try: + if _command_group_running(process): + running.append(process) + except OSError as error: + # EPERM is not evidence of disappearance. Give Darwin's + # communicate owner the grace period to reap; persistent denial + # is reported before escalation, without skipping other groups. + if sys.platform != "darwin" or not isinstance(error, PermissionError) or time.monotonic() >= deadline: + errors.append(error) + running.append(process) + pending = running + remaining = deadline - time.monotonic() + if not pending or remaining <= 0 or errors: + break + time.sleep(min(0.05, remaining)) + for process in pending: + send(process, force=True) + for process in processes: + try: + process.wait(timeout=1) + except (OSError, subprocess.TimeoutExpired) as error: + errors.append(error) + if errors: + raise ExceptionGroup("FPM transport groups could not be fully cleaned up", errors) def terminate_active_commands() -> int: - """TERM every live kubectl child; returns how many were signalled.""" + """Stop registered transport groups with bounded TERM-to-KILL escalation.""" with _ACTIVE_COMMANDS_LOCK: processes = list(_ACTIVE_COMMANDS) - for process in processes: - with suppress(OSError): - process.terminate() + _stop_commands(processes) return len(processes) @@ -245,42 +406,41 @@ def _run_command( check: bool = True, timeout: int | None = None, ) -> subprocess.CompletedProcess[str]: - def _stop(process: subprocess.Popen[str]) -> None: - # Never wait on pipes here and never rely on SIGKILL alone: `tsh - # kubectl` re-execs itself as a pipe-sharing grandchild that SIGKILL - # on the wrapper cannot reach (an unbounded drain would then hang on - # the orphan's open pipe forever), while terminate() IS forwarded. - # Signal politely, give the wrapper a moment, then kill and reap the - # direct child only. - process.terminate() - with suppress(subprocess.TimeoutExpired): - process.wait(timeout=10) - process.kill() - process.wait() - - with subprocess.Popen( - args, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - env=_command_env(), - ) as process: + scope = getattr(_COMMAND_SCOPE, "current", None) + if scope is not None: + scope.reserve() + process = None + try: + process = subprocess.Popen( + args, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + env=_command_env(), + start_new_session=os.name == "posix", + ) with _ACTIVE_COMMANDS_LOCK: _ACTIVE_COMMANDS.add(process) - try: - stdout, stderr = process.communicate(timeout=timeout) - except subprocess.TimeoutExpired as error: - _stop(process) - raise subprocess.TimeoutExpired(args, timeout or 0, output=error.output, stderr=error.stderr) from None - except BaseException: - # Mirror subprocess.run: never abandon a live child (a hung - # kubectl would otherwise block Popen.__exit__'s untimed wait - # forever, unreachable after the registry discard below). - _stop(process) - raise - finally: + if scope is not None and scope.register(process): + raise _CommandCancelled("FPM execution was cancelled during command launch") + stdout, stderr = process.communicate(timeout=timeout) + except BaseException as error: + if process is not None: + try: + _stop_commands([process]) + except Exception as cleanup_error: + error.add_note(f"FPM transport cleanup failed: {cleanup_error!r}") + logger.error("FPM transport cleanup failed", exc_info=True) + raise + finally: + if process is not None: with _ACTIVE_COMMANDS_LOCK: _ACTIVE_COMMANDS.discard(process) + for stream in (process.stdout, process.stderr): + if stream is not None: + stream.close() + if scope is not None: + scope.finish(process) if check and process.returncode != 0: raise subprocess.CalledProcessError(process.returncode, args, output=stdout, stderr=stderr) return subprocess.CompletedProcess(args, process.returncode, stdout, stderr) @@ -712,9 +872,10 @@ def execute(self, pods: list[str], timeout_seconds: int = 14400) -> None: logs_dir = self.cell_dir / "logs" logs_dir.mkdir(parents=True, exist_ok=True) failures = [] + scope = CommandScope() pool = ThreadPoolExecutor(max_workers=len(pods)) try: - futures = {pool.submit(self._run_pod, pod, timeout_seconds): pod for pod in pods} + futures = {pool.submit(scope.run, self._run_pod, pod, timeout_seconds): pod for pod in pods} for future in as_completed(futures): pod = futures[future] try: @@ -725,15 +886,15 @@ def execute(self, pods: list[str], timeout_seconds: int = 14400) -> None: continue (logs_dir / f"{pod}.run.stdout.log").write_text(completed.stdout) (logs_dir / f"{pod}.run.stderr.log").write_text(completed.stderr) - except BaseException: + except BaseException as error: # An interrupt lands here while worker threads sit in kubectl # waits they cannot be signalled out of; kill the children first # so the pool join below returns promptly and salvage/teardown # can start. - terminate_active_commands() + _cancel_preserving_interrupt(scope, error) raise finally: - pool.shutdown(wait=True) + pool.shutdown(wait=True, cancel_futures=True) if failures: raise RuntimeError(f"staged FPM fpm_exec.sh failed: {failures}") @@ -890,6 +1051,85 @@ def cleanup(self) -> None: time.sleep(CLEANUP_PROBE_INTERVAL_SECONDS) +POINTS_FILENAME = "benchmark-points.json" +POINTS_RECEIPT_FILENAME = "point_manifest_receipt.json" + + +def _frozen_points(plan: FPMCollectionPlan) -> str | None: + canonical = getattr(plan.options, "benchmark_points_json", None) + if canonical is None: + return None + expected = getattr(plan.options, "benchmark_points_sha256", None) + if hashlib.sha256(canonical.encode("utf-8")).hexdigest() != expected: + raise ValueError("frozen benchmark-points payload and SHA256 disagree") + return canonical + + +def _stage_points_file(plan: FPMCollectionPlan, cell_dir: Path) -> list[Path]: + canonical = _frozen_points(plan) + if canonical is None: + return [] + path = cell_dir / POINTS_FILENAME + path.write_bytes(canonical.encode("utf-8")) + return [path] + + +def _record_points_receipts(resource, pods, plan, cell, attempt_id: str, *, phase: str) -> None: + if _frozen_points(plan) is None: + return + payload = json.dumps( + { + "plan_sha256": plan.sha256, + "cell_id": cell.cell_id, + "attempt_id": attempt_id, + "sha256": plan.options.benchmark_points_sha256, + "phase": phase, + }, + sort_keys=True, + ) + # The same check runs inside Kubernetes and Slurm containers, before and + # after native execution. It does not modify Generator's emitted scripts. + script = ( + "import hashlib,json,pathlib,sys; receipt=json.loads(sys.argv[1]); " + "actual=hashlib.sha256(pathlib.Path(sys.argv[2]).read_bytes()).hexdigest()\n" + "if actual != receipt['sha256']: raise ValueError('runtime benchmark-points SHA256 mismatch')\n" + "pathlib.Path(sys.argv[3]).write_text(json.dumps(receipt,sort_keys=True))" + ) + execute = getattr(resource, "_exec_checked", None) or resource._exec + for pod in pods: + execute( + pod, + [ + "python3", + "-c", + script, + payload, + f"{REMOTE_WORKDIR}/{POINTS_FILENAME}", + f"{FPM_RESULTS_DIR}/{POINTS_RECEIPT_FILENAME}", + ], + timeout=300, + ) + + +def _validate_points_receipts(plan, cell, raw_root: Path, attempt_id: str) -> None: + if _frozen_points(plan) is None: + return + owners = sorted(raw_root.rglob(COLLECTOR_PROVENANCE_FILENAME)) + if not owners: + raise ValueError("explicit benchmark points have no runtime receipt owners") + expected = { + "plan_sha256": plan.sha256, + "cell_id": cell.cell_id, + "attempt_id": attempt_id, + "sha256": plan.options.benchmark_points_sha256, + "phase": "after", + } + for owner in owners: + path = owner.parent / POINTS_RECEIPT_FILENAME + if json.loads(path.read_text()) != expected: + raise ValueError(f"runtime benchmark-points receipt mismatch: {path}") + + def _cell_generator_overrides( plan: FPMCollectionPlan, cell: FPMCell, @@ -897,6 +1137,10 @@ def _cell_generator_overrides( *, smoke: bool = False, ) -> dict[str, Any]: + explicit_points = _frozen_points(plan) + enforce_eager = bool(getattr(plan.options, "enforce_eager", False)) + if explicit_points is not None and smoke: + raise ValueError("--fpm-benchmark-points-file cannot be combined with --smoke") unsupported_base = set(base) - {"K8sConfig", "generator_dynamo_version"} if unsupported_base: raise ValueError(f"FPM runner accepts deployment-only Generator inputs, got {sorted(unsupported_base)}") @@ -926,6 +1170,15 @@ def _cell_generator_overrides( "--max-model-len", str(plan.options.vllm_max_model_len), ] + if enforce_eager: + scheduler_args.append("--enforce-eager") + if explicit_points is not None: + payload = json.loads(explicit_points) + if not payload[cell.workload_kind]: + raise ValueError(f"benchmark-points manifest has no {cell.workload_kind} points for this cell") + scheduler_args.extend(["--benchmark-points-file", f"{REMOTE_WORKDIR}/{POINTS_FILENAME}"]) + if cell.workload_kind == "decode" and getattr(plan.options, "max_decode_batch_size", None): + scheduler_args.extend(["--max-num-seqs", str(plan.options.max_decode_batch_size)]) if cell.workload_kind == "prefill" and not smoke: profile = plan.options.prefill_sampling compilation_config = { @@ -936,14 +1189,16 @@ def _cell_generator_overrides( [ "--max-num-batched-tokens", str(profile.max_total_prefill_tokens), - "--compilation-config", - json.dumps(compilation_config, sort_keys=True, separators=(",", ":")), "--prefill-max-new-token-samples", str(profile.max_new_token_samples), "--prefill-max-kv-read-token-samples", str(profile.max_kv_read_token_samples), ] ) + if not enforce_eager: + scheduler_args.extend( + ["--compilation-config", json.dumps(compilation_config, sort_keys=True, separators=(",", ":"))] + ) if profile.max_batch_size is not None: scheduler_args.extend(["--max-num-seqs", str(profile.max_batch_size)]) elif smoke: @@ -969,6 +1224,16 @@ def _cell_generator_overrides( ) model_args = [] architecture = getattr(getattr(plan, "capability", None), "architecture", None) + if architecture == "DeepseekV41ForCausalLM": + if plan.options.decoder_replay: + raise NotImplementedError("vLLM DeepSeek-V4.1 true decoder replay is not verified") + model_args.extend( + ["--language-model-only", "--tokenizer-mode=deepseek_v41", '--engram-config={"cpu_offload":false}'] + ) + if cell.workload_kind == "decode" or smoke: + model_args.extend(["--max-num-batched-tokens", str(plan.options.max_prefill_isl)]) + if cell.workload_kind == "prefill" and smoke and plan.options.max_prefill_batch_size: + model_args.extend(["--max-num-seqs", str(plan.options.max_prefill_batch_size)]) if architecture == "GlmMoeDsaForCausalLM": # This is the serving path validated by the pinned GLM-5.2 vLLM image. # The parser does not alter FPM scheduling, but keeping the model's @@ -978,6 +1243,16 @@ def _cell_generator_overrides( {"name": FPM_ENGINE_BENCHMARK_OUTPUT_ENV, "value": f"{FPM_RESULTS_DIR}/benchmark.json"}, {"name": FPM_RUN_ID_ENV, "value": cell.cell_id}, ] + if architecture == "DeepseekV41ForCausalLM": + from aisimulate_core.sdk.deepseek_v41 import MODEL_REVISION + + env.extend( + [ + {"name": "DYN_FPM_DSV41_REAL_KV", "value": "1"}, + {"name": "DYN_FPM_INPUT_TEXT", "value": "/tmp/fpm-bench/fpm_text.txt"}, + {"name": "DYN_FPM_TOKENIZER_REVISION", "value": MODEL_REVISION}, + ] + ) total_gpus = cell.topology.total_gpus generated = { "ServiceConfig": service, @@ -1030,9 +1305,34 @@ def _cell_generator_overrides( if existing is not None and existing != item: raise ValueError(f"conflicting FPM environment value for {name}") resolved_env[name] = copy.deepcopy(item) + if architecture == "DeepseekV41ForCausalLM": + # Image layout belongs to this source-pinned adapter. An explicitly + # configured path uses the existing deployment environment interface; + # both preflight and generated run.sh receive the same resolved value. + configured = resolved_env.get("PYTHONPATH") + if configured is None: + adapter = Path(__file__).parent / "runtime" / "dsv41" / "runtime-paths.json" + python_path = json.loads(adapter.read_text())["python_path"] + else: + python_path = configured.get("value") + if not isinstance(python_path, str) or not python_path: + raise ValueError("V4.1 PYTHONPATH must name explicit absolute runtime paths") + paths = python_path.split(":") + if any(not part or not PurePosixPath(part).is_absolute() or "\n" in part or "\r" in part for part in paths): + raise ValueError("V4.1 PYTHONPATH must name explicit absolute runtime paths") + python_path = ":".join([REMOTE_WORKDIR, *(part for part in paths if part != REMOTE_WORKDIR)]) + resolved_env["PYTHONPATH"] = {"name": "PYTHONPATH", "value": python_path} + readiness = resolved_env.get(READINESS_TIMEOUT_ENV, {}).get("value", DEFAULT_READINESS_TIMEOUT_SECONDS) + if isinstance(readiness, bool) or not re.fullmatch(r"[1-9][0-9]*", str(readiness)) or int(readiness) > 3600: + raise ValueError(f"{READINESS_TIMEOUT_ENV} must be an integer from 1 through 3600") + resolved_env[READINESS_TIMEOUT_ENV] = {"name": READINESS_TIMEOUT_ENV, "value": str(readiness)} merged.setdefault("K8sConfig", {})["extra_env"] = list(resolved_env.values()) policy_args = ((policy.get("params") or {}).get("agg") or {}).get("extra_cli_args") or [] + if any(str(arg).split("=", 1)[0] in {"--enforce-eager", "--no-enforce-eager"} for arg in policy_args): + raise ValueError("eager execution must be supplied through --fpm-enforce-eager") + if any(str(arg).split("=", 1)[0] == "--benchmark-points-file" for arg in policy_args): + raise ValueError("benchmark points must be supplied through --fpm-benchmark-points-file") if cell.workload_kind == "decode": prefix_caching = _decode_prefix_caching_mode(cell) policy_disables = any( @@ -1080,6 +1380,14 @@ def _configured_sampling_metadata( *, smoke: bool, ) -> dict[str, int | str]: + canonical = _frozen_points(plan) + if canonical is not None: + if smoke: + raise ValueError("--fpm-benchmark-points-file cannot be combined with --smoke") + return { + "benchmark_points_sha256": plan.options.benchmark_points_sha256, + "requested_point_count": len(json.loads(canonical)[cell.workload_kind]), + } if cell.workload_kind != "prefill": # Derive checkpoint evidence from the same strategy predicate used to # render the engine flags. _cell_generator_overrides rejects policy @@ -1095,6 +1403,27 @@ def _configured_sampling_metadata( } +def _write_runtime_environment(cell_dir: Path, overrides: dict[str, Any]) -> None: + # Generator owns run.sh and fpm_env.sh. Mirror only startup inputs through + # this Collector-owned file because Slurm does not start a Kubernetes Pod + # with extra_env, and run.sh's exports happen after the preflight process. + names = { + READINESS_TIMEOUT_ENV, + "PYTHONPATH", + "DYN_FPM_DSV41_REAL_KV", + "DYN_FPM_INPUT_TEXT", + "DYN_FPM_TOKENIZER_REVISION", + } + lines = ["# Generated Collector startup environment; engine settings remain in run.sh."] + for item in overrides["K8sConfig"]["extra_env"]: + if item["name"] in names: + value = item.get("value") + if not isinstance(value, str): + raise ValueError(f"Collector startup environment {item['name']} requires a literal string") + lines.append(f"export {item['name']}={shlex.quote(value)}") + (cell_dir / RUNTIME_ENV_FILENAME).write_text("\n".join(lines) + "\n") + + def _render_cell( plan: FPMCollectionPlan, cell: FPMCell, @@ -1149,6 +1478,7 @@ def _render_cell( raise ValueError(f"invalid GeneratorRequest for {cell.cell_id}: {errors}") artifacts = generate_from_request(request, output_dir=str(cell_dir)) _atomic_json(cell_dir / "generator-request.json", params) + _write_runtime_environment(cell_dir, overrides) return artifacts @@ -1241,7 +1571,24 @@ def _runtime_timing_summary(raw_root: Path) -> dict[str, int | float]: } -def _salvage_artifacts(resource: KubernetesCellRunner, cell_id: str) -> None: +def _cell_runner(plan: FPMCollectionPlan, cell: FPMCell, manifest: Path, cell_dir: Path): + executor = getattr(plan.options, "executor", "kubernetes") + if executor == "slurm": + from .slurm import SlurmCellRunner + + return SlurmCellRunner( + manifest, + cell_dir, + image=plan.options.slurm_container_image, + mounts=plan.options.slurm_container_mounts, + total_gpus=cell.topology.total_gpus, + ) + if executor != "kubernetes": + raise ValueError(f"unknown FPM executor {executor!r}") + return KubernetesCellRunner(manifest, cell_dir) + + +def _salvage_artifacts(resource, cell_id: str) -> None: """Best-effort artifact salvage after a failed or interrupted attempt. kubectl-exec disconnects do not stop pod processes, so the runtime may @@ -1311,6 +1658,7 @@ def _recover_completed_attempt( cell_dir = root / "cells" / cell.cell_id try: attempt_id = _required_attempt_id(entry, cell.cell_id) + _validate_points_receipts(plan, cell, cell_dir / "raw", attempt_id) _runtime_collection_summary( cell, cell_dir / "raw", @@ -1342,7 +1690,7 @@ def _recover_completed_attempt( ) return None try: - KubernetesCellRunner(manifest, cell_dir).cleanup() + _cell_runner(plan, cell, manifest, cell_dir).cleanup() except Exception as error: logger.warning( "FPM cell %s recovery refused: teardown of the abandoned workload failed: %s", @@ -1466,6 +1814,8 @@ def _run_collection_impl( database_root: str | None = None, publish_partial: bool = False, ) -> list[dict[str, object]]: + if _frozen_points(plan) is not None and smoke: + raise ValueError("--fpm-benchmark-points-file cannot be combined with --smoke") run_started_at = _utc_now() root = Path(artifact_root).expanduser().resolve() / plan.sha256[:16] if smoke: @@ -1498,6 +1848,12 @@ def _run_collection_impl( Path(str(database_entry.get("parquet", ""))).expanduser(), Path(str(database_entry.get("metadata", ""))).expanduser(), plan, + expected_attempt_ids={ + cell.cell_id: _required_attempt_id(checkpoint["cells"].get(cell.cell_id, {}), cell.cell_id) + for cell in plan.cells + }, + reused_cell_ids=tuple(database_entry.get("skipped_first_publisher_wins", ())), + expected_cell_rows=database_entry.get("cell_rows"), ) except (OSError, RuntimeError, TypeError, ValueError, json.JSONDecodeError) as error: logger.warning("Completed FPM database checkpoint failed validation; rebuilding: %s", error) @@ -1533,6 +1889,12 @@ def _run_collection_impl( entry = checkpoint["cells"].get(cell.cell_id) if not isinstance(entry, dict) or entry.get("status") != "passed": continue + # A validated terminal database no longer depends on retained raw + # receipts. Unpublished cells still need them before publication. + if not formal_database_terminal and _frozen_points(plan) is not None: + _validate_points_receipts( + plan, cell, root / "cells" / cell.cell_id / "raw", _required_attempt_id(entry, cell.cell_id) + ) # This refresh only polishes checkpoint metadata for cells whose # results were already validated and published; raw artifacts that # are no longer readable (disk reclaimed, resume from a different @@ -1576,6 +1938,31 @@ def _run_collection_impl( continue cell_dir = root / "cells" / cell.cell_id + if getattr(plan.options, "executor", "kubernetes") == "slurm": + abandoned_manifest = cell_dir / FPM_MANIFEST_FILENAME + if abandoned_manifest.exists(): + # Shared result mounts stay writable until the old step exits. + # Verify teardown before replacing any part of that directory. + try: + _cell_runner(plan, cell, abandoned_manifest, cell_dir).cleanup() + except Exception as error: + checkpoint["cells"][cell.cell_id] = { + **previous, + "status": "cleanup_failed", + "cleanup_error": str(error), + "artifact_dir": str(cell_dir), + } + errors.append( + { + "module": "fpm_forward", + "cell_id": cell.cell_id, + "error_type": type(error).__name__, + "error_message": str(error), + "classification": "resource_cleanup_failed", + } + ) + _atomic_json(checkpoint_path, checkpoint) + continue if cell_dir.exists() and not resume: shutil.rmtree(cell_dir) cell_dir.mkdir(parents=True, exist_ok=True) @@ -1615,12 +2002,13 @@ def _run_collection_impl( manifest = cell_dir / FPM_MANIFEST_FILENAME run_script = cell_dir / FPM_RUN_SCRIPT_FILENAME env_script = cell_dir / FPM_ENV_FILENAME - if not manifest.exists() or not run_script.exists() or not env_script.exists(): + runtime_env = cell_dir / RUNTIME_ENV_FILENAME + if not manifest.exists() or not run_script.exists() or not env_script.exists() or not runtime_env.exists(): raise RuntimeError( f"Generator FPM target did not emit {FPM_MANIFEST_FILENAME}, " - f"{FPM_ENV_FILENAME}, and {FPM_RUN_SCRIPT_FILENAME}" + f"{FPM_ENV_FILENAME}, {FPM_RUN_SCRIPT_FILENAME}, and {RUNTIME_ENV_FILENAME}" ) - resource = KubernetesCellRunner(manifest, cell_dir) + resource = _cell_runner(plan, cell, manifest, cell_dir) # A prior invocation may have left the same-named workload alive # (cleanup timeout, killed collector host) even when THIS # checkpoint has no record of the cell: workload names derive @@ -1641,8 +2029,18 @@ def _run_collection_impl( [ run_script, env_script, + runtime_env, runtime_exec, runtime_preflight, + *_stage_points_file(plan, cell_dir), + *( + [ + runtime_preflight.parent / "fpm_text.txt", + *sorted(p for p in (runtime_preflight.parent / "dsv41").iterdir() if p.is_file()), + ] + if cell.execution_identity[0] + else [] + ), ], ) resource.prepare_attempt( @@ -1651,13 +2049,16 @@ def _run_collection_impl( plan_sha256=plan.sha256, attempt_id=attempt_id, ) + _record_points_receipts(resource, pods, plan, cell, attempt_id, phase="before") phase_marks["stage_s"] = round(time.monotonic() - mark, 3) mark = time.monotonic() resource.execute(pods) + _record_points_receipts(resource, pods, plan, cell, attempt_id, phase="after") phase_marks["execute_wall_s"] = round(time.monotonic() - mark, 3) mark = time.monotonic() resource.collect(pods) phase_marks["collect_s"] = round(time.monotonic() - mark, 3) + _validate_points_receipts(plan, cell, cell_dir / "raw", attempt_id) runtime_collection = _runtime_collection_summary( cell, cell_dir / "raw", @@ -1772,7 +2173,7 @@ def _run_collection_impl( for cell in plan.cells ) ): - from .database import aggregate_cell, write_formal_database + from .database import aggregate_cell, validate_formal_database_commit, write_formal_database # Explicit partial publication: rows from passed cells only; the # missing cells are recorded so coverage is auditable, never implied. @@ -1802,6 +2203,13 @@ def _run_collection_impl( parquet_path, metadata_path, first_wins_skipped = write_formal_database( plan, formal_rows, systems_root=systems_root ) + committed = validate_formal_database_commit( + parquet_path, + metadata_path, + plan, + reused_cell_ids=first_wins_skipped, + required_cells=tuple(publishable_cells), + ) checkpoint["database"] = { "status": "passed", "parquet": str(parquet_path), @@ -1811,6 +2219,7 @@ def _run_collection_impl( "plan_cells": len(plan.cells), "missing_cells": missing_cells, "skipped_first_publisher_wins": list(first_wins_skipped), + "cell_rows": committed["cell_rows"], } if partial: logger.warning( diff --git a/python/aisimulate/collector/fpm_forward/runtime/dsv41/LICENSE b/python/aisimulate/collector/fpm_forward/runtime/dsv41/LICENSE new file mode 100644 index 000000000..1ed6128b8 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/dsv41/LICENSE @@ -0,0 +1,210 @@ + NOTICE: The test data files under ./lib/llm/tests/data/deepseek-v3.2 are + derived from the DeepSeek-V3.2 model repository (originally developed by + DeepSeek). The original files were obtained from: + https://huggingface.co/deepseek-ai/DeepSeek-V3.2 (commit c69397ecfd1fd142e90e3fbad51f4c7e40b9f3d3) + These files are licensed under the MIT License. The full text of the MIT + License can be found in ./lib/llm/tests/data/deepseek-v3.2/LICENSE. + The rest of this codebase is licensed under the Apache License 2.0 as + described below. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/python/aisimulate/collector/fpm_forward/runtime/dsv41/README.md b/python/aisimulate/collector/fpm_forward/runtime/dsv41/README.md new file mode 100644 index 000000000..78ac88697 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/dsv41/README.md @@ -0,0 +1,216 @@ +# DeepSeek V4.1 native FPM canary producer + +Status: the historical producer completed **126 GB200 calibration points** +(100 prefill, 26 real decode) in eager TP4/DP1 mode. The retained +[collection receipt](https://github.com/ai-dynamo/aisimulate/blob/dd1fa97add17d3d74f580f4c3e0566c1b5f11827/data/experimental/deepseek-v41/gb200-fpm/calibration-v1/collection-receipt.json) +binds producer SHA256 `771d2591eb5893f636a4208d60eb10daf51192eb4062bfec22da2f981cbe8ec0` +and runtime-source manifest `a201055338b2e185db52b44744b0033aca850452580279bde4d3a4d8f8dceaef`. +The earlier allocation-timeout and eager warmup failures remain failed evidence; +they are not the completed calibration run. Subsequent ordinary-serving +[validation](https://github.com/ai-dynamo/aisimulate/blob/dd1fa97add17d3d74f580f4c3e0566c1b5f11827/data/experimental/deepseek-v41/gb200-fpm/ordinary-serving-retention128-v1/README.md) +exposed a large latency mismatch, so completed collection does not establish +serving prediction accuracy or general scheduler qualification. Current changes +to activation, startup configuration and explicit identity attestation have CPU +regression coverage; they have not been collected again on a GPU. This remains +a bounded collection extension, not a general serving scheduler. + +`dsv41_scheduler.py` extends the native Dynamo `InstrumentedScheduler`. It keeps +native `BenchmarkPoint`, FPM messages, schema-v2 rank artifacts, coverage checks, +and second-step decode timing. The producer writes additive provenance directly. +The Collector must continue to reject failed, incomplete, fake, or unmarked rows. + +The DP1 adapter bounds each complete real-KV point at 120 seconds, including +all same-request seed forwards and first-use kernel compilation. This deadline +is independent of Dynamo's short synthetic-point result timeout; that native +timeout, native interval timing and the campaign timeout remain unchanged. +Both timeout values and the adapter source hash are recorded in producer +provenance. Expiration raises a retained failure and never publishes a partial +point or switches to synthetic KV. The first formal-geometry canary exposed +the old eight-second deadline during Engram JIT; that failed attempt is retained. + +## Source and attribution + +The scheduler adapter is modified code derived from NVIDIA's Apache-2.0 Dynamo +implementation: https://github.com/ai-dynamo/dynamo/blob/54960177085413259859c88bd34ed0734d4c2ea9/components/src/dynamo/vllm/instrumented_scheduler.py + +Original copyright: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +All rights reserved. Original SPDX attribution remains in the derived module. +`LICENSE` preserves the upstream Apache-2.0 license. No root NOTICE file exists +in that upstream checkout. The original text fixture and test driver were authored +for this change; they contain no third-party corpus content. The fixture is also +Apache-2.0 licensed under the adjacent license and this attribution. + +Compatibility was initially inspected against vLLM preview commit +79a7108d9aea27ddab99ce1779290d300b17fc23: +https://github.com/vllm-project/vllm/tree/79a7108d9aea27ddab99ce1779290d300b17fc23 + +Relevant original vLLM files carry Apache-2.0 and `Copyright contributors to the +vLLM project`. The adapter calls their request, scheduler-output, and KV-manager +APIs; it does not vendor vLLM implementation files. Current source hashes are in +`runtime-source-sha256.json`. The actual ARM image reports build version +`0.1.dev20904+g179dd0fa9`; that abbreviated SHA does not resolve in the public +upstream repository. Actual installed sources were exported and compared: the six request/scheduler/ +KV-manager files match the inspected preview exactly, as does model_state.py. +The actual model and Engram files add DP-shared table sharding/gathers and per-ubatch +staging; config/engram.py adds the model_has_engram_layers helper. The source +manifest now pins the exported image files. The exported DP helper explicitly returns shard size 1 for DP1, creates no +Engram DP group in that case, and gather_engram_hashes returns its input unchanged. +The exported parallel config confirms num_ubatches=0 for this canary. The actual +ModelConfig preflight passed with outer/text model_type=deepseek_v41, architecture +DeepseekV41ForCausalLM, and resolved Engram cpu_offload=false. Full producer import +passes with the composite runtime below. Completed calibration and subsequent +serving-accuracy limitations are recorded above. Do not substitute the package +version string for source verification. +Artifacts report vllm_revision=null, the actual package version, the inspected API +revision, and the source-manifest digest rather than inventing a global git revision. + +The derived files and source pin are recorded in root THIRD_PARTY_NOTICES.md +and its byte-identical packaged copy. Run scripts/check_packaged_legal_files.py +after changing either notice. + +## Runtime + +Stage this directory on PYTHONPATH with Python `ai-dynamo==1.4.2` and matching +`ai-dynamo-runtime==1.4.2`, then overlay these four unchanged files from Dynamo +`54960177085413259859c88bd34ed0734d4c2ea9`: + +- `components/src/dynamo/vllm/instrumented_scheduler.py` +- `components/src/dynamo/vllm/benchmark_points.py` +- `components/src/dynamo/vllm/gc_policy.py` +- `components/src/dynamo/common/forward_pass_metrics.py` (already identical in 1.4.2) + +Keep both distributions' licenses and the source manifest. This is a composite +Python runtime: the instrumentation revision is not the revision of the entire +worker. Using all Python modules from that newer commit with native runtime 1.4.2 +failed worker import at `update_model_taints`. The verified composition passes +frontend help, worker help, native runtime initialization/shutdown and the +source-checked scheduler import. Generator 1.3/0.24 metadata selects a supported +launch template; it does not describe the deployed backend versions. + +Install the matching AISimulate +wheel, including its native extension, so the shared execution_identity helper +can import. A source-only Python path is insufficient. The image must provide +the pinned vLLM source and native GPU dependencies. + +Set: + +- DYN_FPM_DSV41_REAL_KV=1 +- DYN_FPM_INPUT_TEXT=/staged/path/fpm_text.txt +- DYN_FPM_TOKENIZER_REVISION=fb2764a5cf321eaa5070ca8f9e892818f477c16d + +Use the native Dynamo scheduler class. The lightweight `sitecustomize` hook +defers source verification and subclass activation until that exact scheduler +module is imported. Compiler/helper interpreters do not import Dynamo or vLLM +through this hook. Both native-first and adapter-first imports are tested. An +activation failure raises a chained `RuntimeError` at that later import. It is +outside Python's startup `sitecustomize` exception handler, so ordinary imports +fail with a traceback and allow `finally`/exit cleanup. The preflight catches +activation and shared-SDK import failures, saves a failed audit, and re-raises; +it also rejects a missing scheduler activation before model loading. It never +falls back to the unmodified scheduler. + +Runtime gates require a local pinned V4.1 checkpoint, TP4, DP1, PP1, no EP, no +speculation/DSpark, no ubatching/DBO or context parallelism, no KV or encoder connector, and explicit Engram cpu_offload=false. +The last requirement matters: the preview defaults to CPU UVA offload. All requests +are text-only; no vision inputs or decoder replay are injected. + +This first V4.1 dataset contract supports eager execution only. Pass +`--fpm-enforce-eager`; the frozen plan and rendered arguments record it, the +producer checks the actual model configuration, and the native reader requires +the producer's `execution_mode=eager` evidence. Graph data cannot be published +under this contract. Supporting graph execution later requires a distinct query +identity, as well as separate calibration and validation. The earlier graph +startup attempt is retained as runtime qualification evidence only. + +Native benchmark warmup_iterations must be 0. Native per-shape eager warmups remain +in place and run through the same real-forward path. The token-stream sidecar's +schema 2 labels each history as `warmup` or `measurement`, records the frozen +native warmup IDs, and keeps completed warmup witnesses separately in +`warmup_results`. Only the native measured `results` enter calibration tables. +The reader requires exact coverage of both sets and rejects leaked warmup +timings or unclassified extra histories. This mirrors the eager-replica identity +and save boundaries in [Dynamo's pinned scheduler](https://github.com/ai-dynamo/dynamo/blob/54960177085413259859c88bd34ed0734d4c2ea9/components/src/dynamo/vllm/instrumented_scheduler.py#L2376). +The first eager canary's original mixed sidecar remains failed qualification +evidence; its extra warmup records must not be silently removed. + +An explicit native point +manifest is accepted. For an automatic native grid, use max_model_len<=2050, +max_num_seqs<=2 and max_num_batched_tokens<=512. The complete study requires +2050/2/512 and enough native page capacity. Every grid point is checked against +batch<=2, decode past KV<=2048, prefill prefix-plus-new<=2048 per request, and +total newly scheduled prefill tokens<=512. Unsupported points fail the run +instead of being silently removed. +A native decode point with context<2 fails instead of being relabeled. + +## Startup configuration + +The adapter's image-layout default lives in the staged `runtime-paths.json`. +To use another verified composition, set `PYTHONPATH` through the existing +`K8sConfig.extra_env` deployment input. The Collector prepends its staged adapter +directory and writes the identical resolved path into `run.sh` and +`collector-runtime-env.sh`; the latter also configures preflight in Slurm, where +Kubernetes Pod environment injection is absent. Required source hashes still +apply at the configured paths; changing a path does not bypass verification. + +`FPM_READINESS_TIMEOUT_SECONDS` uses the same `K8sConfig.extra_env` input, defaults +to 900 seconds, and accepts integer budgets from 1 through 3600. This is the +bounded etcd rendezvous wait after each node's preflight, covering staggered +container startup. For example, `{name: FPM_READINESS_TIMEOUT_SECONDS, value: "600"}` +sets a ten-minute budget. The outer execution timeout still bounds the entire +step; this is not an additional native benchmark timing or per-point budget. +Slurm `wait_ready` separately waits for the exact existing allocation to become +RUNNING and validates its node count within the caller's single deadline; it does +not claim the container or model has loaded. + +## State and timing contract + +A request is registered once, with the full tokenizer-produced prompt history. +Its prefix is executed in real forward chunks from position zero. The request and +all its block tables remain alive, including the non-prefix-cacheable compressor +ring and Engram lookback state. Each seed forward drains before measurement starts. +Prefill measures the suffix on that same request. Decode performs native admission +at context-1 then measures the following steady step at the exact requested context. +Those two decode steps can pipeline as in the native benchmark. + +No code assigns num_computed_tokens. vLLM advances it only through the normal +post-schedule accounting, and the producer counts seed tokens only after their +model output callbacks complete. It refuses synthetic-prefix/decode entry points, +allocation failure, lost requests, unexpected CoW, timeout, or missing native FPMs. +The `kvwarm_real_kv` point annotation is earned immediately before successful native +save. Each saved result also has completed_seed_tokens and same_request witnesses. + +Input provenance includes the original UTF-8 text SHA, the SHA of compact JSON token +IDs produced by the actual runtime tokenizer, tokenizer revision, token count and +unique count. The producer also writes benchmark.token-streams.jsonl +with actual prompt and sampled output token IDs for every completed point. Its +SHA and record count are in input_provenance.token_stream_manifest; individual +point witnesses include the stream SHA. Per request, offset is 131*request_index+17*benchmark_id, wrapping the +same real-text token stream to the requested length. This fixture is reproducible; +it is not a representative production corpus. Routing/locality sensitivity and +numerical equivalence remain live validation work. + +From the repository root, run +`python/aisimulate/.venv/bin/pytest -p no:timeout -c python/aisimulate/pytest.ini python/aisimulate/tests/unit/collector/test_fpm_dsv41_producer.py` +to execute the 14 isolated CPU lifecycle and producer-consumer contract checks. +The test wrapper sets the producer and native artifact module paths. They cover actual +seed completion before measure, retained request identities, cached-prefill state, +uneven request lengths, decode pipeline timing sequence, failure paths, and earned +annotations. These are state-machine tests, not substitutes for GPU correctness. + +Same-request seed forwards qualify retained KV state for cached-prefill points. +Cross-request prefix-cache reuse needs separate serving verification with a warm +request, a subsequent request sharing its prefix, and a cold control. End-to-end +latency and full Dynamo FPM traces from those verification runs remain separate +from the calibration grid and its published curves. + + +The pinned Linux ARM64 image is +`vllm/vllm-openai@sha256:d84a123255b822fc22508635218000187221794f59c0694c33b0650d1e377d58`. +Mount the verified composite Dynamo tree at `/opt/dsv41-dynamo` read-only. +Install the matching AISimulate wheel, including its native extension, in the +runtime image. Importing its shared SDK identity through a source-only Python +path is insufficient: the package loads the native extension at import time. +Stage the adapter, source manifest, and fixture through the Collector. The adapter, +Dynamo source and AISimulate wheel must be available during runtime preflight. +The installed native `ai-dynamo-runtime` is separately pinned and import-tested; +the Python package version alone does not establish scheduler compatibility. diff --git a/python/aisimulate/collector/fpm_forward/runtime/dsv41/dsv41_scheduler.py b/python/aisimulate/collector/fpm_forward/runtime/dsv41/dsv41_scheduler.py new file mode 100644 index 000000000..0d24485de --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/dsv41/dsv41_scheduler.py @@ -0,0 +1,507 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Scheduler integration adapted from ai-dynamo/dynamo at +# 54960177085413259859c88bd34ed0734d4c2ea9, components/src/dynamo/vllm/instrumented_scheduler.py. +# Modified: bounded same-request real-KV collection for vLLM V4.1 preview. +"""Canary-only native FPM extension; see README.md for immutable API pins. + +Every state-bearing request is filled by model forwards from position zero. +The same request and block tables survive through its measured suffix/decode. +This module never assigns num_computed_tokens or imports synthetic KV. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import time +from pathlib import Path + +import dynamo.vllm.instrumented_scheduler as native +from vllm.sampling_params import SamplingParams +from vllm.tokenizers import get_tokenizer +from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput +from vllm.v1.request import Request, RequestStatus + +DYNAMO_SHA = "54960177085413259859c88bd34ed0734d4c2ea9" +VLLM_SHA = "79a7108d9aea27ddab99ce1779290d300b17fc23" +MODEL_SHA = "fb2764a5cf321eaa5070ca8f9e892818f477c16d" +MAX_BATCH = 2 +MAX_CONTEXT = 2048 +MAX_NEW = 512 +# This DP1 producer executes a real seed chain and first-use kernel compilation +# before measurement. The native short result timeout bounds a synthetic point, +# not this complete real-KV lifecycle. Keep a finite whole-point deadline without +# changing native measurement timing, DP synchronization or campaign timeout. +REAL_POINT_TIMEOUT_SECONDS = 120.0 + + +def token_slice(pool: list[int], length: int, offset: int) -> list[int]: + """Rotate and repeat the tokenizer's real-text stream deterministically.""" + if not pool or length < 0: + raise ValueError("non-empty token pool and nonnegative length required") + return [pool[(offset + index) % len(pool)] for index in range(length)] + + +class DeepseekV41RealKVScheduler(native.InstrumentedScheduler): + def _bench_init(self, config): + self._real_tags = {} + self._real_callback_stage = None + self._real_stage = None + self._real_requests = [] + self._real_stages = [] + self._real_stage_index = 0 + self._real_submitted = set() + self._real_outstanding = 0 + self._real_deadline = 0.0 + self._real_token_streams = [] + self._real_witnesses = {} + self._real_warmup_results = [] + self._real_expected_warmup_ids = [] + self._real_seed_tokens = 0 + self._real_input = None + self._real_identity = None + super()._bench_init(config) + if not self._bench_active: + raise ValueError("V4.1 canary overlay requires native benchmark mode") + if config.model_config.enforce_eager is not True: + raise ValueError("V4.1 real-KV collection requires enforce_eager=true; graph timing is not qualified") + parallel = config.parallel_config + if (parallel.tensor_parallel_size, parallel.pipeline_parallel_size, parallel.data_parallel_size) != (4, 1, 1): + raise ValueError("V4.1 canary requires pure TP4 / PP1 / DP1") + if ( + parallel.use_ubatching + or parallel.prefill_context_parallel_size != 1 + or parallel.decode_context_parallel_size != 1 + ): + raise ValueError("V4.1 canary requires no ubatching/DBO and no context parallelism") + if parallel.enable_expert_parallel or config.speculative_config is not None: + raise ValueError("V4.1 canary does not support EP or speculative decoding") + if self.connector is not None or self.ec_connector is not None: + raise ValueError("V4.1 canary forbids KV/encoder connectors") + engram = config.engram_config + if engram is None or engram.cpu_offload: + raise ValueError("V4.1 canary requires explicit Engram cpu_offload=false") + if self._bench_explicit_points is None and ( + self.max_model_len > 2050 + or self.max_num_running_reqs > MAX_BATCH + or self.max_num_scheduled_tokens > MAX_NEW + ): + raise ValueError("automatic V4.1 canary grid requires model_len<=2050, max_seqs<=2, batched_tokens<=512") + if self._bench_config.warmup_iterations != 0: + raise ValueError("set native warmup_iterations=0; eager shape warmups remain enabled") + model_path = Path(config.model_config.model) + raw_config = json.loads((model_path / "config.json").read_text()) + if ( + config.model_config.architecture != "DeepseekV41ForCausalLM" + or config.model_config.hf_config.model_type != "deepseek_v41" + ): + raise ValueError("loaded model must be DeepSeek V4.1") + from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, execution_identity + + self._real_identity = dict( + zip( + EXECUTION_COLUMNS, + execution_identity( + raw_config, + engram_cpu_offload=engram.cpu_offload, + input_modality="text", # This producer creates only text token-ID requests. + ), + strict=True, + ) + ) + revision = os.environ.get("DYN_FPM_TOKENIZER_REVISION") + if revision != MODEL_SHA: + raise ValueError("tokenizer revision must equal the pinned V4.1 checkpoint revision") + text_bytes = Path(os.environ["DYN_FPM_INPUT_TEXT"]).read_bytes() + tokenizer = get_tokenizer( + config.model_config.tokenizer, + tokenizer_mode=config.model_config.tokenizer_mode, + trust_remote_code=config.model_config.trust_remote_code, + revision=revision, + ) + self._real_tokens = list(tokenizer.encode(text_bytes.decode("utf-8"), add_special_tokens=False)) + if len(set(self._real_tokens)) < 2 or not all(type(x) is int and x >= 0 for x in self._real_tokens): + raise ValueError("tokenizer text stream must contain distinct valid token ids") + self._real_input = { + "source": "tokenizer_text", + "text_sha256": hashlib.sha256(text_bytes).hexdigest(), + "token_ids_sha256": hashlib.sha256( + json.dumps(self._real_tokens, separators=(",", ":")).encode() + ).hexdigest(), + "tokenizer_revision": revision, + "token_count": len(self._real_tokens), + "unique_token_count": len(set(self._real_tokens)), + "sampling": "rotate stream by 131*request_index+17*benchmark_id; repeat to requested length", + } + + def _kvwarm_warm_eligible(self): + return True + + def _bench_blocks_per_req(self, num_tokens, *, has_cache_hit=False, apply_admission_cap=False): + # Use the pinned managers' native capacity semantics, including one-block rings. + return sum( + manager.get_num_blocks_to_allocate( + request_id="__dsv41_capacity_probe__", + num_tokens=num_tokens, + new_computed_blocks=[], + total_computed_tokens=0, + num_local_computed_tokens=0, + num_tokens_main_model=num_tokens, + apply_admission_cap=apply_admission_cap, + ) + for manager in self.kv_cache_manager.coordinator.single_type_managers + ) + + def _bench_build_grid(self): + built = self._bench_grid_built + super()._bench_build_grid() + if not built: + self._real_validate_grid() + + def _real_validate_grid(self): + self._real_expected_warmup_ids = sorted( + point.benchmark_id for point in self._bench_grid if native.EAGER_WARMUP_REASON in point.sample_reasons + ) + for point in self._bench_grid: + prefix, suffix = self._real_lengths(point) + if ( + not 1 <= point.batch_size <= MAX_BATCH + or max(prefix, default=0) + int(point.point_type == "decode") > MAX_CONTEXT + ): + raise ValueError("V4.1 canary point exceeds batch/context bound") + if point.point_type == "prefill": + if sum(suffix) > MAX_NEW: + raise ValueError("V4.1 canary prefill exceeds total new-token bound") + if any(p + q > MAX_CONTEXT for p, q in zip(prefix, suffix, strict=True)): + raise ValueError("V4.1 canary prefill exceeds prefix plus new-token context bound") + if point.point_type == "decode" and min(prefix) < 1: + raise ValueError("V4.1 real decode requires context >=2; no coordinate clamping") + + def _bench_materialize_prefill_candidate(self, candidate, path, *, generated=False): + partition = candidate.partition.model_dump() if candidate.partition is not None else None + if not self._bench_prefill_point_feasible( + candidate.total_prefill_tokens, + candidate.batch_size, + candidate.total_kv_read_tokens, + partition, + candidate.rows, + ): + self._bench_raise_explicit_infeasible(path, candidate) + capture, padding, reasons = self._bench_cudagraph_metadata( + candidate.total_prefill_tokens, + self._bench_prefill_capture_sizes, + self._bench_capacity_limit("max_num_scheduled_tokens"), + ) + return native.BenchmarkPoint( + point_type="prefill", + total_prefill_tokens=candidate.total_prefill_tokens, + total_kv_read_tokens=candidate.total_kv_read_tokens, + batch_size=candidate.batch_size, + expected_cudagraph_mode=self._bench_prefill_cudagraph_mode if capture is not None else "NONE", + expected_capture_size=capture, + padding_tokens=padding, + partition=partition, + rows=candidate.rows, + sample_reasons=[native._bench_origin_reason(generated), *reasons], + ) + + def _real_lengths(self, point): + if point.point_type == "decode": + contexts = self._bench_decode_context_lengths(point.total_kv_read_tokens, point.batch_size) + return [x - 1 for x in contexts], [1] * point.batch_size + return ( + self._bench_prefill_kv_read_lengths( + point.total_kv_read_tokens, point.batch_size, point.partition, point.rows + ), + self._bench_prefill_new_token_lengths( + point.total_prefill_tokens, point.batch_size, point.partition, point.rows + ), + ) + + def _bench_cache_fake_prefixes(self, *args, **kwargs): + raise RuntimeError("synthetic prefix KV is forbidden for V4.1") + + def _bench_inject_fake_decode(self, *args, **kwargs): + raise RuntimeError("synthetic decode KV is forbidden for V4.1") + + def _bench_new_request_counts_as_decode(self, req_id): + return False # Every newly registered request starts at actual position zero. + + def _bench_should_record_scheduled(self, scheduled): + return self._real_callback_stage in {"admission", "measure"} and super()._bench_should_record_scheduled( + scheduled + ) + + def _real_begin(self, point): + prefix, suffix = self._real_lengths(point) + self._bench_current_point = point + self._bench_current_fpms = [] + self._bench_expected_fpms = 2 if point.point_type == "decode" else 1 + self._bench_admission_kv_tokens = sum(prefix) + self._real_seed_tokens = 0 + self._real_expected_seed = sum(prefix) + self._real_deadline = time.monotonic() + REAL_POINT_TIMEOUT_SECONDS + self._real_requests = [] + self._real_submitted = set() + self._real_stages = [] + remaining = list(prefix) + chunk = min(MAX_NEW, self.max_num_scheduled_tokens // point.batch_size) + if chunk < 1: + raise RuntimeError("native token budget cannot admit one token per request") + while any(remaining): + counts = [min(chunk, value) for value in remaining] + self._real_stages.append(("seed", counts)) + remaining = [value - count for value, count in zip(remaining, counts, strict=True)] + self._real_stages.extend( + [("admission", suffix), ("measure", suffix)] if point.point_type == "decode" else [("measure", suffix)] + ) + self._real_stage_index = 0 + for index in range(point.batch_size): + prompt_len = prefix[index] if point.point_type == "decode" else prefix[index] + suffix[index] + req_id = f"__dsv41_real_{point.benchmark_id}_{index}_{self._bench_seq}" + request = Request( + request_id=req_id, + prompt_token_ids=token_slice(self._real_tokens, prompt_len, 131 * index + 17 * point.benchmark_id), + sampling_params=SamplingParams( + max_tokens=3 if point.point_type == "decode" else 1, ignore_eos=True, temperature=0 + ), + pooling_params=None, + block_hasher=self._bench_block_hasher, + cache_salt=req_id, + ) + request.status = RequestStatus.RUNNING + self.requests[req_id] = request + self.running.append(request) + self._bench_active_req_ids.add(req_id) + self._real_requests.append(request) + self._bench_seq += 1 + + def _real_output(self, stage, counts): + if sum(counts) > self.max_num_scheduled_tokens: + raise RuntimeError("requested real-forward point exceeds native scheduling token budget") + output = SchedulerOutput.make_empty() + output.finished_req_ids = self.finished_req_ids + output.num_common_prefix_blocks = [0] * self.kv_cache_manager.num_kv_cache_groups + for request, count in zip(self._real_requests, counts, strict=True): + if count == 0: + continue + if request.request_id not in self.requests: + raise RuntimeError("real warm request was lost before measurement") + blocks = self.kv_cache_manager.allocate_slots(request, count, delay_cache_blocks=True) + if blocks is None: + raise RuntimeError("real KV allocation failed; synthetic fallback is forbidden") + rid = request.request_id + output.num_scheduled_tokens[rid] = count + if rid not in self._real_submitted: + output.scheduled_new_reqs.append( + NewRequestData.from_request( + request, blocks.get_block_ids(), prefill_token_ids=request._all_token_ids + ) + ) + self._real_submitted.add(rid) + else: + cached = output.scheduled_cached_reqs + cached.req_ids.append(rid) + cached.all_token_ids[rid] = request._all_token_ids.copy() + cached.new_block_ids.append(blocks.get_block_ids(allow_none=True)) + cached.num_computed_tokens.append(request.num_computed_tokens) + cached.num_output_tokens.append(request.num_output_tokens + request.num_output_placeholders) + output.total_num_scheduled_tokens = sum(counts) + output.new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids() or None + copies, retained = self.kv_cache_manager.take_kv_cache_block_copies() + if copies or retained: + raise RuntimeError("unexpected copy-on-write in private same-request real KV path") + self._real_tags[id(output)] = stage + self._real_outstanding += 1 + if stage == "admission" or (stage == "measure" and self._bench_current_point.point_type == "prefill"): + self._bench_sync_pending = True + return output + + def _real_step(self, point_type): + if self._real_stage is not None: + if time.monotonic() >= self._real_deadline: + raise RuntimeError("real KV warm/measurement timed out; no fallback") + if self._real_stage_index < len(self._real_stages): + stage, counts = self._real_stages[self._real_stage_index] + # Complete every seed forward before proceeding. Admission and + # steady decode may pipeline, matching native second-step timing. + if self._real_outstanding and stage != "measure": + return None + if self._real_outstanding and self._real_stage == "seed": + return None + if stage != "seed" and self._real_seed_tokens != self._real_expected_seed: + raise RuntimeError("measured forward requested before all real seed tokens completed") + self._real_stage_index += 1 + self._real_stage = stage + return self._real_output(stage, counts) + if self._real_outstanding: + return None + if len(self._bench_current_fpms) != self._bench_expected_fpms: + raise RuntimeError("real KV point did not yield the exact native FPM count") + stream = { + "benchmark_id": self._bench_current_point.benchmark_id, + "sampling_role": "warmup" + if native.EAGER_WARMUP_REASON in self._bench_current_point.sample_reasons + else "measurement", + "requests": [ + { + "request_index": index, + "prompt_token_ids": list(request.prompt_token_ids), + "output_token_ids": list(request._all_token_ids[len(request.prompt_token_ids) :]), + "computed_tokens": request.num_computed_tokens, + } + for index, request in enumerate(self._real_requests) + ], + } + stream_bytes = json.dumps(stream, sort_keys=True, separators=(",", ":")).encode() + self._real_token_streams.append(stream_bytes) + self._bench_current_point.sample_reasons.append("kvwarm_real_kv") + self._real_witnesses[self._bench_current_point.benchmark_id] = { + "completed_seed_tokens": self._real_seed_tokens, + "same_request": True, + "allocated_fake_tokens": 0, + "token_stream_sha256": hashlib.sha256(stream_bytes).hexdigest(), + } + if stream["sampling_role"] == "warmup": + # Native 5496017 prepends eager replicas with IDs after the + # measured range and discards them at save time (2376-2500, + # 4225-4243). Preserve their real histories under an explicit + # role, while only measured rows enter calibration curves. + point = self._bench_current_point + self._real_warmup_results.append( + { + "point": { + name: getattr(point, name) + for name in ( + "benchmark_id", + "point_type", + "batch_size", + "total_prefill_tokens", + "total_kv_read_tokens", + "sample_reasons", + ) + }, + "real_kv_witness": dict(self._real_witnesses[point.benchmark_id]), + "fpms": list(self._bench_current_fpms), + } + ) + self._bench_save_current_point() + self._bench_cleanup_requests() + self._real_stage = None + self._real_requests = [] + return None + if self._bench_stop_at_timeout_boundary(point_type): + return None + point = self._bench_pop_next(point_type) + if point is None: + self._bench_phase = ( + native._BenchPhase.DECODE_SWEEP + if point_type == "prefill" and self._bench_config.mode == "agg" + else native._BenchPhase.DONE + ) + return None + self._real_begin(point) + self._real_stage = "ready" + return self._real_step(point_type) + + def _bench_step_prefill(self): + return self._real_step("prefill") + + def _bench_step_decode(self): + return self._real_step("decode") + + def schedule(self, throttle_prefills=False): + if not self._bench_active: + return super().schedule(throttle_prefills) + try: + self.current_step += 1 + self.kv_cache_manager.new_step_starts() + output = self._bench_step() + if output is None: + output = SchedulerOutput.make_empty() + output.finished_req_ids = self.finished_req_ids + if output.total_num_scheduled_tokens: + self.sched_step_seq += 1 + self._update_after_schedule(output) + self._bench_synchronize_output(output) + if output.total_num_scheduled_tokens: + self._schedule_times.append(time.monotonic()) + return output + except Exception as error: + self._bench_abort(error) + raise + + def _update_from_output(self, scheduler_output, model_runner_output): + stage = self._real_tags.pop(id(scheduler_output), None) + self._real_callback_stage = stage + try: + result = super()._update_from_output(scheduler_output, model_runner_output) + if stage == "seed": + self._real_seed_tokens += scheduler_output.total_num_scheduled_tokens + if stage is not None: + self._real_outstanding -= 1 + return result + finally: + self._real_callback_stage = None + + def _bench_write_results(self): + super()._bench_write_results() + destination = Path(self._bench_config.output_path) + output = json.loads(destination.read_text()) + stream_bytes = b"\n".join(self._real_token_streams) + (b"\n" if self._real_token_streams else b"") + stream_path = destination.with_suffix(".token-streams.jsonl") + stream_tmp = stream_path.with_suffix(".jsonl.tmp") + stream_tmp.write_bytes(stream_bytes) + os.replace(stream_tmp, stream_path) + output["input_provenance"] = dict(self._real_input or {}) + output["input_provenance"]["token_stream_manifest"] = { + "schema_version": 2, + "warmup_benchmark_ids": self._real_expected_warmup_ids, + "file": stream_path.name, + "sha256": hashlib.sha256(stream_bytes).hexdigest(), + "records": len(self._real_token_streams), + } + output["warmup_results"] = self._real_warmup_results + output["execution_identity"] = self._real_identity + output["execution_mode"] = "eager" + output["kvwarm"] = { + "enabled": True, + "warm_eligible": True, + "skip_reason": None, + "method": "same_request_real_forward", + "max_batch": MAX_BATCH, + "max_context": MAX_CONTEXT, + } + output["producer"] = { + "instrumentation_revision": DYNAMO_SHA, + "dynamo_revision": None, + "vllm_revision": None, + "vllm_package_version": __import__("vllm").__version__, + "reviewed_scheduler_api_revision": VLLM_SHA, + "runtime_source_manifest_sha256": hashlib.sha256( + Path(__file__).with_name("runtime-source-sha256.json").read_bytes() + ).hexdigest(), + "overlay_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "collection_timeouts": { + "same_request_seed_and_measure_seconds": REAL_POINT_TIMEOUT_SECONDS, + "native_point_result_seconds": self._bench_point_result_timeout_seconds, + "scope": "DP1 whole real-KV point; native timing and campaign timeout unchanged", + }, + } + for item in output["results"]: + item["kv_seed_regime"] = "real_kv" + item["real_kv_witness"] = self._real_witnesses[item["point"]["benchmark_id"]] + for item in output["iteration_groups"]: + item["kv_seed_regime"] = "real_kv" + temporary = destination.with_suffix(destination.suffix + ".real.tmp") + temporary.write_text(json.dumps(output, indent=2)) + os.replace(temporary, destination) + + +# Spawned workers may import this class by its defining module before the +# configured native scheduler path. Publish only after class creation so the +# lazy source-checking hook can complete either import order without recursion. +if os.environ.get("DYN_FPM_DSV41_REAL_KV") == "1": + native.InstrumentedScheduler = DeepseekV41RealKVScheduler diff --git a/python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-paths.json b/python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-paths.json new file mode 100644 index 000000000..faf9daeb5 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-paths.json @@ -0,0 +1 @@ +{"python_path":"/opt/dsv41-dynamo/components/src"} diff --git a/python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-source-sha256.json b/python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-source-sha256.json new file mode 100644 index 000000000..840c92158 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-source-sha256.json @@ -0,0 +1,20 @@ +{ + "vllm/v1/request.py": "da610d5203e72d2d524adc5fb39df11502a23520f2f31d9a44cf13d343f9c25d", + "vllm/v1/core/sched/output.py": "41cf5780a013a0c731a0409836ad635d1c84a6aad706ad4e68e189d4a6eaf525", + "vllm/v1/core/sched/async_scheduler.py": "e586a0ef3c6778be56a93e7f9bb712d4de9de6e7d7fe7e3e1d51dae83ecfc508", + "vllm/v1/core/sched/scheduler.py": "e5e1c18b1d7a6ea73adbb4921f64a35b00a96abe673691bf6a7f57281a524519", + "vllm/v1/core/kv_cache_manager.py": "44a097c5fafaab70b952434b90f57ddacb224e208f19ecc719fcb441af39290e", + "vllm/v1/core/single_type_kv_cache_manager.py": "128b98a0511f67d32f44767aa1658776a8246374b9eb8461d397e683ff3d984d", + "vllm/models/deepseek_v4_1/nvidia/model.py": "530ed24c8fd2e9daeb5c3d340ef52246618786f19eace8ec217c8d3bdf271110", + "vllm/models/deepseek_v4_1/common/engram.py": "41c5bdf25cf8337088247be768e3d43fa41e6358397800f8b602ebf1ede69480", + "vllm/config/engram.py": "99f8c132d2c7928b0d2b41615611fc5b92819fc038f8414a7ea6649bf3527712", + "vllm/models/deepseek_v4_1/nvidia/model_state.py": "a2475c04990a17532e93618a58e5bb06a20215b3538ef140d01fe72a1fe9e3ab", + "dynamo/vllm/instrumented_scheduler.py": "765586f5891908f074c7dccae3a35edf4fb2629688f2bfdae11ab0de732ae298", + "dynamo/common/forward_pass_metrics.py": "523dfc4e7ccbc5a37dbcf03b331a41b628227ad8a5d9c64fda8abc96dd31b188", + "dynamo/vllm/benchmark_points.py": "b8be74e3bd375a332331d437220fb85e8d569c343bf8bf5e0f568a21a99f3875", + "dynamo/vllm/gc_policy.py": "4221e5a20718525559ad4c9e9bee34928ee22dcf2cfdf8036936be2a360153a4", + "vllm/distributed/parallel_state.py": "33be0ba9795544029ade187c56331c1b21c32bbda5d5cd86761edd14fa8a8259", + "vllm/distributed/__init__.py": "97f28c30b32662af8cab163939e4ee423037f11630752e13b55cbc25a7d0ab71", + "vllm/config/parallel.py": "b58a49f026a9a748d87c005aea0169d0ed23a24697dd0af438d8c4118112fc46", + "vllm/transformers_utils/configs/deepseek_v41.py": "e0c7583efa9baa95302398daa61604e0f656ecf9f1b072c5715e5ed5a5c54599" +} diff --git a/python/aisimulate/collector/fpm_forward/runtime/dsv41/sitecustomize.py b/python/aisimulate/collector/fpm_forward/runtime/dsv41/sitecustomize.py new file mode 100644 index 000000000..764a603f0 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/dsv41/sitecustomize.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Lazily activate the source-checked overlay only in scheduler processes.""" + +import os + +if os.environ.get("DYN_FPM_DSV41_REAL_KV") == "1": + import hashlib + import importlib.abc + import importlib.machinery + import importlib.util + import json + import sys + from pathlib import Path + + _TARGET = "dynamo.vllm.instrumented_scheduler" + + def _verify_sources(): + expected = json.loads(Path(__file__).with_name("runtime-source-sha256.json").read_text()) + for path, digest in expected.items(): + module_path = path.removesuffix("/__init__.py") if path.endswith("/__init__.py") else path[:-3] + module = module_path.replace("/", ".") + spec = importlib.util.find_spec(module) + if spec is None or spec.origin is None: + raise RuntimeError(f"required pinned source is unavailable: {path}") + actual = hashlib.sha256(Path(spec.origin).read_bytes()).hexdigest() + if actual != digest: + raise RuntimeError(f"pinned source mismatch: {path}; expected={digest}; actual={actual}") + + class _SchedulerLoader(importlib.abc.Loader): + def __init__(self, original): + self.original = original + + def create_module(self, spec): + return self.original.create_module(spec) + + def exec_module(self, module): + try: + _verify_sources() + self.original.exec_module(module) + adapter = sys.modules.get("dsv41_scheduler") + if adapter is not None and not hasattr(adapter, "DeepseekV41RealKVScheduler"): + # A spawned worker may unpickle the adapter class first. + # Its import needs this native base before it can finish; + # the adapter publishes its completed class at module end. + return + from dsv41_scheduler import DeepseekV41RealKVScheduler + + module.InstrumentedScheduler = DeepseekV41RealKVScheduler + except Exception as error: + # This loader runs on a later explicit scheduler import, not + # during sitecustomize initialization. Raising fails that import + # (including adapter-first spawn imports) and permits traceback, + # preflight audit and process cleanup instead of skipping them. + raise RuntimeError(f"V4.1 real-KV producer activation failed: {error}") from error + + def __getattr__(self, name): + return getattr(self.original, name) + + class _SchedulerFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname != _TARGET: + return None + spec = importlib.machinery.PathFinder.find_spec(fullname, path, target) + if spec is not None and spec.loader is not None: + spec.loader = _SchedulerLoader(spec.loader) + return spec + + # NVRTC/compiler helpers inherit PYTHONPATH and this activation variable. + # They must not import vLLM/Dynamo merely by starting a Python interpreter. + sys.meta_path.insert(0, _SchedulerFinder()) diff --git a/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh b/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh index 24e1a7cec..366a609c3 100644 --- a/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh +++ b/python/aisimulate/collector/fpm_forward/runtime/fpm_exec.sh @@ -17,6 +17,15 @@ workdir=/tmp/fpm-bench # the same status before any resource is started. source "${workdir}/fpm_env.sh" +# This is staged by the Collector from the same resolved deployment settings +# used by run.sh. It carries startup configuration across both transports. +source "${workdir}/collector-runtime-env.sh" +if [[ ! "${FPM_READINESS_TIMEOUT_SECONDS:-}" =~ ^[1-9][0-9]*$ ]] || + (( ${#FPM_READINESS_TIMEOUT_SECONDS} > 4 || FPM_READINESS_TIMEOUT_SECONDS > 3600 )); then + echo "FPM_READINESS_TIMEOUT_SECONDS must be an integer from 1 through 3600" >&2 + exit 2 +fi + etcd_endpoint="http://${FPM_MASTER_ADDR}:2379" engine_pid="" etcd_pid="" @@ -71,14 +80,16 @@ if [[ "${FPM_NODE_RANK}" == "0" ]]; then etcd_pid=$! fi +# Adapter activation and its import path have already been sourced from the +# frozen startup settings; no image-specific path belongs in this wrapper. python3 "${workdir}/preflight.py" if [[ "${FPM_NODE_RANK}" == "0" ]]; then # The leader owns the etcd process, so its readiness wait is etcd-aware: # a dead etcd (missing binary, bound port, data-dir permissions) fails the - # cell immediately with direct evidence instead of burning the 120s probe + # cell immediately with direct evidence instead of burning the configured probe # budget and reporting only "readiness timeout". - readiness_deadline=$((SECONDS + 120)) + readiness_deadline=$((SECONDS + FPM_READINESS_TIMEOUT_SECONDS)) while ! (exec 3<>"/dev/tcp/${FPM_MASTER_ADDR}/2379") 2>/dev/null; do if ! kill -0 "${etcd_pid}" 2>/dev/null; then set +e @@ -97,16 +108,16 @@ if [[ "${FPM_NODE_RANK}" == "0" ]]; then sleep 1 done else - python3 - "${FPM_MASTER_ADDR}" <<'PY' + python3 - "${FPM_MASTER_ADDR}" "$FPM_READINESS_TIMEOUT_SECONDS" <<'PY' import socket import sys import time host = sys.argv[1] -deadline = time.monotonic() + 120 +deadline = time.monotonic() + float(sys.argv[2]) while time.monotonic() < deadline: try: - with socket.create_connection((host, 2379), timeout=1): + with socket.create_connection((host, 2379), timeout=min(1.0, max(0.001, deadline - time.monotonic()))): break except OSError: time.sleep(0.2) diff --git a/python/aisimulate/collector/fpm_forward/runtime/fpm_text.txt b/python/aisimulate/collector/fpm_forward/runtime/fpm_text.txt new file mode 100644 index 000000000..f020db8f2 --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/runtime/fpm_text.txt @@ -0,0 +1,9 @@ +A field engineer records a quiet morning beside the river. The water level has changed since yesterday, and the team compares three measurements before deciding where to place the next sensor. Each observation includes a time, a location, and the conditions under which it was collected. A useful report explains both the result and the limits of the evidence. + +Two colleagues review a small software service. They follow a request from its arrival to its response, note which records must remain available between steps, and check that a retry produces the intended result. They write clear names for the inputs and preserve the original data. Their next experiment changes one setting so that its effect can be understood. + +一位研究人员整理了一周的实验记录。她先检查输入是否一致,再比较每次测量的时间和结果。不同的样本可能访问不同的数据,因此平均值并不能解释所有差异。团队保留完整的原始记录,并用另一组样本验证结论。 + +The library opens at nine. A student searches for a map, a teacher prepares a lesson about weather, and a visitor reads the history of a bridge. Their questions share a few words but lead to different shelves. The librarian listens carefully, checks the catalogue, and explains where each source can be found. + +A gardener measures the distance between young trees. Some receive sunlight in the morning and others in the afternoon. She labels the plots, writes down the watering schedule, and returns after several days. The plan leaves enough room for roots, paths, and future observations. diff --git a/python/aisimulate/collector/fpm_forward/runtime/preflight.py b/python/aisimulate/collector/fpm_forward/runtime/preflight.py index e8c85fd86..7ca48c756 100644 --- a/python/aisimulate/collector/fpm_forward/runtime/preflight.py +++ b/python/aisimulate/collector/fpm_forward/runtime/preflight.py @@ -6,6 +6,7 @@ from __future__ import annotations import json +import os from pathlib import Path _AUDIT_PATH = Path("/results/runtime-preflight.json") @@ -44,7 +45,18 @@ def main() -> None: # produce it before this process fails the pod. try: from dynamo.vllm.instrumented_scheduler import BenchmarkPoint, InstrumentedScheduler - except ImportError as error: + + if os.environ.get("DYN_FPM_DSV41_REAL_KV") == "1": + from dsv41_scheduler import DeepseekV41RealKVScheduler + + if InstrumentedScheduler is not DeepseekV41RealKVScheduler: + raise RuntimeError("V4.1 source-checked scheduler activation did not occur") + # Include the shared native SDK import in the rejected-image audit. + from aisimulate_core.sdk.fpm_identity import execution_identity + + if not callable(execution_identity): + raise RuntimeError("V4.1 runtime lacks the shared AISimulate identity helper") + except Exception as error: _write_audit( { "schema_version": 1, @@ -58,7 +70,7 @@ def main() -> None: ) raise RuntimeError( "Dynamo runtime lacks the required native FPM/KV-warm contract; " - f"importing dynamo.vllm.instrumented_scheduler failed: {error}. " + f"runtime activation or identity preflight failed: {error}. " "Provide a compatible Dynamo image." ) from error diff --git a/python/aisimulate/collector/fpm_forward/slurm.py b/python/aisimulate/collector/fpm_forward/slurm.py new file mode 100644 index 000000000..0fffe4c6b --- /dev/null +++ b/python/aisimulate/collector/fpm_forward/slurm.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the ordinary FPM runtime in an existing Slurm/Pyxis allocation. + +Only execution transport differs from Kubernetes. Generator scripts, native +result validation, attempt identity and publication remain owned by the common +campaign. The caller owns the allocation; this runner owns its named steps. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import re +import shutil +import subprocess +import time +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from aisimulate.fpm_contract import FPM_BENCHMARK_RESULT_GLOB + + +class SlurmCellRunner: + def __init__(self, manifest: Path, cell_dir: Path, *, image: str, mounts: tuple[str, ...], total_gpus: int): + from .runner import _expected_nodes + + self.cell_dir = cell_dir.resolve() + self.node_count = _expected_nodes(manifest) + if total_gpus % self.node_count: + raise ValueError("FPM GPUs must divide evenly across Slurm nodes") + self.gpus_per_node = total_gpus // self.node_count + self.image = image + if not image or any(char in image for char in ("\n", "\r")): + raise ValueError("Slurm FPM requires an explicit container image") + self.mounts = mounts + self.job_id = os.environ.get("SLURM_JOB_ID", "") + if not re.fullmatch(r"[0-9]+", self.job_id): + raise ValueError("Slurm FPM must run inside an existing sbatch/salloc allocation") + self.step_name = f"fpm-{hashlib.sha256(str(self.cell_dir).encode()).hexdigest()[:20]}" + # Keep ownership outside the replaceable cell payload so a fresh + # invocation can tear down an abandoned allocation's named steps. + self.owner_path = self.cell_dir.parent / ".slurm-owners" / f"{self.step_name}.json" + self.hosts: list[str] = [] + + def _command(self, args: list[str], *, timeout: float = 60, check: bool = True): + from .runner import _run_command + + try: + return _run_command(args, timeout=timeout, check=check) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error: + # Preparation and cleanup also invoke Slurm. Preserve their failure + # streams even when a campaign formats only str(error), and keep + # concurrent failures or retries from overwriting earlier evidence. + logs = self.cell_dir / "logs" / "transport-failures" / uuid.uuid4().hex + try: + logs.mkdir(parents=True) + for stream in ("stdout", "stderr"): + output = getattr(error, stream, None) or "" + if isinstance(output, bytes): + output = output.decode(errors="replace") + (logs / f"{stream}.log").write_text(output) + (logs / "failure.json").write_text( + json.dumps( + { + "executable": Path(args[0]).name, + "exception": type(error).__name__, + "returncode": getattr(error, "returncode", None), + "timeout_seconds": timeout, + }, + sort_keys=True, + ) + + "\n" + ) + except OSError as log_error: + error.add_note(f"Could not preserve Slurm failure streams: {log_error}") + raise + + def apply(self) -> None: + for executable in ("srun", "scontrol", "squeue", "scancel"): + if not shutil.which(executable): + raise RuntimeError(f"Slurm FPM requires {executable}") + self.owner_path.parent.mkdir(parents=True, exist_ok=True) + self.owner_path.write_text(json.dumps({"job_id": self.job_id, "step_name": self.step_name}) + "\n") + + def wait_ready(self, expected_nodes: int, timeout_seconds: float = 900) -> list[str]: + if expected_nodes != self.node_count: + raise ValueError("Slurm expected node count disagrees with the generated manifest") + if isinstance(timeout_seconds, bool) or not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError("Slurm readiness timeout must be finite and positive") + deadline = time.monotonic() + timeout_seconds + last_state = "unobserved" + + def remaining() -> float: + budget = deadline - time.monotonic() + if budget <= 0: + raise TimeoutError(f"Slurm allocation {self.job_id} not ready before deadline: {last_state}") + return budget + + # This method qualifies the existing allocation, not a container or + # model. Pyxis starts later in _exec; its startup skew has a separate + # bounded rendezvous budget in the staged Collector runtime settings. + while True: + snapshot = self._command(["scontrol", "show", "job", self.job_id, "--oneliner"], timeout=remaining()).stdout + fields = dict(re.findall(r"(?:^|\s)([A-Za-z][A-Za-z0-9_]*)=(\S+)", snapshot)) + if fields.get("JobId") != self.job_id: + raise ValueError("Slurm readiness response does not identify the owned allocation") + last_state = fields.get("JobState", "missing") + if last_state == "RUNNING": + nodelist = fields.get("NodeList") + if not nodelist or nodelist in {"(null)", "None"}: + raise ValueError("running Slurm allocation has no node list") + hosts = self._command(["scontrol", "show", "hostnames", nodelist], timeout=remaining()).stdout.split() + remaining() + if len(hosts) != self.node_count or len(set(hosts)) != self.node_count: + raise ValueError(f"FPM Slurm cell requires exactly {self.node_count} allocated nodes, got {hosts}") + self.hosts = hosts + return self.pods() + if last_state not in {"PENDING", "CONFIGURING", "SUSPENDED"}: + raise RuntimeError(f"Slurm allocation {self.job_id} cannot become ready from {last_state}") + time.sleep(min(1.0, remaining())) + + def pods(self, *, include_terminating: bool = True) -> list[str]: + del include_terminating + return [f"node{rank:04d}" for rank in range(len(self.hosts))] + + def stage(self, pods: list[str], files: list[Path]) -> None: + stage = self.cell_dir / "slurm-runtime" + stage.mkdir(exist_ok=True) + for path in files: + shutil.copy2(path, stage / path.name) + for unit in pods: + (self.cell_dir / "raw" / unit).mkdir(parents=True, exist_ok=True) + + def _exec(self, unit: str, command: list[str], *, timeout: int): + rank = self.pods().index(unit) + mounts = [ + *self.mounts, + f"{self.cell_dir / 'slurm-runtime'}:/tmp/fpm-bench", + f"{self.cell_dir / 'raw' / unit}:/results", + ] + if any("\n" in mount or "," in mount for mount in mounts): + raise ValueError("Slurm container mounts cannot contain newlines or commas") + # Each srun is a one-node step in the caller's allocation. The engine + # itself starts the node's TP/DP workers, exactly as in the Pod runtime. + return self._command( + [ + "srun", + f"--jobid={self.job_id}", + f"--job-name={self.step_name}", + "--nodes=1", + "--ntasks=1", + "--ntasks-per-node=1", + "--exclusive", + "--exact", + f"--nodelist={self.hosts[rank]}", + f"--gpus-per-node={self.gpus_per_node}", + f"--container-image={self.image}", + f"--container-mounts={','.join(mounts)}", + "--container-writable", + "--container-workdir=/tmp/fpm-bench", + "env", + f"FPM_NODE_RANK={rank}", + f"FPM_MASTER_ADDR={self.hosts[0]}", + *command, + ], + timeout=timeout, + ) + + def prepare_attempt(self, pods: list[str], *, cell_id: str, plan_sha256: str, attempt_id: str) -> None: + from .native_artifact import COLLECTOR_PROVENANCE_FILENAME + + payload = json.dumps( + { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + "cell_id": cell_id, + "plan_sha256": plan_sha256, + "attempt_id": attempt_id, + } + ) + script = ( + "import importlib.metadata,json,pathlib,sys; p=json.loads(sys.argv[1]); " + "p['runtime']={'backend':'vllm','backend_version':importlib.metadata.version('vllm')}; " + "pathlib.Path('/results',sys.argv[2]).write_text(json.dumps(p,sort_keys=True)+'\\n')" + ) + for unit in pods: + self._exec(unit, ["python3", "-c", script, payload, COLLECTOR_PROVENANCE_FILENAME], timeout=300) + + def execute(self, pods: list[str], timeout_seconds: int = 14400) -> None: + from .runner import CommandScope, _cancel_preserving_interrupt + + logs = self.cell_dir / "logs" + logs.mkdir(exist_ok=True) + + def run(unit: str) -> None: + try: + result = self._exec(unit, ["bash", "/tmp/fpm-bench/fpm_exec.sh"], timeout=timeout_seconds) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error: + for stream in ("stdout", "stderr"): + output = getattr(error, stream, None) or "" + if isinstance(output, bytes): + output = output.decode(errors="replace") + (logs / f"{unit}.{stream}.log").write_text(output) + raise + (logs / f"{unit}.stdout.log").write_text(result.stdout) + (logs / f"{unit}.stderr.log").write_text(result.stderr) + + scope = CommandScope() + pool = ThreadPoolExecutor(max_workers=len(pods)) + try: + futures = [pool.submit(scope.run, run, pod) for pod in pods] + for future in futures: + future.result() + except BaseException as error: + # Worker threads do not receive the main thread's interrupt. + # Stop their srun children before joining so the campaign can + # promptly salvage artifacts and clean up its owned steps. + _cancel_preserving_interrupt(scope, error) + raise + finally: + pool.shutdown(wait=True, cancel_futures=True) + + def _remote_result_manifest(self, unit: str) -> dict: + from .runner import _file_manifest + + return _file_manifest(self.cell_dir / "raw" / unit) + + def collect(self, pods: list[str], *, require_benchmark: bool = True) -> None: + if require_benchmark and not any( + list((self.cell_dir / "raw" / unit).glob(FPM_BENCHMARK_RESULT_GLOB)) for unit in pods + ): + raise RuntimeError("Slurm FPM result set is missing native benchmark artifacts") + + def cleanup(self) -> None: + # Never cancel the allocation or unrelated steps. This also handles + # resume after a collector process died while its named srun survived. + jobs = {self.job_id} + if self.owner_path.exists(): + owner = json.loads(self.owner_path.read_text()) + if owner.get("step_name") != self.step_name or not re.fullmatch(r"[0-9]+", owner.get("job_id", "")): + raise ValueError("Slurm FPM ownership receipt does not match this campaign cell") + jobs.add(owner["job_id"]) + + def owned_steps() -> list[str]: + # Listing the user's steps works even when an old allocation has + # expired, unlike squeue --jobs= on some Slurm versions. + result = self._command(["squeue", "--steps", "--me", "--noheader", "--format=%i|%j"]) + found = [] + for line in result.stdout.splitlines(): + step_id, _, name = line.strip().partition("|") + if name == self.step_name and any(re.fullmatch(re.escape(job) + r"\.[0-9]+", step_id) for job in jobs): + found.append(step_id) + return found + + for step_id in owned_steps(): + self._command(["scancel", step_id]) + deadline = time.monotonic() + 60 + while remaining := owned_steps(): + if time.monotonic() >= deadline: + raise RuntimeError(f"owned FPM Slurm steps remain after cleanup: {remaining}") + time.sleep(1) diff --git a/python/aisimulate/docs/fpm/deepseek-v41.md b/python/aisimulate/docs/fpm/deepseek-v41.md new file mode 100644 index 000000000..33677e7dd --- /dev/null +++ b/python/aisimulate/docs/fpm/deepseek-v41.md @@ -0,0 +1,174 @@ + + +# DeepSeek-V4.1 FPM collection + +The collection path supports the text backbone with pure TP4, native checkpoint +precision, GPU-resident Engram, eager execution, and DSpark disabled. +Historical [GB200 and GB300 experimental evidence](https://github.com/ai-dynamo/aisimulate/tree/dd1fa97add17d3d74f580f4c3e0566c1b5f11827/data/experimental/deepseek-v41) +is preserved separately from the product code. The GB200 calibration is not +admitted for serving prediction: ordinary-serving validation exposed a large +latency mismatch. Experimental systems overlays do not change curated defaults. + +## Execution identity + +Schema 7 appends four strings to the existing exact model/backend/topology +identity: `model_config_sha256`, `execution_profile`, `engram_residency`, and +`input_modality`. The config hash uses canonical JSON after the same quantization +normalization used by the SDK. V4.1 uses `hbm_tp_sharded` and `text`; its profile +is `full` or `decoder_bounded`. Changing the config or profile cannot borrow +another cell. Schema 6 loads with the legacy values `""`, `full`, `none`, `text`. +Publication upgrades a validated schema-6 pair in memory before emitting schema 7. + +The shared model retains the entire resident-weight inventory even when a replay +profile executes fewer decoder tokens. FPM interpolation uses the original +stage-aware SOL graph. Vision and speculative decoding remain outside this +campaign's measurement contract. + +When a native FPM table labels FMHA by its cache precision, select that table +with `fpm_fmha_dtype: "fp8"` in native engine/replay JSON, or +`ModelConfig(fpm_fmha_quant_mode=FMHAQuantMode.fp8, forward_model="fpm")`. +This option requires `forward_model="fpm"` and changes only the exact FPM cell +selector. The checkpoint's analytical attention graph, interpolation SOL +anchors, and memory inventory remain unchanged. `activation_dtype` retains its +existing arithmetic-override meaning; it is not a substitute for this selector. +The historical GB200 comparison uses the table selector with no activation +override, and its optional op-level SOL comparison uses checkpoint precision. + +`--fpm-decoder-replay` describes true bounded decoder execution. The current +vLLM route rejects it because the verified preview executes the full backbone. +Prefix-cache/SWA tail recomputation does not establish true decoder replay. +Replay OFF data must never be relabeled as ON data. + +For Decoder ON, the FPM v1 telemetry API rejects an iteration with multiple +prefill requests and fresh prefill tokens. Its prompt-length variance cannot +prove equal current extends when requests have different cached prefixes or +completed chunks. Single-prefill and decode-only telemetry remain supported; +explicit homogeneous static inputs retain their separate table-query path. +This admission limit also applies to whole-forward FPM engines before lookup. +Historical reports retain their original predictor identities and coverage; +their supported counts do not describe this stricter current admission rule. + +## Slurm transport + +`--fpm-executor slurm` runs inside a caller-owned Slurm allocation with Pyxis. +It requires an explicit `--fpm-slurm-container-image` and accepts repeated +`--fpm-slurm-container-mount SOURCE:TARGET` arguments. The frozen Generator +manifest determines the node count; the allocated count must agree. A shared +campaign directory must be visible at the same path on each node. + +The transport reuses Generator scripts, native result validation, attempt +identity, recovery, and formal publication. It launches one engine per node +with explicit ranks and a common master, and collects under stable `nodeNNNN` +execution-unit names. Cleanup cancels only receipted, campaign-specific steps +and verifies their exit; it never cancels the caller's allocation. Kubernetes +remains the default transport. + +## Qualification and publication gates + +Use a digest-pinned ARM64 image for GB200. The official preview recipe is at +[vLLM recipes commit ce19de141d448df3e739de977da0830e9a175de5](https://github.com/vllm-project/recipes/blob/ce19de141d448df3e739de977da0830e9a175de5/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml). +The model config revision is +[fb2764a5cf321eaa5070ca8f9e892818f477c16d](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/tree/fb2764a5cf321eaa5070ca8f9e892818f477c16d). +Image support for loading the model does not prove that its Dynamo instrumentation +satisfies the Collector's native FPM contract. The ordinary preflight must pass. + +Engram makes token history part of the workload. Qualification therefore needs +a reproducible tokenizer-generated text corpus and real model-computed KV for +both cached prefills and decode. V4.1 cached-prefill and decode publication reject `fake_fallback`, +legacy, and skipped-warm provenance; the consumer rejects them too. A producer +must earn the `real_kv` marker through execution. The initial canary should use +small batch/context limits, then expand the matrix only after native artifacts, +config identity, input provenance, and timing checks pass. + +Record the immutable checkpoint/image/instrumentation revisions, corpus hash, +resolved engine settings, native artifacts and measured coverage with the data. +Keep internal cluster names, account names, filesystem paths, and raw operational +logs in private campaign storage. + + +The pinned producer is packaged under +`collector/fpm_forward/runtime/dsv41/`. It executes seed chunks on the same +request and block tables as the measured step, retaining Engram history and +compressor ring state. Source hashes cover the installed ARM vLLM files and +Dynamo scheduler/point/FPM definitions. `ai-dynamo-runtime==1.4.2` has passed +an actual import check with Dynamo source `54960177085413259859c88bd34ed0734d4c2ea9`. +This import check does not establish GPU numerical or timing correctness. + +Its initial native grid is bounded to batch 2, per-request context 2048, total +scheduled prefill tokens 512, and global benchmark warmup 0. Use model limit +2050 to include the native decode context-2048 endpoint. Every completed point +archives actual prompt/output tokens in an adjacent JSONL file; the Collector +checks the file checksum, point coverage, per-request computed counts, and +completed seed witness. Unsupported or incomplete points fail the campaign. +The original mixed English/Chinese fixture is reproducible and is not a +representative production workload; Engram locality and routing sensitivity +remain limits on generalizing any resulting curve. + +## Pinned Hugging Face profiles + +FPM datasets are external artifacts. A checked-in pin manifest binds a Hugging +Face dataset commit to each system YAML, native parquet table, and metadata +sidecar by SHA256. Do not replace its immutable revision with `main` or a PR ref. +To materialize an overlay for the existing native reader: + +```python +from pathlib import Path +import aisimulate_core +from aisimulate_core.sdk.fpm_dataset import materialize_fpm_profile +from aisimulate_core.sdk.rust_engine_step import ForwardPassPerfModelConfig, RustForwardPassPerfModel + +pin = Path(aisimulate_core.__file__).parent / "systems/dsv41_fpm_hf.json" +systems_path = materialize_fpm_profile(pin, "gb300-tp4-full", allow_unqualified=True) +config = ForwardPassPerfModelConfig( + model="deepseek-ai/DeepSeek-V4.1-Flash", system="gb300", backend="sglang", + worker_type="aggregated", backend_version="0.0.0.dev0", # Historical identity. + systems_paths=(str(systems_path),), tp=4, moe_tp_size=4, moe_ep_size=1, + estimation_mode="fpm_interpolation", fpm_fmha_quant_mode="fp8", + enable_shared_layer=False, strict_provenance=True, +) +model = RustForwardPassPerfModel.best_available(config) +print(model.estimate_forward_pass_time_ms({ + "version": 1, "wall_time": 0.0, + "scheduled_requests": { + "num_prefill_requests": 0, "sum_prefill_tokens": 0, "sum_prefill_kv_tokens": 0, + "num_decode_requests": 1, "sum_decode_kv_tokens": 2048, + "var_prefill_length": 0.0, "var_decode_kv_tokens": 0.0, + }, +})) +``` + +The packaged `systems/dsv41_fpm_hf.json` pins the three principal #158 tables +(GB300 TP4 `full` / `decoder_bounded`, and quarantined GB200 TP4 `full`) to +[HF dataset commit b35883ee5f8b4a82e24844a7e872056aff64287f](https://huggingface.co/datasets/nvidia/aisimulate-fpm-dataset/tree/b35883ee5f8b4a82e24844a7e872056aff64287f). +The dataset change is reviewed in [HF PR #11](https://huggingface.co/datasets/nvidia/aisimulate-fpm-dataset/discussions/11). +Historical component tables and observations remain in that dataset revision; +they are not substituted for the principal prediction tables. + +The equivalent staging command, from this repository, prints the verified overlay: + +```sh +python -m aisimulate_core.sdk.fpm_dataset \ + python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json \ + gb300-tp4-full --allow-unqualified +``` + +Use `--local-files-only` (or the same Python keyword) after the first download to +run without network access. Every cache use rechecks hashes; corrupted files +fail instead of silently selecting another table. The native reader still owns +schema, real-KV provenance, execution-profile, model-config, and exact-cell +validation. This staging step does not alter interpolation or precision. + +Historical #158 profiles retain their original development or quarantined +status, including the archived backend identity. Reproducing these experiments +requires `--allow-unqualified` (Python: `allow_unqualified=True`), which emits a +warning and does not grant serving admission. In particular, moving GB200 data +to Hugging Face does not resolve its ordinary-serving latency mismatch. + +For restricted datasets, staging uses `HF_TOKEN` or the token saved by `hf auth +login` (`HF_HOME` / `HF_TOKEN_PATH` are respected). Credentials are never stored +in manifests or passed to redirected CDN requests. Set +`HF_HUB_DISABLE_IMPLICIT_TOKEN=1` to download anonymously. No additional Hub SDK +is required for this loader. diff --git a/python/aisimulate/pyproject.toml b/python/aisimulate/pyproject.toml index 897ca8dcd..9df74f1e3 100644 --- a/python/aisimulate/pyproject.toml +++ b/python/aisimulate/pyproject.toml @@ -127,6 +127,10 @@ include = [ { path = "collector/cases/**/*.yaml", format = "wheel" }, { path = "collector/fpm_forward/**/*.py", format = "wheel" }, { path = "collector/fpm_forward/runtime/fpm_exec.sh", format = "wheel" }, + { path = "collector/fpm_forward/runtime/fpm_text.txt", format = "wheel" }, + { path = "collector/fpm_forward/runtime/dsv41/*.json", format = "wheel" }, + { path = "collector/fpm_forward/runtime/dsv41/README.md", format = "wheel" }, + { path = "collector/fpm_forward/runtime/dsv41/LICENSE", format = "wheel" }, # Ship estimator models and profiles with their canonical implementation. "src/aisimulate_core/sdk/**/*.md", "src/aisimulate_core/systems/*.yaml", diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 8869bc457..8b6c83c4b 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -336,6 +336,9 @@ def _supported_estimator_policies(self): or self.estimation_mode != "auto" or self.fallback_policy != "deny" or bool(self.estimator_config) + or getattr(self, "decoder_replay", False) + or getattr(self, "enable_shared_layer", None) is not None + or getattr(self, "strict_provenance", None) is not None ) roles = [getattr(workers, role, None) for role in ("aggregated", "prefill", "decode")] unsupported_provider = "afd" in modes or getattr(workers, "encoder", None) is not None diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index e093292b4..337978a78 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -92,6 +92,7 @@ "gemm_dtype": ("gemm_dtype", "aic_gemm_dtype"), "moe_dtype": ("moe_dtype", "aic_moe_dtype"), "fmha_dtype": ("fmha_dtype", "aic_fmha_dtype"), + "fpm_fmha_dtype": ("fpm_fmha_dtype", "aic_fpm_fmha_dtype"), "kv_cache_dtype": ("kv_cache_dtype", "aic_kv_cache_dtype"), "comm_dtype": ("comm_dtype", "aic_comm_dtype"), "systems_path": ("systems_path",), @@ -247,18 +248,20 @@ def measure(self, spec: ReplaySpec) -> AFDCompanionTiming: if not isinstance(model_name, str) or not model_name or not isinstance(hardware, str) or not hardware: raise ValueError(f"{role} AFD companion requires aic_model_path and aic_system") fpm_parquet_path = timing_overrides.get("fpm_parquet_path") + has_fpm_selector = "fpm_fmha_dtype" in timing_overrides metric = "ttft" if role == "prefill" else "tpot" source = "aisimulate.legacy_cli.api.cli_estimate" try: - if fpm_parquet_path is not None: - # Preserve static integration through the canonical model API. + if fpm_parquet_path is not None or has_fpm_selector: + # The canonical API is the only path that preserves FPM selectors. if forward_model != "fpm": - raise ValueError("fpm_parquet_path requires forward_model='fpm'") + raise ValueError("FPM controls require forward_model='fpm'") quantization = {} for field, parameter in ( ("gemm_dtype", "gemm_quant_mode"), ("moe_dtype", "moe_quant_mode"), ("fmha_dtype", "fmha_quant_mode"), + ("fpm_fmha_dtype", "fpm_fmha_quant_mode"), ("kv_cache_dtype", "kvcache_quant_mode"), ("comm_dtype", "comm_quant_mode"), ): @@ -279,7 +282,11 @@ def measure(self, spec: ReplaySpec) -> AFDCompanionTiming: nextn=kwargs.get("nextn", 0), kv_block_size=args.get("block_size"), estimation_mode="fpm_interpolation", - estimator_config={"fpm_interpolation": {"fpm_parquet_path": fpm_parquet_path}}, + estimator_config=( + {"fpm_interpolation": {"fpm_parquet_path": fpm_parquet_path}} + if fpm_parquet_path is not None + else {} + ), systems_paths=(timing_overrides["systems_path"],) if "systems_path" in timing_overrides else (), **quantization, **{ @@ -329,6 +336,7 @@ def measure(self, spec: ReplaySpec) -> AFDCompanionTiming: "backend_version": deployment.backend_version, "forward_model": forward_model, **({"fpm_parquet_path": fpm_parquet_path} if fpm_parquet_path is not None else {}), + **({"fpm_fmha_dtype": timing_overrides["fpm_fmha_dtype"]} if has_fpm_selector else {}), "metric": metric, }, ) diff --git a/python/aisimulate/src/aisimulate/sdk/fpm_dataset.py b/python/aisimulate/src/aisimulate/sdk/fpm_dataset.py new file mode 100644 index 000000000..0810d0477 --- /dev/null +++ b/python/aisimulate/src/aisimulate/sdk/fpm_dataset.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_dataset.""" + +from aisimulate.sdk._compat import alias_module as _alias_module + +_alias_module(__name__, "aisimulate_core.sdk.fpm_dataset") diff --git a/python/aisimulate/src/aisimulate/sdk/fpm_identity.py b/python/aisimulate/src/aisimulate/sdk/fpm_identity.py new file mode 100644 index 000000000..87184b074 --- /dev/null +++ b/python/aisimulate/src/aisimulate/sdk/fpm_identity.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_identity.""" + +from aisimulate.sdk._compat import alias_module as _alias_module + +_alias_module(__name__, "aisimulate_core.sdk.fpm_identity") diff --git a/python/aisimulate/src/aisimulate_core/sdk/config.py b/python/aisimulate/src/aisimulate_core/sdk/config.py index 72d01e9c6..5a64910d4 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config.py @@ -132,6 +132,10 @@ class ModelConfig: system: str | None = None # DeepSeek-V4.1 text AR: use the backend-verified bounded decoder replay profile. decoder_replay: bool = False + # Whole-forward database identity only; never changes arithmetic or memory. + # Selection emits a warning with the original model mode and matched cell IDs. + # Exact table-label matching is not independent runtime-precision proof. + fpm_fmha_quant_mode: common.FMHAQuantMode | None = field(default=None, kw_only=True) def __post_init__(self) -> None: self.moe_backend = normalize_kernel_backend(self.moe_backend, common.MoEBackend, "moe_backend") diff --git a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py index ea413f727..3fa3629e0 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py @@ -43,6 +43,8 @@ def build_model_config( moe_backend: str | None = None, enable_eplb: bool = False, wideep_num_slots: int | None = None, + *, + fpm_fmha_quant_mode: str | None = None, ) -> ModelConfig: """Build a ModelConfig with optional quant mode overrides.""" validate_moe_controls(enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots) @@ -55,6 +57,7 @@ def build_model_config( gemm_quant_mode=GEMMQuantMode[gemm_quant_mode] if gemm_quant_mode else None, kvcache_quant_mode=KVCacheQuantMode[kvcache_quant_mode] if kvcache_quant_mode else None, fmha_quant_mode=FMHAQuantMode[fmha_quant_mode] if fmha_quant_mode else None, + fpm_fmha_quant_mode=FMHAQuantMode[fpm_fmha_quant_mode] if fpm_fmha_quant_mode else None, moe_quant_mode=MoEQuantMode[moe_quant_mode] if moe_quant_mode else None, comm_quant_mode=CommQuantMode[comm_quant_mode] if comm_quant_mode else None, forward_model=forward_model or "op_level", diff --git a/python/aisimulate/src/aisimulate_core/sdk/deepseek_v41.py b/python/aisimulate/src/aisimulate_core/sdk/deepseek_v41.py index 7435edbd9..973e91bcb 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/deepseek_v41.py +++ b/python/aisimulate/src/aisimulate_core/sdk/deepseek_v41.py @@ -105,6 +105,8 @@ def from_text_config(cls, text_config: dict) -> DeepSeekV41Config: return result def validate(self) -> None: + if self.sliding_window <= 0: + raise ValueError("DeepSeek-V4.1 sliding_window must be positive") layers = self.num_hidden_layers if len(self.compress_ratios) != layers or set(self.compress_ratios) - {0, 1, 2}: raise ValueError("DeepSeek-V4.1 requires one backbone compression ratio (0, 1, or 2) per layer") diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index fd296951b..c325434a5 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -162,6 +162,7 @@ def _fpm_spec_dict(op: FPMForwardOp) -> dict: "phase": op._phase, "model_path": op._model_path, "match_identity": list(op._match_identity), + "original_fmha_quant_mode": op._original_fmha_quant_mode, "weight_bytes": op._weight_bytes, # Speculative verify width for the equivalent-AR decode mapping # (1 = plain AR). Set by the fpm hybrid rewrite in models when a @@ -322,6 +323,7 @@ def _engine_config_dict( "backend_version": _literal_backend_version(system, backend, backend_version, systems_path, database), "fpm_parquet_path": fpm_parquet_path, "kv_block_size": kv_block_size, + "forward_model": getattr(model, "forward_model", getattr(cfg, "forward_model", None)), "decoder_replay": bool(getattr(cfg, "decoder_replay", False)), # ParallelMapping (flattened) "tp_size": int(cfg.tp_size or 1), @@ -334,6 +336,7 @@ def _engine_config_dict( "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)), "activation_dtype": _rust_quant_to_dtype(getattr(cfg, "fmha_quant_mode", None)), + "fpm_fmha_dtype": _rust_quant_to_dtype(getattr(cfg, "fpm_fmha_quant_mode", None)), "kv_cache_dtype": _rust_quant_to_dtype(getattr(cfg, "kvcache_quant_mode", None)), # Shared-layer policy bits only (schema v13): the engine resolves # per-op sources itself (`perf_database/source_resolution.rs`), so the @@ -419,6 +422,7 @@ def compile_engine( moe_quant_mode: str | None = None, kvcache_quant_mode: str | None = None, fmha_quant_mode: str | None = None, + fpm_fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None, attention_backend: str | None = None, moe_backend: str | None = None, @@ -473,6 +477,7 @@ def compile_engine( gemm_quant_mode=gemm_quant_mode, kvcache_quant_mode=kvcache_quant_mode, fmha_quant_mode=fmha_quant_mode, + fpm_fmha_quant_mode=fpm_fmha_quant_mode, moe_quant_mode=moe_quant_mode, comm_quant_mode=comm_quant_mode, forward_model=forward_model, diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_dataset.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_dataset.py new file mode 100644 index 000000000..1fc82132e --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_dataset.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Materialize immutable Hugging Face FPM profiles for the native systems loader. + +Storage verification does not replace native schema, provenance or exact-cell +validation. Downloads happen only on this explicit API/CLI call, never during an +ordinary op-level query. Unqualified profiles require an explicit reproduction-only opt-in. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import tempfile +import warnings +from pathlib import Path, PurePosixPath +from urllib.parse import quote, urlsplit +from urllib.request import HTTPRedirectHandler, Request, build_opener + + +class _HubRedirectHandler(HTTPRedirectHandler): + def redirect_request(self, request, fp, code, msg, headers, newurl): + redirected = super().redirect_request(request, fp, code, msg, headers, newurl) + old, new = urlsplit(request.full_url), urlsplit(newurl) + same_https_origin = ( + old.scheme == new.scheme == "https" + and old.hostname == new.hostname + and (old.port or 443) == (new.port or 443) + ) + authorization = request.get_header("Authorization") + if redirected is not None and same_https_origin and authorization: + redirected.add_unredirected_header("Authorization", authorization) + return redirected + + +urlopen = build_opener(_HubRedirectHandler()).open + + +def _unique_object(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate manifest key: {key}") + result[key] = value + return result + + +def _relative_path(value: str) -> str: + if not isinstance(value, str) or not value or "\\" in value: + raise ValueError("dataset paths must be nonempty relative POSIX paths") + path = PurePosixPath(value) + if path.is_absolute() or any(part in ("", ".", "..") for part in value.split("/")): + raise ValueError(f"unsafe dataset path: {value!r}") + return value + + +def _download_request(url: str) -> Request: + """Use HF login credentials without forwarding them to signed CDN redirects.""" + request = Request(url) + if os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN", "").upper() in ("1", "ON", "YES", "TRUE"): + return request + token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + if not token: + default_home = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "huggingface" + hf_home = Path(os.environ.get("HF_HOME", default_home)) + token_path = Path(os.environ.get("HF_TOKEN_PATH", hf_home / "token")) + try: + token = token_path.read_text().strip() + except FileNotFoundError: + token = None + if token: + # Preserve auth on Hub cache redirects only. Signed CDN redirects do not + # need our credential and the handler never copies it across origins. + request.add_unredirected_header("Authorization", f"Bearer {token.strip()}") + return request + + +def _verify(directory: Path, files: list[dict]) -> None: + for item in files: + path = directory / item["target"] + if path.is_symlink() or not path.is_file() or not path.resolve().is_relative_to(directory.resolve()): + raise ValueError(f"missing or unsafe cached FPM file: {item['target']}") + if hashlib.sha256(path.read_bytes()).hexdigest() != item["sha256"]: + raise ValueError(f"FPM SHA256 mismatch: {item['target']}") + + +def materialize_fpm_profile( + manifest_path: str | Path, + profile: str, + *, + cache_dir: str | Path | None = None, + local_files_only: bool = False, + allow_unqualified: bool = False, +) -> Path: + """Return a verified ``systems_path`` pinned by a checked-in manifest. + + The manifest has ``format_version: 1``, ``repo_id``, a 40-character commit + ``revision`` and ``profiles``. Each admitted profile contains ``admission: + serving`` and ``files`` entries with HF ``path``, relative overlay ``target`` + and ``sha256``. Branches, tags, diagnostic profiles and corrupt caches fail + closed; offline mode never attempts a network request. ``allow_unqualified`` + permits historical diagnostic reproduction, without granting serving admission. + """ + manifest = json.loads(Path(manifest_path).read_text(), object_pairs_hook=_unique_object) + if manifest.get("format_version") != 1: + raise ValueError("unsupported FPM dataset manifest format") + repo = manifest.get("repo_id", "") + revision = manifest.get("revision", "") + if not isinstance(repo, str) or not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", repo): + raise ValueError("invalid Hugging Face dataset repository") + if not isinstance(revision, str) or not re.fullmatch(r"[0-9a-f]{40}", revision): + raise ValueError("FPM dataset revision must be an immutable commit SHA") + entry = manifest["profiles"][profile] + admission = entry.get("admission") + if admission not in ("serving", "development", "quarantined"): + raise ValueError(f"unknown FPM admission: {admission!r}") + if admission != "serving": + if not allow_unqualified: + raise ValueError(f"FPM profile {profile!r} is not admitted for serving") + warnings.warn( + f"FPM profile {profile!r} is {admission}: historical reproduction only, not serving-admitted", + UserWarning, + stacklevel=2, + ) + files = entry["files"] + if not isinstance(files, list) or not files: + raise ValueError("FPM profile must contain files") + targets = set() + for item in files: + _relative_path(item["path"]) + target = _relative_path(item["target"]) + if target in targets: + raise ValueError(f"duplicate FPM target: {target}") + targets.add(target) + if not isinstance(item["sha256"], str) or not re.fullmatch(r"[0-9a-f]{64}", item["sha256"]): + raise ValueError(f"invalid SHA256 for {target}") + parquets = {target for target in targets if target.endswith("/fpm_forward_perf.parquet")} + if not parquets or not any(target.endswith(".yaml") for target in targets): + raise ValueError("FPM profile requires a system YAML and a parquet table") + for target in parquets: + if target.removesuffix(".parquet") + ".metadata.json" not in targets: + raise ValueError(f"missing metadata pair for {target}") + identity = json.dumps({"repo": repo, "revision": revision, "profile": entry}, sort_keys=True).encode() + digest = hashlib.sha256(identity).hexdigest() + root = ( + Path(cache_dir) + if cache_dir is not None + else Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "aisimulate" / "fpm" + ) + destination = root / digest + if destination.exists() or destination.is_symlink(): + if destination.is_symlink(): + raise ValueError("FPM cache directory must not be a symbolic link") + _verify(destination, files) + return destination + if local_files_only: + raise FileNotFoundError(f"pinned FPM profile is not cached: {profile}") + root.mkdir(parents=True, exist_ok=True) + temporary = Path(tempfile.mkdtemp(prefix=".download-", dir=root)) + try: + for item in files: + output = temporary / item["target"] + output.parent.mkdir(parents=True, exist_ok=True) + url = f"https://huggingface.co/datasets/{repo}/resolve/{revision}/{quote(item['path'], safe='/')}" + with urlopen(_download_request(url), timeout=120) as response, output.open("wb") as stream: + shutil.copyfileobj(response, stream) + _verify(temporary, files) + try: + temporary.rename(destination) + except OSError: + # A concurrent reader may have completed the same immutable profile. + if not destination.is_dir() or destination.is_symlink(): + raise + _verify(destination, files) + return destination + finally: + if temporary.exists(): + shutil.rmtree(temporary) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("manifest", type=Path) + parser.add_argument("profile") + parser.add_argument("--cache-dir", type=Path) + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument( + "--allow-unqualified", + action="store_true", + help="historical reproduction only; does not grant serving admission", + ) + args = parser.parse_args() + print( + materialize_fpm_profile( + args.manifest, + args.profile, + cache_dir=args.cache_dir, + local_files_only=args.local_files_only, + allow_unqualified=args.allow_unqualified, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py b/python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py new file mode 100644 index 000000000..a77e6282b --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared whole-forward execution identity; schema 6 has the legacy defaults.""" + +from __future__ import annotations + +import copy +import hashlib +import json + +EXECUTION_COLUMNS = ("model_config_sha256", "execution_profile", "engram_residency", "input_modality") +LEGACY_EXECUTION_IDENTITY = ("", "full", "none", "text") + + +def execution_identity( + raw_config: dict, + *, + decoder_replay: bool = False, + backend: str = "vllm", + engram_cpu_offload: bool | None = None, + input_modality: str | None = None, +) -> tuple[str, ...]: + """Bind V4.1 curves to config and execution; leave existing model keys intact. + + Both SDK and Collector normalize inferred quantization fields before hashing, + so loading the same config via HF, a local directory or AIC cache agrees. + V4.1 callers must supply the supported Engram/text execution facts explicitly. + The SDK supplies its prediction contract; producers must pass their checked + runtime configuration. A checkpoint hash does not prove runtime residency. + """ + architectures = raw_config.get("architectures") or [] + if "DeepseekV41ForCausalLM" not in architectures: + if decoder_replay: + raise ValueError("FPM decoder replay requires a DeepSeek-V4.1 model") + return LEGACY_EXECUTION_IDENTITY + if engram_cpu_offload is not False: + raise ValueError("DeepSeek-V4.1 FPM requires explicit engram_cpu_offload=False") + if input_modality != "text": + raise ValueError("DeepSeek-V4.1 FPM requires explicit input_modality='text'") + from .deepseek_v41 import resolve_execution_profile + from .utils import _attach_inferred_quant_fields + + payload = _attach_inferred_quant_fields(copy.deepcopy(raw_config)) + digest = hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode() + ).hexdigest() + return (digest, resolve_execution_profile(decoder_replay, backend).value, "hbm_tp_sharded", "text") diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py index b5bff6580..5d3c162a8 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py @@ -64,7 +64,7 @@ _FORWARD_MODELS = ("op_level", "fpm") -def _apply_forward_model_fpm(model: BaseModel) -> BaseModel: +def _apply_forward_model_fpm(model: BaseModel, backend_name: str = "vllm") -> BaseModel: """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.""" @@ -119,6 +119,19 @@ def _apply_forward_model_fpm(model: BaseModel) -> BaseModel: decode_op._verify_width = int(model.verify_width) model.context_ops = [prefill_op, *draft_context_ops] model.generation_ops = [decode_op, *draft_generation_ops] + from aisimulate_core.sdk.fpm_identity import execution_identity + + identity = execution_identity( + getattr(model, "raw_config", {}), + decoder_replay=getattr(model.config, "decoder_replay", False), + backend=backend_name, + # The SDK supports this prediction contract; the producer separately + # verifies actual runtime residency and token-only requests. + engram_cpu_offload=False, + input_modality="text", + ) + for op in (prefill_op, decode_op): + op._match_identity = (*op._match_identity[:15], *identity) model.forward_model = "fpm" return model @@ -145,6 +158,8 @@ def get_model( raise InvalidEngineConfigurationError( f"Unknown forward_model: {forward_model!r}. Valid values: {', '.join(_FORWARD_MODELS)}" ) + if getattr(model_config, "fpm_fmha_quant_mode", None) is not None and forward_model != "fpm": + raise InvalidEngineConfigurationError("fpm_fmha_quant_mode requires forward_model='fpm'") # Shallow-copy so mutations below don't poison the @cache'd original. model_info = dict(_get_model_info(model_path)) @@ -220,7 +235,7 @@ def get_model( model.spec_scheme.validate(model, backend_name) materialize_spec_scheme(model) if forward_model == "fpm": - model = _apply_forward_model_fpm(model) + model = _apply_forward_model_fpm(model, backend_name) return model diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py index da38f4dfe..8c090cd0f 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py @@ -645,7 +645,8 @@ def _infer_quant_modes_from_raw_config(raw_config: dict, architecture: str | Non if ( architecture == "DeepseekV41ForCausalLM" - and str(raw_config.get("quantization_config", {}).get("expert_dtype", "")).lower() == "fp4" + and str((_get_language_quantization_config(raw_config) or {}).get("expert_dtype", "")).lower() == "fp4" + and overrides.get("moe_quant_mode") != common.MoEQuantMode.nvfp4 ): overrides["moe_quant_mode"] = common.MoEQuantMode.w4a8_mxfp4_mxfp8 @@ -804,7 +805,9 @@ def _is_dsv4_fp4_expert_model(model_path: str) -> bool: return False raw_config = info.get("raw_config", {}) expert_config = ( - raw_config.get("quantization_config", {}) if architecture == "DeepseekV41ForCausalLM" else raw_config + (_get_language_quantization_config(raw_config) or {}) + if architecture == "DeepseekV41ForCausalLM" + else raw_config ) if str(expert_config.get("expert_dtype") or "").lower() != "fp4": return False 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 728f1f7dc..232634f3a 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py +++ b/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py @@ -20,6 +20,13 @@ systems/data////fpm_forward_perf.parquet systems/data////fpm_forward_perf.metadata.json +An explicit ``fpm_fmha_quant_mode`` selects a recorded table label while +preserving the model's arithmetic and memory modes. The compiled selector emits +a WARNING containing those modes and all matched ``cell_ids`` once per loaded +cell/model mode. A different recorded label remains an exact-match miss. Neither +the selector nor that comparison independently proves the runtime's resolved +attention precision; an engine-derived precision contract remains separate. + (The former Python-side query/loader machinery — the per-call ``query()`` family, the parquet/sidecar validators, and the per-op ``DatabaseMode.SOL`` roofline closure — was retired with the Python engine-step path; the Rust @@ -30,6 +37,7 @@ from enum import Enum +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY from aisimulate_core.sdk.operations.base import PythonOperation _PHASES = ("prefill", "decode") @@ -39,7 +47,7 @@ # handled separately (exact-match, never borrowed); ``weight_quantization`` # is redundant with ``gemm_quant_mode`` (the collector falls one back to the # other) so only ``gemm_quant_mode`` participates in matching. The Rust -# loader's cell keying mirrors this order and arity (15). +# loader's cell keying mirrors this order and arity (19), with schema-6 default execution identity. _CELL_MATCH_COLUMNS = ( "gemm_quant_mode", "moe_quant_mode", @@ -61,6 +69,7 @@ "attention_backend", "enable_wideep", "enable_eplb", + *EXECUTION_COLUMNS, ) @@ -126,10 +135,12 @@ def __init__( self._phase = phase self._model_path = str(model_path) self._weight_bytes = float(weight_bytes) + fmha_selector = getattr(model_config, "fpm_fmha_quant_mode", None) + self._original_fmha_quant_mode = None if fmha_selector is None else _norm_identity(model_config.fmha_quant_mode) self._match_identity = ( _norm_identity(model_config.gemm_quant_mode), _norm_identity(model_config.moe_quant_mode), - _norm_identity(model_config.fmha_quant_mode), + _norm_identity(model_config.fmha_quant_mode if fmha_selector is None else fmha_selector), _norm_identity(model_config.comm_quant_mode), _norm_identity(model_config.kvcache_quant_mode), _norm_identity(model_config.tp_size), @@ -145,6 +156,7 @@ def __init__( _norm_identity(bool(getattr(model_config, "enable_wideep", False))), _norm_identity(bool(getattr(model_config, "enable_eplb", False))), ) + self._match_identity += LEGACY_EXECUTION_IDENTITY 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 d0426533f..7480fb056 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -116,6 +116,7 @@ class ForwardPassPerfModelConfig: gemm_quant_mode: str | None = None moe_quant_mode: str | None = None fmha_quant_mode: str | None = None + fpm_fmha_quant_mode: str | None = dataclass_field(default=None, kw_only=True) kvcache_quant_mode: str | None = None comm_quant_mode: str | None = None nextn: int = 0 @@ -1183,6 +1184,7 @@ def _engine_config_json(model: Any, database: Any) -> str: "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)), + "fpm_fmha_dtype": _quant_to_dtype(getattr(model_config, "fpm_fmha_quant_mode", None)), "kv_cache_dtype": _quant_to_dtype(getattr(model_config, "kvcache_quant_mode", None)), "kv_block_size": None, "nextn": int(nextn) if nextn is not None else None, diff --git a/python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json b/python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json new file mode 100644 index 000000000..514ce42f5 --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json @@ -0,0 +1,181 @@ +{ + "format_version": 1, + "repo_id": "nvidia/aisimulate-fpm-dataset", + "revision": "b35883ee5f8b4a82e24844a7e872056aff64287f", + "profiles": { + "gb200-tp4-full": { + "admission": "quarantined", + "identity": { + "attention_backend": "auto", + "backend": "vllm", + "backend_version": "0.1.dev20904+g179dd0fa9", + "comm_quant_mode": "half", + "cp": 1, + "dp": 1, + "enable_eplb": false, + "enable_wideep": false, + "engram_residency": "hbm_tp_sharded", + "execution_profile": "full", + "fmha_quant_mode": "fp8", + "fmha_resolution": "checkpoint_native", + "gemm_quant_mode": "fp8_block", + "input_modality": "text", + "kv_cache_dtype": "fp8", + "model_config_sha256": "22c3912140adeb60ecbd3c7a9b54e62997adf65ad4c69440f8082cf29f5ef0d4", + "model_path": "deepseek-ai/DeepSeek-V4.1-Flash", + "moe_backend": "auto", + "moe_ep": 1, + "moe_quant_mode": "w4a8_mxfp4_mxfp8", + "moe_tp": 4, + "parallel_strategy": "pure_tp", + "pp": 1, + "system": "gb200", + "tp": 4, + "weight_quantization": "fp8_block" + }, + "files": [ + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb200/vllm/0.1.dev20904-g179dd0fa9/pure-tp4/aisim-dd1fa97add17-pr158/fpm/fpm_forward_perf.full.2edc7f78efe5.parquet", + "target": "data/gb200/vllm/0.1.dev20904+g179dd0fa9/fpm_forward_perf.parquet", + "sha256": "2edc7f78efe54b58009151f77591d2ed2dd42246f0e8a228ddebe5df1b6a816e" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb200/vllm/0.1.dev20904-g179dd0fa9/pure-tp4/aisim-dd1fa97add17-pr158/fpm/fpm_forward_perf.full.2edc7f78efe5.metadata.json", + "target": "data/gb200/vllm/0.1.dev20904+g179dd0fa9/fpm_forward_perf.metadata.json", + "sha256": "2d80a8511cc87a52c9d1fc56cb921a38e726480f87c11fbbe5b6f28c1cf374d9" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb200/vllm/0.1.dev20904-g179dd0fa9/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/LICENSE", + "target": "provenance/LICENSE", + "sha256": "e7c584183cb18a44cd56f453e642164d388c4dcefb49aa6dbd0b03ea40a62878" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb200/vllm/0.1.dev20904-g179dd0fa9/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/THIRD_PARTY_NOTICES.md", + "target": "provenance/THIRD_PARTY_NOTICES.md", + "sha256": "237cf18f406a8ac62e824cda4ddf03072bd18a8f0720cb5a87fbd17cc7b15912" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb200/vllm/0.1.dev20904-g179dd0fa9/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/systems.full.yaml", + "target": "gb200.yaml", + "sha256": "257db66d91458c0a6a9fcccbf06815e1b70db76b0c31b54aa54d5dcc9138c95e" + } + ] + }, + "gb300-tp4-full": { + "admission": "development", + "identity": { + "attention_backend": "auto", + "backend": "sglang", + "backend_version": "0.0.0.dev0", + "comm_quant_mode": "half", + "cp": 1, + "dp": 1, + "enable_eplb": false, + "enable_wideep": false, + "engram_residency": "hbm_tp_sharded", + "execution_profile": "full", + "fmha_quant_mode": "fp8", + "fmha_resolution": "checkpoint_native", + "gemm_quant_mode": "fp8_block", + "input_modality": "text", + "kv_cache_dtype": "fp8", + "model_config_sha256": "22c3912140adeb60ecbd3c7a9b54e62997adf65ad4c69440f8082cf29f5ef0d4", + "model_path": "deepseek-ai/DeepSeek-V4.1-Flash", + "moe_backend": "auto", + "moe_ep": 1, + "moe_quant_mode": "w4a8_mxfp4_mxfp8_trtllm", + "moe_tp": 4, + "parallel_strategy": "pure_tp", + "pp": 1, + "system": "gb300", + "tp": 4, + "weight_quantization": "fp8_block" + }, + "files": [ + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/fpm_forward_perf.full.1ba1736776ee.parquet", + "target": "data/gb300/sglang/0.0.0.dev0/fpm_forward_perf.parquet", + "sha256": "1ba1736776ee06679d63a4d888d6548ee94892e45277e1443fcbc308078aa65a" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/fpm_forward_perf.full.1ba1736776ee.metadata.json", + "target": "data/gb300/sglang/0.0.0.dev0/fpm_forward_perf.metadata.json", + "sha256": "a67af1008f490ac02aa56a43adf5c4fbba8ea63a6d29878dacfee474fa18e53c" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/LICENSE", + "target": "provenance/LICENSE", + "sha256": "e7c584183cb18a44cd56f453e642164d388c4dcefb49aa6dbd0b03ea40a62878" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/THIRD_PARTY_NOTICES.md", + "target": "provenance/THIRD_PARTY_NOTICES.md", + "sha256": "237cf18f406a8ac62e824cda4ddf03072bd18a8f0720cb5a87fbd17cc7b15912" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/systems.full.yaml", + "target": "gb300.yaml", + "sha256": "e0a7041d7637068073aa0ec8c08f4fc09a24a45d24e477af4f3f08870fcdfa11" + } + ] + }, + "gb300-tp4-decoder_bounded": { + "admission": "development", + "identity": { + "attention_backend": "auto", + "backend": "sglang", + "backend_version": "0.0.0.dev0", + "comm_quant_mode": "half", + "cp": 1, + "dp": 1, + "enable_eplb": false, + "enable_wideep": false, + "engram_residency": "hbm_tp_sharded", + "execution_profile": "decoder_bounded", + "fmha_quant_mode": "fp8", + "fmha_resolution": "checkpoint_native", + "gemm_quant_mode": "fp8_block", + "input_modality": "text", + "kv_cache_dtype": "fp8", + "model_config_sha256": "22c3912140adeb60ecbd3c7a9b54e62997adf65ad4c69440f8082cf29f5ef0d4", + "model_path": "deepseek-ai/DeepSeek-V4.1-Flash", + "moe_backend": "auto", + "moe_ep": 1, + "moe_quant_mode": "w4a8_mxfp4_mxfp8_trtllm", + "moe_tp": 4, + "parallel_strategy": "pure_tp", + "pp": 1, + "system": "gb300", + "tp": 4, + "weight_quantization": "fp8_block" + }, + "files": [ + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/fpm_forward_perf.decoder_bounded.0da08ce08b15.parquet", + "target": "data/gb300/sglang/0.0.0.dev0/fpm_forward_perf.parquet", + "sha256": "0da08ce08b15c91f9f4bb27d61a7a648d1b0042256b3afdfd8963350600f643a" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/fpm_forward_perf.decoder_bounded.0da08ce08b15.metadata.json", + "target": "data/gb300/sglang/0.0.0.dev0/fpm_forward_perf.metadata.json", + "sha256": "1ecdd4980b6767c5eecb1f679e80a0c30f1144df7511361db2720f1e9bc9cb67" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/LICENSE", + "target": "provenance/LICENSE", + "sha256": "e7c584183cb18a44cd56f453e642164d388c4dcefb49aa6dbd0b03ea40a62878" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/THIRD_PARTY_NOTICES.md", + "target": "provenance/THIRD_PARTY_NOTICES.md", + "sha256": "237cf18f406a8ac62e824cda4ddf03072bd18a8f0720cb5a87fbd17cc7b15912" + }, + { + "path": "data/history/deepseek-ai--DeepSeek-V4.1-Flash/gb300/sglang/0.0.0.dev0/pure-tp4/aisim-dd1fa97add17-pr158/fpm/provenance/systems.decoder_bounded.yaml", + "target": "gb300.yaml", + "sha256": "e0a7041d7637068073aa0ec8c08f4fc09a24a45d24e477af4f3f08870fcdfa11" + } + ] + } + } +} 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 0ec0385a2..a921e8802 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -94,6 +94,7 @@ def test_stable_function_signatures() -> 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, " + "fpm_fmha_quant_mode: 'str | None' = None, " "comm_quant_mode: 'str | None' = None, attention_backend: 'str | None' = None, " "moe_backend: 'str | None' = None, enable_eplb: 'bool' = False, wideep_num_slots: 'int | None' = None, " "nextn: 'int' = 0, " diff --git a/python/aisimulate/tests/cross_package/test_import_contract.py b/python/aisimulate/tests/cross_package/test_import_contract.py index 56b2f16d2..a94643a5a 100644 --- a/python/aisimulate/tests/cross_package/test_import_contract.py +++ b/python/aisimulate/tests/cross_package/test_import_contract.py @@ -31,6 +31,8 @@ "engine", "engine_table_view", "errors", + "fpm_dataset", + "fpm_identity", "inference_summary", "memory", "models.base", diff --git a/python/aisimulate/tests/e2e/cli/test_cli_build_default.py b/python/aisimulate/tests/e2e/cli/test_cli_build_default.py index 34a2cfd0f..3975f4967 100644 --- a/python/aisimulate/tests/e2e/cli/test_cli_build_default.py +++ b/python/aisimulate/tests/e2e/cli/test_cli_build_default.py @@ -227,5 +227,7 @@ def test_v41_sol_default_accepts_automatic_moe_parallelism(total_gpus): assert "Model: deepseek-ai/DeepSeek-V4.1-Flash (is_moe: True)" in output assert f"Total GPUs: {total_gpus}" in output assert "tp4pp1dp1etp1ep4" in output + # Rendered GPU accounting is total (used = replicas x GPUs per replica). + assert f"{total_gpus} ({total_gpus}={total_gpus // 4}x4)" in output assert "Parallelism width mismatch" not in output assert "No successful experiment runs" not in output diff --git a/python/aisimulate/tests/unit/collector/fixtures/dsv41_producer_lifecycle.py b/python/aisimulate/tests/unit/collector/fixtures/dsv41_producer_lifecycle.py new file mode 100644 index 000000000..879c5e24e --- /dev/null +++ b/python/aisimulate/tests/unit/collector/fixtures/dsv41_producer_lifecycle.py @@ -0,0 +1,530 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU state-machine contract tests; live pinned-image preflight remains required.""" + +import ast +import importlib.util +import json +import os +import re +import sys +import tempfile +import types +import unittest +from collections import deque +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock, patch + +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS + + +def module(name, **values): + obj = types.ModuleType(name) + obj.__dict__.update(values) + sys.modules[name] = obj + return obj + + +class Output: + @classmethod + def make_empty(cls): + return SimpleNamespace( + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], new_block_ids=[], num_computed_tokens=[], num_output_tokens=[], all_token_ids={} + ), + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + finished_req_ids=set(), + ) + + +class NewData: + @classmethod + def from_request(cls, request, block_ids, **kwargs): + return SimpleNamespace( + req_id=request.request_id, + prompt_token_ids=request.prompt_token_ids[:], + num_computed_tokens=request.num_computed_tokens, + block_ids=block_ids, + ) + + +class Request: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + self._all_token_ids = self.prompt_token_ids[:] + self.num_computed_tokens = 0 + self.num_output_tokens = 0 + self.num_output_placeholders = 0 + + +class Base: + def _bench_decode_context_lengths(self, total, batch): + assert total % batch == 0 + return [total // batch] * batch + + def _bench_prefill_kv_read_lengths(self, total, batch, *args): + return [total // batch] * batch + + _bench_prefill_new_token_lengths = _bench_prefill_kv_read_lengths + + def _bench_stop_at_timeout_boundary(self, point_type): + return False + + def _bench_pop_next(self, point_type): + return self._bench_grid.popleft() if self._bench_grid else None + + def _bench_save_current_point(self): + self.saved.append((self._bench_current_point, self._bench_current_fpms[:])) + self._bench_current_point = None + self._bench_current_fpms = [] + + def _bench_cleanup_requests(self): + self.requests.clear() + self.running.clear() + self._bench_active_req_ids.clear() + + def _update_from_output(self, output, model_output): + if self._real_callback_stage in {"admission", "measure"}: + self._bench_current_fpms.append(self._real_callback_stage) + + def _bench_should_record_scheduled(self, scheduled): + return True + + def _bench_write_results(self): + payload = { + "schema_version": 2, + "artifact_type": "rank", + "status": "complete", + "valid": True, + "results": [{"point": {"benchmark_id": 1}, "fpms": []}], + "iteration_groups": [{"benchmark_id": 1}], + } + Path(self._bench_config.output_path).write_text(json.dumps(payload)) + + def _bench_build_explicit_grid(self, points, *, generated=False): + self._bench_grid = deque(points) + + +for name in ["dynamo", "dynamo.vllm", "vllm", "vllm.v1", "vllm.v1.core", "vllm.v1.core.sched"]: + module(name) +module( + "dynamo.vllm.instrumented_scheduler", + InstrumentedScheduler=Base, + EAGER_WARMUP_REASON="eager_warmup", + _BenchPhase=SimpleNamespace(DECODE_SWEEP="decode", DONE="done"), +) +module("vllm.sampling_params", SamplingParams=lambda **kw: SimpleNamespace(**kw)) +module("vllm.tokenizers", get_tokenizer=None) +module("vllm.v1.core.sched.output", CachedRequestData=None, NewRequestData=NewData, SchedulerOutput=Output) +module("vllm.v1.request", Request=Request, RequestStatus=SimpleNamespace(RUNNING="running")) +sys.modules["vllm"].__version__ = "test-runtime" +spec = importlib.util.spec_from_file_location("dsv41_tested", Path(os.environ["AIC_FPM_DSV41_PRODUCER"])) +impl = importlib.util.module_from_spec(spec) +spec.loader.exec_module(impl) + + +def point(kind, batch=2, context=512, new=64): + return SimpleNamespace( + point_type=kind, + batch_size=batch, + total_kv_read_tokens=batch * context, + total_prefill_tokens=batch * new if kind == "prefill" else 0, + partition=None, + rows=None, + sample_reasons=[], + benchmark_id=1, + ) + + +def scheduler(pt): + obj = object.__new__(impl.DeepseekV41RealKVScheduler) + obj._real_tags = {} + obj._real_callback_stage = None + obj._real_stage = None + obj._real_outstanding = 0 + obj._real_tokens = [3, 8, 27, 42, 11, 79, 14] + obj._real_token_streams = [] + obj._real_witnesses = {} + obj._real_warmup_results = [] + obj._real_expected_warmup_ids = [] + obj._bench_point_result_timeout_seconds = 30 + obj._bench_seq = 0 + obj._bench_block_hasher = None + obj._bench_active_req_ids = set() + obj._bench_grid = deque([pt]) + obj._bench_config = SimpleNamespace(mode=pt.point_type) + obj.max_num_scheduled_tokens = 1024 + obj.requests = {} + obj.running = [] + obj.finished_req_ids = set() + obj.saved = [] + obj._bench_sync_pending = False + obj.kv_cache_manager = SimpleNamespace( + num_kv_cache_groups=3, take_new_block_ids=lambda: [], take_kv_cache_block_copies=lambda: ([], []) + ) + obj.kv_cache_manager.allocate_slots = Mock( + side_effect=lambda request, count, **kwargs: SimpleNamespace( + get_block_ids=lambda **kw: ([hash(request.request_id) % 101], [7], [11]) + ) + ) + return obj + + +class Driver: + """A FIFO worker stand-in that only fills memory after a forward completes.""" + + def __init__(self, obj): + self.obj = obj + self.memory = {} + self.inflight = [] + self.identities = {} + + def submit(self, output): + if output is None: + return + snapshots = {} + for rid, count in output.num_scheduled_tokens.items(): + request = self.obj.requests[rid] + self.identities.setdefault(rid, id(request)) + assert self.identities[rid] == id(request) + snapshots[rid] = (request.num_computed_tokens, count) + request.num_computed_tokens += count + if request.num_computed_tokens >= len(request.prompt_token_ids): + request.num_output_placeholders += 1 + self.inflight.append((output, snapshots)) + + def finish(self): + output, snapshots = self.inflight.pop(0) + for rid, (start, count) in snapshots.items(): + request = self.obj.requests[rid] + memory = self.memory.setdefault(rid, []) + assert len(memory) == start, "a forward cannot read uninitialized KV" + tokens = request._all_token_ids[start : start + count] + assert len(tokens) == count, "forward tokens must exist in real history" + memory.extend(tokens) + if start + count >= len(request.prompt_token_ids): + request._all_token_ids.append(13 + request.num_output_tokens) + request.num_output_tokens += 1 + request.num_output_placeholders -= 1 + self.obj._update_from_output(output, SimpleNamespace()) + + +class RealKVTests(unittest.TestCase): + def test_initialization_uses_validated_top_level_engram(self): + with tempfile.TemporaryDirectory() as directory: + Path(directory, "config.json").write_text("{}") + config = SimpleNamespace( + model_config=SimpleNamespace( + enforce_eager=True, + model=directory, + architecture="DeepseekV41ForCausalLM", + hf_config=SimpleNamespace(model_type="deepseek_v41"), + ), + parallel_config=SimpleNamespace( + tensor_parallel_size=4, + pipeline_parallel_size=1, + data_parallel_size=1, + use_ubatching=False, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + enable_expert_parallel=False, + ), + engram_config=SimpleNamespace(cpu_offload=False), + speculative_config=None, + ) + obj = object.__new__(impl.DeepseekV41RealKVScheduler) + obj._bench_active = True + obj.connector = obj.ec_connector = None + obj._bench_explicit_points = [point("prefill")] + obj._bench_config = SimpleNamespace(warmup_iterations=0) + identity = ("c" * 64, "full", "hbm_tp_sharded", "text") + with ( + patch.object(Base, "_bench_init", create=True), + patch("aisimulate_core.sdk.fpm_identity.execution_identity", return_value=identity) as identify, + patch.dict(os.environ, {"DYN_FPM_TOKENIZER_REVISION": "invalid"}), + self.assertRaisesRegex(ValueError, "tokenizer revision"), + ): + obj._bench_init(config) + identify.assert_called_once_with({}, engram_cpu_offload=False, input_modality="text") + self.assertEqual(obj._real_identity, dict(zip(EXECUTION_COLUMNS, identity, strict=True))) + + def test_decode_warms_actual_requests_then_pipelines_two_decode_steps(self): + obj = scheduler(point("decode")) + worker = Driver(obj) + seed = obj._real_step("decode") + worker.submit(seed) + self.assertEqual(seed.total_num_scheduled_tokens, 1022) + self.assertIsNone(obj._real_step("decode")) + self.assertEqual(obj._real_seed_tokens, 0) + worker.finish() + admission = obj._real_step("decode") + worker.submit(admission) + self.assertEqual(admission.scheduled_cached_reqs.num_computed_tokens, [511, 511]) + steady = obj._real_step("decode") + worker.submit(steady) + self.assertEqual(steady.scheduled_cached_reqs.num_computed_tokens, [512, 512]) + self.assertIsNone(obj._real_step("decode")) + worker.finish() + worker.finish() + obj._real_step("decode") + self.assertEqual(obj.saved[0][1], ["admission", "measure"]) + self.assertEqual(obj._real_witnesses[1]["completed_seed_tokens"], 1022) + self.assertTrue(all(len(values) == 513 for values in worker.memory.values())) + self.assertTrue(all(len(set(values)) > 2 for values in worker.memory.values())) + manifest = json.loads(obj._real_token_streams[0]) + self.assertEqual(len(manifest["requests"][0]["prompt_token_ids"]), 511) + self.assertEqual(len(manifest["requests"][0]["output_token_ids"]), 3) + self.assertEqual( + obj._real_witnesses[1]["token_stream_sha256"], impl.hashlib.sha256(obj._real_token_streams[0]).hexdigest() + ) + + def test_cached_prefill_preserves_full_prompt_history_and_same_state(self): + obj = scheduler(point("prefill", context=1024, new=64)) + worker = Driver(obj) + for _ in range(2): + output = obj._real_step("prefill") + worker.submit(output) + self.assertIsNone(obj._real_step("prefill")) + worker.finish() + measured = obj._real_step("prefill") + self.assertEqual(measured.scheduled_cached_reqs.num_computed_tokens, [1024, 1024]) + self.assertEqual(measured.scheduled_cached_reqs.num_output_tokens, [0, 0]) + worker.submit(measured) + worker.finish() + obj._real_step("prefill") + self.assertEqual(obj.saved[0][1], ["measure"]) + self.assertEqual(obj._real_witnesses[1]["completed_seed_tokens"], 2048) + self.assertTrue(all(len(values) == 1088 for values in worker.memory.values())) + + def test_zero_prefix_prefill_is_a_real_forward(self): + obj = scheduler(point("prefill", context=0)) + worker = Driver(obj) + output = obj._real_step("prefill") + self.assertEqual(len(output.scheduled_new_reqs), 2) + worker.submit(output) + worker.finish() + obj._real_step("prefill") + self.assertEqual(obj.saved[0][1], ["measure"]) + + def test_native_eager_warmup_keeps_a_separate_history_role(self): + pt = point("prefill", context=0) + pt.benchmark_id = 5 + pt.sample_reasons = ["eager_warmup"] + obj = scheduler(pt) + obj._real_validate_grid() + worker = Driver(obj) + worker.submit(obj._real_step("prefill")) + worker.finish() + obj._real_step("prefill") + stream = json.loads(obj._real_token_streams[0]) + self.assertEqual(stream["sampling_role"], "warmup") + self.assertEqual(stream["benchmark_id"], 5) + self.assertEqual(obj._real_warmup_results[0]["point"]["benchmark_id"], 5) + self.assertEqual(obj._real_warmup_results[0]["real_kv_witness"], obj._real_witnesses[5]) + self.assertEqual(obj._real_warmup_results[0]["fpms"], ["measure"]) + self.assertEqual(obj._real_expected_warmup_ids, [5]) + + def test_allocation_failure_has_no_fallback(self): + obj = scheduler(point("decode")) + obj.kv_cache_manager.allocate_slots.return_value = None + obj.kv_cache_manager.allocate_slots.side_effect = None + with self.assertRaisesRegex(RuntimeError, "fallback is forbidden"): + obj._real_step("decode") + self.assertEqual(obj.saved, []) + + def test_lost_state_fails(self): + obj = scheduler(point("prefill")) + worker = Driver(obj) + worker.submit(obj._real_step("prefill")) + worker.finish() + obj.requests.clear() + with self.assertRaisesRegex(RuntimeError, "was lost"): + obj._real_step("prefill") + + def test_timeout_never_publishes_real_kv(self): + obj = scheduler(point("decode")) + obj._real_step("decode") + obj._real_deadline = 0 + with self.assertRaisesRegex(RuntimeError, "timed out"): + obj._real_step("decode") + self.assertEqual(obj.saved, []) + + def test_real_seed_jit_has_bounded_deadline_independent_of_fake_point_timeout(self): + obj = scheduler(point("prefill", context=512, new=64)) + obj._bench_point_result_timeout_seconds = 8.0 + worker = Driver(obj) + with patch.object(impl.time, "monotonic", return_value=100.0): + worker.submit(obj._real_step("prefill")) + self.assertEqual(obj._real_deadline, 220.0) + with patch.object(impl.time, "monotonic", return_value=109.0): + self.assertIsNone(obj._real_step("prefill")) + worker.finish() + with patch.object(impl.time, "monotonic", return_value=190.0): + worker.submit(obj._real_step("prefill")) + worker.finish() + with patch.object(impl.time, "monotonic", return_value=191.0): + obj._real_step("prefill") + self.assertEqual(obj.saved[0][1], ["measure"]) + self.assertEqual(obj._bench_point_result_timeout_seconds, 8.0) + + obj = scheduler(point("prefill", context=512, new=64)) + with patch.object(impl.time, "monotonic", return_value=100.0): + obj._real_step("prefill") + with ( + patch.object(impl.time, "monotonic", return_value=220.0), + self.assertRaisesRegex(RuntimeError, "timed out"), + ): + obj._real_step("prefill") + self.assertEqual(obj.saved, []) + self.assertEqual(obj._real_witnesses, {}) + + def test_synthetic_entry_points_are_blocked(self): + obj = scheduler(point("decode")) + for method in [obj._bench_cache_fake_prefixes, obj._bench_inject_fake_decode]: + with self.assertRaisesRegex(RuntimeError, "forbidden"): + method() + + def test_canary_limits_fail_instead_of_silently_skipping(self): + for pt in [ + point("decode", batch=3), + point("decode", context=2049), + point("decode", context=1), + point("prefill", new=513), + point("prefill", batch=2, new=257), + point("prefill", context=2000, new=128), + ]: + with self.subTest(point=pt): + obj = scheduler(pt) + with self.assertRaises(ValueError): + obj._real_validate_grid() + for pt in [ + point("decode", context=2048), + point("prefill", batch=1, context=1536, new=512), + point("prefill", batch=2, context=1536, new=256), + ]: + with self.subTest(boundary=pt): + scheduler(pt)._real_validate_grid() + + def test_native_ring_capacity_semantics_are_used(self): + obj = scheduler(point("decode")) + ring = SimpleNamespace(get_num_blocks_to_allocate=Mock(return_value=1)) + dense = SimpleNamespace(get_num_blocks_to_allocate=Mock(return_value=32)) + obj.kv_cache_manager.coordinator = SimpleNamespace(single_type_managers=[ring, dense]) + self.assertEqual(obj._bench_blocks_per_req(2048), 33) + self.assertEqual(ring.get_num_blocks_to_allocate.call_args.kwargs["num_tokens"], 2048) + + def test_marker_is_earned_after_real_seed_and_measure_complete(self): + pt = point("prefill") + obj = scheduler(pt) + obj._real_validate_grid() + self.assertNotIn("kvwarm_real_kv", pt.sample_reasons) + worker = Driver(obj) + worker.submit(obj._real_step("prefill")) + worker.finish() + self.assertNotIn("kvwarm_real_kv", pt.sample_reasons) + worker.submit(obj._real_step("prefill")) + worker.finish() + obj._real_step("prefill") + self.assertIn("kvwarm_real_kv", pt.sample_reasons) + + def test_uneven_native_lengths_preserve_each_request_prefix(self): + pt = point("prefill", context=512, new=64) + obj = scheduler(pt) + obj._real_lengths = lambda _: ([512, 256], [129, 128]) + worker = Driver(obj) + worker.submit(obj._real_step("prefill")) + worker.finish() + measured = obj._real_step("prefill") + self.assertEqual(measured.scheduled_cached_reqs.num_computed_tokens, [512, 256]) + self.assertEqual(list(measured.num_scheduled_tokens.values()), [129, 128]) + worker.submit(measured) + worker.finish() + obj._real_step("prefill") + self.assertEqual(sorted(map(len, worker.memory.values())), [384, 641]) + + def test_shorter_prefix_parks_without_rebuilding_its_state(self): + obj = scheduler(point("prefill")) + obj._real_lengths = lambda _: ([1024, 256], [129, 128]) + worker = Driver(obj) + worker.submit(obj._real_step("prefill")) + worker.finish() + parked = obj._real_step("prefill") + self.assertEqual(len(parked.num_scheduled_tokens), 1) + worker.submit(parked) + worker.finish() + measured = obj._real_step("prefill") + self.assertEqual(measured.scheduled_cached_reqs.num_computed_tokens, [1024, 256]) + self.assertEqual(len(measured.scheduled_cached_reqs.all_token_ids), 2) + worker.submit(measured) + worker.finish() + obj._real_step("prefill") + self.assertEqual(sorted(map(len, worker.memory.values())), [384, 1153]) + + def test_producer_serialized_provenance_matches_actual_collector_reader(self): + # Execute the actual reader function from the concurrently developed + # Collector source, without importing its unrelated planning dependencies. + consumer_path = Path(os.environ["AIC_FPM_NATIVE_ARTIFACT"]) + tree = ast.parse(consumer_path.read_text()) + names = {"_validate_execution_provenance", "_validate_kvwarm_contract"} + functions = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names] + self.assertEqual(len(functions), 2) + columns = EXECUTION_COLUMNS + scope = { + "Path": Path, + "Any": Any, + "FPMCell": SimpleNamespace, + "re": re, + "EXECUTION_COLUMNS": columns, + "KVWARM_STRATEGIES": {"pure_tp"}, + } + exec(compile(ast.Module(body=functions, type_ignores=[]), str(consumer_path), "exec"), scope) + with tempfile.TemporaryDirectory() as directory: + obj = scheduler(point("decode")) + path = Path(directory) / "benchmark_results.json" + obj._bench_config.output_path = str(path) + obj._real_identity = dict(zip(columns, ("a" * 64, "full", "hbm_tp_sharded", "text"), strict=True)) + obj._real_input = { + "source": "tokenizer_text", + "text_sha256": "b" * 64, + "token_ids_sha256": "c" * 64, + "tokenizer_revision": impl.MODEL_SHA, + "token_count": 10, + "unique_token_count": 7, + } + obj._real_witnesses = {1: {"completed_seed_tokens": 1022, "same_request": True, "allocated_fake_tokens": 0}} + obj._bench_write_results() + payload = json.loads(path.read_text()) + self.assertIsNone(payload["producer"]["vllm_revision"]) + self.assertIsNone(payload["producer"]["dynamo_revision"]) + self.assertEqual(payload["producer"]["instrumentation_revision"], impl.DYNAMO_SHA) + self.assertEqual(payload["producer"]["vllm_package_version"], "test-runtime") + self.assertEqual(payload["producer"]["collection_timeouts"]["same_request_seed_and_measure_seconds"], 120.0) + cell = SimpleNamespace( + execution_identity=tuple(obj._real_identity.values()), + input_text_sha256="b" * 64, + workload_kind="decode", + parallel_strategy="pure_tp", + ) + evidence = scope["_validate_execution_provenance"](cell, payload, path) + for field, value in obj._real_input.items(): + self.assertEqual(evidence[field], value) + self.assertEqual(evidence["token_stream_manifest"]["records"], 0) + self.assertTrue((path.parent / evidence["token_stream_manifest"]["file"]).is_file()) + scope["_validate_kvwarm_contract"](cell, payload["kvwarm"], path) + self.assertEqual(payload["results"][0]["kv_seed_regime"], "real_kv") + self.assertEqual(payload["results"][0]["real_kv_witness"]["completed_seed_tokens"], 1022) + payload["execution"] = payload.pop("execution_identity") + with self.assertRaisesRegex(ValueError, "execution identity"): + scope["_validate_execution_provenance"](cell, payload, path) + + def test_token_slices_are_repeatable_and_offset_varies_history(self): + self.assertEqual(impl.token_slice([1, 3, 7], 5, 1), [3, 7, 1, 3, 7]) + self.assertNotEqual(impl.token_slice([1, 3, 7], 8, 0), impl.token_slice([1, 3, 7], 8, 1)) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/aisimulate/tests/unit/collector/test_fpm_dsv41_producer.py b/python/aisimulate/tests/unit/collector/test_fpm_dsv41_producer.py new file mode 100644 index 000000000..fe31ed6ed --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_dsv41_producer.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run producer state-machine stubs in an isolated process.""" + +import hashlib +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + + +def test_real_kv_producer_lifecycle(): + root = Path(__file__).resolve().parents[3] + runtime = root / "collector/fpm_forward/runtime/dsv41" + environment = dict(os.environ) + environment["AIC_FPM_DSV41_PRODUCER"] = str(runtime / "dsv41_scheduler.py") + environment["AIC_FPM_NATIVE_ARTIFACT"] = str(root / "collector/fpm_forward/native_artifact.py") + completed = subprocess.run( + [sys.executable, str(Path(__file__).with_name("fixtures") / "dsv41_producer_lifecycle.py")], + env=environment, + timeout=30, + capture_output=True, + text=True, + ) + assert completed.returncode == 0, f"{completed.stdout}\n{completed.stderr}" + + +@pytest.mark.parametrize("valid_source", [False, True]) +@pytest.mark.parametrize("adapter_first", [False, True]) +def test_bootstrap_is_lazy_in_helpers_and_fail_closed_at_scheduler_import(tmp_path, valid_source, adapter_first): + runtime = Path(__file__).resolve().parents[3] / "collector/fpm_forward/runtime/dsv41" + shutil.copy2(runtime / "sitecustomize.py", tmp_path / "sitecustomize.py") + module = tmp_path / "dynamo/vllm/instrumented_scheduler.py" + module.parent.mkdir(parents=True) + # A scheduler can spawn compiler/helper interpreters that inherit the same + # environment; those children must not recursively import the scheduler. + helper = "import sys; assert 'dynamo.vllm.instrumented_scheduler' not in sys.modules" + module.write_text( + "import subprocess,sys\n" + f"subprocess.run([sys.executable, '-c', {helper!r}], check=True, timeout=5)\n" + "class InstrumentedScheduler: pass\n" + ) + (tmp_path / "dsv41_scheduler.py").write_text( + "import dynamo.vllm.instrumented_scheduler as native\n" + "class DeepseekV41RealKVScheduler(native.InstrumentedScheduler): pass\n" + "native.InstrumentedScheduler = DeepseekV41RealKVScheduler\n" + ) + (tmp_path / "runtime-source-sha256.json").write_text( + json.dumps( + { + "dynamo/vllm/instrumented_scheduler.py": hashlib.sha256(module.read_bytes()).hexdigest() + if valid_source + else "0" * 64 + } + ) + ) + environment = dict(os.environ, PYTHONPATH=str(tmp_path), DYN_FPM_DSV41_REAL_KV="1") + child = subprocess.run([sys.executable, "-c", helper], env=environment, capture_output=True, text=True, timeout=10) + assert child.returncode == 0, child.stderr + result = subprocess.run( + [ + sys.executable, + "-c", + "try:\n" + + (" import dsv41_scheduler\n" if adapter_first else "") + + " from dynamo.vllm.instrumented_scheduler import InstrumentedScheduler\n" + + " assert InstrumentedScheduler.__name__ == 'DeepseekV41RealKVScheduler'\n" + + "finally:\n print('normal process cleanup ran', flush=True)\n", + ], + env=environment, + capture_output=True, + text=True, + timeout=10, + ) + if valid_source: + assert result.returncode == 0, result.stderr + else: + assert result.returncode == 1 + assert "Traceback (most recent call last)" in result.stderr + assert "normal process cleanup ran" in result.stdout + assert "pinned source mismatch" in result.stderr + + +@pytest.mark.parametrize("failure", ["source", "sdk", "activation"]) +def test_preflight_subprocess_retains_failure_audit_traceback_and_cleanup(tmp_path, failure): + runtime = Path(__file__).resolve().parents[3] / "collector/fpm_forward/runtime" + shutil.copy2(runtime / "preflight.py", tmp_path / "preflight.py") + if failure != "activation": + shutil.copy2(runtime / "dsv41/sitecustomize.py", tmp_path / "sitecustomize.py") + native = tmp_path / "dynamo/vllm/instrumented_scheduler.py" + native.parent.mkdir(parents=True) + native.write_text("class BenchmarkPoint: pass\nclass InstrumentedScheduler: pass\n") + (tmp_path / "dsv41_scheduler.py").write_text( + "import dynamo.vllm.instrumented_scheduler as native\n" + "class DeepseekV41RealKVScheduler(native.InstrumentedScheduler): pass\n" + "native.InstrumentedScheduler = DeepseekV41RealKVScheduler\n" + ) + (tmp_path / "runtime-source-sha256.json").write_text( + json.dumps( + { + "dynamo/vllm/instrumented_scheduler.py": "0" * 64 + if failure == "source" + else hashlib.sha256(native.read_bytes()).hexdigest(), + } + ) + ) + # A regular private package shadows the installed SDK so this test proves + # the exact preflight error path without any model/backend import. + (tmp_path / "aisimulate_core").mkdir() + (tmp_path / "aisimulate_core/__init__.py").write_text("raise ImportError('SDK native extension unavailable')\n") + audit = tmp_path / "audit.json" + child = subprocess.run( + [ + sys.executable, + "-c", + "import preflight\nfrom pathlib import Path\n" + f"preflight._AUDIT_PATH = Path({str(audit)!r})\n" + "try:\n preflight.main()\nfinally:\n print('preflight cleanup ran', flush=True)\n", + ], + cwd=tmp_path, + env=dict(os.environ, PYTHONPATH=str(tmp_path), DYN_FPM_DSV41_REAL_KV="1"), + text=True, + capture_output=True, + timeout=10, + ) + assert child.returncode == 1 + assert "Traceback (most recent call last)" in child.stderr + assert "preflight cleanup ran" in child.stdout + receipt = json.loads(audit.read_text()) + assert receipt["status"] == "failed" + expected = { + "source": "pinned source mismatch", + "sdk": "SDK native extension unavailable", + "activation": "source-checked scheduler activation did not occur", + }[failure] + assert expected in receipt["import_error"] and expected in child.stderr diff --git a/python/aisimulate/tests/unit/collector/test_fpm_exec.py b/python/aisimulate/tests/unit/collector/test_fpm_exec.py index 5916f67bf..27090c792 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_exec.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_exec.py @@ -214,6 +214,7 @@ def _stage( values.update(env_overrides or {}) exports = "\n".join(f"export {name}={shlex.quote(str(values[name]))}" for name in FPM_ENV_EXPORTED_VARS) (workdir / "fpm_env.sh").write_text(f"#!/usr/bin/env bash\n{exports}\n") + (workdir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") (workdir / "preflight.py").write_text("") (workdir / "run.sh").write_text(run_script) if etcd_stub is None: @@ -881,14 +882,16 @@ def accept_report() -> None: def test_fpm_exec_consumes_only_contract_environment(): """Everything fpm_exec.sh needs from the Generator travels through the - fpm_env.sh exports; the only extra FPM_* input is the operator-tunable - barrier timeout, which is defaulted in-script and never rendered.""" + fpm_env.sh exports. The Collector separately stages the configurable + readiness budget; completion-barrier timeout remains operator-tunable.""" script = FPM_EXEC.read_text() - consumed = set(re.findall(r"FPM_[A-Z0-9_]+", script)) + # Match complete shell identifiers. DYN_FPM_* exports configure the native + # producer and must not be mistaken for Generator-owned FPM_* inputs. + consumed = set(re.findall(r"\bFPM_[A-Z0-9_]+\b", script)) assert {"FPM_NODE_RANK", "FPM_MASTER_ADDR", "FPM_BENCHMARK_OUTPUT_PATH"} <= consumed - allowed = set(FPM_ENV_EXPORTED_VARS) | {"FPM_COMPLETION_BARRIER_TIMEOUT_SECONDS"} + allowed = set(FPM_ENV_EXPORTED_VARS) | {"FPM_COMPLETION_BARRIER_TIMEOUT_SECONDS", "FPM_READINESS_TIMEOUT_SECONDS"} assert consumed <= allowed, sorted(consumed - allowed) @@ -951,3 +954,67 @@ def test_fpm_exec_result_naming_matches_contract_on_shared_vectors( assert completed.returncode == 0, completed.stderr for path in expected_paths: assert Path(path).is_file(), path + + +@pytest.mark.parametrize("budget, expected_status", [(1, 1), (10, 0)]) +def test_follower_readiness_budget_covers_delayed_leader_start(tmp_path, budget, expected_status): + marker = tmp_path / "engine-started" + staged = _stage( + tmp_path, + run_script=f"#!/bin/bash\ntouch {shlex.quote(str(marker))}\nexit 0\n", + env_overrides={"FPM_NODE_RANK": 1, "FPM_NODE_COUNT": 2}, + ) + (staged.workdir / "collector-runtime-env.sh").write_text(f"export FPM_READINESS_TIMEOUT_SECONDS={budget}\n") + stop = threading.Event() + ready = threading.Event() + readiness_started = tmp_path / "readiness-started" + # Mark the actual follower loop, not the time the shell was spawned. + source = staged.script.read_text() + deadline_line = "deadline = time.monotonic() + float(sys.argv[2])" + assert source.count(deadline_line) == 1 + staged.script.write_text( + source.replace( + deadline_line, + deadline_line + f"\nopen({str(readiness_started)!r}, 'w').close()", + ) + ) + + def delayed_leader(): + while not readiness_started.exists(): + if stop.wait(0.01): + return + if stop.wait(1.4): + return + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server: + server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server.bind(("127.0.0.1", staged.etcd_port)) + server.listen(8) + ready.set() + stop.wait(10) + + thread = threading.Thread(target=delayed_leader) + thread.start() + try: + result = _run(staged, timeout=10) + finally: + stop.set() + thread.join(timeout=5) + assert not thread.is_alive() + assert result.returncode == expected_status, result.stderr + assert marker.exists() == (expected_status == 0) + if expected_status: + assert "etcd readiness timeout" in result.stderr + else: + assert ready.is_set() + + +@pytest.mark.parametrize("value", ["", "0", "-1", "3601", "999999999999999999999999999999999999"]) +def test_invalid_startup_budget_fails_before_etcd_or_engine(tmp_path, value): + staged = _stage(tmp_path, run_script="#!/bin/bash\nexit 0\n") + (staged.workdir / "collector-runtime-env.sh").write_text( + f"export FPM_READINESS_TIMEOUT_SECONDS={shlex.quote(value)}\n" + ) + result = _run(staged, timeout=5) + assert result.returncode == 2 + assert "must be an integer from 1 through 3600" in result.stderr + assert not staged.etcd_trace.exists() diff --git a/python/aisimulate/tests/unit/collector/test_fpm_explicit_points.py b/python/aisimulate/tests/unit/collector/test_fpm_explicit_points.py new file mode 100644 index 000000000..7738a8ed4 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_explicit_points.py @@ -0,0 +1,278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit point transport preserves the native manifest and attempt identity.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import sys +from types import SimpleNamespace + +import pytest +from collector.fpm_forward.config import FPMCollectionOptions, add_fpm_arguments +from collector.fpm_forward.native_artifact import COLLECTOR_PROVENANCE_FILENAME +from collector.fpm_forward.planner import BackendPolicy, FPMCell +from collector.fpm_forward.runner import ( + POINTS_FILENAME, + POINTS_RECEIPT_FILENAME, + REMOTE_WORKDIR, + _cell_generator_overrides, + _configured_sampling_metadata, + _record_points_receipts, + _stage_points_file, + _validate_points_receipts, + run_collection, +) +from collector.fpm_forward.types import ParallelTopology + +pytestmark = pytest.mark.unit + + +def _payload(): + return { + "schema_version": 3, + "prefill": [ + {"batch_size": 2, "total_prefill_tokens": 8, "total_kv_read_tokens": 256, "rows": [[3, 128], [5, 128]]} + ], + "decode": [{"batch_size": 2, "total_kv_read_tokens": 256}], + } + + +def _options(tmp_path, payload=None, **extra): + path = tmp_path / "source.json" + path.write_text(json.dumps(_payload() if payload is None else payload, indent=2)) + options = FPMCollectionOptions.from_args( + argparse.Namespace(fpm_max_gpus=4, fpm_benchmark_points_file=str(path), **extra) + ) + return options, path + + +def _cell(phase="prefill", policy=None): + return FPMCell( + cell_id="cell", + workload_kind=phase, + topology=ParallelTopology(tp=4, pp=1, dp=1, moe_tp=4, moe_ep=1, cp=1), + weight_quantization="bfloat16", + kv_cache_dtype="auto", + backend_policy=policy or BackendPolicy("baseline", {}, {}), + parallel_strategy="pure_tp", + gemm_quant_mode="bfloat16", + moe_quant_mode="bfloat16", + fmha_quant_mode="bfloat16", + comm_quant_mode="half", + ) + + +def _plan(options): + return SimpleNamespace(options=options, sha256="plan-sha", model_path="text/model") + + +def test_cli_freezes_canonical_payload_rows_and_source_mutation(tmp_path): + options, source = _options(tmp_path) + canonical = json.dumps(_payload(), sort_keys=True, separators=(",", ":")) + assert options.benchmark_points_json == canonical + assert options.benchmark_points_sha256 == hashlib.sha256(canonical.encode()).hexdigest() + assert options.to_dict()["benchmark_points"] == { + "payload": _payload(), + "sha256": options.benchmark_points_sha256, + } + source.write_text("{}") + staged = _stage_points_file(_plan(options), tmp_path) + assert len(staged) == 1 + assert staged[0].name == POINTS_FILENAME + assert staged[0].read_text() == canonical + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + args = parser.parse_args(["--fpm-max-gpus", "4", "--fpm-benchmark-points-file", str(staged[0])]) + assert FPMCollectionOptions.from_args(args) == options + + +def test_manifest_formatting_ignores_paths_but_point_order_changes_identity(tmp_path): + options, source = _options(tmp_path) + source.write_text(options.benchmark_points_json) + same = FPMCollectionOptions.from_args(argparse.Namespace(fpm_max_gpus=4, fpm_benchmark_points_file=str(source))) + assert same.to_dict() == options.to_dict() + changed = _payload() + changed["prefill"][0]["rows"].reverse() + other, _ = _options(tmp_path, changed) + assert other.benchmark_points_sha256 != options.benchmark_points_sha256 + assert other.to_dict() != options.to_dict() + + +@pytest.mark.parametrize( + "payload", + [ + [], + {}, + {"schema_version": True, "prefill": [], "decode": []}, + {"schema_version": 4, "prefill": [], "decode": []}, + {"schema_version": 3, "prefill": [], "decode": []}, + {"schema_version": 3, "prefill": [1], "decode": []}, + ], +) +def test_invalid_envelope_rejected_before_planning(tmp_path, payload): + with pytest.raises(ValueError, match="benchmark-points"): + _options(tmp_path, payload) + + +def test_duplicate_json_fields_are_not_silently_canonicalized(tmp_path): + source = tmp_path / "source.json" + source.write_text('{"schema_version":2,"schema_version":3,"prefill":[],"decode":[]}') + with pytest.raises(ValueError, match="duplicate benchmark-points field"): + FPMCollectionOptions.from_args(argparse.Namespace(fpm_max_gpus=4, fpm_benchmark_points_file=str(source))) + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +def test_generator_owns_explicit_argument_rendering(tmp_path, phase): + options, _ = _options(tmp_path) + plan, cell = _plan(options), _cell(phase) + args = _cell_generator_overrides(plan, cell, {})["params"]["agg"]["extra_cli_args"] + assert args.count("--benchmark-points-file") == 1 + assert args[args.index("--benchmark-points-file") + 1] == f"{REMOTE_WORKDIR}/{POINTS_FILENAME}" + assert _configured_sampling_metadata(plan, cell, smoke=False) == { + "benchmark_points_sha256": options.benchmark_points_sha256, + "requested_point_count": 1, + } + + +def test_default_grid_and_smoke_remain_explicitly_separate(tmp_path): + legacy = FPMCollectionOptions.from_args(argparse.Namespace(fpm_max_gpus=4)) + assert "benchmark_points" not in legacy.to_dict() + assert _stage_points_file(_plan(legacy), tmp_path) == [] + args = _cell_generator_overrides(_plan(legacy), _cell(), {}, smoke=True)["params"]["agg"]["extra_cli_args"] + assert "--benchmark-points-file" not in args + with pytest.raises(ValueError, match="cannot be combined with --smoke"): + _options(tmp_path, smoke=True) + options, _ = _options(tmp_path) + with pytest.raises(ValueError, match="cannot be combined with --smoke"): + run_collection( + _plan(options), + generator_overrides={}, + checkpoint_dir=str(tmp_path / "checkpoint"), + artifact_root=str(tmp_path / "artifacts"), + resume=False, + retry_failed=False, + smoke=True, + ) + assert not (tmp_path / "artifacts").exists() + + +def test_backend_policy_cannot_replace_frozen_points(tmp_path): + options, _ = _options(tmp_path) + cell = _cell( + policy=BackendPolicy("bad", {"params": {"agg": {"extra_cli_args": ["--benchmark-points-file=other"]}}}, {}) + ) + with pytest.raises(ValueError, match="must be supplied through"): + _cell_generator_overrides(_plan(options), cell, {}) + + +def test_staging_rejects_modified_frozen_payload(tmp_path): + options, _ = _options(tmp_path) + from dataclasses import replace + + plan = _plan(replace(options, benchmark_points_json="{}")) + with pytest.raises(ValueError, match="payload and SHA256 disagree"): + _stage_points_file(plan, tmp_path) + + +@pytest.mark.parametrize("transport", ["_exec", "_exec_checked"]) +def test_runtime_hash_and_attempt_receipt_are_checked_in_both_transports(tmp_path, transport): + options, _ = _options(tmp_path) + plan, cell = _plan(options), _cell() + staged = _stage_points_file(plan, tmp_path)[0] + raw = tmp_path / "raw" + unit = raw / "unit" + unit.mkdir(parents=True) + (unit / COLLECTOR_PROVENANCE_FILENAME).write_text("{}") + + def execute(pod, command, timeout): + assert pod == "unit" and timeout == 300 + local = [sys.executable, *command[1:-2], str(staged), str(unit / POINTS_RECEIPT_FILENAME)] + return subprocess.run(local, check=True, capture_output=True, text=True) + + resource = SimpleNamespace(**{transport: execute}) + _record_points_receipts(resource, ["unit"], plan, cell, "attempt", phase="before") + with pytest.raises(ValueError, match="receipt mismatch"): + _validate_points_receipts(plan, cell, raw, "attempt") + _record_points_receipts(resource, ["unit"], plan, cell, "attempt", phase="after") + _validate_points_receipts(plan, cell, raw, "attempt") + with pytest.raises(ValueError, match="receipt mismatch"): + _validate_points_receipts(plan, cell, raw, "another-attempt") + staged.write_text("{}") + with pytest.raises(subprocess.CalledProcessError): + _record_points_receipts(resource, ["unit"], plan, cell, "attempt", phase="after") + (unit / POINTS_RECEIPT_FILENAME).unlink() + with pytest.raises(FileNotFoundError): + _validate_points_receipts(plan, cell, raw, "attempt") + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +def test_eager_is_frozen_and_rendered_without_capture_configuration(tmp_path, phase): + options, _ = _options(tmp_path, fpm_enforce_eager=True) + plan, cell = _plan(options), _cell(phase) + assert options.to_dict()["enforce_eager"] is True + args = _cell_generator_overrides(plan, cell, {})["params"]["agg"]["extra_cli_args"] + assert args.count("--enforce-eager") == 1 + assert "--compilation-config" not in args + assert "--benchmark-points-file" in args + parser = argparse.ArgumentParser() + add_fpm_arguments(parser) + parsed = parser.parse_args(["--fpm-max-gpus", "4", "--fpm-enforce-eager"]) + assert FPMCollectionOptions.from_args(parsed).enforce_eager is True + assert "enforce_eager" not in FPMCollectionOptions.from_args(argparse.Namespace(fpm_max_gpus=4)).to_dict() + + +@pytest.mark.parametrize("argument", ["--enforce-eager", "--no-enforce-eager", "--enforce-eager=false"]) +def test_policy_cannot_hide_eager_identity(tmp_path, argument): + options, _ = _options(tmp_path, fpm_enforce_eager=True) + cell = _cell(policy=BackendPolicy("conflict", {"params": {"agg": {"extra_cli_args": [argument]}}}, {})) + with pytest.raises(ValueError, match="must be supplied through --fpm-enforce-eager"): + _cell_generator_overrides(_plan(options), cell, {}) + + +@pytest.mark.parametrize( + "phase,field", + [ + ("prefill", "batch_size"), + ("prefill", "total_prefill_tokens"), + ("prefill", "total_kv_read_tokens"), + ("decode", "batch_size"), + ("decode", "total_kv_read_tokens"), + ], +) +@pytest.mark.parametrize("value", [None, True, 1.5, "2", -1]) +def test_manifest_coordinates_fail_at_freeze(tmp_path, phase, field, value): + payload = _payload() + if value is None: + del payload[phase][0][field] + else: + payload[phase][0][field] = value + with pytest.raises(ValueError, match=f"benchmark-points {phase} {field}"): + _options(tmp_path, payload) + + +@pytest.mark.parametrize( + "field,value", + [("fpm_enforce_eager", True), ("fpm_enforce_eager", False), ("fpm_benchmark_points_file", "points.json")], +) +def test_explicit_fpm_flags_are_rejected_for_operator_collection(field, value): + from collector.fpm_forward.config import reject_fpm_arguments_without_fpm + + with pytest.raises(ValueError, match=field.replace("_", "-")): + reject_fpm_arguments_without_fpm(argparse.Namespace(ops=["gemm"], **{field: value})) + + +@pytest.mark.parametrize("phase,field", [("prefill", "total_prefill_tokens"), ("decode", "total_kv_read_tokens")]) +@pytest.mark.parametrize("tokens", [0, 1, 2]) +def test_manifest_requires_at_least_one_token_per_sequence(tmp_path, phase, field, tokens): + payload = _payload() + payload[phase][0][field] = tokens + if tokens < 2: + with pytest.raises(ValueError, match=f"benchmark-points {phase} {field}"): + _options(tmp_path, payload) + else: + _options(tmp_path, payload) diff --git a/python/aisimulate/tests/unit/collector/test_fpm_forward.py b/python/aisimulate/tests/unit/collector/test_fpm_forward.py index 20446ac38..86b04ca73 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_forward.py @@ -233,12 +233,24 @@ def test_pure_tp_requires_explicit_model_runtime_capability(): enumerate_fpm_topologies(backend="vllm", is_moe=True, options=options) -def test_plan_contains_only_cell_matrix_and_native_point_contract(): +@pytest.mark.parametrize("explicit", [False, True]) +def test_plan_contains_only_cell_matrix_and_native_point_contract(tmp_path, explicit): + points_file = tmp_path / "points.json" + points_file.write_text( + json.dumps( + { + "schema_version": 3, + "prefill": [{"batch_size": 1, "total_prefill_tokens": 128, "total_kv_read_tokens": 0}], + "decode": [], + } + ) + ) options = FPMCollectionOptions.from_args( _args( fpm_parallel_axes=["dp", "moe_ep"], fpm_dp_sizes=[4], fpm_moe_ep_sizes=[4], + fpm_benchmark_points_file=str(points_file) if explicit else None, ) ) kwargs = { @@ -264,7 +276,7 @@ def test_plan_contains_only_cell_matrix_and_native_point_contract(): assert {cell.workload_kind for cell in first.cells} == {"prefill", "decode"} assert {cell.parallel_strategy for cell in first.cells} == {"dep"} payload = first.to_dict() - assert payload["schema_version"] == 10 + assert payload["schema_version"] == 11 assert payload["capability"]["model_config"]["source_kind"] == "aic_cache" assert len(payload["capability"]["model_config"]["sha256"]) == 64 assert payload["capability"]["model_config"]["payload"]["architectures"] == ["GlmMoeDsaForCausalLM"] @@ -273,11 +285,13 @@ def test_plan_contains_only_cell_matrix_and_native_point_contract(): assert point_generation == { "owner": "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler", "method": "native_self_benchmark", + "source": "frozen_explicit_manifest" if explicit else "native_auto_grid", + "manifest_sha256": options.benchmark_points_sha256, "coordinates": ["batch_size", "total_prefill_tokens", "total_kv_read_tokens"], "partition_policy": "balanced_v1", "point_admission": "dynamo_live_scheduler", "precondition": "vllm_engine_initialized", - "planned_point_count": None, + "planned_point_count": 1 if explicit else None, } assert prefill_sampling["cudagraph_capture_size_count"] == 99 assert prefill_sampling["new_token_axis_point_count"] == 199 @@ -1168,6 +1182,7 @@ def _synthetic_plan_and_cell(tmp_path): model_path="org/model", system="b200_sxm", backend="vllm", + cells=(cell,), options=SimpleNamespace(warmup_iterations=0), capability=SimpleNamespace( support_level="exact", @@ -1338,7 +1353,7 @@ def test_native_validation_rejects_sub_batch_token_totals(tmp_path): _expected_scheduled(decode_point) -def test_formal_database_uses_schema_v6_and_rejects_conflicts(tmp_path): +def test_formal_database_uses_schema_v7_and_rejects_conflicts(tmp_path): plan, cell, cell_dir = _synthetic_plan_and_cell(tmp_path) rows = aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt") parquet, metadata, skipped = write_formal_database(plan, rows, systems_root=tmp_path / "systems") @@ -1346,7 +1361,7 @@ def test_formal_database_uses_schema_v6_and_rejects_conflicts(tmp_path): assert parquet.exists() metadata_payload = json.loads(metadata.read_text()) - assert metadata_payload["schema_version"] == 6 + assert metadata_payload["schema_version"] == 7 assert metadata_payload["coordinate_system"] == "iteration_totals_balanced_v1" assert metadata_payload["backend_version"] == "0.24.0" assert metadata_payload["collector_attempt_ids"] == ["attempt"] @@ -1386,14 +1401,14 @@ def test_formal_database_first_publisher_wins_on_rerun_overlap(tmp_path): assert parquet2.read_bytes() == sealed -def test_formal_database_commit_validation_accepts_sealed_schema_v6_pair(tmp_path): +def test_formal_database_commit_validation_accepts_sealed_schema_v7_pair(tmp_path): plan, cell, cell_dir = _synthetic_plan_and_cell(tmp_path) rows = aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt") parquet, metadata, _skipped = write_formal_database(plan, rows, systems_root=tmp_path / "systems") commit = validate_formal_database_commit(parquet, metadata, plan) - assert commit["schema_version"] == 6 + assert commit["schema_version"] == 7 assert commit["row_count"] == len(rows) @@ -1767,7 +1782,7 @@ def test_formal_database_merge_gate_names_missing_row_key_columns(tmp_path): parquet_path = destination / "fpm_forward_perf.parquet" pq.write_table(pa.Table.from_pylist(stale_rows), parquet_path) (destination / "fpm_forward_perf.metadata.json").write_text( - json.dumps({"parquet_sha256": hashlib.sha256(parquet_path.read_bytes()).hexdigest()}) + json.dumps({"schema_version": 7, "parquet_sha256": hashlib.sha256(parquet_path.read_bytes()).hexdigest()}) ) with pytest.raises(ValueError, match=r"missing columns: \['weight_quantization'\]"): @@ -2384,3 +2399,267 @@ def unavailable(*_args, **_kwargs): assert len(plan.topologies) == 3 assert {decision.disposition for decision in plan.topology_memory_admission} == {"unknown"} + + +def _write_v41_token_streams(payload, path): + import hashlib + + lines = [] + for row in payload["results"]: + point = row["point"] + batch = point["batch_size"] + decode = point["point_type"] == "decode" + seed = point["total_kv_read_tokens"] - (batch if decode else 0) + prompt_total = seed + (0 if decode else point["total_prefill_tokens"]) + lengths = [prompt_total // batch + (index < prompt_total % batch) for index in range(batch)] + stream = { + "benchmark_id": point["benchmark_id"], + "requests": [ + { + "request_index": index, + "prompt_token_ids": [11 + index % 2] * length, + "output_token_ids": [37], + "computed_tokens": length + (2 if decode else 0), + } + for index, length in enumerate(lengths) + ], + } + encoded = json.dumps(stream, sort_keys=True, separators=(",", ":")).encode() + lines.append(encoded) + row["real_kv_witness"] = { + "same_request": True, + "allocated_fake_tokens": 0, + "completed_seed_tokens": seed, + "token_stream_sha256": hashlib.sha256(encoded).hexdigest(), + } + raw = b"\n".join(lines) + b"\n" + sidecar = path.with_suffix(".token-streams.jsonl") + sidecar.write_bytes(raw) + payload["input_provenance"]["token_stream_manifest"] = { + "file": sidecar.name, + "sha256": hashlib.sha256(raw).hexdigest(), + "records": len(lines), + } + + +@pytest.mark.parametrize("marker", ["kvwarm_real_kv", "kvwarm_fake_fallback", None]) +def test_v41_cached_prefill_requires_real_computed_state(tmp_path, marker): + from dataclasses import replace + + from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS + + plan, cell, cell_dir = _synthetic_plan_and_cell(tmp_path) + identity = ("c" * 64, "full", "hbm_tp_sharded", "text") + cell = replace(cell, execution_identity=identity, input_text_sha256="a" * 64) + for path in (cell_dir / "raw").glob("*/benchmark*.json"): + payload = json.loads(path.read_text()) + payload["execution_identity"] = dict(zip(EXECUTION_COLUMNS, identity, strict=True)) + payload["execution_mode"] = "eager" + payload["input_provenance"] = { + "source": "tokenizer_text", + "text_sha256": "a" * 64, + "token_ids_sha256": "b" * 64, + "tokenizer_revision": "pinned", + "token_count": 100, + "unique_token_count": 20, + } + payload["kvwarm"] = {"enabled": True, "warm_eligible": True, "skip_reason": None} + + def mark(value): + if isinstance(value, dict): + if "point_type" in value: + value["sample_reasons"] = [marker] if marker else [] + for child in value.values(): + mark(child) + elif isinstance(value, list): + for child in value: + mark(child) + + mark(payload) + _write_v41_token_streams(payload, path) + path.write_text(json.dumps(payload)) + if marker == "kvwarm_real_kv": + assert aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt")[0]["kv_seed_regime"] == "real_kv" + else: + with pytest.raises(ValueError, match="requires real_kv"): + aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt") + + +@pytest.mark.parametrize("mode", [None, "PIECEWISE", "FULL", False]) +def test_v41_reader_rejects_unqualified_graph_or_missing_execution_mode(tmp_path, mode): + from collector.fpm_forward.native_artifact import _validate_execution_provenance + + from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS + + identity = ("c" * 64, "full", "hbm_tp_sharded", "text") + cell = SimpleNamespace(execution_identity=identity) + payload = {"execution_identity": dict(zip(EXECUTION_COLUMNS, identity, strict=True)), "execution_mode": mode} + with pytest.raises(ValueError, match="verified eager"): + _validate_execution_provenance(cell, payload, tmp_path / "rank.json") + + +@pytest.mark.parametrize("v41,eager", [(True, False), (False, True)]) +def test_explicit_eager_collection_admission(v41, eager, monkeypatch): + from collector.fpm_forward import planner + + monkeypatch.setattr(planner, "execution_identity", lambda *args, **kwargs: ("c" * 64 if v41 else "",)) + with pytest.raises(ValueError, match="eager"): + build_collection_plan( + backend="vllm", + model_path="nvidia/GLM-5.2-NVFP4", + system="b200_sxm", + selected_ops={"dsa_context_module", "dsa_generation_module"}, + options=FPMCollectionOptions.from_args(_args(fpm_enforce_eager=eager)), + ) + + +@pytest.mark.parametrize("corruption", ["missing", "tampered", "fake", "seed", "path", "coverage"]) +def test_v41_real_token_stream_validation_rejects_broken_witness(tmp_path, corruption): + from collector.fpm_forward.native_artifact import _validate_token_streams + + _plan, _cell, cell_dir = _synthetic_plan_and_cell(tmp_path) + path = next((cell_dir / "raw").glob("*/benchmark*.json")) + payload = json.loads(path.read_text()) + payload["input_provenance"] = {} + _write_v41_token_streams(payload, path) + _validate_token_streams(payload, path) + manifest = payload["input_provenance"]["token_stream_manifest"] + if corruption == "missing": + del payload["input_provenance"]["token_stream_manifest"] + elif corruption == "tampered": + path.with_name(manifest["file"]).write_text("changed") + elif corruption == "fake": + payload["results"][0]["real_kv_witness"]["allocated_fake_tokens"] = 1 + elif corruption == "seed": + payload["results"][0]["real_kv_witness"]["completed_seed_tokens"] -= 1 + elif corruption == "path": + manifest["file"] = "../outside.token-streams.jsonl" + else: + manifest["records"] += 1 + with pytest.raises(ValueError, match="V4.1"): + _validate_token_streams(payload, path) + + +@pytest.mark.parametrize( + "corruption", [None, "role", "missing_result", "missing_expected", "measured_warmup", "legacy"] +) +def test_v41_eager_warmup_histories_are_preserved_but_not_measured(tmp_path, corruption): + from copy import deepcopy + + from collector.fpm_forward.native_artifact import _validate_token_streams + + _plan, _cell, cell_dir = _synthetic_plan_and_cell(tmp_path) + path = next((cell_dir / "raw").glob("*/benchmark*.json")) + payload = json.loads(path.read_text()) + payload["input_provenance"] = {} + warmup = deepcopy(payload["results"][0]) + warmup_id = len(payload["results"]) + 1 + warmup["point"].update(benchmark_id=warmup_id, sample_reasons=["eager_warmup"]) + payload["results"].append(warmup) + _write_v41_token_streams(payload, path) + payload["warmup_results"] = [payload["results"].pop()] + manifest = payload["input_provenance"]["token_stream_manifest"] + manifest.update(schema_version=2, warmup_benchmark_ids=[warmup_id]) + sidecar = path.with_name(manifest["file"]) + all_rows = payload["results"] + payload["warmup_results"] + encoded = [] + for line, row in zip(sidecar.read_bytes().splitlines(), all_rows, strict=True): + stream = json.loads(line) + stream["sampling_role"] = "warmup" if stream["benchmark_id"] == warmup_id else "measurement" + if corruption == "role" and stream["benchmark_id"] == warmup_id: + stream["sampling_role"] = "measurement" + changed = json.dumps(stream, sort_keys=True, separators=(",", ":")).encode() + encoded.append(changed) + row["real_kv_witness"]["token_stream_sha256"] = hashlib.sha256(changed).hexdigest() + raw = b"\n".join(encoded) + b"\n" + sidecar.write_bytes(raw) + manifest["sha256"] = hashlib.sha256(raw).hexdigest() + if corruption == "missing_result": + payload["warmup_results"] = [] + elif corruption == "missing_expected": + manifest["warmup_benchmark_ids"] = [] + elif corruption == "measured_warmup": + payload["results"].append(payload["warmup_results"].pop()) + elif corruption == "legacy": + manifest["schema_version"] = 1 + if corruption is None: + _validate_token_streams(payload, path) + assert len(payload["results"]) + 1 == manifest["records"] + else: + with pytest.raises(ValueError, match="V4.1"): + _validate_token_streams(payload, path) + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +@pytest.mark.parametrize("smoke", [False, True]) +def test_v41_native_grid_bounds_reach_both_runtime_phases(phase, smoke): + from collector.fpm_forward.runner import _cell_generator_overrides + + plan = _args_plan() + plan.capability = SimpleNamespace(architecture="DeepseekV41ForCausalLM") + plan.options = FPMCollectionOptions.from_args( + _args( + fpm_max_prefill_isl=64, + fpm_max_prefill_batch_size=1, + fpm_max_decode_batch_size=1, + fpm_max_model_len=258, + fpm_warmup_iterations=0, + ) + ) + generated = _cell_generator_overrides(plan, _args_cell(phase, "pure_tp"), {}, smoke=smoke) + args = generated["params"]["agg"]["extra_cli_args"] + assert args[args.index("--max-num-batched-tokens") + 1] == "64" + assert args[args.index("--max-num-seqs") + 1] == "1" + assert '--engram-config={"cpu_offload":false}' in args + + +@pytest.mark.parametrize( + "corruption", + [ + "results_missing", + "results_null", + "row_null", + "point_missing", + "point_null", + "groups_missing", + "provenance_missing", + ], +) +def test_v41_truncated_native_artifact_raises_actionable_value_error(tmp_path, corruption): + from dataclasses import replace + + from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS + + plan, cell, cell_dir = _synthetic_plan_and_cell(tmp_path) + identity = ("c" * 64, "full", "hbm_tp_sharded", "text") + cell = replace(cell, execution_identity=identity, input_text_sha256="a" * 64) + for path in (cell_dir / "raw").glob("*/benchmark*.json"): + payload = json.loads(path.read_text()) + payload["execution_identity"] = dict(zip(EXECUTION_COLUMNS, identity, strict=True)) + payload["execution_mode"] = "eager" + payload["input_provenance"] = { + "source": "tokenizer_text", + "text_sha256": "a" * 64, + "token_ids_sha256": "b" * 64, + "tokenizer_revision": "pinned", + "token_count": 100, + "unique_token_count": 20, + } + _write_v41_token_streams(payload, path) + if corruption == "results_missing": + del payload["results"] + elif corruption == "results_null": + payload["results"] = None + elif corruption == "row_null": + payload["results"][0] = None + elif corruption == "point_missing": + del payload["results"][0]["point"] + elif corruption == "point_null": + payload["results"][0]["point"] = None + elif corruption == "groups_missing": + del payload["iteration_groups"] + else: + del payload["input_provenance"] + path.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="native"): + aggregate_cell(plan, cell, cell_dir, expected_attempt_id="attempt") diff --git a/python/aisimulate/tests/unit/collector/test_fpm_runner.py b/python/aisimulate/tests/unit/collector/test_fpm_runner.py index a04e6137f..676839dc9 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_runner.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_runner.py @@ -11,6 +11,7 @@ import shutil import signal import subprocess +import sys import time from pathlib import Path from types import SimpleNamespace @@ -1256,6 +1257,7 @@ def render_cell(*args, **kwargs): (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") (cell_dir / "run.sh").write_text("#!/bin/sh\n") (cell_dir / "fpm_env.sh").write_text("#!/bin/sh\n") + (cell_dir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") class FakeResource: def __init__(self, _manifest, _cell_dir): @@ -1314,9 +1316,9 @@ def cleanup(self): # fresh checkpoint proves nothing about the cluster) and one after. assert events.count("cleanup") == 2 assert events.index("cleanup") < events.index("apply") - # Contract: the staged set is exactly the two rendered runtime artifacts + # Contract: the staged set includes the rendered runtime artifacts # plus the collector's own in-pod runtime and preflight. - assert set(staged_names) == {"run.sh", "fpm_env.sh", "fpm_exec.sh", "preflight.py"} + assert set(staged_names) == {"run.sh", "fpm_env.sh", "collector-runtime-env.sh", "fpm_exec.sh", "preflight.py"} assert "cases.json" not in staged_names assert "fpm_scheduler.py" not in staged_names assert "run_with_etcd.sh" not in staged_names @@ -1341,6 +1343,7 @@ def render_cell(*args, **_kwargs): (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") (cell_dir / "run.sh").write_text("#!/bin/sh\n") (cell_dir / "fpm_env.sh").write_text("#!/bin/sh\n") + (cell_dir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") class FakeResource: def __init__(self, _manifest, _cell_dir): @@ -1548,6 +1551,7 @@ def render_cell(*args, **kwargs): (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") (cell_dir / "run.sh").write_text("#!/bin/sh\n") (cell_dir / "fpm_env.sh").write_text("#!/bin/sh\n") + (cell_dir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") class FakeResource: def __init__(self, _manifest, _cell_dir): @@ -1621,6 +1625,7 @@ def test_typed_generator_render_uses_collector_prefill_axis(tmp_path): base = { "K8sConfig": { "k8s_image": "nvcr.io/nvidia/ai-dynamo/vllm-runtime:test", + "extra_env": [{"name": "FPM_READINESS_TIMEOUT_SECONDS", "value": "600"}], "k8s_pvc_mount_path": "/model-cache", "k8s_model_path_in_pvc": "models--nvidia--GLM-5.2-NVFP4", } @@ -1632,7 +1637,12 @@ def test_typed_generator_render_uses_collector_prefill_axis(tmp_path): for artifact in ("k8s_deploy.yaml", "fpm_env.sh", "run.sh"): assert (tmp_path / artifact).exists(), artifact + startup = (tmp_path / "collector-runtime-env.sh").read_text() + assert "export FPM_READINESS_TIMEOUT_SECONDS=600" in startup + # The keepalive Pod intentionally has no engine environment. Staging + # supplies the startup file before either transport launches preflight. script = (tmp_path / "run.sh").read_text() + assert "export FPM_READINESS_TIMEOUT_SECONDS=600" in script assert "--benchmark-mode prefill" in script assert "--benchmark-warmup-iterations 3" in script assert "--scheduler-cls fpm_scheduler" not in script @@ -2025,6 +2035,7 @@ def _running_cell_fixture(tmp_path, plan, cell): (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") (cell_dir / "run.sh").write_text("#!/bin/sh\n") (cell_dir / "fpm_env.sh").write_text("#!/bin/sh\n") + (cell_dir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") checkpoint_dir = tmp_path / "checkpoints" checkpoint_dir.mkdir() @@ -2226,6 +2237,7 @@ def render_cell(*args, **kwargs): (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") (cell_dir / "run.sh").write_text("#!/bin/sh\n") (cell_dir / "fpm_env.sh").write_text("#!/bin/sh\n") + (cell_dir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") applied = [] @@ -2368,7 +2380,7 @@ def test_completed_formal_database_is_terminal_after_commit_validation(monkeypat ) validated = [] - def validate(parquet_path, metadata_path, plan_arg): + def validate(parquet_path, metadata_path, plan_arg, **_kwargs): validated.append((parquet_path, metadata_path, plan_arg.sha256)) return {"schema_version": 6} @@ -2396,6 +2408,236 @@ def forbid_republication(*_args, **_kwargs): assert manifest["attempts"] == [] +@pytest.fixture +def completed_explicit_campaign(monkeypatch, tmp_path): + """Complete real receipt validation, native aggregation and publication on CPU.""" + cell = _cell() + plan = _plan(cell) + points = { + "schema_version": 3, + "prefill": [{"batch_size": 4, "total_prefill_tokens": 257, "total_kv_read_tokens": 128}], + "decode": [], + } + plan.options.benchmark_points_json = json.dumps(points, sort_keys=True, separators=(",", ":")) + plan.options.benchmark_points_sha256 = hashlib.sha256(plan.options.benchmark_points_json.encode()).hexdigest() + runs = [] + + def render_cell(_plan, _cell, cell_dir, _overrides, **_kwargs): + (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") + for name in ("run.sh", "fpm_env.sh", "collector-runtime-env.sh"): + (cell_dir / name).write_text("#!/bin/sh\n") + + class LocalResource: + def __init__(self, _manifest, cell_dir): + self.cell_dir = cell_dir + self.raw = cell_dir / "raw" / "pod-0" + + def apply(self): + self.raw.mkdir(parents=True) + + def wait_ready(self, _expected_nodes): + return ["pod-0"] + + def stage(self, _pods, _files): + pass + + def prepare_attempt(self, _pods, **identity): + _write_provenance(self.raw / "collector-provenance.json", **identity) + + def _exec(self, _pod, command, *, timeout): + # Run the actual before/after receipt producer with local paths. + return subprocess.run( + [ + sys.executable, + *command[1:-2], + str(self.cell_dir / fpm_runner.POINTS_FILENAME), + str(self.raw / fpm_runner.POINTS_RECEIPT_FILENAME), + ], + timeout=timeout, + check=True, + capture_output=True, + text=True, + ) + + def execute(self, _pods): + runs.append("execute") + (self.raw / "benchmark.json").write_text(json.dumps(_native_payload(phase="prefill", rank=0, dp=1))) + + def collect(self, _pods, *, require_benchmark=True): + pass + + def cleanup(self): + pass + + monkeypatch.setattr(fpm_runner, "_render_cell", render_cell) + monkeypatch.setattr(fpm_runner, "KubernetesCellRunner", LocalResource) + checkpoint_dir = tmp_path / "checkpoints" + checkpoint_dir.mkdir() + kwargs = { + "generator_overrides": {}, + "checkpoint_dir": str(checkpoint_dir), + "artifact_root": str(tmp_path / "artifacts"), + "retry_failed": False, + "database_root": str(tmp_path / "db"), + } + assert run_collection(plan, resume=False, **kwargs) == [] + checkpoint_path = checkpoint_dir / "fpm_forward.json" + checkpoint = json.loads(checkpoint_path.read_text()) + assert checkpoint["cells"][cell.cell_id]["status"] == "passed" + assert checkpoint["database"]["status"] == "passed" + from collector.fpm_forward.database import validate_formal_database_commit + + database = checkpoint["database"] + parquet, metadata = Path(database["parquet"]), Path(database["metadata"]) + assert validate_formal_database_commit(parquet, metadata, plan)["row_count"] == 1 + return SimpleNamespace( + plan=plan, + cell=cell, + kwargs=kwargs, + checkpoint_path=checkpoint_path, + parquet=parquet, + metadata=metadata, + runs=runs, + ) + + +@pytest.mark.parametrize("different_artifact_root", [False, True]) +def test_explicit_terminal_database_resumes_after_raw_reclamation(completed_explicit_campaign, different_artifact_root): + campaign = completed_explicit_campaign + committed = (campaign.parquet.read_bytes(), campaign.metadata.read_bytes()) + original_checkpoint = campaign.checkpoint_path.read_bytes() + root = Path(campaign.kwargs["artifact_root"]) + shutil.rmtree(root / campaign.plan.sha256[:16] / "cells" / campaign.cell.cell_id / "raw") + if different_artifact_root: + root = root.with_name("new-artifacts") + campaign.kwargs["artifact_root"] = str(root) + + assert run_collection(campaign.plan, resume=True, **campaign.kwargs) == [] + + assert campaign.runs == ["execute"] + assert campaign.checkpoint_path.read_bytes() == original_checkpoint + assert (campaign.parquet.read_bytes(), campaign.metadata.read_bytes()) == committed + assert json.loads((root / campaign.plan.sha256[:16] / "run-manifest.json").read_text())["attempts"] == [] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("cell_id", "different-cell"), + ("source_plan_sha256", "different-plan"), + ("collector_attempt_id", "different-attempt"), + ("tp", 8), + ("total_prefill_tokens", 258), + ("latency_ms", 999.0), + ], +) +def test_terminal_resume_rejects_resealed_foreign_rows(completed_explicit_campaign, field, value, caplog): + import pyarrow as pa + import pyarrow.parquet as pq + + campaign = completed_explicit_campaign + rows = pq.read_table(campaign.parquet).to_pylist() + rows[0][field] = value + pq.write_table(pa.Table.from_pylist(rows), campaign.parquet) + metadata = json.loads(campaign.metadata.read_text()) + metadata["parquet_sha256"] = hashlib.sha256(campaign.parquet.read_bytes()).hexdigest() + campaign.metadata.write_text(json.dumps(metadata)) + raw = Path(campaign.kwargs["artifact_root"]) / campaign.plan.sha256[:16] / "cells" / campaign.cell.cell_id / "raw" + shutil.rmtree(raw) + + with pytest.raises(ValueError, match="explicit benchmark points have no runtime receipt owners"): + run_collection(campaign.plan, resume=True, **campaign.kwargs) + assert "Completed FPM database checkpoint failed validation" in caplog.text + assert campaign.runs == ["execute"] + + +def test_terminal_resume_preserves_complete_first_publisher_reuse(completed_explicit_campaign): + campaign = completed_explicit_campaign + committed = (campaign.parquet.read_bytes(), campaign.metadata.read_bytes()) + original = json.loads(campaign.checkpoint_path.read_text()) + campaign.plan.sha256 = "another-plan-sha" + assert run_collection(campaign.plan, resume=False, **campaign.kwargs) == [] + repeated = json.loads(campaign.checkpoint_path.read_text()) + assert repeated["database"]["skipped_first_publisher_wins"] == [campaign.cell.cell_id] + assert ( + repeated["cells"][campaign.cell.cell_id]["attempt_id"] != original["cells"][campaign.cell.cell_id]["attempt_id"] + ) + raw = Path(campaign.kwargs["artifact_root"]) / campaign.plan.sha256[:16] / "cells" / campaign.cell.cell_id / "raw" + shutil.rmtree(raw) + + assert run_collection(campaign.plan, resume=True, **campaign.kwargs) == [] + assert campaign.runs == ["execute", "execute"] + assert (campaign.parquet.read_bytes(), campaign.metadata.read_bytes()) == committed + + +def test_legacy_terminal_explicit_checkpoint_checks_requested_coordinates(completed_explicit_campaign, caplog): + import pyarrow as pa + import pyarrow.parquet as pq + + campaign = completed_explicit_campaign + checkpoint = json.loads(campaign.checkpoint_path.read_text()) + del checkpoint["database"]["cell_rows"] + campaign.checkpoint_path.write_text(json.dumps(checkpoint)) + raw = Path(campaign.kwargs["artifact_root"]) / campaign.plan.sha256[:16] / "cells" / campaign.cell.cell_id / "raw" + shutil.rmtree(raw) + assert run_collection(campaign.plan, resume=True, **campaign.kwargs) == [] + + rows = pq.read_table(campaign.parquet).to_pylist() + rows[0]["total_prefill_tokens"] += 1 + pq.write_table(pa.Table.from_pylist(rows), campaign.parquet) + metadata = json.loads(campaign.metadata.read_text()) + metadata["parquet_sha256"] = hashlib.sha256(campaign.parquet.read_bytes()).hexdigest() + campaign.metadata.write_text(json.dumps(metadata)) + with pytest.raises(ValueError, match="explicit benchmark points have no runtime receipt owners"): + run_collection(campaign.plan, resume=True, **campaign.kwargs) + assert "does not cover requested coordinates" in caplog.text + + +@pytest.mark.parametrize("database_status", ["unpublished", "invalid_commit"]) +def test_explicit_resume_still_requires_raw_receipts_without_valid_commit(completed_explicit_campaign, database_status): + campaign = completed_explicit_campaign + if database_status == "unpublished": + checkpoint = json.loads(campaign.checkpoint_path.read_text()) + del checkpoint["database"] + campaign.checkpoint_path.write_text(json.dumps(checkpoint)) + else: + metadata = json.loads(campaign.metadata.read_text()) + metadata["parquet_sha256"] = "invalid" + campaign.metadata.write_text(json.dumps(metadata)) + raw = Path(campaign.kwargs["artifact_root"]) / campaign.plan.sha256[:16] / "cells" / campaign.cell.cell_id / "raw" + shutil.rmtree(raw) + + with pytest.raises(ValueError, match="explicit benchmark points have no runtime receipt owners"): + run_collection(campaign.plan, resume=True, **campaign.kwargs) + assert campaign.runs == ["execute"] + + +@pytest.mark.parametrize("status", ["passed", "failed"]) +def test_explicit_recovery_and_publication_reject_wrong_attempt_receipt(completed_explicit_campaign, status): + campaign = completed_explicit_campaign + checkpoint = json.loads(campaign.checkpoint_path.read_text()) + del checkpoint["database"] + checkpoint["cells"][campaign.cell.cell_id]["status"] = status + campaign.checkpoint_path.write_text(json.dumps(checkpoint)) + cell_dir = Path(campaign.kwargs["artifact_root"]) / campaign.plan.sha256[:16] / "cells" / campaign.cell.cell_id + receipt_path = cell_dir / "raw" / "pod-0" / fpm_runner.POINTS_RECEIPT_FILENAME + receipt = json.loads(receipt_path.read_text()) + receipt["attempt_id"] = "another-attempt" + receipt_path.write_text(json.dumps(receipt)) + + if status == "passed": + with pytest.raises(ValueError, match="runtime benchmark-points receipt mismatch"): + run_collection(campaign.plan, resume=True, **campaign.kwargs) + else: + errors = run_collection(campaign.plan, resume=True, **campaign.kwargs) + assert [error["classification"] for error in errors] == ["campaign_incomplete"] + resumed = json.loads(campaign.checkpoint_path.read_text()) + assert resumed["cells"][campaign.cell.cell_id]["status"] == status + assert "artifact_recovery" not in resumed["cells"][campaign.cell.cell_id] + assert resumed.get("database", {}).get("status") != "passed" + assert campaign.runs == ["execute"] + + def test_user_extra_env_reaches_the_render_request_alongside_collector_identities(): cell = _cell() plan = _plan(cell) @@ -2458,18 +2700,15 @@ def test_publish_partial_ships_passed_cells_and_records_the_missing(tmp_path, mo ) ) - # The database writer has its own coverage; stub it so this test isolates - # the publication gate and the aggregation over passed cells. + # Keep the publication gate observable while exercising its sealed + # per-cell checkpoint proof with the real database writer. + from collector.fpm_forward.database import write_formal_database + published = {} def fake_writer(plan_arg, rows, *, systems_root=None): published["rows"] = rows - parquet = tmp_path / "db" / "fpm_forward_perf.parquet" - parquet.parent.mkdir(parents=True, exist_ok=True) - parquet.write_bytes(b"parquet") - metadata = parquet.with_suffix(".metadata.json") - metadata.write_text("{}") - return parquet, metadata, () + return write_formal_database(plan_arg, rows, systems_root=systems_root) monkeypatch.setattr("collector.fpm_forward.database.write_formal_database", fake_writer) @@ -2603,6 +2842,7 @@ def render_cell(*args, **_kwargs): (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") (cell_dir / "run.sh").write_text("#!/bin/sh\n") (cell_dir / "fpm_env.sh").write_text("#!/bin/sh\n") + (cell_dir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") class FakeResource: def __init__(self, _manifest, _cell_dir): @@ -2683,6 +2923,7 @@ def render_cell(*args, **_kwargs): (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: cell\n") (cell_dir / "run.sh").write_text("#!/bin/sh\n") (cell_dir / "fpm_env.sh").write_text("#!/bin/sh\n") + (cell_dir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") class FakeResource: def __init__(self, _manifest, _cell_dir): @@ -2720,3 +2961,83 @@ def wait_ready(self, _expected_nodes): assert attempt["collector_phase_seconds"].keys() == {"render_s"} checkpoint = json.loads((checkpoint_dir / "fpm_forward_smoke.json").read_text()) assert checkpoint["cells"][cell.cell_id]["collector_phase_seconds"].keys() == {"render_s"} + + +@pytest.mark.parametrize("seconds", ["0", "-1", "1.5", "3601", "nan", True]) +def test_runtime_readiness_configuration_rejects_invalid_budget(seconds): + cell = _cell() + with pytest.raises(ValueError, match="FPM_READINESS_TIMEOUT_SECONDS"): + _cell_generator_overrides( + _plan(cell), + cell, + {"K8sConfig": {"extra_env": [{"name": "FPM_READINESS_TIMEOUT_SECONDS", "value": seconds}]}}, + ) + + +def test_runtime_environment_uses_configured_path_and_budget_before_engine(tmp_path): + cell = _cell() + plan = _plan(cell) + plan.capability = SimpleNamespace(architecture="DeepseekV41ForCausalLM") + plan.options.decoder_replay = False + settings = { + "K8sConfig": { + "extra_env": [ + {"name": "PYTHONPATH", "value": "/custom/runtime components/src"}, + {"name": "FPM_READINESS_TIMEOUT_SECONDS", "value": "600"}, + ] + } + } + overrides = _cell_generator_overrides(plan, cell, settings) + fpm_runner._write_runtime_environment(tmp_path, overrides) + actual = subprocess.run( + [ + "bash", + "-c", + 'source "$1"; printf "%s\\n%s\\n" "$PYTHONPATH" "$FPM_READINESS_TIMEOUT_SECONDS"', + "bash", + str(tmp_path / fpm_runner.RUNTIME_ENV_FILENAME), + ], + check=True, + capture_output=True, + text=True, + ) + assert actual.stdout.splitlines() == ["/tmp/fpm-bench:/custom/runtime components/src", "600"] + env = {item["name"]: item["value"] for item in overrides["K8sConfig"]["extra_env"]} + assert env["PYTHONPATH"] == actual.stdout.splitlines()[0] + assert env["DYN_FPM_DSV41_REAL_KV"] == "1" + + +def test_runtime_default_path_is_adapter_owned(): + cell = _cell() + plan = _plan(cell) + plan.capability = SimpleNamespace(architecture="DeepseekV41ForCausalLM") + plan.options.decoder_replay = False + overrides = _cell_generator_overrides(plan, cell, {}) + env = {item["name"]: item["value"] for item in overrides["K8sConfig"]["extra_env"]} + adapter = Path(fpm_runner.__file__).parent / "runtime/dsv41/runtime-paths.json" + assert env["PYTHONPATH"] == "/tmp/fpm-bench:" + json.loads(adapter.read_text())["python_path"] + assert env["FPM_READINESS_TIMEOUT_SECONDS"] == "900" + + +def test_kubernetes_stages_exact_startup_configuration(tmp_path): + runner = _runner(tmp_path) + source = tmp_path / "collector-runtime-env.sh" + source.write_text("export FPM_READINESS_TIMEOUT_SECONDS=600\n") + remote = tmp_path / "remote" + remote.mkdir() + + def kubectl(*args, **kwargs): + assert args[0] == "cp" + shutil.copy2(args[1], remote / Path(args[1]).name) + return subprocess.CompletedProcess(args, 0, stdout="", stderr="") + + def execute(pod, command, *, timeout): + if command[0] == "mkdir": + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + rewritten = [arg.replace(REMOTE_WORKDIR, str(remote)) for arg in command] + return subprocess.run(rewritten, check=True, capture_output=True, text=True, timeout=timeout) + + runner._kubectl = kubectl + runner._exec_checked = execute + runner.stage(["pod-0"], [source]) + assert (remote / source.name).read_bytes() == source.read_bytes() diff --git a/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py b/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py index 9419f7e52..9e4f86cd8 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_runtime_wrapper.py @@ -112,6 +112,7 @@ def _stage( exports = "\n".join(f"export {name}={shlex.quote(str(values[name]))}" for name in FPM_ENV_EXPORTED_VARS) env_script = f"#!/usr/bin/env bash\n{exports}\n" (workdir / "fpm_env.sh").write_text(env_script) + (workdir / "collector-runtime-env.sh").write_text("export FPM_READINESS_TIMEOUT_SECONDS=900\n") (workdir / "preflight.py").write_text("") (workdir / "run.sh").write_text(run_script) _write_executable(bin_dir / "etcd", _ETCD_STUB) @@ -149,7 +150,7 @@ def test_fpm_exec_starts_leader_etcd_before_preflight(): script = FPM_EXEC.read_text() assert script.index("etcd_pid=$!") < script.index('python3 "${workdir}/preflight.py"') - assert "time.monotonic() + 120" in script + assert "time.monotonic() + float(sys.argv[2])" in script def test_fpm_exec_leader_starts_etcd_and_cleanup_stops_it(tmp_path): @@ -304,3 +305,26 @@ def test_fpm_exec_propagates_fail_closed_env_source(tmp_path): assert "requires rank and leader discovery" in completed.stderr assert not staged.etcd_trace.exists() assert not engine_trace.exists() + + +@pytest.mark.parametrize("failure", [RuntimeError("pinned source mismatch"), ValueError("adapter activation failed")]) +def test_preflight_preserves_non_importerror_runtime_failure(tmp_path, monkeypatch, failure): + import builtins + + from collector.fpm_forward.runtime import preflight + + original_import = builtins.__import__ + + def fail(name, *args, **kwargs): + if name == "dynamo.vllm.instrumented_scheduler": + raise failure + return original_import(name, *args, **kwargs) + + audit = tmp_path / "runtime-preflight.json" + monkeypatch.setattr(preflight, "_AUDIT_PATH", audit) + monkeypatch.setattr(builtins, "__import__", fail) + with pytest.raises(RuntimeError, match="runtime activation or identity preflight failed") as caught: + preflight.main() + assert caught.value.__cause__ is failure + receipt = json.loads(audit.read_text()) + assert receipt["status"] == "failed" and receipt["import_error"] == str(failure) diff --git a/python/aisimulate/tests/unit/collector/test_fpm_slurm.py b/python/aisimulate/tests/unit/collector/test_fpm_slurm.py new file mode 100644 index 000000000..fea09b50e --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_slurm.py @@ -0,0 +1,550 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +import threading +import time +from types import SimpleNamespace + +import pytest +from collector.fpm_forward.slurm import SlurmCellRunner + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def runner(tmp_path, monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "1234") + monkeypatch.setenv("SLURM_JOB_NODELIST", "test-node") + manifest = tmp_path / "manifest.yaml" + manifest.write_text( + json.dumps( + { + "kind": "LeaderWorkerSet", + "metadata": {"name": "cell"}, + "spec": {"replicas": 1, "leaderWorkerTemplate": {"size": 1}}, + } + ) + ) + return SlurmCellRunner(manifest, tmp_path, image="image@sha256:abc", mounts=("/cache:/cache",), total_gpus=4) + + +def test_slurm_requires_scheduler_allocation(runner, monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID") + with pytest.raises(ValueError, match="existing sbatch/salloc"): + SlurmCellRunner(runner.cell_dir / "manifest.yaml", runner.cell_dir, image=runner.image, mounts=(), total_gpus=4) + + +def test_slurm_stage_and_argv_keep_shared_result_unit_identity(runner, monkeypatch): + commands = [] + + def command(args, **kwargs): + commands.append(args) + return SimpleNamespace( + stdout="JobId=1234 JobState=RUNNING NodeList=test-node" if "job" in args else "test-node\n", stderr="" + ) + + monkeypatch.setattr(runner, "_command", command) + units = runner.wait_ready(1) + source = runner.cell_dir / "fpm_exec.sh" + source.write_text("exit 0\n") + startup = runner.cell_dir / "collector-runtime-env.sh" + startup.write_text("export FPM_READINESS_TIMEOUT_SECONDS=600\n") + runner.stage(units, [source, startup]) + assert (runner.cell_dir / "slurm-runtime" / startup.name).read_bytes() == startup.read_bytes() + runner._exec(units[0], ["bash", "/tmp/fpm-bench/fpm_exec.sh"], timeout=10) + assert units == ["node0000"] + assert (runner.cell_dir / "slurm-runtime" / source.name).read_text() == source.read_text() + argv = commands[-1] + assert "--jobid=1234" in argv and "--gpus-per-node=4" in argv + assert "FPM_NODE_RANK=0" in argv and "FPM_MASTER_ADDR=test-node" in argv + assert f"{runner.cell_dir}/raw/node0000:/results" in next(a for a in argv if a.startswith("--container-mounts=")) + + +def test_slurm_cleanup_cancels_only_receipted_steps_and_verifies_exit(runner, monkeypatch): + runner.owner_path.parent.mkdir(parents=True, exist_ok=True) + runner.owner_path.write_text(json.dumps({"job_id": "1233", "step_name": runner.step_name})) + commands = [] + snapshots = iter( + [f"1234.2|{runner.step_name}\n1233.1|{runner.step_name}\n1234.4|other\n1234.batch|{runner.step_name}\n", ""] + ) + + def command(args, **kwargs): + commands.append(args) + return SimpleNamespace(stdout=next(snapshots) if args[0] == "squeue" else "", stderr="") + + monkeypatch.setattr(runner, "_command", command) + runner.cleanup() + assert [args for args in commands if args[0] == "scancel"] == [["scancel", "1234.2"], ["scancel", "1233.1"]] + assert len([args for args in commands if args[0] == "squeue"]) == 2 + + +def test_slurm_cleanup_reports_leaked_steps(runner, monkeypatch): + monkeypatch.setattr( + runner, "_command", lambda *a, **k: SimpleNamespace(stdout=f"1234.2|{runner.step_name}", stderr="") + ) + clock = iter([0, 61]) + monkeypatch.setattr("collector.fpm_forward.slurm.time.monotonic", lambda: next(clock)) + with pytest.raises(RuntimeError, match="remain after cleanup"): + runner.cleanup() + + +def test_slurm_refuses_allocation_geometry_mismatch(runner, monkeypatch): + monkeypatch.setattr( + runner, + "_command", + lambda args, **k: SimpleNamespace( + stdout="JobId=1234 JobState=RUNNING NodeList=node-a,node-b" if "job" in args else "node-a node-b" + ), + ) + with pytest.raises(ValueError, match="exactly 1 allocated nodes"): + runner.wait_ready(1) + + +def test_preparation_preserves_slurm_failure_streams_across_retries(runner, monkeypatch): + runner.hosts = ["test-node"] + errors = [ + subprocess.CalledProcessError(1, ["srun"], output="preparation started", stderr="task resource conflict"), + subprocess.TimeoutExpired(["srun"], 300, output=b"waiting", stderr=b"container startup stalled"), + ] + for failure in errors: + + def fail(*args, **kwargs): + raise failure + + monkeypatch.setattr("collector.fpm_forward.runner._run_command", fail) + with pytest.raises(type(failure)) as caught: + runner.prepare_attempt(runner.pods(), cell_id="cell", plan_sha256="plan", attempt_id="attempt") + assert caught.value is failure + records = list((runner.cell_dir / "logs" / "transport-failures").iterdir()) + assert len(records) == 2 + assert {path.joinpath("stderr.log").read_text() for path in records} == { + "task resource conflict", + "container startup stalled", + } + assert {path.joinpath("stdout.log").read_text() for path in records} == {"preparation started", "waiting"} + assert all(json.loads(path.joinpath("failure.json").read_text())["executable"] == "srun" for path in records) + + +@pytest.mark.parametrize("timeout", [0, -1, float("inf"), float("nan"), True]) +def test_wait_ready_rejects_unbounded_or_invalid_timeout(runner, timeout): + with pytest.raises(ValueError, match="finite and positive"): + runner.wait_ready(1, timeout_seconds=timeout) + + +def test_wait_ready_waits_for_allocation_and_shares_one_deadline(runner, monkeypatch): + clock = [0.0] + commands = [] + states = iter(["PENDING", "CONFIGURING", "RUNNING"]) + monkeypatch.setattr("collector.fpm_forward.slurm.time.monotonic", lambda: clock[0]) + monkeypatch.setattr( + "collector.fpm_forward.slurm.time.sleep", lambda seconds: clock.__setitem__(0, clock[0] + seconds) + ) + + def command(args, *, timeout): + commands.append((args, timeout)) + clock[0] += 0.25 + text = f"JobId=1234 JobState={next(states)} NodeList=test-node" if "job" in args else "test-node" + return SimpleNamespace(stdout=text) + + monkeypatch.setattr(runner, "_command", command) + assert runner.wait_ready(1, timeout_seconds=5) == ["node0000"] + assert [timeout for _, timeout in commands] == [5, 3.75, 2.5, 2.25] + assert commands[-1][0] == ["scontrol", "show", "hostnames", "test-node"] + + +def test_wait_ready_pending_allocation_expires_without_adopting_nodes(runner, monkeypatch): + clock = [0.0] + monkeypatch.setattr("collector.fpm_forward.slurm.time.monotonic", lambda: clock[0]) + monkeypatch.setattr( + "collector.fpm_forward.slurm.time.sleep", lambda seconds: clock.__setitem__(0, clock[0] + seconds) + ) + monkeypatch.setattr( + runner, "_command", lambda *args, **kwargs: SimpleNamespace(stdout="JobId=1234 JobState=PENDING") + ) + with pytest.raises(TimeoutError, match="not ready before deadline: PENDING"): + runner.wait_ready(1, timeout_seconds=0.5) + assert clock[0] == 0.5 and runner.hosts == [] + + +@pytest.mark.parametrize("snapshot", ["JobId=1234 JobState=FAILED", "JobId=4567 JobState=RUNNING NodeList=test-node"]) +def test_wait_ready_rejects_terminal_or_foreign_allocation(runner, monkeypatch, snapshot): + monkeypatch.setattr(runner, "_command", lambda *args, **kwargs: SimpleNamespace(stdout=snapshot)) + with pytest.raises((ValueError, RuntimeError)): + runner.wait_ready(1) + assert runner.hosts == [] + + +@pytest.mark.parametrize("interrupt", [signal.SIGINT, signal.SIGTERM]) +@pytest.mark.parametrize("child_mode", ["cooperative", "ignore_term", "pipe_descendant"]) +@pytest.mark.skipif(os.name != "posix", reason="POSIX transport process groups") +def test_execute_interrupt_stops_live_children_before_join(runner, monkeypatch, interrupt, child_mode): + from collector.fpm_forward import runner as campaign + + stopped = threading.Event() + children = [] + signalled_at = [] + readiness = runner.cell_dir / "child-ready" + monkeypatch.setattr(campaign, "_COMMAND_TERMINATION_GRACE_SECONDS", 0.2) + + def live_command(_unit, _command, *, timeout): + child = ( + "import os,pathlib,signal,sys,time; " + + ("signal.signal(signal.SIGTERM,signal.SIG_IGN); " if child_mode != "cooperative" else "") + + "pathlib.Path(sys.argv[1]).write_text(str(os.getpid())); time.sleep(30)" + ) + if child_mode == "pipe_descendant": + # The direct transport exits on TERM; its descendant ignores TERM + # and retains stdout/stderr, which would block communicate/join. + child = ( + "import subprocess,sys,time; " + f"subprocess.Popen([sys.executable,'-c',{child!r},sys.argv[1]]); time.sleep(30)" + ) + return campaign._run_command([sys.executable, "-c", child, str(readiness)], timeout=timeout) + + def force_cleanup(): + # Keep a broken implementation bounded without trusting its helper. + for process in children: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.kill() + if readiness.exists(): + try: + os.kill(int(readiness.read_text()), signal.SIGKILL) + except ProcessLookupError: + pass + + def interrupt_after_spawn(): + deadline = time.monotonic() + 5 + while not stopped.wait(0.01): + with campaign._ACTIVE_COMMANDS_LOCK: + children[:] = campaign._ACTIVE_COMMANDS + if children and readiness.exists() and readiness.read_text(): + signalled_at.append(time.monotonic()) + os.kill(os.getpid(), interrupt) + # Bound the failing regression as well: the broken executor + # otherwise waits for the child's entire 30-second lifetime. + if not stopped.wait(5): + force_cleanup() + return + if time.monotonic() >= deadline: + return + + monkeypatch.setattr(runner, "_exec", live_command) + interrupter = threading.Thread(target=interrupt_after_spawn) + try: + with pytest.raises(KeyboardInterrupt), campaign._sigterm_as_interrupt(): + interrupter.start() + runner.execute(["node0000"]) + finally: + returned_at = time.monotonic() + stopped.set() + interrupter.join(timeout=6) + campaign.terminate_active_commands() + force_cleanup() + + assert not interrupter.is_alive() + assert len(children) == 1 and len(signalled_at) == 1 + assert returned_at - signalled_at[0] < 4 + expected_signal = signal.SIGKILL if child_mode == "ignore_term" else signal.SIGTERM + assert children[0].returncode == -expected_signal + with campaign._ACTIVE_COMMANDS_LOCK: + assert not campaign._ACTIVE_COMMANDS + + +@pytest.mark.parametrize("backend", ["slurm", "kubernetes"]) +@pytest.mark.parametrize("interrupt", [signal.SIGINT, signal.SIGTERM]) +@pytest.mark.parametrize("window", ["prelaunch", "inflight", "registered"]) +@pytest.mark.skipif(os.name != "posix", reason="POSIX transport process groups") +def test_execute_closes_concurrent_launch_admission(runner, monkeypatch, backend, interrupt, window): + from collector.fpm_forward import runner as campaign + + arrived = threading.Barrier(3) + released = threading.Event() + finished = threading.Event() + children = [] + scopes = [] + original_popen = subprocess.Popen + original_init = campaign.CommandScope.__init__ + original_cancel = campaign.CommandScope.cancel + monkeypatch.setattr(campaign, "_COMMAND_TERMINATION_GRACE_SECONDS", 0.1) + + def scope_init(self): + original_init(self) + scopes.append(self) + + def cancel(self): + try: + return original_cancel(self) + finally: + released.set() + + def popen(*args, **kwargs): + child = original_popen(*args, **kwargs) + children.append(child) + if window == "inflight": + arrived.wait(timeout=5) + assert released.wait(5) + return child + + def command(unit, *args, **kwargs): + if window == "prelaunch": + arrived.wait(timeout=5) + assert released.wait(5) + result = campaign._run_command([sys.executable, "-c", "import time; time.sleep(30)"], check=False) + return (unit, result) if backend == "kubernetes" else result + + def interrupt_workers(): + try: + if window == "registered": + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + if scopes: + with scopes[0].lock: + if len(scopes[0].processes) == 2: + break + time.sleep(0.005) + else: + return + else: + arrived.wait(timeout=5) + os.kill(os.getpid(), interrupt) + if not finished.wait(5): + released.set() + for child in children: + child.kill() + except threading.BrokenBarrierError: + released.set() + + monkeypatch.setattr(campaign.CommandScope, "__init__", scope_init) + monkeypatch.setattr(campaign.CommandScope, "cancel", cancel) + monkeypatch.setattr(subprocess, "Popen", popen) + if backend == "kubernetes": + cell_dir = runner.cell_dir + runner = object.__new__(campaign.KubernetesCellRunner) + runner.cell_dir = cell_dir + monkeypatch.setattr(runner, "_run_pod", command) + else: + monkeypatch.setattr(runner, "_exec", command) + interrupter = threading.Thread(target=interrupt_workers) + try: + with pytest.raises(KeyboardInterrupt), campaign._sigterm_as_interrupt(): + interrupter.start() + runner.execute(["node0000", "node0001"]) + finally: + finished.set() + released.set() + interrupter.join(timeout=6) + for child in children: + if child.poll() is None: + child.kill() + child.wait(timeout=2) + assert not interrupter.is_alive() + assert len(children) == (0 if window == "prelaunch" else 2) + assert all(child.returncode != 0 for child in children) + assert scopes[0].cancelled and scopes[0].inflight == 0 and not scopes[0].processes + with campaign._ACTIVE_COMMANDS_LOCK: + assert not campaign._ACTIVE_COMMANDS + # Cancellation cannot poison salvage, scancel, or a subsequent invocation. + monkeypatch.setattr(subprocess, "Popen", original_popen) + assert campaign._run_command([sys.executable, "-c", "pass"]).returncode == 0 + assert campaign.CommandScope().run(campaign._run_command, [sys.executable, "-c", "pass"]).returncode == 0 + + +def test_cleanup_permission_error_is_reported_without_losing_interrupt(monkeypatch): + from collector.fpm_forward import runner as campaign + + stopped = [] + + class Child: + def wait(self, timeout): + stopped.append("wait") + + child = Child() + monkeypatch.setattr(campaign, "_signal_command", lambda process, force: stopped.append(force)) + + def denied(process): + raise PermissionError("injected group probe denied") + + monkeypatch.setattr(campaign, "_command_group_running", denied) + scope = campaign.CommandScope() + scope.processes.add(child) + error = KeyboardInterrupt() + campaign._cancel_preserving_interrupt(scope, error) + assert stopped == [False, True, "wait"] + assert scope.cancelled + assert any("PermissionError" in note for note in error.__notes__) + + +@pytest.mark.skipif( + not (hasattr(os, "waitid") and hasattr(os, "WNOWAIT")), + reason="Requires waitid without reaping to stage the real exit race", +) +def test_darwin_probe_reaps_child_that_exits_between_poll_and_killpg(monkeypatch): + from collector.fpm_forward import runner as campaign + + child = subprocess.Popen( + [sys.executable, "-c", "import sys; sys.stdin.buffer.read(1)"], + stdin=subprocess.PIPE, + start_new_session=True, + ) + poll, killpg = child.poll, os.killpg + probes = [] + polls = [] + + def exit_after_live_poll(): + status = poll() + polls.append(status) + if len(polls) == 1: + assert status is None # The actual process is alive at the first poll. + child.stdin.write(b"x") + child.stdin.flush() + os.waitid(os.P_PID, child.pid, os.WEXITED | os.WNOWAIT) + # It is now an unreaped zombie, before the group probe runs. + return status + + def darwin_probe(pid, sig): + assert pid == child.pid and sig == 0 + probes.append(child.returncode) + if child.returncode is None: + raise PermissionError("simulated Darwin zombie-only group") + return killpg(pid, sig) # Real ESRCH after the direct child was reaped. + + monkeypatch.setattr(campaign, "sys", SimpleNamespace(platform="darwin")) + monkeypatch.setattr(child, "poll", exit_after_live_poll) + monkeypatch.setattr(os, "killpg", darwin_probe) + try: + assert campaign._command_group_running(child) is False + assert polls == [None, 0] + assert probes == [None, 0] + finally: + child.stdin.close() + child.wait(timeout=3) + + +@pytest.mark.parametrize("force", [False, True]) +@pytest.mark.skipif(not (hasattr(os, "waitid") and hasattr(os, "WNOWAIT")), reason="Requires a real unreaped child") +def test_darwin_signal_accepts_only_reaped_disappeared_group(monkeypatch, force): + from collector.fpm_forward import runner as campaign + + child = subprocess.Popen([sys.executable, "-c", "pass"], start_new_session=True) + os.waitid(os.P_PID, child.pid, os.WEXITED | os.WNOWAIT) + killpg = os.killpg + signals = [] + + def darwin_signal(pid, sig): + signals.append(sig) + if child.returncode is None: + raise PermissionError("simulated Darwin zombie-only group") + return killpg(pid, sig) + + monkeypatch.setattr(campaign, "sys", SimpleNamespace(platform="darwin")) + monkeypatch.setattr(os, "killpg", darwin_signal) + try: + campaign._signal_command(child, force=force) + assert signals == [signal.SIGKILL if force else signal.SIGTERM, 0] + assert child.returncode == 0 + finally: + child.wait(timeout=3) + + +@pytest.mark.parametrize("platform", ["darwin", "linux"]) +@pytest.mark.parametrize("reprobe", ["denied", "surviving_group"]) +@pytest.mark.parametrize("operation", ["probe", "term", "kill"]) +def test_group_permission_denial_is_not_hidden_by_direct_child_exit(monkeypatch, platform, reprobe, operation): + from collector.fpm_forward import runner as campaign + + calls = [] + child = SimpleNamespace(pid=123456789, poll=lambda: 0) + + def denied(pid, sig): + calls.append(sig) + if len(calls) == 1 or reprobe == "denied": + raise PermissionError("permission genuinely denied") + # Direct child is gone, but another process still occupies the group. + return None + + monkeypatch.setattr(campaign, "sys", SimpleNamespace(platform=platform)) + monkeypatch.setattr(os, "killpg", denied) + with pytest.raises(PermissionError, match="genuinely denied"): + if operation == "probe": + campaign._command_group_running(child) + else: + campaign._signal_command(child, force=operation == "kill") + assert len(calls) == (2 if platform == "darwin" else 1) + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX transport process groups") +def test_darwin_cleanup_waits_for_communicate_owner_to_reap(monkeypatch): + """Portable reproduction of Darwin EPERM while communicate owns waitpid.""" + from collector.fpm_forward import runner as campaign + + owner_waiting = threading.Event() + release_owner = threading.Event() + children, results, failures = [], [], [] + real_popen, real_killpg = subprocess.Popen, os.killpg + denied = [] + + def popen(*args, **kwargs): + child = real_popen(*args, **kwargs) + children.append(child) + real_wait = child._wait + + def wait(timeout): + # communicate() calls this wait. Keep its owner's waitpid lock + # across the first cleanup signal/probe without using waitid. + with child._waitpid_lock: + owner_waiting.set() + assert release_owner.wait(5) + return real_wait(timeout) + + child._wait = wait + return child + + def killpg(pid, sig): + if not release_owner.is_set(): + if sig == signal.SIGTERM: + real_killpg(pid, sig) + elif sig == 0: + release_owner.set() + denied.append(sig) + raise PermissionError("Darwin zombie awaiting communicate owner") + return real_killpg(pid, sig) + + def run(): + try: + results.append( + campaign._run_command( + [sys.executable, "-c", "import os, time; os.close(1); os.close(2); time.sleep(30)"], check=False + ) + ) + except BaseException as error: + failures.append(error) + + monkeypatch.setattr(subprocess, "Popen", popen) + monkeypatch.setattr(os, "killpg", killpg) + monkeypatch.setattr(campaign, "sys", SimpleNamespace(platform="darwin")) + monkeypatch.setattr(campaign, "_COMMAND_TERMINATION_GRACE_SECONDS", 1) + worker = threading.Thread(target=run) + worker.start() + try: + assert owner_waiting.wait(5) + campaign._stop_commands(children) + finally: + release_owner.set() + for child in children: + if child.poll() is None: + real_killpg(child.pid, signal.SIGKILL) + worker.join(5) + assert not worker.is_alive() + assert not failures + assert results[0].returncode == -signal.SIGTERM + assert denied[:2] == [signal.SIGTERM, 0] 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 be3bc872b..d02b2b992 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_eighteen(): +def test_engine_spec_schema_version_is_twenty(): """Verification width fields extend the positional operation payload.""" from aisimulate.sdk import engine - assert engine.ENGINE_SPEC_SCHEMA_VERSION == 19 + assert engine.ENGINE_SPEC_SCHEMA_VERSION == 20 def test_lanes_outside_the_known_vocabulary_stay_reachable(): diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_dataset.py b/python/aisimulate/tests/unit/sdk/test_fpm_dataset.py new file mode 100644 index 000000000..18c0ec378 --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_fpm_dataset.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Storage trust boundary tests; native schema/cell checks remain in Rust.""" + +import hashlib +import io +import json + +import pytest + +from aisimulate_core.sdk import fpm_dataset + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def dataset(tmp_path, monkeypatch): + bodies = { + "gb300.yaml": b"data_dir: data/gb300\n", + "data/gb300/sglang/dev-example/fpm_forward_perf.parquet": b"test parquet bytes", + "data/gb300/sglang/dev-example/fpm_forward_perf.metadata.json": b'{"schema_version":7}', + } + manifest = { + "format_version": 1, + "repo_id": "nvidia/aisimulate-fpm-dataset", + "revision": "a" * 40, + "profiles": { + "gb300-full": { + "admission": "serving", + "files": [ + {"path": "archive/" + path, "target": path, "sha256": hashlib.sha256(body).hexdigest()} + for path, body in bodies.items() + ], + } + }, + } + path = tmp_path / "pin.json" + path.write_text(json.dumps(manifest)) + calls = [] + + def fetch(request, **kwargs): + url = request.full_url + calls.append(url) + prefix = f"https://huggingface.co/datasets/{manifest['repo_id']}/resolve/{'a' * 40}/archive/" + assert url.startswith(prefix) + return io.BytesIO(bodies[url.removeprefix(prefix)]) + + monkeypatch.setattr(fpm_dataset, "urlopen", fetch) + return path, manifest, bodies, calls, tmp_path / "cache" + + +def test_pinned_fetch_and_offline_cache(dataset): + path, _, bodies, calls, cache = dataset + staged = fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + assert len(calls) == 3 + for target, body in bodies.items(): + assert (staged / target).read_bytes() == body + assert fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache, local_files_only=True) == staged + assert len(calls) == 3 + + +def test_offline_missing_never_downloads(dataset): + path, _, _, calls, cache = dataset + with pytest.raises(FileNotFoundError): + fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache, local_files_only=True) + assert not calls + + +@pytest.mark.parametrize("mutation", ["revision", "path", "target", "duplicate", "diagnostic", "metadata"]) +def test_bad_manifests_fail_before_network(dataset, mutation): + path, manifest, _, calls, cache = dataset + entry = manifest["profiles"]["gb300-full"] + if mutation == "revision": + manifest["revision"] = "main" + elif mutation in ("path", "target"): + entry["files"][0][mutation] = "../escape" + elif mutation == "duplicate": + entry["files"].append(entry["files"][0]) + elif mutation == "diagnostic": + entry["admission"] = "quarantined" + else: + entry["files"].pop() + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError): + fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + assert not calls + + +def test_corrupt_download_is_not_published(dataset): + path, _, bodies, _, cache = dataset + bodies["gb300.yaml"] = b"corruption" + with pytest.raises(ValueError, match="SHA256"): + fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + assert list(cache.iterdir()) == [] + + +def test_corrupt_cache_fails_without_redownload(dataset): + path, _, _, calls, cache = dataset + staged = fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + (staged / "gb300.yaml").write_bytes(b"corruption") + with pytest.raises(ValueError, match="SHA256"): + fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + assert len(calls) == 3 + + +def test_symlink_cache_file_rejected(dataset, tmp_path): + path, _, bodies, _, cache = dataset + staged = fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + outside = tmp_path / "outside.yaml" + outside.write_bytes(bodies["gb300.yaml"]) + (staged / "gb300.yaml").unlink() + (staged / "gb300.yaml").symlink_to(outside) + with pytest.raises(ValueError, match="unsafe"): + fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + + +@pytest.mark.parametrize("admission", ["development", "quarantined"]) +def test_historical_reproduction_requires_explicit_opt_in(dataset, admission): + path, manifest, _, calls, cache = dataset + manifest["profiles"]["gb300-full"]["admission"] = admission + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="not admitted"): + fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + assert not calls + with pytest.warns(UserWarning, match="historical reproduction only"): + staged = fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache, allow_unqualified=True) + assert staged.is_dir() + with pytest.raises(ValueError, match="not admitted"): + fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache) + + +def test_hf_login_and_environment_credentials_do_not_follow_redirects(tmp_path, monkeypatch): + monkeypatch.delenv("HF_TOKEN", raising=False) + monkeypatch.delenv("HUGGING_FACE_HUB_TOKEN", raising=False) + monkeypatch.delenv("HF_HUB_DISABLE_IMPLICIT_TOKEN", raising=False) + token_path = tmp_path / "token" + token_path.write_text("test-login-secret\n") + monkeypatch.setenv("HF_TOKEN_PATH", str(token_path)) + url = "https://huggingface.co/datasets/org/repo/resolve/" + "a" * 40 + "/file" + request = fpm_dataset._download_request(url) + assert request.get_header("Authorization") == "Bearer test-login-secret" + redirected = fpm_dataset._HubRedirectHandler().redirect_request( + request, None, 302, "Found", {}, "https://cdn.example/file" + ) + assert redirected.get_header("Authorization") is None + monkeypatch.setenv("HF_TOKEN", "test-environment-secret") + assert fpm_dataset._download_request(url).get_header("Authorization") == "Bearer test-environment-secret" + monkeypatch.setenv("HF_HUB_DISABLE_IMPLICIT_TOKEN", "1") + assert fpm_dataset._download_request(url).get_header("Authorization") is None + + +@pytest.mark.parametrize( + ("target", "authenticated"), + [ + ("https://huggingface.co/api/resolve-cache/file", True), + ("https://huggingface.co:443/api/resolve-cache/file", True), + ("https://cdn.example/file", False), + ("https://huggingface.co:444/file", False), + ("http://huggingface.co/file", False), + ], +) +def test_hub_cache_redirects_preserve_auth_only_within_same_https_origin(monkeypatch, target, authenticated): + monkeypatch.setenv("HF_TOKEN", "test-private-dataset-token") + monkeypatch.delenv("HF_HUB_DISABLE_IMPLICIT_TOKEN", raising=False) + request = fpm_dataset._download_request("https://huggingface.co/datasets/org/repo/resolve/commit/file") + redirected = fpm_dataset._HubRedirectHandler().redirect_request(request, None, 302, "Found", {}, target) + assert redirected.get_header("Authorization") == ("Bearer test-private-dataset-token" if authenticated else None) + # A second cross-origin hop must drop same-origin retained auth too. + second = fpm_dataset._HubRedirectHandler().redirect_request( + redirected, None, 302, "Found", {}, "https://cdn.example/second" + ) + assert second.get_header("Authorization") is None + + +def test_concurrent_downloaders_publish_one_complete_profile(dataset, monkeypatch): + from concurrent.futures import ThreadPoolExecutor + from threading import Barrier, local + + path, _, bodies, _, cache = dataset + download = fpm_dataset.urlopen + barrier = Barrier(4) + state = local() + + def simultaneous_download(request, **kwargs): + if not getattr(state, "started", False): + state.started = True + barrier.wait(timeout=10) + return download(request, **kwargs) + + monkeypatch.setattr(fpm_dataset, "urlopen", simultaneous_download) + with ThreadPoolExecutor(max_workers=4) as workers: + results = list( + workers.map(lambda _: fpm_dataset.materialize_fpm_profile(path, "gb300-full", cache_dir=cache), range(4)) + ) + assert len(set(results)) == 1 + assert list(cache.iterdir()) == [results[0]] + for target, body in bodies.items(): + assert (results[0] / target).read_bytes() == body diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py b/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py new file mode 100644 index 000000000..5ab01d5ae --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy +import hashlib +import json +import shutil +from pathlib import Path + +import pytest + +from aisimulate_core.sdk.fpm_identity import EXECUTION_COLUMNS, LEGACY_EXECUTION_IDENTITY, execution_identity +from aisimulate_core.sdk.utils import _attach_inferred_quant_fields, _get_model_config_path + +pytestmark = pytest.mark.unit + + +def config(): + return json.loads((Path(_get_model_config_path()) / "deepseek-ai--DeepSeek-V4.1-Flash_config.json").read_text()) + + +def test_v41_config_and_execution_cannot_borrow_a_table(): + raw = config() + off = execution_identity(raw, engram_cpu_offload=False, input_modality="text") + on = execution_identity(raw, decoder_replay=True, backend="sglang", engram_cpu_offload=False, input_modality="text") + assert off[1:] == ("full", "hbm_tp_sharded", "text") + assert on[0] == off[0] and on[1] == "decoder_bounded" + altered = copy.deepcopy(raw) + altered["text_config"]["kv_source_layer_ids"] = [2, 8, 14] + assert execution_identity(altered, engram_cpu_offload=False, input_modality="text")[0] != off[0] + assert ( + execution_identity( + _attach_inferred_quant_fields(copy.deepcopy(raw)), engram_cpu_offload=False, input_modality="text" + ) + == off + ) + assert raw == config() + with pytest.raises(NotImplementedError, match="not verified for vllm"): + execution_identity(raw, decoder_replay=True, engram_cpu_offload=False, input_modality="text") + + +def test_existing_model_identity_stays_legacy(): + assert execution_identity({"architectures": ["LlamaForCausalLM"]}) == LEGACY_EXECUTION_IDENTITY + + +def test_native_v41_requires_measured_execution_and_text_evidence(): + from types import SimpleNamespace + + from collector.fpm_forward.native_artifact import _validate_execution_provenance + + identity = execution_identity(config(), engram_cpu_offload=False, input_modality="text") + cell = SimpleNamespace(execution_identity=identity, input_text_sha256="a" * 64) + fields = EXECUTION_COLUMNS + payload = { + "execution_mode": "eager", + "execution_identity": dict(zip(fields, identity, strict=True)), + "input_provenance": { + "source": "tokenizer_text", + "text_sha256": "a" * 64, + "token_ids_sha256": "b" * 64, + "tokenizer_revision": "pinned", + "token_count": 100, + "unique_token_count": 20, + }, + } + assert _validate_execution_provenance(cell, payload, Path("artifact")) == payload["input_provenance"] + for mode in (None, "cuda_graph"): + corrupt = copy.deepcopy(payload) + corrupt["execution_mode"] = mode + with pytest.raises(ValueError, match="verified eager execution"): + _validate_execution_provenance(cell, corrupt, Path("artifact")) + corrupt = copy.deepcopy(payload) + corrupt["execution_identity"]["execution_profile"] = "decoder_bounded" + with pytest.raises(ValueError, match="execution identity"): + _validate_execution_provenance(cell, corrupt, Path("artifact")) + corrupt = copy.deepcopy(payload) + corrupt["input_provenance"]["unique_token_count"] = 1 + with pytest.raises(ValueError, match="multiple tokenizer-generated"): + _validate_execution_provenance(cell, corrupt, Path("artifact")) + assert ( + _validate_execution_provenance( + SimpleNamespace(execution_identity=LEGACY_EXECUTION_IDENTITY, input_text_sha256=""), {}, Path("legacy") + ) + is None + ) + + +@pytest.mark.parametrize("replay,backend", [(False, "vllm"), (True, "sglang")]) +def test_v41_fpm_wrap_retains_resident_inventory_and_serialized_stages(replay, backend): + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.deepseek_v41 import MODEL_PATH + from aisimulate_core.sdk.engine import build_ops_json + from aisimulate_core.sdk.models import get_model + + kwargs = dict(tp_size=4, pp_size=1, attention_dp_size=1, moe_tp_size=4, moe_ep_size=1, decoder_replay=replay) + granular = get_model(MODEL_PATH, ModelConfig(**kwargs), backend) + wrapped = get_model(MODEL_PATH, ModelConfig(**kwargs, forward_model="fpm"), backend) + for ops in (wrapped.context_ops, wrapped.generation_ops): + assert len(ops) == 1 + assert ops[0].get_weights() == granular.get_resident_weights_bytes() + assert ops[0]._match_identity[-4:] == execution_identity( + config(), decoder_replay=replay, backend=backend, engram_cpu_offload=False, input_modality="text" + ) + native = json.loads(build_ops_json(ops))[0]["FpmForward"] + assert len(native["match_identity"]) == 19 + stages = [item["Dsv41Stage"] for item in native["sol_ops"] if "Dsv41Stage" in item] + assert len(stages) == 40 + assert all(stage["decoder_replay"] == replay for stage in stages) + assert '"Dsv41Linear"' in json.dumps(native["sol_ops"]) + + +def test_table_selector_preserves_checkpoint_graph_residency_and_cache_identity(): + from types import SimpleNamespace + + from aisimulate_core.sdk.common import FMHAQuantMode + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.deepseek_v41 import MODEL_PATH + from aisimulate_core.sdk.engine import build_ops_json + from aisimulate_core.sdk.models import get_model + from aisimulate_core.sdk.rust_engine_step import _engine_config_json + + kwargs = dict(tp_size=4, pp_size=1, attention_dp_size=1, moe_tp_size=4, moe_ep_size=1) + granular = get_model(MODEL_PATH, ModelConfig(**kwargs), "vllm") + native = get_model(MODEL_PATH, ModelConfig(**kwargs, forward_model="fpm"), "vllm") + selected = get_model( + MODEL_PATH, + ModelConfig(**kwargs, forward_model="fpm", fpm_fmha_quant_mode=FMHAQuantMode.fp8), + "vllm", + ) + overridden = get_model( + MODEL_PATH, ModelConfig(**kwargs, forward_model="fpm", fmha_quant_mode=FMHAQuantMode.fp8), "vllm" + ) + assert native.config.fmha_quant_mode == selected.config.fmha_quant_mode == FMHAQuantMode.bfloat16 + assert overridden.config.fmha_quant_mode == FMHAQuantMode.fp8 + for phase in ("context_ops", "generation_ops"): + direct = json.loads(build_ops_json(getattr(granular, phase))) + baseline = json.loads(build_ops_json(getattr(native, phase)))[0]["FpmForward"] + query = json.loads(build_ops_json(getattr(selected, phase)))[0]["FpmForward"] + arithmetic = json.loads(build_ops_json(getattr(overridden, phase)))[0]["FpmForward"] + assert query["sol_ops"] == baseline["sol_ops"] == direct + assert query["sol_ops"] != arithmetic["sol_ops"] + assert query["match_identity"][2] == arithmetic["match_identity"][2] == "fp8" + assert baseline["match_identity"][2] == "bfloat16" + assert query["original_fmha_quant_mode"] == "bfloat16" + assert baseline["original_fmha_quant_mode"] is None + for model in (native, selected): + assert model.get_resident_weights_bytes() == granular.get_resident_weights_bytes() + assert model.get_additional_activation_bytes(512) == granular.get_additional_activation_bytes(512) + assert model.get_kvcache_bytes_per_sequence(2048) == granular.get_kvcache_bytes_per_sequence(2048) + # Same arithmetic with a different table selector must not reuse a compiled + # handle that selected another FPM identity. + database = SimpleNamespace(system="gb200", backend="vllm", version="test") + assert _engine_config_json(native, database) != _engine_config_json(selected, database) + + +@pytest.mark.parametrize("forward_model", [None, "op_level"]) +def test_table_selector_requires_fpm_before_model_resolution(forward_model): + from aisimulate_core.sdk.common import FMHAQuantMode + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.models import get_model + + with pytest.raises(ValueError, match="requires (forward_model='fpm'|estimation_mode='fpm_interpolation')"): + get_model( + "deliberately-unresolved-model", + ModelConfig(forward_model=forward_model, fpm_fmha_quant_mode=FMHAQuantMode.fp8), + "vllm", + ) + + +def _synthetic_v41_fpm(tmp_path, *, replay=False): + """Exercise the native loader with four invented timings, without campaign files.""" + import pyarrow as pa + import pyarrow.parquet as pq + + from aisimulate_core.sdk.deepseek_v41 import MODEL_PATH + + backend = "sglang" if replay else "vllm" + version = "test-v41-fpm" + systems = tmp_path / "systems" + data = systems / "data" / "gb200" / backend / version + data.mkdir(parents=True) + packaged_systems = Path(_get_model_config_path()).parent / "systems" + shutil.copyfile(packaged_systems / "gb200.yaml", systems / "gb200.yaml") + fields = EXECUTION_COLUMNS + identity = execution_identity( + config(), decoder_replay=replay, backend=backend, engram_cpu_offload=False, input_modality="text" + ) + rows = [] + for phase, new, kv, latency in ( + ("prefill", 32, 128, 10.0), + ("prefill", 64, 128, 12.0), + ("decode", 0, 128, 6.0), + ("decode", 0, 256, 8.0), + ): + rows.append( + dict( + cell_id=f"synthetic-{phase}", + model_path=MODEL_PATH, + system="gb200", + backend=backend, + backend_version=version, + weight_quantization="fp8_block", + gemm_quant_mode="fp8_block", + moe_quant_mode="w4a8_mxfp4_mxfp8_trtllm" if backend == "sglang" else "w4a8_mxfp4_mxfp8", + fmha_quant_mode="fp8", + comm_quant_mode="half", + kv_cache_dtype="fp8", + tp=4, + pp=1, + dp=1, + moe_tp=4, + moe_ep=1, + cp=1, + moe_backend="auto", + attention_backend="auto", + enable_wideep=False, + enable_eplb=False, + workload_kind=phase, + batch_size=1, + total_prefill_tokens=new, + total_kv_read_tokens=kv, + partition_policy="balanced_v1", + kv_seed_regime="real_kv", + latency_ms=latency, + ) + | dict(zip(fields, identity, strict=True)) + ) + parquet = data / "fpm_forward_perf.parquet" + pq.write_table(pa.Table.from_pylist(rows), parquet) + metadata = dict( + schema_name="aic_fpm_forward_perf", + schema_version=7, + coordinate_system="iteration_totals_balanced_v1", + measurement_policy="dynamo_native_single_sample_v1", + row_count=len(rows), + parquet_sha256=hashlib.sha256(parquet.read_bytes()).hexdigest(), + system="gb200", + backend=backend, + backend_version=version, + ) + parquet.with_suffix(".metadata.json").write_text(json.dumps(metadata)) + native = dict( + schema_version=1, + model_name=MODEL_PATH, + system_name="gb200", + backend=backend, + backend_version=version, + systems_path=str(systems), + enable_shared_layer=False, + strict_provenance=True, + tp_size=4, + pp_size=1, + moe_tp_size=4, + moe_ep_size=1, + attention_dp_size=1, + database_mode="SILICON", + decoder_replay=replay, + forward_model="fpm", + fpm_fmha_dtype="fp8", + ) + return native, rows + + +def test_v41_fpm_selector_roundtrip_and_interpolation(tmp_path): + from aisimulate_core.sdk.rust_engine_step import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + selected, rows = _synthetic_v41_fpm(tmp_path) + legacy = dict(selected) + legacy["activation_dtype"] = legacy.pop("fpm_fmha_dtype") + old = RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig.from_legacy_engine_config(legacy, "aggregated") + ) + canonical = ForwardPassPerfModelConfig.from_legacy_engine_config(selected, "aggregated") + assert canonical.fpm_fmha_quant_mode == "fp8" + assert canonical.fmha_quant_mode is None + assert ForwardPassPerfModelConfig(**canonical.to_dict()).to_dict() == canonical.to_dict() + new = RustForwardPassPerfModel.best_available(canonical) + # Two exact endpoints and one interpolation point per phase. The fixture + # values are arbitrary test data, not calibration or validation measurements. + for phase in ("prefill", "decode"): + exact_rows = [row for row in rows if row["workload_kind"] == phase] + points = [(row["total_prefill_tokens"], row["total_kv_read_tokens"], row["latency_ms"]) for row in exact_rows] + points.append((48, 128, None) if phase == "prefill" else (0, 192, None)) + for q, k, expected in points: + scheduled = dict( + num_prefill_requests=1 if phase == "prefill" else 0, + num_decode_requests=1 if phase == "decode" else 0, + sum_prefill_tokens=q, + sum_prefill_kv_tokens=k if phase == "prefill" else 0, + sum_decode_kv_tokens=k if phase == "decode" else 0, + var_prefill_length=0.0, + var_decode_kv_tokens=0.0, + ) + fpm = dict(version=1, wall_time=1.0, scheduled_requests=scheduled) + prediction = new.estimate_forward_pass_time_ms(fpm) + assert prediction == old.estimate_forward_pass_time_ms(fpm) + if expected is not None: + assert prediction == expected + else: + assert exact_rows[0]["latency_ms"] < prediction < exact_rows[1]["latency_ms"] + with pytest.raises(ValueError, match="requires (forward_model='fpm'|estimation_mode='fpm_interpolation')"): + RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig.from_legacy_engine_config(selected | {"forward_model": "op_level"}, "aggregated") + ) + + +def test_v41_fpm_rejects_ambiguous_aggregates_but_keeps_identifiable_inputs(tmp_path): + """The public whole-forward path must reject before a balanced lookup.""" + from aisimulate_core.sdk.rust_engine_step import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + native, rows = _synthetic_v41_fpm(tmp_path, replay=True) + predictor = RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig.from_legacy_engine_config(native, "aggregated") + ) + + # Equal complete prompts can have unequal current extends: (1, 1023) + # and (1023, 1) in (new, prefix) coordinates. Prompt variance is zero. + scheduled = dict( + num_prefill_requests=2, + sum_prefill_tokens=1024, + sum_prefill_kv_tokens=1024, + var_prefill_length=0.0, + num_decode_requests=0, + sum_decode_kv_tokens=0, + var_decode_kv_tokens=0.0, + ) + for decode_batch in (0, 1): + metrics = dict( + version=1, + wall_time=1.0, + scheduled_requests=scheduled + | dict(num_decode_requests=decode_batch, sum_decode_kv_tokens=decode_batch * 129), + ) + with pytest.raises(ValueError, match="multiple prefill requests"): + predictor.estimate_forward_pass_time_ms(metrics) + + for phase in ("prefill", "decode"): + row = next(row for row in rows if row["workload_kind"] == phase and row["batch_size"] == 1) + batch, new, kv = row["batch_size"], row["total_prefill_tokens"], row["total_kv_read_tokens"] + exact = dict( + num_prefill_requests=batch if phase == "prefill" else 0, + sum_prefill_tokens=new, + sum_prefill_kv_tokens=kv if phase == "prefill" else 0, + var_prefill_length=0.0, + num_decode_requests=batch if phase == "decode" else 0, + sum_decode_kv_tokens=kv if phase == "decode" else 0, + var_decode_kv_tokens=0.0, + ) + metrics = dict(version=1, wall_time=1.0, scheduled_requests=exact) + assert predictor.estimate_forward_pass_time_ms(metrics) == row["latency_ms"] + if phase == "decode": + # Fully cached prefill metadata schedules no new prefill compute. + exact.update(num_prefill_requests=2, sum_prefill_kv_tokens=1024) + assert predictor.estimate_forward_pass_time_ms(metrics) == row["latency_ms"] + + +@pytest.mark.parametrize("offload", [None, True, 0, "false"]) +def test_v41_identity_rejects_missing_or_unverified_residency(offload): + with pytest.raises(ValueError, match="explicit engram_cpu_offload=False"): + execution_identity(config(), engram_cpu_offload=offload, input_modality="text") + + +@pytest.mark.parametrize("modality", [None, "image", "multimodal", ""]) +def test_v41_identity_rejects_missing_or_nontext_input(modality): + with pytest.raises(ValueError, match="explicit input_modality='text'"): + execution_identity(config(), engram_cpu_offload=False, input_modality=modality) + + +@pytest.mark.parametrize("quant_config", [None, {}, {"expert_dtype": "fp4"}]) +def test_v41_null_or_native_quantization_config(monkeypatch, quant_config): + from aisimulate_core.sdk import common + from aisimulate_core.sdk.models import helpers + + raw = {"quantization_config": quant_config} + architecture = "DeepseekV41ForCausalLM" + modes = helpers._infer_quant_modes_from_raw_config(raw, architecture) + monkeypatch.setattr(helpers, "_get_model_info", lambda _: {"architecture": architecture, "raw_config": raw}) + native = quant_config == {"expert_dtype": "fp4"} + assert helpers._is_dsv4_fp4_expert_model("test/model") is native + if native: + assert modes["moe_quant_mode"] == common.MoEQuantMode.w4a8_mxfp4_mxfp8 + else: + assert "moe_quant_mode" not in modes + + +def test_v41_preserves_explicit_nvfp4_experts(monkeypatch): + from aisimulate_core.sdk import common + from aisimulate_core.sdk.models import helpers + + raw = { + "quant_algo": "mixed_precision", + "quantization_config": { + "expert_dtype": "fp4", + "quantized_layers": {"model.layers.0.mlp.experts.0.gate_proj": {"quant_algo": "NVFP4"}}, + }, + } + architecture = "DeepseekV41ForCausalLM" + modes = helpers._infer_quant_modes_from_raw_config(raw, architecture) + assert modes["moe_quant_mode"] == common.MoEQuantMode.nvfp4 + monkeypatch.setattr(helpers, "_get_model_info", lambda _: {"architecture": architecture, "raw_config": raw}) + assert not helpers._is_dsv4_fp4_expert_model("test/model") + + +@pytest.mark.parametrize("window", [0, -1]) +def test_v41_rejects_nonpositive_sliding_window(window): + from aisimulate_core.sdk.deepseek_v41 import DeepSeekV41Config + + raw = config()["text_config"] + raw["sliding_window"] = window + with pytest.raises(ValueError, match="sliding_window must be positive"): + DeepSeekV41Config.from_text_config(raw) + + +def test_model_config_builder_preserves_existing_positional_controls(): + from aisimulate_core.sdk.common import FMHAQuantMode + from aisimulate_core.sdk.config_builders import build_model_config + + result = build_model_config( + 4, + 1, + 1, + 4, + 1, + None, + None, + None, + None, + None, + None, + True, + None, + None, + "deepep_moe", + True, + 64, + fpm_fmha_quant_mode="fp8", + ) + assert result.speculation is None + assert result.moe_backend == "deepep_moe" + assert result.enable_eplb is True + assert result.wideep_num_slots == 64 + assert result.fpm_fmha_quant_mode == FMHAQuantMode.fp8 diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 4242d4bef..0462fe15e 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -31,13 +31,13 @@ from aisimulate.compiler import prediction_to_replay_spec from aisimulate.config import CorePredictionConfig from aisimulate.main import main -from aisimulate.runner import EngineReplayRunnerFactory +from aisimulate.runner import AICAFDCompanionPerformanceModel, EngineReplayRunnerFactory from aisimulate.sdk import common, models from aisimulate.sdk import config as sdk_config from aisimulate.sdk.backends.factory import get_backend from aisimulate.sdk.operations import FPMForwardOp from aisimulate.sdk.perf_database import PerfDatabase -from aisimulate.sweeper import AFDLayerTimes +from aisimulate.sweeper import AFDLayerTimes, AFDTopology, BackendDeploymentSpec, ReplaySpec from aisimulate.sweeper.replay import ReplayOutputRequirements from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel from aisimulate_core.sdk.engine import EngineHandle, compile_engine @@ -408,6 +408,11 @@ def fpm_session(tmp_path): _row("decode", 8, 0, 1026, 6.5, model_path=model.model_path, identity=identity), _row("decode", 8, 0, 16400, 9.5, model_path=model.model_path, identity=identity), ] + # Separate selector lane makes swapped warning values observable. + rows += [ + dict(row, fmha_quant_mode="fp8", cell_id=row["cell_id"] + "-fp8", latency_ms=row["latency_ms"] + 1.0) + for row in rows + ] # data_dir comes from the system yaml ("data/h200_sxm"). data_dir = os.path.join(systems_root, "data", SYSTEM, BACKEND, VERSION) _write_pair(data_dir, rows) @@ -417,6 +422,54 @@ def fpm_session(tmp_path): return model, database, backend, isl, osl +def test_afd_companion_packaged_fpm_selector_reaches_native_loader(fpm_session): + _, database, _, isl, osl = fpm_session + topology = AFDTopology( + n_a_nodes=1, + n_f_nodes=1, + gpus_per_node=1, + tp_a=1, + a_batch_size=1, + num_microbatches=1, + phase="decode", + combined_with_pd=True, + ) + spec = ReplaySpec( + backend_deployment=BackendDeploymentSpec( + deployment_mode=topology.adapter_topology, + backend=BACKEND, + backend_version=VERSION, + parallel_config={ + "prefill_tp": 1, + "prefill_pp": 1, + "prefill_attention_dp": 1, + "prefill_moe_tp": 1, + "prefill_moe_ep": 1, + }, + prefill_engine_args={ + "max_num_batched_tokens": isl, + "max_num_seqs": 1, + "aic_model_path": "Qwen/Qwen3-0.6B", + "aic_system": SYSTEM, + "aic_forward_model": "fpm", + "aic_fpm_fmha_dtype": "fp8", + "systems_path": database.systems_root, + }, + num_prefill_workers=1, + ), + workload={"isl": isl, "osl": osl}, + goal={"target": "throughput", "sla": None}, + concurrency=1, + ) + + timing = AICAFDCompanionPerformanceModel().measure(spec) + + # The data tree has a separate fp8 selector row with latency 22.0 + 1.0. + assert timing.latency_ms == pytest.approx(23.0) + assert timing.provenance["source"] == "aisimulate_core.sdk.rust_engine_step.RustForwardPassPerfModel" + assert timing.provenance["fpm_fmha_dtype"] == "fp8" + + class TestFPMStaticAndMixed: @pytest.mark.parametrize("ctx_tokens", [256, 512]) @pytest.mark.parametrize("gen_requests", [0, 2]) @@ -641,6 +694,33 @@ def test_genonly_step_works_with_single_op(self, fpm_session): assert total == pytest.approx(7.0) +def test_explicit_selector_emits_matched_cell_warning_once(fpm_session, capfd): + from aisimulate_core.sdk.rust_engine_step import _cached_engine_handle + + baseline, database, _backend, _isl, _osl = fpm_session + selected = models.get_model( + baseline.model_path, + _model_config(forward_model="fpm", fpm_fmha_quant_mode=common.FMHAQuantMode.fp8), + BACKEND, + ) + original = baseline.config.fmha_quant_mode.name + selector = common.FMHAQuantMode.fp8.name + assert original != selector + capfd.readouterr() + handle = _cached_engine_handle(selected, database) + first = handle.evaluate_context_ops([0], batch_size=1, s=512) + warning = capfd.readouterr().err + assert first[0][1] == 23.0 + assert "WARNING: FPM table FMHA selector" in warning + assert f'original_model_mode="{original}"' in warning + assert f'selector="{selector}"' in warning + assert "matched_cell_ids=" in warning + assert "fpm-test-prefill-fp8" in warning and "fpm-test-decode-fp8" in warning + assert "does not independently verify runtime attention precision" in warning + assert handle.evaluate_context_ops([0], batch_size=1, s=512) == first + assert "FPM table FMHA selector" not in capfd.readouterr().err + + @pytest.mark.parametrize( "forward_model,path", [(None, "/missing/fpm.parquet"), ("op_level", "/missing/fpm.parquet"), ("fpm", "")] ) @@ -886,6 +966,7 @@ def measure(self, request): gemm_quant_mode=common.GEMMQuantMode.fp8, moe_quant_mode=common.MoEQuantMode.fp8, fmha_quant_mode=common.FMHAQuantMode.fp8, + fpm_fmha_quant_mode=common.FMHAQuantMode.bfloat16, kvcache_quant_mode=common.KVCacheQuantMode.fp8, comm_quant_mode=common.CommQuantMode.fp8, ), @@ -910,6 +991,7 @@ def measure(self, request): args.update(aic_system=custom_system, systems_path=str(systems_root), aic_fpm_parquet_path=expected_path) prefix = "aic_" if identity_fields == "prefixed" else "" args.update({f"{prefix}{field}_dtype": "fp8" for field in ("gemm", "moe", "fmha", "kv_cache", "comm")}) + args[f"{prefix}fpm_fmha_dtype"] = "bfloat16" args.update( { "aic_pp_size": 1, @@ -951,3 +1033,51 @@ def test_fpm_detail_distinguishes_memory_budget_from_runtime_capacity(external_f assert sections["source"]["status"] == "unavailable" assert "whole-model FPM" in sections["source"]["unavailable_reason"] assert sections["time"]["serving_metrics"]["mean_ttft_ms"] > 0 + + +def test_fpm_selector_allows_fallback_to_untrained_regression(tmp_path): + systems = tmp_path / "systems" + systems.mkdir() + shutil.copy(Path(_CORE_SYSTEMS) / f"{SYSTEM}.yaml", systems / f"{SYSTEM}.yaml") + model = RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig( + model="Qwen/Qwen3-0.6B", + system=SYSTEM, + backend=BACKEND, + backend_version=VERSION, + worker_type="aggregated", + systems_paths=(str(systems),), + estimation_mode="fpm_interpolation", + fallback_policy="allow", + fpm_fmha_quant_mode="fp8", + ) + ) + diagnostics = model.diagnostics() + assert diagnostics["provenance"]["selected_estimation_mode"] == "fpm_regression" + model.close() + + +def test_canonical_config_preserves_positional_quantization_fields(): + config = ForwardPassPerfModelConfig( + "model", + "system", + "sglang", + "aggregated", + "version", + 4, + 1, + 1, + 4, + 1, + "fp8", + "fp8", + "bfloat16", + "fp8", + "bfloat16", + 0, + fpm_fmha_quant_mode="fp8", + ) + assert config.kvcache_quant_mode == "fp8" + assert config.comm_quant_mode == "bfloat16" + assert config.nextn == 0 + assert config.fpm_fmha_quant_mode == "fp8" diff --git a/tests/test_afd_runner.py b/tests/test_afd_runner.py index 7d39ee2ec..4eac4f198 100644 --- a/tests/test_afd_runner.py +++ b/tests/test_afd_runner.py @@ -481,6 +481,44 @@ def estimator(model, hardware, **kwargs): assert calls[0]["forward_model"] == "fpm" +@pytest.mark.parametrize(("phase", "companion_role"), [("decode", "prefill"), ("prefill", "decode")]) +def test_aic_companion_routes_packaged_fpm_selector_through_canonical_model(phase, companion_role, monkeypatch): + spec = _spec(_topology(phase=phase, combined_with_pd=True), companion_role=companion_role) + engine_args = getattr(spec.backend_deployment, f"{companion_role}_engine_args") + engine_args.pop("timing_model") + engine_args.update( + aic_model_path="test-model", + aic_system="test-system", + aic_forward_model="fpm", + aic_fpm_fmha_dtype="fp8", + ) + configs = [] + + class CanonicalModel: + def static_phase_latency(self, **_kwargs): + return 2.0 + + def close(self): + pass + + def best_available(config): + configs.append(config) + return CanonicalModel() + + monkeypatch.setattr("aisimulate.runner.RustForwardPassPerfModel.best_available", best_available) + + def unexpected_legacy_estimation(*args, **kwargs): + pytest.fail("packaged FPM selector must not be dropped by the legacy estimator path") + + timing = AICAFDCompanionPerformanceModel(unexpected_legacy_estimation).measure(spec) + + assert timing.latency_ms == pytest.approx(2.0 if companion_role == "prefill" else 1.0) + assert len(configs) == 1 + assert configs[0].estimation_mode == "fpm_interpolation" + assert configs[0].fpm_fmha_quant_mode == "fp8" + assert timing.provenance["fpm_fmha_dtype"] == "fp8" + + def test_afd_runner_rejects_unresolved_measurement_before_execution(): spec = _spec(_topology()) spec.backend_deployment.performance_model_metadata["afd"]["measurement_required"] = True diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index eb086c4cb..7147b745e 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -1503,3 +1503,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") + + +@pytest.mark.parametrize( + "option", + [ + {"decoder_replay": True}, + {"enable_shared_layer": True}, + {"enable_shared_layer": False}, + {"strict_provenance": True}, + {"strict_provenance": False}, + ], +) +@pytest.mark.parametrize( + "timing", + [{"type": "fixed", "prefill_ms": 1, "decode_ms": 1}, {"type": "polynomial"}], +) +def test_execution_options_reject_nondefault_timing(option, timing): + engine = _engine() | {"model": "deepseek-ai/DeepSeek-V4.1-Flash", "backend": "sglang"} | option + engine["workers"]["aggregated"] = {"timing": timing} + with pytest.raises(ValidationError, match="estimator policies require.*default timing"): + CorePredictionConfig.model_validate({"engine": engine}) + + +@pytest.mark.parametrize( + "option", + [ + {"decoder_replay": True}, + {"enable_shared_layer": False}, + {"strict_provenance": True}, + ], +) +def test_execution_options_accept_default_timing(option): + engine = _engine() | {"model": "deepseek-ai/DeepSeek-V4.1-Flash", "backend": "sglang"} | option + config = CorePredictionConfig.model_validate({"engine": engine}) + for name, value in option.items(): + assert getattr(config.engine, name) == value + + +@pytest.mark.parametrize("custom_role", ["prefill", "decode"]) +def test_execution_options_reject_mixed_worker_timing(custom_role): + engine = _engine() | { + "mode": "disaggregated", + "strict_provenance": True, + "workers": {"prefill": {}, "decode": {}}, + } + engine["workers"][custom_role] = {"timing": {"type": "polynomial"}} + with pytest.raises(ValidationError, match="default timing in every role"): + CorePredictionConfig.model_validate({"engine": engine})