@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0b7f11a1ee

[AuxOutput] Add block-keyed storage for routed-expert outputs (#45635)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Codex <noreply@openai.com>

Hongxin Xu committed Sep 21, 2026 at 18:25 UTC 0b7f11a1eece3dc9c689f1b5ab90d60c03b10793
32 files changed +3600 -873
tests/config/test_aux_output_config.py new
+119
@@ -0,0 +1,119 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +from typing import Any
6 +
7 +import pytest
8 +
9 +from vllm.config import AuxOutputConfig, VllmConfig
10 +from vllm.engine.arg_utils import EngineArgs
11 +
12 +pytestmark = pytest.mark.cpu_test
13 +
14 +
15 +def _config(
16 + *,
17 + use_v2: bool = True,
18 + pp: int = 1,
19 + dcp: int = 1,
20 + pcp: int = 1,
21 + connector: str | None = None,
22 + runner_type: str = "generate",
23 + is_moe: bool = True,
24 + sliding_window: int | None = None,
25 + attention_chunk_size: int | None = None,
26 + enable_prefix_caching: bool = True,
27 + adaptive_verification: bool = False,
28 +):
29 + return SimpleNamespace(
30 + model_config=SimpleNamespace(
31 + runner_type=runner_type,
32 + is_moe=is_moe,
33 + get_sliding_window=lambda: sliding_window,
34 + attention_chunk_size=attention_chunk_size,
35 + ),
36 + use_v2_model_runner=use_v2,
37 + parallel_config=SimpleNamespace(
38 + pipeline_parallel_size=pp,
39 + decode_context_parallel_size=dcp,
40 + prefill_context_parallel_size=pcp,
41 + ),
42 + aux_output_config=AuxOutputConfig(enable_return_routed_experts=True),
43 + cache_config=SimpleNamespace(
44 + enable_prefix_caching=enable_prefix_caching,
45 + ),
46 + speculative_config=(
47 + SimpleNamespace(enable_adaptive_verification=True)
48 + if adaptive_verification
49 + else None
50 + ),
51 + kv_transfer_config=(
52 + None
53 + if connector is None
54 + else SimpleNamespace(
55 + is_kv_transfer_instance=True,
56 + kv_connector=connector,
57 + )
58 + ),
59 + )
60 +
61 +
62 +def test_aux_output_config_defaults():
63 + config = AuxOutputConfig()
64 +
65 + assert not config.enabled
66 + assert not config.enable_return_routed_experts
67 + assert config.max_bytes is None
68 +
69 +
70 +def test_aux_output_capture_changes_compilation_hash():
71 + disabled = AuxOutputConfig()
72 + enabled = AuxOutputConfig(enable_return_routed_experts=True)
73 +
74 + assert disabled.compute_hash() != enabled.compute_hash()
75 +
76 +
77 +def test_legacy_routed_experts_flag_updates_aux_output_config():
78 + args = EngineArgs(enable_return_routed_experts=True)
79 +
80 + assert args.aux_output_config.enabled
81 + assert args.aux_output_config.enable_return_routed_experts
82 +
83 +
84 +@pytest.mark.parametrize(
85 + ("kwargs", "error"),
86 + [
87 + ({"use_v2": False}, "requires Model Runner V2"),
88 + ({"runner_type": "pooling"}, "only supports generate runners"),
89 + ({"is_moe": False}, "only supports MoE models"),
90 + ({"enable_prefix_caching": False}, "requires prefix caching"),
91 + (
92 + {"adaptive_verification": True},
93 + "adaptive speculative verification",
94 + ),
95 + ({"pp": 2}, "pipeline parallelism"),
96 + ({"dcp": 2}, "context parallelism"),
97 + ({"pcp": 2}, "context parallelism"),
98 + ({"connector": "MooncakeConnector"}, "incompatible with KV connectors"),
99 + ],
100 +)
101 +def test_aux_output_connector_rejects_unsupported_configuration(kwargs, error):
102 + with pytest.raises(ValueError, match=error):
103 + VllmConfig._verify_aux_output_compatibility(_config(**kwargs))
104 +
105 +
106 +@pytest.mark.parametrize(
107 + "kwargs",
108 + [{"sliding_window": 4096}, {"attention_chunk_size": 4096}],
109 +)
110 +def test_aux_output_config_defers_attention_layout_to_kv_config(kwargs):
111 + VllmConfig._verify_aux_output_compatibility(_config(**kwargs))
112 +
113 +
114 +def test_aux_output_guards_are_inactive_when_capture_is_disabled():
115 + config: Any = VllmConfig.__new__(VllmConfig)
116 + config.model_config = SimpleNamespace()
117 + config.aux_output_config = AuxOutputConfig()
118 +
119 + config._verify_aux_output_compatibility()
tests/distributed/aux_output_connector/test_store.py new
+1854
@@ -0,0 +1,1854 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import threading
5 +from contextlib import nullcontext
6 +from dataclasses import dataclass, field
7 +from types import SimpleNamespace
8 +from unittest.mock import Mock
9 +
10 +import numpy as np
11 +import pytest
12 +import torch
13 +
14 +from vllm.distributed.aux_output_connector.connector import (
15 + AuxOutputConnectorMetadata,
16 + AuxOutputSchedulerConnector,
17 + AuxRequestOutput,
18 + PackedBlockHashes,
19 +)
20 +from vllm.distributed.aux_output_connector.routed_experts import (
21 + RoutedExpertsBuffer,
22 + materialize_routed_experts,
23 + publish_routed_experts,
24 + routed_experts_keys,
25 +)
26 +from vllm.distributed.aux_output_connector.store import (
27 + BackgroundBlockObjectStore,
28 + BlockObject,
29 + BlockObjectStore,
30 + BlockObjectStoreError,
31 +)
32 +from vllm.distributed.aux_output_connector.worker import (
33 + AuxOutputWorkerConnector,
34 +)
35 +from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
36 +from vllm.v1.outputs import ModelRunnerOutput
37 +from vllm.v1.worker.gpu import async_utils
38 +from vllm.v1.worker.gpu.sample.output import SamplerOutput
39 +
40 +pytestmark = pytest.mark.cpu_test
41 +
42 +_SHAPE = (3, 2)
43 +_DTYPE = np.dtype("uint8")
44 +_BLOCK_SIZE = 4
45 +
46 +
47 +def test_background_store_publishes_without_blocking_caller():
48 + started = threading.Event()
49 + release = threading.Event()
50 + underlying = Mock()
51 +
52 + def put(*_args, **_kwargs):
53 + started.set()
54 + release.wait()
55 +
56 + underlying.put.side_effect = put
57 + store = BackgroundBlockObjectStore(underlying, max_pending_batches=1)
58 +
59 + store.put([BlockObject("key", b"value")])
60 + assert started.wait(timeout=1)
61 + release.set()
62 + store.close()
63 +
64 + underlying.close.assert_called_once_with()
65 +
66 +
67 +def _make_connector():
68 + return AuxOutputSchedulerConnector()
69 +
70 +
71 +@dataclass
72 +class _Execution:
73 + request_id: str
74 + token_start: int
75 + num_tokens: int
76 + emit_start: int
77 + block_hashes: list[bytes] | PackedBlockHashes
78 +
79 +
80 +@dataclass
81 +class _WorkerStep:
82 + metadata: AuxOutputConnectorMetadata
83 + executions: list[_Execution]
84 +
85 +
86 +@dataclass
87 +class _SchedulerRequest:
88 + request_id: str
89 + block_hashes: list[bytes]
90 + num_tokens: int
91 + num_output_tokens: int = 1
92 + num_computed_tokens: int = 0
93 + num_in_flight_tokens: int = 0
94 + finished: bool = False
95 + sampling_params: SimpleNamespace = field(
96 + default_factory=lambda: SimpleNamespace(routed_experts_prompt_start=0)
97 + )
98 +
99 + def is_finished(self) -> bool:
100 + return self.finished
101 +
102 +
103 +def _request_metadata(
104 + request_id: str,
105 + token_start: int,
106 + num_tokens: int,
107 + emit_start: int,
108 + block_hashes,
109 +) -> _Execution:
110 + return _Execution(request_id, token_start, num_tokens, emit_start, block_hashes)
111 +
112 +
113 +def _metadata(
114 + generation: int,
115 + requests: list[_Execution],
116 + finished_requests: dict[str, list[bytes] | PackedBlockHashes],
117 + hash_updates: dict[str, list[bytes] | PackedBlockHashes] | None = None,
118 +) -> _WorkerStep:
119 + block_hashes = {
120 + request.request_id: request.block_hashes
121 + for request in requests
122 + if request.block_hashes
123 + }
124 + if hash_updates:
125 + block_hashes.update(
126 + (request_id, hashes)
127 + for request_id, hashes in hash_updates.items()
128 + if hashes
129 + )
130 + block_hashes.update(
131 + (request_id, hashes)
132 + for request_id, hashes in finished_requests.items()
133 + if hashes
134 + )
135 + return _WorkerStep(
136 + AuxOutputConnectorMetadata(
137 + generation,
138 + {request.request_id: request.emit_start for request in requests},
139 + block_hashes,
140 + tuple(finished_requests),
141 + ),
142 + requests,
143 + )
144 +
145 +
146 +def _execution_ranges(metadata, request_ids):
147 + by_request = {request.request_id: request for request in metadata.executions}
148 + token_starts = tuple(
149 + by_request[request_id].token_start for request_id in request_ids
150 + )
151 + num_tokens = tuple(by_request[request_id].num_tokens for request_id in request_ids)
152 + return token_starts, num_tokens
153 +
154 +
155 +def _begin_step(worker: AuxOutputWorkerConnector, step: _WorkerStep) -> None:
156 + worker.begin_step(step.metadata)
157 +
158 +
159 +def _make_worker(
160 + max_num_seqs: int,
161 + max_num_batched_tokens: int | None = None,
162 + max_concurrent_batches: int = 2,
163 + max_store_blocks: int | None = None,
164 +) -> AuxOutputWorkerConnector:
165 + object_nbytes = _BLOCK_SIZE * int(np.prod(_SHAPE))
166 + store = _make_store(
167 + max_bytes=(
168 + 1 << 20 if max_store_blocks is None else max_store_blocks * object_nbytes
169 + ),
170 + object_nbytes=object_nbytes,
171 + )
172 + worker = object.__new__(AuxOutputWorkerConnector)
173 + worker._store = store
174 + worker._buffer = RoutedExpertsBuffer(
175 + _DTYPE,
176 + _SHAPE,
177 + _BLOCK_SIZE,
178 + max_num_seqs,
179 + max_num_batched_tokens or max_num_seqs * _BLOCK_SIZE,
180 + max_concurrent_batches,
181 + )
182 + worker._requests = {}
183 + worker._generation = 0
184 + worker._step_metadata = None
185 + worker._pending_outputs = []
186 + worker._lock = threading.Lock()
187 + worker._max_concurrent_batches = max_concurrent_batches
188 + return worker
189 +
190 +
191 +def test_worker_rejects_metadata_generation_rollback():
192 + worker = _make_worker(1)
193 + _begin_step(worker, _metadata(1, [], {}))
194 +
195 + with pytest.raises(AssertionError, match="generation moved backwards"):
196 + _begin_step(worker, _metadata(0, [], {}))
197 +
198 + worker.close()
199 +
200 +
201 +def test_worker_rejects_run_and_finish_in_one_step():
202 + worker = _make_worker(1)
203 + request = _request_metadata("request", 0, 1, 0, [])
204 + metadata = _metadata(
205 + 0,
206 + [request],
207 + {"request": []},
208 + )
209 +
210 + with pytest.raises(AssertionError, match="cannot run and finish"):
211 + _begin_step(worker, metadata)
212 +
213 + worker.close()
214 +
215 +
216 +def _input_batch(request_ids, token_starts, query_start_loc):
217 + return SimpleNamespace(
218 + req_ids=request_ids,
219 + num_computed_tokens_np=token_starts,
220 + query_start_loc_np=query_start_loc,
221 + )
222 +
223 +
224 +def test_non_output_rank_skips_capture_snapshot():
225 + worker = object.__new__(AuxOutputWorkerConnector)
226 + worker._store = None
227 + worker._buffer = None
228 + worker._capturer = Mock()
229 + worker._step_metadata = Mock()
230 + assert worker.prepare_output(_input_batch([], np.array([]), np.array([]))) is None
231 + worker._capturer.snapshot_routing_data.assert_not_called()
232 +
233 +
234 +def test_worker_skips_aux_outputs_for_internal_warmup_step():
235 + worker = _make_worker(1)
236 + worker._capturer = Mock()
237 +
238 + worker.begin_step(None)
239 +
240 + assert worker.prepare_output(_input_batch([], np.array([]), np.array([]))) is None
241 + worker._capturer.snapshot_routing_data.assert_not_called()
242 + worker.close()
243 +
244 +
245 +def test_next_step_does_not_consume_pending_output(monkeypatch):
246 + """begin_step must not wait for or consume an unconsumed step output."""
247 + event = Mock()
248 + monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: event)
249 + monkeypatch.setattr(async_utils, "stream", lambda *args: nullcontext())
250 + monkeypatch.setattr(
251 + async_utils, "async_copy_to_np", lambda tensor: tensor.numpy().copy()
252 + )
253 + worker = _make_worker(1)
254 + worker.begin_step(
255 + _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata
256 + )
257 + rows = torch.ones((1, *_SHAPE), dtype=torch.uint8)
258 + worker._capturer = Mock()
259 + worker._capturer.snapshot_routing_data.return_value = rows
260 + process_output = Mock(wraps=worker.process_output)
261 + monkeypatch.setattr(worker, "process_output", process_output)
262 + pending = worker.prepare_output(
263 + _input_batch(["request"], np.array([0]), np.array([0, 1]))
264 + )
265 + assert pending is not None
266 + output = async_utils.AsyncOutput(
267 + ModelRunnerOutput(["request"], {"request": 0}),
268 + SamplerOutput(
269 + torch.tensor([[7]]), None, None, None, num_rejected=torch.tensor([0])
270 + ),
271 + torch.tensor([1]),
272 + Mock(),
273 + Mock(),
274 + check_ep_fault=False,
275 + pending_aux_output=pending,
276 + )
277 + event.record.assert_called_once()
278 + assert worker._pending_outputs == [pending]
279 +
280 + worker.begin_step(_metadata(0, [], {}).metadata)
281 + process_output.assert_not_called()
282 + assert worker._pending_outputs == [pending]
283 +
284 + result = output.get_output()
285 +
286 + assert result.sampled_token_ids == [[7]]
287 + np.testing.assert_array_equal(
288 + result.aux_output_connector_output["request"].rows, rows.numpy()
289 + )
290 + process_output.assert_called_once()
291 + assert worker._pending_outputs == []
292 + worker.close()
293 +
294 +
295 +@pytest.mark.parametrize("num_pending", [1, 2])
296 +def test_finished_request_teardown_waits_for_pending_output(num_pending):
297 + """Terminal cleanup waits for every outstanding output, not just the first."""
298 + worker = _make_worker(1)
299 + rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, *_SHAPE)
300 + worker._capturer = Mock()
301 + worker._capturer.snapshot_routing_data.return_value = torch.from_numpy(rows)
302 + pending_outputs = []
303 + for i in range(num_pending):
304 + request = _request_metadata("request", 4 * i, 4, 4 * i, [bytes([i]) * 4])
305 + worker.begin_step(_metadata(0, [request], {}).metadata)
306 + pending = worker.prepare_output(
307 + _input_batch(["request"], np.array([4 * i]), np.array([0, 4]))
308 + )
309 + assert pending is not None
310 + pending.enqueue_cpu_copy(np.array([1]), np.array([0]))
311 + pending_outputs.append(pending)
312 +
313 + # The next step finishes the request while its output is unconsumed.
314 + worker.begin_step(_metadata(0, [], {"request": []}).metadata)
315 + assert "request" in worker._requests
316 +
317 + for i, pending in enumerate(pending_outputs):
318 + output = worker.process_output(pending)
319 + np.testing.assert_array_equal(output["request"].rows, rows)
320 + if i + 1 < num_pending:
321 + assert "request" in worker._requests
322 + assert worker._store._references
323 + assert "request" not in worker._requests
324 + # The finished request's store keys must be released after commit.
325 + assert not worker._store._references
326 + worker.close()
327 +
328 +
329 +def test_pending_outputs_bounded_by_max_concurrent_batches():
330 + worker = _make_worker(1, max_concurrent_batches=2)
331 + worker._capturer = Mock()
332 + worker._capturer.snapshot_routing_data.side_effect = lambda num_rows: torch.zeros(
333 + (num_rows, *_SHAPE), dtype=torch.uint8
334 + )
335 + metadata = _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata
336 + batch = _input_batch(["request"], np.array([0]), np.array([0, 1]))
337 + for _ in range(2):
338 + worker.begin_step(metadata)
339 + assert worker.prepare_output(batch) is not None
340 +
341 + worker.begin_step(metadata)
342 + with pytest.raises(AssertionError, match="not consumed in order"):
343 + worker.prepare_output(batch)
344 +
345 + worker.close()
346 +
347 +
348 +def test_worker_rejects_invalid_rejected_token_count():
349 + worker = _make_worker(1)
350 + metadata = _metadata(
351 + 0,
352 + [_request_metadata("request", 0, 1, 0, [])],
353 + {},
354 + )
355 + with pytest.raises(AssertionError, match="rejected-token count is invalid"):
356 + _process_output(
357 + worker,
358 + metadata,
359 + np.zeros((1, *_SHAPE), dtype=_DTYPE),
360 + ["request"],
361 + np.array([-1]),
362 + token_starts=(0,),
363 + num_tokens=(1,),
364 + )
365 + assert worker._pending_outputs == []
366 +
367 + worker.close()
368 +
369 +
370 +def _process_output(
371 + worker,
372 + step,
373 + rows,
374 + request_ids,
375 + num_rejected,
376 + num_sampled=None,
377 + *,
378 + token_starts=None,
379 + num_tokens=None,
380 + query_start_loc=None,
381 +):
382 + if num_sampled is None:
383 + num_sampled = np.ones(len(request_ids), dtype=np.int32)
384 + if token_starts is None or num_tokens is None:
385 + assert isinstance(step, _WorkerStep)
386 + token_starts, num_tokens = _execution_ranges(step, request_ids)
387 + worker._capturer = Mock()
388 + worker._capturer.snapshot_routing_data.return_value = torch.from_numpy(rows)
389 + metadata = step.metadata if isinstance(step, _WorkerStep) else step
390 + worker.begin_step(metadata)
391 + if query_start_loc is None:
392 + query_start_loc = np.concatenate(
393 + (np.zeros(1, dtype=np.int32), np.cumsum(num_tokens, dtype=np.int32))
394 + )
395 + pending = worker.prepare_output(
396 + _input_batch(request_ids, np.asarray(token_starts), query_start_loc)
397 + )
398 + assert pending is not None
399 + pending.enqueue_cpu_copy(num_sampled, num_rejected)
400 + return worker.process_output(pending)
401 +
402 +
403 +def test_worker_ignores_cudagraph_query_padding():
404 + worker = _make_worker(2)
405 + logical = np.arange(2 * 3 * 2, dtype=np.uint8).reshape(2, 3, 2)
406 + metadata = _metadata(
407 + 0,
408 + [
409 + _request_metadata("first", 0, 1, 0, []),
410 + _request_metadata("second", 0, 1, 0, []),
411 + ],
412 + {},
413 + )
414 +
415 + output = _process_output(
416 + worker,
417 + metadata,
418 + logical,
419 + ["first", "second"],
420 + np.zeros(2, dtype=np.int32),
421 + query_start_loc=np.array([0, 1, 2, 2], dtype=np.int32),
422 + )
423 +
424 + np.testing.assert_array_equal(output["first"].rows, logical[:1])
425 + np.testing.assert_array_equal(output["second"].rows, logical[1:])
426 + worker.close()
427 +
428 +
429 +def test_worker_rejects_scheduler_emit_cursor_ahead():
430 + worker = _make_worker(1)
431 + _begin_step(
432 + worker,
433 + _metadata(
434 + 0,
435 + [_request_metadata("request", 0, 1, 0, [])],
436 + {},
437 + ),
438 + )
439 +
440 + with pytest.raises(AssertionError, match="Scheduler emit cursor moved ahead"):
441 + _begin_step(
442 + worker,
443 + _metadata(
444 + 0,
445 + [_request_metadata("request", 1, 1, 1, [])],
446 + {},
447 + ),
448 + )
449 +
450 + worker.close()
451 +
452 +
453 +@pytest.mark.parametrize(
454 + "routing",
455 + [
456 + np.zeros((1, _SHAPE[0], _SHAPE[1] + 1), dtype=_DTYPE),
457 + np.zeros((1, *_SHAPE), dtype=np.int32),
458 + ],
459 + ids=["shape", "dtype"],
460 +)
461 +def test_worker_rejects_mismatched_capture_profile(routing):
462 + worker = _make_worker(1)
463 + metadata = _metadata(
464 + 0,
465 + [_request_metadata("request", 0, 1, 0, [])],
466 + {},
467 + )
468 +
469 + with pytest.raises(AssertionError, match="capture profile changed"):
470 + _process_output(
471 + worker,
472 + metadata,
473 + routing,
474 + ["request"],
475 + np.array([0]),
476 + )
477 + worker.close()
478 +
479 +
480 +@pytest.mark.parametrize(
481 + ("initial_start", "initial_rows", "invalid_start"),
482 + [
483 + (4, [4, 5, 6], 6),
484 + (0, [0], 2),
485 + (None, [], 3),
486 + ],
487 + ids=["overlap", "gap", "unaligned-initial-capture"],
488 +)
489 +def test_logical_buffer_rejects_noncontiguous_capture(
490 + initial_start, initial_rows, invalid_start
491 +):
492 + buffer = RoutedExpertsBuffer(np.dtype("uint8"), (1,), 4, 1, 4, 1)
493 + if initial_start is not None:
494 + rows = np.asarray(initial_rows, dtype=np.uint8).reshape(-1, 1)
495 + assert buffer.capture("request", initial_start, rows) == []
496 +
497 + with pytest.raises(AssertionError, match="not contiguous"):
498 + buffer.capture(
499 + "request", invalid_start, np.array([[invalid_start]], dtype=np.uint8)
500 + )
501 +
502 +
503 +def test_logical_buffer_captures_one_row_per_decode_step():
504 + buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 8, 2)
505 + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=_DTYPE).reshape(_BLOCK_SIZE, 3, 2)
506 +
507 + completed = []
508 + for step in range(_BLOCK_SIZE):
509 + completed += buffer.capture("request", step, logical[step : step + 1])
510 +
511 + assert len(completed) == 1
512 + np.testing.assert_array_equal(completed[0][1], logical)
513 +
514 +
515 +def test_logical_buffer_copies_borrowed_block_when_retained():
516 + buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 4, 2)
517 + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=_DTYPE).reshape(_BLOCK_SIZE, *_SHAPE)
518 +
519 + expected = logical.copy()
520 + completed = buffer.capture("request", 0, logical)
521 + retained = buffer.retain_block(completed[0][1])
522 + logical[:] = 0
523 +
524 + assert not np.shares_memory(retained, logical)
525 + np.testing.assert_array_equal(retained, expected)
526 + buffer.release_block(retained)
527 +
528 +
529 +def _make_store(
530 + *,
531 + max_bytes: int = 1 << 20,
532 + object_nbytes: int = 4,
533 +):
534 + return BlockObjectStore(
535 + max_bytes=max_bytes,
536 + object_nbytes=object_nbytes,
537 + )
538 +
539 +
540 +def test_publish_routed_experts_publishes_full_blocks():
541 + store = _make_store(object_nbytes=_BLOCK_SIZE * int(np.prod(_SHAPE)))
542 + buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 8, 2)
543 + logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2)
544 + hashes = [b"a" * 32, b"b" * 32]
545 + keys = routed_experts_keys(hashes, "0")
546 + blocks = buffer.capture("request", 0, logical)
547 + publish_routed_experts(
548 + store,
549 + batches=[(keys, blocks)],
550 + block_size=_BLOCK_SIZE,
551 + )
552 + expected = logical.copy()
553 + logical[:] = 0
554 +
555 + np.testing.assert_array_equal(
556 + materialize_routed_experts(
557 + store,
558 + keys,
559 + shape_per_token=_SHAPE,
560 + dtype=_DTYPE,
561 + ),
562 + expected,
563 + )
564 + store.close()
565 +
566 +
567 +def test_worker_data_plane_publishes_blocks_and_reuses_prefix():
568 + worker = _make_worker(2)
569 +
570 + hashes = [b"a" * 32, b"b" * 32, b"c" * 32]
571 + logical = np.arange(10 * 3 * 2, dtype=np.uint8).reshape(10, 3, 2)
572 + first = _metadata(
573 + generation=0,
574 + requests=[_request_metadata("first", 0, 8, 0, hashes)],
575 + finished_requests={},
576 + )
577 + output = _process_output(worker, first, logical[:8], ["first"], np.array([0]))
578 + assert output is not None
579 + np.testing.assert_array_equal(output["first"].rows, logical[:8])
580 +
581 + second = _metadata(
582 + generation=0,
583 + requests=[_request_metadata("second", 8, 2, 0, hashes)],
584 + finished_requests={"first": []},
585 + )
586 + output = _process_output(worker, second, logical[8:], ["second"], np.array([0]))
587 + assert output is not None
588 + np.testing.assert_array_equal(output["second"].rows, logical)
589 + worker.close()
590 +
591 +
592 +def test_worker_matches_kv_eviction_order_after_requests_finish():
593 + worker = _make_worker(2, max_store_blocks=3)
594 + hashes = [bytes([value]) * 32 for value in range(5)]
595 + blocks = [
596 + np.full((_BLOCK_SIZE, *_SHAPE), value, dtype=_DTYPE) for value in range(5)
597 + ]
598 +
599 + def publish_and_finish(request_id, block_indexes):
600 + rows = np.concatenate([blocks[index] for index in block_indexes])
601 + step = _metadata(
602 + 0,
603 + [
604 + _request_metadata(
605 + request_id,
606 + 0,
607 + len(rows),
608 + 0,
609 + [hashes[index] for index in block_indexes],
610 + )
611 + ],
612 + {},
613 + )
614 + _process_output(worker, step, rows, [request_id], np.array([0]))
615 + _begin_step(worker, _metadata(0, [], {request_id: []}))
616 +
617 + publish_and_finish("s", [0])
618 + publish_and_finish("ab", [1, 2])
619 + publish_and_finish("c", [3])
620 + publish_and_finish("d", [4])
621 +
622 + assert worker._store is not None
623 + np.testing.assert_array_equal(
624 + materialize_routed_experts(
625 + worker._store,
626 + routed_experts_keys([hashes[1], hashes[3], hashes[4]], "0"),
627 + shape_per_token=_SHAPE,
628 + dtype=_DTYPE,
629 + ),
630 + np.concatenate([blocks[1], blocks[3], blocks[4]]),
631 + )
632 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
633 + materialize_routed_experts(
634 + worker._store,
635 + routed_experts_keys([hashes[2]], "0"),
636 + shape_per_token=_SHAPE,
637 + dtype=_DTYPE,
638 + )
639 + worker.close()
640 +
641 +
642 +def test_worker_rejects_unbacked_capture_gap():
643 + worker = _make_worker(1)
644 + logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
645 + first = _metadata(
646 + 0,
647 + [_request_metadata("request", 0, 1, 0, [])],
648 + {},
649 + )
650 + _process_output(
651 + worker,
652 + first,
653 + logical[:1],
654 + ["request"],
655 + np.array([0]),
656 + np.zeros(1, dtype=np.int32),
657 + )
658 + second = _metadata(
659 + 0,
660 + [_request_metadata("request", 2, 1, 0, [])],
661 + {},
662 + )
663 +
664 + with pytest.raises(AssertionError, match="unbacked token gap"):
665 + _process_output(
666 + worker,
667 + second,
668 + logical[2:3],
669 + ["request"],
670 + np.array([0]),
671 + np.zeros(1, dtype=np.int32),
672 + )
673 + worker.close()
674 +
675 +
676 +def test_worker_replays_rejected_speculative_gap():
677 + worker = _make_worker(1)
678 + logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2)
679 + first = _metadata(
680 + 0,
681 + [_request_metadata("request", 0, 5, 0, [])],
682 + {},
683 + )
684 + first_output = _process_output(
685 + worker, first, logical[:5], ["request"], np.array([4])
686 + )
687 +
688 + queued = _metadata(
689 + 0,
690 + [_request_metadata("request", 5, 4, 0, [])],
691 + {},
692 + )
693 + queued_output = _process_output(
694 + worker,
695 + queued,
696 + logical[1:],
697 + ["request"],
698 + np.array([3]),
699 + )
700 +
701 + rolled_back = _metadata(
702 + 0,
703 + [_request_metadata("request", 6, 3, 0, [b"a" * 32])],
704 + {},
705 + )
706 +
707 + rolled_back_output = _process_output(
708 + worker,
709 + rolled_back,
710 + logical[2:],
711 + ["request"],
712 + np.array([0]),
713 + )
714 +
715 + np.testing.assert_array_equal(first_output["request"].rows, logical[:1])
716 + assert queued_output["request"].token_start == 1
717 + np.testing.assert_array_equal(
718 + queued_output["request"].rows,
719 + logical[1:2],
720 + )
721 + np.testing.assert_array_equal(
722 + rolled_back_output["request"].rows,
723 + logical[2:],
724 + )
725 + assert worker._requests["request"].capture_cursor == len(logical)
726 + worker.close()
727 +
728 +
729 +def test_worker_rejects_overlapping_capture_rows():
730 + worker = _make_worker(1)
731 + logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
732 + first = _metadata(
733 + 0,
734 + [_request_metadata("request", 0, 3, 0, [])],
735 + {},
736 + )
737 + _process_output(
738 + worker,
739 + first,
740 + logical[:3],
741 + ["request"],
742 + np.array([0]),
743 + np.zeros(1, dtype=np.int32),
744 + )
745 + stale_overlap = logical[2:].copy()
746 + stale_overlap[0] += 100
747 + second = _metadata(
748 + 0,
749 + [_request_metadata("request", 2, 2, 0, [b"a" * 32])],
750 + {},
751 + )
752 + with pytest.raises(AssertionError, match="capture moved backwards"):
753 + _process_output(
754 + worker,
755 + second,
756 + stale_overlap,
757 + ["request"],
758 + np.array([0]),
759 + )
760 + worker.close()
761 +
762 +
763 +def test_worker_publishes_entire_batch_before_materializing_prefix():
764 + worker = _make_worker(2)
765 + hashes = [b"a" * 32, b"b" * 32]
766 + logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2)
767 + metadata = _metadata(
768 + generation=0,
769 + requests=[
770 + _request_metadata("consumer", 4, 4, 0, hashes),
771 + _request_metadata("producer", 0, 4, 0, hashes),
772 + ],
773 + finished_requests={},
774 + )
775 +
776 + output = _process_output(
777 + worker,
778 + metadata,
779 + np.concatenate((logical[4:], logical[:4])),
780 + ["consumer", "producer"],
781 + np.array([0, 0]),
782 + )
783 +
784 + assert output is not None
785 + np.testing.assert_array_equal(output["consumer"].rows, logical)
786 + np.testing.assert_array_equal(output["producer"].rows, logical[:4])
787 + worker.close()
788 +
789 +
790 +def test_worker_defers_full_block_until_kv_hash_arrives():
791 + worker = _make_worker(1)
792 + logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2)
793 + first = _metadata(
794 + 0,
795 + [_request_metadata("request", 0, 4, 0, [])],
796 + {},
797 + )
798 +
799 + output = _process_output(
800 + worker,
801 + first,
802 + logical[:4],
803 + ["request"],
804 + np.array([0]),
805 + )
806 +
807 + assert output is not None
808 + np.testing.assert_array_equal(output["request"].rows, logical[:4])
809 + assert worker._requests["request"].pending_blocks[0][0] == 0
810 +
811 + block_hash = b"a" * 32
812 + second = _metadata(
813 + 0,
814 + [_request_metadata("request", 4, 1, 4, [block_hash])],
815 + {},
816 + )
817 + output = _process_output(worker, second, logical[4:], ["request"], np.array([0]))
818 +
819 + assert output is not None
820 + np.testing.assert_array_equal(output["request"].rows, logical[4:])
821 + assert not worker._requests["request"].pending_blocks
822 + assert worker._store is not None
823 + np.testing.assert_array_equal(
824 + materialize_routed_experts(
825 + worker._store,
826 + routed_experts_keys([block_hash], "0"),
827 + shape_per_token=_SHAPE,
828 + dtype=_DTYPE,
829 + ),
830 + logical[:4],
831 + )
832 + worker.close()
833 +
834 +
835 +def test_worker_releases_newly_keyed_blocks_before_capture():
836 + worker = _make_worker(1, max_num_batched_tokens=16)
837 + connector = _make_connector()
838 + block_hashes = [b"a" * 32]
839 + request = _scheduler_request("request", block_hashes, num_tokens=100)
840 + logical = np.arange(52 * 3 * 2, dtype=np.uint8).reshape(52, 3, 2)
841 +
842 + for token_start, num_tokens in [(0, 4), (4, 16), (20, 16)]:
843 + request.num_computed_tokens = token_start
844 + metadata = connector.build_connector_meta(
845 + _step_output([request.request_id], [token_start], [num_tokens]),
846 + {request.request_id: request},
847 + )
848 + _process_output(
849 + worker,
850 + metadata,
851 + logical[token_start : token_start + num_tokens],
852 + [request.request_id],
853 + np.array([0]),
854 + token_starts=(token_start,),
855 + num_tokens=(num_tokens,),
856 + )
857 +
858 + block_hashes.extend(bytes([value]) * 32 for value in range(1, 5))
859 + request.num_computed_tokens = 36
860 + metadata = connector.build_connector_meta(
861 + _step_output([request.request_id], [36], [16]),
862 + {request.request_id: request},
863 + )
864 + _process_output(
865 + worker,
866 + metadata,
867 + logical[36:52],
868 + [request.request_id],
869 + np.array([0]),
870 + token_starts=(36,),
871 + num_tokens=(16,),
872 + )
873 +
874 + state = worker._requests[request.request_id]
875 + assert [start for start, _ in state.pending_blocks] == [
876 + 20,
877 + 24,
878 + 28,
879 + 32,
880 + 36,
881 + 40,
882 + 44,
883 + 48,
884 + ]
885 + worker.close()
886 +
887 +
888 +def test_worker_publishes_pending_block_without_forward():
889 + worker = _make_worker(1)
890 + _begin_step(
891 + worker,
892 + _metadata(
893 + 0,
894 + [_request_metadata("request", 0, 1, 0, [])],
895 + {},
896 + ),
897 + )
898 + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(_BLOCK_SIZE, 3, 2)
899 + assert worker._buffer is not None
900 + worker._requests["request"].pending_blocks = [
901 + (0, worker._buffer.retain_block(logical))
902 + ]
903 + block_hash = b"a" * 32
904 +
905 + _begin_step(worker, _metadata(0, [], {}, {"request": [block_hash]}))
906 +
907 + assert not worker._requests["request"].pending_blocks
908 + assert worker._store is not None
909 + np.testing.assert_array_equal(
910 + materialize_routed_experts(
911 + worker._store,
912 + routed_experts_keys([block_hash], "0"),
913 + shape_per_token=_SHAPE,
914 + dtype=_DTYPE,
915 + ),
916 + logical,
917 + )
918 + worker.close()
919 +
920 +
921 +def test_worker_releases_unscheduled_request_blocks_before_capture():
922 + worker = _make_worker(2, 20, max_concurrent_batches=1)
923 + logical = np.arange(48 * 3 * 2, dtype=np.uint8).reshape(48, 3, 2)
924 + prompt_hashes = [b"a" * 32, b"b" * 32]
925 + prompt = _metadata(
926 + 0,
927 + [
928 + _request_metadata("first", 0, 4, 0, prompt_hashes[:1]),
929 + _request_metadata("second", 0, 4, 0, prompt_hashes[1:]),
930 + ],
931 + {},
932 + )
933 + _process_output(
934 + worker,
935 + prompt,
936 + logical[:8],
937 + ["first", "second"],
938 + np.zeros(2),
939 + np.zeros(2, dtype=np.int32),
940 + )
941 +
942 + first_decode = _metadata(
943 + 0,
944 + [_request_metadata("first", 4, 20, 0, [])],
945 + {},
946 + )
947 + _process_output(
948 + worker,
949 + first_decode,
950 + logical[8:28],
951 + ["first"],
952 + np.zeros(1),
953 + np.zeros(1, dtype=np.int32),
954 + )
955 + assert len(worker._requests["first"].pending_blocks) == 5
956 +
957 + first_hashes = [bytes([value]) * 32 for value in range(2, 7)]
958 + second_decode = _metadata(
959 + 0,
960 + [_request_metadata("second", 4, 20, 0, [])],
961 + {},
962 + {"first": first_hashes},
963 + )
964 + _process_output(
965 + worker,
966 + second_decode,
967 + logical[28:48],
968 + ["second"],
969 + np.zeros(1),
970 + np.zeros(1, dtype=np.int32),
971 + )
972 +
973 + assert not worker._requests["first"].pending_blocks
974 + assert len(worker._requests["second"].pending_blocks) == 5
975 + worker.close()
976 +
977 +
978 +def test_worker_does_not_rematerialize_emitted_rows():
979 + worker = _make_worker(1)
980 + hashes = [b"a" * 32, b"b" * 32]
981 + logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2)
982 +
983 + first = _metadata(
984 + 0,
985 + [_request_metadata("request", 0, 4, 0, hashes)],
986 + {},
987 + )
988 + output = _process_output(worker, first, logical[:4], ["request"], np.array([0]))
989 + assert output is not None
990 +
991 + assert worker._store is not None
992 + worker._store.get_concatenated = Mock(wraps=worker._store.get_concatenated)
993 + second = _metadata(
994 + 0,
995 + # Async scheduling may build this before the scheduler consumes first.
996 + [_request_metadata("request", 4, 4, 0, [])],
997 + {},
998 + )
999 + output = _process_output(worker, second, logical[4:], ["request"], np.array([0]))
1000 +
1001 + assert output is not None
1002 + np.testing.assert_array_equal(output["request"].rows, logical[4:])
1003 + worker._store.get_concatenated.assert_not_called()
1004 + worker.close()
1005 +
1006 +
1007 +def test_worker_resumes_from_scheduler_emit_boundary():
1008 + worker = _make_worker(1)
1009 + logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2)
1010 + resumed = _metadata(
1011 + 0,
1012 + [_request_metadata("request", 4, 1, 4, [b"a" * 32])],
1013 + {},
1014 + )
1015 +
1016 + output = _process_output(worker, resumed, logical[4:], ["request"], np.array([0]))
1017 +
1018 + assert output["request"].token_start == 4
1019 + np.testing.assert_array_equal(output["request"].rows, logical[4:])
1020 + assert worker._requests["request"].emit_cursor == 5
1021 + worker.close()
1022 +
1023 +
1024 +def test_worker_emits_mid_block_chunked_prefill():
1025 + worker = _make_worker(1)
1026 + logical = np.arange(3 * 3 * 2, dtype=np.uint8).reshape(3, 3, 2)
1027 +
1028 + first = _metadata(
1029 + 0,
1030 + [_request_metadata("request", 0, 2, 0, [])],
1031 + {},
1032 + )
1033 + output = _process_output(
1034 + worker,
1035 + first,
1036 + logical[:2],
1037 + ["request"],
1038 + np.array([0]),
1039 + np.zeros(1, dtype=np.int32),
1040 + )
1041 + assert output is not None and not output
1042 +
1043 + second = _metadata(
1044 + 0,
1045 + [_request_metadata("request", 2, 1, 0, [])],
1046 + {},
1047 + )
1048 + output = _process_output(worker, second, logical[2:], ["request"], np.array([0]))
1049 +
1050 + np.testing.assert_array_equal(output["request"].rows, logical)
1051 + worker.close()
1052 +
1053 +
1054 +def test_worker_emits_published_block_and_mid_block_tail():
1055 + worker = _make_worker(1)
1056 + logical = np.arange(7 * 3 * 2, dtype=np.uint8).reshape(7, 3, 2)
1057 +
1058 + first = _metadata(
1059 + 0,
1060 + [_request_metadata("request", 0, 6, 0, [b"a" * 32])],
1061 + {},
1062 + )
1063 + output = _process_output(
1064 + worker,
1065 + first,
1066 + logical[:6],
1067 + ["request"],
1068 + np.array([0]),
1069 + np.zeros(1, dtype=np.int32),
1070 + )
1071 + assert output is not None and not output
1072 +
1073 + second = _metadata(
1074 + 0,
1075 + [_request_metadata("request", 6, 1, 0, [])],
1076 + {},
1077 + )
1078 + output = _process_output(worker, second, logical[6:], ["request"], np.array([0]))
1079 +
1080 + np.testing.assert_array_equal(output["request"].rows, logical)
1081 + worker.close()
1082 +
1083 +
1084 +def test_worker_output_does_not_alias_released_tail_buffer():
1085 + worker = _make_worker(1)
1086 + logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
1087 + first = _metadata(
1088 + 0,
1089 + [_request_metadata("request", 0, 2, 0, [b"a" * 32])],
1090 + {},
1091 + )
1092 + _process_output(
1093 + worker,
1094 + first,
1095 + logical[:2],
1096 + ["request"],
1097 + np.array([0]),
1098 + np.zeros(1, dtype=np.int32),
1099 + )
1100 + second = _metadata(
1101 + 0,
1102 + [_request_metadata("request", 2, 2, 0, [])],
1103 + {},
1104 + )
1105 +
1106 + output = _process_output(worker, second, logical[2:], ["request"], np.array([0]))
1107 + emitted = output["request"].rows
1108 + expected = logical.copy()
1109 +
1110 + replacement = np.full_like(logical, 99)
1111 + reuse = _metadata(
1112 + 0,
1113 + [_request_metadata("reuse", 0, 4, 0, [b"b" * 32])],
1114 + {},
1115 + )
1116 + _process_output(
1117 + worker,
1118 + reuse,
1119 + replacement,
1120 + ["reuse"],
1121 + np.array([0]),
1122 + np.zeros(1, dtype=np.int32),
1123 + )
1124 +
1125 + np.testing.assert_array_equal(emitted, expected)
1126 + worker.close()
1127 +
1128 +
1129 +def test_worker_merges_block_hash_deltas():
1130 + worker = _make_worker(1)
1131 + _begin_step(
1132 + worker,
1133 + _metadata(
1134 + 0,
1135 + [_request_metadata("request", 0, 1, 0, [b"a" * 32, b"b" * 32])],
1136 + {},
1137 + ),
1138 + )
1139 + _begin_step(
1140 + worker,
1141 + _metadata(
1142 + 0,
1143 + [
1144 + _request_metadata(
1145 + "request",
1146 + 1,
1147 + 1,
1148 + 0,
1149 + [b"c" * 32],
1150 + )
1151 + ],
1152 + {},
1153 + ),
1154 + )
1155 +
1156 + assert worker._requests["request"].aux_output_keys == [
1157 + "vllm-artifact/0/" + (b"a" * 32).hex(),
1158 + "vllm-artifact/0/" + (b"b" * 32).hex(),
1159 + "vllm-artifact/0/" + (b"c" * 32).hex(),
1160 + ]
1161 + worker.close()
1162 +
1163 +
1164 +def test_worker_finish_publishes_keyed_block_and_discards_request_state():
1165 + worker = _make_worker(1)
1166 + logical = np.arange(2 * _BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(
1167 + 2 * _BLOCK_SIZE, 3, 2
1168 + )
1169 + running = _metadata(
1170 + 0,
1171 + [_request_metadata("request", 0, len(logical), 0, [])],
1172 + {},
1173 + )
1174 + _process_output(
1175 + worker,
1176 + running,
1177 + logical,
1178 + ["request"],
1179 + np.array([0]),
1180 + np.array([0]),
1181 + )
1182 + block_hash = b"a" * 32
1183 +
1184 + _begin_step(
1185 + worker,
1186 + _metadata(
1187 + 0,
1188 + [],
1189 + {"request": [block_hash]},
1190 + ),
1191 + )
1192 +
1193 + assert "request" not in worker._requests
1194 + assert worker._buffer is not None
1195 + assert "request" not in worker._buffer._requests
1196 + assert worker._store is not None
1197 + np.testing.assert_array_equal(
1198 + materialize_routed_experts(
1199 + worker._store,
1200 + routed_experts_keys([block_hash], "0"),
1201 + shape_per_token=_SHAPE,
1202 + dtype=_DTYPE,
1203 + ),
1204 + logical[:_BLOCK_SIZE],
1205 + )
1206 + worker.close()
1207 +
1208 +
1209 +def test_worker_fails_when_cached_aux_output_is_missing():
1210 + worker = _make_worker(1)
1211 + worker._generation = 0
1212 + metadata = _metadata(
1213 + 0,
1214 + [
1215 + _request_metadata(
1216 + "request",
1217 + 8,
1218 + 1,
1219 + 0,
1220 + [b"a" * 32, b"b" * 32],
1221 + )
1222 + ],
1223 + {},
1224 + )
1225 +
1226 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
1227 + _process_output(
1228 + worker,
1229 + metadata,
1230 + np.zeros((1, *_SHAPE), dtype=_DTYPE),
1231 + ["request"],
1232 + np.array([0]),
1233 + )
1234 + worker.close()
1235 +
1236 +
1237 +def test_store_rejects_oversized_batch_without_partial_write():
1238 + store = _make_store(max_bytes=6, object_nbytes=3)
1239 + store.put([BlockObject("retained", b"rrr")])
1240 +
1241 + with pytest.raises(BlockObjectStoreError, match="cannot retain"):
1242 + store.put(
1243 + [
1244 + BlockObject("first", b"111"),
1245 + BlockObject("second", b"222"),
1246 + BlockObject("third", b"333"),
1247 + ]
1248 + )
1249 +
1250 + assert store.get_concatenated(["retained"]) == b"rrr"
1251 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
1252 + store.get_concatenated(["first"])
1253 + store.close()
1254 +
1255 +
1256 +def test_store_lru_and_immutable_put():
1257 + store = _make_store(max_bytes=8)
1258 + store.put([BlockObject("first", b"1111"), BlockObject("second", b"2222")])
1259 + assert store.get_concatenated(["first"]) == b"1111"
1260 +
1261 + store.put([BlockObject("first", b"1111"), BlockObject("third", b"3333")])
1262 +
1263 + assert store.get_concatenated(["first", "third"]) == b"11113333"
1264 + with pytest.raises(BlockObjectStoreError, match="Increase aux_output_config"):
1265 + store.get_concatenated(["second"])
1266 + store.close()
1267 +
1268 +
1269 +def test_store_does_not_evict_referenced_aux_outputs():
1270 + store = _make_store(max_bytes=8)
1271 + store.put([BlockObject("first", b"1111"), BlockObject("second", b"2222")])
1272 + store.put([], retain_keys=["first"])
1273 + store.put([], retain_keys=["first"], release_keys=["first"])
1274 +
1275 + store.put([BlockObject("third", b"3333")])
1276 +
1277 + assert store.get_concatenated(["first", "third"]) == b"11113333"
1278 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
1279 + store.get_concatenated(["second"])
1280 + store.put([], release_keys=["first"])
1281 + store.close()
1282 +
1283 +
1284 +def test_store_terminal_put_preserves_tail_first_eviction_order():
1285 + store = _make_store(max_bytes=8)
1286 + keys = ["first", "second", "third"]
1287 + store.put(
1288 + [BlockObject("first", b"1111"), BlockObject("second", b"2222")],
1289 + retain_keys=keys,
1290 + )
1291 +
1292 + store.put(
1293 + [BlockObject("third", b"3333")],
1294 + release_keys=reversed(keys),
1295 + )
1296 +
1297 + assert store.get_concatenated(keys[:2]) == b"11112222"
1298 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
1299 + store.get_concatenated(["third"])
1300 + store.close()
1301 +
1302 +
1303 +def test_store_terminal_put_orders_survivors_for_later_eviction():
1304 + store = _make_store(max_bytes=12)
1305 + keys = ["first", "second", "third"]
1306 + store.put(
1307 + [BlockObject("first", b"1111"), BlockObject("second", b"2222")],
1308 + retain_keys=keys,
1309 + )
1310 +
1311 + store.put(
1312 + [BlockObject("third", b"3333")],
1313 + release_keys=reversed(keys),
1314 + )
1315 + store.put([BlockObject("fourth", b"4444")])
1316 +
1317 + assert store.get_concatenated(["first", "second", "fourth"]) == b"111122224444"
1318 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
1319 + store.get_concatenated(["third"])
1320 + store.close()
1321 +
1322 +
1323 +def test_store_terminal_put_keeps_a_shared_incoming_key():
1324 + store = _make_store(max_bytes=8)
1325 + keys = ["first", "second", "third"]
1326 + store.put(
1327 + [BlockObject("first", b"1111"), BlockObject("second", b"2222")],
1328 + retain_keys=(*keys, "third"),
1329 + )
1330 +
1331 + store.put(
1332 + [BlockObject("third", b"3333")],
1333 + release_keys=reversed(keys),
1334 + )
1335 +
1336 + assert store.get_concatenated(["first", "third"]) == b"11113333"
1337 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
1338 + store.get_concatenated(["second"])
1339 + store.put([], release_keys=["third"])
1340 + store.close()
1341 +
1342 +
1343 +def test_store_rejects_access_after_close():
1344 + store = _make_store(max_bytes=8)
1345 + store.close()
1346 +
1347 + with pytest.raises(RuntimeError, match="closed"):
1348 + store.put([BlockObject("first", b"1111")])
1349 + with pytest.raises(RuntimeError, match="closed"):
1350 + store.get_concatenated(["first"])
1351 +
1352 +
1353 +def test_store_reuses_evicted_slot_without_moving_live_objects():
1354 + store = _make_store(max_bytes=12)
1355 + store.put(
1356 + [
1357 + BlockObject("first", b"1111"),
1358 + BlockObject("second", b"2222"),
1359 + BlockObject("third", b"3333"),
1360 + ]
1361 + )
1362 + assert store.get_concatenated(["second"]) == b"2222"
1363 + store.put([BlockObject("fourth", b"4444")])
1364 +
1365 + assert store.get_concatenated(["second", "third", "fourth"]) == b"222233334444"
1366 + store.close()
1367 +
1368 +
1369 +def _scheduler_request(
1370 + request_id: str,
1371 + block_hashes: list[bytes],
1372 + *,
1373 + num_tokens: int = 10,
1374 + num_output_tokens: int = 1,
1375 + prompt_start: int = 0,
1376 +):
1377 + return _SchedulerRequest(
1378 + request_id=request_id,
1379 + block_hashes=block_hashes,
1380 + num_tokens=num_tokens,
1381 + num_output_tokens=num_output_tokens,
1382 + num_computed_tokens=num_tokens,
1383 + sampling_params=SimpleNamespace(routed_experts_prompt_start=prompt_start),
1384 + )
1385 +
1386 +
1387 +def _step_output(
1388 + request_ids: list[str], token_starts: list[int], token_counts: list[int]
1389 +) -> SchedulerOutput:
1390 + output = SchedulerOutput.make_empty()
1391 + output.scheduled_cached_reqs = CachedRequestData(
1392 + req_ids=request_ids,
1393 + resumed_req_ids=set(),
1394 + new_token_ids=[[] for _ in request_ids],
1395 + all_token_ids={},
1396 + new_block_ids=[None for _ in request_ids],
1397 + num_computed_tokens=token_starts,
1398 + num_output_tokens=[0 for _ in request_ids],
1399 + )
1400 + output.num_scheduled_tokens = dict(zip(request_ids, token_counts, strict=True))
1401 + output.total_num_scheduled_tokens = sum(token_counts)
1402 + return output
1403 +
1404 +
1405 +def test_scheduler_connector_builds_worker_metadata_and_forwards_output():
1406 + connector = _make_connector()
1407 + request = _scheduler_request("request", [b"a" * 32], num_tokens=5)
1408 + request.num_computed_tokens = 0
1409 + scheduler_output = _step_output([request.request_id], [0], [4])
1410 +
1411 + metadata = connector.build_connector_meta(
1412 + scheduler_output, {request.request_id: request}
1413 + )
1414 +
1415 + assert metadata.generation == 0
1416 + assert metadata.requests == {request.request_id: 4}
1417 + assert list(metadata.block_hashes[request.request_id]) == [b"a" * 32]
1418 +
1419 + routing = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
1420 + output = {"request": AuxRequestOutput(0, routing)}
1421 + request.num_computed_tokens = 4
1422 + np.testing.assert_array_equal(connector.take_output(request, output), routing)
1423 +
1424 +
1425 +def test_scheduler_starts_worker_output_at_requested_prompt_token():
1426 + connector = _make_connector()
1427 + request = _scheduler_request(
1428 + "request",
1429 + [b"a" * 32],
1430 + num_tokens=4,
1431 + num_output_tokens=0,
1432 + prompt_start=2,
1433 + )
1434 +
1435 + metadata = connector.build_connector_meta(
1436 + _step_output([request.request_id], [0], [4]),
1437 + {request.request_id: request},
1438 + )
1439 +
1440 + assert metadata.requests == {request.request_id: 2}
1441 +
1442 +
1443 +def test_scheduler_connector_preserves_request_finish_order():
1444 + connector = _make_connector()
1445 + first = _scheduler_request("first", [b"a" * 32])
1446 + second = _scheduler_request("second", [b"b" * 32])
1447 + connector.build_connector_meta(
1448 + _step_output([first.request_id, second.request_id], [0, 0], [0, 0]),
1449 + {first.request_id: first, second.request_id: second},
1450 + )
1451 + connector.request_finished(second)
1452 + connector.request_finished(first)
1453 +
1454 + metadata = connector.build_connector_meta(_step_output([], [], []), {})
1455 +
1456 + assert metadata.finished_requests == (second.request_id, first.request_id)
1457 + assert not metadata.block_hashes
1458 +
1459 +
1460 +def test_worker_preemption_terminal_keeps_inflight_hashes_for_publish():
1461 + worker = _make_worker(1)
1462 + block_hashes = [b"a" * 32, b"b" * 32]
1463 + logical = np.arange(2 * _BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(
1464 + 2 * _BLOCK_SIZE, 3, 2
1465 + )
1466 + first = _metadata(
1467 + 0,
1468 + [
1469 + _request_metadata(
1470 + "request",
1471 + 0,
1472 + _BLOCK_SIZE,
1473 + 0,
1474 + block_hashes[:1],
1475 + )
1476 + ],
1477 + {},
1478 + )
1479 + _process_output(
1480 + worker,
1481 + first,
1482 + logical[:_BLOCK_SIZE],
1483 + ["request"],
1484 + np.array([0]),
1485 + )
1486 + inflight = _metadata(
1487 + 0,
1488 + [
1489 + _request_metadata(
1490 + "request",
1491 + _BLOCK_SIZE,
1492 + _BLOCK_SIZE,
1493 + 0,
1494 + block_hashes[1:],
1495 + )
1496 + ],
1497 + {},
1498 + )
1499 + _process_output(
1500 + worker,
1501 + inflight,
1502 + logical[_BLOCK_SIZE:],
1503 + ["request"],
1504 + np.array([0]),
1505 + np.zeros(1, dtype=np.int32),
1506 + token_starts=(_BLOCK_SIZE,),
1507 + num_tokens=(_BLOCK_SIZE,),
1508 + )
1509 + _begin_step(
1510 + worker,
1511 + _metadata(
1512 + 0,
1513 + [],
1514 + {"request": []},
1515 + ),
1516 + )
1517 +
1518 + assert worker._store is not None
1519 + np.testing.assert_array_equal(
1520 + materialize_routed_experts(
1521 + worker._store,
1522 + routed_experts_keys(block_hashes[:1], "0"),
1523 + shape_per_token=_SHAPE,
1524 + dtype=_DTYPE,
1525 + ),
1526 + logical[:_BLOCK_SIZE],
1527 + )
1528 + np.testing.assert_array_equal(
1529 + materialize_routed_experts(
1530 + worker._store,
1531 + routed_experts_keys(block_hashes[1:], "0"),
1532 + shape_per_token=_SHAPE,
1533 + dtype=_DTYPE,
1534 + ),
1535 + logical[_BLOCK_SIZE:],
1536 + )
1537 + worker.close()
1538 +
1539 +
1540 +def test_worker_releases_terminal_pins_before_same_step_publish():
1541 + worker = _make_worker(2, max_store_blocks=1)
1542 + victim_hash = b"a" * 32
1543 + producer_hash = b"b" * 32
1544 + victim_rows = np.zeros((_BLOCK_SIZE, *_SHAPE), dtype=_DTYPE)
1545 + producer_rows = np.ones((_BLOCK_SIZE, *_SHAPE), dtype=_DTYPE)
1546 + running = _metadata(
1547 + 0,
1548 + [
1549 + _request_metadata("victim", 0, _BLOCK_SIZE, 0, [victim_hash]),
1550 + _request_metadata("producer", 0, _BLOCK_SIZE, 0, []),
1551 + ],
1552 + {},
1553 + )
1554 + _process_output(
1555 + worker,
1556 + running,
1557 + np.concatenate((victim_rows, producer_rows)),
1558 + ["victim", "producer"],
1559 + np.zeros(2, dtype=np.int32),
1560 + np.zeros(2, dtype=np.int32),
1561 + )
1562 +
1563 + _begin_step(
1564 + worker,
1565 + _metadata(
1566 + 0,
1567 + [],
1568 + {"victim": []},
1569 + {"producer": [producer_hash]},
1570 + ),
1571 + )
1572 +
1573 + assert worker._store is not None
1574 + np.testing.assert_array_equal(
1575 + materialize_routed_experts(
1576 + worker._store,
1577 + routed_experts_keys([producer_hash], "0"),
1578 + shape_per_token=_SHAPE,
1579 + dtype=_DTYPE,
1580 + ),
1581 + producer_rows,
1582 + )
1583 + with pytest.raises(BlockObjectStoreError, match="does not exist"):
1584 + worker._store.get_concatenated(routed_experts_keys([victim_hash], "0"))
1585 + worker.close()
1586 +
1587 +
1588 +def test_terminal_request_publishes_hash_discovered_after_last_schedule():
1589 + connector = _make_connector()
1590 + worker = _make_worker(1)
1591 + block_hashes: list[bytes] = []
1592 + request = _scheduler_request("request", block_hashes, num_tokens=5)
1593 + request.num_computed_tokens = 0
1594 + metadata = connector.build_connector_meta(
1595 + _step_output([request.request_id], [0], [4]),
1596 + {request.request_id: request},
1597 + )
1598 + logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(_BLOCK_SIZE, 3, 2)
1599 +
1600 + _process_output(
1601 + worker,
1602 + metadata,
1603 + logical,
1604 + [request.request_id],
1605 + np.array([0]),
1606 + token_starts=(0,),
1607 + num_tokens=(_BLOCK_SIZE,),
1608 + )
1609 + assert worker._requests[request.request_id].pending_blocks
1610 +
1611 + block_hash = b"a" * 32
1612 + block_hashes.append(block_hash)
1613 + request.num_computed_tokens = _BLOCK_SIZE
1614 + connector.request_finished(request)
1615 + cleanup = connector.build_connector_meta(_step_output([], [], []), {})
1616 + worker.begin_step(cleanup)
1617 +
1618 + assert worker._store is not None
1619 + np.testing.assert_array_equal(
1620 + materialize_routed_experts(
1621 + worker._store,
1622 + routed_experts_keys([block_hash], "0"),
1623 + shape_per_token=_SHAPE,
1624 + dtype=_DTYPE,
1625 + ),
1626 + logical,
1627 + )
1628 + worker.close()
1629 +
1630 +
1631 +def test_scheduler_connector_recreates_preempted_request_state():
1632 + connector = _make_connector()
1633 + request = _scheduler_request("request", [b"a" * 32], num_tokens=5)
1634 + connector.build_connector_meta(
1635 + _step_output([request.request_id], [0], [4]),
1636 + {request.request_id: request},
1637 + )
1638 + connector.request_finished(request)
1639 + cleanup = connector.build_connector_meta(
1640 + _step_output([], [], []),
1641 + {request.request_id: request},
1642 + )
1643 + assert request.request_id in cleanup.finished_requests
1644 + assert request.request_id not in cleanup.block_hashes
1645 +
1646 + resumed = connector.build_connector_meta(
1647 + _step_output([request.request_id], [0], [4]),
1648 + {request.request_id: request},
1649 + )
1650 + assert list(resumed.block_hashes[request.request_id]) == [b"a" * 32]
1651 +
1652 +
1653 +def test_scheduler_consumes_ordered_stale_aux_outputs():
1654 + connector = _make_connector()
1655 + request = _scheduler_request("request", [], num_tokens=6)
1656 + connector.build_connector_meta(
1657 + _step_output([request.request_id], [0], [4]),
1658 + {request.request_id: request},
1659 + )
1660 + first_rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
1661 + second_rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + 40
1662 + first = {"request": AuxRequestOutput(0, first_rows)}
1663 + second = {"request": AuxRequestOutput(4, second_rows)}
1664 + connector.request_finished(request)
1665 +
1666 + request.num_tokens = 5
1667 + np.testing.assert_array_equal(connector.take_output(request, first), first_rows)
1668 + request.num_tokens = 6
1669 + np.testing.assert_array_equal(
1670 + connector.take_output(request, second), second_rows[:1]
1671 + )
1672 +
1673 +
1674 +def test_scheduler_rejects_stale_output_without_aux_outputs():
1675 + connector = _make_connector()
1676 + request = _scheduler_request("request", [], num_tokens=1)
1677 + connector.build_connector_meta(
1678 + _step_output([request.request_id], [0], [1]),
1679 + {request.request_id: request},
1680 + )
1681 +
1682 + with pytest.raises(
1683 + AssertionError, match="auxiliary output worker output is missing"
1684 + ):
1685 + connector.take_output(
1686 + request,
1687 + {},
1688 + )
1689 +
1690 +
1691 +def test_scheduler_rejects_missing_accepted_aux_output_rows():
1692 + connector = _make_connector()
1693 + request = _scheduler_request("request", [], num_tokens=2)
1694 + connector.build_connector_meta(
1695 + _step_output([request.request_id], [0], [1]),
1696 + {request.request_id: request},
1697 + )
1698 + output = {
1699 + "request": AuxRequestOutput(
1700 + 0,
1701 + np.empty((0, *_SHAPE), dtype=_DTYPE),
1702 + )
1703 + }
1704 +
1705 + with pytest.raises(AssertionError, match="invalid token range"):
1706 + connector.take_output(request, output)
1707 +
1708 +
1709 +def test_scheduler_rejects_aux_output_past_finished_request():
1710 + connector = _make_connector()
1711 + request = _scheduler_request("request", [], num_tokens=1)
1712 + request.finished = True
1713 + connector.build_connector_meta(
1714 + _step_output([request.request_id], [0], [1]),
1715 + {request.request_id: request},
1716 + )
1717 + output = {
1718 + "request": AuxRequestOutput(
1719 + 1,
1720 + np.empty((0, *_SHAPE), dtype=_DTYPE),
1721 + )
1722 + }
1723 +
1724 + with pytest.raises(
1725 + AssertionError, match="finished auxiliary output output has no accepted"
1726 + ):
1727 + connector.take_output(request, output)
1728 +
1729 +
1730 +@pytest.mark.parametrize("chunk_size", [2, 4])
1731 +@pytest.mark.parametrize("max_tokens", [1, 2])
1732 +def test_prompt_end_offset_through_worker_and_scheduler(chunk_size, max_tokens):
1733 + """Omit prompt R3 without losing the empty first output or later decode rows."""
1734 + connector = _make_connector()
1735 + worker = _make_worker(1)
1736 + request = _scheduler_request(
1737 + "request", [b"a" * 32], num_tokens=4, num_output_tokens=0, prompt_start=4
1738 + )
1739 + rows = np.arange(5 * 3 * 2, dtype=_DTYPE).reshape(5, *_SHAPE)
1740 + try:
1741 + for start in range(0, 4, chunk_size):
1742 + metadata = connector.build_connector_meta(
1743 + _step_output([request.request_id], [start], [chunk_size]),
1744 + {request.request_id: request},
1745 + )
1746 + output = _process_output(
1747 + worker,
1748 + metadata,
1749 + rows[start : start + chunk_size],
1750 + [request.request_id],
1751 + np.array([0]),
1752 + num_sampled=np.array([int(start + chunk_size == 4)]),
1753 + token_starts=(start,),
1754 + num_tokens=(chunk_size,),
1755 + )
1756 + if start + chunk_size < 4:
1757 + assert output == {}
1758 + request.num_tokens = 5
1759 + request.num_output_tokens = 1
1760 + request.finished = max_tokens == 1
1761 + result = connector.take_output(request, output)
1762 + np.testing.assert_array_equal(result, rows[:0])
1763 +
1764 + if max_tokens == 2:
1765 + metadata = connector.build_connector_meta(
1766 + _step_output([request.request_id], [4], [1]),
1767 + {request.request_id: request},
1768 + )
1769 + output = _process_output(
1770 + worker,
1771 + metadata,
1772 + rows[4:],
1773 + [request.request_id],
1774 + np.array([0]),
1775 + token_starts=(4,),
1776 + num_tokens=(1,),
1777 + )
1778 + request.num_tokens = 6
1779 + request.finished = True
1780 + np.testing.assert_array_equal(
1781 + connector.take_output(request, output), rows[4:]
1782 + )
1783 + finally:
1784 + worker.close()
1785 +
1786 +
1787 +def test_scheduler_connector_sends_each_block_hash_once():
1788 + connector = _make_connector()
1789 + request = _scheduler_request("request", [b"a" * 32, b"b" * 32], num_tokens=12)
1790 + scheduler_output = _step_output([request.request_id], [0], [4])
1791 +
1792 + first = connector.build_connector_meta(
1793 + scheduler_output, {request.request_id: request}
1794 + )
1795 + request.block_hashes.append(b"c" * 32)
1796 + second = connector.build_connector_meta(
1797 + _step_output([request.request_id], [4], [4]),
1798 + {request.request_id: request},
1799 + )
1800 + third = connector.build_connector_meta(
1801 + _step_output([request.request_id], [8], [4]),
1802 + {request.request_id: request},
1803 + )
1804 +
1805 + assert list(first.block_hashes[request.request_id]) == [b"a" * 32, b"b" * 32]
1806 + assert list(second.block_hashes[request.request_id]) == [b"c" * 32]
1807 + assert request.request_id not in third.block_hashes
1808 +
1809 +
1810 +def test_scheduler_connector_sends_unscheduled_hash_update():
1811 + connector = _make_connector()
1812 + request = _scheduler_request("request", [b"a" * 32], num_tokens=4)
1813 + connector.build_connector_meta(
1814 + _step_output([request.request_id], [0], [4]),
1815 + {request.request_id: request},
1816 + )
1817 +
1818 + request.block_hashes.append(b"b" * 32)
1819 + request.num_tokens = 9
1820 + request.num_computed_tokens = 8
1821 + metadata = connector.build_connector_meta(
1822 + _step_output([], [], []),
1823 + {request.request_id: request},
1824 + )
1825 +
1826 + assert not metadata.requests
1827 + assert list(metadata.block_hashes[request.request_id]) == [b"b" * 32]
1828 +
1829 +
1830 +def test_scheduler_connector_reset_derives_emit_start_from_request():
1831 + connector = _make_connector()
1832 + request = _scheduler_request("request", [b"a" * 32], num_tokens=5)
1833 + request.num_computed_tokens = 0
1834 + before_reset = connector.build_connector_meta(
1835 + _step_output([request.request_id], [0], [4]),
1836 + {request.request_id: request},
1837 + )
1838 + assert list(before_reset.block_hashes[request.request_id]) == [b"a" * 32]
1839 +
1840 + connector.reset()
1841 + empty = connector.build_connector_meta(
1842 + _step_output([], [], []),
1843 + {request.request_id: request},
1844 + )
1845 + assert not empty.block_hashes
1846 + request.num_computed_tokens = 4
1847 + scheduler_output = _step_output([request.request_id], [4], [1])
1848 + metadata = connector.build_connector_meta(
1849 + scheduler_output, {request.request_id: request}
1850 + )
1851 +
1852 + assert metadata.generation == 1
1853 + assert metadata.requests[request.request_id] == 4
1854 + assert list(metadata.block_hashes[request.request_id]) == [b"a" * 32]
tests/entrypoints/openai/test_return_routed_experts.py
+13 -6
@@ -12,15 +12,13 @@ from ...utils import RemoteOpenAIServer
12 MODEL_NAME = "TitanML/tiny-mixtral"
13
14 # tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers.
15 -# The published config has sliding_window=4096, which produces
16 -# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a
17 -# FullAttentionSpec group, so we override sliding_window=null below.
15 +# Its published sliding window is incompatible with AuxOutput retention.
16 NUM_LOCAL_EXPERTS = 8
17 NUM_EXPERTS_PER_TOK = 2
18 NUM_HIDDEN_LAYERS = 2
19
20
23 -def assert_valid_routed_experts(encoded: str | None) -> None:
21 +def decode_routed_experts(encoded: str | None) -> np.ndarray:
22 assert encoded is not None
23 routed_experts = np.load(io.BytesIO(base64.b64decode(encoded)))
24 assert routed_experts.ndim == 3
@@ -30,6 +28,7 @@ def assert_valid_routed_experts(encoded: str | None) -> None:
28 assert topk == NUM_EXPERTS_PER_TOK
29 assert (routed_experts >= 0).all()
30 assert (routed_experts < NUM_LOCAL_EXPERTS).all()
31 + return routed_experts
32
33
34 @pytest.fixture(scope="module")
@@ -44,7 +43,11 @@ def server():
43 "--hf-overrides",
44 '{"sliding_window": null}',
45 ]
47 - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
46 + with RemoteOpenAIServer(
47 + MODEL_NAME,
48 + args,
49 + env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"},
50 + ) as remote_server:
51 yield remote_server
52
53
@@ -63,4 +66,8 @@ async def test_routed_experts(server):
66 choice = result.model_dump()["choices"][0]
67
68 assert choice["token_ids"] is not None
66 - assert_valid_routed_experts(choice["routed_experts"])
69 + assert result.usage is not None
70 + routed_experts = decode_routed_experts(choice["routed_experts"])
71 + assert len(routed_experts) == (
72 + result.usage.prompt_tokens + len(choice["token_ids"]) - 1
73 + )
tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py
+11 -5
@@ -15,9 +15,7 @@ MODEL_NAME = "TitanML/tiny-mixtral"
15 GEN_ENDPOINT = "/inference/v1/generate"
16
17 # tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers.
18 -# The published config has sliding_window=4096, which produces
19 -# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a
20 -# FullAttentionSpec group, so we override sliding_window=null below.
18 +# Its published sliding window is incompatible with AuxOutput retention.
19 NUM_LOCAL_EXPERTS = 8
20 NUM_EXPERTS_PER_TOK = 2
21 NUM_HIDDEN_LAYERS = 2
@@ -36,7 +34,13 @@ def server():
34 '{"sliding_window": null}',
35 "--enable-scale-out",
36 ]
39 - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
37 + with RemoteOpenAIServer(
38 + MODEL_NAME,
39 + args,
40 + env_dict={
41 + "VLLM_USE_V2_MODEL_RUNNER": "1",
42 + },
43 + ) as remote_server:
44 yield remote_server
45
46
@@ -56,9 +60,10 @@ async def client(server: RemoteOpenAIServer):
60 @pytest.mark.asyncio
61 async def test_generate_routed_experts(client):
62 """Test that /inference/v1/generate returns routed_experts when enabled."""
63 + prompt_token_ids = [1, 2, 3]
64 payload = {
65 "model": MODEL_NAME,
61 - "token_ids": [1, 2, 3],
66 + "token_ids": prompt_token_ids,
67 "sampling_params": {"max_tokens": 10, "temperature": 0.0},
68 "stream": False,
69 }
@@ -80,3 +85,4 @@ async def test_generate_routed_experts(client):
85 assert topk == NUM_EXPERTS_PER_TOK
86 assert (routed_experts >= 0).all()
87 assert (routed_experts < NUM_LOCAL_EXPERTS).all()
88 + assert num_tokens == len(prompt_token_ids) + len(choice["token_ids"]) - 1
tests/model_executor/test_routed_experts_capture.py
+218 -194
@@ -7,26 +7,19 @@ from unittest.mock import Mock, patch
7 import pytest
8 import torch
9
10 -from vllm.config import ModelConfig, VllmConfig
10 +from vllm.config import VllmConfig
11 from vllm.config.compilation import CompilationMode
12 from vllm.distributed.eplb.eplb_state import EplbLayerState
13 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
14 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
15 RoutedExpertsCapturer,
16 - RoutedExpertsManager,
16 bind_routed_experts_capturer,
18 - get_routed_experts_attn_gid,
17 )
18 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
21 -from vllm.transformers_utils.model_arch_config_convertor import (
22 - ModelArchConfigConvertorBase,
23 -)
19 from vllm.v1.kv_cache_interface import (
20 FullAttentionSpec,
26 - KVCacheConfig,
21 KVCacheGroupSpec,
22 MLAAttentionSpec,
29 - SlidingWindowMLASpec,
23 UniformTypeKVCacheSpecs,
24 )
25
@@ -42,6 +35,7 @@ def _capturer_with_buffer(
35 num_experts_per_tok: int = 2,
36 dp_rank: int = 0,
37 tp_size: int = 1,
38 + dtype: torch.dtype = torch.int32,
39 ) -> RoutedExpertsCapturer:
40 # Bypass __init__ so the test can use a CPU buffer and skip the
41 # VllmConfig dependency. The CUDA device-tensor allocation in the
@@ -49,10 +43,11 @@ def _capturer_with_buffer(
43 c = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
44 c.dp_rank = dp_rank
45 c.tp_size = tp_size
46 + c.output_dtype = torch.uint8
47 c.device_buffer = torch.full(
48 (max_tokens, num_layers, num_experts_per_tok),
49 -1,
55 - dtype=torch.int32,
50 + dtype=dtype,
51 )
52 return c
53
@@ -88,74 +83,36 @@ def _make_modular_routed_experts():
83 )
84
85
91 -def _make_model_config(hf_config):
92 - num_experts_per_token = ModelArchConfigConvertorBase(
93 - hf_config, hf_config
94 - ).get_num_experts_per_token()
95 - model_config = SimpleNamespace(
96 - hf_text_config=hf_config,
97 - model_arch_config=SimpleNamespace(
98 - num_experts_per_token=num_experts_per_token,
99 - ),
100 - )
101 - model_config.get_num_experts = lambda: hf_config.num_experts
102 - model_config.get_num_experts_per_tok = lambda: (
103 - ModelConfig.get_num_experts_per_tok(model_config)
104 - )
105 - model_config.get_total_num_hidden_layers = lambda: hf_config.num_hidden_layers
106 - return model_config
107 -
108 -
109 -def test_routed_experts_manager_uses_gemma4_top_k_experts():
110 - hf_config = SimpleNamespace(
111 - num_experts=8,
112 - top_k_experts=2,
113 - num_hidden_layers=3,
114 - )
115 - vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
116 - kv_cache_spec = FullAttentionSpec(
117 - block_size=4,
118 - num_kv_heads=1,
119 - head_size=1,
120 - dtype=torch.float32,
121 - )
122 - kv_cache_config = KVCacheConfig(
123 - num_blocks=2,
124 - kv_cache_tensors=[],
125 - kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
126 - )
127 -
128 - manager = RoutedExpertsManager(vllm_config, kv_cache_config)
129 -
130 - assert manager.routed_experts_by_slot.shape == (8, 3, 2)
131 -
132 -
133 -def test_routed_experts_manager_uses_kimi_k3_experts_per_token():
134 - hf_config = SimpleNamespace(
135 - num_experts=8,
136 - num_experts_per_token=2,
137 - num_hidden_layers=3,
138 - )
139 - vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
140 - kv_cache_spec = FullAttentionSpec(
141 - block_size=4,
86 +def _full_attention_kv_group(
87 + spec_type: type[FullAttentionSpec] = FullAttentionSpec,
88 +) -> KVCacheGroupSpec:
89 + full_attention = spec_type(
90 + block_size=16,
91 num_kv_heads=1,
92 head_size=1,
93 dtype=torch.float32,
94 )
146 - kv_cache_config = KVCacheConfig(
147 - num_blocks=2,
148 - kv_cache_tensors=[],
149 - kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
95 + return KVCacheGroupSpec(
96 + ["layer"],
97 + UniformTypeKVCacheSpecs(
98 + block_size=16,
99 + kv_cache_specs={"layer": full_attention},
100 + ),
101 )
102
152 - manager = RoutedExpertsManager(vllm_config, kv_cache_config)
153 -
154 - assert manager.routed_experts_by_slot.shape == (8, 3, 2)
103
104 +@pytest.mark.parametrize("eplb_enabled", [False, True])
105 +def test_base_router_capture_pre_eplb_mapping(eplb_enabled):
106 + eplb_state = None
107 + if eplb_enabled:
108 + eplb_state = EplbLayerState()
109 + eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
110 + eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
111 + eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
112 + eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
113 + eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)]
114 + router = _make_router(eplb_state)
115
157 -def test_base_router_capture_pre_eplb_mapping():
158 - router = _make_router()
116 captured = []
117
118 def capture_fn(ids):
@@ -173,34 +130,7 @@ def test_base_router_capture_pre_eplb_mapping():
130 assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
131
132
176 -def test_base_router_capture_with_eplb_enabled():
177 - eplb_state = EplbLayerState()
178 - eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
179 - eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
180 - eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
181 - eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
182 - eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)]
183 - router = _make_router(eplb_state=eplb_state)
184 -
185 - captured = []
186 -
187 - def capture_fn(ids):
188 - captured.append(ids.clone())
189 -
190 - router.set_capture_fn(capture_fn)
191 - _, topk_ids = router.select_experts(
192 - hidden_states=torch.empty(1),
193 - router_logits=torch.empty(1),
194 - )
195 -
196 - assert len(captured) == 1
197 - # Capture should see logical ids pre-EPLB mapping.
198 - assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
199 - # Our DummyRouter mapping adds +10.
200 - assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
201 -
202 -
203 -def test_public_binding_only_visits_target_model(monkeypatch):
133 +def test_public_binding_binds_target_model_router(monkeypatch):
134 class DummyFusedMoE:
135 def __init__(self, layer_id):
136 self.layer_id = layer_id
@@ -208,7 +138,6 @@ def test_public_binding_only_visits_target_model(monkeypatch):
138 self._quant_method = _make_modular_routed_experts().quant_method
139
140 target_module = DummyFusedMoE(layer_id=7)
211 - draft_module = DummyFusedMoE(layer_id=0)
141
142 import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
143
@@ -221,12 +150,30 @@ def test_public_binding_only_visits_target_model(monkeypatch):
150 )
151
152 assert target_module.router.capture_fn is not None
224 - assert draft_module.router.capture_fn is None
153 topk_ids = torch.tensor([[5, 6]])
154 target_module.router.capture_fn(topk_ids)
155 assert calls == [(7, topk_ids)]
156
157
158 +def test_public_binding_supports_direct_capture_source():
159 + class DirectCaptureSource(torch.nn.Module):
160 + def __init__(self):
161 + super().__init__()
162 + self.layer_id = 3
163 + self.capture_fn = None
164 +
165 + source = DirectCaptureSource()
166 + calls = []
167 + capturer = types.SimpleNamespace(capture=lambda *args: calls.append(args))
168 +
169 + bind_routed_experts_capturer(source, capturer)
170 +
171 + assert source.capture_fn is not None
172 + topk_ids = torch.tensor([[1, 2]])
173 + source.capture_fn(topk_ids)
174 + assert calls == [(3, topk_ids)]
175 +
176 +
177 def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch):
178 class DummyFusedMoE:
179 def __init__(self):
@@ -277,6 +224,51 @@ def test_routed_experts_capturer_single_dp_no_metadata():
224 assert capturer.device_buffer[3, 0, 0].item() == -1
225
226
227 +@pytest.mark.parametrize(
228 + ("num_experts", "dtype_name", "torch_dtype"),
229 + [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)],
230 +)
231 +def test_routed_experts_capturer_exposes_output_profile(
232 + monkeypatch, num_experts, dtype_name, torch_dtype
233 +):
234 + import vllm.model_executor.layers.fused_moe.routed_experts_capturer as module
235 +
236 + monkeypatch.setattr(module, "current_platform", SimpleNamespace(device_type="cpu"))
237 + config = SimpleNamespace(
238 + model_config=SimpleNamespace(
239 + get_total_num_hidden_layers=lambda: 3,
240 + get_num_experts=lambda: num_experts,
241 + get_num_experts_per_tok=lambda: 2,
242 + hf_text_config=SimpleNamespace(model_type="test"),
243 + ),
244 + parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1),
245 + )
246 +
247 + capturer = RoutedExpertsCapturer(8, config)
248 +
249 + assert capturer.shape_per_token == (3, 2)
250 + assert capturer.output_dtype_name == dtype_name
251 + assert capturer.output_dtype == torch_dtype
252 + assert capturer.device_buffer.shape == (8, 3, 2)
253 +
254 +
255 +@pytest.mark.parametrize("output_dtype", [torch.uint8, torch.uint16])
256 +def test_routed_experts_capturer_narrows_snapshot(output_dtype):
257 + capturer = _capturer_with_buffer(dtype=torch.int32)
258 + capturer.output_dtype = output_dtype
259 + topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64)
260 + ctx = SimpleNamespace(dp_metadata=None)
261 + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
262 + capturer.capture(layer_id=0, topk_ids=topk)
263 +
264 + output = capturer.snapshot_routing_data(2)
265 + assert capturer.device_buffer.dtype == torch.int32
266 + assert capturer.device_buffer[:2, 0, :].tolist() == topk.tolist()
267 + assert output.dtype == output_dtype
268 + assert output[:, 0, :].tolist() == topk.tolist()
269 + assert output.data_ptr() != capturer.device_buffer.data_ptr()
270 +
271 +
272 def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
273 """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
274 capturer = _capturer_with_buffer(dp_rank=1)
@@ -307,6 +299,47 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
299 assert torch.equal(capturer.device_buffer[:3, 0, :], topk)
300
301
302 +def test_routed_experts_capturer_dp_ep_gathered_shards():
303 + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
304 + num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32)
305 + ctx = SimpleNamespace(
306 + dp_metadata=SimpleNamespace(
307 + num_tokens_across_dp_cpu=num_tokens_dp,
308 + local_sizes=[2, 2, 1, 1],
309 + )
310 + )
311 + topk = torch.tensor(
312 + [[0, 1], [2, 3], [4, 5], [-1, -1], [10, 11], [12, 13]],
313 + dtype=torch.int32,
314 + )
315 +
316 + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
317 + capturer.capture(layer_id=0, topk_ids=topk)
318 +
319 + assert torch.equal(capturer.device_buffer[:2, 0], topk[4:])
320 +
321 +
322 +def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
323 + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
324 + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
325 + ctx = SimpleNamespace(
326 + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
327 + )
328 + local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
329 + gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32)
330 + tp_group = Mock()
331 + tp_group.all_gather.return_value = gathered
332 +
333 + with (
334 + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
335 + patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
336 + ):
337 + capturer.capture(layer_id=0, topk_ids=local_shard)
338 +
339 + tp_group.all_gather.assert_called_once_with(local_shard, dim=0)
340 + assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3])
341 +
342 +
343 def test_routed_experts_capturer_dp_unexpected_batch_raises():
344 """Mismatch between topk batch dim and DP layout: fail fast."""
345 capturer = _capturer_with_buffer(dp_rank=0)
@@ -324,128 +357,119 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises():
357 assert capturer.device_buffer[0, 0, 0].item() == -1
358
359
327 -def test_routed_experts_attention_group_is_shared_and_fail_closed():
328 - """Both sides key routing data by this gid, so it must skip non-full-attention
329 - groups rather than defaulting to 0, and fail closed when none exists."""
330 - common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
331 - config = SimpleNamespace(
332 - kv_cache_groups=[
333 - KVCacheGroupSpec(
334 - ["swa_layer"],
335 - SlidingWindowMLASpec(block_size=4, sliding_window=8, **common),
336 - ),
337 - KVCacheGroupSpec(["full_layer"], FullAttentionSpec(block_size=4, **common)),
338 - ]
339 - )
340 - assert get_routed_experts_attn_gid(config) == 1
341 -
342 - with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
343 - get_routed_experts_attn_gid(SimpleNamespace(kv_cache_groups=[]))
344 -
345 -
346 -def test_routed_experts_attention_group_unwraps_uniform_type_specs():
347 - """DeepSeek-V4-shaped groups wrap their specs in ``UniformTypeKVCacheSpecs``.
360 +def test_get_aux_output_connector_passes_config(monkeypatch):
361 + import vllm.distributed.aux_output_connector.worker as aux_output_worker
362
349 - The wrapper is not a ``FullAttentionSpec``, so a bare isinstance check finds
350 - no group and fails closed on every worker. Unwrap semantics themselves are
351 - covered by ``test_is_full_attention_spec_*`` in tests/v1/core.
352 - """
353 - common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
354 - swa_spec = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common)
355 - mla_spec = MLAAttentionSpec(block_size=4, **common)
363 + connector = Mock()
364 + constructor = Mock(return_value=connector)
365 + monkeypatch.setattr(aux_output_worker, "AuxOutputWorkerConnector", constructor)
366 config = SimpleNamespace(
357 - kv_cache_groups=[
358 - KVCacheGroupSpec(
359 - ["swa_layer"],
360 - UniformTypeKVCacheSpecs(
361 - block_size=4, kv_cache_specs={"swa_layer": swa_spec}
362 - ),
363 - ),
364 - KVCacheGroupSpec(
365 - ["mla_layer"],
366 - UniformTypeKVCacheSpecs(
367 - block_size=4, kv_cache_specs={"mla_layer": mla_spec}
368 - ),
369 - ),
370 - ]
367 + scheduler_config=SimpleNamespace(max_num_batched_tokens=32)
368 )
369 + model = Mock()
370 + kv_cache_config = Mock()
371
373 - assert get_routed_experts_attn_gid(config) == 1
372 + result = aux_output_worker.get_aux_output_connector(model, config, kv_cache_config)
373
375 - swa_only = SimpleNamespace(kv_cache_groups=[config.kv_cache_groups[0]])
376 - with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
377 - get_routed_experts_attn_gid(swa_only)
378 -
379 -
380 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
381 -def test_mrv2_async_output_returns_existing_routed_experts_field():
382 - from vllm.v1.outputs import ModelRunnerOutput, RoutedExpertsTensors
383 - from vllm.v1.worker.gpu.async_utils import AsyncOutput
384 - from vllm.v1.worker.gpu.sample.output import SamplerOutput
385 -
386 - routed_experts = RoutedExpertsTensors(
387 - routing_data=torch.arange(6, dtype=torch.int32, device="cuda").reshape(3, 1, 2),
388 - slot_mapping=torch.tensor([11, 12, 13], device="cuda"),
389 - )
390 - num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda")
391 - sampler_output = SamplerOutput(
392 - sampled_token_ids=torch.tensor([[1]], device="cuda"),
393 - logprobs_tensors=None,
394 - num_nans=None,
395 - num_sampled=num_sampled,
396 - num_rejected=torch.tensor([0], dtype=torch.int32, device="cuda"),
374 + constructor.assert_called_once_with(
375 + model=model,
376 + kv_cache_config=kv_cache_config,
377 + vllm_config=config,
378 )
398 - output = AsyncOutput(
399 - model_runner_output=ModelRunnerOutput(req_ids=["req"], req_id_to_index={}),
400 - sampler_output=sampler_output,
401 - num_sampled_tokens=num_sampled,
402 - main_stream=torch.cuda.current_stream(),
403 - copy_stream=torch.cuda.Stream(),
404 - routed_experts=routed_experts,
405 - ).get_output()
406 -
407 - assert output.routed_experts is not None
408 - assert output.routed_experts.routing_data[:, 0, 0].tolist() == [0, 2, 4]
409 - assert output.routed_experts.slot_mapping.tolist() == [11, 12, 13]
379 + assert result is connector
380
381
412 -@pytest.mark.parametrize("rank", [0, 1])
413 -def test_all_tp_ranks_initialize_capture(monkeypatch, rank):
414 - pytest.importorskip("vllm.vllm_flash_attn", exc_type=ImportError)
415 - import vllm.v1.worker.gpu.model_runner as model_runner
382 +def test_aux_output_worker_connector_binds_capture_on_non_output_rank(monkeypatch):
383 + import vllm.distributed.aux_output_connector.worker as aux_output_worker
384
385 capturer = Mock()
386 constructor = Mock(return_value=capturer)
387 bind = Mock()
420 - monkeypatch.setattr(model_runner, "RoutedExpertsCapturer", constructor)
421 - monkeypatch.setattr(model_runner, "bind_routed_experts_capturer", bind)
422 -
423 - runner = model_runner.GPUModelRunner.__new__(model_runner.GPUModelRunner)
424 - runner.max_num_tokens = 32
425 - runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(rank=rank))
426 - runner.kv_cache_config = SimpleNamespace()
427 - runner.model = Mock()
388 + monkeypatch.setattr(aux_output_worker, "RoutedExpertsCapturer", constructor)
389 + monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", bind)
390 + monkeypatch.setattr(
391 + aux_output_worker,
392 + "get_tp_group",
393 + lambda: SimpleNamespace(is_first_rank=False, world_size=1),
394 + )
395
429 - runner.init_routed_experts_capturer()
396 + config = SimpleNamespace(
397 + aux_output_config=SimpleNamespace(enable_return_routed_experts=True),
398 + kv_transfer_config=None,
399 + max_concurrent_batches=2,
400 + scheduler_config=SimpleNamespace(max_num_batched_tokens=32),
401 + )
402 + model = Mock()
403 + connector = aux_output_worker.AuxOutputWorkerConnector(
404 + vllm_config=config,
405 + model=model,
406 + kv_cache_config=SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()]),
407 + )
408
409 constructor.assert_called_once_with(
410 max_num_batched_tokens=32,
433 - vllm_config=runner.vllm_config,
434 - kv_cache_config=runner.kv_cache_config,
411 + vllm_config=config,
412 + )
413 + bind.assert_called_once_with(model, capturer)
414 + connector.begin_step(Mock())
415 + assert connector.prepare_output(Mock()) is None
416 + capturer.snapshot_routing_data.assert_not_called()
417 +
418 +
419 +def test_aux_output_worker_connector_default_capacity(monkeypatch):
420 + import vllm.distributed.aux_output_connector.worker as aux_output_worker
421 +
422 + tp_group = SimpleNamespace(is_first_rank=True, world_size=1)
423 + store_constructor = Mock()
424 + background_store_constructor = Mock(side_effect=lambda store, **_: store)
425 + capturer = SimpleNamespace(shape_per_token=(2,), output_dtype_name="int32")
426 + monkeypatch.setattr(aux_output_worker, "get_tp_group", lambda: tp_group)
427 + monkeypatch.setattr(
428 + aux_output_worker, "RoutedExpertsCapturer", Mock(return_value=capturer)
429 + )
430 + monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", Mock())
431 + monkeypatch.setattr(
432 + aux_output_worker,
433 + "resolve_kv_cache_block_sizes",
434 + lambda *_: (32, 16),
435 + )
436 + monkeypatch.setattr(aux_output_worker, "BlockObjectStore", store_constructor)
437 + monkeypatch.setattr(
438 + aux_output_worker, "BackgroundBlockObjectStore", background_store_constructor
439 )
436 - bind.assert_called_once_with(runner.model, capturer)
437 - assert runner.routed_experts_capturer is capturer
440 + monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", Mock())
441 +
442 + config = SimpleNamespace(
443 + aux_output_config=SimpleNamespace(max_bytes=None),
444 + kv_transfer_config=None,
445 + cache_config=SimpleNamespace(enable_prefix_caching=True),
446 + scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=32),
447 + max_concurrent_batches=2,
448 + )
449 + kwargs = dict(
450 + vllm_config=config,
451 + model=Mock(),
452 + kv_cache_config=SimpleNamespace(
453 + num_blocks=10,
454 + kv_cache_groups=[_full_attention_kv_group(MLAAttentionSpec)],
455 + ),
456 + )
457 +
458 + aux_output_worker.AuxOutputWorkerConnector(**kwargs)
459 + assert store_constructor.call_args.kwargs["max_bytes"] == 2560
460 + assert store_constructor.call_args.kwargs["object_nbytes"] == 128
461 + assert background_store_constructor.call_args.kwargs["max_pending_batches"] == 16
462
463
464 def test_v2_model_runner_accepts_routed_experts(monkeypatch):
465 monkeypatch.setattr("importlib.metadata.entry_points", lambda **_: ())
466 config = SimpleNamespace(
467 model_config=SimpleNamespace(
444 - enable_return_routed_experts=True,
468 use_mla=False,
469 logits_processors=None,
470 enable_prompt_embeds=False,
471 ),
472 + aux_output_config=SimpleNamespace(enable_return_routed_experts=True),
473 speculative_config=None,
474 parallel_config=SimpleNamespace(
475 prefill_context_parallel_size=1,
tests/v1/core/test_async_scheduler.py
+1 -2
@@ -354,8 +354,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
354 scheduler.finished_req_ids_dict = None
355 scheduler.grammar_compile_error_reqs = set()
356 scheduler.vllm_config = Mock()
357 - scheduler.vllm_config.model_config.enable_return_routed_experts = False
358 - scheduler.enable_return_routed_experts = False
357 + scheduler.aux_output_connector = None
358 scheduler.return_sampling_mask = False
359 scheduler.recompute_kv_load_failures = False
360 scheduler.defer_block_free = False
tests/v1/core/test_scheduler.py
+71 -23
@@ -19,6 +19,10 @@ from vllm.config import (
19 SpeculativeConfig,
20 VllmConfig,
21 )
22 +from vllm.distributed.aux_output_connector.connector import (
23 + AuxOutputSchedulerConnector,
24 + AuxRequestOutput,
25 +)
26 from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats
27 from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
28 HiSparseConnector,
@@ -55,7 +59,6 @@ from vllm.v1.outputs import (
59 ECConnectorOutput,
60 KVConnectorOutput,
61 ModelRunnerOutput,
58 - RoutedExpertsLists,
62 SamplingMaskLists,
63 make_empty_encoder_model_runner_output,
64 )
@@ -136,21 +139,18 @@ def test_add_requests():
139 assert len(scheduler.waiting) == i + 1
140
141
139 -def test_routed_experts_prompt_start_at_prompt_end():
142 +@pytest.mark.parametrize("max_tokens", [1, 2])
143 +def test_routed_experts_prompt_start_at_prompt_end(max_tokens):
144 """A prompt-end offset should return empty routing data, not crash."""
145 scheduler = create_scheduler()
142 - (request,) = create_requests(num_requests=1)
146 + (request,) = create_requests(num_requests=1, max_tokens=max_tokens)
147 request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens
148 + scheduler.aux_output_connector = AuxOutputSchedulerConnector()
149 scheduler.add_request(request)
150 scheduler_output = scheduler.schedule()
146 - num_tokens = scheduler_output.num_scheduled_tokens[request.request_id]
147 -
148 - manager = Mock()
149 - manager.routed_experts_by_slot = np.empty((0, 1, 1), dtype=np.uint8)
150 - manager.get.return_value = np.empty((0, 1, 1), dtype=np.uint8)
151 - scheduler.enable_return_routed_experts = True
152 - scheduler.routed_experts_mgr = manager
153 - scheduler._re_block_ids = {request.request_id: []}
151 + assert scheduler_output.aux_output_connector_metadata.requests == {
152 + request.request_id: request.num_prompt_tokens
153 + }
154
155 outputs = scheduler.update_from_output(
156 scheduler_output,
@@ -161,21 +161,23 @@ def test_routed_experts_prompt_start_at_prompt_end():
161 logprobs=None,
162 prompt_logprobs_dict={},
163 pooler_output=[],
164 - routed_experts=RoutedExpertsLists(
165 - routing_data=np.zeros((num_tokens, 1, 1), dtype=np.uint8),
166 - slot_mapping=np.arange(num_tokens),
167 - ),
164 + aux_output_connector_output={
165 + request.request_id: AuxRequestOutput(
166 + request.num_prompt_tokens,
167 + np.empty((0, 1, 1), dtype=np.uint8),
168 + )
169 + },
170 ),
171 )
172
171 - manager.get.assert_called_once_with(
172 - [], request.num_prompt_tokens, token_start=request.num_prompt_tokens
173 - )
174 - assert outputs[request.client_index].outputs[0].routed_experts.shape == (0, 1, 1)
173 + output = outputs[request.client_index].outputs[0]
174 + assert output.new_token_ids == [0]
175 + assert output.routed_experts.shape == (0, 1, 1)
176
177
178 def test_finish_request():
179 scheduler = create_scheduler()
180 + scheduler.aux_output_connector = Mock()
181 requests = create_requests(num_requests=10)
182 for request in requests:
183 scheduler.add_request(request)
@@ -184,6 +186,7 @@ def test_finish_request():
186 scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
187 assert request.request_id not in scheduler.requests
188 assert len(scheduler.waiting) == 9 - i
189 + scheduler.aux_output_connector.request_finished.assert_called_with(request)
190
191
192 def test_get_num_unfinished_requests():
@@ -1373,6 +1376,31 @@ def test_preemption_re_records_prefix_cache_query():
1376 assert stats.preempted_requests == 1
1377
1378
1379 +def test_preemption_processes_stale_aux_output():
1380 + scheduler = create_scheduler(enable_prefix_caching=True)
1381 + request = create_requests(num_requests=1)[0]
1382 + scheduler.add_request(request)
1383 + scheduler_output = scheduler.schedule()
1384 + scheduler.running.remove(request)
1385 + scheduler.aux_output_connector = Mock()
1386 + scheduler._preempt_request(request, 0.0)
1387 + scheduler.aux_output_connector.request_finished.assert_called_once_with(request)
1388 +
1389 + scheduler.update_from_output(
1390 + scheduler_output,
1391 + ModelRunnerOutput(
1392 + req_ids=[request.request_id],
1393 + req_id_to_index={request.request_id: 0},
1394 + sampled_token_ids=[[1000]],
1395 + logprobs=None,
1396 + prompt_logprobs_dict={},
1397 + pooler_output=[],
1398 + ),
1399 + )
1400 +
1401 + scheduler.aux_output_connector.take_output.assert_called_once_with(request, None)
1402 +
1403 +
1404 def test_prefix_cache_stats_not_recorded_when_caching_disabled():
1405 """With prefix caching off there is no local lookup, so admitting a request
1406 records no phantom miss."""
@@ -1461,6 +1489,7 @@ def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request():
1489
1490 def test_scheduler_reset_prefix_cache():
1491 scheduler = create_scheduler(enable_prefix_caching=True)
1492 + scheduler.aux_output_connector = Mock()
1493 requests = create_requests(num_requests=10)
1494 for request in requests:
1495 scheduler.add_request(request)
@@ -1477,10 +1506,20 @@ def test_scheduler_reset_prefix_cache():
1506 # Reset prefix cache should fail since there are still running requests
1507 # and they are taking KV cache
1508 assert not scheduler.reset_prefix_cache()
1509 + scheduler.aux_output_connector.reset.assert_not_called()
1510 +
1511 + with pytest.raises(RuntimeError, match=r"pause\(mode='keep'\)"):
1512 + scheduler.reset_prefix_cache(reset_running_requests=True)
1513 +
1514 + # pause(mode="keep") also waits for scheduled model outputs to drain.
1515 + scheduler.set_pause_state(PauseState.PAUSED_ALL)
1516 + with pytest.raises(RuntimeError, match="model output is in flight"):
1517 + scheduler.reset_prefix_cache(reset_running_requests=True)
1518 + for request in requests:
1519 + request.num_in_flight_tokens = 0
1520
1481 - # Reset prefix cache with reset_running_requests=True. All running requests
1482 - # Should be pushed back to the waiting queue and kv cache should be freed
1521 assert scheduler.reset_prefix_cache(reset_running_requests=True)
1522 + scheduler.aux_output_connector.reset.assert_called_once_with()
1523
1524 # Verify requests moved from running to waiting
1525 assert len(scheduler.waiting) == len(requests)
@@ -1490,6 +1529,16 @@ def test_scheduler_reset_prefix_cache():
1529 assert scheduler.waiting[i] == request
1530
1531
1532 +@pytest.mark.parametrize("reset_successful", [False, True])
1533 +def test_aux_output_reset_follows_kv_reset_result(reset_successful: bool):
1534 + scheduler = create_scheduler(enable_prefix_caching=True)
1535 + scheduler.aux_output_connector = Mock()
1536 + scheduler.kv_cache_manager.reset_prefix_cache = Mock(return_value=reset_successful)
1537 +
1538 + assert scheduler.reset_prefix_cache() is reset_successful
1539 + assert scheduler.aux_output_connector.reset.call_count == int(reset_successful)
1540 +
1541 +
1542 def test_kv_cache_release_after_keep_pause_preserves_requests():
1543 """A completed keep pause permits release and recomputation of live requests."""
1544 scheduler = create_scheduler(max_num_seqs=1, enable_prefix_caching=True)
@@ -3791,10 +3840,9 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
3840 scheduler.kv_event_publisher = Mock()
3841 scheduler.finished_req_ids = set()
3842 scheduler.finished_req_ids_dict = None
3843 + scheduler.aux_output_connector = None
3844 scheduler.grammar_compile_error_reqs = set()
3845 scheduler.vllm_config = Mock()
3796 - scheduler.vllm_config.model_config.enable_return_routed_experts = False
3797 - scheduler.enable_return_routed_experts = False
3846 scheduler.return_sampling_mask = False
3847 scheduler.recompute_kv_load_failures = False
3848 scheduler.defer_block_free = False
tests/v1/e2e/general/test_mamba_prefix_cache.py
+1
@@ -1122,6 +1122,7 @@ def _run_mamba_prefix_cache_mrv2(
1122 logprobs_tensors=None,
1123 num_nans=None,
1124 num_sampled=num_sampled,
1125 + num_rejected=num_rejected,
1126 )
1127 return sampler_output, num_sampled, num_rejected
1128
tests/v1/engine/test_output_processor.py
+47 -1
@@ -3,8 +3,9 @@
3
4 import math
5 import time
6 -from unittest.mock import MagicMock
6 +from unittest.mock import MagicMock, Mock
7
8 +import numpy as np
9 import pytest
10
11 from tests.v1.engine.utils import (
@@ -1307,6 +1308,7 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors):
1308 async def test_request_output_collector():
1309 NUM_REQS = 3
1310 TEXT = "a"
1311 + routed_experts = np.arange(12, dtype=np.uint8).reshape(2, 3, 2)
1312
1313 def make_outputs() -> list[RequestOutput]:
1314 return [
@@ -1322,6 +1324,9 @@ async def test_request_output_collector():
1324 token_ids=[idx],
1325 cumulative_logprob=(idx + 1 * 1.0),
1326 logprobs=[{"a": idx, "b": idx}],
1327 + routed_experts=(
1328 + routed_experts if idx == NUM_REQS - 1 else None
1329 + ),
1330 finish_reason="length" if (idx == NUM_REQS - 1) else None,
1331 )
1332 ],
@@ -1374,6 +1379,7 @@ async def test_request_output_collector():
1379
1380 assert output.finished
1381 assert output.outputs[0].finish_reason == "length"
1382 + np.testing.assert_array_equal(output.outputs[0].routed_experts, routed_experts)
1383 # Text, token_ids, and logprobs should get merged.
1384 assert output.outputs[0].text == TEXT * num_to_put
1385 for tok_0, tok_1 in zip(output.outputs[0].token_ids, list(range(num_to_put))):
@@ -1385,6 +1391,46 @@ async def test_request_output_collector():
1391 assert output.outputs[0].cumulative_logprob == cumulative_logprob_expected
1392
1393
1394 +def test_routed_experts_are_accumulated_until_finish():
1395 + state = RequestState(
1396 + request_id="request-int",
1397 + external_req_id="request",
1398 + parent_req=None,
1399 + request_index=0,
1400 + lora_request=None,
1401 + output_kind=RequestOutputKind.DELTA,
1402 + prompt="prompt",
1403 + prompt_token_ids=[1, 2, 3, 4],
1404 + prompt_embeds=None,
1405 + logprobs_processor=Mock(
1406 + logprobs=None,
1407 + cumulative_logprob=0.0,
1408 + pop_prompt_logprobs=Mock(return_value=None),
1409 + ),
1410 + detokenizer=Mock(get_next_output_text=Mock(return_value="")),
1411 + max_tokens_param=2,
1412 + arrival_time=0.0,
1413 + queue=None,
1414 + log_stats=False,
1415 + stream_interval=1,
1416 + )
1417 + prompt_chunk = np.arange(24, dtype=np.uint8).reshape(4, 3, 2)
1418 + decode_chunk = np.arange(12, dtype=np.uint8).reshape(2, 3, 2) + 24
1419 + state.routed_experts_chunks.append(prompt_chunk)
1420 +
1421 + partial = state.make_request_output([5], None, None, None)
1422 + assert partial is not None
1423 + assert partial.outputs[0].routed_experts is None
1424 +
1425 + state.routed_experts_chunks.append(decode_chunk)
1426 + finished = state.make_request_output([6], None, FinishReason.LENGTH, None)
1427 + assert finished is not None
1428 + np.testing.assert_array_equal(
1429 + finished.outputs[0].routed_experts,
1430 + np.concatenate((prompt_chunk, decode_chunk)),
1431 + )
1432 +
1433 +
1434 @pytest.mark.asyncio
1435 async def test_cumulative_output_collector_n():
1436 """Test collector correctly handles multiple outputs by index."""
tests/v1/executor/test_ray_utils.py
+12 -17
@@ -3,12 +3,14 @@
3
4 import numpy as np
5
6 +from vllm.distributed.aux_output_connector.connector import (
7 + AuxRequestOutput,
8 +)
9 from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output
10 from vllm.v1.outputs import (
11 LogprobsLists,
12 LogprobsTensors,
13 ModelRunnerOutput,
11 - RoutedExpertsLists,
14 )
15
16
@@ -59,26 +61,19 @@ def test_detach_zero_copy_from_model_runner_output_copies_only_numpy_views():
61 assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs
62
63
62 -def test_detach_zero_copy_routed_experts_without_logprobs():
64 +def test_detach_zero_copy_aux_output_without_logprobs():
65 + rows = _make_readonly(np.arange(12, dtype=np.uint8).reshape(2, 3, 2))
66 output = ModelRunnerOutput(
67 req_ids=["req-0"],
68 req_id_to_index={"req-0": 0},
66 - routed_experts=RoutedExpertsLists(
67 - routing_data=_make_readonly(np.arange(12, dtype=np.int32).reshape(2, 3, 2)),
68 - slot_mapping=_make_readonly(np.array([7, 8], dtype=np.int64)),
69 - ),
69 + aux_output_connector_output={"req-0": AuxRequestOutput(0, rows)},
70 )
71 - original = output.routed_experts
72 - assert output.logprobs is None
71
72 detach_zero_copy_from_model_runner_output(output)
73
76 - detached = output.routed_experts
77 - assert detached is not None
78 - assert detached is not original
79 - assert detached.routing_data is not original.routing_data
80 - assert detached.slot_mapping is not original.slot_mapping
81 - assert detached.routing_data.flags.writeable
82 - assert detached.slot_mapping.flags.writeable
83 - np.testing.assert_array_equal(detached.routing_data, original.routing_data)
84 - np.testing.assert_array_equal(detached.slot_mapping, original.slot_mapping)
74 + aux_output = output.aux_output_connector_output
75 + assert aux_output is not None
76 + detached = aux_output["req-0"].rows
77 + assert detached is not rows
78 + assert detached.flags.writeable
79 + np.testing.assert_array_equal(detached, rows)
vllm/config/__init__.py
+3
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from vllm.config.attention import AttentionConfig, HiSparseConfig
5 +from vllm.config.aux_output import AuxOutputConfig
6 from vllm.config.cache import CacheConfig
7 from vllm.config.compilation import (
8 CompilationConfig,
@@ -66,6 +67,8 @@ from vllm.config.weight_transfer import WeightTransferConfig
67 # __all__ should only contain classes and functions.
68 # Types and globals should be imported from their respective modules.
69 __all__ = [
70 + # From vllm.config.aux_output
71 + "AuxOutputConfig",
72 # From vllm.config.attention
73 "AttentionConfig",
74 "HiSparseConfig",
vllm/config/aux_output.py new
+31
@@ -0,0 +1,31 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Configuration for execution auxiliary outputs."""
4 +
5 +from pydantic import Field
6 +
7 +from vllm.config.utils import config
8 +
9 +
10 +@config
11 +class AuxOutputConfig:
12 + """Configuration for auxiliary-output delivery."""
13 +
14 + enable_return_routed_experts: bool = False
15 + """Capture and return routed-experts auxiliary outputs."""
16 +
17 + max_bytes: int | None = Field(default=None, gt=0)
18 + """LRU capacity, or ``None`` to derive it from the KV cache capacity."""
19 +
20 + @property
21 + def enabled(self) -> bool:
22 + """Whether any execution auxiliary output is enabled."""
23 + return self.enable_return_routed_experts
24 +
25 + def compute_hash(self) -> str:
26 + """Hash AuxOutput settings that alter the model forward graph."""
27 + from vllm.config.utils import hash_factors
28 +
29 + return hash_factors(
30 + {"enable_return_routed_experts": self.enable_return_routed_experts}
31 + )
vllm/config/model.py
-2
@@ -247,8 +247,6 @@ class ModelConfig:
247
248 NOTE: This disables both `torch.compile` and CUDA graphs, and is
249 equivalent to setting `-cc.mode=none -cc.cudagraph_mode=none`."""
250 - enable_return_routed_experts: bool = False
251 - """Whether to return routed experts."""
250 return_sampling_mask: bool = False
251 """Whether to return the post-processing token support for each sample."""
252 max_logprobs: int = Field(default=20, ge=-1)
vllm/config/vllm.py
+53 -33
@@ -28,6 +28,7 @@ from vllm.utils import random_uuid
28 from vllm.utils.hashing import safe_hash
29
30 from .attention import AttentionConfig, HiSparseConfig
31 +from .aux_output import AuxOutputConfig
32 from .cache import CacheConfig
33 from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode
34 from .device import DeviceConfig
@@ -376,6 +377,8 @@ class VllmConfig:
377 """Model weight offloading configuration."""
378 attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
379 """Attention configuration."""
380 + aux_output_config: AuxOutputConfig = Field(default_factory=AuxOutputConfig)
381 + """Execution auxiliary output configuration."""
382 engram_config: EngramConfig | None = None
383 """N-gram embedding storage and sharding settings."""
384 mamba_config: MambaConfig = Field(default_factory=MambaConfig)
@@ -554,6 +557,7 @@ class VllmConfig:
557 vllm_factors.append(self.ec_transfer_config.compute_hash())
558 else:
559 vllm_factors.append("None")
560 + vllm_factors.append(self.aux_output_config.compute_hash())
561 if self.additional_config:
562 if isinstance(additional_config := self.additional_config, dict):
563 additional_config_hash = safe_hash(
@@ -1084,6 +1088,53 @@ class VllmConfig:
1088 # This is the same for all backends
1089 self.kv_transfer_config.kv_role = "kv_both"
1090
1091 + def _verify_aux_output_compatibility(self) -> None:
1092 + """Reject configurations unsupported by enabled auxiliary outputs."""
1093 + if not self.aux_output_config.enabled:
1094 + return
1095 + if not self.use_v2_model_runner:
1096 + raise ValueError(
1097 + "AuxOutput Connector requires Model Runner V2; set "
1098 + "VLLM_USE_V2_MODEL_RUNNER=1."
1099 + )
1100 + if self.model_config.runner_type != "generate":
1101 + raise ValueError("AuxOutput Connector only supports generate runners.")
1102 + if not self.model_config.is_moe:
1103 + raise ValueError("AuxOutput Connector only supports MoE models.")
1104 + if not self.cache_config.enable_prefix_caching:
1105 + raise ValueError("AuxOutput Connector requires prefix caching.")
1106 + if (
1107 + self.speculative_config is not None
1108 + and self.speculative_config.enable_adaptive_verification
1109 + ):
1110 + raise ValueError(
1111 + "--enable-return-routed-experts is incompatible with "
1112 + "adaptive speculative verification."
1113 + )
1114 + if self.parallel_config.pipeline_parallel_size > 1:
1115 + raise ValueError(
1116 + "--enable-return-routed-experts is incompatible with "
1117 + "pipeline parallelism (PP > 1)."
1118 + )
1119 + if (
1120 + self.parallel_config.decode_context_parallel_size > 1
1121 + or self.parallel_config.prefill_context_parallel_size > 1
1122 + ):
1123 + raise ValueError(
1124 + "--enable-return-routed-experts is incompatible with "
1125 + "context parallelism (DCP/PCP > 1)."
1126 + )
1127 +
1128 + kv_transfer_config = self.kv_transfer_config
1129 + if (
1130 + kv_transfer_config is not None
1131 + and kv_transfer_config.is_kv_transfer_instance
1132 + ):
1133 + raise ValueError(
1134 + "--enable-return-routed-experts is incompatible with KV "
1135 + "connectors (PD disaggregation and KV cache offload)."
1136 + )
1137 +
1138 def _verify_kv_transfer_compat(self) -> None:
1139 """Reject configurations that silently corrupt KV transfers."""
1140 if (
@@ -1294,38 +1345,6 @@ class VllmConfig:
1345
1346 self.parallel_config.is_moe_model = self.model_config.is_moe
1347
1297 - if (
1298 - self.model_config is not None
1299 - and self.model_config.enable_return_routed_experts
1300 - ):
1301 - if self.parallel_config.pipeline_parallel_size > 1:
1302 - raise ValueError(
1303 - "--enable-return-routed-experts is incompatible with "
1304 - "pipeline parallelism (PP > 1)."
1305 - )
1306 - if (
1307 - self.parallel_config.decode_context_parallel_size > 1
1308 - or self.parallel_config.prefill_context_parallel_size > 1
1309 - ):
1310 - raise ValueError(
1311 - "--enable-return-routed-experts is incompatible with context "
1312 - "parallelism (DCP > 1 or PCP > 1)."
1313 - )
1314 -
1315 - # Incompatible with any KV connector — covers both PD disaggregation
1316 - # (kv_producer/kv_consumer: routing captured on P can't reach D) and
1317 - # single-instance KV offload/sharing (kv_both: slot_mapping semantics
1318 - # change when KV blocks live outside local GPU memory, breaking the
1319 - # slot-indexed routed_experts buffer).
1320 - if (
1321 - self.kv_transfer_config is not None
1322 - and self.kv_transfer_config.is_kv_transfer_instance
1323 - ):
1324 - raise ValueError(
1325 - "--enable-return-routed-experts is incompatible with KV "
1326 - "connectors (PD disaggregation, KV cache offload)."
1327 - )
1328 -
1348 if (
1349 self.model_config is not None
1350 and self.model_config.multimodal_config is not None
@@ -2032,6 +2051,7 @@ class VllmConfig:
2051 # Resolve kv_offloading-derived connector name into kv_transfer_config
2052 # before the HMA check below, which inspects the connector class.
2053 self._post_init_kv_transfer_config()
2054 + self._verify_aux_output_compatibility()
2055
2056 if self.is_mm_encoder_only and self.cache_config.enable_prefix_caching:
2057 # Such an instance publishes encoder embeddings and runs no language
@@ -2689,7 +2709,7 @@ class VllmConfig:
2709 f"quantization={self.model_config.quantization}, "
2710 f"quantization_config={self.model_config.quantization_config}, " # noqa
2711 f"enforce_eager={self.model_config.enforce_eager}, "
2692 - f"enable_return_routed_experts={self.model_config.enable_return_routed_experts}, " # noqa
2712 + f"aux_output_config={self.aux_output_config!r}, "
2713 f"kv_cache_dtype={self.cache_config.cache_dtype}, "
2714 f"device_config={self.device_config.device}, "
2715 f"structured_outputs_config={self.structured_outputs_config!r}, "
vllm/distributed/aux_output_connector/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/distributed/aux_output_connector/connector.py new
+157
@@ -0,0 +1,157 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Scheduler-side control plane for execution auxiliary outputs."""
4 +
5 +from __future__ import annotations
6 +
7 +from collections.abc import Iterator, Sequence
8 +from dataclasses import dataclass
9 +from typing import TYPE_CHECKING
10 +
11 +import numpy as np
12 +
13 +if TYPE_CHECKING:
14 + from vllm.v1.core.sched.output import SchedulerOutput
15 + from vllm.v1.request import Request
16 +
17 +
18 +@dataclass
19 +class PackedBlockHashes:
20 + """Contiguous, self-contained block hashes for scheduler-worker IPC."""
21 +
22 + data: bytes
23 + item_size: int
24 +
25 + def __iter__(self) -> Iterator[bytes]:
26 + for start in range(0, len(self.data), self.item_size):
27 + yield self.data[start : start + self.item_size]
28 +
29 +
30 +@dataclass
31 +class AuxOutputConnectorMetadata:
32 + generation: int
33 + requests: dict[str, int]
34 + block_hashes: dict[str, PackedBlockHashes]
35 + finished_requests: tuple[str, ...]
36 +
37 +
38 +@dataclass
39 +class AuxRequestOutput:
40 + token_start: int
41 + rows: np.ndarray
42 +
43 +
44 +class AuxOutputSchedulerConnector:
45 + """Build worker metadata without owning auxiliary output payloads or stores."""
46 +
47 + def __init__(self) -> None:
48 + # Number of hashes already sent to the worker for each active request.
49 + self._sent_hash_counts: dict[str, int] = {}
50 + # Terminal events are delivered with the next connector metadata.
51 + self._finished_requests: dict[str, PackedBlockHashes | None] = {}
52 + self._generation = 0
53 +
54 + def build_connector_meta(
55 + self,
56 + scheduler_output: SchedulerOutput,
57 + requests: dict[str, Request],
58 + ) -> AuxOutputConnectorMetadata:
59 + """Build one step's incremental worker metadata."""
60 + scheduled_requests: dict[str, int] = {}
61 + block_hashes_by_request: dict[str, PackedBlockHashes] = {}
62 + for request_id in scheduler_output.num_scheduled_tokens:
63 + num_sent = self._sent_hash_counts.setdefault(request_id, 0)
64 + request = requests[request_id]
65 + packed = self._pack_new_hashes(request.block_hashes, num_sent)
66 + if packed is not None:
67 + block_hashes_by_request[request_id] = packed
68 + self._sent_hash_counts[request_id] = len(request.block_hashes)
69 + assert request.sampling_params is not None
70 + scheduled_requests[request_id] = max(
71 + request.sampling_params.routed_experts_prompt_start,
72 + 0 if request.num_output_tokens == 0 else request.num_tokens - 1,
73 + )
74 + # A settled token can complete a hash block after the next async schedule
75 + # was built. Send a hash-only update if the request was not rescheduled.
76 + if len(self._sent_hash_counts) != len(scheduled_requests):
77 + for request_id, num_sent in self._sent_hash_counts.items():
78 + if request_id in scheduled_requests:
79 + continue
80 + request = requests[request_id]
81 + packed = self._pack_new_hashes(request.block_hashes, num_sent)
82 + if packed is None:
83 + continue
84 + block_hashes_by_request[request_id] = packed
85 + self._sent_hash_counts[request_id] = len(request.block_hashes)
86 + # Sending transfers ownership of these one-shot events.
87 + finished_requests = tuple(self._finished_requests)
88 + block_hashes_by_request.update(
89 + (request_id, block_hashes)
90 + for request_id, block_hashes in self._finished_requests.items()
91 + if block_hashes is not None
92 + )
93 + self._finished_requests = {}
94 + return AuxOutputConnectorMetadata(
95 + self._generation,
96 + scheduled_requests,
97 + block_hashes_by_request,
98 + finished_requests,
99 + )
100 +
101 + def take_output(
102 + self,
103 + request: Request,
104 + output: dict[str, AuxRequestOutput] | None,
105 + ) -> np.ndarray | None:
106 + """Return the accepted R3 rows for one scheduled request."""
107 + request_id = request.request_id
108 + assert output is not None and request_id in output, (
109 + f"auxiliary output worker output is missing {request_id}"
110 + )
111 + request_output = output[request_id]
112 + token_end = request.num_tokens - 1
113 + local_end = token_end - request_output.token_start
114 + if local_end < 0:
115 + assert not request.is_finished(), (
116 + "finished auxiliary output output has no accepted token range: "
117 + f"request={request_id}, token_end={token_end}, "
118 + f"output_start={request_output.token_start}, "
119 + "output_end="
120 + f"{request_output.token_start + len(request_output.rows)}"
121 + )
122 + return None
123 + assert local_end <= len(request_output.rows), (
124 + "auxiliary output worker output has an invalid token range: "
125 + f"request={request_id}, token_end={token_end}, "
126 + f"output_start={request_output.token_start}, "
127 + f"output_end={request_output.token_start + len(request_output.rows)}"
128 + )
129 + return request_output.rows[:local_end]
130 +
131 + def request_finished(self, request: Request) -> None:
132 + """Queue a request's terminal event and final block hashes."""
133 + request_id = request.request_id
134 + num_sent = self._sent_hash_counts.pop(request_id, None)
135 + if num_sent is None:
136 + return
137 + # The next metadata delivers this terminal event to the worker.
138 + self._finished_requests[request_id] = self._pack_new_hashes(
139 + request.block_hashes, num_sent
140 + )
141 +
142 + @staticmethod
143 + def _pack_new_hashes(
144 + block_hashes: Sequence[bytes], num_sent: int
145 + ) -> PackedBlockHashes | None:
146 + assert num_sent <= len(block_hashes), "KV block-hash history shrank"
147 + new_hashes = block_hashes[num_sent:]
148 + if not new_hashes:
149 + return None
150 + return PackedBlockHashes(b"".join(new_hashes), len(new_hashes[0]))
151 +
152 + def reset(self) -> None:
153 + """Start a new auxiliary output namespace after a prefix-cache reset."""
154 + # The worker drops temporary state on generation changes; resend hashes.
155 + self._sent_hash_counts.clear()
156 + self._finished_requests.clear()
157 + self._generation += 1
vllm/distributed/aux_output_connector/routed_experts.py new
+219
@@ -0,0 +1,219 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Routed-experts auxiliary output keys, publication, and materialization."""
4 +
5 +from __future__ import annotations
6 +
7 +from collections.abc import Hashable, Iterable, Sequence
8 +from dataclasses import dataclass
9 +from typing import Any
10 +
11 +import numpy as np
12 +
13 +from vllm.distributed.aux_output_connector.store import (
14 + BackgroundBlockObjectStore,
15 + BlockObject,
16 + BlockObjectStore,
17 +)
18 +
19 +
20 +@dataclass
21 +class _RequestTail:
22 + slot: int
23 + block_start: int
24 + length: int = 0
25 +
26 +
27 +class RoutedExpertsBuffer:
28 + """Own incomplete and unkeyed full routed-experts blocks."""
29 +
30 + def __init__(
31 + self,
32 + dtype: np.dtype[Any],
33 + shape_per_token: tuple[int, ...],
34 + block_size: int,
35 + max_num_seqs: int,
36 + max_num_batched_tokens: int,
37 + max_concurrent_batches: int,
38 + ) -> None:
39 + self.dtype = dtype
40 + self.shape_per_token = shape_per_token
41 + self.block_size = block_size
42 + max_step_blocks = (
43 + max_num_batched_tokens + block_size - 1
44 + ) // block_size + max_num_seqs
45 + # In-flight full blocks plus one incomplete tail per active request.
46 + max_blocks = max_concurrent_batches * max_step_blocks + max_num_seqs
47 + self._rows: np.ndarray = np.empty(
48 + (max_blocks, block_size, *shape_per_token), dtype=dtype
49 + )
50 + self._free_slots = list(range(len(self._rows) - 1, -1, -1))
51 + self._owned_slots: dict[int, int] = {}
52 + self._requests: dict[Hashable, _RequestTail] = {}
53 +
54 + def _tail(self, request_id: Hashable, block_start: int) -> _RequestTail:
55 + tail = self._requests.get(request_id)
56 + if tail is not None:
57 + assert tail.block_start == block_start, (
58 + "auxiliary output capture skipped an incomplete block: "
59 + f"request={request_id}, expected={tail.block_start}, "
60 + f"actual={block_start}"
61 + )
62 + return tail
63 + assert self._free_slots, "auxiliary output block pool is exhausted"
64 + tail = _RequestTail(self._free_slots.pop(), block_start)
65 + self._requests[request_id] = tail
66 + return tail
67 +
68 + def capture(
69 + self, request_id: Hashable, token_start: int, rows: np.ndarray
70 + ) -> list[tuple[int, np.ndarray]]:
71 + """Stage rows and return completed blocks without retaining them."""
72 + rows = np.asarray(rows)
73 + assert rows.shape[1:] == self.shape_per_token and rows.dtype == self.dtype, (
74 + "routed-experts capture profile changed"
75 + )
76 + if token_start < 0:
77 + raise ValueError("auxiliary output token start must be non-negative")
78 +
79 + completed: list[tuple[int, np.ndarray]] = []
80 + offset = 0
81 + while offset < len(rows):
82 + position = token_start + offset
83 + block_start = position // self.block_size * self.block_size
84 + local_start = position - block_start
85 +
86 + # Full aligned input blocks do not need tail staging.
87 + if (
88 + request_id not in self._requests
89 + and local_start == 0
90 + and len(rows) - offset >= self.block_size
91 + ):
92 + completed.append((block_start, rows[offset : offset + self.block_size]))
93 + offset += self.block_size
94 + continue
95 +
96 + tail = self._tail(request_id, block_start)
97 + assert local_start == tail.length, (
98 + "auxiliary output capture is not contiguous: "
99 + f"request={request_id}, expected={block_start + tail.length}, "
100 + f"actual={position}"
101 + )
102 + count = min(self.block_size - local_start, len(rows) - offset)
103 + if count == 1:
104 + self._rows[tail.slot, local_start] = rows[offset]
105 + else:
106 + self._rows[tail.slot, local_start : local_start + count] = rows[
107 + offset : offset + count
108 + ]
109 + tail.length += count
110 + offset += count
111 + if tail.length == self.block_size:
112 + block = self._rows[tail.slot]
113 + del self._requests[request_id]
114 + self._owned_slots[id(block)] = tail.slot
115 + completed.append((block_start, block))
116 +
117 + return completed
118 +
119 + def read(
120 + self, request_id: Hashable, token_start: int, token_end: int
121 + ) -> np.ndarray:
122 + tail = self._requests.get(request_id)
123 + assert tail is not None, (
124 + f"auxiliary output buffer is missing request {request_id}"
125 + )
126 + local_start = token_start - tail.block_start
127 + local_end = token_end - tail.block_start
128 + assert 0 <= local_start < local_end <= tail.length, (
129 + "auxiliary output range is unavailable: "
130 + f"request={request_id}, range=[{token_start}, {token_end}), "
131 + f"available=[{tail.block_start}, {tail.block_start + tail.length})"
132 + )
133 + return self._rows[tail.slot, local_start:local_end].copy()
134 +
135 + def retain_block(self, rows: np.ndarray) -> np.ndarray:
136 + """Retain one unkeyed block after the current capture call."""
137 + if id(rows) in self._owned_slots:
138 + return rows
139 + assert self._free_slots, "auxiliary output block pool is exhausted"
140 + slot = self._free_slots.pop()
141 + retained = self._rows[slot]
142 + retained[...] = rows
143 + self._owned_slots[id(retained)] = slot
144 + return retained
145 +
146 + def release_block(self, rows: np.ndarray) -> None:
147 + slot = self._owned_slots.pop(id(rows), None)
148 + if slot is not None:
149 + self._free_slots.append(slot)
150 +
151 + def discard(self, request_id: Hashable) -> None:
152 + tail = self._requests.pop(request_id, None)
153 + if tail is not None:
154 + self._free_slots.append(tail.slot)
155 +
156 + def reset(self) -> None:
157 + self._requests.clear()
158 + self._owned_slots.clear()
159 + self._free_slots[:] = range(len(self._rows) - 1, -1, -1)
160 +
161 +
162 +def routed_experts_keys(
163 + block_hashes: Iterable[bytes], aux_output_namespace: str
164 +) -> list[str]:
165 + prefix = f"vllm-artifact/{aux_output_namespace}/"
166 + return [prefix + block_hash.hex() for block_hash in block_hashes]
167 +
168 +
169 +def materialize_routed_experts(
170 + store: BackgroundBlockObjectStore | BlockObjectStore,
171 + aux_output_keys: list[str],
172 + *,
173 + shape_per_token: tuple[int, ...],
174 + dtype: np.dtype[Any],
175 +) -> np.ndarray:
176 + if not aux_output_keys:
177 + raise ValueError("routed-experts auxiliary output key list must not be empty")
178 + payload = store.get_concatenated(aux_output_keys)
179 + return np.frombuffer(payload, dtype=dtype).reshape((-1, *shape_per_token))
180 +
181 +
182 +def publish_routed_experts(
183 + store: BackgroundBlockObjectStore | BlockObjectStore,
184 + *,
185 + batches: Sequence[tuple[Sequence[str], list[tuple[int, np.ndarray]]]],
186 + block_size: int,
187 + retain_keys: Sequence[str] = (),
188 + release_keys: Sequence[str] = (),
189 +) -> None:
190 + """Publish immutable full R3 blocks."""
191 + objects = []
192 + for aux_output_keys, blocks in batches:
193 + for block_start, array in blocks:
194 + if block_start < 0 or block_start % block_size:
195 + raise ValueError(
196 + "auxiliary output block start is not hash-block aligned"
197 + )
198 + block_index = block_start // block_size
199 + if block_index >= len(aux_output_keys):
200 + raise ValueError(
201 + "auxiliary output block has no corresponding key: "
202 + f"start={block_start}, index={block_index}, "
203 + f"keys={len(aux_output_keys)}, block_size={block_size}"
204 + )
205 + if len(array) != block_size:
206 + raise ValueError(
207 + "auxiliary output block length does not match hash block size"
208 + )
209 + objects.append(
210 + BlockObject(
211 + key=aux_output_keys[block_index],
212 + payload=array.tobytes(order="C"),
213 + )
214 + )
215 + store.put(
216 + objects,
217 + retain_keys=retain_keys,
218 + release_keys=release_keys,
219 + )
vllm/distributed/aux_output_connector/store.py new
+234
@@ -0,0 +1,234 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""In-process auxiliary-output storage."""
4 +
5 +from __future__ import annotations
6 +
7 +import mmap
8 +import queue
9 +import threading
10 +from collections import OrderedDict
11 +from collections.abc import Iterable
12 +from dataclasses import dataclass
13 +
14 +
15 +@dataclass(frozen=True)
16 +class BlockObject:
17 + """One immutable auxiliary output object."""
18 +
19 + key: str
20 + payload: bytes
21 +
22 +
23 +class BlockObjectStoreError(RuntimeError):
24 + """AuxOutput storage or retrieval failed."""
25 +
26 +
27 +class BackgroundBlockObjectStore:
28 + """Serialize store mutations on a background thread."""
29 +
30 + def __init__(self, store: BlockObjectStore, *, max_pending_batches: int) -> None:
31 + self._store = store
32 + self._queue: queue.Queue[
33 + tuple[list[BlockObject], tuple[str, ...], tuple[str, ...]] | None
34 + ] = queue.Queue(maxsize=max_pending_batches)
35 + self._error: BaseException | None = None
36 + self._closed = False
37 + self._thread = threading.Thread(
38 + target=self._run,
39 + daemon=True,
40 + name="vllm-aux-output-writer",
41 + )
42 + self._thread.start()
43 +
44 + def _run(self) -> None:
45 + while True:
46 + update = self._queue.get()
47 + try:
48 + if update is None:
49 + return
50 + if self._error is None:
51 + objects, retain_keys, release_keys = update
52 + self._store.put(
53 + objects,
54 + retain_keys=retain_keys,
55 + release_keys=release_keys,
56 + )
57 + except BaseException as error:
58 + self._error = error
59 + finally:
60 + self._queue.task_done()
61 +
62 + def _raise_if_failed(self) -> None:
63 + if self._error is not None:
64 + raise BlockObjectStoreError(
65 + "auxiliary output publication failed"
66 + ) from self._error
67 +
68 + def put(
69 + self,
70 + objects: list[BlockObject],
71 + *,
72 + retain_keys: Iterable[str] = (),
73 + release_keys: Iterable[str] = (),
74 + ) -> None:
75 + retains = tuple(retain_keys)
76 + releases = tuple(release_keys)
77 + if not objects and not retains and not releases:
78 + return
79 + if self._closed:
80 + raise RuntimeError("auxiliary output store is closed")
81 + self._queue.put((objects, retains, releases))
82 + self._raise_if_failed()
83 +
84 + def get_concatenated(self, keys: list[str]) -> bytes:
85 + self._queue.join()
86 + self._raise_if_failed()
87 + return self._store.get_concatenated(keys)
88 +
89 + def close(self) -> None:
90 + if self._closed:
91 + return
92 + self._closed = True
93 + self._queue.join()
94 + self._queue.put(None)
95 + self._thread.join()
96 + self._store.close()
97 + self._raise_if_failed()
98 +
99 +
100 +class BlockObjectStore:
101 + """Single-owner bounded store that fails closed after an eviction."""
102 +
103 + _UNALLOCATED_SLOT = -1
104 +
105 + def __init__(
106 + self,
107 + *,
108 + max_bytes: int,
109 + object_nbytes: int,
110 + ) -> None:
111 + if object_nbytes <= 0:
112 + raise ValueError("auxiliary output object size must be positive")
113 + if max_bytes < object_nbytes:
114 + raise ValueError("auxiliary output store must fit at least one object")
115 + self.object_nbytes = object_nbytes
116 + self.num_slots = max_bytes // object_nbytes
117 + self._lru: OrderedDict[str, int] = OrderedDict()
118 + self._references: dict[str, int] = {}
119 + self._free_slots: list[int] = []
120 + self._next_slot = 0
121 + self._arena: mmap.mmap | None = mmap.mmap(
122 + -1, self.num_slots * self.object_nbytes
123 + )
124 +
125 + def _evict_to_fit(self, protected: set[str]) -> None:
126 + excess = len(self._lru) - self.num_slots
127 + if excess <= 0:
128 + return
129 + victims = []
130 + for key in self._lru:
131 + if key not in self._references and key not in protected:
132 + victims.append(key)
133 + if len(victims) == excess:
134 + break
135 + if len(victims) != excess:
136 + raise BlockObjectStoreError(
137 + "auxiliary output store cannot retain the requested batch: "
138 + f"limit={self.num_slots} objects"
139 + )
140 + for victim in victims:
141 + slot = self._lru.pop(victim)
142 + if slot != self._UNALLOCATED_SLOT:
143 + self._free_slots.append(slot)
144 +
145 + def _allocate_slot(self) -> int:
146 + if self._free_slots:
147 + return self._free_slots.pop()
148 + slot = self._next_slot
149 + self._next_slot += 1
150 + return slot
151 +
152 + def put(
153 + self,
154 + objects: list[BlockObject],
155 + *,
156 + retain_keys: Iterable[str] = (),
157 + release_keys: Iterable[str] = (),
158 + ) -> None:
159 + arena = self._arena
160 + if arena is None:
161 + raise RuntimeError("auxiliary output store is closed")
162 + unique = {obj.key: obj for obj in objects}
163 + self._retain(retain_keys)
164 + terminal_order = self._release(release_keys)
165 + for key in unique:
166 + self._lru.setdefault(key, self._UNALLOCATED_SLOT)
167 + self._lru.move_to_end(key)
168 + for key in terminal_order:
169 + if key in self._lru:
170 + self._lru.move_to_end(key)
171 + try:
172 + self._evict_to_fit(set(unique) - set(terminal_order))
173 + except BlockObjectStoreError:
174 + for key in unique:
175 + if self._lru.get(key) == self._UNALLOCATED_SLOT:
176 + del self._lru[key]
177 + raise
178 + for object_id, obj in unique.items():
179 + slot = self._lru.get(object_id)
180 + if slot is None:
181 + continue
182 + if slot != self._UNALLOCATED_SLOT:
183 + continue
184 + slot = self._allocate_slot()
185 + offset = slot * self.object_nbytes
186 + arena[offset : offset + self.object_nbytes] = obj.payload
187 + self._lru[object_id] = slot
188 +
189 + def _retain(self, keys: Iterable[str]) -> None:
190 + for key in keys:
191 + references = self._references.get(key, 0)
192 + self._references[key] = references + 1
193 +
194 + def _release(self, keys: Iterable[str]) -> list[str]:
195 + terminal_order = []
196 + for key in keys:
197 + references = self._references[key] - 1
198 + if references:
199 + self._references[key] = references
200 + continue
201 + del self._references[key]
202 + terminal_order.append(key)
203 + return terminal_order
204 +
205 + def get_concatenated(self, keys: list[str]) -> bytes:
206 + arena_obj = self._arena
207 + if arena_obj is None:
208 + raise RuntimeError("auxiliary output store is closed")
209 + try:
210 + entries = [self._lru[key] for key in keys]
211 + except KeyError as error:
212 + raise BlockObjectStoreError(
213 + "auxiliary output object does not exist; the object may have been "
214 + f"evicted (used={len(self._lru)}, "
215 + f"limit={self.num_slots} objects). Increase "
216 + "aux_output_config.max_bytes when a KV cache hit requires it."
217 + ) from error
218 + arena = memoryview(arena_obj)
219 + try:
220 + payload = b"".join(
221 + arena[slot * self.object_nbytes : (slot + 1) * self.object_nbytes]
222 + for slot in entries
223 + )
224 + finally:
225 + arena.release()
226 + for key in keys:
227 + self._lru.move_to_end(key)
228 + return payload
229 +
230 + def close(self) -> None:
231 + arena = self._arena
232 + if arena is not None:
233 + self._arena = None
234 + arena.close()
vllm/distributed/aux_output_connector/worker.py new
+412
@@ -0,0 +1,412 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Worker-side auxiliary-output data plane."""
4 +
5 +from __future__ import annotations
6 +
7 +from collections.abc import Iterable, Sequence
8 +from dataclasses import dataclass, field
9 +from threading import Lock
10 +from typing import TYPE_CHECKING, Any
11 +
12 +import numpy as np
13 +import torch
14 +
15 +from vllm.config import VllmConfig
16 +from vllm.distributed.aux_output_connector.connector import (
17 + AuxOutputConnectorMetadata,
18 + AuxRequestOutput,
19 +)
20 +from vllm.distributed.aux_output_connector.routed_experts import (
21 + RoutedExpertsBuffer,
22 + materialize_routed_experts,
23 + publish_routed_experts,
24 + routed_experts_keys,
25 +)
26 +from vllm.distributed.aux_output_connector.store import (
27 + BackgroundBlockObjectStore,
28 + BlockObjectStore,
29 +)
30 +from vllm.distributed.parallel_state import get_tp_group
31 +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
32 + RoutedExpertsCapturer,
33 + bind_routed_experts_capturer,
34 +)
35 +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
36 +
37 +if TYPE_CHECKING:
38 + from vllm.v1.kv_cache_interface import KVCacheConfig
39 + from vllm.v1.worker.gpu.input_batch import InputBatch
40 +
41 +
42 +@dataclass
43 +class _WorkerRequestState:
44 + aux_output_keys: list[str] = field(default_factory=list)
45 + pending_blocks: list[tuple[int, np.ndarray]] = field(default_factory=list)
46 + capture_cursor: int | None = None
47 + scheduled_cursor: int = 0
48 + emit_cursor: int = 0
49 + pending_outputs: int = 0
50 + # Terminal event received; teardown waits for in-flight step outputs.
51 + finished: bool = False
52 +
53 +
54 +@dataclass
55 +class PendingAuxOutput:
56 + """Own one step's R3 snapshot until its asynchronous copy is consumed."""
57 +
58 + connector: AuxOutputWorkerConnector
59 + request_ids: list[str]
60 + token_starts: np.ndarray
61 + query_start_loc: np.ndarray
62 + # GPU snapshot taken on the main stream.
63 + routed_experts_gpu: torch.Tensor
64 + # Filled by enqueue_cpu_copy on the output copy stream.
65 + routed_experts: np.ndarray | None = None
66 + num_sampled: np.ndarray | None = None
67 + num_rejected: np.ndarray | None = None
68 +
69 + def enqueue_cpu_copy(
70 + self, num_sampled: np.ndarray, num_rejected: np.ndarray
71 + ) -> None:
72 + """Enqueue asynchronous D2H copies; call on the output copy stream."""
73 + self.num_sampled = num_sampled
74 + self.num_rejected = num_rejected
75 + self.routed_experts = self.routed_experts_gpu.to(
76 + "cpu", non_blocking=True
77 + ).numpy()
78 +
79 + def process_output(self) -> dict[str, AuxRequestOutput]:
80 + return self.connector.process_output(self)
81 +
82 +
83 +class AuxOutputWorkerConnector:
84 + """Own capture, request tails, and backend resources on the output worker."""
85 +
86 + def __init__(
87 + self,
88 + *,
89 + vllm_config: VllmConfig,
90 + model: torch.nn.Module,
91 + kv_cache_config: KVCacheConfig,
92 + ) -> None:
93 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
94 + capturer = RoutedExpertsCapturer(
95 + max_num_batched_tokens=max_num_batched_tokens, vllm_config=vllm_config
96 + )
97 + bind_routed_experts_capturer(model, capturer)
98 + self._capturer = capturer
99 + self._store: BackgroundBlockObjectStore | None = None
100 + self._buffer: RoutedExpertsBuffer | None = None
101 + self._requests: dict[str, _WorkerRequestState] = {}
102 + self._generation = 0
103 + self._step_metadata: AuxOutputConnectorMetadata | None = None
104 + # Steps whose asynchronous copy has not been consumed yet. The engine
105 + # consumes step outputs in order, so at most max_concurrent_batches
106 + # are outstanding; guarded by _lock against the output thread.
107 + self._pending_outputs: list[PendingAuxOutput] = []
108 + self._lock = Lock()
109 + self._max_concurrent_batches = vllm_config.max_concurrent_batches
110 + # Every TP rank participates in capture collectives, but only the
111 + # executor output rank owns the auxiliary output data plane.
112 + if not get_tp_group().is_first_rank:
113 + return
114 +
115 + shape_per_token = self._capturer.shape_per_token
116 + dtype: np.dtype[Any] = np.dtype(self._capturer.output_dtype_name)
117 + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
118 + kv_cache_config, vllm_config
119 + )
120 + hashes_per_kv_block = scheduler_block_size // hash_block_size
121 + block_nbytes = hash_block_size * int(np.prod(shape_per_token)) * dtype.itemsize
122 + max_bytes = vllm_config.aux_output_config.max_bytes
123 + if max_bytes is None:
124 + max_bytes = kv_cache_config.num_blocks * hashes_per_kv_block * block_nbytes
125 + self._store = BackgroundBlockObjectStore(
126 + BlockObjectStore(max_bytes=max_bytes, object_nbytes=block_nbytes),
127 + max_pending_batches=2 * vllm_config.scheduler_config.max_num_seqs,
128 + )
129 + self._buffer = RoutedExpertsBuffer(
130 + dtype,
131 + shape_per_token,
132 + hash_block_size,
133 + vllm_config.scheduler_config.max_num_seqs,
134 + max_num_batched_tokens,
135 + vllm_config.max_concurrent_batches,
136 + )
137 +
138 + def prepare_output(self, input_batch: InputBatch) -> PendingAuxOutput | None:
139 + """Snapshot one step's R3 tensor for asynchronous CPU transfer."""
140 + if self._buffer is None or self._step_metadata is None:
141 + return None
142 +
143 + request_ids = list(input_batch.req_ids)
144 + query_start_loc = input_batch.query_start_loc_np[: len(request_ids) + 1]
145 + num_rows = int(query_start_loc[-1])
146 + pending_output = PendingAuxOutput(
147 + connector=self,
148 + request_ids=request_ids,
149 + token_starts=input_batch.num_computed_tokens_np,
150 + query_start_loc=query_start_loc,
151 + routed_experts_gpu=self._capturer.snapshot_routing_data(num_rows),
152 + )
153 + with self._lock:
154 + assert len(self._pending_outputs) < self._max_concurrent_batches, (
155 + "auxiliary output step outputs are not consumed in order"
156 + )
157 + self._pending_outputs.append(pending_output)
158 + for request_id in pending_output.request_ids:
159 + self._requests[request_id].pending_outputs += 1
160 + return pending_output
161 +
162 + def process_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]:
163 + """Commit one consumed R3 snapshot and build request outputs.
164 +
165 + Request teardown deferred by begin_step is completed here once no
166 + in-flight step covers the request.
167 + """
168 + with self._lock:
169 + teardown = []
170 + release_keys: list[str] = []
171 + try:
172 + outputs = self._commit_output(pending)
173 + finally:
174 + self._pending_outputs.remove(pending)
175 + for request_id in pending.request_ids:
176 + state = self._requests[request_id]
177 + state.pending_outputs -= 1
178 + if state.pending_outputs == 0 and state.finished:
179 + teardown.append(request_id)
180 + release_keys.extend(reversed(state.aux_output_keys))
181 + if teardown:
182 + self._publish_blocks([], release_keys=release_keys)
183 + self._teardown(teardown)
184 + return outputs
185 +
186 + def _commit_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]:
187 + buffer = self._buffer
188 + store = self._store
189 + assert buffer is not None and store is not None
190 + block_size = buffer.block_size
191 +
192 + routed_experts = pending.routed_experts
193 + num_sampled = pending.num_sampled
194 + num_rejected = pending.num_rejected
195 + assert (
196 + routed_experts is not None
197 + and num_sampled is not None
198 + and num_rejected is not None
199 + ), "auxiliary output CPU copy was not enqueued"
200 +
201 + # Publish the whole batch before materializing any consumer output.
202 + materialize_outputs: list[tuple[str, int, int]] = []
203 + block_batches = []
204 + outputs: dict[str, AuxRequestOutput] = {}
205 +
206 + # Use the ModelRunner's actual batch boundaries rather than rebuilding them.
207 + for request_id, token_start, start, end, sampled, rejected in zip(
208 + pending.request_ids,
209 + pending.token_starts,
210 + pending.query_start_loc[:-1],
211 + pending.query_start_loc[1:],
212 + num_sampled,
213 + num_rejected,
214 + strict=True,
215 + ):
216 + request_num_tokens = end - start
217 + assert request_num_tokens > 0, (
218 + "auxiliary output request token count must be positive"
219 + )
220 + state = self._requests[request_id]
221 +
222 + # Capture precedes speculative acceptance, so discard the rejected
223 + # suffix. Batch boundaries still span the full executed range.
224 + rejected = int(rejected)
225 + assert 0 <= rejected <= request_num_tokens, (
226 + "auxiliary output rejected-token count is invalid"
227 + )
228 + rows = routed_experts[start : end - rejected]
229 +
230 + capture_start = token_start
231 + capture_cursor = state.capture_cursor
232 + if capture_cursor is None:
233 + capture_cursor = capture_start
234 +
235 + assert capture_start >= capture_cursor, (
236 + "auxiliary output capture moved backwards"
237 + )
238 + if capture_start > capture_cursor:
239 + # Reattach after an optimistically scheduled suffix was rejected.
240 + assert capture_cursor < state.scheduled_cursor, (
241 + "auxiliary output capture has an unbacked token gap"
242 + )
243 + capture_start = capture_cursor
244 +
245 + emit_start = state.emit_cursor
246 + # Complete blocks without keys remain pending until a hash update.
247 + completed = buffer.capture(request_id, capture_start, rows)
248 + token_end = capture_start + len(rows)
249 + state.capture_cursor = token_end
250 + state.scheduled_cursor = token_start + request_num_tokens
251 + block_batches.append((state, completed))
252 +
253 + if sampled > 0 and emit_start <= token_end:
254 + if emit_start >= capture_start:
255 + outputs[request_id] = AuxRequestOutput(
256 + emit_start, rows[emit_start - capture_start :]
257 + )
258 + state.emit_cursor = token_end
259 + else:
260 + materialize_outputs.append((request_id, emit_start, token_end))
261 +
262 + # A consumer may reuse a block produced earlier in the same batch.
263 + self._publish_blocks(block_batches)
264 +
265 + for request_id, emit_start, token_end in materialize_outputs:
266 + state = self._requests[request_id]
267 + stored_end = (
268 + min(token_end // block_size, len(state.aux_output_keys)) * block_size
269 + )
270 + if emit_start < stored_end:
271 + first_block = emit_start // block_size
272 + stored = materialize_routed_experts(
273 + store,
274 + state.aux_output_keys[first_block : stored_end // block_size],
275 + shape_per_token=buffer.shape_per_token,
276 + dtype=buffer.dtype,
277 + )
278 + local_start = emit_start % block_size
279 + rows = stored[local_start : local_start + stored_end - emit_start]
280 + if stored_end < token_end:
281 + rows = np.concatenate(
282 + (rows, buffer.read(request_id, stored_end, token_end))
283 + )
284 + else:
285 + rows = buffer.read(request_id, emit_start, token_end)
286 + outputs[request_id] = AuxRequestOutput(emit_start, rows)
287 + state.emit_cursor = token_end
288 + return outputs
289 +
290 + def _publish_blocks(
291 + self,
292 + batches: list[tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]]],
293 + retain_keys: Sequence[str] = (),
294 + release_keys: Sequence[str] = (),
295 + ) -> None:
296 + store = self._store
297 + buffer = self._buffer
298 + assert store is not None and buffer is not None
299 + ready_batches = []
300 + for state, completed in batches:
301 + blocks = state.pending_blocks + completed
302 + keyed_end = len(state.aux_output_keys) * buffer.block_size
303 + ready = [(start, rows) for start, rows in blocks if start < keyed_end]
304 + state.pending_blocks = [
305 + (start, buffer.retain_block(rows))
306 + for start, rows in blocks
307 + if start >= keyed_end
308 + ]
309 + if ready:
310 + ready_batches.append((state.aux_output_keys, ready))
311 + if ready_batches or retain_keys or release_keys:
312 + publish_routed_experts(
313 + store,
314 + batches=ready_batches,
315 + block_size=buffer.block_size,
316 + retain_keys=retain_keys,
317 + release_keys=release_keys,
318 + )
319 + for _, blocks in ready_batches:
320 + for _, rows in blocks:
321 + buffer.release_block(rows)
322 +
323 + def _teardown(self, request_ids: Iterable[str]) -> None:
324 + """Drop finished requests' temporary state after publishing key releases."""
325 + buffer = self._buffer
326 + assert buffer is not None
327 + for request_id in request_ids:
328 + state = self._requests.pop(request_id)
329 + for _, rows in state.pending_blocks:
330 + buffer.release_block(rows)
331 + buffer.discard(request_id)
332 +
333 + def begin_step(self, metadata: AuxOutputConnectorMetadata | None) -> None:
334 + """Apply one scheduler step's request and block-hash updates.
335 +
336 + While an unconsumed step covers a finished request, its teardown is
337 + deferred to that step's process_output.
338 + """
339 + self._step_metadata = metadata
340 + if self._buffer is None or metadata is None:
341 + return
342 + assert not metadata.requests.keys() & metadata.finished_requests, (
343 + "auxiliary output request cannot run and finish in one step"
344 + )
345 + assert metadata.generation >= self._generation, (
346 + "auxiliary output metadata generation moved backwards"
347 + )
348 + with self._lock:
349 + release_keys: list[str] = []
350 + if metadata.generation > self._generation:
351 + # A generation change follows a prefix-cache reset, which the
352 + # scheduler only performs once all model output is consumed.
353 + assert not self._pending_outputs, (
354 + "auxiliary output generation changed with output in flight"
355 + )
356 + release_keys.extend(
357 + key
358 + for state in self._requests.values()
359 + for key in reversed(state.aux_output_keys)
360 + )
361 + self._buffer.reset()
362 + self._requests.clear()
363 + self._generation = metadata.generation
364 + for request_id, emit_start in metadata.requests.items():
365 + state = self._requests.setdefault(
366 + request_id, _WorkerRequestState(emit_cursor=emit_start)
367 + )
368 + if state.pending_outputs == 0:
369 + # In-flight steps leave the worker cursor behind the
370 + # scheduler's optimistic view; only check settled requests.
371 + assert emit_start <= state.emit_cursor, (
372 + "auxiliary output Scheduler emit cursor moved ahead"
373 + )
374 + block_batches: list[
375 + tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]]
376 + ] = []
377 + retained_keys: list[str] = []
378 + for request_id, block_hashes in metadata.block_hashes.items():
379 + state = self._requests[request_id]
380 + keys = routed_experts_keys(block_hashes, str(self._generation))
381 + state.aux_output_keys.extend(keys)
382 + retained_keys.extend(keys)
383 + block_batches.append((state, []))
384 + finished_now: list[str] = []
385 + for request_id in metadata.finished_requests:
386 + state = self._requests[request_id]
387 + state.finished = True
388 + if state.pending_outputs:
389 + continue
390 + finished_now.append(request_id)
391 + release_keys.extend(reversed(state.aux_output_keys))
392 + self._publish_blocks(block_batches, retained_keys, release_keys)
393 + self._teardown(finished_now)
394 +
395 + def close(self) -> None:
396 + with self._lock:
397 + self._pending_outputs.clear()
398 + if self._store is not None:
399 + try:
400 + self._store.close()
401 + finally:
402 + self._store = None
403 +
404 +
405 +def get_aux_output_connector(
406 + model: torch.nn.Module, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
407 +) -> AuxOutputWorkerConnector:
408 + return AuxOutputWorkerConnector(
409 + model=model,
410 + kv_cache_config=kv_cache_config,
411 + vllm_config=vllm_config,
412 + )
vllm/engine/arg_utils.py
+15 -3
@@ -35,6 +35,7 @@ from typing_extensions import TypeIs
35 import vllm.envs as envs
36 from vllm.config import (
37 AttentionConfig,
38 + AuxOutputConfig,
39 CacheConfig,
40 CompilationConfig,
41 ConfigType,
@@ -447,7 +448,9 @@ class EngineArgs:
448 """Arguments for vLLM engine."""
449
450 model: str = ModelConfig.model
450 - enable_return_routed_experts: bool = ModelConfig.enable_return_routed_experts
451 + # Public compatibility argument. Canonical runtime state lives in
452 + # AuxOutputConfig.
453 + enable_return_routed_experts: bool = False
454 return_sampling_mask: bool = ModelConfig.return_sampling_mask
455 model_weights: str = ModelConfig.model_weights
456 served_model_name: str | list[str] | None = ModelConfig.served_model_name
@@ -674,6 +677,7 @@ class EngineArgs:
677 structured_outputs_config: StructuredOutputsConfig = get_field(
678 VllmConfig, "structured_outputs_config"
679 )
680 + aux_output_config: AuxOutputConfig = get_field(VllmConfig, "aux_output_config")
681 reasoning_parser: str = StructuredOutputsConfig.reasoning_parser
682 reasoning_parser_plugin: str | None = None
683
@@ -816,6 +820,10 @@ class EngineArgs:
820 self.compilation_config = CompilationConfig(**self.compilation_config)
821 if isinstance(self.attention_config, dict):
822 self.attention_config = AttentionConfig(**self.attention_config)
823 + if isinstance(self.aux_output_config, dict):
824 + self.aux_output_config = AuxOutputConfig(**self.aux_output_config)
825 + if self.enable_return_routed_experts:
826 + self.aux_output_config.enable_return_routed_experts = True
827 if isinstance(self.engram_config, dict):
828 self.engram_config = EngramConfig(**self.engram_config)
829 if isinstance(self.mamba_config, dict):
@@ -884,6 +892,7 @@ class EngineArgs:
892 """Shared CLI arguments for vLLM engine."""
893 # Model arguments
894 model_kwargs = get_kwargs(ModelConfig)
895 + aux_output_kwargs = get_kwargs(AuxOutputConfig)
896 model_group = parser.add_argument_group(
897 title="ModelConfig",
898 description=ModelConfig.__doc__,
@@ -923,7 +932,7 @@ class EngineArgs:
932 model_group.add_argument("--enforce-eager", **model_kwargs["enforce_eager"])
933 model_group.add_argument(
934 "--enable-return-routed-experts",
926 - **model_kwargs["enable_return_routed_experts"],
935 + **aux_output_kwargs["enable_return_routed_experts"],
936 )
937 model_group.add_argument(
938 "--return-sampling-mask",
@@ -1749,6 +1758,9 @@ class EngineArgs:
1758 vllm_group.add_argument(
1759 "--structured-outputs-config", **vllm_kwargs["structured_outputs_config"]
1760 )
1761 + vllm_group.add_argument(
1762 + "--aux-output-config", **vllm_kwargs["aux_output_config"]
1763 + )
1764 vllm_group.add_argument("--profiler-config", **vllm_kwargs["profiler_config"])
1765 vllm_group.add_argument(
1766 "--optimization-level", **vllm_kwargs["optimization_level"]
@@ -1856,7 +1868,6 @@ class EngineArgs:
1868 quantization_config=self.quantization_config,
1869 allow_deprecated_quantization=self.allow_deprecated_quantization,
1870 enforce_eager=self.enforce_eager,
1859 - enable_return_routed_experts=self.enable_return_routed_experts,
1871 return_sampling_mask=self.return_sampling_mask,
1872 max_logprobs=self.max_logprobs,
1873 logprobs_mode=self.logprobs_mode,
@@ -2669,6 +2680,7 @@ class EngineArgs:
2680 load_config=load_config,
2681 offload_config=offload_config,
2682 attention_config=attention_config,
2683 + aux_output_config=self.aux_output_config,
2684 engram_config=self.engram_config,
2685 mamba_config=mamba_config,
2686 kernel_config=kernel_config,
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+37 -237
@@ -10,15 +10,12 @@ from collections.abc import Callable
10 from functools import partial
11 from typing import Protocol, runtime_checkable
12
13 -import numpy as np
13 import torch
14
15 from vllm.config import VllmConfig
16 from vllm.distributed.parallel_state import get_tp_group
17 from vllm.forward_context import get_forward_context
18 from vllm.platforms import current_platform
20 -from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec
21 -from vllm.v1.outputs import RoutedExpertsTensors
19
20 logger = logging.getLogger(__name__)
21
@@ -46,37 +43,29 @@ def _get_routed_experts_shape(vllm_config: VllmConfig) -> tuple[int, int, int]:
43 class RoutedExpertsCapturer:
44 """Worker-side capturer for routed experts, lives on GPU.
45
49 - Layer-level hooks call :meth:`capture` from inside the forward pass
50 - with the per-layer ``topk_ids`` tensor. The tensor is sliced to the
51 - tokens owned by this DP rank and written into a preallocated device
52 - buffer. At the end of the step, :class:`GPUModelRunner` reads the
53 - device buffer, issues a D2H copy into a pinned CPU buffer, and hands
54 - the result to the scheduler via :class:`RoutedExpertsLists`.
46 + Layer-level hooks call :meth:`capture` inside the forward pass. Routing
47 + rows owned by this DP rank are written into a preallocated device buffer.
48
56 - The device / pinned-CPU transit buffers use ``torch.int32`` (not a
57 - narrow ``uint8``/``uint16`` sized by ``num_experts``). This keeps the
58 - SP all-gather path free of dtype casts, matches the router's native
59 - ``topk_ids`` indices dtype more closely, and costs only a few MB per
60 - worker (``max_num_batched_tokens * num_layers * top_k * 4`` bytes).
61 - The scheduler-side slot buffer
62 - (``RoutedExpertsManager.routed_experts_by_slot``) still uses the
63 - narrow dtype -- numpy fancy-index assignment in ``store_batch``
64 - narrows the data on the way in.
49 + The device buffer uses ``int32``. Stable snapshots use the narrowest dtype
50 + that can represent every logical expert ID.
51
52 Invariants:
53 - One instance per worker; shape is fixed at init and covers the
54 worst-case step (``max_num_batched_tokens`` tokens).
55 - Every routed layer overwrites the current step's token rows.
70 - - ``device_buffer.dtype`` is ``torch.int32``.
56 """
57
58 def __init__(
59 self,
60 max_num_batched_tokens: int,
61 vllm_config: VllmConfig,
77 - kv_cache_config: KVCacheConfig,
62 ) -> None:
79 - num_layers, _, num_experts_per_tok = _get_routed_experts_shape(vllm_config)
63 + num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape(
64 + vllm_config
65 + )
66 + self.shape_per_token = (num_layers, num_experts_per_tok)
67 + self.output_dtype_name = "uint8" if num_experts <= 256 else "uint16"
68 + self.output_dtype = getattr(torch, self.output_dtype_name)
69 logger.info(
70 "RoutedExpertsCapturer: allocating buffer with "
71 "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d "
@@ -87,23 +76,12 @@ class RoutedExpertsCapturer:
76 vllm_config.model_config.hf_text_config.model_type,
77 )
78 self.device_buffer = torch.zeros(
90 - (
91 - max_num_batched_tokens,
92 - num_layers,
93 - num_experts_per_tok,
94 - ),
95 - # Use int32 for the device / host transit buffers: it
96 - # matches the router's native topk_ids dtype, is universally
97 - # supported by NCCL (uint8/uint16 are version-dependent),
98 - # and the extra bytes are small (few MB per worker). The
99 - # big scheduler-side slot buffer stays narrow.
79 + (max_num_batched_tokens, *self.shape_per_token),
80 dtype=torch.int32,
81 device=current_platform.device_type,
82 )
83 self.dp_rank = vllm_config.parallel_config.data_parallel_rank
84 self.tp_size = vllm_config.parallel_config.tensor_parallel_size
105 - # KV cache group whose slot layout the routing data is keyed by.
106 - self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
85
86 def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
87 """Capture expert routing decisions for a specific layer.
@@ -139,31 +117,21 @@ class RoutedExpertsCapturer:
117
118 """
119 ctx = get_forward_context()
142 - if ctx.dp_metadata is None: # single dp
143 - start_loc = 0
144 - end_loc = topk_ids.shape[0]
145 - token_num_per_dp = topk_ids.shape[0]
146 - else: # multi dp
120 + if ctx.dp_metadata is None:
121 + local_topk_ids = topk_ids
122 + else:
123 num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu
148 - token_num_per_dp = int(num_tokens_dp[self.dp_rank].item())
124 + num_local_tokens = int(num_tokens_dp[self.dp_rank].item())
125 total = int(num_tokens_dp.sum().item())
126 n = topk_ids.shape[0]
127 shard_sizes = getattr(ctx.dp_metadata, "local_sizes", None)
128 gathered_size = sum(shard_sizes) if shard_sizes is not None else -1
129
130 if n == total:
155 - # Naive dispatch: all DP ranks' tokens concatenated
156 - # before routing. This rank owns tokens
157 - # [end_loc - token_num_per_dp, end_loc).
158 - cumsum = torch.cumsum(num_tokens_dp, dim=0)
159 - end_loc = int(cumsum[self.dp_rank].item())
160 - start_loc = end_loc - token_num_per_dp
161 - elif n == token_num_per_dp:
162 - # Modular-kernel path: DP combine happens inside
163 - # quant_method.apply; select_experts only sees this
164 - # rank's tokens, take the whole tensor.
165 - start_loc = 0
166 - end_loc = token_num_per_dp
131 + end = int(num_tokens_dp[: self.dp_rank + 1].sum().item())
132 + local_topk_ids = topk_ids[end - num_local_tokens : end]
133 + elif n == num_local_tokens:
134 + local_topk_ids = topk_ids
135 elif shard_sizes is not None and n == gathered_size:
136 # Naive DP+EP dispatch gathers the sequence-parallel shards in
137 # flattened DP-then-TP order. ``local_sizes`` is the exact
@@ -173,12 +141,11 @@ class RoutedExpertsCapturer:
141 assert len(shard_sizes) % num_dp_ranks == 0
142 shards_per_dp_rank = len(shard_sizes) // num_dp_ranks
143 first_shard = self.dp_rank * shards_per_dp_rank
176 - start_loc = sum(shard_sizes[:first_shard])
177 - end_loc = start_loc + token_num_per_dp
144 + start = sum(shard_sizes[:first_shard])
145 + local_topk_ids = topk_ids[start : start + num_local_tokens]
146 elif (
147 self.tp_size > 1
180 - and n != token_num_per_dp
181 - and n == (token_num_per_dp + self.tp_size - 1) // self.tp_size
148 + and n == (num_local_tokens + self.tp_size - 1) // self.tp_size
149 ):
150 # SP + modular-kernel path. All-gather across the TP
151 # group along dim=0 to reconstruct the full per-DP-rank
@@ -193,60 +160,30 @@ class RoutedExpertsCapturer:
160 #
161 # ``topk_ids`` is already whatever the router produced
162 # (typically int32/int64, both supported by NCCL); the
196 - # downstream ``device_buffer[...] = topk_ids[...]``
197 - # setitem narrows into int32 automatically.
198 - topk_ids = get_tp_group().all_gather(topk_ids, dim=0)
199 - start_loc = 0
200 - end_loc = token_num_per_dp
163 + # downstream buffer assignment narrows to the capture dtype.
164 + local_topk_ids = get_tp_group().all_gather(topk_ids, dim=0)[
165 + :num_local_tokens
166 + ]
167 else:
202 - sp_expected = (
203 - (token_num_per_dp + self.tp_size - 1) // self.tp_size
204 - if self.tp_size > 0
205 - else -1
206 - )
168 + sp_expected = (num_local_tokens + self.tp_size - 1) // self.tp_size
169 raise AssertionError(
170 "RoutedExpertsCapturer: unexpected topk_ids batch "
209 - f"dim {n} (expected {total}, {token_num_per_dp}, "
171 + f"dim {n} (expected {total}, {num_local_tokens}, "
172 f"{gathered_size}, or {sp_expected} for "
173 f"dp_rank={self.dp_rank}, tp_size={self.tp_size})"
174 )
175
214 - if layer_id >= self.device_buffer.shape[1]:
176 + if not 0 <= layer_id < self.device_buffer.shape[1]:
177 raise IndexError(
178 f"routed-experts layer {layer_id} exceeds capture buffer "
179 f"layer count {self.device_buffer.shape[1]}"
180 )
181
220 - self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[
221 - start_loc:end_loc, :
222 - ]
223 -
224 - def get_device_buffer(self) -> torch.Tensor:
225 - """Return the underlying device buffer so the model runner can
226 - issue the D2H copy. The tensor is shared; callers must either
227 - clone or fully drain it before the next forward pass overwrites it.
228 - """
229 - return self.device_buffer
230 -
231 - def get_routed_experts(
232 - self, slot_mappings: torch.Tensor, num_tokens: int
233 - ) -> RoutedExpertsTensors:
234 - """Snapshot this step's routing data and its attention slot mapping.
235 -
236 - Both tensors are cloned since the capture buffer and the shared
237 - ``slot_mappings`` are overwritten by the next step, which may race
238 - with an in-flight D2H copy.
182 + self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids
183
240 - Args:
241 - slot_mappings: Per-KV-cache-group slot mappings for this step,
242 - shape ``(num_kv_cache_groups, max_num_batched_tokens)``.
243 - num_tokens: Total number of tokens scheduled in this step.
244 -
245 - """
246 - return RoutedExpertsTensors(
247 - routing_data=self.device_buffer[:num_tokens].clone(),
248 - slot_mapping=slot_mappings[self.attn_gid, :num_tokens].clone(),
249 - )
184 + def snapshot_routing_data(self, num_tokens: int) -> torch.Tensor:
185 + """Return a stable snapshot of the current routing data."""
186 + return self.device_buffer[:num_tokens].to(self.output_dtype)
187
188
189 def bind_routed_experts_capturer(
@@ -268,20 +205,12 @@ def bind_routed_experts_capturer(
205 continue
206 if not isinstance(module, MoERunner):
207 continue
271 - layer_id = module.layer_id
272 -
273 - def capture_fn(
274 - topk_ids: torch.Tensor,
275 - layer_id: int = layer_id,
276 - capturer: RoutedExpertsCapturer = capturer,
277 - ) -> None:
278 - capturer.capture(layer_id, topk_ids)
279 -
208 + capture_fn = partial(capturer.capture, module.layer_id)
209 quant_method = module._quant_method
281 - moe_kernel = getattr(quant_method, "moe_kernel", None)
282 - impl = getattr(moe_kernel, "impl", None)
283 - fused_experts = getattr(impl, "fused_experts", None)
210 if quant_method.is_monolithic:
211 + moe_kernel = getattr(quant_method, "moe_kernel", None)
212 + impl = getattr(moe_kernel, "impl", None)
213 + fused_experts = getattr(impl, "fused_experts", None)
214 if not (
215 isinstance(fused_experts, FusedMoEExpertsMonolithic)
216 and fused_experts.supports_routing_replay_capture()
@@ -303,132 +232,3 @@ def bind_routed_experts_capturer(
232
233 if num_bound == 0:
234 raise ValueError("No supported MoE router found for routed-experts capture.")
306 -
307 -
308 -def get_routed_experts_attn_gid(kv_cache_config: KVCacheConfig) -> int:
309 - """Return the full-attention KV cache group used for routed experts."""
310 - for gid, group in enumerate(kv_cache_config.kv_cache_groups):
311 - if is_full_attention_spec(group.kv_cache_spec):
312 - return gid
313 - raise ValueError("Routed-experts capture requires a full-attention KV cache group.")
314 -
315 -
316 -class RoutedExpertsManager:
317 - """Scheduler-side slot-indexed buffer for routed experts.
318 -
319 - Lives on CPU in the scheduler process. Each slot corresponds to
320 - ``block_id * block_size + offset_in_block`` where ``block_id`` is
321 - drawn from the physical KV-cache block pool, so routing data is
322 - tied to physical blocks and naturally survives preemption for
323 - prefix-cached blocks (prefix hits re-expose the same slots).
324 -
325 - Data flow per step:
326 - 1. Worker D2Hs its device capture buffer into
327 - :class:`RoutedExpertsLists` and returns it via
328 - :class:`ModelRunnerOutput`.
329 - 2. Scheduler calls :meth:`store_batch` with that step's
330 - ``(routing_data, slot_mapping)`` — a single CPU->CPU
331 - fancy-index assign, ~few MB per step.
332 - 3. On request completion / abort / preemption, the scheduler
333 - calls :meth:`get` with the request's block IDs to recover
334 - the full per-token routing.
335 -
336 - Memory: ``routed_experts_by_slot`` is sized for the whole block
337 - pool (``num_blocks * block_size`` slots). For large block pools
338 - this can reach multiple GB; see the init log for the exact size.
339 - """
340 -
341 - def __init__(
342 - self,
343 - vllm_config: VllmConfig,
344 - kv_cache_config: KVCacheConfig,
345 - ) -> None:
346 - # Hybrid models also have KV groups whose slot layout differs.
347 - self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
348 - attn_group = kv_cache_config.kv_cache_groups[self.attn_gid]
349 - self.block_size = attn_group.kv_cache_spec.block_size
350 -
351 - # All kv_cache_groups share the same physical block pool, so
352 - # block IDs span [0, num_blocks) regardless of how many groups
353 - # exist. Sizing to the full pool avoids index-out-of-range
354 - # when different groups happen to land on the same block.
355 - num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape(
356 - vllm_config
357 - )
358 - max_num_slots = kv_cache_config.num_blocks * self.block_size
359 - # Expert IDs are 0..num_experts-1; uint8 fits 256 distinct
360 - # values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping
361 - # this narrow matters because the slot buffer is sized for the
362 - # whole block pool and can reach multiple GB.
363 - expert_id_dtype = np.uint8 if num_experts <= 256 else np.uint16
364 - self.routed_experts_by_slot = np.zeros(
365 - (
366 - max_num_slots,
367 - num_layers,
368 - num_experts_per_tok,
369 - ),
370 - dtype=expert_id_dtype,
371 - )
372 - logger.info(
373 - "RoutedExpertsManager CPU buffer: %.2f GB "
374 - "(slots=%d, layers=%d, top_k=%d, dtype=%s)",
375 - self.routed_experts_by_slot.nbytes / 1e9,
376 - max_num_slots,
377 - num_layers,
378 - num_experts_per_tok,
379 - self.routed_experts_by_slot.dtype.name,
380 - )
381 -
382 - def store_batch(self, data: np.ndarray, slot_mapping: np.ndarray) -> None:
383 - """Persist one step's routed experts into the slot buffer.
384 -
385 - Equivalent to ``slot_buffer[slot_mapping] = data``; numpy fancy
386 - indexing handles repeated / out-of-order indices. Called once
387 - per scheduler step in ``update_from_output``.
388 - """
389 - self.routed_experts_by_slot[slot_mapping] = data
390 -
391 - def get(
392 - self,
393 - block_ids: list[int],
394 - num_tokens: int,
395 - token_start: int = 0,
396 - ) -> np.ndarray:
397 - """Read routed experts data for a completed / preempted request.
398 -
399 - Reconstructs a per-token slot_mapping from the request's block
400 - IDs and returns the routing slice. Because numpy fancy indexing
401 - returns a **copy** (not a view), the returned ndarray is safe
402 - to hold across subsequent :meth:`store_batch` calls — do not
403 - replace the fancy index with a slice without re-verifying.
404 -
405 - Args:
406 - block_ids: Block IDs from the attention KV-cache group.
407 - num_tokens: Number of tokens that have gone through a forward
408 - pass and therefore have routing data written to their
409 - slots (typically ``request.num_tokens - 1``; the last
410 - sampled token has not been forwarded yet). Slots beyond
411 - ``request.num_computed_tokens`` are zero-initialized.
412 - token_start: Skip the first ``token_start`` tokens from the
413 - result. The slot_mapping is sliced before the fancy-index
414 - read, so only the requested slots are fetched — no large
415 - intermediate array is allocated. Clamped to
416 - ``[0, num_tokens]`` automatically.
417 -
418 - Returns:
419 - Array of shape (num_tokens - token_start, num_layers,
420 - num_experts_per_tok).
421 -
422 - """
423 - bs = self.block_size
424 - block_ids_array = np.array(block_ids, dtype=np.int32)
425 - block_offsets = np.arange(bs)
426 - # slot = block_id * block_size + offset_in_block; flatten the
427 - # (num_blocks, block_size) grid and trim to num_tokens, then
428 - # skip the first token_start entries so only the requested
429 - # range is fetched in a single fancy-index read.
430 - slot_mapping = (
431 - block_ids_array.reshape(-1, 1) * bs + block_offsets.reshape(1, -1)
432 - ).flatten()[:num_tokens]
433 - slot_mapping = slot_mapping[token_start:]
434 - return self.routed_experts_by_slot[slot_mapping]
vllm/outputs.py
+4
@@ -194,6 +194,10 @@ class RequestOutput:
194 completion.cumulative_logprob = (
195 next_completion.cumulative_logprob
196 )
197 + # R3 is returned on the terminal output and must survive
198 + # aggregation with earlier chunks that have no R3.
199 + if next_completion.routed_experts is not None:
200 + completion.routed_experts = next_completion.routed_experts
201 completion.finish_reason = next_completion.finish_reason
202 completion.stop_reason = next_completion.stop_reason
203 else:
vllm/v1/core/sched/output.py
+7
@@ -14,6 +14,9 @@ if TYPE_CHECKING:
14 import numpy.typing as npt
15 import torch
16
17 + from vllm.distributed.aux_output_connector.connector import (
18 + AuxOutputConnectorMetadata,
19 + )
20 from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
21 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
22 from vllm.lora.request import LoRARequest
@@ -23,6 +26,7 @@ if TYPE_CHECKING:
26 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
27 from vllm.v1.request import Request
28 else:
29 + AuxOutputConnectorMetadata = object
30 ECConnectorMetadata = object
31 KVConnectorMetadata = object
32 KVCacheBlockCopy = object
@@ -285,6 +289,9 @@ class SchedulerOutput:
289 # synchronously during this step (load_async=False).
290 has_sync_kv_loads: bool = False
291
292 + # Execution-auxiliary output control metadata consumed by the worker connector.
293 + aux_output_connector_metadata: AuxOutputConnectorMetadata | None = None
294 +
295 # EC Cache Connector metadata
296 ec_connector_metadata: ECConnectorMetadata | None = None
297 # EC Cache Manager metadata
vllm/v1/core/sched/scheduler.py
+33 -98
@@ -9,6 +9,7 @@ from typing import Any
9
10 from vllm.compilation.cuda_graph import CUDAGraphStat
11 from vllm.config import KVEventsConfig, VllmConfig
12 +from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
13 from vllm.distributed.ec_transfer.ec_connector.base import (
14 ECConnectorBase,
15 ECConnectorMetadata,
@@ -26,9 +27,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
27 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
28 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
29 from vllm.logger import init_logger
29 -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
30 - RoutedExpertsManager,
31 -)
30 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
31 from vllm.multimodal.encoder_budget import MultiModalBudget
32 from vllm.multimodal.utils import get_mm_features_in_window
@@ -393,26 +391,13 @@ class Scheduler(SchedulerInterface):
391 if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
392 self.perf_metrics = ModelMetrics(vllm_config)
393
396 - self.enable_return_routed_experts = (
397 - vllm_config.model_config.enable_return_routed_experts
394 + self.aux_output_connector = (
395 + AuxOutputSchedulerConnector()
396 + if vllm_config.aux_output_config.enabled
397 + else None
398 )
399 self.return_sampling_mask = vllm_config.model_config.return_sampling_mask
400
401 - if self.enable_return_routed_experts:
402 - assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (
403 - "enable_return_routed_experts does not support context parallelism "
404 - "(dcp_world_size > 1 or pcp_world_size > 1)"
405 - )
406 -
407 - self.routed_experts_mgr = RoutedExpertsManager(
408 - vllm_config=vllm_config,
409 - kv_cache_config=kv_cache_config,
410 - )
411 - # Block-ID snapshot taken at schedule time (before forward),
412 - # so update_from_output can read slot data even if a later
413 - # schedule() frees the blocks (async scheduling race).
414 - self._re_block_ids: dict[str, list[int]] = {}
415 -
401 self._pause_state: PauseState = PauseState.UNPAUSED
402
403 # In-flight requests still prefilling (prefill chunks + in-progress
@@ -1487,6 +1472,13 @@ class Scheduler(SchedulerInterface):
1472 meta = self._build_kv_connector_meta(self.connector, scheduler_output)
1473 scheduler_output.kv_connector_metadata = meta
1474
1475 + if self.aux_output_connector is not None:
1476 + scheduler_output.aux_output_connector_metadata = (
1477 + self.aux_output_connector.build_connector_meta(
1478 + scheduler_output, self.requests
1479 + )
1480 + )
1481 +
1482 # Build the connector meta for ECConnector
1483 if self.ec_connector is not None:
1484 ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta(
@@ -1541,6 +1533,8 @@ class Scheduler(SchedulerInterface):
1533 assert request.status == RequestStatus.RUNNING, (
1534 "Only running requests can be preempted"
1535 )
1536 + if self.aux_output_connector is not None:
1537 + self.aux_output_connector.request_finished(request)
1538 self._free_request_blocks(request)
1539 self.encoder_cache_manager.free(request)
1540 self._inflight_prefills.discard(request)
@@ -1596,22 +1590,6 @@ class Scheduler(SchedulerInterface):
1590 if not request.is_prefill_chunk:
1591 self._inflight_prefills.discard(request)
1592
1599 - # Snapshot block IDs for routed experts before forward starts.
1600 - # A concurrent schedule() may preempt requests and free blocks
1601 - # before update_from_output runs; the snapshot survives that.
1602 - # Use update() to preserve entries from the previous step that
1603 - # have not yet been consumed by update_from_output (async
1604 - # scheduling may call _update_after_schedule again before the
1605 - # prior update_from_output runs).
1606 - if self.enable_return_routed_experts:
1607 - gid = self.routed_experts_mgr.attn_gid
1608 - self._re_block_ids.update(
1609 - {
1610 - rid: self.kv_cache_manager.get_blocks(rid).get_block_ids()[gid]
1611 - for rid in num_scheduled_tokens
1612 - }
1613 - )
1614 -
1593 # Clear the finished and preempted request IDs.
1594 # NOTE: We shouldn't just clear() here because it will also affect
1595 # the scheduler output.
@@ -1962,28 +1940,6 @@ class Scheduler(SchedulerInterface):
1940 num_scheduled_tokens,
1941 )
1942 )
1965 -
1966 - # Persist per-step routed experts into the scheduler-side slot
1967 - # buffer (CPU->CPU fancy-index assign; ~few MB per step).
1968 - # MUST precede the per-request routing reads below: stopped
1969 - # requests may terminate on tokens generated in this very step,
1970 - # whose routing was just D2H'd into model_runner_output.
1971 - routing_data = None
1972 - routing_offsets: dict[str, int] = {}
1973 - if model_runner_output.routed_experts is not None:
1974 - re = model_runner_output.routed_experts
1975 - self.routed_experts_mgr.store_batch(re.routing_data, re.slot_mapping)
1976 - routing_data = re.routing_data.astype(
1977 - self.routed_experts_mgr.routed_experts_by_slot.dtype,
1978 - copy=False,
1979 - )
1980 - # Build offset map using model runner's request order
1981 - # (input_batch ordering), NOT scheduler dict order.
1982 - offset = 0
1983 - for rid in model_runner_output.req_ids:
1984 - routing_offsets[rid] = offset
1985 - offset += num_scheduled_tokens[rid]
1986 -
1943 # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
1944 # the below loop can be a performance bottleneck. We should do our best
1945 # to avoid expensive operations inside the loop.
@@ -2076,7 +2032,6 @@ class Scheduler(SchedulerInterface):
2032 ec_transfer_params = None
2033 prefill_stats = None
2034 status_before_stop = request.status
2079 - num_output_tokens_before = len(request._output_token_ids)
2035
2036 # Check for stop and update request status.
2037 if new_token_ids:
@@ -2113,48 +2068,13 @@ class Scheduler(SchedulerInterface):
2068 stopped = True
2069
2070 routed_experts = None
2116 - if (
2117 - self.enable_return_routed_experts
2118 - and routing_data is not None
2119 - and new_token_ids
2120 - ):
2121 - req_offset = routing_offsets[req_id]
2122 - end = req_offset + num_tokens_scheduled
2123 - block_ids = self._re_block_ids.pop(req_id, [])
2124 - if num_output_tokens_before == 0:
2125 - # Prefill completed: read full prompt routing from
2126 - # slot buffer using the block-ID snapshot taken at
2127 - # schedule time (immune to async preemption).
2128 - if (
2129 - request.sampling_params is not None
2130 - and request.sampling_params.routed_experts_prompt_start
2131 - is not None
2132 - ):
2133 - prompt_start = (
2134 - request.sampling_params.routed_experts_prompt_start
2135 - )
2136 - assert prompt_start <= request.num_prompt_tokens
2137 - else:
2138 - prompt_start = 0
2139 - routed_experts = self.routed_experts_mgr.get(
2140 - block_ids,
2141 - request.num_prompt_tokens,
2142 - token_start=prompt_start,
2143 - )
2144 - else:
2145 - if scheduled_spec_token_ids:
2146 - # Spec decode: accepted tokens at the START of
2147 - # the scheduled range, rejected at the end.
2148 - routed_experts = routing_data[
2149 - req_offset : req_offset + len(new_token_ids)
2150 - ]
2151 - else:
2152 - # Normal decode / re-prefill: token(s) at the END.
2153 - routed_experts = routing_data[end - len(new_token_ids) : end]
2154 -
2071 should_emit_output = bool(
2072 new_token_ids or pooler_output is not None or stopped
2073 )
2074 + if self.aux_output_connector is not None and should_emit_output:
2075 + routed_experts = self.aux_output_connector.take_output(
2076 + request, model_runner_output.aux_output_connector_output
2077 + )
2078 if should_emit_output:
2079 prefill_stats = request.take_prefill_stats()
2080 if prefill_stats is not None:
@@ -2641,6 +2561,8 @@ class Scheduler(SchedulerInterface):
2561 ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
2562 assert request.is_finished()
2563
2564 + if self.aux_output_connector is not None:
2565 + self.aux_output_connector.request_finished(request)
2566 self._inflight_prefills.discard(request)
2567 connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
2568
@@ -2773,6 +2695,16 @@ class Scheduler(SchedulerInterface):
2695 Otherwise, this method will only reset the KV prefix cache when there
2696 is no running requests taking KV cache.
2697 """
2698 + if reset_running_requests and self.aux_output_connector is not None:
2699 + if self._pause_state != PauseState.PAUSED_ALL:
2700 + raise RuntimeError(
2701 + "AuxOutput Connector only supports resetting running requests "
2702 + "after pause(mode='keep')."
2703 + )
2704 + if any(request.num_in_flight_tokens for request in self.requests.values()):
2705 + raise RuntimeError(
2706 + "AuxOutput Connector cannot reset while model output is in flight."
2707 + )
2708 if reset_running_requests:
2709 # For logging.
2710 timestamp = time.monotonic()
@@ -2803,6 +2735,9 @@ class Scheduler(SchedulerInterface):
2735 if reset_connector:
2736 reset_successful = self.reset_connector_cache() and reset_successful
2737
2738 + if reset_successful and self.aux_output_connector is not None:
2739 + self.aux_output_connector.reset()
2740 +
2741 return reset_successful
2742
2743 def reset_connector_cache(self) -> bool:
vllm/v1/executor/ray_utils.py
+4 -8
@@ -233,14 +233,10 @@ def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> No
233 token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
234 )
235
236 - if output.routed_experts is not None:
237 - routing_data, slot_mapping = output.routed_experts
238 - routing_data_c = _copy_if_readonly(routing_data)
239 - slot_mapping_c = _copy_if_readonly(slot_mapping)
240 - if routing_data_c is not routing_data or slot_mapping_c is not slot_mapping:
241 - output.routed_experts = type(output.routed_experts)(
242 - routing_data_c, slot_mapping_c
243 - )
236 + aux_output = output.aux_output_connector_output
237 + if aux_output is not None:
238 + for request_output in aux_output.values():
239 + request_output.rows = _copy_if_readonly(request_output.rows)
240
241
242 class FutureWrapper(Future):
vllm/v1/outputs.py
+3 -77
@@ -15,6 +15,7 @@ from vllm.utils.torch_utils import PIN_MEMORY
15 from vllm.v1.core.sched.output import SchedulerOutput
16
17 if TYPE_CHECKING:
18 + from vllm.distributed.aux_output_connector.connector import AuxRequestOutput
19 from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata
20 from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats
21 from vllm.distributed.kv_events import KVConnectorKVEvents
@@ -23,6 +24,7 @@ if TYPE_CHECKING:
24 )
25 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
26 else:
27 + AuxRequestOutput = object
28 KVConnectorStats = object
29 KVConnectorWorkerMetadata = object
30 KVConnectorKVEvents = object
@@ -188,73 +190,6 @@ class LogprobsTensors(NamedTuple):
190 )
191
192
191 -class RoutedExpertsTensors(NamedTuple):
192 - """Device-side snapshot of routed experts data, pending async D2H.
193 -
194 - Produced by :class:`GPUModelRunner` at the end of each async-scheduled
195 - step. The copy stream waits on the default stream, then issues
196 - non-blocking D2H via :meth:`to_cpu_nonblocking` into a pinned CPU
197 - buffer; :class:`AsyncGPUModelRunnerOutput.get_output` synchronizes
198 - the copy before the scheduler reads it.
199 -
200 - Sliced to ``total_num_scheduled_tokens`` (step-level, across all
201 - requests — NOT per-request). Both ``routing_data`` and
202 - ``slot_mapping`` must be private clones when sourced from shared
203 - capturer / prepare-input buffers, so the next forward pass /
204 - ``_prepare_inputs`` on the default stream does not race with a
205 - D2H still pending on the copy stream.
206 - """
207 -
208 - # (num_scheduled_tokens, num_layers, num_experts_per_tok)
209 - routing_data: torch.Tensor
210 - # (num_scheduled_tokens,)
211 - slot_mapping: torch.Tensor
212 -
213 - def to_cpu_nonblocking(self) -> "RoutedExpertsTensors":
214 - """Issue non-blocking D2H on the current stream.
215 -
216 - NOTE: ``non_blocking=True`` only delivers true overlap when the
217 - CPU target is pinned. The current fallback here allocates a
218 - new pageable CPU tensor per call, which silently degrades to a
219 - synchronous copy; acceptable because the sync happens on the
220 - dedicated copy stream, not the default stream.
221 - """
222 - if self.routing_data.device.type == "cpu":
223 - return self
224 - return RoutedExpertsTensors(
225 - self.routing_data.to("cpu", non_blocking=True),
226 - self.slot_mapping.to("cpu", non_blocking=True),
227 - )
228 -
229 - def tolists(self) -> "RoutedExpertsLists":
230 - """Convert to the numpy-backed form consumed by the scheduler.
231 -
232 - ``.cpu()`` is a no-op when the tensor is already on CPU, so this
233 - is cheap for the post-D2H case; for raw device tensors it will
234 - synchronously block, which is only reached in tests.
235 - """
236 - return RoutedExpertsLists(
237 - self.routing_data.cpu().numpy(),
238 - self.slot_mapping.cpu().numpy(),
239 - )
240 -
241 -
242 -class RoutedExpertsLists(NamedTuple):
243 - """CPU-side routed experts, the form :meth:`RoutedExpertsManager.store_batch`
244 - consumes.
245 -
246 - Batched per scheduler step: the leading dim is the number of tokens
247 - scheduled across all requests in this step (``total_num_scheduled_tokens``),
248 - not per-request tokens. ``slot_mapping[i]`` tells the scheduler which
249 - physical KV-cache slot row ``i`` of ``routing_data`` belongs to.
250 - """
251 -
252 - # (num_scheduled_tokens, num_layers, num_experts_per_tok)
253 - routing_data: np.ndarray
254 - # (num_scheduled_tokens,)
255 - slot_mapping: np.ndarray
256 -
257 -
193 # [num_reqs, <dynamic>]
194 # The shape of each element depends on the pooler used
195 PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None]
@@ -361,16 +296,7 @@ class ModelRunnerOutput:
296 # information related to cudagraph execution
297 cudagraph_stats: CUDAGraphStat | None = None
298
364 - # Per-step routed experts data captured by the worker.
365 - # ``routing_data`` shape: (num_scheduled_tokens, num_layers,
366 - # num_experts_per_tok); expert IDs as uint8/uint16.
367 - # ``slot_mapping`` shape: (num_scheduled_tokens,); physical KV-cache
368 - # slot for each row of routing_data.
369 - # ``num_scheduled_tokens`` is step-level (total across all requests
370 - # in this step), not per-request. The scheduler persists this into
371 - # its slot buffer via ``slot_buffer[slot_mapping] = routing_data``.
372 - # ``None`` when ``enable_return_routed_experts`` is off.
373 - routed_experts: RoutedExpertsLists | None = None
299 + aux_output_connector_output: dict[str, AuxRequestOutput] | None = None
300
301 # ``None`` when ``return_sampling_mask`` is off.
302 sampling_masks: SamplingMaskLists | None = None
vllm/v1/worker/gpu/async_utils.py
+13 -9
@@ -15,12 +15,12 @@ from vllm.v1.outputs import (
15 LogprobsTensors,
16 ModelRunnerOutput,
17 PoolerOutput,
18 - RoutedExpertsTensors,
18 )
19 from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
20 from vllm.v1.worker.utils import raise_if_nan_logits
21
22 if TYPE_CHECKING:
23 + from vllm.distributed.aux_output_connector.worker import PendingAuxOutput
24 from vllm.v1.worker.gpu.input_batch import InputBatch
25
26
@@ -120,8 +120,8 @@ class AsyncOutput(AsyncModelRunnerOutput):
120 num_sampled_tokens: torch.Tensor,
121 main_stream: torch.cuda.Stream,
122 copy_stream: torch.cuda.Stream,
123 - check_ep_fault: bool = False,
124 - routed_experts: RoutedExpertsTensors | None = None,
123 + check_ep_fault: bool,
124 + pending_aux_output: "PendingAuxOutput | None",
125 ):
126 # NOTE(woosuk): We must retain references to the GPU tensors,
127 # as the copy operations are performed on a different CUDA stream than
@@ -129,7 +129,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
129 self.model_runner_output = model_runner_output
130 self.sampler_output = sampler_output
131 self.num_sampled_tokens = num_sampled_tokens
132 - self.routed_experts = routed_experts
132 + self.pending_aux_output = pending_aux_output
133 # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
134 self.copy_event = torch.cuda.Event(blocking=True)
135 self._has_fault: torch.Tensor | None = None
@@ -152,13 +152,15 @@ class AsyncOutput(AsyncModelRunnerOutput):
152 self.sampling_mask_tensors = (
153 sampler_output.sampling_mask_tensors.to_cpu_nonblocking()
154 )
155 - self.routed_experts_cpu: RoutedExpertsTensors | None = None
156 - if routed_experts is not None:
157 - self.routed_experts_cpu = routed_experts.to_cpu_nonblocking()
155 self.prompt_logprobs_dict = {
156 k: v.to_cpu_nonblocking() if v is not None else None
157 for k, v in self.model_runner_output.prompt_logprobs_dict.items()
158 }
159 + if self.pending_aux_output is not None:
160 + self.pending_aux_output.enqueue_cpu_copy(
161 + num_sampled=self.num_sampled_tokens_np,
162 + num_rejected=async_copy_to_np(sampler_output.num_rejected),
163 + )
164 if check_ep_fault:
165 has_fault = get_ep_all2all_manager().query_fault()
166 self._has_fault = has_fault.to("cpu", non_blocking=True)
@@ -192,8 +194,10 @@ class AsyncOutput(AsyncModelRunnerOutput):
194 if self.logprobs_tensors is not None:
195 self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
196 self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
195 - if self.routed_experts_cpu is not None:
196 - self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists()
197 + if self.pending_aux_output is not None:
198 + self.model_runner_output.aux_output_connector_output = (
199 + self.pending_aux_output.process_output()
200 + )
201
202 if self._has_fault is not None and self._has_fault.item():
203 mask = get_ep_all2all_manager().query_active_mask()
vllm/v1/worker/gpu/model_runner.py
+24 -24
@@ -32,6 +32,10 @@ from vllm.compilation.counter import compilation_counter
32 from vllm.compilation.cuda_graph import CUDAGraphStat
33 from vllm.config import VllmConfig
34 from vllm.config.compilation import CUDAGraphMode
35 +from vllm.distributed.aux_output_connector.worker import (
36 + AuxOutputWorkerConnector,
37 + get_aux_output_connector,
38 +)
39 from vllm.distributed.parallel_state import (
40 get_dcp_group,
41 get_pp_group,
@@ -39,10 +43,6 @@ from vllm.distributed.parallel_state import (
43 from vllm.forward_context import BatchDescriptor, set_forward_context
44 from vllm.logger import init_logger
45 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
42 -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
43 - RoutedExpertsCapturer,
44 - bind_routed_experts_capturer,
45 -)
46 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
47 initialize_mamba_ssu_backend,
48 )
@@ -73,7 +73,6 @@ from vllm.v1.outputs import (
73 DraftTokenIds,
74 ECConnectorOutput,
75 ModelRunnerOutput,
76 - RoutedExpertsTensors,
76 )
77 from vllm.v1.watermarking import create_watermarker
78 from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
@@ -348,7 +347,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
347
348 # Expert parallelism load balancer.
349 self.eplb = EPLBController(self.parallel_config, self.device)
351 - self.routed_experts_capturer: RoutedExpertsCapturer | None = None
350 + # The AuxOutput Connector owns R3 capture, copying, and storage.
351 + self.aux_output_connector: AuxOutputWorkerConnector | None = None
352
353 set_offloader(create_offloader(self.vllm_config.offload_config))
354
@@ -356,15 +356,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
356 self.max_model_len = max_model_len
357 self.req_states.max_model_len = max_model_len
358
359 - def init_routed_experts_capturer(self) -> None:
360 - """Initialize target-model capture on every participating worker."""
361 - self.routed_experts_capturer = RoutedExpertsCapturer(
362 - max_num_batched_tokens=self.max_num_tokens,
363 - vllm_config=self.vllm_config,
364 - kv_cache_config=self.kv_cache_config,
365 - )
366 - bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
367 -
359 def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
360 tasks: list[SupportedTask] = []
361 if self.model_config.runner_type == "generate":
@@ -756,6 +747,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
747 else:
748 self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
749
750 + # AuxOutput connector requires resolved kv_cache_config.
751 + if self.vllm_config.aux_output_config.enabled:
752 + self.aux_output_connector = get_aux_output_connector(
753 + self.model, self.vllm_config, kv_cache_config
754 + )
755 +
756 def _init_kv_zero_meta(self) -> None:
757 """Build KV-block zeroing metadata; invoked from gpu_worker."""
758 self.kv_block_zeroer = KVBlockZeroer(
@@ -1648,6 +1645,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1645 self.add_requests(scheduler_output)
1646 self.update_requests(scheduler_output)
1647 self.block_tables.apply_staged_writes()
1648 + if self.aux_output_connector is not None:
1649 + # Register this step before the GPU forward.
1650 + self.aux_output_connector.begin_step(
1651 + scheduler_output.aux_output_connector_metadata
1652 + )
1653 if scheduler_output.total_num_scheduled_tokens == 0:
1654 # No need to run the model.
1655 empty_output = self.kv_connector.no_forward(scheduler_output)
@@ -1969,11 +1971,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1971 aux_hidden_states = None
1972 output_intermediate_tensors = model_output
1973
1972 - routed_experts = None
1973 - if not dummy_run and (capturer := self.routed_experts_capturer) is not None:
1974 - assert slot_mappings is not None
1975 - routed_experts = capturer.get_routed_experts(slot_mappings, num_toks)
1976 -
1974 finished_req_ids = scheduler_output.finished_req_ids
1975 self.execute_model_state = ExecuteModelState(
1976 input_batch=input_batch,
@@ -1984,7 +1981,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1981 dp_sync=dp_sync,
1982 finished_req_ids=finished_req_ids,
1983 ec_connector_output=ec_connector_output,
1987 - routed_experts=routed_experts,
1984 cudagraph_stats=cudagraph_stats,
1985 )
1986
@@ -2014,7 +2010,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2010 dp_sync = self.execute_model_state.dp_sync
2011 finished_req_ids = self.execute_model_state.finished_req_ids
2012 ec_connector_output = self.execute_model_state.ec_connector_output
2017 - routed_experts = self.execute_model_state.routed_experts
2013 cudagraph_stats = self.execute_model_state.cudagraph_stats
2014 self.execute_model_state = None
2015
@@ -2083,6 +2078,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2078 prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]
2079 cudagraph_stats=cudagraph_stats,
2080 )
2081 + pending_aux_output = None
2082 + if self.aux_output_connector is not None:
2083 + pending_aux_output = self.aux_output_connector.prepare_output(input_batch)
2084 +
2085 # Start async output copy here so that it can overlap with speculator proposal.
2086 async_output = AsyncOutput(
2087 model_runner_output=model_runner_output,
@@ -2091,7 +2090,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2090 main_stream=self.main_stream,
2091 copy_stream=self.output_copy_stream,
2092 check_ep_fault=self.check_ep_fault,
2094 - routed_experts=routed_experts,
2093 + pending_aux_output=pending_aux_output,
2094 )
2095
2096 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
@@ -2244,6 +2243,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2243 """Release GPU tensors (model weights, KV caches, workspace) so that
2244 memory is reclaimable when running in the same process."""
2245 torch.accelerator.synchronize()
2246 + if self.aux_output_connector is not None:
2247 + self.aux_output_connector.close()
2248 self.cudagraph_manager = None
2249 self.fast_prefill = None
2250 if hasattr(self, "kv_caches"):
@@ -2316,7 +2317,6 @@ class ExecuteModelState(NamedTuple):
2317 dp_sync: DPSyncState | None
2318 finished_req_ids: set[str]
2319 ec_connector_output: ECConnectorOutput | None
2319 - routed_experts: RoutedExpertsTensors | None
2320 cudagraph_stats: CUDAGraphStat | None
2321
2322
vllm/v1/worker/gpu/sample/batch_shard.py
-2
@@ -617,8 +617,6 @@ def gather_sampler_output(
617 device=device,
618 )
619 if local_output is not None:
620 - assert local_output.num_sampled is not None
621 - assert local_output.num_rejected is not None
620 assert not gather_num_nans or local_output.num_nans is not None
621 num_src_cols = min(
622 local_output.sampled_token_ids.shape[1], max_num_logits_per_req
vllm/v1/worker/gpu/sample/output.py
+2 -2
@@ -17,8 +17,8 @@ class SamplerOutput:
17 sampled_token_ids: torch.Tensor
18 logprobs_tensors: LogprobsTensors | None
19 num_nans: torch.Tensor | None
20 - num_sampled: torch.Tensor | None
21 - num_rejected: torch.Tensor | None = None
20 + num_sampled: torch.Tensor
21 + num_rejected: torch.Tensor
22 sampling_mask_tensors: SamplingMaskTensors | None = None
23
24
vllm/v1/worker/gpu_model_runner.py
-127
@@ -61,10 +61,6 @@ from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType
61 from vllm.model_executor.layers.attention import Attention
62 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
63 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
64 -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
65 - RoutedExpertsCapturer,
66 - bind_routed_experts_capturer,
67 -)
64 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType
65 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
66 initialize_mamba_ssu_backend,
@@ -176,8 +172,6 @@ from vllm.v1.outputs import (
172 LogprobsTensors,
173 ModelRunnerOutput,
174 PoolerOutput,
179 - RoutedExpertsLists,
180 - RoutedExpertsTensors,
175 SamplerOutput,
176 make_empty_encoder_model_runner_output,
177 )
@@ -293,7 +287,6 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
287 invalid_req_indices: list[int],
288 async_output_copy_stream: torch.cuda.Stream,
289 vocab_size: int,
296 - routed_experts: RoutedExpertsTensors | None = None,
290 check_ep_fault: bool = False,
291 num_nans: torch.Tensor | None = None,
292 ):
@@ -309,7 +302,6 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
302 self._sampled_token_ids = sampled_token_ids
303 self.vocab_size = vocab_size
304 self._logprobs_tensors = logprobs_tensors
312 - self._routed_experts = routed_experts
305 self._num_nans = num_nans
306 self._has_fault: torch.Tensor | None = None
307
@@ -325,11 +317,6 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
317 if self._logprobs_tensors
318 else None
319 )
328 - self._routed_experts_cpu = (
329 - self._routed_experts.to_cpu_nonblocking()
330 - if self._routed_experts is not None
331 - else None
332 - )
320 self._num_nans_cpu = (
321 self._num_nans.to("cpu", non_blocking=True)
322 if self._num_nans is not None
@@ -370,10 +357,6 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
357 output.sampled_token_ids = valid_sampled_token_ids
358 output.logprobs = logprobs_lists
359
373 - if self._routed_experts_cpu is not None:
374 - output.routed_experts = self._routed_experts_cpu.tolists()
375 - del self._routed_experts
376 -
360 if self._num_nans_cpu is not None:
361 output.num_nans_in_logits = nans_to_dict(
362 self._num_nans_cpu.tolist(), output.req_id_to_index
@@ -537,9 +520,6 @@ class GPUModelRunner(
520 # These will be overridden in load_model()
521 self.is_multimodal_pruning_enabled = False
522 self.requires_sequential_video_encoding = False
540 - # Set to True after init_routed_experts_capturer() completes.
541 - # Prevents routed experts code from running during profiling/dummy run.
542 - self.routed_experts_initialized = False
523 self.max_model_len = model_config.max_model_len
524
525 # Always set to false after the first forward pass
@@ -2360,18 +2340,6 @@ class GPUModelRunner(
2340 block_table_gid_0 = _get_block_table(0)
2341 slot_mapping_gid_0 = slot_mappings[0]
2342
2363 - if self.routed_experts_initialized:
2364 - # Copy this step's attention slot_mapping into our private
2365 - # device buffer. The shared ``slot_mappings[attn_gid]`` is
2366 - # owned by the attention block table and will be overwritten
2367 - # by the next ``_prepare_inputs``; we need a stable snapshot
2368 - # because the async D2H may still be in flight on the copy
2369 - # stream when the next step runs.
2370 - slot_mapping_attn = slot_mappings[self.routed_experts_capturer.attn_gid]
2371 - self.routed_experts_slot_mapping_device[:num_tokens].copy_(
2372 - slot_mapping_attn[:num_tokens]
2373 - )
2374 -
2343 num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[
2344 :num_reqs_padded
2345 ]
@@ -3745,20 +3713,6 @@ class GPUModelRunner(
3713 invalid_req_indices = []
3714 logprobs_lists = None
3715 if not self.use_async_scheduling:
3748 - # Sync scheduling: issue routed experts D2H into the pinned
3749 - # CPU buffer BEFORE ``_to_list`` below. ``_to_list`` does
3750 - # ``event.synchronize()`` on the async copy stream which
3751 - # waits for every D2H queued on the default stream since
3752 - # the last sync, so this enqueue is naturally covered
3753 - # without requiring its own synchronize.
3754 - if self.routed_experts_initialized:
3755 - buf = self.routed_experts_capturer.get_device_buffer()
3756 - total = scheduler_output.total_num_scheduled_tokens
3757 - self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
3758 - self.routed_experts_slot_mapping_cpu[:total].copy_(
3759 - self.routed_experts_slot_mapping_device[:total],
3760 - non_blocking=True,
3761 - )
3716 with gpu_sync_allowed():
3717 # Get the valid generated tokens.
3718 max_gen_len = sampled_token_ids.shape[-1]
@@ -4802,39 +4756,14 @@ class GPUModelRunner(
4756 else None,
4757 num_nans_in_logits=num_nans_in_logits,
4758 cudagraph_stats=cudagraph_stats,
4805 - routed_experts=None,
4759 )
4760
4761 if not self.use_async_scheduling:
4809 - if self.routed_experts_initialized:
4810 - # Sync path: D2H was issued in ``_bookkeeping_sync`` and
4811 - # synchronized by ``_to_list``'s event.synchronize(), so
4812 - # the pinned buffers are ready to be wrapped as numpy.
4813 - total = scheduler_output.total_num_scheduled_tokens
4814 - output.routed_experts = RoutedExpertsLists(
4815 - routing_data=self.routed_experts_cpu[:total].numpy(),
4816 - slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
4817 - )
4762 return output
4763
4764 with record_function_or_nullcontext(
4765 "gpu_model_runner: AsyncGPUModelRunnerOutput"
4766 ):
4823 - # Async path: produce a device-side snapshot that the async
4824 - # copy stream can D2H later. Both tensors must be private
4825 - # clones because:
4826 - # - ``routing_data`` source is the shared capturer buffer,
4827 - # which the next forward overwrites on the default stream.
4828 - # - ``slot_mapping`` source is our own
4829 - # ``routed_experts_slot_mapping_device``, which the
4830 - # next ``_prepare_inputs`` overwrites on the default
4831 - # stream while the D2H is still pending on the copy
4832 - # stream.
4833 - # Without clones, the copy stream would read torn data.
4834 - routed_experts_snapshot = self.get_routed_experts(
4835 - scheduler_output.total_num_scheduled_tokens
4836 - )
4837 -
4767 async_output = AsyncGPUModelRunnerOutput(
4768 model_runner_output=output,
4769 sampled_token_ids=sampler_output.sampled_token_ids,
@@ -4842,7 +4771,6 @@ class GPUModelRunner(
4771 invalid_req_indices=invalid_req_indices,
4772 async_output_copy_stream=self._get_or_create_async_output_copy_stream(),
4773 vocab_size=self.input_batch.vocab_size,
4845 - routed_experts=routed_experts_snapshot,
4774 check_ep_fault=self.check_ep_fault,
4775 num_nans=num_nans_device,
4776 )
@@ -7468,61 +7396,6 @@ class GPUModelRunner(
7396 kv_transfer_group.register_kv_caches(kv_caches)
7397 kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
7398
7471 - def get_routed_experts(
7472 - self,
7473 - num_tokens: int,
7474 - ) -> RoutedExpertsTensors | None:
7475 - if not self.routed_experts_initialized:
7476 - return None
7477 -
7478 - device_buffer = self.routed_experts_capturer.get_device_buffer()
7479 - return RoutedExpertsTensors(
7480 - routing_data=device_buffer[:num_tokens].clone(),
7481 - slot_mapping=self.routed_experts_slot_mapping_device[:num_tokens].clone(),
7482 - )
7483 -
7484 - def init_routed_experts_capturer(self):
7485 - logger.info(
7486 - "Initializing routed experts capturer, enable_return_routed_experts: %s",
7487 - self.model_config.enable_return_routed_experts,
7488 - )
7489 - self.routed_experts_capturer = RoutedExpertsCapturer(
7490 - max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
7491 - vllm_config=self.vllm_config,
7492 - kv_cache_config=self.kv_cache_config,
7493 - )
7494 - bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
7495 -
7496 - # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on
7497 - # the sync scheduling path. Shape / dtype mirror the device
7498 - # capturer exactly so ``copy_`` is a straight memcpy.
7499 - self.routed_experts_cpu = torch.empty(
7500 - self.routed_experts_capturer.device_buffer.shape,
7501 - dtype=self.routed_experts_capturer.device_buffer.dtype,
7502 - device="cpu",
7503 - pin_memory=PIN_MEMORY,
7504 - )
7505 - # ``slot_mapping`` dtype is fixed to int64 by
7506 - # ``block_table.slot_mapping``; we mirror that here.
7507 - max_tokens = self.scheduler_config.max_num_batched_tokens
7508 - self.routed_experts_slot_mapping_cpu = torch.empty(
7509 - (max_tokens,),
7510 - dtype=torch.int64,
7511 - device="cpu",
7512 - pin_memory=PIN_MEMORY,
7513 - )
7514 - # Private device buffer so the shared ``block_table.slot_mapping``
7515 - # can be overwritten by the next ``_prepare_inputs`` while the
7516 - # D2H is still pending on the copy stream. Written in
7517 - # ``_prepare_inputs``, read in ``_bookkeeping_sync`` (sync path)
7518 - # or cloned into a snapshot (async path).
7519 - self.routed_experts_slot_mapping_device = torch.empty(
7520 - (max_tokens,),
7521 - dtype=torch.int64,
7522 - device=self.device,
7523 - )
7524 - self.routed_experts_initialized = True
7525 -
7399 def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
7400 """Add encoder-only layers to the KV cache config."""
7401 block_size = self.vllm_config.cache_config.block_size
vllm/v1/worker/gpu_worker.py
-3
@@ -761,9 +761,6 @@ class Worker(WorkerBase):
761 ),
762 )
763
764 - if self.model_config.enable_return_routed_experts:
765 - self.model_runner.init_routed_experts_capturer()
766 -
764 # Build KV-zero metadata outside the CuMem pool so the bookkeeping
765 # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch
766 # allocator and are not discarded during sleep/wake cycles.