feat: simulate grouped prefix caching in the vLLM scheduler and plan it from probe evidence

THIRD_PARTY_NOTICES.md

@@ -457,10 +457,11 @@ The identified upstream revision has no root NOTICE file. The repository's
 Apache-2.0 LICENSE supplies the license text; no upstream source file is vendored.
 
 
-The repository-root file `crates/core/src/engine/kv_manager/grouped_prefix.rs`
-adapts rules of vLLM v0.31.0, immutable commit
-`db9527a46873454610df6dbedf79a36d6bf1a7f6`, and of vLLM v0.32.0, immutable
-commit `5df63a7203cdb9328d824a27b682e77c7b9928be`.
+The repository-root files `crates/core/src/engine/kv_manager/grouped_prefix.rs`
+and `crates/core/src/engine/kv_manager/grouped_prefix_tests.rs` (its test
+module, with cases ported from the upstream tests) adapt rules of vLLM v0.31.0,
+immutable commit `db9527a46873454610df6dbedf79a36d6bf1a7f6`, and of vLLM
+v0.32.0, immutable commit `5df63a7203cdb9328d824a27b682e77c7b9928be`.
 One rule table keyed by the vLLM version holds both rows. From
 `kv_cache_coordinator.py`, `grouped_prefix.rs` adapts the hybrid cache-hit
 lookup over the spec groups, the hit alignment, the scheduler block, the pinning
@@ -482,6 +483,21 @@ it adapts the block of the align split, which is the smallest block of the
 prefix-cacheable groups. From `config/vllm.py`, it adapts the bound on in-flight
 tokens that the admission cap uses. The test module ports cases from the
 upstream tests listed below.
+From the same two revisions, `crates/core/src/engine/kv_manager/vllm_backend.rs`
+adapts, from `kv_cache_manager.py`, the whole-prompt admission cost of
+`allocate_slots` and its test of the first chunk before the hit blocks are
+pinned, and `grouped_prefix.rs` counts the blocks of that first chunk.
+`crates/core/src/engine/scheduler/vllm/policy.rs` adapts the whole-prompt test
+of `allocate_slots` and, from `sched/scheduler.py`, the token budget of a prompt
+without chunked prefill and the stop of the waiting loop at a request that does
+not fit, whose hit is searched again at the next step.
+`crates/core/src/engine/config.rs` and `grouped_prefix.rs` adapt the startup
+check of `kv_cache_utils.py` that the pool holds one request of `max_model_len`
+tokens, with the size of each group from `kv_cache_interface.py`, summed per
+group.
+From `sched/scheduler.py` and `kv_cache_interface.py`, `grouped_prefix.rs` also
+adapts the rule that only a rank with a Mamba layer in align mode splits its
+prefill.
 
 - https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/v1/core/kv_cache_coordinator.py
 - https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/v1/core/single_type_kv_cache_manager.py

crates/core/src/engine/common/protocols.rs

@@ -235,6 +235,10 @@ pub(crate) struct MockEngineArgs {
     pub kv_cache_groups: Vec<crate::perfmodel::FpmCacheGroup>,
     #[builder(default = "None")]
     pub kv_cache_capacity_bytes: Option<u64>,
+    /// The vLLM release whose grouped-cache rules apply; see
+    /// `EngineConfig::vllm_version`.
+    #[builder(default = "None")]
+    pub vllm_version: Option<String>,
     #[builder(default = "None")]
     pub native_host_offload: Option<NativeHostOffloadConfig>,
     #[builder(default = "None")]

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

@@ -12,6 +12,7 @@ use crate::engine::belady::BeladyOracle;
 use crate::engine::common::protocols::{KvEventPublishers, PrefillCost};
 use crate::engine::common::sequence::RequestSequence;
 
+use super::grouped_prefix::{AdmissionLimits, GroupedPrefixRules};
 use super::vllm_backend::{
     BlockRequestLease, DecodeBlockReservation as VllmDecodeBlockReservation,
     DestinationReservation as VllmDestinationReservation,
@@ -78,6 +79,24 @@ impl G1Manager {
         self.inner.grouped()
     }
 
+    pub(crate) fn set_grouped_prefix_rules(
+        &mut self,
+        rules: GroupedPrefixRules,
+        limits: AdmissionLimits,
+    ) {
+        self.inner.set_grouped_prefix_rules(rules, limits);
+    }
+
+    /// The rule table of a grouped cache with prefix caching on.
+    pub(crate) fn grouped_prefix_rules(&self) -> Option<&GroupedPrefixRules> {
+        self.inner.grouped_prefix_rules()
+    }
+
+    /// The whole-prompt blocks of `known_tokens` without a prefix hit.
+    pub(crate) fn grouped_cold_admission_blocks(&self, known_tokens: usize) -> usize {
+        self.inner.grouped_cold_admission_blocks(known_tokens)
+    }
+
     pub(crate) fn set_belady_oracle(&mut self, oracle: BeladyOracle) {
         self.inner.set_belady_oracle(oracle);
     }
@@ -257,6 +276,16 @@ impl G1Manager {
         self.inner.finish_lease(owner, lease);
     }
 
+    /// The prefill cost of a waiting candidate. A grouped prefix cache keeps
+    /// the hit on the lease for the allocation, so it is searched once.
+    pub(crate) fn lookup_native_prefill_cost(
+        &self,
+        sequence: &RequestSequence,
+        lease: &mut BlockRequestLease,
+    ) -> PrefillCost {
+        self.inner.lookup_prefill_cost(sequence, lease)
+    }
+
     pub(crate) fn get_native_prefill_cost(
         &self,
         sequence: &RequestSequence,

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

@@ -4,10 +4,6 @@
 //! Request-owned cache groups sharing one rank-local physical budget. Byte
 //! mode replays them cold. Block mode keeps the blocks of every group in one
 //! block pool, so that cached blocks stay until eviction.
-#![cfg_attr(
-    not(test),
-    expect(dead_code, reason = "the scheduler does not use block mode yet")
-)]
 
 use crate::engine::cache::vllm_block_pool::{
     BlockCopyId, BlockReservation, CacheKey, PrefixHit, VllmBlockPool,
@@ -112,13 +108,33 @@ pub(crate) struct GroupStateLease {
 pub(crate) struct GroupedBlockPool {
     pool: VllmBlockPool,
     page_size_bytes: u64,
+    /// The most active blocks seen since the last reset, for tests that
+    /// observe a step's peak after the step.
+    #[cfg(test)]
+    peak_active_blocks: std::cell::Cell<usize>,
 }
 
 impl GroupedBlockPool {
     pub(crate) fn capacity_blocks(&self) -> usize {
         self.pool.capacity()
     }
 
+    #[cfg(test)]
+    pub(crate) fn peak_active_blocks(&self) -> usize {
+        self.peak_active_blocks.get()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn reset_peak_active_blocks(&self) {
+        self.peak_active_blocks.set(self.pool.num_active());
+    }
+
+    fn note_peak(&self) {
+        #[cfg(test)]
+        self.peak_active_blocks
+            .set(self.peak_active_blocks.get().max(self.pool.num_active()));
+    }
+
     pub(crate) fn key_hit(&self, group: u16, hash: SequenceHash) -> Option<PrefixHit> {
         self.pool.key_hit(CacheKey::Grouped { group, hash })
     }
@@ -135,14 +151,19 @@ impl GroupedBlockPool {
     /// Turn the pinned hits into request references. The copies follow the
     /// order of `hits`.
     pub(crate) fn activate(&mut self, reservation: &mut BlockReservation) -> Vec<BlockCopyId> {
-        self.pool
+        let copies = self
+            .pool
             .activate_keys(reservation)
             .map(|(_, id)| id)
-            .collect()
+            .collect();
+        self.note_peak();
+        copies
     }
 
     pub(crate) fn take_private(&mut self, reservation: &mut BlockReservation) -> BlockCopyId {
-        self.pool.allocate_private(reservation)
+        let id = self.pool.allocate_private(reservation);
+        self.note_peak();
+        id
     }
 
     /// Publish a private block. Returns whether `key` became visible.
@@ -233,11 +254,40 @@ impl GroupedKvPool {
             mode: GroupedMode::Blocks(GroupedBlockPool {
                 pool,
                 page_size_bytes,
+                #[cfg(test)]
+                peak_active_blocks: std::cell::Cell::new(0),
             }),
             ..Self::new(groups, capacity_bytes)
         }
     }
 
+    fn debug_assert_byte_mode(&self) {
+        debug_assert!(
+            matches!(self.mode, GroupedMode::Bytes),
+            "the byte-mode API is not called in block mode"
+        );
+    }
+
+    /// Pool blocks that the groups hold for a forward over
+    /// `computed_before..computed_after`, in either mode.
+    pub(crate) fn projected_blocks(
+        &self,
+        computed_before: usize,
+        computed_after: usize,
+    ) -> Option<usize> {
+        self.project(computed_before, computed_after)
+            .map(|lease| lease.ranges.iter().map(std::ops::Range::len).sum())
+    }
+
+    /// Pool blocks that the state groups hold for every running request.
+    pub(crate) fn fixed_request_blocks(&self) -> usize {
+        self.groups
+            .iter()
+            .filter(|group| group.kind == FpmCacheKind::State)
+            .map(|group| group.blocks_per_request.unwrap_or(0) as usize)
+            .sum()
+    }
+
     fn project(&self, computed_before: usize, computed_after: usize) -> Option<GroupedLease> {
         if computed_after < computed_before {
             return None;
@@ -267,13 +317,15 @@ impl GroupedKvPool {
         computed_before: usize,
         computed_after: usize,
     ) -> Option<u64> {
+        self.debug_assert_byte_mode();
         self.project(computed_before, computed_after)
             .map(|lease| lease.bytes)
     }
 
     /// Bytes that the state groups hold for every running request, whatever
     /// its context length.
     pub(crate) fn fixed_request_bytes(&self) -> u64 {
+        self.debug_assert_byte_mode();
         self.groups
             .iter()
             .filter(|group| group.kind == FpmCacheKind::State)
@@ -292,6 +344,7 @@ impl GroupedKvPool {
         computed_before: usize,
         computed_after: usize,
     ) -> bool {
+        self.debug_assert_byte_mode();
         let Some(projected) = self.project(computed_before, computed_after) else {
             return false;
         };
@@ -308,6 +361,7 @@ impl GroupedKvPool {
     }
 
     pub(crate) fn release(&mut self, lease: &mut Option<GroupedLease>) {
+        self.debug_assert_byte_mode();
         if let Some(lease) = lease.take() {
             self.used_bytes -= lease.bytes;
             self.active_blocks -= lease.blocks();