feat: import vLLM 0.33.0 hybrid cache groups through the runtime probe

crates/core/src/engine/kv_manager/grouped.rs

@@ -85,6 +85,13 @@ impl GroupedKvPool {
             .map(|lease| lease.bytes)
     }
 
+    /// Bytes that the state groups hold for one running request, whatever
+    /// its context length: the state charge computed when the pool is built,
+    /// or `u64::MAX` when those bytes overflow.
+    pub(crate) fn fixed_request_bytes(&self) -> u64 {
+        self.state.as_ref().map_or(u64::MAX, |state| state.bytes)
+    }
+
     /// Reserve every group atomically. Expired blocks become reusable in this
     /// transition, but failed attempts leave both the lease and pool unchanged.
     pub(crate) fn allocate(
@@ -144,6 +151,8 @@ mod tests {
                     page_size_bytes: 8,
                     sliding_window: None,
                     blocks_per_request: None,
+                    prefill_checkpoint_blocks: None,
+                    prefill_checkpoint_alignment: None,
                 },
                 FpmCacheGroup {
                     name: "window".into(),
@@ -153,10 +162,13 @@ mod tests {
                     page_size_bytes: 6,
                     sliding_window: Some(3),
                     blocks_per_request: None,
+                    prefill_checkpoint_blocks: None,
+                    prefill_checkpoint_alignment: None,
                 },
             ],
             42,
         );
+        assert_eq!(pool.fixed_request_bytes(), 0);
         let mut first = None;
         let mut second = None;
         assert!(pool.allocate(&mut first, 0, 4)); // 8 + 2 * 6 = 20
@@ -189,6 +201,8 @@ mod tests {
                     page_size_bytes: 8,
                     sliding_window: None,
                     blocks_per_request: None,
+                    prefill_checkpoint_blocks: None,
+                    prefill_checkpoint_alignment: None,
                 },
                 FpmCacheGroup {
                     name: "kda".into(),
@@ -198,10 +212,14 @@ mod tests {
                     page_size_bytes: 10,
                     sliding_window: None,
                     blocks_per_request: Some(2),
+                    prefill_checkpoint_blocks: None,
+                    prefill_checkpoint_alignment: None,
                 },
             ],
             100,
         );
+        // Only the state group counts: 2 blocks of 10 bytes.
+        assert_eq!(pool.fixed_request_bytes(), 20);
         let mut lease = None;
         assert!(pool.allocate(&mut lease, 0, 4)); // 8 + 2 * 10 = 28
         assert_eq!(lease.as_ref().unwrap().blocks(), 3); // 0..1 and the state's 0..2
@@ -225,6 +243,8 @@ mod tests {
             page_size_bytes,
             sliding_window: None,
             blocks_per_request: Some(blocks),
+            prefill_checkpoint_blocks: None,
+            prefill_checkpoint_alignment: None,
         };
         let mut pool = GroupedKvPool::new(vec![state("kda", 10, 2), state("tail", 6, 1)], 100);
         assert_eq!(pool.required_bytes(0, 0), Some(0));

crates/core/src/engine/scheduler/vllm/core.rs

@@ -2488,6 +2488,18 @@ impl VllmCore {
                             .map(|request| request.sequence.num_input_tokens()),
                         max_model_len = self.args.max_model_len,
                         num_gpu_blocks = self.args.num_gpu_blocks,
+                        grouped_minimum_forward_bytes = self
+                            .kv_manager
+                            .grouped()
+                            .zip(self.state.requests.get(&uuid))
+                            .and_then(|(pool, request)| {
+                                let known_tokens = request.sequence.len();
+                                pool.required_bytes(known_tokens.saturating_sub(1), known_tokens)
+                            }),
+                        grouped_capacity_bytes = self
+                            .kv_manager
+                            .grouped()
+                            .map(GroupedKvPool::capacity_bytes),
                         "rejecting request that exceeds a worker admission limit"
                     );
                     rejected_uuids.push(uuid);
@@ -3186,8 +3198,14 @@ impl VllmCore {
                 .ok_or_else(|| anyhow::anyhow!("grouped KV forward byte calculation overflowed"))?;
             anyhow::ensure!(
                 required <= pool.capacity_bytes(),
-                "request {uuid} grouped KV forward needs {required} bytes for {desired_tokens} scheduled tokens after {effective_computed_before} computed tokens, exceeding the entire {}-byte cache pool; reduce max_num_batched_tokens or increase cache memory",
-                pool.capacity_bytes()
+                "request {uuid} grouped KV forward needs {required} bytes for {desired_tokens} scheduled tokens after {effective_computed_before} computed tokens, exceeding the entire {}-byte cache pool{}; reduce max_num_batched_tokens or increase cache memory",
+                pool.capacity_bytes(),
+                match pool.fixed_request_bytes() {
+                    0 => String::new(),
+                    fixed => format!(
+                        "; the fixed state blocks of one request take {fixed} bytes of this pool"
+                    ),
+                }
             );
         }
         let mut actual_computed_after = desired_computed_after;

crates/core/src/engine/scheduler/vllm/grouped_tests.rs

@@ -44,6 +44,8 @@ fn group(name: &str, block: u32, bytes: u64, window: Option<u32>) -> FpmCacheGro
         page_size_bytes: bytes,
         sliding_window: window,
         blocks_per_request: None,
+        prefill_checkpoint_blocks: None,
+        prefill_checkpoint_alignment: None,
     }
 }
 
@@ -56,6 +58,8 @@ fn state(name: &str, bytes: u64, blocks: u32) -> FpmCacheGroup {
         page_size_bytes: bytes,
         sliding_window: None,
         blocks_per_request: Some(blocks),
+        prefill_checkpoint_blocks: None,
+        prefill_checkpoint_alignment: None,
     }
 }
 
@@ -241,6 +245,8 @@ fn grouped_oversized_chunk_fails_without_partial_allocation_or_endless_preemptio
             .contains("exceeding the entire 8-byte cache pool")
     );
     assert!(error.to_string().contains("reduce max_num_batched_tokens"));
+    // Without a state group the message has no sentence about fixed bytes.
+    assert!(!error.to_string().contains("fixed state blocks"));
     assert_eq!(core.state.requests[&id].num_computed_tokens, 8);
     assert_eq!(core.state.requests[&id].sequence.num_allocated_tokens(), 8);
     assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(8));
@@ -250,6 +256,47 @@ fn grouped_oversized_chunk_fails_without_partial_allocation_or_endless_preemptio
     assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0));
 }
 
+#[test]
+fn grouped_state_blocks_exceeding_the_pool_reject_at_admission() {
+    // The 16-byte state block alone is larger than the 12-byte pool, so even
+    // the one-token forward cannot fit and the request never reaches a forward.
+    let mut core = core(vec![group("full", 4, 4, None), state("kda", 16, 1)], 12, 8);
+    let id = submit(&mut core, 115, 4, 2);
+    let mut trace = TraceCollector::default();
+    let pass = core.execute_pass(&mut trace, 0.0);
+    assert!(pass.admissions.is_empty());
+    assert!(
+        pass.output_signals
+            .iter()
+            .any(|output| output.uuid == id && output.rejected)
+    );
+    assert_eq!(pass.mocker_metrics.kv_cache_used_bytes, Some(0));
+}
+
+#[test]
+fn grouped_oversized_chunk_names_the_fixed_state_bytes() {
+    let mut core = core(
+        vec![group("window", 4, 4, Some(4)), state("kda", 2, 1)],
+        9,
+        8,
+    );
+    submit(&mut core, 116, 16, 2);
+    let mut trace = TraceCollector::default();
+    // The one-token forward needs one 4-byte window page and the 2-byte state
+    // block, 6 bytes, so admission passes. The first chunk of 8 tokens needs
+    // two window pages (8 bytes) and the state block, 10 bytes, above the pool.
+    let error = core
+        .try_execute_pass(&mut trace, 0.0)
+        .unwrap_err()
+        .to_string();
+    assert!(error.contains("needs 10 bytes for 8 scheduled tokens after 0 computed tokens"));
+    assert!(error.contains("exceeding the entire 9-byte cache pool"));
+    assert!(error.contains("the fixed state blocks of one request take 2 bytes"));
+    assert!(error.contains("reduce max_num_batched_tokens"));
+    // The sentence comes before the remedy.
+    assert!(error.contains("2 bytes of this pool; reduce max_num_batched_tokens"));
+}
+
 #[test]
 fn grouped_handoff_commands_fail_before_acquiring_state() {
     let mut core = core(vec![group("window", 4, 4, Some(4))], 24, 8);

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

@@ -64,6 +64,16 @@ pub struct FpmCacheGroup {
     /// for every other kind.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub blocks_per_request: Option<u32>,
+    /// The `num_prefill_checkpoint_blocks` of a Mamba spec: 0 or 1, and only
+    /// in a `state` group. In vLLM 0.33.0, a KDA layer that prefills with
+    /// FlashKDA has one checkpoint block.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub prefill_checkpoint_blocks: Option<u32>,
+    /// The `prefill_checkpoint_alignment` of a Mamba spec, in tokens: 16 for
+    /// FlashKDA in vLLM 0.33.0. Positive, and only with
+    /// `prefill_checkpoint_blocks` of 1.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub prefill_checkpoint_alignment: Option<u32>,
 }
 
 impl FpmCacheGroup {
@@ -96,6 +106,22 @@ impl FpmCacheGroup {
                 "blocks_per_request applies only to state cache groups",
             ));
         }
+        if self
+            .prefill_checkpoint_blocks
+            .is_some_and(|blocks| self.kind != FpmCacheKind::State || blocks > 1)
+        {
+            return Err(invalid(
+                "prefill_checkpoint_blocks requires a state group and must be 0 or 1",
+            ));
+        }
+        if self
+            .prefill_checkpoint_alignment
+            .is_some_and(|alignment| self.prefill_checkpoint_blocks != Some(1) || alignment == 0)
+        {
+            return Err(invalid(
+                "prefill_checkpoint_alignment requires prefill_checkpoint_blocks=1 and must be positive",
+            ));
+        }
         Ok(())
     }
 
@@ -750,6 +776,8 @@ mod tests {
             page_size_bytes: 160,
             sliding_window: window,
             blocks_per_request: None,
+            prefill_checkpoint_blocks: None,
+            prefill_checkpoint_alignment: None,
         }
     }
 
@@ -1375,6 +1403,38 @@ mod tests {
         }
     }
 
+    #[test]
+    fn checkpoint_fields_need_a_state_group_and_each_other() {
+        let blocks_text = "prefill_checkpoint_blocks requires a state group and must be 0 or 1";
+        let alignment_text = "prefill_checkpoint_alignment requires prefill_checkpoint_blocks=1 and must be positive";
+        let checkpoint = |blocks, alignment| FpmCacheGroup {
+            prefill_checkpoint_blocks: blocks,
+            prefill_checkpoint_alignment: alignment,
+            ..state_group(Some(1))
+        };
+        for (blocks, alignment) in [(Some(0), None), (Some(1), None), (Some(1), Some(16))] {
+            checkpoint(blocks, alignment).validate().unwrap();
+        }
+        let attention = FpmCacheGroup {
+            prefill_checkpoint_blocks: Some(0),
+            ..group(None)
+        };
+        assert_eq!(rejection(&attention), blocks_text);
+        assert_eq!(rejection(&checkpoint(Some(2), None)), blocks_text);
+        for (blocks, alignment) in [(None, Some(16)), (Some(0), Some(16)), (Some(1), Some(0))] {
+            assert_eq!(rejection(&checkpoint(blocks, alignment)), alignment_text);
+        }
+        // The kind checks run first, so an invalid group keeps its first error.
+        let unusable = FpmCacheGroup {
+            prefill_checkpoint_blocks: Some(2),
+            ..state_group(None)
+        };
+        assert_eq!(
+            rejection(&unusable),
+            "state cache groups require a positive blocks_per_request"
+        );
+    }
+
     #[test]
     fn state_group_holds_a_fixed_range_for_any_length() {
         let state = state_group(Some(2));
@@ -1415,6 +1475,8 @@ mod tests {
             page_size_bytes: 13_045_120,
             sliding_window: None,
             blocks_per_request,
+            prefill_checkpoint_blocks: None,
+            prefill_checkpoint_alignment: None,
         };
         let attention = pool_group("mla_indexer", FpmCacheKind::Attention, 22, None);
         let states = [
@@ -1458,7 +1520,28 @@ mod tests {
         let saved = serde_json::to_value(&state).unwrap();
         assert_eq!(saved["kind"], "state");
         assert_eq!(saved["blocks_per_request"], 1);
-        for original in [attention, state] {
+        // A group saved without the checkpoint fields loads with both absent and saves unchanged.
+        let earlier = r#"{"name":"state","kind":"state","num_layers":2,"block_size_tokens":16,"page_size_bytes":160,"sliding_window":null,"blocks_per_request":1}"#;
+        let loaded = serde_json::from_str::<FpmCacheGroup>(earlier).unwrap();
+        assert_eq!(loaded, state);
+        assert_eq!(
+            (
+                loaded.prefill_checkpoint_blocks,
+                loaded.prefill_checkpoint_alignment
+            ),
+            (None, None)
+        );
+        assert_eq!(serde_json::to_string(&loaded).unwrap(), earlier);
+        // Set fields follow the block count in declaration order.
+        let marked = FpmCacheGroup {
+            prefill_checkpoint_blocks: Some(1),
+            prefill_checkpoint_alignment: Some(16),
+            ..state_group(Some(1))
+        };
+        assert!(serde_json::to_string(&marked).unwrap().ends_with(
+            r#""blocks_per_request":1,"prefill_checkpoint_blocks":1,"prefill_checkpoint_alignment":16}"#
+        ));
+        for original in [attention, state, marked] {
             let saved = serde_json::to_string(&original).unwrap();
             assert_eq!(
                 serde_json::from_str::<FpmCacheGroup>(&saved).unwrap(),

docs/perf-model/fpm-self-service/README.md

@@ -119,7 +119,7 @@ For online learning from per-iteration telemetry, continue with
 | Parallelism | Whole-pass timing can capture intra-iteration computation and communication. Exact topology/resource identity and query coverage remain required; see [current guided CLI coverage](implementation.md#parallelism-support-and-guided-cli-coverage) for collection automation limits. |
 | Self-benchmark serving roles | Aggregated workers collect both prefill and decode. Prefill and decode can also be collected independently with separate accepted configurations. Their role exports must be composed separately for serving; `onboard validate-fpm` currently validates aggregated replay. |
 | Self-benchmark GPU execution | Kubernetes by default, or Slurm/Pyxis inside a caller-owned allocation. Both have multi-node launch paths; placement and GPU requirements come from the generated plan. The workflow does not provision the cluster or acquire a Slurm allocation. |
-| KV cache | Supports standard per-token KV caches, as used by Llama, Qwen3 and MiniMax-M2. Models with linear-attention state, such as Kimi K3's KDA, use the [measured-profile path](examples.md#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile), with the KV cache and state-management limits described in that example. Some configs leave the cache geometry to the [runtime probe](implementation.md#cache-geometry-pending-the-runtime-probe): linear-attention and recurrent-state configs, and architectures that intake does not recognize. For them, `onboard init --model-config` stops with guidance to run the probe, which measures the geometry. A cache geometry override for these configs is a claim that `onboard import-observations` checks against that measurement. Importing cache groups other than full or sliding attention (for example recurrent state) is planned. Other KV cache types may also work, but may need changes to the self-benchmarker to initialize their state and measure the forward pass correctly. |
+| KV cache | Supports standard per-token KV caches, as used by Llama, Qwen3 and MiniMax-M2. Models with linear-attention state, such as Kimi K3's KDA, use the [measured-profile path](examples.md#example-a-onboard-the-collected-kimi-k3-tp8dcp8-profile), with the KV cache and state-management limits described in that example. Some configs leave the cache geometry to the [runtime probe](implementation.md#cache-geometry-pending-the-runtime-probe): linear-attention and recurrent-state configs, and architectures that intake does not recognize. For them, `onboard init --model-config` stops with guidance to run the probe, which measures the geometry. A cache geometry override for these configs is a claim that `onboard import-observations` checks against that measurement. On vLLM 0.33.0 with the Dynamo scheduler file of the `release/1.7.0` branch that the bundle pins ([source notes](../../../python/aisim/collector/fpm_forward/runtime/instrumentation/vllm-0.33.0.md)), the probe imports recurrent-state and tail caches as `state` groups. Other KV cache types may also work, but may need changes to the self-benchmarker to initialize their state and measure the forward pass correctly. |
 | Multimodality | For multimodal models, this workflow covers text generation only; image/audio processing time is excluded. |
 
 Use a source revision containing the required CLI and consumer features; a