feat: price grouped KV handoff transfers by vLLM group rules

THIRD_PARTY_NOTICES.md

@@ -498,6 +498,8 @@ 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.
+`crates/core/src/engine/scheduler/vllm/core.rs` and `vllm_backend.rs` adapt the
+release of the skipped blocks of a completed source.
 
 - 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
@@ -532,6 +534,44 @@ 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_transfer.rs`
+adapts the P/D transfer rules of vLLM's NIXL connector and KV cache manager
+from vLLM v0.31.0, immutable commit `db9527a46873454610df6dbedf79a36d6bf1a7f6`,
+and from vLLM v0.32.0, immutable commit
+`5df63a7203cdb9328d824a27b682e77c7b9928be`: the block clip of each cache
+group, the whole-block reads of a full-attention group, the in-window block
+count of a decode rank for a sliding-window group, the prefill backoff and
+truncation, the blocks that a decode rank pulls, the state pairing and the
+copy on each TP rank. The payload is equal in both versions, so one set of
+rules covers both.
+
+- https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+- https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+- https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+- https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py
+- https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/v1/core/kv_cache_manager.py
+- https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/v1/core/single_type_kv_cache_manager.py
+- https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm/v1/kv_cache_interface.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/distributed/kv_transfer/kv_connector/utils.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/v1/core/kv_cache_manager.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/v1/core/single_type_kv_cache_manager.py
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/v1/kv_cache_interface.py
+
+Copyright contributors to the vLLM project. Licensed under Apache-2.0:
+https://github.com/vllm-project/vllm/blob/db9527a46873454610df6dbedf79a36d6bf1a7f6/LICENSE
+https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/LICENSE
+
+This is a modified Rust simulation adaptation, restricted to no speculative
+decoding, no HiSparse, no bounded replay window and no `all` Mamba cache mode.
+No upstream method body is copied. The identified upstream revisions have no
+root NOTICE file. The repository's Apache-2.0 LICENSE supplies the license
+text, and no upstream source file is vendored.
+
+
 The inference-mode scope and MSA query-position metadata integration in
 `collector/vllm/collect_mla_module.py` and
 `collector/vllm/collect_msa_module.py` are adapted (modified) from serving

crates/core/src/engine/config.rs

@@ -19,7 +19,9 @@ 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::{GroupedPrefixRules, SourceReleaseOrder, VllmRuleVersion};
+use crate::engine::kv_manager::{
+    GroupedPrefixRules, GroupedTransferRules, SourceReleaseOrder, VllmRuleVersion,
+};
 use crate::engine::timing::{TimingModel, TimingModelConfig, built_in_timing_model};
 use crate::perfmodel::FpmCacheGroup;
 
@@ -922,8 +924,9 @@ pub struct EngineConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub kv_cache_capacity_bytes: Option<u64>,
     /// 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`.
+    /// `kv_cache_groups` and `enable_prefix_caching` must set it. A grouped
+    /// prefill or decode rank must set it too. Accepted values: `0.31.0` and
+    /// `0.32.0`.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub vllm_version: Option<String>,
     /// Optional framework-native host-offload simulation.
@@ -1183,6 +1186,34 @@ impl EngineConfig {
         }
     }
 
+    /// The gate of a grouped prefill or decode rank. `VllmCore` construction
+    /// builds `GroupedTransferRules` for such a rank to price its handoff. The
+    /// rules follow vLLM 0.31.0 and 0.32.0, read the marker and the transfer
+    /// bytes of every group, and count the handoff in full-attention blocks of
+    /// one size.
+    fn validate_grouped_transfer(&self) -> Result<()> {
+        ensure!(
+            self.vllm_version
+                .as_deref()
+                .and_then(VllmRuleVersion::parse)
+                .is_some(),
+            "kv_cache_groups with worker_type=prefill or decode require vllm_version 0.31.0 or 0.32.0"
+        );
+        ensure!(
+            self.kv_cache_groups.iter().all(|group| {
+                group.prefix_cacheable.is_some() && group.transfer_bytes_per_block.is_some()
+            }),
+            "kv_cache_groups with worker_type=prefill or decode require every group to declare prefix_cacheable and transfer_bytes_per_block; re-import the runtime probe"
+        );
+        // The groups passed `validate` and declare both fields, so the rules
+        // can fail only on two different full-attention blocks.
+        ensure!(
+            GroupedTransferRules::new(&self.kv_cache_groups).is_ok(),
+            "kv_cache_groups with worker_type=prefill or decode require one block_size_tokens across full-attention groups"
+        );
+        Ok(())
+    }
+
     /// 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
@@ -1430,8 +1461,8 @@ impl EngineConfig {
             );
         } else {
             ensure!(
-                self.backend == Backend::Vllm && self.worker_type == WorkerType::Aggregated,
-                "kv_cache_groups currently require backend=vllm and worker_type=aggregated"
+                self.backend == Backend::Vllm,
+                "kv_cache_groups currently require backend=vllm"
             );
             ensure!(
                 self.native_host_offload.is_none() && self.g3_offload.is_none(),
@@ -1447,8 +1478,7 @@ impl EngineConfig {
             );
             ensure!(
                 self.kv_cache_bytes_per_token.is_none()
-                    && self.kv_transfer_bytes_per_token.is_none()
-                    && self.kv_transfer_bandwidth.is_none(),
+                    && self.kv_transfer_bytes_per_token.is_none(),
                 "kv_cache_groups cannot use linear cache or KV transfer byte geometry"
             );
             ensure!(
@@ -1461,6 +1491,9 @@ impl EngineConfig {
                 group.validate()?;
                 ensure!(names.insert(&group.name), "duplicate KV cache group name");
             }
+            if self.worker_type != WorkerType::Aggregated {
+                self.validate_grouped_transfer()?;
+            }
             if self.enable_prefix_caching {
                 self.validate_grouped_prefix_caching()?;
             }
@@ -1570,12 +1603,14 @@ impl EngineConfig {
         );
         // Only the prefill source prices the handoff, so a positive bandwidth
         // needs payload bytes rather than the topology's fallback latency.
+        // Cache groups carry their own transfer bytes.
         ensure!(
             self.worker_type != WorkerType::Prefill
                 || self
                     .kv_transfer_bandwidth
                     .is_none_or(|bandwidth| bandwidth <= 0.0)
-                || self.kv_transfer_bytes_per_token.is_some(),
+                || self.kv_transfer_bytes_per_token.is_some()
+                || !self.kv_cache_groups.is_empty(),
             "kv_transfer_bandwidth requires kv_transfer_bytes_per_token on prefill workers"
         );
         match &self.timing_model {
@@ -1974,7 +2009,7 @@ mod tests {
             (
                 "worker_type",
                 serde_json::json!("decode"),
-                "worker_type=aggregated",
+                "kv_cache_groups with worker_type=prefill or decode require vllm_version 0.31.0 or 0.32.0",
             ),
             (
                 "enable_prefix_caching",
@@ -2312,6 +2347,139 @@ mod tests {
         validates(wire(false, "", 8));
     }
 
+    /// `glm_groups_wire` with the GLM-5.3-Flash TP 4 transfer bytes:
+    /// `11 * 1_114_112 + 11 * 71_808` for the MLA latent and the indexer keys,
+    /// `11 * 2_048` for the tail ring, and `9 * 1_085_440` or `8 * 1_085_440`
+    /// for a KDA state. The bytes count the layers of the model, not the
+    /// `num_layers` of this wire, which the gate does not read.
+    fn glm_transfer_groups_wire() -> serde_json::Value {
+        let mut groups = glm_groups_wire(None);
+        let bytes = [
+            13_045_120, 22_528, 9_768_960, 9_768_960, 8_683_520, 8_683_520,
+        ];
+        for (group, bytes) in groups.as_array_mut().unwrap().iter_mut().zip(bytes) {
+            group["transfer_bytes_per_block"] = bytes.into();
+        }
+        groups
+    }
+
+    /// A GLM-shaped rank of `role` with prefix caching off. The groups give
+    /// the handoff payload, so it sets a bandwidth and no
+    /// `kv_transfer_bytes_per_token`.
+    fn glm_pd_wire(role: &str) -> serde_json::Value {
+        serde_json::json!({
+            "backend": "vllm",
+            "worker_type": role,
+            "enable_prefix_caching": false,
+            "vllm_version": "0.31.0",
+            "kv_cache_capacity_bytes": 128 * GLM_PAGE,
+            "kv_cache_groups": glm_transfer_groups_wire(),
+            "kv_transfer_bandwidth": 50.0,
+        })
+    }
+
+    #[test]
+    fn grouped_pd_configuration_validates_its_transfer_geometry() {
+        let version = "kv_cache_groups with worker_type=prefill or decode require vllm_version 0.31.0 or 0.32.0";
+        let declare = "kv_cache_groups with worker_type=prefill or decode require every group to declare prefix_cacheable and transfer_bytes_per_block; re-import the runtime probe";
+        let one_block = "kv_cache_groups with worker_type=prefill or decode require one block_size_tokens across full-attention groups";
+        let byte_geometry = "kv_cache_groups cannot use linear cache or KV transfer byte geometry";
+        type Refusal = (fn(&mut serde_json::Value), &'static str);
+        let refusals: [Refusal; 10] = [
+            (
+                |wire| {
+                    wire.as_object_mut().unwrap().remove("vllm_version");
+                },
+                version,
+            ),
+            (|wire| wire["vllm_version"] = "0.30.0".into(), version),
+            // With prefix caching on, the rule of the role fires before the
+            // version rule of prefix caching.
+            (
+                |wire| {
+                    wire["enable_prefix_caching"] = true.into();
+                    wire["prefix_match_unit"] = 2176.into();
+                    wire["max_model_len"] = 131_072.into();
+                    wire.as_object_mut().unwrap().remove("vllm_version");
+                },
+                version,
+            ),
+            // The per-group check runs first, so a malformed group keeps its
+            // own text.
+            (
+                |wire| {
+                    wire["kv_cache_groups"][0]["transfer_bytes_per_block"] = (GLM_PAGE + 1).into()
+                },
+                "invalid engine config: transfer_bytes_per_block must be positive and at most page_size_bytes",
+            ),
+            (
+                |wire| {
+                    wire["kv_cache_groups"][1]
+                        .as_object_mut()
+                        .unwrap()
+                        .remove("transfer_bytes_per_block");
+                },
+                declare,
+            ),
+            (
+                |wire| {
+                    wire["kv_cache_groups"][2]
+                        .as_object_mut()
+                        .unwrap()
+                        .remove("prefix_cacheable");
+                },
+                declare,
+            ),
+            // The handoff timing counts in one full-attention block.
+            (
+                |wire| {
+                    wire["kv_cache_groups"]
+                        .as_array_mut()
+                        .unwrap()
+                        .push(serde_json::json!({
+                            "name": "mla_64", "kind": "attention", "num_layers": 1,
+                            "block_size_tokens": 64, "page_size_bytes": GLM_PAGE,
+                            "prefix_cacheable": true, "transfer_bytes_per_block": GLM_PAGE
+                        }));
+                },
+                one_block,
+            ),
+            (
+                |wire| wire["kv_transfer_bytes_per_token"] = 5995.into(),
+                byte_geometry,
+            ),
+            (
+                |wire| wire["kv_cache_bytes_per_token"] = 1.into(),
+                byte_geometry,
+            ),
+            (
+                |wire| wire["backend"] = "sglang".into(),
+                "kv_cache_groups currently require backend=vllm",
+            ),
+        ];
+        for role in ["prefill", "decode"] {
+            // A positive bandwidth needs no kv_transfer_bytes_per_token.
+            validates(glm_pd_wire(role));
+            let mut v32 = glm_pd_wire(role);
+            v32["vllm_version"] = "0.32.0".into();
+            validates(v32);
+            let mut prefix_caching = glm_pd_wire(role);
+            prefix_caching["enable_prefix_caching"] = true.into();
+            prefix_caching["prefix_match_unit"] = 2176.into();
+            prefix_caching["max_model_len"] = 131_072.into();
+            validates(prefix_caching);
+            for (index, (mutate, expected)) in refusals.iter().enumerate() {
+                let mut wire = glm_pd_wire(role);
+                mutate(&mut wire);
+                assert_eq!(validation_error(wire), *expected, "{role} case {index}");
+            }
+        }
+        // A bandwidth is harmless without a prefill or decode role.
+        let mut aggregated = glm_pd_wire("aggregated");
+        aggregated["kv_transfer_bandwidth"] = 1.0.into();
+        validates(aggregated);
+    }
+
     #[test]
     fn native_host_offload_deserializes_with_default_bandwidths() {
         assert_eq!(DEFAULT_HOST_OFFLOAD_BANDWIDTH_GBPS, 32.0);

crates/core/src/engine/handoff.rs

@@ -61,6 +61,11 @@ pub enum TransferTimingMode {
 /// Schedulers publish this value when source KV ownership becomes held.
 /// The Replayer combines it with the destination's missing-token observation
 /// and schedules the virtual transfer completion.
+///
+/// A grouped KV source counts in units of one full-attention block.
+/// `full_prompt_tokens` is then the number of such blocks of the transferred
+/// prompt. `kv_bytes_per_token` is the payload of one block, summed over the
+/// full-attention groups. `state_bytes` is every fixed byte of the request.
 #[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
 pub struct HandoffTransferTiming {
     /// Which prompt footprint contributes to transfer time.

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

@@ -13,6 +13,7 @@ use crate::engine::common::protocols::{KvEventPublishers, PrefillCost};
 use crate::engine::common::sequence::RequestSequence;
 
 use super::grouped_prefix::{AdmissionLimits, GroupedPrefixRules};
+use super::grouped_transfer::GroupedTransferRules;
 use super::vllm_backend::{
     BlockRequestLease, DecodeBlockReservation as VllmDecodeBlockReservation,
     DestinationReservation as VllmDestinationReservation,
@@ -92,6 +93,15 @@ impl G1Manager {
         self.inner.grouped_prefix_rules()
     }
 
+    pub(crate) fn set_grouped_transfer_rules(&mut self, rules: GroupedTransferRules) {
+        self.inner.set_grouped_transfer_rules(rules);
+    }
+
+    /// The transfer rules of a grouped cache on a prefill or decode rank.
+    pub(crate) fn grouped_transfer(&self) -> Option<&GroupedTransferRules> {
+        self.inner.grouped_transfer()
+    }
+
     /// 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)
@@ -272,6 +282,16 @@ impl G1Manager {
         self.inner.preempt_lease(owner, lease);
     }
 
+    /// Free the grouped blocks that a completed source no longer needs.
+    pub(crate) fn release_skipped_blocks(
+        &mut self,
+        owner: Uuid,
+        lease: &mut BlockRequestLease,
+        computed: usize,
+    ) {
+        self.inner.release_skipped_blocks(owner, lease, computed);
+    }
+
     pub(crate) fn finish_native(&mut self, owner: Uuid, lease: BlockRequestLease) {
         self.inner.finish_lease(owner, lease);
     }