A high-throughput and memory-efficient inference and serving engine for LLMs
[PD][Core] Fix Mamba prefix cache hit rate in PD disaggregation (#44243)
Co-authored-by: lHrHenry233 <2381623149@qq.com> Co-authored-by: underfituu <hzhucong@163.com> Signed-off-by: Zhanqiu Hu <zhu@redhat.com>
zhanqiuhu committed
Jun 11, 2026 at 10:10 UTC
55911db5802d4fa782cf8f19b2842597a36f5814
5 files changed
+534
-3
.buildkite/test_areas/disaggregated.yaml
+14
@@ -61,6 +61,20 @@ steps:
61
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
62
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
63
64
+- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
65
+ key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
66
+ timeout_in_minutes: 25
67
+ working_dir: "/vllm-workspace/tests"
68
+ num_devices: 2
69
+ source_file_dependencies:
70
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
71
+ - vllm/v1/core/sched/
72
+ - vllm/v1/core/kv_cache_coordinator.py
73
+ - tests/v1/kv_connector/nixl_integration/
74
+ commands:
75
+ - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
76
+ - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
77
+
78
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
79
key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
80
timeout_in_minutes: 30
tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
new
+96
@@ -0,0 +1,96 @@
1
+#!/bin/bash
2
+set -xe
3
+
4
+# E2E test: Mamba hybrid prefix cache hits in PD disaggregation.
5
+# Spins up a 1P1D setup with a Mamba hybrid model and verifies
6
+# repeated prompts yield non-zero D-side prefix cache hits.
7
+
8
+PREFILL_GPU_ID=${PREFILL_GPU_ID:-0}
9
+DECODE_GPU_ID=${DECODE_GPU_ID:-1}
10
+MODEL=${MODEL:-"ibm-granite/granite-4.0-h-tiny"}
11
+GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8}
12
+
13
+echo "Running Mamba prefix cache test (GPUs: P=$PREFILL_GPU_ID, D=$DECODE_GPU_ID, model=$MODEL)"
14
+
15
+KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
16
+
17
+# Resolve repository root
18
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
19
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
20
+
21
+trap 'kill $(jobs -pr) 2>/dev/null' SIGINT SIGTERM EXIT
22
+
23
+wait_for_server() {
24
+ local port=$1
25
+ timeout 600 bash -c "
26
+ until curl -s localhost:${port}/v1/completions > /dev/null; do
27
+ sleep 1
28
+ done" && return 0 || return 1
29
+}
30
+
31
+cleanup_instances() {
32
+ echo "Cleaning up any running vLLM instances..."
33
+ pkill -f "vllm serve" || true
34
+ sleep 2
35
+}
36
+
37
+cleanup_instances
38
+
39
+# Start prefill instance
40
+PREFILL_PORT=8001
41
+CUDA_VISIBLE_DEVICES=$PREFILL_GPU_ID \
42
+VLLM_SSM_CONV_STATE_LAYOUT=DS \
43
+VLLM_KV_CACHE_LAYOUT=HND \
44
+VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
45
+vllm serve $MODEL \
46
+ --port $PREFILL_PORT \
47
+ --enforce-eager \
48
+ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
49
+ --max-model-len 16384 \
50
+ --block-size 128 \
51
+ --trust-remote-code \
52
+ --enable-prefix-caching \
53
+ --mamba-cache-mode all \
54
+ --kv-transfer-config "$KV_CONFIG" &
55
+
56
+# Start decode instance
57
+DECODE_PORT=8002
58
+CUDA_VISIBLE_DEVICES=$DECODE_GPU_ID \
59
+VLLM_SSM_CONV_STATE_LAYOUT=DS \
60
+VLLM_KV_CACHE_LAYOUT=HND \
61
+VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
62
+vllm serve $MODEL \
63
+ --port $DECODE_PORT \
64
+ --enforce-eager \
65
+ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
66
+ --max-model-len 16384 \
67
+ --block-size 128 \
68
+ --trust-remote-code \
69
+ --enable-prefix-caching \
70
+ --mamba-cache-mode all \
71
+ --kv-transfer-config "$KV_CONFIG" &
72
+
73
+echo "Waiting for prefill instance on port $PREFILL_PORT..."
74
+wait_for_server "$PREFILL_PORT"
75
+echo "Waiting for decode instance on port $DECODE_PORT..."
76
+wait_for_server "$DECODE_PORT"
77
+
78
+# Start proxy
79
+PROXY_PORT=8192
80
+python3 "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
81
+ --port $PROXY_PORT \
82
+ --prefiller-ports $PREFILL_PORT \
83
+ --decoder-ports $DECODE_PORT &
84
+
85
+sleep 5
86
+
87
+echo "Running Mamba prefix cache test..."
88
+PREFILL_PORT=$PREFILL_PORT \
89
+DECODE_PORT=$DECODE_PORT \
90
+PROXY_PORT=$PROXY_PORT \
91
+python3 -m pytest -s -v \
92
+ "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_mamba_prefix_cache.py"
93
+
94
+echo "Mamba prefix cache test passed!"
95
+
96
+cleanup_instances
tests/v1/kv_connector/nixl_integration/test_mamba_prefix_cache.py
new
+346
@@ -0,0 +1,346 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Verify D-side prefix cache hits reduce transfer for Mamba hybrid PD.
4
+
5
+Sends the same long prompt twice through P/D and asserts that the second
6
+request transfers fewer bytes (because cached blocks are skipped).
7
+"""
8
+
9
+import os
10
+import time
11
+
12
+import openai
13
+import regex as re
14
+import requests
15
+
16
+PREFILL_HOST = os.getenv("PREFILL_HOST", "localhost")
17
+PREFILL_PORT = os.environ["PREFILL_PORT"]
18
+DECODE_HOST = os.getenv("DECODE_HOST", "localhost")
19
+DECODE_PORT = os.environ["DECODE_PORT"]
20
+PROXY_HOST = os.getenv("PROXY_HOST", "localhost")
21
+PROXY_PORT = os.environ["PROXY_PORT"]
22
+
23
+# Long prompt (~9000 tokens) to span many blocks so prefix caching kicks in.
24
+_BASE_PROMPT = """\
25
+The following is a comprehensive overview of distributed systems, covering \
26
+their history, design principles, and modern applications.
27
+
28
+Distributed systems emerged from the need to connect multiple computers to \
29
+work together on shared tasks. In the 1960s, ARPANET demonstrated that \
30
+geographically dispersed machines could communicate through packet switching. \
31
+This laid the groundwork for decades of research into fault tolerance, \
32
+consistency, and performance.
33
+
34
+Leslie Lamport's 1978 paper on logical clocks introduced the concept of \
35
+causal ordering in distributed systems. His later work on the Paxos algorithm \
36
+provided a practical solution to the consensus problem, enabling multiple \
37
+nodes to agree on a single value despite failures. The Byzantine Generals \
38
+Problem, also formulated by Lamport, addressed the challenge of reaching \
39
+agreement when some participants may be malicious.
40
+
41
+The CAP theorem, proposed by Eric Brewer in 2000 and formally proved by Seth \
42
+Gilbert and Nancy Lynch in 2002, states that a distributed system cannot \
43
+simultaneously provide Consistency, Availability, and Partition tolerance. \
44
+This fundamental trade-off has guided the design of distributed databases and \
45
+storage systems ever since. Systems like Google's Bigtable chose consistency \
46
+and partition tolerance, while Amazon's Dynamo prioritized availability and \
47
+partition tolerance.
48
+
49
+Google's MapReduce framework, published in 2004, popularized the concept of \
50
+processing large datasets across clusters of commodity hardware. The \
51
+programming model was simple: users specified a map function to process \
52
+key-value pairs and a reduce function to merge intermediate values. The \
53
+framework handled distribution, fault tolerance, and load balancing \
54
+automatically. This inspired the open-source Hadoop ecosystem, which became \
55
+the foundation for big data processing throughout the 2010s.
56
+
57
+The Google File System (GFS) and its open-source counterpart HDFS provided \
58
+the distributed storage layer beneath MapReduce. These systems replicated \
59
+data across multiple nodes, using a single master for metadata management \
60
+and chunk servers for actual data storage. The master maintained a mapping \
61
+from files to chunks and tracked which chunk servers held each replica.
62
+
63
+Apache Kafka, developed at LinkedIn and open-sourced in 2011, introduced a \
64
+distributed commit log that could handle millions of messages per second. \
65
+Its design separated producers from consumers through topic-based \
66
+publish-subscribe semantics. Partitioning allowed horizontal scaling, while \
67
+replication ensured durability. Kafka's exactly-once semantics, achieved \
68
+through idempotent producers and transactional writes, made it suitable for \
69
+financial and mission-critical applications.
70
+
71
+Raft, published by Diego Ongaro and John Ousterhout in 2014, provided an \
72
+understandable alternative to Paxos for consensus. Its key insight was \
73
+decomposing consensus into leader election, log replication, and safety. A \
74
+leader would be elected through randomized timeouts, then would replicate \
75
+its log entries to followers. Committed entries were guaranteed to be present \
76
+on a majority of servers. Raft's clarity led to its adoption in systems like \
77
+etcd, CockroachDB, and TiKV.
78
+
79
+Container orchestration systems like Kubernetes, released by Google in 2014, \
80
+brought distributed systems concepts to application deployment. Kubernetes \
81
+managed clusters of machines, scheduling containers across nodes while \
82
+maintaining desired state. Its control plane used etcd for consistent state \
83
+storage, an API server for client communication, a scheduler for placement \
84
+decisions, and controllers for reconciliation loops.
85
+
86
+Service meshes emerged to handle the networking complexity of microservices \
87
+architectures. Istio, Linkerd, and Envoy provided transparent proxying, load \
88
+balancing, circuit breaking, and observability without requiring application \
89
+code changes. They implemented the sidecar pattern, deploying a proxy \
90
+alongside each service instance to intercept all network traffic.
91
+
92
+Modern distributed databases like CockroachDB, TiDB, and YugabyteDB combine \
93
+the SQL interface that developers expect with the horizontal scalability of \
94
+NoSQL systems. They use Raft for consensus, multi-version concurrency control \
95
+for transactions, and range-based sharding for data distribution. These \
96
+systems can span multiple data centers while providing serializable isolation.
97
+
98
+Stream processing frameworks evolved from batch-oriented MapReduce to \
99
+real-time systems. Apache Flink provided exactly-once processing with \
100
+event-time semantics, handling out-of-order data through watermarks. Its \
101
+checkpoint mechanism, based on Chandy-Lamport distributed snapshots, allowed \
102
+recovery without data loss. Google's Dataflow model unified batch and \
103
+streaming under a single programming model.
104
+
105
+The rise of machine learning at scale introduced new distributed systems \
106
+challenges. Training large neural networks required distributing computation \
107
+across hundreds or thousands of GPUs. Data parallelism split batches across \
108
+workers, while model parallelism partitioned the network itself. Pipeline \
109
+parallelism overlapped computation stages to maximize utilization. \
110
+Ring-allreduce and parameter server architectures provided different \
111
+trade-offs for gradient synchronization.
112
+
113
+Inference serving systems like vLLM, TensorRT-LLM, and SGLang optimized the \
114
+deployment of large language models. They introduced techniques like \
115
+continuous batching to maximize GPU utilization, PagedAttention for efficient \
116
+KV cache memory management, and speculative decoding to reduce latency. \
117
+Prefill-decode disaggregation separated the compute-intensive prefill phase \
118
+from the memory-bound decode phase across different GPU pools.
119
+
120
+KV cache transfer in disaggregated serving requires careful coordination \
121
+between prefill and decode nodes. The prefill node computes the full KV cache \
122
+for a request's prompt and transfers it to the decode node via high-bandwidth \
123
+interconnects like NVLink, InfiniBand, or RDMA. The decode node then uses \
124
+this transferred cache to generate tokens autoregressively without \
125
+recomputing the prefix.
126
+
127
+Prefix caching optimizes this further by recognizing that multiple requests \
128
+often share common prefixes, such as system prompts or few-shot examples. \
129
+When a decode node receives a new request whose prefix matches a previously \
130
+transferred KV cache, it can skip the transfer for those shared blocks and \
131
+only fetch the new, unique portion. This dramatically reduces both network \
132
+bandwidth consumption and time-to-first-token latency.
133
+
134
+For hybrid architectures combining attention mechanisms with state-space \
135
+models like Mamba, prefix caching becomes more complex. Attention layers \
136
+maintain a KV cache that can be trivially split into independent blocks, \
137
+making prefix matching straightforward. However, Mamba layers maintain a \
138
+recurrent hidden state that represents the entire sequence history in a \
139
+single fixed-size tensor. This state cannot be meaningfully split into \
140
+prefix-aligned blocks the way attention KV caches can.
141
+
142
+The challenge in disaggregated serving of hybrid models is that the cache \
143
+coordination logic must handle these heterogeneous cache types simultaneously. \
144
+A naive approach that requires all cache groups to agree on a single prefix \
145
+hit length will always report zero hits for the Mamba group on a cold decode \
146
+node, dragging the entire prefix cache hit rate to zero even when the \
147
+attention layers have perfect cache hits.
148
+
149
+The solution is to evaluate each cache group independently, allowing the \
150
+attention groups to report their actual cache hits while the Mamba group \
151
+reports zero. The transfer logic then only fetches the blocks that each group \
152
+actually needs: for attention, only the new uncached blocks; for Mamba, \
153
+always the full state. This per-group evaluation preserves the prefix caching \
154
+benefits for attention layers while correctly handling the all-or-nothing \
155
+nature of Mamba state.
156
+
157
+Consistency models in distributed systems range from strong linearizability \
158
+to weak eventual consistency. Linearizability requires that operations appear \
159
+to occur atomically at some point between their invocation and response. \
160
+Sequential consistency relaxes this by only requiring that operations from \
161
+each process appear in program order. Causal consistency preserves causal \
162
+relationships between operations. Eventual consistency only guarantees that \
163
+all replicas will eventually converge to the same state.
164
+
165
+Vector clocks extend Lamport timestamps to capture causality precisely. Each \
166
+process maintains a vector of logical clocks, one per process in the system. \
167
+When a process performs a local event, it increments its own entry. When \
168
+sending a message, it attaches its current vector. Upon receiving a message, \
169
+a process takes the element-wise maximum of its vector and the received \
170
+vector, then increments its own entry. Two events are concurrent if and only \
171
+if neither vector dominates the other.
172
+
173
+Conflict-free replicated data types (CRDTs) provide eventual consistency \
174
+without coordination. They achieve this through mathematical properties: \
175
+either operations are commutative and idempotent (operation-based CRDTs), or \
176
+states form a join-semilattice where merging always produces a valid result \
177
+(state-based CRDTs). Examples include grow-only counters, positive-negative \
178
+counters, grow-only sets, observed-remove sets, and last-writer-wins registers.
179
+
180
+Distributed hash tables (DHTs) like Chord, Kademlia, and Pastry provide \
181
+decentralized key-value lookup. Chord arranges nodes on a circular identifier \
182
+space, using finger tables for O(log n) routing. Kademlia uses XOR distance \
183
+for routing, enabling parallel lookups and natural load balancing. These \
184
+systems underpin peer-to-peer networks, content distribution, and \
185
+decentralized storage.
186
+
187
+Leader election algorithms ensure that exactly one node acts as coordinator \
188
+at any time. The Bully algorithm selects the node with the highest \
189
+identifier. Ring-based algorithms pass election messages around a logical \
190
+ring. In practice, systems often use lease-based leadership where a leader \
191
+must periodically renew its lease, allowing automatic failover when a leader \
192
+becomes unresponsive.
193
+
194
+Distributed transactions spanning multiple partitions require coordination \
195
+protocols. Two-phase commit (2PC) provides atomicity but blocks if the \
196
+coordinator fails. Three-phase commit (3PC) adds a prepare-to-commit phase \
197
+to avoid blocking but does not handle network partitions. Saga patterns \
198
+decompose long-running transactions into compensable sub-transactions, \
199
+providing eventual consistency without global locks.
200
+
201
+Load balancing in distributed systems takes many forms. Round-robin \
202
+distributes requests evenly but ignores server capacity. Weighted round-robin \
203
+accounts for heterogeneous servers. Least-connections routes to the server \
204
+with fewest active requests. Consistent hashing minimizes redistribution when \
205
+servers join or leave. Power-of-two-choices selects the less loaded of two \
206
+randomly chosen servers, providing near-optimal balance with minimal \
207
+coordination.
208
+
209
+Observability in distributed systems requires correlated telemetry across \
210
+service boundaries. Distributed tracing, pioneered by Google's Dapper and \
211
+standardized through OpenTelemetry, propagates trace context through request \
212
+chains. Each service adds spans representing its processing, creating a tree \
213
+structure that reveals latency bottlenecks and error sources. Combined with \
214
+metrics and structured logs, traces provide the visibility needed to operate \
215
+complex distributed systems reliably."""
216
+
217
+# Pad to ~23000 chars (~9000 tokens) to fill many blocks.
218
+PROMPT = _BASE_PROMPT
219
+while len(PROMPT) < 23000:
220
+ n = len(PROMPT)
221
+ PROMPT += f" The value at position {n} is {n * 7 % 9973}."
222
+
223
+
224
+METRICS_OF_INTEREST = [
225
+ "vllm:nixl_bytes_transferred_sum",
226
+ "vllm:nixl_bytes_transferred_count",
227
+ "vllm:nixl_num_descriptors_sum",
228
+ "vllm:nixl_num_descriptors_count",
229
+ "vllm:prefix_cache_hits",
230
+ "vllm:prefix_cache_queries",
231
+]
232
+
233
+
234
+def get_metric(host: str, port: str, metric_name: str) -> float:
235
+ """Scrape a single Prometheus metric from /metrics."""
236
+ url = f"http://{host}:{port}/metrics"
237
+ resp = requests.get(url, timeout=10)
238
+ resp.raise_for_status()
239
+ total = 0.0
240
+ for line in resp.text.splitlines():
241
+ if line.startswith("#"):
242
+ continue
243
+ if line.startswith(metric_name):
244
+ match = re.search(r"[\d.eE+\-]+$", line)
245
+ if match:
246
+ total += float(match.group())
247
+ return total
248
+
249
+
250
+def get_all_metrics(host: str, port: str) -> dict[str, float]:
251
+ """Scrape all metrics of interest."""
252
+ return {name: get_metric(host, port, name) for name in METRICS_OF_INTEREST}
253
+
254
+
255
+def print_metrics(label: str, metrics: dict[str, float]) -> None:
256
+ print(f"\n [{label}]")
257
+ for name, val in metrics.items():
258
+ print(f" {name} = {val}")
259
+
260
+
261
+def test_mamba_prefix_cache_hit():
262
+ """Repeated prompts through PD should transfer fewer bytes on D-side."""
263
+ proxy_client = openai.OpenAI(
264
+ api_key="MY_KEY",
265
+ base_url=f"http://{PROXY_HOST}:{PROXY_PORT}/v1",
266
+ )
267
+ decode_client = openai.OpenAI(
268
+ api_key="MY_KEY",
269
+ base_url=f"http://{DECODE_HOST}:{DECODE_PORT}/v1",
270
+ )
271
+
272
+ models = decode_client.models.list()
273
+ MODEL = models.data[0].id
274
+ print(f"\nModel: {MODEL}")
275
+ print(f"Prompt length: {len(PROMPT)} chars")
276
+
277
+ # Baseline
278
+ m_baseline = get_all_metrics(DECODE_HOST, DECODE_PORT)
279
+ print_metrics("D-side baseline", m_baseline)
280
+
281
+ # Request 1: cold, primes the D-side cache
282
+ print("\n--- Request 1 (cold) ---")
283
+ resp1 = proxy_client.completions.create(
284
+ model=MODEL, prompt=PROMPT, max_tokens=10, temperature=0, seed=42
285
+ )
286
+ output1 = resp1.choices[0].text
287
+ print(f" Output: {output1!r}")
288
+ time.sleep(2)
289
+
290
+ m_after_req1 = get_all_metrics(DECODE_HOST, DECODE_PORT)
291
+ print_metrics("D-side after req1", m_after_req1)
292
+
293
+ transfer_req1 = (
294
+ m_after_req1["vllm:nixl_bytes_transferred_sum"]
295
+ - m_baseline["vllm:nixl_bytes_transferred_sum"]
296
+ )
297
+ descs_req1 = (
298
+ m_after_req1["vllm:nixl_num_descriptors_sum"]
299
+ - m_baseline["vllm:nixl_num_descriptors_sum"]
300
+ )
301
+ print(f" Transfer: {transfer_req1 / 1e6:.2f} MB, {descs_req1:.0f} descs")
302
+
303
+ # Request 2: same prompt, should hit D-side prefix cache
304
+ print("\n--- Request 2 (warm, same prompt) ---")
305
+ resp2 = proxy_client.completions.create(
306
+ model=MODEL, prompt=PROMPT, max_tokens=10, temperature=0, seed=42
307
+ )
308
+ output2 = resp2.choices[0].text
309
+ print(f" Output: {output2!r}")
310
+ time.sleep(2)
311
+
312
+ m_after_req2 = get_all_metrics(DECODE_HOST, DECODE_PORT)
313
+ print_metrics("D-side after req2", m_after_req2)
314
+
315
+ transfer_req2 = (
316
+ m_after_req2["vllm:nixl_bytes_transferred_sum"]
317
+ - m_after_req1["vllm:nixl_bytes_transferred_sum"]
318
+ )
319
+ descs_req2 = (
320
+ m_after_req2["vllm:nixl_num_descriptors_sum"]
321
+ - m_after_req1["vllm:nixl_num_descriptors_sum"]
322
+ )
323
+ print(f" Transfer: {transfer_req2 / 1e6:.2f} MB, {descs_req2:.0f} descs")
324
+
325
+ # P-side metrics (informational)
326
+ m_prefill = get_all_metrics(PREFILL_HOST, PREFILL_PORT)
327
+ print_metrics("P-side final", m_prefill)
328
+
329
+ # Summary
330
+ print("\n--- Summary ---")
331
+ print(f" Req 1: {transfer_req1 / 1e6:.2f} MB ({descs_req1:.0f} descs)")
332
+ print(f" Req 2: {transfer_req2 / 1e6:.2f} MB ({descs_req2:.0f} descs)")
333
+ if transfer_req1 > 0:
334
+ reduction_pct = (1 - transfer_req2 / transfer_req1) * 100
335
+ print(f" Reduction: {reduction_pct:.1f}%")
336
+
337
+ # Assertions
338
+ assert transfer_req1 > 0, (
339
+ f"First request should transfer data, got {transfer_req1} bytes"
340
+ )
341
+ assert transfer_req2 < transfer_req1, (
342
+ f"Second request should transfer fewer bytes due to D-side prefix "
343
+ f"cache hits. Got req1={transfer_req1 / 1e6:.2f} MB, "
344
+ f"req2={transfer_req2 / 1e6:.2f} MB (no reduction)."
345
+ )
346
+ assert output1 == output2, f"Outputs differ: {output1!r} vs {output2!r}"
vllm/v1/core/kv_cache_coordinator.py
+39
@@ -691,6 +691,45 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
691
blocks if blocks is not None else [] for blocks in hit_blocks_by_group
692
), hit_length
693
694
+ def find_longest_cache_hit_per_group(
695
+ self,
696
+ block_hashes: list[BlockHash],
697
+ max_cache_hit_length: int,
698
+ ) -> tuple[tuple[list[KVCacheBlock], ...], tuple[int, ...]]:
699
+ """Like find_longest_cache_hit but evaluates each group independently.
700
+
701
+ Returns:
702
+ (blocks_per_group, hit_lengths_per_group)
703
+ """
704
+
705
+ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
706
+ if kv_cache_spec.block_size == self.hash_block_size:
707
+ return block_hashes
708
+ return BlockHashListWithBlockSize(
709
+ block_hashes, self.hash_block_size, kv_cache_spec.block_size
710
+ )
711
+
712
+ num_groups = len(self.kv_cache_config.kv_cache_groups)
713
+ hit_blocks: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
714
+ hit_lengths: list[int] = [0] * num_groups
715
+
716
+ for spec, group_ids, manager_cls, use_eagle in self.attention_groups:
717
+ blocks = manager_cls.find_longest_cache_hit(
718
+ block_hashes=_get_block_hashes(spec),
719
+ max_length=max_cache_hit_length,
720
+ kv_cache_group_ids=group_ids,
721
+ block_pool=self.block_pool,
722
+ kv_cache_spec=spec,
723
+ drop_eagle_block=use_eagle,
724
+ alignment_tokens=self.scheduler_block_size,
725
+ )
726
+ group_hit = len(blocks[0]) * spec.block_size
727
+ for gid, blks in zip(group_ids, blocks):
728
+ hit_blocks[gid] = blks
729
+ hit_lengths[gid] = group_hit
730
+
731
+ return tuple(hit_blocks), tuple(hit_lengths)
732
+
733
734
def get_kv_cache_coordinator(
735
kv_cache_config: KVCacheConfig,
vllm/v1/core/sched/scheduler.py
+39
-3
@@ -34,6 +34,7 @@ from vllm.v1.core.encoder_cache_manager import (
34
EncoderCacheManager,
35
EncoderDecoderCacheManager,
36
)
37
+from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
38
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
39
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
40
from vllm.v1.core.sched.interface import PauseState, SchedulerInterface
@@ -620,9 +621,44 @@ class Scheduler(SchedulerInterface):
621
# Get already-cached tokens.
622
if request.num_computed_tokens == 0:
623
# Get locally-cached tokens.
623
- new_computed_blocks, num_new_local_computed_tokens = (
624
- self.kv_cache_manager.get_computed_blocks(request)
625
- )
624
+ if (
625
+ self.connector is not None
626
+ and self.has_mamba_layers
627
+ and isinstance(
628
+ self.kv_cache_manager.coordinator,
629
+ HybridKVCacheCoordinator,
630
+ )
631
+ ):
632
+ computed, per_group_hits = (
633
+ self.kv_cache_manager.coordinator.find_longest_cache_hit_per_group(
634
+ request.block_hashes,
635
+ request.num_tokens - 1,
636
+ )
637
+ )
638
+ new_computed_blocks = (
639
+ self.kv_cache_manager.create_kv_cache_blocks(computed)
640
+ )
641
+ # NOTE(ZhanqiuHu): For Mamba hybrid models,
642
+ # num_new_local_computed_tokens should be the FA hit
643
+ # length. This value is passed to the connector's
644
+ # get_num_new_matched_tokens which computes:
645
+ # external = total - local_computed.
646
+ # Using the FA hit skips re-transferring FA blocks
647
+ # already cached on D-side. The Mamba state (always
648
+ # the last block) is transferred unconditionally by
649
+ # _apply_prefix_caching in nixl/worker.py.
650
+ num_new_local_computed_tokens = max(per_group_hits)
651
+ if self.kv_cache_manager.log_stats:
652
+ assert self.kv_cache_manager.prefix_cache_stats is not None
653
+ self.kv_cache_manager.prefix_cache_stats.record(
654
+ num_tokens=request.num_tokens,
655
+ num_hits=num_new_local_computed_tokens,
656
+ preempted=request.num_preemptions > 0,
657
+ )
658
+ else:
659
+ new_computed_blocks, num_new_local_computed_tokens = (
660
+ self.kv_cache_manager.get_computed_blocks(request)
661
+ )
662
663
# In case of hybrid models, obtain hint for Marconi-style APC logic
664
if self.has_mamba_layers: