diff --git a/crates/core/src/perfmodel/config.rs b/crates/core/src/perfmodel/config.rs index 2cc4e44b3c..38e8525060 100644 --- a/crates/core/src/perfmodel/config.rs +++ b/crates/core/src/perfmodel/config.rs @@ -122,7 +122,13 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // - 25 (FPM decoupling): FpmForwardOp also carries the SOL/direct interpolation // selector. The decoupling branch independently used 21 for this positional // field; the combined layout differs from every prior schema, including 24. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 25; +// - 26 (decode context parallelism): the context/generation attention, MLA, +// MLA-module, wide-EP MLA and DSA ops gained a tail-appended `dcp_size` +// (gathered query heads over a 1/dcp KV stripe; striped-context gather). +// Claimed 19 through 25 on its own branch while the DeepSeek-V4.1, MoE +// kernel-source, VR200 pilot, FPM DCP and FPM decoupling changes landed; +// renumbered at each merge (precedent: 15, 18). +pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 26; /// Static engine identity and setup information carried by an /// [`crate::perfmodel::engine::spec::EngineSpec`]. @@ -269,6 +275,12 @@ pub struct ParallelMapping { /// re-derived from this field. #[serde(default)] pub cp_size: Option, + /// Decode-context-parallel size (vLLM `-dcp` / SGLang `--dcp-size`): the + /// decode KV cache is striped by token position across ranks inside the + /// attention group. Part of the engine identity so dcp variants get + /// distinct compiled handles. `None`/1 means no DCP. Like `cp_size`, the + /// per-op math is carried on the ops themselves, not re-derived here. + /// Additive-optional: absent in older payloads. #[serde(default)] pub dcp_size: Option, } diff --git a/crates/core/src/perfmodel/engine/runtime.rs b/crates/core/src/perfmodel/engine/runtime.rs index 9b2c7300f2..27e8d48a79 100644 --- a/crates/core/src/perfmodel/engine/runtime.rs +++ b/crates/core/src/perfmodel/engine/runtime.rs @@ -2612,6 +2612,7 @@ mod tests { cp_size: 1, lane_order: crate::operators::attention::b200_vllm_context_lane_order(), apply_rope: true, + dcp_size: 1, }), ] } @@ -2637,6 +2638,7 @@ mod tests { use_qk_norm: false, scale_num_tokens: 1, verify_query_tokens: 0, + dcp_size: 1, }), ] } @@ -2657,13 +2659,13 @@ mod tests { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { - dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(1), moe_tp_size: Some(1), moe_ep_size: Some(8), cp_size: None, + dcp_size: None, }, quantization: QuantizationConfig { weight_dtype: None, @@ -2728,13 +2730,13 @@ mod tests { config.backend = BackendKind::Sglang; config.backend_version = Some("0.5.18+nvinternal.rubin.0.8full.66997102".into()); config.parallel = ParallelMapping { - dcp_size: None, tp_size: 4, pp_size: 1, attention_dp_size: Some(1), moe_tp_size: Some(4), moe_ep_size: Some(1), cp_size: Some(1), + dcp_size: None, }; config.quantization = QuantizationConfig { weight_dtype: Some(DataType::Bfloat16), diff --git a/crates/core/src/perfmodel/engine/spec.rs b/crates/core/src/perfmodel/engine/spec.rs index 7e62af7eb3..28126bf67e 100644 --- a/crates/core/src/perfmodel/engine/spec.rs +++ b/crates/core/src/perfmodel/engine/spec.rs @@ -232,6 +232,7 @@ mod tests { // survives, not just a single-element degenerate case. lane_order: vec!["trtllm_mha".into(), "flashinfer".into(), "default".into()], apply_rope: false, + dcp_size: 3, } } @@ -248,6 +249,7 @@ mod tests { use_qk_norm: true, scale_num_tokens: 8, verify_query_tokens: 7, + dcp_size: 4, } } @@ -270,6 +272,7 @@ mod tests { kv_cache_dtype: KvCacheQuantMode::Bfloat16, fmha_quant_mode: FmhaQuantMode::Bfloat16, cp_size: 1, + dcp_size: 4, } } @@ -279,6 +282,7 @@ mod tests { scale_factor: 1.0, num_heads: 128, kv_cache_dtype: KvCacheQuantMode::Fp8, + dcp_size: 8, } } @@ -291,6 +295,7 @@ mod tests { fmha_quant_mode: FmhaQuantMode::Fp8, gemm_quant_mode: GemmQuantMode::Fp8Block, native_num_heads: Some(128), + dcp_size: 2, } } @@ -409,6 +414,7 @@ mod tests { cp_size: 1, full_frac: 1.0, attn_projection_quant_modes: None, + dcp_size: 2, } } @@ -564,6 +570,7 @@ mod tests { kv_cache_dtype: KvCacheQuantMode::Fp8, fmha_quant_mode: FmhaQuantMode::Fp8, attn_backend: "flashinfer".into(), + dcp_size: 4, } } @@ -852,13 +859,13 @@ mod tests { moe_kernel_source: None, kv_block_size: Some(64), parallel: ParallelMapping { - dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(8), moe_tp_size: Some(1), moe_ep_size: Some(8), cp_size: None, + dcp_size: None, }, quantization: QuantizationConfig { weight_dtype: Some(DataType::Fp8), @@ -1146,9 +1153,10 @@ mod tests { let selector = bytes.len() - 14; assert_eq!(&bytes[selector..selector + 4], &1u32.to_le_bytes()); // Recreate both branches' incompatible payloads, not just their stamps. - for version in [20u32, 21, 22, 23, 24] { + for version in [20u32, 21, 22, 23, 24, 25] { let mut stale = bytes.clone(); - if version != 21 { + // The decoupling branch's 21 and the merged 25 both carried the selector. + if !matches!(version, 21 | 25) { stale.drain(selector..selector + 4); } if version < 24 { @@ -1372,21 +1380,24 @@ mod tests { .as_object_mut() .unwrap() .remove("verify_query_tokens"); + attention_json.as_object_mut().unwrap().remove("dcp_size"); let attention: GenerationAttentionOp = serde_json::from_value(attention_json).unwrap(); assert_eq!(attention.scale_num_tokens, 1); assert_eq!(attention.verify_query_tokens, 0); + // Pre-v26 producers never emitted decode CP: it defaults to "off". + assert_eq!(attention.dcp_size, 1); let mut fpm_json = serde_json::to_value(fpm_forward()).unwrap(); fpm_json.as_object_mut().unwrap().remove("verify_width"); let fpm: crate::operators::FpmForwardOp = serde_json::from_value(fpm_json).unwrap(); assert_eq!(fpm.verify_width, 1); let mut bytes = handshake_spec().to_bincode().unwrap(); - bytes[..4].copy_from_slice(&17u32.to_le_bytes()); + bytes[..4].copy_from_slice(&18u32.to_le_bytes()); bytes.truncate(4); assert!(matches!( EngineSpec::from_bincode(&bytes), Err(AicError::UnsupportedSchemaVersion { - got: 17, + got: 18, expected: ENGINE_SPEC_SCHEMA_VERSION, .. }) @@ -1424,7 +1435,7 @@ mod tests { for spec in [fpm_spec, pilot_spec] { let bytes = spec.to_bincode().unwrap(); assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec); - for previous_version in [20u32, 21, 22, 23, 24] { + for previous_version in [20u32, 21, 22, 23, 24, 25] { let mut stale = bytes.clone(); stale[..4].copy_from_slice(&previous_version.to_le_bytes()); // DCP and the pilot claimed 22 for different layouts. diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 7a7e65e9bd..a56e37df6e 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -104,6 +104,11 @@ pub struct ForwardPassPerfModelConfig { pub moe_tp_size: Option, #[serde(default)] pub moe_ep_size: Option, + /// Prefill context parallelism (SGLang `--attn-cp-size`, vLLM `-pcp`): extra + /// attention ranks that split prefill tokens, so it widens the attention side + /// (`tp * attention_dp * cp_size == moe_tp_size * moe_ep_size`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cp_size: Option, #[serde(default, alias = "gemm_dtype")] pub gemm_quant_mode: Option, @@ -181,6 +186,7 @@ impl ForwardPassPerfModelConfig { dcp: None, moe_tp_size: None, moe_ep_size: None, + cp_size: None, gemm_quant_mode: None, moe_quant_mode: None, fmha_quant_mode: None, @@ -363,6 +369,11 @@ impl ForwardPassPerfModelConfig { if self.tp == 0 || self.pp == 0 || self.attention_dp == 0 { return Err(invalid_config("tp, pp, and attention_dp must be positive")); } + if self.cp_size == Some(0) { + return Err(invalid_config("cp_size must be positive")); + } + // Decode CP stripes the KV inside the TP group: it must divide tp and + // never enters the width identity below. if self.dcp.is_some_and(|dcp| dcp == 0 || self.tp % dcp != 0) { return Err(invalid_config("dcp must be positive and divide tp")); } @@ -385,14 +396,18 @@ impl ForwardPassPerfModelConfig { "moe_tp_size and moe_ep_size must be positive", )); } - // Profile schema also admits dense TP with unused MoE dimensions - // equal to one, and checks the exact requested deployment below. + // Prefill CP adds attention ranks; decode CP reuses them, so only + // `cp_size` folds into the attention width (same identity as the + // replay-side `AicTimingConfig::validate_parallel_shape`). Profile + // schema also admits dense TP with unused MoE dimensions equal to + // one, and checks the exact requested deployment below. + let cp = u64::from(self.cp_size.unwrap_or(1)); if self.fpm_profile.is_none() - && u64::from(self.tp) * u64::from(self.attention_dp) + && u64::from(self.tp) * u64::from(self.attention_dp) * cp != u64::from(moe_tp) * u64::from(moe_ep) { return Err(invalid_config( - "topology requires tp * attention_dp == moe_tp_size * moe_ep_size", + "topology requires tp * attention_dp * cp_size == moe_tp_size * moe_ep_size", )); } } @@ -880,6 +895,38 @@ mod tests { ); } + #[test] + fn context_parallel_knobs_fold_prefill_cp_only_and_round_trip() { + // Unset knobs stay out of the serialized identity (pre-CP configs are + // byte-identical) and deserialize back to None. + let plain = config(serde_json::json!({})); + let json = serde_json::to_string(&plain).unwrap(); + assert!(!json.contains("cp_size"), "{json}"); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + plain + ); + + // Prefill CP widens the attention side to match the MoE width ... + let mut cfg = config(serde_json::json!({ + "moe_tp_size": 1, "moe_ep_size": 8, "cp_size": 8, "dcp": 1 + })); + cfg.validate().unwrap(); + let json = serde_json::to_string(&cfg).unwrap(); + assert!(json.contains("\"cp_size\":8") && json.contains("\"dcp\":1")); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + cfg + ); + // ... decode CP does not: without prefill CP the widths no longer match. + cfg.cp_size = None; + assert!(cfg.validate().is_err()); + // Zero is rejected like every other parallel size. + cfg.cp_size = Some(8); + cfg.dcp = Some(0); + assert!(cfg.validate().is_err()); + } + #[test] fn role_and_typed_estimator_controls_are_validated() { assert!( diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index 522cc7551d..690812fcb3 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -354,9 +354,13 @@ impl ForwardPassPerfModel { last_error = Some(error); continue; } + // Op-level DCP is priced by the sharded attention ops and their + // collectives; the whole-forward paths need a measured vLLM cell + // (the regression tables carry no DCP column). if config.dcp.is_some_and(|dcp| dcp > 1) - && (mode != EstimationMode::FpmInterpolation - || config.backend != crate::BackendKind::Vllm) + && (mode == EstimationMode::FpmRegression + || (mode == EstimationMode::FpmInterpolation + && config.backend != crate::BackendKind::Vllm)) { let error = AicError::UnsupportedModel( "DCP timing requires measured vLLM FPM interpolation".into(), @@ -429,8 +433,17 @@ impl ForwardPassPerfModel { Err(err) => return Err(err), } } - Err(last_error - .unwrap_or_else(|| AicError::UnsupportedModel("no estimator candidates".into()))) + // Every candidate mode failed. Surface all of them: the last one is + // usually a whole-forward gate (for instance the DCP FPM rule) that + // would otherwise mask the op-level reason the user actually hit. + Err(match failures.len() { + 0 => AicError::UnsupportedModel("no estimator candidates".into()), + 1 => last_error.expect("one recorded failure"), + _ => AicError::UnsupportedModel(format!( + "no estimator mode accepted this configuration: {}", + failures.join("; ") + )), + }) } /// API: diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index 9e4bec62de..63b29d03c2 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -89,6 +89,7 @@ fn context_ops() -> Vec { cp_size: 1, lane_order: crate::operators::attention::b200_vllm_context_lane_order(), apply_rope: true, + dcp_size: 1, }), ] } @@ -114,6 +115,7 @@ fn generation_ops() -> Vec { use_qk_norm: false, scale_num_tokens: 1, verify_query_tokens: 0, + dcp_size: 1, }), ] } @@ -134,13 +136,13 @@ fn fixture_engine_config() -> EngineConfig { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { - dcp_size: None, tp_size: 8, pp_size: 1, attention_dp_size: Some(1), moe_tp_size: Some(1), moe_ep_size: Some(8), cp_size: None, + dcp_size: None, }, quantization: QuantizationConfig { weight_dtype: None, diff --git a/crates/core/src/perfmodel/memory.rs b/crates/core/src/perfmodel/memory.rs index 20938c69ed..b0f9b58637 100644 --- a/crates/core/src/perfmodel/memory.rs +++ b/crates/core/src/perfmodel/memory.rs @@ -543,13 +543,13 @@ mod tests { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { - dcp_size: None, tp_size: 1, pp_size: 1, attention_dp_size: Some(1), moe_tp_size: None, moe_ep_size: None, cp_size: None, + dcp_size: None, }, quantization: QuantizationConfig { weight_dtype: None, diff --git a/crates/core/src/perfmodel/operators/attention.rs b/crates/core/src/perfmodel/operators/attention.rs index 4594edfbc1..53f4960557 100644 --- a/crates/core/src/perfmodel/operators/attention.rs +++ b/crates/core/src/perfmodel/operators/attention.rs @@ -193,6 +193,43 @@ pub struct ContextAttentionOp { /// Appended at the struct tail because bincode payloads are positional. #[serde(default = "default_apply_rope")] pub apply_rope: bool, + /// Decode context parallelism on the SAME engine (aggregated serving): + /// the persistent KV is striped across `dcp_size` ranks, so a prefill + /// with cached context (`prefix > 0`) must first all-gather the other + /// ranks' stripes of that context (vLLM `_context_parallel_compute_prefill_context` + /// / `dcp_manager.kv_gather`, SGLang `all_gather_kv_cache_for_*_extend`). + /// The new-token attention itself is unchanged: every rank computes it for + /// its own heads and only WRITES its stripe. Defaults to 1; tail-appended + /// (schema v26). + #[serde( + default = "crate::operators::gemm::default_seq_split", + deserialize_with = "crate::operators::gemm::deserialize_positive_split" + )] + pub dcp_size: u32, +} + +/// Prefill-side cost of a DCP-striped KV: all-gather `batch * prefix` cached +/// tokens (`kv_elems_per_token`, in comm half-elements) over `dcp` ranks. +/// Zero when there is no DCP or no cached context. +pub(crate) fn dcp_context_gather( + db: &PerfDatabase, + name: &str, + kv_elems_per_token: f64, + dcp: u32, + batch_size: u32, + prefix: u32, +) -> Result, AicError> { + if dcp <= 1 || prefix == 0 || batch_size == 0 { + return Ok(None); + } + let gather = crate::operators::NcclOp::new( + format!("{name}_dcp_context_gather"), + 1.0, + kv_elems_per_token, + dcp, + "all_gather", + ); + Ok(Some(gather.query(db, batch_size.saturating_mul(prefix))?)) } /// Lane precedence for ops built without an explicit order (Rust-side @@ -230,6 +267,7 @@ impl ContextAttentionOp { cp_size: 1, lane_order: default_lane_order(), apply_rope: true, + dcp_size: 1, } } @@ -366,6 +404,18 @@ impl ContextAttentionOp { result = result.scaled(seq_imbalance_correction_scale); } + // Decode CP on the same engine: gather the cached context stripes + // (per-rank K+V of `prefix` tokens, in comm half-elements). Added after + // the zigzag imbalance scaling, like the MLA and DSA context ops: the + // collective's latency does not grow with the prefill imbalance. + let kv_elems = + 2.0 * (self.n_kv * self.head_size) as f64 * self.kv_cache_dtype.mapping().memory / 2.0; + if let Some(gather) = + dcp_context_gather(db, &self.name, kv_elems, self.dcp_size, batch_size, prefix)? + { + result = result.plus(gather); + } + Ok(result.clamp_non_negative().scaled(self.scale_factor)) } } @@ -415,6 +465,27 @@ pub struct GenerationAttentionOp { /// and the ideal-peak guard are both physical bounds. #[serde(default)] pub verify_query_tokens: u32, + /// Decode context parallelism (vLLM `-dcp` / SGLang `--dcp-size`). The + /// KV cache is striped by token position across `dcp_size` ranks of the + /// attention group; each rank all-gathers the group's query heads, attends + /// them over its own `ceil(kv / dcp)` stripe, and the partial outputs are + /// merged by LSE (the merge collectives are separate NCCL ops emitted by + /// the model class). So the kernel prices `n * dcp` query heads against a + /// `1/dcp` KV read; `n_kv` (per-rank, TP-replicated) and the batch are + /// unchanged. Defaults to 1 (no DCP). Appended at the struct tail because + /// bincode payloads are positional (schema v26). + #[serde( + default = "crate::operators::gemm::default_seq_split", + deserialize_with = "crate::operators::gemm::deserialize_positive_split" + )] + pub dcp_size: u32, +} + +/// Decode-CP kernel geometry shared by every decode attention operator: +/// `(query heads seen by the kernel, per-rank KV tokens)` for `dcp` ranks. +pub(crate) fn dcp_geometry(n_local: u32, kv_tokens: u32, dcp: u32) -> (u32, u32) { + let dcp = dcp.max(1); + (n_local * dcp, kv_tokens.div_ceil(dcp)) } impl GenerationAttentionOp { @@ -437,6 +508,7 @@ impl GenerationAttentionOp { use_qk_norm: false, scale_num_tokens: 1, verify_query_tokens: 0, + dcp_size: 1, } } @@ -457,15 +529,25 @@ impl GenerationAttentionOp { ))); } let seq_batch = batch_size / width; + // Decode CP: the kernel sees the whole DCP group's query heads over + // this rank's 1/dcp KV stripe (see the `dcp_size` field docs). + let (n_kernel, kv_local) = dcp_geometry(self.n, kv_seq_tokens, self.dcp_size); + // The stripe holds every dcp-th position, so a sliding window of `w` + // occupies ~`w / dcp` local tokens. + let window_local = if self.window_size > 0 && self.dcp_size > 1 { + self.window_size.div_ceil(self.dcp_size) + } else { + self.window_size + }; let mut result = query_generation_attention_table( db, &self.lane_order, seq_batch, - kv_seq_tokens, - self.n, + kv_local, + n_kernel, self.n_kv, self.head_size, - self.window_size, + window_local, self.kv_cache_dtype, )?; if self.verify_query_tokens > 1 { @@ -477,11 +559,11 @@ impl GenerationAttentionOp { &db.system_spec, self.n_kv, self.head_size, - self.window_size, + window_local, self.kv_cache_dtype, - self.n as f64, + n_kernel as f64, seq_batch as f64, - kv_seq_tokens as f64, + kv_local as f64, attn_flops, ); let block_math_ms = sol.math_ms * self.verify_query_tokens as f64; @@ -1354,6 +1436,38 @@ mod tests { PerfDatabase::load(&systems_root, "b200_sxm", "vllm", "0.24.0").expect("db must load") } + /// The striped-KV context gather is owed only when something is striped + /// (dcp > 1) AND something is cached (prefix > 0); it is an all-gather over + /// `batch * prefix` tokens, so it grows with the cached context. + #[test] + fn dcp_context_gather_prices_only_striped_cached_context() { + let db = b200_vllm_db(); + let kv_elems = 576.0; + assert!( + dcp_context_gather(&db, "context_mla", kv_elems, 1, 4, 30_720) + .unwrap() + .is_none() + ); + assert!( + dcp_context_gather(&db, "context_mla", kv_elems, 8, 4, 0) + .unwrap() + .is_none() + ); + assert!( + dcp_context_gather(&db, "context_mla", kv_elems, 8, 0, 30_720) + .unwrap() + .is_none() + ); + let small = dcp_context_gather(&db, "context_mla", kv_elems, 8, 4, 4_096) + .unwrap() + .expect("striped cached context is gathered"); + let large = dcp_context_gather(&db, "context_mla", kv_elems, 8, 4, 30_720) + .unwrap() + .expect("striped cached context is gathered"); + assert!(small.latency_ms > 0.0); + assert!(large.latency_ms > small.latency_ms); + } + /// The walk order Python serializes for a no-override op on /// b200_sxm/vllm/0.24.0 — see [`b200_vllm_context_lane_order`] and /// [`b200_vllm_generation_lane_order`]. Every pre-lane @@ -2529,4 +2643,53 @@ mod tests { assert_eq!(result.source, Source::Sol); assert_eq!(result.latency_ms, result.sol.unwrap().time_ms()); } + + /// Decode CP: a rank with `n` local heads and `dcp` stripes prices exactly + /// like an un-striped kernel over `n * dcp` heads and `kv / dcp` tokens, + /// and strictly cheaper than the same rank without DCP (the KV read is + /// the decode bottleneck and shrinks by `dcp`). + #[test] + fn generation_attention_dcp_prices_gathered_heads_over_kv_stripe() { + let db = b200_vllm_db().with_mode(DatabaseMode::Sol, Default::default()); + let mut striped = GenerationAttentionOp::new("gen", 16, 1, 128, KvCacheQuantMode::Fp8); + striped.dcp_size = 4; + let gathered = GenerationAttentionOp::new("gen", 64, 1, 128, KvCacheQuantMode::Fp8); + let plain = GenerationAttentionOp::new("gen", 16, 1, 128, KvCacheQuantMode::Fp8); + let a = striped.query(&db, 32, 8192, 1.0).unwrap(); + let b = gathered.query(&db, 32, 2048, 1.0).unwrap(); + let c = plain.query(&db, 32, 8192, 1.0).unwrap(); + assert!((a.latency_ms - b.latency_ms).abs() < 1e-12); + assert!(a.latency_ms < c.latency_ms); + assert_eq!(dcp_geometry(16, 8193, 4), (64, 2049)); + assert_eq!(dcp_geometry(16, 8192, 1), (16, 8192)); + } + + /// Prefill on a DCP-striped engine pays a cached-context gather that + /// scales with the prefix and vanishes without one. + #[test] + fn context_attention_dcp_gathers_cached_context_only_when_prefix_present() { + let db = b200_vllm_db().with_mode(DatabaseMode::Sol, Default::default()); + let plain = with_vllm_lanes_ctx(ContextAttentionOp::new( + "ctx", + 16, + 1, + 128, + KvCacheQuantMode::Fp8, + FmhaQuantMode::Bfloat16, + )); + let mut striped = plain.clone(); + striped.dcp_size = 4; + let no_prefix = ( + plain.query(&db, 8, 512, 0, 1.0).unwrap().latency_ms, + striped.query(&db, 8, 512, 0, 1.0).unwrap().latency_ms, + ); + assert!((no_prefix.0 - no_prefix.1).abs() < 1e-12); + let with_prefix = ( + plain.query(&db, 8, 512, 8192, 1.0).unwrap().latency_ms, + striped.query(&db, 8, 512, 8192, 1.0).unwrap().latency_ms, + ); + assert!(with_prefix.1 > with_prefix.0); + let longer = striped.query(&db, 8, 512, 16384, 1.0).unwrap().latency_ms; + assert!(longer - with_prefix.1 > 0.0); + } } diff --git a/crates/core/src/perfmodel/operators/dsa.rs b/crates/core/src/perfmodel/operators/dsa.rs index ac625bb112..c8d19e0828 100644 --- a/crates/core/src/perfmodel/operators/dsa.rs +++ b/crates/core/src/perfmodel/operators/dsa.rs @@ -62,6 +62,19 @@ pub struct DsaModuleOp { /// default. Weight-estimation only; the latency path never reads it. #[serde(default)] pub attn_projection_quant_modes: Option, + /// Decode context parallelism for the sparse-MLA decode module (vLLM + /// `flashmla_sparse` / `flashinfer_mla_sparse` and SGLang's DSA DCP path): + /// the indexer and the top-k attention run over this rank's `ceil(s / dcp)` + /// KV stripe for the DCP group's `num_heads * dcp` gathered query heads. + /// `index_topk` is kept whole (each rank's selected set is bounded by the + /// global top-k), so the sparse-attention term is an upper bound. Only the + /// generation query reads it; context DSA CP is `cp_size`. Defaults to 1; + /// appended at the struct tail (schema v26). + #[serde( + default = "default_cp_size", + deserialize_with = "crate::operators::gemm::deserialize_positive_split" + )] + pub dcp_size: u32, } /// The four DSA projection groups' quant modes (weight bytes only). @@ -77,6 +90,41 @@ fn default_cp_size() -> u32 { 1 } +/// DSA indexer key width per cached token (bf16, `index_head_dim` = 128): +/// 128 x 2 bytes = 128 comm half-elements. +const DSA_INDEXER_K_ELEMS: f64 = 128.0; + +/// Per-token volume (comm half-elements) a striped DSA worker gathers per +/// cached token: the compressed latent KV plus the indexer K cache weighted by +/// the effective full-indexer fraction (skip-indexer layers never read it, so +/// GLM-5.2's shared-index amortization pays only its `w` share; full-only +/// tables use `w = 1.0`). +pub(crate) fn dsa_cached_context_gather_elems( + kv_cache_dtype: KvCacheQuantMode, + full_fraction: f64, +) -> f64 { + crate::operators::mla::MLA_LATENT_KV_ELEMS * kv_cache_dtype.mapping().memory / 2.0 + + full_fraction.clamp(0.0, 1.0) * DSA_INDEXER_K_ELEMS +} + +/// Add the (optional) striped-KV context gather to an already-scaled CP +/// composition: `None` (dcp = 1 or nothing cached) leaves the CP total +/// untouched; `Some` adds exactly that one collective, scaled like the +/// composition it joins. Kept separate so the arithmetic is testable without +/// the sparse tables the CP path itself needs. +fn compose_cp_result_with_dcp_gather( + result: PerformanceResult, + gather: Option, + scale_factor: f64, +) -> PerformanceResult { + match gather { + None => result, + Some(gather) => result + .plus(gather.scaled(scale_factor)) + .clamp_non_negative(), + } +} + fn default_full_frac() -> f64 { 1.0 } @@ -104,6 +152,7 @@ impl DsaModuleOp { cp_size: 1, full_frac: 1.0, attn_projection_quant_modes: None, + dcp_size: 1, } } @@ -207,17 +256,25 @@ impl DsaModuleOp { // already-scaled results is exact — Python `_amortize`). if self.cp_size > 1 { let full = self.query_context_cp(db, batch_size, isl, prefix, false)?; - if w >= 1.0 { - return Ok(full); - } - let skip = self.query_context_cp(db, batch_size, isl, prefix, true)?; - let mut result = PerformanceResult::new( - w * full.latency_ms + (1.0 - w) * skip.latency_ms, - full.source, - ); - if let Some(components) = blend_sol(w, full.sol, skip.sol) { - result = result.with_sol(components); - } + let mut result = if w >= 1.0 { + full + } else { + let skip = self.query_context_cp(db, batch_size, isl, prefix, true)?; + let mut blended = PerformanceResult::new( + w * full.latency_ms + (1.0 - w) * skip.latency_ms, + full.source, + ); + if let Some(components) = blend_sol(w, full.sol, skip.sol) { + blended = blended.with_sol(components); + } + blended + }; + // A worker that carries both knobs (disaggregated roles do not + // cross-check them) still holds its cached context as 1/dcp + // stripes: the gather is owed on the CP path too. query_context_cp + // already applies scale_factor, so scale only the added gather. + let gather = self.dcp_context_gather(db, batch_size, prefix, w)?; + result = compose_cp_result_with_dcp_gather(result, gather, self.scale_factor); return Ok(result); } // Query at `isl` (new-token count) for the exact `prefix` slice — NOT @@ -247,9 +304,36 @@ impl DsaModuleOp { } blended }; + // Decode CP on the same engine: gather the cached latent-KV stripes + // plus the indexer K cache the full-indexer layers read (weighted by w). + let mut result = result; + if let Some(gather) = self.dcp_context_gather(db, batch_size, prefix, w)? { + result = result.plus(gather); + } Ok(result.clamp_non_negative().scaled(self.scale_factor)) } + /// The cached-context all-gather a striped (dcp > 1) worker pays before a + /// prefill can attend to its context: latent KV plus the indexer K cache + /// weighted by the effective full-indexer fraction; `None` when nothing is + /// striped or nothing is cached (see `attention::dcp_context_gather`). + fn dcp_context_gather( + &self, + db: &PerfDatabase, + batch_size: u32, + prefix: u32, + full_fraction: f64, + ) -> Result, AicError> { + crate::operators::attention::dcp_context_gather( + db, + &self.name, + dsa_cached_context_gather_elems(self.kv_cache_dtype, full_fraction), + self.dcp_size, + batch_size, + prefix, + ) + } + /// Context-Parallel (CP) prefill — GLM-5/DSA sparse composition. /// /// Wires the real data dependencies and delegates to [`Self::query_cp_with`] @@ -431,10 +515,22 @@ impl DsaModuleOp { } else { self.effective_full_frac(db.dsa.has_generation_skip_rows()?) }; + // Decode CP: price the DCP group's gathered heads over this rank's KV + // stripe (see the `dcp_size` field docs). A geometry-adjusted copy keeps + // the table dispatch below unchanged. + let (heads, s) = + crate::operators::attention::dcp_geometry(self.num_heads, s, self.dcp_size); + let geometry = if self.dcp_size > 1 { + let mut op = self.clone(); + op.num_heads = heads; + op + } else { + self.clone() + }; // `dsa_backend="trtllm"` mirrors Python's generation default // (`_query_generation_dsa_module_table(dsa_backend="trtllm")`). let q = |skip_indexer: bool| { - query_generation_table(db, self, batch_size, s, "trtllm", skip_indexer) + query_generation_table(db, &geometry, batch_size, s, "trtllm", skip_indexer) }; let full = q(false)?; let result = if w >= 1.0 { @@ -545,6 +641,49 @@ fn query_generation_table( s: u32, dsa_backend: &str, skip_indexer: bool, +) -> Result { + let result = query_generation_table_raw(db, op, b, s, dsa_backend, skip_indexer)?; + if op.dcp_size <= 1 || matches!(db.database_mode, DatabaseMode::Sol | DatabaseMode::SolFull) { + return Ok(result); + } + // The module tables are measured with the whole top-k. Under DCP the + // global top-k is filtered to this rank's owned slots (vLLM + // `triton_filter_and_convert_dcp_index` + compact), so the sparse + // attention reads ~`topk / dcp` tokens while the indexer still scores the + // full local stripe. Apportion the measured module time with the analytic + // roofline ratio of the two geometries (the SOL leaf already carries it). + let spec = &db.system_spec; + let dims = dsa_dims(&op.architecture); + let flops = dsa_generation_sol_flops(spec, op.gemm_quant_mode)?; + let sol = |topk_divisor: i64| { + dsa_generation_sol_ms( + spec, + dims, + op.kv_cache_dtype, + op.gemm_quant_mode, + b as i64, + s as i64, + op.num_heads as i64, + topk_divisor, + flops, + ) + }; + let (full, striped) = (sol(1), sol(op.dcp_size as i64)); + let ratio = if full > 0.0 { + (striped / full).clamp(0.0, 1.0) + } else { + 1.0 + }; + Ok(result.scaled(ratio)) +} + +fn query_generation_table_raw( + db: &PerfDatabase, + op: &DsaModuleOp, + b: u32, + s: u32, + dsa_backend: &str, + skip_indexer: bool, ) -> Result { let silicon = || { db.dsa @@ -578,6 +717,7 @@ fn query_generation_table( b as i64, s as i64, op.num_heads as i64, + op.dcp_size.max(1) as i64, flops, ))) } @@ -910,6 +1050,7 @@ fn generation_empirical( b as i64, s as i64, num_heads as i64, + op.dcp_size.max(1) as i64, flops, ) }; @@ -1109,6 +1250,7 @@ mod tests { cp_size, full_frac: 1.0, attn_projection_quant_modes: None, + dcp_size: 1, } } @@ -1120,6 +1262,58 @@ mod tests { g } + /// The striped DSA worker gathers latent KV plus the indexer K cache; the + /// indexer share follows the effective full-indexer fraction. + #[test] + fn dsa_cached_context_gather_volume_adds_weighted_indexer_k() { + let latent = crate::operators::mla::MLA_LATENT_KV_ELEMS + * KvCacheQuantMode::Bfloat16.mapping().memory + / 2.0; + assert_eq!( + dsa_cached_context_gather_elems(KvCacheQuantMode::Bfloat16, 1.0), + latent + 128.0 + ); + assert_eq!( + dsa_cached_context_gather_elems(KvCacheQuantMode::Bfloat16, 0.25), + latent + 32.0 + ); + assert_eq!( + dsa_cached_context_gather_elems(KvCacheQuantMode::Bfloat16, 0.0), + latent + ); + } + + /// The CP path owes the striped-KV context gather exactly once: `None` + /// (dcp = 1 or nothing cached) leaves the already-scaled CP total + /// untouched, `Some` adds that one collective scaled like the composition. + #[test] + fn cp_total_gains_exactly_one_scaled_dcp_gather() { + let cp_total = PerformanceResult::new(4575.0, Source::Estimated); + let unchanged = compose_cp_result_with_dcp_gather(cp_total.clone(), None, 0.5); + assert_eq!(unchanged.latency_ms, 4575.0); + assert_eq!(unchanged.source, Source::Estimated); + + let gather = PerformanceResult::new(13.0, Source::Silicon); + let striped = compose_cp_result_with_dcp_gather(cp_total, Some(gather), 0.5); + assert!((striped.latency_ms - (4575.0 + 13.0 * 0.5)).abs() < 1e-9); + } + + /// An explicit zero `dcp_size` in a serialized spec must fail at + /// deserialization instead of being normalized to non-DCP geometry. + #[test] + fn serialized_zero_dcp_size_is_rejected_and_positive_is_kept() { + let mut value = serde_json::to_value(glm_cp_op(1)).unwrap(); + value["dcp_size"] = serde_json::json!(0); + assert!(serde_json::from_value::(value.clone()).is_err()); + value["dcp_size"] = serde_json::json!(4); + assert_eq!( + serde_json::from_value::(value) + .unwrap() + .dcp_size, + 4 + ); + } + /// Composition parity with Python /// `tests/unit/sdk/test_cp_dsa_modeling.py::test_query_cp_composition`, /// same synthetic inputs (cp=8, isl=16384, prefix=0, b=1; base 4300, @@ -1612,6 +1806,7 @@ mod tests { 8, 4096, op.num_heads as i64, + 1, flops, ); assert_eq!(result.latency_ms, expected); diff --git a/crates/core/src/perfmodel/operators/fpm_sol.rs b/crates/core/src/perfmodel/operators/fpm_sol.rs index 824fbd4e7e..0959a524d4 100644 --- a/crates/core/src/perfmodel/operators/fpm_sol.rs +++ b/crates/core/src/perfmodel/operators/fpm_sol.rs @@ -291,22 +291,39 @@ fn context_attention_sol( } let fq_mem = op.fmha_quant_mode.mapping().memory; extra += mem_op_sol_ms(spec, k_num * fq_mem) + mem_op_sol_ms(spec, k_num * fq_mem); // kv write (k_num == v_num) - (fmha + extra * 1.1) * op.scale_factor + // Decode CP on the same engine: all-gather the cached-context KV stripes + // (per-rank K+V of `p` tokens, comm half-elements) over the DCP group. + let gather = if op.dcp_size > 1 && p > 0.0 { + let kv_elems = + 2.0 * (op.n_kv * op.head_size) as f64 * op.kv_cache_dtype.mapping().memory / 2.0; + nccl_sol(spec, op.dcp_size, "all_gather", b * p * kv_elems, 2.0) + } else { + 0.0 + }; + (fmha + extra * 1.1 + gather) * op.scale_factor } /// Generation-attention SOL plus the optional Q/K RMSNorm fused extra. There /// is no 5-sample smoothing and no prefix in SOL mode. fn generation_attention_sol(op: &GenerationAttentionOp, spec: &SystemSpec, b: f64, s: f64) -> f64 { + // Decode CP: the kernel sees the DCP group's `n * dcp` gathered query + // heads over this rank's `ceil(s / dcp)` KV stripe (op-level + // `GenerationAttentionOp::dcp_size` docs); the Q/K norm extra below stays + // on the rank-local heads because it runs before the query gather. + let dcp = op.dcp_size.max(1) as f64; let (n, n_kv, h, w) = ( - op.n as f64, + op.n as f64 * dcp, op.n_kv as f64, op.head_size as f64, op.window_size as f64, ); + let s_local = if dcp > 1.0 { ceil_div(s, dcp) } else { s }; + // The stripe holds every dcp-th position, so a sliding window of `w` + // occupies ~`w / dcp` local tokens. let kv_len = if op.window_size > 0 { - (s - 1.0).min(w) + (s_local - 1.0).min(ceil_div(w, dcp)) } else { - s - 1.0 + s_local - 1.0 }; // fp8 KV -> fp8 compute; everything else (incl. int8 KV) -> bf16 compute. let compute = if op.kv_cache_dtype == crate::common::enums::KvCacheQuantMode::Fp8 { @@ -322,7 +339,7 @@ fn generation_attention_sol(op: &GenerationAttentionOp, spec: &SystemSpec, b: f6 let sol_mem = mem / spec.gpu.mem_bw * 1000.0; let mut latency = sol_math.max(sol_mem); if op.use_qk_norm { - let q_num = n * h; + let q_num = op.n as f64 * h; let k_num = n_kv * h; let qk_norm = 2.0 * mem_op_sol_ms(spec, q_num * 2.0) + 2.0 * mem_op_sol_ms(spec, k_num * 2.0); @@ -400,7 +417,22 @@ fn dsa_context_module_sol( } else { w * sol(false) + (1.0 - w) * sol(true) }; - Ok(ms.max(0.0) * op.scale_factor) + // Decode CP on the same engine: all-gather the cached latent-KV stripes plus + // the indexer K cache weighted by the full-indexer fraction (mirrors + // `DsaModuleOp::query_context`'s `dcp_context_gather`). + let gather = if op.dcp_size > 1 && p > 0 { + let kv_elems = crate::operators::dsa::dsa_cached_context_gather_elems(op.kv_cache_dtype, w); + nccl_sol( + spec, + op.dcp_size, + "all_gather", + b as f64 * p as f64 * kv_elems, + 2.0, + ) + } else { + 0.0 + }; + Ok((ms.max(0.0) + gather) * op.scale_factor) } /// Whole-forward SOL leaf for the DSA generation module (`Op::DsaGeneration`): the @@ -415,14 +447,20 @@ fn dsa_generation_module_sol( ) -> Result { let dims = dsa_dims(&op.architecture); let flops = dsa_generation_sol_flops(spec, op.gemm_quant_mode)?; + // Decode CP geometry mirrors `DsaModuleOp::query_generation`: gathered + // heads over this rank's KV stripe, sparse attention over the rank's + // `ceil(topk / dcp)` owned slots. + let dcp = op.dcp_size.max(1) as f64; + let s_local = if dcp > 1.0 { ceil_div(s, dcp) } else { s }; let ms = dsa_generation_sol_ms( spec, dims, op.kv_cache_dtype, op.gemm_quant_mode, b.round().max(1.0) as i64, - s.round().max(1.0) as i64, - op.num_heads as i64, + s_local.round().max(1.0) as i64, + (op.num_heads as f64 * dcp) as i64, + dcp as i64, flops, ); Ok(ms.max(0.0) * op.scale_factor) @@ -812,6 +850,7 @@ mod tests { cp_size: 1, lane_order: crate::operators::attention::b200_vllm_context_lane_order(), apply_rope: true, + dcp_size: 1, }; let (b, sq, p) = (4.0, 682.6666666666666_f64, 128.5_f64); let (n, n_kv, h) = (48.0, 8.0, 128.0); @@ -854,6 +893,7 @@ mod tests { use_qk_norm: false, scale_num_tokens: 1, verify_query_tokens: 0, + dcp_size: 1, }; let (b, sq) = (256.0, 8441.75_f64); let kv_len = sq - 1.0; @@ -1054,6 +1094,122 @@ mod tests { ); } + /// Decode CP in the whole-forward SOL leaf: `n * dcp` gathered heads over + /// a `ceil(s / dcp)` KV stripe, identical to the un-striped equivalent. + #[test] + fn generation_attention_fpm_sol_applies_dcp_geometry() { + let d = db(); + let mut striped = + GenerationAttentionOp::new("generation_attention", 12, 1, 128, KvCacheQuantMode::Fp8); + striped.dcp_size = 4; + let gathered = + GenerationAttentionOp::new("generation_attention", 48, 1, 128, KvCacheQuantMode::Fp8); + let plain = + GenerationAttentionOp::new("generation_attention", 12, 1, 128, KvCacheQuantMode::Fp8); + let a = op_sol_latency_ms( + &Op::GenerationAttention(striped), + &d, + 256.0, + 256.0, + 8192.0, + 0.0, + ) + .unwrap(); + let b = op_sol_latency_ms( + &Op::GenerationAttention(gathered), + &d, + 256.0, + 256.0, + 2048.0, + 0.0, + ) + .unwrap(); + let c = op_sol_latency_ms( + &Op::GenerationAttention(plain), + &d, + 256.0, + 256.0, + 8192.0, + 0.0, + ) + .unwrap(); + assert!((a - b).abs() < 1e-9, "striped {a} vs gathered {b}"); + assert!( + a < c, + "dcp must shrink the KV-read-bound decode roofline: {a} vs {c}" + ); + } + + /// A sliding window of `w` occupies ~`w / dcp` positions of the stripe, so a + /// striped windowed op prices like the gathered op over a `w / dcp` window. + #[test] + fn generation_attention_fpm_sol_stripes_the_sliding_window() { + let d = db(); + let mut striped = + GenerationAttentionOp::new("generation_attention", 12, 1, 128, KvCacheQuantMode::Fp8); + striped.window_size = 4096; + striped.dcp_size = 4; + let mut gathered = + GenerationAttentionOp::new("generation_attention", 48, 1, 128, KvCacheQuantMode::Fp8); + gathered.window_size = 1024; + let mut unstriped_window = + GenerationAttentionOp::new("generation_attention", 48, 1, 128, KvCacheQuantMode::Fp8); + unstriped_window.window_size = 4096; + let a = op_sol_latency_ms( + &Op::GenerationAttention(striped), + &d, + 256.0, + 256.0, + 32768.0, + 0.0, + ) + .unwrap(); + let b = op_sol_latency_ms( + &Op::GenerationAttention(gathered), + &d, + 256.0, + 256.0, + 8192.0, + 0.0, + ) + .unwrap(); + let c = op_sol_latency_ms( + &Op::GenerationAttention(unstriped_window), + &d, + 256.0, + 256.0, + 8192.0, + 0.0, + ) + .unwrap(); + assert!( + (a - b).abs() < 1e-9, + "striped window {a} vs gathered w/dcp {b}" + ); + assert!(a < c, "the window must shrink with the stripe: {a} vs {c}"); + } + + /// DSA decode under DCP: the sparse attention reads ~`topk / dcp` owned + /// slots, so the striped leaf must undercut the gathered-heads op that keeps + /// the whole top-k (the previous upper bound). + #[test] + fn dsa_generation_fpm_sol_stripes_the_top_k() { + let d = db(); + let mut striped = glm_dsa_op("generation_attention"); + striped.dcp_size = 4; + let mut gathered = glm_dsa_op("generation_attention"); + gathered.num_heads = 256; + let a = + op_sol_latency_ms(&Op::DsaGeneration(striped), &d, 8.0, 8.0, 100000.0, 0.0).unwrap(); + let b = + op_sol_latency_ms(&Op::DsaGeneration(gathered), &d, 8.0, 8.0, 25000.0, 0.0).unwrap(); + assert!(a.is_finite() && a > 0.0, "{a}"); + assert!( + a < b, + "striped top-k must undercut the whole-top-k bound: {a} vs {b}" + ); + } + fn glm_dsa_op(name: &str) -> DsaModuleOp { // nvidia/GLM-5.2-NVFP4: 64 heads, fp8 KV cache, bf16 context FMHA, nvfp4 GEMMs, index_topk 2048 DsaModuleOp::new( @@ -1140,6 +1296,7 @@ mod tests { 8, 100000, 64, + 1, flops, ); assert!(got.is_finite() && got > 0.0, "{got}"); diff --git a/crates/core/src/perfmodel/operators/gemm.rs b/crates/core/src/perfmodel/operators/gemm.rs index 5256d87ad8..bd0221abf4 100644 --- a/crates/core/src/perfmodel/operators/gemm.rs +++ b/crates/core/src/perfmodel/operators/gemm.rs @@ -51,6 +51,23 @@ pub(crate) fn default_seq_split() -> u32 { 1 } +/// Deserialize a parallel-split width that must be at least one. The +/// tail-appended `dcp_size` fields default to 1 when absent, but an explicit +/// zero in a serialized spec must fail at deserialization instead of being +/// normalized away by the geometry helpers (`dcp_geometry`'s `.max(1)`). +pub(crate) fn deserialize_positive_split<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let value = ::deserialize(deserializer)?; + if value == 0 { + return Err(serde::de::Error::custom( + "parallel split size must be a positive integer, got 0", + )); + } + Ok(value) +} + impl GemmOp { /// Convenience constructor for the most common case (no token scaling, /// standard input precision). diff --git a/crates/core/src/perfmodel/operators/mla.rs b/crates/core/src/perfmodel/operators/mla.rs index d5d052b662..f809f18a3a 100644 --- a/crates/core/src/perfmodel/operators/mla.rs +++ b/crates/core/src/perfmodel/operators/mla.rs @@ -50,8 +50,21 @@ pub struct ContextMlaOp { /// pre-CP specs -> 1 (no sharding). #[serde(default = "crate::operators::gemm::default_seq_split")] pub cp_size: u32, + /// Decode context parallelism on the same engine: prefill with cached + /// context all-gathers the other ranks' latent-KV stripes first. See + /// `ContextAttentionOp::dcp_size`. Defaults to 1; tail-appended (schema v26). + #[serde( + default = "crate::operators::gemm::default_seq_split", + deserialize_with = "crate::operators::gemm::deserialize_positive_split" + )] + pub dcp_size: u32, } +/// Absorbed-MLA latent KV per token: `kv_lora_rank (512) + qk_rope_head_dim (64)` +/// elements, the DeepSeek-V3-family geometry every MLA op in this crate assumes +/// (the same constant the Python models spell as `kv_lora_rank + qk_rope_head_dim`). +pub(crate) const MLA_LATENT_KV_ELEMS: f64 = 576.0; + impl ContextMlaOp { pub fn new( name: impl Into, @@ -66,6 +79,7 @@ impl ContextMlaOp { kv_cache_dtype, fmha_quant_mode, cp_size: 1, + dcp_size: 1, } } @@ -94,12 +108,24 @@ impl ContextMlaOp { // balanced chunks, c = ceil(isl / 2cp). Mirrors Python // `ContextMLA.query` and `operators/attention.rs::ContextAttentionOp`. // Latency and energy both sum across the chunks (Python `__add__`). - let result = if self.cp_size > 1 { + let mut result = if self.cp_size > 1 { let c = isl.div_ceil(2 * self.cp_size).max(1); ctx(c, prefix)?.plus(ctx(c, prefix + isl - c)?) } else { ctx(isl, prefix)? }; + // Decode CP on the same engine: gather the cached latent-KV stripes. + let kv_elems = MLA_LATENT_KV_ELEMS * self.kv_cache_dtype.mapping().memory / 2.0; + if let Some(gather) = crate::operators::attention::dcp_context_gather( + db, + &self.name, + kv_elems, + self.dcp_size, + batch_size, + prefix, + )? { + result = result.plus(gather); + } Ok(result.clamp_non_negative().scaled(self.scale_factor)) } } @@ -110,6 +136,16 @@ pub struct GenerationMlaOp { pub scale_factor: f64, pub num_heads: u32, pub kv_cache_dtype: KvCacheQuantMode, + /// Decode context parallelism: the absorbed-MLA decode kernel prices + /// `num_heads * dcp` query heads (the DCP group's gathered queries) over + /// this rank's `ceil(s / dcp)` latent-KV stripe. See + /// [`GenerationAttentionOp::dcp_size`](crate::operators::GenerationAttentionOp). + /// Defaults to 1; appended at the struct tail (schema v26). + #[serde( + default = "crate::operators::gemm::default_seq_split", + deserialize_with = "crate::operators::gemm::deserialize_positive_split" + )] + pub dcp_size: u32, } impl GenerationMlaOp { @@ -119,6 +155,7 @@ impl GenerationMlaOp { scale_factor: 1.0, num_heads, kv_cache_dtype, + dcp_size: 1, } } @@ -128,8 +165,10 @@ impl GenerationMlaOp { batch_size: u32, s: u32, ) -> Result { + let (heads, s_local) = + crate::operators::attention::dcp_geometry(self.num_heads, s, self.dcp_size); let result = - query_generation_mla_table(db, batch_size, s, self.num_heads, self.kv_cache_dtype)?; + query_generation_mla_table(db, batch_size, s_local, heads, self.kv_cache_dtype)?; Ok(result.clamp_non_negative().scaled(self.scale_factor)) } } @@ -151,6 +190,16 @@ pub struct MlaModuleOp { /// it — the layout change is gated by the ENGINE_SPEC_SCHEMA_VERSION bump. #[serde(default)] pub native_num_heads: Option, + /// Decode context parallelism; read by the GENERATION variant only: + /// the fused decode module prices `num_heads * dcp` gathered heads over + /// this rank's `ceil(s / dcp)` latent-KV stripe (the module's projection + /// GEMMs stay rank-local: the frameworks shard those by TP and gather + /// the query afterwards). Defaults to 1; tail-appended (schema v26). + #[serde( + default = "crate::operators::gemm::default_seq_split", + deserialize_with = "crate::operators::gemm::deserialize_positive_split" + )] + pub dcp_size: u32, } impl MlaModuleOp { @@ -169,6 +218,7 @@ impl MlaModuleOp { fmha_quant_mode, gemm_quant_mode, native_num_heads: None, + dcp_size: 1, } } @@ -179,7 +229,7 @@ impl MlaModuleOp { isl: u32, prefix: u32, ) -> Result { - let result = query_context_mla_module_table( + let mut result = query_context_mla_module_table( db, batch_size, isl, @@ -190,6 +240,19 @@ impl MlaModuleOp { self.gemm_quant_mode, self.native_num_heads, )?; + // Decode CP on the same engine: gather the cached latent-KV stripes + // (the fused module carries one `dcp_size` for both phases). + let kv_elems = MLA_LATENT_KV_ELEMS * self.kv_cache_dtype.mapping().memory / 2.0; + if let Some(gather) = crate::operators::attention::dcp_context_gather( + db, + &self.name, + kv_elems, + self.dcp_size, + batch_size, + prefix, + )? { + result = result.plus(gather); + } Ok(result.clamp_non_negative().scaled(self.scale_factor)) } @@ -201,11 +264,13 @@ impl MlaModuleOp { ) -> Result { // No fmha arg: the generation module table has no fmha axis (decode // compute dtype follows the kv-cache dtype). + let (heads, s) = + crate::operators::attention::dcp_geometry(self.num_heads, s, self.dcp_size); let result = query_generation_mla_module_table( db, batch_size, s, - self.num_heads, + heads, self.kv_cache_dtype, self.gemm_quant_mode, self.native_num_heads, @@ -845,6 +910,22 @@ mod tests { const REPO_ROOT_HINT: &str = env!("CARGO_MANIFEST_DIR"); + /// Decode CP on absorbed MLA: `heads * dcp` gathered heads over a `s / dcp` + /// latent-KV stripe, cheaper than the un-striped rank. + #[test] + fn generation_mla_dcp_prices_gathered_heads_over_kv_stripe() { + let db = b200_vllm_db().with_mode(DatabaseMode::Sol, Default::default()); + let mut striped = GenerationMlaOp::new("mla", 16, KvCacheQuantMode::Fp8); + striped.dcp_size = 8; + let gathered = GenerationMlaOp::new("mla", 128, KvCacheQuantMode::Fp8); + let plain = GenerationMlaOp::new("mla", 16, KvCacheQuantMode::Fp8); + let a = striped.query(&db, 32, 16384).unwrap(); + let b = gathered.query(&db, 32, 2048).unwrap(); + let c = plain.query(&db, 32, 16384).unwrap(); + assert!((a.latency_ms - b.latency_ms).abs() < 1e-12); + assert!(a.latency_ms < c.latency_ms); + } + fn b200_vllm_db() -> PerfDatabase { let systems_root = PathBuf::from(REPO_ROOT_HINT) .join("../..") diff --git a/crates/core/src/perfmodel/operators/msa.rs b/crates/core/src/perfmodel/operators/msa.rs index a8cc0b523e..e53d19b8e1 100644 --- a/crates/core/src/perfmodel/operators/msa.rs +++ b/crates/core/src/perfmodel/operators/msa.rs @@ -375,6 +375,7 @@ impl MsaModuleOp { b as i64, s as i64, self.num_heads as i64, + 1, flops, ); let probe = self.dsa_probe(dims.index_topk); diff --git a/crates/core/src/perfmodel/operators/wideep_mla.rs b/crates/core/src/perfmodel/operators/wideep_mla.rs index 8575a503cd..3fb117cd11 100644 --- a/crates/core/src/perfmodel/operators/wideep_mla.rs +++ b/crates/core/src/perfmodel/operators/wideep_mla.rs @@ -545,6 +545,15 @@ pub struct WideEpGenerationMlaOp { /// keep the struct shape close to the Python class. pub fmha_quant_mode: FmhaQuantMode, pub attn_backend: String, + /// Decode context parallelism (SGLang `--dcp-size` composes with DP + /// attention / large EP): `num_heads * dcp` gathered query heads over + /// this rank's `ceil(s / dcp)` latent-KV stripe. Defaults to 1; appended + /// at the struct tail (schema v26). + #[serde( + default = "crate::operators::gemm::default_seq_split", + deserialize_with = "crate::operators::gemm::deserialize_positive_split" + )] + pub dcp_size: u32, } impl WideEpGenerationMlaOp { @@ -561,6 +570,7 @@ impl WideEpGenerationMlaOp { kv_cache_dtype, fmha_quant_mode, attn_backend: "flashinfer".to_string(), + dcp_size: 1, } } @@ -570,11 +580,13 @@ impl WideEpGenerationMlaOp { batch_size: u32, s: u32, ) -> Result { + let (heads, s_local) = + crate::operators::attention::dcp_geometry(self.num_heads, s, self.dcp_size); let (latency, source) = query_wideep_generation_mla_table( db, batch_size, - s, - self.num_heads, + s_local, + heads, self.kv_cache_dtype, &self.attn_backend, )?; diff --git a/crates/core/src/perfmodel/perf_database/dsa.rs b/crates/core/src/perfmodel/perf_database/dsa.rs index 19869cc0e5..1b0e8a802d 100644 --- a/crates/core/src/perfmodel/perf_database/dsa.rs +++ b/crates/core/src/perfmodel/perf_database/dsa.rs @@ -473,6 +473,7 @@ impl DsaTable { c[1] as i64, // b c[2] as i64, // s c[0] as i64, // num_heads + 1, // tables are measured at the full top-k flops, ) }; @@ -912,9 +913,14 @@ pub(crate) fn dsa_generation_sol( b: i64, s: i64, num_heads: i64, + topk_divisor: i64, flops: DsaSolFlops, ) -> SolComponents { let (b, s, num_heads) = (b as i128, s as i128, num_heads as i128); + // Decode CP: the global top-k is filtered to this rank's owned slots + // (vLLM `triton_filter_and_convert_dcp_index` + compact), so the sparse + // attention reads ~`topk / dcp` tokens; the indexer still scores `s`. + let topk_divisor = topk_divisor.max(1) as i128; let (hidden, q_lora, kv_lora) = ( dims.hidden_size as i128, dims.q_lora_rank as i128, @@ -935,7 +941,7 @@ pub(crate) fn dsa_generation_sol( let tokens = b; let proj_out = q_lora + kv_lora + qk_rope + ihd; - let effective_kv = s.min(topk); + let effective_kv = s.min((topk + topk_divisor - 1) / topk_divisor); let gemm_group_ops = 2 * tokens * hidden * proj_out + 2 * tokens * q_lora * num_heads * qk_head_dim @@ -982,9 +988,21 @@ pub(crate) fn dsa_generation_sol_ms( b: i64, s: i64, num_heads: i64, + topk_divisor: i64, flops: DsaSolFlops, ) -> f64 { - dsa_generation_sol(spec, dims, kv_quant, gemm_quant, b, s, num_heads, flops).time_ms() + dsa_generation_sol( + spec, + dims, + kv_quant, + gemm_quant, + b, + s, + num_heads, + topk_divisor, + flops, + ) + .time_ms() } /// Load a DSA module table from an ordered, priority-sorted source list, with diff --git a/crates/core/src/perfmodel/py.rs b/crates/core/src/perfmodel/py.rs index 5bf2de12db..a86d8ed516 100644 --- a/crates/core/src/perfmodel/py.rs +++ b/crates/core/src/perfmodel/py.rs @@ -1104,6 +1104,7 @@ struct EngineBuildRequest { fpm_options: crate::FpmInterpolationConfig, moe_tp_size: Option, moe_ep_size: Option, + cp_size: Option, gemm_quant_mode: Option, moe_quant_mode: Option, kvcache_quant_mode: Option, @@ -1123,7 +1124,6 @@ struct EngineBuildRequest { fpm_profile: Option, worker_type: Option, fpm_interpolation: Option, - cp_size: u32, decode_workload_distribution: Option, prefill_graph_profile: Option, fpm_parquet_path: Option, @@ -1164,6 +1164,7 @@ impl AicEngineBuilder { fpm_options: crate::FpmInterpolationConfig::default(), moe_tp_size: None, moe_ep_size: None, + cp_size: None, gemm_quant_mode: None, moe_quant_mode: None, kvcache_quant_mode: None, @@ -1183,7 +1184,6 @@ impl AicEngineBuilder { fpm_profile: None, worker_type: None, fpm_interpolation: None, - cp_size: 1, decode_workload_distribution: None, prefill_graph_profile: None, fpm_parquet_path: None, @@ -1270,6 +1270,15 @@ impl AicEngineBuilder { self } + /// Set optional prefill context parallelism (`cp_size`, widens the attention + /// side) and decode context parallelism (`dcp_size`, stripes the decode KV + /// across the existing TP ranks). Unset keeps both at one. + pub fn context_parallelism(mut self, cp_size: Option, dcp_size: Option) -> Self { + self.request.cp_size = cp_size; + self.request.dcp_size = dcp_size; + self + } + /// Override the GEMM quantization mode. pub fn gemm_quant_mode(mut self, value: impl Into) -> Self { self.request.gemm_quant_mode = Some(value.into()); @@ -1457,7 +1466,7 @@ mod builder_tests { assert_eq!(request.strict_provenance, Some(true)); assert_eq!(request.fpm_profile.as_deref(), Some(r#"{"model":"model"}"#)); assert_eq!(request.fpm_interpolation.as_deref(), Some("direct")); - assert_eq!(request.cp_size, 1); + assert_eq!(request.cp_size, None); assert_eq!( request.fpm_parquet_path.as_deref(), Some("/artifacts/reviewed-fpm.parquet") @@ -1513,10 +1522,9 @@ fn build_engine_from_request(request: EngineBuildRequest) -> Result Result { - if request.cp_size != 1 { + if request.cp_size == Some(0) { return Err(AicError::InvalidEngineConfig( - "cp_size must be the integer 1; this SDK entry point does not support context parallelism" - .to_string(), + "cp_size must be a positive integer, got 0".to_string(), )); } if request.fpm_fmha_quant_mode.is_some() && request.forward_model.as_deref() != Some("fpm") { @@ -1562,6 +1570,7 @@ fn compile_engine_from_request(request: EngineBuildRequest) -> Result Result>, apply_rope: bool, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { + if dcp_size == 0 { + return Err(PyValueError::new_err( + "ContextAttention dcp_size must be positive", + )); + } let inner = Op::ContextAttention(ContextAttentionOp { name, scale_factor, @@ -1047,6 +1053,7 @@ impl PyContextAttention { cp_size, lane_order: lane_order.unwrap_or_else(default_lane_order), apply_rope, + dcp_size, }); Ok((PyContextAttention, PyOperation { inner })) } @@ -1071,7 +1078,9 @@ impl PyContextAttention { o.apply_rope, ) .into_pyobject(py)?; - Ok((args, PyDict::new(py))) + let kwargs = PyDict::new(py); + kwargs.set_item("dcp_size", o.dcp_size)?; + Ok((args, kwargs)) } #[getter(_n)] @@ -1079,6 +1088,24 @@ impl PyContextAttention { Ok(slf.as_super().context_attention()?.n) } + /// Decode CP on the same engine (cached-context gather); see + /// `ContextAttentionOp::dcp_size`. + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().context_attention()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err( + "ContextAttention dcp_size must be positive", + )); + } + slf.as_super().context_attention_mut()?.dcp_size = value; + Ok(()) + } + #[getter(_n_kv)] fn n_kv(slf: PyRef<'_, Self>) -> PyResult { Ok(slf.as_super().context_attention()?.n_kv) @@ -1168,7 +1195,7 @@ impl PyGenerationAttention { const _ENGINE_QUERY_SHAPE: &'static str = "generation"; #[new] - #[pyo3(signature = (name, scale_factor, n, n_kv, kv_cache_dtype, window_size=0, head_size=128, use_qk_norm=false, lane_order=None, *, scale_num_tokens=1, verify_query_tokens=0))] + #[pyo3(signature = (name, scale_factor, n, n_kv, kv_cache_dtype, window_size=0, head_size=128, use_qk_norm=false, lane_order=None, *, scale_num_tokens=1, verify_query_tokens=0, dcp_size=1))] #[allow(clippy::too_many_arguments)] fn new( name: String, @@ -1182,7 +1209,13 @@ impl PyGenerationAttention { lane_order: Option>, scale_num_tokens: u32, verify_query_tokens: u32, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { + if dcp_size == 0 { + return Err(PyValueError::new_err( + "GenerationAttention dcp_size must be positive", + )); + } let inner = Op::GenerationAttention(GenerationAttentionOp { name, scale_factor, @@ -1195,6 +1228,7 @@ impl PyGenerationAttention { use_qk_norm, scale_num_tokens, verify_query_tokens, + dcp_size, }); Ok((PyGenerationAttention, PyOperation { inner })) } @@ -1221,6 +1255,7 @@ impl PyGenerationAttention { kwargs.set_item("lane_order", o.lane_order.clone())?; kwargs.set_item("scale_num_tokens", o.scale_num_tokens)?; kwargs.set_item("verify_query_tokens", o.verify_query_tokens)?; + kwargs.set_item("dcp_size", o.dcp_size)?; Ok((args, kwargs)) } @@ -1229,6 +1264,23 @@ impl PyGenerationAttention { Ok(slf.as_super().generation_attention()?.n) } + /// Decode context parallelism (see `GenerationAttentionOp::dcp_size`). + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().generation_attention()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err( + "GenerationAttention dcp_size must be positive", + )); + } + slf.as_super().generation_attention_mut()?.dcp_size = value; + Ok(()) + } + #[getter(_n_kv)] fn n_kv(slf: PyRef<'_, Self>) -> PyResult { Ok(slf.as_super().generation_attention()?.n_kv) @@ -1408,7 +1460,8 @@ impl PyContextMLA { const _ENGINE_QUERY_SHAPE: &'static str = "context"; #[new] - #[pyo3(signature = (name, scale_factor, num_heads, kvcache_quant_mode, fmha_quant_mode, cp_size=1))] + #[pyo3(signature = (name, scale_factor, num_heads, kvcache_quant_mode, fmha_quant_mode, cp_size=1, dcp_size=1))] + #[allow(clippy::too_many_arguments)] fn new( name: String, scale_factor: f64, @@ -1416,7 +1469,13 @@ impl PyContextMLA { kvcache_quant_mode: &Bound<'_, PyAny>, fmha_quant_mode: &Bound<'_, PyAny>, cp_size: u32, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { + if dcp_size == 0 { + return Err(PyValueError::new_err( + "ContextMLA dcp_size must be positive", + )); + } let inner = Op::ContextMla(ContextMlaOp { name, scale_factor, @@ -1424,10 +1483,29 @@ impl PyContextMLA { kv_cache_dtype: kv_quant(kvcache_quant_mode)?, fmha_quant_mode: fmha_quant(fmha_quant_mode)?, cp_size, + dcp_size, }); Ok((PyContextMLA, PyOperation { inner })) } + /// Decode CP on the same engine (cached-context gather); see + /// `ContextMlaOp::dcp_size`. + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().context_mla()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err( + "ContextMLA dcp_size must be positive", + )); + } + slf.as_super().context_mla_mut()?.dcp_size = value; + Ok(()) + } + fn __getnewargs_ex__<'py>( slf: PyRef<'py, Self>, py: Python<'py>, @@ -1442,7 +1520,9 @@ impl PyContextMLA { o.cp_size, ) .into_pyobject(py)?; - Ok((args, PyDict::new(py))) + let kwargs = PyDict::new(py); + kwargs.set_item("dcp_size", o.dcp_size)?; + Ok((args, kwargs)) } #[getter(_num_heads)] @@ -1498,18 +1578,25 @@ impl PyGenerationMLA { const _ENGINE_QUERY_SHAPE: &'static str = "generation"; #[new] - #[pyo3(signature = (name, scale_factor, num_heads, kv_cache_dtype))] + #[pyo3(signature = (name, scale_factor, num_heads, kv_cache_dtype, dcp_size=1))] fn new( name: String, scale_factor: f64, num_heads: u32, kv_cache_dtype: &Bound<'_, PyAny>, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { + if dcp_size == 0 { + return Err(PyValueError::new_err( + "GenerationMLA dcp_size must be positive", + )); + } let inner = Op::GenerationMla(GenerationMlaOp { name, scale_factor, num_heads, kv_cache_dtype: kv_quant(kv_cache_dtype)?, + dcp_size, }); Ok((PyGenerationMLA, PyOperation { inner })) } @@ -1526,7 +1613,9 @@ impl PyGenerationMLA { enum_token(&o.kv_cache_dtype), ) .into_pyobject(py)?; - Ok((args, PyDict::new(py))) + let kwargs = PyDict::new(py); + kwargs.set_item("dcp_size", o.dcp_size)?; + Ok((args, kwargs)) } #[getter(_num_heads)] @@ -1534,6 +1623,23 @@ impl PyGenerationMLA { Ok(slf.as_super().generation_mla()?.num_heads) } + /// Decode context parallelism (see `GenerationMlaOp::dcp_size`). + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().generation_mla()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err( + "GenerationMLA dcp_size must be positive", + )); + } + slf.as_super().generation_mla_mut()?.dcp_size = value; + Ok(()) + } + #[getter(_kv_cache_dtype)] fn kv_cache_dtype<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult> { py_enum_member( @@ -1561,7 +1667,7 @@ impl PyMLAModule { const _ENGINE_QUERY_SHAPE: &'static str = "module"; #[new] - #[pyo3(signature = (name, scale_factor, is_context, num_heads, kvcache_quant_mode, fmha_quant_mode, gemm_quant_mode, native_num_heads=None))] + #[pyo3(signature = (name, scale_factor, is_context, num_heads, kvcache_quant_mode, fmha_quant_mode, gemm_quant_mode, native_num_heads=None, dcp_size=1))] #[allow(clippy::too_many_arguments)] fn new( name: String, @@ -1572,7 +1678,11 @@ impl PyMLAModule { fmha_quant_mode: &Bound<'_, PyAny>, gemm_quant_mode: &Bound<'_, PyAny>, native_num_heads: Option, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { + if dcp_size == 0 { + return Err(PyValueError::new_err("MLAModule dcp_size must be positive")); + } let module = MlaModuleOp { name, scale_factor, @@ -1581,6 +1691,7 @@ impl PyMLAModule { fmha_quant_mode: fmha_quant(fmha_quant_mode)?, gemm_quant_mode: gemm_quant(gemm_quant_mode)?, native_num_heads, + dcp_size, }; let inner = if is_context { Op::MlaModuleContext(module) @@ -1608,6 +1719,7 @@ impl PyMLAModule { .into_pyobject(py)?; let kwargs = PyDict::new(py); kwargs.set_item("native_num_heads", o.native_num_heads)?; + kwargs.set_item("dcp_size", o.dcp_size)?; Ok((args, kwargs)) } @@ -1616,6 +1728,21 @@ impl PyMLAModule { matches!(slf.as_super().op(), Op::MlaModuleContext(_)) } + /// Decode context parallelism (generation variant; see `MlaModuleOp::dcp_size`). + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().mla_module()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err("MLAModule dcp_size must be positive")); + } + slf.as_super().mla_module_mut()?.dcp_size = value; + Ok(()) + } + #[setter(_is_context)] fn set_is_context(mut slf: PyRefMut<'_, Self>, value: bool) -> PyResult<()> { let base = slf.as_super(); @@ -3183,7 +3310,8 @@ impl PyWideEPGenerationMLA { const _ENGINE_QUERY_SHAPE: &'static str = "generation"; #[new] - #[pyo3(signature = (name, scale_factor, tp_size, kvcache_quant_mode, fmha_quant_mode, attn_backend="flashinfer"))] + #[pyo3(signature = (name, scale_factor, tp_size, kvcache_quant_mode, fmha_quant_mode, attn_backend="flashinfer", dcp_size=1))] + #[allow(clippy::too_many_arguments)] fn new( name: String, scale_factor: f64, @@ -3191,12 +3319,18 @@ impl PyWideEPGenerationMLA { kvcache_quant_mode: &Bound<'_, PyAny>, fmha_quant_mode: &Bound<'_, PyAny>, attn_backend: &str, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { if tp_size == 0 || 128 % tp_size != 0 { return Err(PyValueError::new_err(format!( "WideEPGenerationMLA tp_size must divide 128, got {tp_size}" ))); } + if dcp_size == 0 { + return Err(PyValueError::new_err( + "WideEPGenerationMLA dcp_size must be positive", + )); + } let inner = Op::WideEpGenerationMla(crate::operators::WideEpGenerationMlaOp { name, scale_factor, @@ -3204,6 +3338,7 @@ impl PyWideEPGenerationMLA { kv_cache_dtype: kv_quant(kvcache_quant_mode)?, fmha_quant_mode: fmha_quant(fmha_quant_mode)?, attn_backend: attn_backend.to_string(), + dcp_size, }); Ok((PyWideEPGenerationMLA, PyOperation { inner })) } @@ -3222,7 +3357,9 @@ impl PyWideEPGenerationMLA { o.attn_backend.clone(), ) .into_pyobject(py)?; - Ok((args, PyDict::new(py))) + let kwargs = PyDict::new(py); + kwargs.set_item("dcp_size", o.dcp_size)?; + Ok((args, kwargs)) } #[getter(_tp_size)] @@ -3230,6 +3367,23 @@ impl PyWideEPGenerationMLA { Ok(128 / slf.as_super().wideep_generation_mla()?.num_heads.max(1)) } + /// Decode context parallelism (see `WideEpGenerationMlaOp::dcp_size`). + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().wideep_generation_mla()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err( + "WideEPGenerationMLA dcp_size must be positive", + )); + } + slf.as_super().wideep_generation_mla_mut()?.dcp_size = value; + Ok(()) + } + #[getter(_kvcache_quant_mode)] fn kvcache_quant_mode<'py>( slf: PyRef<'py, Self>, @@ -3560,7 +3714,7 @@ impl PyContextDSAModule { const _ENGINE_QUERY_SHAPE: &'static str = "context"; #[new] - #[pyo3(signature = (name, scale_factor, num_heads, kvcache_quant_mode, fmha_quant_mode, gemm_quant_mode, architecture="DeepseekV32ForCausalLM", cp_size=1, index_topk_freq=1, dsa_full_layer_fraction=None, attn_projection_quant_modes=None))] + #[pyo3(signature = (name, scale_factor, num_heads, kvcache_quant_mode, fmha_quant_mode, gemm_quant_mode, architecture="DeepseekV32ForCausalLM", cp_size=1, index_topk_freq=1, dsa_full_layer_fraction=None, attn_projection_quant_modes=None, dcp_size=1))] #[allow(clippy::too_many_arguments)] fn new( name: String, @@ -3574,7 +3728,13 @@ impl PyContextDSAModule { index_topk_freq: i64, dsa_full_layer_fraction: Option, attn_projection_quant_modes: Option<&Bound<'_, PyDict>>, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { + if dcp_size == 0 { + return Err(PyValueError::new_err( + "ContextDSAModule dcp_size must be positive", + )); + } let gemm = gemm_quant(gemm_quant_mode)?; let freq = index_topk_freq.max(1) as f64; let inner = Op::DsaContext(crate::operators::DsaModuleOp { @@ -3592,10 +3752,29 @@ impl PyContextDSAModule { attn_projection_quant_modes, gemm, )?), + dcp_size, }); Ok((PyContextDSAModule, PyOperation { inner })) } + /// Decode CP on the same engine (cached-context gather); see + /// `DsaModuleOp::dcp_size`. + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().dsa()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err( + "ContextDSAModule dcp_size must be positive", + )); + } + slf.as_super().dsa_mut()?.dcp_size = value; + Ok(()) + } + fn __getnewargs_ex__<'py>( slf: PyRef<'py, Self>, py: Python<'py>, @@ -3620,6 +3799,7 @@ impl PyContextDSAModule { dsa_projection_dict(py, quants)?, )?; } + kwargs.set_item("dcp_size", o.dcp_size)?; Ok((args, kwargs)) } @@ -3707,7 +3887,7 @@ impl PyGenerationDSAModule { const _ENGINE_QUERY_SHAPE: &'static str = "generation"; #[new] - #[pyo3(signature = (name, scale_factor, num_heads, kv_cache_dtype, gemm_quant_mode, architecture="DeepseekV32ForCausalLM", index_topk_freq=1, dsa_full_layer_fraction=None, attn_projection_quant_modes=None))] + #[pyo3(signature = (name, scale_factor, num_heads, kv_cache_dtype, gemm_quant_mode, architecture="DeepseekV32ForCausalLM", index_topk_freq=1, dsa_full_layer_fraction=None, attn_projection_quant_modes=None, dcp_size=1))] #[allow(clippy::too_many_arguments)] fn new( name: String, @@ -3719,7 +3899,13 @@ impl PyGenerationDSAModule { index_topk_freq: i64, dsa_full_layer_fraction: Option, attn_projection_quant_modes: Option<&Bound<'_, PyDict>>, + dcp_size: u32, ) -> PyResult<(Self, PyOperation)> { + if dcp_size == 0 { + return Err(PyValueError::new_err( + "GenerationDSAModule dcp_size must be positive", + )); + } let gemm = gemm_quant(gemm_quant_mode)?; let freq = index_topk_freq.max(1) as f64; let inner = Op::DsaGeneration(crate::operators::DsaModuleOp { @@ -3737,6 +3923,7 @@ impl PyGenerationDSAModule { attn_projection_quant_modes, gemm, )?), + dcp_size, }); Ok((PyGenerationDSAModule, PyOperation { inner })) } @@ -3763,9 +3950,27 @@ impl PyGenerationDSAModule { dsa_projection_dict(py, quants)?, )?; } + kwargs.set_item("dcp_size", o.dcp_size)?; Ok((args, kwargs)) } + /// Decode context parallelism (see `DsaModuleOp::dcp_size`). + #[getter(_dcp_size)] + fn dcp_size(slf: PyRef<'_, Self>) -> PyResult { + Ok(slf.as_super().dsa()?.dcp_size) + } + + #[setter(_dcp_size)] + fn set_dcp_size(mut slf: PyRefMut<'_, Self>, value: u32) -> PyResult<()> { + if value == 0 { + return Err(PyValueError::new_err( + "GenerationDSAModule dcp_size must be positive", + )); + } + slf.as_super().dsa_mut()?.dcp_size = value; + Ok(()) + } + #[getter(_num_heads)] fn num_heads(slf: PyRef<'_, Self>) -> PyResult { Ok(slf.as_super().dsa()?.num_heads) diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 07cb6827e5..4690310e3e 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -177,6 +177,10 @@ struct AicTimingConfig { moe_tp_size: Option, #[serde(default)] moe_ep_size: Option, + /// Prefill context parallelism (SGLang attn-cp / vLLM PCP): folds into the + /// attention width like attention_dp. `None` means 1. + #[serde(default)] + cp_size: Option, #[serde(default, alias = "gemm_quant_mode")] gemm_dtype: Option, #[serde(default, alias = "moe_quant_mode")] @@ -322,6 +326,7 @@ impl AicTimingConfig { dcp: self.dcp, moe_tp_size: self.moe_tp_size, moe_ep_size: self.moe_ep_size, + cp_size: self.cp_size, gemm_quant_mode: self.gemm_dtype.clone(), moe_quant_mode: self.moe_dtype.clone(), fmha_quant_mode: self.fmha_dtype.clone(), @@ -414,9 +419,11 @@ impl AicTimingConfig { && self.pp > 0 && self.attention_dp > 0 && self.moe_tp_size != Some(0) - && self.moe_ep_size != Some(0), - "AIC timing parallel sizes tp, pp, attention_dp, moe_tp_size, and \ - moe_ep_size must be positive" + && self.moe_ep_size != Some(0) + && self.cp_size != Some(0) + && self.dcp != Some(0), + "AIC timing parallel sizes tp, pp, attention_dp, moe_tp_size, \ + moe_ep_size, cp_size, and dcp must be positive" ); ensure!(self.nextn <= 5, "AIC nextn must be in 0..=5"); self.speculative_depth()?; @@ -425,11 +432,15 @@ impl AicTimingConfig { "AIC moe_tp_size and moe_ep_size must be configured together" ); if let (Some(moe_tp), Some(moe_ep)) = (self.moe_tp_size, self.moe_ep_size) { + // Prefill CP widens the attention side (mirrors ModelConfig's + // `tp * attention_dp * cp == moe_tp * moe_ep`); decode CP reuses + // ranks inside the attention group and is deliberately absent. + let cp = u64::from(self.cp_size.unwrap_or(1)); ensure!( self.fpm_profile.is_some() - || u64::from(self.tp) * u64::from(self.attention_dp) + || u64::from(self.tp) * u64::from(self.attention_dp) * cp == u64::from(moe_tp) * u64::from(moe_ep), - "AIC topology requires tp * attention_dp == moe_tp_size * moe_ep_size" + "AIC topology requires tp * attention_dp * cp_size == moe_tp_size * moe_ep_size" ); } Ok(()) @@ -963,6 +974,9 @@ fn aic_capacity_kwargs<'py>( kwargs.set_item("attention_dp_size", config.attention_dp)?; kwargs.set_item("moe_tp_size", config.moe_tp_size)?; kwargs.set_item("moe_ep_size", config.moe_ep_size)?; + kwargs.set_item("cp_size", config.cp_size.unwrap_or(1))?; + // Capacity estimation prices the 1/dcp KV stripe; unrecorded DCP is 1. + kwargs.set_item("dcp_size", config.dcp.unwrap_or(1))?; kwargs.set_item("gemm_quant_mode", config.gemm_dtype.as_deref())?; kwargs.set_item("moe_quant_mode", config.moe_dtype.as_deref())?; kwargs.set_item("fmha_quant_mode", config.fmha_dtype.as_deref())?; @@ -1126,10 +1140,6 @@ fn materialize_aic_capacity( if capacity_is_explicit || role.rank.state_cache.is_some() { return Ok(()); } - ensure!( - config.dcp.is_none_or(|dcp| dcp == 1), - "DCP FPM replay requires explicit KV block capacity; automatic DCP/hybrid sizing is unsupported" - ); let blocks = estimate(config, role)?; ensure!(blocks > 0, "AIC estimated zero KV-cache blocks"); role.rank.num_gpu_blocks = blocks; @@ -3233,6 +3243,24 @@ mod tests { assert!(infer_aic_timing_topology(&mut engine).is_err()); } + #[test] + fn estimator_request_carries_both_context_parallel_knobs() { + // The native estimator path builds the engine from this request, so a + // dropped knob would silently price a dcp=1 engine for a dcp=8 worker. + let mut config = aic_config(); + config.cp_size = Some(2); + config.dcp = Some(4); + let request = config + .estimator_request(ForwardPassWorkerType::Aggregated) + .unwrap(); + assert_eq!(request.cp_size, Some(2)); + assert_eq!(request.dcp, Some(4)); + let unset = aic_config() + .estimator_request(ForwardPassWorkerType::Aggregated) + .unwrap(); + assert_eq!((unset.cp_size, unset.dcp), (None, None)); + } + #[test] fn canonical_and_legacy_default_selection_are_distinct() { let mut config = aic_config(); @@ -3314,6 +3342,7 @@ mod tests { dcp: None, moe_tp_size: None, moe_ep_size: None, + cp_size: None, gemm_dtype: None, moe_dtype: None, fmha_dtype: None, @@ -3351,6 +3380,26 @@ mod tests { } } + #[test] + fn parallel_shape_folds_prefill_cp_but_not_decode_cp_into_width() { + let mut config = aic_config(); + config.tp = 1; + config.attention_dp = 1; + config.moe_tp_size = Some(1); + config.moe_ep_size = Some(8); + // Prefill CP widens the attention side to match the MoE width ... + config.cp_size = Some(8); + config.dcp = Some(8); + config.validate_parallel_shape().unwrap(); + // ... decode CP does not: without prefill CP the widths no longer match. + config.cp_size = None; + assert!(config.validate_parallel_shape().is_err()); + // Zero is rejected like every other parallel size. + config.cp_size = Some(8); + config.dcp = Some(0); + assert!(config.validate_parallel_shape().is_err()); + } + #[test] fn fixed_capacity_does_not_require_vllm_pp1_profile_resources() { for (backend, pp) in [("sglang", 1), ("vllm", 2)] { @@ -3451,7 +3500,10 @@ mod tests { } #[test] - fn dcp_capacity_requires_explicit_blocks() { + fn dcp_capacity_is_estimated_on_the_kv_stripe() { + // The KV-capacity estimator prices the 1/dcp stripe (memory.py's + // `_cp_kv_memory_divisor`), so automatic sizing stays available under + // DCP; explicit capacity still bypasses the estimator. let mut config = aic_config(); config.tp = 4; config.dcp = Some(4); @@ -3459,22 +3511,15 @@ mod tests { role.tensor_parallel_size = 4; role.rank.num_gpu_blocks = 17; - let error = materialize_aic_capacity(&config, &mut role, false, |_, _| { - panic!("DCP must be rejected before estimating capacity") - }) - .unwrap_err(); - assert!(error.to_string().contains("explicit KV block capacity")); - assert_eq!(role.rank.num_gpu_blocks, 17); + materialize_aic_capacity(&config, &mut role, false, |_, _| Ok(4 * 321)).unwrap(); + assert_eq!(role.rank.num_gpu_blocks, 4 * 321); + role.rank.num_gpu_blocks = 17; materialize_aic_capacity(&config, &mut role, true, |_, _| { panic!("explicit DCP capacity must not invoke the estimator") }) .unwrap(); assert_eq!(role.rank.num_gpu_blocks, 17); - - config.dcp = Some(1); - materialize_aic_capacity(&config, &mut role, false, |_, _| Ok(321)).unwrap(); - assert_eq!(role.rank.num_gpu_blocks, 321); } #[test] diff --git a/crates/core/tests/perfmodel/memory_round_trip.rs b/crates/core/tests/perfmodel/memory_round_trip.rs index b4ef917e11..0b68f06a99 100644 --- a/crates/core/tests/perfmodel/memory_round_trip.rs +++ b/crates/core/tests/perfmodel/memory_round_trip.rs @@ -93,13 +93,13 @@ fn request(tolerance_fraction: Option) -> KvCacheEstimateRequest { moe_kernel_source: None, kv_block_size: None, parallel: ParallelMapping { - dcp_size: None, tp_size: 1, pp_size: 1, attention_dp_size: Some(1), moe_tp_size: None, moe_ep_size: None, cp_size: None, + dcp_size: None, }, quantization: QuantizationConfig { weight_dtype: None, diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index a39f099814..aaedd5e9db 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -169,7 +169,8 @@ mod tests { // v23: ParallelMapping gained optional recorded DCP identity. // v24: FpmForwardOp carries typed DCP separately from matching strings. // v25: SOL/direct interpolation selector combined with typed DCP. - assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 25); + // v26: attention / MLA / DSA ops gained tail-appended decode-CP dcp_size. + assert_eq!(ENGINE_SPEC_SCHEMA_VERSION, 26); assert_eq!(FPM_VERSION, 1); assert_eq!(ForwardPassMetrics::default().version, FPM_VERSION); } diff --git a/docs/cli/migrate-from-aiconfigurator.md b/docs/cli/migrate-from-aiconfigurator.md index 99681f2f03..da9b3dd4c8 100644 --- a/docs/cli/migrate-from-aiconfigurator.md +++ b/docs/cli/migrate-from-aiconfigurator.md @@ -1560,10 +1560,14 @@ and pins CP=1. #### 5.6.2 Context parallelism (CP) -**The unified AISimulate CLI has no CP configuration field.** AIC exposes per-role -`agg_cp_candidates`, `prefill_cp_candidates`, and `decode_cp_candidates`. CP>1 support depends on -the model family and backend; the dense-model PP example above does not establish CP>1 support. -Keep supported CP workflows on AIC. See +**The unified `predict` CLI exposes CP as two per-phase knobs; `recommend` does not sweep them.** +`parallelism.prefill_context` is prefill context parallelism (AIC `*_cp_candidates`: SGLang +`--attn-cp-size`, vLLM `-pcp`) and `parallelism.decode_context` is decode context parallelism +(AIC `dcp_size` / `prefill_dcp_size` / `decode_dcp_size`: vLLM `-dcp`, SGLang `--dcp-size`). +Aggregated workers accept at most one of the two above 1; disaggregated roles carry each knob +independently. Support above 1 depends on the model family and backend (`supports_cp` / +`supports_dcp` on the model class); the dense-model PP example above does not establish CP>1 +support. For CP candidate sweeps keep the workflow on AIC. See [advanced AIC search controls](../../python/aisimulate/docs/advanced_tuning.md). diff --git a/docs/cli/user-guide.md b/docs/cli/user-guide.md index 4d30b65f45..260067141d 100644 --- a/docs/cli/user-guide.md +++ b/docs/cli/user-guide.md @@ -1021,6 +1021,8 @@ engine: | `engine.workers..parallelism.attention_data` | `1` | Feasible registry values | `parallelism` | Positive and model/backend compatible. | | `engine.workers..parallelism.moe_tensor` | `1` | Feasible registry values | `parallelism` | Positive and model/backend compatible. | | `engine.workers..parallelism.moe_expert` | `1` | Feasible registry values | `parallelism` | Positive and model/backend compatible. | +| `engine.workers..parallelism.prefill_context` | `1` | `x` | `-` | `predict` only. Prefill context parallelism (SGLang `--attn-cp-size`, vLLM `-pcp`): splits prefill tokens across extra attention ranks; decode stays replicated on them. Requires model/backend CP support. | +| `engine.workers..parallelism.decode_context` | `1` | `x` | `-` | `predict` only. Decode context parallelism (vLLM `-dcp`, SGLang `--dcp-size`): stripes the decode KV cache across ranks inside the attention group without adding GPUs. Aggregated workers accept at most one of `prefill_context` / `decode_context` above 1; disaggregated roles carry each knob independently. Modeled for DeepSeek-V3-class MLA (vLLM, SGLang), DeepSeek-V3.2 / GLM-5 DSA (vLLM, SGLang) and dense / Qwen-MoE GQA (vLLM, `decode_context <= tensor / kv_heads`); other families fail loud. | | `engine.workers..scheduler.max_batched_tokens` | Aggregated/prefill/decode: `8192` | Prefill/aggregated: `{choices: [8192, 16384, 32768]}`; decode: `-` | `-` | Positive. | | `engine.workers..scheduler.max_sequences` | Aggregated `256`; prefill `1`; decode `256` | Prefill: `{choices: [1, 2, 4, 8, 16, 32, 64, 128, 256]}`; aggregated/decode: `{choices: [256, 512, 1024]}` | `-` | Positive. | | `engine.workers..scheduler.prefill_schedule_interval` | `1` | `x` | `-` | `predict` only. Positive. Values above one throttle prefill admission only for vLLM attention-DP groups. | diff --git a/docs/core-api.md b/docs/core-api.md index fd3ab37cff..02dfb9b8bc 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -275,6 +275,13 @@ Use `RustForwardPassPerfModel.best_available(config)` from Python or `ForwardPassPerfModelConfig` owns model, hardware, backend, topology, data policies, a required immutable `worker_type`, and the complete nested `estimator_config`. Worker roles are `prefill`, `decode`, and `aggregated`. +Topology carries two optional context-parallel knobs next to `tp`, `pp`, +`attention_dp`, `moe_tp_size`, and `moe_ep_size`: `cp_size` (prefill context +parallelism, SGLang `--attn-cp-size` / vLLM `-pcp`; extra attention ranks, so +`tp * attention_dp * cp_size == moe_tp_size * moe_ep_size`) and `dcp_size` +(decode context parallelism, vLLM `-dcp` / SGLang `--dcp-size`; stripes the +decode KV cache across the existing TP ranks and adds no GPUs). Both default to +one and stay out of the serialized identity when unset. ```python from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel diff --git a/python/aisimulate/docs/generator_overview.md b/python/aisimulate/docs/generator_overview.md index 3e9b3fd5c2..875c4c8545 100644 --- a/python/aisimulate/docs/generator_overview.md +++ b/python/aisimulate/docs/generator_overview.md @@ -56,6 +56,8 @@ flowchart TD ``` DSL rules users can extend to influence generated configs. Field names come from `backend_config_mapping.yaml` and `deployment_config.yaml`; prefixes like `agg_`, `prefill_`, and `decode_` scope the impact to that role’s generated outputs. + **Context parallelism**: two orthogonal per-worker knobs flow through the same chain. `context_parallel_size` is prefill CP (vLLM `--prefill-context-parallel-size`, SGLang `--enable-prefill-cp --attn-cp-size --cp-strategy`); it adds attention ranks, so `gpus_per_worker` multiplies by it (the SGLang rule folds it into `--tensor-parallel-size`, because SGLang's CP ranks live inside `--tp`). `decode_context_parallel_size` is decode CP (vLLM `--decode-context-parallel-size`, SGLang `--dcp-size`); it stripes the KV cache across ranks already in the attention group and adds no GPUs. `generator/context_parallel.py` is the single decision point: it version-gates the flags (vLLM PCP and DCP >= 0.14.1, the first `cli_args` template that renders them, `--dcp-comm-backend` >= 0.18.0; SGLang `--attn-cp-size`/`--dcp-size` >= 0.5.15, `--dcp-comm-backend` >= 0.5.17), picks `--cp-strategy interleave` for DSA architectures and `zigzag` otherwise, and rejects TRT-LLM (no CP perf data). Both request paths use it: `module_bridge` reads the swept `cp`/`(p)cp` column plus the task's `dcp_size`/`prefill_dcp_size`/`decode_dcp_size`, and `from_sweeper_candidate` reads the optional `cp`/`dcp` (or `prefill_`/`decode_`-prefixed) candidate columns. The `aisimulate sweep`/`recommend` search spaces do not enumerate either knob yet; they are CLI-driven inputs. + **Rule selection**: Use `--generator-set rule=benchmark` to switch to a different rule plugin folder under `src/aisimulate/generator/rule_plugin/`. If `rule` is not provided, the default production rules are used (tuned for deployment, including max batch size and CUDA graph batch size adjustments). The `benchmark` rules are designed to align generated configs with AIC simulation, using broader CUDA graph batch sizes and a stricter max batch size derived from the simulated batch size. You can add your own rule sets by creating a folder under `rule_plugin/` and selecting it via `--generator-set rule=`. - Backend templates (`config/backend_templates//`): @@ -215,7 +217,7 @@ generator mapping fails closed rather than producing an artifact that contradict `--deployment-target fpm` supports a vLLM single aggregated-worker topology with exactly one worker replica. It emits a keepalive Pod when the resolved topology fits on one node. Multinode output defaults to a `LeaderWorkerSet`; set `K8sConfig.fpm_orchestrator: grove` to emit a Grove `PodCliqueSet` instead. Router/planner configurations and invalid FPM topologies fail closed. Other deployment targets keep their existing behavior. -`Workers.agg.gpus_per_worker` is the total GPU count for the worker replica; a value larger than `NodeConfig.num_gpus_per_node` produces a multinode worker. The resolved `TP * PP * DP` must equal that total, which must divide evenly across the resolved node count. Multinode DP must also divide evenly across nodes. Any resolved DP greater than one uses the `mp` data-parallel backend. The default `lws` orchestrator requires LeaderWorkerSet; `grove` requires Grove. Grove output leaves scheduler and queue selection unset unless they are supplied explicitly through `K8sConfig.worker_extra_pod_spec.schedulerName` and `K8sConfig.fpm_resource_labels`. +`Workers.agg.gpus_per_worker` is the total GPU count for the worker replica; a value larger than `NodeConfig.num_gpus_per_node` produces a multinode worker. The resolved `TP * PP * DP * PCP` must equal that total, which must divide evenly across the resolved node count. Multinode DP must also divide evenly across nodes. Any resolved DP greater than one uses the `mp` data-parallel backend. The default `lws` orchestrator requires LeaderWorkerSet; `grove` requires Grove. Grove output leaves scheduler and queue selection unset unless they are supplied explicitly through `K8sConfig.worker_extra_pod_spec.schedulerName` and `K8sConfig.fpm_resource_labels`. The FPM overlay accepts `Workers.agg.extra_cli_args` as a `list[str]` and concrete `K8sConfig.extra_env` entries in `{name, value}` form. Rules, mappings, and versioned templates still produce the base vLLM command; `extra_cli_args` are appended to that resolved command. `--benchmark-mode` is required and accepts `agg`, `prefill`, or `decode`; it selects the runtime collection phase without changing the required single aggregated-worker topology. `valueFrom`, `envFrom`, and Secret-derived environment values are not supported in V1. The generator owns multinode coordination arguments such as `--nnodes`, `--node-rank`, `--headless`, and the local data-parallel rank arguments, so callers must not pass those through the overlay. diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 20bd0a8742..18b3386be2 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -132,6 +132,10 @@ def finish_lowering(value): "aic_pp_size": resolved["pp"], "aic_moe_tp_size": resolved["moe_tp_size"], "aic_moe_ep_size": resolved["moe_ep_size"], + # Rust serializes the context-parallel knobs only when set; + # decode CP must reach the KV-capacity estimate (1/dcp per rank). + "aic_cp_size": resolved.get("cp_size"), + "aic_dcp_size": resolved.get("dcp"), "aic_fpm_profile": resolved.get("fpm_profile"), "worker_type": resolved["worker_type"], }.items() @@ -173,10 +177,6 @@ def finish_lowering(value): return finish_lowering(lowered) if lowered.get("num_gpu_blocks") is not None: return finish_lowering(lowered) - if canonical_result is not None and (resolved.get("dcp") or 1) > 1: - raise ValueError( - "DCP FPM replay requires explicit KV block capacity; automatic DCP/hybrid sizing is unsupported" - ) backend = lowered.get("aic_backend") if backend is None: return finish_lowering(lowered) @@ -225,6 +225,8 @@ def finish_lowering(value): moe_tp_size=lowered.get("aic_moe_tp_size"), moe_ep_size=lowered.get("aic_moe_ep_size"), attention_dp_size=attention_dp, + cp_size=(lowered.get("aic_cp_size") if lowered.get("aic_cp_size") is not None else 1), + dcp_size=(lowered.get("aic_dcp_size") if lowered.get("aic_dcp_size") is not None else 1), gemm_dtype=lowered.get("aic_gemm_dtype"), moe_dtype=lowered.get("aic_moe_dtype"), fmha_dtype=lowered.get("aic_fmha_dtype"), @@ -334,6 +336,8 @@ def estimate_num_gpu_blocks( moe_tp_size: int | None = None, moe_ep_size: int | None = None, attention_dp_size: int | None = None, + cp_size: int = 1, + dcp_size: int = 1, gemm_dtype: str | None = None, moe_dtype: str | None = None, fmha_dtype: str | None = None, @@ -409,6 +413,8 @@ def estimate_num_gpu_blocks( attention_dp_size=(attention_dp_size if attention_dp_size is not None else 1), moe_tp_size=moe_tp_size, moe_ep_size=moe_ep_size, + cp_size=cp_size, + dcp_size=dcp_size, gemm_quant_mode=_quant_mode_name("gemm", gemm_dtype), moe_quant_mode=_quant_mode_name("moe", moe_dtype), fmha_quant_mode=_quant_mode_name("fmha", fmha_dtype), diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index ea7b3de9dd..fa37980369 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -229,6 +229,7 @@ def _deployment( if mode == "agg": assert engine.workers.aggregated is not None worker = engine.workers.aggregated + _require_exclusive_context_parallelism(worker) parallel = _parallel_mapping(worker, prefix="") state_size = resolve_state_size(engine, worker) if worker.kv_cache.state_cache is not None: @@ -259,6 +260,7 @@ def _deployment( assert engine.workers.prefill is not None and engine.workers.decode is not None prefill = engine.workers.prefill decode = engine.workers.decode + _require_disaggregated_context_parallelism(engine.backend, prefill, decode) transfer_bytes_per_token = None if engine.kv_transfer is not None: transfer_bytes_per_token = _resolve_kv_bytes_per_token( @@ -311,6 +313,15 @@ def _afd_deployment( companion: ReplicaParallelConfig | None = None if companion_worker is not None: parallel = companion_worker.parallelism + # The AFD companion's ParallelShape, GPU accounting and provenance carry + # tp/pp/attention_dp/moe only; a CP knob here would price a wider worker + # than the topology reports. Fail closed until AFD models CP explicitly. + if _prefill_cp(parallel) != 1 or _decode_cp(parallel) != 1: + raise ValueError( + f"AFD companion ({companion_role}) workers do not support context parallelism: got " + f"parallelism.prefill_context={parallel.prefill_context}, " + f"parallelism.decode_context={parallel.decode_context}; set both to 1" + ) companion = ReplicaParallelConfig( shape=ParallelShape( tp=parallel.tensor, @@ -387,9 +398,66 @@ def _afd_deployment( ) +def _require_exclusive_context_parallelism(worker: WorkerPredictionConfig) -> None: + """Aggregated workers model prefill CP and decode CP as mutually exclusive. + + This is a deployment-topology rule, deliberately NOT a ModelConfig rule: + disaggregated deployments carry the two knobs on different workers, so the + same ModelConfig fields are valid there without any cross-check. In one + aggregated engine the frameworks either do not compose the two at all + (SGLang's prefill CP and DCP paths never reference each other) or only for + a narrow model class (vLLM: sparse MLA with ``ag_rs``), so the simulator + prices them one at a time. + """ + parallel = worker.parallelism + if _prefill_cp(parallel) > 1 and _decode_cp(parallel) > 1: + raise ValueError( + "aggregated workers support at most one of parallelism.prefill_context and " + f"parallelism.decode_context above 1 (got prefill_context={parallel.prefill_context}, " + f"decode_context={parallel.decode_context}); use a disaggregated deployment to apply " + "prefill CP on the prefill worker and decode CP on the decode worker" + ) + + +def _require_disaggregated_context_parallelism( + backend: str, prefill: WorkerPredictionConfig, decode: WorkerPredictionConfig +) -> None: + """Disaggregated roles carry their own knobs, with two layout constraints. + + A prefill engine gains nothing from striping its KV (the KV only passes + through), so decode CP on the prefill worker is accepted only as the + PCP+DCP layout, i.e. equal to its prefill CP. vLLM's NIXL connector + additionally refuses to pair a replicated-PCP prefill (pcp > 1, dcp = 1) + with a DCP-sharded decode, and requires the two DCP sizes to divide one + another (``nixl/base_worker.py``). SGLang re-lays the KV out per decode + DCP rank on the prefill side and has no such pairing rule. + """ + p, d = prefill.parallelism, decode.parallelism + p_dcp, d_dcp = _decode_cp(p), _decode_cp(d) + if p_dcp not in (1, _prefill_cp(p)): + raise ValueError( + f"prefill workers accept parallelism.decode_context only as 1 or equal to prefill_context " + f"(got decode_context={p.decode_context}, prefill_context={p.prefill_context}); a prefill " + "engine only stripes its KV to match a PCP+DCP layout" + ) + if backend != "vllm" or d_dcp == 1: + return + if _prefill_cp(p) > 1 and p_dcp == 1: + raise ValueError( + f"vLLM cannot pair a replicated-PCP prefill worker (prefill_context={p.prefill_context}, " + f"decode_context=1) with a DCP-sharded decode worker (decode_context={d.decode_context}); " + "set the prefill worker's decode_context equal to its prefill_context or drop one knob" + ) + if d_dcp % p_dcp and p_dcp % d_dcp: + raise ValueError( + f"vLLM requires the prefill and decode DCP sizes to divide one another (got prefill " + f"decode_context={p.decode_context}, decode decode_context={d.decode_context})" + ) + + def _parallel_mapping(worker: WorkerPredictionConfig, *, prefix: str) -> dict[str, JSONValue]: parallel = worker.parallelism - return { + mapping: dict[str, JSONValue] = { f"{prefix}replicas": parallel.replicas, f"{prefix}tp": parallel.tensor, f"{prefix}pp": parallel.pipeline, @@ -397,7 +465,27 @@ def _parallel_mapping(worker: WorkerPredictionConfig, *, prefix: str) -> dict[st f"{prefix}moe_tp": parallel.moe_tensor, f"{prefix}moe_ep": parallel.moe_expert, **({f"{prefix}dcp": parallel.decode_context} if parallel.decode_context is not None else {}), + **_prefill_cp_knob(parallel, f"{prefix}cp"), } + return mapping + + +def _prefill_cp_knob(parallel: Any, cp_key: str) -> dict[str, int]: + """Prefill CP, spelled out only when above one so cp=1 deployments, engine + args and estimator identities stay byte-identical to pre-CP outputs. Decode + CP is carried by ``decode_context`` itself: ``None`` (not requested) is + distinct from an explicit 1 for the FPM cell identity.""" + return {cp_key: _prefill_cp(parallel)} if _prefill_cp(parallel) != 1 else {} + + +def _prefill_cp(parallel: Any) -> int: + """The priced prefill-CP size: unset (``None``) means one.""" + return parallel.prefill_context or 1 + + +def _decode_cp(parallel: Any) -> int: + """The priced decode-CP size: unrecorded (``None``) means one.""" + return parallel.decode_context or 1 def _worker_performance_model_metadata( @@ -449,6 +537,7 @@ def _worker_performance_model_metadata( config["speculation"] = engine.speculation.cost_config() if worker.timing.fpm_parquet_path is not None: config["fpm_parquet_path"] = worker.timing.fpm_parquet_path + config.update(_prefill_cp_knob(parallel, "cp_size")) return { "provider": "aic", "config": config, @@ -484,10 +573,6 @@ def _worker_engine_args( raise ValueError("decode_context must divide tensor parallelism") if worker.timing.type != "default" or engine.mode == "afd" or engine.workers.encoder is not None: raise ValueError("DCP requires the canonical forward-pass timing provider") - if parallel.decode_context > 1 and capacity.type != "fixed": - raise ValueError( - "DCP FPM replay requires explicit KV block capacity; automatic DCP/hybrid sizing is unsupported" - ) if capacity.type == "default" and memory_fraction is None: memory_fraction = 0.88 if backend == "sglang" else 0.9 block_size = cache.block_size @@ -525,6 +610,9 @@ def _worker_engine_args( payload[f"aic_{field}"] = value if parallel.pipeline != 1: payload["aic_pp_size"] = parallel.pipeline + payload.update(_prefill_cp_knob(parallel, "aic_cp_size")) + if _decode_cp(parallel) != 1: + payload["aic_dcp_size"] = parallel.decode_context if parallel.moe_tensor * parallel.moe_expert > 1: payload["aic_moe_tp_size"] = parallel.moe_tensor payload["aic_moe_ep_size"] = parallel.moe_expert @@ -580,6 +668,8 @@ def _worker_engine_args( "aic_model_path", "aic_moe_tp_size", "aic_moe_ep_size", + "aic_cp_size", + "aic_dcp_size", ): payload.pop(name, None) if worker.timing.type == "default" and engine.mode != "afd" and engine.workers.encoder is None: @@ -605,6 +695,7 @@ def _worker_engine_args( dcp=parallel.decode_context, moe_tp_size=parallel.moe_tensor if sharded_moe else None, moe_ep_size=parallel.moe_expert if sharded_moe else None, + **_prefill_cp_knob(parallel, "cp_size"), kv_block_size=block_size, nextn=engine.nextn, **{ @@ -678,8 +769,6 @@ def _resolve_kv_bytes_per_token( if configured != "auto": return configured parallel = worker.parallelism - if (parallel.decode_context or 1) > 1: - raise ValueError("DCP KV transfer sizing requires explicit bytes_per_token") return estimate_kv_bytes_per_token( engine.model, tp_size=parallel.tensor, diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 207ea77421..58e33b88e6 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -84,9 +84,19 @@ class ParallelismPresetConfig(StrictModel): attention_data: PositiveInt = 1 moe_tensor: PositiveInt = 1 moe_expert: PositiveInt = 1 + # Prefill context parallelism (SGLang ``--attn-cp-size`` / vLLM ``-pcp``): + # splits prefill tokens across extra attention ranks; decode stays replicated + # on them. Widens the worker like attention_data does. ``None`` (not + # requested) means 1 and keeps default dumps free of the CP knobs. + prefill_context: Annotated[int, Field(strict=True, gt=0)] | None = None class ParallelismPredictionConfig(ParallelismPresetConfig): + # Decode context parallelism (vLLM ``-dcp`` / SGLang ``--dcp-size``): stripes the + # decode KV cache across ranks that already belong to the attention group, so it + # adds no GPUs. ``None`` (not requested) is priced as 1 but keeps the FPM cell + # identity on unrecorded-DCP profiles; an explicit 1 selects recorded-DCP1 + # profiles. Aggregated workers accept at most one of the two knobs above 1. decode_context: Annotated[int, Field(strict=True, gt=0)] | None = None @@ -583,13 +593,17 @@ def _validate_custom_preset_entries(cls, value: Any) -> Any: "moe_tensor", "moe_expert", } + # The context-parallel knobs are not searched by recommend, but a + # ParallelismPredictionConfig round-tripped through model_dump carries + # them; accept them here and let recommend reject values above 1. + optional = {"prefill_context", "decode_context"} if not preset: raise ValueError("parallelism preset list must be nonempty") for index, entry in enumerate(preset): if not isinstance(entry, dict): raise ValueError(f"parallelism preset entry {index} must be a mapping") missing = required - set(entry) - unknown = set(entry) - required + unknown = set(entry) - required - optional if missing or unknown: raise ValueError( "parallelism preset entries must cover exactly all knobs; " diff --git a/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py b/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py index a550b6ef57..ab8c98e542 100644 --- a/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py +++ b/python/aisimulate/src/aisimulate/generator/builders/fpm_builder.py @@ -355,12 +355,15 @@ def _resolve_topology(context: dict[str, Any], worker: DGDService, args: list[st tensor_parallel_size = _positive_cli_int(args, "--tensor-parallel-size") pipeline_parallel_size = _positive_cli_int(args, "--pipeline-parallel-size") data_parallel_size = _positive_cli_int(args, "--data-parallel-size") - expected_gpus = tensor_parallel_size * pipeline_parallel_size * data_parallel_size + # vLLM prefill context parallelism (-pcp) expands the world size + # (world = tp * pp * pcp); decode CP (-dcp) reuses the TP ranks and adds nothing. + prefill_context_parallel_size = _positive_cli_int(args, "--prefill-context-parallel-size") + expected_gpus = tensor_parallel_size * pipeline_parallel_size * data_parallel_size * prefill_context_parallel_size if expected_gpus != total_gpus: raise ValueError( "FPM topology does not match the resolved GPU count: " f"tp({tensor_parallel_size}) * pp({pipeline_parallel_size}) * dp({data_parallel_size}) " - f"!= gpus({total_gpus})" + f"* pcp({prefill_context_parallel_size}) != gpus({total_gpus})" ) if data_parallel_size > 1: @@ -371,7 +374,8 @@ def _resolve_topology(context: dict[str, Any], worker: DGDService, args: list[st local_data_parallel_size = 1 if ( data_parallel_size > 1 - and local_data_parallel_size * tensor_parallel_size * pipeline_parallel_size > gpus_per_node + and local_data_parallel_size * tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size + > gpus_per_node ): raise ValueError("FPM local parallel topology exceeds the per-node GPU count") diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml b/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml index 7d25c96063..642ef0c6fb 100644 --- a/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml +++ b/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml @@ -41,6 +41,45 @@ parameters: sglang: null trtllm: null +# Context parallelism. Two orthogonal per-worker knobs; version gates and the +# SGLang layout choice live in generator/context_parallel.py (the single +# decision point both request paths funnel through), so an unsupported knob +# fails there instead of vanishing from a template that has no line for it. +- param_key: context_parallel_size + # Prefill CP. vLLM -pcp (>= 0.12.0; expands the world size). SGLang + # --attn-cp-size: the CP ranks live INSIDE --tp, so sglang.rule folds + # tensor_parallel_size = tp * cp for dense models (MoE is already folded to + # moe_tp * moe_ep, which equals tp * dp * cp); the SGLang template emits the + # inseparable --enable-prefill-cp switch next to it. TRT-LLM: no mapping + # (its Ulysses/Ring CP has no perf data in the simulator). + vllm: prefill-context-parallel-size + sglang: attn-cp-size + trtllm: null + +- param_key: cp_strategy + # SGLang prefill CP layout: "interleave" for DSA architectures (zigzag is + # rejected at startup for them), "zigzag" otherwise. Resolved in + # generator/context_parallel.py::cp_strategy_for. + vllm: null + sglang: cp-strategy + trtllm: null + +- param_key: decode_context_parallel_size + # Decode CP: stripes the decode KV across ranks already in the attention + # group; adds no GPUs. vLLM -dcp (>= 0.10.2), SGLang --dcp-size (>= 0.5.15). + # TRT-LLM Helix is a disaggregated-decode-only layout the simulator does not + # model: null here, rejected in context_parallel.py. + vllm: decode-context-parallel-size + sglang: dcp-size + trtllm: null + +- param_key: dcp_comm_backend + # Partial-output merge collective for DCP: vLLM ag_rs|a2a (>= 0.18.0), + # SGLang ag_rs|a2a|fi_a2a (>= 0.5.17). Unset lets the framework resolve it. + vllm: dcp-comm-backend + sglang: dcp-comm-backend + trtllm: null + # Vision encoder parallelism (VL models). # Guard: SGLang's --mm-enable-dp-encoder only exists >= 0.5.10. - param_key: enable_encoder_dp diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.15.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.15.j2 new file mode 100644 index 0000000000..bacfb777fd --- /dev/null +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.15.j2 @@ -0,0 +1,47 @@ +{# +SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 +-#} +{%- set args = [] -%} +{%- if sglang['tensor-parallel-size'] is defined and sglang['tensor-parallel-size'] is not none -%}{%- set _ = args.append('--tensor-parallel-size "' ~ sglang['tensor-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['pipeline-parallel-size'] is defined and sglang['pipeline-parallel-size'] is not none -%}{%- set _ = args.append('--pipeline-parallel-size "' ~ sglang['pipeline-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['data-parallel-size'] is defined and sglang['data-parallel-size'] is not none -%}{%- set _ = args.append('--data-parallel-size "' ~ sglang['data-parallel-size'] ~ '"') -%}{%- endif -%} +{#- Guard: context-parallel flags appear only when the request carries the knob (>1). + --attn-cp-size is inert without --enable-prefill-cp, so the pair is emitted together; + --cp-strategy must be "interleave" for DSA architectures (generator/context_parallel.py). + --dcp-size is the SGLang decode CP flag (>= 0.5.15); --dcp-comm-backend needs the 0.5.17 template. -#} +{%- if sglang['attn-cp-size'] is defined and sglang['attn-cp-size'] is not none -%}{%- set _ = args.append('--enable-prefill-cp') -%}{%- set _ = args.append('--attn-cp-size "' ~ sglang['attn-cp-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['cp-strategy'] is defined and sglang['cp-strategy'] is not none -%}{%- set _ = args.append('--cp-strategy "' ~ sglang['cp-strategy'] ~ '"') -%}{%- endif -%} +{%- if sglang['dcp-size'] is defined and sglang['dcp-size'] is not none -%}{%- set _ = args.append('--dcp-size "' ~ sglang['dcp-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['page-size'] is defined and sglang['page-size'] is not none -%}{%- set _ = args.append('--page-size "' ~ sglang['page-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['kv-cache-dtype'] is defined and sglang['kv-cache-dtype'] is not none -%}{%- set _ = args.append('--kv-cache-dtype "' ~ sglang['kv-cache-dtype'] ~ '"') -%}{%- endif -%} +{%- if sglang['mem-fraction-static'] is defined and sglang['mem-fraction-static'] is not none -%}{%- set _ = args.append('--mem-fraction-static "' ~ sglang['mem-fraction-static'] ~ '"') -%}{%- endif -%} +{%- if sglang['max-total-tokens'] is defined and sglang['max-total-tokens'] is not none -%}{%- set _ = args.append('--max-total-tokens "' ~ sglang['max-total-tokens'] ~ '"') -%}{%- endif -%} +{%- if sglang['context-length'] is defined and sglang['context-length'] is not none -%}{%- set _ = args.append('--context-length "' ~ sglang['context-length'] ~ '"') -%}{%- endif -%} +{%- if sglang['max-running-requests'] is defined and sglang['max-running-requests'] is not none -%}{%- set _ = args.append('--max-running-requests "' ~ sglang['max-running-requests'] ~ '"') -%}{%- endif -%} +{%- if sglang['chunked-prefill-size'] is defined and sglang['chunked-prefill-size'] is not none -%}{%- set _ = args.append('--chunked-prefill-size "' ~ sglang['chunked-prefill-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['max-prefill-tokens'] is defined and sglang['max-prefill-tokens'] is not none -%}{%- set _ = args.append('--max-prefill-tokens "' ~ sglang['max-prefill-tokens'] ~ '"') -%}{%- endif -%} +{%- if sglang['enable-mixed-chunk'] -%}{%- set _ = args.append('--enable-mixed-chunk') -%}{%- endif -%} +{%- if sglang['skip-tokenizer-init'] -%}{%- set _ = args.append('--skip-tokenizer-init') -%}{%- endif -%} +{%- if sglang['trust-remote-code'] -%}{%- set _ = args.append('--trust-remote-code') -%}{%- endif -%} +{%- if sglang['disable-radix-cache'] -%}{%- set _ = args.append('--disable-radix-cache') -%}{%- endif -%} +{%- if sglang['enable-dp-attention'] -%}{%- set _ = args.append('--enable-dp-attention') -%}{%- endif -%} +{%- if sglang['mm-enable-dp-encoder'] -%}{%- set _ = args.append('--mm-enable-dp-encoder') -%}{%- endif -%} +{%- if sglang['expert-parallel-size'] is defined and sglang['expert-parallel-size'] is not none -%}{%- set _ = args.append('--expert-parallel-size "' ~ sglang['expert-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['moe-runner-backend'] is defined and sglang['moe-runner-backend'] is not none -%}{%- set _ = args.append('--moe-runner-backend "' ~ sglang['moe-runner-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['moe-a2a-backend'] is defined and sglang['moe-a2a-backend'] is not none -%}{%- set _ = args.append('--moe-a2a-backend "' ~ sglang['moe-a2a-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['attention-backend'] is defined and sglang['attention-backend'] is not none -%}{%- set _ = args.append('--attention-backend "' ~ sglang['attention-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['disaggregation-transfer-backend'] is defined and sglang['disaggregation-transfer-backend'] is not none -%}{%- set _ = args.append('--disaggregation-transfer-backend "' ~ sglang['disaggregation-transfer-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['disaggregation-bootstrap-port'] is defined and sglang['disaggregation-bootstrap-port'] is not none -%}{%- set _ = args.append('--disaggregation-bootstrap-port "' ~ sglang['disaggregation-bootstrap-port'] ~ '"') -%}{%- endif -%} +{%- if sglang['elastic-ep-backend'] is defined and sglang['elastic-ep-backend'] is not none -%}{%- set _ = args.append('--elastic-ep-backend "' ~ sglang['elastic-ep-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['disable-cuda-graph'] -%}{%- set _ = args.append('--disable-cuda-graph') -%}{%- endif -%} +{%- if sglang['cuda-graph-bs'] is defined and sglang['cuda-graph-bs'] is not none -%}{%- set _ = args.append('--cuda-graph-bs ' ~ (sglang['cuda-graph-bs'] | join(' '))) -%}{%- endif -%} +{%- if sglang['disable-cuda-graph-padding'] -%}{%- set _ = args.append('--disable-cuda-graph-padding') -%}{%- endif -%} +{%- if sglang['cuda-graph-max-bs'] is defined and sglang['cuda-graph-max-bs'] is not none -%}{%- set _ = args.append('--cuda-graph-max-bs "' ~ sglang['cuda-graph-max-bs'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-algorithm'] is defined and sglang['speculative-algorithm'] is not none -%}{%- set _ = args.append('--speculative-algorithm "' ~ sglang['speculative-algorithm'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-num-steps'] is defined and sglang['speculative-num-steps'] is not none -%}{%- set _ = args.append('--speculative-num-steps "' ~ sglang['speculative-num-steps'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-eagle-topk'] is defined and sglang['speculative-eagle-topk'] is not none -%}{%- set _ = args.append('--speculative-eagle-topk "' ~ sglang['speculative-eagle-topk'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-num-draft-tokens'] is defined and sglang['speculative-num-draft-tokens'] is not none -%}{%- set _ = args.append('--speculative-num-draft-tokens "' ~ sglang['speculative-num-draft-tokens'] ~ '"') -%}{%- endif -%} +{%- if sglang['disable-overlap-schedule'] -%}{%- set _ = args.append('--disable-overlap-schedule') -%}{%- endif -%} +{%- if sglang['load-balance-method'] is defined and sglang['load-balance-method'] is not none -%}{%- set _ = args.append('--load-balance-method "' ~ sglang['load-balance-method'] ~ '"') -%}{%- endif -%} +{{ args | join(' ') }} diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.17.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.17.j2 new file mode 100644 index 0000000000..dd4e0ba846 --- /dev/null +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.17.j2 @@ -0,0 +1,48 @@ +{# +SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 +-#} +{%- set args = [] -%} +{%- if sglang['tensor-parallel-size'] is defined and sglang['tensor-parallel-size'] is not none -%}{%- set _ = args.append('--tensor-parallel-size "' ~ sglang['tensor-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['pipeline-parallel-size'] is defined and sglang['pipeline-parallel-size'] is not none -%}{%- set _ = args.append('--pipeline-parallel-size "' ~ sglang['pipeline-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['data-parallel-size'] is defined and sglang['data-parallel-size'] is not none -%}{%- set _ = args.append('--data-parallel-size "' ~ sglang['data-parallel-size'] ~ '"') -%}{%- endif -%} +{#- Guard: context-parallel flags appear only when the request carries the knob (>1). + --attn-cp-size is inert without --enable-prefill-cp, so the pair is emitted together; + --cp-strategy must be "interleave" for DSA architectures (generator/context_parallel.py). + --dcp-size (>= 0.5.15) and --dcp-comm-backend (>= 0.5.17) are the SGLang decode CP flags. -#} +{%- if sglang['attn-cp-size'] is defined and sglang['attn-cp-size'] is not none -%}{%- set _ = args.append('--enable-prefill-cp') -%}{%- set _ = args.append('--attn-cp-size "' ~ sglang['attn-cp-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['cp-strategy'] is defined and sglang['cp-strategy'] is not none -%}{%- set _ = args.append('--cp-strategy "' ~ sglang['cp-strategy'] ~ '"') -%}{%- endif -%} +{%- if sglang['dcp-size'] is defined and sglang['dcp-size'] is not none -%}{%- set _ = args.append('--dcp-size "' ~ sglang['dcp-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['dcp-comm-backend'] is defined and sglang['dcp-comm-backend'] is not none -%}{%- set _ = args.append('--dcp-comm-backend "' ~ sglang['dcp-comm-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['page-size'] is defined and sglang['page-size'] is not none -%}{%- set _ = args.append('--page-size "' ~ sglang['page-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['kv-cache-dtype'] is defined and sglang['kv-cache-dtype'] is not none -%}{%- set _ = args.append('--kv-cache-dtype "' ~ sglang['kv-cache-dtype'] ~ '"') -%}{%- endif -%} +{%- if sglang['mem-fraction-static'] is defined and sglang['mem-fraction-static'] is not none -%}{%- set _ = args.append('--mem-fraction-static "' ~ sglang['mem-fraction-static'] ~ '"') -%}{%- endif -%} +{%- if sglang['max-total-tokens'] is defined and sglang['max-total-tokens'] is not none -%}{%- set _ = args.append('--max-total-tokens "' ~ sglang['max-total-tokens'] ~ '"') -%}{%- endif -%} +{%- if sglang['context-length'] is defined and sglang['context-length'] is not none -%}{%- set _ = args.append('--context-length "' ~ sglang['context-length'] ~ '"') -%}{%- endif -%} +{%- if sglang['max-running-requests'] is defined and sglang['max-running-requests'] is not none -%}{%- set _ = args.append('--max-running-requests "' ~ sglang['max-running-requests'] ~ '"') -%}{%- endif -%} +{%- if sglang['chunked-prefill-size'] is defined and sglang['chunked-prefill-size'] is not none -%}{%- set _ = args.append('--chunked-prefill-size "' ~ sglang['chunked-prefill-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['max-prefill-tokens'] is defined and sglang['max-prefill-tokens'] is not none -%}{%- set _ = args.append('--max-prefill-tokens "' ~ sglang['max-prefill-tokens'] ~ '"') -%}{%- endif -%} +{%- if sglang['enable-mixed-chunk'] -%}{%- set _ = args.append('--enable-mixed-chunk') -%}{%- endif -%} +{%- if sglang['skip-tokenizer-init'] -%}{%- set _ = args.append('--skip-tokenizer-init') -%}{%- endif -%} +{%- if sglang['trust-remote-code'] -%}{%- set _ = args.append('--trust-remote-code') -%}{%- endif -%} +{%- if sglang['disable-radix-cache'] -%}{%- set _ = args.append('--disable-radix-cache') -%}{%- endif -%} +{%- if sglang['enable-dp-attention'] -%}{%- set _ = args.append('--enable-dp-attention') -%}{%- endif -%} +{%- if sglang['mm-enable-dp-encoder'] -%}{%- set _ = args.append('--mm-enable-dp-encoder') -%}{%- endif -%} +{%- if sglang['expert-parallel-size'] is defined and sglang['expert-parallel-size'] is not none -%}{%- set _ = args.append('--expert-parallel-size "' ~ sglang['expert-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if sglang['moe-runner-backend'] is defined and sglang['moe-runner-backend'] is not none -%}{%- set _ = args.append('--moe-runner-backend "' ~ sglang['moe-runner-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['moe-a2a-backend'] is defined and sglang['moe-a2a-backend'] is not none -%}{%- set _ = args.append('--moe-a2a-backend "' ~ sglang['moe-a2a-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['attention-backend'] is defined and sglang['attention-backend'] is not none -%}{%- set _ = args.append('--attention-backend "' ~ sglang['attention-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['disaggregation-transfer-backend'] is defined and sglang['disaggregation-transfer-backend'] is not none -%}{%- set _ = args.append('--disaggregation-transfer-backend "' ~ sglang['disaggregation-transfer-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['disaggregation-bootstrap-port'] is defined and sglang['disaggregation-bootstrap-port'] is not none -%}{%- set _ = args.append('--disaggregation-bootstrap-port "' ~ sglang['disaggregation-bootstrap-port'] ~ '"') -%}{%- endif -%} +{%- if sglang['elastic-ep-backend'] is defined and sglang['elastic-ep-backend'] is not none -%}{%- set _ = args.append('--elastic-ep-backend "' ~ sglang['elastic-ep-backend'] ~ '"') -%}{%- endif -%} +{%- if sglang['disable-cuda-graph'] -%}{%- set _ = args.append('--disable-cuda-graph') -%}{%- endif -%} +{%- if sglang['cuda-graph-bs'] is defined and sglang['cuda-graph-bs'] is not none -%}{%- set _ = args.append('--cuda-graph-bs ' ~ (sglang['cuda-graph-bs'] | join(' '))) -%}{%- endif -%} +{%- if sglang['disable-cuda-graph-padding'] -%}{%- set _ = args.append('--disable-cuda-graph-padding') -%}{%- endif -%} +{%- if sglang['cuda-graph-max-bs'] is defined and sglang['cuda-graph-max-bs'] is not none -%}{%- set _ = args.append('--cuda-graph-max-bs "' ~ sglang['cuda-graph-max-bs'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-algorithm'] is defined and sglang['speculative-algorithm'] is not none -%}{%- set _ = args.append('--speculative-algorithm "' ~ sglang['speculative-algorithm'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-num-steps'] is defined and sglang['speculative-num-steps'] is not none -%}{%- set _ = args.append('--speculative-num-steps "' ~ sglang['speculative-num-steps'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-eagle-topk'] is defined and sglang['speculative-eagle-topk'] is not none -%}{%- set _ = args.append('--speculative-eagle-topk "' ~ sglang['speculative-eagle-topk'] ~ '"') -%}{%- endif -%} +{%- if sglang['speculative-num-draft-tokens'] is defined and sglang['speculative-num-draft-tokens'] is not none -%}{%- set _ = args.append('--speculative-num-draft-tokens "' ~ sglang['speculative-num-draft-tokens'] ~ '"') -%}{%- endif -%} +{%- if sglang['disable-overlap-schedule'] -%}{%- set _ = args.append('--disable-overlap-schedule') -%}{%- endif -%} +{%- if sglang['load-balance-method'] is defined and sglang['load-balance-method'] is not none -%}{%- set _ = args.append('--load-balance-method "' ~ sglang['load-balance-method'] ~ '"') -%}{%- endif -%} +{{ args | join(' ') }} diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.14.1.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.14.1.j2 index 51ebd247f6..51d7f1d30d 100644 --- a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.14.1.j2 +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.14.1.j2 @@ -6,6 +6,10 @@ SPDX-License-Identifier: Apache-2.0 {%- if vllm['tensor-parallel-size'] is defined -%}{%- set _ = args.append('--tensor-parallel-size "' ~ vllm['tensor-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['pipeline-parallel-size'] is defined -%}{%- set _ = args.append('--pipeline-parallel-size "' ~ vllm['pipeline-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['data-parallel-size'] is defined -%}{%- set _ = args.append('--data-parallel-size "' ~ vllm['data-parallel-size'] ~ '"') -%}{%- endif -%} +{#- Guard: context-parallel flags appear only when the request carries the knob (>1); + generator/context_parallel.py version-gates them (PCP and DCP >= 0.14.1, the first template that renders them). No --dcp-comm-backend before the 0.19.0 template. -#} +{%- if vllm['prefill-context-parallel-size'] is defined and vllm['prefill-context-parallel-size'] is not none -%}{%- set _ = args.append('--prefill-context-parallel-size "' ~ vllm['prefill-context-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if vllm['decode-context-parallel-size'] is defined and vllm['decode-context-parallel-size'] is not none -%}{%- set _ = args.append('--decode-context-parallel-size "' ~ vllm['decode-context-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['enable-expert-parallel'] -%}{%- set _ = args.append('--enable-expert-parallel') -%}{%- endif -%} {%- if vllm['mm-encoder-tp-mode'] is defined -%}{%- set _ = args.append('--mm-encoder-tp-mode "' ~ vllm['mm-encoder-tp-mode'] ~ '"') -%}{%- endif -%} {%- if vllm['block-size'] is defined -%}{%- set _ = args.append('--block-size "' ~ vllm['block-size'] ~ '"') -%}{%- endif -%} diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.16.0.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.16.0.j2 index 51ebd247f6..51d7f1d30d 100644 --- a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.16.0.j2 +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.16.0.j2 @@ -6,6 +6,10 @@ SPDX-License-Identifier: Apache-2.0 {%- if vllm['tensor-parallel-size'] is defined -%}{%- set _ = args.append('--tensor-parallel-size "' ~ vllm['tensor-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['pipeline-parallel-size'] is defined -%}{%- set _ = args.append('--pipeline-parallel-size "' ~ vllm['pipeline-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['data-parallel-size'] is defined -%}{%- set _ = args.append('--data-parallel-size "' ~ vllm['data-parallel-size'] ~ '"') -%}{%- endif -%} +{#- Guard: context-parallel flags appear only when the request carries the knob (>1); + generator/context_parallel.py version-gates them (PCP and DCP >= 0.14.1, the first template that renders them). No --dcp-comm-backend before the 0.19.0 template. -#} +{%- if vllm['prefill-context-parallel-size'] is defined and vllm['prefill-context-parallel-size'] is not none -%}{%- set _ = args.append('--prefill-context-parallel-size "' ~ vllm['prefill-context-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if vllm['decode-context-parallel-size'] is defined and vllm['decode-context-parallel-size'] is not none -%}{%- set _ = args.append('--decode-context-parallel-size "' ~ vllm['decode-context-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['enable-expert-parallel'] -%}{%- set _ = args.append('--enable-expert-parallel') -%}{%- endif -%} {%- if vllm['mm-encoder-tp-mode'] is defined -%}{%- set _ = args.append('--mm-encoder-tp-mode "' ~ vllm['mm-encoder-tp-mode'] ~ '"') -%}{%- endif -%} {%- if vllm['block-size'] is defined -%}{%- set _ = args.append('--block-size "' ~ vllm['block-size'] ~ '"') -%}{%- endif -%} diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.19.0.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.19.0.j2 index 51ebd247f6..57dca748c9 100644 --- a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.19.0.j2 +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.19.0.j2 @@ -6,6 +6,11 @@ SPDX-License-Identifier: Apache-2.0 {%- if vllm['tensor-parallel-size'] is defined -%}{%- set _ = args.append('--tensor-parallel-size "' ~ vllm['tensor-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['pipeline-parallel-size'] is defined -%}{%- set _ = args.append('--pipeline-parallel-size "' ~ vllm['pipeline-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['data-parallel-size'] is defined -%}{%- set _ = args.append('--data-parallel-size "' ~ vllm['data-parallel-size'] ~ '"') -%}{%- endif -%} +{#- Guard: context-parallel flags appear only when the request carries the knob (>1); + generator/context_parallel.py version-gates them (PCP and DCP >= 0.14.1, the first template that renders them; comm backend >= 0.19.0, the first template that renders it). -#} +{%- if vllm['prefill-context-parallel-size'] is defined and vllm['prefill-context-parallel-size'] is not none -%}{%- set _ = args.append('--prefill-context-parallel-size "' ~ vllm['prefill-context-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if vllm['decode-context-parallel-size'] is defined and vllm['decode-context-parallel-size'] is not none -%}{%- set _ = args.append('--decode-context-parallel-size "' ~ vllm['decode-context-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if vllm['dcp-comm-backend'] is defined and vllm['dcp-comm-backend'] is not none -%}{%- set _ = args.append('--dcp-comm-backend "' ~ vllm['dcp-comm-backend'] ~ '"') -%}{%- endif -%} {%- if vllm['enable-expert-parallel'] -%}{%- set _ = args.append('--enable-expert-parallel') -%}{%- endif -%} {%- if vllm['mm-encoder-tp-mode'] is defined -%}{%- set _ = args.append('--mm-encoder-tp-mode "' ~ vllm['mm-encoder-tp-mode'] ~ '"') -%}{%- endif -%} {%- if vllm['block-size'] is defined -%}{%- set _ = args.append('--block-size "' ~ vllm['block-size'] ~ '"') -%}{%- endif -%} diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.20.1.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.20.1.j2 index f95b17032f..69af841f52 100644 --- a/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.20.1.j2 +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/vllm/cli_args.0.20.1.j2 @@ -6,6 +6,11 @@ SPDX-License-Identifier: Apache-2.0 {%- if vllm['tensor-parallel-size'] is defined -%}{%- set _ = args.append('--tensor-parallel-size "' ~ vllm['tensor-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['pipeline-parallel-size'] is defined -%}{%- set _ = args.append('--pipeline-parallel-size "' ~ vllm['pipeline-parallel-size'] ~ '"') -%}{%- endif -%} {%- if vllm['data-parallel-size'] is defined -%}{%- set _ = args.append('--data-parallel-size "' ~ vllm['data-parallel-size'] ~ '"') -%}{%- endif -%} +{#- Guard: context-parallel flags appear only when the request carries the knob (>1); + generator/context_parallel.py version-gates them (PCP and DCP >= 0.14.1, the first template that renders them; comm backend >= 0.19.0, the first template that renders it). -#} +{%- if vllm['prefill-context-parallel-size'] is defined and vllm['prefill-context-parallel-size'] is not none -%}{%- set _ = args.append('--prefill-context-parallel-size "' ~ vllm['prefill-context-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if vllm['decode-context-parallel-size'] is defined and vllm['decode-context-parallel-size'] is not none -%}{%- set _ = args.append('--decode-context-parallel-size "' ~ vllm['decode-context-parallel-size'] ~ '"') -%}{%- endif -%} +{%- if vllm['dcp-comm-backend'] is defined and vllm['dcp-comm-backend'] is not none -%}{%- set _ = args.append('--dcp-comm-backend "' ~ vllm['dcp-comm-backend'] ~ '"') -%}{%- endif -%} {%- if vllm['enable-expert-parallel'] -%}{%- set _ = args.append('--enable-expert-parallel') -%}{%- endif -%} {%- if vllm['mm-encoder-tp-mode'] is defined -%}{%- set _ = args.append('--mm-encoder-tp-mode "' ~ vllm['mm-encoder-tp-mode'] ~ '"') -%}{%- endif -%} {%- if vllm['block-size'] is defined -%}{%- set _ = args.append('--block-size "' ~ vllm['block-size'] ~ '"') -%}{%- endif -%} diff --git a/python/aisimulate/src/aisimulate/generator/context_parallel.py b/python/aisimulate/src/aisimulate/generator/context_parallel.py new file mode 100644 index 0000000000..48652b067b --- /dev/null +++ b/python/aisimulate/src/aisimulate/generator/context_parallel.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Context-parallel knobs for deployment rendering: the single decision point +for which backend/version can launch prefill CP and decode CP, and how. + +The simulator carries two orthogonal per-worker knobs: + +* ``context_parallel_size`` -- prefill context parallelism (SGLang + ``--attn-cp-size`` + ``--enable-prefill-cp``, vLLM ``-pcp``). Extra + attention ranks that split the prefill tokens; decode stays replicated on + them, so the worker grows by ``cp`` GPUs. +* ``decode_context_parallel_size`` -- decode context parallelism (vLLM + ``-dcp``, SGLang ``--dcp-size``). Stripes the decode KV across ranks that + already belong to the attention group; adds no GPUs. + +Both entry points into the generator (the SDK result bridge and the Sweeper +candidate request) funnel through :func:`context_parallel_params` so a knob a +backend/version cannot launch fails here, loudly, instead of being dropped by +a template that has no line for it (the "silently wrong deployment" pattern). + +# Guard: the minimum versions below are the FIRST release tag containing the +# flag in the framework's own history (verified 2026-09-17 against the vLLM and +# SGLang git logs), floored to the versioned cli_args template that renders +# the flag. Re-verify on every backend version bump. +""" + +from __future__ import annotations + +from typing import Any + +from packaging.version import InvalidVersion, Version + +from aisimulate_core.sdk.config import validate_parallel_size + +# DeepSeek Sparse Attention architectures: SGLang's prefill CP requires the +# ``interleave`` layout for them (zigzag is rejected at startup); dense and +# plain-MLA models use ``zigzag``. +DSA_ARCHITECTURES: frozenset[str] = frozenset({"DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"}) +DSA_MODEL_FAMILIES: frozenset[str] = frozenset({"DEEPSEEKV32"}) + +_MIN_VERSIONS: dict[tuple[str, str], Version] = { + # vLLM: DCP landed in v0.10.2 (#23734) and PCP in v0.12.0 (#28718), but the + # generator only renders the flags from cli_args.0.14.1.j2 onwards (the + # 0.10.2 / 0.11.0 / 0.12.0 templates predate them and are frozen), and + # template selection is a floor match. The floor is therefore the first + # template that emits the flags, not the framework release, so a request + # can never resolve to a template that silently drops the knob. The DCP + # communication backend selector arrived in v0.18.0 (#34883), but the first + # template that renders --dcp-comm-backend is cli_args.0.19.0.j2 (0.18.x + # would floor-match the 0.16.0 template and drop the flag). + ("vllm", "decode_context_parallel_size"): Version("0.14.1"), + ("vllm", "context_parallel_size"): Version("0.14.1"), + ("vllm", "dcp_comm_backend"): Version("0.19.0"), + # SGLang: --enable-prefill-cp/--cp-strategy in v0.5.14 (#27312) and + # --dcp-size in v0.5.15 (#25090); the generator renders both from the + # cli_args.0.5.15 template, so the floor is 0.5.15 for either knob. + # --dcp-comm-backend arrived in v0.5.17 (cli_args.0.5.17 template). + ("sglang", "context_parallel_size"): Version("0.5.15"), + ("sglang", "decode_context_parallel_size"): Version("0.5.15"), + ("sglang", "dcp_comm_backend"): Version("0.5.17"), +} + +_DCP_COMM_BACKENDS: dict[str, frozenset[str]] = { + "vllm": frozenset({"ag_rs", "a2a"}), + "sglang": frozenset({"ag_rs", "a2a", "fi_a2a"}), +} + + +class ContextParallelUnsupportedError(ValueError): + """The requested CP knob cannot be launched on this backend/version.""" + + +def _parse_version(backend_version: str | None) -> Version | None: + if backend_version is None: + return None + text = str(backend_version).strip() + if text.lower().startswith("v") and len(text) > 1: + text = text[1:] + try: + return Version(text) + except InvalidVersion: + return None + + +def _positive(value: Any, name: str) -> int: + """Unset (``None``) means 1; anything else must be a positive integer.""" + if value is None: + return 1 + try: + return validate_parallel_size(name, value) + except ValueError as exc: + raise ContextParallelUnsupportedError(str(exc)) from exc + + +def cp_strategy_for(architecture: str | None, model_family: str | None = None) -> str: + """SGLang ``--cp-strategy`` for a model: ``interleave`` for DSA, else ``zigzag``.""" + if architecture in DSA_ARCHITECTURES or (model_family or "").upper() in DSA_MODEL_FAMILIES: + return "interleave" + return "zigzag" + + +def _require_version(backend: str, knob: str, backend_version: str | None, value: Any) -> None: + minimum = _MIN_VERSIONS.get((backend, knob)) + if minimum is None: + raise ContextParallelUnsupportedError( + f"{knob}={value!r} cannot be rendered for backend {backend!r}: the generator has no " + "launch flag for it on this backend" + + ( + " (TensorRT-LLM's Helix context parallelism is a disaggregated-decode-only layout " + "that the simulator does not model)" + if backend == "trtllm" + else "" + ) + ) + parsed = _parse_version(backend_version) + if parsed is None: + raise ContextParallelUnsupportedError( + f"{knob}={value!r} requires a parseable {backend} backend_version to check support " + f"(need >= {minimum}), got {backend_version!r}" + ) + if parsed < minimum: + raise ContextParallelUnsupportedError( + f"{knob}={value!r} requires {backend} >= {minimum} (flag not available in {backend_version})" + ) + + +def context_parallel_params( + *, + backend: str, + backend_version: str | None, + context_parallel_size: Any = 1, + decode_context_parallel_size: Any = 1, + dcp_comm_backend: str | None = None, + architecture: str | None = None, + model_family: str | None = None, +) -> dict[str, Any]: + """Per-worker generator params for the CP knobs, or ``{}`` when both are 1. + + Raises :class:`ContextParallelUnsupportedError` when the backend/version + has no launch flag for a knob that is set, so an evaluated CP deployment + can never degrade into an un-CP'd launch. + """ + cp = _positive(context_parallel_size, "context_parallel_size") + dcp = _positive(decode_context_parallel_size, "decode_context_parallel_size") + params: dict[str, Any] = {} + if cp > 1: + _require_version(backend, "context_parallel_size", backend_version, cp) + params["context_parallel_size"] = cp + if backend == "sglang": + params["cp_strategy"] = cp_strategy_for(architecture, model_family) + if dcp > 1: + _require_version(backend, "decode_context_parallel_size", backend_version, dcp) + params["decode_context_parallel_size"] = dcp + if dcp_comm_backend: + if dcp <= 1: + raise ContextParallelUnsupportedError( + f"dcp_comm_backend={dcp_comm_backend!r} needs decode_context_parallel_size > 1" + ) + allowed = _DCP_COMM_BACKENDS.get(backend, frozenset()) + if dcp_comm_backend not in allowed: + raise ContextParallelUnsupportedError( + f"dcp_comm_backend={dcp_comm_backend!r} is not a {backend} choice; expected one of {sorted(allowed)}" + ) + _require_version(backend, "dcp_comm_backend", backend_version, dcp_comm_backend) + params["dcp_comm_backend"] = dcp_comm_backend + return params diff --git a/python/aisimulate/src/aisimulate/generator/module_bridge.py b/python/aisimulate/src/aisimulate/generator/module_bridge.py index d6fea2c909..b68b494d5f 100644 --- a/python/aisimulate/src/aisimulate/generator/module_bridge.py +++ b/python/aisimulate/src/aisimulate/generator/module_bridge.py @@ -12,6 +12,7 @@ from aisimulate.sdk.task_v2 import Task from .aggregators import collect_generator_params +from .context_parallel import ContextParallelUnsupportedError, context_parallel_params from .rendering import apply_defaults @@ -150,6 +151,11 @@ def _build_worker_params(prefix: str, extra_overrides: dict | None) -> tuple[dic dp = _safe_int(_series_val(result_df, f"{prefix}dp", 1), 1) moe_tp = _safe_int(_series_val(result_df, f"{prefix}moe_tp", 1), 1) moe_ep = _safe_int(_series_val(result_df, f"{prefix}moe_ep", 1), 1) + # Prefill CP is a swept column of the result frame; decode CP is the + # task's per-role scalar (dcp_size / prefill_dcp_size / decode_dcp_size). + cp = _safe_int(_series_val(result_df, f"{prefix}cp", 1), 1) + dcp_attr = {"": "dcp_size", "(p)": "prefill_dcp_size", "(d)": "decode_dcp_size"}.get(prefix, "dcp_size") + dcp = _safe_int(getattr(task_config, dcp_attr, 1), 1) bs = _safe_int(_series_val(result_df, f"{prefix}bs", 1), 1) memory = _safe_float(_series_val(result_df, f"{prefix}memory", None), None) @@ -165,11 +171,14 @@ def _build_worker_params(prefix: str, extra_overrides: dict | None) -> tuple[dic "tensor_parallel_size": tp, "pipeline_parallel_size": pp, "data_parallel_size": dp, - "gpus_per_worker": tp * pp * dp, "moe_tensor_parallel_size": moe_tp, "moe_expert_parallel_size": moe_ep, "max_batch_size": bs, **{k: v for k, v in quant.items() if v is not None}, + # Raw task-derived CP knobs; validated ONCE below, after the + # Workers. overrides have had their say. + **{k: v for k, v in (("context_parallel_size", cp), ("decode_context_parallel_size", dcp)) if v != 1}, + **{k: v for k, v in (("dcp_comm_backend", getattr(task_config, "dcp_comm", None)),) if v is not None}, } if memory is not None: @@ -182,6 +191,37 @@ def _build_worker_params(prefix: str, extra_overrides: dict | None) -> tuple[dic worker_payload["attention_backend"] = attention_backend worker_payload = _deep_merge(worker_payload, extra_overrides) + + # The single decision point sees the EFFECTIVE knobs (task values, then + # overrides) exactly once: it rejects malformed values (2.5, "x", True) + # and knobs the backend version cannot launch, and derives cp_strategy. + effective = context_parallel_params( + backend=task_config.primary_backend_name, + backend_version=getattr(task_config, "primary_backend_version", None), + context_parallel_size=worker_payload.pop("context_parallel_size", 1), + decode_context_parallel_size=worker_payload.pop("decode_context_parallel_size", 1), + dcp_comm_backend=worker_payload.pop("dcp_comm_backend", None), + architecture=getattr(task_config, "architecture", None), + model_family=getattr(task_config, "model_family", None) or None, + ) + # cp_strategy follows the architecture (SGLang rejects zigzag for DSA + # at startup); an override may only restate it. + requested_strategy = worker_payload.pop("cp_strategy", None) + if requested_strategy is not None and requested_strategy != effective.get("cp_strategy"): + raise ContextParallelUnsupportedError( + f"Workers override cp_strategy={requested_strategy!r} does not match the strategy this " + f"model requires ({effective.get('cp_strategy')!r}); drop the override or fix the model" + ) + worker_payload.update(effective) + # Prefill CP ranks are extra attention GPUs (folded into --tp by the + # SGLang rules, an expanded world size on vLLM); decode CP adds none. + if "gpus_per_worker" not in worker_payload: + worker_payload["gpus_per_worker"] = ( + _safe_int(worker_payload.get("tensor_parallel_size"), tp) + * _safe_int(worker_payload.get("pipeline_parallel_size"), pp) + * _safe_int(worker_payload.get("data_parallel_size"), dp) + * effective.get("context_parallel_size", 1) + ) effective_attention_backend = worker_payload.get("attention_backend") if effective_attention_backend is not None: # Task normalization promotes CLI strings to AttentionBackend @@ -267,7 +307,8 @@ def _build_worker_params(prefix: str, extra_overrides: dict | None) -> tuple[dic tp = agg_params.get("tensor_parallel_size", 1) pp = agg_params.get("pipeline_parallel_size", 1) dp = agg_params.get("data_parallel_size", 1) - gpus_per_replica = tp * pp * dp + # gpus_per_worker already counts the prefill-CP ranks. + gpus_per_replica = agg_params.get("gpus_per_worker") or tp * pp * dp agg_workers = effective_total_gpus // gpus_per_replica prefill_params, prefill_workers = None, 0 decode_params, decode_workers = None, 0 @@ -281,12 +322,13 @@ def _build_worker_params(prefix: str, extra_overrides: dict | None) -> tuple[dic p_tp = prefill_params.get("tensor_parallel_size", 1) p_pp = prefill_params.get("pipeline_parallel_size", 1) p_dp = prefill_params.get("data_parallel_size", 1) - prefill_gpus_per_worker = p_tp * p_pp * p_dp + # gpus_per_worker already counts the prefill-CP ranks. + prefill_gpus_per_worker = prefill_params.get("gpus_per_worker") or p_tp * p_pp * p_dp d_tp = decode_params.get("tensor_parallel_size", 1) d_pp = decode_params.get("pipeline_parallel_size", 1) d_dp = decode_params.get("data_parallel_size", 1) - decode_gpus_per_worker = d_tp * d_pp * d_dp + decode_gpus_per_worker = decode_params.get("gpus_per_worker") or d_tp * d_pp * d_dp # Each replica uses prefill_workers_per_replica prefill workers + decode_workers_per_replica decode workers # For simplicity, assume 1:1 prefill:decode ratio per replica diff --git a/python/aisimulate/src/aisimulate/generator/request/sweeper.py b/python/aisimulate/src/aisimulate/generator/request/sweeper.py index c0e8bd3bb5..169ba79c23 100644 --- a/python/aisimulate/src/aisimulate/generator/request/sweeper.py +++ b/python/aisimulate/src/aisimulate/generator/request/sweeper.py @@ -16,6 +16,7 @@ from dataclasses import asdict, is_dataclass, replace from typing import Any +from ..context_parallel import ContextParallelUnsupportedError, context_parallel_params from .schema import ( BackendSpec, EmitTargets, @@ -111,7 +112,14 @@ def _flatten_overrides(overrides: Mapping[str, Any] | None) -> dict[str, Any]: return raw -def _role_sizing(config: Mapping[str, Any], role: str, *, backend: str) -> tuple[RoleSizing, int]: +def _role_sizing( + config: Mapping[str, Any], + role: str, + *, + backend: str, + backend_version: str, + architecture: str | None = None, +) -> tuple[RoleSizing, int]: prefix = "" if role == "agg" else f"{role}_" tp = _positive_int(config.get(f"{prefix}tp"), path=f"candidate.config.{prefix}tp") @@ -128,6 +136,10 @@ def _role_sizing(config: Mapping[str, Any], role: str, *, backend: str) -> tuple config.get(f"{prefix}moe_ep"), path=f"candidate.config.{prefix}moe_ep", ) + # Context-parallel knobs are optional columns: prefill CP (``cp``) adds + # attention ranks, decode CP (``dcp``) stripes KV inside the TP group. + cp = _positive_int(config.get(f"{prefix}cp", 1), path=f"candidate.config.{prefix}cp") + dcp = _positive_int(config.get(f"{prefix}dcp", 1), path=f"candidate.config.{prefix}dcp") workers_key = "replicas" if role == "agg" else f"{role}_replicas" workers = _positive_int(config.get(workers_key), path=f"candidate.config.{workers_key}") @@ -150,11 +162,24 @@ def _role_sizing(config: Mapping[str, Any], role: str, *, backend: str) -> tuple if memory_value is not None and blocks_value is not None: raise SweeperCandidateError(f"{memory_path} and {blocks_path} are mutually exclusive") + try: + cp_params = context_parallel_params( + backend=backend, + backend_version=backend_version, + context_parallel_size=cp, + decode_context_parallel_size=dcp, + architecture=architecture, + ) + except ContextParallelUnsupportedError as exc: + raise SweeperCandidateError(f"candidate.config.{prefix}cp/dcp: {exc}") from exc + extra: dict[str, Any] = { - "gpus_per_worker": tp * pp * dp, + # Prefill CP ranks are extra GPUs; decode CP reuses the TP ranks. + "gpus_per_worker": tp * pp * dp * cp, "max_num_tokens": max_num_tokens, "tokens_per_block": tokens_per_block, "disable_prefix_cache": not bool(config.get(f"{role}_enable_prefix_caching")), + **cp_params, } if blocks_value is not None: num_gpu_blocks = _positive_int(blocks_value, path=blocks_path) @@ -302,12 +327,22 @@ def from_sweeper_candidate( ) hardware_sku = role_hardware["prefill"] + # Resolved before role sizing: SGLang's prefill-CP layout (--cp-strategy) + # depends on the model architecture. + resolved_model_facts = model_facts or _resolve_model_facts(model_path, config) + active_roles = ("agg",) if mode == "agg" else ("prefill", "decode") roles: dict[str, RoleSizing] = {} workers: dict[str, int] = {} expected_gpus = 0 for role in active_roles: - sizing, count = _role_sizing(config, role, backend=backend) + sizing, count = _role_sizing( + config, + role, + backend=backend, + backend_version=backend_version, + architecture=resolved_model_facts.architecture, + ) roles[role] = sizing workers[role] = count expected_gpus += count * int(sizing.extra["gpus_per_worker"]) @@ -356,7 +391,6 @@ def from_sweeper_candidate( raw["BenchConfig.estimated_concurrency"] = _positive_int(concurrency, path="candidate.config.concurrency") _map_adapters(config, raw) - resolved_model_facts = model_facts or _resolve_model_facts(model_path, config) candidate_nextn = config.get("aic_nextn") if candidate_nextn is not None and resolved_model_facts.nextn != candidate_nextn: resolved_model_facts = ModelFacts( diff --git a/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/sglang.rule b/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/sglang.rule index 3b192c3a20..ad8b0bd627 100644 --- a/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/sglang.rule +++ b/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/sglang.rule @@ -19,7 +19,12 @@ when ModelConfig.is_moe and (moe_tensor_parallel_size and moe_expert_parallel_si agg_prefill_decode tensor_parallel_size = moe_tensor_parallel_size * moe_expert_parallel_size # GPUs per worker follow the same TP/PP/DP product that SGLang expects -agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (data_parallel_size or 1) +# Prefill CP ranks live inside --tp on SGLang (see rule_plugin/sglang.rule for +# why this is a per-role expression rather than a `when` block). +agg_prefill_decode tensor_parallel_size = ((tensor_parallel_size or 1) * context_parallel_size) if ((context_parallel_size or 1) > 1 and not ModelConfig.is_moe) else (tensor_parallel_size if tensor_parallel_size is defined else none) + +# MoE fold already equals the attention-group width; --dp partitions --tp (see rule_plugin/sglang.rule). +agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (1 if (ModelConfig.is_moe and (moe_tensor_parallel_size and moe_expert_parallel_size)) else (data_parallel_size or 1)) agg_prefill_decode enable_attention_dp = ((data_parallel_size or 1) > 1) and ModelConfig.is_moe diff --git a/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/vllm.rule b/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/vllm.rule index fdfa297c2c..93a8b168bf 100644 --- a/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/vllm.rule +++ b/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/vllm.rule @@ -7,7 +7,8 @@ when not (preserve_engine_limits or false): decode max_batch_size = (max_batch_size if max_batch_size else 128) -agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (data_parallel_size or 1) +# vLLM prefill CP (-pcp) expands the world size; decode CP adds no GPUs. +agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (data_parallel_size or 1) * (context_parallel_size or 1) agg_prefill_decode enable_expert_parallel = ((moe_expert_parallel_size or 1) > 1) agg_prefill_decode cuda_graph_batch_sizes = ((range(1, max_batch_size + 1) | list) if max_batch_size else []) diff --git a/python/aisimulate/src/aisimulate/generator/rule_plugin/sglang.rule b/python/aisimulate/src/aisimulate/generator/rule_plugin/sglang.rule index 38c5b9306f..a5a076b56c 100644 --- a/python/aisimulate/src/aisimulate/generator/rule_plugin/sglang.rule +++ b/python/aisimulate/src/aisimulate/generator/rule_plugin/sglang.rule @@ -20,8 +20,24 @@ agg enable_mixed_chunk = true when ModelConfig.is_moe and (moe_tensor_parallel_size and moe_expert_parallel_size): agg_prefill_decode tensor_parallel_size = moe_tensor_parallel_size * moe_expert_parallel_size +# Prefill context parallelism: SGLang's attention-CP ranks live INSIDE --tp +# (attn_tp = tp / dp / attn_cp), so the launched TP is the attention-group width +# tp x cp for dense models. MoE models are already folded above to +# moe_tp x moe_ep, which equals tp x dp x cp by the simulator's width identity. +# Guard: keep this AFTER the MoE fold and BEFORE gpus_per_worker. It must stay a +# per-role conditional expression, not a `when` block: `when` is evaluated once +# over the merged agg_prefill_decode scope, so a prefill-only cp would fold the +# decode worker's TP too. (--enable-prefill-cp is emitted by the template +# alongside --attn-cp-size; the pair is inseparable.) +agg_prefill_decode tensor_parallel_size = ((tensor_parallel_size or 1) * context_parallel_size) if ((context_parallel_size or 1) > 1 and not ModelConfig.is_moe) else (tensor_parallel_size if tensor_parallel_size is defined else none) + # GPUs per worker follow the same TP/PP/DP product that SGLang expects -agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (data_parallel_size or 1) +# (prefill CP is already folded into TP above; decode CP adds no GPUs). +# Guard: when the MoE fold above applied, tensor_parallel_size is already +# moe_tp x moe_ep == the whole attention-group width (tp x dp x cp): SGLang's +# --dp-size partitions that --tp, it does not multiply it, so DP must not be +# counted again. +agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (1 if (ModelConfig.is_moe and (moe_tensor_parallel_size and moe_expert_parallel_size)) else (data_parallel_size or 1)) agg_prefill_decode enable_attention_dp = ((data_parallel_size or 1) > 1) and ModelConfig.is_moe diff --git a/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule b/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule index f197b91049..81e690b19f 100644 --- a/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule +++ b/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule @@ -10,11 +10,16 @@ when not (preserve_engine_limits or false): # Enforce vLLM MoE parallelism: vLLM uses unified TP for both attention and MoE, # so tensor_parallel_size must account for moe_tp × moe_ep. # Without this, MoE models get tp=1 and the full model is replicated per GPU, causing OOM. +# The evaluated attention width is tp * dp * cp == moe_tp * moe_ep; prefill CP +# ranks are separate -pcp ranks on vLLM, so they must NOT be folded into --tp +# (gpus_per_worker below multiplies by context_parallel_size once). when ModelConfig.is_moe and (moe_tensor_parallel_size and moe_expert_parallel_size): - agg_prefill_decode tensor_parallel_size = (moe_tensor_parallel_size * moe_expert_parallel_size) // (data_parallel_size or 1) + agg_prefill_decode tensor_parallel_size = (moe_tensor_parallel_size * moe_expert_parallel_size) // ((data_parallel_size or 1) * (context_parallel_size or 1)) -# GPUs per worker follow the same TP/PP/DP product -agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (data_parallel_size or 1) +# GPUs per worker follow the TP/PP/DP product. vLLM's prefill context +# parallelism (-pcp) EXPANDS the world size (world = tp * pp * pcp), so it +# multiplies the worker; decode CP (-dcp) reuses TP ranks and adds no GPUs. +agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_parallel_size or 1) * (data_parallel_size or 1) * (context_parallel_size or 1) agg_prefill_decode enable_expert_parallel = ((moe_expert_parallel_size or 1) > 1) when not (preserve_engine_limits or false): diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index 8997154b0d..d40a0c4ccf 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -520,9 +520,24 @@ def _parallel_entries(role: str, raw: dict[str, Any]) -> tuple[str, Any]: return "independent", {"choices": choices, "log_ranges": log_ranges} +_CONTEXT_PARALLEL_KEYS = ("prefill_context", "decode_context") + + def _parallel_mapping(value: Any, path: str) -> dict[str, int]: if not isinstance(value, dict): raise ValueError(f"{path} entries must be mappings") + # Context parallelism is a predict-only knob: preset entries are full + # ParallelismPredictionConfig dumps, so they may carry the two keys unset + # (None) or at 1. The sweeper does not enumerate CP/DCP yet, so anything + # else is rejected explicitly instead of being silently dropped. + value = dict(value) + for key in _CONTEXT_PARALLEL_KEYS: + leaf = value.pop(key, None) + if leaf not in (None, 1): + raise ValueError( + f"{path}.{key}={leaf!r} is not supported by recommend; context parallelism " + "is a predict-only knob (use `aisimulate predict` with a fixed parallelism)" + ) missing = set(_PARALLEL_KEYS) - set(value) unknown = set(value) - set(_PARALLEL_KEYS) if missing or unknown: diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index ce28a217f3..46d4ac6f65 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -91,6 +91,8 @@ "pp": ("aic_pp_size",), "moe_tp_size": ("moe_tp_size", "aic_moe_tp_size"), "moe_ep_size": ("moe_ep_size", "aic_moe_ep_size"), + "cp_size": ("cp_size", "aic_cp_size"), + "dcp": ("dcp", "dcp_size", "aic_dcp_size"), "gemm_dtype": ("gemm_dtype", "aic_gemm_dtype"), "moe_dtype": ("moe_dtype", "aic_moe_dtype"), "fmha_dtype": ("fmha_dtype", "aic_fmha_dtype"), @@ -1191,7 +1193,7 @@ def _pop_aic_timing_overrides(rank: dict[str, JSONValue], role: str) -> dict[str value = rank.pop(configured[0]) if target == "decode_workload_distribution" and value is None: continue - if target in {"pp", "moe_tp_size", "moe_ep_size", "wideep_num_slots"}: + if target in {"pp", "moe_tp_size", "moe_ep_size", "wideep_num_slots", "cp_size", "dcp"}: value = _positive_int(value, f"engine provider {role} {target}") elif target == "fpm_profile": from aisimulate_core.sdk.fpm_profile import load_fpm_profile @@ -1376,6 +1378,11 @@ def _materialize_engine_role( memory_diagnostics[role] = role_memory capacity_materialized = False num_gpu_blocks_is_explicit = False + # A nested canonical timing config carries the CP knobs too; check them + # against parallel_config BEFORE capacity materialization resolves them + # into AIC inputs, so a mismatch cannot size the KV cache for one topology + # while the deployment reports another. + _require_nested_context_parallel_match(role_config, parallel_config, role) if "rank" not in role_config: state_cache = _manual_state_cache(role_config, deployment_backend, role) num_gpu_blocks_is_explicit = role_config.get("num_gpu_blocks") is not None @@ -1563,6 +1570,16 @@ def _materialize_engine_role( memory_fraction_overrides[memory_field] = float(value) aic_timing_overrides = _pop_aic_timing_overrides(rank, role) + for target in ("cp_size", "dcp"): + if target in aic_timing_overrides: + # Same contract as tp / attention_dp: a directly supplied deployment + # must not price one CP topology while parallel_config reports another. + _require_parallel_match( + parallel_config, + f"{parallel_prefix}{'cp' if target == 'cp_size' else 'dcp'}", + aic_timing_overrides[target], + f"engine provider {role} {target}", + ) timing_model = rank.get("timing_model") uses_aic_timing = timing_model is None or ( @@ -1890,6 +1907,37 @@ def _sample_synthetic_lengths( return [rng.randint(lower, upper) for _ in range(count)] +def _require_nested_context_parallel_match( + role_config: Mapping[str, JSONValue], + parallel_config: Mapping[str, JSONValue], + role: str, +) -> None: + """Reject ``timing_model.config.cp_size/dcp`` that disagree with ``parallel_config``. + + The timing model may sit at the top level (flat CLI/Sweeper form) or under + the nested ``rank`` descriptor (execution-level input); both are checked. + """ + rank_config = role_config.get("rank") + timing_source: Mapping[str, JSONValue] = rank_config if isinstance(rank_config, Mapping) else role_config + timing = timing_source.get("timing_model") + if not isinstance(timing, dict) or timing.get("type") != "external" or timing.get("provider") != "aic": + return + nested = timing.get("config") + if not isinstance(nested, dict): + return + prefix = "" if role == "aggregated" else f"{role}_" + for target, parallel_field in (("cp_size", "cp"), ("dcp", "dcp")): + value = nested.get(target) + if value is None: + continue + _require_parallel_match( + parallel_config, + f"{prefix}{parallel_field}", + _positive_int(value, f"engine provider {role} timing_model.config.{target}"), + f"engine provider {role} timing_model.config.{target}", + ) + + def _require_parallel_match( parallel_config: Mapping[str, JSONValue], field: str, diff --git a/python/aisimulate/src/aisimulate/sdk/task_v2.py b/python/aisimulate/src/aisimulate/sdk/task_v2.py index a24afc41a3..0607b584cc 100644 --- a/python/aisimulate/src/aisimulate/sdk/task_v2.py +++ b/python/aisimulate/src/aisimulate/sdk/task_v2.py @@ -123,7 +123,9 @@ def _default_cp_list_for(model_family: str, backend_name: str) -> list[int]: Capability-derived: any model whose class declares ``supports_cp`` on this backend is auto-swept over cp ∈ {1,2,4,8}. Keying off the registry (not a hardcoded family list) means the sweep policy never drifts from - ``BaseModel.supports_cp``. Decode is always forced to cp=1 by iter_parallel. + ``BaseModel.supports_cp``. Decode defaults to [1] in ``_fill_role_search`` + (prefill CP never runs on a decode worker); decode CP is the separate + per-role ``*_dcp_size`` scalar. """ from aisimulate.sdk.models.base import _MODEL_REGISTRY @@ -586,6 +588,13 @@ class Task: fmha_quant_mode: common.FMHAQuantMode | None = None comm_quant_mode: common.CommQuantMode | None = None + # Decode context parallelism for the agg worker (vLLM ``-dcp`` / SGLang + # ``--dcp-size``). A per-role scalar, not a swept candidate list: the + # (tp, pp, dp, moe_tp, moe_ep, cp) tuple is unpacked positionally by the + # sweep/pareto code, so widening it is a separate change. Combining it with + # a prefill-CP candidate above 1 is rejected in ``iter_parallel("agg")``. + dcp_size: int | None = None + # ====== 3. Agg search space ====== agg_num_gpu_candidates: list[int] | None = None agg_tp_candidates: list[int] | None = None @@ -603,6 +612,8 @@ class Task: prefill_enable_wideep: bool = False prefill_enable_chunked_prefill: bool = False prefill_enable_eplb: bool = False + # Decode CP on a prefill worker is allowed but pointless; carried for symmetry. + prefill_dcp_size: int | None = None prefill_gemm_quant_mode: common.GEMMQuantMode | None = None prefill_moe_quant_mode: common.MoEQuantMode | None = None prefill_kvcache_quant_mode: common.KVCacheQuantMode | None = None @@ -623,6 +634,8 @@ class Task: decode_system_name: str = "" decode_backend_name: str = "trtllm" decode_backend_version: str | None = None + # Decode context parallelism for the decode worker (its natural home). + decode_dcp_size: int | None = None decode_enable_wideep: bool = False decode_enable_eplb: bool = False decode_gemm_quant_mode: common.GEMMQuantMode | None = None @@ -2022,10 +2035,11 @@ def _fill_role_search(self, role: str, src: dict[str, list[int]]) -> None: for k_src, k_attr in map_to_attr.items(): if getattr(self, k_attr) is None: if k_src == "cp_list": - # Decode is always cp=1 (CP is prefill-only). prefill/agg - # auto-sweep cp for CP-validated families (else [1]); an - # explicit worker-config cp_list still wins. A user-supplied - # non-1 decode cp is rejected in iter_parallel. + # Prefill CP never runs on a decode worker, so decode + # DEFAULTS to [1] (an explicit list still wins and is + # priced as given). prefill/agg auto-sweep cp for + # CP-validated families (else [1]); an explicit + # worker-config cp_list still wins. if role == "decode": value = [1] else: @@ -2097,12 +2111,15 @@ def build_model_config( """ num_gpus_per_node = self._num_gpus_per_node(role) model_config = config.ModelConfig( - tp_size=parallel[0] if parallel is not None else 1, + # The template config's TP is a placeholder; ModelConfig requires it to + # be divisible by the role's dcp_size, so the placeholder is the dcp. + tp_size=parallel[0] if parallel is not None else (self._role_attr(role, "dcp_size") or 1), pp_size=parallel[1] if parallel is not None else 1, attention_dp_size=parallel[2] if parallel is not None else 1, moe_tp_size=parallel[3] if parallel is not None else 1, moe_ep_size=parallel[4] if parallel is not None else 1, cp_size=parallel[5] if parallel is not None else 1, + dcp_size=self._role_attr(role, "dcp_size"), gemm_quant_mode=self._role_attr(role, "gemm_quant_mode"), moe_quant_mode=self._role_attr(role, "moe_quant_mode"), kvcache_quant_mode=self._role_attr(role, "kvcache_quant_mode"), @@ -2216,14 +2233,42 @@ def iter_parallel(self, role: Literal["agg", "prefill", "decode"]) -> Iterator[P def _cands(dim: str) -> list[int]: return getattr(self, f"{prefix}{dim}_candidates") - # CP is modeled for context/prefill only; decode must be cp=1. Fail loud - # rather than silently coercing a user-supplied decode cp>1. cp_list = _cands("cp") or [1] - if role == "decode" and any(c != 1 for c in cp_list): + # Topology-layer rules (same as the compiler's): one aggregated engine + # prices prefill CP and decode CP one at a time; a disaggregated prefill + # engine stripes its KV only as the PCP+DCP layout (dcp == cp); vLLM + # refuses a replicated-PCP prefill next to a DCP-sharded decode and + # needs the two DCP sizes to divide one another. Fail loud instead of + # silently dropping the user's prefill-CP candidates. + # None means "not requested" (priced as 1); it stays None on the ModelConfig + # so the FPM cell identity can tell unrecorded DCP from recorded DCP=1. + dcp_size = self._role_attr(role, "dcp_size") or 1 + if role == "agg" and dcp_size > 1 and any(c > 1 for c in cp_list): raise ValueError( - f"decode CP must be 1 (CP is modeled for prefill only); got " - f"decode_cp_candidates={cp_list}. Enable CP via prefill/agg instead." + f"aggregated workers support at most one of prefill CP and decode CP above 1; got " + f"dcp_size={dcp_size} with agg_cp_candidates={cp_list}. Pin agg_cp_candidates=[1] " + "(or set dcp_size=1), or use serving_mode='disagg' to put prefill CP on the prefill " + "worker and decode CP on the decode worker." ) + if role == "prefill": + if dcp_size > 1 and cp_list != [dcp_size]: + raise ValueError( + f"prefill_dcp_size={dcp_size} is only accepted as the PCP+DCP layout, i.e. with " + f"prefill_cp_candidates=[{dcp_size}] (got {cp_list}); a prefill engine gains nothing " + "from striping its KV on its own" + ) + decode_dcp = self.decode_dcp_size or 1 + if self.prefill_backend_name == "vllm" and decode_dcp > 1: + if dcp_size == 1 and any(c > 1 for c in cp_list): + raise ValueError( + f"vLLM cannot pair a replicated-PCP prefill worker (prefill_cp_candidates={cp_list}, " + f"prefill_dcp_size=1) with a DCP-sharded decode worker (decode_dcp_size={decode_dcp})" + ) + if decode_dcp % dcp_size and dcp_size % decode_dcp: + raise ValueError( + f"vLLM requires prefill_dcp_size={dcp_size} and decode_dcp_size={decode_dcp} to divide " + "one another" + ) backend = common.BackendName[self._role_attr(role, "backend_name")] parallel = enumerate_parallel_config( diff --git a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py index d3ca843ecf..2fb35476fc 100644 --- a/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py +++ b/python/aisimulate/src/aisimulate/sweeper/forward_pass_estimator.py @@ -59,6 +59,10 @@ def _request(self, sample: dict[str, Any], role: str) -> ForwardPassPerfModelCon prefix = _role_prefix(role) moe_tp = int(sample[f"{prefix}moe_tp"]) moe_ep = int(sample[f"{prefix}moe_ep"]) + # Context-parallel knobs are optional columns (the compiler spells them + # out only when != 1); keep 1 out of the identity so pre-CP candidates + # resolve to the same estimator as before. + cp_size = int(sample.get(f"{prefix}cp") or 1) block_size = sample[f"{role}_block_size"] if block_size is None: block_size = _DEFAULT_BLOCK_SIZE[backend] @@ -89,6 +93,7 @@ def _request(self, sample: dict[str, Any], role: str) -> ForwardPassPerfModelCon dcp=sample.get(f"{prefix}dcp"), moe_tp_size=moe_tp if moe_tp * moe_ep > 1 else None, moe_ep_size=moe_ep if moe_tp * moe_ep > 1 else None, + cp_size=cp_size if cp_size != 1 else None, nextn=int(nextn or 0), **{name: getattr(self._search_space, name) for name in ENGINE_MODEL_CONTROL_FIELDS}, speculation=self._search_space.speculation.cost_config() diff --git a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py index 6eff2a9b0f..86e38ea96b 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py +++ b/python/aisimulate/src/aisimulate_core/sdk/backends/base_backend.py @@ -1221,8 +1221,8 @@ def _run_encoder(batch_size: int) -> tuple[dict[str, float], dict[str, float], d # KV-per-seq context for capacity probing in CLI detail reports. try: kv_seq_len_used = isl + img_ctx_tokens + beam_width * osl - # CP shards persistent KV across cp ranks (full/cp per rank). - kv_bytes_per_seq = model.get_kvcache_bytes_per_sequence(kv_seq_len_used) / model._cp_kv_memory_divisor() + # Per-rank persistent KV (decode CP stripes the token-linear KV). + kv_bytes_per_seq = model.get_kvcache_rank_bytes_per_sequence(kv_seq_len_used) summary.set_kv_per_seq(kv_bytes_per_seq, kv_seq_len_used) except Exception: # Best-effort; downstream report degrades gracefully when unset. @@ -2292,9 +2292,9 @@ def _get_memory_usage( activations *= 1.0 + self.ACTIVATION_OVERHEAD_FRAC seq_tokens = max_seq_len if max_seq_len is not None else isl + beam_width * osl - # CP shards persistent KV across cp ranks (full/cp per rank); the - # all-gather is a transient compute buffer, not steady-state footprint. - kvcache = batch_size * model.get_kvcache_bytes_per_sequence(seq_tokens) / model._cp_kv_memory_divisor() + # Per-rank persistent KV: decode CP stripes the token-linear KV (prefill + # CP's all-gather is a transient compute buffer, not steady-state footprint). + kvcache = batch_size * model.get_kvcache_rank_bytes_per_sequence(seq_tokens) if has_draft_scheme: kvcache += batch_size * scheme.draft_kv_bytes_per_sequence(model, seq_tokens) # should not be divided by pp_size as you need to hold all kvcache for stages. diff --git a/python/aisimulate/src/aisimulate_core/sdk/config.py b/python/aisimulate/src/aisimulate_core/sdk/config.py index 56f48bbea1..54125c3dd2 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config.py @@ -29,6 +29,13 @@ def normalize_kernel_backend( raise ValueError(f"{field_name} must be one of {choices}, got {value!r}.") from exc +def validate_parallel_size(name: str, value: object) -> int: + """Return ``value`` when it is a positive ``int`` (not ``bool``); raise ``ValueError`` otherwise.""" + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + return value + + def normalize_kernel_source(value: str | None, field_name: str) -> str | None: """Validate an optional exact collected kernel-source label.""" if value is None: @@ -93,6 +100,25 @@ class ModelConfig: # from backend_name when cp_size > 1; default "none". Dense models branch on # this in their op pipeline; GLM-5 DSA ignores it (handled in ContextDSAModule). cp_style: str = "none" + # Decode context parallelism (vLLM ``-dcp`` / SGLang ``--dcp-size``) is the + # keyword-only ``dcp_size`` above: it stripes the DECODE KV cache by token + # position across ranks that already belong to the attention group, so it + # does NOT fold into attn_width or total_gpus_per_worker. ``None`` means not + # requested (priced as 1) and keeps the FPM cell identity on the unrecorded + # profiles; an explicit 1 selects recorded-DCP1 cells. ``cp_size`` (prefill + # CP) and ``dcp_size`` are orthogonal per-phase knobs; whether one deployment + # may set both is a topology-layer decision (agg vs disagg), not a ModelConfig one. + # DCP partial-output merge collective: "ag_rs" (query all-gather + LSE + # all-gather + output reduce-scatter; vLLM default) or "a2a" (query + # all-gather + one packed all-to-all; SGLang default on CUDA). None picks + # the backend default in BaseModel._dcp_comm_style. + dcp_comm: str | None = None + # Replicated query projection under DCP (vLLM ``dcp_q_replicate`` / SGLang + # ``--dcp-replicate-q-proj``): every rank runs the Q up-projection for the + # whole DCP group's heads and skips the per-layer query all-gather. Only + # meaningful with the a2a-style merges. None picks the backend/model default + # in BaseModel._dcp_q_replicate. + dcp_q_replicate: bool | None = None workload_distribution: str = "power_law" # Explicit decode-only profile. Context retains workload_distribution. decode_workload_distribution: str | None = field(default=None, kw_only=True) @@ -159,6 +185,7 @@ class ModelConfig: fpm_fmha_quant_mode: common.FMHAQuantMode | None = field(default=None, kw_only=True) def __post_init__(self) -> None: + validate_parallel_size("cp_size", self.cp_size) if self.dcp_size is not None and ( type(self.dcp_size) is not int or self.dcp_size <= 0 or self.tp_size % self.dcp_size ): diff --git a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py index 6e4452e91a..b065fcee0b 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/config_builders.py +++ b/python/aisimulate/src/aisimulate_core/sdk/config_builders.py @@ -46,8 +46,9 @@ def build_model_config( wideep_num_slots: int | None = None, *, moe_kernel_source: str | None = None, - fpm_fmha_quant_mode: str | None = None, + cp_size: int = 1, dcp_size: int | None = None, + fpm_fmha_quant_mode: str | None = None, decode_workload_distribution: str | None = None, prefill_graph_profile: str | None = None, ) -> ModelConfig: @@ -60,11 +61,12 @@ def build_model_config( ) return ModelConfig( tp_size=tp_size, - dcp_size=dcp_size, pp_size=pp_size, attention_dp_size=attention_dp_size, moe_tp_size=moe_tp_size, moe_ep_size=moe_ep_size, + cp_size=cp_size, + dcp_size=dcp_size, gemm_quant_mode=GEMMQuantMode[gemm_quant_mode] if gemm_quant_mode else None, kvcache_quant_mode=KVCacheQuantMode[kvcache_quant_mode] if kvcache_quant_mode else None, fmha_quant_mode=FMHAQuantMode[fmha_quant_mode] if fmha_quant_mode else None, diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index 630b9da311..9f001a666f 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -48,6 +48,7 @@ from typing import Any import aisimulate_core +from aisimulate_core.sdk.config import validate_parallel_size 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 @@ -153,6 +154,8 @@ # - 23 (DCP identity): optional recorded dcp_size in ParallelMapping. # - 24 (typed FPM DCP): dcp_size moves out of the op matching tuple. # - 25 (FPM decoupling): merge the SOL/direct selector into the typed DCP layout. +# - 26 (decode context parallelism): Generation/Context attention, MLA, DSA +# and MLAModule ops gained a tail-appended `dcp_size` (positional bincode). # Single owner: the Rust crate constant. Python re-exports it for # diagnostics/tests instead of declaring a twin to keep in sync. ENGINE_SPEC_SCHEMA_VERSION = aisimulate_core.engine_spec_schema_version() @@ -361,6 +364,8 @@ def _engine_config_dict( "moe_tp_size": _opt_int(getattr(cfg, "moe_tp_size", None)), "moe_ep_size": _opt_int(getattr(cfg, "moe_ep_size", None)), "cp_size": _opt_int(getattr(cfg, "cp_size", None)), + # Decode CP joins the identity too: a dcp variant prices decode attention + # and KV capacity differently, so it must never share a compiled handle. "dcp_size": _opt_int(getattr(cfg, "dcp_size", None)), # QuantizationConfig (flattened) "weight_dtype": _rust_quant_to_dtype(getattr(cfg, "gemm_quant_mode", None)), @@ -515,8 +520,9 @@ def compile_engine( logical batches are exploratory and queries outside 1..32 fail. Missing or mismatched approved profile data raises ``DecodeMoeProfileError``. - Engine schema 25 combines the FPM interpolation selector, typed DCP, and - exact-profile policy on MoE operators. Older binary specs require + Engine schema 26 adds decode context parallelism on top of the schema-25 + FPM interpolation selector, typed DCP and exact-profile policy on MoE + operators. Older binary specs require recompilation; default representations are therefore not byte-identical. """ if not isinstance(decoder_replay, bool): @@ -525,8 +531,8 @@ def compile_engine( raise InvalidEngineConfigurationError( f"decoder_replay requires model={DEEPSEEK_V41_MODEL_PATH!r} and backend='sglang'" ) - if type(cp_size) is not int or cp_size != 1: - raise ValueError("cp_size must be the integer 1; this SDK entry point does not support context parallelism") + # Before any model construction: the profile and op-level paths both read it. + validate_parallel_size("cp_size", cp_size) if fpm_parquet_path is not None: if not fpm_parquet_path: raise ValueError("fpm_parquet_path cannot be empty") @@ -676,6 +682,7 @@ def compile_engine( moe_backend=moe_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, + cp_size=cp_size, speculation=resolved_speculation, ) model_config.fpm_config = fpm_config @@ -1006,6 +1013,7 @@ def build_database_probe_spec_json( "moe_tp_size": None, "moe_ep_size": None, "cp_size": None, + "dcp_size": None, "weight_dtype": None, "moe_dtype": None, "activation_dtype": None, diff --git a/python/aisimulate/src/aisimulate_core/sdk/memory.py b/python/aisimulate/src/aisimulate_core/sdk/memory.py index 15116bd2ee..daaa063fa1 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/memory.py +++ b/python/aisimulate/src/aisimulate_core/sdk/memory.py @@ -51,6 +51,7 @@ from aisimulate_core.sdk import perf_database from aisimulate_core.sdk.backends.factory import get_backend from aisimulate_core.sdk.common import DefaultHFModels +from aisimulate_core.sdk.config import validate_parallel_size from aisimulate_core.sdk.config_builders import apply_nextn, build_model_config, validate_moe_controls, validate_nextn from aisimulate_core.sdk.fpm_profile import FpmModelProfile, load_fpm_profile from aisimulate_core.sdk.models import get_model @@ -285,6 +286,8 @@ def from_request( attention_dp_size: int = 1, moe_tp_size: int | None = None, moe_ep_size: int | None = None, + cp_size: int = 1, + dcp_size: int = 1, gemm_quant_mode: str | None = None, moe_quant_mode: str | None = None, kvcache_quant_mode: str | None = None, @@ -344,6 +347,8 @@ def from_request( attention_backend=attention_backend, enable_eplb=enable_eplb, wideep_num_slots=wideep_num_slots, + cp_size=cp_size, + dcp_size=dcp_size, ) # Apply nextn/MTP onto the config BEFORE get_model so the built model is # spec-decode aware (e.g. for any draft-module weights). This does NOT scale @@ -406,6 +411,22 @@ def from_request( if backend != "sglang": non_kv_bytes += activations_bytes + # Per-RANK persistent KV. Prefill CP keeps the full KV on every rank; + # decode CP stripes the token-linear KV so each rank holds 1/dcp of every + # token, while rank-local per-request state (Kimi-K3 KDA) stays whole + # (BaseModel.get_kvcache_rank_*). The per-token figure and the + # byte-budget -> token-count inverse must see the same striping. + # Duck-typed model doubles without the rank hooks keep the full KV. + rank_bytes_per_sequence = getattr( + model, "get_kvcache_rank_bytes_per_sequence", model.get_kvcache_bytes_per_sequence + ) + + def tokens_from_kv_bytes(kv_budget_bytes: float) -> int: + # Resolved lazily: some callers never invert (stubs without the + # capacity method). + capacity = getattr(model, "get_kvcache_rank_batch_capacity", None) or model.get_kvcache_batch_capacity + return int(capacity(float(kv_budget_bytes), max_batch_size)) + return cls( { "weights_bytes": weights_bytes, @@ -418,10 +439,9 @@ def from_request( "pre_model_load_overhead_bytes": ( runtime_overhead_bytes + comm_overhead_bytes if backend == "sglang" else 0.0 ), - "kv_size_per_token_bytes": float(model.get_kvcache_bytes_per_sequence(1)), + "kv_size_per_token_bytes": float(rank_bytes_per_sequence(1)), "gpu_memory_capacity_bytes": float(database.system_spec["gpu"]["mem_capacity"]), - # Model's byte-budget -> token-count inverse (KV-curve aware). - "tokens_from_kv_bytes": lambda budget: model.get_kvcache_batch_capacity(budget, max_batch_size), + "tokens_from_kv_bytes": tokens_from_kv_bytes, } ) @@ -1029,6 +1049,8 @@ def estimate_kv_cache( attention_dp_size: int = 1, moe_tp_size: int | None = None, moe_ep_size: int | None = None, + cp_size: int = 1, + dcp_size: int = 1, gemm_quant_mode: str | None = None, moe_quant_mode: str | None = None, kvcache_quant_mode: str | None = None, @@ -1049,7 +1071,6 @@ def estimate_kv_cache( allow_hf_config_download: bool = False, fpm_profile: dict | str | FpmModelProfile | None = None, worker_type: str = "aggregated", - cp_size: int = 1, context_length: int | None = None, request_occupancy_tokens: int | None = None, ) -> dict[str, Any]: @@ -1098,7 +1119,8 @@ def estimate_kv_cache( route never constructs an analytical model or uses naive fallback. Separate CUDA graph bytes apply only to declared profile overheads. worker_type: canonical serving role selecting this deployment's resources. - cp_size: profile context-parallel identity; currently only CP1 is supported. + cp_size: prefill context parallelism (folds into the attention width); also + the FPM profile cell identity, whose recorded cells are CP1 today. context_length: optional context bound for a grouped profile's per-request cache peak, also checked against runtime memory's ``max_model_len``. Defaults to the profile context and does not change its byte budget. @@ -1118,12 +1140,13 @@ def estimate_kv_cache( tolerance, no KV budget, insufficient model metadata, or (with the fallback off) an unsupported model/backend. """ - if type(cp_size) is not int or cp_size != 1: - raise ValueError("cp_size must be the integer 1; this SDK entry point does not support context parallelism") _validate_memory_fraction(backend, memory_fraction_kind, memory_fraction_value) _validate_tolerance(tolerance_fraction) _validate_naive_reservation(naive_kv_reservation) _validate_cuda_graph_reservation(cuda_graph_reserved_bytes) + # Before the model build and the naive fallback, which would swallow the error. + validate_parallel_size("cp_size", cp_size) + validate_parallel_size("dcp_size", dcp_size) fraction = float(memory_fraction_value) is_of_free = memory_fraction_kind == "of_free" @@ -1246,6 +1269,8 @@ def estimate_kv_cache( attention_dp_size=int(attention_dp_size), moe_tp_size=moe_tp_size, moe_ep_size=moe_ep_size, + cp_size=int(cp_size), + dcp_size=int(dcp_size), gemm_quant_mode=gemm_quant_mode, moe_quant_mode=moe_quant_mode, kvcache_quant_mode=kvcache_quant_mode, @@ -1260,6 +1285,14 @@ def estimate_kv_cache( systems_path=systems_path, ) except Exception as exc: # native model build unsupported (model/backend/perf DB) + if isinstance(exc, NotImplementedError) and (int(cp_size) > 1 or int(dcp_size) > 1): + # A context-parallel capability rejection must not degrade into the + # naive estimator, which knows neither knob and would return an + # unstriped (cp=dcp=1) capacity for a request that asked otherwise. + raise ValueError( + f"context parallelism is not supported for KV-cache estimation of this model/backend " + f"(model={model_path}, backend={backend}, cp_size={cp_size}, dcp_size={dcp_size}): {exc}" + ) from exc if ( not allow_naive_fallback or enable_eplb @@ -1319,6 +1352,8 @@ def estimate_num_gpu_blocks( attention_dp_size: int = 1, moe_tp_size: int | None = None, moe_ep_size: int | None = None, + cp_size: int = 1, + dcp_size: int = 1, gemm_quant_mode: str | None = None, moe_quant_mode: str | None = None, kvcache_quant_mode: str | None = None, @@ -1340,7 +1375,6 @@ def estimate_num_gpu_blocks( diagnostics: dict[str, Any] | None = None, fpm_profile: dict | str | FpmModelProfile | None = None, worker_type: str = "aggregated", - cp_size: int = 1, context_length: int | None = None, ) -> int: """Convert the KV-cache token capacity to a scheduler block count. @@ -1386,6 +1420,10 @@ def estimate_num_gpu_blocks( attention_dp_size=attention_dp_size if fpm_profile is not None else int(attention_dp_size), moe_tp_size=moe_tp_size, moe_ep_size=moe_ep_size, + # Raw values: estimate_kv_cache validates them (positive, integral, not + # bool) before any coercion, so True / 1.9 cannot become cp=dcp=1 here. + cp_size=cp_size, + dcp_size=dcp_size, gemm_quant_mode=gemm_quant_mode, moe_quant_mode=moe_quant_mode, kvcache_quant_mode=kvcache_quant_mode, @@ -1406,7 +1444,6 @@ def estimate_num_gpu_blocks( allow_hf_config_download=allow_hf_config_download, fpm_profile=fpm_profile, worker_type=worker_type, - cp_size=cp_size, context_length=context_length, ) diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py index 4ab809be23..273cb98513 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/__init__.py @@ -173,7 +173,10 @@ def get_model( rewrites each phase list to a single whole-model ``FPMForwardOp``. """ forward_model = getattr(model_config, "forward_model", "op_level") or "op_level" - if (model_config.dcp_size or 1) > 1 and (forward_model != "fpm" or backend_name != "vllm"): + # Whole-forward DCP timing needs a measured vLLM cell (the recorded-DCP FPM + # identity); the op-level path prices the sharded attention itself and is + # gated per model class by ``supports_dcp`` below. + if (model_config.dcp_size or 1) > 1 and forward_model == "fpm" and backend_name != "vllm": raise NotImplementedError("DCP timing requires measured vLLM FPM interpolation") if forward_model not in _FORWARD_MODELS: raise InvalidEngineConfigurationError( @@ -248,6 +251,18 @@ def get_model( else: model_config.cp_style = "none" + # Decode context parallelism is a separate modeling capability: the model + # class must price the KV-sharded decode attention (+ LSE merge comm) before + # dcp>1 can be estimated. Deployment policy (whether a role may combine + # prefill CP with DCP) is decided by the topology layer, not here; this only + # guards against silently wrong numbers. + if (model_config.dcp_size or 1) > 1 and forward_model != "fpm" and not cls.supports_dcp(backend_name): + raise NotImplementedError( + f"Decode context parallelism (dcp_size={model_config.dcp_size}) is not supported for " + f"model_family={model_family!r} on backend={backend_name!r}. The model class " + f"must override ``supports_dcp`` and implement the KV-sharded decode path." + ) + # Resolve the speculative scheme BEFORE construction (an explicit mtp # scheme writes its depth back onto nextn, which model families read), # attach it after, and gate unsupported (model, backend) combinations. @@ -264,9 +279,20 @@ def get_model( _validate_decode_moe_profile(model_path, model_config, backend_name) _validate_prefill_graph_profile(model_path, model_config, backend_name) model = cls.create(model_info, model_config, backend_name) + # Backend-specific defaults below construction (e.g. the DCP merge + # collective: a2a on sglang, ag_rs elsewhere) read the backend identity off + # the model; families whose constructors do not record it get it here. + if getattr(model, "_backend_name", None) is None: + model._backend_name = backend_name model.spec_scheme = build_spec_scheme(model_config, spec_config) model.spec_scheme.validate(model, backend_name) materialize_spec_scheme(model) + # Op-level decode CP: rewrite the decode attention (sharded KV stripe + + # merge collectives) after speculation materialized the draft ops. The + # whole-forward path prices DCP from the recorded-DCP measured cell instead + # (the per-rank KV stripe still enters memory via _cp_kv_memory_divisor). + if (model_config.dcp_size or 1) > 1 and forward_model != "fpm": + model._apply_decode_context_parallel() if model_config.prefill_graph_profile is not None: model.apply_prefill_graph_profile() if model_config.moe_kernel_source is not None: diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/base.py b/python/aisimulate/src/aisimulate_core/sdk/models/base.py index b0b21d6328..71e87d0cc3 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/base.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/base.py @@ -197,6 +197,20 @@ def supports_cp(cls, backend_name: str) -> bool: """ return False + @classmethod + def supports_dcp(cls, backend_name: str) -> bool: + """Whether this (model, backend) combo models decode context parallelism. + + Default False. A DCP-capable model class overrides this once its + generation pipeline prices the KV-sharded decode attention (per-rank + KV / dcp, query gather, LSE merge collective). ``get_model`` checks this + BEFORE construction so dcp_size>1 fails loud instead of silently + pricing decode as if the KV were not sharded. This is a modeling + capability check only; deployment policy (which roles may combine + prefill CP with DCP) lives in the topology layer. + """ + return False + @classmethod def _resolve_cp_style(cls, backend_name: str) -> str: """Pick the CP variant for this (model, backend). Called only when cp_size>1.""" @@ -231,8 +245,300 @@ def _cp_attn_comm_ops(self) -> list: ] return [] + # ------------------------------------------------------------------ + # Decode context parallelism (DCP): one rewrite for every model class. + # The decode attention op prices the DCP group's gathered query heads + # over this rank's 1/dcp KV stripe (Rust `dcp_size` on the op); the + # collectives that make that possible are modeled here as NCCL ops: + # * query all-gather over the DCP group (skipped by the frameworks' + # replicate-q-proj option, which is NOT modeled: gathered Q is priced), + # * the partial-output merge -- "ag_rs" (LSE all-gather + fp32 output + # reduce-scatter, vLLM default) or "a2a" (one packed all-to-all, + # SGLang default on CUDA; the tiny LSE exchange rides inside it). + # ------------------------------------------------------------------ + + def _dcp_comm_style(self) -> str: + """``"ag_rs"`` or ``"a2a"``. SGLang's ``fi_a2a`` (FlashInfer MNNVL all-to-all, + its default where the DCP group sits in one MNNVL domain) moves the same + packed buffer as ``a2a`` and is priced as such.""" + style = getattr(self.config, "dcp_comm", None) + if style is None: + style = "a2a" if getattr(self, "_backend_name", None) == "sglang" else "ag_rs" + if style == "fi_a2a": + style = "a2a" + if style not in ("ag_rs", "a2a"): + raise ValueError(f"dcp_comm must be 'ag_rs', 'a2a' or 'fi_a2a', got {style!r}") + return style + + def _dcp_q_replicate(self) -> bool: + """Whether every rank runs the Q up-projection for the whole DCP group. + + vLLM ``dcp_q_replicate`` / SGLang ``--dcp-replicate-q-proj``: the + per-layer query all-gather disappears and the Q-side ops (``q_b`` + projection, absorb BMM) run on ``n_local * dcp`` heads instead. Only + defined for the a2a-style merges. ``None`` means the backend/model + default, which is off unless a model class says otherwise. + """ + explicit = getattr(self.config, "dcp_q_replicate", None) + if explicit is None: + return False + if explicit and self._dcp_comm_style() != "a2a": + raise ValueError("dcp_q_replicate requires the a2a / fi_a2a DCP merge (it removes the query all-gather)") + return bool(explicit) + + def _mla_latent_dims(self) -> tuple[int, int]: + """(q dim gathered per head, output dim merged per head) for absorbed MLA decode.""" + kv_lora_rank, qk_rope_head_dim = 0, 0 + if isinstance(self.extra_params, dict): + kv_lora_rank = int(self.extra_params.get("kv_lora_rank") or 0) + qk_rope_head_dim = int(self.extra_params.get("qk_rope_head_dim") or 0) + kv_lora_rank = kv_lora_rank or 512 + qk_rope_head_dim = qk_rope_head_dim or 64 + return kv_lora_rank + qk_rope_head_dim, kv_lora_rank + + def _decode_attention_dcp_dims(self, op) -> tuple[int, int, int] | None: + """``(rank-local query heads, q dim per head, output dim per head)`` for a + decode attention op, or ``None`` when ``op`` is not one. + + Matches on the engine's core classes: ops nested inside a FallbackOp + come back as bare core instances, not the Python shell subclasses. + """ + import aisimulate_core._native as _core + + if isinstance(op, _core.GenerationAttention): + return int(op._n), int(op._head_size), int(op._head_size) + if isinstance(op, _core.GenerationMLA): + q_dim, v_dim = self._mla_latent_dims() + return int(op._num_heads), q_dim, v_dim + if isinstance(op, _core.WideEPGenerationMLA): + q_dim, v_dim = self._mla_latent_dims() + return 128 // int(op._tp_size), q_dim, v_dim + if isinstance(op, _core.GenerationDSAModule): + q_dim, v_dim = self._mla_latent_dims() + return int(op._num_heads), q_dim, v_dim + if isinstance(op, _core.MLAModule) and not op._is_context: + q_dim, v_dim = self._mla_latent_dims() + return int(op._num_heads), q_dim, v_dim + return None + + @staticmethod + def _through_fallback(op, leaf): + """Apply ``leaf(op) -> (op, hit)`` to an op or to every interior of a ``FallbackOp``. + + ``FallbackOp`` exposes clones of its inner ops, so a block whose interior + was touched is rebuilt around the new primary/fallback list. Returns the + (possibly rebuilt) op and the first truthy ``hit``. + """ + import aisimulate_core._native as _core + import aisimulate_core.sdk.operations as ops + + if isinstance(op, _core.OverlapOp): + groups, hit = [], None + for group in (op._group_a, op._group_b): + rebuilt = [] + for inner in group: + inner, inner_hit = BaseModel._through_fallback(inner, leaf) + rebuilt.append(inner) + hit = hit or inner_hit + groups.append(rebuilt) + if not hit: + return op, hit + return ops.OverlapOp(op._name, groups[0], groups[1], seq_split=int(getattr(op, "_seq_split", 1) or 1)), hit + if not isinstance(op, _core.FallbackOp): + return leaf(op) + primary, hit = BaseModel._through_fallback(op._primary, leaf) + fallback = [] + for inner in op._fallback: + inner, inner_hit = BaseModel._through_fallback(inner, leaf) + fallback.append(inner) + hit = hit or inner_hit + if not hit: + return op, hit + return ops.FallbackOp(op._name, primary=primary, fallback=fallback), hit + + def _rewrite_op_for_dcp(self, op, dcp: int): + """Return ``(op with dcp applied, (dims, scale) | None)``; the collectives + are then priced once for the whole (possibly rebuilt) block.""" + + import aisimulate_core._native as _core + + def leaf(inner): + dims = self._decode_attention_dcp_dims(inner) + if dims is None: + return inner, None + inner._dcp_size = dcp + # DSA: the indexer's per-rank top-k candidates are merged across the + # DCP group before the sparse attention (priced in _dcp_attn_comm_ops). + topk = 0 + if isinstance(inner, _core.GenerationDSAModule): + from aisimulate_core.sdk.operations.dsa import DSA_MODEL_DIMS + + topk = int(DSA_MODEL_DIMS.get(str(inner._architecture), {}).get("index_topk", 0) or 0) + return inner, (dims, float(inner._scale_factor), topk) + + return self._through_fallback(op, leaf) + + def _replicate_q_side_for_dcp(self, op, dcp: int): + """``dcp_q_replicate``: the Q up-projection and the absorb BMM ahead of the + decode attention run for the whole group's ``n_local * dcp`` heads on every + rank (weights included), so their per-layer cost scales by ``dcp``.""" + + def leaf(inner): + name = str(inner._name) + if "q_b" in name or "bmm_pre" in name: + inner._scale_factor = float(inner._scale_factor) * dcp + return inner, True + return inner, False + + return self._through_fallback(op, leaf) + + def _dcp_attn_comm_ops( + self, name: str, scale: float, *, n_local: int, q_dim: int, v_dim: int, dsa_topk: int = 0 + ) -> list: + """The per-layer DCP collectives that accompany one decode attention op.""" + import aisimulate_core.sdk.operations as ops + from aisimulate_core.sdk import common + + dcp = int(self.config.dcp_size or 1) + comm_quant_mode = self.config.comm_quant_mode + gathered_heads = n_local * dcp + + def nccl(suffix: str, kind: str, elements_per_token: int, dtype=comm_quant_mode): + # The NCCL table is keyed by the collective's whole buffer (nccl-tests + # `size`: the all-gather receive buffer, the reduce-scatter input, the + # all-to-all per-rank buffer), matching `context_cp_all_gather`. + return ops.NCCL( + f"{name}_dcp_{suffix}", + scale, + kind, + num_elements_per_token=elements_per_token, + num_gpus=dcp, + comm_quant_mode=dtype, + ) + + collectives = [] + if not self._dcp_q_replicate(): + # Every rank contributes its n_local query heads; the gathered buffer + # holds all of them. With fp8 attention the MQA query is quantized + # BEFORE the gather (vLLM `_decode_concat_quant_fp8_op` -> + # `dcp_manager.query_gather`), so the buffer follows fmha_quant_mode. + fmha = self.config.fmha_quant_mode + q_dtype = common.CommQuantMode.fp8 if fmha is not None and fmha.value.memory == 1 else comm_quant_mode + collectives.append(nccl("q_all_gather", "all_gather", gathered_heads * q_dim, q_dtype)) + if dsa_topk > 0: + # DSA indexer under DCP (vLLM `_merge_dcp_topk_global`): each rank + # packs its (score, id) fp32 candidates, one all-gather over the + # group, then a stable top-k over the dcp * topk merged entries. + packed = dsa_topk * 4 # (fp32 score, int32 id) per candidate, in half-elements + collectives.append(ops.ElementWise(f"{name}_dcp_topk_pack", scale, dsa_topk * 4, packed)) + collectives.append(nccl("topk_all_gather", "all_gather", packed * dcp)) + collectives.append(ops.ElementWise(f"{name}_dcp_topk_merge", scale, packed * dcp, dsa_topk * 2)) + # Partial-output merge (vllm/v1/attention/ops/dcp.py). Both styles pay + # the collectives AND the elementwise passes around them; the latter + # are launch/latency-bound at decode batch sizes but add up over the + # layers, so they are priced explicitly. + if self._dcp_comm_style() == "ag_rs": + # `cp_lse_ag_out_rs`: all-gather the fp32 LSE (2 half-elements per + # gathered head from each of the dcp ranks), `correct_attn_out` + # rescales the partial outputs in place, then reduce-scatter the + # corrected outputs by head. + collectives.append(nccl("lse_all_gather", "all_gather", gathered_heads * 2 * dcp)) + collectives.append( + ops.ElementWise(f"{name}_dcp_lse_correct", scale, gathered_heads * v_dim, gathered_heads * v_dim) + ) + collectives.append(nccl("out_reduce_scatter", "reduce_scatter", gathered_heads * v_dim)) + else: + # `dcp_a2a_lse_reduce`: pack output + LSE (2 half slots per head) + # into one send buffer, one all-to-all, then unpack and combine + # the dcp partials of this rank's heads. + packed = gathered_heads * (v_dim + 2) + collectives.append(ops.ElementWise(f"{name}_dcp_a2a_pack", scale, gathered_heads * v_dim, packed)) + collectives.append(nccl("out_all_to_all", "alltoall", packed)) + collectives.append(ops.ElementWise(f"{name}_dcp_a2a_combine", scale, packed, n_local * v_dim)) + return collectives + + def _dcp_kv_head_replication(self) -> int | None: + """How many DCP ranks can share one KV copy, or ``None`` for "any that + divide TP" (MLA: one latent KV, replicated on every TP rank). GQA + classes return ``tp / kv_heads``: DCP only de-duplicates KV heads that + TP already replicates (vLLM: ``dcp <= tp_size // total_kv_heads``).""" + return None + + def _validate_dcp_topology(self) -> None: + # ModelConfig already requires dcp_size to divide tp_size. + dcp = int(self.config.dcp_size or 1) + limit = self._dcp_kv_head_replication() + if limit is not None and dcp > limit: + raise ValueError( + f"dcp_size={dcp} exceeds the KV-head replication of this layout ({limit}): " + "GQA decode CP only de-duplicates kv heads that TP replicates (dcp <= tp / kv_heads)." + ) + + def _stripe_context_for_dcp(self, op, dcp: int): + """Mark the prefill-side attention ops of a DCP-striped engine. + + With the persistent KV striped, a prefill that reuses cached context + (``prefix > 0``) first all-gathers the other ranks' stripes; the Rust + context ops price that from ``dcp_size``. New-token attention is + unchanged. FallbackOp interiors are clones, so the block is rebuilt. + """ + import aisimulate_core._native as _core + + def leaf(inner): + striped = isinstance(inner, (_core.ContextAttention, _core.ContextMLA, _core.ContextDSAModule)) or ( + isinstance(inner, _core.MLAModule) and inner._is_context + ) + if striped: + inner._dcp_size = dcp + return inner, striped + + return self._through_fallback(op, leaf) + + def _apply_decode_context_parallel(self) -> None: + """Rewrite ``generation_ops`` for ``config.dcp_size > 1``. + + Every top-level decode attention op gets ``_dcp_size`` and is followed + by its merge collectives. Draft ops are left alone (the frameworks + replicate the draft KV on every DCP rank). Called by ``get_model`` + after construction when ``dcp_size > 1``; a class that declares + ``supports_dcp`` but exposes no rewritable decode attention op fails + loud instead of silently pricing decode as unsharded. + """ + dcp = int(self.config.dcp_size or 1) + self._validate_dcp_topology() + q_replicate = self._dcp_q_replicate() + rewritten: list = [] + touched = 0 + for op in self.generation_ops: + if op._name.startswith("draft_"): + rewritten.append(op) + continue + op, found = self._rewrite_op_for_dcp(op, dcp) + if q_replicate: + op, _ = self._replicate_q_side_for_dcp(op, dcp) + rewritten.append(op) + if found is None: + continue + (n_local, q_dim, v_dim), scale, dsa_topk = found + rewritten.extend( + self._dcp_attn_comm_ops(op._name, scale, n_local=n_local, q_dim=q_dim, v_dim=v_dim, dsa_topk=dsa_topk) + ) + touched += 1 + if touched == 0: + raise NotImplementedError( + f"{type(self).__name__} declares supports_dcp but exposes no top-level decode attention op " + "to shard; decode context parallelism cannot be priced for this graph." + ) + self.generation_ops = rewritten + # Prefill side of the same engine: cached-context gather (aggregated + # serving). A decode-only worker never queries context_ops, so this is + # inert there. + self.context_ops = [ + op if op._name.startswith("draft_") else self._stripe_context_for_dcp(op, dcp)[0] for op in self.context_ops + ] + def _cp_kv_memory_divisor(self) -> int: - """Per-rank persistent-KV divisor under CP (always 1: full KV per rank). + """Per-rank persistent-KV divisor: 1 under prefill CP, ``dcp_size`` under DCP. Verified against sglang v0.5.13 that CP gives **no** per-rank KV-memory savings for any family -- each rank holds the FULL KV: @@ -246,10 +552,16 @@ def _cp_kv_memory_divisor(self) -> int: local pool (decode page_table / ``cache_seqlens`` span the full seq_len with no gather) -- so the full KV must reside per rank. - CP therefore saves prefill *compute*, not KV memory. Kept as a method so - a future seq-sliced variant (e.g. Ring) can override. + Prefill CP therefore saves prefill *compute*, not KV memory. + + Decode context parallelism is the knob that DOES shard the persistent + KV: vLLM ``-dcp`` / SGLang ``--dcp-size`` stripe the KV by token + position across the dcp ranks (rank ``r`` owns position ``p`` when + ``p mod dcp == r``), so each rank stores about ``1/dcp`` of every + sequence. Both frameworks widen the logical page by ``dcp`` so the + per-rank stripe stays balanced to within one token. """ - return 1 + return int(self.config.dcp_size or 1) def get_kvcache_elements_per_token(self) -> int: """KV cache size per token (per GPU) summed over all layers, in elements. @@ -306,6 +618,23 @@ def get_kvcache_batch_capacity(self, kv_budget_bytes: float, max_batch_size: int """Total-token capacity; models with per-request state may reserve it here.""" return self.get_kvcache_max_tokens(kv_budget_bytes) + def get_kvcache_rank_bytes_per_sequence(self, seq_len: int) -> float: + """Per-RANK persistent KV for one sequence: the full KV under prefill CP, + the ``1/dcp`` stripe under decode CP (see :meth:`_cp_kv_memory_divisor`). + + Hybrid models whose per-request state is rank-local and never striped + (Kimi-K3's KDA state) override this and + :meth:`get_kvcache_rank_batch_capacity` so the divisor touches only the + token-linear KV. + """ + return self.get_kvcache_bytes_per_sequence(seq_len) / self._cp_kv_memory_divisor() + + def get_kvcache_rank_batch_capacity(self, kv_budget_bytes: float, max_batch_size: int) -> int: + """Token capacity of a RANK-LOCAL byte budget: the inverse of + :meth:`get_kvcache_rank_bytes_per_sequence`. Under decode CP a budget of + ``B`` bytes holds the tokens whose full KV is ``B * dcp`` bytes.""" + return self.get_kvcache_batch_capacity(float(kv_budget_bytes) * self._cp_kv_memory_divisor(), max_batch_size) + def _binary_search_kvcache_max_tokens(self, kv_budget_bytes: float) -> int: """Monotonic-search inverse of :meth:`get_kvcache_bytes_per_sequence`. diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py index 160c9bcfb2..6d8b55af97 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek.py @@ -48,6 +48,17 @@ def supports_cp(cls, backend_name: str) -> bool: # (its Ring CP has no perf data / _cp_attn_comm_ops). return backend_name == "sglang" + @classmethod + def supports_dcp(cls, backend_name: str) -> bool: + # Dense MLA decode CP: vLLM `-dcp` (every MLA backend) and SGLang + # `--dcp-size` (flashinfer / trtllm_mla / cutedsl_mla / flashmla). Both + # stripe the latent KV inside the TP group and LSE-merge the partial + # outputs; BaseModel._apply_decode_context_parallel rewrites the + # GenerationMLA / GenerationAttention / WideEPGenerationMLA op and adds + # the merge collectives. TRT-LLM's Helix is a different layout (cp adds + # ranks, attention weights replicated) and is not modeled. + return backend_name in ("vllm", "sglang") + @classmethod def create(cls, model_info: dict, model_config, backend_name: str) -> BaseModel: moe_args = (model_info["topk"], model_info["num_experts"], model_info["moe_inter_size"]) 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 0a3d9f8179..432e891bd7 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py @@ -299,6 +299,41 @@ def apply_prefill_graph_profile(self): *(original[name] for name in retained), ] + @classmethod + def supports_dcp(cls, backend_name: str) -> bool: + # Sparse-MLA (DSA) decode CP: vLLM `flashmla_sparse` / `flashinfer_mla_sparse` + # + DCP-aware indexer. SGLang's NSA backend has no DCP path yet (the + # DSA-family DCP PRs, #39330 / #36990 / #39117, are open), so it is not + # claimed. The GenerationDSAModule op prices the gathered heads over + # the per-rank KV stripe with the sparse attention over the rank's + # ceil(topk / dcp) owned slots; BaseModel._apply_decode_context_parallel + # adds the merge collectives and the indexer top-k gather. + return backend_name == "vllm" + + def _validate_dcp_topology(self) -> None: + super()._validate_dcp_topology() + # vLLM's FlashMLA sparse backend (the default DSA decode kernel unless + # FLASHINFER_MLA_SPARSE is selected, `attention_backend="flashinfer"` + # here) only serves DCP with the `ag_rs` merge and the fp8_ds_mla + # mixed-batch KV path (`flashmla_sparse.py`: NotImplementedError + # otherwise). Do not price a configuration vLLM refuses to launch. + backend = getattr(self, "_backend_name", None) + attention_backend = str(getattr(self.config, "attention_backend", None) or "").lower() + if backend != "vllm" or "flashinfer" in attention_backend: + return + if self._dcp_comm_style() != "ag_rs": + raise ValueError( + f"{self.architecture} decode context parallelism on vLLM's FlashMLA sparse backend only runs " + f"with dcp_comm='ag_rs' (got {self.config.dcp_comm!r}); select attention_backend='flashinfer' " + "(FLASHINFER_MLA_SPARSE) for the a2a merges." + ) + if self.config.kvcache_quant_mode.value.memory != 1: + raise ValueError( + f"{self.architecture} decode context parallelism on vLLM's FlashMLA sparse backend requires an " + f"fp8 KV cache (got kvcache_quant_mode={self.config.kvcache_quant_mode.name}); the bf16 sparse " + "path returns no LSE for the DCP merge." + ) + @classmethod def create(cls, model_info: dict, model_config, backend_name: str) -> BaseModel: moe_args = (model_info["topk"], model_info["num_experts"], model_info["moe_inter_size"]) 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 8305a48e9b..cad8d0ac51 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/kimi_k3.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/kimi_k3.py @@ -63,6 +63,24 @@ class KimiK3Model(BaseModel): # context — over-charges only short-context high-concurrency mixes. KDA_STATE_SLOTS_PER_REQUEST = 5 + @classmethod + def supports_dcp(cls, backend_name: str) -> bool: + # SGLang ships Kimi-K3 decode CP (`arg_groups/model_overrides/kimi_k3.py`: + # cutedsl_mla / tokenspeed_mla / aiter with `enable_dcp`, + # `models/kimi_k3.py::prepare_context_parallel_metadata_for_dcp`). DCP + # stripes the KV of the MLA layers only; the KDA layers keep no KV and + # ignore it, which is exactly what the op-type rewrite does. vLLM has no + # K3 DCP path (the hybrid KV manager rejects dcp > 1). + return backend_name == "sglang" + + def _dcp_q_replicate(self) -> bool: + # SGLang's K3 override defaults `dcp_replicate_q_proj=True` under the + # a2a / fi_a2a merges (its DCP defaults), dropping the per-layer query + # all-gather in favour of a replicated Q projection. + if getattr(self.config, "dcp_q_replicate", None) is None and self._backend_name == "sglang": + return self._dcp_comm_style() == "a2a" + return super()._dcp_q_replicate() + @classmethod def create(cls, model_info: dict, model_config, backend_name: str) -> BaseModel: return cls( @@ -824,3 +842,20 @@ def get_kvcache_max_tokens(self, kv_budget_bytes: float) -> int: if budget <= 0 or per_token <= 0: return 0 return int(budget // per_token) + + # Decode CP stripes only the MLA layers' token-linear KV; the KDA state is + # per-request and rank-local (every rank runs the full recurrence on its + # own KDA head shard), so it must be reserved whole BEFORE the stripe. + def get_kvcache_rank_bytes_per_sequence(self, seq_len: int) -> float: + token_bytes = ( + max(0, seq_len) * self.config.kvcache_quant_mode.value.memory * self.get_kvcache_elements_per_token() + ) + return token_bytes / self._cp_kv_memory_divisor() + self._kda_state_bytes_per_request() + + def get_kvcache_rank_batch_capacity(self, kv_budget_bytes: float, max_batch_size: int) -> int: + token_budget = float(kv_budget_bytes) - self._kda_state_bytes_per_request() + if token_budget <= 0: + return 0 + return self.get_kvcache_batch_capacity( + token_budget * self._cp_kv_memory_divisor() + self._kda_state_bytes_per_request(), max_batch_size + ) diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/llama.py b/python/aisimulate/src/aisimulate_core/sdk/models/llama.py index 5d4cbcee91..cd28c69a4a 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/llama.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/llama.py @@ -23,9 +23,20 @@ class LLAMAModel(BaseModel): @classmethod def supports_cp(cls, backend_name: str) -> bool: # Dense GQA prefill CP: SGLang AllGather (zigzag in-seq split). vLLM CP - # is decode-time DCP (a different concept) -- intentionally excluded. + # is decode-time DCP (a different concept) -- see supports_dcp. return backend_name == "sglang" + @classmethod + def supports_dcp(cls, backend_name: str) -> bool: + # GQA decode CP: vLLM FlashAttention / FlashInfer backends (dcp <= tp / + # kv_heads, the TP-replicated kv heads are what DCP de-duplicates). + # SGLang's GQA DCP exists only in its Triton backend, so it is not + # claimed here. + return backend_name == "vllm" + + def _dcp_kv_head_replication(self) -> int | None: + return int(self.config.tp_size) // max(1, int(self._num_kv_heads)) + @classmethod def create(cls, model_info: dict, model_config, backend_name: str) -> BaseModel: return cls( diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/moe.py b/python/aisimulate/src/aisimulate_core/sdk/models/moe.py index c2f269ac9c..50a4ec41d1 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/moe.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/moe.py @@ -36,6 +36,14 @@ def supports_cp(cls, backend_name: str) -> bool: # ContextAttention cp_size, token-major seq_split, MoEDispatch attn_cp_size). return backend_name == "sglang" + @classmethod + def supports_dcp(cls, backend_name: str) -> bool: + # GQA decode CP on vLLM (FlashAttention / FlashInfer); see llama.py. + return backend_name == "vllm" + + def _dcp_kv_head_replication(self) -> int | None: + return int(self.config.tp_size) // max(1, int(self._num_kv_heads)) + @classmethod def create(cls, model_info: dict, model_config, backend_name: str) -> BaseModel: moe_args = (model_info["topk"], model_info["num_experts"], model_info["moe_inter_size"]) 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 4826d484f8..c6ab193c1a 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -133,6 +133,10 @@ class ForwardPassPerfModelConfig: dcp: int | None = dataclass_field(default=None, kw_only=True) moe_tp_size: int | None = None moe_ep_size: int | None = None + # Prefill context parallelism (widens the attention side) and decode + # context parallelism (stripes the decode KV across the TP ranks); None + # keeps both at one and out of the serialized identity. + cp_size: int | None = dataclass_field(default=None, kw_only=True) gemm_quant_mode: str | None = None moe_quant_mode: str | None = None fmha_quant_mode: str | None = None @@ -1379,6 +1383,11 @@ def _engine_config_json(model: Any, database: Any) -> str: "fpm_config": fpm_config.cache_identity() if fpm_config is not None else None, "decoder_replay": bool(getattr(model_config, "decoder_replay", False)), "cp_style": getattr(model_config, "cp_style", None), + # DCP op-shaping overrides: the merge collective + # (ag_rs vs a2a) and the replicated-Q variant compile + # different op graphs for one (tp, dcp) identity. + "dcp_comm": getattr(model_config, "dcp_comm", None), + "dcp_q_replicate": getattr(model_config, "dcp_q_replicate", None), "workload_distribution": getattr(model_config, "workload_distribution", None), "decode_workload_distribution": getattr(model_config, "decode_workload_distribution", None), "prefill_graph_profile": getattr(model_config, "prefill_graph_profile", None), diff --git a/python/aisimulate/tests/unit/generator/test_context_parallel.py b/python/aisimulate/tests/unit/generator/test_context_parallel.py new file mode 100644 index 0000000000..498d1d7e04 --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_context_parallel.py @@ -0,0 +1,691 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prefill CP / decode CP rendering through the generator. + +Covers the single decision point (``generator/context_parallel.py``), the +rule + mapping + versioned-template chain on both request paths (Sweeper +candidate and SDK result bridge), and the FPM world-size check. +""" + +from __future__ import annotations + +import copy +import shlex +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest +from jinja2 import Environment, FileSystemLoader + +from aisimulate.fpm_contract import FPM_RUN_SCRIPT_FILENAME +from aisimulate.generator.api import generate_backend_artifacts, generate_from_request +from aisimulate.generator.context_parallel import ( + ContextParallelUnsupportedError, + context_parallel_params, + cp_strategy_for, +) +from aisimulate.generator.module_bridge import task_config_to_generator_config +from aisimulate.generator.rendering.engine import render_backend_templates +from aisimulate.generator.request import ModelFacts, SweeperCandidateError, from_sweeper_candidate + +from .test_fpm_artifacts import _params as _fpm_params +from .test_sweeper_request import _agg_candidate, _cli_flag_value, _disagg_candidate + +pytestmark = pytest.mark.unit + +_TEMPLATE_ROOT = ( + Path(__file__).resolve().parents[3] / "src" / "aisimulate" / "generator" / "config" / "backend_templates" +) +_WORKLOAD = {"isl": 4096, "osl": 512} + + +def _flags(cli_args: str) -> list[str]: + return shlex.split(cli_args) + + +def _render_template(backend: str, name: str, **context) -> str: + env = Environment(loader=FileSystemLoader(str(_TEMPLATE_ROOT / backend))) + return env.get_template(name).render(**context) + + +# --------------------------------------------------------------------------- # +# The decision point +# --------------------------------------------------------------------------- # + + +def test_both_knobs_at_one_render_nothing(): + assert context_parallel_params(backend="vllm", backend_version="0.20.1") == {} + assert context_parallel_params(backend="trtllm", backend_version="1.3.0rc14") == {} + + +def test_decode_cp_alone_renders_only_its_own_knob(): + params = context_parallel_params(backend="vllm", backend_version="0.20.1", decode_context_parallel_size=8) + assert params == {"decode_context_parallel_size": 8} + + +@pytest.mark.parametrize( + ("backend", "version", "kwargs", "expected"), + [ + # vLLM floors are the first cli_args template that renders the flags (0.14.1), + # not the framework release that introduced them. + ("vllm", "0.14.1", {"context_parallel_size": 2}, {"context_parallel_size": 2}), + ("vllm", "0.14.1", {"decode_context_parallel_size": 4}, {"decode_context_parallel_size": 4}), + ( + "vllm", + "0.19.0", + {"decode_context_parallel_size": 4, "dcp_comm_backend": "a2a"}, + {"decode_context_parallel_size": 4, "dcp_comm_backend": "a2a"}, + ), + ( + "sglang", + "0.5.15", + {"context_parallel_size": 2, "architecture": "Qwen3ForCausalLM"}, + {"context_parallel_size": 2, "cp_strategy": "zigzag"}, + ), + ("sglang", "0.5.15", {"decode_context_parallel_size": 2}, {"decode_context_parallel_size": 2}), + ( + "sglang", + "0.5.17", + {"decode_context_parallel_size": 2, "dcp_comm_backend": "fi_a2a"}, + {"decode_context_parallel_size": 2, "dcp_comm_backend": "fi_a2a"}, + ), + ], +) +def test_minimum_versions_are_inclusive(backend, version, kwargs, expected): + assert context_parallel_params(backend=backend, backend_version=version, **kwargs) == expected + + +@pytest.mark.parametrize( + ("backend", "version", "kwargs", "needle"), + [ + ("vllm", "0.13.0", {"context_parallel_size": 2}, ">= 0.14.1"), + ("vllm", "0.12.0", {"decode_context_parallel_size": 2}, ">= 0.14.1"), + # 0.18.0 has the flag upstream but floor-matches the 0.16.0 template, which + # does not render it; the floor is the first template that does. + ("vllm", "0.18.0", {"decode_context_parallel_size": 2, "dcp_comm_backend": "a2a"}, ">= 0.19.0"), + ("sglang", "0.5.14", {"context_parallel_size": 2}, ">= 0.5.15"), + ("sglang", "0.5.11", {"decode_context_parallel_size": 2}, ">= 0.5.15"), + ("sglang", "0.5.16", {"decode_context_parallel_size": 2, "dcp_comm_backend": "a2a"}, ">= 0.5.17"), + ("trtllm", "1.3.0rc14", {"decode_context_parallel_size": 2}, "Helix"), + ("trtllm", "1.3.0rc14", {"context_parallel_size": 2}, "no launch flag"), + ("vllm", None, {"decode_context_parallel_size": 2}, "parseable"), + ], +) +def test_unsupported_backend_or_version_fails_loudly(backend, version, kwargs, needle): + with pytest.raises(ContextParallelUnsupportedError, match=needle): + context_parallel_params(backend=backend, backend_version=version, **kwargs) + + +def test_dcp_comm_backend_is_validated_per_backend(): + with pytest.raises(ContextParallelUnsupportedError, match="not a vllm choice"): + context_parallel_params( + backend="vllm", backend_version="0.20.1", decode_context_parallel_size=2, dcp_comm_backend="fi_a2a" + ) + with pytest.raises(ContextParallelUnsupportedError, match="needs decode_context_parallel_size > 1"): + context_parallel_params(backend="vllm", backend_version="0.20.1", dcp_comm_backend="a2a") + + +@pytest.mark.parametrize( + ("architecture", "family", "expected"), + [ + ("Qwen3ForCausalLM", None, "zigzag"), + ("DeepseekV3ForCausalLM", "DEEPSEEK", "zigzag"), + ("DeepseekV32ForCausalLM", None, "interleave"), + ("GlmMoeDsaForCausalLM", None, "interleave"), + (None, "DEEPSEEKV32", "interleave"), + ], +) +def test_sglang_cp_strategy_follows_dsa(architecture, family, expected): + assert cp_strategy_for(architecture, family) == expected + + +# --------------------------------------------------------------------------- # +# Versioned templates +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("version", ["0.14.1", "0.16.0", "0.19.0", "0.20.1"]) +def test_vllm_templates_render_both_cp_flags(version): + out = _render_template( + "vllm", + f"cli_args.{version}.j2", + vllm={"tensor-parallel-size": 8, "prefill-context-parallel-size": 2, "decode-context-parallel-size": 8}, + speculative_config={}, + ) + tokens = _flags(out) + assert _cli_flag_value(out, "--prefill-context-parallel-size") == "2" + assert _cli_flag_value(out, "--decode-context-parallel-size") == "8" + assert "--dcp-comm-backend" not in tokens + + +@pytest.mark.parametrize(("version", "rendered"), [("0.16.0", False), ("0.19.0", True), ("0.20.1", True)]) +def test_vllm_dcp_comm_backend_only_from_0_19_template(version, rendered): + out = _render_template( + "vllm", + f"cli_args.{version}.j2", + vllm={"decode-context-parallel-size": 8, "dcp-comm-backend": "a2a"}, + speculative_config={}, + ) + assert ("--dcp-comm-backend" in _flags(out)) is rendered + + +@pytest.mark.parametrize("version", ["0.5.15", "0.5.17"]) +def test_sglang_templates_render_prefill_and_decode_cp(version): + out = _render_template( + "sglang", + f"cli_args.{version}.j2", + sglang={ + "tensor-parallel-size": 8, + "attn-cp-size": 4, + "cp-strategy": "zigzag", + "dcp-size": 2, + "dcp-comm-backend": "a2a", + }, + ) + tokens = _flags(out) + assert "--enable-prefill-cp" in tokens + assert _cli_flag_value(out, "--attn-cp-size") == "4" + assert _cli_flag_value(out, "--cp-strategy") == "zigzag" + assert _cli_flag_value(out, "--dcp-size") == "2" + assert ("--dcp-comm-backend" in tokens) is (version == "0.5.17") + + +# --------------------------------------------------------------------------- # +# Sweeper candidate path (rule + mapping + template, per backend) +# --------------------------------------------------------------------------- # + + +def test_vllm_agg_decode_cp_renders_flag_without_extra_gpus(): + candidate = _agg_candidate(tp=8, replicas=1, used_gpus=8, dcp=8) + request = from_sweeper_candidate( + candidate, + workload=_WORKLOAD, + model_facts=ModelFacts(is_moe=False, architecture="Qwen3ForCausalLM"), + ) + assert request.topology.roles["agg"].extra["gpus_per_worker"] == 8 + assert request.topology.roles["agg"].extra["decode_context_parallel_size"] == 8 + + cli = generate_from_request(request)["cli_args_agg"] + assert _cli_flag_value(cli, "--tensor-parallel-size") == "8" + assert _cli_flag_value(cli, "--decode-context-parallel-size") == "8" + assert "--prefill-context-parallel-size" not in _flags(cli) + + +def test_vllm_agg_prefill_cp_multiplies_worker_gpus(): + candidate = _agg_candidate(tp=2, replicas=4, used_gpus=16, cp=2) + request = from_sweeper_candidate( + candidate, + workload=_WORKLOAD, + model_facts=ModelFacts(is_moe=False, architecture="Qwen3ForCausalLM"), + ) + assert request.topology.roles["agg"].extra["gpus_per_worker"] == 4 + + artifacts = generate_from_request(request) + cli = artifacts["cli_args_agg"] + # vLLM's PCP expands the world size; TP stays the weight shard. + assert _cli_flag_value(cli, "--tensor-parallel-size") == "2" + assert _cli_flag_value(cli, "--prefill-context-parallel-size") == "2" + assert "--decode-context-parallel-size" not in _flags(cli) + + +def test_vllm_prefill_cp_gpu_mismatch_is_rejected(): + # used_gpus still counts tp*replicas only: the candidate is inconsistent. + candidate = _agg_candidate(tp=2, replicas=4, used_gpus=8, cp=2) + with pytest.raises(SweeperCandidateError): + from_sweeper_candidate(candidate, workload=_WORKLOAD, model_facts=ModelFacts(is_moe=False)) + + +def test_vllm_too_old_for_prefill_cp_is_rejected(): + candidate = _agg_candidate(backend_version="0.11.0", tp=2, replicas=4, used_gpus=16, cp=2) + with pytest.raises(SweeperCandidateError, match=">= 0.14.1"): + from_sweeper_candidate(candidate, workload=_WORKLOAD, model_facts=ModelFacts(is_moe=False)) + + +def _sglang_dense_disagg(**overrides): + base = { + "model_name": "Qwen/Qwen3-32B-FP8", + "backend": "sglang", + "backend_version": "0.5.16", + "prefill_tp": 2, + "prefill_moe_tp": 1, + "prefill_moe_ep": 1, + "prefill_replicas": 1, + "prefill_cp": 4, + "decode_tp": 4, + "decode_moe_tp": 1, + "decode_moe_ep": 1, + "decode_replicas": 1, + "decode_dcp": 4, + "used_gpus": 12, + # The shared disagg fixture carries a TRT-LLM-only KVBM adapter. + "adapters": {"dynamo.router": {"mode": "kv"}}, + } + base.update(overrides) + return _disagg_candidate(**base) + + +def test_sglang_disagg_folds_prefill_cp_into_tp_and_stripes_decode(): + request = from_sweeper_candidate( + _sglang_dense_disagg(), + workload=_WORKLOAD, + model_facts=ModelFacts(is_moe=False, architecture="Qwen3ForCausalLM"), + ) + prefill = request.topology.roles["prefill"].extra + decode = request.topology.roles["decode"].extra + assert prefill["gpus_per_worker"] == 8 + assert prefill["context_parallel_size"] == 4 + assert prefill["cp_strategy"] == "zigzag" + assert decode["gpus_per_worker"] == 4 + assert decode["decode_context_parallel_size"] == 4 + + artifacts = generate_from_request(request) + p_cli = artifacts["cli_args_prefill"] + d_cli = artifacts["cli_args_decode"] + # SGLang's attention-CP ranks live inside --tp: tp(2) x cp(4) -> 8. + assert _cli_flag_value(p_cli, "--tensor-parallel-size") == "8" + assert "--enable-prefill-cp" in _flags(p_cli) + assert _cli_flag_value(p_cli, "--attn-cp-size") == "4" + assert _cli_flag_value(p_cli, "--cp-strategy") == "zigzag" + assert "--dcp-size" not in _flags(p_cli) + + assert _cli_flag_value(d_cli, "--tensor-parallel-size") == "4" + assert _cli_flag_value(d_cli, "--dcp-size") == "4" + for flag in ("--enable-prefill-cp", "--attn-cp-size", "--cp-strategy"): + assert flag not in _flags(d_cli) + + +def test_sglang_moe_prefill_cp_uses_moe_fold_and_interleave_for_dsa(): + candidate = _agg_candidate( + model_name="deepseek-ai/DeepSeek-V3.2", + backend="sglang", + backend_version="0.5.16", + hardware_sku="gb200", + tp=1, + moe_tp=1, + moe_ep=8, + cp=8, + replicas=1, + used_gpus=8, + ) + request = from_sweeper_candidate( + candidate, + workload=_WORKLOAD, + model_facts=ModelFacts(is_moe=True, architecture="DeepseekV32ForCausalLM"), + ) + assert request.topology.roles["agg"].extra["cp_strategy"] == "interleave" + + cli = generate_from_request(request)["cli_args_agg"] + # MoE: TP = moe_tp * moe_ep already equals the attention-group width; the + # dense-only CP fold must not double it. + assert _cli_flag_value(cli, "--tensor-parallel-size") == "8" + assert _cli_flag_value(cli, "--attn-cp-size") == "8" + assert _cli_flag_value(cli, "--cp-strategy") == "interleave" + assert "--enable-prefill-cp" in _flags(cli) + + +def test_sglang_too_old_for_decode_cp_is_rejected(): + with pytest.raises(SweeperCandidateError, match=">= 0.5.15"): + from_sweeper_candidate( + _sglang_dense_disagg(backend_version="0.5.11"), + workload=_WORKLOAD, + model_facts=ModelFacts(is_moe=False, architecture="Qwen3ForCausalLM"), + ) + + +def test_trtllm_decode_cp_is_rejected(): + candidate = _disagg_candidate(decode_dcp=4) + with pytest.raises(SweeperCandidateError, match="Helix"): + from_sweeper_candidate( + candidate, + workload=_WORKLOAD, + model_facts=ModelFacts(is_moe=True, architecture="DeepseekV3ForCausalLM"), + ) + + +def test_candidates_without_cp_columns_are_unchanged(): + request = from_sweeper_candidate( + _agg_candidate(), + workload=_WORKLOAD, + model_facts=ModelFacts(is_moe=False, architecture="Qwen3ForCausalLM"), + ) + extra = request.topology.roles["agg"].extra + assert extra["gpus_per_worker"] == 2 + for key in ("context_parallel_size", "decode_context_parallel_size", "cp_strategy", "dcp_comm_backend"): + assert key not in extra + + +# --------------------------------------------------------------------------- # +# SGLang GPU accounting under the MoE fold +# --------------------------------------------------------------------------- # + + +def _sglang_rule_params(*, is_moe: bool, rule: str | None = None, **agg) -> dict: + params = { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "moe_tensor_parallel_size": 1, + "moe_expert_parallel_size": 1, + "max_batch_size": 64, + } + params.update(agg) + payload = { + "params": {"agg": params}, + "ModelConfig": {"is_moe": is_moe, "prefix": 0, "nextn": 0}, + "SlaConfig": {"isl": 1024, "osl": 128}, + "DynConfig": {"mode": "agg"}, + "WorkerConfig": {"agg_workers": 1}, + "BenchConfig": {}, + } + if rule is not None: + payload["rule"] = rule + return payload + + +def test_vllm_moe_fold_keeps_prefill_cp_ranks_out_of_tp(): + from aisimulate.generator.rendering.rule_engine import apply_rule_plugins + + # Sweeper candidate tp=8, dp=1, cp=2 with moe_tp=1 x moe_ep=16: the attention + # width tp*dp*cp == moe_tp*moe_ep is 16 GPUs. vLLM's -pcp ranks are separate + # from --tp, so the fold must yield --tp 8 and a 16-GPU worker, not --tp 16 + # with 32 GPUs. + moe = apply_rule_plugins( + _sglang_rule_params( + is_moe=True, + tensor_parallel_size=8, + context_parallel_size=2, + moe_tensor_parallel_size=1, + moe_expert_parallel_size=16, + ), + "vllm", + )["params"]["agg"] + assert moe["tensor_parallel_size"] == 8 + assert moe["gpus_per_worker"] == 16 + + +@pytest.mark.parametrize("rule", [None, "benchmark"]) +def test_sglang_moe_fold_does_not_count_attention_dp_twice(rule): + from aisimulate.generator.rendering.rule_engine import apply_rule_plugins + + # DeepSeek-style DEP on SGLang: attention tp=1 x dp=8, experts moe_ep=8. + # --tp is folded to moe_tp x moe_ep (8) and --dp partitions that --tp, so + # the worker is 8 GPUs, not 8 x 8. + moe = apply_rule_plugins( + _sglang_rule_params(is_moe=True, rule=rule, data_parallel_size=8, moe_expert_parallel_size=8), + "sglang", + )["params"]["agg"] + assert moe["tensor_parallel_size"] == 8 + assert moe["gpus_per_worker"] == 8 + + # Dense attention-DP keeps the plain tp x pp x dp product. + dense = apply_rule_plugins( + _sglang_rule_params(is_moe=False, rule=rule, tensor_parallel_size=2, data_parallel_size=2), + "sglang", + )["params"]["agg"] + assert dense["tensor_parallel_size"] == 2 + assert dense["gpus_per_worker"] == 4 + + # MoE without moe_tp/moe_ep configured (no fold) behaves like dense. + unfolded = apply_rule_plugins( + _sglang_rule_params( + is_moe=True, + rule=rule, + tensor_parallel_size=2, + data_parallel_size=2, + moe_tensor_parallel_size=None, + moe_expert_parallel_size=None, + ), + "sglang", + )["params"]["agg"] + assert unfolded["gpus_per_worker"] == 4 + + +# --------------------------------------------------------------------------- # +# SDK result bridge path +# --------------------------------------------------------------------------- # + + +def _task(**overrides) -> SimpleNamespace: + fields = { + "primary_backend_name": "sglang", + "primary_system_name": "gb200", + "primary_backend_version": "0.5.16", + "primary_model_path": "Qwen/Qwen3-32B-FP8", + "prefix": 0, + "is_moe": False, + "nextn": 0, + "nextn_accepted": None, + "serving_mode": "agg", + "total_gpus": 0, + "system_name": "gb200", + "prefill_system_name": "gb200", + "decode_system_name": "gb200", + "isl": 1024, + "osl": 256, + "ttft": 2000.0, + "tpot": 50.0, + "attention_backend": None, + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test_bridge_reads_prefill_cp_from_result_row(): + row = pd.Series({"workers": 1, "tp": 2, "pp": 1, "dp": 1, "cp": 4, "bs": 64}) + + result = task_config_to_generator_config(_task(), row, num_gpus_per_node=8) + agg = result["params"]["agg"] + assert agg["context_parallel_size"] == 4 + assert agg["cp_strategy"] == "zigzag" + assert agg["gpus_per_worker"] == 8 + assert "decode_context_parallel_size" not in agg + + artifacts = generate_backend_artifacts(result, "sglang", backend_version="0.5.16", deployment_target="dynamo-j2") + cli = artifacts["cli_args_agg"] + assert _cli_flag_value(cli, "--tensor-parallel-size") == "8" + assert _cli_flag_value(cli, "--attn-cp-size") == "4" + assert "--enable-prefill-cp" in _flags(cli) + + +def test_bridge_reads_decode_cp_from_task(): + row = pd.Series({"workers": 1, "tp": 4, "pp": 1, "dp": 1, "bs": 64}) + + task = _task(primary_backend_version="0.5.17", dcp_size=4, dcp_comm="a2a") + result = task_config_to_generator_config(task, row, num_gpus_per_node=8) + agg = result["params"]["agg"] + assert agg["decode_context_parallel_size"] == 4 + assert agg["dcp_comm_backend"] == "a2a" + assert agg["gpus_per_worker"] == 4 + assert "context_parallel_size" not in agg + + artifacts = generate_backend_artifacts(result, "sglang", backend_version="0.5.17", deployment_target="dynamo-j2") + cli = artifacts["cli_args_agg"] + assert _cli_flag_value(cli, "--tensor-parallel-size") == "4" + assert _cli_flag_value(cli, "--dcp-size") == "4" + assert _cli_flag_value(cli, "--dcp-comm-backend") == "a2a" + + +def test_vllm_bridge_renders_dcp_comm_backend_at_its_minimum_version(): + """Complete path at the floor: the accepted minimum for ``dcp_comm_backend`` + must select a template that renders it. 0.18.0 has the flag upstream but + floor-matches cli_args.0.16.0.j2, which has no line for it, so the floor is + 0.19.0 and an 0.18.0 request is rejected rather than silently launched on + the backend-default merge.""" + row = pd.Series({"workers": 1, "tp": 8, "pp": 1, "dp": 1, "bs": 64}) + + task = _task(primary_backend_name="vllm", primary_backend_version="0.19.0", dcp_size=4, dcp_comm="a2a") + result = task_config_to_generator_config(task, row, num_gpus_per_node=8) + assert result["params"]["agg"]["dcp_comm_backend"] == "a2a" + cli = generate_backend_artifacts(result, "vllm", backend_version="0.19.0", deployment_target="dynamo-j2")[ + "cli_args_agg" + ] + assert _cli_flag_value(cli, "--decode-context-parallel-size") == "4" + assert _cli_flag_value(cli, "--dcp-comm-backend") == "a2a" + + with pytest.raises(ContextParallelUnsupportedError, match=">= 0.19.0"): + task_config_to_generator_config( + _task(primary_backend_name="vllm", primary_backend_version="0.18.0", dcp_size=4, dcp_comm="a2a"), + row, + num_gpus_per_node=8, + ) + + +def test_bridge_disagg_uses_per_role_decode_cp_and_prefill_cp_column(): + row = pd.Series( + { + "(p)workers": 1, + "(p)tp": 1, + "(p)pp": 1, + "(p)dp": 1, + "(p)cp": 2, + "(p)bs": 8, + "(d)workers": 1, + "(d)tp": 4, + "(d)pp": 1, + "(d)dp": 1, + "(d)bs": 64, + } + ) + task = _task( + primary_backend_name="vllm", + primary_backend_version="0.20.1", + serving_mode="disagg", + prefill_dcp_size=1, + decode_dcp_size=4, + ) + + result = task_config_to_generator_config(task, row, num_gpus_per_node=8) + prefill = result["params"]["prefill"] + decode = result["params"]["decode"] + assert prefill["context_parallel_size"] == 2 + assert prefill["gpus_per_worker"] == 2 + assert "decode_context_parallel_size" not in prefill + assert decode["decode_context_parallel_size"] == 4 + assert decode["gpus_per_worker"] == 4 + assert "context_parallel_size" not in decode + + +def test_bridge_without_cp_is_unchanged(): + row = pd.Series({"workers": 1, "tp": 2, "pp": 1, "dp": 1, "bs": 64}) + agg = task_config_to_generator_config(_task(), row, num_gpus_per_node=8)["params"]["agg"] + assert agg["gpus_per_worker"] == 2 + for key in ("context_parallel_size", "decode_context_parallel_size", "cp_strategy", "dcp_comm_backend"): + assert key not in agg + + +def test_bridge_revalidates_context_parallel_overrides_and_resizes_the_worker(): + row = pd.Series({"workers": 1, "tp": 4, "pp": 1, "dp": 1, "bs": 64}) + task = _task(primary_backend_version="0.5.17") + + # A Workers. override introducing decode CP is validated and does not + # change the GPU count; one introducing prefill CP grows the worker. + striped = task_config_to_generator_config( + task, row, generator_overrides={"Workers": {"agg": {"decode_context_parallel_size": 4}}}, num_gpus_per_node=8 + )["params"]["agg"] + assert striped["decode_context_parallel_size"] == 4 + assert striped["gpus_per_worker"] == 4 + + wider = task_config_to_generator_config( + task, row, generator_overrides={"Workers": {"agg": {"context_parallel_size": 2}}}, num_gpus_per_node=8 + )["params"]["agg"] + assert (wider["context_parallel_size"], wider["cp_strategy"]) == (2, "zigzag") + assert wider["gpus_per_worker"] == 8 + + # Overrides cannot bypass the decision point: a comm backend without DCP, + # or a knob a too-old backend cannot launch, still fail loud. + with pytest.raises(ContextParallelUnsupportedError, match="needs decode_context_parallel_size > 1"): + task_config_to_generator_config( + task, row, generator_overrides={"Workers": {"agg": {"dcp_comm_backend": "a2a"}}}, num_gpus_per_node=8 + ) + with pytest.raises(ContextParallelUnsupportedError, match=">= 0.5.15"): + task_config_to_generator_config( + _task(primary_backend_version="0.5.11"), + row, + generator_overrides={"Workers": {"agg": {"decode_context_parallel_size": 4}}}, + num_gpus_per_node=8, + ) + # cp_strategy is model-derived: an override may restate it but not contradict it. + restated = task_config_to_generator_config( + task, + row, + generator_overrides={"Workers": {"agg": {"context_parallel_size": 2, "cp_strategy": "zigzag"}}}, + num_gpus_per_node=8, + )["params"]["agg"] + assert restated["cp_strategy"] == "zigzag" + with pytest.raises(ContextParallelUnsupportedError, match="cp_strategy='interleave' does not match"): + task_config_to_generator_config( + task, + row, + generator_overrides={"Workers": {"agg": {"context_parallel_size": 2, "cp_strategy": "interleave"}}}, + num_gpus_per_node=8, + ) + # Malformed override values are rejected, never coerced into another topology. + for bad in ("invalid", 2.5, True, 0): + with pytest.raises(ContextParallelUnsupportedError, match="must be a positive integer"): + task_config_to_generator_config( + task, + row, + generator_overrides={"Workers": {"agg": {"decode_context_parallel_size": bad}}}, + num_gpus_per_node=8, + ) + + +def test_bridge_rejects_cp_on_trtllm(): + row = pd.Series({"workers": 1, "tp": 2, "pp": 1, "dp": 1, "bs": 64}) + task = _task(primary_backend_name="trtllm", primary_backend_version="1.3.0rc14", dcp_size=2) + with pytest.raises(ContextParallelUnsupportedError, match="Helix"): + task_config_to_generator_config(task, row, num_gpus_per_node=8) + + +# --------------------------------------------------------------------------- # +# FPM target world-size check +# --------------------------------------------------------------------------- # + + +def _fpm_with(**agg_overrides) -> dict: + params = copy.deepcopy(_fpm_params()) + params["params"]["agg"].update(agg_overrides) + params["WorkerConfig"]["agg_gpus_per_worker"] = params["params"]["agg"]["gpus_per_worker"] + return params + + +def _fpm_engine_command(run_sh: str) -> list[str]: + line = next(line for line in run_sh.splitlines() if line.startswith("engine_command=(")) + return shlex.split(line[len("engine_command=(") : -1]) + + +def test_fpm_counts_prefill_cp_in_world_size(): + params = _fpm_with(tensor_parallel_size=2, gpus_per_worker=4, context_parallel_size=2) + artifacts = render_backend_templates(params, "vllm", version="0.20.1", deployment_target="fpm") + command = _fpm_engine_command(artifacts[FPM_RUN_SCRIPT_FILENAME]) + assert command[command.index("--tensor-parallel-size") + 1] == "2" + assert command[command.index("--prefill-context-parallel-size") + 1] == "2" + + +def test_fpm_decode_cp_adds_no_gpus(): + params = _fpm_with(tensor_parallel_size=4, gpus_per_worker=4, decode_context_parallel_size=4) + artifacts = render_backend_templates(params, "vllm", version="0.20.1", deployment_target="fpm") + command = _fpm_engine_command(artifacts[FPM_RUN_SCRIPT_FILENAME]) + assert command[command.index("--decode-context-parallel-size") + 1] == "4" + assert "--prefill-context-parallel-size" not in command + + +def test_fpm_topology_check_multiplies_prefill_cp(): + from aisimulate.generator.builders.fpm_builder import _resolve_topology + + worker = SimpleNamespace(resources={"limits": {"gpu": "4"}}, multinode=None) + context = {"NodeConfig": {"num_gpus_per_node": 8}} + + topology = _resolve_topology( + context, worker, ["--tensor-parallel-size", "2", "--prefill-context-parallel-size", "2"] + ) + assert topology["total_gpus"] == 4 + + with pytest.raises(ValueError, match=r"pcp\(1\)"): + _resolve_topology(context, worker, ["--tensor-parallel-size", "2"]) + + # Decode CP reuses the TP ranks: it must not be counted. + with pytest.raises(ValueError, match="does not match"): + _resolve_topology(context, worker, ["--tensor-parallel-size", "2", "--decode-context-parallel-size", "2"]) diff --git a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py index d3fcc12ec7..5cd26f86cb 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_base_backend.py @@ -141,7 +141,7 @@ def _model(): _num_experts=0, model_family="test", get_kvcache_bytes_per_sequence=lambda _seq_len: 0.0, - _cp_kv_memory_divisor=lambda: 1, + get_kvcache_rank_bytes_per_sequence=lambda _seq_len: 0.0, ) @staticmethod @@ -542,7 +542,7 @@ def test_run_agg_b1_uses_scheduled_activation_peak( model._num_experts = 0 model.model_family = "test" model.get_kvcache_bytes_per_sequence = lambda _seq_len: 0.0 - model._cp_kv_memory_divisor = lambda: 1 + model.get_kvcache_rank_bytes_per_sequence = lambda _seq_len: 0.0 database.system_spec["misc"] = {"nccl_mem": {1: 0}, "other_mem": 0} monkeypatch.setattr( backend, diff --git a/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py b/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py index 5b33cc30e9..643f003c1d 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_deepseek_v4_workspace_memory.py @@ -83,7 +83,7 @@ def _model(*, tp=4, dp=1, ep=1, hidden_size=7168, heads=128): _topk=6, config=SimpleNamespace(pp_size=1, tp_size=tp, attention_dp_size=dp, moe_ep_size=ep, nextn=0), get_kvcache_bytes_per_sequence=lambda _seq: 1024, - _cp_kv_memory_divisor=lambda: 1, + get_kvcache_rank_bytes_per_sequence=lambda _seq: 1024, ) diff --git a/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py b/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py index 86203a2cfc..52933838d2 100644 --- a/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py +++ b/python/aisimulate/tests/unit/sdk/backends/test_step3p7_memory.py @@ -44,7 +44,7 @@ def test_step3p7_uses_moe_activation_and_dispatch_workspace(backend, moe_coeffic get_kvcache_bytes_per_sequence=lambda _seq: 0, get_resident_weights_bytes=lambda: 0, get_additional_activation_bytes=lambda _tokens: 0, - _cp_kv_memory_divisor=lambda: 1, + get_kvcache_rank_bytes_per_sequence=lambda _seq: 0, ) database = SimpleNamespace(system_spec={"misc": {"nccl_mem": {1: 0}, "other_mem": 0}}) diff --git a/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py b/python/aisimulate/tests/unit/sdk/database/test_attention_lanes.py index 3f84cd61f9..1a4cf73013 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,15 @@ def _fetch(database, attribute, *, shared_layer=None): monkeypatch.setattr(_etv, "fetch_attention_lane_density", _fetch) -def test_engine_spec_schema_version_is_twenty_five(): - """Direct interpolation with typed FPM DCP changes the positional operator payload.""" +def test_engine_spec_schema_version_is_twenty_six(): + """Decode-CP ``dcp_size`` fields extend the positional operation payload + (v26), on top of the FPM-decoupling direct-interpolation selector (v25), + the typed FPM DCP (v24), the recorded-DCP identity (v23), the VR200 pilot + (v22), the exact MoE kernel source (v21), the DeepSeek-V4.1 FPM selector + (v20), the V4.1 KV layout (v19) and the verification width fields (v18).""" from aisimulate.sdk import engine - assert engine.ENGINE_SPEC_SCHEMA_VERSION == 25 + assert engine.ENGINE_SPEC_SCHEMA_VERSION == 26 def test_lanes_outside_the_known_vocabulary_stay_reachable(): diff --git a/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py b/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py index 78a5c04e2a..5d96e9dd6f 100644 --- a/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py +++ b/python/aisimulate/tests/unit/sdk/speculation/test_consumer_equivalence.py @@ -295,7 +295,7 @@ def test_memory_accounting_includes_draft_bytes(self): model._num_experts = 0 model.model_family = "GPT" model.get_kvcache_bytes_per_sequence = lambda seq_len: 2048.0 - model._cp_kv_memory_divisor = lambda: 1 + model.get_kvcache_rank_bytes_per_sequence = lambda seq_len: 2048.0 m_base = backend._get_memory_usage(base, database, **kwargs) m_draft = backend._get_memory_usage(drafted, database, **kwargs) 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 f5521ac867..f81b6efe88 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 @@ -2550,3 +2550,113 @@ def _task(**kwargs): assert _task(engine_step_backend="rust").engine_step_backend == "rust" assert _task(engine_step_backend="RUST").engine_step_backend == "rust" assert _task(engine_step_backend=None).engine_step_backend is None + + +# --------------------------------------------------------------------------- +# Context parallelism: prefill CP candidates vs the per-role decode-CP scalar +# --------------------------------------------------------------------------- + + +def _disagg_task(**overrides) -> Task: + # sglang: the only backend whose prefill CP is modeled (enumerate_parallel_config + # rejects cp>1 elsewhere), so the CP candidates below are actually enumerated. + kwargs = { + "serving_mode": "disagg", + "prefill_model_path": "deepseek-ai/DeepSeek-V3", + "prefill_system_name": "h200_sxm", + "prefill_backend_name": "sglang", + "decode_model_path": "deepseek-ai/DeepSeek-V3", + "decode_system_name": "h200_sxm", + "decode_backend_name": "sglang", + "total_gpus": 32, + } + kwargs.update(overrides) + return Task(**kwargs) + + +def test_decode_prefill_cp_candidates_are_enumerated_not_rejected(): + """The old ``decode CP must be 1`` gate is gone: a decode worker's prefill-CP + list is priced as given (the optimizer, not a gate, decides it is useless).""" + t = _disagg_task(decode_cp_candidates=[1, 2]) + tuples = list(t.iter_parallel("decode")) + assert tuples + assert all(len(tup) == 6 for tup in tuples) + + +def test_decode_prefill_cp_defaults_to_one(): + assert _disagg_task().decode_cp_candidates == [1] + + +def test_role_dcp_size_reaches_model_config(): + # ModelConfig requires dcp_size to divide tp_size, so pin the decode TP. + t = _disagg_task(decode_dcp_size=8, decode_tp_candidates=[8]) + assert t.build_model_config(role="decode").dcp_size == 8 + # Not requested stays None (priced as 1; keeps the unrecorded FPM identity). + assert t.build_model_config(role="prefill").dcp_size is None + + +def test_malformed_role_dcp_size_is_rejected_not_defaulted(): + t = _disagg_task(decode_dcp_size=0) + with pytest.raises(ValueError, match="dcp_size must be positive"): + t.build_model_config(role="decode") + + +@pytest.mark.parametrize( + ("overrides", "needle"), + [ + # A prefill engine stripes its KV only as the PCP+DCP layout. + ({"prefill_dcp_size": 4}, "PCP\\+DCP layout"), + ({"prefill_dcp_size": 2, "prefill_cp_candidates": [1, 2]}, "PCP\\+DCP layout"), + # vLLM NIXL: replicated PCP cannot feed a DCP-sharded decode; DCP sizes divide one another. + ( + {"prefill_backend_name": "vllm", "prefill_cp_candidates": [2], "decode_dcp_size": 8}, + "replicated-PCP", + ), + ( + {"prefill_backend_name": "vllm", "prefill_cp_candidates": [4], "prefill_dcp_size": 4, "decode_dcp_size": 6}, + "divide one another", + ), + ], +) +def test_disagg_prefill_context_parallel_layout_rules(overrides, needle): + # Task validates its parallel space on construction, so the rule fires there. + with pytest.raises(ValueError, match=needle): + list(_disagg_task(**overrides).iter_parallel("prefill")) + + +def test_disagg_allows_prefill_cp_and_decode_dcp_on_different_workers(): + """SGLang re-lays the KV out per decode DCP rank on the prefill side: no pairing rule.""" + t = _disagg_task(prefill_cp_candidates=[1, 2], decode_dcp_size=8, decode_tp_candidates=[8]) + assert list(t.iter_parallel("prefill")) + assert list(t.iter_parallel("decode")) + assert t.build_model_config(role="decode").dcp_size == 8 + + +def test_agg_rejects_prefill_cp_candidates_combined_with_dcp(): + # Task construction already enumerates the agg search space (fmha data + # fallback), so the topology rule fires at construction time. + with pytest.raises(ValueError, match="at most one of prefill CP and decode CP"): + Task( + serving_mode="agg", + model_path="deepseek-ai/DeepSeek-V3", + system_name="h200_sxm", + backend_name="sglang", + total_gpus=8, + dcp_size=8, + agg_cp_candidates=[1, 2], + ) + + +def test_agg_dcp_with_pinned_prefill_cp_enumerates_and_reaches_model_config(): + t = Task( + serving_mode="agg", + model_path="deepseek-ai/DeepSeek-V3", + system_name="h200_sxm", + backend_name="sglang", + total_gpus=8, + dcp_size=8, + agg_cp_candidates=[1], + agg_tp_candidates=[8], + ) + assert list(t.iter_parallel("agg")) + assert t.build_model_config(role="agg").dcp_size == 8 diff --git a/python/aisimulate/tests/unit/sdk/test_dcp_kv_capacity.py b/python/aisimulate/tests/unit/sdk/test_dcp_kv_capacity.py new file mode 100644 index 0000000000..00e9253cba --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_dcp_kv_capacity.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Decode CP stripes the persistent KV: per-rank bytes/token drop by dcp AND the +rank-local budget therefore holds dcp times as many tokens. Both halves must +move together; the token inverse used to see the full per-token size and left +the block count unchanged.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from aisimulate_core.sdk import memory + +pytestmark = pytest.mark.unit + +_GIB = 1024**3 +_PER_TOKEN = 100.0 + + +class _StubModel: + def __init__(self, dcp: int): + self._dcp = dcp + + def get_kvcache_bytes_per_sequence(self, seq_len): + return _PER_TOKEN * seq_len + + def get_kvcache_batch_capacity(self, budget, max_batch_size): + return int(budget // _PER_TOKEN) + + def get_kvcache_rank_bytes_per_sequence(self, seq_len): + return _PER_TOKEN * seq_len / self._dcp + + def get_kvcache_rank_batch_capacity(self, budget, max_batch_size): + return int(budget * self._dcp // _PER_TOKEN) + + +def _estimate(monkeypatch, model) -> dict: + class _StubBackend: + def _get_memory_usage(self, *a, **k): + return {"weights": 10.0, "activations": 1.0, "others": 1.0, "nccl": 1.0, "kvcache": 0.0} + + db = SimpleNamespace(version="0.14.1", system_spec={"gpu": {"mem_capacity": 100 * _GIB}}) + monkeypatch.setattr(memory, "get_model", lambda *a, **k: model) + monkeypatch.setattr(memory, "get_backend", lambda backend: _StubBackend()) + monkeypatch.setattr(memory.perf_database, "get_database", lambda *a, **k: db) + estimator = memory.KVCacheEstimator.from_request( + "deepseek-ai/DeepSeek-V3", + "h200_sxm", + "vllm", + max_num_tokens=8192, + max_batch_size=256, + tp_size=8, + moe_tp_size=8, + moe_ep_size=1, + dcp_size=getattr(model, "_dcp", 1), + ) + return estimator.estimate(is_of_free=False, fraction=0.9, gpu_memory_capacity_bytes_override=None) + + +def test_dcp_stripes_per_token_bytes_and_multiplies_token_capacity(monkeypatch): + base = _estimate(monkeypatch, _StubModel(dcp=1)) + striped = _estimate(monkeypatch, _StubModel(dcp=4)) + + assert base["total_kv_size_bytes"] == striped["total_kv_size_bytes"] + assert base["kv_size_per_token_bytes"] == int(_PER_TOKEN) + assert striped["kv_size_per_token_bytes"] == int(_PER_TOKEN / 4) + # A rank-local budget of B bytes holds tokens whose full KV is 4B bytes + # (both sides floor to whole tokens, so allow the rounding slack). + assert abs(striped["total_kv_size_tokens"] - 4 * base["total_kv_size_tokens"]) <= 4 + # Consistency: bytes / (bytes per token) == tokens on both sides. + assert striped["total_kv_size_tokens"] == striped["total_kv_size_bytes"] // striped["kv_size_per_token_bytes"] + + +def test_model_without_rank_hooks_keeps_full_kv(monkeypatch): + # Duck-typed model doubles (upstream memory tests) carry no rank hooks. + class _Plain: + def get_kvcache_bytes_per_sequence(self, seq_len): + return _PER_TOKEN * seq_len + + def get_kvcache_batch_capacity(self, budget, max_batch_size): + return int(budget // _PER_TOKEN) + + out = _estimate(monkeypatch, _Plain()) + assert out["kv_size_per_token_bytes"] == int(_PER_TOKEN) + assert out["total_kv_size_tokens"] == out["total_kv_size_bytes"] // int(_PER_TOKEN) + + +def test_kimi_k3_dcp_reserves_rank_local_kda_state_before_striping(): + """Hybrid model: only the MLA layers' token-linear KV is striped by dcp; the + per-request KDA state is rank-local and must be reserved whole. A budget + smaller than one request's state admits nothing, however large dcp is.""" + from aisimulate_core.sdk.config import ModelConfig + from aisimulate_core.sdk.models import get_model + + def _k3(dcp): + cfg = ModelConfig(tp_size=8, attention_dp_size=1, moe_tp_size=8, moe_ep_size=1, dcp_size=dcp) + return get_model("moonshotai/Kimi-K3", cfg, "sglang") + + plain, striped = _k3(1), _k3(4) + state = striped._kda_state_bytes_per_request() + per_token = plain.get_kvcache_bytes_per_sequence(1) - state + assert state > 0 and per_token > 0 + + # Per-rank bytes: state whole, token KV / 4. + assert striped.get_kvcache_rank_bytes_per_sequence(0) == pytest.approx(state) + assert striped.get_kvcache_rank_bytes_per_sequence(4096) == pytest.approx(state + 4096 * per_token / 4) + + # Inverse: half a request's state fits zero tokens (the old divisor + # reported ~10k tokens here); state + N striped tokens fits exactly N. + assert striped.get_kvcache_rank_batch_capacity(state / 2, max_batch_size=1) == 0 + assert striped.get_kvcache_rank_batch_capacity(state, max_batch_size=1) == 0 + assert striped.get_kvcache_rank_batch_capacity(state + 1000 * per_token / 4, max_batch_size=1) == 1000 + assert plain.get_kvcache_rank_batch_capacity(state + 1000 * per_token, max_batch_size=1) == 1000 diff --git a/python/aisimulate/tests/unit/sdk/test_dcp_plumbing.py b/python/aisimulate/tests/unit/sdk/test_dcp_plumbing.py new file mode 100644 index 0000000000..fcb743c14e --- /dev/null +++ b/python/aisimulate/tests/unit/sdk/test_dcp_plumbing.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Decode-context-parallel (DCP) field plumbing. + +DCP (vLLM ``-dcp`` / SGLang ``--dcp-size``) is a per-phase knob orthogonal to +prefill CP (``cp_size``): it stripes the DECODE KV cache across ranks that +already belong to the attention group. This file pins the first-version +contract: + +* the field exists on ``ModelConfig`` and does NOT widen the attention side; +* it travels through the shared config builder and the engine identity; +* it shards the per-rank persistent KV (prefill CP never did); +* ``get_model`` gates it on a model-level capability flag. + +Deployment policy (whether one worker may combine prefill CP with DCP) is +deliberately NOT a model-building concern; it lives in the topology layer and +is covered by the compiler / task_v2 tests. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from aisimulate_core.sdk import common, config +from aisimulate_core.sdk.config_builders import build_model_config +from aisimulate_core.sdk.models.base import BaseModel + +pytestmark = pytest.mark.unit + + +def test_model_config_dcp_defaults_to_one_and_does_not_widen_attention(): + cfg = config.ModelConfig( + tp_size=4, + attention_dp_size=1, + cp_size=2, + dcp_size=4, + moe_tp_size=1, + moe_ep_size=8, + ) + assert cfg.dcp_size == 4 + # Prefill CP folds into the attention width; decode CP reuses those ranks. + assert cfg.attn_width == 8 + assert cfg.total_gpus_per_worker == 8 + assert cfg.resolve_moe_parallelism() == (1, 8) + assert config.ModelConfig().dcp_size is None + + +@pytest.mark.parametrize("knob", ["cp_size", "dcp_size"]) +@pytest.mark.parametrize("bad", [0, 1.5, True]) +def test_model_config_rejects_malformed_context_parallel_sizes(knob, bad): + with pytest.raises(ValueError, match=f"{knob} must be (a )?positive"): + config.ModelConfig(**{knob: bad}) + + +def test_build_model_config_carries_both_context_parallel_knobs(): + cfg = build_model_config( + tp_size=8, + pp_size=1, + attention_dp_size=1, + moe_tp_size=16, + moe_ep_size=1, + cp_size=2, + dcp_size=8, + ) + assert (cfg.cp_size, cfg.dcp_size) == (2, 8) + defaults = build_model_config(tp_size=1, pp_size=1, attention_dp_size=1, moe_tp_size=1, moe_ep_size=1) + assert (defaults.cp_size, defaults.dcp_size) == (1, None) + + +def test_base_model_does_not_claim_dcp_support_by_default(): + for backend in ("vllm", "sglang", "trtllm"): + assert BaseModel.supports_dcp(backend) is False + + +def test_persistent_kv_divisor_is_dcp_not_prefill_cp(): + shell = SimpleNamespace(config=SimpleNamespace(cp_size=8, dcp_size=1)) + assert BaseModel._cp_kv_memory_divisor(shell) == 1 + shell.config.dcp_size = 4 + assert BaseModel._cp_kv_memory_divisor(shell) == 4 + + +def test_get_model_gates_dcp_on_model_capability(): + from aisimulate_core.sdk.models import get_model + + # Dense GQA DCP is only modeled on vLLM (SGLang's GQA DCP lives in its + # Triton backend), so the sglang request must fail loud. + model_config = config.ModelConfig(tp_size=16, moe_tp_size=16, moe_ep_size=1, dcp_size=2) + with pytest.raises(NotImplementedError, match="dcp_size=2"): + get_model("meta-llama/Meta-Llama-3.1-70B", model_config, "sglang") + + +# -------------------------------------------------------------------------- +# Model-level rewrite: sharded decode attention + merge collectives +# -------------------------------------------------------------------------- + + +def _generation_op_names(model) -> list[str]: + return [op._name for op in model.generation_ops] + + +def _decode_attention_ops(model): + """``[(top-level container name, decode attention op)]`` incl. FallbackOp interiors.""" + # Core classes: FallbackOp interiors come back as bare core instances. + import aisimulate_core._native as core + + kinds = (core.GenerationAttention, core.GenerationMLA, core.WideEPGenerationMLA, core.GenerationDSAModule) + + def walk(container_name, op): + if isinstance(op, core.FallbackOp): + yield from walk(container_name, op._primary) + for inner in op._fallback: + yield from walk(container_name, inner) + elif isinstance(op, kinds) or (isinstance(op, core.MLAModule) and not op._is_context): + yield container_name, op + + return [pair for top in model.generation_ops if not top._name.startswith("draft_") for pair in walk(top._name, top)] + + +def _context_attention_ops(model): + import aisimulate_core._native as core + + kinds = (core.ContextAttention, core.ContextMLA, core.ContextDSAModule) + + def walk(op): + if isinstance(op, core.FallbackOp): + yield from walk(op._primary) + for inner in op._fallback: + yield from walk(inner) + elif isinstance(op, kinds) or (isinstance(op, core.MLAModule) and op._is_context): + yield op + + return [op for top in model.context_ops if not top._name.startswith("draft_") for op in walk(top)] + + +@pytest.mark.parametrize( + ("backend", "merge_suffixes"), + [ + ("sglang", ("_dcp_a2a_pack", "_dcp_out_all_to_all", "_dcp_a2a_combine")), + ("vllm", ("_dcp_lse_all_gather", "_dcp_lse_correct", "_dcp_out_reduce_scatter")), + ], +) +def test_deepseek_dcp_rewrites_decode_attention_and_adds_merge_collectives(backend, merge_suffixes): + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=8) + model = get_model("deepseek-ai/DeepSeek-V3", model_config, backend) + + attention = _decode_attention_ops(model) + assert attention, "DeepSeek decode graph must expose a decode attention op" + # Every decode attention op -- the fused MLAModule primary AND its granular + # GenerationMLA / GenerationAttention fallback -- is striped. + assert all(op._dcp_size == 8 for _, op in attention) + names = _generation_op_names(model) + for container, _ in {container: None for container, _ in attention}.items(): + index = names.index(container) + # The q gather, then the merge collectives and their elementwise + # passes, ride directly behind the block they serve. + expected = [f"{container}_dcp_q_all_gather", *(f"{container}{suffix}" for suffix in merge_suffixes)] + assert names[index + 1 : index + 1 + len(expected)] == expected + # The prefill graph gains no collectives of its own ... + assert not any("_dcp_" in op._name for op in model.context_ops) + # ... but its attention ops are marked so cached-context prefill pays the + # stripe gather (aggregated serving on the same engine). + context_attention = _context_attention_ops(model) + assert context_attention + assert all(op._dcp_size == 8 for op in context_attention) + + +def test_dcp_comm_override_selects_the_merge_collective(): + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, dcp_comm="a2a") + model = get_model("deepseek-ai/DeepSeek-V3", model_config, "vllm") + names = _generation_op_names(model) + assert any(name.endswith("_dcp_out_all_to_all") for name in names) + assert not any(name.endswith("_dcp_out_reduce_scatter") for name in names) + + +def test_dsa_dcp_rewrites_the_sparse_decode_module(): + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig( + tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, kvcache_quant_mode=common.KVCacheQuantMode.fp8 + ) + model = get_model("deepseek-ai/DeepSeek-V3.2", model_config, "vllm") + attention = _decode_attention_ops(model) + assert attention + assert all(op._dcp_size == 4 for _, op in attention) + names = _generation_op_names(model) + assert any(name.endswith("_dcp_q_all_gather") for name in names) + # The indexer's per-rank top-k candidates are gathered and merged across the + # DCP group before the sparse attention (vLLM `_merge_dcp_topk_global`). + assert any(name.endswith("_dcp_topk_all_gather") for name in names) + assert any(name.endswith("_dcp_topk_merge") for name in names) + + +def test_dsa_dcp_is_not_claimed_on_sglang(): + from aisimulate_core.sdk.models import get_model + + with pytest.raises(NotImplementedError, match="not supported for model_family='DEEPSEEKV32' on backend='sglang'"): + get_model( + "deepseek-ai/DeepSeek-V3.2", + config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4), + "sglang", + ) + + +@pytest.mark.parametrize( + ("overrides", "needle"), + [ + ({"dcp_comm": "a2a"}, "only runs with dcp_comm='ag_rs'"), + ({"kvcache_quant_mode": common.KVCacheQuantMode.bfloat16}, "requires an fp8 KV cache"), + ], +) +def test_dsa_dcp_on_vllm_flashmla_sparse_rejects_what_vllm_refuses(overrides, needle): + from aisimulate_core.sdk.models import get_model + + kwargs = dict(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, kvcache_quant_mode=common.KVCacheQuantMode.fp8) + kwargs.update(overrides) + with pytest.raises(ValueError, match=needle): + get_model("deepseek-ai/DeepSeek-V3.2", config.ModelConfig(**kwargs), "vllm") + # The FlashInfer sparse backend has no such restriction. + get_model( + "deepseek-ai/DeepSeek-V3.2", + config.ModelConfig(**kwargs, attention_backend="flashinfer"), + "vllm", + ) + + +def test_gqa_dcp_is_bounded_by_kv_head_replication(): + from aisimulate_core.sdk.models import get_model + + # Llama-3.1-70B: 8 kv heads. tp=16 replicates each kv head twice -> dcp<=2. + ok = get_model( + "meta-llama/Meta-Llama-3.1-70B", + config.ModelConfig(tp_size=16, moe_tp_size=16, moe_ep_size=1, dcp_size=2), + "vllm", + ) + assert all(op._dcp_size == 2 for _, op in _decode_attention_ops(ok)) + with pytest.raises(ValueError, match="KV-head replication"): + get_model( + "meta-llama/Meta-Llama-3.1-70B", + config.ModelConfig(tp_size=16, moe_tp_size=16, moe_ep_size=1, dcp_size=4), + "vllm", + ) + with pytest.raises(ValueError, match="must be positive and divide tp_size"): + get_model( + "meta-llama/Meta-Llama-3.1-70B", + config.ModelConfig(tp_size=16, moe_tp_size=16, moe_ep_size=1, dcp_size=3), + "vllm", + ) + + +def test_fpm_forward_model_prices_dcp_from_recorded_cells_on_vllm_only(): + from aisimulate_core.sdk.models import get_model + + # Whole-forward DCP timing comes from the recorded-DCP FPM cell identity + # (upstream #284), so the op-level rewrite and supports_dcp gate are skipped + # on that path; off vLLM there are no such cells. + model_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=8, forward_model="fpm") + model = get_model("deepseek-ai/DeepSeek-V3", model_config, "vllm") + assert not any("_dcp_" in op._name for op in model.generation_ops) + assert model._cp_kv_memory_divisor() == 8 + with pytest.raises(NotImplementedError, match="DCP timing requires measured vLLM FPM interpolation"): + get_model("deepseek-ai/DeepSeek-V3", model_config, "sglang") + + +def test_get_model_records_the_backend_for_families_that_do_not(): + # The DCP merge collective defaults per backend (a2a on sglang); dense + # families do not store backend_name themselves, so get_model must. + from aisimulate_core.sdk.models import get_model + + for backend in ("vllm", "sglang"): + model = get_model("meta-llama/Meta-Llama-3.1-70B", config.ModelConfig(tp_size=8), backend) + assert model._backend_name == backend + assert model._dcp_comm_style() == ("a2a" if backend == "sglang" else "ag_rs") + + +def test_dcp_one_leaves_the_decode_graph_untouched(): + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1) + model = get_model("deepseek-ai/DeepSeek-V3", model_config, "sglang") + assert all(op._dcp_size == 1 for _, op in _decode_attention_ops(model)) + assert not any("_dcp_" in name for name in _generation_op_names(model)) + + +def test_engine_identity_includes_dcp_size(): + from aisimulate_core.sdk.rust_engine_step import _engine_config_json + + def make(dcp): + cfg = SimpleNamespace( + tp_size=4, + pp_size=1, + moe_tp_size=1, + moe_ep_size=4, + attention_dp_size=1, + cp_size=None, + dcp_size=dcp, + gemm_quant_mode=SimpleNamespace(name="fp8", value=None), + moe_quant_mode=SimpleNamespace(name="nvfp4", value=None), + fmha_quant_mode=SimpleNamespace(name="bfloat16", value=None), + comm_quant_mode=SimpleNamespace(name="half", value=None), + kvcache_quant_mode=SimpleNamespace(name="fp8", value=None), + ) + model = SimpleNamespace( + config=cfg, + model_path="org/model-a", + architecture="X", + forward_model="op_level", + _nextn=None, + _nextn_accepted=None, + ) + database = SimpleNamespace(system="b200_sxm", backend="vllm", version="0.25.1", systems_root="/tmp/x") + return _engine_config_json(model, database) + + assert make(1) != make(8) + assert json.loads(make(8))["dcp_size"] == 8 + assert json.loads(make(None))["dcp_size"] is None + + +def test_fi_a2a_prices_as_the_packed_all_to_all(): + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, dcp_comm="fi_a2a") + names = _generation_op_names(get_model("deepseek-ai/DeepSeek-V3", model_config, "sglang")) + assert any(name.endswith("_dcp_out_all_to_all") for name in names) + assert not any(name.endswith("_dcp_out_reduce_scatter") for name in names) + + +def test_query_all_gather_follows_the_attention_dtype(): + from aisimulate_core.sdk.models import get_model + + def q_gather_mode(fmha): + model_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, fmha_quant_mode=fmha) + model = get_model("deepseek-ai/DeepSeek-V3", model_config, "vllm") + gathers = [op for op in model.generation_ops if op._name.endswith("_dcp_q_all_gather")] + assert gathers + return {str(op._comm_quant_mode).rsplit(".", 1)[-1] for op in gathers} + + # vLLM quantizes the MQA query to fp8 before the DCP gather when the + # attention runs in fp8; otherwise the buffer follows comm_quant_mode. + assert q_gather_mode(common.FMHAQuantMode.fp8) == {"fp8"} + assert q_gather_mode(common.FMHAQuantMode.bfloat16) == {"half"} + + +def test_query_replication_drops_the_gather_and_widens_the_q_projection(): + from aisimulate_core.sdk.models import get_model + + base_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, dcp_comm="a2a") + base = get_model("deepseek-ai/DeepSeek-V3", base_config, "vllm") + replicated = get_model( + "deepseek-ai/DeepSeek-V3", + config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, dcp_comm="a2a", dcp_q_replicate=True), + "vllm", + ) + assert any(name.endswith("_dcp_q_all_gather") for name in _generation_op_names(base)) + assert not any(name.endswith("_dcp_q_all_gather") for name in _generation_op_names(replicated)) + + def q_side_scale(model): + found = {} + + def walk(op): + import aisimulate_core._native as core + + if isinstance(op, core.FallbackOp): + walk(op._primary) + for inner in op._fallback: + walk(inner) + elif isinstance(op, core.OverlapOp): + for inner in (*op._group_a, *op._group_b): + walk(inner) + elif "q_b" in op._name or "bmm_pre" in op._name: + found[op._name] = float(op._scale_factor) + + for op in model.generation_ops: + walk(op) + return found + + base_scales, replicated_scales = q_side_scale(base), q_side_scale(replicated) + assert base_scales + assert set(replicated_scales) == set(base_scales) + for name, scale in base_scales.items(): + assert replicated_scales[name] == pytest.approx(scale * 4), name + + +def test_query_replication_requires_an_a2a_merge(): + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig(tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, dcp_q_replicate=True) + with pytest.raises(ValueError, match="dcp_q_replicate requires the a2a"): + get_model("deepseek-ai/DeepSeek-V3", model_config, "vllm") + + +def test_kimi_k3_dcp_is_modeled_on_sglang_with_replicated_q_by_default(): + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig(tp_size=8, attention_dp_size=1, moe_tp_size=8, moe_ep_size=1, dcp_size=4) + model = get_model("moonshotai/Kimi-K3", model_config, "sglang") + attention = _decode_attention_ops(model) + assert attention, "K3 exposes its MLA decode attention to the DCP rewrite" + assert all(op._dcp_size == 4 for _, op in attention) + names = _generation_op_names(model) + # SGLang's K3 override: a2a merge with dcp_replicate_q_proj=True, so no + # per-layer query all-gather; the KDA layers carry no KV and stay untouched. + assert any(name.endswith("_dcp_out_all_to_all") for name in names) + assert not any(name.endswith("_dcp_q_all_gather") for name in names) + assert not any("kda" in name and "_dcp_" in name for name in names) + with pytest.raises(NotImplementedError, match="not supported for model_family='KIMIK3' on backend='vllm'"): + get_model("moonshotai/Kimi-K3", model_config, "vllm") + + +def test_striped_context_dsa_module_survives_a_deepcopy(): + import copy + + from aisimulate_core.sdk.models import get_model + + model_config = config.ModelConfig( + tp_size=8, moe_tp_size=8, moe_ep_size=1, dcp_size=4, kvcache_quant_mode=common.KVCacheQuantMode.fp8 + ) + model = get_model("deepseek-ai/DeepSeek-V3.2", model_config, "vllm") + before = [op._dcp_size for op in _context_attention_ops(model)] + assert before and all(value == 4 for value in before) + # ProcessPoolExecutor / copy.deepcopy rebuild ops from __getnewargs_ex__; + # the striped-KV context gather must not silently drop to dcp=1. + copied = copy.deepcopy(model) + assert [op._dcp_size for op in _context_attention_ops(copied)] == before diff --git a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py index 92ace2255f..eddda2028b 100644 --- a/python/aisimulate/tests/unit/sdk/test_fpm_profile.py +++ b/python/aisimulate/tests/unit/sdk/test_fpm_profile.py @@ -422,16 +422,18 @@ def test_profile_engine_controls_fail_before_model_lookup(profile_dict, direct_c profile_memory(**controls) -@pytest.mark.parametrize("cp_size", [2, 0, True, 1.0, "1", None]) +# Prefill CP (cp_size > 1) is a supported knob at these entry points; only malformed +# sizes are rejected, and before any model construction. +@pytest.mark.parametrize("cp_size", [0, True, 1.0, "1", None]) @pytest.mark.parametrize("with_profile", [False, True]) def test_sdk_entry_points_reject_unsupported_cp_before_construction(profile_dict, monkeypatch, cp_size, with_profile): monkeypatch.setattr(engine, "build_model_config", _fail_graph) monkeypatch.setattr(memory.KVCacheEstimator, "from_request", _fail_graph) monkeypatch.setattr(memory.NaiveKVCacheEstimator, "from_model_path", _fail_graph) kwargs = {"cp_size": cp_size, "fpm_profile": profile_dict if with_profile else None} - with pytest.raises(ValueError, match="cp_size must be the integer 1"): + with pytest.raises(ValueError, match="cp_size must be a positive integer"): engine.compile_engine("test/unknown-decoder", "test_gpu", "vllm", "0.25.1", forward_model="fpm", **kwargs) - with pytest.raises(ValueError, match="cp_size must be the integer 1"): + with pytest.raises(ValueError, match="cp_size must be a positive integer"): memory.estimate_kv_cache( "test/unknown-decoder", "test_gpu", @@ -1053,7 +1055,8 @@ def test_profile_block_budget_preserves_resource_provenance(profile_dict, profil ({"max_num_tokens": True}, "positive integer"), ({"fmha_quant_mode": "bfloat16"}, "identity conflict"), ({"attention_backend": "different"}, "identity conflict"), - ({"cp_size": 2}, "cp_size must be the integer 1"), + # Prefill CP is a supported knob; a profile without that cell fails loud. + ({"cp_size": 2}, "no matching FPM deployment profile"), ({"nextn": 1}, "nextn must be 0"), ({"gpu_memory_capacity_bytes_override": 200}, "no KV budget"), ({"memory_fraction_kind": "of_free"}, "incompatible memory fraction"), diff --git a/python/aisimulate/tests/unit/sdk/test_memory_estimation.py b/python/aisimulate/tests/unit/sdk/test_memory_estimation.py index 41a6446505..ee08a40b58 100644 --- a/python/aisimulate/tests/unit/sdk/test_memory_estimation.py +++ b/python/aisimulate/tests/unit/sdk/test_memory_estimation.py @@ -903,6 +903,66 @@ def _boom(*args, **kwargs): assert out["tolerance_adjusted"] is None +@pytest.mark.parametrize("knob", ["cp_size", "dcp_size"]) +def test_estimate_kv_cache_rejects_malformed_context_parallel_sizes_before_any_build(monkeypatch, knob): + # Validation runs before the model build and the naive fallback: neither + # may turn True (int 1) into a cp=dcp=1 estimate. + def _never(*args, **kwargs): + raise AssertionError("model build must not be reached") + + monkeypatch.setattr(memory.KVCacheEstimator, "from_request", classmethod(_never)) + monkeypatch.setattr(memory.NaiveKVCacheEstimator, "from_model_path", classmethod(_never)) + with pytest.raises(ValueError, match=f"{knob} must be a positive integer"): + memory.estimate_kv_cache( + "foo/bar", + "h200_sxm", + "vllm", + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.9, + allow_naive_fallback=True, + **{knob: True}, + ) + + +def test_estimate_kv_cache_never_falls_back_after_a_context_parallel_capability_rejection(monkeypatch): + # get_model rejects dcp/cp>1 for families without the sharded path with + # NotImplementedError; the naive estimator knows neither knob, so falling + # back would silently return an unstriped capacity for a striped request. + def _unsupported(*args, **kwargs): + raise NotImplementedError("Decode context parallelism (dcp_size=2) is not supported") + + monkeypatch.setattr(memory.KVCacheEstimator, "from_request", classmethod(_unsupported)) + monkeypatch.setattr(memory.NaiveKVCacheEstimator, "_load_config", lambda *a, **k: dict(_RAW_UNSUPPORTED)) + with pytest.raises(ValueError, match="context parallelism is not supported"): + memory.estimate_kv_cache( + "foo/bar-unknown-arch", + "h200_sxm", + "vllm", + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.9, + gpu_memory_capacity_bytes_override=200 * _GIB, + dcp_size=2, + allow_naive_fallback=True, + ) + # Without a CP knob the same rejection keeps the documented naive fallback. + out = memory.estimate_kv_cache( + "foo/bar-unknown-arch", + "h200_sxm", + "vllm", + max_num_tokens=8192, + max_batch_size=256, + memory_fraction_kind="of_total", + memory_fraction_value=0.9, + gpu_memory_capacity_bytes_override=200 * _GIB, + allow_naive_fallback=True, + ) + assert out["source"] == "naive_fallback" + + def test_estimate_kv_cache_propagates_when_fallback_disabled(monkeypatch): def _boom(*args, **kwargs): raise RuntimeError("perf DB not available") 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 fe0bf74504..caf006e7ed 100644 --- a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py +++ b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py @@ -955,6 +955,26 @@ def best_available(config_json): assert not hasattr(rust_engine_step.RustForwardPassPerfModel, "from_regression") +def test_forward_pass_config_carries_context_parallel_knobs_through_rust_normalization() -> None: + import aisimulate_core + from aisimulate_core.sdk import ForwardPassPerfModelConfig + + plain = ForwardPassPerfModelConfig(model="m", system="s", backend="vllm", worker_type="decode") + assert (plain.cp_size, plain.dcp) == (None, None) + normalized = json.loads(aisimulate_core.RustForwardPassPerfModel.normalize_config(json.dumps(plain.to_dict()))) + # Unset prefill CP never enters the serialized identity (pre-CP configs stay + # byte-identical); unrecorded decode CP serializes as the canonical null. + assert "cp_size" not in normalized and normalized.get("dcp") is None + + striped = ForwardPassPerfModelConfig( + model="m", system="s", backend="vllm", worker_type="decode", tp=8, cp_size=2, dcp=4 + ) + payload = striped.to_dict() + assert (payload["cp_size"], payload["dcp"]) == (2, 4) + normalized = json.loads(aisimulate_core.RustForwardPassPerfModel.normalize_config(json.dumps(payload))) + assert (normalized["cp_size"], normalized["dcp"]) == (2, 4) + + def test_forward_pass_config_requires_role_and_defaults_to_auto_deny() -> None: from aisimulate_core.sdk import ForwardPassPerfModelConfig @@ -1020,6 +1040,7 @@ def test_forward_pass_config_preserves_existing_positional_arguments(tmp_path: P "dcp": None, "fpm_fmha_quant_mode": None, "moe_kernel_source": None, + "cp_size": None, } source = " source_with_spaces " @@ -1030,6 +1051,7 @@ def test_forward_pass_config_preserves_existing_positional_arguments(tmp_path: P "dcp": None, "fpm_fmha_quant_mode": None, "moe_kernel_source": source, + "cp_size": None, } assert json.loads(json.dumps(pinned.to_dict()))["moe_kernel_source"] == source @@ -2339,3 +2361,36 @@ def test_canonical_native_selection_retries_roots_and_pins_effective_configurati assert resolved["transfer_policy"] == ["xshape"] assert resolved["estimator_config"]["correction"]["enabled"] is False assert config.systems_paths == (str(tmp_path), packaged) + + +@pytest.mark.parametrize( + ("variant_a", "variant_b"), + [ + ({"dcp_comm": "ag_rs"}, {"dcp_comm": "a2a"}), + ({"dcp_comm": "a2a", "dcp_q_replicate": False}, {"dcp_comm": "a2a", "dcp_q_replicate": True}), + ({}, {"dcp_comm": "a2a"}), + ], +) +def test_engine_config_json_separates_dcp_op_shaping_overrides(variant_a, variant_b) -> None: + """``dcp_comm`` / ``dcp_q_replicate`` select different DCP op graphs for one + (tp, dcp) identity; the cache key must not let a warm ``ag_rs`` handle answer + an ``a2a`` request (or vice versa, whichever compiled first).""" + + def _model(**dcp_overrides): + return SimpleNamespace( + model_path="deepseek-ai/DeepSeek-V3", + architecture="DeepseekV3ForCausalLM", + config=ModelConfig( + tp_size=8, pp_size=1, attention_dp_size=1, moe_tp_size=8, moe_ep_size=1, dcp_size=4, **dcp_overrides + ), + ) + + database = SimpleNamespace(system="b200_sxm", backend="vllm", version="0.24.0") + key_a = rust_engine_step._engine_config_json(_model(**variant_a), database) + key_b = rust_engine_step._engine_config_json(_model(**variant_b), database) + assert key_a != key_b + identity = json.loads(json.loads(key_b)["extra"]["identity"])["model_config"] + assert identity["dcp_comm"] == variant_b.get("dcp_comm") + assert identity["dcp_q_replicate"] == variant_b.get("dcp_q_replicate") + # Same overrides -> same key (the memo still hits). + assert rust_engine_step._engine_config_json(_model(**variant_b), database) == key_b diff --git a/python/aisimulate/tests/unit/test_support_plan.py b/python/aisimulate/tests/unit/test_support_plan.py index 57b527c359..4abf3656f8 100644 --- a/python/aisimulate/tests/unit/test_support_plan.py +++ b/python/aisimulate/tests/unit/test_support_plan.py @@ -202,6 +202,7 @@ def no_model_or_collector_import(name, *args, **kwargs): "attention_data": 1, "moe_tensor": 4, "moe_expert": 1, + "prefill_context": None, "decode_context": None, } assert recommendation.engine.workers.aggregated.timing.estimation_mode == "fpm_interpolation" diff --git a/tests/sweeper/test_forward_pass_estimator.py b/tests/sweeper/test_forward_pass_estimator.py index cd94226e07..762a4a3448 100644 --- a/tests/sweeper/test_forward_pass_estimator.py +++ b/tests/sweeper/test_forward_pass_estimator.py @@ -518,6 +518,62 @@ def close(self): assert (seen[-1].tp, seen[-1].kv_block_size, seen[-1].backend_version) == (4, 1, None) +def test_resolver_carries_context_parallel_columns_into_the_estimator_identity(monkeypatch): + # A dcp=8 candidate must not resolve to (or share a cache entry with) the + # dcp=1 estimator: the native engine is built from this request. + import aisimulate_core + from aisimulate.sweeper.config import SearchSpace + from aisimulate.sweeper.forward_pass_estimator import ForwardPassEstimatorResolver + + seen = [] + + class Model: + def __init__(self, config): + seen.append(config) + self.config = json.loads( + aisimulate_core.RustForwardPassPerfModel.normalize_config(json.dumps(config.to_dict())) + ) + self.config["backend_version"] = "0.24.0" + self.config["estimation_mode"] = "op_level" + + def diagnostics(self): + return { + "readiness": "ready", + "provenance": {"config": self.config, "selected_systems_root": self.config["systems_paths"][0]}, + } + + def close(self): + pass + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", Model) + space = SearchSpace(model_name="deepseek-ai/DeepSeek-V3", hardware_sku="h200_sxm", backend=["vllm"]) + resolver = ForwardPassEstimatorResolver(space) + sample = { + "deployment_mode": "agg", + "model_name": space.model_name, + "hardware_sku": space.hardware_sku, + "backend": "vllm", + "tp": 8, + "pp": 1, + "attention_dp": 1, + "moe_tp": 1, + "moe_ep": 8, + "agg_block_size": 64, + } + plain = resolver.resolve_candidate(sample)["agg"] + striped = resolver.resolve_candidate({**sample, "dcp": 8})["agg"] + assert len(seen) == 2 + assert (seen[0].cp_size, seen[0].dcp) == (None, None) + assert (seen[1].cp_size, seen[1].dcp) == (None, 8) + # Unit knobs stay out of the serialized identity; a set knob is carried. + assert plain.config.get("dcp") is None and "cp_size" not in plain.config + assert striped.config["dcp"] == 8 + + # Prefill CP folds into the width identity; the sample carries the wider MoE side. + resolver.resolve_candidate({**sample, "tp": 1, "moe_ep": 8, "cp": 8}) + assert (seen[-1].tp, seen[-1].cp_size, seen[-1].dcp) == (1, 8, None) + + def test_search_rejects_unknown_controls_and_policies_on_custom_timing(): from aisimulate.sweeper.config import SearchSpace diff --git a/tests/test_aic.py b/tests/test_aic.py index fcb7891b7f..9bd03420de 100644 --- a/tests/test_aic.py +++ b/tests/test_aic.py @@ -163,6 +163,52 @@ def estimate(**kwargs): assert calls[0]["free_gpu_memory_fraction"] == 0.0 +def test_materialize_forwards_context_parallel_knobs(monkeypatch) -> None: + calls = [] + + def estimate(**kwargs): + calls.append(kwargs) + return 1 + + monkeypatch.setattr(aic, "estimate_num_gpu_blocks", estimate) + aic.materialize_aic_num_gpu_blocks( + { + "engine_type": "vllm", + "aic_backend": "vllm", + "aic_system": "h200_sxm", + "aic_model_path": "test-model", + "aic_cp_size": 2, + "aic_dcp_size": 4, + "block_size": 64, + } + ) + + assert calls[0]["cp_size"] == 2 + assert calls[0]["dcp_size"] == 4 + + +def test_materialize_defaults_context_parallel_knobs_to_one(monkeypatch) -> None: + calls = [] + + def estimate(**kwargs): + calls.append(kwargs) + return 1 + + monkeypatch.setattr(aic, "estimate_num_gpu_blocks", estimate) + aic.materialize_aic_num_gpu_blocks( + { + "engine_type": "vllm", + "aic_backend": "vllm", + "aic_system": "h200_sxm", + "aic_model_path": "test-model", + "block_size": 64, + } + ) + + assert calls[0]["cp_size"] == 1 + assert calls[0]["dcp_size"] == 1 + + @pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) @pytest.mark.parametrize("timing", ["default", "fixed", "polynomial"]) def test_prediction_omitted_version_uses_current_database_for_capacity(monkeypatch, backend, timing): diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index 4696b15731..1238e04d83 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -311,6 +311,25 @@ def test_recommendation_accepts_domains_and_parallel_preset() -> None: assert isinstance(config.engine, EngineRecommendationConfig) +def test_parallel_preset_entries_round_trip_with_context_knobs() -> None: + # A preset entry dumped back to a mapping carries prefill_context unset + # (decode CP lives on the prediction config only); recommend must accept the + # knobs (unset or at 1) so a config can be re-validated from model_dump, + # while unrelated keys stay rejected. + from aisimulate.config.engine import ParallelismRecommendationConfig + + entry = {"replicas": 1, "tensor": 2, "pipeline": 1, "attention_data": 1, "moe_tensor": 1, "moe_expert": 1} + parallelism = ParallelismRecommendationConfig.model_validate({"preset": [entry]}) + dumped = parallelism.model_dump()["preset"][0] + assert dumped["prefill_context"] is None and "decode_context" not in dumped + + again = ParallelismRecommendationConfig.model_validate({"preset": [dumped]}) + assert again.preset[0].tensor == 2 + + with pytest.raises(ValidationError, match="unknown=\\['mystery'\\]"): + ParallelismRecommendationConfig.model_validate({"preset": [{**entry, "mystery": 1}]}) + + def test_parallel_preset_modes_lower_without_conflating_semantics() -> None: independent = { "preset": False, @@ -1363,6 +1382,178 @@ def test_pd_predict_checks_effective_prefill_hardware_in_router_hook(router_hard prediction_to_replay_spec(config, adapter_specs={"dynamo.router": spec}) +def test_prediction_parallelism_accepts_context_parallel_knobs() -> None: + config = CorePredictionConfig.model_validate( + {"engine": {**_engine(), "workers": {"aggregated": {"parallelism": {"decode_context": 4}}}}} + ) + assert config.engine.workers.aggregated is not None + parallel = config.engine.workers.aggregated.parallelism + assert parallel.prefill_context is None + assert parallel.decode_context == 4 + + +@pytest.mark.parametrize("field", ["prefill_context", "decode_context"]) +def test_prediction_rejects_nonpositive_context_parallel(field: str) -> None: + with pytest.raises(ValidationError): + CorePredictionConfig.model_validate( + {"engine": {**_engine(), "workers": {"aggregated": {"parallelism": {field: 0}}}}} + ) + + +def test_aggregated_deployment_rejects_prefill_and_decode_cp_together() -> None: + from aisimulate.compiler import _deployment + + config = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(), + "backend": "sglang", + "workers": {"aggregated": {"parallelism": {"prefill_context": 2, "decode_context": 2}}}, + } + } + ) + with pytest.raises(ValueError, match="at most one of parallelism.prefill_context"): + _deployment(config.engine, workload={}, afd_performance_model=None) + + +def test_aggregated_deployment_accepts_one_context_parallel_knob() -> None: + from aisimulate.compiler import _deployment + + config = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(), + "backend": "sglang", + "workers": {"aggregated": {"parallelism": {"tensor": 4, "decode_context": 4}}}, + } + } + ) + spec = _deployment(config.engine, workload={}, afd_performance_model=None) + # cp=1 is not spelled out, so pre-CP deployments stay byte-identical. + assert "cp" not in spec.parallel_config + assert spec.parallel_config["dcp"] == 4 + # Default timing lowers the canonical estimator config onto the rank; the + # knob must ride that identity (the native engine is built from it). + timing = spec.agg_engine_args["timing_model"]["config"] + assert timing["dcp"] == 4 + assert timing["cp_size"] is None + assert spec.performance_model_metadata["aggregated"]["config"]["dcp"] == 4 + assert "cp_size" not in spec.performance_model_metadata["aggregated"]["config"] + + +@pytest.mark.parametrize("knob", ["prefill_context", "decode_context"]) +def test_afd_companion_rejects_context_parallelism(knob: str) -> None: + # The companion's ParallelShape / GPU accounting carry tp/pp/dp/moe only, so + # a CP knob would price a wider worker than the topology reports. + from pathlib import Path + + from aisimulate.compiler import prediction_to_replay_spec + + tiny_model = Path(__file__).resolve().parents[1] / "tests/e2e/configs/unified_cli/fixtures/tiny-model" + config = CorePredictionConfig.model_validate( + { + "engine": { + "mode": "afd", + "model": str(tiny_model), + "hardware": "h200_sxm", + "backend": "vllm", + "context_length": 2048, + "afd": { + "phase": "decode", + "combined_with_pd": True, + "n_a_nodes": 1, + "n_f_nodes": 1, + "tp_a": 8, + "a_batch_size": 8, + }, + "workers": {"prefill": {"parallelism": {knob: 2}}}, + } + } + ) + with pytest.raises(ValueError, match="AFD companion .* do not support context parallelism"): + prediction_to_replay_spec(config) + + +def test_disaggregated_deployment_carries_each_context_parallel_knob_per_role() -> None: + from aisimulate.compiler import _deployment + + config = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(), + "mode": "disaggregated", + "backend": "sglang", + "workers": { + "prefill": {"parallelism": {"prefill_context": 2}}, + "decode": {"parallelism": {"tensor": 4, "decode_context": 4}}, + }, + } + } + ) + spec = _deployment(config.engine, workload={}, afd_performance_model=None) + assert spec.parallel_config["prefill_cp"] == 2 + assert "prefill_dcp" not in spec.parallel_config + assert "decode_cp" not in spec.parallel_config + assert spec.parallel_config["decode_dcp"] == 4 + prefill_timing = spec.prefill_engine_args["timing_model"]["config"] + decode_timing = spec.decode_engine_args["timing_model"]["config"] + assert (prefill_timing["cp_size"], prefill_timing.get("dcp")) == (2, None) + assert (decode_timing["cp_size"], decode_timing["dcp"]) == (None, 4) + assert spec.performance_model_metadata["prefill"]["config"]["cp_size"] == 2 + assert spec.performance_model_metadata["decode"]["config"]["dcp"] == 4 + + +@pytest.mark.parametrize( + ("backend", "prefill", "decode", "needle"), + [ + # A prefill engine stripes its KV only as the PCP+DCP layout. + ("sglang", {"decode_context": 4}, {}, "only as 1 or equal to prefill_context"), + ("vllm", {"prefill_context": 2, "decode_context": 4}, {"decode_context": 4}, "only as 1 or equal"), + # vLLM NIXL: replicated PCP cannot feed a DCP-sharded decode ... + ("vllm", {"prefill_context": 2}, {"decode_context": 4}, "replicated-PCP"), + # ... and the two DCP sizes must divide one another. + ("vllm", {"prefill_context": 4, "decode_context": 4}, {"decode_context": 6}, "divide one another"), + ], +) +def test_disaggregated_context_parallel_layout_rules(backend: str, prefill: dict, decode: dict, needle: str) -> None: + from aisimulate.compiler import _deployment + + config = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(), + "mode": "disaggregated", + "backend": backend, + "workers": {"prefill": {"parallelism": prefill}, "decode": {"parallelism": decode}}, + } + } + ) + with pytest.raises(ValueError, match=needle): + _deployment(config.engine, workload={}, afd_performance_model=None) + + +def test_vllm_disaggregated_pcp_plus_dcp_prefill_pairs_with_a_dividing_decode_dcp() -> None: + from aisimulate.compiler import _deployment + + config = CorePredictionConfig.model_validate( + { + "engine": { + **_engine(), + "mode": "disaggregated", + "backend": "vllm", + "context_length": 4096, + "workers": { + "prefill": {"parallelism": {"tensor": 2, "prefill_context": 2, "decode_context": 2}}, + "decode": {"parallelism": {"tensor": 8, "decode_context": 8}}, + }, + } + } + ) + spec = _deployment(config.engine, workload={}, afd_performance_model=None) + assert (spec.parallel_config["prefill_cp"], spec.parallel_config["prefill_dcp"]) == (2, 2) + assert spec.parallel_config["decode_dcp"] == 8 + + @pytest.mark.parametrize("seed", [0, 42, 2**64 - 1]) @pytest.mark.parametrize("warmup", [False, True]) def test_agentic_snapshot_seed_compiles_for_prediction_and_recommendation(seed: int, warmup: bool) -> None: @@ -1599,13 +1790,11 @@ def test_prediction_dcp_identity_roundtrip_and_explicit_capacity(): assert timing["estimation_mode"] == "fpm_interpolation" assert deployment.agg_engine_args["num_gpu_blocks"] == 2175 assert deployment.parallel_config["dcp"] == 8 - raw["engine"]["workers"]["aggregated"]["kv_cache"]["capacity"] = {"type": "default"} - with pytest.raises(ValueError, match="explicit KV block capacity"): - prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)) + # Automatic KV capacity and transfer sizing stay available under DCP: the + # capacity estimate prices the 1/dcp KV stripe and the per-token transfer + # bytes are unchanged by striping. cache = raw["engine"]["workers"]["aggregated"]["kv_cache"] cache.update(capacity={"type": "fixed", "blocks": 2175}, host_offload={"num_host_blocks": 32}) - with pytest.raises(ValueError, match="explicit bytes_per_token"): - prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)) cache["bytes_per_token"] = 1728 deployment = prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)).backend_deployment assert deployment.agg_engine_args["kv_cache_bytes_per_token"] == 1728 diff --git a/tests/test_runner.py b/tests/test_runner.py index 663721a96e..3424b64879 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1258,6 +1258,62 @@ def test_runner_rejects_conflicting_backend_version_in_explicit_aic_timing(): EngineReplayRunnerFactory(runtime=RecordingRuntime()).create(0).run(_spec(deployment=deployment)) +@pytest.mark.parametrize(("field", "alias"), [("dcp", "aic_dcp_size"), ("cp", "aic_cp_size")]) +def test_runner_rejects_parallel_config_that_conflicts_with_context_parallel_args(field, alias): + # Same contract as tp / attention_dp: capacity and timing must not price one + # CP topology while parallel_config reports another. + engine_args = _engine_args() + engine_args[alias] = 4 + deployment = BackendDeploymentSpec( + deployment_mode="agg", + backend="vllm", + backend_version="test", + parallel_config={"tp": 2, "attention_dp": 1, field: 2, "replicas": 1}, + agg_engine_args=engine_args, + num_workers=1, + ) + + with pytest.raises(ValueError, match=f"parallel_config.{field}=2 conflicts"): + EngineReplayRunnerFactory(runtime=RecordingRuntime()).create(0).run(_spec(deployment=deployment)) + + +@pytest.mark.parametrize("nested_rank", [False, True]) +@pytest.mark.parametrize(("field", "target"), [("dcp", "dcp"), ("cp", "cp_size")]) +def test_runner_rejects_nested_timing_context_parallel_that_conflicts_with_parallel_config(field, target, nested_rank): + # The canonical timing config nests the knobs under timing_model.config + # (flat form or under the execution-level `rank` descriptor); they are + # checked against parallel_config BEFORE capacity materialization could + # resolve them into AIC inputs. + timing = { + "type": "external", + "provider": "aic", + "config": { + "model": "test-model", + "system": "test-system", + "backend": "vllm", + "worker_type": "aggregated", + "tp": 2, + "attention_dp": 1, + target: 2, + }, + } + engine_args = _engine_args(timing=timing) + engine_args.pop("num_gpu_blocks") + if nested_rank: + engine_args = {"rank": engine_args} + deployment = BackendDeploymentSpec( + deployment_mode="agg", + backend="vllm", + backend_version="test", + parallel_config={"tp": 2, "attention_dp": 1, field: 4, "replicas": 1}, + agg_engine_args=engine_args, + num_workers=1, + ) + + with pytest.raises(ValueError, match=f"parallel_config.{field}=4 conflicts"): + EngineReplayRunnerFactory(runtime=RecordingRuntime()).create(0).run(_spec(deployment=deployment)) + + def test_runner_rejects_parallel_config_that_conflicts_with_engine_args(): deployment = BackendDeploymentSpec( deployment_mode="agg", @@ -1608,6 +1664,108 @@ def test_memory_detail_with_explicit_blocks_does_not_guess_components(): assert "memory_breakdown" not in data +def test_runner_forwards_decode_cp_to_aic_timing_and_capacity(monkeypatch): + runtime = RecordingRuntime() + engine_args = _engine_args() + engine_args.pop("num_gpu_blocks") + engine_args.pop("timing_model") + engine_args["aic_backend_version"] = "test" + engine_args["aic_dcp_size"] = 4 + engine_args["gpu_memory_utilization"] = 0.8 + calls = [] + + def estimate(**kwargs): + calls.append(kwargs) + return 321 + + monkeypatch.setattr(aic, "estimate_num_gpu_blocks", estimate) + deployment = BackendDeploymentSpec( + deployment_mode="agg", + backend="vllm", + backend_version="test", + agg_engine_args=engine_args, + num_workers=1, + ) + + EngineReplayRunnerFactory(runtime=runtime).create(0).run(_spec(deployment=deployment)) + + timing_config = runtime.execution_spec["engine"]["rank"]["timing_model"]["config"] + assert timing_config["dcp"] == 4 + assert "cp_size" not in timing_config + assert calls[0]["dcp_size"] == 4 + assert calls[0]["cp_size"] == 1 + + +def test_public_decode_cp_reaches_canonical_timing_and_capacity(tmp_path, monkeypatch): + # Default timing lowers the canonical ForwardPassPerfModelConfig onto the + # rank and strips the legacy aic_* args, so the knob must survive that path: + # into the estimator identity (native engine build) AND into the KV capacity. + path = tmp_path / "prediction.yaml" + path.write_text( + """\ +engine: + mode: aggregated + model: example/model + hardware: h200_sxm + backend: vllm + context_length: 4096 + workers: + aggregated: + parallelism: + tensor: 8 + decode_context: 4 + kv_cache: + block_size: 16 + capacity: + type: default + memory_fraction: 0.8 +""", + encoding="utf-8", + ) + calls = [] + + def estimate(**kwargs): + calls.append(kwargs) + return 321 + + monkeypatch.setattr(aic, "estimate_num_gpu_blocks", estimate) + import aisimulate_core + from aisimulate_core.sdk import RustForwardPassPerfModel + + requested = [] + + class Estimator: + def __init__(self, config): + requested.append(config) + self.config = json.loads(aisimulate_core.RustForwardPassPerfModel.normalize_config(json.dumps(config))) + self.config.update(backend_version="0.24.0", estimation_mode="op_level", fallback_policy="deny") + + def diagnostics(self): + return {"readiness": "ready", "provenance": {"config": self.config}} + + def close(self): + pass + + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", Estimator) + public = CorePredictionConfig.from_yaml(path) + spec = prediction_to_replay_spec(public) + assert spec.backend_deployment.agg_engine_args["timing_model"]["config"]["dcp"] == 4 + assert "aic_dcp_size" not in spec.backend_deployment.agg_engine_args + + runtime = RecordingRuntime() + EngineReplayRunnerFactory(runtime=runtime).create(0).run(spec) + + # The Python request carries the dataclass fields (unset -> None); the Rust + # provenance drops unset knobs, which is what lands on the rank below. + assert requested and requested[0]["dcp"] == 4 and requested[0]["cp_size"] is None + rank = runtime.execution_spec["spec"]["engine"]["rank"] + assert rank["timing_model"]["config"]["dcp"] == 4 + assert "cp_size" not in rank["timing_model"]["config"] + assert rank["num_gpu_blocks"] == 321 + assert calls[0]["dcp_size"] == 4 + assert calls[0]["cp_size"] == 1 + + def test_native_report_memory_error_becomes_host_resource_failure(monkeypatch): import sys from types import ModuleType