From fc32cb093b252fc6d7cccdae10ca2f34d8af7f4d Mon Sep 17 00:00:00 2001 From: Yiming Liu Date: Sat, 7 Nov 2026 01:49:44 +0800 Subject: [PATCH 1/4] feat: simulate grouped prefix caching in the vLLM scheduler by vLLM version A grouped KV cache rank can now run with enable_prefix_caching on. The rank names the vLLM release it follows in the new optional field vllm_version (0.31.0 or 0.32.0), so that the grouped rule table of the coordinator applies the matching hit, split and retention rules. The field is a string, so the gate can name a wrong value in its error. The configuration gate replaces the old refusal of prefix caching with groups. With groups and prefix caching on it requires prefix_cacheable on every group, one page size, prefix_match_unit, a supported vllm_version, a rule table that accepts the groups, and a pool that holds the full prompt of max_model_len plus the state blocks of one request, summed per group as vLLM's startup check does. The rule table refuses a convolution group, so a rank with one runs only with prefix caching off. A rank without max_model_len bounds no prompt, so the capacity rule does not apply there. vllm_version without groups and prefix_match_unit without state_cache or groups are refused. A grouped rank with prefix_match_unit and with prefix caching off is refused too, because only prefix caching reads the unit. The scheduler builds the grouped pool in block mode when prefix caching is on and keeps byte mode otherwise. Admission searches the hit once per waiting candidate, stores it on the lease and tests the whole prompt in blocks, where inactive cached blocks stay evictable. Without chunked prefill, admission compares the tokens after the hit with the step budget, as vLLM does. The allocation then tests the blocks of the first chunk without the window cap, as vLLM's allocate_slots does, before it pins the hit and takes the fresh blocks. A waiting request that fails either test waits and preempts no running request, as on vLLM's waiting path, and its hit blocks keep their place in the LRU. Only a prompt whose blocks do not fit the empty pool is rejected. A hit inside a block can need more blocks than the pool holds, because its copy and inactive blocks count. Such a request waits too, because vLLM searches the hit again at the next step and running requests can evict the hit blocks. With no running request, vLLM waits forever and the replay stops with its livelock error. The align split comes from the rule table of the version. Finalization publishes the cached blocks and state keys of the step. Preemption and completion release the request's blocks tail first and a preempted request searches its hit again. Block mode reports kv_cache_cached_bytes next to the used bytes, counts the cached blocks as inactive blocks and includes them in the physical cache usage. The replay dashboard adds these cached bytes to the used bytes of a grouped rank, so that its resident cache keeps the retained prefixes, as the physical cache usage of a linear rank does. Every path with prefix caching off, and every linear path, is unchanged. The byte-mode API of the grouped pool now asserts its mode in debug builds, the dead-code allowances of the two grouped modules go, and the coordinator tests move to grouped_prefix_tests.rs. The prefix_cacheable doc comment states the rule the gate enforces, and the block_size_tokens doc comments say that grouped prefix caching reads the block of a state group. vllm_backend.rs, policy.rs and config.rs adapt rules of vLLM 0.31.0 and 0.32.0, so their headers name both revisions, and the notice entry names the rule that each one adapts. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yiming Liu --- THIRD_PARTY_NOTICES.md | 21 +- crates/core/src/engine/common/protocols.rs | 4 + crates/core/src/engine/config.rs | 315 ++- .../core/src/engine/kv_manager/g1_manager.rs | 29 + crates/core/src/engine/kv_manager/grouped.rs | 68 +- .../src/engine/kv_manager/grouped_prefix.rs | 1857 ++--------------- .../engine/kv_manager/grouped_prefix_tests.rs | 1795 ++++++++++++++++ crates/core/src/engine/kv_manager/mod.rs | 3 + .../src/engine/kv_manager/vllm_backend.rs | 256 ++- crates/core/src/engine/protocol.rs | 3 + crates/core/src/engine/scheduler/mod.rs | 4 + crates/core/src/engine/scheduler/rank.rs | 2 + crates/core/src/engine/scheduler/vllm/core.rs | 145 +- .../engine/scheduler/vllm/grouped_tests.rs | 716 ++++++- .../core/src/engine/scheduler/vllm/policy.rs | 56 +- crates/core/src/perfmodel/fpm/resources.rs | 12 +- crates/core/src/replay/components/engine.rs | 1 + crates/core/src/replay/telemetry.rs | 10 +- crates/core/tests/runtime_contract.rs | 110 + docs/getting-started/understand-results.md | 2 +- python/aisimulate/THIRD_PARTY_NOTICES.md | 21 +- .../src/aisimulate/replay/dashboard.js | 7 +- .../src/aisimulate/replay/dashboard.py | 7 +- .../src/aisimulate_core/fpm_profile.py | 5 +- tests/browser/replay_dashboard.cjs | 6 +- tests/test_replay_dashboard.py | 6 + 26 files changed, 3642 insertions(+), 1819 deletions(-) create mode 100644 crates/core/src/engine/kv_manager/grouped_prefix_tests.rs diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 948ddbf54a..7e92b962c9 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,18 @@ 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. - 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..6cf80f08a4 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)?; + // Without a model length the rank bounds no prompt, so there is no + // full prompt to size the pool for. + if let Some(max_model_len) = self.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,219 @@ 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")); + + // 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..515c1a42e5 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 { @@ -478,6 +481,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 +589,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 +785,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 +816,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 +855,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 +888,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 +1349,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..544b4e3105 --- /dev/null +++ b/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs @@ -0,0 +1,1795 @@ +// 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 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..7304920d19 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.unwrap_or(usize::MAX), + }, + ); + } 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..7f2f01eec8 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, 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..7e92b962c9 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,18 @@ 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. - 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/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_core/fpm_profile.py b/python/aisimulate/src/aisimulate_core/fpm_profile.py index 3864973e4d..c91500bfe8 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`` 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_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} ) From 7c16b83c89cdcefd8b0ff6e9943b8bba7adc21ff Mon Sep 17 00:00:00 2001 From: Yiming Liu Date: Sat, 7 Nov 2026 01:51:45 +0800 Subject: [PATCH 2/4] fix: require max_model_len for grouped prefix caching The grouped prefix-caching gate refused nothing for a rank without max_model_len and skipped its capacity rule, and the scheduler then ran the sliding-window admission cap without a model length. vLLM always sizes the pool with a model length, and unsupported input fails loudly in this repository, so the gate now refuses such a rank with "kv_cache_groups with enable_prefix_caching=true require max_model_len" and the scheduler always passes the real length to the admission limits. Prefix caching off and the linear path are unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yiming Liu --- crates/core/src/engine/config.rs | 33 ++++++++++--------- crates/core/src/engine/scheduler/vllm/core.rs | 2 +- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/crates/core/src/engine/config.rs b/crates/core/src/engine/config.rs index 6cf80f08a4..e2a655e8b7 100644 --- a/crates/core/src/engine/config.rs +++ b/crates/core/src/engine/config.rs @@ -1221,21 +1221,21 @@ impl EngineConfig { ); }; let rules = GroupedPrefixRules::new(groups, unit, version)?; - // Without a model length the rank bounds no prompt, so there is no - // full prompt to size the pool for. - if let Some(max_model_len) = self.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" - ); - } + // 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(()) } @@ -2201,6 +2201,9 @@ mod tests { 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); diff --git a/crates/core/src/engine/scheduler/vllm/core.rs b/crates/core/src/engine/scheduler/vllm/core.rs index 7304920d19..dab8a3a507 100644 --- a/crates/core/src/engine/scheduler/vllm/core.rs +++ b/crates/core/src/engine/scheduler/vllm/core.rs @@ -971,7 +971,7 @@ impl VllmCore { 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.unwrap_or(usize::MAX), + max_model_len: args.max_model_len.expect("validated max_model_len"), }, ); } else { From c37e974653c35d077eb21a517d92464db2fdd8eb Mon Sep 17 00:00:00 2001 From: Yiming Liu Date: Sat, 7 Nov 2026 02:50:04 +0800 Subject: [PATCH 3/4] feat: plan grouped prefix caching from runtime probe evidence and keep validate-fpm cold A grouped FPM profile could only replay cold. The Rust scheduler now models grouped prefix caching for vLLM 0.31.0 and 0.32.0, so Python opens its gate when the profile carries the evidence of a runtime probe that ran with prefix caching and recorded its hash block. One helper in config/engine.py, grouped_prefix_caching_requirements, holds the rule for the prediction, the sweeper and the onboarding plan. It keeps the old cold-replay text for a profile without that evidence. It refuses a group without the prefix_cacheable marker, a convolution group, a prefix-cacheable state group outside mamba_cache_mode=align, a vLLM version other than 0.31.0 and 0.32.0, unequal pages, and a prefix_match_unit other than the probe's hash block. For GLM-5.3-Flash it refuses a hash block or a budget off the 4-token grid of the indexer cache. It returns the hash block that the rank uses. The convolution refusal uses the text of the Rust rule table, because neither side models a convolution group with prefix caching. The compiler gives a grouped rank with prefix caching that hash block as prefix_match_unit and the backend version as vllm_version. It already writes max_model_len for every vLLM rank, which the Rust gate needs. The sweeper builds its candidate ranks without the compiler, so sweeper/deploy.py writes the same three keys and checks each candidate with its own budget before replay. Without that, every candidate of such a profile failed in the Rust gate with "require prefix_match_unit". A grouped deployment with prefix caching may set prefix_match_unit without state_cache. The engine validation checks it, because the cache configuration does not see the profile. Every other case keeps the text "prefix_match_unit currently requires state_cache". The onboarding plan turns prefix caching on only when the deployment passes grouped_prefix_caching_requirements, so the plan and the gate share every check. Otherwise it plans a cold replay, as before, and a profile that fails one requirement does not fail planning. validate-fpm keeps its replay cold, because it tests FPM timing. It also drops an explicit prefix_match_unit from that replay, because only prefix caching reads the unit on a grouped rank. The test profiles with prefix-caching evidence leave out the fictional convolution group of the grouped fixture. The doc comments of the probe's three prefix-caching fields, in Rust and in Python, now say that grouped prefix caching needs them and that replay hashes prompts in blocks of hash_block_size_tokens. With prefix caching off, the payloads and the results do not change. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yiming Liu --- crates/core/src/perfmodel/fpm/resources.rs | 6 +- python/aisimulate/src/aisimulate/compiler.py | 8 + .../src/aisimulate/config/engine.py | 118 +++++- .../aisimulate/src/aisimulate/support/plan.py | 20 +- .../src/aisimulate/support/validation.py | 7 + .../src/aisimulate/sweeper/config.py | 18 +- .../src/aisimulate/sweeper/deploy.py | 39 ++ .../src/aisimulate_core/fpm_profile.py | 6 +- .../tests/unit/test_onboard_runtime.py | 33 ++ tests/test_fpm_grouped_workflow.py | 386 +++++++++++++++++- 10 files changed, 602 insertions(+), 39 deletions(-) diff --git a/crates/core/src/perfmodel/fpm/resources.rs b/crates/core/src/perfmodel/fpm/resources.rs index 7f2f01eec8..7a7e97f98d 100644 --- a/crates/core/src/perfmodel/fpm/resources.rs +++ b/crates/core/src/perfmodel/fpm/resources.rs @@ -238,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/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/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 c91500bfe8..608d5d4ef0 100644 --- a/python/aisimulate/src/aisimulate_core/fpm_profile.py +++ b/python/aisimulate/src/aisimulate_core/fpm_profile.py @@ -127,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/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") From d23b71f3fb9f7d8edfefea378a15592a35a4d647 Mon Sep 17 00:00:00 2001 From: Yiming Liu Date: Sat, 7 Nov 2026 04:33:42 +0800 Subject: [PATCH 4/4] fix: apply the Mamba align split only to ranks with a Mamba state group vLLM runs the align-mode prefill split only when the cache holds a MambaSpec layer in align mode (v0.31 sched:354-356, :695-697, :1127-1129, v0.32 sched:377-379, :724-726, :1173-1175). The grouped scheduler asked the rule table for a split on every rank with grouped prefix rules (core.rs:3026-3035). A rank without a Mamba group, such as full attention with a sliding window, then got prefill chunks that vLLM never makes: a stop at the last block boundary, a floor to the block, a stop at the next boundary and a stop at the junction. align_split now returns the desired tokens unless the rules hold a State rule. A State rule is a prefix-cacheable Mamba group, which the Python gate admits only with mamba_cache_mode align (engine.py:1348-1351). The GLM tail ring is a KpoolTailSpec or a CircularBufferSpec, not a MambaSpec, and is a Ring rule. GLM keeps its split and its schedules. The new test runs on both rule rows, with one full-attention group and with full attention plus a sliding window. The single-group parity test of the next commit found the defect. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yiming Liu --- THIRD_PARTY_NOTICES.md | 3 ++ .../src/engine/kv_manager/grouped_prefix.rs | 17 ++++++++++- .../engine/kv_manager/grouped_prefix_tests.rs | 30 +++++++++++++++++++ python/aisimulate/THIRD_PARTY_NOTICES.md | 3 ++ 4 files changed, 52 insertions(+), 1 deletion(-) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 7e92b962c9..72d166d46f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -495,6 +495,9 @@ not fit, whose hit is searched again at the next step. 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/kv_manager/grouped_prefix.rs b/crates/core/src/engine/kv_manager/grouped_prefix.rs index 515c1a42e5..e6f673d141 100644 --- a/crates/core/src/engine/kv_manager/grouped_prefix.rs +++ b/crates/core/src/engine/kv_manager/grouped_prefix.rs @@ -387,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, @@ -397,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; diff --git a/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs b/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs index 544b4e3105..f7055f11fa 100644 --- a/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs +++ b/crates/core/src/engine/kv_manager/grouped_prefix_tests.rs @@ -1123,6 +1123,36 @@ fn align_split_with_a_checkpoint_block_ports_vllm_0320() { 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, diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 7e92b962c9..72d166d46f 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -495,6 +495,9 @@ not fit, whose hit is searched again at the next step. 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