From 2740c7c265c6f49a207ec73c0a3d661c745e33d4 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Fri, 18 Sep 2026 14:32:15 -0700 Subject: [PATCH 1/8] feat: define DeepSeek V4.1 text architecture and execution contract Reapply reviewed PR159 net changes from 40fdb6039eb5ff5f1a642aef5e43da46722e89ea onto current main. --- THIRD_PARTY_NOTICES.md | 35 + crates/core/src/engine/cache/radix_cache.rs | 20 + crates/core/src/engine/common/perf_model.rs | 66 +- .../src/engine/kv_manager/sglang_backend.rs | 93 +- .../core/src/engine/scheduler/sglang/core.rs | 77 +- .../core/src/engine/scheduler/sglang/tests.rs | 276 ++++++ crates/core/src/engine/timing.rs | 30 + crates/core/src/perfmodel/common/enums.rs | 2 + .../core/src/perfmodel/common/system_spec.rs | 5 + crates/core/src/perfmodel/config.rs | 9 +- crates/core/src/perfmodel/engine/readiness.rs | 55 +- crates/core/src/perfmodel/engine/runtime.rs | 410 +++++++- crates/core/src/perfmodel/engine/spec.rs | 109 ++- crates/core/src/perfmodel/fpm/config.rs | 4 + crates/core/src/perfmodel/fpm/tests.rs | 1 + crates/core/src/perfmodel/memory.rs | 1 + .../core/src/perfmodel/operators/attention.rs | 6 +- crates/core/src/perfmodel/operators/dsv41.rs | 918 ++++++++++++++++++ crates/core/src/perfmodel/operators/mod.rs | 2 + crates/core/src/perfmodel/operators/op.rs | 41 +- .../core/src/perfmodel/perf_database/gemm.rs | 1 + .../core/src/perfmodel/perf_database/mod.rs | 1 + .../perf_database/moe_expert_compute.rs | 1 + crates/core/src/perfmodel/py.rs | 13 + crates/core/src/perfmodel/py_ops.rs | 11 +- crates/core/src/python.rs | 107 +- crates/core/src/replay/agg_tests.rs | 21 +- .../core/tests/perfmodel/memory_round_trip.rs | 1 + crates/tests/public-api/src/lib.rs | 2 +- docs/deepseek-v41-storage.md | 57 ++ docs/deepseek-v41.md | 119 +++ python/aisimulate/THIRD_PARTY_NOTICES.md | 35 + python/aisimulate/src/aisimulate/compiler.py | 13 + .../src/aisimulate/config/engine.py | 6 +- python/aisimulate/src/aisimulate/runner.py | 7 + .../src/aisimulate/sdk/deepseek_v41.py | 8 + .../src/aisimulate/sdk/inference_session.py | 4 + .../src/aisimulate/sdk/models/deepseek_v41.py | 8 + .../aisimulate/src/aisimulate/sdk/task_v2.py | 7 + ...deepseek-ai--DeepSeek-V4.1-Flash_README.md | 11 + ...epseek-ai--DeepSeek-V4.1-Flash_config.json | 170 ++++ .../src/aisimulate_core/sdk/afd_partition.py | 10 + .../sdk/backends/base_backend.py | 15 +- .../sdk/backends/sglang_backend.py | 1 + .../sdk/backends/trtllm_backend.py | 1 + .../src/aisimulate_core/sdk/common.py | 5 + .../src/aisimulate_core/sdk/config.py | 2 + .../src/aisimulate_core/sdk/deepseek_v41.py | 229 +++++ .../src/aisimulate_core/sdk/engine.py | 6 +- .../src/aisimulate_core/sdk/memory.py | 2 +- .../aisimulate_core/sdk/models/__init__.py | 10 +- .../src/aisimulate_core/sdk/models/base.py | 18 + .../sdk/models/deepseek_v41.py | 297 ++++++ .../src/aisimulate_core/sdk/models/helpers.py | 17 +- .../aisimulate_core/sdk/rust_engine_step.py | 2 + .../src/aisimulate_core/sdk/utils.py | 2 + .../src/aisimulate_core/systems/b200_sxm.yaml | 3 +- .../src/aisimulate_core/systems/b300_sxm.yaml | 1 + .../src/aisimulate_core/systems/gb200.yaml | 1 + .../src/aisimulate_core/systems/gb300.yaml | 1 + .../cross_package/test_core_public_api.py | 1 + .../cross_package/test_import_contract.py | 2 + .../tests/e2e/cli/test_cli_build_default.py | 42 + .../unit/cli/test_afd_phase_completion.py | 11 + .../unit/sdk/backends/test_base_backend.py | 4 + .../test_deepseek_v4_workspace_memory.py | 2 + .../unit/sdk/backends/test_step3p7_memory.py | 2 + .../unit/sdk/database/test_attention_lanes.py | 2 +- .../unit/sdk/models/test_deepseek_v41.py | 488 ++++++++++ .../sdk/models/test_deepseek_v41_residency.py | 115 +++ .../unit/sdk/models/test_model_config.py | 1 + .../tests/unit/sdk/models/test_qwen35.py | 55 ++ .../speculation/test_consumer_equivalence.py | 2 + .../tests/unit/sdk/test_v41_native_bridge.py | 126 +++ tests/test_runner.py | 46 + tests/test_unified_traffic_runtime.py | 5 +- 76 files changed, 4220 insertions(+), 70 deletions(-) create mode 100644 crates/core/src/perfmodel/operators/dsv41.rs create mode 100644 docs/deepseek-v41-storage.md create mode 100644 docs/deepseek-v41.md create mode 100644 python/aisimulate/src/aisimulate/sdk/deepseek_v41.py create mode 100644 python/aisimulate/src/aisimulate/sdk/models/deepseek_v41.py create mode 100644 python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_README.md create mode 100644 python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json create mode 100644 python/aisimulate/src/aisimulate_core/sdk/deepseek_v41.py create mode 100644 python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py create mode 100644 python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py create mode 100644 python/aisimulate/tests/unit/sdk/models/test_deepseek_v41_residency.py create mode 100644 python/aisimulate/tests/unit/sdk/test_v41_native_bridge.py diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index dc273cc39..0fdce628f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -437,6 +437,33 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` +## SGLang DeepSeek-V4.1 serving contracts + +The V4.1 execution and performance-model logic in +`src/aiconfigurator_core/sdk/deepseek_v41.py`, `sdk/models/deepseek_v41.py` +(with the same `src/aiconfigurator_core/` prefix), and repository-root +`crates/core/src/perfmodel/operators/dsv41.rs` is informed by and modified from +SGLang's serving architecture at immutable commit +`1aa0e962b206102b7c439a4a0c4981cfec6e87bc`: + +- `python/sglang/srt/models/deepseek_v4.py` and `deepseek_v2.py` +- `python/sglang/srt/layers/engram.py` +- `python/sglang/srt/layers/attention/dsv4/compressor.py` +- `python/sglang/srt/layers/attention/dsv4/dsv41_sparse.py` +- `python/sglang/srt/layers/attention/deepseek_v4_backend.py` +- `python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py` +- `python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py` + +Source: https://github.com/sgl-project/sglang/tree/1aa0e962b206102b7c439a4a0c4981cfec6e87bc +Copyright 2023-2024 SGLang Team and SGLang contributors. Licensed under Apache-2.0; its terms are +reproduced in the repository `LICENSE`. These are analytical adaptations, +not a copy of the model execution implementation. The modified analytical +scoring/storage adaptations and their independently written regression cases +also appear in `python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py`, +Rust operator/spec unit tests, `docs/deepseek-v41.md`, and +`docs/deepseek-v41-storage.md`. They distinguish candidate masking from scoring +and physical FlashMLA cache payload from logical FP4 values. + ## DeepSeek model configuration files The following model configuration files are copied from, or formatting-only @@ -448,8 +475,16 @@ adaptations of, the named DeepSeek model repositories: | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3_config.json` | `deepseek-ai/DeepSeek-V3@e815299b0bcbac849fa540c768ef21845365c9eb` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3.2_config.json` | `deepseek-ai/DeepSeek-V3.2@c69397ecfd1fd142e90e3fbad51f4c7e40b9f3d3` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Flash_config.json` | `deepseek-ai/DeepSeek-V4-Flash@60d8d70770c6776ff598c94bb586a859a38244f1` | +| `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json` | `deepseek-ai/DeepSeek-V4.1-Flash@fb2764a5cf321eaa5070ca8f9e892818f477c16d` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Pro_config.json` | `deepseek-ai/DeepSeek-V4-Pro@b5968e9190ef611bbf34a7229255be88a0e937c1` | +The V4.1 descriptor and performance formulas in `src/aiconfigurator_core/sdk/deepseek_v41.py`, +`src/aiconfigurator_core/sdk/models/deepseek_v41.py`, and repository-root +`crates/core/src/perfmodel/operators/dsv41.rs` are AISimulate performance-model +adaptations of the architecture described by `inference/model.py` and +`DeepSeek_V41_Tech_Report.pdf` at the same V4.1 revision (modified; no model execution code). +Source: https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/tree/fb2764a5cf321eaa5070ca8f9e892818f477c16d + Upstream repositories: https://huggingface.co/deepseek-ai diff --git a/crates/core/src/engine/cache/radix_cache.rs b/crates/core/src/engine/cache/radix_cache.rs index b0e899ea7..07202b36d 100644 --- a/crates/core/src/engine/cache/radix_cache.rs +++ b/crates/core/src/engine/cache/radix_cache.rs @@ -43,6 +43,7 @@ impl KvPageId { /// Manages free / allocated pages for the simulated SGLang KV cache. /// /// SGLang's paged allocator owns and frees whole pages in production. +#[derive(Clone)] pub struct PagePool { next_fresh: usize, free: Vec, @@ -157,6 +158,7 @@ impl PagePool { } /// A single node in the radix tree. +#[derive(Clone)] pub struct TreeNode { /// Children keyed by the first complete page on the child edge. pub children: FxHashMap, @@ -195,6 +197,7 @@ pub struct RadixCache { } /// Optional forecast metadata; physical ownership and leaf eligibility stay in the radix tree. +#[derive(Clone)] struct BeladyLeaves { oracle: BeladyOracle, cursor: usize, @@ -205,6 +208,23 @@ struct BeladyLeaves { } impl RadixCache { + /// Snapshot allocator and radix metadata for a fallible admission transaction. + /// This is not a second owner of live request leases; only one cache state may commit. + pub(crate) fn admission_checkpoint(&self) -> Self { + Self { + nodes: self.nodes.clone(), + root: self.root, + page_pool: self.page_pool.clone(), + page_size: self.page_size, + #[cfg(test)] + test_now: self.test_now, + evictable_leaves: self.evictable_leaves.clone(), + belady: self.belady.clone(), + evictable_size: self.evictable_size, + protected_size: self.protected_size, + } + } + pub fn new(total_tokens: usize, page_size: usize) -> Self { assert!(page_size >= 1, "page_size must be >= 1"); let mut nodes = SlotMap::with_key(); diff --git a/crates/core/src/engine/common/perf_model.rs b/crates/core/src/engine/common/perf_model.rs index f01cc183c..1ab334d36 100644 --- a/crates/core/src/engine/common/perf_model.rs +++ b/crates/core/src/engine/common/perf_model.rs @@ -42,6 +42,22 @@ impl std::fmt::Debug for PerfModel { } impl PerfModel { + pub(crate) fn prefill_batch_validation_can_fail(&self) -> bool { + match self { + Self::Polynomial => false, + Self::External { timing } => timing.prefill_batch_validation_can_fail(), + } + } + + pub fn validate_prefill_batch(&self, requests: &[(usize, usize)]) -> Result<()> { + if let PerfModel::External { timing } = self { + timing + .validate_prefill_batch(requests) + .context("external prefill batch cannot be represented")?; + } + Ok(()) + } + /// Predict prefill time in milliseconds. /// /// Callers always pass all parameters; each variant uses what it needs: @@ -101,17 +117,17 @@ impl PerfModel { ) .context("external decode prediction failed")?, }; - // Validate before the floor: f64::max would hide NaN or -infinity. + // Reject non-finite provider output before it reaches the scheduler. ensure!( time.is_finite(), "decode timing provider returned non-finite duration {time}ms" ); - // Token-emitting decode steps should not collapse onto the same timestamp. - let result = time.max(1.0); tracing::trace!( - "Decode time prediction: batch_size={batch_size}, active_kv_tokens={active_kv_tokens}, context_length={context_length}, time={result:.2}ms" + "Decode time prediction: batch_size={batch_size}, active_kv_tokens={active_kv_tokens}, context_length={context_length}, time={time:.2}ms" ); - Ok(result) + // The polynomial applies its own historical floor. External measurements + // can legitimately be below one millisecond and must retain that value. + Ok(time) } } @@ -206,6 +222,46 @@ mod tests { assert_eq!(model.predict_decode_time(9, 0, 128, 0).unwrap(), 9.0); } + #[test] + fn external_decode_preserves_finite_values_and_rejects_non_finite() { + struct FixedDecode(f64); + impl TimingModel for FixedDecode { + fn predict_prefill_ms(&self, _: usize, _: usize, _: usize) -> anyhow::Result { + Ok(0.0) + } + fn predict_decode_ms( + &self, + _: usize, + _: usize, + _: usize, + _: usize, + ) -> anyhow::Result { + Ok(self.0) + } + } + // Finite values retain the provider's submillisecond resolution; + // modeled_duration_ms still rejects negative durations downstream. + for value in [0.6, 0.0, -1.0] { + let model = PerfModel::External { + timing: Arc::new(FixedDecode(value)), + }; + assert_eq!(model.predict_decode_time(1, 128, 128, 1024).unwrap(), value); + } + // Main's finite-time boundary must reject NaN and both infinities before + // any duration normalization or scheduling can hide invalid output. + for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + let model = PerfModel::External { + timing: Arc::new(FixedDecode(value)), + }; + let error = model.predict_decode_time(1, 128, 128, 1024).unwrap_err(); + assert!( + error + .to_string() + .contains("decode timing provider returned non-finite duration") + ); + } + } + #[test] fn external_prefill_errors_propagate_with_context() { let error = PerfModel::External { diff --git a/crates/core/src/engine/kv_manager/sglang_backend.rs b/crates/core/src/engine/kv_manager/sglang_backend.rs index 00680342e..0445e766a 100644 --- a/crates/core/src/engine/kv_manager/sglang_backend.rs +++ b/crates/core/src/engine/kv_manager/sglang_backend.rs @@ -32,7 +32,38 @@ pub(crate) struct RadixRequestLease { last_node: Option, } +/// Metadata used only with the matching KV-manager admission checkpoint. +pub(crate) struct RadixLeaseCheckpoint { + pages: Vec, + materialized_tokens: usize, + cached_tokens: usize, + admission_reused_tokens: usize, + page_hashes_len: usize, + last_node: Option, +} + impl RadixRequestLease { + pub(crate) fn admission_checkpoint(&self) -> RadixLeaseCheckpoint { + RadixLeaseCheckpoint { + pages: self.pages.clone(), + materialized_tokens: self.materialized_tokens, + cached_tokens: self.cached_tokens, + admission_reused_tokens: self.admission_reused_tokens, + page_hashes_len: self.page_hashes.len(), + last_node: self.last_node, + } + } + + /// Restore only after the matching manager checkpoint has been rolled back. + pub(crate) fn restore_admission(&mut self, checkpoint: RadixLeaseCheckpoint) { + self.pages = checkpoint.pages; + self.materialized_tokens = checkpoint.materialized_tokens; + self.cached_tokens = checkpoint.cached_tokens; + self.admission_reused_tokens = checkpoint.admission_reused_tokens; + self.page_hashes.truncate(checkpoint.page_hashes_len); + self.last_node = checkpoint.last_node; + } + #[cfg(test)] pub(crate) fn pages(&self) -> &[KvPageId] { &self.pages @@ -135,6 +166,15 @@ pub struct SglangKvManager { /// block hash so router events reflect logical block visibility, not /// transient slot ownership. block_hash_refcounts: FxHashMap, + pending_admission_events: Option>, +} + +#[must_use = "admission must be committed or rolled back"] +pub(crate) struct SglangAdmissionCheckpoint { + cache: RadixCache, + next_event_id: u64, + page_to_block_hash: Vec>, + block_hash_refcounts: FxHashMap, } pub struct DecodeTokenReservation { @@ -227,6 +267,7 @@ impl SglangKvManager { next_event_id: 0, page_to_block_hash, block_hash_refcounts: FxHashMap::default(), + pending_admission_events: None, } } @@ -243,6 +284,50 @@ impl SglangKvManager { self.cache.set_belady_oracle(oracle); } + /// Hold speculative events until the timing provider accepts the admitted batch. + pub(crate) fn begin_admission(&mut self) -> SglangAdmissionCheckpoint { + assert!( + self.pending_admission_events.is_none(), + "nested KV admission" + ); + self.pending_admission_events = Some(Vec::new()); + SglangAdmissionCheckpoint { + cache: self.cache.admission_checkpoint(), + next_event_id: self.next_event_id, + page_to_block_hash: self.page_to_block_hash.clone(), + block_hash_refcounts: self.block_hash_refcounts.clone(), + } + } + + pub(crate) fn commit_admission(&mut self, _checkpoint: SglangAdmissionCheckpoint) { + for event in self + .pending_admission_events + .take() + .expect("active KV admission") + { + self.publish_event(event); + } + } + + pub(crate) fn rollback_admission(&mut self, checkpoint: SglangAdmissionCheckpoint) { + self.pending_admission_events + .take() + .expect("active KV admission"); + self.cache = checkpoint.cache; + self.next_event_id = checkpoint.next_event_id; + self.page_to_block_hash = checkpoint.page_to_block_hash; + self.block_hash_refcounts = checkpoint.block_hash_refcounts; + self.log_trace("admission_rollback", 0); + } + + fn publish_event(&mut self, event: KvEvent) { + if let Some(events) = self.pending_admission_events.as_mut() { + events.push(event); + } else if let Err(error) = self.kv_event_publishers.publish(event, None) { + tracing::warn!("Failed to publish SGLang KV event: {error}"); + } + } + #[cfg(test)] pub fn cache_mut(&mut self) -> &mut RadixCache { &mut self.cache @@ -999,9 +1084,7 @@ impl SglangKvManager { }; self.next_event_id += 1; - if let Err(e) = self.kv_event_publishers.publish(event, None) { - tracing::warn!("Failed to publish SGLang KV event: {e}"); - } + self.publish_event(event); hashed_blocks } @@ -1042,9 +1125,7 @@ impl SglangKvManager { }; self.next_event_id += 1; - if let Err(e) = self.kv_event_publishers.publish(event, None) { - tracing::warn!("Failed to publish SGLang KV remove event: {e}"); - } + self.publish_event(event); } } diff --git a/crates/core/src/engine/scheduler/sglang/core.rs b/crates/core/src/engine/scheduler/sglang/core.rs index 4a2055521..616c1c015 100644 --- a/crates/core/src/engine/scheduler/sglang/core.rs +++ b/crates/core/src/engine/scheduler/sglang/core.rs @@ -788,6 +788,27 @@ impl SglangCore { if let Some(oracle) = &self.belady { oracle.retire_requests(rejected.iter().map(|signal| signal.uuid)); } + // 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. + let admission_checkpoint = (!self.waiting.is_empty() + && self.config.perf_model.prefill_batch_validation_can_fail()) + .then(|| { + let waiting = self + .waiting + .iter() + .map(|request| { + ( + request.uuid, + request.materialized_tokens, + request.allocated_tokens, + request.kv_lease.admission_checkpoint(), + ) + }) + .collect::>(); + (self.kv_manager.begin_admission(), waiting) + }); + let running_before_admission = self.running.len(); let mut admissions = self.promote_prebuilt_ready(); let materialized_waiting = !self.prebuilt_ready.is_empty(); apply_schedule_policy(&mut self.waiting, &self.kv_manager, &self.config); @@ -808,6 +829,56 @@ impl SglangCore { AdmissionStage::FreshKv => Default::default(), }; + let batch_size = admit.can_run.len(); + let mean_isl = admit.total_isl.checked_div(batch_size).unwrap_or(0); + let mean_prefix = admit.total_prefix.checked_div(batch_size).unwrap_or(0); + let prefill_time = (|| { + self.config.perf_model.validate_prefill_batch( + &admit + .prefill_fpm + .iter() + .map(|item| (item.tokens_computed, item.prefix_tokens)) + .collect::>(), + )?; + simulate_prefill_duration(batch_size, mean_isl, mean_prefix, &self.config, true) + })(); + let prefill_time = match prefill_time { + Ok(duration) => { + if let Some((checkpoint, _)) = admission_checkpoint { + self.kv_manager.commit_admission(checkpoint); + } + duration + } + Err(error) => { + // A retry is still part of the caller's prepared group round; + // it must not consume the prefill interval a second time. + self.group_pass_prepared = grouped; + if let Some((checkpoint, waiting)) = admission_checkpoint { + self.kv_manager.rollback_admission(checkpoint); + let mut requests = self + .waiting + .drain(..) + .chain(admit.can_run) + .map(|request| (request.uuid, request)) + .collect::>(); + for (uuid, materialized, allocated, lease) in waiting { + let mut request = + requests.remove(&uuid).expect("admission request retained"); + request.kv_lease.restore_admission(lease); + request.materialized_tokens = materialized; + request.allocated_tokens = allocated; + request.debug_assert_invariants(self.config.block_size); + self.waiting.push_back(request); + } + debug_assert!(requests.is_empty()); + } + for request in self.running.drain(running_before_admission..).rev() { + self.prebuilt_ready.push_front(request); + } + return Err(error); + } + }; + admissions.append(&mut admit.admissions); for admission in &admissions { if let Some(collector) = collector.as_deref_mut() { @@ -818,12 +889,6 @@ impl SglangCore { // Capture per-request prefill FPM data before dispersing can_run. let prefill_fpm = admit.prefill_fpm; - let batch_size = admit.can_run.len(); - let mean_isl = admit.total_isl.checked_div(batch_size).unwrap_or(0); - let mean_prefix = admit.total_prefix.checked_div(batch_size).unwrap_or(0); - let prefill_time = - simulate_prefill_duration(batch_size, mean_isl, mean_prefix, &self.config, true)?; - // This committed prefill retires the whole request's input forecast exactly once. // Later chunks and preemption recomputation intentionally do not restore demand: // the oracle estimates global input demand, while native execution remains causal. diff --git a/crates/core/src/engine/scheduler/sglang/tests.rs b/crates/core/src/engine/scheduler/sglang/tests.rs index d0219669f..996004cf4 100644 --- a/crates/core/src/engine/scheduler/sglang/tests.rs +++ b/crates/core/src/engine/scheduler/sglang/tests.rs @@ -2598,3 +2598,279 @@ mod forward_pass_metrics { } } } + +#[test] +fn prefill_provider_validation_receives_actual_chunk_geometry() { + use std::sync::{Arc, Mutex}; + struct GeometryTiming(Arc>>>); + impl crate::engine::TimingModel for GeometryTiming { + fn validate_prefill_batch(&self, requests: &[(usize, usize)]) -> anyhow::Result<()> { + self.0.lock().unwrap().push(requests.to_vec()); + anyhow::ensure!( + requests.windows(2).all(|pair| pair[0] == pair[1]), + "heterogeneous geometry" + ); + Ok(()) + } + fn predict_prefill_ms(&self, _: usize, _: usize, _: usize) -> anyhow::Result { + Ok(0.6) + } + fn predict_decode_ms(&self, _: usize, _: usize, _: usize, _: usize) -> anyhow::Result { + Ok(0.6) + } + } + let observed = Arc::new(Mutex::new(Vec::new())); + let mut args = test_args(128, 4, 8); + args.perf_model = crate::engine::common::perf_model::PerfModel::External { + timing: Arc::new(GeometryTiming(Arc::clone(&observed))), + } + .into(); + let mut core = SglangCore::new(args); + core.receive(direct_request((0..16).collect(), 1)); + core.try_execute_pass_internal(None, 0.0).unwrap(); + core.try_execute_pass_internal(None, 1.0).unwrap(); + assert_eq!(*observed.lock().unwrap(), vec![vec![(8, 0)], vec![(8, 8)]]); + + let mut args = test_args(128, 4, 32); + args.perf_model = crate::engine::common::perf_model::PerfModel::External { + timing: Arc::new(GeometryTiming(Arc::clone(&observed))), + } + .into(); + let mut core = SglangCore::new(args); + core.receive(direct_request((0..4).collect(), 1)); + core.receive(direct_request((100..112).collect(), 1)); + let error = core.try_execute_pass_internal(None, 0.0).unwrap_err(); + assert_eq!(error.root_cause().to_string(), "heterogeneous geometry"); + assert_eq!( + observed.lock().unwrap().last().unwrap(), + &vec![(4, 0), (12, 0)] + ); +} + +mod admission_validation_rollback { + use super::*; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + struct FallibleTiming { + fail: Arc, + fail_in_prediction: bool, + } + + impl crate::engine::TimingModel for FallibleTiming { + fn validate_prefill_batch(&self, _: &[(usize, usize)]) -> anyhow::Result<()> { + anyhow::ensure!( + self.fail_in_prediction || !self.fail.load(Ordering::Relaxed), + "rejected prefill geometry" + ); + Ok(()) + } + + fn predict_prefill_ms(&self, _: usize, _: usize, _: usize) -> anyhow::Result { + anyhow::ensure!( + !self.fail_in_prediction || !self.fail.load(Ordering::Relaxed), + "prefill timing unavailable" + ); + Ok(0.6) + } + + fn predict_decode_ms(&self, _: usize, _: usize, _: usize, _: usize) -> anyhow::Result { + Ok(0.6) + } + } + + fn install_timing(core: &mut SglangCore, fail: &Arc, fail_in_prediction: bool) { + core.config.perf_model = crate::engine::common::perf_model::PerfModel::External { + timing: Arc::new(FallibleTiming { + fail: Arc::clone(fail), + fail_in_prediction, + }), + } + .into(); + } + + fn capacity(core: &SglangCore) -> (usize, usize, usize, usize) { + let cache = core.kv_manager.cache(); + ( + cache.available_tokens(), + cache.evictable_size, + cache.protected_size, + cache.num_nodes(), + ) + } + + #[test] + fn fresh_requests_survive_failed_passes_and_are_admitted_only_on_retry() { + for fail_in_prediction in [false, true] { + let mut core = SglangCore::new_with_kv_capture(test_args(32, 4, 32), 0); + let mut collector = crate::engine::trace::TraceCollector::default(); + let ids = [ + core.receive(direct_request((0..4).collect(), 1)), + core.receive(direct_request((100..112).collect(), 1)), + ]; + for id in ids { + collector.on_arrival(id, 0.0, 0, 1); + } + let fail = Arc::new(AtomicBool::new(true)); + install_timing(&mut core, &fail, fail_in_prediction); + let waiting_before = format!("{:?}", core.waiting); + let capacity_before = capacity(&core); + let ratio_before = core.new_token_ratio; + for now_ms in [1.0, 2.0] { + assert!(core.try_execute_pass(&mut collector, now_ms).is_err()); + assert_eq!(format!("{:?}", core.waiting), waiting_before); + assert!(core.running.is_empty()); + assert_eq!(capacity(&core), capacity_before); + assert_eq!(core.new_token_ratio, ratio_before); + assert!(core.drain_kv_events().is_empty()); + for id in ids { + let snapshot = collector.snapshot(id).unwrap(); + assert_eq!(snapshot.arrival_ms, Some(0.0)); + assert_eq!(snapshot.first_admit_ms, None); + assert_eq!(snapshot.first_token_ms, None); + } + } + fail.store(false, Ordering::Relaxed); + let pass = core.try_execute_pass(&mut collector, 10.0).unwrap(); + assert_eq!(pass.admissions.len(), 2); + assert_eq!(pass.completed_requests, 2); + assert!(core.is_empty()); + assert_eq!(pass.kv_events.first().unwrap().event_id, 0); + for id in ids { + assert_eq!(collector.snapshot(id).unwrap().first_admit_ms, Some(10.0)); + assert!(collector.snapshot(id).unwrap().first_token_ms.unwrap() >= 10.0); + } + } + } + + #[test] + fn continuation_failure_restores_cache_eviction_leases_queue_order_and_trace() { + let mut core = SglangCore::new_with_kv_capture(test_args(12, 4, 32), 0); + let mut collector = crate::engine::trace::TraceCollector::default(); + let seed: Vec = (1000..1016).collect(); + core.receive(direct_request(seed.clone(), 0)); + core.try_execute_pass(&mut collector, 0.0).unwrap(); + let running = core.receive(direct_request((100..104).collect(), 8)); + core.try_execute_pass(&mut collector, 1.0).unwrap(); + core.config.chunked_prefill_size = 8; + let continuation = core.receive(direct_request((0..20).collect(), 1)); + core.try_execute_pass(&mut collector, 2.0).unwrap(); + assert_eq!(core.waiting[0].materialized_tokens, 8); + assert!(core.waiting[0].kv_lease.is_active()); + let fresh = core.receive(direct_request((200..212).collect(), 1)); + // LPM changes this order during the rejected pass. Rollback must restore + // the original order, including the continuation's existing ownership. + core.waiting.swap(0, 1); + core.config.schedule_policy = SchedulePolicy::Lpm; + core.config.chunked_prefill_size = 32; + let fail = Arc::new(AtomicBool::new(true)); + install_timing(&mut core, &fail, false); + let waiting_before = format!("{:?}", core.waiting); + let running_before = format!("{:?}", core.running); + let capacity_before = capacity(&core); + let snapshots = + [running, continuation, fresh].map(|id| format!("{:?}", collector.snapshot(id))); + assert_eq!(core.kv_manager.cache().prefix_match_len(&seed), 16); + assert!(core.drain_kv_events().is_empty()); + + for now_ms in [3.0, 4.0] { + assert!(core.try_execute_pass(&mut collector, now_ms).is_err()); + assert_eq!(format!("{:?}", core.waiting), waiting_before); + assert_eq!(format!("{:?}", core.running), running_before); + assert_eq!(capacity(&core), capacity_before); + assert_eq!(core.kv_manager.cache().prefix_match_len(&seed), 16); + assert!(core.drain_kv_events().is_empty()); + assert_eq!( + [running, continuation, fresh].map(|id| format!("{:?}", collector.snapshot(id))), + snapshots + ); + } + + fail.store(false, Ordering::Relaxed); + let pass = core.try_execute_pass(&mut collector, 10.0).unwrap(); + assert_eq!(pass.admissions.len(), 1, "continuation is not readmitted"); + assert_eq!(pass.admissions[0].uuid, fresh); + assert_eq!(pass.completed_requests, 2); + assert!( + removed_event_count(&pass.kv_events) > 0, + "fixture must force eviction" + ); + assert_eq!( + collector.snapshot(continuation).unwrap().first_admit_ms, + Some(2.0) + ); + assert_eq!( + collector.snapshot(fresh).unwrap().first_admit_ms, + Some(10.0) + ); + assert_eq!(core.running.len(), 1); + assert_eq!(core.running[0].uuid, running); + } + + #[test] + fn failed_validation_restores_promoted_destination_without_duplicate_admission() { + let mut core = SglangCore::new_with_kv_capture(test_args(32, 4, 32), 0); + let mut collector = crate::engine::trace::TraceCollector::default(); + let destination = Uuid::from_u128(90_101); + let handoff_id = HandoffId::from(Uuid::from_u128(90_102)); + core.apply_command(SchedulerCommand::ReserveDestination { + handoff_id, + request: DirectRequest { + tokens: (0..8).collect(), + max_output_tokens: 2, + uuid: Some(destination), + ..Default::default() + }, + }) + .unwrap(); + core.apply_command(SchedulerCommand::ActivateDestination { handoff_id }) + .unwrap(); + core.drain_kv_events(); + let fail = Arc::new(AtomicBool::new(true)); + install_timing(&mut core, &fail, false); + let capacity_before = capacity(&core); + assert!(core.try_execute_pass(&mut collector, 1.0).is_err()); + assert!(core.running.is_empty()); + assert!(core.prebuilt_request(destination).is_some()); + assert_eq!(capacity(&core), capacity_before); + assert!(collector.snapshot(destination).is_none()); + fail.store(false, Ordering::Relaxed); + let pass = core.try_execute_pass(&mut collector, 2.0).unwrap(); + assert_eq!(pass.admissions.len(), 1); + assert_eq!(pass.admissions[0].uuid, destination); + assert_eq!( + collector.snapshot(destination).unwrap().first_admit_ms, + Some(2.0) + ); + } + + #[test] + fn failed_group_pass_retry_does_not_advance_prefill_interval_twice() { + let mut args = test_args(32, 4, 32); + args.prefill_decode_interval = 2; + let mut core = SglangCore::new(args); + core.receive(direct_request((0..4).collect(), 1)); + // A peer's prefill has armed two synchronized decode/idle rounds. + core.prepare_group_pass(); + core.finish_group_pass(true, true); + let fail = Arc::new(AtomicBool::new(true)); + install_timing(&mut core, &fail, false); + + core.prepare_group_pass(); + assert!(core.try_execute_hidden_pass(1.0).is_err()); + fail.store(false, Ordering::Relaxed); + let retried = core.try_execute_hidden_pass(1.0).unwrap(); + assert!(retried.admissions.is_empty()); + core.finish_group_pass(false, false); + + core.prepare_group_pass(); + let second_interval_round = core.try_execute_hidden_pass(2.0).unwrap(); + assert!(second_interval_round.admissions.is_empty()); + core.finish_group_pass(false, false); + + core.prepare_group_pass(); + let admitted = core.try_execute_hidden_pass(3.0).unwrap(); + assert_eq!(admitted.admissions.len(), 1); + assert_eq!(admitted.completed_requests, 1); + } +} diff --git a/crates/core/src/engine/timing.rs b/crates/core/src/engine/timing.rs index cc7ac770e..719ccce75 100644 --- a/crates/core/src/engine/timing.rs +++ b/crates/core/src/engine/timing.rs @@ -374,6 +374,20 @@ 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. + fn prefill_batch_validation_can_fail(&self) -> bool { + true + } + + /// Validate actual (new tokens, cached prefix) pairs before a scheduler + /// reduces them to means. Providers with nonlinear per-request execution + /// policies may reject batches that their aggregate API cannot represent. + fn validate_prefill_batch(&self, _requests: &[(usize, usize)]) -> Result<()> { + Ok(()) + } + /// Predict one prefill batch's latency in milliseconds. fn predict_prefill_ms( &self, @@ -404,6 +418,10 @@ pub trait TimingModel: Send + Sync { struct PolynomialTimingModel; impl TimingModel for PolynomialTimingModel { + fn prefill_batch_validation_can_fail(&self) -> bool { + false + } + fn predict_prefill_ms( &self, batch_size: usize, @@ -436,6 +454,10 @@ struct FixedTimingModel { } impl TimingModel for FixedTimingModel { + fn prefill_batch_validation_can_fail(&self) -> bool { + false + } + fn predict_prefill_ms( &self, batch_size: usize, @@ -527,12 +549,20 @@ mod tests { decode_ms: 3.0, }) .unwrap(); + assert!(!model.prefill_batch_validation_can_fail()); assert_eq!(model.predict_prefill_ms(0, 128, 0).unwrap(), 0.0); assert_eq!(model.predict_decode_ms(0, 128, 64, 1024).unwrap(), 0.0); assert_eq!(model.predict_prefill_ms(2, 128, 0).unwrap(), 7.0); assert_eq!(model.predict_decode_ms(2, 128, 64, 1024).unwrap(), 3.0); } + #[test] + fn polynomial_model_does_not_need_admission_checkpoint() { + let model = built_in_timing_model(&TimingModelConfig::Polynomial).unwrap(); + assert!(!model.prefill_batch_validation_can_fail()); + model.validate_prefill_batch(&[(4, 0), (12, 8)]).unwrap(); + } + #[test] fn external_provider_requires_runner_resolution() { let error = built_in_timing_model(&TimingModelConfig::External { diff --git a/crates/core/src/perfmodel/common/enums.rs b/crates/core/src/perfmodel/common/enums.rs index ed7cbf290..e760765f3 100644 --- a/crates/core/src/perfmodel/common/enums.rs +++ b/crates/core/src/perfmodel/common/enums.rs @@ -539,6 +539,7 @@ pub enum ModelFamily { MinimaxM3, Qwen3Vl, Qwen3VlMoe, + DeepSeekV41, } impl ModelFamily { @@ -552,6 +553,7 @@ impl ModelFamily { Self::DeepSeek => "DEEPSEEK", Self::DeepSeekV32 => "DEEPSEEKV32", Self::DeepSeekV4 => "DEEPSEEKV4", + Self::DeepSeekV41 => "DEEPSEEKV41", Self::KimiK25 => "KIMIK25", Self::NemotronNas => "NEMOTRONNAS", Self::NemotronH => "NEMOTRONH", diff --git a/crates/core/src/perfmodel/common/system_spec.rs b/crates/core/src/perfmodel/common/system_spec.rs index 8067689e3..8c0bb7dee 100644 --- a/crates/core/src/perfmodel/common/system_spec.rs +++ b/crates/core/src/perfmodel/common/system_spec.rs @@ -42,6 +42,9 @@ pub struct GpuSpec { /// HBM capacity in bytes. #[serde(default)] pub mem_capacity: Option, + /// Peak scalar FP32 throughput (not tensor-core throughput). + #[serde(default)] + pub fp32_flops: Option, /// Peak TC-FLOPS at bf16. #[serde(default)] pub bfloat16_tc_flops: Option, @@ -229,6 +232,7 @@ mod tests { let spec = SystemSpec { data_dir: PathBuf::from("data/synthetic"), gpu: GpuSpec { + fp32_flops: None, mem_bw: 1.0, mem_bw_empirical_scaling_factor: 1.0, mem_empirical_constant_latency: 0.0, @@ -267,6 +271,7 @@ mod tests { let spec = SystemSpec { data_dir: PathBuf::from("data/synthetic"), gpu: GpuSpec { + fp32_flops: None, mem_bw: 1.0, mem_bw_empirical_scaling_factor: 1.0, mem_empirical_constant_latency: 0.0, diff --git a/crates/core/src/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index 723b1ca75..7cbb47c04 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -80,7 +80,10 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // batch/query widths and FpmForwardOp gained verify_width. Upstream used // 14/15, already occupied here; these are positional bincode layout changes. // TokenScale was appended to remap draft query widths before op lookup. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 18; +// - 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; /// Static engine identity and setup information carried by an /// [`crate::perfmodel::engine::spec::EngineSpec`]. @@ -115,6 +118,10 @@ pub struct EngineConfig { #[serde(default)] pub forward_model: Option, + /// Use the backend-verified bounded DeepSeek-V4.1 decoder execution profile. + #[serde(default)] + pub decoder_replay: bool, + // KV pub kv_block_size: Option, diff --git a/crates/core/src/perfmodel/engine/readiness.rs b/crates/core/src/perfmodel/engine/readiness.rs index 4001a3e0b..1b39654df 100644 --- a/crates/core/src/perfmodel/engine/readiness.rs +++ b/crates/core/src/perfmodel/engine/readiness.rs @@ -111,6 +111,12 @@ impl Availability<'_> { DatabaseMode::Sol | DatabaseMode::SolFull ); match op { + Dsv41Stage(stage) => { + for child in &stage.children { + self.op(child)?; + } + return Ok(()); + } Overlap(group) => { for child in group.group_a.iter().chain(&group.group_b) { self.op(child)?; @@ -272,7 +278,20 @@ impl Availability<'_> { } Embedding(_) | Elementwise(_) | P2P(_) | CustomAllReduce(_) | Nccl(_) | MoeDispatch(_) => Ok(()), - Overlap(_) | Fallback(_) | TokenScale(_) | FpmForward(_) => Ok(()), + Dsv41Attention(_) | Dsv41Mhc(_) | Dsv41Engram(_) | Dsv41Linear(_) => { + match self.db.database_mode { + DatabaseMode::Silicon => Err(AicError::PerfDatabase(format!( + "DeepSeek-V4.1 {} has no measured SILICON data", + op.name() + ))), + DatabaseMode::Empirical => Err(AicError::EmpiricalNotImplemented(format!( + "DeepSeek-V4.1 {} has no empirical anchor", + op.name() + ))), + _ => Ok(()), + } + } + Overlap(_) | Fallback(_) | TokenScale(_) | FpmForward(_) | Dsv41Stage(_) => Ok(()), } } } @@ -335,6 +354,40 @@ mod tests { write_parquet(path, &[Col::Str("kernel_source", vec!["fixture"])]); } + #[test] + fn dsv41_stage_preserves_analytic_only_readiness() { + use crate::operators::dsv41::{Dsv41LinearOp, Dsv41StageOp}; + let root = systems(); + let stage = Op::Dsv41Stage(Dsv41StageOp { + name: "stage".into(), + is_context: true, + decoder_replay: true, + bounded: true, + window_size: 128, + children: vec![Op::Dsv41Linear(Dsv41LinearOp { + name: "projection".into(), + n: 16, + k: 16, + quant_mode: GemmQuantMode::Fp8Block, + })], + }); + for mode in [ + DatabaseMode::Silicon, + DatabaseMode::Empirical, + DatabaseMode::Hybrid, + DatabaseMode::Sol, + DatabaseMode::SolFull, + ] { + let db = PerfDatabase::load(root.path(), "b200_sxm", "vllm", "0.24.0") + .unwrap() + .with_mode(mode, TransferPolicy::ALL); + assert_eq!( + validate(&db, [&stage].into_iter()).is_ok(), + !matches!(mode, DatabaseMode::Silicon | DatabaseMode::Empirical) + ); + } + } + #[test] fn dsv4_requires_its_phase_module_not_an_auxiliary_table() { for kind in ["Csa", "Hca"] { diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 7d6b96d15..af80434fb 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -887,6 +887,31 @@ impl Engine { marginal_decode_ms, ]); } + if self.has_dsv41_stages() { + let isl = isl.max(1); + if ctx_tokens > 0 && prefix >= isl { + return Err(AicError::InvalidEngineConfig( + "V4.1 prefill requires isl > prefix".into(), + )); + } + // The SDK's packed context count includes prefix for complete + // requests. A remainder describes this iteration's partial extend. + let mut prefills = Vec::with_capacity(2); + if ctx_tokens / isl > 0 { + prefills.push((ctx_tokens / isl, isl - prefix, prefix)); + } + if ctx_tokens % isl > 0 { + prefills.push((1, ctx_tokens % isl, prefix)); + } + return self.dsv41_mixed_workload( + &prefills, + gen_tokens.saturating_mul(self.nextn.saturating_add(1)), + isl.saturating_add(osl / 2).saturating_add(1), + seq_imbalance_correction_scale, + gen_seq_imbalance_correction_scale, + on_op, + ); + } // Python divides by `isl` (`floor(ctx/isl)`, `ceil(ctx/isl)`) without // a guard — callers always pass isl >= 1. Clamp to avoid a Rust // div-by-zero panic on degenerate input Python would crash on. @@ -1005,6 +1030,139 @@ impl Engine { ]) } + fn has_dsv41_stages(&self) -> bool { + self.context_ops + .iter() + .any(|op| matches!(op, Op::Dsv41Stage(_))) + } + + /// Scope every prefill extend before fusing token-major work with decode. + /// Applying a decoder tail to the combined batch would incorrectly discard + /// decode tokens and other requests' tails. The tuple is (batch, new, prefix). + #[allow(clippy::too_many_arguments)] + fn dsv41_mixed_workload( + &self, + prefills: &[(u32, u32, u32)], + decode_batch: u32, + decode_kv: u32, + context_scale: f64, + generation_scale: f64, + mut on_op: impl FnMut(MixedPass, &Op, PerformanceResult), + ) -> Result<[f64; 4], AicError> { + let mut totals = [0.0; 4]; + if prefills.is_empty() { + if decode_batch > 0 { + // Generation can fuse or overlap children differently from + // prefill. Preserve that graph when no prefill is scheduled. + for outer in &self.generation_ops { + let children: &[Op] = match outer { + Op::Dsv41Stage(stage) => &stage.children, + _ => std::slice::from_ref(outer), + }; + for child in children { + let result = query_generation_op( + child, + &self.db, + decode_batch, + 1, + decode_kv, + generation_scale, + 0, + None, + )?; + let (bucket, pass) = if child.is_generation_attention() { + (3, MixedPass::DecodeAttention) + } else { + (1, MixedPass::SharedNonAttention) + }; + totals[bucket] += result.latency_ms; + on_op(pass, child, result); + } + } + } + totals[0] = totals[1] + totals[3]; + return Ok(totals); + } + let prefill_requests: u32 = prefills.iter().map(|(batch, _, _)| batch).sum(); + for outer in &self.context_ops { + let (stage, children): (_, &[Op]) = match outer { + Op::Dsv41Stage(stage) => (Some(stage), &stage.children), + _ => (None, std::slice::from_ref(outer)), + }; + let scopes: Vec<_> = prefills + .iter() + .map(|&(batch, s, prefix)| { + let (s, prefix) = stage.map_or((s as f64, prefix as f64), |stage| { + stage.scope(s as f64, prefix as f64) + }); + (batch, s as u32, prefix as u32) + }) + .collect(); + let tokens = scopes + .iter() + .try_fold(decode_batch, |total, &(batch, s, _)| { + batch.checked_mul(s).and_then(|n| total.checked_add(n)) + }) + .ok_or_else(|| { + AicError::InvalidEngineConfig("V4.1 mixed token count overflow".into()) + })?; + for child in children { + if child.is_context_attention() { + for &(batch, s, prefix) in &scopes { + if batch == 0 || s == 0 { + continue; + } + let result = query_context_op( + child, + &self.db, + batch, + s, + prefix, + context_scale, + None, + )?; + totals[2] += result.latency_ms; + on_op(MixedPass::ContextAttention, child, result); + } + } else if tokens > 0 { + let x = if child.is_logits_gemm() { + prefill_requests.saturating_add(decode_batch) + } else { + tokens + }; + let result = + query_context_op(child, &self.db, 1, tokens, 0, context_scale, Some(x))?; + totals[1] += result.latency_ms; + on_op(MixedPass::SharedNonAttention, child, result); + } + } + } + if decode_batch > 0 { + for outer in &self.generation_ops { + let children: &[Op] = match outer { + Op::Dsv41Stage(stage) => &stage.children, + _ => std::slice::from_ref(outer), + }; + for child in children.iter().filter(|op| op.is_generation_attention()) { + let result = query_generation_op( + child, + &self.db, + decode_batch, + 1, + decode_kv, + generation_scale, + 0, + None, + )?; + totals[3] += result.latency_ms; + on_op(MixedPass::DecodeAttention, child, result); + } + } + } + totals[0] = totals[1] + totals[2] + totals[3]; + Ok(totals) + } + /// One generation-only step latency. LITERAL mirror of Python /// `_get_genonly_step_latency` (`base_backend.py:1040-1100`): /// `run_static(batch=gen_tokens, isl=isl+osl//2, osl=2, mode=static_gen)` @@ -1370,7 +1528,7 @@ impl Engine { }, )?; let mut ctx_attn = ctx_attn.into_values(); - if ctx_tokens > 0 { + if ctx_tokens > 0 && !self.has_dsv41_stages() { // Mirror the scalar bucket and Python's fold-then-single-true- // division (`base_backend.py:1244-1246`): one `/ scale2` per // folded name, never a per-entry reciprocal multiply. @@ -1774,6 +1932,57 @@ impl Engine { return Ok(total); } + 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. + let mut prefills = Vec::new(); + if has_prefill { + let n = sched.num_prefill_requests; + let q = sched.sum_prefill_tokens / n; + let qr = sched.sum_prefill_tokens % n; + let p = sched.sum_prefill_kv_tokens / n; + let pr = sched.sum_prefill_kv_tokens % n; + let mut bounds = vec![0, qr, pr, n]; + bounds.sort_unstable(); + bounds.dedup(); + for pair in bounds.windows(2) { + let count = pair[1] - pair[0]; + let query = q + u32::from(pair[0] < qr); + if query > 0 { + prefills.push((count, query, p + u32::from(pair[0] < pr))); + } + } + } + return self + .dsv41_mixed_workload( + &prefills, + sched.num_decode_requests, + sched.sum_decode_kv_tokens / sched.num_decode_requests.max(1), + 1.0, + 1.0, + |_, _, _| {}, + ) + .map(|parts| parts[0]); + } + if has_prefill && has_decode { // Mix step (continuous batching): compose like Python's // `_get_mix_step_latency`. `sum_prefill_kv_tokens` is exactly the @@ -1932,6 +2141,7 @@ mod tests { backend: BackendKind::Vllm, backend_version: Some("0.24.0".to_string()), forward_model: None, + decoder_replay: false, kv_block_size: None, parallel: ParallelMapping { tp_size: 8, @@ -2218,6 +2428,204 @@ mod tests { )); } + // Linear memory probes isolate the engine's workload orchestration from + // kernel formulas. Their expected token counts are request-level contracts. + fn dsv41_probe_engine(replay: bool) -> Engine { + let leaf = |name: &str| { + Op::Elementwise(ElementwiseOp { + name: name.into(), + scale_factor: 1.0, + bytes_per_token: 8192.0, + scale_num_tokens: 1, + seq_split: 1, + }) + }; + let stage = |is_context, bounded| { + Op::Dsv41Stage(crate::operators::Dsv41StageOp { + name: if bounded { "decoder" } else { "encoder" }.into(), + is_context, + bounded, + decoder_replay: replay, + window_size: 128, + children: vec![ + leaf("norm"), + leaf(if is_context { + "context_attention" + } else { + "generation_attention" + }), + ], + }) + }; + let mut engine = build_engine(None); + engine.db = Arc::new( + PerfDatabase::load(&systems_root(), "b200_sxm", "vllm", "0.24.0") + .unwrap() + .with_mode(DatabaseMode::Sol, TransferPolicy::default()), + ); + engine.context_ops = vec![stage(true, false), stage(true, true)]; + engine.generation_ops = vec![stage(false, false), stage(false, true)]; + engine + } + + fn dsv41_probe_token_ms(engine: &Engine) -> f64 { + let Op::Dsv41Stage(stage) = &engine.context_ops[0] else { + unreachable!() + }; + query_context_op(&stage.children[0], &engine.db, 1, 1, 0, 1.0, None) + .unwrap() + .latency_ms + } + + #[test] + fn dsv41_mixed_scopes_each_request_before_adding_decode() { + let engine = dsv41_probe_engine(true); + let unit = dsv41_probe_token_ms(&engine); + // Two 256-token extends: encoder 512, decoder 2*128; both + // stages also execute all 200 decode requests, not one 128-token tail. + let parts = engine + .mixed_step_breakdown(512, 200, 256, 32, 0, 1.0, 1.0) + .unwrap(); + assert!((parts[1] / unit - (512.0 + 256.0 + 400.0)).abs() < 1e-9); + assert!((parts[2] / unit - 768.0).abs() < 1e-9); + assert!((parts[3] / unit - 400.0).abs() < 1e-9); + } + + #[test] + fn dsv41_partial_extend_and_prefix_do_not_fill_decoder_tail() { + let engine = dsv41_probe_engine(true); + let unit = dsv41_probe_token_ms(&engine); + for q in [1, 127, 128, 129] { + let parts = engine + .mixed_step_breakdown(q, 3, 4096, 32, 2048, 1.0, 1.0) + .unwrap(); + assert!((parts[2] / unit - f64::from(q + q.min(128))).abs() < 1e-9); + assert!((parts[1] / unit - f64::from(q + q.min(128) + 6)).abs() < 1e-9); + let (shared, context, decode) = engine + .mixed_step_breakdown_per_op(q, 3, 4096, 32, 2048, 1.0, 1.0) + .unwrap(); + assert!((shared.iter().map(|v| v.1).sum::() - parts[1]).abs() < 1e-12); + assert!((context.iter().map(|v| v.1).sum::() - parts[2]).abs() < 1e-12); + assert!((decode.iter().map(|v| v.1).sum::() - parts[3]).abs() < 1e-12); + } + } + + #[test] + fn dsv41_replay_never_changes_decode_work() { + for replay in [false, true] { + let mut engine = dsv41_probe_engine(replay); + for outer in &mut engine.generation_ops { + let Op::Dsv41Stage(stage) = outer else { + unreachable!() + }; + let norm = stage.children[0].clone(); + stage.children[0] = Op::Overlap(crate::operators::op::OverlapOp::new( + "decode_fused", + vec![norm.clone(), norm.clone()], + vec![norm], + )); + } + let mixed = engine + .mixed_step_latency(0, 257, 2048, 32, 0, 1.0, 1.0) + .unwrap(); + let decode = engine.decode_step_latency(257, 2048, 32, 1.0).unwrap(); + assert!((mixed - decode).abs() < 1e-12); + } + } + + #[test] + fn dsv41_telemetry_retains_prefill_remainders() { + let engine = dsv41_probe_engine(false); + let unit = dsv41_probe_token_ms(&engine); + let mut metrics = ForwardPassMetrics::default(); + metrics.scheduled_requests.num_prefill_requests = 2; + metrics.scheduled_requests.sum_prefill_tokens = 257; + metrics.scheduled_requests.sum_prefill_kv_tokens = 513; + metrics.scheduled_requests.num_decode_requests = 3; + metrics.scheduled_requests.sum_decode_kv_tokens = 1536; + let result = engine.forward_pass_time_ms(&[metrics]).unwrap(); + assert!((result / unit - 4.0 * 260.0).abs() < 1e-9); + } + + #[test] + fn dsv41_replay_rejects_equal_prompt_heterogeneous_extends() { + // The scheduler observes identical 1024-token prompts, but different + // cached prefixes leave extends of 1 and 1023 tokens. Prompt variance + // is zero although the real bounded tails total 129, not 2 * 128. + let requests = [(1024, 1023, 1), (1024, 1, 1023)]; + assert!( + requests + .iter() + .all(|&(prompt, prefix, query)| { prompt == 1024 && prefix + query == prompt }) + ); + let engine = dsv41_probe_engine(true); + let mut metrics = ForwardPassMetrics::default(); + metrics.scheduled_requests.num_prefill_requests = requests.len() as u32; + metrics.scheduled_requests.sum_prefill_tokens = requests.iter().map(|r| r.2).sum(); + metrics.scheduled_requests.sum_prefill_kv_tokens = requests.iter().map(|r| r.1).sum(); + // Matches build_fpm_snapshot: variance is over prompt, not query. + metrics.scheduled_requests.var_prefill_length = 0.0; + let error = engine.forward_pass_time_ms(&[metrics]).unwrap_err(); + assert!(matches!(error, AicError::InvalidForwardPassMetrics(_))); + assert!(error.to_string().contains("multiple prefill requests")); + } + + #[test] + fn dsv41_replay_rejects_multiple_prefills_without_geometry() { + let engine = dsv41_probe_engine(true); + let mut metrics = ForwardPassMetrics::default(); + metrics.scheduled_requests.num_prefill_requests = 2; + for tokens in [1, 2, 127, 128, 129, 256, 1024] { + for variance in [0.0, 64.0] { + metrics.scheduled_requests.sum_prefill_tokens = tokens; + metrics.scheduled_requests.var_prefill_length = variance; + for decode_batch in [0, 3] { + metrics.scheduled_requests.num_decode_requests = decode_batch; + metrics.scheduled_requests.sum_decode_kv_tokens = decode_batch * 512; + assert!(matches!( + engine.forward_pass_time_ms(std::slice::from_ref(&metrics)), + Err(AicError::InvalidForwardPassMetrics(_)) + )); + } + } + } + } + + #[test] + fn dsv41_replay_telemetry_keeps_single_prefill_and_decode_boundaries() { + let engine = dsv41_probe_engine(true); + let unit = dsv41_probe_token_ms(&engine); + let mut metrics = ForwardPassMetrics::default(); + for query in [0, 1, 127, 128, 129, 1024] { + for prefix in [0, 1024] { + metrics.scheduled_requests.num_prefill_requests = 1; + metrics.scheduled_requests.sum_prefill_tokens = query; + metrics.scheduled_requests.sum_prefill_kv_tokens = prefix; + for decode_batch in [0, 3] { + metrics.scheduled_requests.num_decode_requests = decode_batch; + metrics.scheduled_requests.sum_decode_kv_tokens = decode_batch * 512; + let result = engine + .forward_pass_time_ms(std::slice::from_ref(&metrics)) + .unwrap(); + let expected = 2 * (query + query.min(128) + 2 * decode_batch); + assert!((result / unit - f64::from(expected)).abs() < 1e-9); + } + } + } + // Cached-prefill metadata alone is not fresh prefill work. Keep the + // decode-only path (and an otherwise empty iteration) available. + metrics.scheduled_requests.num_prefill_requests = 2; + metrics.scheduled_requests.sum_prefill_tokens = 0; + for decode_batch in [0, 3] { + metrics.scheduled_requests.num_decode_requests = decode_batch; + metrics.scheduled_requests.sum_decode_kv_tokens = decode_batch * 512; + let result = engine + .forward_pass_time_ms(std::slice::from_ref(&metrics)) + .unwrap(); + assert!((result / unit - f64::from(4 * decode_batch)).abs() < 1e-9); + } + } + #[test] fn both_equals_context_plus_generation() { let engine = build_engine(None); diff --git a/crates/core/src/perfmodel/engine/spec.rs b/crates/core/src/perfmodel/engine/spec.rs index 442525ef4..3ac40705c 100644 --- a/crates/core/src/perfmodel/engine/spec.rs +++ b/crates/core/src/perfmodel/engine/spec.rs @@ -167,10 +167,11 @@ mod tests { use crate::operators::op::{FallbackOp, OverlapOp}; use crate::operators::{ ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4MegaMoeOp, - Dsv4ModuleOp, ElementwiseOp, EmbeddingOp, EncoderAttentionOp, GdnOp, GemmOp, - GenerationAttentionOp, GenerationMlaOp, KdaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, - MlaModuleOp, MoEDispatchOp, MoeAllToAllOp, MoeExpertComputeOp, MoeOp, NcclOp, P2POp, - VisionEncoderOp, WideEpContextMlaOp, WideEpGenerationMlaOp, + Dsv4ModuleOp, Dsv41AttentionOp, Dsv41EngramOp, Dsv41LinearOp, Dsv41MhcOp, Dsv41StageOp, + ElementwiseOp, EmbeddingOp, EncoderAttentionOp, GdnOp, GemmOp, GenerationAttentionOp, + GenerationMlaOp, KdaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, MoEDispatchOp, + MoeAllToAllOp, MoeExpertComputeOp, MoeOp, NcclOp, P2POp, VisionEncoderOp, + WideEpContextMlaOp, WideEpGenerationMlaOp, }; use crate::perf_database::dsv4::AttnKind; use crate::{ @@ -711,6 +712,57 @@ mod tests { numerator: 5, denominator: 6, }), + OpSpec::Dsv41Attention(Dsv41AttentionOp { + name: "v41_attention".into(), + is_context: true, + role: "full".into(), + compress_ratio: 2, + hidden_size: 5120, + num_heads: 16, + head_dim: 512, + q_lora_rank: 1280, + o_lora_rank: 1024, + o_groups: 2, + index_n_heads: 8, + index_head_dim: 128, + index_topk: 512, + window_size: 128, + candidate_limit: 0, + is_candidate_source: false, + bounded_prefill: false, + gemm_quant_mode: GemmQuantMode::Fp8Block, + fmha_quant_mode: FmhaQuantMode::Fp8, + kv_cache_layout: crate::operators::dsv41::Dsv41KvCacheLayout::SglangFp8Bf16, + }), + OpSpec::Dsv41Mhc(Dsv41MhcOp { + name: "v41_mhc".into(), + hidden_size: 5120, + hc_mult: 4, + sinkhorn_iters: 20, + }), + OpSpec::Dsv41Engram(Dsv41EngramOp { + name: "v41_engram".into(), + num_embeddings: 384006168, + head_dim: 256, + hash_columns: 24, + hidden_size: 5120, + 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, + decoder_replay: true, + bounded: true, + window_size: 128, + children: vec![OpSpec::Gemm(gemm())], + }), ]; // Exhaustiveness guard: if a variant is added to `Op`, this match @@ -752,6 +804,11 @@ mod tests { | OpSpec::Kda(_) | OpSpec::MoeAllToAll(_) | OpSpec::MoeExpertCompute(_) + | OpSpec::Dsv41Attention(_) + | OpSpec::Dsv41Mhc(_) + | OpSpec::Dsv41Engram(_) + | OpSpec::Dsv41Stage(_) + | OpSpec::Dsv41Linear(_) | OpSpec::TokenScale(_) => {} } } @@ -767,6 +824,7 @@ mod tests { backend: crate::BackendKind::Trtllm, backend_version: Some("1.0.0rc3".into()), forward_model: None, + decoder_replay: false, kv_block_size: Some(64), parallel: ParallelMapping { tp_size: 8, @@ -842,8 +900,16 @@ mod tests { // Appending is the only safe growth direction. 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(); assert_eq!( - TOKEN_SCALE_INDEX as usize + 1, + appended, + vec![36, 37, 38, 39, 40], + "V41 appended indices moved" + ); + assert_eq!( + TOKEN_SCALE_INDEX as usize + 6, all_op_variants().len(), "all_op_variants() must cover exactly the pinned variant count" ); @@ -1021,6 +1087,37 @@ mod tests { } } + /// The backend layout is positional even though legacy JSON has a default. + #[test] + fn dsv41_layout_round_trip_and_stale_v18_payload_rejection() { + let attention = all_op_variants() + .into_iter() + .find(|op| matches!(op, OpSpec::Dsv41Attention(_))) + .unwrap(); + let spec = EngineSpec::new(sample_engine_config(), vec![], vec![attention]); + let mut bytes = spec.to_bincode().unwrap(); + assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); + // SglangFp8Bf16 is enum1 (four bytes), the final field of the final op. + // Removing it restores the actual prior Dsv41Attention schema18 shape. + assert_eq!(&bytes[bytes.len() - 4..], &1u32.to_le_bytes()); + bytes.truncate(bytes.len() - 4); + bytes[..4].copy_from_slice(&18u32.to_le_bytes()); + assert!(matches!( + EngineSpec::from_bincode(&bytes), + Err(AicError::UnsupportedSchemaVersion { + kind: "EngineSpec", + got: 18, + expected: 19 + }) + )); + // A false schema19 stamp cannot silently use the JSON-only default. + bytes[..4].copy_from_slice(&19u32.to_le_bytes()); + assert!(matches!( + EngineSpec::from_bincode(&bytes), + Err(AicError::EngineSpec(_)) + )); + } + /// v11 -> v12 regression (PR-6): `DsaModuleOp` gained /// `attn_projection_quant_modes`, a positional bincode layout change. A /// pre-PR v11 producer's DSA payload must be rejected by the VERSION GATE @@ -1166,7 +1263,7 @@ mod tests { EngineSpec::from_bincode(&bytes), Err(AicError::UnsupportedSchemaVersion { got: 17, - expected: 18, + expected: ENGINE_SPEC_SCHEMA_VERSION, .. }) )); diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 37822037e..334a870be 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -116,6 +116,9 @@ pub struct ForwardPassPerfModelConfig { pub speculation: Option, #[serde(default)] pub kv_block_size: Option, + /// Preserve DeepSeek V4.1 decoder replay execution identity. + #[serde(default)] + pub decoder_replay: bool, #[serde(default)] #[serde(alias = "forward_model")] pub estimation_mode: EstimationMode, @@ -165,6 +168,7 @@ impl ForwardPassPerfModelConfig { nextn: 0, speculation: None, kv_block_size: None, + decoder_replay: false, estimation_mode: EstimationMode::Auto, database_mode: DatabaseMode::default(), transfer_policy: None, diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index fb4101506..0953bb3d6 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -107,6 +107,7 @@ fn fixture_engine_config() -> EngineConfig { backend: BackendKind::Vllm, backend_version: Some("0.24.0".to_string()), forward_model: None, + decoder_replay: false, kv_block_size: None, parallel: ParallelMapping { tp_size: 8, diff --git a/crates/core/src/perfmodel/memory.rs b/crates/core/src/perfmodel/memory.rs index 8a028b95c..50cb44766 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -507,6 +507,7 @@ mod tests { backend: BackendKind::Vllm, backend_version: Some("test-version".to_string()), forward_model: None, + decoder_replay: false, kv_block_size: None, parallel: ParallelMapping { tp_size: 1, diff --git a/crates/core/src/perfmodel/operators/attention.rs b/crates/core/src/perfmodel/operators/attention.rs index 2860f905d..1f8d77c55 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 18). + /// payloads are positional (current ENGINE_SPEC_SCHEMA_VERSION 19). #[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 18). +/// field — introduced at schema v8, current ENGINE_SPEC_SCHEMA_VERSION 19). /// 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 18). + /// bincode payloads are positional, current ENGINE_SPEC_SCHEMA_VERSION 19). #[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/dsv41.rs b/crates/core/src/perfmodel/operators/dsv41.rs new file mode 100644 index 000000000..e2e20546b --- /dev/null +++ b/crates/core/src/perfmodel/operators/dsv41.rs @@ -0,0 +1,918 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! DeepSeek-V4.1 text AR operator and phase contracts. +//! +//! Architecture source: deepseek-ai/DeepSeek-V4.1-Flash, revision +//! fb2764a5cf321eaa5070ca8f9e892818f477c16d, config.json / inference/model.py +//! and DeepSeek_V41_Tech_Report.pdf (MIT, Copyright (c) 2023 DeepSeek). +//! Serving layout/scoring source: sgl-project/sglang at +//! 1aa0e962b206102b7c439a4a0c4981cfec6e87bc, +//! python/sglang/srt/layers/attention/deepseek_v4_backend.py and +//! python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py (Apache-2.0, +//! Copyright SGLang contributors). Independently expressed analytical +//! adaptations; see THIRD_PARTY_NOTICES.md and docs/deepseek-v41-storage.md. + +use serde::{Deserialize, Serialize}; + +use crate::common::enums::{DatabaseMode, FmhaQuantMode, GemmQuantMode}; +use crate::common::error::AicError; +use crate::common::system_spec::{SystemSpec, quant_tc_flops}; +use crate::operators::base::{PerformanceResult, SolComponents, Source}; +use crate::operators::op::{Op, RuntimeContext}; +use crate::perf_database::PerfDatabase; + +fn leaf(spec: &SystemSpec, flops: f64, bytes: f64, rate: f64) -> PerformanceResult { + PerformanceResult::sol(SolComponents::new( + flops / rate * 1e3, + bytes / spec.gpu.mem_bw * 1e3, + )) +} + +fn zero() -> PerformanceResult { + PerformanceResult::sol(SolComponents::new(0.0, 0.0)) +} + +/// New V41 kernels have no measured lookup in the SOL release. HYBRID's +/// analytic contribution is explicitly SOL, never an invented utilization or +/// a V4 module hit. SILICON must fail until the V41 collector publishes it. +fn analytic_mode(db: &PerfDatabase, name: &str) -> Result<(), AicError> { + match db.database_mode { + DatabaseMode::Silicon => Err(AicError::PerfDatabase(format!( + "DeepSeek-V4.1 {name} has no measured SILICON data" + ))), + DatabaseMode::Empirical => Err(AicError::EmpiricalNotImplemented(format!( + "DeepSeek-V4.1 {name} has no empirical anchor" + ))), + _ => Ok(()), + } +} + +/// Sum floor(t/r) from t=1 through n; fractional tails retain continuous +/// workload weight for FPM interpolation, with integer publication boundaries. +fn floor_prefix(n: f64, ratio: f64) -> f64 { + let n = n.max(0.0); + let whole = n.floor(); + let q = (whole / ratio).floor(); + let rem = whole - q * ratio; + ratio * q * (q - 1.0) / 2.0 + q * (rem + 1.0) + (n - whole) * q +} + +fn limited_pairs(query: f64, prefix: f64, limit: f64) -> f64 { + let antiderivative = |n: f64| { + let ramp = n.min(limit).max(0.0); + ramp * (ramp + 1.0) / 2.0 + (n - limit).max(0.0) * limit + }; + antiderivative(prefix + query) - antiderivative(prefix) +} + +fn compressed_pairs(query: f64, prefix: f64, ratio: f64, topk: f64) -> f64 { + let antiderivative = |n: f64| { + let saturation = ratio * topk; + floor_prefix(n.min(saturation), ratio) + (n - saturation).max(0.0) * topk + }; + antiderivative(prefix + query) - antiderivative(prefix) +} + +/// Persistent payload layout, independent of attention arithmetic precision. +/// LogicalFp4 is the legacy theoretical layout, not a qualified runtime layout. +#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Dsv41KvCacheLayout { + #[default] + LogicalFp4, + SglangFp8Bf16, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Dsv41AttentionOp { + pub name: String, + pub is_context: bool, + pub role: String, + pub compress_ratio: u32, + pub hidden_size: u32, + pub num_heads: u32, + pub head_dim: u32, + pub q_lora_rank: u32, + pub o_lora_rank: u32, + pub o_groups: u32, + pub index_n_heads: u32, + pub index_head_dim: u32, + pub index_topk: u32, + pub window_size: u32, + pub candidate_limit: u32, + pub is_candidate_source: bool, + pub bounded_prefill: bool, + pub gemm_quant_mode: GemmQuantMode, + pub fmha_quant_mode: FmhaQuantMode, + // Appended in schema 19. The default accepts older JSON; bincode is positional. + #[serde(default)] + pub kv_cache_layout: Dsv41KvCacheLayout, +} + +impl Dsv41AttentionOp { + fn validate_role(&self) -> Result<(), AicError> { + match self.role.as_str() { + "swa" | "full" | "reindex" | "reuse" => Ok(()), + _ => Err(AicError::ModelConfig(format!( + "DeepSeek-V4.1 attention role must be swa, full, reindex, or reuse; got {:?}", + self.role + ))), + } + } + + pub fn weight_bytes(&self) -> f64 { + let (h, q, o, n, d, g) = ( + self.hidden_size as f64, + self.q_lora_rank as f64, + self.o_lora_rank as f64, + self.num_heads as f64, + self.head_dim as f64, + self.o_groups as f64, + ); + let w8 = self.gemm_quant_mode.mapping().memory + + if self.gemm_quant_mode == GemmQuantMode::Fp8Block { + 1.0 / 1024.0 + } else { + 0.0 + }; + let mut bytes = (h * q + q * n * d + h * d + g * o * h) * w8 + n * d * o * 2.0; + if self.role == "full" { + bytes += h * d * 2.0; + if self.compress_ratio > 1 { + bytes += h * d * 2.0; + } + bytes += d * self.index_head_dim as f64 * 2.0; + } + if self.role == "full" || self.role == "reindex" { + bytes += q * self.index_n_heads as f64 * self.index_head_dim as f64 * w8; + bytes += h * self.index_n_heads as f64 * 2.0; + } + // Attention sinks, normalization vectors and compressor positional bias. + bytes + n * 4.0 + (q + d) * 2.0 + } + + /// Shared f64 roofline for native queries and whole-model FPM interpolation. + /// Every serial projection/kernel contributes its own max(math,memory). + pub fn sol( + &self, + spec: &SystemSpec, + batch: f64, + s: f64, + prefix: f64, + ) -> Result { + self.validate_role()?; + if self.role == "full" && self.compress_ratio == 0 { + return Err(AicError::ModelConfig( + "DeepSeek-V4.1 full attention requires a positive compress_ratio".into(), + )); + } + if batch <= 0.0 || s <= 0.0 { + return Ok(zero()); + } + let (h, q, o, n, d, g) = ( + self.hidden_size as f64, + self.q_lora_rank as f64, + self.o_lora_rank as f64, + self.num_heads as f64, + self.head_dim as f64, + self.o_groups as f64, + ); + let (window_entry_bytes, main_entry_bytes) = match self.kv_cache_layout { + Dsv41KvCacheLayout::LogicalFp4 => (d, d * 0.5625), + Dsv41KvCacheLayout::SglangFp8Bf16 => { + // Only the pinned V4.1 low-ratio, 512-wide FlashMLA contract. + if self.head_dim != 512 || self.index_head_dim != 128 || self.compress_ratio > 2 { + return Err(AicError::ModelConfig( + "SGLang V4.1 KV layout requires head_dim=512, index=128, ratio=0/1/2" + .into(), + )); + } + (584.0, 584.0) + } + }; + let tokens = if self.is_context { batch * s } else { batch }; + let bf16 = quant_tc_flops(spec, GemmQuantMode::Bfloat16.mapping())?; + let fp8 = quant_tc_flops(spec, GemmQuantMode::Fp8.mapping())?; + let gemm = quant_tc_flops(spec, self.gemm_quant_mode.mapping())?; + let attn = quant_tc_flops(spec, self.fmha_quant_mode.mapping())?; + let w8 = self.gemm_quant_mode.mapping().memory + + if self.gemm_quant_mode == GemmQuantMode::Fp8Block { + 1.0 / 1024.0 + } else { + 0.0 + }; + let mm = |a: f64, b: f64, nt: f64, weight: f64, rate: f64| { + leaf( + spec, + 2.0 * nt * a * b, + a * b * weight + nt * (a + b) * 2.0, + rate, + ) + }; + let mut result = mm(h, q, tokens, w8, gemm) + .plus(mm(q, n * d, tokens, w8, gemm)) + .plus(mm(h, d, tokens, w8, gemm)) + .plus(leaf( + spec, + 2.0 * tokens * n * d * o, + n * d * o * 2.0 + tokens * (n * d + g * o) * 2.0, + bf16, + )) + .plus(mm(g * o, h, tokens, w8, gemm)); + if self.kv_cache_layout == Dsv41KvCacheLayout::SglangFp8Bf16 { + // Fused norm/RoPE/store reads the BF16 projection output and writes + // one physical SWA row. Projection's own BF16 output write is above. + result = result.plus(leaf( + spec, + 0.0, + tokens * (d * 2.0 + window_entry_bytes), + bf16, + )); + } + let end = if self.is_context { prefix + s } else { s }; + let ratio = self.compress_ratio as f64; + let compressed_len = if ratio > 0.0 { + (end / ratio).floor() + } else { + 0.0 + }; + if self.role == "full" { + let mult = if self.compress_ratio > 1 { 2.0 } else { 1.0 }; + result = result.plus(mm(h, d, tokens, 2.0, bf16).scaled(mult)); + let produced = if self.is_context { + (end / ratio).floor() - (prefix / ratio).floor() + } else { + (end / ratio).floor() - ((end - 1.0).max(0.0) / ratio).floor() + }; + let ihd = self.index_head_dim as f64; + result = result.plus(mm(d, ihd, batch * produced, 2.0, bf16)); + // Main FP4 rounding is temporary: SGLang persists only the 584-byte + // FlashMLA row, not an additional packed main cache. Index stays FP4. + let packed = main_entry_bytes + ihd * 0.53125; + result = result.plus(leaf(spec, 0.0, batch * produced * (d * 2.0 + packed), bf16)); + } + if self.role == "full" || self.role == "reindex" { + let (inh, ihd) = (self.index_n_heads as f64, self.index_head_dim as f64); + result = result + .plus(mm(q, inh * ihd, tokens, w8, fp8)) + .plus(mm(h, inh, tokens, 2.0, bf16)); + // Dense/paged index GEMMs score the full compressed context before + // two_level_decode_logits / _mask_topk_scores masks candidates. + // candidate_limit describes eligibility, not a pre-GEMM gather. + let index_len = compressed_len; + let fp4 = quant_tc_flops(spec, GemmQuantMode::Nvfp4.mapping())?; + result = result.plus(leaf( + spec, + 2.0 * tokens * inh * ihd * index_len, + batch * index_len * ihd * 0.53125 + + tokens * inh * ihd * 0.53125 + + tokens * index_len * 4.0, + fp4, + )); + // Materialized scores, top-k positions, and optional coarse candidate blocks. + let candidate_bytes = if self.is_candidate_source { + tokens * compressed_len / 8.0 * 4.0 + } else { + 0.0 + }; + result = result.plus(leaf( + spec, + 0.0, + tokens * (index_len * 4.0 + self.index_topk as f64 * 4.0) + candidate_bytes, + bf16, + )); + } + let wp = if self.is_context { + batch + * limited_pairs( + s, + if self.bounded_prefill { 0.0 } else { prefix }, + self.window_size as f64, + ) + } else { + batch * end.min(self.window_size as f64) + }; + let cp = if self.compress_ratio == 0 { + 0.0 + } else if self.is_context { + batch * compressed_pairs(s, prefix, ratio, self.index_topk as f64) + } else { + batch * compressed_len.min(self.index_topk as f64) + }; + // SOL assumes ideal reuse of overlapping window entries across + // queries. Pair count determines arithmetic, but HBM sees each + // unique window KV row once. Irregular compressed top-k selections + // retain their per-query traffic; no such reuse is guaranteed there. + let window_rows = if self.is_context { + batch + * (s + if self.bounded_prefill { + 0.0 + } else { + prefix.min((self.window_size as f64 - 1.0).max(0.0)) + }) + } else { + batch * end.min(self.window_size as f64) + }; + result = result.plus(leaf( + spec, + 4.0 * n * d * (wp + cp), + (window_rows * window_entry_bytes + cp * main_entry_bytes) + tokens * n * d * 4.0, + attn, + )); + Ok(result) + } + + pub fn query( + &self, + db: &PerfDatabase, + ctx: &RuntimeContext, + ) -> Result { + self.validate_role()?; + analytic_mode(db, "CSA2 attention")?; + self.sol( + &db.system_spec, + ctx.batch_size as f64, + ctx.s as f64, + ctx.prefix as f64, + ) + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Dsv41MhcOp { + pub name: String, + pub hidden_size: u32, + pub hc_mult: u32, + pub sinkhorn_iters: u32, +} +impl Dsv41MhcOp { + pub fn weight_bytes(&self) -> f64 { + let hc = self.hc_mult as f64; + 2.0 * ((hc + 2.0) * hc * (hc * self.hidden_size as f64 + 1.0) + 3.0) * 4.0 + } + pub fn sol(&self, spec: &SystemSpec, tokens: f64) -> Result { + if tokens <= 0.0 { + return Ok(zero()); + } + let (h, hc) = (self.hidden_size as f64, self.hc_mult as f64); + let mixes = (hc + 2.0) * hc; + let ops = 2.0 + * tokens + * (2.0 * hc * h * mixes + + (hc * hc + 2.0 * hc) * self.sinkhorn_iters as f64 + + 2.0 * hc * hc * h + + 2.0 * hc * h); + // Single pass over each attention/FFN residual site. Coefficients are + // FP32; residual activation traffic remains BF16. + let bytes = + self.weight_bytes() + 2.0 * tokens * hc * h * 2.0 * 2.0 + 2.0 * tokens * mixes * 4.0; + let fp32 = spec + .gpu + .fp32_flops + .filter(|v| v.is_finite() && *v > 0.0) + .ok_or_else(|| { + AicError::MissingSystemFlops( + "DeepSeek-V4.1 scalar mHC requires fp32_flops in the system spec".into(), + ) + })?; + Ok(leaf(spec, ops, bytes, fp32)) + } + pub fn query(&self, db: &PerfDatabase, tokens: u32) -> Result { + analytic_mode(db, "single-pass mHC")?; + self.sol(&db.system_spec, tokens as f64) + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Dsv41EngramOp { + pub name: String, + pub num_embeddings: u64, + pub head_dim: u32, + pub hash_columns: u32, + pub hidden_size: u32, + pub hc_mult: u32, + pub tp_size: u32, +} +impl Dsv41EngramOp { + pub fn weight_bytes(&self) -> f64 { + let rows = self.num_embeddings.div_ceil(self.tp_size as u64) as f64; + let d = self.head_dim as f64; + let projection = + self.hash_columns as f64 * d * self.hidden_size as f64 * (self.hc_mult + 1) as f64; + rows * (d + d / 32.0) + + projection * (1.0 + 1.0 / 1024.0) + + 2.0 * self.hc_mult as f64 * self.hidden_size as f64 * 2.0 + } + pub fn sol(&self, spec: &SystemSpec, tokens: f64) -> Result { + if tokens <= 0.0 { + return Ok(zero()); + } + let input = self.hash_columns as f64 * self.head_dim as f64; + let output = self.hidden_size as f64 * (self.hc_mult + 1) as f64; + let fp8 = quant_tc_flops(spec, GemmQuantMode::Fp8.mapping())?; + let lookup = leaf( + spec, + 0.0, + tokens * input * ((1.0 + 1.0 / 32.0) / self.tp_size as f64 + 2.0), + fp8, + ); + let projection = leaf( + spec, + 2.0 * tokens * input * output, + input * output * (1.0 + 1.0 / 1024.0) + tokens * (input + output) * 2.0, + fp8, + ); + let gate = leaf( + spec, + 0.0, + tokens * self.hidden_size as f64 * self.hc_mult as f64 * 8.0, + fp8, + ); + Ok(lookup.plus(projection).plus(gate)) + } + pub fn query(&self, db: &PerfDatabase, tokens: u32) -> Result { + analytic_mode(db, "Engram")?; + self.sol(&db.system_spec, tokens as f64) + } +} + +/// Checkpoint-native dense projection. V41's FP8 block is 32x32, so it +/// must not query the legacy 128x128 FP8Block GEMM table. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Dsv41LinearOp { + pub name: String, + pub n: u32, + pub k: u32, + pub quant_mode: GemmQuantMode, +} +impl Dsv41LinearOp { + pub fn weight_bytes(&self) -> f64 { + let elements = self.n as f64 * self.k as f64; + elements * self.quant_mode.mapping().memory + + if self.quant_mode == GemmQuantMode::Fp8Block { + self.n.div_ceil(32) as f64 * self.k.div_ceil(32) as f64 + } else { + 0.0 + } + } + pub fn sol(&self, spec: &SystemSpec, tokens: f64) -> Result { + if tokens <= 0.0 { + return Ok(zero()); + } + let rate = quant_tc_flops(spec, self.quant_mode.mapping())?; + Ok(leaf( + spec, + 2.0 * tokens * self.n as f64 * self.k as f64, + self.weight_bytes() + tokens * (self.n + self.k) as f64 * 2.0, + rate, + )) + } + pub fn query(&self, db: &PerfDatabase, tokens: u32) -> Result { + analytic_mode(db, "32x32 dense projection")?; + self.sol(&db.system_spec, tokens as f64) + } +} + +/// A decoder layer's token domain. All resident weights stay present when its +/// late-layer prefill inputs are shortened. Decoder replay never shortens decode. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Dsv41StageOp { + pub name: String, + pub is_context: bool, + pub decoder_replay: bool, + pub bounded: bool, + pub window_size: u32, + pub children: Vec, +} +impl Dsv41StageOp { + pub fn scope(&self, s: f64, prefix: f64) -> (f64, f64) { + if self.is_context && self.decoder_replay && self.bounded { + let tail = s.min(self.window_size as f64); + (tail, prefix + s - tail) + } else { + (s, prefix) + } + } + pub fn weight_bytes(&self) -> f64 { + self.children.iter().map(Op::weight_bytes).sum() + } + pub fn query( + &self, + db: &PerfDatabase, + ctx: &RuntimeContext, + ) -> Result { + let (s, prefix) = self.scope(ctx.s as f64, ctx.prefix as f64); + let mut scoped = *ctx; + scoped.s = s as u32; + scoped.prefix = prefix as u32; + if self.is_context { + scoped.num_tokens = ctx.batch_size * scoped.s; + } + let mut total = None; + for op in &self.children { + let mut child = scoped; + if self.is_context && op.is_logits_gemm() { + child.num_tokens = ctx.batch_size; + } + let result = op.query(db, &child)?; + total = Some(match total { + None => result, + Some(prev) => PerformanceResult::plus(prev, result), + }); + } + Ok(total.unwrap_or_else(|| PerformanceResult::new(0.0, Source::Sol))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::enums::TransferPolicy; + use std::path::PathBuf; + + fn test_db(mode: DatabaseMode) -> PerfDatabase { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../python/aisimulate/src/aiconfigurator_core/systems"); + PerfDatabase::load(&root, "gb300", "sglang", "0.5.14") + .unwrap() + .with_mode(mode, TransferPolicy::ALL) + } + + fn unit_spec() -> SystemSpec { + let mut spec = test_db(DatabaseMode::Sol).system_spec.clone(); + spec.gpu.mem_bw = 1e6; + spec.gpu.bfloat16_tc_flops = Some(1e9); + spec.gpu.fp8_tc_flops = Some(1e9); + spec.gpu.fp4_tc_flops = Some(1e9); + spec.gpu.fp32_flops = Some(1e9); + spec + } + + fn attention(role: &str, ratio: u32) -> Dsv41AttentionOp { + Dsv41AttentionOp { + name: "attention".into(), + is_context: true, + role: role.into(), + compress_ratio: ratio, + hidden_size: 8, + num_heads: 2, + head_dim: 4, + q_lora_rank: 4, + o_lora_rank: 2, + o_groups: 1, + index_n_heads: 4, + index_head_dim: 2, + index_topk: 3, + window_size: 128, + candidate_limit: 0, + is_candidate_source: false, + bounded_prefill: false, + gemm_quant_mode: GemmQuantMode::Bfloat16, + fmha_quant_mode: FmhaQuantMode::Bfloat16, + kv_cache_layout: Dsv41KvCacheLayout::LogicalFp4, + } + } + + fn assert_components(result: PerformanceResult, flops: f64, bytes: f64) { + let components = result.sol.unwrap(); + assert!((components.math_ms - flops / 1e6).abs() < 1e-10); + assert!((components.mem_ms - bytes / 1e3).abs() < 1e-10); + assert_eq!(result.source, Source::Sol); + } + + #[test] + fn all_attention_roles_have_independent_numeric_rooflines_and_weights() { + // Hand-expanded B=1,Q=4,P=0 ledger. Base projections: 1024 FLOPs, + // 704 bytes; the SWA kernel has 10 causal pairs but four KV rows. + // Full ratio-two owns TWO 8x4 BF16 compressor matrices, not one. + for (role, ratio, weights, flops, bytes) in [ + ("swa", 0, 280.0, 1344.0, 848.0), + ("reuse", 2, 280.0, 1472.0, 857.0), + ("reindex", 2, 408.0, 2112.0, 1308.125), + ("full", 2, 552.0, 2656.0, 1690.75), + ("full", 1, 488.0, 2720.0, 1654.75), + ] { + let op = attention(role, ratio); + assert_eq!(op.weight_bytes(), weights, "{role}, ratio={ratio}"); + assert_components(op.sol(&unit_spec(), 1.0, 4.0, 0.0).unwrap(), flops, bytes); + } + } + + #[test] + fn serialized_full_attention_rejects_zero_compression_in_both_phases() { + let spec = unit_spec(); + for is_context in [true, false] { + let mut value = serde_json::to_value(attention("full", 2)).unwrap(); + value["compress_ratio"] = serde_json::json!(0); + value["is_context"] = serde_json::json!(is_context); + let malformed: Dsv41AttentionOp = serde_json::from_value(value).unwrap(); + for (batch, sequence) in [(1.0, 128.0), (0.0, 128.0), (1.0, 0.0)] { + let error = malformed.sol(&spec, batch, sequence, 32.0).unwrap_err(); + assert!(matches!(error, AicError::ModelConfig(ref message) + if message.contains("positive compress_ratio"))); + } + // Ratio zero is valid for the uncompressed sliding-window role. + let mut sliding_window = malformed; + sliding_window.role = "swa".into(); + let result = sliding_window.sol(&spec, 1.0, 128.0, 32.0).unwrap(); + assert!(result.latency_ms.is_finite()); + assert!(result.latency_ms > 0.0); + } + } + + #[test] + fn window_flops_use_pairs_but_hbm_uses_unique_rows() { + let spec = unit_spec(); + let mut op = attention("swa", 0); + let short = op.sol(&spec, 1.0, 4.0, 0.0).unwrap().sol.unwrap(); + let prefix = op.sol(&spec, 1.0, 4.0, 1000.0).unwrap().sol.unwrap(); + // Four queries now attend 512 pairs, but load just 127 additional rows. + assert!((prefix.math_ms - short.math_ms - 32.0 * (512.0 - 10.0) / 1e6).abs() < 1e-12); + assert!((prefix.mem_ms - short.mem_ms - 127.0 * 4.0 / 1e3).abs() < 1e-12); + op.bounded_prefill = true; + assert_eq!(op.sol(&spec, 1.0, 4.0, 1000.0).unwrap().sol.unwrap(), short); + // Remove heads/projections so this isolates the 4096-row SWA traffic. + op.hidden_size = 0; + op.q_lora_rank = 0; + op.o_lora_rank = 0; + op.num_heads = 0; + op.o_groups = 0; + op.head_dim = 512; + // The empty hidden->KV projection writes each KV row in BF16; + // attention then reads each row once in FP8: (2+1)*4096*512 bytes. + assert_components(op.sol(&spec, 1.0, 4096.0, 0.0).unwrap(), 0.0, 6_291_456.0); + } + + #[test] + fn candidate_eligibility_never_caps_full_context_scoring() { + let spec = unit_spec(); + let mut op = attention("reindex", 1); + op.is_context = false; + op.index_topk = 3; // Attention selection has saturated well before these inputs. + op.candidate_limit = 16_384; + for context in [16_384.0, 16_385.0, 131_072.0] { + let limited = op.sol(&spec, 1.0, context, 0.0).unwrap(); + op.candidate_limit = 0; + assert_eq!(limited, op.sol(&spec, 1.0, context, 0.0).unwrap()); + op.candidate_limit = 16_384; + } + let lo = op.sol(&spec, 1.0, 16_384.0, 0.0).unwrap().sol.unwrap(); + let hi = op.sol(&spec, 1.0, 131_072.0, 0.0).unwrap().sol.unwrap(); + // Per extra key: 2*4 heads*2 dimensions =16 FLOPs; 2*17/32=1.0625 + // packed K bytes, plus4 score-write and4 score-read bytes. All other + // geometry is saturated and constant in this independent decode ledger. + let extra = 131_072.0 - 16_384.0; + assert!((hi.math_ms - lo.math_ms - 16.0 * extra / 1e6).abs() < 1e-10); + assert!((hi.mem_ms - lo.mem_ms - 9.0625 * extra / 1e3).abs() < 1e-10); + } + + #[test] + fn prefixed_prefill_scores_across_candidate_boundary() { + let spec = unit_spec(); + let mut op = attention("reindex", 1); + op.candidate_limit = 16_384; + // Three actual queries with a long cached prefix. Their sparse/window + // attention is saturated; the full index matrix still grows with end. + for end in [16_384.0, 16_385.0, 131_072.0] { + let limited = op.sol(&spec, 1.0, 3.0, end - 3.0).unwrap(); + op.candidate_limit = 0; + assert_eq!(limited, op.sol(&spec, 1.0, 3.0, end - 3.0).unwrap()); + op.candidate_limit = 16_384; + } + let lo = op.sol(&spec, 1.0, 3.0, 16_381.0).unwrap().sol.unwrap(); + let hi = op.sol(&spec, 1.0, 3.0, 16_382.0).unwrap().sol.unwrap(); + // One extra K row: 3 queries *16 FLOPs; 1.0625 K bytes plus + // 3*(4-byte score write +4-byte score read). Prefix is not re-executed. + assert!((hi.math_ms - lo.math_ms - 48.0 / 1e6).abs() < 1e-10); + assert!((hi.mem_ms - lo.mem_ms - 25.0625 / 1e3).abs() < 1e-10); + } + + #[test] + fn sglang_physical_reads_and_publication_count_one_persistent_record() { + let spec = unit_spec(); + let mut op = attention("full", 2); + op.is_context = false; + op.head_dim = 512; + op.index_head_dim = 128; + // Relative to the separately hand-tested logical ledger: every token + // stores its SWA row (1024-byte source read +584-byte write), the128 + // resident window rows add72 bytes each, and the three selected main + // entries add296 each. Even tokens publish one more main row (+296). + for (context, published) in [(129.0, 0.0), (130.0, 1.0)] { + let logical = op.sol(&spec, 1.0, context, 0.0).unwrap().sol.unwrap(); + op.kv_cache_layout = Dsv41KvCacheLayout::SglangFp8Bf16; + let physical = op.sol(&spec, 1.0, context, 0.0).unwrap().sol.unwrap(); + assert_eq!(physical.math_ms, logical.math_ms); + let added = 1024.0 + 584.0 + 128.0 * 72.0 + (3.0 + published) * 296.0; + assert!((physical.mem_ms - logical.mem_ms - added / 1e3).abs() < 1e-10); + op.kv_cache_layout = Dsv41KvCacheLayout::LogicalFp4; + } + op.role = "swa".into(); + op.compress_ratio = 0; + op.is_context = true; + op.hidden_size = 0; + op.q_lora_rank = 0; + op.o_lora_rank = 0; + op.num_heads = 0; + op.o_groups = 0; + op.kv_cache_layout = Dsv41KvCacheLayout::SglangFp8Bf16; + // Per query: BF16 projection write1024, store read1024/write584, + // attention read584 =3216 bytes. No persistent packed FP4 record. + assert_components(op.sol(&spec, 1.0, 4096.0, 0.0).unwrap(), 0.0, 13_172_736.0); + } + + #[test] + fn physical_layout_rejects_unqualified_geometry_and_json_defaults_are_theoretical() { + let mut op = attention("full", 2); + let mut json = serde_json::to_value(&op).unwrap(); + json.as_object_mut().unwrap().remove("kv_cache_layout"); + let legacy: Dsv41AttentionOp = serde_json::from_value(json.clone()).unwrap(); + assert_eq!(legacy.kv_cache_layout, Dsv41KvCacheLayout::LogicalFp4); + json["kv_cache_layout"] = serde_json::json!("unknown"); + assert!(serde_json::from_value::(json).is_err()); + op.kv_cache_layout = Dsv41KvCacheLayout::SglangFp8Bf16; + assert!(op.sol(&unit_spec(), 1.0, 128.0, 0.0).is_err()); + op.head_dim = 512; + op.index_head_dim = 128; + for ratio in [4, 128] { + op.compress_ratio = ratio; + assert!(op.sol(&unit_spec(), 1.0, 128.0, 0.0).is_err()); + } + } + + #[test] + fn mhc_prices_both_residual_sites_at_scalar_fp32_rate() { + let op = Dsv41MhcOp { + name: "mhc".into(), + hidden_size: 5120, + hc_mult: 4, + sinkhorn_iters: 20, + }; + assert_eq!(op.weight_bytes(), 3_932_376.0); + assert_components(op.sol(&unit_spec(), 2.0).unwrap(), 4_753_280.0, 4_260_440.0); + let mut spec = unit_spec(); + for rate in [ + None, + Some(0.0), + Some(-1.0), + Some(f64::NAN), + Some(f64::INFINITY), + ] { + spec.gpu.fp32_flops = rate; + assert!(matches!( + op.sol(&spec, 1.0), + Err(AicError::MissingSystemFlops(_)) + )); + } + } + + #[test] + fn engram_production_weights_include_exact_tables_and_replicated_projections() { + for (tp_size, expected) in [ + (1, 203_073_076_240.0), + (4, 51_004_552_072.0), + (8, 25_659_798_088.0), + ] { + let total: f64 = [384_006_168, 384_016_682] + .into_iter() + .map(|num_embeddings| { + Dsv41EngramOp { + name: "engram".into(), + num_embeddings, + head_dim: 256, + hash_columns: 24, + hidden_size: 5120, + hc_mult: 4, + tp_size, + } + .weight_bytes() + }) + .sum(); + assert_eq!(total, expected); + } + let op = Dsv41EngramOp { + name: "small".into(), + num_embeddings: 64, + head_dim: 32, + hash_columns: 2, + hidden_size: 8, + hc_mult: 4, + tp_size: 4, + }; + assert_eq!(op.weight_bytes(), 3218.5); + // Lookup 289 bytes + projection 2978.5 bytes + gate 512 bytes. + assert_components(op.sol(&unit_spec(), 2.0).unwrap(), 10240.0, 3779.5); + } + + #[test] + fn block32_linear_counts_partial_scale_tiles_and_numeric_roofline() { + let op = Dsv41LinearOp { + name: "linear".into(), + n: 33, + k: 65, + quant_mode: GemmQuantMode::Fp8Block, + }; + assert_eq!(op.weight_bytes(), 2151.0); // 2145 weights + 2*3 scale bytes + assert_components(op.sol(&unit_spec(), 2.0).unwrap(), 8580.0, 2543.0); + } + + #[test] + fn new_operators_reject_unmeasured_modes_and_empty_work_is_zero() { + let attn = attention("full", 2); + let mhc = Dsv41MhcOp { + name: "mhc".into(), + hidden_size: 8, + hc_mult: 4, + sinkhorn_iters: 20, + }; + let engram = Dsv41EngramOp { + name: "engram".into(), + num_embeddings: 64, + head_dim: 32, + hash_columns: 2, + hidden_size: 8, + hc_mult: 4, + tp_size: 4, + }; + let linear = Dsv41LinearOp { + name: "linear".into(), + n: 32, + k: 32, + quant_mode: GemmQuantMode::Fp8Block, + }; + for mode in [DatabaseMode::Silicon, DatabaseMode::Empirical] { + let db = test_db(mode); + let ctx = RuntimeContext { + batch_size: 1, + s: 4, + prefix: 0, + num_tokens: 4, + ..RuntimeContext::default() + }; + for result in [ + attn.query(&db, &ctx), + mhc.query(&db, 4), + engram.query(&db, 4), + linear.query(&db, 4), + ] { + match mode { + DatabaseMode::Silicon => { + assert!(matches!(result, Err(AicError::PerfDatabase(_)))) + } + _ => assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_)))), + } + } + } + let spec = unit_spec(); + for result in [ + attn.sol(&spec, 1.0, 0.0, 0.0), + mhc.sol(&spec, 0.0), + engram.sol(&spec, 0.0), + linear.sol(&spec, 0.0), + ] { + assert_eq!(result.unwrap().latency_ms, 0.0); + } + } + #[test] + fn deserialized_attention_roles_are_validated_before_empty_work() { + for role in ["ful", "", "FULL", "swa "] { + let mut json = serde_json::to_value(attention("full", 2)).unwrap(); + json["role"] = serde_json::Value::String(role.into()); + let op: Dsv41AttentionOp = serde_json::from_value(json).unwrap(); + for (batch, seq) in [(1.0, 4.0), (0.0, 4.0), (1.0, 0.0)] { + assert!(matches!( + op.sol(&unit_spec(), batch, seq, 0.0), + Err(AicError::ModelConfig(_)) + )); + } + } + for role in ["swa", "full", "reindex", "reuse"] { + let json = serde_json::to_value(attention(role, 2)).unwrap(); + let op: Dsv41AttentionOp = serde_json::from_value(json.clone()).unwrap(); + assert_eq!(serde_json::to_value(&op).unwrap(), json); + assert!(op.sol(&unit_spec(), 1.0, 4.0, 0.0).is_ok()); + assert_eq!(op.sol(&unit_spec(), 0.0, 4.0, 0.0).unwrap().latency_ms, 0.0); + } + } + + #[test] + fn bounded_scope_preserves_absolute_position() { + let stage = Dsv41StageOp { + name: "tail".into(), + is_context: true, + decoder_replay: true, + bounded: true, + window_size: 128, + children: vec![], + }; + assert_eq!(stage.scope(256.0, 1000.0), (128.0, 1128.0)); + assert_eq!(stage.scope(3.0, 1000.0), (3.0, 1000.0)); + } + #[test] + fn compression_publication_and_saturation() { + assert_eq!(compressed_pairs(4.0, 0.0, 2.0, 512.0), 4.0); + assert_eq!(compressed_pairs(2.0, 1024.0, 2.0, 512.0), 1024.0); + assert_eq!(compressed_pairs(1.0, 128.0, 2.0, 512.0), 64.0); + } +} diff --git a/crates/core/src/perfmodel/operators/mod.rs b/crates/core/src/perfmodel/operators/mod.rs index e500ad385..04603f622 100644 --- a/crates/core/src/perfmodel/operators/mod.rs +++ b/crates/core/src/perfmodel/operators/mod.rs @@ -18,6 +18,7 @@ pub mod base; pub mod communication; pub mod dsa; pub mod dsv4; +pub mod dsv41; pub mod elementwise; pub mod embedding; pub mod fpm_forward; @@ -43,6 +44,7 @@ pub use base::{PerformanceResult, Source}; pub use communication::{CustomAllReduceOp, NcclOp, P2POp}; pub use dsa::DsaModuleOp; pub use dsv4::{Dsv4MegaMoeOp, Dsv4ModuleOp}; +pub use dsv41::{Dsv41AttentionOp, Dsv41EngramOp, Dsv41LinearOp, Dsv41MhcOp, Dsv41StageOp}; pub use elementwise::ElementwiseOp; pub use embedding::EmbeddingOp; pub use fpm_forward::{FpmForwardOp, FpmPhase}; diff --git a/crates/core/src/perfmodel/operators/op.rs b/crates/core/src/perfmodel/operators/op.rs index b98840a96..19410e133 100644 --- a/crates/core/src/perfmodel/operators/op.rs +++ b/crates/core/src/perfmodel/operators/op.rs @@ -21,10 +21,11 @@ use serde::{Deserialize, Serialize}; use crate::common::error::AicError; use crate::operators::{ ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4MegaMoeOp, Dsv4ModuleOp, - ElementwiseOp, EmbeddingOp, EncoderAttentionOp, FpmForwardOp, GdnOp, GemmOp, - GenerationAttentionOp, GenerationMlaOp, KdaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, - MoEDispatchOp, MoeAllToAllOp, MoeExpertComputeOp, MoeOp, MsaModuleOp, NcclOp, P2POp, - PerformanceResult, Source, VisionEncoderOp, WideEpContextMlaOp, WideEpGenerationMlaOp, + Dsv41AttentionOp, Dsv41EngramOp, Dsv41LinearOp, Dsv41MhcOp, Dsv41StageOp, ElementwiseOp, + EmbeddingOp, EncoderAttentionOp, FpmForwardOp, GdnOp, GemmOp, GenerationAttentionOp, + GenerationMlaOp, KdaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, MoEDispatchOp, + MoeAllToAllOp, MoeExpertComputeOp, MoeOp, MsaModuleOp, NcclOp, P2POp, PerformanceResult, + Source, VisionEncoderOp, WideEpContextMlaOp, WideEpGenerationMlaOp, }; use crate::perf_database::PerfDatabase; @@ -177,6 +178,11 @@ pub enum Op { /// ratio changes query dimensions before lookup, never the result. /// Appended as part of the speculative-decoding schema-18 migration. TokenScale(TokenScaleOp), + Dsv41Attention(Dsv41AttentionOp), + Dsv41Mhc(Dsv41MhcOp), + Dsv41Engram(Dsv41EngramOp), + Dsv41Stage(Dsv41StageOp), + Dsv41Linear(Dsv41LinearOp), } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -266,6 +272,11 @@ impl Op { /// family multiplies its own scale_factor inside its `weight_bytes`. pub fn weight_bytes(&self) -> f64 { match self { + Op::Dsv41Attention(o) => o.weight_bytes(), + Op::Dsv41Mhc(o) => o.weight_bytes(), + Op::Dsv41Engram(o) => o.weight_bytes(), + Op::Dsv41Stage(o) => o.weight_bytes(), + Op::Dsv41Linear(o) => o.weight_bytes(), Op::TokenScale(o) => o.op.weight_bytes(), Op::Gemm(o) => o.weights_bytes(), Op::Embedding(o) => o.weights_bytes(), @@ -321,6 +332,11 @@ impl Op { /// debugging. pub fn name(&self) -> &str { match self { + Op::Dsv41Attention(o) => &o.name, + Op::Dsv41Mhc(o) => &o.name, + Op::Dsv41Engram(o) => &o.name, + Op::Dsv41Stage(o) => &o.name, + Op::Dsv41Linear(o) => &o.name, Op::TokenScale(o) => o.op.name(), Op::Gemm(o) => &o.name, Op::Embedding(o) => &o.name, @@ -365,6 +381,11 @@ impl Op { /// returns them). Every variant carries `name`. pub fn set_name(&mut self, name: String) { match self { + Op::Dsv41Attention(o) => o.name = name, + Op::Dsv41Mhc(o) => o.name = name, + Op::Dsv41Engram(o) => o.name = name, + Op::Dsv41Stage(o) => o.name = name, + Op::Dsv41Linear(o) => o.name = name, Op::TokenScale(o) => o.op.set_name(name), Op::Gemm(o) => o.name = name, Op::Embedding(o) => o.name = name, @@ -439,6 +460,13 @@ impl Op { Op::FpmForward(_) => {} // no scale_factor on this family (composite/whole-model) Op::Overlap(_) => {} // no scale_factor on this family (composite/whole-model) Op::Fallback(_) => {} // no scale_factor on this family (composite/whole-model) + // V4.1 operators encode individual stages, without a repetition field. + // The Python setter rejects these families before reaching this mutator. + Op::Dsv41Attention(_) + | Op::Dsv41Mhc(_) + | Op::Dsv41Engram(_) + | Op::Dsv41Stage(_) + | Op::Dsv41Linear(_) => {} Op::Dsv4MegaMoe(o) => o.scale_factor = scale_factor, Op::Kda(o) => o.scale_factor = scale_factor, Op::MoeAllToAll(o) => o.scale_factor = scale_factor, @@ -500,6 +528,11 @@ impl Op { ctx: &RuntimeContext, ) -> Result { match self { + Op::Dsv41Attention(op) => op.query(db, ctx), + Op::Dsv41Mhc(op) => op.query(db, ctx.num_tokens), + Op::Dsv41Engram(op) => op.query(db, ctx.num_tokens), + Op::Dsv41Stage(op) => op.query(db, ctx), + Op::Dsv41Linear(op) => op.query(db, ctx.num_tokens), Op::TokenScale(op) => { let scaled = RuntimeContext { batch_size: op.scale_tokens(ctx.batch_size)?, diff --git a/crates/core/src/perfmodel/perf_database/gemm.rs b/crates/core/src/perfmodel/perf_database/gemm.rs index b31cbe778..c7b315a76 100644 --- a/crates/core/src/perfmodel/perf_database/gemm.rs +++ b/crates/core/src/perfmodel/perf_database/gemm.rs @@ -1260,6 +1260,7 @@ mod tests { let mut spec = SystemSpec { data_dir: std::path::PathBuf::from("data/synthetic"), gpu: GpuSpec { + fp32_flops: None, mem_bw: 1.0, mem_bw_empirical_scaling_factor: 1.0, mem_empirical_constant_latency: 0.0, diff --git a/crates/core/src/perfmodel/perf_database/mod.rs b/crates/core/src/perfmodel/perf_database/mod.rs index e68b5f5da..5ad43b733 100644 --- a/crates/core/src/perfmodel/perf_database/mod.rs +++ b/crates/core/src/perfmodel/perf_database/mod.rs @@ -852,6 +852,7 @@ pub(crate) mod energy_test_fixtures { SystemSpec { data_dir: "data".into(), gpu: GpuSpec { + fp32_flops: None, mem_bw: 7.7e12, mem_bw_empirical_scaling_factor: 0.92, mem_empirical_constant_latency: 2e-6, diff --git a/crates/core/src/perfmodel/perf_database/moe_expert_compute.rs b/crates/core/src/perfmodel/perf_database/moe_expert_compute.rs index 301d22b31..3758b4f23 100644 --- a/crates/core/src/perfmodel/perf_database/moe_expert_compute.rs +++ b/crates/core/src/perfmodel/perf_database/moe_expert_compute.rs @@ -762,6 +762,7 @@ mod tests { SystemSpec { data_dir: PathBuf::from("data/synthetic"), gpu: GpuSpec { + fp32_flops: None, mem_bw: 1e9, mem_bw_empirical_scaling_factor: 1.0, mem_empirical_constant_latency: 0.0, diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index cfe0cc340..d81f0654f 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1052,6 +1052,7 @@ struct EngineBuildRequest { kv_block_size: Option, systems_path: Option, forward_model: Option, + decoder_replay: bool, database_mode: Option, shared_layer: Option, transfer_policy: Option>, @@ -1097,6 +1098,7 @@ impl AicEngineBuilder { kv_block_size: None, systems_path: None, forward_model: None, + decoder_replay: false, database_mode: None, shared_layer: None, transfer_policy: None, @@ -1112,6 +1114,12 @@ impl AicEngineBuilder { self } + /// Select the verified V4.1 bounded decoder execution profile. + pub fn decoder_replay(mut self, enabled: bool) -> Self { + self.request.decoder_replay = enabled; + self + } + /// Select the performance-database lookup mode. Unset keeps SILICON. pub fn database_mode(mut self, database_mode: DatabaseMode) -> Self { self.request.database_mode = Some(database_mode.as_str().to_owned()); @@ -1381,6 +1389,7 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result) -> En kv_block_size: config.kv_block_size, systems_path: systems_path.map(str::to_owned), forward_model: config.forward_model.clone(), + decoder_replay: config.decoder_replay, database_mode: Some(config.database_mode.as_str().to_owned()), shared_layer: config.enable_shared_layer, transfer_policy: config.transfer_policy.clone(), @@ -1710,6 +1721,7 @@ impl PyForwardPassPerfModel { nextn: request.nextn, speculation: request.speculation, kv_block_size: request.kv_block_size, + decoder_replay: request.decoder_replay, estimation_mode, database_mode: legacy.database_mode, transfer_policy: request.transfer_policy, @@ -1930,6 +1942,7 @@ mod tests { backend: BackendKind::Vllm, backend_version: Some("0.24.0".to_string()), forward_model: None, + decoder_replay: false, kv_block_size: None, parallel: ParallelMapping { tp_size: 8, diff --git a/crates/core/src/perfmodel/py_ops.rs b/crates/core/src/perfmodel/py_ops.rs index aa6446be4..e651893ee 100644 --- a/crates/core/src/perfmodel/py_ops.rs +++ b/crates/core/src/perfmodel/py_ops.rs @@ -165,9 +165,13 @@ pub(crate) fn wrap_op(py: Python<'_>, op: Op) -> PyResult> { // FpmForward has no family class: FPMForwardOp stays a Python class // (callable slot + pinned signature) whose spec adapter converts to a // BASE-wrapped engine op for list assembly. - Op::FpmForward(_) | Op::TokenScale(_) => { - Ok(Py::new(py, PyOperation { inner: op })?.into_any()) - } + Op::FpmForward(_) + | Op::Dsv41Attention(_) + | Op::Dsv41Mhc(_) + | Op::Dsv41Engram(_) + | Op::Dsv41Stage(_) + | Op::Dsv41Linear(_) + | Op::TokenScale(_) => Ok(Py::new(py, PyOperation { inner: op })?.into_any()), // Vision is never wrapped: compile decomposes it into child ops. other => Err(PyTypeError::new_err(format!( "no Python class wrapper for engine op variant {:?}", @@ -4289,6 +4293,7 @@ pub(crate) fn reject_retired_ops(ops: &[Op]) -> Result<(), String> { reject_retired_ops(&o.fallback)?; } Op::FpmForward(o) => reject_retired_ops(&o.sol_ops)?, + Op::Dsv41Stage(o) => reject_retired_ops(&o.children)?, Op::TokenScale(o) => reject_retired_ops(std::slice::from_ref(&o.op))?, _ => {} } diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 871621bb6..7f798888e 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -192,6 +192,8 @@ struct AicTimingConfig { #[serde(default)] forward_model: Option, #[serde(default)] + decoder_replay: bool, + #[serde(default)] worker_type: Option, #[serde(default)] estimation_mode: Option, @@ -207,7 +209,7 @@ struct AicTimingConfig { systems_paths: Vec, #[serde(default)] attention_backend: Option, - #[serde(default)] + #[serde(default, alias = "shared_layer")] enable_shared_layer: Option, #[serde(default)] strict_provenance: bool, @@ -271,6 +273,7 @@ impl AicTimingConfig { nextn: self.nextn, speculation: self.speculation.clone(), kv_block_size: self.kv_block_size, + decoder_replay: self.decoder_replay, estimation_mode: mode, fallback_policy: self.fallback_policy, estimator_config: self.estimator_config.clone(), @@ -373,6 +376,7 @@ type PhaseEvidenceKey = (u32, u32, u32, u32, bool); struct AicTimingModel { engine: Py, + decoder_replay: bool, use_fpm_decode_totals: bool, fpm_decode_kv_ceiling: Option, evidence: Mutex, @@ -420,6 +424,7 @@ impl AicTimingModel { })?; Ok(Self { engine, + decoder_replay: config.decoder_replay, use_fpm_decode_totals, fpm_decode_kv_ceiling, evidence: Mutex::new(TimingEvidenceSummary::default()), @@ -509,6 +514,21 @@ fn phase_evidence_from_python( } impl TimingModel for AicTimingModel { + fn prefill_batch_validation_can_fail(&self) -> bool { + self.decoder_replay + } + + fn validate_prefill_batch(&self, requests: &[(usize, usize)]) -> Result<()> { + if self.decoder_replay { + ensure!( + requests.windows(2).all(|pair| pair[0] == pair[1]), + "decoder replay requires identical per-request new-token and cached-prefix lengths; \ + heterogeneous prefill cannot be represented by the mean-based replay timing API" + ); + } + Ok(()) + } + fn predict_prefill_ms( &self, batch_size: usize, @@ -564,9 +584,14 @@ impl TimingModel for AicTimingModel { } let batch_size = checked_u32(batch_size, "decode batch size")?; - let mean_context_length = checked_u32(mean_context_length, "mean context length")?; - let evidence = - self.predict_phase_evidence(batch_size, mean_context_length, 2, 0, "static_gen")?; + // Both scheduler implementations include the current input token in + // their sequence length before sampling the next token. The op API's + // isl is past KV; osl=2 adds the current token exactly once. + let mean_past_kv = mean_context_length + .checked_sub(1) + .context("mean decode context must include the current input token")?; + let mean_past_kv = checked_u32(mean_past_kv, "mean past KV length")?; + let evidence = self.predict_phase_evidence(batch_size, mean_past_kv, 2, 0, "static_gen")?; let latency_ms = evidence.latency_ms; self.record_evidence(evidence, false)?; Ok(latency_ms) @@ -1990,6 +2015,7 @@ mod tests { fn timing_model(engine: Py, use_fpm_decode_totals: bool) -> AicTimingModel { AicTimingModel { engine, + decoder_replay: false, use_fpm_decode_totals, fpm_decode_kv_ceiling: None, evidence: Mutex::new(TimingEvidenceSummary::default()), @@ -2068,6 +2094,7 @@ mod tests { strict_provenance: false, systems_path: None, forward_model: None, + decoder_replay: false, } } @@ -2523,6 +2550,36 @@ mod tests { assert_eq!(config.forward_model.as_deref(), Some("fpm")); } + #[test] + fn timing_policy_reaches_canonical_request() { + for replay in [false, true] { + 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", + "enable_shared_layer": false, "strict_provenance": true + })) + .unwrap(); + let request = config + .estimator_request(ForwardPassWorkerType::Aggregated) + .unwrap(); + assert_eq!(request.decoder_replay, replay); + assert_eq!(request.database_mode, crate::DatabaseMode::Silicon); + assert_eq!(request.enable_shared_layer, Some(false)); + assert!(request.strict_provenance); + let round_trip: ForwardPassPerfModelConfig = + serde_json::from_str(&serde_json::to_string(&request).unwrap()).unwrap(); + assert_eq!(round_trip, request); + } + let defaults = serde_json::from_value::(serde_json::json!({ + "model": "test-model", "backend": "sglang", "system": "test-system", "tp": 1 + })) + .unwrap(); + assert!(!defaults.decoder_replay); + assert_eq!(defaults.database_mode, crate::DatabaseMode::default()); + assert!(defaults.enable_shared_layer.is_none()); + assert!(!defaults.strict_provenance); + } + #[test] fn fpm_decode_timing_queries_exact_past_kv_total() { pyo3::prepare_freethreaded_python(); @@ -2550,6 +2607,48 @@ mod tests { assert_eq!(latency, 546_048.0); } + #[test] + fn op_level_decode_timing_converts_inclusive_mean_to_past_kv() { + pyo3::prepare_freethreaded_python(); + let engine = Python::with_gil(|py| Py::new(py, DecodeCoordinateProbe).unwrap().into_any()); + let timing = timing_model(engine, false); + + let latency = timing + .predict_decode_ms(35, 546_081, 15_602, 546_048) + .unwrap(); + + assert_eq!(latency, 546_070.0); + for inclusive_length in [1, 128, 129, 2049] { + assert_eq!( + timing + .predict_decode_ms(2, 2 * inclusive_length, inclusive_length, 8192) + .unwrap(), + (2 * inclusive_length) as f64, + ); + } + assert!(timing.predict_decode_ms(1, 0, 0, 8192).is_err()); + } + + #[test] + fn bounded_replay_rejects_heterogeneous_actual_prefill_geometry() { + pyo3::prepare_freethreaded_python(); + let engine = Python::with_gil(|py| Py::new(py, DecodeCoordinateProbe).unwrap().into_any()); + let mut timing = timing_model(engine, false); + timing.decoder_replay = true; + for geometry in [vec![], vec![(3, 1536)], vec![(129, 128), (129, 128)]] { + timing.validate_prefill_batch(&geometry).unwrap(); + } + for geometry in [vec![(127, 128), (129, 128)], vec![(3, 128), (3, 1536)]] { + assert!(timing.prefill_batch_validation_can_fail()); + let error = timing.validate_prefill_batch(&geometry).unwrap_err(); + assert!(error.to_string().contains("heterogeneous prefill")); + timing.decoder_replay = false; + assert!(!timing.prefill_batch_validation_can_fail()); + timing.validate_prefill_batch(&geometry).unwrap(); + timing.decoder_replay = true; + } + } + #[test] fn op_level_timing_exposes_typed_python_evidence() { pyo3::prepare_freethreaded_python(); diff --git a/crates/core/src/replay/agg_tests.rs b/crates/core/src/replay/agg_tests.rs index 225b7b9d2..39bf7dd10 100644 --- a/crates/core/src/replay/agg_tests.rs +++ b/crates/core/src/replay/agg_tests.rs @@ -473,8 +473,9 @@ fn g3_zero_time_restore_finishes_without_future_arrival_in_both_scopes() { ); let report = run_engine_replay(input).unwrap(); assert_eq!(report.request_counts.completed_requests, 3); - // Native decode timing has a 1ms floor even for fixed decode_ms=0. - assert_eq!(report.throughput.duration_ms, 201.0); + // Restore, prefill and fixed decode are all zero-time, so the final + // request completes at its 200ms arrival without a synthetic floor. + assert_eq!(report.throughput.duration_ms, 200.0); let g3 = report.g3_offload.unwrap(); assert!(g3.read.completed_jobs > 0); assert!(g3.lookup_hits <= g3.lookup_probes); @@ -607,17 +608,18 @@ fn g3_zero_duration_promotions_under_temporary_g2_pressure_fall_back_and_drain() rank["g3_offload"]["write_bandwidth_gbps"] = json!(0.001); rank["g3_offload"]["shared_write_bandwidth_gbps"] = json!(0.001); }); - // First prefix writes finish at2001ms. A generation-only full block + // First prefix writes finish at2000ms. A generation-only full block // evicts one G1 block without a G2 prompt store. X evicts the other G1 - // block and starts a slow write at2008ms, pinning one G2 slot. + // block and starts a slow write at2007ms, pinning one G2 slot. // The final request's zero-time reads must not // alternate forever in the other slot or fabricate a protected hit. let report = run_engine_replay(input).unwrap(); assert_eq!(report.request_counts.completed_requests, 4); - assert_eq!(report.throughput.duration_ms, 2010.0); + // The final restore and fixed decode take zero time after arrival. + assert_eq!(report.throughput.duration_ms, 2009.0); let stats = report.g3_offload.unwrap(); // Both reads start at the current lookup time2009; neither is backdated - // to X's2008 boundary, and the A/B cycle terminates at this timestamp. + // to X's2007 boundary, and the A/B cycle terminates at this timestamp. assert_eq!(stats.read.completed_jobs, 2); assert_eq!(stats.read.completed_bytes, 2_000_000); assert_eq!(stats.pending_blocks, 0); @@ -652,8 +654,8 @@ fn g3_post_lookup_touch_can_evict_a_later_prefix_block() { assert_eq!(report.request_counts.completed_requests, 4); // At100ms A is missing. Reserving it precedes the full-prefix touch, so // it evicts the older B; reserving B then evicts X. Both blocks must be - // read in one batch:2ms first-byte +2ms data, followed by1ms decode. - assert_eq!(report.throughput.duration_ms, 105.0); + // read in one batch:2ms first-byte +2ms data, followed by zero-time fixed decode. + assert_eq!(report.throughput.duration_ms, 100.0 + 2.0 + 2.0); let stats = report.g3_offload.unwrap(); assert_eq!(stats.read.completed_jobs, 1); assert_eq!(stats.read.completed_bytes, 2_000_000); @@ -688,7 +690,8 @@ fn g3_idle_gap_does_not_consume_first_byte_or_transfer_time_before_arrival() { }); let report = run_engine_replay(input).unwrap(); assert_eq!(report.request_counts.completed_requests, 4); - assert_eq!(report.throughput.duration_ms, arrival + 4.0); + // One block pays 2ms first-byte +1ms data; fixed decode adds zero. + assert_eq!(report.throughput.duration_ms, arrival + 2.0 + 1.0); let stats = report.g3_offload.unwrap(); assert_eq!(stats.read.completed_jobs, 1); assert_eq!(stats.read.completed_bytes, 1_000_000); diff --git a/crates/core/tests/perfmodel/memory_round_trip.rs b/crates/core/tests/perfmodel/memory_round_trip.rs index a804b34d4..2e9abb6e2 100644 --- a/crates/core/tests/perfmodel/memory_round_trip.rs +++ b/crates/core/tests/perfmodel/memory_round_trip.rs @@ -86,6 +86,7 @@ fn request(tolerance_fraction: Option) -> KvCacheEstimateRequest { backend: BackendKind::Trtllm, backend_version: Some("1.3.0rc10".to_string()), forward_model: None, + decoder_replay: false, kv_block_size: None, parallel: ParallelMapping { tp_size: 1, diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index a73ee5944..1ae17c3a1 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, 18); + assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 19); 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 new file mode 100644 index 000000000..e3d4d86b9 --- /dev/null +++ b/docs/deepseek-v41-storage.md @@ -0,0 +1,57 @@ +# DeepSeek-V4.1 scoring and physical KV payload + +These independently expressed analytical formulas use SGLang at immutable +[1aa0e962b206102b7c439a4a0c4981cfec6e87bc](https://github.com/sgl-project/sglang/tree/1aa0e962b206102b7c439a4a0c4981cfec6e87bc). +The upstream sources are Apache-2.0, Copyright SGLang contributors; see +`THIRD_PARTY_NOTICES.md`. No SGLang execution code is included here. + +## Index scoring precedes candidate masking + +In [`deepseek_v4_backend.py`](https://github.com/sgl-project/sglang/blob/1aa0e962b206102b7c439a4a0c4981cfec6e87bc/python/sglang/srt/layers/attention/deepseek_v4_backend.py), +`_low_ratio_index_topk_extend` calls the dense index GEMM over all compressed +keys before publishing/consuming candidate masks. `_low_ratio_index_topk_decode` +likewise calls the paged index GEMM before `two_level_decode_logits` masks its +result. Neither inspected SM100 path gathers candidates before the GEMM. +Thus `candidate_limit=16384` limits eligibility; it cannot cap scoring FLOPs, +key reads, or the materialized score array at contexts 16384/131072. Selection +still respects `index_topk`. No runtime-specific pre-GEMM optimization is assumed +for the other backends' theoretical analytical graph. + +## Persistent layout and ownership + +[`deepseek_v4_memory_pool.py`](https://github.com/sgl-project/sglang/blob/1aa0e962b206102b7c439a4a0c4981cfec6e87bc/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py) +`DeepSeekV4SingleKVPool.get_bytes_per_token/create_buffer` gives +448 FP8 NoPE bytes +128 BF16 RoPE bytes +7 scale bytes +1 scale-padding byte = 584. +Both SWA and compressed main use this layout. `get_dsv4_indexer_bytes_per_token` +forces low-ratio FP4 index storage: 128/2 + 128/32 = 68 bytes. Main FP4 rounding +before FlashMLA storage changes values; it does not allocate a second 288-byte +persistent main record. The exact layout is scoped to 512-wide, 64-RoPE, +128-index, ratio 0/1/2 V4.1; other geometries are rejected. High-ratio and unified +BF16 pool alternatives are outside this contract. + +Only four Full layers own compressed main/index pools; Reindex and Reuse share +them. One ratio-one owner and three ratio-two owners yield 652*(1+3/2)=1630 bytes +per token after the windows fill. At 131072 tokens, 40*128*584 window bytes plus +three ratio-two FP32 pair states plus compressed pools total 216662016 bytes +(206.625 MiB). Odd/even publication boundaries are preserved by forward and +inverse capacity APIs. This is physical **payload**, not allocator consumption: +576-byte page rounding, spare pages, fragmentation and other workspaces are not +included. Runtime measurements are still needed to qualify total capacity. + +## Ideal read/write traffic and compatibility + +Sparse attention reads 584 bytes per selected main row; SWA reads 584 per unique +window row with the existing ideal reuse assumption. A fused SWA norm/RoPE/store +reads the BF16 projection output and writes one physical row per token. The +projection already accounts for producing its BF16 output. Full owners publish +one compressed main/index record per completed group; no extra persistent +FP4 main record or extra fused main store is counted. Low-ratio incomplete-group +or padding writes in fallback kernels and intermediate traffic are not measured; +these remain SOL lower bounds, not an exact kernel traffic trace. + +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 +the new positional field. Model observations and calibration data are unchanged. diff --git a/docs/deepseek-v41.md b/docs/deepseek-v41.md new file mode 100644 index 000000000..b92fb0c06 --- /dev/null +++ b/docs/deepseek-v41.md @@ -0,0 +1,119 @@ + + +# DeepSeek-V4.1 Flash text modeling + +`deepseek-ai/DeepSeek-V4.1-Flash` uses the pinned checkpoint configuration +`fb2764a5cf321eaa5070ca8f9e892818f477c16d`. The model covers the 40-layer text +backbone with ordinary autoregressive decoding (`nextn=0`). The vision encoder, +three DSpark layers, encoder replay, and host-resident Engram tables require +separate execution contracts. + +The SOL graph can be constructed for SGLang, vLLM, and TensorRT-LLM. This is an +analytical model capability; it does not certify that a released version of +those runtimes loads or executes this checkpoint. Initial parallel mappings use +attention DP=1, PP=1, CP=1, and TP dividing eight output groups; attention TP must +match MoE TP times EP. The decomposed shared expert is TP-sharded. Fused shared +expert slots, replicated shared experts, MegaMoE, and alternative collective +implementations require measurements and explicit execution identities. + +## Decoder replay + +`ModelConfig(decoder_replay=False)` and `compile_engine(..., decoder_replay=False)` +select `full`: every actual extend token traverses all 40 layers. The SGLang-only +`decoder_replay=True` selects `decoder_bounded`, matching the inspected source at +[`1aa0e962b206102b7c439a4a0c4981cfec6e87bc`](https://github.com/sgl-project/sglang/tree/1aa0e962b206102b7c439a4a0c4981cfec6e87bc). +Layers 0–20 process the complete actual extend; layers 21–39 process each +request's last `min(extend_length, 128)` tokens. The absolute sequence endpoint +and shared global KV remain unchanged. A three-token extend after a long cached +prefix still has three late-layer query tokens. The bounded late-layer SWA +window starts at that extend tail. Decode always traverses all layers. + +The native mixed path scopes explicitly grouped requests before combining their +non-attention work with decode tokens. For bounded replay, the op-level FPM v1 +consumer accepts fresh prefill work only when it describes one prefill request. Its +variance field measures full prompt lengths; equal prompts can have different +cached prefixes or completed chunks, so zero variance cannot establish identical +extend tails. Multiple-prefill aggregates are therefore rejected, including +balanced batches and small total token counts. Explicit static/mixed geometry, +single-prefill telemetry, decode-only work, and Decoder OFF retain their existing +paths. Replay never removes resident weights or changes cache-capacity inventory. + +## Operators and memory + +CSA2 has two SWA layers, four Full owners (2, 8, 14, 20), four Reindex layers +(24, 28, 32, 36), and 30 Reuse layers. Full owners alone publish compressed KV; +Reindex layers rebuild selection and Reuse layers share it. Projection, packing, +index scoring/selection, and sparse-attention work remain explicit in the +analytical module. Single-pass mHC uses the system's scalar `fp32_flops` field; +BF16 compressor projections use tensor-core throughput. GB200 and GB300 provide +an explicit nominal FP32 rate. A missing rate is an error rather than a BF16 +substitution. HGX B200 and HGX B300 use 75 TFLOPS per GPU, from the +[NVIDIA HGX specification](https://www.nvidia.com/en-us/data-center/hgx/)'s +600 TFLOPS FP32 for each eight-GPU baseboard (accessed September 10, 2026). + +The indexer is replicated across attention TP: every rank owns all 32 index +heads, their projections and the full scoring/selection workload. This matches +[SGLang's pinned V4.1 indexer](https://github.com/sgl-project/sglang/blob/1aa0e962b206102b7c439a4a0c4981cfec6e87bc/python/sglang/srt/layers/attention/dsv4/dsv41_sparse.py#L203). +The [DeepSeek reference](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/fb2764a5cf321eaa5070ca8f9e892818f477c16d/inference/model.py#L507) +instead shards index heads and all-reduces their scores before selection. +Those two execution strategies must not be mixed. The TRT-LLM graph currently +uses the same replicated analytical baseline; runtime qualification is pending. + +SWA arithmetic counts causal query-key pairs. Its SOL HBM traffic counts unique +window KV rows with ideal reuse across queries: `Q + min(P, W-1)` rows per +prefill request, or `min(S, W)` for decode. Bounded decoder prefill uses `P=0` +for this local window. Actual kernel tiling and cache misses can read more; +this is a lower-bound assumption. Compressed sparse reads retain per-query +traffic because selected positions can differ. Both independent BF16 compressor +matrices are read by ratio-two Full owners. + +The V4.1 graph explicitly prices attention output reduction and omits its +redundant pre-MLP dispatch under the required DP=CP=1 topology. Shared +`MoEDispatch` behavior is unchanged: Qwen3.5's SGLang attention-DP path retains +the folded TP reduce-scatter plus DP all-gather, and its unqualified TRT-LLM +path retains the previously documented collective behavior. + +SGLang's pinned CUDA FlashMLA layout stores 584 bytes per main/SWA entry +(FP8 NoPE, BF16 RoPE, scales/padding), and its low-ratio index stores 68 bytes. +Three half-rate owners and one full-rate owner give a 1,630-byte global slope. +All 40 layers retain their 128-token window; ratio-two owners retain FP32 +pooling state. Reindex/reuse layers share the compressed pools. Per-sequence +capacity follows exact publication boundaries, and batch capacity reserves +complete window/state buffers before applying the slope. The serialized +`sglang_fp8_bf16` layout separates physical storage from attention precision. + +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. + +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 +traffic assumes uniformly distributed hash ownership; hotspots and cache reuse +need measurement. Replicated projection/gate weights, mHC weights, MoE MXFP4 +scales, dispatch workspace, expanded residual buffers, and Engram temporary +buffers are included. Backend activation coefficients remain heuristic, and +runtime allocator measurements are still required to qualify capacity. V4.1 +uses the existing MoE coefficient family (SGLang TP4: 13; vLLM/TRT-LLM TP4: 10), +with the mHC and Engram buffers added separately, rather than the dense default. + +## Result provenance + +SOL returns analytical bounds. HYBRID may combine existing measured/empirical +BF16 GEMM, MoE, and collective data with **SOL** contributions for CSA2, Engram, +single-pass mHC, and 32x32-block FP8 shared-expert projections; those new components are uncalibrated. SILICON fails for missing +V4.1 data and never substitutes V4 attention or mHC tables. EMPIRICAL similarly +requires a V4.1 anchor. Measured support is a separate dependent change. + +The independent FPM consumer uses the same model descriptor, resident inventory, +execution profile, and Rust `f64` SOL methods. Its checkpoint/profile/residency +identity and full-model measurements are owned by that dependent implementation. + +The [historical SOL validation bundle](https://github.com/ai-dynamo/aisimulate/tree/71343738bace213caa7d6966b21c287af03ac37f/data/experimental/deepseek-v41/prediction-refresh-20260914/sol) +preserves observed inputs, source identities, replay tools, and result coverage +at that commit. Its results remain specific to the recorded source and runtime +identities. diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index dc273cc39..0fdce628f 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -437,6 +437,33 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` +## SGLang DeepSeek-V4.1 serving contracts + +The V4.1 execution and performance-model logic in +`src/aiconfigurator_core/sdk/deepseek_v41.py`, `sdk/models/deepseek_v41.py` +(with the same `src/aiconfigurator_core/` prefix), and repository-root +`crates/core/src/perfmodel/operators/dsv41.rs` is informed by and modified from +SGLang's serving architecture at immutable commit +`1aa0e962b206102b7c439a4a0c4981cfec6e87bc`: + +- `python/sglang/srt/models/deepseek_v4.py` and `deepseek_v2.py` +- `python/sglang/srt/layers/engram.py` +- `python/sglang/srt/layers/attention/dsv4/compressor.py` +- `python/sglang/srt/layers/attention/dsv4/dsv41_sparse.py` +- `python/sglang/srt/layers/attention/deepseek_v4_backend.py` +- `python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py` +- `python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py` + +Source: https://github.com/sgl-project/sglang/tree/1aa0e962b206102b7c439a4a0c4981cfec6e87bc +Copyright 2023-2024 SGLang Team and SGLang contributors. Licensed under Apache-2.0; its terms are +reproduced in the repository `LICENSE`. These are analytical adaptations, +not a copy of the model execution implementation. The modified analytical +scoring/storage adaptations and their independently written regression cases +also appear in `python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py`, +Rust operator/spec unit tests, `docs/deepseek-v41.md`, and +`docs/deepseek-v41-storage.md`. They distinguish candidate masking from scoring +and physical FlashMLA cache payload from logical FP4 values. + ## DeepSeek model configuration files The following model configuration files are copied from, or formatting-only @@ -448,8 +475,16 @@ adaptations of, the named DeepSeek model repositories: | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3_config.json` | `deepseek-ai/DeepSeek-V3@e815299b0bcbac849fa540c768ef21845365c9eb` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3.2_config.json` | `deepseek-ai/DeepSeek-V3.2@c69397ecfd1fd142e90e3fbad51f4c7e40b9f3d3` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Flash_config.json` | `deepseek-ai/DeepSeek-V4-Flash@60d8d70770c6776ff598c94bb586a859a38244f1` | +| `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json` | `deepseek-ai/DeepSeek-V4.1-Flash@fb2764a5cf321eaa5070ca8f9e892818f477c16d` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Pro_config.json` | `deepseek-ai/DeepSeek-V4-Pro@b5968e9190ef611bbf34a7229255be88a0e937c1` | +The V4.1 descriptor and performance formulas in `src/aiconfigurator_core/sdk/deepseek_v41.py`, +`src/aiconfigurator_core/sdk/models/deepseek_v41.py`, and repository-root +`crates/core/src/perfmodel/operators/dsv41.rs` are AISimulate performance-model +adaptations of the architecture described by `inference/model.py` and +`DeepSeek_V41_Tech_Report.pdf` at the same V4.1 revision (modified; no model execution code). +Source: https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/tree/fb2764a5cf321eaa5070ca8f9e892818f477c16d + Upstream repositories: https://huggingface.co/deepseek-ai diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index d9fc8869d..9d85114d2 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -375,6 +375,12 @@ def _worker_performance_model_metadata( "nextn": None, **({"speculation": engine.speculation.cost_config()} if engine.speculation is not None else {}), "forward_model": worker.timing.forward_model, + **({"decoder_replay": True} if engine.decoder_replay else {}), + **{ + field: getattr(engine, field) + for field in ("database_mode", "enable_shared_layer", "strict_provenance") + if getattr(engine, field) is not None + }, }, } @@ -416,6 +422,12 @@ def _worker_engine_args( payload["speculation"] = engine.speculation.model_dump(mode="json") if engine.backend_version is not None: payload["aic_backend_version"] = engine.backend_version + if engine.decoder_replay: + payload["aic_decoder_replay"] = True + for field in ("database_mode", "enable_shared_layer", "strict_provenance"): + value = getattr(engine, field) + if value is not None: + payload[f"aic_{field}"] = value if parallel.pipeline != 1: payload["aic_pp_size"] = parallel.pipeline if parallel.moe_tensor * parallel.moe_expert > 1: @@ -475,6 +487,7 @@ def _worker_engine_args( backend=backend, backend_version=engine.backend_version, worker_type=role, + decoder_replay=engine.decoder_replay, tp=parallel.tensor, pp=parallel.pipeline, attention_dp=parallel.attention_data, diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 4407fe8c6..d6ce2a944 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -7,7 +7,7 @@ from typing import Annotated, Any, Literal -from pydantic import Field, field_validator, model_validator +from pydantic import Field, StrictBool, field_validator, model_validator from .common import Choices, IntegerRange, NumericRange, StrictModel @@ -335,6 +335,10 @@ class EnginePredictionConfig(EstimatorPolicyConfig): hardware: str backend: Backend = "vllm" backend_version: str | None = None + decoder_replay: StrictBool = False + database_mode: Literal["SILICON", "SOL", "HYBRID", "EMPIRICAL"] | None = None + enable_shared_layer: StrictBool | None = None + strict_provenance: StrictBool | None = None context_length: PositiveInt | Literal["max"] = "max" speculation: NgramSpeculationConfig | None = None workers: WorkersPredictionConfig = Field(default_factory=WorkersPredictionConfig) diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index 3a73d806b..0b235c535 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -89,6 +89,10 @@ "comm_dtype": ("comm_dtype", "aic_comm_dtype"), "systems_path": ("systems_path",), "forward_model": ("forward_model", "aic_forward_model"), + "decoder_replay": ("decoder_replay", "aic_decoder_replay"), + "database_mode": ("database_mode", "aic_database_mode"), + "enable_shared_layer": ("enable_shared_layer", "shared_layer", "aic_enable_shared_layer"), + "strict_provenance": ("strict_provenance", "aic_strict_provenance"), } _AIC_FORWARD_MODELS = frozenset({"op_level", "fpm"}) @@ -1298,6 +1302,9 @@ def _materialize_engine_role( value = rank.pop(configured[0]) if target in {"pp", "moe_tp_size", "moe_ep_size"}: value = _positive_int(value, f"engine provider {role} {target}") + elif target in {"decoder_replay", "enable_shared_layer", "strict_provenance"}: + if not isinstance(value, bool): + raise ValueError(f"engine provider {role} {target} must be a boolean") elif not isinstance(value, str) or not value: raise ValueError(f"engine provider {role} {target} must be a string") if target == "forward_model" and value not in _AIC_FORWARD_MODELS: diff --git a/python/aisimulate/src/aisimulate/sdk/deepseek_v41.py b/python/aisimulate/src/aisimulate/sdk/deepseek_v41.py new file mode 100644 index 000000000..4fa2bf2d6 --- /dev/null +++ b/python/aisimulate/src/aisimulate/sdk/deepseek_v41.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.deepseek_v41.""" + +from aisimulate.sdk._compat import alias_module as _alias_module + +_alias_module(__name__, "aisimulate_core.sdk.deepseek_v41") diff --git a/python/aisimulate/src/aisimulate/sdk/inference_session.py b/python/aisimulate/src/aisimulate/sdk/inference_session.py index 8700e5c67..68a25f6e4 100644 --- a/python/aisimulate/src/aisimulate/sdk/inference_session.py +++ b/python/aisimulate/src/aisimulate/sdk/inference_session.py @@ -1015,6 +1015,10 @@ def __init__( backend: BaseBackend, afd_config: config.AFDConfig, ) -> None: + from aisimulate_core.sdk.afd_partition import validate_afd_model_architecture + from aisimulate_core.sdk.utils import get_model_config_from_model_path + + validate_afd_model_architecture(get_model_config_from_model_path(model_path)["architecture"]) self._model_path = model_path self._a_model_config = a_model_config self._f_model_config = f_model_config diff --git a/python/aisimulate/src/aisimulate/sdk/models/deepseek_v41.py b/python/aisimulate/src/aisimulate/sdk/models/deepseek_v41.py new file mode 100644 index 000000000..5d83fd607 --- /dev/null +++ b/python/aisimulate/src/aisimulate/sdk/models/deepseek_v41.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.models.deepseek_v41.""" + +from aisimulate.sdk._compat import alias_module as _alias_module + +_alias_module(__name__, "aisimulate_core.sdk.models.deepseek_v41") diff --git a/python/aisimulate/src/aisimulate/sdk/task_v2.py b/python/aisimulate/src/aisimulate/sdk/task_v2.py index 72cf86155..00ad270b0 100644 --- a/python/aisimulate/src/aisimulate/sdk/task_v2.py +++ b/python/aisimulate/src/aisimulate/sdk/task_v2.py @@ -885,6 +885,10 @@ def __post_init__(self) -> None: self.nextn, self.nextn_accepted = normalize_speculative_decoding(self.nextn, self.nextn_accepted) self._validate_deepseek_v4_hardware() self._resolve_model_identity() + if self.serving_mode == "afd": + from aisimulate_core.sdk.afd_partition import validate_afd_model_architecture + + validate_afd_model_architecture(self._architecture) if self.nextn == "auto": raise ValueError("nextn='auto' requires a model path to resolve num_nextn_predict_layers.") self._resolve_backend_version() @@ -2354,6 +2358,9 @@ def _validate_disagg(self) -> None: # per-role DB check in _validate_database_quant_modes governs support. def _validate_afd(self) -> None: + from aisimulate_core.sdk.afd_partition import validate_afd_model_architecture + + validate_afd_model_architecture(self._architecture) if not self.model_path: raise ValueError("afd mode requires model_path") if not self.system_name: diff --git a/python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_README.md b/python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_README.md new file mode 100644 index 000000000..e9adb4beb --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_README.md @@ -0,0 +1,11 @@ +# DeepSeek-V4.1 configuration provenance + +`deepseek-ai--DeepSeek-V4.1-Flash_config.json` is copied without modification from +https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/fb2764a5cf321eaa5070ca8f9e892818f477c16d/config.json. + +Copyright (c) 2023 DeepSeek. MIT license; the complete upstream license is +reproduced in the packaged `THIRD_PARTY_NOTICES.md`. + +The text AR performance model uses the 40 backbone layers. The supplied +vision and DSpark configuration remains in this source snapshot but is not +enabled by the text AR model. diff --git a/python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json b/python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json new file mode 100644 index 000000000..09917a913 --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json @@ -0,0 +1,170 @@ +{ + "architectures": [ + "DeepseekV41ForCausalLM" + ], + "model_type": "deepseek_v41", + "dtype": "bfloat16", + "transformers_version": "5.6.0", + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "image_token_id": 129264, + "quantization_config": { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [ + 32, + 32 + ], + "scale_fmt": "ue8m0", + "expert_dtype": "fp4" + }, + "text_config": { + "model_type": "deepseek_v41_text", + "vocab_size": 129280, + "hidden_size": 5120, + "moe_intermediate_size": 2304, + "num_hidden_layers": 40, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "qk_rope_head_dim": 64, + "q_lora_rank": 1280, + "o_lora_rank": 1024, + "o_groups": 8, + "hidden_act": "silu", + "swiglu_limit": 10.0, + "rms_norm_eps": 1e-20, + "attention_bias": false, + "attention_dropout": 0.0, + "initializer_range": 0.02, + "use_cache": true, + "tie_word_embeddings": false, + "max_position_embeddings": 1048576, + "rope_theta": 10000, + "rope_scaling": { + "rope_type": "yarn", + "factor": 16, + "beta_fast": 32, + "beta_slow": 1, + "original_max_position_embeddings": 65536 + }, + "n_routed_experts": 384, + "n_shared_experts": 1, + "num_experts_per_tok": 6, + "scoring_func": "sqrtsoftplus", + "topk_method": "noaux_tc", + "norm_topk_prob": true, + "routed_scaling_factor": 1.5, + "sliding_window": 128, + "compress_ratios": [ + 0, + 0, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0 + ], + "compress_rope_theta": 160000, + "kv_source_layer_ids": [ + 2, + 8, + 14, + 20 + ], + "index_source_layer_ids": [ + 2, + 8, + 14, + 20, + 24, + 28, + 32, + 36 + ], + "index_n_heads": 32, + "index_head_dim": 128, + "index_topk": 512, + "candidate_source_layer_id": 20, + "candidate_topk_blocks": 2048, + "candidate_block_size": 8, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "hc_eps": 1e-06, + "engram_layer_ids": [ + 1, + 14 + ], + "engram_num_embeddings": [ + 384006168, + 384016682 + ], + "engram_max_ngram_size": 4, + "engram_vocab_size": 16000000, + "engram_n_heads": 8, + "engram_head_dim": 256, + "engram_pad_token_id": 2, + "engram_compressed_vocab_size": 99092, + "num_nextn_predict_layers": 3, + "dspark_block_size": 5, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [ + 37, + 38, + 39 + ], + "dspark_markov_rank": 256, + "dspark_n_routed_experts": 128, + "dspark_num_experts_per_tok": 3 + }, + "vision_config": { + "model_type": "deepseek_v41_vision", + "num_hidden_layers": 32, + "hidden_size": 1024, + "num_attention_heads": 16, + "intermediate_size": 2816, + "patch_size": 14, + "rope_theta": 10000, + "downsample_ratio": 3, + "max_image_tokens": 1024, + "min_pixels": 295936, + "max_wh_ratio": null + } +} diff --git a/python/aisimulate/src/aisimulate_core/sdk/afd_partition.py b/python/aisimulate/src/aisimulate_core/sdk/afd_partition.py index 144b9bcd3..2103026fd 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/afd_partition.py +++ b/python/aisimulate/src/aisimulate_core/sdk/afd_partition.py @@ -27,6 +27,16 @@ class AFDPartitionError(ValueError): """Raised when an op sequence cannot be safely partitioned for AFD.""" +def validate_afd_model_architecture(architecture: str) -> None: + """Reject architectures whose stages/residency cannot be split across A/F pools.""" + if architecture == "DeepseekV41ForCausalLM": + raise NotImplementedError( + "AFD does not support DeepSeek-V4.1: mixed attention/FFN stages and " + "their weight, scale, Engram, and cache residency are not partitioned. " + "Use agg, disagg, or static estimation." + ) + + @dataclass class AFDOpsPartition: """Partitioned AFD operation lists. diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py index 296724df9..6eff2a9b0 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py @@ -112,6 +112,7 @@ class BaseBackend: "DEEPSEEK", "DEEPSEEKV32", "DEEPSEEKV4", + "DEEPSEEKV41", "KIMIK25", ) @@ -2224,15 +2225,10 @@ def _get_memory_usage( isinstance(scheme, SpecSchemeBase) and not isinstance(scheme, MTPScheme) and type(scheme) is not NullScheme ) - weights = 0.0 - for op in model.context_ops: - # Materialized draft ops are excluded here: the scheme's own - # byte-exact accounting below is the single source of truth - # (the op-list subset under-counts aliased/owned embed and - # sampling heads unevenly across schemes). - if has_draft_scheme and op._name.startswith("draft_"): - continue - weights += op.get_weights() + # Target residency is independent of the speculative scheme and can + # include scales or skipped-layer weights outside the executed ops. + # Its inventory excludes scheme-owned draft weights, added below. + weights = model.get_resident_weights_bytes() # count weights on a single GPU weights /= model.config.pp_size if has_draft_scheme: @@ -2265,6 +2261,7 @@ def _get_memory_usage( * 4 ) + activations += model.get_additional_activation_bytes(num_tokens) activations = max(activations, self.MIN_ACTIVATION_BYTES) # MTP correction: speculative decoding verifies nextn+1 tokens per decode step, diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/sglang_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/sglang_backend.py index ae8755bad..3b28100b1 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/sglang_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/sglang_backend.py @@ -79,6 +79,7 @@ class SGLANGBackend(BaseBackend): "DEEPSEEK": {1: 28, 2: 17, 4: 13, 8: 13}, "DEEPSEEKV32": {1: 28, 2: 17, 4: 13, 8: 13}, "DEEPSEEKV4": {1: 28, 2: 17, 4: 13, 8: 13}, + "DEEPSEEKV41": {1: 28, 2: 17, 4: 13, 8: 13}, "KIMIK25": {1: 28, 2: 17, 4: 13, 8: 13}, "KIMIK3": {1: 28, 2: 17, 4: 13, 8: 13}, "default": {1: 13, 2: 8, 4: 6.5, 8: 6.5}, diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/trtllm_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/trtllm_backend.py index b7f0fab99..6e70897ee 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/trtllm_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/trtllm_backend.py @@ -49,6 +49,7 @@ class TRTLLMBackend(BaseBackend): "DEEPSEEK": {1: 22, 2: 13, 4: 10, 8: 10}, "DEEPSEEKV32": {1: 22, 2: 13, 4: 10, 8: 10}, "DEEPSEEKV4": {1: 22, 2: 13, 4: 10, 8: 10}, + "DEEPSEEKV41": {1: 22, 2: 13, 4: 10, 8: 10}, "KIMIK25": {1: 22, 2: 13, 4: 10, 8: 10}, # 4+6/TP, fp8 will have relatively low act, but ignore here. need more experiments "default": {1: 10, 2: 6, 4: 5, 8: 5}, diff --git a/python/aisimulate/src/aisimulate_core/sdk/common.py b/python/aisimulate/src/aisimulate_core/sdk/common.py index c19f26685..5c27ad60c 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/common.py +++ b/python/aisimulate/src/aisimulate_core/sdk/common.py @@ -12,6 +12,7 @@ from packaging.version import InvalidVersion, Version +from aisimulate_core.sdk.deepseek_v41 import DeepSeekV41Config as DeepSeekV41Config from aisimulate_core.sdk.errors import MissingSystemFlopsError @@ -656,6 +657,7 @@ def get_default_models() -> set[str]: and this set via get_default_models() """ DefaultHFModels = { + "deepseek-ai/DeepSeek-V4.1-Flash", # Llama 3.1 Models "meta-llama/Meta-Llama-3.1-8B", "meta-llama/Meta-Llama-3.1-70B", @@ -812,6 +814,7 @@ def get_default_models() -> set[str]: "DEEPSEEK", "DEEPSEEKV32", "DEEPSEEKV4", + "DEEPSEEKV41", "KIMIK25", "KIMIK3", "NEMOTRONNAS", @@ -837,6 +840,7 @@ def get_default_models() -> set[str]: "DeepseekV32ForCausalLM": "DEEPSEEKV32", "GlmMoeDsaForCausalLM": "DEEPSEEKV32", "DeepseekV4ForCausalLM": "DEEPSEEKV4", + "DeepseekV41ForCausalLM": "DEEPSEEKV41", "KimiK25ForConditionalGeneration": "KIMIK25", "KimiK3ForConditionalGeneration": "KIMIK3", "NemotronForCausalLM": "NEMOTRONNAS", @@ -871,6 +875,7 @@ def get_default_models() -> set[str]: # Multimodal architectures whose LLM config lives under a nested key (e.g. "text_config"). # _parse_hf_config_json will flatten these before parsing. MULTIMODAL_TEXT_CONFIG_KEY = { + "DeepseekV41ForCausalLM": "text_config", "KimiK25ForConditionalGeneration": "text_config", "KimiK3ForConditionalGeneration": "text_config", # Step-3.7/3.5-Flash ship a vision tower and nest the whole decoder under diff --git a/python/aisimulate/src/aisimulate_core/sdk/config.py b/python/aisimulate/src/aisimulate_core/sdk/config.py index 67eefd6e4..72d01e9c6 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config.py @@ -130,6 +130,8 @@ class ModelConfig: # Internal system identity used by phase/quantization-specific communication # dtype selection. It travels with ModelConfig through sweep replacements. system: str | None = None + # DeepSeek-V4.1 text AR: use the backend-verified bounded decoder replay profile. + decoder_replay: bool = False 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/deepseek_v41.py b/python/aisimulate/src/aisimulate_core/sdk/deepseek_v41.py new file mode 100644 index 000000000..9815dd119 --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/sdk/deepseek_v41.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared DeepSeek-V4.1 text architecture and execution contract. + +Structural source: deepseek-ai/DeepSeek-V4.1-Flash at +fb2764a5cf321eaa5070ca8f9e892818f477c16d (config.json and inference/model.py; +MIT, Copyright (c) 2023 DeepSeek). See THIRD_PARTY_NOTICES.md. +This module describes performance-model inputs; it does not execute model code. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from enum import StrEnum + +MODEL_PATH = "deepseek-ai/DeepSeek-V4.1-Flash" +MODEL_REVISION = "fb2764a5cf321eaa5070ca8f9e892818f477c16d" + + +class DeepSeekV41ExecutionProfile(StrEnum): + FULL = "full" + DECODER_BOUNDED = "decoder_bounded" + + +class DeepSeekV41KVCacheLayout(StrEnum): + """Storage contracts, separate from the precision of the logical values. + + The packed layout is a theoretical inventory, not a qualified runtime + layout for vLLM or TRT-LLM. The SGLang layout is pinned in the source + references and THIRD_PARTY_NOTICES.md. + """ + + LOGICAL_FP4 = "logical_fp4" + SGLANG_FP8_BF16 = "sglang_fp8_bf16" + + +def resolve_kv_cache_layout(backend_name: str) -> DeepSeekV41KVCacheLayout: + """Only the pinned SGLang backend has a physical storage contract here.""" + if backend_name == "sglang": + return DeepSeekV41KVCacheLayout.SGLANG_FP8_BF16 + if backend_name in ("vllm", "trtllm"): + return DeepSeekV41KVCacheLayout.LOGICAL_FP4 + raise ValueError(f"unknown DeepSeek-V4.1 KV backend: {backend_name}") + + +def resolve_execution_profile(decoder_replay: bool, backend_name: str) -> DeepSeekV41ExecutionProfile: + """Only SGLang has a verified bounded-decoder execution contract.""" + if not isinstance(decoder_replay, bool): + raise ValueError("decoder_replay must be a boolean") + if decoder_replay and backend_name != "sglang": + raise NotImplementedError(f"DeepSeek-V4.1 decoder replay is not verified for {backend_name}") + return DeepSeekV41ExecutionProfile.DECODER_BOUNDED if decoder_replay else DeepSeekV41ExecutionProfile.FULL + + +@dataclass(frozen=True) +class DeepSeekV41Config: + """Backbone geometry; the separate three DSpark layers are not AR layers.""" + + num_hidden_layers: int + hidden_size: int + num_attention_heads: int + head_dim: int + qk_rope_head_dim: int + q_lora_rank: int + o_lora_rank: int + o_groups: int + moe_intermediate_size: int + n_routed_experts: int + n_shared_experts: int + num_experts_per_tok: int + sliding_window: int + compress_ratios: tuple[int, ...] + kv_source_layer_ids: tuple[int, ...] + index_source_layer_ids: tuple[int, ...] + index_n_heads: int + index_head_dim: int + index_topk: int + candidate_source_layer_id: int + candidate_topk_blocks: int + candidate_block_size: int + hc_mult: int + hc_sinkhorn_iters: int + hc_eps: float + engram_layer_ids: tuple[int, ...] + engram_num_embeddings: tuple[int, ...] + engram_max_ngram_size: int + engram_n_heads: int + engram_head_dim: int + + @classmethod + def from_text_config(cls, text_config: dict) -> DeepSeekV41Config: + values = {field.name: text_config[field.name] for field in fields(cls)} + for name in ( + "compress_ratios", + "kv_source_layer_ids", + "index_source_layer_ids", + "engram_layer_ids", + "engram_num_embeddings", + ): + values[name] = tuple(values[name]) + values["compress_ratios"] = values["compress_ratios"][: values["num_hidden_layers"]] + result = cls(**values) + result.validate() + return result + + def validate(self) -> None: + 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") + for ids in (self.kv_source_layer_ids, self.index_source_layer_ids, self.engram_layer_ids): + if tuple(sorted(set(ids))) != ids or any(i < 0 or i >= layers for i in ids): + raise ValueError("DeepSeek-V4.1 source/layer IDs must be unique, ascending backbone indices") + if not set(self.kv_source_layer_ids) <= set(self.index_source_layer_ids): + raise ValueError("DeepSeek-V4.1 KV owners must also own an indexer") + if any(self.compress_ratios[i] <= 0 for i in self.kv_source_layer_ids): + raise ValueError("DeepSeek-V4.1 KV owners require a positive compression ratio") + if self.candidate_source_layer_id not in self.kv_source_layer_ids: + raise ValueError("DeepSeek-V4.1 candidate source must own compressed KV") + if len(self.engram_layer_ids) != len(self.engram_num_embeddings): + raise ValueError("DeepSeek-V4.1 Engram layer/table counts differ") + source = None + for layer, ratio in enumerate(self.compress_ratios): + if layer in self.kv_source_layer_ids: + source = layer + if ratio and (source is None or self.compress_ratios[source] != ratio): + raise ValueError(f"DeepSeek-V4.1 layer {layer} has no compatible preceding KV owner") + + @property + def decoder_start_layer(self) -> int: + """The first layer eligible for bounded replay; layer 20 builds global KV.""" + return max(self.kv_source_layer_ids) + 1 + + def layer_role(self, layer: int) -> str: + if not self.compress_ratios[layer]: + return "swa" + if layer in self.kv_source_layer_ids: + return "full" + return "reindex" if layer in self.index_source_layer_ids else "reuse" + + @property + def compressed_entry_bytes(self) -> float: + """Theoretical packed-FP4 value and scale bytes, not SGLang storage.""" + return self.head_dim / 2 + (self.head_dim + 15) // 16 + + @property + def index_entry_bytes(self) -> float: + return self.index_head_dim / 2 + (self.index_head_dim + 31) // 32 + + def kv_entry_bytes(self, layout: DeepSeekV41KVCacheLayout) -> tuple[float, float, float]: + """Window, compressed-main and index payload bytes, excluding page padding. + + SGLang 1aa0e962, deepseek_v4_memory_pool.py:123-145 and + deepseek_v4_backend.py:3012-3016: FP4-rounded main values are stored + as 448 FP8 bytes + 64 BF16 RoPE values + 7 scales + 1 scale pad. + Low-ratio (1/2) index keys stay packed FP4: 64 payload + 4 scales. + """ + if layout == DeepSeekV41KVCacheLayout.SGLANG_FP8_BF16: + if (self.head_dim, self.qk_rope_head_dim, self.index_head_dim) != (512, 64, 128): + raise ValueError("pinned SGLang V4.1 KV storage requires head_dim=512, RoPE=64, index=128") + return 584.0, 584.0, 68.0 + if layout == DeepSeekV41KVCacheLayout.LOGICAL_FP4: + return float(self.head_dim), self.compressed_entry_bytes, self.index_entry_bytes + raise ValueError(f"unknown DeepSeek-V4.1 KV storage layout: {layout}") + + def kvcache_bytes( + self, + sequence_length: int, + layout: DeepSeekV41KVCacheLayout = DeepSeekV41KVCacheLayout.LOGICAL_FP4, + ) -> float: + """Unique pools and FP32 state; default is the theoretical packed inventory.""" + window_entry, main_entry, index_entry = self.kv_entry_bytes(layout) + length = max(0, sequence_length) + if not length: + return 0.0 + window = self.num_hidden_layers * min(length, self.sliding_window) * window_entry + global_cache = 0.0 + state = 0.0 + for owner in self.kv_source_layer_ids: + ratio = self.compress_ratios[owner] + global_cache += (length // ratio) * (main_entry + index_entry) + if ratio > 1: + state += 2 * ratio * self.head_dim * 4 + return window + global_cache + state + + def engram_table_bytes(self, tp_size: int) -> float: + """GPU-resident row-sharded FP8 tables including one UE8M0 scale per 32.""" + row_bytes = self.engram_head_dim + (self.engram_head_dim + 31) // 32 + return float(sum((n + tp_size - 1) // tp_size for n in self.engram_num_embeddings) * row_bytes) + + +@dataclass(frozen=True) +class V41RequestWorkload: + """One actual forward's uncached tokens and original absolute KV position.""" + + query_tokens: int + prefix_tokens: int = 0 + + def __post_init__(self) -> None: + if self.query_tokens < 0 or self.prefix_tokens < 0: + raise ValueError("DeepSeek-V4.1 token counts must be nonnegative") + + +@dataclass(frozen=True) +class V41StageWorkload: + query_tokens: int + prefix_tokens: int + + +def stage_workloads( + descriptor: DeepSeekV41Config, + profile: DeepSeekV41ExecutionProfile | str, + requests: tuple[V41RequestWorkload, ...], + layer: int, +) -> tuple[V41StageWorkload, ...]: + """Bound per request, retaining its absolute attention position. + + Inputs describe an actual extend call, not the entire logical prompt. + Cached prefixes are not new inputs and do not fill a short extend's tail. + """ + profile = DeepSeekV41ExecutionProfile(profile) + bounded = profile == DeepSeekV41ExecutionProfile.DECODER_BOUNDED and layer >= descriptor.decoder_start_layer + return tuple( + V41StageWorkload( + min(r.query_tokens, descriptor.sliding_window) if bounded else r.query_tokens, + r.prefix_tokens + (max(r.query_tokens - descriptor.sliding_window, 0) if bounded else 0), + ) + for r in requests + ) diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index 854b6a3aa..ff3c6e7e3 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -47,7 +47,7 @@ from aisimulate_core.sdk.config_builders import apply_nextn, build_model_config from aisimulate_core.sdk.errors import InvalidEngineConfigurationError as InvalidEngineConfigurationError from aisimulate_core.sdk.models import get_model -from aisimulate_core.sdk.models.helpers import resolve_sglang_mla_compute +from aisimulate_core.sdk.models.helpers import resolve_dsv4_moe_arch, resolve_sglang_mla_compute from aisimulate_core.sdk.operations import FPMForwardOp from aisimulate_core.sdk.operations.base import Operation from aisimulate_core.sdk.perf_database import load_system_spec @@ -319,6 +319,7 @@ def _engine_config_dict( # Rust side reloads the perf database from this string verbatim. "backend_version": _literal_backend_version(system, backend, backend_version, systems_path, database), "kv_block_size": kv_block_size, + "decoder_replay": bool(getattr(cfg, "decoder_replay", False)), # ParallelMapping (flattened) "tp_size": int(cfg.tp_size or 1), "pp_size": int(cfg.pp_size or 1), @@ -422,6 +423,7 @@ def compile_engine( kv_block_size: int | None = None, systems_path: str | None = None, forward_model: str | None = None, + decoder_replay: bool = False, database_mode: str | None = None, shared_layer: bool | None = None, transfer_policy: str | list[str] | None = None, @@ -463,6 +465,8 @@ def compile_engine( apply_nextn(model_config, nextn) except (ValueError, TypeError, KeyError) as exc: raise InvalidEngineConfigurationError(str(exc)) from exc + model_config.decoder_replay = decoder_replay + resolve_dsv4_moe_arch(model_config, model_path, system_name=system, backend_name=backend) # Slot policy FIRST, tolerance second: resolve the requested version to a # literal (raising on unlisted versions / unpopulated aliases) before the diff --git a/python/aisimulate/src/aisimulate_core/sdk/memory.py b/python/aisimulate/src/aisimulate_core/sdk/memory.py index 819fb02b9..291e7a726 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory.py @@ -388,7 +388,7 @@ def from_request( "kv_size_per_token_bytes": float(model.get_kvcache_bytes_per_sequence(1)), "gpu_memory_capacity_bytes": float(database.system_spec["gpu"]["mem_capacity"]), # Model's byte-budget -> token-count inverse (KV-curve aware). - "tokens_from_kv_bytes": model.get_kvcache_max_tokens, + "tokens_from_kv_bytes": lambda budget: model.get_kvcache_batch_capacity(budget, max_batch_size), } ) diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py index 9905d3243..b5bff6580 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py @@ -68,7 +68,13 @@ def _apply_forward_model_fpm(model: BaseModel) -> 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.""" - from aisimulate_core.sdk.operations.fpm_forward import FPMForwardOp + from aisimulate_core.sdk.operations.fpm_forward import _CELL_MATCH_COLUMNS, FPMForwardOp + + if getattr(model.config, "decoder_replay", False) and "execution_profile" not in _CELL_MATCH_COLUMNS: + # Rewriting the staged graph is safe only when measured curves are + # keyed by execution profile. Legacy FPM cells cannot distinguish the + # bounded decoder tail from a full forward at the same coordinates. + raise NotImplementedError("decoder_replay requires FPM tables with execution_profile identity") if model.encoder_ops: raise NotImplementedError( @@ -104,7 +110,7 @@ def _apply_forward_model_fpm(model: BaseModel) -> BaseModel: generation_ops = [op for op in model.generation_ops if not op._name.startswith("draft_")] draft_context_ops = [op for op in model.context_ops if op._name.startswith("draft_")] draft_generation_ops = [op for op in model.generation_ops if op._name.startswith("draft_")] - weight_bytes = float(sum(op.get_weights() for op in context_ops)) + weight_bytes = model.get_resident_weights_bytes() prefill_op = FPMForwardOp("prefill", model.config, model.model_path, sol_ops=context_ops, weight_bytes=weight_bytes) decode_op = FPMForwardOp( "decode", model.config, model.model_path, sol_ops=generation_ops, weight_bytes=weight_bytes diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/base.py b/python/aisimulate/src/aisimulate_core/sdk/models/base.py index 6381211f1..b0b21d632 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/base.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/base.py @@ -162,6 +162,20 @@ def verify_width(self) -> int: def activation_hidden_size(self) -> int: return self._num_heads * self._head_size + def get_additional_activation_bytes(self, num_tokens: int) -> float: + """Architecture-specific buffers beyond the backend's generic workspace.""" + return 0.0 + + def get_resident_weights_bytes(self) -> float: + """Resident target weights per TP/EP rank, before PP division. + + Models with phase-dependent execution can override this inventory; + skipping token work must never remove resident decoder weights. + Scheme-owned draft weights are accounted for separately, including + any draft weights absent from the materialized context-op subset. + """ + return float(sum(op.get_weights() for op in self.context_ops if not op._name.startswith("draft_"))) + # ------------------------------------------------------------------ # Context parallelism (CP) declaration + comm factory (1145-style). # GLM-5 DSA does NOT use these -- it handles CP inside ContextDSAModule. @@ -288,6 +302,10 @@ def get_kvcache_max_tokens(self, kv_budget_bytes: float) -> int: return 0 return int(budget // per_token) + def get_kvcache_batch_capacity(self, kv_budget_bytes: float, max_batch_size: int) -> int: + """Total-token capacity; models with per-request state may reserve it here.""" + return self.get_kvcache_max_tokens(kv_budget_bytes) + def _binary_search_kvcache_max_tokens(self, kv_budget_bytes: float) -> int: """Monotonic-search inverse of :meth:`get_kvcache_bytes_per_sequence`. diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py new file mode 100644 index 000000000..d0e75babe --- /dev/null +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DeepSeek-V4.1 text AR, with independently scoped prefill decoder stages. + +Architecture source: deepseek-ai/DeepSeek-V4.1-Flash at +fb2764a5cf321eaa5070ca8f9e892818f477c16d. See sdk/deepseek_v41.py and +THIRD_PARTY_NOTICES.md. Vision and DSpark are deliberately separate capabilities. +""" + +from __future__ import annotations + +import json + +import aisimulate_core._native as core +import aisimulate_core.sdk.operations as ops +from aisimulate_core.sdk import common +from aisimulate_core.sdk.deepseek_v41 import ( + DeepSeekV41Config, + resolve_execution_profile, + resolve_kv_cache_layout, +) +from aisimulate_core.sdk.models.base import BaseModel, register_model + + +def _native(kind: str, **values): + return core.op_from_spec_json(json.dumps({kind: values})) + + +def _spec(op): + return json.loads(op._spec_json()) + + +@register_model("DEEPSEEKV41") +class DeepSeekV41Model(BaseModel): + """40-layer text backbone; resident weights do not depend on replay mode.""" + + @classmethod + def create(cls, info: dict, model_config, backend_name: str) -> BaseModel: + return cls(info, model_config, backend_name) + + @property + def activation_hidden_size(self) -> int: + return self._hidden_size + + def __init__(self, info: dict, model_config, backend_name: str): + super().__init__( + info["model_path"], + info["model_family"], + info["architecture"], + info["layers"], + info["n"], + info["n_kv"], + info["d"], + info["hidden_size"], + info["inter_size"], + info["vocab"], + info["context"], + model_config, + info["extra_params"], + ) + d = self.extra_params + if not isinstance(d, DeepSeekV41Config): + raise TypeError("DeepSeekV41Model requires the shared DeepSeekV41Config descriptor") + if self._nextn: + raise NotImplementedError( + "DeepSeek-V4.1 text AR requires nextn=0; DSpark has a separate execution contract" + ) + if self._num_layers != d.num_hidden_layers: + raise ValueError("DeepSeek-V4.1 layer overrides cannot preserve cross-layer KV ownership") + if model_config.moe_backend == "megamoe": + raise NotImplementedError( + "DeepSeek-V4.1 SOL/Hybrid uses decomposed MoE; MegaMoE requires separate measurement" + ) + if model_config.attention_dp_size != 1 or model_config.pp_size != 1 or model_config.cp_size != 1: + raise NotImplementedError( + "DeepSeek-V4.1 text baseline requires attention_dp_size=1, pp_size=1 and cp_size=1; " + "DP Engram collectives and PP cache ownership need separate contracts" + ) + tp = model_config.tp_size + mtp, ep = model_config.resolve_moe_parallelism() + if tp * model_config.attention_dp_size != mtp * ep: + raise ValueError("attention TP * DP must equal MoE TP * EP for DeepSeek-V4.1") + if d.n_routed_experts % ep or d.index_n_heads % tp or d.o_groups % tp: + raise ValueError("DeepSeek-V4.1 EP must divide experts; TP must divide index heads and output groups") + self._topk, self._num_experts, self._moe_inter_size = info["topk"], info["num_experts"], info["moe_inter_size"] + self.raw_config = info["raw_config"] + self.text_config = self.raw_config["text_config"] + self.execution_profile = resolve_execution_profile(model_config.decoder_replay, backend_name).value + self.kv_cache_layout = resolve_kv_cache_layout(backend_name) + # Validate physical dimensions before constructing any native operators. + d.kv_entry_bytes(self.kv_cache_layout) + self.engram_residency = "hbm_tp_sharded" + h = self._hidden_size + distribution = ( + "power_law_1.01" + if model_config.workload_distribution == "power_law" + else model_config.workload_distribution + ) + + def attention(layer: int, context: bool): + return _native( + "Dsv41Attention", + name="context_attention" if context else "generation_attention", + is_context=context, + role=d.layer_role(layer), + compress_ratio=d.compress_ratios[layer], + hidden_size=h, + num_heads=d.num_attention_heads // tp, + head_dim=d.head_dim, + q_lora_rank=d.q_lora_rank, + o_lora_rank=d.o_lora_rank, + o_groups=max(1, d.o_groups // tp), + # SGLang's pinned V4.1 indexer replicates both projections + # and all index heads across TP (dsv41_sparse.py:203-224). + # The DeepSeek reference instead shards heads and reduces + # scores; that different execution contract is not used here. + index_n_heads=d.index_n_heads, + index_head_dim=d.index_head_dim, + index_topk=d.index_topk, + window_size=d.sliding_window, + candidate_limit=d.candidate_topk_blocks * d.candidate_block_size + if layer > d.candidate_source_layer_id + else 0, + is_candidate_source=layer == d.candidate_source_layer_id, + bounded_prefill=context and model_config.decoder_replay and layer >= d.decoder_start_layer, + gemm_quant_mode=model_config.gemm_quant_mode.name, + fmha_quant_mode=model_config.fmha_quant_mode.name, + kv_cache_layout=self.kv_cache_layout.value, + ) + + def stage(layer: int, context: bool): + phase = "context" if context else "generation" + children = [] + if layer in d.engram_layer_ids: + engram_index = d.engram_layer_ids.index(layer) + children.append( + _native( + "Dsv41Engram", + name=f"{phase}_engram", + num_embeddings=d.engram_num_embeddings[engram_index], + head_dim=d.engram_head_dim, + hash_columns=(d.engram_max_ngram_size - 1) * d.engram_n_heads, + hidden_size=h, + hc_mult=d.hc_mult, + tp_size=tp, + ) + ) + children.append( + ops.NCCL( + f"{phase}_engram_allreduce", + 1, + "all_reduce", + (d.engram_max_ngram_size - 1) * d.engram_n_heads * d.engram_head_dim, + tp, + common.CommQuantMode.half, + ) + ) + children.extend( + [ + _native( + "Dsv41Mhc", + name=f"{phase}_mhc", + hidden_size=h, + hc_mult=d.hc_mult, + sinkhorn_iters=d.hc_sinkhorn_iters, + ), + ops.ElementWise(f"{phase}_attn_norm", 1, h, h, 0.8), + attention(layer, context), + ops.NCCL(f"{phase}_attention_allreduce", 1, "all_reduce", h, tp, common.CommQuantMode.half), + ops.ElementWise(f"{phase}_ffn_norm", 1, h, h, 0.8), + ] + ) + local_inter = self._moe_inter_size * d.n_shared_experts // tp + shared = [ + _native( + "Dsv41Linear", + name=f"{phase}_shared_gate_up_gemm", + n=2 * local_inter, + k=h, + quant_mode=model_config.gemm_quant_mode.name, + ), + ops.ElementWise(f"{phase}_shared_act_gate", 1, 2 * local_inter, local_inter, 0.8), + _native( + "Dsv41Linear", + name=f"{phase}_shared_ffn2_gemm", + n=h, + k=local_inter, + quant_mode=model_config.gemm_quant_mode.name, + ), + ] + # DP=CP=1 is required above and the attention output was already + # reduced. Omit this model's redundant pre-MLP dispatch; retain + # the legacy backend-specific semantics for all other models. + routed = [ + ops.GEMM(f"{phase}_router_gemm", 1, self._num_experts, h, common.GEMMQuantMode.bfloat16), + ops.MoE( + f"{phase}_moe", + 1, + h, + self._moe_inter_size, + self._topk, + self._num_experts, + mtp, + ep, + model_config.moe_quant_mode, + distribution, + model_config.attention_dp_size, + ), + ] + # The post-expert reduction consumes both routed and shared + # partials, and therefore follows the optional compute overlap. + combine = ops.MoEDispatch( + f"{phase}_moe_post_dispatch", + 1, + h, + self._topk, + self._num_experts, + mtp, + ep, + model_config.attention_dp_size, + False, + quant_mode=model_config.moe_quant_mode, + backend=backend_name, + is_context=context, + attn_ar_modeled=True, + ) + # SGLang's qualified TP eager path executes forward_normal on + # one stream. Dual-stream shared/routed work requires capture or + # graph/SBO dispatch (sglang@1aa0e962 deepseek_v2.py:885-960, + # 1107-1126,1191-1222). Other backend/EP modes remain assumptions. + if context or (ep == 1 and backend_name == "sglang"): + children.extend(shared + routed) + else: + children.append(ops.OverlapOp(f"{phase}_moe_overlap", group_a=routed, group_b=shared)) + children.append(combine) + return _native( + "Dsv41Stage", + name=f"{phase}_v41_layer_{layer}", + is_context=context, + decoder_replay=model_config.decoder_replay, + bounded=layer >= d.decoder_start_layer, + window_size=d.sliding_window, + children=[_spec(op) for op in children], + ) + + for context in (True, False): + phase = "context" if context else "generation" + target = self.context_ops if context else self.generation_ops + target.append(ops.Embedding(f"{phase}_embedding", 1, self._vocab_size // tp, h, 0.3)) + target.append(ops.NCCL(f"{phase}_embedding_allreduce", 1, "all_reduce", h, tp, common.CommQuantMode.half)) + target.extend(stage(i, context) for i in range(d.num_hidden_layers)) + target.extend( + [ + ops.GEMM(f"{phase}_logits_gemm", 1, self._vocab_size // tp, h, common.GEMMQuantMode.bfloat16), + ops.P2P(f"{phase}_p2p", model_config.pp_size - 1, h, model_config.pp_size), + ] + ) + self._resident_weight_bytes = float(sum(op.get_weights() for op in self.context_ops)) + # Native MXFP4 scales are one UE8M0 byte per 32 weights. The generic + # legacy MoE op stores only packed-weight bytes, so preserve the scale + # inventory here independently of either phase's execution count. + if model_config.moe_quant_mode.name.startswith("w4") and "mxfp4" in model_config.moe_quant_mode.name: + expert_elements = d.num_hidden_layers * 3 * h * self._moe_inter_size * self._num_experts / (mtp * ep) + self._resident_weight_bytes += expert_elements / 32 + + def get_additional_activation_bytes(self, num_tokens: int) -> float: + d = self.extra_params + # Two expanded residual buffers plus the largest Engram lookup, + # projection and hash workspace; layer-local buffers are reusable. + residual = 2 * d.hc_mult * d.hidden_size * 2 + hash_columns = (d.engram_max_ngram_size - 1) * d.engram_n_heads + engram = 2 * (hash_columns * d.engram_head_dim + (d.hc_mult + 1) * d.hidden_size) + statistics = 2 * (d.hc_mult + 2) * d.hc_mult * 4 + hashes = hash_columns * len(d.engram_layer_ids) * 8 + return float(num_tokens * (residual + engram + statistics + hashes)) + + def get_resident_weights_bytes(self) -> float: + return self._resident_weight_bytes + + def get_kvcache_bytes_per_sequence(self, seq_len: int) -> float: + return self.extra_params.kvcache_bytes(seq_len, self.kv_cache_layout) + + def get_kvcache_max_tokens(self, kv_budget_bytes: float) -> int: + return self._binary_search_kvcache_max_tokens(kv_budget_bytes) + + def get_kvcache_batch_capacity(self, kv_budget_bytes: float, max_batch_size: int) -> int: + d = self.extra_params + window_entry, main_entry, index_entry = d.kv_entry_bytes(self.kv_cache_layout) + fixed = d.num_hidden_layers * d.sliding_window * window_entry + fixed += sum( + 2 * d.compress_ratios[i] * d.head_dim * 4 for i in d.kv_source_layer_ids if d.compress_ratios[i] > 1 + ) + slope = sum((main_entry + index_entry) / d.compress_ratios[i] for i in d.kv_source_layer_ids) + # Reserve complete ring/state buffers for every scheduler slot. Charging + # the asymptotic global slope also covers odd ratio-two publication tails. + return max(0, int((kv_budget_bytes - max_batch_size * fixed) // slope)) diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py index 5117a8703..da38f4dfe 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py @@ -33,6 +33,7 @@ "DEEPSEEK", "DEEPSEEKV32", "DEEPSEEKV4", + "DEEPSEEKV41", "KIMIK25", "KIMIK3", "HYBRIDMOE", @@ -642,6 +643,12 @@ def _infer_quant_modes_from_raw_config(raw_config: dict, architecture: str | Non ): overrides["moe_quant_mode"] = common.MoEQuantMode.w4a8_mxfp4_mxfp8 + if ( + architecture == "DeepseekV41ForCausalLM" + and str(raw_config.get("quantization_config", {}).get("expert_dtype", "")).lower() == "fp4" + ): + overrides["moe_quant_mode"] = common.MoEQuantMode.w4a8_mxfp4_mxfp8 + # KVCache quant mode # TODO: support fp4 kv cache if kv_cache_algo == "fp8": @@ -652,7 +659,7 @@ def _infer_quant_modes_from_raw_config(raw_config: dict, architecture: str | Non raise ValueError(f"Unsupported kv cache algorithm: {kv_cache_algo}") # DSV4 sparse attention requires FP8 KV cache across all backends. - if architecture == "DeepseekV4ForCausalLM": + if architecture in {"DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"}: overrides["kvcache_quant_mode"] = common.KVCacheQuantMode.fp8 # FMHA quant mode @@ -792,10 +799,14 @@ def _is_dsv4_fp4_expert_model(model_path: str) -> bool: ``expert_dtype`` and return False. """ info = _get_model_info(model_path) - if info.get("architecture") != "DeepseekV4ForCausalLM": + architecture = info.get("architecture") + if architecture not in {"DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"}: return False raw_config = info.get("raw_config", {}) - if str(raw_config.get("expert_dtype") or "").lower() != "fp4": + expert_config = ( + raw_config.get("quantization_config", {}) if architecture == "DeepseekV41ForCausalLM" else raw_config + ) + if str(expert_config.get("expert_dtype") or "").lower() != "fp4": return False _gemm_algos, moe_algos = _collect_mixed_precision_layer_algos(raw_config) return "nvfp4" not in moe_algos 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 9c995d738..2667f49d7 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -121,6 +121,7 @@ class ForwardPassPerfModelConfig: nextn: int = 0 speculation: dict[str, Any] | None = None kv_block_size: int | None = None + decoder_replay: bool = False estimation_mode: str = "auto" database_mode: str = "SILICON" transfer_policy: str | tuple[str, ...] | None = None @@ -1202,6 +1203,7 @@ def _engine_config_json(model: Any, database: Any) -> str: "comm": _raw_quant_name(getattr(model_config, "comm_quant_mode", None)), }, "model_config": { + "decoder_replay": bool(getattr(model_config, "decoder_replay", False)), "cp_style": getattr(model_config, "cp_style", None), "workload_distribution": getattr(model_config, "workload_distribution", None), "overwrite_num_layers": getattr(model_config, "overwrite_num_layers", None), diff --git a/python/aisimulate/src/aisimulate_core/sdk/utils.py b/python/aisimulate/src/aisimulate_core/sdk/utils.py index 1bae9c393..94d6251ce 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/utils.py +++ b/python/aisimulate/src/aisimulate_core/sdk/utils.py @@ -1191,6 +1191,8 @@ def positive_int(value, name): "index_n_heads": config["index_n_heads"], "index_topk": config["index_topk"], } + elif architecture == "DeepseekV41ForCausalLM": + extra_params = common.DeepSeekV41Config.from_text_config(config) elif architecture == "DeepseekV4ForCausalLM": compress_ratios = tuple(config["compress_ratios"]) if len(compress_ratios) < layers: diff --git a/python/aisimulate/src/aisimulate_core/systems/b200_sxm.yaml b/python/aisimulate/src/aisimulate_core/systems/b200_sxm.yaml index 8ed09ff5b..244a16d8e 100644 --- a/python/aisimulate/src/aisimulate_core/systems/b200_sxm.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/b200_sxm.yaml @@ -5,6 +5,7 @@ data_dir: data/b200_sxm # relative to systems_dir gpu: + fp32_flops: 75000000000000 # https://www.nvidia.com/en-us/data-center/hgx/ (600 TFLOPS FP32 / 8 GPUs) mem_bw: 7700000000000 # 7.7TB/s mem_bw_empirical_scaling_factor: 0.8 # some nonofficial correction based on observations, you should try to modify based on your own observations mem_empirical_constant_latency: 0.000003 # 3us some nonofficial correction based on observations, you should try to modify based on your own observations @@ -31,4 +32,4 @@ misc: 4: 411041792 # 392MB 8: 411041792 # 392MB other_mem: 3758096384 # increase from 551MB to 3.5GB for safer deployment, this will cover part of the inaccurate mem calc. - nccl_version: '2.27.3' \ No newline at end of file + nccl_version: '2.27.3' diff --git a/python/aisimulate/src/aisimulate_core/systems/b300_sxm.yaml b/python/aisimulate/src/aisimulate_core/systems/b300_sxm.yaml index 144af7638..1ea5ecd32 100644 --- a/python/aisimulate/src/aisimulate_core/systems/b300_sxm.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/b300_sxm.yaml @@ -5,6 +5,7 @@ data_dir: data/b300_sxm # relative to systems_dir gpu: + fp32_flops: 75000000000000 # https://www.nvidia.com/en-us/data-center/hgx/ (600 TFLOPS FP32 / 8 GPUs) mem_bw: 7750000000000 # 7.75 TB/s (Total mem_bw: 62TB/s / 8 GPUs) mem_bw_empirical_scaling_factor: 0.8 # some nonofficial correction based on observations, you should try to modify based on your own observations mem_empirical_constant_latency: 0.000003 # 3us some nonofficial correction based on observations, you should try to modify based on your own observations diff --git a/python/aisimulate/src/aisimulate_core/systems/gb200.yaml b/python/aisimulate/src/aisimulate_core/systems/gb200.yaml index e721d2daf..fe33a950b 100644 --- a/python/aisimulate/src/aisimulate_core/systems/gb200.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/gb200.yaml @@ -5,6 +5,7 @@ data_dir: data/gb200 # relative to systems_dir gpu: + fp32_flops: 80000000000000 # https://www.nvidia.com/pt-br/data-center/gb200-nvl72/ (5760 TFLOPS FP32 / 72 GPUs) mem_bw: 8000000000000 # 8TB/s mem_bw_empirical_scaling_factor: 0.8 # some nonofficial correction based on observations, you should try to modify based on your own observations mem_empirical_constant_latency: 0.000003 # 3us some nonofficial correction based on observations, you should try to modify based on your own observations diff --git a/python/aisimulate/src/aisimulate_core/systems/gb300.yaml b/python/aisimulate/src/aisimulate_core/systems/gb300.yaml index a93085823..1f93ba748 100644 --- a/python/aisimulate/src/aisimulate_core/systems/gb300.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/gb300.yaml @@ -5,6 +5,7 @@ data_dir: data/gb300 # relative to systems_dir gpu: + fp32_flops: 80000000000000 # https://lenovopress.lenovo.com/lp2357.pdf (GB300, 80 TFLOPS FP32 per GPU) mem_bw: 8000000000000 # 8TB/s mem_bw_empirical_scaling_factor: 0.8 # some nonofficial correction based on observations, you should try to modify based on your own observations mem_empirical_constant_latency: 0.000003 # 3us some nonofficial correction based on observations, you should try to modify based on your own observations 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 b05f73880..5998e4cda 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -100,6 +100,7 @@ def test_stable_function_signatures() -> None: "kv_block_size: 'int | None' = None, " "systems_path: 'str | None' = None, " "forward_model: 'str | None' = None, " + "decoder_replay: 'bool' = False, " "database_mode: 'str | None' = None, shared_layer: 'bool | None' = None, " "transfer_policy: 'str | list[str] | None' = None, " "strict_provenance: 'bool | None' = None) -> 'bytes'" diff --git a/python/aisimulate/tests/cross_package/test_import_contract.py b/python/aisimulate/tests/cross_package/test_import_contract.py index 239c0b6d5..56b2f16d2 100644 --- a/python/aisimulate/tests/cross_package/test_import_contract.py +++ b/python/aisimulate/tests/cross_package/test_import_contract.py @@ -25,6 +25,7 @@ "backends.trtllm_backend", "backends.vllm_backend", "common", + "deepseek_v41", "config", "config_builders", "engine", @@ -38,6 +39,7 @@ "models.deepseek", "models.deepseek_v32", "models.deepseek_v4", + "models.deepseek_v41", "models.gemma4", "models.gpt", "models.helpers", 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 b612f7f54..82485e19c 100644 --- a/python/aisimulate/tests/e2e/cli/test_cli_build_default.py +++ b/python/aisimulate/tests/e2e/cli/test_cli_build_default.py @@ -183,3 +183,45 @@ def test_cli_default_build_subset(case: dict): # TODO: remove try/except around save_results assert "Failed to save results" not in combined_output + + +@pytest.mark.parametrize("total_gpus", [4, 32]) +def test_v41_sol_default_accepts_automatic_moe_parallelism(total_gpus): + """Exercise the actual default sweep: model construction alone missed this bug.""" + completed = sp.run( + [ + "aiconfigurator", + "cli", + "default", + "--model", + "deepseek-ai/DeepSeek-V4.1-Flash", + "--system", + "gb300", + "--backend", + "sglang", + "--total-gpus", + str(total_gpus), + "--serving-mode", + "agg", + "--database-mode", + "SOL", + "--isl", + "1024", + "--osl", + "128", + "--nextn", + "0", + ], + capture_output=True, + text=True, + ) + output = f"{completed.stdout}\n{completed.stderr}" + assert completed.returncode == 0, output + # The CLI currently logs sweep failures without a nonzero exit code, so + # require actual ranked results and the correct topology, not exit zero. + assert "AIConfigurator Final Results" in output, output + 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 + assert "Parallelism width mismatch" not in output + assert "No successful experiment runs" not in output diff --git a/python/aisimulate/tests/unit/cli/test_afd_phase_completion.py b/python/aisimulate/tests/unit/cli/test_afd_phase_completion.py index f13dddcf4..da75656c0 100644 --- a/python/aisimulate/tests/unit/cli/test_afd_phase_completion.py +++ b/python/aisimulate/tests/unit/cli/test_afd_phase_completion.py @@ -20,6 +20,17 @@ pytestmark = pytest.mark.unit +@pytest.fixture(autouse=True) +def _fake_model_architecture(monkeypatch): + """Keep the synthetic model local while exercising the AFD capability guard.""" + + def resolve_model(model_path): + assert model_path == "test-model" + return {"architecture": "DeepseekV3ForCausalLM"} + + monkeypatch.setattr("aisimulate_core.sdk.utils.get_model_config_from_model_path", resolve_model) + + def _fake_phase_metrics( *, t_a_layer: float, diff --git a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py index 20fb5933f..d3fcc12ec 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py @@ -89,6 +89,8 @@ def model(): model.model_name = "test-model" model.forward_model = "op_level" model._nextn = 0 + model.get_resident_weights_bytes.return_value = 0.0 + model.get_additional_activation_bytes.return_value = 0.0 model.encoder_ops = [] model.context_ops = [ _StaticOp("context_attention", latency_ms=11.0, energy_wms=110.0), @@ -125,6 +127,8 @@ class TestMTPActivationMemoryScaling: def _model(): return SimpleNamespace( context_ops=[SimpleNamespace(get_weights=lambda: 0.0)], + get_resident_weights_bytes=lambda: 0.0, + get_additional_activation_bytes=lambda _tokens: 0.0, config=ModelConfig( tp_size=1, pp_size=1, diff --git a/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py b/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py index 70967f563..5b33cc30e 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py @@ -73,6 +73,8 @@ def backends(): def _model(*, tp=4, dp=1, ep=1, hidden_size=7168, heads=128): return SimpleNamespace( context_ops=[SimpleNamespace(get_weights=lambda: 2 * GIB)], + get_resident_weights_bytes=lambda: 2 * GIB, + get_additional_activation_bytes=lambda _tokens: 0, model_family="DEEPSEEKV4", _num_heads=heads, _head_size=512, diff --git a/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py b/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py index 247e57133..86203a2cf 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py @@ -42,6 +42,8 @@ def test_step3p7_uses_moe_activation_and_dispatch_workspace(backend, moe_coeffic nextn=0, ), get_kvcache_bytes_per_sequence=lambda _seq: 0, + get_resident_weights_bytes=lambda: 0, + get_additional_activation_bytes=lambda _tokens: 0, _cp_kv_memory_divisor=lambda: 1, ) database = SimpleNamespace(system_spec={"misc": {"nccl_mem": {1: 0}, "other_mem": 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 bdb9a9367..be3bc872b 100644 --- a/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py +++ b/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py @@ -247,7 +247,7 @@ def test_engine_spec_schema_version_is_eighteen(): """Verification width fields extend the positional operation payload.""" from aisimulate.sdk import engine - assert engine.ENGINE_SPEC_SCHEMA_VERSION == 18 + assert engine.ENGINE_SPEC_SCHEMA_VERSION == 19 def test_lanes_outside_the_known_vocabulary_stay_reachable(): diff --git a/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py b/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py new file mode 100644 index 000000000..b06c5acf2 --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py @@ -0,0 +1,488 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections import Counter +from dataclasses import replace + +import pytest + +from aisimulate_core.sdk import common +from aisimulate_core.sdk.deepseek_v41 import ( + MODEL_PATH, + DeepSeekV41Config, + DeepSeekV41KVCacheLayout, + V41RequestWorkload, + resolve_execution_profile, + resolve_kv_cache_layout, + stage_workloads, +) +from aisimulate_core.sdk.models.helpers import _infer_quant_modes_from_raw_config +from aisimulate_core.sdk.utils import get_model_config_from_model_path + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def descriptor(): + return get_model_config_from_model_path(MODEL_PATH)["extra_params"] + + +def test_real_v41_config_and_quant(descriptor): + info = get_model_config_from_model_path(MODEL_PATH) + assert isinstance(descriptor, DeepSeekV41Config) + assert (info["layers"], info["hidden_size"], info["n"]) == (40, 5120, 64) + assert len(descriptor.compress_ratios) == 40 + assert Counter(descriptor.layer_role(i) for i in range(40)) == {"swa": 2, "full": 4, "reindex": 4, "reuse": 30} + quant = _infer_quant_modes_from_raw_config(info["raw_config"], info["architecture"]) + assert quant["moe_quant_mode"] == common.MoEQuantMode.w4a8_mxfp4_mxfp8 + + +@pytest.mark.parametrize("system", ["gb200", "gb300"]) +def test_sglang_blackwell_experts_select_measured_kernel_lane(system): + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.models.helpers import resolve_dsv4_moe_arch, resolve_dsv4_moe_arch_mode + + assert resolve_dsv4_moe_arch_mode(MODEL_PATH, system, "sglang") == common.MoEQuantMode.w4a8_mxfp4_mxfp8_trtllm + assert resolve_dsv4_moe_arch_mode(MODEL_PATH, system, "vllm") is None + explicit = ModelConfig(moe_quant_mode=common.MoEQuantMode.bfloat16) + resolve_dsv4_moe_arch(explicit, MODEL_PATH, system_name=system, backend_name="sglang") + assert explicit.moe_quant_mode == common.MoEQuantMode.bfloat16 + + +@pytest.mark.parametrize("backend", ["sglang", "vllm", "trtllm"]) +@pytest.mark.parametrize("total_gpus", [4, 32]) +def test_default_task_enumerates_matching_moe_parallel_widths(backend, total_gpus): + from aisimulate.sdk.task_v2 import Task + + task = Task( + model_path=MODEL_PATH, + system_name="gb300", + backend_name=backend, + total_gpus=total_gpus, + database_mode="SOL", + isl=1024, + osl=128, + nextn=0, + ) + assert task.model_family == "DEEPSEEKV41" + assert task.is_moe + parallel = [tuple(choice) for choice in task.iter_parallel("agg")] + assert (4, 1, 1, 1, 4, 1) in parallel + assert all(tp * dp * cp == moe_tp * moe_ep for tp, _, dp, moe_tp, moe_ep, cp in parallel) + + +def test_shared_pool_memory_slope_and_odd_decode(descriptor): + assert descriptor.compressed_entry_bytes == 288 + assert descriptor.index_entry_bytes == 68 + assert descriptor.kvcache_bytes(130) - descriptor.kvcache_bytes(128) == 2 * 890 + # Only the ratio-one owner grows on odd tokens; the three ratio-two owners + # publish their next compressed entry on the following even token. + assert descriptor.kvcache_bytes(129) - descriptor.kvcache_bytes(128) == 356 + assert descriptor.engram_table_bytes(1) == 202758032400 + + +def test_sglang_physical_payload_owners_and_publication_boundaries(descriptor): + layout = DeepSeekV41KVCacheLayout.SGLANG_FP8_BF16 + assert descriptor.kv_entry_bytes(layout) == (584, 584, 68) + # 40 SWA pools; one full-rate and three half-rate shared KV/index owners. + # Reindex layers share both physical pools, rather than allocating copies. + assert descriptor.kvcache_bytes(131072, layout) == 216662016 == 206.625 * 2**20 + assert descriptor.kvcache_bytes(129, layout) - descriptor.kvcache_bytes(128, layout) == 652 + assert descriptor.kvcache_bytes(130, layout) - descriptor.kvcache_bytes(129, layout) == 4 * 652 + assert descriptor.kvcache_bytes(130, layout) - descriptor.kvcache_bytes(128, layout) == 3260 + for invalid in (replace(descriptor, head_dim=256), replace(descriptor, qk_rope_head_dim=32)): + with pytest.raises(ValueError, match="pinned SGLang"): + invalid.kv_entry_bytes(layout) + with pytest.raises(ValueError, match="unknown"): + descriptor.kv_entry_bytes("unknown") + with pytest.raises(ValueError, match="unknown"): + resolve_kv_cache_layout("unknown") + + +@pytest.mark.parametrize("backend", ["sglang", "vllm", "trtllm"]) +@pytest.mark.parametrize("tokens", [127, 128, 129, 130, 16384, 131072]) +def test_physical_and_theoretical_capacity_inverse(backend, tokens, descriptor): + model = _build_model(backend=backend) + expected = DeepSeekV41KVCacheLayout.SGLANG_FP8_BF16 if backend == "sglang" else DeepSeekV41KVCacheLayout.LOGICAL_FP4 + assert model.kv_cache_layout == expected + capacity = descriptor.kvcache_bytes(tokens, expected) + assert model.get_kvcache_bytes_per_sequence(tokens) == capacity + assert model.get_kvcache_max_tokens(capacity) == tokens + assert model.get_kvcache_max_tokens(capacity - 1) < tokens + + +def test_replay_tail_is_bounded_per_actual_request(descriptor): + requests = (V41RequestWorkload(256, 1024), V41RequestWorkload(3, 4096)) + full = stage_workloads(descriptor, "full", requests, 39) + assert [(r.query_tokens, r.prefix_tokens) for r in full] == [(256, 1024), (3, 4096)] + tail = stage_workloads(descriptor, "decoder_bounded", requests, 21) + assert [(r.query_tokens, r.prefix_tokens) for r in tail] == [(128, 1152), (3, 4096)] + assert stage_workloads(descriptor, "decoder_bounded", requests, 20) == full + + +@pytest.mark.parametrize("backend", ["vllm", "trtllm"]) +def test_unverified_replay_backend_fails(backend): + assert resolve_execution_profile(False, backend) == "full" + with pytest.raises(NotImplementedError, match="not verified"): + resolve_execution_profile(True, backend) + + +def _build_model(*, replay=False, backend="sglang", tp=4): + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.models import get_model + + return get_model( + MODEL_PATH, + ModelConfig(tp_size=tp, pp_size=1, attention_dp_size=1, moe_tp_size=1, moe_ep_size=tp, decoder_replay=replay), + backend, + ) + + +@pytest.mark.parametrize("backend", ["sglang", "vllm", "trtllm"]) +def test_v41_text_graph_full_profile_all_backends(backend): + import json + + model = _build_model(backend=backend) + assert model.execution_profile == "full" + assert not model.encoder_ops + specs = [json.loads(op._spec_json()) for op in model.context_ops] + stages = [op["Dsv41Stage"] for op in specs if "Dsv41Stage" in op] + assert len(stages) == 40 + layouts = { + child["Dsv41Attention"]["kv_cache_layout"] + for stage in stages + for child in stage["children"] + if "Dsv41Attention" in child + } + assert layouts == {resolve_kv_cache_layout(backend).value} + assert sum(stage["bounded"] for stage in stages) == 19 + assert all(not stage["decoder_replay"] for stage in stages) + roles = Counter( + next(c["Dsv41Attention"]["role"] for c in stage["children"] if "Dsv41Attention" in c) for stage in stages + ) + assert roles == {"swa": 2, "full": 4, "reindex": 4, "reuse": 30} + + +def test_replay_keeps_inventory_and_kv_pool_capacity(): + full, replay = _build_model(), _build_model(replay=True) + assert replay.execution_profile == "decoder_bounded" + assert full.get_resident_weights_bytes() == replay.get_resident_weights_bytes() + assert full.get_resident_weights_bytes() > full.extra_params.engram_table_bytes(4) + assert full.get_kvcache_bytes_per_sequence(4096) == replay.get_kvcache_bytes_per_sequence(4096) + assert replay.get_kvcache_max_tokens(replay.get_kvcache_bytes_per_sequence(4096)) == 4096 + + +def test_batch_capacity_reserves_each_request_window(descriptor): + model = _build_model() + fixed = 40 * 128 * 584 + 3 * 2 * 2 * 512 * 4 + assert model.get_kvcache_batch_capacity(4 * fixed + 4096 * 1630, 4) == 4096 + assert model.get_kvcache_batch_capacity(4 * fixed - 1, 4) == 0 + assert model.get_additional_activation_bytes(1024) > 1024 * 4 * 5120 * 2 + + +@pytest.mark.parametrize(("dp", "pp"), [(2, 1), (1, 2)]) +def test_unmodeled_parallel_cache_ownership_fails(dp, pp): + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.models import get_model + + with pytest.raises(NotImplementedError, match="DP Engram collectives and PP cache ownership"): + get_model( + MODEL_PATH, + ModelConfig(tp_size=4, pp_size=pp, attention_dp_size=dp, moe_tp_size=1, moe_ep_size=4 * dp), + "sglang", + ) + + +def test_native_sol_replay_decode_and_short_extend_contract(): + from aisimulate_core.sdk.engine import EngineHandle, compile_engine + + def engine(replay): + return EngineHandle( + compile_engine( + MODEL_PATH, + "gb300", + "sglang", + tp_size=4, + moe_tp_size=1, + moe_ep_size=4, + decoder_replay=replay, + database_mode="SOL", + ) + ) + + full, replay = engine(False), engine(True) + assert replay.predict_prefill_latency(1, 1024) < full.predict_prefill_latency(1, 1024) + assert replay.predict_decode_latency(1, 1024) == full.predict_decode_latency(1, 1024) + # Both modes process the same three actual query tokens; the bounded SWA + # intentionally does not read the pre-existing ring prefix. + assert replay.predict_prefill_latency(1, 4099, 4096) > 0 + assert replay.mixed_step_latency(2048, 1, 1024, 2) > 0 + + +def test_native_block32_shared_projections_have_distinct_perf_identity(): + import json + + model = _build_model() + first_layer = next( + json.loads(op._spec_json())["Dsv41Stage"] + for op in model.context_ops + if "Dsv41Stage" in json.loads(op._spec_json()) + ) + linears = [c["Dsv41Linear"] for c in first_layer["children"] if "Dsv41Linear" in c] + assert {(op["n"], op["k"]) for op in linears} == {(1152, 5120), (5120, 576)} + assert all(c["Gemm"]["quant_mode"] == "bfloat16" for c in first_layer["children"] if "Gemm" in c) + + +@pytest.mark.parametrize("first_replay", [False, True]) +def test_engine_cache_keeps_replay_profiles_separate(first_replay): + from aisimulate_core.sdk import rust_engine_step + from aisimulate_core.sdk.perf_database import get_database_view + + database = get_database_view("gb300", "sglang", "current", allow_missing_data=True, database_mode="SOL") + rust_engine_step._engine_handle_cache_clear() + try: + models = {replay: _build_model(replay=replay) for replay in (first_replay, not first_replay)} + handles = {replay: rust_engine_step._cached_engine_handle(model, database) for replay, model in models.items()} + assert handles[False] is not handles[True] + assert handles[True].predict_prefill_latency(1, 1024) < handles[False].predict_prefill_latency(1, 1024) + assert handles[True].predict_decode_latency(1, 1024) == handles[False].predict_decode_latency(1, 1024) + assert ( + rust_engine_step._cached_engine_handle(_build_model(replay=first_replay), database) is handles[first_replay] + ) + finally: + rust_engine_step._engine_handle_cache_clear() + + +def test_sglang_pure_tp_eager_shared_and_routed_costs_are_sequential(): + import json + + import aisimulate_core._native as native + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.engine import _evaluate_single_op + from aisimulate_core.sdk.models import get_model + from aisimulate_core.sdk.perf_database import get_database_view + + model = get_model(MODEL_PATH, ModelConfig(tp_size=4, moe_tp_size=4, moe_ep_size=1), "sglang") + stages = [ + json.loads(op._spec_json())["Dsv41Stage"] + for op in model.generation_ops + if "Dsv41Stage" in json.loads(op._spec_json()) + ] + assert len(stages) == 40 + assert all(not any("Overlap" in child for child in stage["children"]) for stage in stages) + children = stages[0]["children"] + shared = [child for child in children if "shared_" in next(iter(child.values())).get("name", "")] + routed = [ + child + for child in children + if next(iter(child.values())).get("name", "") + in ("generation_router_gemm", "generation_moe_pre_dispatch", "generation_moe") + ] + assert len(shared) == 3 and len(routed) == 2 + db = get_database_view("gb300", "sglang", "current", allow_missing_data=True, database_mode="SOL") + + def cost(children): + stage = stages[0] | {"children": children} + op = native.op_from_spec_json(json.dumps({"Dsv41Stage": stage})) + return float(_evaluate_single_op(db, op, is_context=False, batch_size=2, s=129, prefix=0, x=2)) + + assert cost(shared + routed) == pytest.approx(cost(shared) + cost(routed)) + assert cost(shared + routed) > max(cost(shared), cost(routed)) + + +def test_unqualified_ep_generation_overlap_is_not_changed_by_tp_eager_fix(): + import json + + model = _build_model() + stages = [ + json.loads(op._spec_json())["Dsv41Stage"] + for op in model.generation_ops + if "Dsv41Stage" in json.loads(op._spec_json()) + ] + assert all(any("Overlap" in child for child in stage["children"]) for stage in stages) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"compress_ratios": (0,)}, "one backbone compression ratio"), + ({"compress_ratios": (3,) * 40}, "one backbone compression ratio"), + ({"kv_source_layer_ids": (8, 2)}, "unique, ascending"), + ({"index_source_layer_ids": (2, 8, 14)}, "must also own an indexer"), + ({"candidate_source_layer_id": 3}, "candidate source must own"), + ({"engram_num_embeddings": (1,)}, "layer/table counts differ"), + ( + {"kv_source_layer_ids": (0, 2, 8, 14, 20), "index_source_layer_ids": (0, 2, 8, 14, 20, 24, 28, 32, 36)}, + "positive compression ratio", + ), + ], +) +def test_descriptor_rejects_malformed_ownership(descriptor, changes, message): + with pytest.raises(ValueError, match=message): + replace(descriptor, **changes).validate() + + +def test_descriptor_rejects_ownerless_and_incompatible_compression(descriptor): + for layer, ratio in [(0, 1), (3, 1)]: + ratios = list(descriptor.compress_ratios) + ratios[layer] = ratio + with pytest.raises(ValueError, match="no compatible preceding KV owner"): + replace(descriptor, compress_ratios=tuple(ratios)).validate() + + +@pytest.mark.parametrize( + ("changes", "error", "message"), + [ + ({"nextn": 1}, NotImplementedError, "nextn=0"), + ({"overwrite_num_layers": 2}, ValueError, "layer overrides"), + ({"moe_backend": "megamoe"}, NotImplementedError, "decomposed MoE"), + ({"cp_size": 2, "moe_ep_size": 8}, NotImplementedError, "Context parallelism"), + ({"tp_size": 16, "moe_ep_size": 16}, ValueError, "TP must divide"), + ], +) +def test_model_constructor_rejects_unmodeled_configuration(changes, error, message): + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.models import get_model + + config = ModelConfig(**({"tp_size": 4, "moe_tp_size": 1, "moe_ep_size": 4} | changes)) + with pytest.raises(error, match=message): + get_model(MODEL_PATH, config, "sglang") + + +@pytest.mark.parametrize("tp", [1, 4, 8]) +def test_replicated_indexer_preserves_weights_and_score_cost_across_tp(tp): + import json + + import aisimulate_core._native as native + from aisimulate_core.sdk.engine import _evaluate_single_op + from aisimulate_core.sdk.perf_database import get_database_view + + model = _build_model(tp=tp) + specs = [json.loads(op._spec_json()) for op in model.context_ops] + attention = next( + child["Dsv41Attention"] + for stage in specs + if "Dsv41Stage" in stage + for child in stage["Dsv41Stage"]["children"] + if "Dsv41Attention" in child and child["Dsv41Attention"]["role"] == "reindex" + ) + assert attention["index_n_heads"] == 32 + indexed = native.op_from_spec_json(json.dumps({"Dsv41Attention": attention})) + reused = native.op_from_spec_json(json.dumps({"Dsv41Attention": attention | {"role": "reuse"}})) + # 1280*32*128 FP8 weights plus 32x32 block scales, and 5120*32 BF16 gates. + assert indexed.get_weights() - reused.get_weights() == 5_575_680 + db = get_database_view("gb300", "sglang", "current", allow_missing_data=True, database_mode="SOL") + + def cost(op): + return float(_evaluate_single_op(db, op, is_context=True, batch_size=1, s=4096, prefix=0, x=4096)) + + index_cost = cost(indexed) - cost(reused) + # Every TP uses the identical replicated indexer. Changing main attention + # heads must not divide this complete scoring/selection contribution. + reference = attention | {"num_heads": 64, "o_groups": 8} + full = native.op_from_spec_json(json.dumps({"Dsv41Attention": reference})) + full_reuse = native.op_from_spec_json(json.dumps({"Dsv41Attention": reference | {"role": "reuse"}})) + assert index_cost == pytest.approx(cost(full) - cost(full_reuse), rel=1e-12) + assert index_cost > 0 + + +@pytest.mark.parametrize( + ("backend_name", "coefficient", "overhead"), [("sglang", 13, 1.15), ("vllm", 10, 1.0), ("trtllm", 10, 1.0)] +) +def test_v41_actual_activation_memory_uses_moe_coefficient(backend_name, coefficient, overhead): + from aisimulate_core.sdk.backends.factory import get_backend + from aisimulate_core.sdk.perf_database import get_database_view + + model = _build_model(backend=backend_name) + backend = get_backend(backend_name) + db = get_database_view("gb300", backend_name, "current", allow_missing_data=True, database_mode="SOL") + memory = backend._get_memory_usage(model, db, 1, 1, 8192, 1, num_tokens=8192) + # Generic MoE workspace plus the separately owned mHC/Engram buffers. + workspace_width = 5120 if backend_name == "sglang" else 64 * 512 + workspace = 8192 * workspace_width * 384 * 6 / 4 / 128 * 4 + expanded = model.get_additional_activation_bytes(8192) + expected = (2 * 8192 * 64 * 512 * coefficient + workspace + expanded) * overhead + assert memory["activations"] * (1 << 30) == pytest.approx(expected) + assert memory["weights"] * (1 << 30) == model.get_resident_weights_bytes() + + +def test_nested_stage_rejects_retired_dispatch_at_serialization(): + import json + + import aisimulate_core._native as native + + model = _build_model() + stage = next(json.loads(op._spec_json()) for op in model.context_ops if "Dsv41Stage" in json.loads(op._spec_json())) + dispatch = next(child for child in stage["Dsv41Stage"]["children"] if "MoeDispatch" in child) + dispatch["MoeDispatch"]["flavor"] = "RetiredDeepEp" + inner = stage | {"Dsv41Stage": stage["Dsv41Stage"] | {"children": [dispatch]}} + outer = stage | {"Dsv41Stage": stage["Dsv41Stage"] | {"children": [inner]}} + op = native.op_from_spec_json(json.dumps(outer)) + with pytest.raises(ValueError, match="retired"): + native.ops_json_from_ops([op]) + + +@pytest.mark.parametrize("system", ["b200_sxm", "b300_sxm"]) +def test_hgx_blackwell_has_published_scalar_fp32_rate(system): + from aisimulate_core.sdk.engine import EngineHandle, compile_engine + from aisimulate_core.sdk.perf_database import get_database_view + + db = get_database_view(system, "sglang", "current", allow_missing_data=True, database_mode="SOL") + assert db.system_spec["gpu"]["fp32_flops"] == 75e12 + engine = EngineHandle( + compile_engine(MODEL_PATH, system, "sglang", tp_size=4, moe_tp_size=4, moe_ep_size=1, database_mode="SOL") + ) + assert engine.predict_prefill_latency(1, 128) > 0 + + +@pytest.mark.parametrize("decoder_replay", [False, True]) +def test_afd_rejects_v41_before_search_or_session_construction(decoder_replay): + from aisimulate.sdk.inference_session import AFDInferenceSession + from aisimulate.sdk.task_v2 import Task + from aisimulate_core.sdk.config import ModelConfig + + with pytest.raises(NotImplementedError, match="AFD does not support DeepSeek-V4.1"): + Task( + model_path=MODEL_PATH, + system_name="gb300", + backend_name="sglang", + total_gpus=8, + database_mode="SOL", + serving_mode="afd", + isl=1024, + osl=128, + nextn=0, + ) + with pytest.raises(NotImplementedError, match="AFD does not support DeepSeek-V4.1"): + AFDInferenceSession( + MODEL_PATH, + ModelConfig(decoder_replay=decoder_replay), + ModelConfig(decoder_replay=decoder_replay), + None, + None, + None, + ) + + +@pytest.mark.parametrize("is_context", [False, True]) +@pytest.mark.parametrize("batch,seq", [(1, 4), (0, 4), (1, 0)]) +def test_native_attention_rejects_unknown_role_before_zero_work(is_context, batch, seq): + import json + + import aisimulate_core._native as native + from aisimulate_core.sdk.engine import _evaluate_single_op + from aisimulate_core.sdk.perf_database import get_database_view + + model = _build_model() + specs = [json.loads(op._spec_json()) for op in model.context_ops] + attention = next( + child["Dsv41Attention"] + for stage in specs + if "Dsv41Stage" in stage + for child in stage["Dsv41Stage"]["children"] + if "Dsv41Attention" in child + ) + op = native.op_from_spec_json(json.dumps({"Dsv41Attention": attention | {"role": "ful", "is_context": is_context}})) + db = get_database_view("gb300", "sglang", "current", allow_missing_data=True, database_mode="SOL") + with pytest.raises(ValueError, match="attention role must be"): + _evaluate_single_op(db, op, is_context=is_context, batch_size=batch, s=seq, prefix=0, x=batch * seq) diff --git a/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41_residency.py b/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41_residency.py new file mode 100644 index 000000000..565749499 --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41_residency.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Target residency survives speculative and whole-forward graph rewrites.""" + +from types import SimpleNamespace + +import pytest + +from aisimulate_core.sdk.backends.factory import get_backend +from aisimulate_core.sdk.config import ModelConfig +from aisimulate_core.sdk.deepseek_v41 import MODEL_PATH +from aisimulate_core.sdk.models import _apply_forward_model_fpm, get_model +from aisimulate_core.sdk.speculation import NullScheme, SpeculationConfig +from aisimulate_core.sdk.speculation.draft_model import DraftModelScheme +from aisimulate_core.sdk.speculation.ngram import NgramScheme + +pytestmark = pytest.mark.unit + + +def _model(*, model_path=MODEL_PATH, speculation=None, forward_model="op_level", pp=1, backend="sglang", replay=False): + return get_model( + model_path, + ModelConfig( + tp_size=4, + pp_size=pp, + attention_dp_size=1, + moe_tp_size=4, + moe_ep_size=1, + speculation=speculation, + forward_model=forward_model, + decoder_replay=replay, + ), + backend, + ) + + +def _weights_bytes(model, backend): + # Zero runtime overhead isolates model/scheme residency from device capacity. + database = SimpleNamespace(system_spec={"misc": {"nccl_mem": {4: 0.0}, "other_mem": 0.0}}) + memory = get_backend(backend)._get_memory_usage(model, database, batch_size=1, beam_width=1, isl=128, osl=1) + return memory["weights"] * (1 << 30) + + +@pytest.mark.parametrize("forward_model", ["op_level", "fpm"]) +def test_ngram_keeps_complete_native_mxfp4_target_inventory(forward_model): + plain = _model() + ngram = _model( + forward_model=forward_model, + speculation=SpeculationConfig(kind="ngram", params={"num_speculative_tokens": 3}), + ) + assert type(plain.spec_scheme) is NullScheme + assert isinstance(ngram.spec_scheme, NgramScheme) + # Checkpoint geometry: 40 layers, three matrices per expert, 384 experts, + # TP4, with one native MXFP4 scale byte for every 32 expert elements. + expected_scales = 40 * 3 * 5120 * 2304 * 384 // 4 // 32 + assert expected_scales == 4_246_732_800 + target_op_bytes = sum(op.get_weights() for op in plain.context_ops) + expected_target = target_op_bytes + expected_scales + assert plain.get_resident_weights_bytes() == expected_target + assert _weights_bytes(plain, "sglang") == expected_target + assert ngram.spec_scheme.draft_weights_bytes(ngram) == 0 + assert ngram.get_resident_weights_bytes() == expected_target + assert _weights_bytes(ngram, "sglang") == expected_target + if forward_model == "fpm": + assert ngram.context_ops[0].get_weights() == expected_target + + +@pytest.mark.parametrize("model_path,pp", [(MODEL_PATH, 1), ("Qwen/Qwen3-8B", 2)]) +@pytest.mark.parametrize("forward_model", ["op_level", "fpm"]) +def test_owned_draft_weights_are_added_once_after_target_pp_division(model_path, pp, forward_model): + plain = _model(model_path=model_path, pp=pp, backend="vllm") + drafted = _model( + model_path=model_path, + pp=pp, + backend="vllm", + forward_model=forward_model, + speculation=SpeculationConfig( + kind="draft_model", + draft_model_path="Qwen/Qwen3-0.6B", + params={"num_speculative_tokens": 3}, + ), + ) + scheme = drafted.spec_scheme + assert isinstance(scheme, DraftModelScheme) + owned_draft = scheme._draft_model + expected_draft = owned_draft.get_resident_weights_bytes() + assert expected_draft > 0 + assert sum(op.get_weights() for op in drafted.context_ops if op._name.startswith("draft_")) > 0 + assert scheme.draft_weights_bytes(drafted) == expected_draft + assert drafted.get_resident_weights_bytes() == plain.get_resident_weights_bytes() + expected_target_per_stage = plain.get_resident_weights_bytes() / pp + assert _weights_bytes(plain, "vllm") == expected_target_per_stage + assert _weights_bytes(drafted, "vllm") == expected_target_per_stage + expected_draft + if forward_model == "fpm": + assert drafted.context_ops[0].get_weights() == plain.get_resident_weights_bytes() + + +def test_legacy_fpm_rejects_decoder_replay_before_rewriting_graph(monkeypatch): + from aisimulate_core.sdk.operations import fpm_forward + + # A legacy selector has no way to distinguish OFF and ON measurements. + # Retain this boundary regression when newer FPM schemas add that axis. + monkeypatch.setattr( + fpm_forward, + "_CELL_MATCH_COLUMNS", + tuple(column for column in fpm_forward._CELL_MATCH_COLUMNS if column != "execution_profile"), + ) + model = _model(replay=True) + original_context = list(model.context_ops) + original_generation = list(model.generation_ops) + with pytest.raises(NotImplementedError, match="execution_profile identity"): + _apply_forward_model_fpm(model) + assert model.context_ops == original_context + assert model.generation_ops == original_generation diff --git a/python/aisimulate/tests/unit/sdk/models/test_model_config.py b/python/aisimulate/tests/unit/sdk/models/test_model_config.py index fe14681f4..c86350ddf 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_model_config.py +++ b/python/aisimulate/tests/unit/sdk/models/test_model_config.py @@ -107,6 +107,7 @@ def test_model_configs_have_correct_structure(self): ("deepseek-ai/DeepSeek-V3", True), ("deepseek-ai/DeepSeek-V3.2", True), ("deepseek-ai/DeepSeek-V4-Flash", True), + ("deepseek-ai/DeepSeek-V4.1-Flash", True), ("deepseek-ai/DeepSeek-V4-Pro", True), ("sgl-project/DeepSeek-V4-Flash-FP8", True), ("sgl-project/DeepSeek-V4-Pro-FP8", True), diff --git a/python/aisimulate/tests/unit/sdk/models/test_qwen35.py b/python/aisimulate/tests/unit/sdk/models/test_qwen35.py index 689e4e10d..6a8da0d75 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_qwen35.py +++ b/python/aisimulate/tests/unit/sdk/models/test_qwen35.py @@ -47,6 +47,61 @@ def _flatten_ops(phase_ops): yield op +@pytest.mark.parametrize("is_context", [True, False]) +def test_sglang_attention_dp_prices_folded_tp_reduction_numerically(is_context): + """Qwen omits its attention AR, so dispatch must retain BOTH collectives.""" + from aisimulate_core.sdk.engine import _evaluate_single_op + from aisimulate_core.sdk.operations import NCCL + from aisimulate_core.sdk.perf_database import get_database_view + + model = models.get_model( + "Qwen/Qwen3.5-397B-A17B", + _model_config(tp_size=4, attention_dp_size=2, moe_tp_size=1, moe_ep_size=8), + "sglang", + ) + db = get_database_view("gb300", "sglang", "current", allow_missing_data=True, database_mode="SOL") + ops = list(_flatten_ops(model.context_ops if is_context else model.generation_ops)) + phase = "context" if is_context else "generation" + assert not any(op._name in (f"{phase}_gdn_ar", f"{phase}_full_ar") for op in ops) + + def cost(op, tokens): + return float(_evaluate_single_op(db, op, is_context=is_context, batch_size=1, s=64, prefix=0, x=tokens)) + + dispatches = [op for op in ops if op._name.endswith("_moe_pre_dispatch")] + assert len(dispatches) == 2 + for op in dispatches: + spec = json.loads(op._spec_json())["MoeDispatch"] + assert spec["attn_ar_modeled"] + scale, h = spec["scale_factor"], spec["hidden_size"] + reduction = NCCL("reference", scale, "reduce_scatter", h, 4, common.CommQuantMode.half) + gather = NCCL("reference", scale, "all_gather", h, 8, common.CommQuantMode.half) + assert cost(op, 64) == pytest.approx(cost(reduction, 64) + cost(gather, 128), rel=1e-12) + assert cost(op, 64) > cost(gather, 128) + + +@pytest.mark.parametrize("is_context", [True, False]) +def test_trtllm_qwen_dispatch_retains_legacy_allreduce_cost(is_context): + """The unqualified TRT-LLM path keeps its documented pre-existing behavior.""" + from aisimulate_core.sdk.engine import _evaluate_single_op + from aisimulate_core.sdk.perf_database import get_database_view + + model = models.get_model("Qwen/Qwen3.5-397B-A17B", _model_config(tp_size=8), "trtllm") + db = get_database_view("gb200", "trtllm", "current", allow_missing_data=True, database_mode="SOL") + + def cost(op): + return float(_evaluate_single_op(db, op, is_context=is_context, batch_size=1, s=64, prefix=0, x=64)) + + ops = list(_flatten_ops(model.context_ops if is_context else model.generation_ops)) + dispatches = [op for op in ops if op._name.endswith("_moe_pre_dispatch")] + assert len(dispatches) == 2 + for op in dispatches: + spec = json.loads(op._spec_json())["MoeDispatch"] + assert spec["flavor"] == "TrtllmAlltoall" and spec["attn_ar_modeled"] + reference = CustomAllReduce("reference", spec["scale_factor"], spec["hidden_size"], 8) + assert cost(op) == pytest.approx(cost(reference), rel=1e-12) + assert cost(op) > 0 + + @pytest.mark.parametrize( ("model_name", "expected_out_hidden", "is_moe"), [ diff --git a/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py b/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py index 216704949..78a5c04e2 100644 --- a/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py +++ b/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py @@ -85,6 +85,8 @@ def _fake_model(nextn: int = 0, spec_scheme=None, speculation=None): model.encoder_ops = [] model.context_ops = [_RecordingOp("context_attention", 11.0)] model.generation_ops = [_RecordingOp("generation_attention", 2.0)] + model.get_resident_weights_bytes = lambda: sum(op.get_weights() for op in model.context_ops) + model.get_additional_activation_bytes = lambda num_tokens: 0.0 model.config = _model_config(speculation=speculation) model.config.nextn = nextn if spec_scheme is None: diff --git a/python/aisimulate/tests/unit/sdk/test_v41_native_bridge.py b/python/aisimulate/tests/unit/sdk/test_v41_native_bridge.py new file mode 100644 index 000000000..e8368b630 --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_v41_native_bridge.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from aisimulate_core.sdk.deepseek_v41 import MODEL_PATH +from aisimulate_core.sdk.rust_engine_step import ForwardPassPerfModelConfig, RustForwardPassPerfModel + +pytestmark = pytest.mark.unit + + +def _native_model(config): + return RustForwardPassPerfModel.best_available( + ForwardPassPerfModelConfig.from_legacy_engine_config(config, "aggregated") + ) + + +def _config(replay=False, backend="sglang"): + return { + "schema_version": 1, + "model_name": MODEL_PATH, + "system_name": "gb300", + "backend": backend, + "backend_version": "0.5.14" if backend == "sglang" else "0.24.0", + "tp_size": 4, + "pp_size": 1, + "moe_tp_size": 4, + "moe_ep_size": 1, + "attention_dp_size": 1, + "database_mode": "SOL", + "decoder_replay": replay, + } + + +def test_native_builder_preserves_replay_and_backend_guard(): + sample = { + "version": 1, + "scheduled_requests": { + "num_prefill_requests": 1, + "sum_prefill_tokens": 256, + "sum_prefill_kv_tokens": 0, + }, + } + full = _native_model(_config(False)) + bounded = _native_model(_config(True)) + assert 0 < bounded.estimate_forward_pass_time_ms(sample) < full.estimate_forward_pass_time_ms(sample) + with pytest.raises(ValueError, match="not verified"): + _native_model(_config(True, "vllm")) + + +@pytest.mark.parametrize("extends", [((1, 1023), (1023, 1)), ((512, 512), (512, 512)), ((1, 1023), (2, 1022))]) +def test_replay_fpm_rejects_multiple_prefills_even_with_equal_prompts(extends): + # FPM v1 reports variance of full prompt lengths, not actual extends. + # All three valid batches have equal 1024-token prompts. Their aggregate + # does not identify the individual tails, even if the caller knows them. + assert {query + prefix for query, prefix in extends} == {1024} + sample = { + "version": 1, + "scheduled_requests": { + "num_prefill_requests": len(extends), + "sum_prefill_tokens": sum(query for query, _ in extends), + "sum_prefill_kv_tokens": sum(prefix for _, prefix in extends), + "var_prefill_length": 0.0, + }, + } + bounded = _native_model(_config(True)) + with pytest.raises(ValueError, match="multiple prefill requests"): + bounded.estimate_forward_pass_time_ms(sample) + # Decoder OFF retains the existing aggregate approximation. + full = _native_model(_config(False)) + assert full.estimate_forward_pass_time_ms(sample) > 0 + + +@pytest.mark.parametrize("query", [1, 127, 128, 129]) +def test_replay_fpm_single_prefill_matches_explicit_static_geometry(query): + from aisimulate_core.sdk.engine import EngineHandle + + bounded = _native_model(_config(True)) + static = EngineHandle.compile( + MODEL_PATH, + "gb300", + "sglang", + backend_version="0.5.14", + tp_size=4, + moe_tp_size=4, + moe_ep_size=1, + decoder_replay=True, + database_mode="SOL", + ) + sample = { + "version": 1, + "scheduled_requests": { + "num_prefill_requests": 1, + "sum_prefill_tokens": query, + "sum_prefill_kv_tokens": 1024, + }, + } + assert bounded.estimate_forward_pass_time_ms(sample) == pytest.approx( + static.predict_prefill_latency(1, 1024 + query, 1024), rel=1e-12 + ) + + +def test_compile_engine_resolves_v41_native_expert_lane(monkeypatch): + from aisimulate_core.sdk import engine + + original = engine.get_model + captured = [] + + def build(*args, **kwargs): + model = original(*args, **kwargs) + captured.append(model.config.moe_quant_mode.name) + return model + + monkeypatch.setattr(engine, "get_model", build) + handle = engine.EngineHandle.compile( + MODEL_PATH, + "gb300", + "sglang", + backend_version="0.5.14", + tp_size=4, + moe_tp_size=4, + moe_ep_size=1, + database_mode="SOL", + ) + assert captured == ["w4a8_mxfp4_mxfp8_trtllm"] + assert handle.predict_prefill_latency(1, 256) > 0 diff --git a/tests/test_runner.py b/tests/test_runner.py index e9f4ad34e..371afdc58 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1284,6 +1284,52 @@ def test_runner_rejects_unknown_forward_model(value): EngineReplayRunnerFactory(runtime=RecordingRuntime()).create(0).run(_spec(deployment=deployment)) +@pytest.mark.parametrize("replay", [False, True]) +def test_public_replay_keeps_decoder_profile_and_database_policy(replay): + public = CorePredictionConfig.model_validate( + { + "engine": { + "model": "example/model", + "hardware": "gb300", + "backend": "sglang", + "decoder_replay": replay, + "database_mode": "SILICON", + "enable_shared_layer": False, + "strict_provenance": True, + "workers": {"aggregated": {"kv_cache": {"capacity": {"type": "fixed", "blocks": 128}}}}, + } + } + ) + runtime = RecordingRuntime() + spec = prediction_to_replay_spec(public) + EngineReplayRunnerFactory(runtime=runtime).create(0).run(spec) + config = runtime.execution_spec["spec"]["engine"]["rank"]["timing_model"]["config"] + assert config.get("decoder_replay", False) is replay + assert config["database_mode"] == "SILICON" + assert config["enable_shared_layer"] is False + assert config["strict_provenance"] is True + metadata = spec.backend_deployment.performance_model_metadata["aggregated"]["config"] + assert metadata.get("decoder_replay", False) is replay + assert metadata["database_mode"] == "SILICON" + + +@pytest.mark.parametrize("field", ["aic_decoder_replay", "aic_enable_shared_layer", "aic_strict_provenance"]) +def test_replay_policy_alias_requires_a_boolean(field): + engine_args = _engine_args() + engine_args.pop("timing_model") + engine_args[field] = "false" + deployment = BackendDeploymentSpec( + deployment_mode="agg", + backend="vllm", + backend_version="test", + parallel_config={"tp": 2, "attention_dp": 1, "replicas": 2}, + agg_engine_args=engine_args, + num_workers=2, + ) + with pytest.raises(ValueError, match="must be a boolean"): + EngineReplayRunnerFactory(runtime=RecordingRuntime()).create(0).run(_spec(deployment=deployment)) + + def test_memory_detail_reuses_capacity_calculation_without_changing_execution(monkeypatch): from aisimulate_core.sdk import memory diff --git a/tests/test_unified_traffic_runtime.py b/tests/test_unified_traffic_runtime.py index 545f4755e..79fff7f86 100644 --- a/tests/test_unified_traffic_runtime.py +++ b/tests/test_unified_traffic_runtime.py @@ -264,7 +264,10 @@ def test_b200_power_survives_native_json_and_runner_normalization() -> None: assert report.metrics["completed_requests"] == 100 native_summary = report.metadata["native_report"] - for name, expected in {"power_w": 655.9411158961074, "power_coverage": 0.9070317503277924}.items(): + # Decode converts scheduler-inclusive length to past KV before pricing the + # current token. Reverting only that conversion reproduces the older + # section 4.11 capture (655.9411158961074 W, coverage 0.9070317503277924). + for name, expected in {"power_w": 655.957349601573, "power_coverage": 0.907023184956731}.items(): assert native_summary[name] == pytest.approx(expected) assert report.metrics[name] == native_summary[name] From b32dd0d58cbfbd210a4aa37ba8384cb636f6897d Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Fri, 18 Sep 2026 14:36:46 -0700 Subject: [PATCH 2/8] fix: use consolidated package paths for V4.1 fixtures and notices --- THIRD_PARTY_NOTICES.md | 10 +++++----- crates/core/src/perfmodel/operators/dsv41.rs | 2 +- python/aisimulate/THIRD_PARTY_NOTICES.md | 10 +++++----- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 0fdce628f..d311585c9 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -440,8 +440,8 @@ SOFTWARE. ## SGLang DeepSeek-V4.1 serving contracts The V4.1 execution and performance-model logic in -`src/aiconfigurator_core/sdk/deepseek_v41.py`, `sdk/models/deepseek_v41.py` -(with the same `src/aiconfigurator_core/` prefix), and repository-root +`src/aisimulate_core/sdk/deepseek_v41.py`, `sdk/models/deepseek_v41.py` +(with the same `src/aisimulate_core/` prefix), and repository-root `crates/core/src/perfmodel/operators/dsv41.rs` is informed by and modified from SGLang's serving architecture at immutable commit `1aa0e962b206102b7c439a4a0c4981cfec6e87bc`: @@ -475,11 +475,11 @@ adaptations of, the named DeepSeek model repositories: | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3_config.json` | `deepseek-ai/DeepSeek-V3@e815299b0bcbac849fa540c768ef21845365c9eb` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3.2_config.json` | `deepseek-ai/DeepSeek-V3.2@c69397ecfd1fd142e90e3fbad51f4c7e40b9f3d3` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Flash_config.json` | `deepseek-ai/DeepSeek-V4-Flash@60d8d70770c6776ff598c94bb586a859a38244f1` | -| `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json` | `deepseek-ai/DeepSeek-V4.1-Flash@fb2764a5cf321eaa5070ca8f9e892818f477c16d` | +| `src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json` | `deepseek-ai/DeepSeek-V4.1-Flash@fb2764a5cf321eaa5070ca8f9e892818f477c16d` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Pro_config.json` | `deepseek-ai/DeepSeek-V4-Pro@b5968e9190ef611bbf34a7229255be88a0e937c1` | -The V4.1 descriptor and performance formulas in `src/aiconfigurator_core/sdk/deepseek_v41.py`, -`src/aiconfigurator_core/sdk/models/deepseek_v41.py`, and repository-root +The V4.1 descriptor and performance formulas in `src/aisimulate_core/sdk/deepseek_v41.py`, +`src/aisimulate_core/sdk/models/deepseek_v41.py`, and repository-root `crates/core/src/perfmodel/operators/dsv41.rs` are AISimulate performance-model adaptations of the architecture described by `inference/model.py` and `DeepSeek_V41_Tech_Report.pdf` at the same V4.1 revision (modified; no model execution code). diff --git a/crates/core/src/perfmodel/operators/dsv41.rs b/crates/core/src/perfmodel/operators/dsv41.rs index e2e20546b..e4eb2e82b 100644 --- a/crates/core/src/perfmodel/operators/dsv41.rs +++ b/crates/core/src/perfmodel/operators/dsv41.rs @@ -533,7 +533,7 @@ mod tests { fn test_db(mode: DatabaseMode) -> PerfDatabase { let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("../../python/aisimulate/src/aiconfigurator_core/systems"); + .join("../../python/aisimulate/src/aisimulate_core/systems"); PerfDatabase::load(&root, "gb300", "sglang", "0.5.14") .unwrap() .with_mode(mode, TransferPolicy::ALL) diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 0fdce628f..d311585c9 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -440,8 +440,8 @@ SOFTWARE. ## SGLang DeepSeek-V4.1 serving contracts The V4.1 execution and performance-model logic in -`src/aiconfigurator_core/sdk/deepseek_v41.py`, `sdk/models/deepseek_v41.py` -(with the same `src/aiconfigurator_core/` prefix), and repository-root +`src/aisimulate_core/sdk/deepseek_v41.py`, `sdk/models/deepseek_v41.py` +(with the same `src/aisimulate_core/` prefix), and repository-root `crates/core/src/perfmodel/operators/dsv41.rs` is informed by and modified from SGLang's serving architecture at immutable commit `1aa0e962b206102b7c439a4a0c4981cfec6e87bc`: @@ -475,11 +475,11 @@ adaptations of, the named DeepSeek model repositories: | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3_config.json` | `deepseek-ai/DeepSeek-V3@e815299b0bcbac849fa540c768ef21845365c9eb` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V3.2_config.json` | `deepseek-ai/DeepSeek-V3.2@c69397ecfd1fd142e90e3fbad51f4c7e40b9f3d3` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Flash_config.json` | `deepseek-ai/DeepSeek-V4-Flash@60d8d70770c6776ff598c94bb586a859a38244f1` | -| `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json` | `deepseek-ai/DeepSeek-V4.1-Flash@fb2764a5cf321eaa5070ca8f9e892818f477c16d` | +| `src/aisimulate_core/model_configs/deepseek-ai--DeepSeek-V4.1-Flash_config.json` | `deepseek-ai/DeepSeek-V4.1-Flash@fb2764a5cf321eaa5070ca8f9e892818f477c16d` | | `src/aiconfigurator_core/model_configs/deepseek-ai--DeepSeek-V4-Pro_config.json` | `deepseek-ai/DeepSeek-V4-Pro@b5968e9190ef611bbf34a7229255be88a0e937c1` | -The V4.1 descriptor and performance formulas in `src/aiconfigurator_core/sdk/deepseek_v41.py`, -`src/aiconfigurator_core/sdk/models/deepseek_v41.py`, and repository-root +The V4.1 descriptor and performance formulas in `src/aisimulate_core/sdk/deepseek_v41.py`, +`src/aisimulate_core/sdk/models/deepseek_v41.py`, and repository-root `crates/core/src/perfmodel/operators/dsv41.rs` are AISimulate performance-model adaptations of the architecture described by `inference/model.py` and `DeepSeek_V41_Tech_Report.pdf` at the same V4.1 revision (modified; no model execution code). From 3213af70d14968a191c1ad0f62b1b29c09fe0563 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Fri, 18 Sep 2026 14:44:10 -0700 Subject: [PATCH 3/8] fix: preserve canonical estimator policy in V4.1 replay --- python/aisimulate/src/aisimulate/compiler.py | 2 ++ .../aisimulate/src/aisimulate/config/engine.py | 1 - tests/test_runner.py | 17 ++++++++++++++++- tests/test_unified_traffic_runtime.py | 1 + 4 files changed, 19 insertions(+), 2 deletions(-) diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index 9d85114d2..04e71e78b 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -488,6 +488,8 @@ def _worker_engine_args( backend_version=engine.backend_version, worker_type=role, decoder_replay=engine.decoder_replay, + enable_shared_layer=engine.enable_shared_layer, + strict_provenance=bool(engine.strict_provenance), tp=parallel.tensor, pp=parallel.pipeline, attention_dp=parallel.attention_data, diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index d6ce2a944..9acf0597d 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -336,7 +336,6 @@ class EnginePredictionConfig(EstimatorPolicyConfig): backend: Backend = "vllm" backend_version: str | None = None decoder_replay: StrictBool = False - database_mode: Literal["SILICON", "SOL", "HYBRID", "EMPIRICAL"] | None = None enable_shared_layer: StrictBool | None = None strict_provenance: StrictBool | None = None context_length: PositiveInt | Literal["max"] = "max" diff --git a/tests/test_runner.py b/tests/test_runner.py index 371afdc58..4a7c079e5 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1285,7 +1285,22 @@ def test_runner_rejects_unknown_forward_model(value): @pytest.mark.parametrize("replay", [False, True]) -def test_public_replay_keeps_decoder_profile_and_database_policy(replay): +def test_public_replay_keeps_decoder_profile_and_database_policy(replay, monkeypatch): + from aisimulate_core.sdk.rust_engine_step import RustForwardPassPerfModel + + class ReadyEstimator: + def __init__(self, config): + self.config = config + + def diagnostics(self): + return {"readiness": "ready", "provenance": {"config": self.config}} + + def close(self): + pass + + # This exercises configuration transport with a recording runtime, not + # readiness or prediction for the deliberately synthetic example model. + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", ReadyEstimator) public = CorePredictionConfig.model_validate( { "engine": { diff --git a/tests/test_unified_traffic_runtime.py b/tests/test_unified_traffic_runtime.py index 79fff7f86..98537a2c5 100644 --- a/tests/test_unified_traffic_runtime.py +++ b/tests/test_unified_traffic_runtime.py @@ -132,6 +132,7 @@ def test_prediction_spec_separates_perf_identity_from_fixed_timing() -> None: "moe_ep_size": None, "nextn": None, "forward_model": "op_level", + "database_mode": "SILICON", }, } } From da60013b4a61d20ddd3a076f79822925d69f1c7b Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Fri, 18 Sep 2026 15:55:40 -0700 Subject: [PATCH 4/8] fix: preserve EPD default policy equivalence after V4.1 integration Signed-off-by: Harry Lee --- .../aisimulate/src/aisimulate/config/epd.py | 5 ++++- .../tests/unit/sdk/test_memory_estimation.py | 4 ++++ tests/test_epd_cli.py | 21 +++++++++++++++++++ 3 files changed, 29 insertions(+), 1 deletion(-) diff --git a/python/aisimulate/src/aisimulate/config/epd.py b/python/aisimulate/src/aisimulate/config/epd.py index 72ea3e488..ec6d2a2f1 100644 --- a/python/aisimulate/src/aisimulate/config/epd.py +++ b/python/aisimulate/src/aisimulate/config/epd.py @@ -86,7 +86,7 @@ def _language_execution(spec: ReplaySpec) -> dict: from ..capacity import DEFAULT_BACKEND_VERSIONS from ..runner import _materialize_engine_role - from .engine import SchedulerPredictionConfig + from .engine import EstimatorPolicyConfig, SchedulerPredictionConfig deployment = spec.backend_deployment roles = ( @@ -110,6 +110,9 @@ def _language_execution(spec: ReplaySpec) -> dict: timing.get("backend_version") or DEFAULT_BACKEND_VERSIONS[deployment.backend], ) timing.setdefault("cuda_graph_reserved_bytes", 0) + # Legacy Sweeper descriptors omit the default policy; the compiler + # serializes it explicitly. Compare their resolved execution meaning. + timing.setdefault("database_mode", EstimatorPolicyConfig().database_mode) # HandoffTransferTiming::delay_ms uses the same fallback for either mode # when a complete byte-count/bandwidth transfer model is unavailable. if rank.get("kv_transfer_bytes_per_token") is None or rank.get("kv_transfer_bandwidth") is None: diff --git a/python/aisimulate/tests/unit/sdk/test_memory_estimation.py b/python/aisimulate/tests/unit/sdk/test_memory_estimation.py index 20111aa8d..41a644650 100644 --- a/python/aisimulate/tests/unit/sdk/test_memory_estimation.py +++ b/python/aisimulate/tests/unit/sdk/test_memory_estimation.py @@ -286,6 +286,10 @@ def get_kvcache_bytes_per_sequence(self, seq_len): def get_kvcache_max_tokens(self, budget): return int(budget // 1024) + def get_kvcache_batch_capacity(self, budget, max_batch_size): + assert max_batch_size == 128 + return self.get_kvcache_max_tokens(budget) + class Backend: def _get_memory_usage(self, *args, num_tokens, **kwargs): return {"weights": 59, "activations": num_tokens / 1024, "others": 2, "nccl": 1} diff --git a/tests/test_epd_cli.py b/tests/test_epd_cli.py index 32cc6707b..18d054768 100644 --- a/tests/test_epd_cli.py +++ b/tests/test_epd_cli.py @@ -84,6 +84,27 @@ def _recommendation(mode="aggregated"): return raw +@pytest.mark.parametrize("mode", ["aggregated", "disaggregated"]) +def test_epd_language_policy_normalizes_only_equivalent_default(mode): + from aisimulate.config.epd import _language_execution + + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate(_prediction(mode))) + expected = _language_execution(spec) + legacy = deepcopy(spec) + deployment = legacy.backend_deployment + arguments = ( + [deployment.agg_engine_args] + if mode == "aggregated" + else [deployment.prefill_engine_args, deployment.decode_engine_args] + ) + for args in arguments: + assert args.pop("aic_database_mode") == "SILICON" + assert _language_execution(legacy) == expected + # A genuinely different estimator policy must still reject a callback. + arguments[0]["aic_database_mode"] = "SOL" + assert _language_execution(legacy) != expected + + @pytest.mark.parametrize("mode", ["aggregated", "disaggregated", "heterogeneous"]) @pytest.mark.parametrize("relative_stop", [False, True]) def test_native_cli_epd_recommend_yaml_predict(tmp_path, capsys, mode, relative_stop): From bc1ad8dfa5fc665f00941c8c8d38a435024cea25 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Fri, 18 Sep 2026 14:33:26 -0700 Subject: [PATCH 5/8] feat: preserve DeepSeek V4.1 FPM support and pinned HF datasets --- THIRD_PARTY_NOTICES.md | 13 + .../perfmodel/test_engine_step_parity.py | 4 +- crates/core/src/perfmodel/config.rs | 6 +- crates/core/src/perfmodel/engine/runtime.rs | 298 +++++++++- crates/core/src/perfmodel/engine/spec.rs | 25 + crates/core/src/perfmodel/fpm/config.rs | 43 ++ crates/core/src/perfmodel/fpm/tests.rs | 1 + crates/core/src/perfmodel/memory.rs | 1 + .../src/perfmodel/operators/fpm_forward.rs | 95 +++- .../core/src/perfmodel/operators/fpm_sol.rs | 64 +++ .../perfmodel/perf_database/fpm_forward.rs | 189 ++++++- crates/core/src/perfmodel/py.rs | 43 ++ crates/core/src/python.rs | 6 + .../core/tests/perfmodel/memory_round_trip.rs | 1 + python/aisimulate/THIRD_PARTY_NOTICES.md | 13 + .../collector/fpm_forward/capabilities.py | 33 ++ .../collector/fpm_forward/config.py | 129 ++++- .../collector/fpm_forward/database.py | 132 ++++- .../collector/fpm_forward/native_artifact.py | 130 +++++ .../collector/fpm_forward/planner.py | 47 +- .../collector/fpm_forward/runner.py | 501 +++++++++++++++-- .../fpm_forward/runtime/dsv41/LICENSE | 210 ++++++++ .../fpm_forward/runtime/dsv41/README.md | 216 ++++++++ .../runtime/dsv41/dsv41_scheduler.py | 507 ++++++++++++++++++ .../runtime/dsv41/runtime-paths.json | 1 + .../runtime/dsv41/runtime-source-sha256.json | 20 + .../runtime/dsv41/sitecustomize.py | 71 +++ .../collector/fpm_forward/runtime/fpm_exec.sh | 21 +- .../fpm_forward/runtime/fpm_text.txt | 9 + .../fpm_forward/runtime/preflight.py | 16 +- .../aisimulate/collector/fpm_forward/slurm.py | 266 +++++++++ python/aisimulate/docs/fpm/deepseek-v41.md | 174 ++++++ python/aisimulate/pyproject.toml | 4 + python/aisimulate/src/aisimulate/runner.py | 1 + .../src/aisimulate/sdk/fpm_identity.py | 8 + .../src/aisimulate_core/sdk/config.py | 4 + .../aisimulate_core/sdk/config_builders.py | 2 + .../src/aisimulate_core/sdk/engine.py | 5 + .../src/aisimulate_core/sdk/fpm_dataset.py | 210 ++++++++ .../src/aisimulate_core/sdk/fpm_identity.py | 48 ++ .../aisimulate_core/sdk/models/__init__.py | 19 +- .../sdk/operations/fpm_forward.py | 16 +- .../aisimulate_core/sdk/rust_engine_step.py | 2 + .../aisimulate_core/systems/dsv41_fpm_hf.json | 181 +++++++ .../cross_package/test_core_public_api.py | 1 + .../cross_package/test_import_contract.py | 1 + .../fixtures/dsv41_producer_lifecycle.py | 490 +++++++++++++++++ .../unit/collector/test_fpm_dsv41_producer.py | 143 +++++ .../tests/unit/collector/test_fpm_exec.py | 61 ++- .../collector/test_fpm_explicit_points.py | 234 ++++++++ .../tests/unit/collector/test_fpm_forward.py | 237 +++++++- .../tests/unit/collector/test_fpm_runner.py | 343 +++++++++++- .../collector/test_fpm_runtime_wrapper.py | 26 +- .../tests/unit/collector/test_fpm_slurm.py | 479 +++++++++++++++++ .../unit/sdk/models/test_deepseek_v41.py | 5 +- .../tests/unit/sdk/test_fpm_dataset.py | 199 +++++++ .../unit/sdk/test_fpm_execution_identity.py | 366 +++++++++++++ .../tests/unit/sdk/test_fpm_forward.py | 25 + 58 files changed, 6250 insertions(+), 145 deletions(-) create mode 100644 python/aisimulate/collector/fpm_forward/runtime/dsv41/LICENSE create mode 100644 python/aisimulate/collector/fpm_forward/runtime/dsv41/README.md create mode 100644 python/aisimulate/collector/fpm_forward/runtime/dsv41/dsv41_scheduler.py create mode 100644 python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-paths.json create mode 100644 python/aisimulate/collector/fpm_forward/runtime/dsv41/runtime-source-sha256.json create mode 100644 python/aisimulate/collector/fpm_forward/runtime/dsv41/sitecustomize.py create mode 100644 python/aisimulate/collector/fpm_forward/runtime/fpm_text.txt create mode 100644 python/aisimulate/collector/fpm_forward/slurm.py create mode 100644 python/aisimulate/docs/fpm/deepseek-v41.md create mode 100644 python/aisimulate/src/aisimulate/sdk/fpm_identity.py create mode 100644 python/aisimulate/src/aisimulate_core/sdk/fpm_dataset.py create mode 100644 python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py create mode 100644 python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json create mode 100644 python/aisimulate/tests/unit/collector/fixtures/dsv41_producer_lifecycle.py create mode 100644 python/aisimulate/tests/unit/collector/test_fpm_dsv41_producer.py create mode 100644 python/aisimulate/tests/unit/collector/test_fpm_explicit_points.py create mode 100644 python/aisimulate/tests/unit/collector/test_fpm_slurm.py create mode 100644 python/aisimulate/tests/unit/sdk/test_fpm_dataset.py create mode 100644 python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index d311585c9..5c44e8710 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -64,6 +64,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/`. 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/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index 7cbb47c04..b9dabfc58 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -82,7 +82,8 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // TokenScale was appended to remap draft query widths before op lookup. // - 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. +// FpmForwardOp also carries original_fmha_quant_mode for table-selection +// diagnostics. Both append positional fields; defaults support old JSON only. pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 19; /// Static engine identity and setup information carried by an @@ -226,6 +227,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 af80434fb..4d21216a0 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 { @@ -1867,6 +1874,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 @@ -1933,26 +1969,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; @@ -2155,6 +2175,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) }), @@ -2986,6 +3007,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, sol_ops: vec![], }) }; @@ -3013,6 +3035,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( @@ -3070,6 +3093,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, sol_ops: vec![], }) }; @@ -3081,6 +3105,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 { @@ -3379,6 +3638,7 @@ mod tests { match_identity: default_identity(4), weight_bytes: 0.0, verify_width: 1, + original_fmha_quant_mode: None, sol_ops: vec![], })], vec![], @@ -3696,6 +3956,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 ... @@ -3741,6 +4002,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 3ac40705c..5795eb694 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()), @@ -838,6 +839,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) }), @@ -1268,4 +1270,27 @@ mod tests { }) )); } + #[test] + fn fpm_selector_json_default_and_schema18_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(&18u32.to_le_bytes()); + bytes.truncate(4); + assert!(matches!( + EngineSpec::from_bincode(&bytes), + Err(AicError::UnsupportedSchemaVersion { + got: 18, + expected: ENGINE_SPEC_SCHEMA_VERSION, + + .. + }) + )); + } } diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 334a870be..0699bbce5 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")] @@ -163,6 +166,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, @@ -269,6 +273,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() { @@ -307,6 +323,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 auto_selects_all_modes_even_when_fallback_is_denied() { 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 0953bb3d6..87c7b29e6 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -121,6 +121,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 50cb44766..50d764783 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -521,6 +521,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/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 ac3548c37..91e7201e3 100644 --- a/crates/core/src/perfmodel/perf_database/fpm_forward.rs +++ b/crates/core/src/perfmodel/perf_database/fpm_forward.rs @@ -19,7 +19,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 @@ -39,9 +39,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}; @@ -51,7 +51,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 @@ -88,7 +88,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", @@ -104,15 +104,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, @@ -143,6 +148,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 @@ -272,7 +295,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 @@ -334,7 +365,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: {}. \ @@ -365,7 +396,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"), @@ -415,7 +447,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 @@ -473,7 +508,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)?; @@ -499,6 +534,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", @@ -655,6 +695,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) } @@ -692,6 +737,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(), @@ -704,6 +753,14 @@ 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].is_empty() + && (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; @@ -852,6 +909,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) { @@ -1104,6 +1162,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 { @@ -1126,6 +1185,7 @@ pub(crate) mod tests { system: "b200_sxm", backend: "vllm", kv_seed_regime: None, + execution: None, } } } @@ -1144,7 +1204,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![ @@ -1163,10 +1223,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, |_| {}) } @@ -1186,6 +1250,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); @@ -1223,6 +1288,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 = @@ -1302,6 +1379,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"); @@ -1365,7 +1454,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(), @@ -1411,6 +1504,72 @@ 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 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 d81f0654f..fd7a577b0 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, nextn: u32, @@ -1091,6 +1092,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, nextn: 0, @@ -1193,6 +1195,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()); @@ -1312,6 +1320,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, @@ -1324,6 +1334,8 @@ mod builder_tests { let request = engine_build_request(&config, None); 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(), @@ -1344,6 +1356,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. @@ -1359,6 +1380,12 @@ 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(), + )); + } + 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" @@ -1386,6 +1413,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)) } @@ -1529,6 +1568,8 @@ fn engine_build_request(config: &EngineConfig, systems_path: Option<&str>) -> En .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. @@ -1716,6 +1757,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, @@ -1956,6 +1998,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 7f798888e..de215704e 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -169,6 +169,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")] @@ -268,6 +270,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, @@ -2073,6 +2076,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, @@ -2556,6 +2560,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(); @@ -2563,6 +2568,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 2e9abb6e2..60dbb1c2b 100644 --- a/crates/core/tests/perfmodel/memory_round_trip.rs +++ b/crates/core/tests/perfmodel/memory_round_trip.rs @@ -100,6 +100,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/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index d311585c9..5c44e8710 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -64,6 +64,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/`. 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..dcd9f38a6 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,31 @@ 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 transport content; native Dynamo owns point/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") + 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 +229,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 +283,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 +328,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 +366,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 +389,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 +569,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 +653,12 @@ 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_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..03f8ee107 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,120 @@ 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"] + 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) + 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,10 +368,22 @@ 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]] = [] for path, payload in rank_payloads: + 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}") if ( payload.get("schema_version") != FPM_NATIVE_BENCHMARK_RESULT_SCHEMA_VERSION or payload.get("artifact_type") != "rank" @@ -439,4 +568,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..04b19d7c0 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,174 @@ 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: + 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. Still attempt KILL, + # report the observation, and continue cleaning other groups. + 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 +400,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 +866,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 +880,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 +1045,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 +1131,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 +1164,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 +1183,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 +1218,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 +1237,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 +1299,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 +1374,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 +1397,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 +1472,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 +1565,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 +1652,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 +1684,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 +1808,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 +1842,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 +1883,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 +1932,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 +1996,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 +2023,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 +2043,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 +2167,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 +2197,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 +2213,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..1f7c75425 --- /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_results.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..dee620aa4 --- /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=config.model_config.engram_config.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..cb208a297 --- /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 757ff1f99f716a1d21985723ae534cef4afbb5ab](https://huggingface.co/datasets/nvidia/aisimulate-fpm-dataset/tree/757ff1f99f716a1d21985723ae534cef4afbb5ab). +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 eccb06d44..6d4407870 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/runner.py b/python/aisimulate/src/aisimulate/runner.py index 0b235c535..d225c04bf 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -85,6 +85,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",), 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 ec4dbe909..e893126cb 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py @@ -39,6 +39,7 @@ def build_model_config( forward_model: str | None = None, enable_encoder_dp: bool = True, attention_backend: str | None = None, + fpm_fmha_quant_mode: str | None = None, speculation=None, ) -> ModelConfig: """Build a ModelConfig with optional quant mode overrides.""" @@ -51,6 +52,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/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index ff3c6e7e3..1ff67a7ff 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -161,6 +161,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 @@ -319,6 +320,7 @@ def _engine_config_dict( # Rust side reloads the perf database from this string verbatim. "backend_version": _literal_backend_version(system, backend, backend_version, systems_path, database), "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), @@ -331,6 +333,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 @@ -416,6 +419,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, nextn: int = 0, @@ -454,6 +458,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/operations/fpm_forward.py b/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py index 4b0a621e3..d454812b2 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py +++ b/python/aisimulate/src/aisimulate_core/sdk/operations/fpm_forward.py @@ -19,6 +19,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 @@ -29,6 +36,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") @@ -38,7 +46,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", @@ -60,6 +68,7 @@ "attention_backend", "enable_wideep", "enable_eplb", + *EXECUTION_COLUMNS, ) @@ -125,10 +134,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), @@ -144,6 +155,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 2667f49d7..5acf3d9ee 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 = None kvcache_quant_mode: str | None = None comm_quant_mode: str | None = None nextn: int = 0 @@ -1160,6 +1161,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..49bba04ff --- /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": "757ff1f99f716a1d21985723ae534cef4afbb5ab", + "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": "bd1ac729bedbec8c2d1cdf9594b362662e808010428345990b218e5c55dc06d1" + }, + { + "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": "a888f6e14bfea669dd514bb919d281af33cdf75b68f912f3321d898081d038aa" + }, + { + "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": "821121b14b92bc4ccaa1a16bb8b0efeb1c5f03b7509cab8527a1b72637bdc396" + }, + { + "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 5998e4cda..114d26d69 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, " "nextn: 'int' = 0, " "speculation: 'dict | None' = None, " diff --git a/python/aisimulate/tests/cross_package/test_import_contract.py b/python/aisimulate/tests/cross_package/test_import_contract.py index 56b2f16d2..c77d5465a 100644 --- a/python/aisimulate/tests/cross_package/test_import_contract.py +++ b/python/aisimulate/tests/cross_package/test_import_contract.py @@ -31,6 +31,7 @@ "engine", "engine_table_view", "errors", + "fpm_identity", "inference_summary", "memory", "models.base", 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..8c3d73394 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/fixtures/dsv41_producer_lifecycle.py @@ -0,0 +1,490 @@ +# 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 + + +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_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 = ("model_config_sha256", "execution_profile", "engram_residency", "input_modality") + 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..70b6bddf2 --- /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") + subprocess.run( + [sys.executable, str(Path(__file__).with_name("fixtures") / "dsv41_producer_lifecycle.py")], + env=environment, + check=True, + timeout=30, + capture_output=True, + text=True, + ) + + +@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..faadac49f 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,53 @@ 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), (3, 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() + + def delayed_leader(): + 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..562009407 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_explicit_points.py @@ -0,0 +1,234 @@ +# 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, {}) diff --git a/python/aisimulate/tests/unit/collector/test_fpm_forward.py b/python/aisimulate/tests/unit/collector/test_fpm_forward.py index bc0ac9eb3..6ff153904 100644 --- a/python/aisimulate/tests/unit/collector/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/collector/test_fpm_forward.py @@ -233,12 +233,18 @@ 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}], "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 +270,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 +279,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 @@ -1157,6 +1165,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", @@ -1327,7 +1336,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") @@ -1335,7 +1344,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"] @@ -1375,14 +1384,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) @@ -1756,7 +1765,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'\]"): @@ -2373,3 +2382,215 @@ 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 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..bb99234ef --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_fpm_slurm.py @@ -0,0 +1,479 @@ +# 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, "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, "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) diff --git a/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py b/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py index b06c5acf2..78361357e 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py +++ b/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py @@ -466,7 +466,8 @@ def test_afd_rejects_v41_before_search_or_session_construction(decoder_replay): @pytest.mark.parametrize("is_context", [False, True]) @pytest.mark.parametrize("batch,seq", [(1, 4), (0, 4), (1, 0)]) -def test_native_attention_rejects_unknown_role_before_zero_work(is_context, batch, seq): +@pytest.mark.parametrize("database_mode", ["SOL", "SILICON", "HYBRID", "EMPIRICAL"]) +def test_native_attention_rejects_unknown_role_before_zero_work(is_context, batch, seq, database_mode): import json import aisimulate_core._native as native @@ -483,6 +484,6 @@ def test_native_attention_rejects_unknown_role_before_zero_work(is_context, batc if "Dsv41Attention" in child ) op = native.op_from_spec_json(json.dumps({"Dsv41Attention": attention | {"role": "ful", "is_context": is_context}})) - db = get_database_view("gb300", "sglang", "current", allow_missing_data=True, database_mode="SOL") + db = get_database_view("gb300", "sglang", "current", allow_missing_data=True, database_mode=database_mode) with pytest.raises(ValueError, match="attention role must be"): _evaluate_single_op(db, op, is_context=is_context, batch_size=batch, s=seq, prefix=0, x=batch * seq) 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..b8e319f8e --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py @@ -0,0 +1,366 @@ +# 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 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 = ("model_config_sha256", "execution_profile", "engram_residency", "input_modality") + 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 = ("model_config_sha256", "execution_profile", "engram_residency", "input_modality") + 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) diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index bf3374e2d..26e752b1c 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -603,3 +603,28 @@ def test_genonly_step_works_with_single_op(self, fpm_session): ) assert per_op["fpm_forward_decode"] == pytest.approx(7.0) assert total == pytest.approx(7.0) + + +def test_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=baseline.config.fmha_quant_mode), + BACKEND, + ) + original = baseline.config.fmha_quant_mode.name + 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] == 22.0 + assert "WARNING: FPM table FMHA selector" in warning + assert f'original_model_mode="{original}"' in warning + assert f'selector="{original}"' in warning + assert "matched_cell_ids=" in warning + assert "fpm-test-prefill" in warning and "fpm-test-decode" 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 From 25a0e3c96414754fb387fadc0e6af27333e5a36b Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Fri, 18 Sep 2026 14:50:51 -0700 Subject: [PATCH 6/8] Expose pinned FPM dataset staging through public SDK facade --- python/aisimulate/src/aisimulate/sdk/fpm_dataset.py | 8 ++++++++ .../tests/cross_package/test_import_contract.py | 1 + 2 files changed, 9 insertions(+) create mode 100644 python/aisimulate/src/aisimulate/sdk/fpm_dataset.py 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/tests/cross_package/test_import_contract.py b/python/aisimulate/tests/cross_package/test_import_contract.py index c77d5465a..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,7 @@ "engine", "engine_table_view", "errors", + "fpm_dataset", "fpm_identity", "inference_summary", "memory", From 438d735f7648e0be9b197062fa31996104effc71 Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Fri, 18 Sep 2026 14:52:25 -0700 Subject: [PATCH 7/8] Pin reviewed HF archive provenance without changing observed data --- python/aisimulate/docs/fpm/deepseek-v41.md | 2 +- .../src/aisimulate_core/systems/dsv41_fpm_hf.json | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/python/aisimulate/docs/fpm/deepseek-v41.md b/python/aisimulate/docs/fpm/deepseek-v41.md index cb208a297..33677e7dd 100644 --- a/python/aisimulate/docs/fpm/deepseek-v41.md +++ b/python/aisimulate/docs/fpm/deepseek-v41.md @@ -142,7 +142,7 @@ print(model.estimate_forward_pass_time_ms({ 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 757ff1f99f716a1d21985723ae534cef4afbb5ab](https://huggingface.co/datasets/nvidia/aisimulate-fpm-dataset/tree/757ff1f99f716a1d21985723ae534cef4afbb5ab). +[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. diff --git a/python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json b/python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json index 49bba04ff..514ce42f5 100644 --- a/python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json +++ b/python/aisimulate/src/aisimulate_core/systems/dsv41_fpm_hf.json @@ -1,7 +1,7 @@ { "format_version": 1, "repo_id": "nvidia/aisimulate-fpm-dataset", - "revision": "757ff1f99f716a1d21985723ae534cef4afbb5ab", + "revision": "b35883ee5f8b4a82e24844a7e872056aff64287f", "profiles": { "gb200-tp4-full": { "admission": "quarantined", @@ -42,7 +42,7 @@ { "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": "bd1ac729bedbec8c2d1cdf9594b362662e808010428345990b218e5c55dc06d1" + "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", @@ -100,7 +100,7 @@ { "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": "a888f6e14bfea669dd514bb919d281af33cdf75b68f912f3321d898081d038aa" + "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", @@ -158,7 +158,7 @@ { "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": "821121b14b92bc4ccaa1a16bb8b0efeb1c5f03b7509cab8527a1b72637bdc396" + "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", From 34a72058ce9366cc3548789786da20dd4d1e2832 Mon Sep 17 00:00:00 2001 From: Yiming Liu Date: Mon, 21 Sep 2026 18:02:14 +0800 Subject: [PATCH 8/8] fix: preserve packaged FPM selector in AFD companions Signed-off-by: Yiming Liu --- python/aisimulate/src/aisimulate/runner.py | 14 +++-- .../tests/unit/sdk/test_fpm_forward.py | 52 ++++++++++++++++++- tests/test_afd_runner.py | 38 ++++++++++++++ 3 files changed, 98 insertions(+), 6 deletions(-) diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index 29f704c31..337978a78 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -248,13 +248,14 @@ 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"), @@ -281,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, **{ @@ -331,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/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index fa363ed8e..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 @@ -422,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]) 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