@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6c2eef5a5d
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[P/D] KVConnector for decode benchmarking (#25986)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

Tyler Michael Smith committed Oct 21, 2025 at 19:30 UTC 6c2eef5a5d95b3d213856626a40bafb2c5cb1e18
4 files changed +838 -1
tests/v1/kv_connector/unit/test_decode_bench_connector.py new
+415
@@ -0,0 +1,415 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for DecodeBenchConnector.
5 +
6 +Tests the functionality of the DecodeBenchConnector which fills KV cache
7 +with dummy values for decode performance benchmarking.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm import SamplingParams
14 +from vllm.config import KVTransferConfig
15 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
16 +
17 +# ruff: noqa: E501
18 +from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import (
19 + DecodeBenchConnector,
20 + DecodeBenchConnectorMetadata,
21 +)
22 +from vllm.forward_context import ForwardContext
23 +from vllm.utils.hashing import sha256
24 +from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
25 +from vllm.v1.core.sched.scheduler import Scheduler
26 +from vllm.v1.request import Request
27 +
28 +from .utils import (
29 + EOS_TOKEN_ID,
30 + create_model_runner_output,
31 + create_scheduler,
32 + create_vllm_config,
33 +)
34 +
35 +
36 +class DecodeBenchTestRunner:
37 + """Test runner for DecodeBenchConnector."""
38 +
39 + def __init__(self, block_size: int, num_gpu_blocks: int):
40 + self.block_size = block_size
41 + self.num_gpu_blocks = num_gpu_blocks
42 +
43 + self.req_id = -1
44 +
45 + # Create vllm config with DecodeBenchConnector
46 + vllm_config = create_vllm_config(
47 + block_size=block_size, max_num_batched_tokens=1000
48 + )
49 + vllm_config.kv_transfer_config = KVTransferConfig(
50 + kv_connector="DecodeBenchConnector",
51 + kv_role="kv_both",
52 + )
53 +
54 + self.vllm_config = vllm_config
55 + self.scheduler: Scheduler = create_scheduler(
56 + vllm_config, num_blocks=num_gpu_blocks
57 + )
58 +
59 + # Create worker-side connector
60 + self.worker_connector = DecodeBenchConnector(
61 + vllm_config, KVConnectorRole.WORKER
62 + )
63 +
64 + # Create dummy KV caches for testing
65 + # Shape: [num_blocks, 2, num_heads, block_size, head_dim]
66 + # Using simplified shape for testing
67 + num_heads = 4
68 + head_dim = 64
69 + self.kv_caches = {
70 + f"layer_{i}": torch.zeros(
71 + num_gpu_blocks, 2, num_heads, block_size, head_dim
72 + )
73 + for i in range(2) # 2 layers for testing
74 + }
75 +
76 + # Register KV caches with worker connector
77 + self.worker_connector.register_kv_caches(self.kv_caches)
78 +
79 + # Extract scheduler-side connector
80 + scheduler_connector = self.scheduler.connector
81 + assert scheduler_connector is not None
82 + assert isinstance(scheduler_connector, DecodeBenchConnector)
83 + self.scheduler_connector: DecodeBenchConnector = scheduler_connector
84 +
85 + init_none_hash(sha256)
86 + self._block_hasher = get_request_block_hasher(block_size, sha256)
87 +
88 + self._dummy_ctx: ForwardContext = ForwardContext(
89 + no_compile_layers={}, attn_metadata={}, virtual_engine=0
90 + )
91 +
92 + def new_request(self, token_ids: list[int]) -> Request:
93 + """Create a new request with given token IDs."""
94 + self.req_id += 1
95 +
96 + req = Request(
97 + request_id=str(self.req_id),
98 + prompt_token_ids=token_ids,
99 + sampling_params=SamplingParams(max_tokens=100),
100 + pooling_params=None,
101 + eos_token_id=EOS_TOKEN_ID,
102 + block_hasher=self._block_hasher,
103 + )
104 +
105 + self.scheduler.add_request(req)
106 + return req
107 +
108 + def run_single_step(self, token_id: int = 0):
109 + """Run a single scheduler + worker step."""
110 + scheduler_output = self.scheduler.schedule()
111 +
112 + # Get connector metadata
113 + kv_connector_metadata = scheduler_output.kv_connector_metadata
114 + assert kv_connector_metadata is not None
115 + assert isinstance(kv_connector_metadata, DecodeBenchConnectorMetadata)
116 +
117 + # Bind metadata and load KV
118 + self.worker_connector.bind_connector_metadata(kv_connector_metadata)
119 + self.worker_connector.start_load_kv(self._dummy_ctx)
120 +
121 + if scheduler_output.total_num_scheduled_tokens > 0:
122 + self.worker_connector.wait_for_save()
123 +
124 + self.worker_connector.clear_connector_metadata()
125 +
126 + # Create model runner output
127 + model_runner_output = create_model_runner_output(
128 + reqs=self.scheduler.running,
129 + token_id=token_id,
130 + )
131 +
132 + self.scheduler.update_from_output(scheduler_output, model_runner_output)
133 +
134 + return scheduler_output, kv_connector_metadata
135 +
136 +
137 +def test_decode_bench_connector_basic():
138 + """Test basic functionality of DecodeBenchConnector."""
139 + block_size = 16
140 + num_gpu_blocks = 100
141 +
142 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
143 +
144 + # Create a request with multiple blocks worth of tokens
145 + num_tokens = block_size * 3 # 3 blocks
146 + token_ids = [1] * num_tokens
147 +
148 + req = runner.new_request(token_ids)
149 +
150 + # Run first step - should fill KV cache with dummy values
151 + scheduler_output, metadata = runner.run_single_step()
152 +
153 + # Check that get_num_new_matched_tokens returned correct value
154 + # Should be num_tokens - 1 (all except the last token for decode)
155 + expected_fill_tokens = num_tokens - 1
156 +
157 + # Check metadata has the request to fill
158 + assert len(metadata.reqs_to_fill) == 1
159 + assert req.request_id in metadata.reqs_to_fill
160 +
161 + block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
162 + assert num_tokens_to_fill == expected_fill_tokens
163 +
164 + # For standard attention, there's only one group
165 + assert len(block_ids_per_group) == 1
166 + block_ids = block_ids_per_group[0]
167 +
168 + # Calculate expected number of blocks
169 + expected_num_blocks = (expected_fill_tokens + block_size - 1) // block_size
170 + assert len(block_ids) == expected_num_blocks
171 +
172 + # Verify KV caches were filled with constant value
173 + for layer_name, kv_cache in runner.kv_caches.items():
174 + for block_id in block_ids:
175 + # Check that the block was filled
176 + block_data = kv_cache[block_id]
177 + # Should be filled with constant value 0.015
178 + assert torch.allclose(block_data, torch.tensor(0.015))
179 +
180 +
181 +def test_decode_bench_connector_no_refill():
182 + """Test that DecodeBenchConnector only fills once per request."""
183 + block_size = 16
184 + num_gpu_blocks = 100
185 +
186 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
187 +
188 + # Create a request
189 + num_tokens = block_size * 2
190 + token_ids = [1] * num_tokens
191 +
192 + runner.new_request(token_ids)
193 +
194 + # Run first step - should fill KV cache
195 + _, metadata1 = runner.run_single_step()
196 + assert len(metadata1.reqs_to_fill) == 1
197 +
198 + # Run second step - should NOT fill again (already filled)
199 + _, metadata2 = runner.run_single_step()
200 + assert len(metadata2.reqs_to_fill) == 0
201 +
202 +
203 +def test_decode_bench_connector_single_token():
204 + """Test DecodeBenchConnector with single token request."""
205 + block_size = 16
206 + num_gpu_blocks = 100
207 +
208 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
209 +
210 + # Create a request with just 1 token
211 + # Should not fill anything (need at least 2 tokens: 1 to fill, 1 to decode)
212 + token_ids = [1]
213 +
214 + runner.new_request(token_ids)
215 +
216 + # Run step - should NOT fill KV cache
217 + _, metadata = runner.run_single_step()
218 + assert len(metadata.reqs_to_fill) == 0
219 +
220 +
221 +def test_decode_bench_connector_two_tokens():
222 + """Test DecodeBenchConnector with two token request."""
223 + block_size = 16
224 + num_gpu_blocks = 100
225 +
226 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
227 +
228 + # Create a request with 2 tokens
229 + # Should fill 1 token (first token), decode the second
230 + token_ids = [1, 2]
231 +
232 + req = runner.new_request(token_ids)
233 +
234 + # Run step
235 + _, metadata = runner.run_single_step()
236 +
237 + assert len(metadata.reqs_to_fill) == 1
238 + assert req.request_id in metadata.reqs_to_fill
239 +
240 + block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
241 + assert num_tokens_to_fill == 1
242 + # For standard attention, there's only one group
243 + assert len(block_ids_per_group) == 1
244 + assert len(block_ids_per_group[0]) == 1 # 1 token needs 1 block
245 +
246 +
247 +def test_decode_bench_connector_large_context():
248 + """Test DecodeBenchConnector with large context size."""
249 + block_size = 16
250 + num_gpu_blocks = 1000
251 +
252 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
253 +
254 + # Create a request with many blocks
255 + num_blocks = 20
256 + num_tokens = block_size * num_blocks
257 + token_ids = list(range(num_tokens))
258 +
259 + req = runner.new_request(token_ids)
260 +
261 + # Run step
262 + _, metadata = runner.run_single_step()
263 +
264 + assert len(metadata.reqs_to_fill) == 1
265 + assert req.request_id in metadata.reqs_to_fill
266 +
267 + block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
268 +
269 + # Should fill all tokens except the last one
270 + expected_fill_tokens = num_tokens - 1
271 + assert num_tokens_to_fill == expected_fill_tokens
272 +
273 + # For standard attention, there's only one group
274 + assert len(block_ids_per_group) == 1
275 + block_ids = block_ids_per_group[0]
276 +
277 + # Calculate expected number of blocks
278 + expected_num_blocks = (expected_fill_tokens + block_size - 1) // block_size
279 + assert len(block_ids) == expected_num_blocks
280 +
281 + # Verify blocks were filled
282 + for layer_name, kv_cache in runner.kv_caches.items():
283 + for block_id in block_ids:
284 + block_data = kv_cache[block_id]
285 + assert torch.allclose(block_data, torch.tensor(0.015))
286 +
287 +
288 +def test_decode_bench_connector_multiple_requests():
289 + """Test DecodeBenchConnector with multiple sequential requests."""
290 + block_size = 16
291 + num_gpu_blocks = 100
292 +
293 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
294 +
295 + # First request
296 + req1 = runner.new_request([1] * (block_size * 2))
297 + _, metadata1 = runner.run_single_step()
298 +
299 + assert len(metadata1.reqs_to_fill) == 1
300 + assert req1.request_id in metadata1.reqs_to_fill
301 +
302 + # Complete first request
303 + while runner.scheduler.running:
304 + runner.run_single_step()
305 +
306 + # Add EOS to finish
307 + scheduler_output = runner.scheduler.schedule()
308 + model_runner_output = create_model_runner_output(
309 + reqs=runner.scheduler.running,
310 + token_id=EOS_TOKEN_ID,
311 + use_eos=True,
312 + )
313 + runner.scheduler.update_from_output(scheduler_output, model_runner_output)
314 +
315 + # Second request - should also get filled
316 + req2 = runner.new_request([2] * (block_size * 3))
317 + _, metadata2 = runner.run_single_step()
318 +
319 + assert len(metadata2.reqs_to_fill) == 1
320 + assert req2.request_id in metadata2.reqs_to_fill
321 +
322 + # Different request should have different metadata
323 + _, num_tokens1 = metadata1.reqs_to_fill[req1.request_id]
324 + _, num_tokens2 = metadata2.reqs_to_fill[req2.request_id]
325 +
326 + assert num_tokens1 == block_size * 2 - 1
327 + assert num_tokens2 == block_size * 3 - 1
328 +
329 +
330 +def test_decode_bench_connector_partial_block():
331 + """Test DecodeBenchConnector with partial block filling."""
332 + block_size = 16
333 + num_gpu_blocks = 100
334 +
335 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
336 +
337 + # Create a request that doesn't align to block boundaries
338 + # e.g., 2.5 blocks worth of tokens
339 + num_tokens = block_size * 2 + block_size // 2
340 + token_ids = [1] * num_tokens
341 +
342 + req = runner.new_request(token_ids)
343 +
344 + # Run step
345 + _, metadata = runner.run_single_step()
346 +
347 + assert len(metadata.reqs_to_fill) == 1
348 + assert req.request_id in metadata.reqs_to_fill
349 +
350 + block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
351 +
352 + # Should fill all tokens except the last one
353 + expected_fill_tokens = num_tokens - 1
354 + assert num_tokens_to_fill == expected_fill_tokens
355 +
356 + # For standard attention, there's only one group
357 + assert len(block_ids_per_group) == 1
358 + block_ids = block_ids_per_group[0]
359 +
360 + # Should allocate 3 blocks to hold the partial data
361 + expected_num_blocks = 3
362 + assert len(block_ids) == expected_num_blocks
363 +
364 +
365 +def test_decode_bench_connector_concurrent_requests():
366 + """Test DecodeBenchConnector with multiple concurrent requests in the same batch."""
367 + block_size = 16
368 + num_gpu_blocks = 1000
369 +
370 + runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
371 +
372 + # Create multiple requests that will be batched together
373 + req1 = runner.new_request([1] * (block_size * 2))
374 + req2 = runner.new_request([2] * (block_size * 3))
375 + req3 = runner.new_request([3] * (block_size * 1))
376 +
377 + # Run first step - all requests should be filled concurrently
378 + _, metadata = runner.run_single_step()
379 +
380 + # All three requests should be in the metadata
381 + assert len(metadata.reqs_to_fill) == 3
382 + assert req1.request_id in metadata.reqs_to_fill
383 + assert req2.request_id in metadata.reqs_to_fill
384 + assert req3.request_id in metadata.reqs_to_fill
385 +
386 + # Verify each request has correct fill info
387 + block_ids_per_group1, num_tokens1 = metadata.reqs_to_fill[req1.request_id]
388 + block_ids_per_group2, num_tokens2 = metadata.reqs_to_fill[req2.request_id]
389 + block_ids_per_group3, num_tokens3 = metadata.reqs_to_fill[req3.request_id]
390 +
391 + # Verify token counts (all tokens except last one)
392 + assert num_tokens1 == block_size * 2 - 1
393 + assert num_tokens2 == block_size * 3 - 1
394 + assert num_tokens3 == block_size * 1 - 1
395 +
396 + # Verify block counts for each request
397 + assert len(block_ids_per_group1[0]) == 2 # 2 blocks
398 + assert len(block_ids_per_group2[0]) == 3 # 3 blocks
399 + assert len(block_ids_per_group3[0]) == 1 # 1 block
400 +
401 + # Verify all blocks are filled in KV cache
402 + for req_id, (block_ids_per_group, _) in metadata.reqs_to_fill.items():
403 + block_ids = block_ids_per_group[0]
404 + for layer_name, kv_cache in runner.kv_caches.items():
405 + for block_id in block_ids:
406 + block_data = kv_cache[block_id]
407 + assert torch.allclose(block_data, torch.tensor(0.015))
408 +
409 + # Run second step - should NOT fill again (already filled)
410 + _, metadata2 = runner.run_single_step()
411 + assert len(metadata2.reqs_to_fill) == 0
412 +
413 +
414 +if __name__ == "__main__":
415 + pytest.main([__file__, "-v"])
vllm/distributed/kv_transfer/kv_connector/factory.py
+6
@@ -130,3 +130,9 @@ KVConnectorFactory.register_connector(
130 "vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector",
131 "OffloadingConnector",
132 )
133 +
134 +KVConnectorFactory.register_connector(
135 + "DecodeBenchConnector",
136 + "vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector",
137 + "DecodeBenchConnector",
138 +)
vllm/distributed/kv_transfer/kv_connector/v1/__init__.py
+4 -1
@@ -4,5 +4,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
4 KVConnectorBase_V1,
5 KVConnectorRole,
6 )
7 +from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import ( # noqa E:501
8 + DecodeBenchConnector,
9 +)
10
8 -__all__ = ["KVConnectorRole", "KVConnectorBase_V1"]
11 +__all__ = ["KVConnectorRole", "KVConnectorBase_V1", "DecodeBenchConnector"]
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py new
+413
@@ -0,0 +1,413 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +DecodeBenchConnector: A KV Connector for decode instance performance testing.
5 +
6 +This connector emulates a prefill-decode disaggregated setting by filling
7 +the KV cache with dummy values, allowing measurement of decoder performance
8 +under larger input sequence lengths (ISL) in resource-limited environments.
9 +
10 +Usage:
11 + To use this connector for benchmarking, configure it in the kv_transfer_config:
12 +
13 + Example:
14 + vllm serve <model> --kv-transfer-config '{
15 + "kv_connector": "DecodeBenchConnector",
16 + "kv_role": "kv_both",
17 + "kv_connector_extra_config": {
18 + "fill_mean": 0.015,
19 + "fill_std": 0.0
20 + }
21 + }'
22 +
23 + Then run your benchmark with desired input/output lengths:
24 + vllm bench serve --base-url http://127.0.0.1:8000 --model <model> \\
25 + --dataset-name random --random-input-len 40000 \\
26 + --random-output-len 100 --max-concurrency 10
27 +
28 + Configuration options (via kv_connector_extra_config):
29 + - fill_mean (float): Mean value for random normal fill (default: 0.015)
30 + - fill_std (float): Standard deviation for random fill (default: 0.0)
31 + Set to 0 for constant values, >0 for random sampling
32 +"""
33 +
34 +from dataclasses import dataclass
35 +from typing import TYPE_CHECKING, Any
36 +
37 +import torch
38 +
39 +from vllm.distributed.kv_transfer.kv_connector.v1 import (
40 + KVConnectorBase_V1,
41 + KVConnectorRole,
42 +)
43 +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
44 +from vllm.logger import init_logger
45 +from vllm.utils import cdiv
46 +
47 +if TYPE_CHECKING:
48 + from vllm.attention.backends.abstract import AttentionMetadata
49 + from vllm.config import VllmConfig
50 + from vllm.forward_context import ForwardContext
51 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
52 + from vllm.v1.core.sched.output import SchedulerOutput
53 + from vllm.v1.request import Request
54 +
55 +logger = init_logger(__name__)
56 +
57 +
58 +@dataclass
59 +class DecodeBenchConnectorMetadata(KVConnectorMetadata):
60 + """Metadata for DecodeBenchConnector.
61 +
62 + Contains information about which requests need their KV cache filled
63 + with dummy values for benchmarking purposes.
64 + """
65 +
66 + # request_id -> (block_ids_per_group, num_tokens_to_fill)
67 + # block_ids_per_group is a tuple of lists, one per KV cache group
68 + # For standard attention: single group, e.g., ([1, 2, 3],)
69 + # For MLA: multiple groups, e.g., ([1, 2], [1, 2])
70 + reqs_to_fill: dict[str, tuple[tuple[list[int], ...], int]]
71 +
72 +
73 +class DecodeBenchConnector(KVConnectorBase_V1):
74 + """
75 + A KV Connector for decode instance performance testing.
76 +
77 + This connector fills the KV cache with dummy (non-zero) values to
78 + emulate a prefill-decode disaggregated setting, enabling performance
79 + testing of the decoder with larger input sequence lengths.
80 + """
81 +
82 + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
83 + super().__init__(vllm_config, role)
84 +
85 + self.connector_scheduler: DecodeBenchConnectorScheduler | None = None
86 + self.connector_worker: DecodeBenchConnectorWorker | None = None
87 +
88 + if role == KVConnectorRole.SCHEDULER:
89 + self.connector_scheduler = DecodeBenchConnectorScheduler(vllm_config)
90 + elif role == KVConnectorRole.WORKER:
91 + self.connector_worker = DecodeBenchConnectorWorker(vllm_config)
92 +
93 + # ==============================
94 + # Worker-side methods
95 + # ==============================
96 +
97 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
98 + assert self.connector_worker is not None
99 + self.connector_worker.register_kv_caches(kv_caches)
100 +
101 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
102 + assert self.connector_worker is not None
103 + assert isinstance(self._connector_metadata, DecodeBenchConnectorMetadata)
104 + self.connector_worker.start_fill_kv(self._connector_metadata)
105 +
106 + def wait_for_layer_load(self, layer_name: str) -> None:
107 + # All operations are synchronous, so nothing to wait for
108 + pass
109 +
110 + def save_kv_layer(
111 + self,
112 + layer_name: str,
113 + kv_layer: torch.Tensor,
114 + attn_metadata: "AttentionMetadata",
115 + **kwargs: Any,
116 + ) -> None:
117 + # This connector doesn't save KV cache (benchmarking only)
118 + pass
119 +
120 + def wait_for_save(self):
121 + # This connector doesn't save KV cache (benchmarking only)
122 + pass
123 +
124 + # ==============================
125 + # Scheduler-side methods
126 + # ==============================
127 +
128 + def get_num_new_matched_tokens(
129 + self,
130 + request: "Request",
131 + num_computed_tokens: int,
132 + ) -> tuple[int | None, bool]:
133 + assert self.connector_scheduler is not None
134 + return self.connector_scheduler.get_num_new_matched_tokens(
135 + request, num_computed_tokens
136 + )
137 +
138 + def update_state_after_alloc(
139 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
140 + ):
141 + assert self.connector_scheduler is not None
142 + return self.connector_scheduler.update_state_after_alloc(
143 + request, blocks, num_external_tokens
144 + )
145 +
146 + def build_connector_meta(
147 + self, scheduler_output: "SchedulerOutput"
148 + ) -> KVConnectorMetadata:
149 + assert self.connector_scheduler is not None
150 + return self.connector_scheduler.build_connector_meta(scheduler_output)
151 +
152 + def request_finished(
153 + self,
154 + request: "Request",
155 + block_ids: list[int],
156 + ) -> tuple[bool, dict[str, Any] | None]:
157 + assert self.connector_scheduler is not None
158 + self.connector_scheduler.request_finished(request)
159 + return False, None
160 +
161 +
162 +class DecodeBenchConnectorScheduler:
163 + """Scheduler-side implementation for DecodeBenchConnector."""
164 +
165 + def __init__(self, vllm_config: "VllmConfig"):
166 + self.vllm_config = vllm_config
167 + self.block_size = vllm_config.cache_config.block_size
168 +
169 + # Track which requests have already been filled
170 + self._filled_requests: set[str] = set()
171 +
172 + # Track pending fills for the current scheduler step
173 + # request_id -> (block_ids_per_group, num_tokens_to_fill)
174 + # Note: _pending_fills doesn't need explicit cleanup - it's cleared
175 + # after build_connector_meta() is called in the same scheduler step
176 + self._pending_fills: dict[str, tuple[tuple[list[int], ...], int]] = {}
177 +
178 + def get_num_new_matched_tokens(
179 + self,
180 + request: "Request",
181 + num_computed_tokens: int,
182 + ) -> tuple[int, bool]:
183 + """
184 + For new requests, return the number of tokens that should be filled
185 + with dummy KV cache values.
186 +
187 + Returns:
188 + (num_tokens_to_fill, is_async)
189 + - num_tokens_to_fill: number of uncomputed tokens minus 1
190 + (we fill everything except the last token for decode)
191 + - is_async: False (synchronous filling)
192 + """
193 + req_id = request.request_id
194 +
195 + # Only fill once per request on first scheduling
196 + if req_id in self._filled_requests:
197 + return 0, False
198 +
199 + # Calculate how many tokens we need to fill
200 + # Fill all uncomputed tokens except the last one (which will be decoded)
201 + # This simulates having processed a long prefill
202 + num_uncomputed_tokens = request.num_tokens - num_computed_tokens
203 + num_tokens_to_fill = max(0, num_uncomputed_tokens - 1)
204 +
205 + if num_tokens_to_fill == 0:
206 + return 0, False
207 +
208 + # Return False for synchronous operation - the fill is fast enough
209 + # that async overhead isn't worth it
210 + return num_tokens_to_fill, False
211 +
212 + def update_state_after_alloc(
213 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
214 + ):
215 + """
216 + Called after blocks are allocated. Store the block IDs so we can
217 + fill them with dummy values.
218 +
219 + Supports both standard attention (single KV cache group) and MLA
220 + (multiple KV cache groups).
221 + """
222 + req_id = request.request_id
223 +
224 + if num_external_tokens == 0:
225 + return
226 +
227 + # Get the block IDs that were allocated
228 + # block_groups is a tuple of lists, one per KV cache group
229 + # For standard attention: 1 group
230 + # For MLA: multiple groups (one per attention type)
231 + block_groups = blocks.get_block_ids()
232 +
233 + # Calculate how many blocks we need to fill
234 + # num_external_tokens are the tokens we said we'd provide
235 + num_blocks_to_fill = cdiv(num_external_tokens, self.block_size)
236 +
237 + # Extract the first num_blocks_to_fill blocks from each group
238 + # All groups should have the same block IDs for the same request
239 + block_ids_per_group = tuple(
240 + group_blocks[:num_blocks_to_fill] for group_blocks in block_groups
241 + )
242 +
243 + # Store the blocks to fill for all group. _pending_fills doesn't need cleanup
244 + # as it's cleared after build_connector_meta
245 + self._pending_fills[req_id] = (
246 + block_ids_per_group,
247 + num_external_tokens,
248 + )
249 + self._filled_requests.add(req_id)
250 +
251 + logger.debug(
252 + "DecodeBenchConnector: Allocated %d blocks across %d KV cache groups "
253 + "for request %s",
254 + num_blocks_to_fill,
255 + len(block_groups),
256 + req_id,
257 + )
258 +
259 + def build_connector_meta(
260 + self, scheduler_output: "SchedulerOutput"
261 + ) -> KVConnectorMetadata:
262 + """
263 + Build metadata containing information about which blocks to fill
264 + with dummy KV values.
265 + """
266 + meta = DecodeBenchConnectorMetadata(reqs_to_fill=self._pending_fills.copy())
267 +
268 + # Clear pending fills after building metadata
269 + self._pending_fills.clear()
270 +
271 + return meta
272 +
273 + def request_finished(self, request: "Request"):
274 + """
275 + Called when a request has finished. Clean up any state.
276 + """
277 + self._filled_requests.discard(request.request_id)
278 +
279 +
280 +class DecodeBenchConnectorWorker:
281 + """Worker-side implementation for DecodeBenchConnector."""
282 +
283 + def __init__(self, vllm_config: "VllmConfig"):
284 + self.vllm_config = vllm_config
285 + self.block_size = vllm_config.cache_config.block_size
286 +
287 + # Get fill parameters from extra config
288 + kv_transfer_config = vllm_config.kv_transfer_config
289 + assert kv_transfer_config is not None
290 + self.fill_mean = kv_transfer_config.get_from_extra_config("fill_mean", 0.015)
291 + self.fill_std = kv_transfer_config.get_from_extra_config("fill_std", 0.0)
292 +
293 + # Will be populated via register_kv_caches
294 + self.kv_caches: dict[str, torch.Tensor] | None = None
295 +
296 + # Mapping from KV cache group index to list of layer names in that group
297 + self.group_to_layers: dict[int, list[str]] | None = None
298 +
299 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
300 + """Store references to the KV cache tensors and build group mapping."""
301 + self.kv_caches = kv_caches
302 +
303 + # For simplicity, assume all layers belong to group 0 (standard attention)
304 + # For MLA models with multiple groups, the metadata will handle the mapping
305 + # We just need to fill the blocks specified in the metadata
306 + self.group_to_layers = {0: list(kv_caches.keys())}
307 +
308 + logger.debug(
309 + "DecodeBenchConnector: Registered %d KV cache layers",
310 + len(kv_caches),
311 + )
312 +
313 + def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata):
314 + """
315 + Fill the allocated KV cache blocks with dummy (non-zero) values.
316 +
317 + This simulates having a populated KV cache from a prefill phase,
318 + allowing decode performance testing with larger context sizes.
319 +
320 + Supports both standard attention (single group) and MLA (multiple groups).
321 + """
322 + if not metadata.reqs_to_fill:
323 + return
324 +
325 + assert self.kv_caches is not None, "KV caches must be registered before filling"
326 + assert self.group_to_layers is not None, "Group mapping must be initialized"
327 +
328 + for req_id, (block_ids_per_group, num_tokens) in metadata.reqs_to_fill.items():
329 + # Fill blocks for each KV cache group
330 + for group_idx, block_ids in enumerate(block_ids_per_group):
331 + self._fill_blocks(group_idx, block_ids, num_tokens)
332 +
333 + logger.debug(
334 + "DecodeBenchConnector: Filled %d blocks (%d tokens) across %d groups "
335 + "for request %s",
336 + len(block_ids_per_group[0]) if block_ids_per_group else 0,
337 + num_tokens,
338 + len(block_ids_per_group),
339 + req_id,
340 + )
341 +
342 + def _fill_blocks(self, group_idx: int, block_ids: list[int], num_tokens: int):
343 + """
344 + Fill specified blocks with dummy non-zero values for a specific KV cache group.
345 +
346 + Args:
347 + group_idx: The KV cache group index to fill
348 + block_ids: List of block IDs to fill in this group
349 + num_tokens: Total number of tokens to fill across these blocks
350 + """
351 + if not block_ids:
352 + return
353 +
354 + assert self.kv_caches is not None
355 + assert self.group_to_layers is not None
356 +
357 + # Get the layers that belong to this group
358 + layer_names = self.group_to_layers.get(group_idx, [])
359 +
360 + # Fill only the layers in this group
361 + for layer_name in layer_names:
362 + if layer_name not in self.kv_caches:
363 + logger.warning(
364 + "DecodeBenchConnector: Layer %s not found in KV caches", layer_name
365 + )
366 + continue
367 +
368 + kv_cache = self.kv_caches[layer_name]
369 +
370 + # Convert block_ids to tensor on device
371 + block_ids_tensor = torch.tensor(
372 + block_ids, dtype=torch.long, device=kv_cache.device
373 + )
374 +
375 + # Filter invalid block IDs
376 + valid_mask = block_ids_tensor < kv_cache.shape[0]
377 + valid_block_ids = block_ids_tensor[valid_mask]
378 +
379 + if len(valid_block_ids) == 0:
380 + continue
381 +
382 + # Create fill values - either constant or random
383 + block_shape = kv_cache.shape[1:]
384 + if self.fill_std > 0:
385 + # Random normal sampling
386 + fill_values = torch.normal(
387 + mean=self.fill_mean,
388 + std=self.fill_std,
389 + size=(len(valid_block_ids),) + block_shape,
390 + dtype=kv_cache.dtype,
391 + device=kv_cache.device,
392 + )
393 + else:
394 + # Constant fill value
395 + fill_values = torch.full(
396 + (len(valid_block_ids),) + block_shape,
397 + self.fill_mean,
398 + dtype=kv_cache.dtype,
399 + device=kv_cache.device,
400 + )
401 +
402 + # Batch fill operation
403 + kv_cache[valid_block_ids] = fill_values
404 +
405 + logger.debug(
406 + "DecodeBenchConnector: Filled %d blocks in group %d with %s values "
407 + "(mean=%.3f, std=%.3f)",
408 + len(block_ids),
409 + group_idx,
410 + "random" if self.fill_std > 0 else "constant",
411 + self.fill_mean,
412 + self.fill_std,
413 + )