diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 948ddbf54a..72d166d46f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/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 diff --git a/crates/core/src/engine/common/protocols.rs b/crates/core/src/engine/common/protocols.rs index e8060134b1..98d58765a2 100644 --- a/crates/core/src/engine/common/protocols.rs +++ b/crates/core/src/engine/common/protocols.rs @@ -235,6 +235,10 @@ pub(crate) struct MockEngineArgs { pub kv_cache_groups: Vec, #[builder(default = "None")] pub kv_cache_capacity_bytes: Option, + /// The vLLM release whose grouped-cache rules apply; see + /// `EngineConfig::vllm_version`. + #[builder(default = "None")] + pub vllm_version: Option, #[builder(default = "None")] pub native_host_offload: Option, #[builder(default = "None")] diff --git a/crates/core/src/engine/config.rs b/crates/core/src/engine/config.rs index 9d5ee0a790..e2a655e8b7 100644 --- a/crates/core/src/engine/config.rs +++ b/crates/core/src/engine/config.rs @@ -3,18 +3,23 @@ // SPDX-License-Identifier: Apache-2.0 // Modified behavioral adaptation of SGLang v0.5.22; see THIRD_PARTY_NOTICES.md. // Source: https://github.com/sgl-project/sglang/tree/baa9befa381511cd4d56f83980951e5a45fffc6d +// Copyright contributors to the vLLM project. +// Modified Rust simulation of the vLLM v0.31.0 +// (db9527a46873454610df6dbedf79a36d6bf1a7f6) and v0.32.0 +// (5df63a7203cdb9328d824a27b682e77c7b9928be) grouped KV rules. +// Immutable sources and license: THIRD_PARTY_NOTICES.md. //! Serializable engine configuration. use std::sync::Arc; -use anyhow::{Result, ensure}; +use anyhow::{Result, bail, ensure}; use serde::{Deserialize, Deserializer, Serialize}; use crate::engine::common::speculative::normalize_conditional_accept_rates; use crate::engine::handoff::TransferTimingMode; use crate::engine::host_offload::HostPolicy; -use crate::engine::kv_manager::SourceReleaseOrder; +use crate::engine::kv_manager::{GroupedPrefixRules, SourceReleaseOrder, VllmRuleVersion}; use crate::engine::timing::{TimingModel, TimingModelConfig, built_in_timing_model}; use crate::perfmodel::FpmCacheGroup; @@ -916,6 +921,11 @@ pub struct EngineConfig { /// Shared physical budget for `kv_cache_groups`, including page padding. #[serde(skip_serializing_if = "Option::is_none")] pub kv_cache_capacity_bytes: Option, + /// vLLM release whose grouped-cache rules this rank follows. A rank with + /// `kv_cache_groups` and `enable_prefix_caching` must set it. Accepted + /// values: `0.31.0` and `0.32.0`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub vllm_version: Option, /// Optional framework-native host-offload simulation. #[serde(skip_serializing_if = "Option::is_none")] pub native_host_offload: Option, @@ -1013,6 +1023,8 @@ struct EngineConfigWire { #[serde(default)] kv_cache_capacity_bytes: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + vllm_version: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] native_host_offload: Option, #[serde(default, skip_serializing_if = "Option::is_none")] state_cache: Option, @@ -1082,6 +1094,7 @@ impl<'de> Deserialize<'de> for EngineConfig { kv_cache_bytes_per_token: wire.kv_cache_bytes_per_token, kv_cache_groups: wire.kv_cache_groups, kv_cache_capacity_bytes: wire.kv_cache_capacity_bytes, + vllm_version: wire.vllm_version, native_host_offload: wire.native_host_offload, g3_offload: wire.g3_offload, kvcr: wire.kvcr, @@ -1129,6 +1142,7 @@ impl Default for EngineConfig { kv_cache_bytes_per_token: None, kv_cache_groups: Vec::new(), kv_cache_capacity_bytes: None, + vllm_version: None, native_host_offload: None, g3_offload: None, kvcr: None, @@ -1169,15 +1183,80 @@ impl EngineConfig { } } + /// The grouped prefix-caching gate. Every group declares whether it + /// caches, the pool has one page size, the hash unit and the vLLM + /// version are set, the rule table accepts the groups, and the pool + /// holds the full prompt of `max_model_len` plus the state blocks of one + /// request, summed per group as vLLM's startup check does (v0.31 + /// `kv_cache_utils.py:938-942`, v0.32 `:1043-1047`). + fn validate_grouped_prefix_caching(&self) -> Result<()> { + let groups = &self.kv_cache_groups; + ensure!( + groups.iter().all(|group| group.prefix_cacheable.is_some()), + "kv_cache_groups with enable_prefix_caching=true require every group to declare prefix_cacheable" + ); + ensure!( + groups + .iter() + .any(|group| group.prefix_cacheable == Some(true)), + "kv_cache_groups with enable_prefix_caching=true require at least one prefix_cacheable group" + ); + let page_size_bytes = groups[0].page_size_bytes; + ensure!( + groups + .iter() + .all(|group| group.page_size_bytes == page_size_bytes), + "kv_cache_groups with enable_prefix_caching=true require equal page_size_bytes across groups" + ); + let Some(unit) = self.prefix_match_unit else { + bail!("kv_cache_groups with enable_prefix_caching=true require prefix_match_unit"); + }; + let Some(version) = self + .vllm_version + .as_deref() + .and_then(VllmRuleVersion::parse) + else { + bail!( + "kv_cache_groups with enable_prefix_caching=true require vllm_version 0.31.0 or 0.32.0" + ); + }; + let rules = GroupedPrefixRules::new(groups, unit, version)?; + // vLLM always sizes the pool with a model length. + let Some(max_model_len) = self.max_model_len else { + bail!("kv_cache_groups with enable_prefix_caching=true require max_model_len"); + }; + let capacity_blocks = self + .kv_cache_capacity_bytes + .expect("validated grouped cache capacity") + / page_size_bytes; + let max_in_flight_tokens = self + .max_concurrent_batches() + .saturating_mul(self.max_num_batched_tokens); + ensure!( + capacity_blocks >= rules.startup_blocks(max_model_len, max_in_flight_tokens) as u64, + "kv_cache_groups with enable_prefix_caching=true need capacity for the full prompt of max_model_len plus the state blocks of one request" + ); + Ok(()) + } + fn validate_state_cache(&self) -> Result<()> { if let Some(unit) = self.prefix_match_unit { + // A grouped cache hashes at the unit itself. The grouped gate + // checks it against the group blocks. ensure!( - unit > 0 && self.block_size > 0 && self.block_size.is_multiple_of(unit), + unit > 0 + && (!self.kv_cache_groups.is_empty() + || (self.block_size > 0 && self.block_size.is_multiple_of(unit))), "prefix_match_unit must be a positive divisor of block_size" ); ensure!( - self.state_cache.is_some(), - "prefix_match_unit currently requires state_cache" + self.state_cache.is_some() || !self.kv_cache_groups.is_empty(), + "prefix_match_unit currently requires state_cache or kv_cache_groups" + ); + // Only grouped prefix caching reads the unit, as in Python. + ensure!( + self.kv_cache_groups.is_empty() || self.enable_prefix_caching, + "prefix_match_unit on kv_cache_groups requires enable_prefix_caching=true" ); ensure!( !self.emit_kv_events, @@ -1345,15 +1424,15 @@ impl EngineConfig { self.kv_cache_capacity_bytes.is_none(), "kv_cache_capacity_bytes requires kv_cache_groups" ); + ensure!( + self.vllm_version.is_none(), + "vllm_version requires kv_cache_groups" + ); } else { ensure!( self.backend == Backend::Vllm && self.worker_type == WorkerType::Aggregated, "kv_cache_groups currently require backend=vllm and worker_type=aggregated" ); - ensure!( - !self.enable_prefix_caching, - "kv_cache_groups currently require enable_prefix_caching=false" - ); ensure!( self.native_host_offload.is_none() && self.g3_offload.is_none(), "kv_cache_groups currently support only HBM without host or G3 offload" @@ -1382,6 +1461,9 @@ impl EngineConfig { group.validate()?; ensure!(names.insert(&group.name), "duplicate KV cache group name"); } + if self.enable_prefix_caching { + self.validate_grouped_prefix_caching()?; + } } if let Some(g3) = &self.g3_offload { g3.validate()?; @@ -1897,7 +1979,7 @@ mod tests { ( "enable_prefix_caching", serde_json::json!(true), - "enable_prefix_caching=false", + "declare prefix_cacheable", ), ("aic_nextn", serde_json::json!(1), "speculative"), ("emit_kv_events", serde_json::json!(true), "KV event"), @@ -2014,6 +2096,222 @@ mod tests { } } + const GLM_PAGE: u64 = 13_045_120; + + /// The six GLM-5.3-Flash TP 4 groups that the runtime probe recorded on + /// vLLM 0.31.0. `checkpoint` is the `prefill_checkpoint_blocks` of G2 to + /// G5, with alignment 16 for 1. + fn glm_groups_wire(checkpoint: Option) -> serde_json::Value { + let mut groups = vec![ + serde_json::json!({ + "name": "mla", "kind": "attention", "num_layers": 22, + "block_size_tokens": 2176, "page_size_bytes": GLM_PAGE, + "prefix_cacheable": true + }), + serde_json::json!({ + "name": "kpool_tail", "kind": "state", "num_layers": 11, + "block_size_tokens": 4, "page_size_bytes": GLM_PAGE, + "blocks_per_request": 1, "prefix_cacheable": false + }), + ]; + for i in 0..4 { + let mut group = serde_json::json!({ + "name": format!("kda_{i}"), "kind": "state", "num_layers": 4, + "block_size_tokens": 2176, "page_size_bytes": GLM_PAGE, + "blocks_per_request": 1, "prefix_cacheable": true + }); + if let Some(blocks) = checkpoint { + group["prefill_checkpoint_blocks"] = blocks.into(); + if blocks == 1 { + group["prefill_checkpoint_alignment"] = 16.into(); + } + } + groups.push(group); + } + serde_json::Value::Array(groups) + } + + fn glm_prefix_wire(version: Option<&str>, checkpoint: Option) -> serde_json::Value { + let mut wire = serde_json::json!({ + "backend": "vllm", + "enable_prefix_caching": true, + "prefix_match_unit": 2176, + "max_model_len": 131_072, + "kv_cache_capacity_bytes": 74 * GLM_PAGE, + "kv_cache_groups": glm_groups_wire(checkpoint), + }); + if let Some(version) = version { + wire["vllm_version"] = version.into(); + } + wire + } + + fn validation_error(wire: serde_json::Value) -> String { + let config: EngineConfig = serde_json::from_value(wire).unwrap(); + config.validate().unwrap_err().to_string() + } + + fn validates(wire: serde_json::Value) { + let config: EngineConfig = serde_json::from_value(wire).unwrap(); + config.validate().unwrap(); + } + + #[test] + fn grouped_prefix_caching_validates_its_geometry() { + validates(glm_prefix_wire(Some("0.31.0"), None)); + validates(glm_prefix_wire(Some("0.31.0"), Some(0))); + let round_trip: EngineConfig = + serde_json::from_value(glm_prefix_wire(Some("0.31.0"), None)).unwrap(); + assert_eq!(round_trip.vllm_version.as_deref(), Some("0.31.0")); + let encoded = serde_json::to_value(&round_trip).unwrap(); + assert_eq!(encoded["vllm_version"], "0.31.0"); + assert_eq!( + serde_json::from_value::(encoded).unwrap(), + round_trip + ); + + // The version rules. + assert!( + validation_error(glm_prefix_wire(None, None)) + .contains("require vllm_version 0.31.0 or 0.32.0") + ); + assert!( + validation_error(glm_prefix_wire(Some("0.30.0"), None)) + .contains("require vllm_version 0.31.0 or 0.32.0") + ); + validates(glm_prefix_wire(Some("0.32.0"), Some(0))); + assert!(validation_error(glm_prefix_wire(Some("0.32.0"), None)).contains("re-import")); + validates(glm_prefix_wire(Some("0.32.0"), Some(1))); + assert!( + validation_error(glm_prefix_wire(Some("0.31.0"), Some(1))).contains("do not model") + ); + + // The capacity rule: cdiv(131072, 2176) + (2 + checkpoint) * 4 + 1. + for (version, checkpoint, bound) in [("0.31.0", Some(0), 70), ("0.32.0", Some(1), 74)] { + let mut wire = glm_prefix_wire(Some(version), checkpoint); + wire["kv_cache_capacity_bytes"] = ((bound - 1) * GLM_PAGE).into(); + assert!( + validation_error(wire.clone()).contains("need capacity for the full prompt"), + "{version} at {} blocks", + bound - 1 + ); + wire["kv_cache_capacity_bytes"] = (bound * GLM_PAGE).into(); + validates(wire); + } + let mut small = glm_prefix_wire(Some("0.31.0"), None); + small["kv_cache_capacity_bytes"] = (10 * GLM_PAGE).into(); + assert!(validation_error(small).contains("need capacity for the full prompt")); + let mut unbounded = glm_prefix_wire(Some("0.31.0"), None); + unbounded.as_object_mut().unwrap().remove("max_model_len"); + assert!(validation_error(unbounded).contains("require max_model_len")); + + // The geometry rules, in gate order. + let mut undeclared = glm_prefix_wire(Some("0.31.0"), None); + undeclared["kv_cache_groups"][3] + .as_object_mut() + .unwrap() + .remove("prefix_cacheable"); + assert!(validation_error(undeclared).contains("declare prefix_cacheable")); + let mut none_cacheable = glm_prefix_wire(Some("0.31.0"), None); + for group in none_cacheable["kv_cache_groups"].as_array_mut().unwrap() { + group["prefix_cacheable"] = false.into(); + } + assert!(validation_error(none_cacheable).contains("at least one prefix_cacheable group")); + let mut unequal = glm_prefix_wire(Some("0.31.0"), None); + unequal["kv_cache_groups"][1]["page_size_bytes"] = (GLM_PAGE - 1).into(); + assert!(validation_error(unequal).contains("equal page_size_bytes")); + let mut no_unit = glm_prefix_wire(Some("0.31.0"), None); + no_unit.as_object_mut().unwrap().remove("prefix_match_unit"); + assert!(validation_error(no_unit).contains("require prefix_match_unit")); + let mut bad_unit = glm_prefix_wire(Some("0.31.0"), None); + bad_unit["prefix_match_unit"] = 100.into(); + assert!( + validation_error(bad_unit) + .contains("prefix_match_unit must divide the block_size_tokens") + ); + let mut with_state_cache = glm_prefix_wire(Some("0.31.0"), None); + with_state_cache["state_cache"] = serde_json::json!({"bytes_per_request": 64}); + with_state_cache["num_gpu_blocks"] = 64.into(); + with_state_cache["block_size"] = 4.into(); + with_state_cache["kv_cache_bytes_per_token"] = 16.into(); + assert!( + validation_error(with_state_cache) + .contains("cannot use linear cache or KV transfer byte geometry") + ); + + // Prefix caching off keeps the byte mode, with or without the version. + // Only prefix caching reads the unit, so a rank with the unit and + // without prefix caching is refused. + let mut off = glm_prefix_wire(Some("0.31.0"), None); + off["enable_prefix_caching"] = false.into(); + assert!( + serde_json::from_value::(off.clone()) + .unwrap_err() + .to_string() + .contains( + "prefix_match_unit on kv_cache_groups requires enable_prefix_caching=true" + ) + ); + off.as_object_mut().unwrap().remove("prefix_match_unit"); + validates(off.clone()); + off.as_object_mut().unwrap().remove("vllm_version"); + validates(off); + // A linear rank has no grouped-cache rules to follow. + let linear: EngineConfig = serde_json::from_value(serde_json::json!({ + "backend": "vllm", "vllm_version": "0.31.0" + })) + .unwrap(); + assert!( + linear + .validate() + .unwrap_err() + .to_string() + .contains("vllm_version requires kv_cache_groups") + ); + assert!( + serde_json::to_value(EngineConfig::default()) + .unwrap() + .get("vllm_version") + .is_none() + ); + } + + #[test] + fn grouped_prefix_caching_refuses_convolution_groups() { + // The grouped rule table has no rule for a convolution group, so the + // gate refuses it with prefix caching on, whatever the capacity. Cold + // replay keeps the group. + let wire = |prefix_caching: bool, version: &str, pages: u64| { + let mut wire = serde_json::json!({ + "backend": "vllm", + "enable_prefix_caching": prefix_caching, + "max_model_len": 128, + "max_num_batched_tokens": 8, + "kv_cache_capacity_bytes": pages * 64, + "kv_cache_groups": [{ + "name": "convolution", "kind": "convolution", "num_layers": 1, + "block_size_tokens": 4, "page_size_bytes": 64, "sliding_window": 4, + "prefix_cacheable": true + }], + }); + if prefix_caching { + wire["prefix_match_unit"] = 4.into(); + wire["vllm_version"] = version.into(); + } + wire + }; + for version in ["0.31.0", "0.32.0"] { + for pages in [8, 32] { + assert_eq!( + validation_error(wire(true, version, pages)), + "grouped prefix caching does not support convolution cache groups; turn prefix caching off to replay them cold", + "{version} at {pages} pages" + ); + } + } + validates(wire(false, "", 8)); + } + #[test] fn native_host_offload_deserializes_with_default_bandwidths() { assert_eq!(DEFAULT_HOST_OFFLOAD_BANDWIDTH_GBPS, 32.0); diff --git a/crates/core/src/engine/kv_manager/g1_manager.rs b/crates/core/src/engine/kv_manager/g1_manager.rs index cc81c78549..8ca0ec283a 100644 --- a/crates/core/src/engine/kv_manager/g1_manager.rs +++ b/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, diff --git a/crates/core/src/engine/kv_manager/grouped.rs b/crates/core/src/engine/kv_manager/grouped.rs index 97a43c9412..a88f4a3333 100644 --- a/crates/core/src/engine/kv_manager/grouped.rs +++ b/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,6 +108,10 @@ 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, } impl GroupedBlockPool { @@ -119,6 +119,22 @@ impl GroupedBlockPool { 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 { 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 { - 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 { + 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 { if computed_after < computed_before { return None; @@ -267,6 +317,7 @@ impl GroupedKvPool { computed_before: usize, computed_after: usize, ) -> Option { + self.debug_assert_byte_mode(); self.project(computed_before, computed_after) .map(|lease| lease.bytes) } @@ -274,6 +325,7 @@ impl GroupedKvPool { /// 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) { + self.debug_assert_byte_mode(); if let Some(lease) = lease.take() { self.used_bytes -= lease.bytes; self.active_blocks -= lease.blocks(); diff --git a/crates/core/src/engine/kv_manager/grouped_prefix.rs b/crates/core/src/engine/kv_manager/grouped_prefix.rs index eca84a6856..e6f673d141 100644 --- a/crates/core/src/engine/kv_manager/grouped_prefix.rs +++ b/crates/core/src/engine/kv_manager/grouped_prefix.rs @@ -27,19 +27,6 @@ //! - https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/v1/core/block_pool.py //! - https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/v1/engine/core.py //! - https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/config/vllm.py -//! -//! The test module ports cases from the vLLM tests below. -//! -//! - https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py -//! - https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py -//! - https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/tests/v1/core/test_mamba_align_chunk_split.py -#![cfg_attr( - not(test), - expect( - dead_code, - reason = "the scheduler does not use grouped prefix caching yet" - ) -)] use super::grouped::{GroupLease, GroupStateLease, GroupedBlockLease, GroupedBlockPool}; use super::vllm_backend::BlockRequestLease; @@ -63,13 +50,6 @@ impl VllmRuleVersion { _ => None, } } - - pub(crate) fn as_str(self) -> &'static str { - match self { - Self::V0_31_0 => "0.31.0", - Self::V0_32_0 => "0.32.0", - } - } } /// The vLLM 0.32.0 FlashKDA prefill checkpoint block (v0.32 `iface:1040-1043`). @@ -113,10 +93,10 @@ pub(crate) struct GroupedPrefixRules { /// group order (`coord:775-779`). specs: Vec, hash_block: usize, - /// The hit alignment `A` (`coord:687-735`). + /// The hit alignment `A` (`coord:687-735`): the hash block with + /// fine-grained hits, else the scheduler block, the LCM of every group + /// block (`coord:87-92`). alignment: usize, - /// The LCM of every group block (`coord:87-92`). - scheduler_block: usize, /// The smallest cacheable group block, vLLM's `cache_config.block_size` /// (`core:336-350`). cacheable_block: usize, @@ -305,13 +285,36 @@ impl GroupedPrefixRules { specs, hash_block, alignment, - scheduler_block, cacheable_block, version, checkpoint, }) } + /// The pool blocks that vLLM's startup check reserves for one request of + /// `max_model_len` tokens, summed per group (`kv_cache_utils.py:938-942`): + /// `cdiv(max_model_len, block)` for full attention (`iface:578-583`), the + /// admission cap of a sliding window (`iface:859-866`), `2 + checkpoint` + /// for a state in align mode (`iface:1083-1092`) and 1 for a ring. + pub(crate) fn startup_blocks( + &self, + max_model_len: usize, + max_in_flight_tokens: usize, + ) -> usize { + let checkpoint = usize::from(self.checkpoint.is_some()); + self.groups + .iter() + .map(|rule| match rule.kind { + GroupRuleKind::FullAttention => max_model_len.div_ceil(rule.block), + GroupRuleKind::SlidingWindow { window } => { + window_admission_blocks(window, rule.block, max_in_flight_tokens, max_model_len) + } + GroupRuleKind::State => 2 + checkpoint, + GroupRuleKind::Ring => 1, + }) + .sum() + } + /// Whether hits land on hash boundaries inside a block (`coord:703`, /// `sched:375-379`). fn fine_grained(&self) -> bool { @@ -384,7 +387,9 @@ impl GroupedPrefixRules { /// Modified simulation of the align split (v0.31 `sched:414-527`, v0.32 /// `sched:443-556`). `known` is the request's token count and `budget` - /// the step budget. `desired` is clamped to the tokens the request has. + /// the step budget. Without a State rule the split is off and returns + /// `desired`. During the prefill, `desired` is clamped to the tokens the + /// request has. pub(crate) fn align_split( &self, start: usize, @@ -394,6 +399,19 @@ impl GroupedPrefixRules { junction: usize, budget: usize, ) -> usize { + // vLLM splits only with a MambaSpec layer in align mode (v0.31 + // `sched:354-356`, `iface:1568-1574`, v0.32 `sched:377-379`, + // `iface:1563-1569`), for running and waiting requests (v0.31 + // `sched:695-697`, `:1127-1129`, v0.32 `sched:724-726`, `:1173-1175`). + // A State rule is a prefix-cacheable Mamba group, admitted only in + // align mode. A tail ring is not a MambaSpec. + if !self + .groups + .iter() + .any(|rule| rule.kind == GroupRuleKind::State) + { + return desired; + } let prefill_end = prompt.max(known.saturating_sub(1)); if start >= prefill_end { return desired; @@ -478,6 +496,22 @@ pub(crate) struct AdmissionLimits { pub(crate) max_model_len: usize, } +/// `SlidingWindowSpec.max_admission_blocks_per_request` (v0.31 +/// `iface:859-866`, v0.32 `iface:878-885`), without multi-module speculative +/// decoding. +fn window_admission_blocks( + window: usize, + block: usize, + max_in_flight_tokens: usize, + max_model_len: usize, +) -> usize { + (window - 1) + .saturating_add(max_in_flight_tokens) + .min(max_model_len) + .div_ceil(block) + + 1 +} + /// The first end, a multiple of `step` at most `max`, with `needed` /// contiguous resident blocks before it, or with every block from 0, and the /// run length (`single:960-1069`). `resident(i)` tells whether block `i` @@ -570,6 +604,15 @@ impl GroupedPrefixCoordinator { } } + /// The hit of a request that takes no part in prefix caching. + pub(crate) fn empty_hit(&self) -> GroupedHit { + GroupedHit { + tokens: 0, + junction_tokens: 0, + per_group: vec![GroupHitBlocks::None; self.rules.groups.len()], + } + } + /// The fixed point over the spec groups, full attention first /// (`coord:842-977`, `coord:775-779`), capped at `num_tokens - 1` /// (`mgr:289-295`). The longest hit, which sets the junction, is the @@ -757,6 +800,30 @@ impl GroupedPrefixCoordinator { 0 } + /// The tokens of the first full-attention group's hit blocks that another + /// request holds active. They cost no capacity at admission. + pub(crate) fn active_full_attention_hit_tokens( + &self, + pool: &GroupedBlockPool, + hit: &GroupedHit, + ) -> usize { + self.rules + .groups + .iter() + .zip(&hit.per_group) + .find(|(rule, _)| rule.kind == GroupRuleKind::FullAttention && rule.cacheable) + .map_or(0, |(rule, blocks)| match blocks { + GroupHitBlocks::Attention(keys) => { + keys.iter() + .flatten() + .filter(|key| key_is_resident(pool, **key) && !key_is_inactive(pool, **key)) + .count() + * rule.block + } + _ => 0, + }) + } + /// The blocks a new request needs for the whole prompt, counting inactive /// hit blocks again (`single:216-268`, `single:1716-1816`, `mgr:519-536`). pub(crate) fn admission_blocks( @@ -764,6 +831,34 @@ impl GroupedPrefixCoordinator { pool: &GroupedBlockPool, hit: &GroupedHit, prompt_tokens: usize, + ) -> usize { + self.blocks_to_allocate(pool, hit, prompt_tokens, prompt_tokens, true) + } + + /// The blocks that the first allocation of a request takes for its tokens + /// up to `target`, counting inactive hit blocks again: vLLM's second test + /// in `allocate_slots`, on the chunk and without the admission cap of a + /// sliding window (`mgr:558-576`). + pub(crate) fn first_allocation_blocks( + &self, + pool: &GroupedBlockPool, + hit: &GroupedHit, + prompt: usize, + target: usize, + ) -> usize { + self.blocks_to_allocate(pool, hit, target, prompt, false) + } + + /// vLLM's `get_num_blocks_to_allocate` for the first `tokens` tokens of a + /// request without blocks. Only a chunk that reaches `prompt` takes a + /// checkpoint block (v0.32 `single:1643-1674`). + fn blocks_to_allocate( + &self, + pool: &GroupedBlockPool, + hit: &GroupedHit, + tokens: usize, + prompt: usize, + admission_cap: bool, ) -> usize { let rules = &self.rules; let mut total = 0; @@ -775,15 +870,20 @@ impl GroupedPrefixCoordinator { GroupHitBlocks::Attention(keys) => keys, _ => &[], }; - let mut required = prompt_tokens.div_ceil(block); + let mut required = tokens.div_ceil(block); let mut skipped_blocks = 0; if let GroupRuleKind::SlidingWindow { window } = rule.kind { // The admission cap of `SlidingWindowSpec` // (`iface:859-866`), and the blocks below the window // at the hit (`single:1153-1188`). - let retained = (window - 1 + self.limits.max_in_flight_tokens) - .min(self.limits.max_model_len); - required = required.min(retained.div_ceil(block) + 1); + if admission_cap { + required = required.min(window_admission_blocks( + window, + block, + self.limits.max_in_flight_tokens, + self.limits.max_model_len, + )); + } skipped_blocks = hit.tokens.saturating_sub(window - 1) / block; } let new = required.saturating_sub(skipped_blocks.max(keys.len())); @@ -803,23 +903,24 @@ impl GroupedPrefixCoordinator { _ => None, }; let inactive = key.is_some_and(|key| key_is_inactive(pool, key)); - // `cdiv(prompt, B) - len(computed)`, plus the copy block of + // `cdiv(tokens, B) - len(computed)`, plus the copy block of // a hit inside a block, capped at `1 + partial + checkpoint` // (`single:1757-1777`, `:1806-1811`, v0.32 `single:1770`). // A valid checkpoint needs two new blocks or more, so one // new block and the checkpoint block below reach the cap. let copy = key.is_some() && !hit.tokens.is_multiple_of(block); - let new = prompt_tokens + let new = tokens .div_ceil(block) .saturating_sub(hit.tokens.div_ceil(block)) .min(1) + usize::from(copy); // v0.32 `single:1750-1758`, `:1770`. - let checkpoint = rules.checkpoint_valid( - hit.tokens, - prompt_tokens, - rules.checkpoint_position(prompt_tokens), - ); + let checkpoint = tokens >= prompt + && rules.checkpoint_valid( + hit.tokens, + tokens, + rules.checkpoint_position(tokens), + ); new + usize::from(inactive) + usize::from(checkpoint) } GroupRuleKind::Ring => 1, @@ -1263,1702 +1364,5 @@ enum Plan { } #[cfg(test)] -mod tests { - use super::*; - use crate::engine::cache::vllm_block_pool::CacheKey; - use crate::engine::common::hashing::SequenceHash; - use crate::engine::kv_manager::grouped::{GroupLease, GroupedBlockPool, GroupedKvPool}; - use crate::engine::kv_manager::vllm_backend::BlockRequestLease; - use crate::perfmodel::{FpmCacheGroup, FpmCacheKind}; - use uuid::Uuid; - - const BUDGET: usize = 8192; - const GLM_BLOCK: u32 = 2176; - - fn group( - name: &str, - kind: FpmCacheKind, - block: u32, - window: Option, - state: Option<(bool, u32)>, - ) -> FpmCacheGroup { - FpmCacheGroup { - name: name.into(), - kind, - num_layers: 1, - block_size_tokens: block, - page_size_bytes: 1, - sliding_window: window, - blocks_per_request: state.map(|_| 1), - prefix_cacheable: Some(state.is_none_or(|(cacheable, _)| cacheable)), - transfer_bytes_per_block: None, - prefill_checkpoint_blocks: state - .and_then(|(cacheable, ckpt)| cacheable.then_some(ckpt)), - prefill_checkpoint_alignment: state - .and_then(|(cacheable, ckpt)| (cacheable && ckpt == 1).then_some(16)), - } - } - - fn attention(name: &str, block: u32, window: Option) -> FpmCacheGroup { - group(name, FpmCacheKind::Attention, block, window, None) - } - - fn state(name: &str, block: u32, cacheable: bool, checkpoint: bool) -> FpmCacheGroup { - group( - name, - FpmCacheKind::State, - block, - None, - Some((cacheable, u32::from(checkpoint))), - ) - } - - /// The six GLM-5.3-Flash TP 4 groups that the runtime probe recorded on - /// vLLM 0.31.0. With `checkpoint`, G2 to G5 carry - /// `prefill_checkpoint_blocks: Some(1)` and alignment 16. - fn groups_glm(checkpoint: bool) -> Vec { - let mut groups = vec![ - attention("mla", GLM_BLOCK, None), - state("kpool_tail", 4, false, false), - ]; - for i in 0..4 { - groups.push(state(&format!("kda_{i}"), GLM_BLOCK, true, checkpoint)); - } - groups - } - - fn rules_glm_hash( - version: VllmRuleVersion, - checkpoint: bool, - hash_block: usize, - ) -> GroupedPrefixRules { - GroupedPrefixRules::new(&groups_glm(checkpoint), hash_block, version).unwrap() - } - - fn rules_glm(version: VllmRuleVersion, checkpoint: bool) -> GroupedPrefixRules { - rules_glm_hash(version, checkpoint, GLM_BLOCK as usize) - } - - fn rules(groups: &[FpmCacheGroup], hash_block: usize) -> GroupedPrefixRules { - GroupedPrefixRules::new(groups, hash_block, VllmRuleVersion::V0_31_0).unwrap() - } - - /// A block pool with `blocks` pool blocks for `groups`. - fn kv_pool(groups: &[FpmCacheGroup], blocks: u64) -> GroupedKvPool { - GroupedKvPool::with_prefix_cache(groups.to_vec(), blocks, 1) - } - - /// A lease with chained hashes of the tokens `0..tokens` at `hash_block`. - fn lease(tokens: usize, hash_block: usize) -> BlockRequestLease { - let tokens: Vec = (0..tokens as u32).collect(); - let mut lease = BlockRequestLease::new(Uuid::nil(), Vec::new()); - lease.configure_prefix_hashes(&tokens, hash_block); - lease - } - - fn hash_at(hash_block: usize, tokens: usize) -> SequenceHash { - lease(tokens, hash_block).prefix_hash_at(tokens).unwrap() - } - - fn key(group: u16, hash_block: usize, tokens: usize) -> CacheKey { - CacheKey::Grouped { - group, - hash: hash_at(hash_block, tokens), - } - } - - /// Publish one block under `key` and leave it in the inactive LRU, as a - /// finished request leaves its cached blocks. - fn seed(pool: &mut GroupedBlockPool, key: CacheKey) { - let mut reservation = pool.reserve(&[], 1).unwrap(); - let id = pool.take_private(&mut reservation); - assert!(pool.cache(id, key)); - pool.release(id); - } - - fn seed_many(pool: &mut GroupedBlockPool, group: u16, hash_block: usize, ends: &[usize]) { - for &end in ends { - seed(pool, key(group, hash_block, end)); - } - } - - fn seed_states(pool: &mut GroupedBlockPool, groups: &[u16], hash_block: usize, end: usize) { - for &group in groups { - seed(pool, key(group, hash_block, end)); - } - } - - fn hit_exists(pool: &GroupedBlockPool, group: u16, hash_block: usize, tokens: usize) -> bool { - pool.key_hit(group, hash_at(hash_block, tokens)).is_some() - } - - fn state_lease(lease: &GroupedBlockLease, group: usize) -> &GroupStateLease { - match &lease.groups[group] { - GroupLease::State(state) => state, - other => panic!("group {group} is not a state lease: {other:?}"), - } - } - - fn attention_copies(lease: &GroupedBlockLease, group: usize) -> usize { - match &lease.groups[group] { - GroupLease::Attention { copies, .. } => copies.iter().flatten().count(), - other => panic!("group {group} is not an attention lease: {other:?}"), - } - } - - fn split( - rules: &GroupedPrefixRules, - start: usize, - prompt: usize, - known: usize, - desired: usize, - junction: usize, - ) -> usize { - rules.align_split(start, prompt, known, desired, junction, BUDGET) - } - - /// Run a whole prefill of `prompt` tokens from `hit`: the chunks of - /// `align_split` at `budget`, one step each. Returns the chunk ends. - fn prefill( - coordinator: &mut GroupedPrefixCoordinator, - pool: &mut GroupedBlockPool, - lease: &mut GroupedBlockLease, - request: &BlockRequestLease, - prompt: usize, - hit: &GroupedHit, - budget: usize, - ) -> Vec { - let mut ends = Vec::new(); - let mut computed = hit.tokens; - let mut step = 1; - while computed < prompt { - coordinator.begin_step(); - let chunk = coordinator.rules().align_split( - computed, - prompt, - prompt, - budget, - hit.tokens + hit.junction_tokens, - budget, - ); - assert!(chunk > 0, "the split must make progress at {computed}"); - let target = computed + chunk; - assert!(coordinator.allocate(pool, lease, prompt, computed, target, step)); - coordinator.finalize(pool, lease, request, prompt, target, hit, step); - computed = target; - ends.push(target); - step += 1; - } - coordinator.begin_step(); - ends - } - - const GLM_STATES: [u16; 4] = [2, 3, 4, 5]; - - /// One batch of `BUDGET` tokens in flight and no model-length cap. - const LIMITS: AdmissionLimits = AdmissionLimits { - max_in_flight_tokens: BUDGET, - max_model_len: usize::MAX, - }; - - fn with_limits(rules: GroupedPrefixRules) -> GroupedPrefixCoordinator { - GroupedPrefixCoordinator::new(rules, LIMITS) - } - - #[test] - fn version_rows_differ_in_the_last_cacheable_stop_and_the_partial_tail() { - // v0.31 sched:444 `num_tokens - num_tokens % block_size`; v0.32 - // sched:476-480 with iface:1110-1122 `(n - 1) // B * B`. - let v31 = rules_glm(VllmRuleVersion::V0_31_0, false); - let v32 = rules_glm(VllmRuleVersion::V0_32_0, false); - assert_eq!(v31.last_cacheable_stop(9000), 8704); - assert_eq!(v32.last_cacheable_stop(9000), 8704); - assert_eq!(v31.last_cacheable_stop(8704), 8704); - assert_eq!(v32.last_cacheable_stop(8704), 6528); - assert_eq!(v32.last_cacheable_stop(0), 0); - // v0.31 sched:482-487 `prompt // hash * hash`; v0.32 sched:517-525 - // `(prompt - 1) // hash * hash`. - let v31 = rules_glm_hash(VllmRuleVersion::V0_31_0, false, 4); - let v32 = rules_glm_hash(VllmRuleVersion::V0_32_0, false, 4); - assert_eq!(v31.partial_tail(9000), 9000); - assert_eq!(v32.partial_tail(9000), 8996); - assert_eq!(v31.partial_tail(4356), 4356); - assert_eq!(v32.partial_tail(4356), 4352); - assert_eq!( - VllmRuleVersion::parse("0.31.0"), - Some(VllmRuleVersion::V0_31_0) - ); - assert_eq!( - VllmRuleVersion::parse("0.32.0"), - Some(VllmRuleVersion::V0_32_0) - ); - assert_eq!(VllmRuleVersion::parse("0.30.0"), None); - assert_eq!(VllmRuleVersion::V0_31_0.as_str(), "0.31.0"); - assert_eq!(VllmRuleVersion::V0_32_0.as_str(), "0.32.0"); - } - - #[test] - fn rules_refuse_unsupported_checkpoint_fields() { - // v0.32 sched:390-394 runs checkpoint mode only when every Mamba - // group has a positive count. - let err = GroupedPrefixRules::new(&groups_glm(true), 2176, VllmRuleVersion::V0_31_0) - .unwrap_err() - .to_string(); - assert_eq!( - err, - "kv_cache_groups with vllm_version=0.31.0 do not model prefill checkpoint blocks" - ); - let mut groups = groups_glm(false); - groups[2].prefill_checkpoint_blocks = None; - let err = GroupedPrefixRules::new(&groups, 2176, VllmRuleVersion::V0_32_0) - .unwrap_err() - .to_string(); - assert_eq!( - err, - "kv_cache_groups with vllm_version=0.32.0 require prefill_checkpoint_blocks on every prefix-cacheable state group; re-import the runtime probe" - ); - let mut groups = groups_glm(true); - groups[3].prefill_checkpoint_alignment = Some(32); - let err = GroupedPrefixRules::new(&groups, 2176, VllmRuleVersion::V0_32_0) - .unwrap_err() - .to_string(); - assert_eq!( - err, - "kv_cache_groups with vllm_version=0.32.0 require one prefill_checkpoint_alignment across state groups" - ); - let zero = rules_glm(VllmRuleVersion::V0_32_0, false); - assert!(zero.checkpoint.is_none()); - let one = rules_glm(VllmRuleVersion::V0_32_0, true); - assert_eq!(one.checkpoint, Some(CheckpointRule { alignment: 16 })); - // Block mode holds one running block per state group, so another count - // is refused. - let mut groups = groups_glm(false); - groups[4].blocks_per_request = Some(2); - let err = GroupedPrefixRules::new(&groups, 2176, VllmRuleVersion::V0_31_0) - .unwrap_err() - .to_string(); - assert_eq!( - err, - "kv_cache_groups with enable_prefix_caching=true model one state block per request; blocks_per_request must be 1" - ); - // utils:799-805 and coord:87-92. - assert!(GroupedPrefixRules::new(&groups_glm(false), 3, VllmRuleVersion::V0_31_0).is_err()); - } - - #[test] - fn rules_refuse_convolution_groups() { - // A convolution group keeps a window, and the table has no rule for it. - let mut convolution = attention("convolution", 4, Some(4)); - convolution.kind = FpmCacheKind::Convolution; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - for groups in [ - vec![convolution.clone()], - vec![attention("full", 4, None), convolution.clone()], - ] { - let err = GroupedPrefixRules::new(&groups, 4, version) - .unwrap_err() - .to_string(); - assert_eq!( - err, - "grouped prefix caching does not support convolution cache groups; turn prefix caching off to replay them cold", - "{version:?}" - ); - } - } - } - - #[test] - fn alignment_is_the_hash_block_only_with_fine_grained_hits() { - // coord:687-735. G2 to G5 are state rules with block 2176 and G0 is - // full attention. At hash 2176 no state block exceeds the hash block. - let glm = rules_glm(VllmRuleVersion::V0_31_0, false); - assert_eq!(glm.alignment, 2176); - assert_eq!(glm.scheduler_block, 2176); - assert_eq!(glm.cacheable_block, 2176); - let fine = rules_glm_hash(VllmRuleVersion::V0_31_0, false, 4); - assert_eq!(fine.alignment, 4); - assert_eq!(fine.scheduler_block, 2176); - // Fine-grained hits need a Mamba group in align mode whose block exceeds - // the hash block. Two attention groups have none, so they give the - // scheduler block, with or without a window (coord:691-704). - let windowed = [attention("a", 64, Some(100)), attention("b", 16, None)]; - assert_eq!(rules(&windowed, 16).alignment, 64); - let plain = [attention("a", 64, None), attention("b", 16, None)]; - assert_eq!(rules(&plain, 16).alignment, 64); - // A state group above the hash block turns fine-grained hits on when - // every other cacheable rule is full attention or at the hash block. - let with_state = [ - attention("a", 64, None), - attention("b", 16, None), - state("s", 64, true, false), - ]; - assert_eq!(rules(&with_state, 16).alignment, 16); - // A window at another block than the hash block turns them off again - // (single:960 has no sub-block lookup). - let with_window = [ - attention("a", 64, Some(100)), - attention("b", 16, None), - state("s", 64, true, false), - ]; - assert_eq!(rules(&with_window, 16).alignment, 64); - } - - #[test] - fn full_attention_hit_is_capped_below_the_prompt() { - // mgr:289-295 caps the lookup at `num_tokens - 1`. - let groups = [attention("full", 16, None)]; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let coordinator = with_limits(GroupedPrefixRules::new(&groups, 16, version).unwrap()); - let mut kv = kv_pool(&groups, 16); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, 16, &[16, 32, 48, 64]); - let hit = coordinator.find_longest_hit(pool, &lease(64, 16), 64); - assert_eq!(hit.tokens, 48); - assert_eq!(hit.junction_tokens, 0); - assert_eq!(hit.per_group.len(), 1); - assert!( - matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 3) - ); - let hit = coordinator.find_longest_hit(pool, &lease(65, 16), 65); - assert_eq!(hit.tokens, 64); - } - } - - #[test] - fn sliding_window_hits_need_a_contiguous_run_before_an_aligned_end() { - // single:960-1069: `cdiv(32, 16) == 2` contiguous blocks before an - // aligned end, or a run from block 0. - let groups = [ - attention("full", 16, None), - attention("window", 16, Some(33)), - ]; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let r = GroupedPrefixRules::new(&groups, 16, version).unwrap(); - assert_eq!(r.alignment, 16); - let coordinator = with_limits(r); - let mut kv = kv_pool(&groups, 32); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, 16, &[16, 32, 48, 64, 80, 96]); - seed_many(pool, 1, 16, &[32, 48, 64, 96]); - let hit = coordinator.find_longest_hit(pool, &lease(96, 16), 96); - // The end 64 has blocks 2 and 3 before it. The end 80 lacks block 4. - assert_eq!(hit.tokens, 64); - match &hit.per_group[1] { - GroupHitBlocks::Attention(keys) => { - assert_eq!(keys.len(), 4); - assert!(keys[0].is_none() && keys[1].is_none()); - assert_eq!(keys[2], Some(key(1, 16, 48))); - assert_eq!(keys[3], Some(key(1, 16, 64))); - } - other => panic!("{other:?}"), - } - // Junction: the full-attention hit before the loop (80) minus 64. - assert_eq!(hit.junction_tokens, 16); - - let mut kv = kv_pool(&groups, 32); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, 16, &[16, 32, 48, 64, 80, 96]); - seed_many(pool, 1, 16, &[16]); - let hit = coordinator.find_longest_hit(pool, &lease(96, 16), 96); - assert_eq!(hit.tokens, 16); - } - } - - #[test] - fn one_token_window_hits_only_at_a_cached_block() { - // A window of one token needs `cdiv(0, 4) == 0` contiguous blocks - // (v0.31 single:961-971, v0.32 :943-953). vLLM counts a block only - // after it finds the block cached, so a hit still needs the cached - // block at its end (v0.31 single:1021-1044, v0.32 :1003-1026). - let groups = [attention("window", 4, Some(1))]; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let coordinator = with_limits(GroupedPrefixRules::new(&groups, 4, version).unwrap()); - let mut kv = kv_pool(&groups, 32); - let pool = kv.block_pool_mut().unwrap(); - let hit = coordinator.find_longest_hit(pool, &lease(100, 4), 100); - assert_eq!((hit.tokens, hit.junction_tokens), (0, 0), "{version:?}"); - seed_many(pool, 0, 4, &[48]); - let hit = coordinator.find_longest_hit(pool, &lease(100, 4), 100); - assert_eq!((hit.tokens, hit.junction_tokens), (48, 0), "{version:?}"); - match &hit.per_group[0] { - GroupHitBlocks::Attention(keys) => { - assert_eq!(keys.len(), 12, "{version:?}"); - assert!(keys[..11].iter().all(Option::is_none), "{version:?}"); - assert_eq!(keys[11], Some(key(0, 4, 48)), "{version:?}"); - } - other => panic!("{other:?}"), - } - } - } - - #[test] - fn sliding_window_admission_reserves_the_window_plus_the_in_flight_tokens() { - // iface:859-866: `cdiv(min(window - 1 + max_in_flight_tokens, - // max_model_len), block) + 1`, with `max_in_flight_tokens = - // max_concurrent_batches * max_num_batched_tokens` (vllm.py:605-613). - let groups = [ - attention("full", 16, None), - attention("window", 16, Some(33)), - ]; - let prompt: usize = 100_000; - let hit = GroupedHit { - tokens: 0, - junction_tokens: 0, - per_group: vec![ - GroupHitBlocks::Attention(Vec::new()), - GroupHitBlocks::Attention(Vec::new()), - ], - }; - let full_blocks = prompt.div_ceil(16); - let kv = kv_pool(&groups, 16); - let pool = kv.block_pool().unwrap(); - // One batch of 8,192: (32 + 8192) / 16 + 1 = 515. - let one_batch = with_limits(rules(&groups, 16)); - assert_eq!( - one_batch.admission_blocks(pool, &hit, prompt), - full_blocks + 515 - ); - // Two batches of 8,192: (32 + 16384) / 16 + 1 = 1027. - let two_batches = GroupedPrefixCoordinator::new( - rules(&groups, 16), - AdmissionLimits { - max_in_flight_tokens: 2 * BUDGET, - max_model_len: usize::MAX, - }, - ); - assert_eq!( - two_batches.admission_blocks(pool, &hit, prompt), - full_blocks + 1027 - ); - // The model length caps it: cdiv(4096, 16) + 1 = 257. - let capped = GroupedPrefixCoordinator::new( - rules(&groups, 16), - AdmissionLimits { - max_in_flight_tokens: 2 * BUDGET, - max_model_len: 4096, - }, - ); - assert_eq!( - capped.admission_blocks(pool, &hit, prompt), - full_blocks + 257 - ); - // A short prompt needs fewer blocks than the cap. - assert_eq!(one_batch.admission_blocks(pool, &hit, 100), 7 + 7); - } - - #[test] - fn state_group_hits_only_at_a_kept_boundary() { - // single:1487-1564 scans aligned ends right to left for one cached - // state block. pool:198-224: every group of the spec group must hold - // the key. - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let groups = groups_glm(false); - let coordinator = with_limits(rules_glm(version, false)); - let hash = 2176; - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); - seed_states(pool, &GLM_STATES, hash, 8704); - let request = lease(9000, hash); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - assert_eq!(hit.tokens, 8704); - assert_eq!(hit.junction_tokens, 0); - assert_eq!( - hit.per_group[2], - GroupHitBlocks::State(Some(key(2, hash, 8704))) - ); - assert_eq!(hit.per_group[1], GroupHitBlocks::None); - - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); - seed_states(pool, &[2, 4, 5], hash, 8704); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - assert_eq!(hit.tokens, 0); - assert_eq!(hit.per_group[0], GroupHitBlocks::Attention(Vec::new())); - assert_eq!(hit.per_group[3], GroupHitBlocks::State(None)); - - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); - seed_states(pool, &GLM_STATES, hash, 6528); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - assert_eq!(hit.tokens, 6528); - assert!( - matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 3) - ); - } - } - - #[test] - fn ring_group_never_hits_and_holds_one_block() { - // single:1214-1230 claims one block per request; single:498-499 - // caches nothing for a group that is not prefix cacheable. - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let groups = groups_glm(false); - let mut coordinator = with_limits(rules_glm(version, false)); - let hash = 2176; - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(2500, hash); - let hit = coordinator.find_longest_hit(pool, &request, 2500); - assert_eq!(hit.tokens, 0); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 2500, 0, 2176, 1)); - assert_eq!( - block_lease.groups[1], - GroupLease::Ring { - copy: Some(ring_copy(&block_lease)) - } - ); - coordinator.finalize(pool, &mut block_lease, &request, 2500, 2176, &hit, 1); - assert_eq!(pool.inactive_cached_blocks(), 0); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 2500, 2176, 2500, 2)); - coordinator.finalize(pool, &mut block_lease, &request, 2500, 2500, &hit, 2); - coordinator.release_request(pool, &mut block_lease); - coordinator.begin_step(); - - let hit = coordinator.find_longest_hit(pool, &lease(2500, hash), 2500); - assert_eq!(hit.tokens, 2176); - assert_eq!(hit.per_group[1], GroupHitBlocks::None); - assert!( - matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 1) - ); - assert_eq!( - hit.per_group[2], - GroupHitBlocks::State(Some(key(2, hash, 2176))) - ); - } - } - - fn ring_copy(lease: &GroupedBlockLease) -> crate::engine::cache::vllm_block_pool::BlockCopyId { - match &lease.groups[1] { - GroupLease::Ring { copy: Some(copy) } => *copy, - other => panic!("{other:?}"), - } - } - - #[test] - fn hybrid_hit_is_the_longest_prefix_every_cacheable_group_serves() { - // coord:842-977: the fixed point shortens to 4,352; coord:970-977 - // reports the uncached common prefix. - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let groups = groups_glm(false); - let coordinator = with_limits(rules_glm(version, false)); - let hash = 2176; - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704, 10880]); - seed_states(pool, &GLM_STATES, hash, 4352); - let hit = coordinator.find_longest_hit(pool, &lease(11000, hash), 11000); - assert_eq!(hit.tokens, 4352); - assert_eq!(hit.junction_tokens, 10880 - 4352); - assert!( - matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 2) - ); - } - } - - #[test] - fn hybrid_partial_hash_truncates_full_attention_hit_length() { - // Ported from v0.31 partial_test:1596. Values hand-derived: the - // fine-grained probe of single:806-825 extends the two full blocks to - // the partial entry at 4,356, and the state hit there keeps it. - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let groups = groups_glm(false); - let coordinator = with_limits(rules_glm_hash(version, false, 4)); - let hash = 4; - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 4356]); - seed_states(pool, &GLM_STATES, hash, 4356); - let request = lease(4500, hash); - let hit = coordinator.find_longest_hit(pool, &request, 4500); - assert_eq!(hit.tokens, 4356); - assert_eq!(hit.junction_tokens, 0); - match &hit.per_group[0] { - GroupHitBlocks::Attention(keys) => { - assert_eq!(keys.len(), 3); - assert_eq!(keys[2], Some(key(0, hash, 4356))); - } - other => panic!("{other:?}"), - } - - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 4356]); - seed_states(pool, &GLM_STATES, hash, 4352); - let hit = coordinator.find_longest_hit(pool, &request, 4500); - assert_eq!(hit.tokens, 4352); - // coord:970-977: the full-attention hit of 4,356 minus 4,352. - assert_eq!(hit.junction_tokens, 4); - assert!( - matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 2) - ); - } - } - - /// Two full-attention groups of one spec and a state group, all at block 64. - fn two_full_groups_and_a_state() -> [FpmCacheGroup; 3] { - [ - attention("full_a", 64, None), - attention("full_b", 64, None), - state("state", 64, true, false), - ] - } - - #[test] - fn groups_of_one_spec_hit_jointly_and_set_the_junction_together() { - // coord:737-779 (v0.32 coord:708-750) looks up the groups of one spec - // together, and pool:198-224 finds a block only when every group of the - // spec holds it. The longest hit is the largest result of one spec's - // lookup (coord:947, :970-977). Group a holds 192 tokens, group b 64 and - // the state group the state at 64, so vLLM hits 64 with no junction. - let groups = two_full_groups_and_a_state(); - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let coordinator = with_limits(GroupedPrefixRules::new(&groups, 64, version).unwrap()); - let mut kv = kv_pool(&groups, 16); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, 64, &[64, 128, 192]); - seed_many(pool, 1, 64, &[64]); - seed_many(pool, 2, 64, &[64]); - let hit = coordinator.find_longest_hit(pool, &lease(401, 64), 401); - assert_eq!((hit.tokens, hit.junction_tokens), (64, 0), "{version:?}"); - assert_eq!( - hit.per_group[0], - GroupHitBlocks::Attention(vec![Some(key(0, 64, 64))]) - ); - assert_eq!( - hit.per_group[1], - GroupHitBlocks::Attention(vec![Some(key(1, 64, 64))]) - ); - assert_eq!( - hit.per_group[2], - GroupHitBlocks::State(Some(key(2, 64, 64))) - ); - // So the first chunk runs to the last cacheable stop 384 (v0.31 - // sched:444, v0.32 sched:476-480). A junction at 192 would end it - // there. - let junction = hit.tokens + hit.junction_tokens; - assert_eq!( - split(coordinator.rules(), 64, 401, 401, BUDGET, junction), - 320, - "{version:?}" - ); - } - } - - #[test] - fn groups_of_one_spec_hit_inside_a_block_only_jointly() { - // With fine-grained hits, single:794-828 probes the hash boundaries of the - // first partial block for every group of the spec at once. Group a holds - // two full blocks, group b one full block and the partial entry at 96. - // No boundary inside the second block is cached in both, so vLLM hits 64. - let groups = two_full_groups_and_a_state(); - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let rules = GroupedPrefixRules::new(&groups, 16, version).unwrap(); - assert_eq!(rules.alignment, 16); - let coordinator = with_limits(rules); - let mut kv = kv_pool(&groups, 16); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, 16, &[64, 128]); - seed_many(pool, 1, 16, &[64, 96]); - seed_many(pool, 2, 16, &[64, 96]); - let hit = coordinator.find_longest_hit(pool, &lease(200, 16), 200); - assert_eq!((hit.tokens, hit.junction_tokens), (64, 0), "{version:?}"); - assert_eq!( - hit.per_group[0], - GroupHitBlocks::Attention(vec![Some(key(0, 16, 64))]) - ); - assert_eq!( - hit.per_group[1], - GroupHitBlocks::Attention(vec![Some(key(1, 16, 64))]) - ); - assert_eq!( - hit.per_group[2], - GroupHitBlocks::State(Some(key(2, 16, 64))) - ); - } - } - - #[test] - fn state_groups_of_one_spec_hit_jointly() { - // Without a full-attention group, the first spec looked up is the first - // one in group order (coord:775-779), here the two state groups. They - // hold the states at 64 and 192, and at 64. The joint lookup - // (single:1540-1564) hits 64 and leaves no junction. - let groups = [ - state("state_a", 64, true, false), - state("state_b", 64, true, false), - attention("window", 64, Some(65)), - ]; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let coordinator = with_limits(GroupedPrefixRules::new(&groups, 64, version).unwrap()); - let mut kv = kv_pool(&groups, 16); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, 64, &[64, 192]); - seed_many(pool, 1, 64, &[64]); - seed_many(pool, 2, 64, &[64]); - let hit = coordinator.find_longest_hit(pool, &lease(401, 64), 401); - assert_eq!((hit.tokens, hit.junction_tokens), (64, 0), "{version:?}"); - assert_eq!( - hit.per_group[0], - GroupHitBlocks::State(Some(key(0, 64, 64))) - ); - assert_eq!( - hit.per_group[1], - GroupHitBlocks::State(Some(key(1, 64, 64))) - ); - } - // On GLM-5.3-Flash the full-attention group G0 comes first, so its hit - // sets the junction, as in vLLM. The extra state of G2 at 8,704 changes - // nothing, and the four KDA groups hit 4,352 together. - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let groups = groups_glm(false); - let coordinator = with_limits(rules_glm(version, false)); - let hash = 2176; - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); - seed_states(pool, &GLM_STATES, hash, 4352); - seed_states(pool, &[2], hash, 8704); - let hit = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); - assert_eq!( - (hit.tokens, hit.junction_tokens), - (4352, 8704 - 4352), - "{version:?}" - ); - for group in GLM_STATES { - assert_eq!( - hit.per_group[usize::from(group)], - GroupHitBlocks::State(Some(key(group, hash, 4352))) - ); - } - } - } - - #[test] - fn admission_counts_inactive_hits_and_state_copies() { - // single:216-268 counts inactive hit blocks again. The keys are seeded - // directly, because a finished request of 4,352 tokens keeps its state - // key only at its replay boundary, 2,176. - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let groups = groups_glm(false); - let coordinator = with_limits(rules_glm(version, false)); - let hash = 2176; - let mut kv = kv_pool(&groups, 16); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352]); - seed_states(pool, &GLM_STATES, hash, 4352); - let hit = coordinator.find_longest_hit(pool, &lease(4400, hash), 4400); - assert_eq!(hit.tokens, 4352); - // 2 inactive G0 blocks, 1 new G0 block, 1 ring, 4 x (running + - // inactive cached state block). - assert_eq!(coordinator.admission_blocks(pool, &hit, 4400), 12); - // Another running request holds the two G0 blocks active. - let other = [key(0, hash, 2176), key(0, hash, 4352)]; - let mut reservation = pool.reserve(&other, 0).unwrap(); - let held = pool.activate(&mut reservation); - assert_eq!(held.len(), 2); - assert_eq!(coordinator.admission_blocks(pool, &hit, 4400), 10); - pool.cancel(reservation); - } - } - - #[test] - fn resumed_state_holds_cached_and_running_blocks_during_the_first_step() { - // single:1858-1868: a resumed request holds the hit block and the - // running block. single:1658-1675 frees the hit block once - // `cdiv(processed, block) - 1 > hit_index`: `cdiv(6000, 2176) - 1 == - // 2 > 1`. - let groups = groups_glm(false); - let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); - let hash = 2176; - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352]); - seed_states(pool, &GLM_STATES, hash, 4352); - let request = lease(9000, hash); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - assert_eq!(hit.tokens, 4352); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); - for group in GLM_STATES { - let state = state_lease(&block_lease, usize::from(group)); - assert_eq!(state.hit.map(|(_, index)| index), Some(1)); - assert!(state.running.is_none()); - } - assert!(coordinator.allocate(pool, &mut block_lease, 9000, 4352, 6000, 1)); - let state_blocks: usize = GLM_STATES - .iter() - .map(|&group| { - let state = state_lease(&block_lease, usize::from(group)); - assert!(state.hit.is_some() && state.running.is_some()); - assert!(state.previous.is_none() && state.copy.is_none()); - 2 - }) - .sum(); - assert_eq!(state_blocks, 8); - // G0: 2 hits + 1 new, ring 1, states 8. - assert_eq!(pool.active_blocks(), 12); - coordinator.finalize(pool, &mut block_lease, &request, 9000, 6000, &hit, 1); - for group in GLM_STATES { - let state = state_lease(&block_lease, usize::from(group)); - assert!(state.hit.is_none() && state.running.is_some()); - } - assert_eq!(pool.active_blocks(), 8); - assert_eq!(pool.inactive_cached_blocks(), 4); - } - - #[test] - fn boundary_crossing_keeps_the_previous_running_block_for_one_step() { - // single:1658-1675 and mamba_utils:576-584: the block of the earlier - // boundary goes at the next allocation. One pass, as on the linear path, - // where vLLM's async scheduling holds it one more. - let groups = groups_glm(false); - let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); - let hash = 2176; - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(9000, hash); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 9000, 0, 2176, 1)); - coordinator.finalize(pool, &mut block_lease, &request, 9000, 2176, &hit, 1); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 9000, 2176, 4352, 2)); - let state_and_ring = |lease: &GroupedBlockLease| -> usize { - lease.groups[1..] - .iter() - .map(|g| match g { - GroupLease::Attention { .. } => 0, - GroupLease::Ring { copy } => usize::from(copy.is_some()), - GroupLease::State(s) => [ - s.running, - s.hit.map(|h| h.0), - s.previous.map(|p| p.0), - s.copy, - s.checkpoint.map(|c| c.0), - ] - .iter() - .flatten() - .count(), - }) - .sum() - }; - for group in GLM_STATES { - let state = state_lease(&block_lease, usize::from(group)); - assert!(state.running.is_some()); - assert_eq!(state.previous.map(|(_, step)| step), Some(2)); - } - assert_eq!(state_and_ring(&block_lease), 9); - coordinator.finalize(pool, &mut block_lease, &request, 9000, 4352, &hit, 2); - assert_eq!(state_and_ring(&block_lease), 9); - coordinator.begin_step(); - let before = pool.active_blocks(); - assert!(coordinator.allocate(pool, &mut block_lease, 9000, 4352, 6528, 3)); - // Four previous blocks go, four new running blocks and one G0 block come. - assert_eq!(pool.active_blocks(), before + 1); - for group in GLM_STATES { - let state = state_lease(&block_lease, usize::from(group)); - assert_eq!(state.previous.map(|(_, step)| step), Some(3)); - } - assert_eq!(state_and_ring(&block_lease), 9); - } - - #[test] - fn partial_hit_takes_a_copy_block_for_the_shared_block() { - // single:387-409 and single:1770-1777: a hit that ends inside a block - // takes one copy block. single:800-805: the full-attention probe needs - // the full blocks below the partial entry, so blocks 0 and 1 are - // cached too. - let groups = [attention("full", 16, None), state("s", 32, true, false)]; - let r = rules(&groups, 4); - assert_eq!(r.alignment, 4); - let mut coordinator = with_limits(r); - let hash = 4; - let mut kv = kv_pool(&groups, 32); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[16, 32, 36]); - seed_many(pool, 1, hash, &[36]); - let request = lease(50, hash); - let hit = coordinator.find_longest_hit(pool, &request, 50); - assert_eq!(hit.tokens, 36); - // G0: 4 live - 3 hits + 3 inactive + 1 copy = 5. State - // (single:1757-1777): `cdiv(50, 32) - cdiv(36, 32) == 0` new blocks - // plus the copy block, plus the inactive hit = 2. - assert_eq!(coordinator.admission_blocks(pool, &hit, 50), 7); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); - assert_eq!(pool.active_blocks(), 4); - assert!(coordinator.allocate(pool, &mut block_lease, 50, 36, 50, 1)); - // G0 holds blocks 0 to 3, with block 2 as a private copy. The shared - // block 2 is inactive again. The target 50 stays in the hit's block, - // so the state's copy-on-write block is its running block - // (single:1891-1895, `_apply_cow` at single:454-474); the hit block - // stays until the copy has run (sched:1434-1442). - assert_eq!(attention_copies(&block_lease, 0), 4); - let state = state_lease(&block_lease, 1); - assert!(state.hit.is_some() && state.running.is_some() && state.copy.is_none()); - assert_eq!(pool.active_blocks(), 4 + 2); - assert_eq!(pool.inactive_cached_blocks(), 1); - coordinator.finalize(pool, &mut block_lease, &request, 50, 50, &hit, 1); - let state = state_lease(&block_lease, 1); - assert!(state.hit.is_none() && state.running.is_some()); - assert_eq!(pool.active_blocks(), 4 + 1); - assert_eq!(pool.inactive_cached_blocks(), 2); - - // A target beyond the hit's block: the copy sits at the hit's index - // and a running block follows. `remove_skipped_blocks` frees the - // copy at the next allocation (single:1658-1675, index set at :1868). - let mut kv = kv_pool(&groups, 32); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[16, 32, 36]); - seed_many(pool, 1, hash, &[36]); - let request = lease(70, hash); - let hit = coordinator.find_longest_hit(pool, &request, 70); - assert_eq!(hit.tokens, 36); - // G0: 5 live - 3 hits + 3 inactive + 1 copy = 6. State: `cdiv(70, 32) - // - cdiv(36, 32) == 1` new block plus the copy block plus the - // inactive hit = 3. - assert_eq!(coordinator.admission_blocks(pool, &hit, 70), 9); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); - assert!(coordinator.allocate(pool, &mut block_lease, 70, 36, 70, 1)); - let state = state_lease(&block_lease, 1); - assert!(state.hit.is_some() && state.running.is_some() && state.copy.is_some()); - assert_eq!(pool.active_blocks(), 5 + 3); - coordinator.finalize(pool, &mut block_lease, &request, 70, 70, &hit, 1); - let state = state_lease(&block_lease, 1); - assert!(state.hit.is_none() && state.copy.is_some()); - assert_eq!(pool.active_blocks(), 5 + 2); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 70, 70, 71, 2)); - let state = state_lease(&block_lease, 1); - assert!(state.copy.is_none() && state.running.is_some()); - assert_eq!(pool.active_blocks(), 5 + 1); - } - - #[test] - fn finalize_publishes_states_at_the_replay_boundary_and_the_junction() { - // single:1614-1622 with retention 0 (cache:160-166): the replay - // boundary `(9000 - 1) // 2176 * 2176 == 8704` and the junction. - let groups = groups_glm(false); - let hash = 2176; - let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(9000, hash); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - let mut block_lease = coordinator.new_lease(); - let ends = prefill( - &mut coordinator, - pool, - &mut block_lease, - &request, - 9000, - &hit, - 2176, - ); - assert_eq!(ends, vec![2176, 4352, 6528, 8704, 9000]); - coordinator.release_request(pool, &mut block_lease); - for group in GLM_STATES { - for end in [2176, 4352, 6528] { - assert!( - !hit_exists(pool, group, hash, end), - "group {group} at {end}" - ); - } - assert!(hit_exists(pool, group, hash, 8704)); - } - // G0 blocks 0 to 3 and four state blocks stay cached. - assert_eq!(pool.inactive_cached_blocks(), 8); - assert_eq!(pool.active_blocks(), 0); - - // A hit of 4,352 with a junction of 2,176. - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - seed_many(pool, 0, hash, &[2176, 4352, 6528]); - seed_states(pool, &GLM_STATES, hash, 4352); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - assert_eq!((hit.tokens, hit.junction_tokens), (4352, 2176)); - let mut block_lease = coordinator.new_lease(); - assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); - let ends = prefill( - &mut coordinator, - pool, - &mut block_lease, - &request, - 9000, - &hit, - 2176, - ); - assert_eq!(ends, vec![6528, 8704, 9000]); - coordinator.release_request(pool, &mut block_lease); - for group in GLM_STATES { - assert!(hit_exists(pool, group, hash, 6528), "group {group}"); - assert!(hit_exists(pool, group, hash, 8704), "group {group}"); - assert!(!hit_exists(pool, group, hash, 2176), "group {group}"); - } - } - - #[test] - fn align_split_follows_vllm_0310() { - // sched:414-527 at v0.31. - let r = rules_glm(VllmRuleVersion::V0_31_0, false); - // The floor of 8,192 to 2,176 is 6,528 (sched:460-468). - assert_eq!(split(&r, 0, 9000, 9000, 8192, 0), 6528); - // The last cacheable stop `9000 - 9000 % 2176 == 8704` (sched:444). - assert_eq!(split(&r, 6528, 9000, 9000, 8192, 0), 2176); - assert_eq!(split(&r, 8704, 9000, 9000, 8192, 0), 296); - // A prompt of exactly 8,704 tokens has no stop at 6,528, because its - // last cacheable stop `8704 - 8704 % 2176` is the prompt itself - // (sched:444). - assert_eq!(split(&r, 6528, 8704, 8704, 8192, 0), 2176); - // With hash 4 the tail 9,000 equals the prompt and adds no stop - // (sched:482-487, :512-514). With prompt 9,002 the tail is 9,000. - let fine = rules_glm_hash(VllmRuleVersion::V0_31_0, false, 4); - assert_eq!(split(&fine, 8704, 9000, 9000, 8192, 0), 296); - assert_eq!(split(&fine, 8704, 9002, 9002, 8192, 0), 296); - // The junction stop (sched:488-519). - assert_eq!(split(&r, 4352, 9000, 9000, 8192, 6528), 2176); - // Past the prefill the split returns the request (sched:434-436). - assert_eq!(split(&r, 9000, 9000, 9001, 1, 0), 1); - } - - #[test] - fn align_split_follows_vllm_0320() { - // v0.32 sched:443-556. - let r = rules_glm(VllmRuleVersion::V0_32_0, false); - assert_eq!(split(&r, 0, 9000, 9000, 8192, 0), 6528); - assert_eq!(split(&r, 6528, 9000, 9000, 2472, 0), 2176); - // The last stop of 8,704 tokens is `(8704 - 1) // 2176 * 2176 == 6528` - // (v0.32 sched:476-480). - assert_eq!(split(&r, 4352, 8704, 8704, 4352, 0), 2176); - let v31 = rules_glm(VllmRuleVersion::V0_31_0, false); - assert_eq!(split(&v31, 4352, 8704, 8704, 4352, 0), 4352); - // Checkpoint mode (v0.32 sched:482-501, iface:1125-1146). For - // (6528, 9000, pos 8704): 6528 % 2176 == 0; cdiv(9000, 2176) - 2 == 3 - // > 6527 / 2176 == 2; 6528 + 2176 <= 8704; 6528 < 8704 < 9000; - // 2176 % 16 == 0. - let ckpt = rules_glm(VllmRuleVersion::V0_32_0, true); - assert_eq!(ckpt.checkpoint_position(9000), 8704); - assert!(ckpt.checkpoint_valid(6528, 9000, 8704)); - assert_eq!(split(&ckpt, 6528, 9000, 9000, 2472, 0), 2472); - // For (4352, 8704, pos 6528): 4352 % 2176 == 0; cdiv(8704, 2176) - 2 - // == 2 > 4351 / 2176 == 1; 4352 + 2176 <= 6528; 4352 < 6528 < 8704; - // 2176 % 16 == 0. - assert_eq!(ckpt.checkpoint_position(8704), 6528); - assert!(ckpt.checkpoint_valid(4352, 8704, 6528)); - assert_eq!(split(&ckpt, 4352, 8704, 8704, 4352, 0), 4352); - // Without a checkpoint rule the check is false. - assert!(!r.checkpoint_valid(6528, 9000, 8704)); - } - - #[test] - fn align_split_with_a_checkpoint_block_ports_vllm_0320() { - // Ported from v0.32 split_test:127-180, which stubs `block_size` to - // 1,600 and `hash_block_size` to 16 (split_test:139-153). The engine - // shape with the same numbers: one full-attention rule and one state - // rule at block 1,600, hash 16, checkpoint 1 at alignment 16. The - // state block is above the hash block and the other rule is full - // attention, so fine-grained hits are on and A is 16. - let groups = [attention("full", 1600, None), state("s", 1600, true, true)]; - let r = GroupedPrefixRules::new(&groups, 16, VllmRuleVersion::V0_32_0).unwrap(); - assert_eq!(r.alignment, 16); - assert_eq!(r.cacheable_block, 1600); - // The engine's checkpoint position is `(2002 - 1) / 16 * 16 == 2000`. - // The vLLM test stubs partial hits off and gets 1,600. Both positions - // pass the five conditions (0 % 16 == 0; cdiv(2002, 1600) - 2 == 0 > - // -1; 0 + 16 <= pos; 0 < pos < 2002; pos % 16 == 0), so both give - // 2,002. - assert_eq!(r.checkpoint_position(2002), 2000); - assert!(r.checkpoint_valid(0, 2002, 2000)); - assert!(r.checkpoint_valid(0, 2002, 1600)); - assert_eq!(r.align_split(0, 2002, 2002, 2002, 0, 16384), 2002); - // The end 2,000 is below the prompt, so the floor to 1,600 applies - // and no checkpoint is taken. - assert_eq!(r.align_split(0, 3602, 3602, 2000, 0, 16384), 1600); - let groups = [attention("full", 1600, None), state("s", 1600, true, false)]; - let v31 = GroupedPrefixRules::new(&groups, 16, VllmRuleVersion::V0_31_0).unwrap(); - assert_eq!(v31.align_split(0, 2002, 2002, 2002, 0, 16384), 1600); - assert_eq!(v31.align_split(0, 3602, 3602, 2000, 0, 16384), 1600); - } - - #[test] - fn partial_tail_moves_one_unit_down_on_vllm_0320() { - // At hash 4 the 0.32.0 tail, `(prompt - 1) / hash * hash`, is 8,996, - // between the last stop 8,704 and the prompt (v0.32 sched:517-525, - // single:860-868, single:2069-2073). The 0.31.0 tail is 9,000. - let groups = groups_glm(false); - let hash = 4; - for (version, expected_ends, expected_hit) in [ - (VllmRuleVersion::V0_32_0, vec![6528, 8704, 8996, 9000], 8996), - (VllmRuleVersion::V0_31_0, vec![6528, 8704, 9000], 8704), - ] { - let mut coordinator = with_limits(rules_glm_hash(version, false, hash)); - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(9000, hash); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - let mut block_lease = coordinator.new_lease(); - let ends = prefill( - &mut coordinator, - pool, - &mut block_lease, - &request, - 9000, - &hit, - BUDGET, - ); - assert_eq!(ends, expected_ends, "{version:?}"); - if version == VllmRuleVersion::V0_32_0 { - for group in GLM_STATES { - assert!(hit_exists(pool, group, hash, 8704)); - assert!(hit_exists(pool, group, hash, 8996)); - } - } - let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); - assert_eq!(resend.tokens, expected_hit, "{version:?}"); - coordinator.release_request(pool, &mut block_lease); - let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); - assert_eq!(resend.tokens, expected_hit, "{version:?}"); - } - } - - #[test] - fn checkpoint_block_is_taken_in_the_final_chunk_and_published_at_the_prompt_checkpoint() { - // v0.32 single:1643-1674, :1745-1773 (one more block in the chunk - // that reaches the prompt), :2017-2066 (published at the prompt - // checkpoint under retention 0). - let groups = groups_glm(true); - let hash = 2176; - let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_32_0, true)); - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(9000, hash); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - assert_eq!(hit.tokens, 0); - // 5 for G0, 1 for the ring, 2 for each KDA group. - assert_eq!(coordinator.admission_blocks(pool, &hit, 9000), 14); - let v31 = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); - assert_eq!(v31.admission_blocks(pool, &hit, 9000), 10); - - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert_eq!(split(coordinator.rules(), 0, 9000, 9000, BUDGET, 0), 6528); - assert!(coordinator.allocate(pool, &mut block_lease, 9000, 0, 6528, 1)); - for group in GLM_STATES { - let state = state_lease(&block_lease, usize::from(group)); - assert!(state.running.is_some()); - assert!(state.previous.is_none() && state.checkpoint.is_none() && state.hit.is_none()); - } - coordinator.finalize(pool, &mut block_lease, &request, 9000, 6528, &hit, 1); - coordinator.begin_step(); - assert_eq!( - split(coordinator.rules(), 6528, 9000, 9000, BUDGET, 0), - 2472 - ); - assert!(coordinator.allocate(pool, &mut block_lease, 9000, 6528, 9000, 2)); - for group in GLM_STATES { - let state = state_lease(&block_lease, usize::from(group)); - assert!(state.running.is_some() && state.previous.is_some()); - assert_eq!(state.checkpoint.map(|(_, position)| position), Some(8704)); - } - // G0 5, ring 1, 4 x (running + previous + checkpoint). - assert_eq!(pool.active_blocks(), 5 + 1 + 12); - coordinator.finalize(pool, &mut block_lease, &request, 9000, 9000, &hit, 2); - for group in GLM_STATES { - assert!(hit_exists(pool, group, hash, 8704), "group {group}"); - let state = state_lease(&block_lease, usize::from(group)); - let (checkpoint, _) = state.checkpoint.unwrap(); - assert!(!pool.is_private(checkpoint)); - assert!(pool.is_private(state.running.unwrap())); - } - coordinator.begin_step(); - let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); - assert_eq!(resend.tokens, 8704); - coordinator.release_request(pool, &mut block_lease); - // Four G0 blocks and the four checkpoint copies stay in the LRU. - assert_eq!(pool.inactive_cached_blocks(), 8); - assert_eq!(pool.active_blocks(), 0); - let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); - assert_eq!(resend.tokens, 8704); - } - - #[test] - fn checkpoint_block_goes_at_the_next_step_and_stays_cached_until_eviction() { - // A Mamba manager skips every token before the last computed one (v0.32 - // single:1954-1959). So the request's next allocation first frees every - // block below its last block, the checkpoint block included (v0.32 - // kv_cache_manager.py:546-559, single:672-710, :1592-1641). A freed - // block keeps its key in the LRU (v0.32 block_pool.py:774-803). - let groups = groups_glm(true); - let hash = 2176; - let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_32_0, true)); - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(9000, hash); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - let mut block_lease = coordinator.new_lease(); - let ends = prefill( - &mut coordinator, - pool, - &mut block_lease, - &request, - 9000, - &hit, - BUDGET, - ); - assert_eq!(ends, [6528, 9000]); - // G0 5, ring 1, 4 x (running + previous + checkpoint). - assert_eq!(pool.active_blocks(), 5 + 1 + 12); - // The prefill ran steps 1 and 2. Each decode step holds G0 5, the ring - // and the four running blocks only, and a resend still hits the - // checkpoint keys. - for (step, target) in [(3, 9001), (4, 9002)] { - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 9000, target - 1, target, step)); - assert_eq!(pool.active_blocks(), 5 + 1 + 4, "step {step}"); - for group in GLM_STATES { - let state = state_lease(&block_lease, usize::from(group)); - assert!(state.checkpoint.is_none() && state.previous.is_none()); - let cached = pool.key_hit(group, hash_at(hash, 8704)); - assert!(cached.is_some_and(|hit| !hit.is_active), "group {group}"); - } - coordinator.finalize(pool, &mut block_lease, &request, 9000, target, &hit, step); - let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); - assert_eq!(resend.tokens, 8704, "step {step}"); - } - // The checkpoint keys go only when a new block evicts them, oldest first - // in the order the groups freed them (v0.32 coord:447-450). - assert_eq!(pool.inactive_cached_blocks(), 4); - let free = pool.capacity_blocks() - pool.active_blocks() - 4; - let mut reservation = pool.reserve(&[], free + 1).unwrap(); - for _ in 0..=free { - pool.take_private(&mut reservation); - } - pool.cancel(reservation); - assert!(!hit_exists(pool, 2, hash, 8704)); - for group in [3, 4, 5] { - assert!(hit_exists(pool, group, hash, 8704), "group {group}"); - } - } - - #[test] - fn checkpoint_block_publishes_the_prompt_checkpoint_under_retention_zero() { - // Ported from v0.32 partial_test:775-807: the prompt-end checkpoint - // `(prompt - 1) / 32 * 32` stays published under retention 0. - let groups = [attention("full", 32, None), state("s", 32, true, true)]; - let hash = 32; - for (prompt, expected) in [(120, 96), (104, 96), (96, 64)] { - let mut coordinator = with_limits( - GroupedPrefixRules::new(&groups, hash, VllmRuleVersion::V0_32_0).unwrap(), - ); - let mut kv = kv_pool(&groups, 16); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(prompt, hash); - let hit = coordinator.find_longest_hit(pool, &request, prompt); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, prompt, 0, prompt, 1)); - assert_eq!( - state_lease(&block_lease, 1).checkpoint.map(|(_, p)| p), - Some(expected) - ); - coordinator.finalize(pool, &mut block_lease, &request, prompt, prompt, &hit, 1); - assert!(hit_exists(pool, 1, hash, expected), "prompt {prompt}"); - coordinator.release_request(pool, &mut block_lease); - coordinator.begin_step(); - let resend = coordinator.find_longest_hit(pool, &lease(prompt, hash), prompt); - assert_eq!(resend.tokens, expected, "prompt {prompt}"); - } - } - - #[test] - fn state_key_published_in_a_step_is_not_a_hit_in_that_step() { - // single:1726-1734: a state block cached in this step is not a hit - // until the next step. The block pool stores no step, so the coordinator - // keeps the set. - let groups = groups_glm(false); - let hash = 2176; - let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); - let mut kv = kv_pool(&groups, 64); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(2500, hash); - let hit = coordinator.find_longest_hit(pool, &request, 2500); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 2500, 0, 2176, 1)); - coordinator.finalize(pool, &mut block_lease, &request, 2500, 2176, &hit, 1); - for group in GLM_STATES { - assert!(hit_exists(pool, group, hash, 2176)); - } - let same_step = coordinator.find_longest_hit(pool, &lease(2500, hash), 2500); - assert_eq!(same_step.tokens, 0); - coordinator.begin_step(); - let next_step = coordinator.find_longest_hit(pool, &lease(2500, hash), 2500); - assert_eq!(next_step.tokens, 2176); - } - - #[test] - fn allocate_fails_atomically() { - let groups = groups_glm(false); - let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); - // Room for the first chunk (G0 1, ring 1, 4 running) but not the - // second (G0 1 more, 4 new running blocks while 4 previous stay). - let mut kv = kv_pool(&groups, 9); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(9000, 2176); - let hit = coordinator.find_longest_hit(pool, &request, 9000); - let mut block_lease = coordinator.new_lease(); - coordinator.begin_step(); - assert!(coordinator.allocate(pool, &mut block_lease, 9000, 0, 2176, 1)); - coordinator.finalize(pool, &mut block_lease, &request, 9000, 2176, &hit, 1); - assert_eq!(pool.active_blocks(), 6); - coordinator.begin_step(); - let before = format!("{block_lease:?}"); - assert!(!coordinator.allocate(pool, &mut block_lease, 9000, 2176, 4352, 2)); - assert_eq!(pool.active_blocks(), 6); - assert_eq!(format!("{block_lease:?}"), before); - coordinator.release_request(pool, &mut block_lease); - assert_eq!(pool.active_blocks(), 0); - } - - #[test] - fn release_frees_groups_in_profile_order_each_tail_first() { - // coord:387-395 frees the managers in group order, and each manager frees - // its blocks tail first (single:581-589). Cached blocks join the LRU tail - // in that order (pool:777-808). So the cached block of a state group that - // the profile lists first is evicted before the attention blocks. - let groups = [state("s", 64, true, false), attention("full", 64, None)]; - let hash = 64; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let mut coordinator = - with_limits(GroupedPrefixRules::new(&groups, hash, version).unwrap()); - let mut kv = kv_pool(&groups, 6); - let pool = kv.block_pool_mut().unwrap(); - let request = lease(200, hash); - let hit = coordinator.find_longest_hit(pool, &request, 200); - let mut block_lease = coordinator.new_lease(); - let ends = prefill( - &mut coordinator, - pool, - &mut block_lease, - &request, - 200, - &hit, - 256, - ); - assert_eq!(ends, vec![192, 200], "{version:?}"); - coordinator.release_request(pool, &mut block_lease); - // The state at the replay boundary 192 and the attention blocks that - // end at 64, 128 and 192 stay cached. Two blocks are free. - assert_eq!(pool.inactive_cached_blocks(), 4, "{version:?}"); - assert_eq!(pool.active_blocks(), 0, "{version:?}"); - // Three fresh blocks evict one cached block: the state block. - let reservation = pool.reserve(&[], 3).unwrap(); - pool.cancel(reservation); - assert!(!hit_exists(pool, 0, hash, 192), "{version:?}"); - for end in [64, 128, 192] { - assert!(hit_exists(pool, 1, hash, end), "{version:?} at {end}"); - } - // Four fresh blocks evict one more: the attention tail block. - let reservation = pool.reserve(&[], 4).unwrap(); - pool.cancel(reservation); - assert!(!hit_exists(pool, 1, hash, 192), "{version:?}"); - for end in [64, 128] { - assert!(hit_exists(pool, 1, hash, end), "{version:?} at {end}"); - } - } - } - - #[test] - fn checkpoint_rule_needs_state_groups_at_the_smallest_cacheable_block() { - // vLLM 0.32.0 tests a checkpoint against `cache_config.block_size`, the - // smallest prefix-cacheable block, in the scheduler (v0.32 sched:469, - // :496) and against the Mamba block in the manager (v0.32 single:1652, - // :1662). The rule table holds one block for both. - let small_attention = [attention("full", 16, None), state("s", 64, true, true)]; - let err = GroupedPrefixRules::new(&small_attention, 16, VllmRuleVersion::V0_32_0) - .unwrap_err() - .to_string(); - assert_eq!( - err, - "kv_cache_groups with vllm_version=0.32.0 require every state group with prefill checkpoint blocks to have the smallest block_size_tokens of the prefix_cacheable groups" - ); - // Without checkpoint blocks the same shape is modeled on both rows. - let no_checkpoint = [attention("full", 16, None), state("s", 64, true, false)]; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - assert!(GroupedPrefixRules::new(&no_checkpoint, 16, version).is_ok()); - } - // A larger attention block keeps the state block the smallest. The - // checkpoint position then moves in units of the scheduler block, the LCM - // 128, in the manager (v0.32 coord:699-706) and in the scheduler (v0.32 - // sched:473-475, :209): `(300 - 1) / 128 * 128 == 256`. - let large_attention = [attention("full", 128, None), state("s", 64, true, true)]; - let r = GroupedPrefixRules::new(&large_attention, 64, VllmRuleVersion::V0_32_0).unwrap(); - assert_eq!((r.alignment, r.cacheable_block), (128, 64)); - assert_eq!(r.checkpoint_position(300), 256); - } - - #[test] - fn retired_window_blocks_are_released_tail_first() { - // v0.31 single:660-685 and v0.32 single:645-670 free a retired window - // range backward, and free_blocks appends its cached blocks to the free - // queue in that order (v0.31 pool:777-806, v0.32 pool:774-803). So the - // tail block of the range is evicted first. - let groups = [attention("window", 4, Some(8))]; - let hash = 4; - for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { - let mut coordinator = - with_limits(GroupedPrefixRules::new(&groups, hash, version).unwrap()); - let mut kv = kv_pool(&groups, 8); - let pool = kv.block_pool_mut().unwrap(); - // The run that a finished request of 100 tokens keeps before its - // replay boundary 96: the blocks that end at 92 and 96. - seed_many(pool, 0, hash, &[92, 96]); - let request = lease(108, hash); - let hit = coordinator.find_longest_hit(pool, &request, 108); - assert_eq!(hit.tokens, 96, "{version:?}"); - let mut block_lease = coordinator.new_lease(); - coordinator.pin_hit(pool, &mut block_lease, &hit).unwrap(); - assert!(coordinator.allocate(pool, &mut block_lease, 108, 96, 104, 1)); - coordinator.finalize(pool, &mut block_lease, &request, 108, 104, &hit, 1); - // The next allocation retires both hit blocks: (104 - 7) / 4 = 24. - assert!(coordinator.allocate(pool, &mut block_lease, 108, 104, 108, 2)); - assert_eq!(pool.inactive_cached_blocks(), 2, "{version:?}"); - assert_eq!(pool.active_blocks(), 3, "{version:?}"); - // Three blocks are free, so four fresh blocks evict one cached block: - // the tail block of the retired range, which ends at 96. - let reservation = pool.reserve(&[], 4).unwrap(); - pool.cancel(reservation); - assert!(!hit_exists(pool, 0, hash, 96), "{version:?}"); - assert!(hit_exists(pool, 0, hash, 92), "{version:?}"); - } - } - - /// The sliding-window search before the single traversal: every candidate - /// end scans its own run again. - fn rescanning_window_run( - max: usize, - block: usize, - step: usize, - needed: usize, - mut resident: impl FnMut(usize) -> bool, - ) -> (usize, usize) { - let mut end = max / step * step; - while end > 0 { - let blocks = end / block; - let mut run = 0; - while run < blocks && run < needed && resident(blocks - run - 1) { - run += 1; - } - if run >= needed || run == blocks { - return (end, run); - } - end -= step; - } - (0, 0) - } - - #[test] - fn window_search_looks_up_each_block_at_most_once() { - // single:1018-1044 finds a sliding-window hit in one backward traversal. - // 4,096 candidate blocks of which the last 2,048 are resident, and a - // window that needs 4,096 contiguous blocks: no hit. - let (block, needed) = (16, 4096); - let resident = |i: usize| i >= 2048; - let mut lookups = 0; - let found = window_run(4096 * block, block, block, needed, |i| { - lookups += 1; - resident(i) - }); - let mut rescans = 0; - let rescanned = rescanning_window_run(4096 * block, block, block, needed, |i| { - rescans += 1; - resident(i) - }); - assert_eq!(found, (0, 0)); - assert_eq!(rescanned, found); - assert_eq!(rescans, 2_102_272); - assert_eq!(lookups, 4096); - } - - #[test] - fn window_search_finds_the_hits_of_the_rescanning_search() { - // A fixed xorshift sequence: random residency, blocks, alignments, - // windows and lengths. Each block is looked up at most once. - let mut state: u64 = 0x9E37_79B9_7F4A_7C15; - let mut next = |bound: u64| { - state ^= state << 13; - state ^= state >> 7; - state ^= state << 17; - state % bound - }; - for case in 0..20_000 { - let block = 1 + next(8) as usize; - let step = block * (1 + next(4) as usize); - let needed = next(7) as usize; - let max = next(400) as usize; - let density = next(101); - let blocks = max / block; - let resident: Vec = (0..blocks).map(|_| next(100) < density).collect(); - let mut lookups = vec![0; blocks]; - let found = window_run(max, block, step, needed, |i| { - lookups[i] += 1; - resident[i] - }); - let expected = rescanning_window_run(max, block, step, needed, |i| resident[i]); - assert_eq!( - found, expected, - "case {case}: max {max} block {block} step {step} needed {needed} {resident:?}" - ); - assert!(lookups.iter().all(|&count| count <= 1), "case {case}"); - } - } - - /// Whether a scan of every completed attention block, as finalization did - /// before it kept a cursor, would still find a block to publish. - fn rescan_finds_a_block( - coordinator: &GroupedPrefixCoordinator, - pool: &GroupedBlockPool, - lease: &GroupedBlockLease, - request: &BlockRequestLease, - prompt: usize, - computed_after: usize, - hit: &GroupedHit, - ) -> bool { - use crate::engine::cache::vllm_block_pool::BlockCopyId; - let rules = coordinator.rules(); - let junction = (hit.junction_tokens > 0).then_some(hit.tokens + hit.junction_tokens); - let publishable = |g: usize, copy: Option<&Option>, end: usize| { - copy.is_some_and(|copy| copy.is_some_and(|id| pool.is_private(id))) - && grouped_key(g, request.prefix_hash_at(end)).is_some() - }; - rules.groups.iter().enumerate().any(|(g, rule)| { - let GroupLease::Attention { copies, .. } = &lease.groups[g] else { - return false; - }; - let block = rule.block; - match rule.kind { - GroupRuleKind::FullAttention => (0..computed_after / block) - .any(|i| publishable(g, copies.get(i), (i + 1) * block)), - GroupRuleKind::SlidingWindow { window } => [ - Some(rules.replay_boundary(prompt, block)), - junction.map(|junction| rules.block_floor(junction, block)), - ] - .into_iter() - .flatten() - .filter(|&pos| pos > 0 && pos <= computed_after) - .any(|pos| { - let top = pos / block; - (top.saturating_sub((window - 1).div_ceil(block))..top) - .any(|i| publishable(g, copies.get(i), (i + 1) * block)) - }), - _ => false, - } - }) - } - - #[test] - fn finalize_publishes_what_a_full_rescan_would() { - // Two requests share a prefix on a full-attention and a sliding-window - // group, with chunked prefills and decode steps. After every step a scan - // of all completed blocks finds nothing more to publish, so the cursor - // skips only blocks that are cached or have no key. - let mut state: u64 = 0xD1B5_4A32_D192_ED03; - let mut next = |bound: u64| { - state ^= state << 13; - state ^= state >> 7; - state ^= state << 17; - (state % bound) as usize - }; - let mut hits = 0; - for case in 0..400 { - let hash = [4, 8][next(2)]; - let groups = [ - attention("full", (hash * (1 + next(2))) as u32, None), - attention( - "window", - (hash * (1 + next(2))) as u32, - Some(2 + next(40) as u32), - ), - ]; - let version = [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0][next(2)]; - let mut coordinator = - with_limits(GroupedPrefixRules::new(&groups, hash, version).unwrap()); - let mut kv = kv_pool(&groups, 4096); - let pool = kv.block_pool_mut().unwrap(); - let first = 1 + next(300); - // Half of the cases resend the whole first prompt, which reaches its - // replay boundary. - let shared = if next(2) == 0 { - first - } else { - next(first as u64 + 1) - }; - let second: Vec = (0..shared as u32) - .chain(10_000..10_000 + (1 + next(300)) as u32) - .collect(); - let first: Vec = (0..first as u32).collect(); - let mut step = 1; - for tokens in [first, second] { - let mut request = BlockRequestLease::new(Uuid::nil(), Vec::new()); - request.configure_prefix_hashes(&tokens, hash); - let prompt = tokens.len(); - let hit = coordinator.find_longest_hit(pool, &request, prompt); - hits += usize::from(hit.tokens > 0); - let mut block_lease = coordinator.new_lease(); - coordinator.pin_hit(pool, &mut block_lease, &hit).unwrap(); - // A chunked prefill, then up to eight decode steps. - let (budget, end) = (1 + next(64), prompt + 1 + next(8)); - let mut computed = hit.tokens; - while computed < end { - coordinator.begin_step(); - let target = (computed + budget).min(prompt).max(computed + 1); - assert!(coordinator.allocate( - pool, - &mut block_lease, - prompt, - computed, - target, - step - )); - coordinator.finalize( - pool, - &mut block_lease, - &request, - prompt, - target, - &hit, - step, - ); - assert!( - !rescan_finds_a_block( - &coordinator, - pool, - &block_lease, - &request, - prompt, - target, - &hit - ), - "case {case} {groups:?} prompt {prompt} hit {} at {target}", - hit.tokens - ); - computed = target; - step += 1; - } - coordinator.release_request(pool, &mut block_lease); - } - } - assert!(hits > 150, "{hits} hits"); - } -} +#[path = "grouped_prefix_tests.rs"] +mod tests; diff --git a/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs b/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs new file mode 100644 index 0000000000..f7055f11fa --- /dev/null +++ b/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs @@ -0,0 +1,1825 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Rule and coordinator cases of the grouped prefix cache, with the vLLM +//! 0.31.0 and 0.32.0 literals re-derived from the pinned trees. Some cases +//! are ported from the vLLM tests at `db9527a46873454610df6dbedf79a36d6bf1a7f6` +//! and `5df63a7203cdb9328d824a27b682e77c7b9928be`. Copyright contributors to +//! the vLLM project. See THIRD_PARTY_NOTICES.md. Apache-2.0. +//! +//! - https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +//! - https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +//! - https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/tests/v1/core/test_mamba_align_chunk_split.py + +use super::*; +use crate::engine::cache::vllm_block_pool::CacheKey; +use crate::engine::common::hashing::SequenceHash; +use crate::engine::kv_manager::grouped::{GroupLease, GroupedBlockPool, GroupedKvPool}; +use crate::engine::kv_manager::vllm_backend::BlockRequestLease; +use crate::perfmodel::{FpmCacheGroup, FpmCacheKind}; +use uuid::Uuid; + +const BUDGET: usize = 8192; +const GLM_BLOCK: u32 = 2176; + +fn group( + name: &str, + kind: FpmCacheKind, + block: u32, + window: Option, + state: Option<(bool, u32)>, +) -> FpmCacheGroup { + FpmCacheGroup { + name: name.into(), + kind, + num_layers: 1, + block_size_tokens: block, + page_size_bytes: 1, + sliding_window: window, + blocks_per_request: state.map(|_| 1), + prefix_cacheable: Some(state.is_none_or(|(cacheable, _)| cacheable)), + transfer_bytes_per_block: None, + prefill_checkpoint_blocks: state.and_then(|(cacheable, ckpt)| cacheable.then_some(ckpt)), + prefill_checkpoint_alignment: state + .and_then(|(cacheable, ckpt)| (cacheable && ckpt == 1).then_some(16)), + } +} + +fn attention(name: &str, block: u32, window: Option) -> FpmCacheGroup { + group(name, FpmCacheKind::Attention, block, window, None) +} + +fn state(name: &str, block: u32, cacheable: bool, checkpoint: bool) -> FpmCacheGroup { + group( + name, + FpmCacheKind::State, + block, + None, + Some((cacheable, u32::from(checkpoint))), + ) +} + +/// The six GLM-5.3-Flash TP 4 groups that the runtime probe recorded on +/// vLLM 0.31.0. With `checkpoint`, G2 to G5 carry +/// `prefill_checkpoint_blocks: Some(1)` and alignment 16. +fn groups_glm(checkpoint: bool) -> Vec { + let mut groups = vec![ + attention("mla", GLM_BLOCK, None), + state("kpool_tail", 4, false, false), + ]; + for i in 0..4 { + groups.push(state(&format!("kda_{i}"), GLM_BLOCK, true, checkpoint)); + } + groups +} + +fn rules_glm_hash( + version: VllmRuleVersion, + checkpoint: bool, + hash_block: usize, +) -> GroupedPrefixRules { + GroupedPrefixRules::new(&groups_glm(checkpoint), hash_block, version).unwrap() +} + +fn rules_glm(version: VllmRuleVersion, checkpoint: bool) -> GroupedPrefixRules { + rules_glm_hash(version, checkpoint, GLM_BLOCK as usize) +} + +fn rules(groups: &[FpmCacheGroup], hash_block: usize) -> GroupedPrefixRules { + GroupedPrefixRules::new(groups, hash_block, VllmRuleVersion::V0_31_0).unwrap() +} + +/// A block pool with `blocks` pool blocks for `groups`. +fn kv_pool(groups: &[FpmCacheGroup], blocks: u64) -> GroupedKvPool { + GroupedKvPool::with_prefix_cache(groups.to_vec(), blocks, 1) +} + +/// A lease with chained hashes of the tokens `0..tokens` at `hash_block`. +fn lease(tokens: usize, hash_block: usize) -> BlockRequestLease { + let tokens: Vec = (0..tokens as u32).collect(); + let mut lease = BlockRequestLease::new(Uuid::nil(), Vec::new()); + lease.configure_prefix_hashes(&tokens, hash_block); + lease +} + +fn hash_at(hash_block: usize, tokens: usize) -> SequenceHash { + lease(tokens, hash_block).prefix_hash_at(tokens).unwrap() +} + +fn key(group: u16, hash_block: usize, tokens: usize) -> CacheKey { + CacheKey::Grouped { + group, + hash: hash_at(hash_block, tokens), + } +} + +/// Publish one block under `key` and leave it in the inactive LRU, as a +/// finished request leaves its cached blocks. +fn seed(pool: &mut GroupedBlockPool, key: CacheKey) { + let mut reservation = pool.reserve(&[], 1).unwrap(); + let id = pool.take_private(&mut reservation); + assert!(pool.cache(id, key)); + pool.release(id); +} + +fn seed_many(pool: &mut GroupedBlockPool, group: u16, hash_block: usize, ends: &[usize]) { + for &end in ends { + seed(pool, key(group, hash_block, end)); + } +} + +fn seed_states(pool: &mut GroupedBlockPool, groups: &[u16], hash_block: usize, end: usize) { + for &group in groups { + seed(pool, key(group, hash_block, end)); + } +} + +fn hit_exists(pool: &GroupedBlockPool, group: u16, hash_block: usize, tokens: usize) -> bool { + pool.key_hit(group, hash_at(hash_block, tokens)).is_some() +} + +fn state_lease(lease: &GroupedBlockLease, group: usize) -> &GroupStateLease { + match &lease.groups[group] { + GroupLease::State(state) => state, + other => panic!("group {group} is not a state lease: {other:?}"), + } +} + +fn attention_copies(lease: &GroupedBlockLease, group: usize) -> usize { + match &lease.groups[group] { + GroupLease::Attention { copies, .. } => copies.iter().flatten().count(), + other => panic!("group {group} is not an attention lease: {other:?}"), + } +} + +fn split( + rules: &GroupedPrefixRules, + start: usize, + prompt: usize, + known: usize, + desired: usize, + junction: usize, +) -> usize { + rules.align_split(start, prompt, known, desired, junction, BUDGET) +} + +/// Run a whole prefill of `prompt` tokens from `hit`: the chunks of +/// `align_split` at `budget`, one step each. Returns the chunk ends. +fn prefill( + coordinator: &mut GroupedPrefixCoordinator, + pool: &mut GroupedBlockPool, + lease: &mut GroupedBlockLease, + request: &BlockRequestLease, + prompt: usize, + hit: &GroupedHit, + budget: usize, +) -> Vec { + let mut ends = Vec::new(); + let mut computed = hit.tokens; + let mut step = 1; + while computed < prompt { + coordinator.begin_step(); + let chunk = coordinator.rules().align_split( + computed, + prompt, + prompt, + budget, + hit.tokens + hit.junction_tokens, + budget, + ); + assert!(chunk > 0, "the split must make progress at {computed}"); + let target = computed + chunk; + assert!(coordinator.allocate(pool, lease, prompt, computed, target, step)); + coordinator.finalize(pool, lease, request, prompt, target, hit, step); + computed = target; + ends.push(target); + step += 1; + } + coordinator.begin_step(); + ends +} + +const GLM_STATES: [u16; 4] = [2, 3, 4, 5]; + +/// One batch of `BUDGET` tokens in flight and no model-length cap. +const LIMITS: AdmissionLimits = AdmissionLimits { + max_in_flight_tokens: BUDGET, + max_model_len: usize::MAX, +}; + +fn with_limits(rules: GroupedPrefixRules) -> GroupedPrefixCoordinator { + GroupedPrefixCoordinator::new(rules, LIMITS) +} + +#[test] +fn version_rows_differ_in_the_last_cacheable_stop_and_the_partial_tail() { + // v0.31 sched:444 `num_tokens - num_tokens % block_size`; v0.32 + // sched:476-480 with iface:1110-1122 `(n - 1) // B * B`. + let v31 = rules_glm(VllmRuleVersion::V0_31_0, false); + let v32 = rules_glm(VllmRuleVersion::V0_32_0, false); + assert_eq!(v31.last_cacheable_stop(9000), 8704); + assert_eq!(v32.last_cacheable_stop(9000), 8704); + assert_eq!(v31.last_cacheable_stop(8704), 8704); + assert_eq!(v32.last_cacheable_stop(8704), 6528); + assert_eq!(v32.last_cacheable_stop(0), 0); + // v0.31 sched:482-487 `prompt // hash * hash`; v0.32 sched:517-525 + // `(prompt - 1) // hash * hash`. + let v31 = rules_glm_hash(VllmRuleVersion::V0_31_0, false, 4); + let v32 = rules_glm_hash(VllmRuleVersion::V0_32_0, false, 4); + assert_eq!(v31.partial_tail(9000), 9000); + assert_eq!(v32.partial_tail(9000), 8996); + assert_eq!(v31.partial_tail(4356), 4356); + assert_eq!(v32.partial_tail(4356), 4352); + assert_eq!( + VllmRuleVersion::parse("0.31.0"), + Some(VllmRuleVersion::V0_31_0) + ); + assert_eq!( + VllmRuleVersion::parse("0.32.0"), + Some(VllmRuleVersion::V0_32_0) + ); + assert_eq!(VllmRuleVersion::parse("0.30.0"), None); +} + +#[test] +fn rules_refuse_unsupported_checkpoint_fields() { + // v0.32 sched:390-394 runs checkpoint mode only when every Mamba + // group has a positive count. + let err = GroupedPrefixRules::new(&groups_glm(true), 2176, VllmRuleVersion::V0_31_0) + .unwrap_err() + .to_string(); + assert_eq!( + err, + "kv_cache_groups with vllm_version=0.31.0 do not model prefill checkpoint blocks" + ); + let mut groups = groups_glm(false); + groups[2].prefill_checkpoint_blocks = None; + let err = GroupedPrefixRules::new(&groups, 2176, VllmRuleVersion::V0_32_0) + .unwrap_err() + .to_string(); + assert_eq!( + err, + "kv_cache_groups with vllm_version=0.32.0 require prefill_checkpoint_blocks on every prefix-cacheable state group; re-import the runtime probe" + ); + let mut groups = groups_glm(true); + groups[3].prefill_checkpoint_alignment = Some(32); + let err = GroupedPrefixRules::new(&groups, 2176, VllmRuleVersion::V0_32_0) + .unwrap_err() + .to_string(); + assert_eq!( + err, + "kv_cache_groups with vllm_version=0.32.0 require one prefill_checkpoint_alignment across state groups" + ); + let zero = rules_glm(VllmRuleVersion::V0_32_0, false); + assert!(zero.checkpoint.is_none()); + let one = rules_glm(VllmRuleVersion::V0_32_0, true); + assert_eq!(one.checkpoint, Some(CheckpointRule { alignment: 16 })); + // Block mode holds one running block per state group, so another count + // is refused. + let mut groups = groups_glm(false); + groups[4].blocks_per_request = Some(2); + let err = GroupedPrefixRules::new(&groups, 2176, VllmRuleVersion::V0_31_0) + .unwrap_err() + .to_string(); + assert_eq!( + err, + "kv_cache_groups with enable_prefix_caching=true model one state block per request; blocks_per_request must be 1" + ); + // utils:799-805 and coord:87-92. + assert!(GroupedPrefixRules::new(&groups_glm(false), 3, VllmRuleVersion::V0_31_0).is_err()); +} + +#[test] +fn rules_refuse_convolution_groups() { + // A convolution group keeps a window, and the table has no rule for it. + let mut convolution = attention("convolution", 4, Some(4)); + convolution.kind = FpmCacheKind::Convolution; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + for groups in [ + vec![convolution.clone()], + vec![attention("full", 4, None), convolution.clone()], + ] { + let err = GroupedPrefixRules::new(&groups, 4, version) + .unwrap_err() + .to_string(); + assert_eq!( + err, + "grouped prefix caching does not support convolution cache groups; turn prefix caching off to replay them cold", + "{version:?}" + ); + } + } +} + +#[test] +fn alignment_is_the_hash_block_only_with_fine_grained_hits() { + // coord:687-735. G2 to G5 are state rules with block 2176 and G0 is + // full attention. At hash 2176 no state block exceeds the hash block. + let glm = rules_glm(VllmRuleVersion::V0_31_0, false); + assert_eq!(glm.alignment, 2176); + assert_eq!(glm.cacheable_block, 2176); + let fine = rules_glm_hash(VllmRuleVersion::V0_31_0, false, 4); + assert_eq!(fine.alignment, 4); + // Fine-grained hits need a Mamba group in align mode whose block exceeds + // the hash block. Two attention groups have none, so they give the + // scheduler block, with or without a window (coord:691-704). + let windowed = [attention("a", 64, Some(100)), attention("b", 16, None)]; + assert_eq!(rules(&windowed, 16).alignment, 64); + let plain = [attention("a", 64, None), attention("b", 16, None)]; + assert_eq!(rules(&plain, 16).alignment, 64); + // A state group above the hash block turns fine-grained hits on when + // every other cacheable rule is full attention or at the hash block. + let with_state = [ + attention("a", 64, None), + attention("b", 16, None), + state("s", 64, true, false), + ]; + assert_eq!(rules(&with_state, 16).alignment, 16); + // A window at another block than the hash block turns them off again + // (single:960 has no sub-block lookup). + let with_window = [ + attention("a", 64, Some(100)), + attention("b", 16, None), + state("s", 64, true, false), + ]; + assert_eq!(rules(&with_window, 16).alignment, 64); +} + +#[test] +fn full_attention_hit_is_capped_below_the_prompt() { + // mgr:289-295 caps the lookup at `num_tokens - 1`. + let groups = [attention("full", 16, None)]; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let coordinator = with_limits(GroupedPrefixRules::new(&groups, 16, version).unwrap()); + let mut kv = kv_pool(&groups, 16); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, 16, &[16, 32, 48, 64]); + let hit = coordinator.find_longest_hit(pool, &lease(64, 16), 64); + assert_eq!(hit.tokens, 48); + assert_eq!(hit.junction_tokens, 0); + assert_eq!(hit.per_group.len(), 1); + assert!(matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 3)); + let hit = coordinator.find_longest_hit(pool, &lease(65, 16), 65); + assert_eq!(hit.tokens, 64); + } +} + +#[test] +fn sliding_window_hits_need_a_contiguous_run_before_an_aligned_end() { + // single:960-1069: `cdiv(32, 16) == 2` contiguous blocks before an + // aligned end, or a run from block 0. + let groups = [ + attention("full", 16, None), + attention("window", 16, Some(33)), + ]; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let r = GroupedPrefixRules::new(&groups, 16, version).unwrap(); + assert_eq!(r.alignment, 16); + let coordinator = with_limits(r); + let mut kv = kv_pool(&groups, 32); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, 16, &[16, 32, 48, 64, 80, 96]); + seed_many(pool, 1, 16, &[32, 48, 64, 96]); + let hit = coordinator.find_longest_hit(pool, &lease(96, 16), 96); + // The end 64 has blocks 2 and 3 before it. The end 80 lacks block 4. + assert_eq!(hit.tokens, 64); + match &hit.per_group[1] { + GroupHitBlocks::Attention(keys) => { + assert_eq!(keys.len(), 4); + assert!(keys[0].is_none() && keys[1].is_none()); + assert_eq!(keys[2], Some(key(1, 16, 48))); + assert_eq!(keys[3], Some(key(1, 16, 64))); + } + other => panic!("{other:?}"), + } + // Junction: the full-attention hit before the loop (80) minus 64. + assert_eq!(hit.junction_tokens, 16); + + let mut kv = kv_pool(&groups, 32); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, 16, &[16, 32, 48, 64, 80, 96]); + seed_many(pool, 1, 16, &[16]); + let hit = coordinator.find_longest_hit(pool, &lease(96, 16), 96); + assert_eq!(hit.tokens, 16); + } +} + +#[test] +fn one_token_window_hits_only_at_a_cached_block() { + // A window of one token needs `cdiv(0, 4) == 0` contiguous blocks + // (v0.31 single:961-971, v0.32 :943-953). vLLM counts a block only + // after it finds the block cached, so a hit still needs the cached + // block at its end (v0.31 single:1021-1044, v0.32 :1003-1026). + let groups = [attention("window", 4, Some(1))]; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let coordinator = with_limits(GroupedPrefixRules::new(&groups, 4, version).unwrap()); + let mut kv = kv_pool(&groups, 32); + let pool = kv.block_pool_mut().unwrap(); + let hit = coordinator.find_longest_hit(pool, &lease(100, 4), 100); + assert_eq!((hit.tokens, hit.junction_tokens), (0, 0), "{version:?}"); + seed_many(pool, 0, 4, &[48]); + let hit = coordinator.find_longest_hit(pool, &lease(100, 4), 100); + assert_eq!((hit.tokens, hit.junction_tokens), (48, 0), "{version:?}"); + match &hit.per_group[0] { + GroupHitBlocks::Attention(keys) => { + assert_eq!(keys.len(), 12, "{version:?}"); + assert!(keys[..11].iter().all(Option::is_none), "{version:?}"); + assert_eq!(keys[11], Some(key(0, 4, 48)), "{version:?}"); + } + other => panic!("{other:?}"), + } + } +} + +#[test] +fn sliding_window_admission_reserves_the_window_plus_the_in_flight_tokens() { + // iface:859-866: `cdiv(min(window - 1 + max_in_flight_tokens, + // max_model_len), block) + 1`, with `max_in_flight_tokens = + // max_concurrent_batches * max_num_batched_tokens` (vllm.py:605-613). + let groups = [ + attention("full", 16, None), + attention("window", 16, Some(33)), + ]; + let prompt: usize = 100_000; + let hit = GroupedHit { + tokens: 0, + junction_tokens: 0, + per_group: vec![ + GroupHitBlocks::Attention(Vec::new()), + GroupHitBlocks::Attention(Vec::new()), + ], + }; + let full_blocks = prompt.div_ceil(16); + let kv = kv_pool(&groups, 16); + let pool = kv.block_pool().unwrap(); + // One batch of 8,192: (32 + 8192) / 16 + 1 = 515. + let one_batch = with_limits(rules(&groups, 16)); + assert_eq!( + one_batch.admission_blocks(pool, &hit, prompt), + full_blocks + 515 + ); + // Two batches of 8,192: (32 + 16384) / 16 + 1 = 1027. + let two_batches = GroupedPrefixCoordinator::new( + rules(&groups, 16), + AdmissionLimits { + max_in_flight_tokens: 2 * BUDGET, + max_model_len: usize::MAX, + }, + ); + assert_eq!( + two_batches.admission_blocks(pool, &hit, prompt), + full_blocks + 1027 + ); + // The model length caps it: cdiv(4096, 16) + 1 = 257. + let capped = GroupedPrefixCoordinator::new( + rules(&groups, 16), + AdmissionLimits { + max_in_flight_tokens: 2 * BUDGET, + max_model_len: 4096, + }, + ); + assert_eq!( + capped.admission_blocks(pool, &hit, prompt), + full_blocks + 257 + ); + // A short prompt needs fewer blocks than the cap. + assert_eq!(one_batch.admission_blocks(pool, &hit, 100), 7 + 7); +} + +#[test] +fn state_group_hits_only_at_a_kept_boundary() { + // single:1487-1564 scans aligned ends right to left for one cached + // state block. pool:198-224: every group of the spec group must hold + // the key. + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let groups = groups_glm(false); + let coordinator = with_limits(rules_glm(version, false)); + let hash = 2176; + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); + seed_states(pool, &GLM_STATES, hash, 8704); + let request = lease(9000, hash); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + assert_eq!(hit.tokens, 8704); + assert_eq!(hit.junction_tokens, 0); + assert_eq!( + hit.per_group[2], + GroupHitBlocks::State(Some(key(2, hash, 8704))) + ); + assert_eq!(hit.per_group[1], GroupHitBlocks::None); + + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); + seed_states(pool, &[2, 4, 5], hash, 8704); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + assert_eq!(hit.tokens, 0); + assert_eq!(hit.per_group[0], GroupHitBlocks::Attention(Vec::new())); + assert_eq!(hit.per_group[3], GroupHitBlocks::State(None)); + + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); + seed_states(pool, &GLM_STATES, hash, 6528); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + assert_eq!(hit.tokens, 6528); + assert!(matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 3)); + } +} + +#[test] +fn ring_group_never_hits_and_holds_one_block() { + // single:1214-1230 claims one block per request; single:498-499 + // caches nothing for a group that is not prefix cacheable. + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let groups = groups_glm(false); + let mut coordinator = with_limits(rules_glm(version, false)); + let hash = 2176; + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(2500, hash); + let hit = coordinator.find_longest_hit(pool, &request, 2500); + assert_eq!(hit.tokens, 0); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 2500, 0, 2176, 1)); + assert_eq!( + block_lease.groups[1], + GroupLease::Ring { + copy: Some(ring_copy(&block_lease)) + } + ); + coordinator.finalize(pool, &mut block_lease, &request, 2500, 2176, &hit, 1); + assert_eq!(pool.inactive_cached_blocks(), 0); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 2500, 2176, 2500, 2)); + coordinator.finalize(pool, &mut block_lease, &request, 2500, 2500, &hit, 2); + coordinator.release_request(pool, &mut block_lease); + coordinator.begin_step(); + + let hit = coordinator.find_longest_hit(pool, &lease(2500, hash), 2500); + assert_eq!(hit.tokens, 2176); + assert_eq!(hit.per_group[1], GroupHitBlocks::None); + assert!(matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 1)); + assert_eq!( + hit.per_group[2], + GroupHitBlocks::State(Some(key(2, hash, 2176))) + ); + } +} + +fn ring_copy(lease: &GroupedBlockLease) -> crate::engine::cache::vllm_block_pool::BlockCopyId { + match &lease.groups[1] { + GroupLease::Ring { copy: Some(copy) } => *copy, + other => panic!("{other:?}"), + } +} + +#[test] +fn hybrid_hit_is_the_longest_prefix_every_cacheable_group_serves() { + // coord:842-977: the fixed point shortens to 4,352; coord:970-977 + // reports the uncached common prefix. + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let groups = groups_glm(false); + let coordinator = with_limits(rules_glm(version, false)); + let hash = 2176; + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704, 10880]); + seed_states(pool, &GLM_STATES, hash, 4352); + let hit = coordinator.find_longest_hit(pool, &lease(11000, hash), 11000); + assert_eq!(hit.tokens, 4352); + assert_eq!(hit.junction_tokens, 10880 - 4352); + assert!(matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 2)); + } +} + +#[test] +fn hybrid_partial_hash_truncates_full_attention_hit_length() { + // Ported from v0.31 partial_test:1596. Values hand-derived: the + // fine-grained probe of single:806-825 extends the two full blocks to + // the partial entry at 4,356, and the state hit there keeps it. + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let groups = groups_glm(false); + let coordinator = with_limits(rules_glm_hash(version, false, 4)); + let hash = 4; + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 4356]); + seed_states(pool, &GLM_STATES, hash, 4356); + let request = lease(4500, hash); + let hit = coordinator.find_longest_hit(pool, &request, 4500); + assert_eq!(hit.tokens, 4356); + assert_eq!(hit.junction_tokens, 0); + match &hit.per_group[0] { + GroupHitBlocks::Attention(keys) => { + assert_eq!(keys.len(), 3); + assert_eq!(keys[2], Some(key(0, hash, 4356))); + } + other => panic!("{other:?}"), + } + + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 4356]); + seed_states(pool, &GLM_STATES, hash, 4352); + let hit = coordinator.find_longest_hit(pool, &request, 4500); + assert_eq!(hit.tokens, 4352); + // coord:970-977: the full-attention hit of 4,356 minus 4,352. + assert_eq!(hit.junction_tokens, 4); + assert!(matches!(&hit.per_group[0], GroupHitBlocks::Attention(keys) if keys.len() == 2)); + } +} + +/// Two full-attention groups of one spec and a state group, all at block 64. +fn two_full_groups_and_a_state() -> [FpmCacheGroup; 3] { + [ + attention("full_a", 64, None), + attention("full_b", 64, None), + state("state", 64, true, false), + ] +} + +#[test] +fn groups_of_one_spec_hit_jointly_and_set_the_junction_together() { + // coord:737-779 (v0.32 coord:708-750) looks up the groups of one spec + // together, and pool:198-224 finds a block only when every group of the + // spec holds it. The longest hit is the largest result of one spec's + // lookup (coord:947, :970-977). Group a holds 192 tokens, group b 64 and + // the state group the state at 64, so vLLM hits 64 with no junction. + let groups = two_full_groups_and_a_state(); + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let coordinator = with_limits(GroupedPrefixRules::new(&groups, 64, version).unwrap()); + let mut kv = kv_pool(&groups, 16); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, 64, &[64, 128, 192]); + seed_many(pool, 1, 64, &[64]); + seed_many(pool, 2, 64, &[64]); + let hit = coordinator.find_longest_hit(pool, &lease(401, 64), 401); + assert_eq!((hit.tokens, hit.junction_tokens), (64, 0), "{version:?}"); + assert_eq!( + hit.per_group[0], + GroupHitBlocks::Attention(vec![Some(key(0, 64, 64))]) + ); + assert_eq!( + hit.per_group[1], + GroupHitBlocks::Attention(vec![Some(key(1, 64, 64))]) + ); + assert_eq!( + hit.per_group[2], + GroupHitBlocks::State(Some(key(2, 64, 64))) + ); + // So the first chunk runs to the last cacheable stop 384 (v0.31 + // sched:444, v0.32 sched:476-480). A junction at 192 would end it + // there. + let junction = hit.tokens + hit.junction_tokens; + assert_eq!( + split(coordinator.rules(), 64, 401, 401, BUDGET, junction), + 320, + "{version:?}" + ); + } +} + +#[test] +fn groups_of_one_spec_hit_inside_a_block_only_jointly() { + // With fine-grained hits, single:794-828 probes the hash boundaries of the + // first partial block for every group of the spec at once. Group a holds + // two full blocks, group b one full block and the partial entry at 96. + // No boundary inside the second block is cached in both, so vLLM hits 64. + let groups = two_full_groups_and_a_state(); + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let rules = GroupedPrefixRules::new(&groups, 16, version).unwrap(); + assert_eq!(rules.alignment, 16); + let coordinator = with_limits(rules); + let mut kv = kv_pool(&groups, 16); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, 16, &[64, 128]); + seed_many(pool, 1, 16, &[64, 96]); + seed_many(pool, 2, 16, &[64, 96]); + let hit = coordinator.find_longest_hit(pool, &lease(200, 16), 200); + assert_eq!((hit.tokens, hit.junction_tokens), (64, 0), "{version:?}"); + assert_eq!( + hit.per_group[0], + GroupHitBlocks::Attention(vec![Some(key(0, 16, 64))]) + ); + assert_eq!( + hit.per_group[1], + GroupHitBlocks::Attention(vec![Some(key(1, 16, 64))]) + ); + assert_eq!( + hit.per_group[2], + GroupHitBlocks::State(Some(key(2, 16, 64))) + ); + } +} + +#[test] +fn state_groups_of_one_spec_hit_jointly() { + // Without a full-attention group, the first spec looked up is the first + // one in group order (coord:775-779), here the two state groups. They + // hold the states at 64 and 192, and at 64. The joint lookup + // (single:1540-1564) hits 64 and leaves no junction. + let groups = [ + state("state_a", 64, true, false), + state("state_b", 64, true, false), + attention("window", 64, Some(65)), + ]; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let coordinator = with_limits(GroupedPrefixRules::new(&groups, 64, version).unwrap()); + let mut kv = kv_pool(&groups, 16); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, 64, &[64, 192]); + seed_many(pool, 1, 64, &[64]); + seed_many(pool, 2, 64, &[64]); + let hit = coordinator.find_longest_hit(pool, &lease(401, 64), 401); + assert_eq!((hit.tokens, hit.junction_tokens), (64, 0), "{version:?}"); + assert_eq!( + hit.per_group[0], + GroupHitBlocks::State(Some(key(0, 64, 64))) + ); + assert_eq!( + hit.per_group[1], + GroupHitBlocks::State(Some(key(1, 64, 64))) + ); + } + // On GLM-5.3-Flash the full-attention group G0 comes first, so its hit + // sets the junction, as in vLLM. The extra state of G2 at 8,704 changes + // nothing, and the four KDA groups hit 4,352 together. + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let groups = groups_glm(false); + let coordinator = with_limits(rules_glm(version, false)); + let hash = 2176; + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 6528, 8704]); + seed_states(pool, &GLM_STATES, hash, 4352); + seed_states(pool, &[2], hash, 8704); + let hit = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); + assert_eq!( + (hit.tokens, hit.junction_tokens), + (4352, 8704 - 4352), + "{version:?}" + ); + for group in GLM_STATES { + assert_eq!( + hit.per_group[usize::from(group)], + GroupHitBlocks::State(Some(key(group, hash, 4352))) + ); + } + } +} + +#[test] +fn admission_counts_inactive_hits_and_state_copies() { + // single:216-268 counts inactive hit blocks again. The keys are seeded + // directly, because a finished request of 4,352 tokens keeps its state + // key only at its replay boundary, 2,176. + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let groups = groups_glm(false); + let coordinator = with_limits(rules_glm(version, false)); + let hash = 2176; + let mut kv = kv_pool(&groups, 16); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352]); + seed_states(pool, &GLM_STATES, hash, 4352); + let hit = coordinator.find_longest_hit(pool, &lease(4400, hash), 4400); + assert_eq!(hit.tokens, 4352); + // 2 inactive G0 blocks, 1 new G0 block, 1 ring, 4 x (running + + // inactive cached state block). + assert_eq!(coordinator.admission_blocks(pool, &hit, 4400), 12); + // Another running request holds the two G0 blocks active. + let other = [key(0, hash, 2176), key(0, hash, 4352)]; + let mut reservation = pool.reserve(&other, 0).unwrap(); + let held = pool.activate(&mut reservation); + assert_eq!(held.len(), 2); + assert_eq!(coordinator.admission_blocks(pool, &hit, 4400), 10); + pool.cancel(reservation); + } +} + +#[test] +fn resumed_state_holds_cached_and_running_blocks_during_the_first_step() { + // single:1858-1868: a resumed request holds the hit block and the + // running block. single:1658-1675 frees the hit block once + // `cdiv(processed, block) - 1 > hit_index`: `cdiv(6000, 2176) - 1 == + // 2 > 1`. + let groups = groups_glm(false); + let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); + let hash = 2176; + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352]); + seed_states(pool, &GLM_STATES, hash, 4352); + let request = lease(9000, hash); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + assert_eq!(hit.tokens, 4352); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); + for group in GLM_STATES { + let state = state_lease(&block_lease, usize::from(group)); + assert_eq!(state.hit.map(|(_, index)| index), Some(1)); + assert!(state.running.is_none()); + } + assert!(coordinator.allocate(pool, &mut block_lease, 9000, 4352, 6000, 1)); + let state_blocks: usize = GLM_STATES + .iter() + .map(|&group| { + let state = state_lease(&block_lease, usize::from(group)); + assert!(state.hit.is_some() && state.running.is_some()); + assert!(state.previous.is_none() && state.copy.is_none()); + 2 + }) + .sum(); + assert_eq!(state_blocks, 8); + // G0: 2 hits + 1 new, ring 1, states 8. + assert_eq!(pool.active_blocks(), 12); + coordinator.finalize(pool, &mut block_lease, &request, 9000, 6000, &hit, 1); + for group in GLM_STATES { + let state = state_lease(&block_lease, usize::from(group)); + assert!(state.hit.is_none() && state.running.is_some()); + } + assert_eq!(pool.active_blocks(), 8); + assert_eq!(pool.inactive_cached_blocks(), 4); +} + +#[test] +fn boundary_crossing_keeps_the_previous_running_block_for_one_step() { + // single:1658-1675 and mamba_utils:576-584: the block of the earlier + // boundary goes at the next allocation. One pass, as on the linear path, + // where vLLM's async scheduling holds it one more. + let groups = groups_glm(false); + let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); + let hash = 2176; + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(9000, hash); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 9000, 0, 2176, 1)); + coordinator.finalize(pool, &mut block_lease, &request, 9000, 2176, &hit, 1); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 9000, 2176, 4352, 2)); + let state_and_ring = |lease: &GroupedBlockLease| -> usize { + lease.groups[1..] + .iter() + .map(|g| match g { + GroupLease::Attention { .. } => 0, + GroupLease::Ring { copy } => usize::from(copy.is_some()), + GroupLease::State(s) => [ + s.running, + s.hit.map(|h| h.0), + s.previous.map(|p| p.0), + s.copy, + s.checkpoint.map(|c| c.0), + ] + .iter() + .flatten() + .count(), + }) + .sum() + }; + for group in GLM_STATES { + let state = state_lease(&block_lease, usize::from(group)); + assert!(state.running.is_some()); + assert_eq!(state.previous.map(|(_, step)| step), Some(2)); + } + assert_eq!(state_and_ring(&block_lease), 9); + coordinator.finalize(pool, &mut block_lease, &request, 9000, 4352, &hit, 2); + assert_eq!(state_and_ring(&block_lease), 9); + coordinator.begin_step(); + let before = pool.active_blocks(); + assert!(coordinator.allocate(pool, &mut block_lease, 9000, 4352, 6528, 3)); + // Four previous blocks go, four new running blocks and one G0 block come. + assert_eq!(pool.active_blocks(), before + 1); + for group in GLM_STATES { + let state = state_lease(&block_lease, usize::from(group)); + assert_eq!(state.previous.map(|(_, step)| step), Some(3)); + } + assert_eq!(state_and_ring(&block_lease), 9); +} + +#[test] +fn partial_hit_takes_a_copy_block_for_the_shared_block() { + // single:387-409 and single:1770-1777: a hit that ends inside a block + // takes one copy block. single:800-805: the full-attention probe needs + // the full blocks below the partial entry, so blocks 0 and 1 are + // cached too. + let groups = [attention("full", 16, None), state("s", 32, true, false)]; + let r = rules(&groups, 4); + assert_eq!(r.alignment, 4); + let mut coordinator = with_limits(r); + let hash = 4; + let mut kv = kv_pool(&groups, 32); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[16, 32, 36]); + seed_many(pool, 1, hash, &[36]); + let request = lease(50, hash); + let hit = coordinator.find_longest_hit(pool, &request, 50); + assert_eq!(hit.tokens, 36); + // G0: 4 live - 3 hits + 3 inactive + 1 copy = 5. State + // (single:1757-1777): `cdiv(50, 32) - cdiv(36, 32) == 0` new blocks + // plus the copy block, plus the inactive hit = 2. + assert_eq!(coordinator.admission_blocks(pool, &hit, 50), 7); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); + assert_eq!(pool.active_blocks(), 4); + assert!(coordinator.allocate(pool, &mut block_lease, 50, 36, 50, 1)); + // G0 holds blocks 0 to 3, with block 2 as a private copy. The shared + // block 2 is inactive again. The target 50 stays in the hit's block, + // so the state's copy-on-write block is its running block + // (single:1891-1895, `_apply_cow` at single:454-474); the hit block + // stays until the copy has run (sched:1434-1442). + assert_eq!(attention_copies(&block_lease, 0), 4); + let state = state_lease(&block_lease, 1); + assert!(state.hit.is_some() && state.running.is_some() && state.copy.is_none()); + assert_eq!(pool.active_blocks(), 4 + 2); + assert_eq!(pool.inactive_cached_blocks(), 1); + coordinator.finalize(pool, &mut block_lease, &request, 50, 50, &hit, 1); + let state = state_lease(&block_lease, 1); + assert!(state.hit.is_none() && state.running.is_some()); + assert_eq!(pool.active_blocks(), 4 + 1); + assert_eq!(pool.inactive_cached_blocks(), 2); + + // A target beyond the hit's block: the copy sits at the hit's index + // and a running block follows. `remove_skipped_blocks` frees the + // copy at the next allocation (single:1658-1675, index set at :1868). + let mut kv = kv_pool(&groups, 32); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[16, 32, 36]); + seed_many(pool, 1, hash, &[36]); + let request = lease(70, hash); + let hit = coordinator.find_longest_hit(pool, &request, 70); + assert_eq!(hit.tokens, 36); + // G0: 5 live - 3 hits + 3 inactive + 1 copy = 6. State: `cdiv(70, 32) + // - cdiv(36, 32) == 1` new block plus the copy block plus the + // inactive hit = 3. + assert_eq!(coordinator.admission_blocks(pool, &hit, 70), 9); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); + assert!(coordinator.allocate(pool, &mut block_lease, 70, 36, 70, 1)); + let state = state_lease(&block_lease, 1); + assert!(state.hit.is_some() && state.running.is_some() && state.copy.is_some()); + assert_eq!(pool.active_blocks(), 5 + 3); + coordinator.finalize(pool, &mut block_lease, &request, 70, 70, &hit, 1); + let state = state_lease(&block_lease, 1); + assert!(state.hit.is_none() && state.copy.is_some()); + assert_eq!(pool.active_blocks(), 5 + 2); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 70, 70, 71, 2)); + let state = state_lease(&block_lease, 1); + assert!(state.copy.is_none() && state.running.is_some()); + assert_eq!(pool.active_blocks(), 5 + 1); +} + +#[test] +fn finalize_publishes_states_at_the_replay_boundary_and_the_junction() { + // single:1614-1622 with retention 0 (cache:160-166): the replay + // boundary `(9000 - 1) // 2176 * 2176 == 8704` and the junction. + let groups = groups_glm(false); + let hash = 2176; + let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(9000, hash); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + let mut block_lease = coordinator.new_lease(); + let ends = prefill( + &mut coordinator, + pool, + &mut block_lease, + &request, + 9000, + &hit, + 2176, + ); + assert_eq!(ends, vec![2176, 4352, 6528, 8704, 9000]); + coordinator.release_request(pool, &mut block_lease); + for group in GLM_STATES { + for end in [2176, 4352, 6528] { + assert!( + !hit_exists(pool, group, hash, end), + "group {group} at {end}" + ); + } + assert!(hit_exists(pool, group, hash, 8704)); + } + // G0 blocks 0 to 3 and four state blocks stay cached. + assert_eq!(pool.inactive_cached_blocks(), 8); + assert_eq!(pool.active_blocks(), 0); + + // A hit of 4,352 with a junction of 2,176. + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + seed_many(pool, 0, hash, &[2176, 4352, 6528]); + seed_states(pool, &GLM_STATES, hash, 4352); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + assert_eq!((hit.tokens, hit.junction_tokens), (4352, 2176)); + let mut block_lease = coordinator.new_lease(); + assert!(coordinator.pin_hit(pool, &mut block_lease, &hit).is_ok()); + let ends = prefill( + &mut coordinator, + pool, + &mut block_lease, + &request, + 9000, + &hit, + 2176, + ); + assert_eq!(ends, vec![6528, 8704, 9000]); + coordinator.release_request(pool, &mut block_lease); + for group in GLM_STATES { + assert!(hit_exists(pool, group, hash, 6528), "group {group}"); + assert!(hit_exists(pool, group, hash, 8704), "group {group}"); + assert!(!hit_exists(pool, group, hash, 2176), "group {group}"); + } +} + +#[test] +fn align_split_follows_vllm_0310() { + // sched:414-527 at v0.31. + let r = rules_glm(VllmRuleVersion::V0_31_0, false); + // The floor of 8,192 to 2,176 is 6,528 (sched:460-468). + assert_eq!(split(&r, 0, 9000, 9000, 8192, 0), 6528); + // The last cacheable stop `9000 - 9000 % 2176 == 8704` (sched:444). + assert_eq!(split(&r, 6528, 9000, 9000, 8192, 0), 2176); + assert_eq!(split(&r, 8704, 9000, 9000, 8192, 0), 296); + // A prompt of exactly 8,704 tokens has no stop at 6,528, because its + // last cacheable stop `8704 - 8704 % 2176` is the prompt itself + // (sched:444). + assert_eq!(split(&r, 6528, 8704, 8704, 8192, 0), 2176); + // With hash 4 the tail 9,000 equals the prompt and adds no stop + // (sched:482-487, :512-514). With prompt 9,002 the tail is 9,000. + let fine = rules_glm_hash(VllmRuleVersion::V0_31_0, false, 4); + assert_eq!(split(&fine, 8704, 9000, 9000, 8192, 0), 296); + assert_eq!(split(&fine, 8704, 9002, 9002, 8192, 0), 296); + // The junction stop (sched:488-519). + assert_eq!(split(&r, 4352, 9000, 9000, 8192, 6528), 2176); + // Past the prefill the split returns the request (sched:434-436). + assert_eq!(split(&r, 9000, 9000, 9001, 1, 0), 1); +} + +#[test] +fn align_split_follows_vllm_0320() { + // v0.32 sched:443-556. + let r = rules_glm(VllmRuleVersion::V0_32_0, false); + assert_eq!(split(&r, 0, 9000, 9000, 8192, 0), 6528); + assert_eq!(split(&r, 6528, 9000, 9000, 2472, 0), 2176); + // The last stop of 8,704 tokens is `(8704 - 1) // 2176 * 2176 == 6528` + // (v0.32 sched:476-480). + assert_eq!(split(&r, 4352, 8704, 8704, 4352, 0), 2176); + let v31 = rules_glm(VllmRuleVersion::V0_31_0, false); + assert_eq!(split(&v31, 4352, 8704, 8704, 4352, 0), 4352); + // Checkpoint mode (v0.32 sched:482-501, iface:1125-1146). For + // (6528, 9000, pos 8704): 6528 % 2176 == 0; cdiv(9000, 2176) - 2 == 3 + // > 6527 / 2176 == 2; 6528 + 2176 <= 8704; 6528 < 8704 < 9000; + // 2176 % 16 == 0. + let ckpt = rules_glm(VllmRuleVersion::V0_32_0, true); + assert_eq!(ckpt.checkpoint_position(9000), 8704); + assert!(ckpt.checkpoint_valid(6528, 9000, 8704)); + assert_eq!(split(&ckpt, 6528, 9000, 9000, 2472, 0), 2472); + // For (4352, 8704, pos 6528): 4352 % 2176 == 0; cdiv(8704, 2176) - 2 + // == 2 > 4351 / 2176 == 1; 4352 + 2176 <= 6528; 4352 < 6528 < 8704; + // 2176 % 16 == 0. + assert_eq!(ckpt.checkpoint_position(8704), 6528); + assert!(ckpt.checkpoint_valid(4352, 8704, 6528)); + assert_eq!(split(&ckpt, 4352, 8704, 8704, 4352, 0), 4352); + // Without a checkpoint rule the check is false. + assert!(!r.checkpoint_valid(6528, 9000, 8704)); +} + +#[test] +fn align_split_with_a_checkpoint_block_ports_vllm_0320() { + // Ported from v0.32 split_test:127-180, which stubs `block_size` to + // 1,600 and `hash_block_size` to 16 (split_test:139-153). The engine + // shape with the same numbers: one full-attention rule and one state + // rule at block 1,600, hash 16, checkpoint 1 at alignment 16. The + // state block is above the hash block and the other rule is full + // attention, so fine-grained hits are on and A is 16. + let groups = [attention("full", 1600, None), state("s", 1600, true, true)]; + let r = GroupedPrefixRules::new(&groups, 16, VllmRuleVersion::V0_32_0).unwrap(); + assert_eq!(r.alignment, 16); + assert_eq!(r.cacheable_block, 1600); + // The engine's checkpoint position is `(2002 - 1) / 16 * 16 == 2000`. + // The vLLM test stubs partial hits off and gets 1,600. Both positions + // pass the five conditions (0 % 16 == 0; cdiv(2002, 1600) - 2 == 0 > + // -1; 0 + 16 <= pos; 0 < pos < 2002; pos % 16 == 0), so both give + // 2,002. + assert_eq!(r.checkpoint_position(2002), 2000); + assert!(r.checkpoint_valid(0, 2002, 2000)); + assert!(r.checkpoint_valid(0, 2002, 1600)); + assert_eq!(r.align_split(0, 2002, 2002, 2002, 0, 16384), 2002); + // The end 2,000 is below the prompt, so the floor to 1,600 applies + // and no checkpoint is taken. + assert_eq!(r.align_split(0, 3602, 3602, 2000, 0, 16384), 1600); + let groups = [attention("full", 1600, None), state("s", 1600, true, false)]; + let v31 = GroupedPrefixRules::new(&groups, 16, VllmRuleVersion::V0_31_0).unwrap(); + assert_eq!(v31.align_split(0, 2002, 2002, 2002, 0, 16384), 1600); + assert_eq!(v31.align_split(0, 3602, 3602, 2000, 0, 16384), 1600); +} + +#[test] +fn align_split_needs_a_mamba_state_group() { + // vLLM splits only with Mamba layers in align mode (v0.31 `sched:354-356`, + // `:695-697`, v0.32 `sched:377-379`, `:724-726`). Without a state rule + // every chunk is the desired count, on both rows. + let full = [attention("full", 64, None)]; + let windowed = [attention("full", 64, None), attention("swa", 64, Some(128))]; + // A tail ring, such as the GLM `KpoolTailSpec`, is not a MambaSpec. + let ring = [attention("full", 64, None), state("ring", 4, false, false)]; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + for groups in [&full[..], &windowed[..], &ring[..]] { + let r = GroupedPrefixRules::new(groups, 64, version).unwrap(); + // The last block boundary of the prompt is no stop. + assert_eq!(split(&r, 0, 300, 300, 300, 0), 300, "{version:?}"); + // A chunk that ends below the prompt keeps its length. + assert_eq!(split(&r, 0, 1000, 1000, 500, 0), 500, "{version:?}"); + // A start inside a block runs past the next boundary. + assert_eq!(split(&r, 100, 1000, 1000, 500, 0), 500, "{version:?}"); + // The junction is no stop. + assert_eq!(split(&r, 64, 1000, 1000, 800, 640), 800, "{version:?}"); + } + // A state group keeps the split: the floor of 8,192 to 2,176. + assert_eq!( + split(&rules_glm(version, false), 0, 9000, 9000, 8192, 0), + 6528, + "{version:?}" + ); + } +} + +#[test] +fn partial_tail_moves_one_unit_down_on_vllm_0320() { + // At hash 4 the 0.32.0 tail, `(prompt - 1) / hash * hash`, is 8,996, + // between the last stop 8,704 and the prompt (v0.32 sched:517-525, + // single:860-868, single:2069-2073). The 0.31.0 tail is 9,000. + let groups = groups_glm(false); + let hash = 4; + for (version, expected_ends, expected_hit) in [ + (VllmRuleVersion::V0_32_0, vec![6528, 8704, 8996, 9000], 8996), + (VllmRuleVersion::V0_31_0, vec![6528, 8704, 9000], 8704), + ] { + let mut coordinator = with_limits(rules_glm_hash(version, false, hash)); + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(9000, hash); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + let mut block_lease = coordinator.new_lease(); + let ends = prefill( + &mut coordinator, + pool, + &mut block_lease, + &request, + 9000, + &hit, + BUDGET, + ); + assert_eq!(ends, expected_ends, "{version:?}"); + if version == VllmRuleVersion::V0_32_0 { + for group in GLM_STATES { + assert!(hit_exists(pool, group, hash, 8704)); + assert!(hit_exists(pool, group, hash, 8996)); + } + } + let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); + assert_eq!(resend.tokens, expected_hit, "{version:?}"); + coordinator.release_request(pool, &mut block_lease); + let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); + assert_eq!(resend.tokens, expected_hit, "{version:?}"); + } +} + +#[test] +fn checkpoint_block_is_taken_in_the_final_chunk_and_published_at_the_prompt_checkpoint() { + // v0.32 single:1643-1674, :1745-1773 (one more block in the chunk + // that reaches the prompt), :2017-2066 (published at the prompt + // checkpoint under retention 0). + let groups = groups_glm(true); + let hash = 2176; + let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_32_0, true)); + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(9000, hash); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + assert_eq!(hit.tokens, 0); + // 5 for G0, 1 for the ring, 2 for each KDA group. + assert_eq!(coordinator.admission_blocks(pool, &hit, 9000), 14); + let v31 = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); + assert_eq!(v31.admission_blocks(pool, &hit, 9000), 10); + + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert_eq!(split(coordinator.rules(), 0, 9000, 9000, BUDGET, 0), 6528); + assert!(coordinator.allocate(pool, &mut block_lease, 9000, 0, 6528, 1)); + for group in GLM_STATES { + let state = state_lease(&block_lease, usize::from(group)); + assert!(state.running.is_some()); + assert!(state.previous.is_none() && state.checkpoint.is_none() && state.hit.is_none()); + } + coordinator.finalize(pool, &mut block_lease, &request, 9000, 6528, &hit, 1); + coordinator.begin_step(); + assert_eq!( + split(coordinator.rules(), 6528, 9000, 9000, BUDGET, 0), + 2472 + ); + assert!(coordinator.allocate(pool, &mut block_lease, 9000, 6528, 9000, 2)); + for group in GLM_STATES { + let state = state_lease(&block_lease, usize::from(group)); + assert!(state.running.is_some() && state.previous.is_some()); + assert_eq!(state.checkpoint.map(|(_, position)| position), Some(8704)); + } + // G0 5, ring 1, 4 x (running + previous + checkpoint). + assert_eq!(pool.active_blocks(), 5 + 1 + 12); + coordinator.finalize(pool, &mut block_lease, &request, 9000, 9000, &hit, 2); + for group in GLM_STATES { + assert!(hit_exists(pool, group, hash, 8704), "group {group}"); + let state = state_lease(&block_lease, usize::from(group)); + let (checkpoint, _) = state.checkpoint.unwrap(); + assert!(!pool.is_private(checkpoint)); + assert!(pool.is_private(state.running.unwrap())); + } + coordinator.begin_step(); + let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); + assert_eq!(resend.tokens, 8704); + coordinator.release_request(pool, &mut block_lease); + // Four G0 blocks and the four checkpoint copies stay in the LRU. + assert_eq!(pool.inactive_cached_blocks(), 8); + assert_eq!(pool.active_blocks(), 0); + let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); + assert_eq!(resend.tokens, 8704); +} + +#[test] +fn checkpoint_block_goes_at_the_next_step_and_stays_cached_until_eviction() { + // A Mamba manager skips every token before the last computed one (v0.32 + // single:1954-1959). So the request's next allocation first frees every + // block below its last block, the checkpoint block included (v0.32 + // kv_cache_manager.py:546-559, single:672-710, :1592-1641). A freed + // block keeps its key in the LRU (v0.32 block_pool.py:774-803). + let groups = groups_glm(true); + let hash = 2176; + let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_32_0, true)); + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(9000, hash); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + let mut block_lease = coordinator.new_lease(); + let ends = prefill( + &mut coordinator, + pool, + &mut block_lease, + &request, + 9000, + &hit, + BUDGET, + ); + assert_eq!(ends, [6528, 9000]); + // G0 5, ring 1, 4 x (running + previous + checkpoint). + assert_eq!(pool.active_blocks(), 5 + 1 + 12); + // The prefill ran steps 1 and 2. Each decode step holds G0 5, the ring + // and the four running blocks only, and a resend still hits the + // checkpoint keys. + for (step, target) in [(3, 9001), (4, 9002)] { + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 9000, target - 1, target, step)); + assert_eq!(pool.active_blocks(), 5 + 1 + 4, "step {step}"); + for group in GLM_STATES { + let state = state_lease(&block_lease, usize::from(group)); + assert!(state.checkpoint.is_none() && state.previous.is_none()); + let cached = pool.key_hit(group, hash_at(hash, 8704)); + assert!(cached.is_some_and(|hit| !hit.is_active), "group {group}"); + } + coordinator.finalize(pool, &mut block_lease, &request, 9000, target, &hit, step); + let resend = coordinator.find_longest_hit(pool, &lease(9000, hash), 9000); + assert_eq!(resend.tokens, 8704, "step {step}"); + } + // The checkpoint keys go only when a new block evicts them, oldest first + // in the order the groups freed them (v0.32 coord:447-450). + assert_eq!(pool.inactive_cached_blocks(), 4); + let free = pool.capacity_blocks() - pool.active_blocks() - 4; + let mut reservation = pool.reserve(&[], free + 1).unwrap(); + for _ in 0..=free { + pool.take_private(&mut reservation); + } + pool.cancel(reservation); + assert!(!hit_exists(pool, 2, hash, 8704)); + for group in [3, 4, 5] { + assert!(hit_exists(pool, group, hash, 8704), "group {group}"); + } +} + +#[test] +fn checkpoint_block_publishes_the_prompt_checkpoint_under_retention_zero() { + // Ported from v0.32 partial_test:775-807: the prompt-end checkpoint + // `(prompt - 1) / 32 * 32` stays published under retention 0. + let groups = [attention("full", 32, None), state("s", 32, true, true)]; + let hash = 32; + for (prompt, expected) in [(120, 96), (104, 96), (96, 64)] { + let mut coordinator = + with_limits(GroupedPrefixRules::new(&groups, hash, VllmRuleVersion::V0_32_0).unwrap()); + let mut kv = kv_pool(&groups, 16); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(prompt, hash); + let hit = coordinator.find_longest_hit(pool, &request, prompt); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, prompt, 0, prompt, 1)); + assert_eq!( + state_lease(&block_lease, 1).checkpoint.map(|(_, p)| p), + Some(expected) + ); + coordinator.finalize(pool, &mut block_lease, &request, prompt, prompt, &hit, 1); + assert!(hit_exists(pool, 1, hash, expected), "prompt {prompt}"); + coordinator.release_request(pool, &mut block_lease); + coordinator.begin_step(); + let resend = coordinator.find_longest_hit(pool, &lease(prompt, hash), prompt); + assert_eq!(resend.tokens, expected, "prompt {prompt}"); + } +} + +#[test] +fn state_key_published_in_a_step_is_not_a_hit_in_that_step() { + // single:1726-1734: a state block cached in this step is not a hit + // until the next step. The block pool stores no step, so the coordinator + // keeps the set. + let groups = groups_glm(false); + let hash = 2176; + let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); + let mut kv = kv_pool(&groups, 64); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(2500, hash); + let hit = coordinator.find_longest_hit(pool, &request, 2500); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 2500, 0, 2176, 1)); + coordinator.finalize(pool, &mut block_lease, &request, 2500, 2176, &hit, 1); + for group in GLM_STATES { + assert!(hit_exists(pool, group, hash, 2176)); + } + let same_step = coordinator.find_longest_hit(pool, &lease(2500, hash), 2500); + assert_eq!(same_step.tokens, 0); + coordinator.begin_step(); + let next_step = coordinator.find_longest_hit(pool, &lease(2500, hash), 2500); + assert_eq!(next_step.tokens, 2176); +} + +#[test] +fn allocate_fails_atomically() { + let groups = groups_glm(false); + let mut coordinator = with_limits(rules_glm(VllmRuleVersion::V0_31_0, false)); + // Room for the first chunk (G0 1, ring 1, 4 running) but not the + // second (G0 1 more, 4 new running blocks while 4 previous stay). + let mut kv = kv_pool(&groups, 9); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(9000, 2176); + let hit = coordinator.find_longest_hit(pool, &request, 9000); + let mut block_lease = coordinator.new_lease(); + coordinator.begin_step(); + assert!(coordinator.allocate(pool, &mut block_lease, 9000, 0, 2176, 1)); + coordinator.finalize(pool, &mut block_lease, &request, 9000, 2176, &hit, 1); + assert_eq!(pool.active_blocks(), 6); + coordinator.begin_step(); + let before = format!("{block_lease:?}"); + assert!(!coordinator.allocate(pool, &mut block_lease, 9000, 2176, 4352, 2)); + assert_eq!(pool.active_blocks(), 6); + assert_eq!(format!("{block_lease:?}"), before); + coordinator.release_request(pool, &mut block_lease); + assert_eq!(pool.active_blocks(), 0); +} + +#[test] +fn release_frees_groups_in_profile_order_each_tail_first() { + // coord:387-395 frees the managers in group order, and each manager frees + // its blocks tail first (single:581-589). Cached blocks join the LRU tail + // in that order (pool:777-808). So the cached block of a state group that + // the profile lists first is evicted before the attention blocks. + let groups = [state("s", 64, true, false), attention("full", 64, None)]; + let hash = 64; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let mut coordinator = with_limits(GroupedPrefixRules::new(&groups, hash, version).unwrap()); + let mut kv = kv_pool(&groups, 6); + let pool = kv.block_pool_mut().unwrap(); + let request = lease(200, hash); + let hit = coordinator.find_longest_hit(pool, &request, 200); + let mut block_lease = coordinator.new_lease(); + let ends = prefill( + &mut coordinator, + pool, + &mut block_lease, + &request, + 200, + &hit, + 256, + ); + assert_eq!(ends, vec![192, 200], "{version:?}"); + coordinator.release_request(pool, &mut block_lease); + // The state at the replay boundary 192 and the attention blocks that + // end at 64, 128 and 192 stay cached. Two blocks are free. + assert_eq!(pool.inactive_cached_blocks(), 4, "{version:?}"); + assert_eq!(pool.active_blocks(), 0, "{version:?}"); + // Three fresh blocks evict one cached block: the state block. + let reservation = pool.reserve(&[], 3).unwrap(); + pool.cancel(reservation); + assert!(!hit_exists(pool, 0, hash, 192), "{version:?}"); + for end in [64, 128, 192] { + assert!(hit_exists(pool, 1, hash, end), "{version:?} at {end}"); + } + // Four fresh blocks evict one more: the attention tail block. + let reservation = pool.reserve(&[], 4).unwrap(); + pool.cancel(reservation); + assert!(!hit_exists(pool, 1, hash, 192), "{version:?}"); + for end in [64, 128] { + assert!(hit_exists(pool, 1, hash, end), "{version:?} at {end}"); + } + } +} + +#[test] +fn checkpoint_rule_needs_state_groups_at_the_smallest_cacheable_block() { + // vLLM 0.32.0 tests a checkpoint against `cache_config.block_size`, the + // smallest prefix-cacheable block, in the scheduler (v0.32 sched:469, + // :496) and against the Mamba block in the manager (v0.32 single:1652, + // :1662). The rule table holds one block for both. + let small_attention = [attention("full", 16, None), state("s", 64, true, true)]; + let err = GroupedPrefixRules::new(&small_attention, 16, VllmRuleVersion::V0_32_0) + .unwrap_err() + .to_string(); + assert_eq!( + err, + "kv_cache_groups with vllm_version=0.32.0 require every state group with prefill checkpoint blocks to have the smallest block_size_tokens of the prefix_cacheable groups" + ); + // Without checkpoint blocks the same shape is modeled on both rows. + let no_checkpoint = [attention("full", 16, None), state("s", 64, true, false)]; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + assert!(GroupedPrefixRules::new(&no_checkpoint, 16, version).is_ok()); + } + // A larger attention block keeps the state block the smallest. The + // checkpoint position then moves in units of the scheduler block, the LCM + // 128, in the manager (v0.32 coord:699-706) and in the scheduler (v0.32 + // sched:473-475, :209): `(300 - 1) / 128 * 128 == 256`. + let large_attention = [attention("full", 128, None), state("s", 64, true, true)]; + let r = GroupedPrefixRules::new(&large_attention, 64, VllmRuleVersion::V0_32_0).unwrap(); + assert_eq!((r.alignment, r.cacheable_block), (128, 64)); + assert_eq!(r.checkpoint_position(300), 256); +} + +#[test] +fn retired_window_blocks_are_released_tail_first() { + // v0.31 single:660-685 and v0.32 single:645-670 free a retired window + // range backward, and free_blocks appends its cached blocks to the free + // queue in that order (v0.31 pool:777-806, v0.32 pool:774-803). So the + // tail block of the range is evicted first. + let groups = [attention("window", 4, Some(8))]; + let hash = 4; + for version in [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0] { + let mut coordinator = with_limits(GroupedPrefixRules::new(&groups, hash, version).unwrap()); + let mut kv = kv_pool(&groups, 8); + let pool = kv.block_pool_mut().unwrap(); + // The run that a finished request of 100 tokens keeps before its + // replay boundary 96: the blocks that end at 92 and 96. + seed_many(pool, 0, hash, &[92, 96]); + let request = lease(108, hash); + let hit = coordinator.find_longest_hit(pool, &request, 108); + assert_eq!(hit.tokens, 96, "{version:?}"); + let mut block_lease = coordinator.new_lease(); + coordinator.pin_hit(pool, &mut block_lease, &hit).unwrap(); + assert!(coordinator.allocate(pool, &mut block_lease, 108, 96, 104, 1)); + coordinator.finalize(pool, &mut block_lease, &request, 108, 104, &hit, 1); + // The next allocation retires both hit blocks: (104 - 7) / 4 = 24. + assert!(coordinator.allocate(pool, &mut block_lease, 108, 104, 108, 2)); + assert_eq!(pool.inactive_cached_blocks(), 2, "{version:?}"); + assert_eq!(pool.active_blocks(), 3, "{version:?}"); + // Three blocks are free, so four fresh blocks evict one cached block: + // the tail block of the retired range, which ends at 96. + let reservation = pool.reserve(&[], 4).unwrap(); + pool.cancel(reservation); + assert!(!hit_exists(pool, 0, hash, 96), "{version:?}"); + assert!(hit_exists(pool, 0, hash, 92), "{version:?}"); + } +} + +/// The sliding-window search before the single traversal: every candidate +/// end scans its own run again. +fn rescanning_window_run( + max: usize, + block: usize, + step: usize, + needed: usize, + mut resident: impl FnMut(usize) -> bool, +) -> (usize, usize) { + let mut end = max / step * step; + while end > 0 { + let blocks = end / block; + let mut run = 0; + while run < blocks && run < needed && resident(blocks - run - 1) { + run += 1; + } + if run >= needed || run == blocks { + return (end, run); + } + end -= step; + } + (0, 0) +} + +#[test] +fn window_search_looks_up_each_block_at_most_once() { + // single:1018-1044 finds a sliding-window hit in one backward traversal. + // 4,096 candidate blocks of which the last 2,048 are resident, and a + // window that needs 4,096 contiguous blocks: no hit. + let (block, needed) = (16, 4096); + let resident = |i: usize| i >= 2048; + let mut lookups = 0; + let found = window_run(4096 * block, block, block, needed, |i| { + lookups += 1; + resident(i) + }); + let mut rescans = 0; + let rescanned = rescanning_window_run(4096 * block, block, block, needed, |i| { + rescans += 1; + resident(i) + }); + assert_eq!(found, (0, 0)); + assert_eq!(rescanned, found); + assert_eq!(rescans, 2_102_272); + assert_eq!(lookups, 4096); +} + +#[test] +fn window_search_finds_the_hits_of_the_rescanning_search() { + // A fixed xorshift sequence: random residency, blocks, alignments, + // windows and lengths. Each block is looked up at most once. + let mut state: u64 = 0x9E37_79B9_7F4A_7C15; + let mut next = |bound: u64| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state % bound + }; + for case in 0..20_000 { + let block = 1 + next(8) as usize; + let step = block * (1 + next(4) as usize); + let needed = next(7) as usize; + let max = next(400) as usize; + let density = next(101); + let blocks = max / block; + let resident: Vec = (0..blocks).map(|_| next(100) < density).collect(); + let mut lookups = vec![0; blocks]; + let found = window_run(max, block, step, needed, |i| { + lookups[i] += 1; + resident[i] + }); + let expected = rescanning_window_run(max, block, step, needed, |i| resident[i]); + assert_eq!( + found, expected, + "case {case}: max {max} block {block} step {step} needed {needed} {resident:?}" + ); + assert!(lookups.iter().all(|&count| count <= 1), "case {case}"); + } +} + +/// Whether a scan of every completed attention block, as finalization did +/// before it kept a cursor, would still find a block to publish. +fn rescan_finds_a_block( + coordinator: &GroupedPrefixCoordinator, + pool: &GroupedBlockPool, + lease: &GroupedBlockLease, + request: &BlockRequestLease, + prompt: usize, + computed_after: usize, + hit: &GroupedHit, +) -> bool { + use crate::engine::cache::vllm_block_pool::BlockCopyId; + let rules = coordinator.rules(); + let junction = (hit.junction_tokens > 0).then_some(hit.tokens + hit.junction_tokens); + let publishable = |g: usize, copy: Option<&Option>, end: usize| { + copy.is_some_and(|copy| copy.is_some_and(|id| pool.is_private(id))) + && grouped_key(g, request.prefix_hash_at(end)).is_some() + }; + rules.groups.iter().enumerate().any(|(g, rule)| { + let GroupLease::Attention { copies, .. } = &lease.groups[g] else { + return false; + }; + let block = rule.block; + match rule.kind { + GroupRuleKind::FullAttention => { + (0..computed_after / block).any(|i| publishable(g, copies.get(i), (i + 1) * block)) + } + GroupRuleKind::SlidingWindow { window } => [ + Some(rules.replay_boundary(prompt, block)), + junction.map(|junction| rules.block_floor(junction, block)), + ] + .into_iter() + .flatten() + .filter(|&pos| pos > 0 && pos <= computed_after) + .any(|pos| { + let top = pos / block; + (top.saturating_sub((window - 1).div_ceil(block))..top) + .any(|i| publishable(g, copies.get(i), (i + 1) * block)) + }), + _ => false, + } + }) +} + +#[test] +fn finalize_publishes_what_a_full_rescan_would() { + // Two requests share a prefix on a full-attention and a sliding-window + // group, with chunked prefills and decode steps. After every step a scan + // of all completed blocks finds nothing more to publish, so the cursor + // skips only blocks that are cached or have no key. + let mut state: u64 = 0xD1B5_4A32_D192_ED03; + let mut next = |bound: u64| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state % bound) as usize + }; + let mut hits = 0; + for case in 0..400 { + let hash = [4, 8][next(2)]; + let groups = [ + attention("full", (hash * (1 + next(2))) as u32, None), + attention( + "window", + (hash * (1 + next(2))) as u32, + Some(2 + next(40) as u32), + ), + ]; + let version = [VllmRuleVersion::V0_31_0, VllmRuleVersion::V0_32_0][next(2)]; + let mut coordinator = with_limits(GroupedPrefixRules::new(&groups, hash, version).unwrap()); + let mut kv = kv_pool(&groups, 4096); + let pool = kv.block_pool_mut().unwrap(); + let first = 1 + next(300); + // Half of the cases resend the whole first prompt, which reaches its + // replay boundary. + let shared = if next(2) == 0 { + first + } else { + next(first as u64 + 1) + }; + let second: Vec = (0..shared as u32) + .chain(10_000..10_000 + (1 + next(300)) as u32) + .collect(); + let first: Vec = (0..first as u32).collect(); + let mut step = 1; + for tokens in [first, second] { + let mut request = BlockRequestLease::new(Uuid::nil(), Vec::new()); + request.configure_prefix_hashes(&tokens, hash); + let prompt = tokens.len(); + let hit = coordinator.find_longest_hit(pool, &request, prompt); + hits += usize::from(hit.tokens > 0); + let mut block_lease = coordinator.new_lease(); + coordinator.pin_hit(pool, &mut block_lease, &hit).unwrap(); + // A chunked prefill, then up to eight decode steps. + let (budget, end) = (1 + next(64), prompt + 1 + next(8)); + let mut computed = hit.tokens; + while computed < end { + coordinator.begin_step(); + let target = (computed + budget).min(prompt).max(computed + 1); + assert!(coordinator.allocate( + pool, + &mut block_lease, + prompt, + computed, + target, + step + )); + coordinator.finalize(pool, &mut block_lease, &request, prompt, target, &hit, step); + assert!( + !rescan_finds_a_block( + &coordinator, + pool, + &block_lease, + &request, + prompt, + target, + &hit + ), + "case {case} {groups:?} prompt {prompt} hit {} at {target}", + hit.tokens + ); + computed = target; + step += 1; + } + coordinator.release_request(pool, &mut block_lease); + } + } + assert!(hits > 150, "{hits} hits"); +} + +#[test] +fn first_allocation_test_counts_the_blocks_that_the_first_allocation_takes() { + // The test before a first allocation (mgr:558-576) counts the inactive + // hit blocks and the fresh blocks of the chunk. pin_hit and the first + // allocate hold exactly these at their peak, so the test never refuses a + // chunk that fits and never admits one that does not. + let layouts = [ + ( + vec![ + attention("full", 16, None), + attention("window", 16, Some(40)), + ], + 16, + VllmRuleVersion::V0_31_0, + ), + ( + vec![attention("full", 64, None), state("s", 64, true, false)], + 16, + VllmRuleVersion::V0_31_0, + ), + ( + vec![attention("full", 64, None), state("s", 64, true, false)], + 16, + VllmRuleVersion::V0_32_0, + ), + ( + vec![ + attention("full", 64, None), + state("ring", 4, false, false), + state("s", 64, true, true), + ], + 64, + VllmRuleVersion::V0_32_0, + ), + ]; + let mut state: u64 = 0x2545_F491_4F6C_DD1D; + let mut next = |bound: u64| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state % bound) as usize + }; + let mut hits = 0; + for (groups, hash, version) in &layouts { + for case in 0..150 { + // A small step budget, so that the admission cap of the window + // binds. The allocation test leaves it out. + let mut coordinator = GroupedPrefixCoordinator::new( + GroupedPrefixRules::new(groups, *hash, *version).unwrap(), + AdmissionLimits { + max_in_flight_tokens: 16, + max_model_len: usize::MAX, + }, + ); + let mut kv = kv_pool(groups, 512); + let pool = kv.block_pool_mut().unwrap(); + let first = 1 + next(400); + let request = lease(first, *hash); + let hit = coordinator.find_longest_hit(pool, &request, first); + let mut block_lease = coordinator.new_lease(); + let budget = 64 * (1 + next(4)); + prefill( + &mut coordinator, + pool, + &mut block_lease, + &request, + first, + &hit, + budget, + ); + coordinator.release_request(pool, &mut block_lease); + // A second request shares a prefix of the first one. + let shared = if next(2) == 0 { + first + } else { + next(first as u64 + 1) + }; + let tokens: Vec = (0..shared as u32) + .chain(10_000..10_000 + (1 + next(300)) as u32) + .collect(); + let mut request = BlockRequestLease::new(Uuid::nil(), Vec::new()); + request.configure_prefix_hashes(&tokens, *hash); + let prompt = tokens.len(); + coordinator.begin_step(); + let hit = coordinator.find_longest_hit(pool, &request, prompt); + hits += usize::from(hit.tokens > 0); + let target = if next(2) == 0 { + prompt + } else { + hit.tokens + 1 + next((prompt - hit.tokens) as u64) + }; + let needed = coordinator.first_allocation_blocks(pool, &hit, prompt, target); + let before = pool.active_blocks(); + pool.reset_peak_active_blocks(); + let mut block_lease = coordinator.new_lease(); + coordinator.pin_hit(pool, &mut block_lease, &hit).unwrap(); + assert!(coordinator.allocate(pool, &mut block_lease, prompt, hit.tokens, target, 100)); + // The copy of a shared block is taken while the hit block is + // still held, so the peak counts both. + assert_eq!( + pool.peak_active_blocks() - before, + needed, + "{version:?} case {case}: prompt {prompt} hit {} target {target}", + hit.tokens + ); + coordinator.release_request(pool, &mut block_lease); + } + } + assert!(hits > 200, "{hits} hits"); +} diff --git a/crates/core/src/engine/kv_manager/mod.rs b/crates/core/src/engine/kv_manager/mod.rs index 529264c063..f70d871078 100644 --- a/crates/core/src/engine/kv_manager/mod.rs +++ b/crates/core/src/engine/kv_manager/mod.rs @@ -46,7 +46,10 @@ mod vllm_firewall_tests; pub(crate) use g1_manager::{ DestinationReservation, G1Manager, NativeAllocation, SourceReleaseOrder, SourceReuseDependency, }; +#[cfg(test)] +pub(crate) use grouped::GroupedBlockPool; pub(crate) use grouped::GroupedKvPool; +pub(crate) use grouped_prefix::{AdmissionLimits, GroupedPrefixRules, VllmRuleVersion}; pub(crate) use sglang_backend::SglangKvManager; pub(crate) use vllm_backend::BlockRequestLease; diff --git a/crates/core/src/engine/kv_manager/vllm_backend.rs b/crates/core/src/engine/kv_manager/vllm_backend.rs index e7214b4721..5d30f70239 100644 --- a/crates/core/src/engine/kv_manager/vllm_backend.rs +++ b/crates/core/src/engine/kv_manager/vllm_backend.rs @@ -1,5 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project. +// Modified Rust simulation of the vLLM v0.29.0 +// (98dff2a81d747d1dba01a47f939f48c3526d4206) align-mode state cache and of +// the v0.31.0 (db9527a46873454610df6dbedf79a36d6bf1a7f6) and v0.32.0 +// (5df63a7203cdb9328d824a27b682e77c7b9928be) grouped KV rules. +// Immutable sources and license: THIRD_PARTY_NOTICES.md. //! vLLM G1 manager over a minimal physical block-pool model. //! @@ -11,6 +17,9 @@ use uuid::Uuid; use super::AllocationRequirement; use super::grouped::{GroupedKvPool, GroupedLease}; +use super::grouped_prefix::{ + AdmissionLimits, GroupedHit, GroupedPrefixCoordinator, GroupedPrefixRules, +}; use super::state_cache_manager::{StateCacheManager, StateRequest}; use crate::engine::belady::BeladyOracle; use crate::engine::cache::vllm_block_pool::{ @@ -130,10 +139,15 @@ pub(crate) struct BlockRequestLease { entries: Vec, allocated_tokens: usize, prefix_hashes: Option<(usize, Vec)>, + /// The prompt length behind `prefix_hashes`. + prompt_tokens: usize, /// Present only while newly acquired capacity awaits a scheduler-installed /// write fence. The default path does not allocate this state. pending_capacity_writes: Option>, grouped: Option, + /// The grouped prefix hit found at admission, so that the allocation + /// does not search again. Cleared by preemption. + grouped_hit: Option, state: Option>, pub(crate) session_id: Option, } @@ -152,14 +166,17 @@ impl BlockRequestLease { entries, allocated_tokens: 0, prefix_hashes: None, + prompt_tokens: 0, pending_capacity_writes: None, grouped: None, + grouped_hit: None, state: None, session_id: None, } } pub(crate) fn configure_prefix_hashes(&mut self, tokens: &[u32], unit: usize) { + self.prompt_tokens = tokens.len(); let mut parent = None; let hashes = block_hashes(tokens, unit) .map(|local| { @@ -203,6 +220,14 @@ impl BlockRequestLease { self.owner } + /// The end of the uncached common prefix of the grouped hit, `hit + + /// junction` (`mgr:312-321`). Zero without a hit. + pub(crate) fn grouped_junction(&self) -> usize { + self.grouped_hit + .as_ref() + .map_or(0, |hit| hit.tokens + hit.junction_tokens) + } + #[cfg(test)] pub(crate) fn computed_state_tokens(&self) -> Option { self.state.as_ref().map(|state| state.computed_tokens) @@ -380,20 +405,47 @@ pub(crate) struct VllmKvManager { dp_rank: u32, next_event_id: u64, grouped: Option, + /// The prefix cache rules of a grouped cache in block mode. + grouped_prefix: Option, state_cache: Option, prefix_match_unit: Option, + /// Scheduler steps begun, for the one-step block lifetimes of the + /// grouped coordinator. + step: u64, } impl VllmKvManager { pub(crate) fn set_grouped_cache(&mut self, grouped: GroupedKvPool) { - assert!(!self.enable_prefix_caching); self.grouped = Some(grouped); } + /// Turn grouped prefix caching on. The pool must be in block mode. + pub(crate) fn set_grouped_prefix_rules( + &mut self, + rules: GroupedPrefixRules, + limits: AdmissionLimits, + ) { + assert!( + self.enable_prefix_caching + && self + .grouped + .as_ref() + .is_some_and(|grouped| grouped.block_pool().is_some()), + "grouped prefix rules need prefix caching and a block-mode pool" + ); + self.grouped_prefix = Some(GroupedPrefixCoordinator::new(rules, limits)); + } + pub(crate) fn grouped(&self) -> Option<&GroupedKvPool> { self.grouped.as_ref() } + pub(crate) fn grouped_prefix_rules(&self) -> Option<&GroupedPrefixRules> { + self.grouped_prefix + .as_ref() + .map(GroupedPrefixCoordinator::rules) + } + pub(crate) fn set_belady_oracle(&mut self, oracle: BeladyOracle) { self.pool.set_belady_oracle(oracle); } @@ -423,8 +475,10 @@ impl VllmKvManager { dp_rank, next_event_id: 0, grouped: None, + grouped_prefix: None, state_cache: None, prefix_match_unit: None, + step: 0, } } @@ -456,8 +510,8 @@ impl VllmKvManager { pub(crate) fn configure_prefix_match_unit(&mut self, unit: Option) { assert!( - unit.is_none() || self.state_cache.is_some(), - "validated state_cache" + unit.is_none() || self.state_cache.is_some() || self.grouped.is_some(), + "validated state_cache or kv_cache_groups" ); self.prefix_match_unit = unit; if let Some(state) = &mut self.state_cache { @@ -479,8 +533,9 @@ impl VllmKvManager { requires_logits: bool, cost: PrefillCost, ) -> PrefillCost { - if self.state_cache.is_some() { + if self.state_cache.is_some() || self.grouped_prefix.is_some() { // Joint lookup already applied bounds before selecting a real state. + // The grouped hit is capped below the prompt in `find_longest_hit`. cost } else { apply_prefix_recompute( @@ -572,6 +627,10 @@ impl VllmKvManager { if let Some(manager) = &mut self.state_cache { manager.begin_step(&mut self.pool); } + if let Some(coordinator) = &mut self.grouped_prefix { + coordinator.begin_step(); + } + self.step = self.step.checked_add(1).expect("scheduler step overflow"); } pub(crate) fn state_cache_blocks(&self) -> usize { @@ -632,9 +691,66 @@ impl VllmKvManager { ) -> NativeAllocation { lease.debug_assert_owner(owner); if let Some(grouped) = &mut self.grouped { - assert_eq!(reusable_prefix_blocks, 0); + assert!(self.enable_prefix_caching || reusable_prefix_blocks == 0); let previous_blocks = lease.resident_block_count(); - if !grouped.allocate( + if let Some(coordinator) = &self.grouped_prefix { + let pool = grouped.block_pool_mut().expect("block mode"); + let hit = lease + .grouped_hit + .as_ref() + .expect("a grouped prefill cost precedes the allocation"); + let prompt = lease.prompt_tokens; + let block_lease = lease + .grouped + .get_or_insert_with(GroupedLease::default) + .block_lease + .get_or_insert_with(|| coordinator.new_lease()); + let first = !block_lease.allocated; + if first { + // The scheduler passes the hit in prefix hash units; the + // stored hit is the same lookup. + debug_assert_eq!( + reusable_prefix_blocks * self.prefix_match_unit.unwrap_or(self.block_size), + hit.tokens + ); + // vLLM tests the blocks of the chunk before it touches the + // hit (`mgr:558-576`), so a chunk that does not fit leaves + // the hit blocks where they are in the inactive LRU. + let needed = + coordinator.first_allocation_blocks(pool, hit, prompt, cumulative_tokens); + if needed > pool.capacity_blocks() - pool.active_blocks() { + return NativeAllocation::CapacityExhausted; + } + if coordinator.pin_hit(pool, block_lease, hit).is_err() { + return NativeAllocation::CapacityExhausted; + } + } else if cumulative_tokens <= lease.allocated_tokens { + return NativeAllocation::Ready { + value: 0, + dependencies: Vec::new(), + }; + } + let computed_before = if first { + hit.tokens + } else { + lease.allocated_tokens + }; + if !coordinator.allocate( + pool, + block_lease, + prompt, + computed_before, + cumulative_tokens, + self.step, + ) { + if first { + // Undo the pin, so that the hit blocks return to the + // inactive LRU. + coordinator.release_request(pool, block_lease); + } + return NativeAllocation::CapacityExhausted; + } + } else if !grouped.allocate( &mut lease.grouped, lease.allocated_tokens, cumulative_tokens, @@ -997,7 +1113,10 @@ impl VllmKvManager { if lease.entries[position].identity.sequence_hash.is_none() { lease.entries[position].identity = sequence.complete_block_identity(position, parent_hash); - lease.entries[position].pending_cache = self.enable_prefix_caching; + // A grouped lease holds no linear copies; its coordinator + // caches the group blocks below. + lease.entries[position].pending_cache = + self.enable_prefix_caching && self.grouped.is_none(); } let entry = &mut lease.entries[position]; @@ -1075,6 +1194,38 @@ impl VllmKvManager { } manager.commit(&mut self.pool, state, computed_after, hash); } + if let (Some(grouped), Some(coordinator)) = (&mut self.grouped, &mut self.grouped_prefix) + && publish_state + && self.enable_prefix_caching + && sequence.enable_prefix_caching() + && let Some(mut block_lease) = lease + .grouped + .as_mut() + .and_then(|grouped| grouped.block_lease.take()) + { + // The block lease leaves the request lease for the call, which + // reads the request's prefix hashes. + if block_lease.allocated { + let hit = lease + .grouped_hit + .as_ref() + .expect("a grouped prefill cost precedes the allocation"); + coordinator.finalize( + grouped.block_pool_mut().expect("block mode"), + &mut block_lease, + lease, + sequence.num_input_tokens(), + computed_after, + hit, + self.step, + ); + } + lease + .grouped + .as_mut() + .expect("grouped lease present") + .block_lease = Some(block_lease); + } #[cfg(debug_assertions)] sequence.debug_assert_finalized_range( lease.entries.len(), @@ -1337,6 +1488,8 @@ impl VllmKvManager { } self.release_lease_entries(lease); lease.allocated_tokens = 0; + // A re-admission searches the hit again. + lease.grouped_hit = None; } pub(crate) fn finish_lease(&mut self, owner: Uuid, mut lease: BlockRequestLease) { @@ -1347,7 +1500,21 @@ impl VllmKvManager { fn release_lease_entries(&mut self, lease: &mut BlockRequestLease) { if let Some(grouped) = &mut self.grouped { - grouped.release(&mut lease.grouped); + if let Some(coordinator) = &self.grouped_prefix { + if let Some(block_lease) = lease + .grouped + .as_mut() + .and_then(|grouped| grouped.block_lease.as_mut()) + { + coordinator.release_request( + grouped.block_pool_mut().expect("block mode"), + block_lease, + ); + } + lease.grouped = None; + } else { + grouped.release(&mut lease.grouped); + } } if let (Some(manager), Some(state)) = (&mut self.state_cache, &mut lease.state) { manager.release(&mut self.pool, state); @@ -1495,11 +1662,73 @@ impl VllmKvManager { } } + /// The prefill cost of a waiting request, and for a grouped prefix cache + /// the hit behind it, stored on the lease for the allocation. + pub(crate) fn lookup_prefill_cost( + &self, + sequence: &RequestSequence, + lease: &mut BlockRequestLease, + ) -> PrefillCost { + if self.grouped_prefix.is_some() { + let (cost, hit) = self.grouped_prefill_cost(sequence, lease); + lease.grouped_hit = Some(hit); + return cost; + } + self.get_lease_prefill_cost(sequence, lease) + } + + /// The whole-prompt admission cost of a grouped prefix cache + /// (`mgr:519-536`): the blocks of `admission_blocks` and the tokens of + /// `find_longest_hit`. A request without prefix caching hits nothing. + fn grouped_prefill_cost( + &self, + sequence: &RequestSequence, + lease: &BlockRequestLease, + ) -> (PrefillCost, GroupedHit) { + let coordinator = self.grouped_prefix.as_ref().expect("grouped prefix rules"); + let pool = self + .grouped + .as_ref() + .and_then(GroupedKvPool::block_pool) + .expect("block mode"); + let known_tokens = sequence.len(); + let hit = if self.enable_prefix_caching && sequence.enable_prefix_caching() { + coordinator.find_longest_hit(pool, lease, known_tokens) + } else { + coordinator.empty_hit() + }; + let active_cached_tokens = coordinator + .active_full_attention_hit_tokens(pool, &hit) + .min(hit.tokens); + let cost = PrefillCost { + new_blocks: coordinator.admission_blocks(pool, &hit, known_tokens), + new_tokens: known_tokens - hit.tokens, + cached_tokens: hit.tokens, + active_cached_tokens, + }; + (cost, hit) + } + + /// The whole-prompt blocks of `known_tokens` without a hit: what a grouped + /// prefix cache request needs from the empty pool. + pub(crate) fn grouped_cold_admission_blocks(&self, known_tokens: usize) -> usize { + let coordinator = self.grouped_prefix.as_ref().expect("grouped prefix rules"); + let pool = self + .grouped + .as_ref() + .and_then(GroupedKvPool::block_pool) + .expect("block mode"); + coordinator.admission_blocks(pool, &coordinator.empty_hit(), known_tokens) + } + pub(crate) fn get_lease_prefill_cost( &self, sequence: &RequestSequence, lease: &BlockRequestLease, ) -> PrefillCost { + if self.grouped_prefix.is_some() { + return self.grouped_prefill_cost(sequence, lease).0; + } let (mut overlap_blocks, mut active_overlap_blocks): (usize, usize) = if self.enable_prefix_caching && sequence.enable_prefix_caching() { let mut overlap = 0; @@ -1767,8 +1996,17 @@ impl VllmKvManager { .map_or_else(|| self.pool.num_active(), GroupedKvPool::active_blocks) } + /// Cached blocks that no request holds. A grouped cache keeps them in its + /// block pool, and byte mode caches none. pub(crate) fn num_inactive_blocks(&self) -> usize { - self.pool.num_inactive() + self.grouped.as_ref().map_or_else( + || self.pool.num_inactive(), + |grouped| { + grouped + .block_pool() + .map_or(0, |pool| pool.inactive_cached_blocks()) + }, + ) } pub(crate) fn max_capacity(&self) -> usize { diff --git a/crates/core/src/engine/protocol.rs b/crates/core/src/engine/protocol.rs index 4165eda8bd..fd7b505e58 100644 --- a/crates/core/src/engine/protocol.rs +++ b/crates/core/src/engine/protocol.rs @@ -363,6 +363,9 @@ pub struct Metrics { pub kv_cache_used_bytes: Option, /// Physical grouped-cache capacity; absent for the legacy linear pool. pub kv_cache_capacity_bytes: Option, + /// Cached grouped-cache bytes that no request holds; present only with + /// grouped prefix caching. + pub kv_cache_cached_bytes: Option, /// Byte occupancy for grouped caches; otherwise `active_blocks / total_blocks`. pub cache_usage: f64, /// Physical resident fraction. This includes inactive reusable vLLM blocks diff --git a/crates/core/src/engine/scheduler/mod.rs b/crates/core/src/engine/scheduler/mod.rs index 481e3ddc4c..cad594b02a 100644 --- a/crates/core/src/engine/scheduler/mod.rs +++ b/crates/core/src/engine/scheduler/mod.rs @@ -184,6 +184,9 @@ pub struct MockerMetrics { pub total_blocks: u64, pub kv_cache_used_bytes: Option, pub kv_cache_capacity_bytes: Option, + /// Bytes of cached grouped-cache blocks that no request holds. Only a + /// grouped cache with prefix caching reports it. + pub kv_cache_cached_bytes: Option, pub gpu_cache_usage_perc: f64, pub physical_gpu_cache_usage_perc: f64, pub running_requests: u64, @@ -250,6 +253,7 @@ impl MockerMetrics { total_blocks, kv_cache_used_bytes: None, kv_cache_capacity_bytes: None, + kv_cache_cached_bytes: None, gpu_cache_usage_perc, physical_gpu_cache_usage_perc, running_requests, diff --git a/crates/core/src/engine/scheduler/rank.rs b/crates/core/src/engine/scheduler/rank.rs index 3ff6380880..c9bb030097 100644 --- a/crates/core/src/engine/scheduler/rank.rs +++ b/crates/core/src/engine/scheduler/rank.rs @@ -535,6 +535,7 @@ fn core_args(config: &EngineConfig, timing: Arc) -> MockEngineA kv_cache_bytes_per_token: config.kv_cache_bytes_per_token, kv_cache_groups: config.kv_cache_groups.clone(), kv_cache_capacity_bytes: config.kv_cache_capacity_bytes, + vllm_version: config.vllm_version.clone(), native_host_offload: config.native_host_offload.clone(), state_cache: config.state_cache, prefix_match_unit: config.prefix_match_unit, @@ -706,6 +707,7 @@ fn map_metrics(metrics: MockerMetrics) -> Metrics { total_blocks: metrics.total_blocks, kv_cache_used_bytes: metrics.kv_cache_used_bytes, kv_cache_capacity_bytes: metrics.kv_cache_capacity_bytes, + kv_cache_cached_bytes: metrics.kv_cache_cached_bytes, cache_usage: metrics.gpu_cache_usage_perc, physical_cache_usage: metrics.physical_gpu_cache_usage_perc, running_requests: metrics.running_requests, diff --git a/crates/core/src/engine/scheduler/vllm/core.rs b/crates/core/src/engine/scheduler/vllm/core.rs index 7f0e366426..dab8a3a507 100644 --- a/crates/core/src/engine/scheduler/vllm/core.rs +++ b/crates/core/src/engine/scheduler/vllm/core.rs @@ -25,14 +25,16 @@ use crate::engine::common::utils::{ }; use crate::engine::host_offload::{G2Binding, HostTierHandle}; use crate::engine::kv_manager::AllocationRequirement; +use crate::engine::kv_manager::{ + AdmissionLimits, G1Manager, GroupedKvPool, GroupedPrefixRules, VllmRuleVersion, +}; use crate::engine::kv_manager::{DestinationReservation, G1Acquire, NativeAllocation}; -use crate::engine::kv_manager::{G1Manager, GroupedKvPool}; use crate::engine::scheduler::queue_metrics::{QueueStats, QueuedLength}; use crate::engine::scheduler::vllm::host_offload::{ CompletedLoad, HostLookup, HostTransferProgress, StartLoad, VllmHostOffloadAdapter, VllmHostRequestState, }; -use crate::engine::scheduler::vllm::policy::{self, AdmissionDecision, PolicySequence}; +use crate::engine::scheduler::vllm::policy::{self, AdmissionDecision}; use crate::engine::scheduler::vllm::request::RequestKvState; use crate::engine::scheduler::{ ActiveHandoffRequests, AdmissionEvent, AdmissionInvariant, AdmissionStage, @@ -936,8 +938,7 @@ impl VllmCore { args.state_cache, args.kv_cache_bytes_per_token, args.aic_nextn.is_some(), - ) - .with_prefix_match_unit(args.prefix_match_unit); + ); if let Some(order) = native_host_offload .as_ref() .and_then(VllmHostOffloadAdapter::store_source_hold) @@ -945,12 +946,42 @@ impl VllmCore { kv_manager.hold_native_store_sources(order); } if !args.kv_cache_groups.is_empty() { - kv_manager.set_grouped_cache(GroupedKvPool::new( - args.kv_cache_groups.clone(), - args.kv_cache_capacity_bytes - .expect("validated grouped cache capacity"), - )); + let capacity_bytes = args + .kv_cache_capacity_bytes + .expect("validated grouped cache capacity"); + if args.enable_prefix_caching { + let unit = args + .prefix_match_unit + .expect("validated grouped prefix_match_unit"); + let version = args + .vllm_version + .as_deref() + .and_then(VllmRuleVersion::parse) + .expect("validated vllm_version"); + let rules = GroupedPrefixRules::new(&args.kv_cache_groups, unit, version) + .expect("validated grouped prefix rules"); + kv_manager.set_grouped_cache(GroupedKvPool::with_prefix_cache( + args.kv_cache_groups.clone(), + capacity_bytes, + args.kv_cache_groups[0].page_size_bytes, + )); + kv_manager.set_grouped_prefix_rules( + rules, + AdmissionLimits { + max_in_flight_tokens: args + .max_concurrent_batches + .saturating_mul(args.max_num_batched_tokens.unwrap_or(usize::MAX)), + max_model_len: args.max_model_len.expect("validated max_model_len"), + }, + ); + } else { + kv_manager.set_grouped_cache(GroupedKvPool::new( + args.kv_cache_groups.clone(), + capacity_bytes, + )); + } } + let kv_manager = kv_manager.with_prefix_match_unit(args.prefix_match_unit); Ok(Self { kv_manager, belady_oracle: None, @@ -1917,8 +1948,11 @@ impl VllmCore { metrics.total_blocks = 0; metrics.kv_cache_used_bytes = Some(pool.used_bytes()); metrics.kv_cache_capacity_bytes = Some(pool.capacity_bytes()); + metrics.kv_cache_cached_bytes = pool.block_pool().map(|_| pool.cached_bytes()); metrics.gpu_cache_usage_perc = pool.used_bytes() as f64 / pool.capacity_bytes() as f64; - metrics.physical_gpu_cache_usage_perc = metrics.gpu_cache_usage_perc; + // Cached blocks that no request holds stay resident. + metrics.physical_gpu_cache_usage_perc = + (pool.used_bytes() + pool.cached_bytes()) as f64 / pool.capacity_bytes() as f64; } metrics } @@ -2223,7 +2257,10 @@ impl VllmCore { .requests .get_mut(&uuid) .expect("waiting request missing from state"); - let mut cost = request.sequence.prefill_cost(&self.kv_manager); + let mut cost = self.kv_manager.lookup_native_prefill_cost( + &request.sequence.sequence, + &mut request.sequence.lease, + ); if request.sequence.sequence.enable_prefix_caching() && let Some(counts) = &mut self.g1_prefix_cache_lookup { @@ -2295,6 +2332,7 @@ impl VllmCore { raw_prefill_cost .as_ref() .expect("fresh admission must retain its G1 lookup"), + self.kv_manager.grouped_prefix_rules().is_some(), ) { decision } else { @@ -2350,6 +2388,7 @@ impl VllmCore { grouped_minimum_forward_bytes = self .kv_manager .grouped() + .filter(|pool| pool.block_pool().is_none()) .zip(self.state.requests.get(&uuid)) .and_then(|(pool, request)| { let known_tokens = request.sequence.len(); @@ -2927,14 +2966,17 @@ impl VllmCore { let request = self .state .requests - .get(&uuid) + .get_mut(&uuid) .unwrap_or_else(|| panic!("schedule_request: {uuid} missing from state.requests")); request.debug_assert_invariants(uuid); let cached_prefix_tokens = if request.num_computed_tokens == 0 { match prefill_cost { Some(cost) => cost.cached_tokens, None => { - let cost = request.sequence.prefill_cost(&self.kv_manager); + let cost = self.kv_manager.lookup_native_prefill_cost( + &request.sequence.sequence, + &mut request.sequence.lease, + ); if request.sequence.sequence.enable_prefix_caching() && let Some(counts) = &mut self.g1_prefix_cache_lookup { @@ -2957,6 +2999,7 @@ impl VllmCore { } else { 0 }; + let request = &self.state.requests[&uuid]; let effective_computed_before = request.num_computed_tokens + cached_prefix_tokens; let prompt_len = request.sequence.num_input_tokens(); let prompt_before = effective_computed_before.min(prompt_len); @@ -2980,12 +3023,23 @@ impl VllmCore { return Ok(ScheduleOutcome::Blocked); } - let desired_tokens = self.align_state_prefill( - effective_computed_before, - prompt_len, - request.sequence.len(), - remaining_known_tokens.min(*token_budget), - ); + let desired_tokens = match self.kv_manager.grouped_prefix_rules() { + // The grouped rule table owns the split of its vLLM version. + Some(rules) => rules.align_split( + effective_computed_before, + prompt_len, + request.sequence.len(), + remaining_known_tokens.min(*token_budget), + request.sequence.lease.grouped_junction(), + self.args.max_num_batched_tokens.unwrap_or(usize::MAX), + ), + None => self.align_state_prefill( + effective_computed_before, + prompt_len, + request.sequence.len(), + remaining_known_tokens.min(*token_budget), + ), + }; // Speculative positions consume scheduler budget but are not committed // sequence/KV state until verification accepts them. let scheduled_drafts = draft_tokens.min(token_budget.saturating_sub(desired_tokens)); @@ -2999,20 +3053,41 @@ impl VllmCore { let desired_computed_after = effective_computed_before + desired_tokens; if let Some(pool) = self.kv_manager.grouped() { - let required = pool - .required_bytes(effective_computed_before, desired_computed_after) - .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(), - match pool.fixed_request_bytes() { - 0 => String::new(), - fixed => format!( - "; the fixed state blocks of one request take {fixed} bytes of this pool" - ), - } - ); + if let Some(blocks) = pool.block_pool() { + let required = pool + .projected_blocks(effective_computed_before, desired_computed_after) + .ok_or_else(|| { + anyhow::anyhow!("grouped KV forward block calculation overflowed") + })?; + anyhow::ensure!( + required <= blocks.capacity_blocks(), + "request {uuid} grouped KV forward needs {required} blocks for {desired_tokens} scheduled tokens after {effective_computed_before} computed tokens, exceeding the entire {}-block cache pool{}; reduce max_num_batched_tokens or increase cache memory", + blocks.capacity_blocks(), + match pool.fixed_request_blocks() { + 0 => String::new(), + fixed => format!( + "; the fixed state blocks of one request take {fixed} blocks of this pool" + ), + } + ); + } else { + let required = pool + .required_bytes(effective_computed_before, desired_computed_after) + .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(), + 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; let host_offload_enabled = self.native_host_offload.is_some(); @@ -3077,6 +3152,10 @@ impl VllmCore { NativeAllocation::CapacityExhausted => {} } + if from_waiting && !policy::waiting_allocation_may_preempt(&self.kv_manager) { + actual_computed_after = effective_computed_before; + break; + } let Some(preempted) = self.policy_preempt(pressure_at_ms) else { actual_computed_after = effective_computed_before; break; diff --git a/crates/core/src/engine/scheduler/vllm/grouped_tests.rs b/crates/core/src/engine/scheduler/vllm/grouped_tests.rs index 54d5a81082..04240cff12 100644 --- a/crates/core/src/engine/scheduler/vllm/grouped_tests.rs +++ b/crates/core/src/engine/scheduler/vllm/grouped_tests.rs @@ -6,20 +6,39 @@ use uuid::Uuid; use crate::engine::common::perf_model::PerfModel; use crate::engine::common::protocols::{DirectRequest, MockEngineArgs}; +use crate::engine::kv_manager::BlockRequestLease; use crate::engine::scheduler::SchedulerCommand; use crate::engine::trace::TraceCollector; use crate::perfmodel::{FpmCacheGroup, FpmCacheKind}; use super::core::VllmCore; +/// Decode queries `(active_tokens, mean_context, token_bound)` and prefill +/// queries `(batch, new_tokens, prefix_tokens)`. #[derive(Default)] -struct CapturedDecodeQueries(Mutex>); +struct CapturedDecodeQueries( + Mutex>, + Mutex>, +); impl crate::engine::TimingModel for CapturedDecodeQueries { fn predict_prefill_ms(&self, _: usize, _: usize, _: usize) -> anyhow::Result { Ok(0.0) } + fn predict_prefill_totals_ms( + &self, + batch_size: usize, + new_tokens: usize, + prefix_tokens: usize, + ) -> anyhow::Result { + self.1 + .lock() + .unwrap() + .push((batch_size, new_tokens, prefix_tokens)); + Ok(0.0) + } + fn predict_decode_ms( &self, _: usize, @@ -86,14 +105,109 @@ fn core(groups: Vec, capacity: u64, chunk: usize) -> VllmCore { } fn submit(core: &mut VllmCore, id: u128, prompt: usize, output: usize) -> Uuid { + submit_tokens(core, id, (0..prompt as u32).collect(), output) +} + +fn submit_tokens(core: &mut VllmCore, id: u128, tokens: Vec, output: usize) -> Uuid { core.receive(DirectRequest { uuid: Some(Uuid::from_u128(id)), - tokens: (0..prompt as u32).collect(), + tokens, max_output_tokens: output, ..Default::default() }) } +const GLM_PAGE: u64 = 13_045_120; +const GLM_BLOCK: u32 = 2176; + +fn glm_state(name: &str, block: u32, cacheable: bool, checkpoint: bool) -> FpmCacheGroup { + FpmCacheGroup { + name: name.into(), + kind: FpmCacheKind::State, + num_layers: 1, + block_size_tokens: block, + page_size_bytes: GLM_PAGE, + sliding_window: None, + blocks_per_request: Some(1), + prefix_cacheable: Some(cacheable), + transfer_bytes_per_block: None, + prefill_checkpoint_blocks: cacheable.then_some(u32::from(checkpoint)), + prefill_checkpoint_alignment: (cacheable && checkpoint).then_some(16), + } +} + +/// The six GLM-5.3-Flash TP 4 groups that the runtime probe recorded on +/// vLLM 0.31.0, at one page size. With `checkpoint`, G2 to G5 carry one +/// FlashKDA checkpoint block at alignment 16. +fn glm_groups(checkpoint: bool) -> Vec { + let mut mla = group("mla", GLM_BLOCK, GLM_PAGE, None); + mla.prefix_cacheable = Some(true); + let mut groups = vec![mla, glm_state("kpool_tail", 4, false, false)]; + for i in 0..4 { + groups.push(glm_state(&format!("kda_{i}"), GLM_BLOCK, true, checkpoint)); + } + groups +} + +/// Three GLM cores: vLLM 0.31.0, and vLLM 0.32.0 without and with checkpoint +/// blocks, which GLM-5.3-Flash takes only under the FlashKDA prefill backend. +const GLM_CORES: [(&str, bool); 3] = [("0.31.0", false), ("0.32.0", false), ("0.32.0", true)]; + +fn prefix_core( + groups: Vec, + capacity: u64, + chunk: usize, + unit: usize, + version: &str, +) -> VllmCore { + VllmCore::new( + MockEngineArgs::builder() + .block_size(4) + .num_gpu_blocks(1) + .max_num_batched_tokens(Some(chunk)) + .max_num_seqs(Some(3)) + .max_model_len(Some(131_072)) + .enable_prefix_caching(true) + .prefix_match_unit(Some(unit)) + .vllm_version(Some(version.to_string())) + .kv_cache_groups(groups) + .kv_cache_capacity_bytes(Some(capacity)) + .speedup_ratio(0.0) + .build() + .unwrap(), + ) +} + +/// Step until the core drains. Returns `(uuid, reused_input_tokens)` of every +/// admission in order. +fn drain( + core: &mut VllmCore, + trace: &mut TraceCollector, + from_step: &mut usize, +) -> Vec<(Uuid, usize)> { + let mut admissions = Vec::new(); + for _ in 0..200 { + let pass = core.execute_pass(trace, *from_step as f64); + *from_step += 1; + admissions.extend( + pass.admissions + .iter() + .map(|admission| (admission.uuid, admission.reused_input_tokens)), + ); + for output in pass.output_signals { + assert!(!output.rejected, "{} rejected", output.uuid); + } + if core.is_drained() { + return admissions; + } + } + panic!("the core did not drain"); +} + +fn block_pool(core: &VllmCore) -> &crate::engine::kv_manager::GroupedBlockPool { + core.kv_manager.grouped().unwrap().block_pool().unwrap() +} + #[test] fn grouped_window_replays_long_context_with_bounded_storage_and_logical_fpm_lengths() { let mut core = core(vec![group("window", 4, 8, Some(4))], 24, 8); @@ -332,6 +446,7 @@ fn grouped_state_blocks_stay_constant_through_prefill_and_decode() { } // Every pass adds one 8-byte page; the two 10-byte state blocks stay. assert_eq!(bytes, vec![28, 36, 44, 52, 60]); + assert!(core.kv_manager.grouped().unwrap().block_pool().is_none()); core.apply_command(SchedulerCommand::CancelRequest { request_id: id }) .unwrap(); assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0)); @@ -475,3 +590,600 @@ fn grouped_state_request_waits_when_only_attention_pages_fit() { assert!(pass.output_signals.iter().all(|output| !output.rejected)); } } + +#[test] +fn grouped_prefix_caching_off_keeps_byte_mode_accounting() { + let mut core = core(glm_groups(false), 64 * GLM_PAGE, 8192); + assert!(core.kv_manager.grouped().unwrap().block_pool().is_none()); + let mut trace = TraceCollector::default(); + let mut step = 0; + let first = submit(&mut core, 301, 9000, 4); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(first, 0)]); + // A resend of the same prompt replays cold: nothing is cached in byte mode. + let second = submit(&mut core, 302, 9000, 4); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(second, 0)]); + let metrics = core.mocker_metrics(); + assert_eq!(metrics.kv_cache_used_bytes, Some(0)); + assert_eq!(metrics.kv_cache_cached_bytes, None); + assert_eq!(core.kv_manager.grouped().unwrap().cached_bytes(), 0); + assert_eq!(core.kv_manager.num_active_blocks(), 0); +} + +#[test] +fn grouped_glm_shaped_multi_turn_replay_reports_hit_tokens_and_block_counts() { + for (version, checkpoint) in GLM_CORES { + let mut core = prefix_core(glm_groups(checkpoint), 64 * GLM_PAGE, 8192, 2176, version); + let mut trace = TraceCollector::default(); + let mut step = 0; + let turn_one: Vec = (0..9000).collect(); + let first = submit_tokens(&mut core, 401, turn_one.clone(), 4); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(first, 0)]); + assert_eq!(block_pool(&core).active_blocks(), 0); + + // Turn two: the first 9,000 tokens again plus 4,000 new ones. The hit + // is the state at request one's replay boundary (9000 - 1) / 2176 * + // 2176 = 8704, where G0 holds four cached blocks too. + let mut turn_two = turn_one.clone(); + turn_two.extend(100_000..104_000); + let second = submit_tokens(&mut core, 402, turn_two.clone(), 4); + block_pool(&core).reset_peak_active_blocks(); + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert_eq!(pass.admissions.len(), 1, "{version} {checkpoint}"); + assert_eq!(pass.admissions[0].uuid, second); + assert_eq!( + pass.admissions[0].reused_input_tokens, 8704, + "{version} {checkpoint}" + ); + // During the first step: G0 holds the four pinned blocks and the fresh + // ones of the chunk, the ring one, each KDA group the cached and the + // running block, plus the checkpoint block on the third core. + let (peak, after) = if checkpoint { (19, 15) } else { (14, 10) }; + assert_eq!( + block_pool(&core).peak_active_blocks(), + peak, + "{version} {checkpoint}" + ); + // After the step the four cached state blocks are released. + assert_eq!( + block_pool(&core).active_blocks(), + after, + "{version} {checkpoint}" + ); + assert_eq!(block_pool(&core).inactive_cached_blocks(), 4); + assert_eq!( + pass.mocker_metrics.kv_cache_used_bytes, + Some(after as u64 * GLM_PAGE) + ); + assert_eq!( + pass.mocker_metrics.kv_cache_cached_bytes, + Some(4 * GLM_PAGE) + ); + assert!(drain(&mut core, &mut trace, &mut step).is_empty()); + + // A request that shares only the first 2,176 tokens finds a G0 block + // but no state key at 2,176 (retention 0), so it hits nothing. + let mut other: Vec = (0..2176).collect(); + other.extend(200_000..202_824); + let third = submit_tokens(&mut core, 403, other, 4); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(third, 0)]); + + // The whole second turn plus 100 tokens hits request two's replay + // boundary (13000 - 1) / 2176 * 2176 = 10880. + let mut turn_three = turn_two; + turn_three.extend(300_000..300_100); + let fourth = submit_tokens(&mut core, 404, turn_three, 4); + assert_eq!( + drain(&mut core, &mut trace, &mut step), + vec![(fourth, 10880)] + ); + assert_eq!(block_pool(&core).active_blocks(), 0); + assert_eq!(core.mocker_metrics().kv_cache_used_bytes, Some(0)); + } +} + +#[test] +fn grouped_prefix_hit_lowers_prefill_tokens_and_raises_prefix_tokens() { + for (version, checkpoint) in GLM_CORES { + let mut core = prefix_core(glm_groups(checkpoint), 64 * GLM_PAGE, 8192, 2176, version); + let captured = Arc::new(CapturedDecodeQueries::default()); + core.args.perf_model = Arc::new(PerfModel::External { + timing: captured.clone(), + }); + let mut trace = TraceCollector::default(); + let mut step = 0; + let turn_one: Vec = (0..9000).collect(); + submit_tokens(&mut core, 501, turn_one.clone(), 1); + drain(&mut core, &mut trace, &mut step); + // The first prefill stops at the block boundary 6528 inside the 8192 + // budget, then at the last cacheable stop 8704, then at the prompt. + // In checkpoint mode the final chunk runs from 6528 to the prompt. + let first_turn: Vec<(usize, usize, usize)> = if checkpoint { + vec![(1, 6528, 0), (1, 2472, 6528)] + } else { + vec![(1, 6528, 0), (1, 2176, 6528), (1, 296, 8704)] + }; + assert_eq!( + *captured.1.lock().unwrap(), + first_turn, + "{version} {checkpoint}" + ); + captured.1.lock().unwrap().clear(); + + let mut turn_two = turn_one; + turn_two.extend(100_000..104_000); + submit_tokens(&mut core, 502, turn_two, 1); + drain(&mut core, &mut trace, &mut step); + // From the hit at 8704 the split ends the first chunk at the last + // cacheable stop 10880 and the second at the prompt. The checkpoint + // branch runs to the prompt end in one chunk. + let second_turn: Vec<(usize, usize, usize)> = if checkpoint { + vec![(1, 4296, 8704)] + } else { + vec![(1, 2176, 8704), (1, 2120, 10880)] + }; + assert_eq!( + *captured.1.lock().unwrap(), + second_turn, + "{version} {checkpoint}" + ); + } +} + +#[test] +fn grouped_prefix_admission_tests_the_whole_prompt() { + for (version, checkpoint) in GLM_CORES { + // A prompt of 40,000 tokens needs 19 G0 blocks, 1 ring block and one + // running block per KDA group, plus one checkpoint block each on the + // third core: 24 or 28, above the 20-block pool. + let mut core = prefix_core(glm_groups(checkpoint), 20 * GLM_PAGE, 8192, 2176, version); + let id = submit(&mut core, 601, 40_000, 2); + let request = &core.state.requests[&id].sequence; + let cost = core + .kv_manager + .get_native_prefill_cost(&request.sequence, &request.lease); + assert_eq!( + cost.new_blocks, + if checkpoint { 28 } else { 24 }, + "{version}" + ); + assert_eq!(cost.cached_tokens, 0); + let mut trace = TraceCollector::default(); + let pass = core.execute_pass(&mut trace, 0.0); + assert!(pass.admissions.is_empty(), "{version} {checkpoint}"); + assert!( + pass.output_signals + .iter() + .any(|output| output.uuid == id && output.rejected) + ); + assert_eq!(block_pool(&core).active_blocks(), 0); + assert!(core.is_drained()); + + // Under contention the whole prompt decides, as in vLLM's default + // admission (mgr:519-536). A running request of 13,000 tokens holds + // 3 G0 pages, the ring and four state blocks after its first chunk, + // which the split ends at 6,528, so 22 of 30 blocks are free. The + // 40,000-token prompt needs 24 and waits without preempting anything. + let mut core = prefix_core(glm_groups(checkpoint), 30 * GLM_PAGE, 8192, 2176, version); + let running = submit(&mut core, 603, 13_000, 2); + let waiting = submit_tokens(&mut core, 604, (100_000..140_000).collect(), 2); + let pass = core.execute_pass(&mut trace, 0.0); + assert_eq!(pass.admissions.len(), 1, "{version} {checkpoint}"); + assert_eq!(pass.admissions[0].uuid, running); + assert_eq!(block_pool(&core).active_blocks(), 8); + // The second pass has budget left after the running request's + // chunk, and the waiting request still waits for the whole prompt. + let pass = core.execute_pass(&mut trace, 1.0); + assert!(pass.admissions.is_empty(), "{version} {checkpoint}"); + assert_eq!(pass.mocker_metrics.vllm_preemptions_total, 0); + assert_eq!(core.state.requests[&waiting].num_computed_tokens, 0); + } + // Prefix caching off keeps the first-chunk test of byte mode. The first + // pass spends its budget on the running request. In the second pass the + // running request takes its last 4,808 tokens (6 G0 pages, the ring and + // the states: 11 pages) and the waiting request's first chunk of the + // 3,384 tokens left (2 pages, the ring and the states: 7 pages) fits. + let mut off = core(glm_groups(false), 30 * GLM_PAGE, 8192); + let running = submit(&mut off, 605, 13_000, 2); + let waiting = submit_tokens(&mut off, 606, (100_000..140_000).collect(), 2); + let mut trace = TraceCollector::default(); + let pass = off.execute_pass(&mut trace, 0.0); + assert_eq!(pass.admissions.len(), 1); + assert_eq!(pass.admissions[0].uuid, running); + let pass = off.execute_pass(&mut trace, 1.0); + assert_eq!(pass.admissions.len(), 1); + assert_eq!(pass.admissions[0].uuid, waiting); + assert_eq!(pass.mocker_metrics.kv_cache_used_bytes, Some(18 * GLM_PAGE)); + // Byte mode also refuses a final forward that can never fit, as + // grouped_state_blocks_exceeding_the_pool_reject_at_admission shows: the + // 40,000-token prompt's last forward needs 24 pages of a 20-page pool. + let mut off = core(glm_groups(false), 20 * GLM_PAGE, 8192); + let id = submit(&mut off, 602, 40_000, 2); + let pass = off.execute_pass(&mut trace, 0.0); + assert!(pass.admissions.is_empty()); + assert!( + pass.output_signals + .iter() + .any(|output| output.uuid == id && output.rejected) + ); +} + +#[test] +fn grouped_prefix_waiting_request_keeps_fifo() { + // Each 20,000-token prompt needs 10 G0 blocks, the ring block and four + // running blocks: 15 of 20. The second request waits for the first one + // to finish and preempts nothing. + let mut core = prefix_core(glm_groups(false), 20 * GLM_PAGE, 8192, 2176, "0.31.0"); + let first = submit_tokens(&mut core, 701, (0..20_000).collect(), 2); + let second = submit_tokens(&mut core, 702, (100_000..120_000).collect(), 2); + let mut trace = TraceCollector::default(); + let mut step = 0; + let initial = core.execute_pass(&mut trace, step as f64); + step += 1; + assert_eq!(initial.admissions.len(), 1); + assert_eq!(initial.admissions[0].uuid, first); + let mut saw_first_complete = false; + let mut admitted_second_at = None; + for _ in 0..40 { + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert_eq!(pass.mocker_metrics.vllm_preemptions_total, 0); + if pass + .admissions + .iter() + .any(|admission| admission.uuid == second) + { + admitted_second_at = Some(step); + assert!( + saw_first_complete, + "the second request waited for the first" + ); + } + for output in pass.output_signals { + assert!(!output.rejected); + saw_first_complete |= output.uuid == first && output.completed; + } + if core.is_drained() { + break; + } + } + assert!(admitted_second_at.is_some()); + assert!(core.is_drained()); + assert_eq!(block_pool(&core).active_blocks(), 0); +} + +#[test] +fn grouped_prefix_waiting_hit_waits_when_its_first_chunk_does_not_fit() { + // One sliding-window group of 4-token blocks with an 8-token window, a + // 9-block pool and a step budget of 8 tokens. + for version in ["0.31.0", "0.32.0"] { + let mut window = group("window", 4, 8, Some(8)); + window.prefix_cacheable = Some(true); + let mut core = prefix_core(vec![window], 9 * 8, 8, 4, version); + let mut trace = TraceCollector::default(); + let mut step = 0; + // A finished request of 100 tokens keeps the two blocks before its + // replay boundary 96 cached. + submit_tokens(&mut core, 1101, (0..100).collect(), 1); + drain(&mut core, &mut trace, &mut step); + assert_eq!(block_pool(&core).inactive_cached_blocks(), 2, "{version}"); + submit_tokens(&mut core, 1102, (1000..1008).collect(), 30); + submit_tokens(&mut core, 1103, (2000..2008).collect(), 30); + for _ in 0..3 { + core.execute_pass(&mut trace, step as f64); + step += 1; + } + assert_eq!(block_pool(&core).active_blocks(), 5, "{version}"); + // A request of 104 tokens hits those blocks at 96. The whole-prompt + // test counts only the two inactive hit blocks, because the window + // caps it (mgr:519-536). After their decode steps the running + // requests hold 6 of the 9 blocks, which leaves 3, but the first + // chunk of 6 tokens needs the two hit blocks and two fresh blocks + // (mgr:558-576). vLLM then stops its waiting loop (v0.31 + // sched:1223-1230, v0.32 sched:1275-1282), so the request waits and + // preempts nothing. + let waiting = submit_tokens(&mut core, 1104, (0..104).collect(), 1); + let request = &core.state.requests[&waiting].sequence; + let cost = core + .kv_manager + .get_native_prefill_cost(&request.sequence, &request.lease); + assert_eq!((cost.new_blocks, cost.cached_tokens), (2, 96), "{version}"); + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert_eq!(pass.mocker_metrics.vllm_preemptions_total, 0, "{version}"); + assert!(pass.admissions.is_empty(), "{version}"); + assert_eq!( + core.state.requests[&waiting].num_computed_tokens, 0, + "{version}" + ); + assert_eq!(core.state.running.len(), 2, "{version}"); + // The waiting request starts with its hit once the decode steps of the + // running requests leave room. + let admissions = drain(&mut core, &mut trace, &mut step); + assert_eq!(admissions.first(), Some(&(waiting, 96)), "{version}"); + } +} + +#[test] +fn grouped_prefix_waiting_hit_keeps_its_lru_place_when_its_first_chunk_does_not_fit() { + // vLLM tests the blocks of the first chunk before it touches the hit + // (mgr:558-589). So a waiting request that fails that test leaves its hit + // blocks where they are in the LRU, ahead of a block cached after them. + // Pinning and then releasing them would move them behind that block. + let cached = |core: &VllmCore, tokens: &[u32], end: usize| { + let mut lease = BlockRequestLease::new(Uuid::nil(), Vec::new()); + lease.configure_prefix_hashes(tokens, 4); + block_pool(core) + .key_hit(0, lease.prefix_hash_at(end).unwrap()) + .is_some() + }; + for version in ["0.31.0", "0.32.0"] { + let mut window = group("window", 4, 8, Some(8)); + window.prefix_cacheable = Some(true); + let mut core = prefix_core(vec![window], 9 * 8, 8, 4, version); + let mut trace = TraceCollector::default(); + let mut step = 0; + // The LRU holds the blocks that end at 96 and 92 of a finished request + // of 100 tokens, then the block that ends at 4 of a finished request + // of 8 tokens. + let first: Vec = (0..100).collect(); + let second: Vec = (3000..3008).collect(); + submit_tokens(&mut core, 1401, first.clone(), 1); + drain(&mut core, &mut trace, &mut step); + submit_tokens(&mut core, 1402, second.clone(), 1); + drain(&mut core, &mut trace, &mut step); + assert_eq!(block_pool(&core).inactive_cached_blocks(), 3, "{version}"); + // As in grouped_prefix_waiting_hit_waits_when_its_first_chunk_does_not_fit, + // two running requests hold 6 of the 9 blocks, and the request of 104 + // tokens fails the first-chunk test of its hit at 96. + submit_tokens(&mut core, 1403, (1000..1008).collect(), 30); + submit_tokens(&mut core, 1404, (2000..2008).collect(), 30); + for _ in 0..3 { + core.execute_pass(&mut trace, step as f64); + step += 1; + } + let waiting = submit_tokens(&mut core, 1405, (0..104).collect(), 1); + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert!(pass.admissions.is_empty(), "{version}"); + assert_eq!(block_pool(&core).active_blocks(), 6, "{version}"); + assert_eq!(block_pool(&core).inactive_cached_blocks(), 3, "{version}"); + core.apply_command(SchedulerCommand::CancelRequest { + request_id: waiting, + }) + .unwrap(); + // No block is free. A request of 4 new tokens takes one fresh block, + // which evicts the oldest cached block: the one that ends at 96. + let fresh = submit_tokens(&mut core, 1406, (4000..4004).collect(), 1); + let pass = core.execute_pass(&mut trace, step as f64); + assert_eq!(pass.admissions.len(), 1, "{version}"); + assert_eq!(pass.admissions[0].uuid, fresh, "{version}"); + assert!(!cached(&core, &first, 96), "{version}"); + assert!(cached(&core, &first, 92), "{version}"); + assert!(cached(&core, &second, 4), "{version}"); + } +} + +#[test] +fn grouped_prefix_without_chunked_prefill_budgets_the_uncached_tokens() { + // With chunked prefill off, vLLM compares the tokens after the prefix hit + // with the step budget (v0.31 sched:1084, :1115-1121, v0.32 sched:1129, + // :1160-1167), as the linear pool does. Two running requests decode one + // token each and leave 14 of the 16 tokens. A prompt of 15 tokens with 12 + // cached tokens needs 3 of them. + for version in ["0.31.0", "0.32.0"] { + let args = || { + MockEngineArgs::builder() + .block_size(4) + .max_num_batched_tokens(Some(16)) + .max_num_seqs(Some(4)) + .max_model_len(Some(4096)) + .enable_prefix_caching(true) + .enable_chunked_prefill(false) + .speedup_ratio(0.0) + }; + let linear = VllmCore::new(args().num_gpu_blocks(64).build().unwrap()); + let mut full = group("full", 4, 1, None); + full.prefix_cacheable = Some(true); + let grouped = VllmCore::new( + args() + .num_gpu_blocks(1) + .prefix_match_unit(Some(4)) + .vllm_version(Some(version.to_string())) + .kv_cache_groups(vec![full]) + .kv_cache_capacity_bytes(Some(64)) + .build() + .unwrap(), + ); + for (pool, mut core) in [("linear", linear), ("grouped", grouped)] { + let mut trace = TraceCollector::default(); + let mut step = 0; + submit_tokens(&mut core, 1301, (0..15).collect(), 1); + drain(&mut core, &mut trace, &mut step); + submit_tokens(&mut core, 1302, (1000..1004).collect(), 10); + submit_tokens(&mut core, 1303, (2000..2004).collect(), 10); + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert_eq!(pass.admissions.len(), 2, "{version} {pool}"); + let waiting = submit_tokens(&mut core, 1304, (0..15).collect(), 1); + let pass = core.execute_pass(&mut trace, step as f64); + let admitted: Vec<(Uuid, usize)> = pass + .admissions + .iter() + .map(|admission| (admission.uuid, admission.reused_input_tokens)) + .collect(); + assert_eq!(admitted, vec![(waiting, 12)], "{version} {pool}"); + } + } +} + +/// Full-attention and state groups of 16-token blocks, a 4-token hash unit +/// and a pool of five blocks: the blocks of one 34-token request, as vLLM's +/// startup check counts them (3 + 2). +fn fine_grained_core(version: &str) -> VllmCore { + let mut full = group("full", 16, 8, None); + full.prefix_cacheable = Some(true); + let mut state = state("state", 8, 1); + state.block_size_tokens = 16; + state.prefix_cacheable = Some(true); + state.prefill_checkpoint_blocks = (version == "0.32.0").then_some(0); + VllmCore::new( + MockEngineArgs::builder() + .block_size(4) + .num_gpu_blocks(1) + .max_num_batched_tokens(Some(64)) + .max_num_seqs(Some(4)) + .max_model_len(Some(34)) + .enable_prefix_caching(true) + .prefix_match_unit(Some(4)) + .vllm_version(Some(version.to_string())) + .kv_cache_groups(vec![full, state]) + .kv_cache_capacity_bytes(Some(5 * 8)) + .speedup_ratio(0.0) + .build() + .unwrap(), + ) +} + +#[test] +fn grouped_prefix_hit_that_needs_more_blocks_than_the_pool_waits() { + // A finished 13-token request leaves its full-attention block and its + // state keyed at the hash boundary 12. A 33-token prompt that repeats it + // hits 12 tokens inside a block. That hit needs 7 blocks: 2 new, the + // inactive hit block and a copy for full attention (single:216-268), and + // the running block, a copy and the inactive hit block for the state + // (single:1716-1816). Cold the prompt needs 4 of the 5 blocks. vLLM keeps + // such a request waiting and searches its hit again at the next step + // (v0.31 sched:930-937, :1223-1230, v0.32 sched:960-967, :1275-1282). + for version in ["0.31.0", "0.32.0"] { + // With nothing running, nothing changes the cache, so the request + // waits in every pass. The replay reports that as a livelock + // (runtime_contract.rs grouped_prefix_hit_above_the_pool_is_a_livelock). + let mut core = fine_grained_core(version); + let mut trace = TraceCollector::default(); + let mut step = 0; + submit_tokens(&mut core, 1501, (0..13).collect(), 1); + drain(&mut core, &mut trace, &mut step); + assert_eq!(block_pool(&core).inactive_cached_blocks(), 2, "{version}"); + let waiting = submit_tokens(&mut core, 1502, (0..33).collect(), 1); + let request = &core.state.requests[&waiting].sequence; + let cost = core + .kv_manager + .get_native_prefill_cost(&request.sequence, &request.lease); + assert_eq!((cost.new_blocks, cost.cached_tokens), (7, 12), "{version}"); + for _ in 0..3 { + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert!(pass.admissions.is_empty(), "{version}"); + assert!(pass.output_signals.is_empty(), "{version}"); + assert_eq!(pass.mocker_metrics.vllm_preemptions_total, 0, "{version}"); + } + assert!(core.state.requests.contains_key(&waiting), "{version}"); + + // A running request can evict the hit blocks. Its decode crosses a + // block boundary at 16 tokens and takes two blocks while one is free, + // which evicts the cached full-attention block of the hit. The + // waiting request then hits nothing and starts cold once the running + // request finishes. + let mut core = fine_grained_core(version); + let mut trace = TraceCollector::default(); + let mut step = 0; + submit_tokens(&mut core, 1503, (0..13).collect(), 1); + drain(&mut core, &mut trace, &mut step); + let running = submit_tokens(&mut core, 1504, (1000..1004).collect(), 20); + let waiting = submit_tokens(&mut core, 1505, (0..33).collect(), 1); + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert_eq!(pass.admissions.len(), 1, "{version}"); + assert_eq!(pass.admissions[0].uuid, running, "{version}"); + assert!(pass.output_signals.iter().all(|output| !output.rejected)); + let mut finished = Vec::new(); + let mut admissions = Vec::new(); + for _ in 0..40 { + let pass = core.execute_pass(&mut trace, step as f64); + step += 1; + assert_eq!(pass.mocker_metrics.vllm_preemptions_total, 0, "{version}"); + admissions.extend( + pass.admissions + .iter() + .map(|admission| (admission.uuid, admission.reused_input_tokens)), + ); + for output in pass.output_signals { + assert!(!output.rejected, "{version}"); + if output.completed { + finished.push(output.uuid); + } + } + if core.is_drained() { + break; + } + } + assert_eq!(admissions, vec![(waiting, 0)], "{version}"); + assert_eq!(finished, vec![running, waiting], "{version}"); + } +} + +#[test] +fn grouped_prefix_physical_usage_counts_cached_blocks() { + // Metrics::physical_cache_usage includes the inactive reusable blocks and + // Metrics::inactive_blocks counts them (protocol.rs). Active usage counts + // only the blocks that requests hold. + for (version, checkpoint) in GLM_CORES { + let mut core = prefix_core(glm_groups(checkpoint), 64 * GLM_PAGE, 8192, 2176, version); + let mut trace = TraceCollector::default(); + let mut step = 0; + submit_tokens(&mut core, 1201, (0..9000).collect(), 4); + drain(&mut core, &mut trace, &mut step); + // G0 keeps its four full blocks, and each KDA group the state at the + // replay boundary 8704. + assert_eq!(block_pool(&core).inactive_cached_blocks(), 8, "{version}"); + let metrics = core.mocker_metrics(); + assert_eq!(metrics.kv_cache_used_bytes, Some(0), "{version}"); + assert_eq!( + metrics.kv_cache_cached_bytes, + Some(8 * GLM_PAGE), + "{version}" + ); + assert_eq!(metrics.gpu_cache_usage_perc, 0.0, "{version}"); + assert_eq!( + metrics.physical_gpu_cache_usage_perc, + 8.0 / 64.0, + "{version}" + ); + assert_eq!(metrics.active_decode_blocks, 0, "{version}"); + assert_eq!(metrics.inactive_decode_blocks, 8, "{version}"); + } +} + +#[test] +fn grouped_prefix_eviction_frees_cached_blocks_oldest_first() { + // One full-attention group of 4-token blocks in a pool of 16 blocks. + let mut full = group("full", 4, 8, None); + full.prefix_cacheable = Some(true); + let mut core = prefix_core(vec![full], 16 * 8, 64, 4, "0.31.0"); + let mut trace = TraceCollector::default(); + let mut step = 0; + // Three finished requests of 16 tokens leave 12 cached blocks. + for (id, base) in [(801, 0), (802, 100), (803, 200)] { + submit_tokens(&mut core, id, (base..base + 16).collect(), 1); + drain(&mut core, &mut trace, &mut step); + } + assert_eq!(block_pool(&core).inactive_cached_blocks(), 12); + // A fourth request of 40 tokens needs 10 fresh blocks: the 4 free ones + // and the 6 evicted oldest first, in release order (tail first): all of + // request one and the two last blocks of request two. + let fourth = submit_tokens(&mut core, 804, (300..340).collect(), 1); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(fourth, 0)]); + assert_eq!(block_pool(&core).inactive_cached_blocks(), 16); + // A fifth request equal to the third one still hits, capped below the + // prompt: (16 - 1) / 4 * 4 = 12. + let fifth = submit_tokens(&mut core, 805, (200..216).collect(), 1); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(fifth, 12)]); + // The fifth request's one fresh block evicted the older of request + // two's two surviving blocks, so request two hits its first block only. + let sixth = submit_tokens(&mut core, 806, (100..116).collect(), 1); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(sixth, 4)]); + // Request one is gone. + let seventh = submit_tokens(&mut core, 807, (0..16).collect(), 1); + assert_eq!(drain(&mut core, &mut trace, &mut step), vec![(seventh, 0)]); +} diff --git a/crates/core/src/engine/scheduler/vllm/policy.rs b/crates/core/src/engine/scheduler/vllm/policy.rs index ddd17cc0f0..2a649e6bd9 100644 --- a/crates/core/src/engine/scheduler/vllm/policy.rs +++ b/crates/core/src/engine/scheduler/vllm/policy.rs @@ -2,7 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright contributors to the vLLM project. // Modified Rust simulation of vLLM b3563fc65ae0f5359802593d78e7ea097e1fed31 -// scheduling contracts; immutable sources and license: THIRD_PARTY_NOTICES.md. +// scheduling contracts and of the v0.31.0 +// (db9527a46873454610df6dbedf79a36d6bf1a7f6) and v0.32.0 +// (5df63a7203cdb9328d824a27b682e77c7b9928be) grouped KV rules; +// immutable sources and license: THIRD_PARTY_NOTICES.md. //! Engine-specific policy for the shared vLLM/TRT-LLM scheduler core. //! @@ -84,17 +87,54 @@ pub(super) enum AdmissionDecision { Reject, } -/// Grouped cold-cache admission uses the pending forward's physical bytes. -/// Logical prompt length remains independent of bounded cache residency. +/// Grouped admission. In byte mode (prefix caching off) the pending +/// forward's physical bytes decide, and logical prompt length remains +/// independent of bounded cache residency. With `whole_prompt` (prefix +/// caching on) the blocks of the whole prompt decide, as vLLM's +/// `allocate_slots` does by default (`mgr:519-536`). Only a prompt whose +/// blocks do not fit the empty pool is rejected. pub(super) fn decide_grouped_waiting_admission( sequence: &S, kv_manager: &G1Manager, token_budget: usize, enable_chunked_prefill: bool, prefill_cost: &PrefillCost, + whole_prompt: bool, ) -> Option { let pool = kv_manager.grouped()?; let required_tokens = sequence.len(); + if whole_prompt { + let blocks = pool.block_pool().expect("block mode"); + let capacity = blocks.capacity_blocks(); + // A prompt whose blocks never fit the pool is never admitted. The test + // leaves the hit out: the copy and inactive blocks of a hit depend on + // the cache, which changes while the request waits. + if kv_manager.grouped_cold_admission_blocks(required_tokens) > capacity { + return Some(AdmissionDecision::Reject); + } + // Without chunked prefill, vLLM compares the tokens after the prefix + // hit with the budget (v0.31 `sched:1084`, `:1115-1121`, v0.32 + // `sched:1129`, `:1160-1167`). + if token_budget == 0 || (!enable_chunked_prefill && prefill_cost.new_tokens > token_budget) + { + return Some(AdmissionDecision::Wait); + } + // Inactive cached blocks are evictable, so only active blocks count + // against the request. Contention keeps FIFO and never preempts. A + // hit that needs more blocks than the pool waits too. vLLM searches + // the hit again at the next step (v0.31 `sched:930-937`, + // `:1223-1230`, v0.32 `sched:960-967`, `:1275-1282`), and running + // requests can evict the hit blocks in between. + return Some( + if prefill_cost.new_blocks > capacity - blocks.active_blocks() { + AdmissionDecision::Wait + } else { + AdmissionDecision::Admit { + prefill_cost: prefill_cost.clone(), + } + }, + ); + } // A known context whose smallest possible final forward cannot fit will // never be admitted, regardless of chunking or competing requests. let minimum = pool.required_bytes(required_tokens.saturating_sub(1), required_tokens); @@ -120,6 +160,16 @@ pub(super) fn decide_grouped_waiting_admission( } } +/// Whether a waiting request whose allocation fails may preempt running +/// work. With grouped prefix caching it may not: admission tests the blocks +/// of the whole prompt, and the allocation then tests the blocks of the +/// chunk, as vLLM's `allocate_slots` does (`mgr:519-576`). A waiting request +/// that fails either test stays queued, because vLLM stops its waiting loop +/// there (v0.31 `sched:1223-1230`, v0.32 `sched:1275-1282`). +pub(super) fn waiting_allocation_may_preempt(kv_manager: &G1Manager) -> bool { + kv_manager.grouped_prefix_rules().is_none() +} + #[derive(Debug, Clone, Copy)] pub(super) struct WaitingAdmissionConfig { pub(super) policy: SchedulingPolicy, diff --git a/crates/core/src/perfmodel/fpm/resources.rs b/crates/core/src/perfmodel/fpm/resources.rs index a9942443a0..7a7e97f98d 100644 --- a/crates/core/src/perfmodel/fpm/resources.rs +++ b/crates/core/src/perfmodel/fpm/resources.rs @@ -40,8 +40,10 @@ pub enum FpmCacheKind { /// can be more than the bytes of the group's own layers. Its /// `block_size_tokens` is the block size of vLLM's state spec, as the runtime /// probe records it: the Mamba block (the scheduler block in `align` mode, -/// `max_model_len` in mode `none`) or the ring block. Sizing and replay do not -/// read it yet. +/// `max_model_len` in mode `none`) or the ring block. Sizing and replay with +/// prefix caching off do not read it. Grouped prefix caching reads it as the +/// Mamba block of `align` mode, and every group's block enters the scheduler +/// block, the least common multiple of the group blocks. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct FpmCacheGroup { @@ -65,9 +67,9 @@ pub struct FpmCacheGroup { #[serde(default, skip_serializing_if = "Option::is_none")] pub blocks_per_request: Option, /// Whether vLLM's manager for this group takes part in prefix caching. A - /// ring spec does not. Replay with prefix caching on needs it on every - /// group, and does not read it yet. `None` in a profile imported before - /// this field existed. + /// ring spec does not. A rank with `kv_cache_groups` and + /// `enable_prefix_caching` refuses a group that leaves it unset. `None` + /// in a profile imported before this field existed. #[serde(default, skip_serializing_if = "Option::is_none")] pub prefix_cacheable: Option, /// The payload that a P/D transfer moves for one block of this group, @@ -236,8 +238,10 @@ pub enum FpmMambaCacheMode { /// non-KV memory limit or a capacity that can be scaled to another setting. /// /// The three optional prefix-caching fields record the settings that the probe -/// ran with. They are evidence only: selection, sizing and replay do not read -/// them, and the simulated deployment decides prefix caching. +/// ran with. Selection and sizing do not read them. Grouped prefix caching +/// needs a probe that ran with prefix caching, and replay then hashes prompts +/// in blocks of `hash_block_size_tokens`. For a linear profile, the simulated +/// deployment decides prefix caching. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct FpmRuntimeMemoryConfig { diff --git a/crates/core/src/replay/components/engine.rs b/crates/core/src/replay/components/engine.rs index 450ab8b368..44a0203f86 100644 --- a/crates/core/src/replay/components/engine.rs +++ b/crates/core/src/replay/components/engine.rs @@ -127,6 +127,7 @@ impl RankMetricsState { total_blocks: self.latest.total_blocks, kv_cache_used_bytes: self.latest.kv_cache_used_bytes, kv_cache_capacity_bytes: self.latest.kv_cache_capacity_bytes, + kv_cache_cached_bytes: self.latest.kv_cache_cached_bytes, active_cache_usage: self.latest.cache_usage, physical_cache_usage: self.latest.physical_cache_usage, running_requests: self.latest.running_requests, diff --git a/crates/core/src/replay/telemetry.rs b/crates/core/src/replay/telemetry.rs index 4a00f075a5..e7952730e6 100644 --- a/crates/core/src/replay/telemetry.rs +++ b/crates/core/src/replay/telemetry.rs @@ -52,8 +52,13 @@ pub struct ReplaySchedulerMetricsSnapshot { pub kv_cache_used_bytes: Option, #[serde(skip_serializing_if = "Option::is_none")] pub kv_cache_capacity_bytes: Option, - /// Legacy/backend-native block utilization, or physical byte utilization - /// for grouped caches (which have no scalar total block capacity). + /// Cached grouped-cache bytes that no request holds; present only with + /// grouped prefix caching. + #[serde(skip_serializing_if = "Option::is_none")] + pub kv_cache_cached_bytes: Option, + /// Legacy/backend-native block utilization, or active byte utilization for + /// grouped caches. With grouped prefix caching, cached bytes that no request + /// holds count in `physical_cache_usage` only. pub active_cache_usage: f64, /// Physical resident utilization; equal to active utilization for SGLang. pub physical_cache_usage: f64, @@ -916,6 +921,7 @@ mod tests { total_blocks: 0, kv_cache_used_bytes: Some(1024), kv_cache_capacity_bytes: Some(2048), + kv_cache_cached_bytes: None, active_cache_usage: 0.5, physical_cache_usage: 0.5, running_requests: 1, diff --git a/crates/core/tests/runtime_contract.rs b/crates/core/tests/runtime_contract.rs index 7af590d281..b83eb4b3dc 100644 --- a/crates/core/tests/runtime_contract.rs +++ b/crates/core/tests/runtime_contract.rs @@ -177,6 +177,116 @@ fn grouped_cache_replay_emits_physical_bytes_without_scalar_block_capacity() { ); } +#[test] +fn grouped_prefix_cache_replay_reports_reuse() { + const GLM_PAGE: u64 = 13_045_120; + for (version, checkpoint) in [("0.31.0", None), ("0.32.0", Some(0))] { + let samples = Arc::new(Mutex::new(Vec::new())); + let turn_one: Vec = (0..9000).collect(); + let mut turn_two = turn_one.clone(); + turn_two.extend(100_000..104_000); + let mut first = request("turn-1", 0.0, turn_one.len(), 4); + first.input_token_ids = Some(turn_one); + let mut second = request("turn-2", 1_000.0, turn_two.len(), 4); + second.input_token_ids = Some(turn_two); + let mut spec = aggregated_spec(Backend::Vllm, 1, 0.0, vec![first, second]); + let rank = &mut spec.engine["rank"]; + rank["enable_prefix_caching"] = true.into(); + rank["prefix_match_unit"] = 2176.into(); + rank["vllm_version"] = version.into(); + rank["max_model_len"] = 131_072.into(); + rank["max_num_batched_tokens"] = 8192.into(); + rank["kv_cache_capacity_bytes"] = (74 * GLM_PAGE).into(); + let mut groups = vec![ + serde_json::json!({"name":"mla", "kind":"attention", "num_layers":22, "block_size_tokens":2176, "page_size_bytes":GLM_PAGE, "prefix_cacheable":true}), + serde_json::json!({"name":"kpool_tail", "kind":"state", "num_layers":11, "block_size_tokens":4, "page_size_bytes":GLM_PAGE, "blocks_per_request":1, "prefix_cacheable":false}), + ]; + for i in 0..4 { + let mut group = serde_json::json!({"name":format!("kda_{i}"), "kind":"state", "num_layers":4, "block_size_tokens":2176, "page_size_bytes":GLM_PAGE, "blocks_per_request":1, "prefix_cacheable":true}); + if let Some(blocks) = checkpoint { + group["prefill_checkpoint_blocks"] = blocks.into(); + } + groups.push(group); + } + rank["kv_cache_groups"] = serde_json::Value::Array(groups); + let report = Replayer::new(spec, ReplayEngineFactory::new()) + .unwrap() + .with_telemetry_observer( + 1.0, + Box::new(RecordingTelemetryObserver { + samples: Arc::clone(&samples), + }), + ) + .unwrap() + .run() + .unwrap(); + assert_eq!(report.request_counts.completed_requests, 2, "{version}"); + // The second turn reuses 8,704 of its 13,000 prompt tokens. + let expected_ratio = 8704.0 / (9000.0 + 13_000.0); + assert!( + (report.prefix_cache_reused_ratio - expected_ratio).abs() < 1e-9, + "{version}: {}", + report.prefix_cache_reused_ratio + ); + let samples = samples.lock().unwrap(); + let mut saw_cached_bytes = false; + for sample in samples.iter() { + for row in &sample.decode_scheduler_metrics { + assert_eq!(row.total_blocks, 0); + assert_eq!(row.kv_cache_capacity_bytes, Some(74 * GLM_PAGE)); + let cached = row + .kv_cache_cached_bytes + .expect("block mode reports cached bytes"); + assert_eq!(cached % GLM_PAGE, 0); + saw_cached_bytes |= cached > 0; + let emitted = serde_json::to_value(row).unwrap(); + assert_eq!(emitted["kv_cache_cached_bytes"], cached); + } + } + assert!(saw_cached_bytes, "{version}"); + } +} + +#[test] +fn grouped_prefix_hit_above_the_pool_is_a_livelock() { + // The geometry of grouped_prefix_hit_that_needs_more_blocks_than_the_pool_waits + // in grouped_tests.rs. vLLM's startup check counts no block of a hit + // inside a block, so five blocks validate for a max_model_len of 34. The + // second prompt hits 12 tokens inside a block and then needs seven + // blocks. vLLM keeps it waiting, and with nothing running nothing evicts + // the hit blocks. Replay reports the livelock and drops no request. + for (version, checkpoint) in [("0.31.0", None), ("0.32.0", Some(0))] { + let mut first = request("first", 0.0, 13, 1); + first.input_token_ids = Some((0..13).collect()); + let mut second = request("second", 1_000.0, 33, 1); + second.input_token_ids = Some((0..33).collect()); + let mut spec = aggregated_spec(Backend::Vllm, 1, 0.0, vec![first, second]); + let rank = &mut spec.engine["rank"]; + rank["enable_prefix_caching"] = true.into(); + rank["prefix_match_unit"] = 4.into(); + rank["vllm_version"] = version.into(); + rank["max_model_len"] = 34.into(); + rank["kv_cache_capacity_bytes"] = (5 * 64).into(); + let mut state = serde_json::json!({"name":"state", "kind":"state", "num_layers":1, "block_size_tokens":16, "page_size_bytes":64, "blocks_per_request":1, "prefix_cacheable":true}); + if let Some(blocks) = checkpoint { + state["prefill_checkpoint_blocks"] = blocks.into(); + } + rank["kv_cache_groups"] = serde_json::json!([ + {"name":"full", "kind":"attention", "num_layers":1, "block_size_tokens":16, "page_size_bytes":64, "prefix_cacheable":true}, + state, + ]); + let error = Replayer::new(spec, ReplayEngineFactory::new()) + .unwrap() + .run() + .unwrap_err(); + assert_eq!( + error.to_string(), + "replay invariant violated: offline replay detected an effect-free zero-duration pass with 1 in-flight requests remaining", + "{version}" + ); + } +} + #[test] fn replay_accept_length_counts_decode_work_before_output_truncation() { for backend in [Backend::Vllm, Backend::Sglang] { diff --git a/docs/getting-started/understand-results.md b/docs/getting-started/understand-results.md index 145687d55d..6addbb439c 100644 --- a/docs/getting-started/understand-results.md +++ b/docs/getting-started/understand-results.md @@ -166,7 +166,7 @@ The exports provide the following inputs for warmup and profiling views: | Request/conversation/lane timeline | Absolute request timestamps, lineage, lane, and lifecycle events; retain canceled/failed requests | | ISL/OSL, TTFT, E2E, interactivity | Request lengths and timing; `itl_ms` is per-request average token gap and is null with fewer than two delivered tokens | | Queue depth | `prefill_scheduler_metrics` / `decode_scheduler_metrics`: `running_requests`, `waiting_requests`, worker/rank; plus `router_pending_*_requests` | -| KV use | Scheduler `active_blocks`, `inactive_blocks`, `total_blocks`, `active_cache_usage`, `physical_cache_usage`; grouped caches use `kv_cache_used_bytes` / `kv_cache_capacity_bytes` | +| KV use | Scheduler `active_blocks`, `inactive_blocks`, `total_blocks`, `active_cache_usage`, `physical_cache_usage`; grouped caches use (`kv_cache_used_bytes` + `kv_cache_cached_bytes`) / `kv_cache_capacity_bytes`, and an absent `kv_cache_cached_bytes` counts as 0 | | Prefill/decode throughput | `traffic.committed_prefill_tokens` / `traffic.delivered_output_tokens`, divided by `traffic.duration_s` | | Prefix lookup reuse | Role interval `g1_prefix_cache_lookup`: `matched_tokens`, `queried_tokens`, `query_count`; actual retries count, even if admission later blocks | | Prompt token sources | `traffic.first_admission`: `hbm_reused_tokens`, `host_reused_tokens`, `uncached_tokens`; distinct from executed work and readmission recomputation | diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 948ddbf54a..72d166d46f 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/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 diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index c128529210..590ede8c31 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -25,6 +25,7 @@ HostWorkerPredictionConfig, VisionPredictionConfig, WorkerPredictionConfig, + fpm_worker_deployment, resolve_block_size, ) from .config.traffic import SyntheticSessionSource, SyntheticSource, TraceSource @@ -818,6 +819,13 @@ def _worker_engine_args( payload["max_model_len"] = effective_context_length if cache.prefix_match_unit is not None: payload["prefix_match_unit"] = cache.prefix_match_unit + deployment = fpm_worker_deployment(engine, worker, role) if cache.prefix_caching else None + if deployment is not None and deployment.resources.cache_layout == "grouped": + # The rank hashes at the probe's hash block and follows the grouped rules of its vLLM version. + payload["prefix_match_unit"] = ( + cache.prefix_match_unit or deployment.resources.runtime_memory.hash_block_size_tokens + ) + payload["vllm_version"] = engine.backend_version if cache.state_cache is not None: payload["state_cache"] = {"bytes_per_request": cache.state_cache.bytes_per_request} payload["kv_cache_bytes_per_token"] = cache.bytes_per_token diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 6232d57cb8..63378e8b00 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -327,11 +327,14 @@ def _validate_kvcr(self): @model_validator(mode="after") def _validate_manual_geometry(self) -> KvCachePredictionConfig: infer_state = self.state_cache is not None and self.state_cache.bytes_per_request is None - if self.prefix_match_unit is not None: - if self.state_cache is None: - raise ValueError("prefix_match_unit currently requires state_cache") - if not infer_state and (self.block_size is None or self.block_size % self.prefix_match_unit): - raise ValueError("prefix_match_unit must be a positive divisor of block_size") + # Without state_cache, the engine validation checks the unit, because only it sees a grouped FPM profile. + if ( + self.prefix_match_unit is not None + and self.state_cache is not None + and not infer_state + and (self.block_size is None or self.block_size % self.prefix_match_unit) + ): + raise ValueError("prefix_match_unit must be a positive divisor of block_size") if self.state_cache is not None: if self.capacity.type != "fixed": raise ValueError("state_cache requires fixed capacity (blocks or bytes)") @@ -930,6 +933,7 @@ def _reject_auto_hardware(cls, value: str) -> str: @model_validator(mode="after") def _validate_roles(self) -> EnginePredictionConfig: _validate_fpm_profile(self, {self.mode}, {self.backend}) + _validate_prediction_prefix_match_unit(self) if self.decoder_replay: # Keep the pre-supervision configuration path lightweight. Import # the canonical model identity only when this runtime feature is @@ -1264,19 +1268,7 @@ def _validate_fpm_profile(engine, modes: set[str], backends: set[str]) -> None: if worker.timing.type != "default": raise ValueError("engine.fpm_profile requires default timing for every worker") if isinstance(worker, WorkerPredictionConfig): - parallel = worker.parallelism - deployment = profile.select( - model=engine.model, - system=worker.hardware or engine.hardware, - backend=engine.backend, - backend_version=backend_version, - worker_type=role, - tp_size=parallel.tensor, - pp_size=parallel.pipeline, - attention_dp_size=parallel.attention_data, - moe_tp_size=parallel.moe_tensor, - moe_ep_size=parallel.moe_expert, - ) + deployment = fpm_worker_deployment(engine, worker, role) deployment.resources.validate_envelope( max_num_tokens=worker.scheduler.max_batched_tokens, max_batch_size=worker.scheduler.max_sequences, @@ -1286,7 +1278,13 @@ def _validate_fpm_profile(engine, modes: set[str], backends: set[str]) -> None: raise ValueError("grouped FPM cache supports only aggregated vLLM without speculative decoding") cache = worker.kv_cache if cache.prefix_caching: - raise ValueError("grouped FPM cache requires kv_cache.prefix_caching=false for cold replay") + grouped_prefix_caching_requirements( + deployment, + backend_version=backend_version, + architecture=profile.architecture, + prefix_match_unit=cache.prefix_match_unit, + max_batched_tokens=worker.scheduler.max_batched_tokens, + ) if cache.host_offload is not None or cache.g3_offload is not None: raise ValueError("grouped FPM cache supports only HBM; host and G3 offload are unsupported") if cache.capacity.type == "fixed" or cache.bytes_per_token != "auto": @@ -1297,6 +1295,88 @@ def _validate_fpm_profile(engine, modes: set[str], backends: set[str]) -> None: # lowering to SearchSpace, which validates reachable grouped workers. +def fpm_worker_deployment(engine, worker, role: str): + """The FPM profile deployment that a prediction worker runs, or None without a profile.""" + if engine.fpm_profile is None: + return None + parallel = worker.parallelism + return engine.fpm_profile.select( + model=engine.model, + system=worker.hardware or engine.hardware, + backend=engine.backend, + backend_version=requested_backend_version(engine.backend_version, "vllm"), + worker_type=role, + tp_size=parallel.tensor, + pp_size=parallel.pipeline, + attention_dp_size=parallel.attention_data, + moe_tp_size=parallel.moe_tensor, + moe_ep_size=parallel.moe_expert, + ) + + +def grouped_prefix_caching_requirements( + deployment, + *, + backend_version: str | None, + architecture: str, + prefix_match_unit: int | None, + max_batched_tokens: int, + field: str = "kv_cache.prefix_caching", +) -> int: + """Check one grouped FPM deployment for prefix caching and return the hash block of its rank. + + The prediction and the sweeper share this rule, and planning turns prefix caching on only when it passes. + ``field`` names the setting that turns prefix caching off. The Rust gate also checks the cache group geometry + and the pool capacity. + """ + memory = deployment.resources.runtime_memory + # The evidence: a runtime probe ran with prefix caching and recorded the hash block that replay must follow. + if memory is None or memory.enable_prefix_caching is not True or memory.hash_block_size_tokens is None: + raise ValueError(f"grouped FPM cache requires {field}=false for cold replay") + groups = deployment.resources.cache_groups + if any(group.prefix_cacheable is None for group in groups): + raise ValueError( + "grouped FPM cache prefix caching requires every cache group to declare prefix_cacheable; " + "re-import the runtime probe" + ) + # The Rust rule table has no rule for a convolution group. + if any(group.kind == "convolution" for group in groups): + raise ValueError( + "grouped prefix caching does not support convolution cache groups; turn prefix caching off to replay " + "them cold" + ) + if memory.mamba_cache_mode != "align" and any(group.kind == "state" and group.prefix_cacheable for group in groups): + raise ValueError( + "grouped FPM cache prefix caching requires mamba_cache_mode=align for a prefix-cacheable state group" + ) + if backend_version not in ("0.31.0", "0.32.0"): + raise ValueError("grouped FPM cache prefix caching is modeled for vLLM 0.31.0 and 0.32.0 only") + if len({group.page_size_bytes for group in groups}) > 1: + raise ValueError("grouped FPM cache prefix caching requires equal page_size_bytes across cache groups") + hash_block = memory.hash_block_size_tokens + if prefix_match_unit is not None and prefix_match_unit != hash_block: + raise ValueError( + "kv_cache.prefix_match_unit must equal the profile's hash_block_size_tokens for a grouped FPM cache" + ) + # The GLM indexer keeps one state for every 4 tokens, so hits and chunk starts stay on that grid. + if architecture == "Glm5NextForConditionalGeneration" and (hash_block % 4 or max_batched_tokens % 4): + raise ValueError( + "GLM-5.3-Flash prefix caching requires prefix_match_unit and max_batched_tokens that are multiples of 4" + ) + return hash_block + + +def _validate_prediction_prefix_match_unit(engine: EnginePredictionConfig) -> None: + for role in ("aggregated", "prefill", "decode"): + worker = getattr(engine.workers, role) + if worker is None or worker.kv_cache.prefix_match_unit is None or worker.kv_cache.state_cache is not None: + continue + # A grouped FPM cache with prefix caching hashes at the unit without state_cache. + deployment = fpm_worker_deployment(engine, worker, role) if worker.kv_cache.prefix_caching else None + if deployment is None or deployment.resources.cache_layout != "grouped": + raise ValueError("prefix_match_unit currently requires state_cache") + + def _validate_speculation(engine, *, modes: set[str], backends: set[str]) -> None: if engine.speculation is None: return diff --git a/python/aisimulate/src/aisimulate/replay/dashboard.js b/python/aisimulate/src/aisimulate/replay/dashboard.js index c1ea615153..6616807ad6 100644 --- a/python/aisimulate/src/aisimulate/replay/dashboard.js +++ b/python/aisimulate/src/aisimulate/replay/dashboard.js @@ -833,7 +833,12 @@ window.AISimReplay = { mount(root, report, options = {}) { finite(pool?.kv_cache_capacity_bytes) && pool.kv_cache_capacity_bytes > 0 ? percent( - pool.kv_cache_used_bytes, + finite(pool.kv_cache_used_bytes) + ? pool.kv_cache_used_bytes + + (finite(pool.kv_cache_cached_bytes) + ? pool.kv_cache_cached_bytes + : 0) + : null, pool.kv_cache_capacity_bytes, ) : finite(pool?.physical_cache_usage) diff --git a/python/aisimulate/src/aisimulate/replay/dashboard.py b/python/aisimulate/src/aisimulate/replay/dashboard.py index bb94e577e1..7cd47e2c85 100644 --- a/python/aisimulate/src/aisimulate/replay/dashboard.py +++ b/python/aisimulate/src/aisimulate/replay/dashboard.py @@ -87,7 +87,11 @@ def _cache_percent(workers: list[dict[str, Any]]) -> float | None: capacities = [worker.get("kv_cache_capacity_bytes") for worker in workers] if all(value is not None and value > 0 for value in capacities): used = [worker.get("kv_cache_used_bytes") for worker in workers] - return 100 * sum(used) / sum(capacities) if all(value is not None for value in used) else None + if any(value is None for value in used): + return None + # Cached blocks that no request holds stay resident. A pool without prefix caching reports none. + cached = sum(worker.get("kv_cache_cached_bytes") or 0 for worker in workers) + return 100 * (sum(used) + cached) / sum(capacities) capacities = [worker.get("total_blocks") for worker in workers] usage = [worker.get("physical_cache_usage") for worker in workers] if all(value is not None and value > 0 for value in capacities) and all(value is not None for value in usage): @@ -124,6 +128,7 @@ def _sample(row: Mapping[str, Any], *, aggregated: bool) -> dict[str, Any]: "total_blocks", "kv_cache_used_bytes", "kv_cache_capacity_bytes", + "kv_cache_cached_bytes", "active_cache_usage", "physical_cache_usage", "running_requests", diff --git a/python/aisimulate/src/aisimulate/support/plan.py b/python/aisimulate/src/aisimulate/support/plan.py index 978395e356..b5c58e9588 100644 --- a/python/aisimulate/src/aisimulate/support/plan.py +++ b/python/aisimulate/src/aisimulate/support/plan.py @@ -21,6 +21,7 @@ import yaml from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.config.engine import grouped_prefix_caching_requirements from .config_profile import _cache_group_shapes, linear_block_size, linear_runtime_resources from .fpm import fpm_cli_args @@ -148,8 +149,23 @@ def _configs( if request.fpm_profile is not None: engine["fpm_profile"] = request.fpm_profile.model_dump(mode="json") worker["timing"]["estimator_config"] = {"fpm_interpolation": {"method": "direct"}} - if request.profile_deployment().resources.cache_layout == "grouped": - worker.setdefault("kv_cache", {})["prefix_caching"] = False + deployment = request.profile_deployment() + resources = deployment.resources + if resources.cache_layout == "grouped": + # Replay reuses prefixes only when the deployment passes the grouped prefix-caching requirements that the + # prediction checks. Otherwise the plan keeps the cold replay. + try: + grouped_prefix_caching_requirements( + deployment, + backend_version=deployment.backend_version, + architecture=request.fpm_profile.architecture, + prefix_match_unit=None, + max_batched_tokens=worker["scheduler"]["max_batched_tokens"], + ) + prefix_caching = True + except ValueError: + prefix_caching = False + worker.setdefault("kv_cache", {})["prefix_caching"] = prefix_caching else: # Replay of a linear profile with runtime probe evidence uses the measured block size. observed = _probe_observed_resources(request) diff --git a/python/aisimulate/src/aisimulate/support/validation.py b/python/aisimulate/src/aisimulate/support/validation.py index 22a2ea4ccc..62c50dad68 100644 --- a/python/aisimulate/src/aisimulate/support/validation.py +++ b/python/aisimulate/src/aisimulate/support/validation.py @@ -77,6 +77,13 @@ def _prediction_config(request: SupportRequest, plan: Path, trace: Path) -> dict if engine["mode"] != "aggregated" or engine.get("speculation") is not None: raise ValueError("FPM replay validation requires aggregated serving with speculation disabled") worker = engine["workers"]["aggregated"] + deployment = request.profile_deployment() + if deployment is not None and deployment.resources.cache_layout == "grouped": + # validate-fpm tests FPM timing, so its replay stays cold even when the plan reuses prefixes. On a grouped + # rank only prefix caching reads prefix_match_unit, so the unit goes too. + cache = worker.setdefault("kv_cache", {}) + cache["prefix_caching"] = False + cache.pop("prefix_match_unit", None) timing = worker["timing"] timing["estimation_mode"] = "fpm_interpolation" timing["fallback_policy"] = "deny" diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index 64a5c3686f..0dd23263d4 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -36,7 +36,11 @@ is_active_engine_model_control, requested_backend_version, ) -from ..config.engine import NgramSpeculationConfig, ReplayEstimationMode +from ..config.engine import ( + NgramSpeculationConfig, + ReplayEstimationMode, + grouped_prefix_caching_requirements, +) class OptimizationTarget(str, Enum): @@ -1250,7 +1254,7 @@ def _validate_grouped_cache(self) -> SearchSpace: profile = load_fpm_profile(self.fpm_profile) if not any(deployment.resources.cache_layout == "grouped" for deployment in profile.deployments): return self - from .search_space import _fpm_parallel_configs, _parallel_role + from .search_space import _fpm_parallel_configs, _parallel_role, _role_runtime for mode in self.deployment_mode: configs = _fpm_parallel_configs(self, mode) @@ -1273,7 +1277,15 @@ def _validate_grouped_cache(self) -> SearchSpace: if role != "agg" or self.aic_nextn or self.speculation is not None: raise ValueError("grouped FPM cache requires aggregated vLLM without speculative decoding") if self.agg_enable_prefix_caching: - raise ValueError("grouped FPM cache requires agg_enable_prefix_caching=false for cold replay") + # The search has no prefix_match_unit. Each candidate is checked again with its own budget. + grouped_prefix_caching_requirements( + deployment, + backend_version=self.requested_backend_version("vllm"), + architecture=profile.architecture, + prefix_match_unit=None, + max_batched_tokens=_role_runtime(self, "vllm", "agg")[0], + field="agg_enable_prefix_caching", + ) if self.agg_native_host_offload is not None: raise ValueError("grouped FPM cache supports only HBM without host offload") if self.agg_num_gpu_blocks is not None or self.agg_kv_bytes_per_token != "auto": diff --git a/python/aisimulate/src/aisimulate/sweeper/deploy.py b/python/aisimulate/src/aisimulate/sweeper/deploy.py index 8e8c7b0564..62638e5982 100644 --- a/python/aisimulate/src/aisimulate/sweeper/deploy.py +++ b/python/aisimulate/src/aisimulate/sweeper/deploy.py @@ -117,6 +117,8 @@ def _engine_args_payload( context_length = sample.get(f"{role}_context_length") or sample.get("context_length") if context_length is not None: payload["max_model_len"] = int(context_length) + if role == "agg" and payload["enable_prefix_caching"] and sample.get("fpm_profile") is not None: + _add_grouped_prefix_caching(payload, sample, backend_version=backend_version) if moe_tp * moe_ep > 1: payload["aic_moe_tp_size"] = moe_tp payload["aic_moe_ep_size"] = moe_ep @@ -213,6 +215,43 @@ def _engine_args_payload( return payload +def _add_grouped_prefix_caching(payload: dict[str, Any], sample: dict[str, Any], *, backend_version: str) -> None: + """Give a grouped FPM rank with prefix caching the keys that the Rust gate needs, after the shared check.""" + from aisimulate_core.sdk.fpm_profile import load_fpm_profile + + from ..config.engine import grouped_prefix_caching_requirements + + profile = load_fpm_profile(sample["fpm_profile"]) + # A profile without grouped deployments keeps its candidate path as it was. + if all(deployment.resources.cache_layout != "grouped" for deployment in profile.deployments): + return + deployment = profile.select( + model=str(sample["model_name"]), + system=_role_hardware_sku(sample, "agg"), + backend=str(sample["backend"]), + backend_version=backend_version, + tp_size=int(sample["tp"]), + pp_size=int(sample["pp"]), + attention_dp_size=int(sample["attention_dp"]), + moe_tp_size=int(sample["moe_tp"]), + moe_ep_size=int(sample["moe_ep"]), + ) + if deployment.resources.cache_layout != "grouped": + return + # The search gate checks the largest budget of the domain. Each candidate is checked with its own, before replay. + payload["prefix_match_unit"] = grouped_prefix_caching_requirements( + deployment, + backend_version=backend_version, + architecture=profile.architecture, + prefix_match_unit=None, + max_batched_tokens=int(sample["agg_max_num_batched_tokens"]), + field="agg_enable_prefix_caching", + ) + payload["vllm_version"] = deployment.backend_version + # vLLM always sizes its pool with a model length, and the Rust gate refuses a rank without one. + payload.setdefault("max_model_len", profile.context_length) + + def _profile_kv_args(sample: dict[str, Any], role: str, backend_version: str) -> dict[str, Any]: if sample.get("fpm_profile") is None: return {} diff --git a/python/aisimulate/src/aisimulate_core/fpm_profile.py b/python/aisimulate/src/aisimulate_core/fpm_profile.py index 3864973e4d..608d5d4ef0 100644 --- a/python/aisimulate/src/aisimulate_core/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/fpm_profile.py @@ -41,7 +41,10 @@ class FpmCacheGroup(_ProfileModel): ``block_size_tokens`` is the block size of vLLM's state spec, as the runtime probe records it: the Mamba block (the scheduler block in ``align`` mode, ``max_model_len`` in mode ``none``) or the ring block. - Sizing and replay do not read it yet. + Sizing and replay with prefix caching off do not read it. Grouped prefix + caching reads it as the Mamba block of ``align`` mode, and every group's + block enters the scheduler block, the least common multiple of the group + blocks. ``blocks_per_request`` counts the pool blocks that a running request holds between block boundaries. For a Mamba spec, that is ``1 + k``, where ``k`` @@ -124,8 +127,10 @@ class FpmRuntimeMemoryProfile(_ProfileModel): This capacity already accounts for runtime graph and workspace reservations. It is not a non-KV memory limit and must not be rescaled to another setting. The three optional prefix-caching fields record the settings that the probe - ran with. They are evidence only: selection, sizing and replay do not read - them, and the simulated deployment decides prefix caching. + ran with. Selection and sizing do not read them. Grouped prefix caching + needs a probe that ran with prefix caching, and replay then hashes prompts + in blocks of ``hash_block_size_tokens``. For a linear profile, the + simulated deployment decides prefix caching. """ kv_cache_bytes: Annotated[int, Field(gt=0, le=2**53)] diff --git a/python/aisimulate/tests/unit/test_onboard_runtime.py b/python/aisimulate/tests/unit/test_onboard_runtime.py index 2242d4709d..749c846e01 100644 --- a/python/aisimulate/tests/unit/test_onboard_runtime.py +++ b/python/aisimulate/tests/unit/test_onboard_runtime.py @@ -1510,6 +1510,39 @@ def test_finalized_linear_runtime_profile_keeps_the_measured_block_size(tmp_path check_plan(resolved, target) +def test_finalized_grouped_profile_with_prefix_caching_evidence_plans_and_recommends_with_it(tmp_path, capsys): + import yaml + + from aisimulate.support.finalization import finalize + from aisimulate.sweeper import SweepResult + from aisimulate.sweeper.result import CandidateStatus + + # The vLLM 0.31.0 records of the hybrid probe ran with prefix caching in align mode, and the import marks + # every cache group. Finalization keeps that evidence. + _, request, _, root = _completed_runtime_collection(tmp_path, capsys, layout="mamba_hybrid") + target = tmp_path / "finalized" + finalize(request, root, target) + resources = SupportRequest.from_yaml(target / "request.yaml").profile_deployment().resources + assert (resources.cache_layout, resources.runtime_memory.enable_prefix_caching) == ("grouped", True) + assert all(group.prefix_cacheable is not None for group in resources.cache_groups) + for name in ("predict/pilot.yaml", "recommend/pilot.yaml"): + worker = yaml.safe_load((target / name).read_text())["engine"]["workers"]["aggregated"] + assert worker["kv_cache"]["prefix_caching"] is True + # The synthetic timings measure 16-token prefills and decodes that read 8 or 32 tokens, so the pilot runs its + # planned engine with a request that they cover. + pilot = yaml.safe_load((target / "recommend/pilot.yaml").read_text()) + pilot["traffic"]["source"].update(input_tokens=16, output_tokens=4) + path = tmp_path / "recommend-pilot.yaml" + path.write_text(yaml.safe_dump(pilot)) + output = tmp_path / "recommendation" + assert cli.main(["recommend", "--config", str(path), "--output-dir", str(output), "--format", "json"]) == 0 + capsys.readouterr() + # The sweeper's candidate rank carries the hash block and the vLLM version, so the Rust gate admits it. + result = SweepResult.from_json((output / "recommendation.json").read_text()) + assert result.candidates + assert all(candidate.status is CandidateStatus.FEASIBLE for candidate in result.candidates) + + @pytest.mark.parametrize( "corruption", ["missing-observed-resources", "probe-not-a-dict", "observed-resources-not-a-dict", "missing-cache-groups"], diff --git a/tests/browser/replay_dashboard.cjs b/tests/browser/replay_dashboard.cjs index bdd8d13beb..380dfb9e6d 100644 --- a/tests/browser/replay_dashboard.cjs +++ b/tests/browser/replay_dashboard.cjs @@ -1108,6 +1108,8 @@ async function main() { dp_rank: 0, kv_cache_capacity_bytes: 4 * 2 ** 30, kv_cache_used_bytes: (index ? 3 : 1) * 2 ** 30, + // Retained prefixes stay resident. The first sample has none. + ...(index ? { kv_cache_cached_bytes: 2 ** 29 } : {}), }, ]; sample.host_pools = [ @@ -1192,8 +1194,8 @@ async function main() { ); close( (await traces("cache"))[0].y.at(-1), - 65, - "grouped byte utilization", + 75, + "grouped byte utilization with retained prefixes", ); await mode("queue", "preemptions"); assert.deepEqual( diff --git a/tests/test_fpm_grouped_workflow.py b/tests/test_fpm_grouped_workflow.py index 3338f8fdea..e89d354f60 100644 --- a/tests/test_fpm_grouped_workflow.py +++ b/tests/test_fpm_grouped_workflow.py @@ -112,7 +112,7 @@ def no_graph(*_args, **_kwargs): return profile, root -def _write_timings(root, profile, *, decode_ceiling=8192): +def _write_timings(root, profile, *, decode_ceiling=8192, version="0.25.1"): import pyarrow as pa import pyarrow.parquet as pq @@ -147,7 +147,7 @@ def _write_timings(root, profile, *, decode_ceiling=8192): "kv_seed_regime": "real_kv", } ) - path = root / "data/h200_sxm/vllm/0.25.1/fpm_forward_perf.parquet" + path = root / f"data/h200_sxm/vllm/{version}/fpm_forward_perf.parquet" path.parent.mkdir(parents=True, exist_ok=True) pq.write_table(pa.Table.from_pylist(rows), path) path.with_suffix(".metadata.json").write_text( @@ -159,7 +159,7 @@ def _write_timings(root, profile, *, decode_ceiling=8192): "measurement_policy": "dynamo_native_single_sample_v1", "system": "h200_sxm", "backend": "vllm", - "backend_version": "0.25.1", + "backend_version": version, "row_count": len(rows), "parquet_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), } @@ -238,6 +238,70 @@ def _with_state_group(profile): return profile +# A convolution group as the probe import writes one for Inkling: a 4-token window, marked cacheable. Its block and +# page fit _with_prefix_evidence. +_CONVOLUTION_GROUP = { + "name": "conv", + "kind": "convolution", + "num_layers": 3, + "block_size_tokens": 64, + "page_size_bytes": 64, + "sliding_window": 4, + "prefix_cacheable": True, +} + + +def _with_prefix_evidence(profile, *, hash_block=64, version="0.31.0"): + """Runtime memory from a probe that ran with prefix caching on, and every group marked cacheable. + + The Rust gate needs one page size, a hash block that divides every cacheable block and one running block + per state group, so every page is 64 bytes and every block is one hash block. Grouped prefix caching refuses + convolution groups, so the fictional convolution group goes. + """ + profile = deepcopy(profile) + for deployment in profile["deployments"]: + deployment["backend_version"] = version + resources = deployment["resources"] + resources["cache_groups"] = [group for group in resources["cache_groups"] if group["kind"] != "convolution"] + for name in ("weights_bytes", "activations_bytes", "runtime_overhead_bytes", "comm_overhead_bytes"): + resources.pop(name) + # Runtime memory requires the exact recorded scheduler envelope of _prediction. + resources["max_batch_size"] = 4 + resources["runtime_memory"] = { + "kv_cache_bytes": 3456, + "gpu_memory_utilization": 0.9, + # The Rust gate admits this pool at 1024 tokens, as vLLM's startup check would. + "max_model_len": 1024, + "provenance": "Synthetic pool of 54 blocks of 64 bytes from a fictional probe. No GPU qualification.", + "enable_prefix_caching": True, + "mamba_cache_mode": "align", + "hash_block_size_tokens": hash_block, + } + for group in resources["cache_groups"]: + group.update(block_size_tokens=hash_block, page_size_bytes=64, prefix_cacheable=True) + if group["kind"] == "state": + group["blocks_per_request"] = 1 + if version == "0.32.0": + group["prefill_checkpoint_blocks"] = 0 + return profile + + +def _prefix_caching_prediction(profile, root, version): + """Two identical 520-token requests in sequence on a grouped rank with prefix caching.""" + raw = _prediction(profile, root) + raw["engine"]["backend_version"] = version + # The rank runs at the probe's max_model_len. The Rust gate needs a pool for one prompt of max_model_len plus + # the state blocks of one request: 16 + 13 + 2 = 31 of the 54 blocks at 1024 tokens. + raw["engine"]["context_length"] = profile["deployments"][0]["resources"]["runtime_memory"]["max_model_len"] + raw["engine"]["workers"]["aggregated"]["kv_cache"]["prefix_caching"] = True + raw["traffic"] = { + "source": {"type": "synthetic", "input_tokens": 520, "output_tokens": 4, "cached_prefix_tokens": 520}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 2}, + } + return raw + + @pytest.mark.parametrize("unrelated", ["topology", "hardware", "version"]) def test_mixed_profile_predict_preserves_selected_linear_prefix_cache(grouped_case, tmp_path, unrelated): profile, root = grouped_case @@ -852,18 +916,24 @@ def render_without_launch(_args, resolved): assert collector_cli.main([arg for arg in command if arg != "--plan-only"]) == 0 -@pytest.mark.parametrize("dep", [False, True]) -def test_onboard_validation_replays_windowed_weka_with_full_context(grouped_case, tmp_path, monkeypatch, dep): +def _plan_and_validate_weka(profile, tmp_path, monkeypatch, *, dep=False, version="0.25.1", context_length=2048): + """Plan the onboarding of a grouped profile and replay a windowed Weka trace with validate-fpm. + + A shorter context, such as a probe's max_model_len, gets a workload and a second request that fit in it. The + second request still crosses the 512-token window. + """ from aisimulate import supervision - profile, _root = grouped_case + # The default workload of 1,024 input and 128 output tokens needs a context of 1,152 tokens. + workload = {} if context_length >= 1152 else {"workload": {"input_tokens": 512, "output_tokens": 4}} + second = min(1152, context_length - 64) request = SupportRequest.model_validate( { "identity": { "model": profile["model"], "model_revision": profile["model_revision"], "model_kind": "moe", - "framework_version": "0.25.1", + "framework_version": version, "gpu": "h200_sxm", "interconnect": "nvswitch", }, @@ -872,15 +942,16 @@ def test_onboard_validation_replays_windowed_weka_with_full_context(grouped_case "attention_data_parallel": 2 if dep else 1, "moe_tensor_parallel": 1 if dep else 2, "moe_expert_parallel": 2 if dep else 1, - "context_length": 2048, + "context_length": context_length, }, + **workload, "collection": {"max_num_tokens": 128, "max_batch_size": 4}, "fpm_profile": profile, } ) plan = tmp_path / "collection" create_plan(request, plan) - _write_timings(plan / "systems", profile) + _write_timings(plan / "systems", profile, version=version) trace = tmp_path / "trace.jsonl" trace.write_text( json.dumps( @@ -902,9 +973,9 @@ def test_onboard_validation_replays_windowed_weka_with_full_context(grouped_case "t": 0.1, "type": "s", "model": "trace/source-model", - "in": 1152, + "in": second, "out": 4, - "hash_ids": list(range(18)), + "hash_ids": list(range(second // 64)), }, ], } @@ -930,6 +1001,13 @@ def test_onboard_validation_replays_windowed_weka_with_full_context(grouped_case ) == 0 ) + return plan, output + + +@pytest.mark.parametrize("dep", [False, True]) +def test_onboard_validation_replays_windowed_weka_with_full_context(grouped_case, tmp_path, monkeypatch, dep): + profile, _root = grouped_case + _plan, output = _plan_and_validate_weka(profile, tmp_path, monkeypatch, dep=dep) validation = json.loads((output / "validation.json").read_text()) assert validation["status"] == "covered" assert validation["accuracy"] == "not_assessed" @@ -992,3 +1070,289 @@ def test_aggregated_only_grouped_profile_skips_unavailable_disagg(grouped_case): with pytest.warns(UserWarning, match="disagg.*skipped"): branches = enumerate_branches(config, max_seq_len=2048) assert [branch.deployment_mode for branch in branches] == ["agg"] + + +def _prefix_search(profile, root, version): + return {**_mixed_search(profile, root), "fpm_profile": profile, "backend_version": version} + + +@pytest.mark.parametrize( + "case,error", + [ + ("no_evidence", "prefix_caching=false for cold replay"), + ("probe_without_prefix_caching", "prefix_caching=false for cold replay"), + ("unmarked_group", "every cache group to declare prefix_cacheable; re-import the runtime probe"), + ("convolution_group", "grouped prefix caching does not support convolution cache groups"), + ("state_without_align", "requires mamba_cache_mode=align for a prefix-cacheable state group"), + ("vllm_0320", None), + ("vllm_0300", "is modeled for vLLM 0.31.0 and 0.32.0 only"), + ("unequal_pages", "requires equal page_size_bytes across cache groups"), + ("other_unit", "must equal the profile's hash_block_size_tokens for a grouped FPM cache"), + ], +) +def test_grouped_prefix_caching_needs_vllm_evidence_with_align_mode(grouped_case, case, error): + from aisimulate.sweeper.config import SearchSpace + + profile, root = grouped_case + version = {"no_evidence": "0.25.1", "vllm_0320": "0.32.0", "vllm_0300": "0.30.0"}.get(case, "0.31.0") + profile = _with_state_group(profile) + if case != "no_evidence": + profile = _with_prefix_evidence(profile, version=version) + for deployment in profile["deployments"]: + resources = deployment["resources"] + if case == "probe_without_prefix_caching": + resources["runtime_memory"].update(enable_prefix_caching=False, mamba_cache_mode="none") + elif case == "unmarked_group": + del resources["cache_groups"][1]["prefix_cacheable"] + elif case == "convolution_group": + resources["cache_groups"].append(dict(_CONVOLUTION_GROUP)) + elif case == "state_without_align": + resources["runtime_memory"]["mamba_cache_mode"] = "none" + elif case == "unequal_pages": + resources["cache_groups"][1]["page_size_bytes"] = 16 + raw = _prediction(profile, root) + raw["engine"]["backend_version"] = version + cache = raw["engine"]["workers"]["aggregated"]["kv_cache"] + cache["prefix_caching"] = True + if case == "other_unit": + cache["prefix_match_unit"] = 32 + # The search space has no prefix_match_unit, so only the prediction can state another unit. + search = None if case == "other_unit" else _prefix_search(profile, root, version) + if error is None: + CorePredictionConfig.model_validate(raw) + SearchSpace.model_validate(search) + return + with pytest.raises(ValidationError, match=error): + CorePredictionConfig.model_validate(raw) + if search is not None: + with pytest.raises(ValidationError, match=error): + SearchSpace.model_validate(search) + + +@pytest.mark.parametrize("version", ["0.31.0", "0.32.0"]) +def test_grouped_prefix_caching_replays_with_reuse(grouped_case, tmp_path, version): + profile, root = grouped_case + profile = _with_prefix_evidence(_with_state_group(profile), version=version) + _write_timings(root, profile, version=version) + raw = _prefix_caching_prediction(profile, root, version) + args = prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)).backend_deployment.agg_engine_args + assert (args["enable_prefix_caching"], args["prefix_match_unit"], args["vllm_version"]) == (True, 64, version) + path = tmp_path / "predict.yaml" + path.write_text(yaml.safe_dump(raw)) + output = tmp_path / "prediction" + assert cli.main(["predict", "-c", str(path), "--output-dir", str(output), "--capture-per-request"]) == 0 + report = json.loads((output / "prediction.json").read_text()) + assert report["completed_requests"] == 2 + # Both versions keep the state at 512, the last 64-token boundary below the 520-token prompt, so the + # identical second request reuses 512 tokens. + assert [row["reused_input_tokens"] for row in report["per_request"]] == [0, 512] + + +def test_grouped_prefix_caching_rank_takes_the_probe_hash_block(grouped_case): + profile, root = grouped_case + raw = _prefix_caching_prediction(_with_prefix_evidence(_with_state_group(profile)), root, "0.31.0") + + def engine_args(): + return prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)).backend_deployment.agg_engine_args + + implicit = engine_args() + raw["engine"]["workers"]["aggregated"]["kv_cache"]["prefix_match_unit"] = 64 + assert engine_args() == implicit + # The Rust gate also refuses a rank without max_model_len. The compiler writes it for every vLLM rank. + assert (implicit["prefix_match_unit"], implicit["vllm_version"], implicit["max_model_len"]) == (64, "0.31.0", 1024) + raw["engine"]["workers"]["aggregated"]["kv_cache"] = {"prefix_caching": False} + assert not {"prefix_match_unit", "vllm_version"} & engine_args().keys() + + +@pytest.mark.parametrize( + "hash_block,budget,refused", + [(2176, 8190, True), (2176, 8192, False), (2174, 8192, True)], +) +def test_glm_grid_refuses_off_grid_units_and_budgets(grouped_case, hash_block, budget, refused): + from aisimulate.sweeper.config import SearchSpace + + profile, root = grouped_case + profile = _with_prefix_evidence(profile, hash_block=hash_block) + profile["architecture"] = "Glm5NextForConditionalGeneration" + for deployment in profile["deployments"]: + # Runtime memory requires the exact recorded budget. + deployment["resources"]["max_num_tokens"] = budget + raw = _prediction(profile, root) + raw["engine"]["backend_version"] = "0.31.0" + worker = raw["engine"]["workers"]["aggregated"] + worker["kv_cache"]["prefix_caching"] = True + worker["scheduler"]["max_batched_tokens"] = budget + search = {**_prefix_search(profile, root, "0.31.0"), "agg_max_num_batched_tokens": [budget]} + if not refused: + CorePredictionConfig.model_validate(raw) + SearchSpace.model_validate(search) + return + error = "GLM-5.3-Flash prefix caching requires prefix_match_unit and max_batched_tokens that are multiples of 4" + with pytest.raises(ValidationError, match=error): + CorePredictionConfig.model_validate(raw) + with pytest.raises(ValidationError, match=error): + SearchSpace.model_validate(search) + + +def test_grouped_plan_keeps_prefix_caching_on_with_evidence_and_validation_cold(grouped_case, tmp_path, monkeypatch): + profile, _root = grouped_case + profile = _with_prefix_evidence(profile) + # Onboarding keeps the context within the probe's max_model_len (support/cli.py:317-320). + plan, output = _plan_and_validate_weka(profile, tmp_path, monkeypatch, version="0.31.0", context_length=1024) + for name in ("predict/pilot.yaml", "recommend/pilot.yaml"): + worker = yaml.safe_load((plan / name).read_text())["engine"]["workers"]["aggregated"] + assert worker["kv_cache"]["prefix_caching"] is True + validation = json.loads((output / "validation.json").read_text()) + assert validation["status"] == "covered" + # validate-fpm tests FPM timing, so its replay stays cold. + saved = CorePredictionConfig.from_yaml(output / "predict.yaml") + assert saved.engine.workers.aggregated.kv_cache.prefix_caching is False + + +def test_cold_validation_drops_the_unit_that_only_prefix_caching_reads(grouped_case, tmp_path, monkeypatch): + from aisimulate.support.validation import _prediction_config + + profile, _root = grouped_case + plan, _output = _plan_and_validate_weka( + _with_prefix_evidence(profile), tmp_path, monkeypatch, version="0.31.0", context_length=1024 + ) + # The probe's hash block is a valid prefix_match_unit for a grouped rank with prefix caching. The plan does not + # write it, and validate-fpm refuses a changed pilot, so the test gives such a pilot to _prediction_config. + pilot = yaml.safe_load((plan / "predict/pilot.yaml").read_text()) + cache = pilot["engine"]["workers"]["aggregated"]["kv_cache"] + cache["prefix_match_unit"] = 64 + CorePredictionConfig.model_validate(pilot) + edited = tmp_path / "edited" + (edited / "predict").mkdir(parents=True) + (edited / "predict/pilot.yaml").write_text(yaml.safe_dump(pilot)) + # On a grouped rank only prefix caching reads the unit, so the cold replay turns prefix caching off and drops it. + cold = _prediction_config(SupportRequest.from_yaml(plan / "request.yaml"), edited, tmp_path / "trace.jsonl") + del cache["prefix_match_unit"] + assert cold["engine"]["workers"]["aggregated"]["kv_cache"] == {**cache, "prefix_caching": False} + + +def test_grouped_plan_without_the_probe_hash_block_stays_cold(grouped_case, tmp_path, monkeypatch): + # Replay hashes prompts in blocks of the probe's hash block, so a probe record without it is no evidence. + profile, _root = grouped_case + profile = _with_prefix_evidence(profile) + for deployment in profile["deployments"]: + del deployment["resources"]["runtime_memory"]["hash_block_size_tokens"] + plan, _output = _plan_and_validate_weka(profile, tmp_path, monkeypatch, version="0.31.0", context_length=1024) + for name in ("predict/pilot.yaml", "recommend/pilot.yaml"): + worker = yaml.safe_load((plan / name).read_text())["engine"]["workers"]["aggregated"] + assert worker["kv_cache"]["prefix_caching"] is False + + +@pytest.mark.parametrize("version", ["0.30.0", "0.25.1"]) +def test_grouped_plan_on_an_unmodeled_vllm_version_stays_cold(grouped_case, tmp_path, monkeypatch, version): + # Replay models grouped prefix caching for vLLM 0.31.0 and 0.32.0 only, so a probe of another version with + # evidence and markers plans the cold replay that it planned before. + profile, _root = grouped_case + profile = _with_prefix_evidence(profile, version=version) + plan, output = _plan_and_validate_weka(profile, tmp_path, monkeypatch, version=version, context_length=1024) + for name in ("predict/pilot.yaml", "recommend/pilot.yaml"): + worker = yaml.safe_load((plan / name).read_text())["engine"]["workers"]["aggregated"] + assert worker["kv_cache"]["prefix_caching"] is False + assert json.loads((output / "validation.json").read_text())["status"] == "covered" + + +@pytest.mark.parametrize("unsupported", ["state_without_align", "convolution_group"]) +def test_grouped_plan_that_fails_a_prefix_caching_requirement_stays_cold( + grouped_case, tmp_path, monkeypatch, unsupported +): + # A probe with evidence and markers on a modeled version still plans the cold replay that it planned before + # when the deployment fails another grouped prefix-caching requirement, and that replay runs. + profile, _root = grouped_case + profile = _with_prefix_evidence(_with_state_group(profile)) + for deployment in profile["deployments"]: + resources = deployment["resources"] + if unsupported == "state_without_align": + resources["runtime_memory"]["mamba_cache_mode"] = "none" + else: + resources["cache_groups"].append(dict(_CONVOLUTION_GROUP)) + plan, output = _plan_and_validate_weka(profile, tmp_path, monkeypatch, version="0.31.0", context_length=1024) + for name in ("predict/pilot.yaml", "recommend/pilot.yaml"): + worker = yaml.safe_load((plan / name).read_text())["engine"]["workers"]["aggregated"] + assert worker["kv_cache"]["prefix_caching"] is False + assert json.loads((output / "validation.json").read_text())["status"] == "covered" + assert json.loads((output / "prediction/prediction.json").read_text())["completed_requests"] == 2 + + +@pytest.mark.parametrize("prefix_caching", [False, True]) +def test_grouped_sweeper_candidate_rank_takes_the_probe_hash_block(grouped_case, prefix_caching): + from aisimulate.sweeper.config import SearchSpace + from aisimulate.sweeper.deploy import build_backend_deployment + from aisimulate.sweeper.parallel_enum import ParallelShape, ReplicaParallelConfig + from aisimulate.sweeper.sample import unroll_sample + + profile, root = grouped_case + profile = _with_prefix_evidence(_with_state_group(profile)) + raw = {**_prefix_search(profile, root, "0.31.0"), "agg_enable_prefix_caching": prefix_caching} + sample = unroll_sample( + search_space=SearchSpace.model_validate(raw), + selection={ + "deployment_mode": "agg", + "backend": "vllm", + "agg_max_num_batched_tokens": 128, + "agg_max_num_seqs": 4, + }, + parallel_config=ReplicaParallelConfig(ParallelShape(tp=2, dp=1, moe_tp=2, moe_ep=1), replicas=1), + ) + args = build_backend_deployment(sample, backend_version="0.31.0").agg_engine_args + if prefix_caching: + # Without a search context length, the rank takes the profile's. + assert (args["prefix_match_unit"], args["vllm_version"], args["max_model_len"]) == (64, "0.31.0", 2048) + else: + assert not {"prefix_match_unit", "vllm_version", "max_model_len"} & args.keys() + + +def test_grouped_prefix_caching_recommendation_replays_every_candidate(grouped_case, tmp_path, monkeypatch): + from aisimulate.recommend import _run_recommendation + from aisimulate.sweeper import SweepResult + from aisimulate.sweeper.result import CandidateStatus + + monkeypatch.setattr("aisimulate.recommend.run_recommendation", _run_recommendation) + profile, root = grouped_case + profile = _with_prefix_evidence(_with_state_group(profile)) + _write_timings(root, profile, version="0.31.0") + raw = _prefix_caching_prediction(profile, root, "0.31.0") + raw["engine"]["workers"]["aggregated"]["parallelism"] = {"preset": "default"} + raw["optimization"] = {"constraints": {"max_candidate_gpus": 2}} + raw["optimizer"] = {"algorithm": "random", "max_trials": 2, "parallelism": 1} + path = tmp_path / "recommend.yaml" + path.write_text(yaml.safe_dump(raw)) + output = tmp_path / "recommendation" + assert cli.main(["recommend", "-c", str(path), "--output-dir", str(output)]) == 0 + result = SweepResult.from_json((output / "recommendation.json").read_text()) + # Both topologies replay. A candidate that the Rust gate refused would be failed with "replay failed". + assert len(result.candidates) == 2 + assert all(candidate.status is CandidateStatus.FEASIBLE for candidate in result.candidates) + exported = CorePredictionConfig.from_yaml(output / "recommendations/0001.yaml") + assert exported.engine.workers.aggregated.kv_cache.prefix_caching is True + + +def test_grouped_sweeper_checks_each_candidate_budget_before_replay(grouped_case): + from aisimulate.sweeper.config import SearchSpace + from aisimulate.sweeper.deploy import build_backend_deployment + from aisimulate.sweeper.parallel_enum import ParallelShape, ReplicaParallelConfig + from aisimulate.sweeper.sample import unroll_sample + + profile, root = grouped_case + profile = _with_prefix_evidence(profile, hash_block=2176) + profile["architecture"] = "Glm5NextForConditionalGeneration" + # The search gate checks the largest budget, 8192. The 8190 candidate is refused when its rank is built. + search = SearchSpace.model_validate( + {**_prefix_search(profile, root, "0.31.0"), "agg_max_num_batched_tokens": [8190, 8192]} + ) + sample = unroll_sample( + search_space=search, + selection={ + "deployment_mode": "agg", + "backend": "vllm", + "agg_max_num_batched_tokens": 8190, + "agg_max_num_seqs": 4, + }, + parallel_config=ReplicaParallelConfig(ParallelShape(tp=2, dp=1, moe_tp=2, moe_ep=1), replicas=1), + ) + with pytest.raises(ValueError, match="max_batched_tokens that are multiples of 4"): + build_backend_deployment(sample, backend_version="0.31.0") diff --git a/tests/test_replay_dashboard.py b/tests/test_replay_dashboard.py index bed4311b2d..bff82100e4 100644 --- a/tests/test_replay_dashboard.py +++ b/tests/test_replay_dashboard.py @@ -184,6 +184,12 @@ def test_byte_capacities_and_unknown_metrics_are_preserved(capture): assert data["samples"][0]["cache_pct"] == 30 assert data["samples"][0]["workers"][0]["total_blocks"] is None assert data["samples"][0]["workers"][0]["kv_cache_capacity_bytes"] == 800 + # A grouped prefix-cached rank keeps cached blocks that no request holds. They stay resident, and an + # absent field counts as none. + row["decode_scheduler_metrics"][1]["kv_cache_cached_bytes"] = 100 + data = build_replay_dashboard_data({}, capture) + assert data["samples"][0]["cache_pct"] == 40 + assert data["samples"][0]["workers"][1]["kv_cache_cached_bytes"] == 100 row["decode_scheduler_metrics"].append( {"worker_id": 1, "dp_rank": 0, "total_blocks": 4, "physical_cache_usage": 0.5} )