diff --git a/crates/core/src/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index d73f824a9..840b3f130 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -108,7 +108,10 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // - 20 (DeepSeek-V4.1 FPM): FpmForwardOp gained original_fmha_quant_mode // for selector diagnostics. This appends a positional field after the schema-19 // release; serde defaults support legacy JSON, not legacy bincode. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 20; +// - 21 (AIC-1781): EngineConfig and MoeOp gained exact `moe_kernel_source` +// identity. Renumbered from the branch's concurrent v20 claim after the +// DeepSeek-V4.1 FPM layout landed first. +pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 21; /// Static engine identity and setup information carried by an /// [`crate::perfmodel::engine::spec::EngineSpec`]. @@ -152,6 +155,10 @@ pub struct EngineConfig { /// Use the backend-verified bounded DeepSeek-V4.1 decoder execution profile. #[serde(default)] pub decoder_replay: bool, + /// Exact collected MoE compute kernel-source lane. Unlike + /// `moe_backend`, this selects one measured MoE table lane. + #[serde(default)] + pub moe_kernel_source: Option, // KV pub kv_block_size: Option, diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index ea894c537..ca67b169c 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -2323,6 +2323,7 @@ mod tests { forward_model: None, fpm_parquet_path: None, decoder_replay: false, + moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { tp_size: 8, diff --git a/crates/core/src/perfmodel/engine/spec.rs b/crates/core/src/perfmodel/engine/spec.rs index 856a6a44b..17ece3757 100644 --- a/crates/core/src/perfmodel/engine/spec.rs +++ b/crates/core/src/perfmodel/engine/spec.rs @@ -319,6 +319,7 @@ mod tests { workload_distribution: "power_law_1.2".into(), is_gated: true, moe_backend: None, + moe_kernel_source: Some("sglang_flashinfer_trtllm_moe".into()), enable_eplb: false, is_context: false, } @@ -827,6 +828,7 @@ mod tests { forward_model: None, fpm_parquet_path: None, decoder_replay: false, + moe_kernel_source: None, kv_block_size: Some(64), parallel: ParallelMapping { tp_size: 8, @@ -986,6 +988,32 @@ mod tests { assert_eq!(spec, decoded); } + #[test] + fn moe_and_fpm_fields_round_trip_and_reject_concurrent_schema20() { + // These independent positional fields were both introduced as schema20 + // on separate branches. The merged wire layout needs its own version. + let spec = EngineSpec::new( + sample_engine_config(), + vec![OpSpec::FpmForward(fpm_forward())], + vec![OpSpec::Moe(moe())], + ); + assert_eq!(spec.schema_version, 21); + let mut bytes = spec.to_bincode().unwrap(); + assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); + + bytes[..4].copy_from_slice(&20u32.to_le_bytes()); + // Reject the stale version before even trying to decode its payload. + bytes.truncate(4); + assert!(matches!( + EngineSpec::from_bincode(&bytes), + Err(AicError::UnsupportedSchemaVersion { + kind: "EngineSpec", + got: 20, + expected: ENGINE_SPEC_SCHEMA_VERSION, + }) + )); + } + /// A version skew combined with an op-layout change must surface as a clear /// [`AicError::UnsupportedSchemaVersion`], NOT a generic bincode I/O error. /// diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 3b802c9b8..98008a3d0 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -139,6 +139,10 @@ pub struct ForwardPassPerfModelConfig { pub estimator_config: EstimatorConfig, #[serde(default)] pub attention_backend: Option, + /// Exact collected MoE compute kernel-source lane. This is separate from + /// `moe_backend`, which describes topology/runtime behavior. + #[serde(default)] + pub moe_kernel_source: Option, #[serde(default)] pub moe_backend: Option, #[serde(default)] @@ -189,6 +193,7 @@ impl ForwardPassPerfModelConfig { moe_backend: None, enable_eplb: false, wideep_num_slots: None, + moe_kernel_source: None, enable_shared_layer: None, strict_provenance: false, } @@ -231,6 +236,13 @@ impl ForwardPassPerfModelConfig { { return Err(invalid_config("backend_version cannot be empty")); } + if self + .moe_kernel_source + .as_ref() + .is_some_and(|value| value.trim().is_empty()) + { + return Err(invalid_config("moe_kernel_source cannot be empty")); + } if self.tp == 0 || self.pp == 0 || self.attention_dp == 0 { return Err(invalid_config("tp, pp, and attention_dp must be positive")); } @@ -268,6 +280,7 @@ impl ForwardPassPerfModelConfig { } if self.estimation_mode == EstimationMode::FpmInterpolation && (self.enable_eplb + || self.moe_kernel_source.is_some() || self.wideep_num_slots.is_some() || self .moe_backend @@ -275,7 +288,7 @@ impl ForwardPassPerfModelConfig { .is_some_and(|value| value != "default")) { return Err(invalid_config( - "fpm_interpolation does not support EPLB, slots or moe_backend overrides", + "fpm_interpolation does not support EPLB, slots, moe_backend or moe_kernel_source overrides", )); } if self.nextn > 5 { @@ -416,6 +429,64 @@ mod tests { cfg.validate().unwrap(); } + #[test] + fn moe_kernel_source_round_trips_without_pinning_the_default() { + let default = config(serde_json::json!({})); + assert_eq!(default.moe_kernel_source, None); + + let pinned = config(serde_json::json!({ + "moe_kernel_source": "sglang_flashinfer_trtllm_moe" + })); + let round_trip: ForwardPassPerfModelConfig = + serde_json::from_str(&serde_json::to_string(&pinned).unwrap()).unwrap(); + assert_eq!(round_trip, pinned); + } + + #[test] + fn interpolation_rejects_exact_moe_sources_before_fallback() { + for fallback in ["deny", "allow"] { + for fpm_selector in [None, Some("fp8")] { + let cfg = config(serde_json::json!({ + "estimation_mode": "fpm_interpolation", + "fallback_policy": fallback, + "moe_kernel_source": "source_that_does_not_exist", + "fpm_fmha_quant_mode": fpm_selector + })); + let error = crate::ForwardPassPerfModel::best_available(cfg).unwrap_err(); + assert!(error.to_string().contains("moe_kernel_source overrides")); + } + } + config(serde_json::json!({"estimation_mode": "fpm_interpolation"})) + .validate() + .unwrap(); + config(serde_json::json!({"moe_kernel_source": "sglang_flashinfer_trtllm_moe"})) + .validate() + .unwrap(); + } + + #[test] + fn blank_moe_kernel_source_is_rejected_before_estimator_selection() { + for source in ["", " \t"] { + let error = config(serde_json::json!({"moe_kernel_source": source})) + .validate() + .unwrap_err(); + assert_eq!( + error.to_string(), + "invalid engine config: invalid forward pass perf model config: moe_kernel_source cannot be empty" + ); + } + + for estimation_mode in [EstimationMode::Auto, EstimationMode::FpmRegression] { + let mut cfg = config(serde_json::json!({"moe_kernel_source": " \t"})); + cfg.estimation_mode = estimation_mode; + let error = crate::ForwardPassPerfModel::best_available(cfg).unwrap_err(); + assert_eq!( + error.to_string(), + "invalid engine config: invalid forward pass perf model config: moe_kernel_source cannot be empty" + ); + } + } + #[test] fn prompt_lookup_is_distinct_from_mtp_and_validated_before_selection() { let mut cfg = config(serde_json::json!({ diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index a4dee04a0..ecfbc4ecc 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -705,6 +705,7 @@ fn build_native_candidate( ) -> Result<(Engine, PathBuf), AicError> { if config.estimation_mode == EstimationMode::FpmInterpolation && (config.enable_eplb + || config.moe_kernel_source.is_some() || config.wideep_num_slots.is_some() || config .moe_backend @@ -712,7 +713,7 @@ fn build_native_candidate( .is_some_and(|value| value != "default")) { return Err(AicError::UnsupportedModel( - "FPM interpolation does not support EPLB, slots or moe_backend overrides".into(), + "FPM interpolation does not support EPLB, slots, moe_backend or moe_kernel_source overrides".into(), )); } if config.estimation_mode == EstimationMode::FpmInterpolation && config.nextn != 0 { @@ -1115,6 +1116,21 @@ fn can_fallback_to_regression(err: &AicError) -> bool { mod fallback_errors { use super::*; + #[cfg(feature = "python")] + #[test] + fn interpolation_candidate_does_not_discard_exact_moe_source() { + let mut config = ForwardPassPerfModelConfig::new( + "model", + "system", + crate::BackendKind::Sglang, + crate::ForwardPassWorkerType::Decode, + ); + config.moe_kernel_source = Some("source_that_does_not_exist".into()); + config.estimation_mode = EstimationMode::FpmInterpolation; + let error = build_native_candidate(&config).err().unwrap(); + assert!(error.to_string().contains("moe_kernel_source overrides")); + } + #[test] fn corruption_is_not_a_coverage_gap() { let yaml = serde_yaml::from_str::("broken: [").unwrap_err(); diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index e75d36ba0..ec631752e 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -109,6 +109,7 @@ fn fixture_engine_config() -> EngineConfig { forward_model: None, fpm_parquet_path: None, decoder_replay: false, + moe_kernel_source: None, 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 820973e23..e63a89da0 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -509,6 +509,7 @@ mod tests { forward_model: None, fpm_parquet_path: None, decoder_replay: false, + moe_kernel_source: None, 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 604e77ace..cf13f5cab 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 20). + /// payloads are positional (current ENGINE_SPEC_SCHEMA_VERSION 21). #[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 20). +/// field — introduced at schema v8, current ENGINE_SPEC_SCHEMA_VERSION 21). /// 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 20). + /// bincode payloads are positional, current ENGINE_SPEC_SCHEMA_VERSION 21). #[serde(default = "default_lane_order")] pub lane_order: Vec, /// Per-head RMSNorm on Q and K before decode attention. Appended at the diff --git a/crates/core/src/perfmodel/operators/fpm_sol.rs b/crates/core/src/perfmodel/operators/fpm_sol.rs index 2c84c2b5a..5091f9fe8 100644 --- a/crates/core/src/perfmodel/operators/fpm_sol.rs +++ b/crates/core/src/perfmodel/operators/fpm_sol.rs @@ -895,6 +895,7 @@ mod tests { workload_distribution: "uniform".into(), is_gated: true, moe_backend: None, + moe_kernel_source: None, enable_eplb: false, is_context: true, }; diff --git a/crates/core/src/perfmodel/operators/moe.rs b/crates/core/src/perfmodel/operators/moe.rs index 35dbef00f..89a68c838 100644 --- a/crates/core/src/perfmodel/operators/moe.rs +++ b/crates/core/src/perfmodel/operators/moe.rs @@ -44,7 +44,7 @@ use crate::operators::base::{PerformanceResult, SolComponents, Source}; use crate::operators::util_empirical::{self, UtilGrid}; use crate::perf_database::PerfDatabase; use crate::perf_database::gemm::quant_tc_flops; -use crate::perf_database::moe::{MoeKernel, MoeSiblingSlice}; +use crate::perf_database::moe::MoeSiblingSlice; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -144,30 +144,24 @@ pub(crate) fn policy_fingerprint(policy: TransferPolicy) -> String { /// to the wideep context/generation MoE tables retired with AIC-1601 — see /// the typed error raised at each selector point.) #[derive(Clone, Copy, PartialEq)] -enum MoeTableSel { +enum MoeTableSel<'a> { Standard, LowLatency, + ExactKernelSource(&'a str), } -impl MoeTableSel { - /// Grid cache-key tag. Python folds `kernel_tag` ("std" / "ll") plus - /// `id(node)` into the key. - fn tag(self) -> &'static str { +impl MoeTableSel<'_> { + /// Grid cache-key tag. Exact source selection must be included so a grid + /// built from one kernel lane can never answer for another. + fn tag(self) -> String { match self { - Self::Standard => "std", - Self::LowLatency => "ll", + Self::Standard => "std".to_owned(), + Self::LowLatency => "ll".to_owned(), + Self::ExactKernelSource(source) => format!("kernel_source={source}"), } } } -/// The perf-DB kernel grid behind a selector. -fn moe_kernel(table: MoeTableSel) -> MoeKernel { - match table { - MoeTableSel::Standard => MoeKernel::Standard, - MoeTableSel::LowLatency => MoeKernel::LowLatency, - } -} - /// A sibling slice the transfer ladder may borrow: the reference slice's /// shape + token curve, the quant whose SOL reconstructs its util /// (`sol_quant`: QUERY quant for same-profile tiers, REFERENCE quant for @@ -214,6 +208,11 @@ pub struct MoeOp { /// pre-existing specs -> None -> the regular table. #[serde(default)] pub moe_backend: Option, + /// Exact collected MoE kernel-source lane. `None` preserves the + /// framework/default lookup; `Some` selects only rows carrying this exact + /// `kernel_source` value and fails on a missing or ineligible lane. + #[serde(default)] + pub moe_kernel_source: Option, /// EPLB enabled (Python `MoE._enable_eplb`). On the sglang branch the /// PREFILL token count is corrected to `int(num_tokens * 0.8)` /// (operations/moe.py: expert-parallel load balancing evens the @@ -266,6 +265,7 @@ impl MoeOp { workload_distribution: workload_distribution.into(), is_gated: true, moe_backend: None, + moe_kernel_source: None, enable_eplb: false, is_context: false, } @@ -317,6 +317,21 @@ impl MoeOp { } } + fn validate_kernel_source(&self, num_tokens: u32) -> Result<(), AicError> { + if self.moe_kernel_source.as_deref() == Some("moe_torch_flow_min_latency") + && (!self.is_gated || self.quant_mode != MoeQuantMode::Nvfp4 || num_tokens > 128) + { + return Err(AicError::InvalidEngineConfig(format!( + "MoE kernel source moe_torch_flow_min_latency requires gated nvfp4 with at most \ + 128 tokens; op {} has is_gated={}, quant={}, num_tokens={num_tokens}", + self.name, + self.is_gated, + self.quant_mode.name(), + ))); + } + Ok(()) + } + /// SILICON resolution (retired-deepep gate + low-latency probe + the default /// grid, scale/clamp applied per branch — the audit-PR body, unchanged). pub(crate) fn silicon_pr( @@ -324,6 +339,7 @@ impl MoeOp { db: &PerfDatabase, num_tokens: u32, ) -> Result { + self.validate_kernel_source(num_tokens)?; let is_sglang = db.backend == "sglang"; // SGLang EPLB prefill correction — INSIDE get_silicon only (Python // operations/moe.py:684, sglang branch: `num_tokens_corrected = @@ -361,6 +377,29 @@ impl MoeOp { ))); } + if let Some(kernel_source) = self.moe_kernel_source.as_deref() { + let value = db.moe.query_kernel_source( + kernel_source, + num_tokens, + self.hidden_size, + self.inter_size, + self.topk, + self.num_experts, + self.moe_tp_size, + self.moe_ep_size, + self.quant_mode, + &self.workload_distribution, + &sol, + )?; + return Ok(PerformanceResult::with_energy( + value.latency, + value.energy, + Source::Silicon, + ) + .clamp_non_negative() + .scaled(self.scale_factor)); + } + // Mirrors Python's MoE._query_moe_table TRT-LLM gate: for nvfp4 // gated MoE at num_tokens <= 128, probe the // `moe_torch_flow_min_latency` grid first and fall back to the @@ -414,6 +453,7 @@ impl MoeOp { /// `MoE._query_moe_table::get_empirical` (`operations/moe.py:327-572`): /// own-shape grid → xshape → xquant → xprofile → typed empirical miss. fn empirical_latency(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + self.validate_kernel_source(num_tokens)?; let spec = &db.system_spec; let quant = self.quant_mode; let num_gemms: u64 = if self.is_gated { 3 } else { 2 }; @@ -449,18 +489,20 @@ impl MoeOp { self.name ))); } - let table = if num_tokens <= 128 - && quant == MoeQuantMode::Nvfp4 - && self.is_gated - && db.moe.low_latency_available()? - { - match self.slice_points(db, MoeTableSel::LowLatency) { - Ok(_) => MoeTableSel::LowLatency, - Err(err) if err.is_missing_perf_data() => MoeTableSel::Standard, - Err(err) => return Err(err), + let table = match self.moe_kernel_source.as_deref() { + Some(kernel_source) => MoeTableSel::ExactKernelSource(kernel_source), + None if num_tokens <= 128 + && quant == MoeQuantMode::Nvfp4 + && self.is_gated + && db.moe.low_latency_available()? => + { + match self.slice_points(db, MoeTableSel::LowLatency) { + Ok(_) => MoeTableSel::LowLatency, + Err(err) if err.is_missing_perf_data() => MoeTableSel::Standard, + Err(err) => return Err(err), + } } - } else { - MoeTableSel::Standard + None => MoeTableSel::Standard, }; let kernel_tag = table.tag(); @@ -535,7 +577,7 @@ impl MoeOp { } } if let Some(reference) = - self.reference_grid(db, kernel_tag, num_gemms, policy, &candidates)? + self.reference_grid(db, &kernel_tag, num_gemms, policy, &candidates)? { grid = Some(reference); } @@ -563,7 +605,7 @@ impl MoeOp { &mut candidates, )?; if let Some(reference) = - self.reference_grid(db, kernel_tag, num_gemms, policy, &candidates)? + self.reference_grid(db, &kernel_tag, num_gemms, policy, &candidates)? { if !reference.is_empty() { grid = Some(reference); @@ -594,19 +636,43 @@ impl MoeOp { fn slice_points( &self, db: &PerfDatabase, - table: MoeTableSel, + table: MoeTableSel<'_>, ) -> Result, AicError> { - db.moe.slice_points( - moe_kernel(table), - self.quant_mode.name(), - &self.workload_distribution, - self.topk, - self.num_experts, - self.hidden_size, - self.inter_size, - self.moe_tp_size, - self.moe_ep_size, - ) + match table { + MoeTableSel::Standard => db.moe.slice_points( + crate::perf_database::moe::MoeKernel::Standard, + self.quant_mode.name(), + &self.workload_distribution, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + ), + MoeTableSel::LowLatency => db.moe.slice_points( + crate::perf_database::moe::MoeKernel::LowLatency, + self.quant_mode.name(), + &self.workload_distribution, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + ), + MoeTableSel::ExactKernelSource(kernel_source) => db.moe.slice_points_for_kernel_source( + kernel_source, + self.quant_mode.name(), + &self.workload_distribution, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + ), + } } /// Distinct quant names of the selected table, in first-seen (file row) @@ -614,9 +680,19 @@ impl MoeOp { fn table_available_quants( &self, db: &PerfDatabase, - table: MoeTableSel, + table: MoeTableSel<'_>, ) -> Result, AicError> { - db.moe.available_quants(moe_kernel(table)) + match table { + MoeTableSel::Standard => db + .moe + .available_quants(crate::perf_database::moe::MoeKernel::Standard), + MoeTableSel::LowLatency => db + .moe + .available_quants(crate::perf_database::moe::MoeKernel::LowLatency), + MoeTableSel::ExactKernelSource(kernel_source) => { + db.moe.available_quants_for_kernel_source(kernel_source) + } + } } /// Enumerate `source_quant`'s collected sibling slices (same table, @@ -627,19 +703,37 @@ impl MoeOp { fn collect_candidates( &self, db: &PerfDatabase, - table: MoeTableSel, + table: MoeTableSel<'_>, source_quant: MoeQuantMode, sol_quant: MoeQuantMode, provenance: &'static str, out: &mut Vec, ) -> Result<(), AicError> { - let slices = db.moe.sibling_slices( - moe_kernel(table), - source_quant.name(), - &self.workload_distribution, - self.moe_tp_size, - self.moe_ep_size, - ); + let slices = match table { + MoeTableSel::Standard => db.moe.sibling_slices( + crate::perf_database::moe::MoeKernel::Standard, + source_quant.name(), + &self.workload_distribution, + self.moe_tp_size, + self.moe_ep_size, + ), + MoeTableSel::LowLatency => db.moe.sibling_slices( + crate::perf_database::moe::MoeKernel::LowLatency, + source_quant.name(), + &self.workload_distribution, + self.moe_tp_size, + self.moe_ep_size, + ), + MoeTableSel::ExactKernelSource(kernel_source) => { + db.moe.sibling_slices_for_kernel_source( + kernel_source, + source_quant.name(), + &self.workload_distribution, + self.moe_tp_size, + self.moe_ep_size, + ) + } + }; let slices = match slices { Ok(slices) => slices, Err(err) if err.is_missing_perf_data() => return Ok(()), @@ -874,6 +968,7 @@ mod tests { attention_dp_size, is_gated: true, moe_backend: None, + moe_kernel_source: None, enable_eplb: false, is_context: false, } @@ -911,6 +1006,7 @@ mod tests { attention_dp_size: 1, is_gated: true, moe_backend: None, + moe_kernel_source: None, enable_eplb: false, is_context: false, } @@ -929,6 +1025,161 @@ mod tests { assert_eq!(result.source, source, "{label}: wrong source"); } + fn exact_source_fixture(mode: DatabaseMode) -> (tempfile::TempDir, PerfDatabase, MoeOp) { + use crate::perf_database::energy_test_fixtures::{ + Col, write_energy_systems_root, write_parquet, + }; + + // Synthetic values distinguish the restricted NVFP4 lane from the + // unrestricted FlashInfer FP8 lane; these are not measured timings. + let tmp = tempfile::tempdir().expect("tmpdir"); + let data = write_energy_systems_root(tmp.path()); + write_parquet( + &data.join("moe_perf.parquet"), + &[ + Col::Str( + "moe_dtype", + vec!["nvfp4", "nvfp4", "fp8_block", "fp8_block"], + ), + Col::I64("num_tokens", vec![64, 128, 64, 256]), + Col::I64("hidden_size", vec![4096; 4]), + Col::I64("inter_size", vec![2048; 4]), + Col::I64("topk", vec![2; 4]), + Col::I64("num_experts", vec![8; 4]), + Col::I64("moe_tp_size", vec![1; 4]), + Col::I64("moe_ep_size", vec![1; 4]), + Col::Str("distribution", vec!["uniform"; 4]), + Col::Str( + "kernel_source", + vec![ + "moe_torch_flow_min_latency", + "moe_torch_flow_min_latency", + "sglang_flashinfer_trtllm_moe", + "sglang_flashinfer_trtllm_moe", + ], + ), + Col::F64("latency", vec![1.0, 2.0, 3.0, 6.0]), + ], + ); + let db = PerfDatabase::load(tmp.path(), "testsys", "vllm", "1.0") + .expect("synthetic database") + .with_mode(mode, TransferPolicy::ALL); + let mut op = MoeOp::new( + "moe", + 4096, + 2048, + 2, + 8, + 1, + 1, + MoeQuantMode::Nvfp4, + "uniform", + ); + op.moe_kernel_source = Some("moe_torch_flow_min_latency".into()); + (tmp, db, op) + } + + fn assert_explicit_low_latency_rejects_ineligible_queries(mode: DatabaseMode) { + let (_tmp, db, op) = exact_source_fixture(mode); + for (is_gated, quant, tokens, attention_dp_size) in [ + (false, MoeQuantMode::Nvfp4, 64, 1), + (true, MoeQuantMode::Nvfp4, 129, 1), + (true, MoeQuantMode::Fp8Block, 64, 1), + (true, MoeQuantMode::Nvfp4Wo, 64, 1), + (true, MoeQuantMode::Nvfp4, 65, 2), + ] { + let mut invalid = op.clone(); + invalid.is_gated = is_gated; + invalid.quant_mode = quant; + invalid.attention_dp_size = attention_dp_size; + let result = invalid.query(&db, tokens); + assert!( + matches!(&result, Err(AicError::InvalidEngineConfig(message)) + if message.contains("moe_torch_flow_min_latency") + && message.contains("gated nvfp4") + && message.contains("128")), + "ineligible exact lane must fail before lookup or transfer: {mode:?}, \ + gated={is_gated}, quant={quant:?}, tokens={tokens}, \ + attention_dp={attention_dp_size}, got {result:?}" + ); + } + } + + #[test] + fn moe_explicit_low_latency_rejects_ineligible_silicon_queries() { + assert_explicit_low_latency_rejects_ineligible_queries(DatabaseMode::Silicon); + } + + #[test] + fn moe_explicit_low_latency_rejects_ineligible_empirical_queries() { + assert_explicit_low_latency_rejects_ineligible_queries(DatabaseMode::Empirical); + } + + #[test] + fn moe_explicit_low_latency_rejects_ineligible_hybrid_queries() { + assert_explicit_low_latency_rejects_ineligible_queries(DatabaseMode::Hybrid); + } + + #[test] + fn moe_explicit_low_latency_accepts_gated_nvfp4_through_128_tokens() { + for mode in [ + DatabaseMode::Silicon, + DatabaseMode::Empirical, + DatabaseMode::Hybrid, + ] { + let (_tmp, db, mut op) = exact_source_fixture(mode); + let source = if mode == DatabaseMode::Empirical { + Source::Empirical + } else { + Source::Silicon + }; + // Endpoints are the synthetic rows. Linear interpolation at 127 + // is 1 + 63/64 of the 1 ms interval: 1.984375 ms. + for (tokens, expected_ms) in [(64, 1.0), (127, 1.984375), (128, 2.0)] { + let result = op.query(&db, tokens).expect("eligible exact lane"); + assert_routing(&result, source, "low latency"); + if source == Source::Silicon { + assert_eq!(result.latency_ms, expected_ms); + } + } + op.attention_dp_size = 2; + let result = op.query(&db, 64).expect("128 gathered tokens"); + assert_routing(&result, source, "low latency with dp"); + if source == Source::Silicon { + assert_eq!(result.latency_ms, 2.0); + } + } + } + + #[test] + fn moe_other_exact_sources_allow_non_gated_fp8_above_128_tokens() { + for mode in [ + DatabaseMode::Silicon, + DatabaseMode::Empirical, + DatabaseMode::Hybrid, + ] { + let (_tmp, db, mut op) = exact_source_fixture(mode); + op.moe_kernel_source = Some("sglang_flashinfer_trtllm_moe".into()); + op.is_gated = false; + op.quant_mode = MoeQuantMode::Fp8Block; + let result = op + .query(&db, 256) + .expect("other source is not low-latency restricted"); + assert_routing( + &result, + if mode == DatabaseMode::Empirical { + Source::Empirical + } else { + Source::Silicon + }, + "other exact source", + ); + if result.source == Source::Silicon { + assert_eq!(result.latency_ms, 6.0); // exact synthetic FlashInfer row + } + } + } + /// Oracle values generated from the Python reference on the same data /// (shared layer pinned OFF so Python reads exactly the primary parquet /// the Rust table loads): @@ -1127,6 +1378,7 @@ mod tests { attention_dp_size: 1, is_gated: true, moe_backend: None, + moe_kernel_source: None, enable_eplb: false, is_context: false, }; @@ -1156,6 +1408,7 @@ mod tests { attention_dp_size: 1, is_gated: true, moe_backend: None, + moe_kernel_source: None, enable_eplb: false, is_context: false, }; diff --git a/crates/core/src/perfmodel/operators/moe_expert_compute.rs b/crates/core/src/perfmodel/operators/moe_expert_compute.rs index bdb6d9c2d..15bf976c9 100644 --- a/crates/core/src/perfmodel/operators/moe_expert_compute.rs +++ b/crates/core/src/perfmodel/operators/moe_expert_compute.rs @@ -202,6 +202,7 @@ impl MoeExpertComputeOp { workload_distribution: self.workload_distribution.clone(), is_gated: self.is_gated, moe_backend: None, + moe_kernel_source: None, enable_eplb: false, is_context: self.inference_phase == "context", }; diff --git a/crates/core/src/perfmodel/perf_database/moe.rs b/crates/core/src/perfmodel/perf_database/moe.rs index 8a5587f3e..52e501368 100644 --- a/crates/core/src/perfmodel/perf_database/moe.rs +++ b/crates/core/src/perfmodel/perf_database/moe.rs @@ -74,14 +74,13 @@ pub struct MoeSiblingSlice { pub points: Vec<(u32, f64)>, } -/// Two parallel grids split by `kernel_source`. Mirrors Python's split in -/// `aisimulate.sdk.operations.moe.MoE.load_data`, where rows tagged -/// `kernel_source == "moe_torch_flow_min_latency"` route to a separate -/// accumulator that the TRT-LLM SILICON path probes first for small-token -/// nvfp4 gated MoE queries. +/// Parallel grids split by `kernel_source`. `default` and `low_latency` +/// preserve the legacy automatic lookup behavior, while `by_kernel_source` +/// retains every named lane for an exact user-selected lookup. struct LoadedMoeGrids { default: MoeGrids, low_latency: MoeGrids, + by_kernel_source: BTreeMap, } struct MoeGrids { @@ -94,6 +93,29 @@ struct MoeGrids { quants_in_load_order: Vec, } +fn insert_moe_row( + index: &mut MoeIndex>, + quants_in_load_order: &mut Vec, + quant: &str, + distribution: &str, + shape: MoeShapeKey, + num_tokens: u32, + value: LeafValue, +) { + if !quants_in_load_order + .iter() + .any(|existing| existing == quant) + { + quants_in_load_order.push(quant.to_owned()); + } + // First-seen value wins within both the legacy merged grid and an exact + // lane, including across priority-ordered shared-layer sources. + index + .entry(quant.to_owned(), distribution.to_owned(), shape) + .entry(num_tokens) + .or_insert(value); +} + #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] struct MoeKey { quant: String, @@ -167,8 +189,70 @@ impl MoeTable { workload_distribution: &str, sol: &dyn Fn(f64) -> f64, ) -> Result { - let loaded = self.load()?; - let grids = &loaded.default; + self.query_for_kernel_source( + None, + num_tokens, + hidden_size, + inter_size, + topk, + num_experts, + moe_tp_size, + moe_ep_size, + quant, + workload_distribution, + sol, + ) + } + + /// Query one exact collected kernel-source lane. The requested lane is + /// never merged with or substituted by another lane: an unavailable lane + /// is a typed missing-data error instead of a silent fallback. + #[allow(clippy::too_many_arguments)] + pub fn query_kernel_source( + &self, + kernel_source: &str, + num_tokens: u32, + hidden_size: u32, + inter_size: u32, + topk: u32, + num_experts: u32, + moe_tp_size: u32, + moe_ep_size: u32, + quant: MoeQuantMode, + workload_distribution: &str, + sol: &dyn Fn(f64) -> f64, + ) -> Result { + self.query_for_kernel_source( + Some(kernel_source), + num_tokens, + hidden_size, + inter_size, + topk, + num_experts, + moe_tp_size, + moe_ep_size, + quant, + workload_distribution, + sol, + ) + } + + #[allow(clippy::too_many_arguments)] + fn query_for_kernel_source( + &self, + kernel_source: Option<&str>, + num_tokens: u32, + hidden_size: u32, + inter_size: u32, + topk: u32, + num_experts: u32, + moe_tp_size: u32, + moe_ep_size: u32, + quant: MoeQuantMode, + workload_distribution: &str, + sol: &dyn Fn(f64) -> f64, + ) -> Result { + let grids = self.grids_for_kernel_source(kernel_source)?; let quant_name = quant.name(); let shape = MoeShapeKey { @@ -301,6 +385,63 @@ impl MoeTable { moe_ep_size: u32, ) -> Result, AicError> { let grids = self.grids_for(kernel)?; + self.slice_points_in_grids( + grids, + format!("{kernel:?}"), + quant_name, + workload_distribution, + topk, + num_experts, + hidden_size, + inter_size, + moe_tp_size, + moe_ep_size, + ) + } + + /// Own-slice token curve from an exact collected kernel-source lane. + #[allow(clippy::too_many_arguments)] + pub fn slice_points_for_kernel_source( + &self, + kernel_source: &str, + quant_name: &str, + workload_distribution: &str, + topk: u32, + num_experts: u32, + hidden_size: u32, + inter_size: u32, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { + let grids = self.grids_for_kernel_source(Some(kernel_source))?; + self.slice_points_in_grids( + grids, + format!("kernel_source={kernel_source:?}"), + quant_name, + workload_distribution, + topk, + num_experts, + hidden_size, + inter_size, + moe_tp_size, + moe_ep_size, + ) + } + + #[allow(clippy::too_many_arguments)] + fn slice_points_in_grids( + &self, + grids: &MoeGrids, + grid_label: String, + quant_name: &str, + workload_distribution: &str, + topk: u32, + num_experts: u32, + hidden_size: u32, + inter_size: u32, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { let shape = MoeShapeKey { topk, num_experts, @@ -316,7 +457,7 @@ impl MoeTable { let by_tokens = by_tokens.filter(|curve| !curve.is_empty()).ok_or_else(|| { let key = MoeKey::from_shape(quant_name, dist, shape); AicError::PerfDatabase(format!( - "MoE data missing for {key:?} ({kernel:?}) at {}", + "MoE data missing for {key:?} ({grid_label}) at {}", self.data_root.display() )) })?; @@ -343,12 +484,47 @@ impl MoeTable { moe_ep_size: u32, ) -> Result, AicError> { let grids = self.grids_for(kernel)?; + Ok(Self::sibling_slices_in_grids( + grids, + quant_name, + workload_distribution, + moe_tp_size, + moe_ep_size, + )) + } + + /// Collected sibling slices from an exact kernel-source lane. + pub fn sibling_slices_for_kernel_source( + &self, + kernel_source: &str, + quant_name: &str, + workload_distribution: &str, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { + let grids = self.grids_for_kernel_source(Some(kernel_source))?; + Ok(Self::sibling_slices_in_grids( + grids, + quant_name, + workload_distribution, + moe_tp_size, + moe_ep_size, + )) + } + + fn sibling_slices_in_grids( + grids: &MoeGrids, + quant_name: &str, + workload_distribution: &str, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Vec { let (_, by_shape) = grids .index .resolve_uniform_shapes(quant_name, workload_distribution); let mut slices = Vec::new(); let Some(by_shape) = by_shape else { - return Ok(slices); + return slices; }; for (shape, curve) in by_shape { if shape.moe_tp_size != moe_tp_size @@ -365,7 +541,7 @@ impl MoeTable { points: curve.iter().map(|(t, leaf)| (t, leaf.latency)).collect(), }); } - Ok(slices) + slices } /// Distinct quant names present in the kernel grid, in first-seen @@ -376,6 +552,17 @@ impl MoeTable { Ok(self.grids_for(kernel)?.quants_in_load_order.clone()) } + /// Distinct quant names in an exact collected kernel-source lane. + pub fn available_quants_for_kernel_source( + &self, + kernel_source: &str, + ) -> Result, AicError> { + Ok(self + .grids_for_kernel_source(Some(kernel_source))? + .quants_in_load_order + .clone()) + } + fn grids_for(&self, kernel: MoeKernel) -> Result<&MoeGrids, AicError> { let loaded = self.load()?; Ok(match kernel { @@ -384,6 +571,19 @@ impl MoeTable { }) } + fn grids_for_kernel_source(&self, kernel_source: Option<&str>) -> Result<&MoeGrids, AicError> { + let loaded = self.load()?; + match kernel_source { + None => Ok(&loaded.default), + Some(kernel_source) => loaded.by_kernel_source.get(kernel_source).ok_or_else(|| { + AicError::PerfDatabase(format!( + "MoE kernel_source {kernel_source:?} is unavailable at {}", + self.data_root.display() + )) + }), + } + } + fn load(&self) -> Result<&LoadedMoeGrids, AicError> { let cell = self.moe.get_or_init(|| load_moe_parquet(&self.moe_sources)); cell.as_ref().map_err(clone_err) @@ -417,6 +617,11 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result MoeIndex::default(); let mut default_quants: Vec = Vec::new(); let mut low_latency_quants: Vec = Vec::new(); + let mut kernel_source_indices: BTreeMap< + String, + MoeIndex>, + > = BTreeMap::new(); + let mut kernel_source_quants: BTreeMap> = BTreeMap::new(); let mut any_source = false; for source in sources { let path = source.path(); @@ -474,22 +679,34 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result } else { (&mut default_index, &mut default_quants) }; - // First-seen (file row) quant order — Python's dict insertion - // order, consumed by `available_quants`. - if !target_quants.iter().any(|q| q == &quant) { - target_quants.push(quant.clone()); - } let latency = row.f64(latency_col)?; let power = row.f64_optional(power_col)?.unwrap_or(0.0); - // Python's `load_moe_data` wraps the leaf insert in a try/except KeyError - // and skips on conflict, i.e. it keeps the FIRST occurrence of each - // (shape, num_tokens) tuple. Some perf files contain duplicate rows - // (same kernel_source, same shape) — preserving first-wins parity here, - // extended across shared-layer sources (earlier source wins). - target - .entry(quant, distribution, shape) - .entry(row.u32(num_tokens_col)?) - .or_insert(LeafValue::with_power(latency, power)); + let value = LeafValue::with_power(latency, power); + let num_tokens = row.u32(num_tokens_col)?; + insert_moe_row( + target, + target_quants, + &quant, + &distribution, + shape, + num_tokens, + value, + ); + if !kernel_source.is_empty() { + let source_index = kernel_source_indices + .entry(kernel_source.clone()) + .or_default(); + let source_quants = kernel_source_quants.entry(kernel_source).or_default(); + insert_moe_row( + source_index, + source_quants, + &quant, + &distribution, + shape, + num_tokens, + value, + ); + } } } if !any_source || (default_index.is_empty() && low_latency_index.is_empty()) { @@ -502,6 +719,21 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result .unwrap_or_default() ))); } + let by_kernel_source = kernel_source_indices + .into_iter() + .map(|(kernel_source, index)| { + let quants_in_load_order = kernel_source_quants + .remove(&kernel_source) + .expect("kernel-source quant order is inserted with its index"); + ( + kernel_source, + MoeGrids { + index: index.map_values(|curve| LeafAxisCurve::from_map("num_tokens", curve)), + quants_in_load_order, + }, + ) + }) + .collect(); Ok(LoadedMoeGrids { default: MoeGrids { index: default_index.map_values(|curve| LeafAxisCurve::from_map("num_tokens", curve)), @@ -512,6 +744,7 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result .map_values(|curve| LeafAxisCurve::from_map("num_tokens", curve)), quants_in_load_order: low_latency_quants, }, + by_kernel_source, }) } @@ -678,6 +911,85 @@ mod tests { ); } + #[test] + fn moe_exact_kernel_source_selects_the_requested_sglang_lane() { + // Both lanes cover the same FP8 MoE shape. Unset retains the legacy + // first-row/default result; an exact selector must use only its lane. + use crate::perf_database::energy_test_fixtures::{Col, write_parquet}; + + let tmp = tempfile::tempdir().expect("tmpdir"); + write_parquet( + &tmp.path().join("moe_perf.parquet"), + &[ + Col::Str("moe_dtype", vec!["fp8_block", "fp8_block"]), + Col::I64("num_tokens", vec![1024, 1024]), + Col::I64("hidden_size", vec![8192, 8192]), + Col::I64("inter_size", vec![2048, 2048]), + Col::I64("topk", vec![10, 10]), + Col::I64("num_experts", vec![512, 512]), + Col::I64("moe_tp_size", vec![1, 1]), + Col::I64("moe_ep_size", vec![8, 8]), + Col::Str("distribution", vec!["power_law_1.2", "power_law_1.2"]), + Col::Str( + "kernel_source", + vec!["sglang_fused_moe_triton", "sglang_flashinfer_trtllm_moe"], + ), + Col::F64("latency", vec![1.0, 0.25]), + ], + ); + let table = MoeTable::new(tmp.path().to_path_buf()); + + let legacy = table + .query( + 1024, + 8192, + 2048, + 10, + 512, + 1, + 8, + MoeQuantMode::Fp8Block, + "power_law_1.2", + &proxy_sol, + ) + .expect("default lane resolves"); + let flashinfer = table + .query_kernel_source( + "sglang_flashinfer_trtllm_moe", + 1024, + 8192, + 2048, + 10, + 512, + 1, + 8, + MoeQuantMode::Fp8Block, + "power_law_1.2", + &proxy_sol, + ) + .expect("requested FlashInfer lane resolves"); + + assert_eq!(legacy.latency, 1.0); + assert_eq!(flashinfer.latency, 0.25); + + let error = table + .query_kernel_source( + "sglang_missing_moe", + 1024, + 8192, + 2048, + 10, + 512, + 1, + 8, + MoeQuantMode::Fp8Block, + "power_law_1.2", + &proxy_sol, + ) + .expect_err("an unavailable exact lane must not fall back to the default"); + assert!(error.to_string().contains("sglang_missing_moe")); + } + /// ENERGY oracle on a synthetic power-carrying fixture. Python twin /// (pandas fixture, `energy_test_fixtures` spec): /// diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index e637363eb..38e636567 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1051,6 +1051,7 @@ struct EngineBuildRequest { moe_backend: Option, enable_eplb: bool, wideep_num_slots: Option, + moe_kernel_source: Option, nextn: u32, speculation: Option, kv_block_size: Option, @@ -1102,6 +1103,7 @@ impl AicEngineBuilder { moe_backend: None, enable_eplb: false, wideep_num_slots: None, + moe_kernel_source: None, nextn: 0, speculation: None, kv_block_size: None, @@ -1233,6 +1235,12 @@ impl AicEngineBuilder { self } + /// Select an exact collected MoE compute kernel-source lane. + pub fn moe_kernel_source(mut self, value: impl Into) -> Self { + self.request.moe_kernel_source = Some(value.into()); + self + } + /// Configure speculative decoding. pub fn speculative_decoding(mut self, nextn: u32) -> Self { self.request.nextn = nextn; @@ -1272,6 +1280,7 @@ mod builder_tests { assert!(builder.request.moe_tp_size.is_none()); assert!(builder.request.moe_ep_size.is_none()); assert!(builder.request.attention_backend.is_none()); + assert!(builder.request.moe_kernel_source.is_none()); assert!(builder.request.kv_block_size.is_none()); assert!(builder.request.fpm_parquet_path.is_none()); assert!(builder.request.database_mode.is_none()); @@ -1289,6 +1298,7 @@ mod builder_tests { .attention_dp_size(4) .moe_parallelism(Some(1), Some(8)) .attention_backend("fa3") + .moe_kernel_source("sglang_flashinfer_trtllm_moe") .database_mode(DatabaseMode::Empirical) .shared_layer(true) .transfer_policy(vec!["xshape".to_owned(), "xquant".to_owned()]) @@ -1306,6 +1316,10 @@ mod builder_tests { (Some(1), Some(8)) ); assert_eq!(builder.request.attention_backend.as_deref(), Some("fa3")); + assert_eq!( + builder.request.moe_kernel_source.as_deref(), + Some("sglang_flashinfer_trtllm_moe") + ); assert_eq!(builder.request.database_mode.as_deref(), Some("EMPIRICAL")); assert_eq!(builder.request.shared_layer, Some(true)); assert_eq!( @@ -1466,6 +1480,7 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result, + moe_kernel_source: Option, enable_eplb: bool, seq_split: u32, ) -> PyResult<(Self, PyOperation)> { @@ -1799,6 +1800,7 @@ impl PyMoE { workload_distribution, is_gated, moe_backend, + moe_kernel_source, enable_eplb, is_context, }); @@ -1828,6 +1830,7 @@ impl PyMoE { kwargs.set_item("is_context", o.is_context)?; kwargs.set_item("is_gated", o.is_gated)?; kwargs.set_item("moe_backend", o.moe_backend.clone())?; + kwargs.set_item("moe_kernel_source", o.moe_kernel_source.clone())?; kwargs.set_item("enable_eplb", o.enable_eplb)?; Ok((args, kwargs)) } @@ -1896,6 +1899,11 @@ impl PyMoE { Ok(slf.as_super().moe()?.moe_backend.clone()) } + #[getter(_moe_kernel_source)] + fn moe_kernel_source(slf: PyRef<'_, Self>) -> PyResult> { + Ok(slf.as_super().moe()?.moe_kernel_source.clone()) + } + #[getter(_enable_eplb)] fn enable_eplb(slf: PyRef<'_, Self>) -> PyResult { Ok(slf.as_super().moe()?.enable_eplb) diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 14af520df..a4a40de8b 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -226,6 +226,8 @@ struct AicTimingConfig { #[serde(default)] attention_backend: Option, #[serde(default)] + moe_kernel_source: Option, + #[serde(default)] moe_backend: Option, #[serde(default)] enable_eplb: bool, @@ -322,6 +324,7 @@ impl AicTimingConfig { moe_backend: self.moe_backend.clone(), enable_eplb: self.enable_eplb, wideep_num_slots: self.wideep_num_slots, + moe_kernel_source: self.moe_kernel_source.clone(), enable_shared_layer: self.enable_shared_layer, strict_provenance: self.strict_provenance, }) @@ -2766,6 +2769,21 @@ mod tests { ); } + #[test] + fn canonical_timing_config_preserves_moe_kernel_source() { + let mut config = aic_config(); + config.moe_kernel_source = Some("sglang_flashinfer_trtllm_moe".into()); + + assert_eq!( + config + .estimator_request(ForwardPassWorkerType::Aggregated) + .unwrap() + .moe_kernel_source + .as_deref(), + Some("sglang_flashinfer_trtllm_moe") + ); + } + fn aic_config() -> AicTimingConfig { AicTimingConfig { model: "test-model".into(), @@ -2801,6 +2819,7 @@ mod tests { moe_backend: None, enable_eplb: false, wideep_num_slots: None, + moe_kernel_source: None, enable_shared_layer: None, strict_provenance: false, systems_path: None, diff --git a/crates/core/tests/perfmodel/memory_round_trip.rs b/crates/core/tests/perfmodel/memory_round_trip.rs index d2236e4ba..3880eb025 100644 --- a/crates/core/tests/perfmodel/memory_round_trip.rs +++ b/crates/core/tests/perfmodel/memory_round_trip.rs @@ -88,6 +88,7 @@ fn request(tolerance_fraction: Option) -> KvCacheEstimateRequest { forward_model: None, fpm_parquet_path: None, decoder_replay: false, + moe_kernel_source: None, 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 8f95f8285..60321911b 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -130,7 +130,10 @@ 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, 20); + // v19: Dsv41AttentionOp gained kv_cache_layout. + // v20: FpmForwardOp gained original_fmha_quant_mode for selector diagnostics. + // v21: EngineConfig and MoeOp gained exact moe_kernel_source identity. + assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 21); assert_eq!(FPM_VERSION, 1); assert_eq!(ForwardPassMetrics::default().version, FPM_VERSION); } diff --git a/docs/core-api.md b/docs/core-api.md index 670716234..042c2a114 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -208,18 +208,15 @@ The parquet identity must match the requested model, hardware, backend version, ### Engine identity controls -The canonical configuration also carries quantization overrides and -`attention_backend`, `moe_backend`, `enable_eplb` (default `false`), and -`wideep_num_slots` (default absent). These controls reach model construction, -KV memory sizing, and replay provenance. EPLB/slots and nondefault MoE backend -selection require an MoE model. Collected FPM interpolation cannot represent -EPLB, slots, or MoE backend overrides; it rejects an explicit incompatible -request and is skipped during automatic selection for those identities. - -Rust callers using exhaustive `ForwardPassPerfModelConfig` literals must add -`moe_backend: None`, `enable_eplb: false`, and `wideep_num_slots: None`. -`ForwardPassPerfModelConfig::new(...)` supplies these defaults. This extends -the canonical configuration introduced by #242. +The canonical configuration also carries quantization overrides and `attention_backend`, `moe_backend`, `moe_kernel_source` (default absent), `enable_eplb` (default `false`), and `wideep_num_slots` (default absent). These controls reach model construction, KV memory sizing, and replay provenance. EPLB/slots and nondefault MoE backend or kernel-source selection require an MoE model. Collected FPM interpolation cannot represent EPLB, slots, MoE backend, or kernel-source overrides; it rejects an explicit incompatible request and is skipped during automatic selection for those identities. + +`moe_kernel_source` selects an exact, nonblank collected `kernel_source` label for fused MoE compute. It is distinct from the existing `moe_backend` graph/backend control; source labels are not backend aliases and are preserved without trimming. `None` keeps the existing default source-selection policy, including eligible low-latency NVFP4 selection. `SILICON` reads only the requested source's table; `EMPIRICAL` derives its estimate from that same source; `HYBRID` may fall back to empirical estimation within that source, but does not substitute a different source. Missing source data remains an error. An explicit `moe_torch_flow_min_latency` requires gated NVFP4 and at most 128 tokens after attention-DP gathering. Pure-roofline `SOL` remains table-independent and does not claim measured support for the requested source. + +An explicit source is rejected for dense graphs, MegaMoE modules, large-EP expert-compute graphs, and any constructed timing phase with no compatible fused MoE operator. It is also incompatible with whole-forward FPM, including the legacy Task `forward_model='fpm'` rewrite. Invalid graph/source combinations fail as invalid configuration rather than triggering estimator fallback. An untrained `fpm_regression` model remains not-ready; retaining a source in its configuration is not evidence of source-specific prediction support. + +AFD regular companions currently reject exact-source requests through their legacy estimator and fixed timing paths. The external-FPM companion forwards the source to canonical validation, which rejects the incompatible FPM request. These controls describe standalone AISimulate behavior, not downstream Dynamo planner integration. + +Rust callers using exhaustive `ForwardPassPerfModelConfig` literals must add `moe_backend: None`, `moe_kernel_source: None`, `enable_eplb: false`, and `wideep_num_slots: None`. Direct `EngineConfig` and `MoeOp` literals likewise require the new `moe_kernel_source` field. `ForwardPassPerfModelConfig::new(...)` supplies its default. This extends the canonical configuration introduced by #242. Rust callers constructing `SyntheticTraceSpec` must also add `cached_prefix_tokens: 0` to preserve existing prefix-sharing behavior. A positive diff --git a/docs/deepseek-v41-storage.md b/docs/deepseek-v41-storage.md index 2a7a22fd2..f498ae287 100644 --- a/docs/deepseek-v41-storage.md +++ b/docs/deepseek-v41-storage.md @@ -53,5 +53,5 @@ Other backends retain the explicitly unqualified `logical_fp4` inventory and traffic, avoiding a guessed physical cache precision. The operator serializes `kv_cache_layout` independently of `fmha_quant_mode`. Legacy JSON without this field defaults to the theoretical layout; fresh SGLang graphs always name the -physical layout. Engine wire schema 20 rejects earlier bincode schemas before decoding +physical layout. Engine wire schema 21 rejects earlier bincode schemas before decoding the new positional field. Model observations and calibration data are unchanged. diff --git a/docs/deepseek-v41.md b/docs/deepseek-v41.md index f953bac6e..d18e10934 100644 --- a/docs/deepseek-v41.md +++ b/docs/deepseek-v41.md @@ -94,7 +94,7 @@ The `logical_fp4` layout (288-byte compressed main, 68-byte index, FP8 window, 890-byte global slope) remains the explicit theoretical estimate for vLLM and TRT-LLM; their runtime storage is unqualified. Both inventories exclude allocator page padding and spare pages. See [source proof and limits](deepseek-v41-storage.md) -for full-context scoring, cache read/write accounting, and schema 20 compatibility. +for full-context scoring, cache read/write accounting, and schema 21 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 @@ -128,6 +128,7 @@ preserves observed inputs, source identities, replay tools, and result coverage at that commit. Its results remain specific to the recorded source and runtime identities. -FPM selector diagnostics append a positional field to `FpmForwardOp`, so this -branch uses EngineSpec schema 20. Schema-19 and earlier bincode inputs are rejected -before decoding; legacy JSON may still use the documented field defaults. +FPM selector diagnostics append a positional field to `FpmForwardOp` in schema 20. +Exact MoE kernel-source identity extends the layout in EngineSpec schema 21. +Schema-20 and earlier bincode inputs are rejected before decoding; legacy JSON +may still use the documented field defaults. diff --git a/python/aisimulate/collector/cases/models/Qwen3_5MoeForCausalLM_cases.yaml b/python/aisimulate/collector/cases/models/Qwen3_5MoeForCausalLM_cases.yaml index b9de8e14d..e3e2516fa 100644 --- a/python/aisimulate/collector/cases/models/Qwen3_5MoeForCausalLM_cases.yaml +++ b/python/aisimulate/collector/cases/models/Qwen3_5MoeForCausalLM_cases.yaml @@ -166,7 +166,7 @@ model_case_values: # attention override above, so it applies verbatim regardless of that # gap. # - # No sglang_moe_backends override is declared: the SM100 + # Without the deployment pin, the SM100 # flashinfer_trtllm moe-runner-backend forcer # (`_qwen3_moe_family_overrides`, arg_groups/overrides.py:1283-1290, # decorating :1291) is gated by the SAME kind of architecture-string @@ -184,10 +184,13 @@ model_case_values: # and Fp8MoEMethod.create_moe_runner (layers/quantization/fp8.py:2265-2294, # `is_deepgemm_moe_runner_backend_enabled()` requires a non-none a2a # backend so falls to `else: moe_runner_backend = MoeRunnerBackend.TRITON` - # at :2281) for fp8_block. That is exactly the - # base_ops/moe.yaml:79-81 default ({default: triton} for both - # bfloat16 and fp8_block already), so redeclaring it here would just - # duplicate the base value -- omitted instead. + # at :2281) for fp8_block. Keep that base Triton default for bfloat16. + # The FP8 deployment pins the runner instead: ai-dynamo/dynamo + # 5fc15ee, recipes/qwen3.8-2.4t-a95b/sglang/agg-gb200-chat/deploy.yaml: + # 101-112 (SM100) and agg-gb300-chat/deploy.yaml:12-30 (SM103) set + # fp8 quantization and moe-runner-backend: flashinfer_trtllm. + sglang_moe_backends: + fp8_block: {100: flashinfer_trtllm, 103: flashinfer_trtllm} # # vLLM has no schema field analogous to sglang_moe_backends: # collect_moe.py (collector/vllm/collect_moe.py) constructs FusedMoE diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 00b5f1f5a..60e0de294 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -129,7 +129,7 @@ def finish_lowering(value): ): if resolved[source] is not None: lowered[target] = resolved[source] - for name in ("moe_backend", "attention_backend", "enable_eplb", "wideep_num_slots"): + for name in ("moe_backend", "moe_kernel_source", "attention_backend", "enable_eplb", "wideep_num_slots"): if resolved.get(name) is not None and not (name == "enable_eplb" and resolved[name] is False): lowered[f"aic_{name}"] = resolved[name] if resolved["systems_paths"]: @@ -208,7 +208,7 @@ def finish_lowering(value): comm_dtype=lowered.get("aic_comm_dtype"), **{ name: lowered[f"aic_{name}"] - for name in ("moe_backend", "attention_backend", "enable_eplb", "wideep_num_slots") + for name in ("moe_backend", "moe_kernel_source", "attention_backend", "enable_eplb", "wideep_num_slots") if lowered.get(f"aic_{name}") is not None }, systems_path=capacity_systems_path, @@ -241,6 +241,7 @@ def estimate_num_gpu_blocks( kv_cache_dtype: str | None = None, comm_dtype: str | None = None, moe_backend: str | None = None, + moe_kernel_source: str | None = None, attention_backend: str | None = None, enable_eplb: bool = False, wideep_num_slots: int | None = None, @@ -315,6 +316,7 @@ def estimate_num_gpu_blocks( name: value for name, value in ( ("moe_backend", moe_backend), + ("moe_kernel_source", moe_kernel_source), ("attention_backend", attention_backend), ("enable_eplb", enable_eplb), ("wideep_num_slots", wideep_num_slots), diff --git a/python/aisimulate/src/aisimulate/config/common.py b/python/aisimulate/src/aisimulate/config/common.py index 442798668..3a5d59da2 100644 --- a/python/aisimulate/src/aisimulate/config/common.py +++ b/python/aisimulate/src/aisimulate/config/common.py @@ -224,6 +224,7 @@ def split_config_sections( "enable_eplb", "wideep_num_slots", "moe_backend", + "moe_kernel_source", "attention_backend", "gemm_quant_mode", "moe_quant_mode", @@ -247,7 +248,7 @@ def is_active_engine_model_control(name: str, value: Any) -> bool: def omit_inactive_moe_controls(config: dict[str, Any]) -> dict[str, Any]: """Keep additive defaults out of timing payloads parsed by older runners.""" result = dict(config) - for name in ("moe_backend", "wideep_num_slots", "enable_eplb"): + for name in ("moe_backend", "moe_kernel_source", "wideep_num_slots", "enable_eplb"): if not is_active_engine_model_control(name, result.get(name)): result.pop(name, None) return result diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 7980a74c7..12ec73dff 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -325,6 +325,7 @@ class EstimatorPolicyConfig(StrictModel): enable_eplb: bool = Field(default=False, strict=True) wideep_num_slots: PositiveInt | None = None moe_backend: str | None = None + moe_kernel_source: str | None = None attention_backend: str | None = None gemm_quant_mode: str | None = None moe_quant_mode: str | None = None diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index 2733d5d5e..20a28a8a1 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -101,6 +101,7 @@ "forward_model": ("forward_model", "aic_forward_model"), "fpm_parquet_path": ("fpm_parquet_path", "aic_fpm_parquet_path"), "moe_backend": ("aic_moe_backend",), + "moe_kernel_source": ("moe_kernel_source", "aic_moe_kernel_source"), "attention_backend": ("aic_attention_backend",), "enable_eplb": ("aic_enable_eplb",), "wideep_num_slots": ("aic_wideep_num_slots",), @@ -190,6 +191,8 @@ def measure(self, spec: ReplaySpec) -> AFDCompanionTiming: timing_model = args.get("timing_model") if isinstance(timing_model, Mapping) and timing_model.get("type") == "fixed": + if any(args.get(alias) is not None for alias in _AIC_TIMING_FIELD_ALIASES["moe_kernel_source"]): + raise ValueError("moe_kernel_source is not supported by fixed AFD companion timing") key = "prefill_ms" if role == "prefill" else "decode_ms" latency = _positive_number(timing_model.get(key), f"{role} timing_model.{key}") return AFDCompanionTiming( @@ -296,6 +299,7 @@ def measure(self, spec: ReplaySpec) -> AFDCompanionTiming: for name in ( "attention_backend", "moe_backend", + "moe_kernel_source", "enable_eplb", "wideep_num_slots", "decoder_replay", @@ -318,6 +322,8 @@ def measure(self, spec: ReplaySpec) -> AFDCompanionTiming: raw = {metric: latency} source = "aisimulate_core.sdk.rust_engine_step.RustForwardPassPerfModel" else: + if timing_overrides.get("moe_kernel_source") is not None: + raise ValueError("moe_kernel_source is not supported by the AFD companion legacy estimator") result = estimator(model_name, hardware, **kwargs) raw = getattr(result, "raw", None) except Exception as exc: diff --git a/python/aisimulate/src/aisimulate/sdk/task_v1_compat.py b/python/aisimulate/src/aisimulate/sdk/task_v1_compat.py index f383978eb..7d1158e4f 100644 --- a/python/aisimulate/src/aisimulate/sdk/task_v1_compat.py +++ b/python/aisimulate/src/aisimulate/sdk/task_v1_compat.py @@ -86,13 +86,14 @@ "nextn_accept_rates", "moe_backend", "attention_backend", + "moe_kernel_source", "wideep_num_slots", } ) # Global scalars that may appear inside the V1 ``config:`` block. _CONFIG_SCALARS: frozenset[str] = frozenset( - {"moe_backend", "nextn", "nextn_accept_rates", "attention_backend", "wideep_num_slots"} + {"moe_backend", "nextn", "nextn_accept_rates", "attention_backend", "moe_kernel_source", "wideep_num_slots"} ) # Keys that are part of the V1 format itself (handled structurally, never copied). diff --git a/python/aisimulate/src/aisimulate/sdk/task_v2.py b/python/aisimulate/src/aisimulate/sdk/task_v2.py index 00ad270b0..a24afc41a 100644 --- a/python/aisimulate/src/aisimulate/sdk/task_v2.py +++ b/python/aisimulate/src/aisimulate/sdk/task_v2.py @@ -576,6 +576,9 @@ class Task: # closed; None/default uses the mapped framework default or safe default fallback. SGLang WideEP # maps None/default to flashinfer and also supports fa3. attention_backend: common.AttentionBackend | None = None + # Exact collected MoE compute kernel-source lane; unlike moe_backend this + # does not select topology or a runtime implementation family. + moe_kernel_source: str | None = field(default=None, kw_only=True) wideep_num_slots: int | None = None # EPLB slot count; defaults to num_experts when None gemm_quant_mode: common.GEMMQuantMode | None = None moe_quant_mode: common.MoEQuantMode | None = None @@ -860,6 +863,7 @@ def __post_init__(self) -> None: common.AttentionBackend, "attention_backend", ) + self.moe_kernel_source = config.normalize_kernel_source(self.moe_kernel_source, "moe_kernel_source") # Canonicalize at construction so downstream config and routing see # exactly one spelling for the only supported engine-step backend. self.engine_step_backend = validate_engine_step_backend(self.engine_step_backend) @@ -2125,6 +2129,7 @@ def build_model_config( # this field as the kernel-LANE override — materializing a lane name here # would silently pin every model to the flashinfer lane. attention_backend=self.attention_backend, + moe_kernel_source=self.moe_kernel_source, wideep_num_slots=self.wideep_num_slots, forward_model=self.forward_model or "op_level", moe_comm_backend=None, diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index e8810df9e..2c161c2fb 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -542,6 +542,7 @@ class SearchSpace(BaseModel): enable_eplb: bool = Field(default=False, strict=True) wideep_num_slots: int | None = Field(default=None, strict=True, gt=0) moe_backend: str | None = None + moe_kernel_source: str | None = None attention_backend: str | None = None gemm_quant_mode: str | None = None moe_quant_mode: str | None = None diff --git a/python/aisimulate/src/aisimulate_core/sdk/config.py b/python/aisimulate/src/aisimulate_core/sdk/config.py index 5a64910d4..b529c6780 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config.py @@ -28,6 +28,15 @@ def normalize_kernel_backend( raise ValueError(f"{field_name} must be one of {choices}, got {value!r}.") from exc +def normalize_kernel_source(value: str | None, field_name: str) -> str | None: + """Validate an optional exact collected kernel-source label.""" + if value is None: + return None + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field_name} must be a non-empty string or None, got {value!r}.") + return value + + def has_video_input( *, num_videos: int = 0, @@ -108,6 +117,11 @@ class ModelConfig: # (backend, version, sm_version) wins. Do NOT default this to a lane name — # that would silently pin every model to that lane. attention_backend: common.AttentionBackend | None = None + # Exact collected MoE kernel-source lane. This is intentionally separate + # from ``moe_backend``, which describes topology/runtime behavior rather + # than a single measured compute-kernel lane. ``None`` preserves the + # existing framework/default lookup. + moe_kernel_source: str | None = field(default=None, kw_only=True) # DEPRECATED and ignored (large-EP is selected per tuple via # moe_comm_backend); kept for a compatibility window because ModelConfig # is exported through the supported core SDK facade and removal breaks @@ -144,6 +158,7 @@ def __post_init__(self) -> None: common.AttentionBackend, "attention_backend", ) + self.moe_kernel_source = normalize_kernel_source(self.moe_kernel_source, "moe_kernel_source") def resolve_moe_parallelism(self) -> tuple[int, int]: """Resolve and validate MoE parallelism dimensions in-place. diff --git a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py index 3fa3629e0..4979cc5eb 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py @@ -21,6 +21,7 @@ MoEQuantMode, ) from aisimulate_core.sdk.config import ModelConfig +from aisimulate_core.sdk.errors import InvalidEngineConfigurationError logger = logging.getLogger(__name__) @@ -44,10 +45,16 @@ def build_model_config( enable_eplb: bool = False, wideep_num_slots: int | None = None, *, + moe_kernel_source: str | None = None, fpm_fmha_quant_mode: str | None = None, ) -> ModelConfig: """Build a ModelConfig with optional quant mode overrides.""" - validate_moe_controls(enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots) + validate_moe_controls( + enable_eplb=enable_eplb, + wideep_num_slots=wideep_num_slots, + moe_backend=moe_backend, + moe_kernel_source=moe_kernel_source, + ) return ModelConfig( tp_size=tp_size, pp_size=pp_size, @@ -66,6 +73,7 @@ def build_model_config( moe_backend=None if moe_backend == "default" else moe_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, + moe_kernel_source=moe_kernel_source, speculation=speculation, ) @@ -75,6 +83,7 @@ def validate_moe_controls( enable_eplb: bool = False, wideep_num_slots: int | None = None, moe_backend: str | None = None, + moe_kernel_source: str | None = None, model_path: str | None = None, ) -> None: """Reject invalid MoE identity before model construction or fallback.""" @@ -82,11 +91,18 @@ def validate_moe_controls( raise ValueError("enable_eplb must be a boolean") if wideep_num_slots is not None and (type(wideep_num_slots) is not int or wideep_num_slots <= 0): raise ValueError("wideep_num_slots must be a positive integer") - if model_path is not None and (enable_eplb or wideep_num_slots is not None or moe_backend not in (None, "default")): + if moe_kernel_source is not None and moe_backend == "megamoe": + raise InvalidEngineConfigurationError("moe_kernel_source is not supported by the MegaMoE module") + if model_path is not None and ( + enable_eplb + or wideep_num_slots is not None + or moe_backend not in (None, "default") + or moe_kernel_source is not None + ): from aisimulate_core.sdk.models import check_is_moe if not check_is_moe(model_path): - raise ValueError("EPLB, slots and moe_backend require an MoE model") + raise InvalidEngineConfigurationError("EPLB, slots, moe_backend and moe_kernel_source require an MoE model") def validate_nextn(nextn: int | None) -> int: diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index c325434a5..16b7f5e3b 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -44,7 +44,7 @@ from typing import Any import aisimulate_core -from aisimulate_core.sdk.config_builders import apply_nextn, build_model_config +from aisimulate_core.sdk.config_builders import apply_nextn, build_model_config, validate_moe_controls from aisimulate_core.sdk.deepseek_v41 import MODEL_PATH as DEEPSEEK_V41_MODEL_PATH from aisimulate_core.sdk.errors import InvalidEngineConfigurationError as InvalidEngineConfigurationError from aisimulate_core.sdk.models import get_model @@ -322,6 +322,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), "fpm_parquet_path": fpm_parquet_path, + "moe_kernel_source": getattr(cfg, "moe_kernel_source", None), "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)), @@ -425,6 +426,7 @@ def compile_engine( fpm_fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None, attention_backend: str | None = None, + moe_kernel_source: str | None = None, moe_backend: str | None = None, enable_eplb: bool = False, wideep_num_slots: int | None = None, @@ -467,6 +469,13 @@ def compile_engine( resolved_moe_tp = moe_tp_size if moe_tp_size is not None else 1 resolved_moe_ep = moe_ep_size if moe_ep_size is not None else 1 try: + validate_moe_controls( + model_path=model_path, + enable_eplb=enable_eplb, + wideep_num_slots=wideep_num_slots, + moe_backend=moe_backend, + moe_kernel_source=moe_kernel_source, + ) resolved_speculation = SpeculationConfig(**speculation) if speculation is not None else None model_config = build_model_config( tp_size=tp_size, @@ -482,6 +491,7 @@ def compile_engine( comm_quant_mode=comm_quant_mode, forward_model=forward_model, attention_backend=attention_backend, + moe_kernel_source=moe_kernel_source, moe_backend=moe_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, @@ -492,11 +502,6 @@ def compile_engine( apply_nextn(model_config, nextn) except (ValueError, TypeError, KeyError) as exc: raise InvalidEngineConfigurationError(str(exc)) from exc - if enable_eplb or wideep_num_slots is not None or moe_backend not in (None, "default"): - from aisimulate_core.sdk.models import check_is_moe - - if not check_is_moe(model_path): - raise InvalidEngineConfigurationError("EPLB, slots and moe_backend require an MoE model") model_config.decoder_replay = decoder_replay try: resolve_dsv4_moe_arch(model_config, model_path, system_name=system, backend_name=backend) diff --git a/python/aisimulate/src/aisimulate_core/sdk/memory.py b/python/aisimulate/src/aisimulate_core/sdk/memory.py index 1e0115274..d2f8ead2b 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory.py @@ -279,6 +279,7 @@ def from_request( fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None, moe_backend: str | None = None, + moe_kernel_source: str | None = None, attention_backend: str | None = None, enable_eplb: bool = False, wideep_num_slots: int | None = None, @@ -309,7 +310,11 @@ def from_request( resolved_moe_tp = moe_tp_size if moe_tp_size is not None else 1 resolved_moe_ep = moe_ep_size if moe_ep_size is not None else 1 validate_moe_controls( - model_path=model_path, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, moe_backend=moe_backend + model_path=model_path, + enable_eplb=enable_eplb, + wideep_num_slots=wideep_num_slots, + moe_backend=moe_backend, + moe_kernel_source=moe_kernel_source, ) model_config = build_model_config( tp_size=tp_size, @@ -323,6 +328,7 @@ def from_request( moe_quant_mode=moe_quant_mode, comm_quant_mode=comm_quant_mode, moe_backend=moe_backend, + moe_kernel_source=moe_kernel_source, attention_backend=attention_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, @@ -1017,6 +1023,7 @@ def estimate_kv_cache( fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None, moe_backend: str | None = None, + moe_kernel_source: str | None = None, attention_backend: str | None = None, enable_eplb: bool = False, wideep_num_slots: int | None = None, @@ -1089,7 +1096,11 @@ def estimate_kv_cache( # Validate the compute-side MTP depth before any fallback path. validate_nextn(nextn) validate_moe_controls( - model_path=model_path, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, moe_backend=moe_backend + model_path=model_path, + enable_eplb=enable_eplb, + wideep_num_slots=wideep_num_slots, + moe_backend=moe_backend, + moe_kernel_source=moe_kernel_source, ) try: @@ -1111,6 +1122,7 @@ def estimate_kv_cache( fmha_quant_mode=fmha_quant_mode, comm_quant_mode=comm_quant_mode, moe_backend=moe_backend, + moe_kernel_source=moe_kernel_source, attention_backend=attention_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, @@ -1123,6 +1135,7 @@ def estimate_kv_cache( or enable_eplb or wideep_num_slots is not None or moe_backend not in (None, "default") + or moe_kernel_source is not None or attention_backend is not None ): raise ValueError( @@ -1182,6 +1195,7 @@ def estimate_num_gpu_blocks( fmha_quant_mode: str | None = None, comm_quant_mode: str | None = None, moe_backend: str | None = None, + moe_kernel_source: str | None = None, attention_backend: str | None = None, enable_eplb: bool = False, wideep_num_slots: int | None = None, @@ -1244,6 +1258,7 @@ def estimate_num_gpu_blocks( fmha_quant_mode=fmha_quant_mode, comm_quant_mode=comm_quant_mode, moe_backend=moe_backend, + moe_kernel_source=moe_kernel_source, attention_backend=attention_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py index 5d3c162a8..65d31bf60 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py @@ -28,6 +28,7 @@ import copy import importlib +import json import pkgutil from aisimulate_core.sdk import config @@ -64,6 +65,18 @@ _FORWARD_MODELS = ("op_level", "fpm") +def _uses_moe_kernel_source(spec, source: str) -> bool: + """Inspect native composite children as well as top-level MoE operators.""" + if isinstance(spec, dict): + moe = spec.get("Moe") + if isinstance(moe, dict) and moe.get("moe_kernel_source") == source: + return True + return any(_uses_moe_kernel_source(child, source) for child in spec.values()) + if isinstance(spec, list): + return any(_uses_moe_kernel_source(child, source) for child in spec) + return False + + 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 @@ -158,6 +171,16 @@ def get_model( raise InvalidEngineConfigurationError( f"Unknown forward_model: {forward_model!r}. Valid values: {', '.join(_FORWARD_MODELS)}" ) + if model_config.moe_kernel_source is not None: + from aisimulate_core.sdk.config_builders import validate_moe_controls + + if forward_model == "fpm": + raise InvalidEngineConfigurationError("moe_kernel_source is not supported with forward_model='fpm'") + validate_moe_controls( + model_path=model_path, + moe_backend=model_config.moe_backend, + moe_kernel_source=model_config.moe_kernel_source, + ) 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'") @@ -234,6 +257,14 @@ def get_model( model.spec_scheme = build_spec_scheme(model_config, spec_config) model.spec_scheme.validate(model, backend_name) materialize_spec_scheme(model) + if model_config.moe_kernel_source is not None: + for phase, phase_ops in (("context", model.context_ops), ("generation", model.generation_ops)): + if not any( + _uses_moe_kernel_source(json.loads(op._spec_json()), model_config.moe_kernel_source) for op in phase_ops + ): + raise InvalidEngineConfigurationError( + f"moe_kernel_source is not supported: {phase} graph has no compatible MoE operator" + ) if forward_model == "fpm": model = _apply_forward_model_fpm(model, backend_name) return model diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/blocks/moe.py b/python/aisimulate/src/aisimulate_core/sdk/models/blocks/moe.py index 6cf205cd3..0efc8735c 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/blocks/moe.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/blocks/moe.py @@ -28,6 +28,7 @@ import aisimulate_core.sdk.operations as ops from aisimulate_core.sdk import common +from aisimulate_core.sdk.errors import InvalidEngineConfigurationError from aisimulate_core.sdk.models.helpers import check_is_moe from aisimulate_core.sdk.operations.moe_comm import MOE_A2A_BACKENDS, communication_dtype_for, nodes_for @@ -296,6 +297,10 @@ def _default_moe_block_ops( # attribute keep working; absent/uncovered phase means the fused path below. comm_backend = (getattr(cfg, "moe_comm_backend", None) or {}).get(inference_phase) if comm_backend: + if getattr(cfg, "moe_kernel_source", None) is not None: + raise InvalidEngineConfigurationError( + "moe_kernel_source is unsupported when moe_comm_backend selects the large-EP MoE graph" + ) return _large_ep_block_ops( comm_backend, prefix=prefix, @@ -397,6 +402,7 @@ def _default_moe_block_ops( quant_mode, workload_distribution, cfg.attention_dp_size, + moe_kernel_source=cfg.moe_kernel_source, ), ops.MoEDispatch( f"{prefix}_moe_post_dispatch", diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py index c51984763..160c9bcfb 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py @@ -799,6 +799,7 @@ def _attn_mode(group: str) -> common.GEMMQuantMode: moe_quant_mode, workload_distribution, attention_dp_size, + moe_kernel_source=self.config.moe_kernel_source, ) ] ) @@ -1029,6 +1030,7 @@ def _attn_mode(group: str) -> common.GEMMQuantMode: moe_quant_mode, workload_distribution, attention_dp_size, + moe_kernel_source=self.config.moe_kernel_source, ), ops.MoEDispatch( "generation_moe_post_dispatch", diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py index 2bbb1ddcb..9e076fac7 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py @@ -502,6 +502,7 @@ def __init__(self, topk: int, num_experts: int, moe_inter_size: int, *args, back moe_quant_mode, workload_distribution, attention_dp_size, + moe_kernel_source=self.config.moe_kernel_source, ), ops.MoEDispatch( "context_moe_post_dispatch", @@ -631,6 +632,7 @@ def __init__(self, topk: int, num_experts: int, moe_inter_size: int, *args, back moe_quant_mode, workload_distribution, attention_dp_size, + moe_kernel_source=self.config.moe_kernel_source, ), ops.MoEDispatch( "generation_moe_post_dispatch", diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v4.py b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v4.py index f7e59296d..d82dad131 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v4.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v4.py @@ -203,6 +203,7 @@ def _moe_ops(phase: str, num_layers: float, is_context: bool, attn_cp: int = 1): moe_quant_mode, workload_distribution, attention_dp_size, + moe_kernel_source=self.config.moe_kernel_source, ), ops.MoEDispatch( f"{phase}_moe_post_dispatch", diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py index eeb1664ca..2d997f334 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v41.py @@ -206,6 +206,7 @@ def stage(layer: int, context: bool): model_config.moe_quant_mode, distribution, model_config.attention_dp_size, + moe_kernel_source=model_config.moe_kernel_source, ), ] # The post-expert reduction consumes both routed and shared diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/gemma4.py b/python/aisimulate/src/aisimulate_core/sdk/models/gemma4.py index c6feb7d7b..14460401d 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/gemma4.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/gemma4.py @@ -285,6 +285,7 @@ def _moe_ops( moe_q, wl_dist, attn_dp, + moe_kernel_source=self.config.moe_kernel_source, ), ops.MoEDispatch( f"{prefix}_moe_post_dispatch", diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/kimi_k3.py b/python/aisimulate/src/aisimulate_core/sdk/models/kimi_k3.py index 3d37b3c3b..8305a48e9 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/kimi_k3.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/kimi_k3.py @@ -723,6 +723,7 @@ def _ffn_ops(self, prefix: str, context: bool) -> list: moe_q, workload_dist, attn_dp, + moe_kernel_source=self.config.moe_kernel_source, ), ops.MoEDispatch( f"{prefix}_moe_post_dispatch", diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py b/python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py index 9b9504972..3859d7e2c 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py @@ -305,6 +305,7 @@ def _build_context_ops(self) -> None: moe_quant_mode, workload_distribution, attention_dp_size, + moe_kernel_source=self.config.moe_kernel_source, is_gated=False, # NemotronH uses Relu2 (non-gated) ), ops.MoEDispatch( @@ -597,6 +598,7 @@ def _build_generation_ops(self) -> None: moe_quant_mode, workload_distribution, attention_dp_size, + moe_kernel_source=self.config.moe_kernel_source, is_gated=False, # NemotronH uses Relu2 (non-gated) ), ops.MoEDispatch( diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/qwen35.py b/python/aisimulate/src/aisimulate_core/sdk/models/qwen35.py index 2fa04a7d2..96b790adb 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/qwen35.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/qwen35.py @@ -580,6 +580,7 @@ def _ffn_context_ops( attn_dp, is_context=True, moe_backend=self.config.moe_backend, + moe_kernel_source=self.config.moe_kernel_source, # EPLB is not modeled for Qwen3.5 (the load curve stays 1.2). enable_eplb=False, ) @@ -865,6 +866,7 @@ def _ffn_generation_ops( attn_dp, is_context=False, moe_backend=self.config.moe_backend, + moe_kernel_source=self.config.moe_kernel_source, enable_eplb=False, ) ) 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 7480fb056..8c7b4e912 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -130,6 +130,7 @@ class ForwardPassPerfModelConfig: fallback_policy: str = "deny" estimator_config: dict[str, Any] = dataclass_field(default_factory=dict) attention_backend: str | None = None + moe_kernel_source: str | None = dataclass_field(default=None, kw_only=True) enable_shared_layer: bool | None = None strict_provenance: bool = False moe_backend: str | None = None @@ -1197,6 +1198,7 @@ def _engine_config_json(model: Any, database: Any) -> str: # per phase); without this key they would share a cached handle and # silently answer with the other mode's engine. "forward_model": getattr(model, "forward_model", "op_level"), + "moe_kernel_source": getattr(model_config, "moe_kernel_source", None), # Same identity built against different systems roots reads different # perf trees; the root is part of the engine identity. "systems_root": str(getattr(database, "systems_root", "") or ""), @@ -1235,6 +1237,7 @@ def _engine_config_json(model: Any, database: Any) -> str: "sms": getattr(model_config, "sms", None), "moe_backend": getattr(model_config, "moe_backend", None), "attention_backend": getattr(model_config, "attention_backend", None), + "moe_kernel_source": getattr(model_config, "moe_kernel_source", None), # enable_wideep is gone from the identity: the deprecated # flag is constant False on every Task-built ModelConfig; # moe_comm_backend + num_gpus_per_node below carry the diff --git a/python/aisimulate/src/aisimulate_core/sdk/speculation/dspark.py b/python/aisimulate/src/aisimulate_core/sdk/speculation/dspark.py index ce57eb808..78bdb9a2b 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/speculation/dspark.py +++ b/python/aisimulate/src/aisimulate_core/sdk/speculation/dspark.py @@ -286,6 +286,7 @@ def _v4_block_ops(self, model, *, is_context: bool) -> list: cfg.moe_quant_mode, workload_distribution, cfg.attention_dp_size, + moe_kernel_source=cfg.moe_kernel_source, ), ops.MoEDispatch( f"dspark_{phase}_moe_post_dispatch", 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 a921e8802..d7ece1d63 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -96,6 +96,7 @@ def test_stable_function_signatures() -> None: "kvcache_quant_mode: 'str | None' = None, fmha_quant_mode: 'str | None' = None, " "fpm_fmha_quant_mode: 'str | None' = None, " "comm_quant_mode: 'str | None' = None, attention_backend: 'str | None' = None, " + "moe_kernel_source: 'str | None' = None, " "moe_backend: 'str | None' = None, enable_eplb: 'bool' = False, wideep_num_slots: 'int | None' = None, " "nextn: 'int' = 0, " "speculation: 'dict | None' = None, " diff --git a/python/aisimulate/tests/unit/collector/test_case_generator.py b/python/aisimulate/tests/unit/collector/test_case_generator.py new file mode 100644 index 000000000..13141d13d --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_case_generator.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression coverage for model-specific collector case declarations.""" + +from __future__ import annotations + +import pytest +from collector.case_generator import get_common_moe_test_cases, get_sglang_moe_backend + +pytestmark = pytest.mark.unit + + +QWEN38_MAX_FP8 = "Qwen/Qwen3.8-2.4T-A95B-FP8" + + +def test_qwen38_fp8_recipe_pins_only_blackwell_moe_runner(monkeypatch): + """The FP8 artifact follows its serving recipe; bf16 keeps the base default.""" + monkeypatch.setenv("COLLECTOR_MODEL_PATH", QWEN38_MAX_FP8) + + cases = get_common_moe_test_cases(backend="sglang") + assert cases + + for sm_version in (100, 103): + assert {get_sglang_moe_backend(case, "fp8_block", sm_version) for case in cases} == {"flashinfer_trtllm"} + assert {get_sglang_moe_backend(case, "bfloat16", sm_version) for case in cases} == {"triton"} diff --git a/python/aisimulate/tests/unit/collector/test_model_cases.py b/python/aisimulate/tests/unit/collector/test_model_cases.py index db6544f95..7d89c40af 100644 --- a/python/aisimulate/tests/unit/collector/test_model_cases.py +++ b/python/aisimulate/tests/unit/collector/test_model_cases.py @@ -1967,7 +1967,7 @@ def test_qwen38_max_nvfp4_moe_cases_are_declared_with_correct_shape_and_runner(m base_sample = base_cases[0] assert get_sglang_moe_backend(base_sample, "bfloat16", 100) == "triton" - assert get_sglang_moe_backend(base_sample, "fp8_block", 100) == "triton" + assert get_sglang_moe_backend(base_sample, "fp8_block", 100) == "flashinfer_trtllm" assert moe_model_allows_quantization("sglang", "RadixArk/Qwen3.8-2.4T-A95B-NVFP4", "nvfp4") assert not moe_model_allows_quantization("sglang", "RadixArk/Qwen3.8-2.4T-A95B-NVFP4", "bfloat16") 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 d02b2b992..d2c826498 100644 --- a/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py +++ b/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py @@ -243,11 +243,11 @@ def _fetch(database, attribute, *, shared_layer=None): monkeypatch.setattr(_etv, "fetch_attention_lane_density", _fetch) -def test_engine_spec_schema_version_is_twenty(): - """Verification width fields extend the positional operation payload.""" +def test_engine_spec_schema_version_is_twenty_one(): + """Exact MoE source extends the payload after the schema-20 FPM selector.""" from aisimulate.sdk import engine - assert engine.ENGINE_SPEC_SCHEMA_VERSION == 20 + assert engine.ENGINE_SPEC_SCHEMA_VERSION == 21 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 index 2307f09c1..2a19fb8f8 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py +++ b/python/aisimulate/tests/unit/sdk/models/test_deepseek_v41.py @@ -138,6 +138,36 @@ def _build_model(*, replay=False, backend="sglang", tp=4): ) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("source", [None, "sglang_flashinfer_trtllm_moe"]) +def test_moe_kernel_source_reaches_every_stage_child(replay, source): + import json + + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.models import get_model + + model = get_model( + MODEL_PATH, + ModelConfig(tp_size=4, moe_tp_size=1, moe_ep_size=4, decoder_replay=replay, moe_kernel_source=source), + "sglang", + ) + + def moe_specs(value): + if isinstance(value, dict): + if "Moe" in value: + yield value["Moe"] + for child in value.values(): + yield from moe_specs(child) + elif isinstance(value, list): + for child in value: + yield from moe_specs(child) + + for phase in (model.context_ops, model.generation_ops): + moes = [moe for op in phase for moe in moe_specs(json.loads(op._spec_json()))] + assert len(moes) == 40 + assert all(moe.get("moe_kernel_source") == source for moe in moes) + + @pytest.mark.parametrize("backend", ["sglang", "vllm", "trtllm"]) def test_v41_text_graph_full_profile_all_backends(backend): import json 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 c86350ddf..b1acaaec2 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_model_config.py +++ b/python/aisimulate/tests/unit/sdk/models/test_model_config.py @@ -16,6 +16,8 @@ import aisimulate.sdk.operations as ops from aisimulate.sdk import common, config, models +from aisimulate.sdk.config_builders import build_model_config +from aisimulate.sdk.errors import InvalidEngineConfigurationError from aisimulate.sdk.models import ( LLAMAModel, Qwen3VLModel, @@ -25,6 +27,7 @@ get_model_family, ) from aisimulate.sdk.performance_result import PerformanceResult +from aisimulate.sdk.speculation import SpeculationConfig from aisimulate.sdk.utils import get_model_config_from_model_path pytestmark = pytest.mark.unit @@ -49,6 +52,91 @@ def test_model_config_rejects_unknown_kernel_backend(field, value): config.ModelConfig(**{field: value}) +def test_build_model_config_preserves_existing_positional_arguments(): + speculation = SpeculationConfig(kind="mtp", params={"depth": 2}) + model_config = build_model_config( + 4, + 2, + 2, + 2, + 4, + "fp8", + "fp8", + "bfloat16", + "fp8", + "half", + "fpm", + False, + "fa3", + speculation, + "megamoe", + True, + 64, + ) + + assert (model_config.tp_size, model_config.pp_size, model_config.attention_dp_size) == (4, 2, 2) + assert (model_config.moe_tp_size, model_config.moe_ep_size) == (2, 4) + assert model_config.gemm_quant_mode is common.GEMMQuantMode.fp8 + assert model_config.kvcache_quant_mode is common.KVCacheQuantMode.fp8 + assert model_config.fmha_quant_mode is common.FMHAQuantMode.bfloat16 + assert model_config.moe_quant_mode is common.MoEQuantMode.fp8 + assert model_config.comm_quant_mode is common.CommQuantMode.half + assert model_config.forward_model == "fpm" + assert model_config.enable_encoder_dp is False + assert model_config.attention_backend is common.AttentionBackend.fa3 + assert model_config.speculation is speculation + assert model_config.moe_backend is common.MoEBackend.megamoe + assert model_config.enable_eplb is True + assert model_config.wideep_num_slots == 64 + assert model_config.moe_kernel_source is None + + +@pytest.mark.parametrize("source", [None, "sglang_flashinfer_trtllm_moe", " source_with_spaces "]) +def test_build_model_config_preserves_exact_kernel_source_keyword(source): + model_config = build_model_config(4, 1, 1, 4, 1, moe_kernel_source=source) + + assert model_config.moe_kernel_source == source + + +@pytest.mark.parametrize("source", ["", " ", "\t\n", "\u2003", 1, False, []]) +def test_model_config_rejects_invalid_kernel_source(source): + with pytest.raises(ValueError, match="moe_kernel_source must be a non-empty string"): + config.ModelConfig(moe_kernel_source=source) + + +@pytest.mark.parametrize( + ("model_path", "moe_backend", "message"), + [ + ("Qwen/Qwen3-32B", None, "require an MoE model"), + ("deepseek-ai/DeepSeek-V4-Pro", "megamoe", "MegaMoE"), + ], +) +def test_model_graph_rejects_moe_kernel_source_without_a_compatible_operator(model_path, moe_backend, message): + model_config = config.ModelConfig( + tp_size=1, + attention_dp_size=8, + moe_tp_size=1, + moe_ep_size=8, + moe_backend=moe_backend, + moe_kernel_source="missing_source", + ) + + with pytest.raises(InvalidEngineConfigurationError, match=message): + get_model(model_path, model_config, "sglang") + + +def test_model_graph_rejects_source_when_layer_override_removes_all_moe_operators(): + # Kimi-K3 has a dense first layer; its checkpoint metadata still declares MoE. + model_path = "moonshotai/Kimi-K3" + assert check_is_moe(model_path) + model_config = config.ModelConfig( + overwrite_num_layers=1, moe_tp_size=1, moe_ep_size=1, moe_kernel_source="missing_source" + ) + + with pytest.raises(InvalidEngineConfigurationError, match="moe_kernel_source.*no compatible MoE"): + get_model(model_path, model_config, "sglang") + + class TestSupportedModels: """Test default models configuration from support_matrix.csv.""" diff --git a/python/aisimulate/tests/unit/sdk/models/test_moe_block_builder_followups.py b/python/aisimulate/tests/unit/sdk/models/test_moe_block_builder_followups.py index 81161bf0d..1ffbb4b21 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_moe_block_builder_followups.py +++ b/python/aisimulate/tests/unit/sdk/models/test_moe_block_builder_followups.py @@ -23,6 +23,7 @@ import aisimulate_core.sdk.operations as ops from aisimulate_core.sdk import common, config +from aisimulate_core.sdk.errors import InvalidEngineConfigurationError from aisimulate_core.sdk.models import get_model from aisimulate_core.sdk.models.blocks.moe import MoEBlockShape, build_moe_block_ops @@ -59,6 +60,11 @@ def _dispatches(op_list): class TestGpusPerNodeGuard: + def test_large_ep_rejects_an_explicit_moe_kernel_source(self): + cfg = _cfg(moe_comm_backend={"context": "deepep_ht"}, moe_kernel_source="pinned_source") + with pytest.raises(InvalidEngineConfigurationError, match="moe_kernel_source.*large-EP"): + _build(cfg, gpus_per_node=8) + def test_large_ep_without_gpus_per_node_raises(self): cfg = _cfg(moe_comm_backend={"context": "deepep_ht"}) with pytest.raises(ValueError, match="gpus_per_node"): diff --git a/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py b/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py index 3dc039e4c..f5521ac86 100644 --- a/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py +++ b/python/aisimulate/tests/unit/sdk/task_v2/test_task_config.py @@ -405,6 +405,98 @@ def test_attention_backend_and_wideep_num_slots_reach_model_config(): assert mc.wideep_num_slots == 288 +@pytest.mark.parametrize("source", ["", " ", "\t\n", "\u2003", 1, False, []]) +def test_task_rejects_invalid_moe_kernel_source(source): + with pytest.raises(ValueError, match="moe_kernel_source must be a non-empty string"): + Task(moe_kernel_source=source) + + +@pytest.mark.parametrize("source", [None, "sglang_flashinfer_trtllm_moe", " source_with_spaces "]) +def test_task_round_trip_preserves_exact_moe_kernel_source(source): + task = Task(moe_kernel_source=source) + restored = Task.from_yaml(yaml.safe_load(task.to_yaml())) + + assert task.moe_kernel_source == source + assert restored.moe_kernel_source == source + assert restored.build_model_config(role="agg").moe_kernel_source == source + + +def test_moe_kernel_source_reaches_qwen38_compute_ops_and_compiled_spec(): + """A V2 Task pins only Qwen3.8's compute MoE lane, including the native spec.""" + from aisimulate.sdk import models + from aisimulate.sdk.engine import build_engine_spec_json + from aisimulate.sdk.operations.moe import MoE, MoEDispatch + from aisimulate.sdk.perf_database import get_database + + lane = "sglang_flashinfer_trtllm_moe" + task = Task( + serving_mode="agg", + model_path="Qwen/Qwen3.8-2.4T-A95B-FP8", + system_name="b200_sxm", + backend_name="sglang", + backend_version="0.5.17", + moe_kernel_source=lane, + ) + model_config = task.build_model_config(role="agg", parallel=(8, 1, 1, 1, 8, 1)) + assert model_config.moe_kernel_source == lane + + model = models.get_model(task.model_path, model_config, task.backend_name) + + def all_ops(ops): + pending = list(ops) + while pending: + op = pending.pop() + yield op + for group in ("_group_a", "_group_b", "_fallback"): + pending.extend(getattr(op, group, ()) or ()) + + ops = list(all_ops([*model.context_ops, *model.generation_ops])) + compute_ops = [op for op in ops if isinstance(op, MoE)] + dispatch_ops = [op for op in ops if isinstance(op, MoEDispatch)] + assert compute_ops + assert all(op._moe_kernel_source == lane for op in compute_ops) + assert all(not hasattr(op, "_moe_kernel_source") for op in dispatch_ops) + + database = get_database("b200_sxm", "sglang", "0.5.17") + spec = json.loads( + build_engine_spec_json( + model, + model_path=task.model_path, + system=task.system_name, + backend=task.backend_name, + backend_version=task.backend_version, + kv_block_size=None, + systems_path=None, + nextn=task.nextn, + database=database, + ) + ) + + def all_moe_specs(value): + if isinstance(value, list): + for item in value: + yield from all_moe_specs(item) + elif isinstance(value, dict): + if "Moe" in value: + yield value["Moe"] + for item in value.values(): + yield from all_moe_specs(item) + + moe_specs = list(all_moe_specs([spec["context_ops"], spec["generation_ops"]])) + assert spec["engine"]["moe_kernel_source"] == lane + assert moe_specs + assert all(moe["moe_kernel_source"] == lane for moe in moe_specs) + + default_config = Task( + serving_mode="agg", + model_path="Qwen/Qwen3.8-2.4T-A95B-FP8", + system_name="b200_sxm", + backend_name="sglang", + backend_version="0.5.17", + ).build_model_config(role="agg") + assert default_config.moe_kernel_source is None + + def test_invalid_attention_backend_rejected(): with pytest.raises(ValueError, match="attention_backend"): Task( diff --git a/python/aisimulate/tests/unit/sdk/task_v2/test_v1_compat.py b/python/aisimulate/tests/unit/sdk/task_v2/test_v1_compat.py index 39e82d921..799981a4f 100644 --- a/python/aisimulate/tests/unit/sdk/task_v2/test_v1_compat.py +++ b/python/aisimulate/tests/unit/sdk/task_v2/test_v1_compat.py @@ -227,6 +227,17 @@ def test_attention_backend_and_wideep_num_slots_map(self): assert out["attention_backend"] == "fa3" assert out["wideep_num_slots"] == 288 + def test_moe_kernel_source_maps_from_legacy_config(self): + """The recipe-pinned MoE lane remains explicit after V1 conversion.""" + out = convert_v1_to_v2( + { + "serving_mode": "agg", + "model_path": "Qwen/Qwen3.8-2.4T-A95B-FP8", + "config": {"moe_kernel_source": "sglang_flashinfer_trtllm_moe"}, + } + ) + assert out["moe_kernel_source"] == "sglang_flashinfer_trtllm_moe" + def test_multiple_profiles_raise(self): """Multiple V1 profiles drop all but the first -- a silent semantic change, so reject.""" with pytest.raises(ValueError) as exc: diff --git a/python/aisimulate/tests/unit/sdk/test_compile_engine_mtp.py b/python/aisimulate/tests/unit/sdk/test_compile_engine_mtp.py index 2492adf96..7bedb3b3e 100644 --- a/python/aisimulate/tests/unit/sdk/test_compile_engine_mtp.py +++ b/python/aisimulate/tests/unit/sdk/test_compile_engine_mtp.py @@ -112,6 +112,27 @@ def _capture_spec(model, **_kwargs): assert captured["model"].config.attention_backend == "fa3" +def test_compile_engine_propagates_moe_kernel_source_to_model_config(monkeypatch): + captured = {} + + def _capture_spec(model, **_kwargs): + captured["model"] = model + return "{}" + + monkeypatch.setattr(engine, "build_engine_spec_json", _capture_spec) + monkeypatch.setattr(engine, "_maybe_load_database", lambda *a, **k: None) + monkeypatch.setattr(engine.aisimulate_core, "engine_spec_bincode_from_json", lambda s: b"") + + engine.compile_engine( + "Qwen/Qwen3-30B-A3B", + "h200_sxm", + "trtllm", + moe_kernel_source="sglang_flashinfer_trtllm_moe", + ) + + assert captured["model"].config.moe_kernel_source == "sglang_flashinfer_trtllm_moe" + + def test_compile_engine_propagates_database_mode_to_database_view(monkeypatch): captured = {} diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py b/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py index 5ab01d5ae..46a197679 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_execution_identity.py @@ -434,9 +434,11 @@ def test_model_config_builder_preserves_existing_positional_controls(): True, 64, fpm_fmha_quant_mode="fp8", + moe_kernel_source="sglang_flashinfer_trtllm_moe", ) assert result.speculation is None assert result.moe_backend == "deepep_moe" assert result.enable_eplb is True assert result.wideep_num_slots == 64 assert result.fpm_fmha_quant_mode == FMHAQuantMode.fp8 + assert result.moe_kernel_source == "sglang_flashinfer_trtllm_moe" diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py index 4ab7a07af..dc13974a8 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_forward.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_forward.py @@ -1076,8 +1076,10 @@ def test_canonical_config_preserves_positional_quantization_fields(): "bfloat16", 0, fpm_fmha_quant_mode="fp8", + moe_kernel_source="sglang_flashinfer_trtllm_moe", ) assert config.kvcache_quant_mode == "fp8" assert config.comm_quant_mode == "bfloat16" assert config.nextn == 0 assert config.fpm_fmha_quant_mode == "fp8" + assert config.moe_kernel_source == "sglang_flashinfer_trtllm_moe" diff --git a/python/aisimulate/tests/unit/sdk/test_moe_kernel_source_consumers.py b/python/aisimulate/tests/unit/sdk/test_moe_kernel_source_consumers.py new file mode 100644 index 000000000..0806aea9a --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_moe_kernel_source_consumers.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exact MoE source survives current CLI, Sweeper and Replay boundaries.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest + +from aisimulate.compiler import _worker_engine_args +from aisimulate.config.engine import EnginePredictionConfig +from aisimulate.runner import _materialize_engine_role +from aisimulate.sweeper.config import SearchSpace +from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver +from aisimulate_core.sdk import RustForwardPassPerfModel +from aisimulate_core.sdk.errors import InvalidEngineConfigurationError + +pytestmark = pytest.mark.unit +SOURCE = "sglang_flashinfer_trtllm_moe" + + +@pytest.mark.parametrize("source", [None, SOURCE]) +def test_cli_saved_config_and_replay_preserve_source(source): + engine = EnginePredictionConfig( + model="Qwen/Qwen3-30B-A3B", + hardware="b200_sxm", + backend="sglang", + backend_version="0.5.17", + moe_kernel_source=source, + workers={"aggregated": {"kv_cache": {"capacity": {"type": "fixed", "blocks": 100}}}}, + ) + restored = EnginePredictionConfig.model_validate_json(engine.model_dump_json()) + assert restored.moe_kernel_source == source + payload = _worker_engine_args(restored, restored.workers.aggregated, "aggregated", transfer_bytes_per_token=None) + config = payload["timing_model"]["config"] + assert config.get("moe_kernel_source") == source + assert ("moe_kernel_source" in config) == (source is not None) + role = _materialize_engine_role("sglang", "0.5.17", {}, payload, "aggregated") + assert role["rank"]["timing_model"]["config"].get("moe_kernel_source") == source + + +@pytest.mark.parametrize("field", ["moe_kernel_source", "aic_moe_kernel_source"]) +def test_flat_replay_source_reaches_timing(field): + role = _materialize_engine_role( + "sglang", + "0.5.17", + {}, + {"aic_model_path": "Qwen/Qwen3-30B-A3B", "aic_system": "b200_sxm", "num_gpu_blocks": 100, field: SOURCE}, + "aggregated", + ) + assert role["rank"]["timing_model"]["config"]["moe_kernel_source"] == SOURCE + assert field not in role["rank"] + + +def test_sweeper_source_is_round_tripped_resolved_and_cached(monkeypatch): + space = SearchSpace( + model_name="Qwen/Qwen3-30B-A3B", hardware_sku="b200_sxm", backend=["sglang"], moe_kernel_source=SOURCE + ) + space = SearchSpace.model_validate_json(space.model_dump_json()) + resolver = ForwardPassEstimatorResolver(space) + sample = { + "backend": "sglang", + "deployment_mode": "agg", + "hardware_sku": "b200_sxm", + "tp": 1, + "pp": 1, + "attention_dp": 1, + "moe_tp": 1, + "moe_ep": 1, + "agg_block_size": 1, + } + calls = [] + + def best_available(request): + calls.append(request) + resolved = request.to_dict() | {"backend_version": "0.5.17"} + return SimpleNamespace( + diagnostics=lambda: {"readiness": "ready", "provenance": {"config": resolved}}, close=lambda: None + ) + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", best_available) + request = resolver._request(sample, "agg") + assert request.moe_kernel_source == SOURCE + assert resolver.resolve_candidate(sample)["agg"].config["moe_kernel_source"] == SOURCE + resolver.resolve_candidate(sample) + assert len(calls) == 1 + resolver._resolve(replace(request, moe_kernel_source=None), "agg") + assert len(calls) == 2 + + +def test_source_rejects_custom_timing_in_current_config(): + with pytest.raises(ValueError, match="engine model controls require default timing"): + EnginePredictionConfig( + model="Qwen/Qwen3-30B-A3B", + hardware="b200_sxm", + moe_kernel_source=SOURCE, + workers={"aggregated": {"timing": {"type": "fixed", "prefill_ms": 1, "decode_ms": 1}}}, + ) + + +def test_sweeper_capacity_preserves_source_through_model_config(monkeypatch): + from aisimulate.sweeper import kv_estimate + from aisimulate.sweeper.parallel_enum import ParallelShape + from aisimulate.sweeper.search_space import _engine_memory_kwargs + from aisimulate_core.sdk import memory + + space = SearchSpace(model_name="Qwen/Qwen3-30B-A3B", hardware_sku="b200_sxm", moe_kernel_source=SOURCE) + captured = {} + original = memory.build_model_config + + def build_config(**kwargs): + captured["config"] = original(**kwargs) + raise RuntimeError("stop before database access") + + monkeypatch.setattr(memory, "build_model_config", build_config) + with pytest.raises(ValueError, match="stop before database access"): + kv_estimate.estimate_kv_tokens( + ParallelShape(1, 1, 1, 1), + model_name=space.model_name, + hardware_sku=space.hardware_sku, + backend="sglang", + backend_version="0.5.17", + **_engine_memory_kwargs(space), + ) + assert captured["config"].moe_kernel_source == SOURCE + + +def test_explicit_source_prevents_naive_capacity_fallback(monkeypatch): + from aisimulate_core.sdk import memory + + def unavailable(*args, **kwargs): + assert kwargs["moe_kernel_source"] == SOURCE + raise RuntimeError("unavailable model") + + monkeypatch.setattr(memory.KVCacheEstimator, "from_request", unavailable) + with pytest.raises(ValueError, match="unavailable model"): + memory.estimate_num_gpu_blocks( + "Qwen/Qwen3-30B-A3B", + "b200_sxm", + "sglang", + scheduler_block_size=1, + max_num_tokens=1, + max_batch_size=1, + memory_fraction_kind="of_total", + memory_fraction_value=0.9, + moe_kernel_source=SOURCE, + allow_naive_fallback=True, + ) + + +@pytest.mark.parametrize("entry_point", ["compile", "native_memory", "memory_with_fallback"]) +@pytest.mark.parametrize( + ("model_path", "moe_backend", "message"), + [ + ("Qwen/Qwen3-32B", None, "require an MoE model"), + ("deepseek-ai/DeepSeek-V4-Pro", "megamoe", "MegaMoE"), + ], +) +def test_unsupported_source_is_rejected_before_database_or_capacity_fallback( + entry_point, model_path, moe_backend, message +): + from aisimulate_core.sdk import engine, memory + + controls = {"moe_backend": moe_backend, "moe_kernel_source": SOURCE} + with pytest.raises(InvalidEngineConfigurationError, match=message): + if entry_point == "compile": + engine.compile_engine(model_path, "missing-system", "sglang", **controls) + elif entry_point == "native_memory": + memory.KVCacheEstimator.from_request( + model_path, "missing-system", "sglang", max_num_tokens=1, max_batch_size=1, **controls + ) + else: + memory.estimate_num_gpu_blocks( + model_path, + "missing-system", + "sglang", + scheduler_block_size=1, + max_num_tokens=1, + max_batch_size=1, + memory_fraction_kind="of_total", + memory_fraction_value=0.9, + allow_naive_fallback=True, + **controls, + ) + + +@pytest.mark.parametrize("estimation_mode", ["auto", "op_level"]) +@pytest.mark.parametrize( + ("model_path", "moe_backend", "message"), + [ + ("Qwen/Qwen3-32B", None, "require an MoE model"), + ("deepseek-ai/DeepSeek-V4-Pro", "megamoe", "MegaMoE"), + ], +) +def test_canonical_source_validation_cannot_fall_back(estimation_mode, model_path, moe_backend, message): + with pytest.raises(ValueError, match=f"invalid engine config:.*{message}"): + RustForwardPassPerfModel.best_available( + { + "model": model_path, + "system": "b200_sxm", + "backend": "sglang", + "worker_type": "aggregated", + "estimation_mode": estimation_mode, + "fallback_policy": "allow", + "moe_backend": moe_backend, + "moe_kernel_source": SOURCE, + } + ) + + +def test_task_cannot_rewrite_an_explicit_moe_source_into_whole_forward_fpm(): + from aisimulate.sdk.task_v2 import Task + from aisimulate_core.sdk.models import get_model + + task = Task(model_path="Qwen/Qwen3-30B-A3B", forward_model="fpm", moe_kernel_source=SOURCE) + model_config = task.build_model_config(role="agg") + with pytest.raises(InvalidEngineConfigurationError, match="moe_kernel_source.*forward_model='fpm'"): + get_model(task.model_path, model_config, "sglang") + + +@pytest.mark.parametrize( + ("forward_model", "message"), + [ + ("op_level", "fpm_fmha_quant_mode requires forward_model='fpm'"), + ("fpm", "moe_kernel_source is not supported with forward_model='fpm'"), + ], +) +def test_independent_selectors_cannot_silently_discard_each_other(forward_model, message): + from aisimulate_core.sdk.config_builders import build_model_config + from aisimulate_core.sdk.models import get_model + + config = build_model_config( + 1, + 1, + 1, + 1, + 1, + forward_model=forward_model, + moe_kernel_source=SOURCE, + fpm_fmha_quant_mode="fp8", + ) + with pytest.raises(InvalidEngineConfigurationError, match=message): + get_model("Qwen/Qwen3-30B-A3B", config, "sglang") diff --git a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py index 66a1223d2..509e6de47 100644 --- a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py +++ b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py @@ -935,6 +935,62 @@ def test_forward_pass_config_requires_role_and_defaults_to_auto_deny() -> None: assert config.estimation_mode == "auto" assert config.fallback_policy == "deny" assert config.estimator_config == {} + assert config.moe_kernel_source is None + + pinned = ForwardPassPerfModelConfig( + model="m", + system="s", + backend="vllm", + worker_type="decode", + moe_kernel_source="sglang_flashinfer_trtllm_moe", + ) + assert pinned.to_dict()["moe_kernel_source"] == "sglang_flashinfer_trtllm_moe" + + +def test_forward_pass_config_preserves_existing_positional_arguments(tmp_path: Path) -> None: + from aisimulate_core.sdk import ForwardPassPerfModelConfig + + # Positional order frozen from main at e8828036e3d3, before moe_kernel_source was added. + legacy_fields = { + "model": "Qwen/Qwen3-30B-A3B", + "system": "b200_sxm", + "backend": "sglang", + "worker_type": "decode", + "backend_version": "0.5.17", + "tp": 4, + "pp": 2, + "attention_dp": 2, + "moe_tp_size": 1, + "moe_ep_size": 4, + "gemm_quant_mode": "fp8", + "moe_quant_mode": "fp8_block", + "fmha_quant_mode": "bfloat16", + "kvcache_quant_mode": "fp8", + "comm_quant_mode": "half", + "nextn": 2, + "speculation": {"kind": "mtp", "params": {"depth": 2}}, + "kv_block_size": 64, + "decoder_replay": True, + "estimation_mode": "fpm_regression", + "database_mode": "EMPIRICAL", + "transfer_policy": ("gemm", "moe"), + "systems_paths": (str(tmp_path),), + "fallback_policy": "allow", + "estimator_config": {"correction": {"enabled": False}}, + "attention_backend": "fa3", + "enable_shared_layer": False, + "strict_provenance": True, + "moe_backend": "megamoe", + "enable_eplb": True, + "wideep_num_slots": 64, + } + config = ForwardPassPerfModelConfig(*legacy_fields.values()) + assert vars(config) == {**legacy_fields, "fpm_fmha_quant_mode": None, "moe_kernel_source": None} + + source = " source_with_spaces " + pinned = ForwardPassPerfModelConfig(*legacy_fields.values(), moe_kernel_source=source) + assert vars(pinned) == {**legacy_fields, "fpm_fmha_quant_mode": None, "moe_kernel_source": source} + assert json.loads(json.dumps(pinned.to_dict()))["moe_kernel_source"] == source def _supported_fpm_config() -> dict[str, object]: @@ -1343,14 +1399,14 @@ def test_sparse_cp_ops_emit_cp_fields_in_spec(): assert spec["attn_kind"] == "Csa" -def test_engine_config_json_identity_disambiguates_collapsed_quant_modes(): +def test_engine_config_json_identity_disambiguates_collapsed_quant_modes_and_moe_lanes(): """Two models differing only in a wire-collapsed dtype (sq vs int8_wo both -> "int8") or an identity-omitted ModelConfig field (moe_backend) must get DISTINCT handle-cache keys — sharing one cached handle silently returns the other model's latencies.""" from aisimulate.sdk import common - def _model(gemm_mode, moe_backend=None): + def _model(gemm_mode, moe_backend=None, moe_kernel_source=None): cfg = SimpleNamespace( tp_size=8, pp_size=1, @@ -1365,6 +1421,7 @@ def _model(gemm_mode, moe_backend=None): comm_quant_mode=None, moe_backend=moe_backend, attention_backend=None, + moe_kernel_source=moe_kernel_source, # enable_wideep dropped from the fixture: the deprecated flag left # the engine identity (constant False; moe_comm_backend + # num_gpus_per_node carry the regime). @@ -1387,6 +1444,11 @@ def _model(gemm_mode, moe_backend=None): ) assert key_sq != key_deepep, "moe_backend must participate in the cache identity" + key_flashinfer = rust_engine_step._engine_config_json( + _model(common.GEMMQuantMode.sq, moe_kernel_source="sglang_flashinfer_trtllm_moe"), database + ) + assert key_sq != key_flashinfer, "a pinned MoE lane must not reuse the default handle" + def test_engine_config_json_identity_includes_database_policy(): """Two views of the SAME on-disk identity that differ only in the diff --git a/tests/sweeper/test_engine_request.py b/tests/sweeper/test_engine_request.py index 79f38c29f..689d86353 100644 --- a/tests/sweeper/test_engine_request.py +++ b/tests/sweeper/test_engine_request.py @@ -22,6 +22,7 @@ "enable_eplb": True, "wideep_num_slots": 256, "moe_backend": "deepep_moe", + "moe_kernel_source": "sglang_flashinfer_trtllm_moe", "attention_backend": "fa3", "gemm_quant_mode": "fp8", "moe_quant_mode": "fp8", diff --git a/tests/test_afd_runner.py b/tests/test_afd_runner.py index 4eac4f198..21526af31 100644 --- a/tests/test_afd_runner.py +++ b/tests/test_afd_runner.py @@ -315,11 +315,21 @@ def test_single_output_token_completes_at_prefill_without_decode(phase, companio assert record["tpot_ms"] == 0.0 -def test_default_companion_model_consumes_fixed_timing_without_aic_lookup(): +@pytest.mark.parametrize( + "overrides", + [ + {}, + {"moe_kernel_source": None}, + {"aic_moe_kernel_source": None}, + {"attention_backend": "fa3", "enable_eplb": True}, + ], +) +def test_default_companion_model_consumes_fixed_timing_without_aic_lookup(overrides): spec = _spec( _topology(phase="prefill", combined_with_pd=True), companion_role="decode", ) + spec.backend_deployment.decode_engine_args.update(overrides) timing = AICAFDCompanionPerformanceModel().measure(spec) @@ -329,6 +339,58 @@ def test_default_companion_model_consumes_fixed_timing_without_aic_lookup(): assert timing.provenance["provider"] == "fixed" +@pytest.mark.parametrize(("phase", "companion_role"), [("decode", "prefill"), ("prefill", "decode")]) +@pytest.mark.parametrize("field", ["moe_kernel_source", "aic_moe_kernel_source"]) +@pytest.mark.parametrize("custom_timing", [False, True]) +def test_aic_companion_legacy_estimator_rejects_an_explicit_moe_source(phase, companion_role, field, custom_timing): + 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") + if custom_timing: + engine_args["timing_model"] = {"type": "polynomial"} + engine_args.update(aic_model_path="Qwen/Qwen3-30B-A3B", aic_system="b200_sxm") + engine_args[field] = "sglang_flashinfer_trtllm_moe" + + def estimator(*args, **kwargs): + pytest.fail("an explicit source must not reach an estimator that cannot consume it") + + with pytest.raises(InvalidRunnerError, match="moe_kernel_source.*not supported.*AFD companion"): + AICAFDCompanionPerformanceModel(estimator).measure(spec) + + +@pytest.mark.parametrize(("phase", "companion_role"), [("decode", "prefill"), ("prefill", "decode")]) +@pytest.mark.parametrize("field", ["moe_kernel_source", "aic_moe_kernel_source"]) +def test_aic_companion_external_fpm_preserves_source_for_canonical_rejection(phase, companion_role, field, tmp_path): + spec = _spec(_topology(phase=phase, combined_with_pd=True), companion_role=companion_role) + spec = replace( + spec, backend_deployment=replace(spec.backend_deployment, backend="sglang", backend_version="0.5.17") + ) + spec.backend_deployment.parallel_config[f"{companion_role}_tp"] = 1 + engine_args = getattr(spec.backend_deployment, f"{companion_role}_engine_args") + engine_args.pop("timing_model") + engine_args.update( + aic_model_path="Qwen/Qwen3-30B-A3B", + aic_system="b200_sxm", + aic_forward_model="fpm", + aic_fpm_parquet_path=str(tmp_path / "not-read.parquet"), + ) + engine_args[field] = "sglang_flashinfer_trtllm_moe" + + with pytest.raises(InvalidRunnerError, match="fpm_interpolation.*moe_kernel_source overrides"): + AICAFDCompanionPerformanceModel().measure(spec) + + +@pytest.mark.parametrize(("phase", "companion_role"), [("decode", "prefill"), ("prefill", "decode")]) +@pytest.mark.parametrize("field", ["moe_kernel_source", "aic_moe_kernel_source"]) +def test_aic_companion_fixed_timing_rejects_an_explicit_moe_source(phase, companion_role, field): + 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[field] = "sglang_flashinfer_trtllm_moe" + + with pytest.raises(ValueError, match="moe_kernel_source.*not supported.*fixed AFD companion timing"): + AICAFDCompanionPerformanceModel().measure(spec) + + @pytest.mark.parametrize(("phase", "companion_role"), [("decode", "prefill"), ("prefill", "decode")]) @pytest.mark.parametrize("forward_model", ["fpm", "op_level", None]) @pytest.mark.parametrize("field", ["forward_model", "aic_forward_model"])