feat(core): op-level decode CP and prefill CP on top of the recorded DCP FPM profiles

crates/core/src/perfmodel/config.rs

@@ -119,7 +119,13 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1;
 //   JSON defaults preserve unrecorded DCP.
 // - 24 (typed FPM DCP): FpmForwardOp carries dcp_size separately from the
 //   base matching identity, so control flow never parses the string tuple.
-pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 24;
+// - 25 (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 22 on its own branch while the DeepSeek-V4.1, MoE
+//   kernel-source, VR200 pilot and FPM DCP changes landed; renumbered at each
+//   merge (precedent: 15, 18).
+pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 25;
 
 /// Static engine identity and setup information carried by an
 /// [`crate::perfmodel::engine::spec::EngineSpec`].
@@ -266,6 +272,12 @@ pub struct ParallelMapping {
     /// re-derived from this field.
     #[serde(default)]
     pub cp_size: Option<u32>,
+    /// 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<u32>,
 }

crates/core/src/perfmodel/engine/runtime.rs

@@ -2440,6 +2440,7 @@ mod tests {
                 cp_size: 1,
                 lane_order: crate::operators::attention::b200_vllm_context_lane_order(),
                 apply_rope: true,
+                dcp_size: 1,
             }),
         ]
     }
@@ -2465,6 +2466,7 @@ mod tests {
                 use_qk_norm: false,
                 scale_num_tokens: 1,
                 verify_query_tokens: 0,
+                dcp_size: 1,
             }),
         ]
     }
@@ -2485,13 +2487,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,
@@ -2556,13 +2558,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),

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,
         }
     }
 
@@ -851,13 +858,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),
@@ -1047,7 +1054,7 @@ mod tests {
             vec![OpSpec::FpmForward(fpm_forward())],
             vec![OpSpec::Moe(moe())],
         );
-        assert_eq!(spec.schema_version, 24);
+        assert_eq!(spec.schema_version, ENGINE_SPEC_SCHEMA_VERSION);
         let mut bytes = spec.to_bincode().unwrap();
         assert_eq!(EngineSpec::from_bincode(&bytes).unwrap(), spec);
 
@@ -1328,21 +1335,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-v25 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,
                 ..
             })
@@ -1380,7 +1390,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] {
+            for previous_version in [20u32, 21, 22, 23, 24] {
                 let mut stale = bytes.clone();
                 stale[..4].copy_from_slice(&previous_version.to_le_bytes());
                 // DCP and the pilot claimed 22 for different layouts.
@@ -1391,7 +1401,7 @@ mod tests {
                         Err(AicError::UnsupportedSchemaVersion {
                             kind: "EngineSpec",
                             got,
-                            expected: 24,
+                            expected: ENGINE_SPEC_SCHEMA_VERSION,
                         }) if got == previous_version
                     ));
                 }

crates/core/src/perfmodel/fpm/config.rs

@@ -101,6 +101,11 @@ pub struct ForwardPassPerfModelConfig {
     pub moe_tp_size: Option<u32>,
     #[serde(default)]
     pub moe_ep_size: Option<u32>,
+    /// 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<u32>,
 
     #[serde(default, alias = "gemm_dtype")]
     pub gemm_quant_mode: Option<String>,
@@ -177,6 +182,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,
@@ -281,6 +287,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"));
         }
@@ -303,11 +314,15 @@ impl ForwardPassPerfModelConfig {
                     "moe_tp_size and moe_ep_size must be positive",
                 ));
             }
-            if u64::from(self.tp) * u64::from(self.attention_dp)
+            // 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`).
+            let cp = u64::from(self.cp_size.unwrap_or(1));
+            if 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",
                 ));
             }
         }
@@ -647,6 +662,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::<ForwardPassPerfModelConfig>(&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::<ForwardPassPerfModelConfig>(&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!(

crates/core/src/perfmodel/fpm/model.rs

@@ -331,9 +331,13 @@ impl ForwardPassPerfModel {
         let mut failures = Vec::new();
         let mut last_error = None;
         for mode in config.candidate_modes() {
+            // 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(),

crates/core/src/perfmodel/fpm/tests.rs

@@ -89,6 +89,7 @@ fn context_ops() -> Vec<Op> {
             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<Op> {
             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,

crates/core/src/perfmodel/memory.rs

@@ -515,13 +515,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,