test_scheduler.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 7,350 lines 271 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	import dataclasses
4	from concurrent.futures import Future
5	from types import SimpleNamespace
6	from unittest.mock import Mock
7
8	import numpy as np
9	import pytest
10	import torch
11
12	import vllm.envs as envs
13	from vllm.config import (
14	CacheConfig,
15	ECTransferConfig,
16	KVTransferConfig,
17	ModelConfig,
18	SchedulerConfig,
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,
29	HiSparseConnectorScheduler,
30	)
31	from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
32	from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
33	from vllm.multimodal.inputs import (
34	MultiModalFeatureSpec,
35	MultiModalKwargsItem,
36	PlaceholderRange,
37	)
38	from vllm.sampling_params import SamplingParams, StructuredOutputsParams
39	from vllm.utils.hashing import sha256
40	from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
41	from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
42	from vllm.v1.core.kv_cache_manager import KVCacheManager
43	from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
44	from vllm.v1.core.sched.diffusion_scheduler import (
45	DiffusionAsyncScheduler,
46	DiffusionScheduler,
47	diffusion_canvas_width,
48	)
49	from vllm.v1.core.sched.interface import PauseState
50	from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
51	from vllm.v1.core.sched.scheduler import Scheduler
52	from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
53	from vllm.v1.engine import FinishReason
54	from vllm.v1.engine.core import EngineCore
55	from vllm.v1.executor.uniproc_executor import UniProcExecutor
56	from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
57	from vllm.v1.kv_cache_interface import (
58	FullAttentionSpec,
59	KVCacheConfig,
60	KVCacheGroupSpec,
61	MambaSpec,
62	)
63	from vllm.v1.outputs import (
64	DraftTokenIds,
65	ECConnectorOutput,
66	KVConnectorOutput,
67	ModelRunnerOutput,
68	SamplingMaskLists,
69	make_empty_encoder_model_runner_output,
70	)
71	from vllm.v1.request import Request, RequestStatus
72	from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
73
74	from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
75
76	pytestmark = pytest.mark.cpu_test
77
78
79	def test_make_scheduled_encoder_input_stats_output_embeddings():
80	scheduler = create_scheduler()
81	mm_features = [
82	MultiModalFeatureSpec(
83	data=MultiModalKwargsItem.dummy(),
84	modality="image",
85	identifier="image-0",
86	mm_position=PlaceholderRange(offset=0, length=196),
87	),
88	MultiModalFeatureSpec(
89	data=MultiModalKwargsItem.dummy(),
90	modality="video",
91	identifier="video-0",
92	mm_position=PlaceholderRange(offset=200, length=196),
93	),
94	MultiModalFeatureSpec(
95	data=MultiModalKwargsItem.dummy(),
96	modality="audio",
97	identifier="audio-0",
98	mm_position=PlaceholderRange(offset=400, length=49),
99	),
100	]
101	scheduler.requests["req"] = Mock(mm_features=mm_features)
102
103	stats = scheduler._make_scheduled_encoder_input_stats({"req": [0, 1, 2]})
104
105	assert stats is not None
106	assert stats.num_inputs == 3
107	assert stats.output_tokens == 441
108
109
110	def test_scheduled_encoder_input_stats_disabled_without_iteration_logging(
111	monkeypatch: pytest.MonkeyPatch,
112	):
113	scheduler = create_scheduler()
114	make_stats = Mock(side_effect=AssertionError("stats should not be computed"))
115	monkeypatch.setattr(scheduler, "_make_scheduled_encoder_input_stats", make_stats)
116
117	scheduler_output = scheduler.schedule()
118
119	make_stats.assert_not_called()
120	assert scheduler_output.scheduled_encoder_input_stats is None
121
122
123	def test_scheduled_encoder_input_stats_disabled_without_log_stats(
124	monkeypatch: pytest.MonkeyPatch,
125	):
126	scheduler = create_scheduler()
127	scheduler.log_stats = False
128	scheduler.observability_config.enable_logging_iteration_details = True
129	make_stats = Mock(side_effect=AssertionError("stats should not be computed"))
130	monkeypatch.setattr(scheduler, "_make_scheduled_encoder_input_stats", make_stats)
131
132	scheduler_output = scheduler.schedule()
133
134	make_stats.assert_not_called()
135	assert scheduler_output.scheduled_encoder_input_stats is None
136
137
138	def test_add_requests():
139	scheduler = create_scheduler()
140	requests = create_requests(num_requests=10)
141
142	for i, request in enumerate(requests):
143	scheduler.add_request(request)
144	assert request.request_id in scheduler.requests
145	assert len(scheduler.waiting) == i + 1
146
147
148	@pytest.mark.parametrize("max_tokens", [1, 2])
149	def test_routed_experts_prompt_start_at_prompt_end(max_tokens):
150	"""A prompt-end offset should return empty routing data, not crash."""
151	scheduler = create_scheduler()
152	(request,) = create_requests(num_requests=1, max_tokens=max_tokens)
153	request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens
154	scheduler.aux_output_connector = AuxOutputSchedulerConnector()
155	scheduler.add_request(request)
156	scheduler_output = scheduler.schedule()
157	assert scheduler_output.aux_output_connector_metadata.requests == {
158	request.request_id: request.num_prompt_tokens
159	}
160
161	outputs = scheduler.update_from_output(
162	scheduler_output,
163	ModelRunnerOutput(
164	req_ids=[request.request_id],
165	req_id_to_index={request.request_id: 0},
166	sampled_token_ids=[[0]],
167	logprobs=None,
168	prompt_logprobs_dict={},
169	pooler_output=[],
170	aux_output_connector_output={
171	request.request_id: AuxRequestOutput(
172	request.num_prompt_tokens,
173	np.empty((0, 1, 1), dtype=np.uint8),
174	)
175	},
176	),
177	)
178
179	output = outputs[request.client_index].outputs[0]
180	assert output.new_token_ids == [0]
181	assert output.routed_experts.shape == (0, 1, 1)
182
183
184	def test_finish_request():
185	scheduler = create_scheduler()
186	scheduler.aux_output_connector = Mock()
187	requests = create_requests(num_requests=10)
188	for request in requests:
189	scheduler.add_request(request)
190
191	for i, request in enumerate(requests):
192	scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
193	assert request.request_id not in scheduler.requests
194	assert len(scheduler.waiting) == 9 - i
195	scheduler.aux_output_connector.request_finished.assert_called_with(request)
196
197
198	def test_get_num_unfinished_requests():
199	scheduler = create_scheduler()
200	requests = create_requests(num_requests=10)
201	for request in requests:
202	scheduler.add_request(request)
203
204	for i, request in enumerate(requests):
205	scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_STOPPED)
206	assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1
207
208
209	def test_max_num_active_seqs_caps_admission():
210	"""Admission is capped by max_num_active_seqs, not max_num_seqs."""
211	scheduler = create_scheduler(max_num_seqs=16, max_num_active_seqs=3)
212	requests = create_requests(num_requests=10)
213	for request in requests:
214	scheduler.add_request(request)
215
216	output = scheduler.schedule()
217
218	assert len(output.scheduled_new_reqs) == 3
219	assert len(scheduler.running) == 3
220	assert len(scheduler.waiting) == 7
221	# Execution capacity is untouched.
222	assert scheduler.max_num_running_reqs == 16
223
224
225	def _bind_hisparse_connector(scheduler):
226	connector = object.__new__(HiSparseConnector)
227	connector.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False)
228	from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager
229
230	coordinator = get_hisparse_coordinator(make_hisparse_kv_cache_manager(16, 16))
231	connector.connector_scheduler.bind_coordinator(coordinator)
232	composite = object.__new__(MultiConnector)
233	composite._connectors = (connector,)
234	scheduler.connector = composite
235	return coordinator
236
237
238	def test_pending_hisparse_spill_keeps_scheduler_alive():
239	"""A final host spill must complete after the last request finishes."""
240	scheduler = create_scheduler()
241	coordinator = _bind_hisparse_connector(scheduler)
242	pending = coordinator.pending_spills
243	pending[0] = Mock()
244	assert scheduler.has_requests()
245	pending.clear()
246	assert not scheduler.has_requests()
247
248
249	@pytest.mark.parametrize(
250	"enable_prefix_caching, prompt_logprobs",
251	[
252	(False, None),
253	(True, 5),
254	],
255	)
256	def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int \| None):
257	"""Test scheduling.
258	Two cases: default APC/no prompt logprobs; APC=True + prompt logprobs
259	"""
260	scheduler = create_scheduler(enable_prefix_caching=enable_prefix_caching)
261	requests = create_requests(num_requests=10, prompt_logprobs=prompt_logprobs)
262	for request in requests:
263	scheduler.add_request(request)
264
265	# Test initial scheduling
266	output = scheduler.schedule()
267	assert len(output.scheduled_new_reqs) == len(requests)
268	assert output.scheduled_cached_reqs.num_reqs == 0
269	assert len(output.finished_req_ids) == 0
270	# Verify all requests are scheduled.
271	for req_id, num_tokens in output.num_scheduled_tokens.items():
272	assert num_tokens == len(requests[int(req_id)].prompt_token_ids)
273
274	# Verify requests moved from waiting to running
275	assert len(scheduler.waiting) == 0
276	assert len(scheduler.running) == len(requests)
277	for i, request in enumerate(requests):
278	assert scheduler.running[i] == request
279
280
281	def test_scheduler_stats_route_to_existing_output_client():
282	scheduler = create_scheduler()
283	request = create_requests(num_requests=1)[0]
284	request.client_index = 1
285	scheduler.add_request(request)
286
287	scheduler_output = scheduler.schedule()
288	model_output = ModelRunnerOutput(
289	req_ids=[request.request_id],
290	req_id_to_index={request.request_id: 0},
291	sampled_token_ids=[[1000]],
292	logprobs=None,
293	prompt_logprobs_dict={},
294	pooler_output=[],
295	)
296
297	engine_core_outputs = scheduler.update_from_output(scheduler_output, model_output)
298
299	assert 0 not in engine_core_outputs
300	assert engine_core_outputs[1].scheduler_stats is not None
301	assert len(engine_core_outputs[1].outputs) == 1
302
303
304	def test_schedule_multimodal_requests():
305	scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
306	mm_positions = [[PlaceholderRange(offset=i, length=100)] for i in range(10)]
307	requests = create_requests(
308	num_requests=10,
309	num_tokens=200,
310	mm_positions=mm_positions,
311	)
312	for request in requests:
313	scheduler.add_request(request)
314
315	output = scheduler.schedule()
316	assert len(output.scheduled_new_reqs) == len(requests)
317	assert output.scheduled_cached_reqs.num_reqs == 0
318	assert len(output.finished_req_ids) == 0
319	for req_id, num_tokens in output.num_scheduled_tokens.items():
320	assert num_tokens == len(requests[int(req_id)].prompt_token_ids)
321	assert len(output.scheduled_encoder_inputs) == 10
322	for req_id, encoder_input in output.scheduled_encoder_inputs.items():
323	assert len(encoder_input) == 1
324
325
326	def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
327	"""Async scheduling + PP: allow multi-step in-flight scheduling per request."""
328	scheduler = create_scheduler(async_scheduling=True, pipeline_parallel_size=2)
329	(req,) = create_requests(num_requests=1, num_tokens=8)
330	scheduler.add_request(req)
331
332	_ = scheduler.schedule()
333	assert req.num_output_placeholders > 0
334
335	# before any update_from_output, we still expect the request can be
336	# scheduled again (multi-step in-flight).
337	output = scheduler.schedule()
338	assert req.request_id in output.num_scheduled_tokens
339
340
341	def test_cached_request_data_resumed_all_token_ids_mrv1_only():
342	"""all_token_ids carries a resumed request's token ids to the connector
343	for the V1 model runner, but is skipped entirely for the V2 model runner.
344	"""
345	from vllm.v1.core.kv_cache_manager import KVCacheBlocks
346
347	scheduler = create_scheduler(use_v2_model_runner=False)
348	(req,) = create_requests(num_requests=1, num_tokens=8)
349	req.append_output_token_ids([101, 102, 103])
350
351	# A resumed request was not scheduled in the previous step.
352	assert req.request_id not in scheduler.prev_step_scheduled_req_ids
353
354	empty_blocks = KVCacheBlocks(blocks=((),))
355
356	def make_cached():
357	return scheduler._make_cached_request_data(
358	running_reqs=[],
359	resumed_reqs=[req],
360	num_scheduled_tokens={req.request_id: 1},
361	spec_decode_tokens={},
362	req_to_new_blocks={req.request_id: empty_blocks},
363	)
364
365	# V1 model runner: the full token id list is propagated.
366	assert not scheduler.use_v2_model_runner
367	cached = make_cached()
368	assert req.request_id in cached.resumed_req_ids
369	assert cached.all_token_ids[req.request_id] == list(req.all_token_ids)
370
371	# V2 model runner: all_token_ids is skipped entirely.
372	scheduler.use_v2_model_runner = True
373	cached = make_cached()
374	assert req.request_id in cached.resumed_req_ids
375	assert cached.all_token_ids == {}
376
377
378	def test_schedule_partial_requests():
379	"""Test scheduling behavior with partial requests.
380
381	This test verifies that:
382	1. The scheduler can handle multiple partial requests in a single step when
383	constrained by encoder budget.
384	2. A request in RUNNING state may be unscheduled in subsequent steps if
385	there is insufficient encoder budget.
386	"""
387	scheduler = create_scheduler(
388	model="llava-hf/llava-1.5-7b-hf",
389	max_num_batched_tokens=1024,
390	)
391	mm_positions = [[PlaceholderRange(offset=100, length=600)] for _ in range(3)]
392	requests = create_requests(
393	num_requests=3,
394	num_tokens=800,
395	mm_positions=mm_positions,
396	)
397	for request in requests:
398	scheduler.add_request(request)
399
400	output = scheduler.schedule()
401	assert len(output.scheduled_new_reqs) == 3
402	assert output.scheduled_cached_reqs.num_reqs == 0
403	assert len(output.finished_req_ids) == 0
404
405	assert scheduler.max_num_encoder_input_tokens == 1024
406	# The first request is scheduled fully.
407	assert output.num_scheduled_tokens[requests[0].request_id] == 800
408	# The second request is scheduled partially.
409	# The <img> tokens are not scheduled because of the encoder budget.
410	assert output.num_scheduled_tokens[requests[1].request_id] == 100
411	# The third request is also scheduled partially.
412	# The <img> tokens are not scheduled because of the encoder budget.
413	assert output.num_scheduled_tokens[requests[2].request_id] == 100
414	req_to_index = {request.request_id: i for i, request in enumerate(requests)}
415	model_runner_output = ModelRunnerOutput(
416	req_ids=[request.request_id for request in requests],
417	req_id_to_index=req_to_index,
418	# Only the first request has a sampled token id because
419	# the rest requests are still being prefilled.
420	sampled_token_ids=[[0], [], []],
421	logprobs=None,
422	prompt_logprobs_dict={},
423	pooler_output=[],
424	)
425	scheduler.update_from_output(output, model_runner_output)
426
427	# Schedule the next step.
428	# Only the first and second requests are scheduled.
429	# The third request is in the RUNNING state but not scheduled in this step
430	# because of the encoder budget.
431	output = scheduler.schedule()
432	assert len(scheduler.running) == 3
433	assert len(output.scheduled_new_reqs) == 0
434	assert output.scheduled_cached_reqs.num_reqs == 2
435	assert len(output.finished_req_ids) == 0
436	assert output.num_scheduled_tokens[requests[0].request_id] == 1
437	assert output.num_scheduled_tokens[requests[1].request_id] == 700
438	assert requests[2].request_id not in output.num_scheduled_tokens
439
440
441	def test_encoder_only_prompt_longer_than_budget_is_chunked():
442	"""The engine switches chunked prefill off for an encoder-only instance,
443	but the token budget must still split its prompt across steps instead of
444	leaving the request waiting forever."""
445	scheduler = create_scheduler(
446	model="llava-hf/llava-1.5-7b-hf",
447	max_num_batched_tokens=1024,
448	max_model_len=4096,
449	mm_encoder_only=True,
450	)
451	# EngineCore turns this off after config validation whenever the instance
452	# holds no KV cache, which is every encoder-only instance.
453	scheduler.scheduler_config.enable_chunked_prefill = False
454	(request,) = create_requests(
455	num_requests=1,
456	num_tokens=2500,
457	mm_positions=[[PlaceholderRange(offset=100, length=600)]],
458	)
459	scheduler.add_request(request)
460
461	def advance(output):
462	scheduler.update_from_output(
463	output,
464	ModelRunnerOutput(
465	req_ids=[request.request_id],
466	req_id_to_index={request.request_id: 0},
467	sampled_token_ids=[[]],
468	logprobs=None,
469	prompt_logprobs_dict={},
470	pooler_output=[],
471	),
472	)
473
474	first = scheduler.schedule()
475	assert first.num_scheduled_tokens[request.request_id] == 1024
476	assert request.request_id in first.scheduled_encoder_inputs
477	advance(first)
478
479	second = scheduler.schedule()
480	assert second.num_scheduled_tokens[request.request_id] == 1024
481	advance(second)
482
483	third = scheduler.schedule()
484	assert third.num_scheduled_tokens[request.request_id] == 452
485
486
487	@pytest.mark.parametrize("has_running", [True, False])
488	def test_schedule_prefills_gating(has_running: bool):
489	"""DP prefill-balancing gate: when `throttle_prefills` is True, a new
490	WAITING (prefill) request is deferred ONLY if this rank has running work to
491	protect. With no running requests, the prefill is admitted regardless (so a
492	throttled step is never wasted as a dummy), and running/decode requests are
493	unaffected. Once the cadence allows prefills again, the request is admitted.
494	"""
495	scheduler = create_scheduler(max_num_seqs=16, max_num_batched_tokens=8192)
496
497	if has_running:
498	# Establish a running (decode) request via a prefill + output step.
499	(running_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["run0"])
500	scheduler.add_request(running_req)
501	output = scheduler.schedule()
502	assert len(output.scheduled_new_reqs) == 1
503	scheduler.update_from_output(
504	output,
505	ModelRunnerOutput(
506	req_ids=["run0"],
507	req_id_to_index={"run0": 0},
508	sampled_token_ids=[[0]],
509	logprobs=None,
510	prompt_logprobs_dict={},
511	pooler_output=[],
512	),
513	)
514	assert len(scheduler.running) == 1
515
516	# Add a new WAITING (prefill) request, with prefills gated off.
517	(new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"])
518	scheduler.add_request(new_req)
519	output = scheduler.schedule(throttle_prefills=True)
520
521	if has_running:
522	# There is running work to protect, so the new prefill is deferred...
523	assert "new0" not in output.num_scheduled_tokens
524	assert new_req.status == RequestStatus.WAITING
525	# ...while the running/decode request keeps being scheduled.
526	assert "run0" in output.num_scheduled_tokens
527	# When the cadence allows prefills again, the request is admitted.
528	output = scheduler.schedule()
529
530	# No running work to protect (or cadence now open): the prefill is admitted.
531	assert "new0" in output.num_scheduled_tokens
532	assert any(r.req_id == "new0" for r in output.scheduled_new_reqs)
533
534
535	def _setup_remote_kv_resume(num_prompt_tokens: int, matched_tokens: int):
536	"""Drive a remote-KV request `r2` to the resume point (async load complete)
537	while another request `r1` is already decoding, so the step is throttle-
538	eligible. Returns the scheduler. The connector matches `matched_tokens` of
539	`r2`'s prompt; the rest (if any) is local prefill.
540	"""
541	from tests.v1.kv_connector.unit.utils import create_model_runner_output
542
543	BLOCK_SIZE = 16
544	scheduler = create_scheduler(
545	enable_prefix_caching=True,
546	use_kv_connector=mock_kv(matched_tokens=matched_tokens, is_async=True),
547	block_size=BLOCK_SIZE,
548	)
549	# Distinct prompts so r2 gets no local prefix cache hit from r1, only the
550	# connector's external async load.
551	r1, r2 = create_requests(
552	num_requests=2,
553	num_tokens=num_prompt_tokens,
554	max_tokens=20,
555	block_size=BLOCK_SIZE,
556	req_ids=["r1", "r2"],
557	)
558
559	# r1: drive through its async KV load into the running (decode) state, so
560	# self.running is non-empty (which makes the next step throttle-eligible).
561	scheduler.add_request(r1)
562	_step_until_kv_transfer_finished(scheduler, ["r1"])
563	output = scheduler.schedule() # promote + schedule r1
564	assert "r1" in output.num_scheduled_tokens
565	scheduler.update_from_output(
566	output, create_model_runner_output([r1], token_id=1000)
567	)
568	assert scheduler.running # r1 now decoding
569
570	# r2: a second remote-KV request; complete its async load while r1 decodes.
571	scheduler.add_request(r2)
572	output = scheduler.schedule() # r1 decodes; r2 -> WAITING_FOR_REMOTE_KVS
573	assert r2.status == RequestStatus.WAITING_FOR_REMOTE_KVS
574	scheduler.update_from_output(
575	output, create_model_runner_output([r1], finished_recving={"r2"})
576	)
577	assert "r2" in scheduler.finished_recving_kv_req_ids
578	return scheduler
579
580
581	def test_throttle_prefills_excludes_fully_transferred_remote_kv():
582	"""A remote-KV resume whose whole prompt was transferred (no local prefill
583	left, e.g. the decode side of P/D disaggregation) must NOT be throttled by
584	the DP prefill cadence -- its single-token step has no prefill compute to
585	defer, so delaying it would be pointless.
586	"""
587	block_size = 16
588	num_prompt = block_size * 2
589	# Fully matched: the whole prompt is loaded remotely.
590	scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt)
591
592	output = scheduler.schedule(throttle_prefills=True)
593	assert "r2" in output.num_scheduled_tokens
594	assert "r1" in output.num_scheduled_tokens
595
596
597	def test_throttle_prefills_defers_remote_kv_resume_with_local_prefill():
598	"""A remote-KV resume with local prefill still to compute (the connector
599	only matched part of the prompt) IS throttled by the DP prefill cadence,
600	like any other request doing local prefill compute this step.
601	"""
602	block_size = 16
603	num_prompt = block_size * 4
604	# Half matched: the remaining half is local prefill compute.
605	scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt // 2)
606
607	output = scheduler.schedule(throttle_prefills=True)
608	assert "r2" not in output.num_scheduled_tokens # deferred (has local prefill)
609	assert "r1" in output.num_scheduled_tokens
610
611
612	def test_throttle_defers_inflight_prefill_chunk():
613	"""DP prefill balancing throttles ALL prefill compute on a throttled step,
614	not just new admissions: an in-progress (chunked) prefill already in the
615	running queue is also deferred, so the step runs decode-only, while a
616	separate decode keeps being scheduled."""
617	scheduler = create_scheduler(
618	max_num_seqs=16, max_num_batched_tokens=50, enable_chunked_prefill=True
619	)
620
621	# A short request that finishes prefill in one step -> a running decode.
622	(decode_req,) = create_requests(num_requests=1, num_tokens=4, req_ids=["dec0"])
623	scheduler.add_request(decode_req)
624	output = scheduler.schedule()
625	scheduler.update_from_output(
626	output,
627	ModelRunnerOutput(
628	req_ids=["dec0"],
629	req_id_to_index={"dec0": 0},
630	sampled_token_ids=[[0]],
631	logprobs=None,
632	prompt_logprobs_dict={},
633	pooler_output=[],
634	),
635	)
636	assert decode_req in scheduler.running and not decode_req.is_prefill_chunk
637
638	# A long request (80 tokens, budget 50) -> prefilled in chunks.
639	(chunk_req,) = create_requests(num_requests=1, num_tokens=80, req_ids=["chk0"])
640	scheduler.add_request(chunk_req)
641	output = scheduler.schedule() # first chunk of chk0 + decode of dec0
642	assert output.num_scheduled_tokens["chk0"] > 0
643	scheduler.update_from_output(
644	output,
645	ModelRunnerOutput(
646	req_ids=["dec0", "chk0"],
647	req_id_to_index={"dec0": 0, "chk0": 1},
648	sampled_token_ids=[[0], []], # no token sampled for partial prefill
649	logprobs=None,
650	prompt_logprobs_dict={},
651	pooler_output=[],
652	),
653	)
654	assert chunk_req.is_prefill_chunk # still mid-prefill, in running
655
656	# Throttled step: the in-flight prefill chunk is deferred, the decode runs.
657	output = scheduler.schedule(throttle_prefills=True)
658	assert "chk0" not in output.num_scheduled_tokens
659	assert "dec0" in output.num_scheduled_tokens
660
661	# When the cadence opens again, the prefill chunk resumes.
662	output = scheduler.schedule()
663	assert "chk0" in output.num_scheduled_tokens
664
665
666	def test_throttle_capacity_bound_guard_admits():
667	"""Saturation guard: if a cadence-aligned release step cannot drain the
668	waiting prefill queue (it ran out of token budget), the throttle backs off on
669	the next step so the backlog cannot grow into a TTFT avalanche -- prefills are
670	admitted even though throttle_prefills is set."""
671	scheduler = create_scheduler(
672	max_num_seqs=16, max_num_batched_tokens=200, enable_chunked_prefill=True
673	)
674	a, b = create_requests(num_requests=2, num_tokens=200, req_ids=["a", "b"])
675	scheduler.add_request(a)
676	scheduler.add_request(b)
677
678	# Release step (throttle off): `a` fills the 200-token budget; `b` cannot be
679	# reached, so the waiting queue is not drained -> capacity-bound.
680	output = scheduler.schedule()
681	assert "a" in output.num_scheduled_tokens
682	assert "b" not in output.num_scheduled_tokens
683	assert scheduler.prefill_capacity_bound
684
685	# Throttle. Because the previous release was capacity-bound, the guard backs
686	# off and `b` is admitted rather than stalling the backlog.
687	output = scheduler.schedule(throttle_prefills=True)
688	assert "b" in output.num_scheduled_tokens
689
690
691	def test_no_mm_input_chunking():
692	# Disable multimodal input chunking.
693	scheduler = create_scheduler(
694	model="llava-hf/llava-1.5-7b-hf",
695	max_num_batched_tokens=1024,
696	disable_chunked_mm_input=True,
697	max_model_len=2048,
698	)
699	mm_positions = [[PlaceholderRange(offset=400, length=800)]]
700	requests = create_requests(
701	num_requests=1, num_tokens=1200, mm_positions=mm_positions
702	)
703	for request in requests:
704	scheduler.add_request(request)
705
706	output = scheduler.schedule()
707	assert len(output.scheduled_new_reqs) == 1
708	assert output.scheduled_cached_reqs.num_reqs == 0
709	assert len(output.finished_req_ids) == 0
710	# We want to only see the 400 text tokens at the start scheduled
711	assert output.num_scheduled_tokens[requests[0].request_id] == 400
712
713	req_to_index = {request.request_id: i for i, request in enumerate(requests)}
714	model_runner_output = ModelRunnerOutput(
715	req_ids=[request.request_id for request in requests],
716	req_id_to_index=req_to_index,
717	sampled_token_ids=[[] for _ in range(len(requests))],
718	logprobs=None,
719	prompt_logprobs_dict={},
720	pooler_output=[],
721	)
722	scheduler.update_from_output(output, model_runner_output)
723
724	output = scheduler.schedule()
725	assert len(scheduler.running) == 1
726	assert len(output.scheduled_new_reqs) == 0
727	assert output.scheduled_cached_reqs.num_reqs == 1
728	assert len(output.finished_req_ids) == 0
729	assert output.num_scheduled_tokens[requests[0].request_id] == 800
730
731	# Test that we fail if we disable chunked mm input and use too small
732	# of a max_num_batched_tokens for the mm input.
733	with pytest.raises(ValueError):
734	_ = create_scheduler(
735	model="llava-hf/llava-1.5-7b-hf",
736	max_num_batched_tokens=100,
737	disable_chunked_mm_input=True,
738	)
739
740
741	@pytest.mark.parametrize("enable_prefix_caching", [True, False])
742	def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
743	"""Test scheduling behavior with concurrent partial requests.
744
745	This test verifies that: there are multiple long prefill requests in the
746	RUNNING state, and we can schedule them together.
747
748	"""
749	scheduler = create_scheduler(
750	model="facebook/opt-125m",
751	max_num_batched_tokens=1024,
752	long_prefill_token_threshold=400,
753	enable_prefix_caching=enable_prefix_caching,
754	)
755	requests = create_requests(
756	num_requests=3,
757	num_tokens=800,
758	)
759	for request in requests:
760	scheduler.add_request(request)
761
762	output = scheduler.schedule()
763	assert len(output.scheduled_new_reqs) == 3
764	assert output.scheduled_cached_reqs.num_reqs == 0
765	assert len(output.finished_req_ids) == 0
766
767	# The first request is scheduled partially - 400.
768	assert output.num_scheduled_tokens[requests[0].request_id] == 400
769	# The second request is scheduled partially - 400.
770	assert output.num_scheduled_tokens[requests[1].request_id] == 400
771	# The third request is also scheduled partially - 1024 - 400 - 400 = 224.
772	assert output.num_scheduled_tokens[requests[2].request_id] == 224
773	req_to_index = {request.request_id: i for i, request in enumerate(requests)}
774	model_runner_output = ModelRunnerOutput(
775	req_ids=[request.request_id for request in requests],
776	req_id_to_index=req_to_index,
777	sampled_token_ids=[[] for _ in range(len(requests))],
778	logprobs=None,
779	prompt_logprobs_dict={},
780	pooler_output=[],
781	)
782	scheduler.update_from_output(output, model_runner_output)
783
784	# Schedule the next step. All three requests are running.
785	# Processed the remaining prefills of the first and second requests.
786	output1 = scheduler.schedule()
787	assert len(scheduler.running) == 3
788	assert len(output1.scheduled_new_reqs) == 0
789	assert output1.scheduled_cached_reqs.num_reqs == 3
790	assert len(output1.finished_req_ids) == 0
791	assert output1.num_scheduled_tokens[requests[0].request_id] == 400
792	assert output1.num_scheduled_tokens[requests[1].request_id] == 400
793	assert output1.num_scheduled_tokens[requests[2].request_id] == 224
794
795	# Schedule the third step. All three requests are running.
796	# First and second requests are in the decode stage.
797	# All the remaining tokens in the third request are processed.
798	model_runner_output = ModelRunnerOutput(
799	req_ids=[request.request_id for request in requests],
800	req_id_to_index=req_to_index,
801	sampled_token_ids=[[0], [0]] + [[] for _ in range(len(requests) - 2)],
802	logprobs=None,
803	prompt_logprobs_dict={},
804	pooler_output=[],
805	)
806	scheduler.update_from_output(output1, model_runner_output)
807	output2 = scheduler.schedule()
808	assert len(scheduler.running) == 3
809	assert len(output2.scheduled_new_reqs) == 0
810	assert output2.scheduled_cached_reqs.num_reqs == 3
811	assert len(output2.finished_req_ids) == 0
812	assert output2.num_scheduled_tokens[requests[0].request_id] == 1
813	assert output2.num_scheduled_tokens[requests[1].request_id] == 1
814	assert output2.num_scheduled_tokens[requests[2].request_id] == 800 - 224 - 224
815
816
817	def test_long_prefill_threshold_ignored_when_alone():
818	"""A lone long prefill is not capped by the threshold: it has no other
819	request to starve."""
820	scheduler = create_scheduler(
821	max_num_batched_tokens=1024,
822	long_prefill_token_threshold=400,
823	)
824	request = create_requests(num_requests=1, num_tokens=2000)[0]
825	scheduler.add_request(request)
826
827	output = scheduler.schedule()
828	assert output.num_scheduled_tokens[request.request_id] == 1024
829
830
831	def test_long_prefill_threshold_applies_with_other_requests():
832	"""The threshold caps the prefill as soon as another request is queued."""
833	scheduler = create_scheduler(
834	max_num_batched_tokens=1024,
835	long_prefill_token_threshold=400,
836	)
837	long_req = create_requests(num_requests=1, num_tokens=2000)[0]
838	short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0]
839	for request in [long_req, short_req]:
840	scheduler.add_request(request)
841
842	output = scheduler.schedule()
843	assert output.num_scheduled_tokens[long_req.request_id] == 400
844	assert output.num_scheduled_tokens[short_req.request_id] == 10
845
846
847	def test_long_prefill_threshold_floored_by_fair_share():
848	"""With the adaptive flag, the effective threshold never falls below the
849	fair share of the token budget: max_num_batched_tokens / num queued +
850	running requests."""
851	scheduler = create_scheduler(
852	max_num_batched_tokens=1024,
853	long_prefill_token_threshold=100,
854	long_prefill_token_threshold_adaptive=True,
855	)
856	long_req = create_requests(num_requests=1, num_tokens=2000)[0]
857	short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0]
858	for request in [long_req, short_req]:
859	scheduler.add_request(request)
860
861	output = scheduler.schedule()
862	# 100 is below the fair share (1024 // 2 = 512), so the floor binds.
863	assert output.num_scheduled_tokens[long_req.request_id] == 512
864	assert output.num_scheduled_tokens[short_req.request_id] == 10
865
866
867	def test_update_from_output_routes_sampling_masks_by_request():
868	"""Each request receives the sampler row at its own batch index."""
869	scheduler = create_scheduler()
870	scheduler.return_sampling_mask = True
871	requests = create_requests(num_requests=3, max_tokens=10)
872	for req in requests:
873	req.num_computed_tokens = req.num_tokens
874	scheduler.requests[req.request_id] = req
875	scheduler.running.append(req)
876	req.status = RequestStatus.RUNNING
877
878	scheduler_output = SchedulerOutput(
879	scheduled_new_reqs=[],
880	scheduled_cached_reqs=CachedRequestData.make_empty(),
881	num_scheduled_tokens={req.request_id: 1 for req in requests},
882	total_num_scheduled_tokens=3,
883	scheduled_encoder_inputs={},
884	scheduled_spec_decode_tokens={},
885	num_common_prefix_blocks=[],
886	finished_req_ids=set(),
887	free_encoder_mm_hashes=[],
888	)
889	model_output = ModelRunnerOutput(
890	req_ids=[req.request_id for req in requests],
891	req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
892	sampled_token_ids=[[1], [3], [4]],
893	logprobs=None,
894	prompt_logprobs_dict={},
895	pooler_output=[],
896	sampling_masks=SamplingMaskLists(
897	np.array([1, 2, 3, 4, 5, 6], dtype=np.int32), np.array([0, 2, 3, 6])
898	),
899	)
900
901	outputs = scheduler.update_from_output(scheduler_output, model_output)[0].outputs
902
903	assert [out.request_id for out in outputs] == [req.request_id for req in requests]
904	assert [out.new_sampling_mask.token_ids.tolist() for out in outputs] == [
905	[1, 2],
906	[3],
907	[4, 5, 6],
908	]
909	assert all(out.new_sampling_mask.offsets is None for out in outputs)
910
911
912	def test_update_from_output_routes_multi_position_sampling_masks():
913	scheduler = create_scheduler()
914	scheduler.return_sampling_mask = True
915	requests = create_requests(num_requests=2, max_tokens=10)
916	for req in requests:
917	req.num_computed_tokens = req.num_tokens
918	scheduler.requests[req.request_id] = req
919	scheduler.running.append(req)
920	req.status = RequestStatus.RUNNING
921
922	scheduler_output = SchedulerOutput(
923	scheduled_new_reqs=[],
924	scheduled_cached_reqs=CachedRequestData.make_empty(),
925	num_scheduled_tokens={req.request_id: 1 for req in requests},
926	total_num_scheduled_tokens=2,
927	scheduled_encoder_inputs={},
928	scheduled_spec_decode_tokens={},
929	num_common_prefix_blocks=[],
930	finished_req_ids=set(),
931	free_encoder_mm_hashes=[],
932	)
933	model_output = ModelRunnerOutput(
934	req_ids=[req.request_id for req in requests],
935	req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
936	sampled_token_ids=[[1, 2], [3, 4, 5]],
937	logprobs=None,
938	prompt_logprobs_dict={},
939	pooler_output=[],
940	sampling_masks=SamplingMaskLists(
941	token_ids=np.array([1, 6, 2, 3, 7, 8, 4, 5, 9], dtype=np.int32),
942	offsets=np.array([0, 2, 3, 6, 7, 9]),
943	cu_num_generated_tokens=[0, 2, 5],
944	),
945	)
946
947	outputs = scheduler.update_from_output(scheduler_output, model_output)[0].outputs
948
949	assert [out.new_sampling_mask.to_nested_list() for out in outputs] == [
950	[[1, 6], [2]],
951	[[3, 7, 8], [4], [5, 9]],
952	]
953
954
955	def test_stop_via_update_from_output():
956	"""Test stopping behavior through update_from_output."""
957	scheduler = create_scheduler(num_speculative_tokens=1)
958
959	# Test case 1: Stop on EOS token
960	requests = create_requests(num_requests=2, max_tokens=10)
961	for req in requests:
962	req.num_computed_tokens = req.num_tokens
963	scheduler.requests[req.request_id] = req
964	scheduler.running.append(req)
965	req.status = RequestStatus.RUNNING
966
967	scheduler_output = SchedulerOutput(
968	scheduled_new_reqs=[],
969	scheduled_cached_reqs=CachedRequestData.make_empty(),
970	num_scheduled_tokens={requests[0].request_id: 1, requests[1].request_id: 2},
971	total_num_scheduled_tokens=3,
972	scheduled_encoder_inputs={},
973	scheduled_spec_decode_tokens={
974	requests[0].request_id: [],
975	requests[1].request_id: [10],
976	},
977	num_common_prefix_blocks=[],
978	finished_req_ids=set(),
979	free_encoder_mm_hashes=[],
980	)
981
982	model_output = ModelRunnerOutput(
983	req_ids=[req.request_id for req in requests],
984	req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
985	sampled_token_ids=[
986	[EOS_TOKEN_ID],
987	[10, 11],
988	], # First request hits EOS, second continues
989	logprobs=None,
990	prompt_logprobs_dict={},
991	pooler_output=[],
992	)
993
994	scheduler.update_from_output(scheduler_output, model_output)
995
996	# Verify first request stopped, second continues
997	assert len(scheduler.running) == 1
998	assert scheduler.running[0].request_id == requests[1].request_id
999	assert requests[0].status == RequestStatus.FINISHED_STOPPED
1000	assert requests[0].request_id in scheduler.finished_req_ids
1001	assert list(requests[0].output_token_ids) == [EOS_TOKEN_ID]
1002	assert list(requests[1].output_token_ids) == [10, 11]
1003
1004	# Test case 2: Stop on custom stop token
1005	scheduler = create_scheduler(num_speculative_tokens=2)
1006	requests = create_requests(num_requests=2, max_tokens=10, stop_token_ids=[42, 43])
1007	for req in requests:
1008	req.num_computed_tokens = req.num_tokens
1009	scheduler.requests[req.request_id] = req
1010	scheduler.running.append(req)
1011	req.status = RequestStatus.RUNNING
1012
1013	scheduler_output = SchedulerOutput(
1014	scheduled_new_reqs=[],
1015	scheduled_cached_reqs=CachedRequestData.make_empty(),
1016	num_scheduled_tokens={requests[0].request_id: 3, requests[1].request_id: 2},
1017	total_num_scheduled_tokens=5,
1018	scheduled_encoder_inputs={},
1019	scheduled_spec_decode_tokens={
1020	requests[0].request_id: [10, 42],
1021	requests[1].request_id: [13],
1022	},
1023	num_common_prefix_blocks=[],
1024	finished_req_ids=set(),
1025	free_encoder_mm_hashes=[],
1026	)
1027
1028	model_output = ModelRunnerOutput(
1029	req_ids=[req.request_id for req in requests],
1030	req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
1031	sampled_token_ids=[[10, 42, 12], [13, 14]], # First request hits stop token
1032	logprobs=None,
1033	prompt_logprobs_dict={},
1034	pooler_output=[],
1035	)
1036
1037	scheduler.update_from_output(scheduler_output, model_output)
1038
1039	# Verify first request stopped on custom token
1040	assert len(scheduler.running) == 1
1041	assert scheduler.running[0].request_id == requests[1].request_id
1042	assert requests[0].status == RequestStatus.FINISHED_STOPPED
1043	assert requests[0].stop_reason == 42
1044	assert requests[0].request_id in scheduler.finished_req_ids
1045	assert list(requests[0].output_token_ids) == [10, 42]
1046	assert list(requests[1].output_token_ids) == [13, 14]
1047
1048	# Test case 3: Stop on max tokens
1049	scheduler = create_scheduler(num_speculative_tokens=2)
1050	requests = create_requests(num_requests=2, max_tokens=2)
1051	for req in requests:
1052	req.num_computed_tokens = req.num_tokens
1053	scheduler.requests[req.request_id] = req
1054	scheduler.running.append(req)
1055	req.status = RequestStatus.RUNNING
1056
1057	scheduler_output = SchedulerOutput(
1058	scheduled_new_reqs=[],
1059	scheduled_cached_reqs=CachedRequestData.make_empty(),
1060	num_scheduled_tokens={requests[0].request_id: 3, requests[1].request_id: 1},
1061	total_num_scheduled_tokens=4,
1062	scheduled_encoder_inputs={},
1063	scheduled_spec_decode_tokens={
1064	requests[0].request_id: [10, 11],
1065	requests[1].request_id: [],
1066	},
1067	num_common_prefix_blocks=[],
1068	finished_req_ids=set(),
1069	free_encoder_mm_hashes=[],
1070	)
1071
1072	model_output = ModelRunnerOutput(
1073	req_ids=[req.request_id for req in requests],
1074	req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
1075	sampled_token_ids=[[10, 11, 12], [13]], # First request exceeds max_tokens
1076	logprobs=None,
1077	prompt_logprobs_dict={},
1078	pooler_output=[],
1079	)
1080
1081	scheduler.update_from_output(scheduler_output, model_output)
1082
1083	# Verify first request stopped due to length
1084	assert len(scheduler.running) == 1
1085	assert scheduler.running[0].request_id == requests[1].request_id
1086	assert requests[0].status == RequestStatus.FINISHED_LENGTH_CAPPED
1087	assert requests[0].request_id in scheduler.finished_req_ids
1088	assert list(requests[0].output_token_ids) == [10, 11] # Truncated to max_tokens
1089	assert list(requests[1].output_token_ids) == [13]
1090
1091	# Test case 4: Ignore EOS flag
1092	scheduler = create_scheduler(num_speculative_tokens=2)
1093	requests = create_requests(num_requests=1, max_tokens=10, ignore_eos=True)
1094	requests[0].num_computed_tokens = requests[0].num_tokens
1095	scheduler.requests[requests[0].request_id] = requests[0]
1096	scheduler.running.append(requests[0])
1097
1098	scheduler_output = SchedulerOutput(
1099	scheduled_new_reqs=[],
1100	scheduled_cached_reqs=CachedRequestData.make_empty(),
1101	num_scheduled_tokens={requests[0].request_id: 3},
1102	total_num_scheduled_tokens=3,
1103	scheduled_encoder_inputs={},
1104	scheduled_spec_decode_tokens={requests[0].request_id: [EOS_TOKEN_ID, 10]},
1105	num_common_prefix_blocks=[],
1106	finished_req_ids=set(),
1107	free_encoder_mm_hashes=[],
1108	)
1109
1110	model_output = ModelRunnerOutput(
1111	req_ids=[requests[0].request_id],
1112	req_id_to_index={requests[0].request_id: 0},
1113	sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
1114	logprobs=None,
1115	prompt_logprobs_dict={},
1116	pooler_output=[],
1117	)
1118
1119	scheduler.update_from_output(scheduler_output, model_output)
1120
1121	# Verify request continues past EOS
1122	assert len(scheduler.running) == 1
1123	assert not requests[0].is_finished()
1124	assert list(requests[0].output_token_ids) == [EOS_TOKEN_ID, 10, 11]
1125
1126
1127	def test_check_stop_min_tokens():
1128	"""Test that sampled stop tokens still finish under min_tokens."""
1129	from vllm.v1.core.sched.utils import check_stop
1130
1131	# Test case 1: regular tokens below min_tokens should keep running.
1132	sampling_params = SamplingParams(
1133	ignore_eos=False,
1134	max_tokens=20,
1135	min_tokens=5,
1136	)
1137	sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
1138	request = Request(
1139	request_id="0",
1140	prompt_token_ids=[0, 1, 2],
1141	sampling_params=sampling_params,
1142	pooling_params=None,
1143	)
1144	request.append_output_token_ids([10, 11, 12])
1145
1146	result = check_stop(request, max_model_len=100)
1147	assert result is False, "Should not stop when num_output_tokens<min_tokens"
1148
1149	# Test case 2: if EOS is sampled, it should stop even below min_tokens.
1150	request_eos = Request(
1151	request_id="eos",
1152	prompt_token_ids=[0, 1, 2],
1153	sampling_params=sampling_params,
1154	pooling_params=None,
1155	)
1156	request_eos.append_output_token_ids([10, 11, EOS_TOKEN_ID])
1157	result = check_stop(request_eos, max_model_len=100)
1158	assert result is True, "Sampled EOS should stop even below min_tokens"
1159	assert request_eos.status == RequestStatus.FINISHED_STOPPED
1160
1161	# Test case 3: ignore_eos should still prevent EOS-based stopping.
1162	sampling_params_ignore_eos = SamplingParams(
1163	ignore_eos=True,
1164	max_tokens=20,
1165	min_tokens=5,
1166	)
1167	sampling_params_ignore_eos.update_from_generation_config({}, EOS_TOKEN_ID)
1168	request_ignore_eos = Request(
1169	request_id="ignore-eos",
1170	prompt_token_ids=[0, 1, 2],
1171	sampling_params=sampling_params_ignore_eos,
1172	pooling_params=None,
1173	)
1174	request_ignore_eos.append_output_token_ids([EOS_TOKEN_ID])
1175	result = check_stop(request_ignore_eos, max_model_len=100)
1176	assert result is False, "ignore_eos should still suppress EOS stopping"
1177
1178	# Test case 4: min_tokens = 0 should follow normal stopping logic.
1179	sampling_params_no_min = SamplingParams(
1180	ignore_eos=False,
1181	max_tokens=20,
1182	min_tokens=0,
1183	)
1184	sampling_params_no_min.update_from_generation_config({}, EOS_TOKEN_ID)
1185	request_no_min = Request(
1186	request_id="1",
1187	prompt_token_ids=[0, 1, 2],
1188	sampling_params=sampling_params_no_min,
1189	pooling_params=None,
1190	)
1191	request_no_min.append_output_token_ids([10, EOS_TOKEN_ID])
1192
1193	result = check_stop(request_no_min, max_model_len=100)
1194	assert result is True, "Should stop on EOS when min_tokens=0"
1195	assert request_no_min.status == RequestStatus.FINISHED_STOPPED
1196
1197	# Test case 5: if a stop token is sampled, it should stop immediately.
1198	sampling_params_stop = SamplingParams(
1199	ignore_eos=False,
1200	max_tokens=20,
1201	min_tokens=5,
1202	stop_token_ids=[42],
1203	)
1204	sampling_params_stop.update_from_generation_config({}, EOS_TOKEN_ID)
1205	request_stop = Request(
1206	request_id="2",
1207	prompt_token_ids=[0, 1, 2],
1208	sampling_params=sampling_params_stop,
1209	pooling_params=None,
1210	)
1211	request_stop.append_output_token_ids([10, 11, 42])
1212	result = check_stop(request_stop, max_model_len=100)
1213	assert result is True, "Sampled stop tokens should stop even below min_tokens"
1214	assert request_stop.status == RequestStatus.FINISHED_STOPPED
1215	assert request_stop.stop_reason == 42
1216
1217
1218	def test_check_stop_length_cap_below_min_tokens():
1219	"""Length caps must finish a request even below min_tokens.
1220
1221	Guards against requests hanging forever when min_tokens cannot be
1222	satisfied: min_tokens only defers EOS/stop tokens, so a request that
1223	reaches max_model_len (or a max_tokens clamped below min_tokens) must
1224	still finish as FINISHED_LENGTH_CAPPED instead of spinning in the
1225	scheduler with zero new tokens to schedule.
1226	"""
1227	from vllm.v1.core.sched.utils import check_stop
1228
1229	# Reaching max_model_len below min_tokens must stop the request.
1230	sampling_params = SamplingParams(
1231	ignore_eos=True,
1232	max_tokens=500,
1233	min_tokens=500,
1234	)
1235	sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
1236	request = Request(
1237	request_id="0",
1238	prompt_token_ids=[0, 1, 2],
1239	sampling_params=sampling_params,
1240	pooling_params=None,
1241	)
1242	# 7 output tokens: num_tokens == 10 == max_model_len, min_tokens unmet.
1243	request.append_output_token_ids([10, 11, 12, 13, 14, 15, 16])
1244
1245	result = check_stop(request, max_model_len=10)
1246	assert result is True, "Should stop at max_model_len even below min_tokens"
1247	assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
1248
1249	# Reaching max_tokens below min_tokens must stop the request. This can
1250	# happen when max_tokens is clamped to the context window after
1251	# SamplingParams validation (e.g. max_tokens=None on a long prompt).
1252	sampling_params_capped = SamplingParams(
1253	ignore_eos=True,
1254	max_tokens=500,
1255	min_tokens=500,
1256	)
1257	sampling_params_capped.max_tokens = 3
1258	sampling_params_capped.update_from_generation_config({}, EOS_TOKEN_ID)
1259	request_capped = Request(
1260	request_id="1",
1261	prompt_token_ids=[0, 1, 2],
1262	sampling_params=sampling_params_capped,
1263	pooling_params=None,
1264	)
1265	request_capped.append_output_token_ids([10, 11, 12])
1266
1267	result = check_stop(request_capped, max_model_len=100)
1268	assert result is True, "Should stop at max_tokens even below min_tokens"
1269	assert request_capped.status == RequestStatus.FINISHED_LENGTH_CAPPED
1270
1271
1272	@pytest.mark.parametrize(
1273	"enable_prefix_caching, prompt_logprobs",
1274	[
1275	(False, None),
1276	(True, 5),
1277	],
1278	)
1279	def test_schedule_concurrent_batches(
1280	enable_prefix_caching: bool, prompt_logprobs: int \| None
1281	):
1282	scheduler = create_scheduler(
1283	max_num_batched_tokens=1024,
1284	max_num_seqs=2,
1285	enable_prefix_caching=enable_prefix_caching,
1286	)
1287	requests = create_requests(
1288	num_requests=2,
1289	num_tokens=512,
1290	prompt_logprobs=prompt_logprobs,
1291	)
1292
1293	# Schedule the first request.
1294	scheduler.add_request(requests[0])
1295	scheduler_output0 = scheduler.schedule()
1296	assert len(scheduler_output0.scheduled_new_reqs) == 1
1297	assert scheduler_output0.num_scheduled_tokens[requests[0].request_id] == 512
1298
1299	# The first request is still running, so only schedule the second request.
1300	scheduler.add_request(requests[1])
1301	scheduler_output1 = scheduler.schedule()
1302	assert len(scheduler_output1.scheduled_new_reqs) == 1
1303	assert scheduler_output1.num_scheduled_tokens[requests[1].request_id] == 512
1304
1305	# Model output of the first request.
1306	model_runner_output = ModelRunnerOutput(
1307	req_ids=[requests[0].request_id],
1308	req_id_to_index={requests[0].request_id: 0},
1309	sampled_token_ids=[[0]],
1310	logprobs=None,
1311	prompt_logprobs_dict={},
1312	pooler_output=[],
1313	)
1314	scheduler.update_from_output(scheduler_output0, model_runner_output)
1315
1316	# Schedule the next step.
1317	# The first request can be scheduled again while the second
1318	# request is still running.
1319	scheduler_output2 = scheduler.schedule()
1320	assert scheduler_output2.num_scheduled_tokens[requests[0].request_id] == 1
1321
1322	# Model output of the second request.
1323	model_runner_output = ModelRunnerOutput(
1324	req_ids=[requests[1].request_id],
1325	req_id_to_index={requests[1].request_id: 0},
1326	sampled_token_ids=[[0]],
1327	logprobs=None,
1328	prompt_logprobs_dict={},
1329	pooler_output=[],
1330	)
1331	scheduler.update_from_output(scheduler_output1, model_runner_output)
1332
1333
1334	@pytest.mark.parametrize("enable_chunked_prefill", [True, False])
1335	def test_schedule_order(enable_chunked_prefill: bool):
1336	scheduler = create_scheduler(
1337	max_num_batched_tokens=1024,
1338	max_num_seqs=3,
1339	enable_chunked_prefill=enable_chunked_prefill,
1340	)
1341
1342	# long requests
1343	requests = create_requests(num_requests=2, num_tokens=800, req_ids=["1", "2"])
1344	# short requests
1345	requests += create_requests(num_requests=2, num_tokens=10, req_ids=["3", "4"])
1346
1347	for request in requests:
1348	scheduler.add_request(request)
1349
1350	scheduler_output1 = scheduler.schedule()
1351
1352	if enable_chunked_prefill:
1353	# When enable chunked prefill, long requests will be chunked.
1354	assert len(scheduler_output1.scheduled_new_reqs) == 2
1355	else:
1356	# When disable chunked prefill, should not skip the long requests,
1357	# and scheduling subsequent short requests in advance,
1358	# even though there is still token budgets remaining.
1359	assert len(scheduler_output1.scheduled_new_reqs) == 1
1360
1361
1362	def test_preempt_during_execution():
1363	# NOTE(woosuk): The actual number of available blocks is 10 instead of 11
1364	# because block 0 is reserved as the null block.
1365	scheduler = create_scheduler(
1366	max_num_batched_tokens=100,
1367	block_size=16,
1368	num_blocks=11,
1369	enable_prefix_caching=False,
1370	)
1371	requests = create_requests(num_requests=2, num_tokens=80, block_size=16)
1372
1373	# Schedule the first request.
1374	scheduler.add_request(requests[0])
1375	scheduler_output0 = scheduler.schedule()
1376	assert len(scheduler_output0.num_scheduled_tokens) == 1
1377	assert len(scheduler_output0.scheduled_new_reqs[0].block_ids[0]) == 5
1378
1379	# Schedule the second request while the first request is still running.
1380	# This scenario can occur in certain cases, when max_concurrent_batches > 1
1381	# (e.g., when pipeline parallelism is used).
1382	scheduler.add_request(requests[1])
1383	scheduler_output1 = scheduler.schedule()
1384	assert len(scheduler_output1.num_scheduled_tokens) == 1
1385	assert len(scheduler_output1.scheduled_new_reqs[0].block_ids[0]) == 5
1386
1387	# Get the output of the first request.
1388	model_runner_output0 = ModelRunnerOutput(
1389	req_ids=[requests[0].request_id],
1390	req_id_to_index={requests[0].request_id: 0},
1391	sampled_token_ids=[[0]],
1392	logprobs=None,
1393	prompt_logprobs_dict={},
1394	pooler_output=[],
1395	)
1396	scheduler.update_from_output(scheduler_output0, model_runner_output0)
1397
1398	# Schedule the first request again. This will cause the preemption
1399	# of the second request because the KV cache is full.
1400	_ = scheduler.schedule()
1401	assert len(scheduler.running) == 1
1402	assert scheduler.running[0] == requests[0]
1403	assert requests[1].status == RequestStatus.PREEMPTED
1404
1405	model_runner_output1 = ModelRunnerOutput(
1406	req_ids=[requests[1].request_id],
1407	req_id_to_index={requests[1].request_id: 0},
1408	sampled_token_ids=[[42]],
1409	logprobs=None,
1410	prompt_logprobs_dict={},
1411	pooler_output=[],
1412	)
1413	scheduler.update_from_output(scheduler_output1, model_runner_output1)
1414
1415	# The second request (that is preempted) should be updated with the
1416	# sampled token id.
1417	assert len(requests[1].output_token_ids) == 1
1418	assert requests[1].output_token_ids[0] == 42
1419
1420
1421	def test_pending_hisparse_reclamation_defers_preemption(monkeypatch):
1422	"""Pending reclamation must stall rather than preempt its owning request."""
1423	scheduler = create_scheduler(enable_prefix_caching=False)
1424	request = create_requests(num_requests=1, num_tokens=16, block_size=16)[0]
1425	scheduler.add_request(request)
1426	scheduler_output = scheduler.schedule()
1427	scheduler.update_from_output(
1428	scheduler_output,
1429	ModelRunnerOutput(
1430	req_ids=[request.request_id],
1431	req_id_to_index={request.request_id: 0},
1432	sampled_token_ids=[[0]],
1433	logprobs=None,
1434	prompt_logprobs_dict={},
1435	pooler_output=[],
1436	),
1437	)
1438
1439	monkeypatch.setattr(
1440	scheduler.kv_cache_manager, "allocate_slots", Mock(return_value=None)
1441	)
1442	coordinator = _bind_hisparse_connector(scheduler)
1443	monkeypatch.setattr(
1444	scheduler.connector, "build_connector_meta", Mock(return_value=None)
1445	)
1446	monkeypatch.setattr(
1447	coordinator,
1448	"has_pending_reclamation",
1449	Mock(return_value=True),
1450	)
1451
1452	deferred_output = scheduler.schedule()
1453
1454	assert not deferred_output.num_scheduled_tokens
1455	assert scheduler.running == [request]
1456	assert request.status == RequestStatus.RUNNING
1457	assert request.num_preemptions == 0
1458
1459
1460	def test_prefix_cache_query_not_inflated_by_connector_defer():
1461	"""The GPU prefix-cache query is recorded at admission, so a request the
1462	connector defers several times is counted once, not once per retry."""
1463	num_defers_before_matching = 3
1464	scheduler = create_scheduler(
1465	enable_prefix_caching=True,
1466	use_kv_connector=mock_kv(
1467	matched_tokens=0,
1468	is_async=False,
1469	num_defers_before_matching=num_defers_before_matching,
1470	),
1471	)
1472	request = create_requests(num_requests=1, num_tokens=32, block_size=16)[0]
1473	scheduler.add_request(request)
1474
1475	# Each deferred step re-runs the lookup but records nothing.
1476	for _ in range(num_defers_before_matching):
1477	assert not scheduler.schedule().scheduled_new_reqs
1478
1479	output = scheduler.schedule()
1480	assert any(r.req_id == request.request_id for r in output.scheduled_new_reqs)
1481
1482	stats = scheduler.kv_cache_manager.prefix_cache_stats
1483	assert stats is not None
1484	assert stats.requests == 1
1485	assert stats.queries == request.num_tokens
1486
1487
1488	def test_preemption_re_records_prefix_cache_query():
1489	"""A preempted request re-enters the lookup on resume, so its recomputation
1490	is counted again into the preempted stats."""
1491	scheduler = create_scheduler(enable_prefix_caching=True)
1492	request = create_requests(num_requests=1)[0]
1493	scheduler.add_request(request)
1494
1495	scheduler_output = scheduler.schedule()
1496	stats = scheduler.kv_cache_manager.prefix_cache_stats
1497	assert stats is not None
1498	assert (stats.requests, stats.preempted_requests) == (1, 0)
1499
1500	scheduler.running.remove(request)
1501	scheduler._preempt_request(request, 0.0)
1502	assert request.status == RequestStatus.PREEMPTED
1503
1504	scheduler.update_from_output(
1505	scheduler_output,
1506	ModelRunnerOutput(
1507	req_ids=[request.request_id],
1508	req_id_to_index={request.request_id: 0},
1509	sampled_token_ids=[[1000]],
1510	logprobs=None,
1511	prompt_logprobs_dict={},
1512	pooler_output=[],
1513	),
1514	)
1515	assert request.num_stale_output_tokens == 0
1516	stats = scheduler.kv_cache_manager.prefix_cache_stats
1517	assert stats is not None
1518
1519	scheduler.schedule()
1520	assert request.status == RequestStatus.RUNNING
1521	assert stats.preempted_requests == 1
1522
1523
1524	def test_preemption_processes_stale_aux_output():
1525	scheduler = create_scheduler(enable_prefix_caching=True)
1526	request = create_requests(num_requests=1)[0]
1527	scheduler.add_request(request)
1528	scheduler_output = scheduler.schedule()
1529	scheduler.running.remove(request)
1530	scheduler.aux_output_connector = Mock()
1531	scheduler._preempt_request(request, 0.0)
1532	scheduler.aux_output_connector.request_finished.assert_called_once_with(request)
1533
1534	scheduler.update_from_output(
1535	scheduler_output,
1536	ModelRunnerOutput(
1537	req_ids=[request.request_id],
1538	req_id_to_index={request.request_id: 0},
1539	sampled_token_ids=[[1000]],
1540	logprobs=None,
1541	prompt_logprobs_dict={},
1542	pooler_output=[],
1543	),
1544	)
1545
1546	scheduler.aux_output_connector.take_output.assert_called_once_with(request, None)
1547
1548
1549	def test_prefix_cache_stats_not_recorded_when_caching_disabled():
1550	"""With prefix caching off there is no local lookup, so admitting a request
1551	records no phantom miss."""
1552	scheduler = create_scheduler(enable_prefix_caching=False)
1553	for request in create_requests(num_requests=2):
1554	scheduler.add_request(request)
1555
1556	scheduler.schedule()
1557
1558	stats = scheduler.kv_cache_manager.prefix_cache_stats
1559	assert stats is not None
1560	assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
1561
1562
1563	def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request():
1564	"""A real cache hit rejected once by allocate_slots and admitted on the next
1565	step is counted exactly once, hits included."""
1566	block_size = 16
1567	scheduler = create_scheduler(
1568	enable_prefix_caching=True,
1569	enable_chunked_prefill=False,
1570	block_size=block_size,
1571	)
1572
1573	# Seed the cache so the next request with the same prompt hits it.
1574	seed = create_requests(
1575	num_requests=1,
1576	num_tokens=block_size * 2,
1577	max_tokens=2,
1578	same_prompt=True,
1579	block_size=block_size,
1580	req_ids=["seed"],
1581	)[0]
1582	scheduler.add_request(seed)
1583	_step_until_done(
1584	scheduler,
1585	scheduler.schedule(),
1586	ModelRunnerOutput(
1587	req_ids=["seed"],
1588	req_id_to_index={"seed": 0},
1589	sampled_token_ids=[[1000]],
1590	logprobs=None,
1591	prompt_logprobs_dict={},
1592	pooler_output=[],
1593	),
1594	)
1595
1596	# The seeding step swapped in a fresh accumulator, so re-read it; the
1597	# retried request must be the only thing recorded from here on.
1598	stats = scheduler.kv_cache_manager.prefix_cache_stats
1599	assert stats is not None
1600	assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
1601
1602	retried = create_requests(
1603	num_requests=1,
1604	num_tokens=block_size * 3,
1605	max_tokens=1,
1606	same_prompt=True,
1607	block_size=block_size,
1608	req_ids=["retried"],
1609	)[0]
1610	scheduler.add_request(retried)
1611
1612	# Reject the first allocation attempt, then delegate to the real one.
1613	orig_allocate_slots = scheduler.kv_cache_manager.allocate_slots
1614	allocate_results: list = []
1615
1616	def spy_allocate_slots(args, *kwargs):
1617	result = None if not allocate_results else orig_allocate_slots(args, *kwargs)
1618	allocate_results.append(result)
1619	return result
1620
1621	scheduler.kv_cache_manager.allocate_slots = spy_allocate_slots
1622
1623	assert not scheduler.schedule().scheduled_new_reqs
1624	assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
1625
1626	assert "retried" in scheduler.schedule().num_scheduled_tokens
1627	assert allocate_results[0] is None and allocate_results[1] is not None
1628	assert (stats.requests, stats.queries, stats.hits) == (
1629	1,
1630	retried.num_tokens,
1631	block_size * 2,
1632	)
1633
1634
1635	def test_scheduler_reset_prefix_cache():
1636	scheduler = create_scheduler(enable_prefix_caching=True)
1637	scheduler.aux_output_connector = Mock()
1638	requests = create_requests(num_requests=10)
1639	for request in requests:
1640	scheduler.add_request(request)
1641
1642	# Initial scheduling, requests should be at the running state now
1643	_ = scheduler.schedule()
1644
1645	# Verify requests moved from waiting to running
1646	assert len(scheduler.waiting) == 0
1647	assert len(scheduler.running) == len(requests)
1648	for i, request in enumerate(requests):
1649	assert scheduler.running[i] == request
1650
1651	# Reset prefix cache should fail since there are still running requests
1652	# and they are taking KV cache
1653	assert not scheduler.reset_prefix_cache()
1654	scheduler.aux_output_connector.reset.assert_not_called()
1655
1656	# Pause completes pending model outputs before the caller resets the scheduler.
1657	for request in requests:
1658	request.num_in_flight_tokens = 0
1659
1660	assert scheduler.reset_prefix_cache(reset_running_requests=True)
1661	scheduler.aux_output_connector.reset.assert_called_once_with()
1662
1663	# Verify requests moved from running to waiting
1664	assert len(scheduler.waiting) == len(requests)
1665	assert len(scheduler.running) == 0
1666
1667	for i, request in enumerate(requests):
1668	assert scheduler.waiting[i] == request
1669
1670
1671	@pytest.mark.parametrize("reset_successful", [False, True])
1672	def test_aux_output_reset_follows_kv_reset_result(reset_successful: bool):
1673	scheduler = create_scheduler(enable_prefix_caching=True)
1674	scheduler.aux_output_connector = Mock()
1675	scheduler.kv_cache_manager.reset_prefix_cache = Mock(return_value=reset_successful)
1676
1677	assert scheduler.reset_prefix_cache() is reset_successful
1678	assert scheduler.aux_output_connector.reset.call_count == int(reset_successful)
1679
1680
1681	def test_kv_cache_release_after_keep_pause_preserves_requests():
1682	"""A completed keep pause permits release and recomputation of live requests."""
1683	scheduler = create_scheduler(max_num_seqs=1, enable_prefix_caching=True)
1684	running, waiting = create_requests(num_requests=2)
1685	for request in (running, waiting):
1686	scheduler.add_request(request)
1687	scheduler_output = scheduler.schedule()
1688	scheduler.update_from_output(
1689	scheduler_output,
1690	ModelRunnerOutput(
1691	req_ids=[running.request_id],
1692	req_id_to_index={running.request_id: 0},
1693	sampled_token_ids=[[1]],
1694	),
1695	)
1696	assert scheduler.kv_cache_manager.get_block_ids(running.request_id)[0]
1697
1698	core = object.__new__(EngineCore)
1699	core.scheduler = scheduler
1700	core.model_executor = Mock(is_sleeping=False)
1701	core.mm_receiver_cache = None
1702	core.batch_queue = None
1703
1704	assert core.pause_scheduler(mode="keep", clear_cache=False) is None
1705	assert scheduler.running == [running]
1706	assert list(scheduler.waiting) == [waiting]
1707	assert not scheduler.has_requests()
1708
1709	core.release_kv_cache_memory()
1710
1711	core.model_executor.discard.assert_called_once_with(("kv_cache",))
1712	assert scheduler.requests == {
1713	request.request_id: request for request in (running, waiting)
1714	}
1715	assert not scheduler.running
1716	assert list(scheduler.waiting) == [running, waiting]
1717	assert running.status == RequestStatus.PREEMPTED
1718	assert running.num_computed_tokens == 0
1719	assert list(running.output_token_ids) == [1]
1720	assert not scheduler.kv_cache_manager.get_block_ids(running.request_id)[0]
1721	assert scheduler.schedule().total_num_scheduled_tokens == 0
1722
1723	assert core.wake_up(tags=["kv_cache"])
1724	resumed = scheduler.schedule()
1725	assert resumed.num_scheduled_tokens == {running.request_id: running.num_tokens}
1726
1727
1728	@pytest.mark.parametrize(
1729	"sleeping_tags",
1730	[{"weights", "kv_cache"}, {"weights"}, {"kv_cache"}],
1731	ids=["full-sleep", "weights-only", "kv-only-or-repeated-release"],
1732	)
1733	def test_kv_cache_release_rejects_nonresident_memory(sleeping_tags):
1734	"""Reject release without resetting caches or discarding more memory."""
1735	core = object.__new__(EngineCore)
1736	core.scheduler = create_scheduler()
1737	core.scheduler.set_pause_state(PauseState.PAUSED_ALL)
1738	core.batch_queue = None
1739	core._reset_caches = Mock()
1740	core.model_executor = object.__new__(UniProcExecutor)
1741	core.model_executor.sleeping_tags = sleeping_tags.copy()
1742	core.model_executor.collective_rpc = Mock()
1743
1744	with pytest.raises(
1745	RuntimeError, match="requires all executor memory to be resident"
1746	):
1747	core.release_kv_cache_memory()
1748
1749	core._reset_caches.assert_not_called()
1750	core.model_executor.collective_rpc.assert_not_called()
1751	assert core.model_executor.sleeping_tags == sleeping_tags
1752
1753
1754	def test_reset_connector_cache_no_connector_is_no_op_success():
1755	"""``reset_connector_cache`` must return True when no connector is
1756	configured.
1757
1758	Without this, ``reset_prefix_cache(reset_connector=True)`` returns
1759	``False`` on every engine that doesn't have a KV connector configured —
1760	even when the local prefix cache reset succeeded — and any caller that
1761	interprets the return value as "did the reset I asked for succeed?"
1762	sees a spurious failure.
1763	"""
1764	scheduler = create_scheduler(enable_prefix_caching=True)
1765	assert scheduler.connector is None
1766
1767	# No-connector reset is treated as success.
1768	assert scheduler.reset_connector_cache() is True
1769
1770	# End-to-end: reset_prefix_cache(reset_connector=True) on an idle
1771	# scheduler succeeds with or without a connector.
1772	assert scheduler.reset_prefix_cache(reset_connector=True) is True
1773
1774
1775	def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch):
1776	monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
1777	(tmp_path / "config.json").write_text(
1778	'{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
1779	)
1780	scheduler = create_scheduler(
1781	model=str(tmp_path),
1782	max_num_seqs=16,
1783	max_num_batched_tokens=20,
1784	num_speculative_tokens=4,
1785	parallel_drafting=True,
1786	skip_tokenizer_init=True,
1787	)
1788	speculative_config = scheduler.vllm_config.speculative_config
1789	assert speculative_config is not None
1790	assert scheduler.max_num_scheduled_tokens == 20
1791	assert speculative_config.max_num_new_slots_for_drafting == 3
1792
1793	for request in create_requests(num_requests=2, num_tokens=10):
1794	scheduler.add_request(request)
1795
1796	assert scheduler.schedule().num_scheduled_tokens == {"0": 10, "1": 4}
1797
1798
1799	# Note - these test cases mirror some of those in test_rejection_sampler.py
1800	@pytest.mark.parametrize(
1801	"spec_tokens,output_tokens,expected,expected_per_req",
1802	[
1803	([[1, 2, 3]], [[1, 2, 3, 4]], (1, 3, 3, [1, 1, 1]), [[0, 0, 0, 1]]), # perfect
1804	([[1, 2, 3]], [[1, 5]], (1, 3, 1, [1, 0, 0]), [[0, 1, 0, 0]]), # early mismatch
1805	(
1806	[[1, 2], [3]],
1807	[[1, 2, 5], [3, 4]],
1808	(2, 3, 3, [2, 1]),
1809	[[0, 0, 1], [0, 1, 0]],
1810	), # multiple sequences
1811	([[1]], [[1, 2]], (1, 1, 1, [1]), [[0, 1]]), # single token sequence
1812	([[]], [[5]], (0, 0, 0, [0]), [[0, 0]]), # empty sequence -> empty accumulator
1813	(
1814	[[1, 2, 3], [4, 5, 6]],
1815	[[1, 2, 7], [4, 8]],
1816	(2, 6, 3, [2, 1, 0]),
1817	[[0, 0, 1, 0], [0, 1, 0, 0]],
1818	), # multiple mismatches
1819	],
1820	)
1821	def test_schedule_spec_decoding_stats(
1822	spec_tokens, output_tokens, expected, expected_per_req
1823	):
1824	"""Test scheduling behavior with speculative decoding.
1825
1826	This test verifies that:
1827	1. Speculated tokens get scheduled correctly
1828	2. The aggregate SpecDecodingStats count draft and accepted tokens
1829	3. The per-request accumulator (enabled via per_request_spec_decode_metrics)
1830	buckets the same acceptance by accepted draft count (j)
1831	"""
1832	num_spec_tokens = max(1, max(len(t) for t in spec_tokens))
1833	scheduler = create_scheduler(
1834	num_speculative_tokens=num_spec_tokens,
1835	per_request_spec_decode_metrics="summary",
1836	)
1837	requests = create_requests(num_requests=len(spec_tokens), num_tokens=1)
1838	req_ids = []
1839	req_to_index = {}
1840	for i, request in enumerate(requests):
1841	scheduler.add_request(request)
1842	req_ids.append(request.request_id)
1843	req_to_index[request.request_id] = i
1844
1845	# Schedule a decode, which will also draft speculative tokens
1846	output = scheduler.schedule()
1847	assert len(output.scheduled_new_reqs) == len(requests)
1848	assert output.total_num_scheduled_tokens == len(requests)
1849	for i in range(len(requests)):
1850	req_id = requests[i].request_id
1851	assert output.num_scheduled_tokens[req_id] == 1
1852	assert req_id not in output.scheduled_spec_decode_tokens
1853
1854	model_runner_output = ModelRunnerOutput(
1855	req_ids=req_ids,
1856	req_id_to_index=req_to_index,
1857	sampled_token_ids=[[0] for _ in range(len(requests))],
1858	logprobs=None,
1859	prompt_logprobs_dict={},
1860	pooler_output=[],
1861	)
1862	engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
1863	draft_token_ids = DraftTokenIds(req_ids, spec_tokens)
1864	scheduler.update_draft_token_ids(draft_token_ids)
1865
1866	for i in range(len(requests)):
1867	running_req = scheduler.running[i]
1868	# The prompt token
1869	assert running_req.num_computed_tokens == 1
1870	# The prompt token and the sampled token
1871	assert running_req.num_tokens == 2
1872	# The prompt token, the sampled token, and the speculated tokens
1873	assert running_req.num_tokens_with_spec == 2 + len(spec_tokens[i])
1874
1875	# No draft or accepted tokens counted yet
1876	assert not engine_core_outputs or (
1877	engine_core_outputs[0].scheduler_stats.spec_decoding_stats is None
1878	)
1879
1880	# Schedule the speculated tokens for validation
1881	output = scheduler.schedule()
1882	assert len(output.scheduled_new_reqs) == 0
1883	# The sampled token and speculated tokens
1884	assert output.total_num_scheduled_tokens == len(requests) + sum(
1885	len(ids) for ids in spec_tokens
1886	)
1887	for i in range(len(requests)):
1888	req_id = requests[i].request_id
1889	assert output.num_scheduled_tokens[req_id] == 1 + len(spec_tokens[i])
1890	if spec_tokens[i]:
1891	assert len(output.scheduled_spec_decode_tokens[req_id]) == len(
1892	spec_tokens[i]
1893	)
1894	else:
1895	assert req_id not in output.scheduled_spec_decode_tokens
1896
1897	model_runner_output = ModelRunnerOutput(
1898	req_ids=req_ids,
1899	req_id_to_index=req_to_index,
1900	sampled_token_ids=output_tokens,
1901	logprobs=None,
1902	prompt_logprobs_dict={},
1903	pooler_output=[],
1904	)
1905	engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
1906
1907	scheduler_stats = (
1908	engine_core_outputs[0].scheduler_stats if engine_core_outputs else None
1909	)
1910	if expected[0] == 0:
1911	assert scheduler_stats is not None
1912	assert scheduler_stats.spec_decoding_stats is None
1913	else:
1914	assert scheduler_stats is not None
1915	assert scheduler_stats.spec_decoding_stats is not None
1916	stats = scheduler_stats.spec_decoding_stats
1917	assert stats.num_drafts == expected[0]
1918	assert stats.num_draft_tokens == expected[1]
1919	assert stats.num_accepted_tokens == expected[2]
1920	assert stats.num_accepted_tokens_per_pos == expected[3]
1921
1922	# Per-request accumulator: the same acceptance, bucketed by accepted draft
1923	# count (j) on each request rather than summed across the batch. The
1924	# accumulator is created eagerly on add_request, so every request has one
1925	# (an empty histogram when it drafted nothing).
1926	for i, req_id in enumerate(req_ids):
1927	payload = scheduler.requests[req_id].spec_decode_metrics.to_dict()
1928	assert payload["acceptance_histogram"] == expected_per_req[i]
1929	assert payload["num_draft_tokens"] == len(spec_tokens[i])
1930	assert "per_step_accepted" not in payload # summary level
1931
1932
1933	def _run_spec_verify_steps(scheduler, rounds, num_invalid_per_round=None):
1934	"""Drive prefill + one draft/verify step per round for a single request.
1935
1936	``rounds`` is a list of ``(spec_token_ids, output_token_ids)`` -- one verify
1937	step each. ``num_invalid_per_round`` optionally injects the grammar-invalid
1938	draft count structured-output decoding would set on the verify output.
1939	Returns the request.
1940	"""
1941	[req] = create_requests(num_requests=1, num_tokens=1)
1942	scheduler.add_request(req)
1943	rid = req.request_id
1944	req_to_index = {rid: 0}
1945
1946	def _mk_output(sampled):
1947	return ModelRunnerOutput(
1948	req_ids=[rid],
1949	req_id_to_index=req_to_index,
1950	sampled_token_ids=sampled,
1951	logprobs=None,
1952	prompt_logprobs_dict={},
1953	pooler_output=[],
1954	)
1955
1956	# Prefill: sample one token.
1957	scheduler.update_from_output(scheduler.schedule(), _mk_output([[0]]))
1958	for i, (spec, out) in enumerate(rounds):
1959	scheduler.update_draft_token_ids(DraftTokenIds([rid], [spec]))
1960	output = scheduler.schedule()
1961	if num_invalid_per_round is not None and num_invalid_per_round[i]:
1962	output.num_invalid_spec_tokens = {rid: num_invalid_per_round[i]}
1963	scheduler.update_from_output(output, _mk_output([out]))
1964	return req
1965
1966
1967	def test_per_request_spec_decode_detailed_records_per_step():
1968	scheduler = create_scheduler(
1969	num_speculative_tokens=3,
1970	per_request_spec_decode_metrics="detailed",
1971	)
1972	# Three verify steps for one request; num_accepted = len(output) - 1, so the
1973	# outputs below accept 3, 1, then 0 drafts across the steps.
1974	req = _run_spec_verify_steps(
1975	scheduler,
1976	[
1977	([1, 2, 3], [1, 2, 3, 4]), # accept 3
1978	([5, 6, 7], [5, 8]), # accept 1
1979	([9, 10, 11], [12]), # accept 0
1980	],
1981	)
1982	payload = scheduler.requests[req.request_id].spec_decode_metrics.to_dict()
1983	assert payload["per_step_accepted"] == [3, 1, 0]
1984	assert payload["per_step_drafted"] == [3, 3, 3]
1985	assert payload["num_spec_steps"] == 3
1986
1987
1988	def test_per_request_spec_decode_subtracts_invalid_drafts():
1989	# Grammar-invalidated drafts (num_invalid_spec_tokens, set by structured
1990	# output) are excluded from the proposed count, mirroring the aggregate.
1991	scheduler = create_scheduler(
1992	num_speculative_tokens=3,
1993	per_request_spec_decode_metrics="summary",
1994	)
1995	# One verify step: 3 drafted, 1 grammar-invalid, output accepts 2.
1996	req = _run_spec_verify_steps(
1997	scheduler,
1998	[([1, 2, 3], [1, 2, 5])],
1999	num_invalid_per_round=[1],
2000	)
2001	payload = scheduler.requests[req.request_id].spec_decode_metrics.to_dict()
2002	assert payload["num_draft_tokens"] == 2 # 3 drafted - 1 invalid
2003	assert payload["num_accepted_draft_tokens"] == 2 # len([1,2,5]) - 1
2004	assert payload["acceptance_histogram"] == [0, 0, 1, 0]
2005
2006
2007	def test_per_request_spec_decode_acceptance_disabled_by_default():
2008	scheduler = create_scheduler(num_speculative_tokens=3)
2009	assert scheduler.spec_decode_metrics_level == "none"
2010	req = _run_spec_verify_steps(scheduler, [([1, 2, 3], [1, 2, 3, 4])])
2011	assert scheduler.requests[req.request_id].spec_decode_metrics is None
2012
2013
2014	def test_spec_decoding_stats_empty_output():
2015	"""Test that spec decoding stats handle empty output tokens gracefully.
2016
2017	This is a regression test for a bug where empty sampled_token_ids
2018	would cause num_accepted = len([]) - 1 = -1, leading to a
2019	ValueError when incrementing a Prometheus counter with a negative value.
2020	"""
2021	num_spec_tokens = 3
2022	scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens)
2023	requests = create_requests(num_requests=1, num_tokens=1)
2024	request = requests[0]
2025	req_id = request.request_id
2026
2027	scheduler.add_request(request)
2028
2029	# Initial schedule (prefill)
2030	output = scheduler.schedule()
2031	assert len(output.scheduled_new_reqs) == 1
2032
2033	# Complete the prefill with a sampled token
2034	model_runner_output = ModelRunnerOutput(
2035	req_ids=[req_id],
2036	req_id_to_index={req_id: 0},
2037	sampled_token_ids=[[0]],
2038	logprobs=None,
2039	prompt_logprobs_dict={},
2040	pooler_output=[],
2041	)
2042	scheduler.update_from_output(output, model_runner_output)
2043
2044	# Add draft tokens for speculation
2045	draft_token_ids = DraftTokenIds([req_id], [[1, 2, 3]])
2046	scheduler.update_draft_token_ids(draft_token_ids)
2047
2048	# Schedule the speculated tokens for validation
2049	output = scheduler.schedule()
2050	assert req_id in output.scheduled_spec_decode_tokens
2051	assert len(output.scheduled_spec_decode_tokens[req_id]) == 3
2052
2053	# Simulate empty output tokens (e.g., due to request abortion or error)
2054	# This would previously cause num_accepted = -1 and crash
2055	model_runner_output = ModelRunnerOutput(
2056	req_ids=[req_id],
2057	req_id_to_index={req_id: 0},
2058	sampled_token_ids=[[]], # Empty output tokens
2059	logprobs=None,
2060	prompt_logprobs_dict={},
2061	pooler_output=[],
2062	)
2063
2064	# This should not raise an error
2065	engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
2066
2067	# Spec decoding stats should be None since no tokens were generated
2068	scheduler_stats = (
2069	engine_core_outputs[0].scheduler_stats if engine_core_outputs else None
2070	)
2071	assert scheduler_stats is None or scheduler_stats.spec_decoding_stats is None
2072
2073
2074	def test_no_spec_tokens_scheduled_for_prefill_chunks():
2075	"""Test that draft tokens are ignored for prefill chunk requests.
2076
2077	When a request is being prefilled in chunks (chunked prefill), draft tokens
2078	from `update_draft_token_ids` should be ignored until the prefill is complete.
2079
2080	The bug manifests when:
2081	- A prefill chunk is scheduled
2082	- Draft tokens are provided via update_draft_token_ids
2083	- The next schedule has enough budget to include spec tokens
2084
2085	Without the fix, spec tokens would incorrectly be scheduled with the
2086	remaining prefill tokens. With the fix, draft tokens are ignored for
2087	prefill chunks.
2088	"""
2089	num_spec_tokens = 3
2090	# Use budget of 50, with 80 token prompt:
2091	# - First chunk: 50 tokens
2092	# - Second chunk: 30 remaining + potentially 3 spec tokens = 33
2093	# Without fix: num_scheduled_spec_tokens = 33 + 50 - 80 = 3 (BUG!)
2094	# With fix: spec_token_ids cleared, so no spec tokens scheduled
2095	scheduler = create_scheduler(
2096	num_speculative_tokens=num_spec_tokens,
2097	max_num_batched_tokens=50,
2098	enable_chunked_prefill=True,
2099	)
2100	requests = create_requests(num_requests=1, num_tokens=80)
2101	req = requests[0]
2102	scheduler.add_request(req)
2103
2104	# First schedule - prefill chunk (50 of 80 tokens)
2105	output = scheduler.schedule()
2106	assert len(output.scheduled_new_reqs) == 1
2107	assert output.num_scheduled_tokens[req.request_id] == 50
2108
2109	# Update from output (no sampled token since still prefilling)
2110	req_to_index = {req.request_id: 0}
2111	model_runner_output = ModelRunnerOutput(
2112	req_ids=[req.request_id],
2113	req_id_to_index=req_to_index,
2114	sampled_token_ids=[[]],
2115	logprobs=None,
2116	prompt_logprobs_dict={},
2117	pooler_output=[],
2118	)
2119	scheduler.update_from_output(output, model_runner_output)
2120
2121	# Provide draft tokens while request is still in prefill.
2122	# The fix ensures these are ignored for prefill chunks.
2123	draft_token_ids = DraftTokenIds([req.request_id], [[1, 2, 3]])
2124	scheduler.update_draft_token_ids(draft_token_ids)
2125
2126	# Second schedule - remaining 30 tokens of prefill
2127	output = scheduler.schedule()
2128	# KEY ASSERTION: Should schedule exactly the remaining 30 prefill tokens,
2129	# NOT 33 (30 + 3 spec). Without the fix, this would be 33.
2130	assert output.num_scheduled_tokens[req.request_id] == 30, (
2131	f"Expected 30 tokens (remaining prefill only), "
2132	f"got {output.num_scheduled_tokens[req.request_id]}. "
2133	"Spec tokens should not be scheduled with prefill chunks."
2134	)
2135	# No spec tokens should be in the output
2136	assert req.request_id not in output.scheduled_spec_decode_tokens, (
2137	"Spec tokens should not be scheduled with prefill chunks"
2138	)
2139
2140	# Update from output with a sampled token (prefill complete)
2141	model_runner_output = ModelRunnerOutput(
2142	req_ids=[req.request_id],
2143	req_id_to_index=req_to_index,
2144	sampled_token_ids=[[42]],
2145	logprobs=None,
2146	prompt_logprobs_dict={},
2147	pooler_output=[],
2148	)
2149	scheduler.update_from_output(output, model_runner_output)
2150
2151	# Now provide draft tokens - should be accepted since prefill is complete
2152	draft_token_ids = DraftTokenIds([req.request_id], [[1, 2, 3]])
2153	scheduler.update_draft_token_ids(draft_token_ids)
2154
2155	# spec_token_ids SHOULD be set after prefill is complete
2156	assert req.spec_token_ids == [1, 2, 3], (
2157	f"spec_token_ids should be set after prefill, got {req.spec_token_ids}"
2158	)
2159
2160	# Third schedule - decode phase with spec tokens
2161	output = scheduler.schedule()
2162	# 1 new token + 3 spec tokens = 4
2163	assert output.num_scheduled_tokens[req.request_id] == 4
2164	assert req.request_id in output.scheduled_spec_decode_tokens
2165	assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens
2166
2167
2168	def _model_output(scheduler, output, sampled):
2169	"""Feed `sampled` (per-request list) back to the scheduler."""
2170	req_ids = list(output.num_scheduled_tokens.keys())
2171	scheduler.update_from_output(
2172	output,
2173	ModelRunnerOutput(
2174	req_ids=req_ids,
2175	req_id_to_index={r: i for i, r in enumerate(req_ids)},
2176	sampled_token_ids=sampled,
2177	logprobs=None,
2178	prompt_logprobs_dict={},
2179	pooler_output=[],
2180	),
2181	)
2182
2183
2184	def test_spec_decode_padding_first_decode_step():
2185	"""A request taking its first decode step (whole prompt already computed via
2186	a prefix-cache hit) is padded with placeholder (-1) spec tokens so it enters
2187	the worker with the same 1 + num_spec_tokens shape as the other speculative
2188	decodes, keeping the batch uniform.
2189	"""
2190	num_spec = 3
2191	scheduler = create_scheduler(
2192	num_speculative_tokens=num_spec,
2193	enable_prefix_caching=True,
2194	block_size=16,
2195	)
2196	# Two identical 33-token prompts: 2 full blocks (32 tokens) get cached, so a
2197	# second identical request hits num_computed == num_prompt_tokens - 1.
2198	r1, r2 = create_requests(
2199	num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16
2200	)
2201
2202	# Drive r1 through prefill so its prompt blocks are cached, then give it real
2203	# drafts so it is a running speculative decode (1 + num_spec shape).
2204	scheduler.add_request(r1)
2205	out = scheduler.schedule()
2206	assert out.num_scheduled_tokens[r1.request_id] == 33
2207	_model_output(scheduler, out, [[100]])
2208	scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))
2209
2210	# r2 arrives; its whole prompt is a prefix-cache hit -> first decode step.
2211	scheduler.add_request(r2)
2212	out = scheduler.schedule()
2213
2214	# r1 verifies its real drafts.
2215	assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]
2216	# r2 is padded to the 1 + num_spec shape with placeholder (-1) drafts.
2217	assert out.num_scheduled_tokens[r2.request_id] == 1 + num_spec
2218	assert out.scheduled_spec_decode_tokens[r2.request_id] == [-1] * num_spec
2219
2220
2221	def test_spec_decode_padding_resumed_request_without_running_requests():
2222	"""Pad a synchronously resumed request even without local running work."""
2223	num_spec = 3
2224	scheduler = create_scheduler(
2225	num_speculative_tokens=num_spec,
2226	use_kv_connector=mock_kv(matched_tokens=32, is_async=False),
2227	)
2228	(resumed,) = create_requests(num_requests=1, num_tokens=33, max_tokens=16)
2229
2230	scheduler.add_request(resumed)
2231	out = scheduler.schedule()
2232
2233	assert out.num_scheduled_tokens[resumed.request_id] == 1 + num_spec
2234	assert out.scheduled_spec_decode_tokens[resumed.request_id] == [-1] * num_spec
2235
2236
2237	def test_spec_decode_padding_skipped_for_diffusion():
2238	"""Diffusion spec tokens are the fixed-size denoising canvas, not
2239	rejectable drafts: a first-decode-step request must keep its 1-token span
2240	instead of being padded to 1 + num_spec_tokens, which would overflow the
2241	canvas.
2242	"""
2243	num_spec = 3
2244	scheduler = create_scheduler(
2245	num_speculative_tokens=num_spec,
2246	enable_prefix_caching=True,
2247	block_size=16,
2248	)
2249	# Diffusion schedulers initialize this to 0 (model_config.is_diffusion).
2250	scheduler.num_sampled_tokens_per_step = 0
2251	r1, r2 = create_requests(
2252	num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16
2253	)
2254
2255	scheduler.add_request(r1)
2256	out = scheduler.schedule()
2257	assert out.num_scheduled_tokens[r1.request_id] == 33
2258	_model_output(scheduler, out, [[100]])
2259	scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))
2260
2261	# r2 arrives; its whole prompt is a prefix-cache hit -> needs exactly
2262	# 1 token while r1 is a running speculative decode.
2263	scheduler.add_request(r2)
2264	out = scheduler.schedule()
2265
2266	assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]
2267	# r2 keeps its true 1-token span; no placeholder drafts are attached.
2268	assert out.num_scheduled_tokens[r2.request_id] == 1
2269	assert r2.request_id not in out.scheduled_spec_decode_tokens
2270
2271
2272	@pytest.mark.parametrize("aligned_tokens", [1, 2, 3])
2273	def test_spec_decode_padding_dropped_when_recurrent_alignment_clips(
2274	aligned_tokens: int,
2275	):
2276	"""A clipped padded request must drop its padding, not shorten it.
2277
2278	The sampler derives a request's row window from its draft count, so a
2279	padded request that keeps fewer than 1 + num_spec query rows would select
2280	the preceding request's hidden states.
2281	"""
2282	num_spec = 3
2283	scheduler = create_scheduler(
2284	num_speculative_tokens=num_spec,
2285	enable_prefix_caching=True,
2286	block_size=16,
2287	)
2288	r1, r2 = create_requests(
2289	num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16
2290	)
2291
2292	scheduler.add_request(r1)
2293	out = scheduler.schedule()
2294	assert out.num_scheduled_tokens[r1.request_id] == 33
2295	_model_output(scheduler, out, [[100]])
2296	scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))
2297
2298	# Model the recurrent align split that clips a padded 1+3 prompt tail at
2299	# the next cache boundary. The real QSA block size produces all three cases
2300	# depending on prompt length modulo four.
2301	scheduler.need_mamba_block_aligned_split = True
2302	original_split = scheduler._mamba_block_aligned_split
2303
2304	def align_prompt_tail(request, num_new_tokens, *args):
2305	if request is r2:
2306	return aligned_tokens
2307	return original_split(request, num_new_tokens, *args)
2308
2309	scheduler._mamba_block_aligned_split = align_prompt_tail
2310	scheduler.add_request(r2)
2311	out = scheduler.schedule()
2312
2313	assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]
2314	assert out.num_scheduled_tokens[r2.request_id] == 1
2315	assert r2.request_id not in out.scheduled_spec_decode_tokens
2316
2317
2318	def test_spec_decode_padding_skipped_with_prefill_in_batch():
2319	"""Padding is skipped when the batch contains a prefill chunk: the batch is
2320	already mixed/non-uniform, so padding a new decode request buys nothing.
2321	"""
2322	num_spec = 3
2323	scheduler = create_scheduler(
2324	num_speculative_tokens=num_spec,
2325	enable_prefix_caching=True,
2326	block_size=16,
2327	max_num_batched_tokens=64,
2328	)
2329	# r_warm + r_candidate share a prompt so r_candidate gets a full prefix hit.
2330	r_warm, r_candidate = create_requests(
2331	num_requests=2, num_tokens=33, same_prompt=True, max_tokens=1
2332	)
2333	# r_long has a different, long prompt that prefills over multiple chunks.
2334	(r_long,) = create_requests(num_requests=1, num_tokens=100, max_tokens=16)
2335
2336	# Warm the prefix cache with r_warm's prompt (it finishes; blocks stay cached).
2337	scheduler.add_request(r_warm)
2338	out = scheduler.schedule()
2339	assert out.num_scheduled_tokens[r_warm.request_id] == 33
2340	_model_output(scheduler, out, [[100]])
2341	assert r_warm.request_id in scheduler.finished_req_ids
2342
2343	# Start r_long; after one chunk it remains a prefill chunk in the running queue.
2344	scheduler.add_request(r_long)
2345	out = scheduler.schedule()
2346	_model_output(scheduler, out, [[]]) # still prefilling, no sampled token
2347	assert r_long.is_prefill_chunk
2348
2349	# r_candidate arrives (prefix-cache hit -> first decode step) alongside the
2350	# in-flight prefill chunk.
2351	scheduler.add_request(r_candidate)
2352	out = scheduler.schedule()
2353
2354	# The batch has a prefill chunk, so r_candidate is NOT padded.
2355	assert r_long.request_id in out.num_scheduled_tokens
2356	assert out.num_scheduled_tokens[r_candidate.request_id] == 1
2357	assert r_candidate.request_id not in out.scheduled_spec_decode_tokens
2358
2359
2360	def test_scheduler_stats_waiting_queues():
2361	"""Test that scheduler stats correctly report capacity-bound and blocked
2362	waiting requests."""
2363	# Create scheduler with limited capacity so we can have waiting requests
2364	scheduler = create_scheduler(max_num_batched_tokens=100)
2365
2366	# Create requests: some will be scheduled, some will wait on capacity,
2367	# and some will be blocked by constraints
2368	all_requests = create_requests(num_requests=5, num_tokens=50)
2369
2370	# Add 3 requests - only 2 can be scheduled (2 * 50 = 100 tokens)
2371	# The 3rd will remain in waiting queue (capacity constraint)
2372	for request in all_requests[:3]:
2373	scheduler.add_request(request)
2374
2375	# Add 2 more that are blocked on grammar compilation
2376	for request in all_requests[3:]:
2377	request.status = RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR
2378	scheduler.add_request(request)
2379
2380	# Schedule - this will schedule 2 requests, leaving 1 in waiting
2381	output = scheduler.schedule()
2382
2383	# Verify: 2 scheduled, 1 still waiting on capacity, 2 blocked by constraints
2384	assert len(output.scheduled_new_reqs) == 2
2385	assert len(scheduler.waiting) == 3
2386	assert len(scheduler.deferred_waiting) == 2
2387
2388	# Call update_from_output() to get frontend-facing stat
2389	scheduled_req_ids = list(output.num_scheduled_tokens.keys())
2390	model_runner_output = ModelRunnerOutput(
2391	req_ids=scheduled_req_ids,
2392	req_id_to_index={req_id: i for i, req_id in enumerate(scheduled_req_ids)},
2393	sampled_token_ids=[[1]] * len(scheduled_req_ids),
2394	logprobs=None,
2395	prompt_logprobs_dict={},
2396	pooler_output=[],
2397	)
2398	engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
2399	assert engine_core_outputs and len(engine_core_outputs) > 0
2400	stats = engine_core_outputs[0].scheduler_stats
2401	assert stats is not None
2402
2403	# Verify stats match queue lengths after scheduling
2404	assert stats.num_running_reqs == 2 # 2 were scheduled
2405	assert stats.num_waiting_reqs == 1 # 1 waiting on capacity
2406	assert stats.num_skipped_waiting_reqs == 2 # 2 blocked by constraints
2407
2408
2409	def _assert_right_scheduler_output(
2410	output: SchedulerOutput,
2411	num_requests: int,
2412	expected_num_scheduled_tokens: int,
2413	):
2414	"""Check if SchedulerOutput is correct after remote KV cache hit."""
2415	# We should inject the kv_connector_metadata.
2416	assert len(output.kv_connector_metadata.requests) == num_requests
2417
2418	# Only num_tokens - matched_num_new_tokens should be scheduled.
2419	for _, num_scheduled_tokens in output.num_scheduled_tokens.items():
2420	assert num_scheduled_tokens == expected_num_scheduled_tokens
2421
2422
2423	def _assert_right_kv_cache_manager(
2424	scheduler: Scheduler,
2425	requests: list[Request],
2426	num_tokens: int,
2427	block_size: int,
2428	num_requests: int,
2429	num_total_blocks: int,
2430	):
2431	"""Check whether KVCacheManager is correct after allocate."""
2432	# Make sure the request stats are right.
2433	EXPECTED_TOTAL_BLOCKS = num_tokens // block_size
2434	for req in requests:
2435	blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
2436	0
2437	].req_to_blocks[req.request_id]
2438	hashes = req.block_hashes
2439	assert (
2440	scheduler.kv_cache_manager.coordinator.single_type_managers[
2441	0
2442	].num_cached_block[req.request_id]
2443	== EXPECTED_TOTAL_BLOCKS
2444	)
2445	assert len(blocks) == EXPECTED_TOTAL_BLOCKS
2446	assert len(hashes) == EXPECTED_TOTAL_BLOCKS
2447
2448	# Make sure we actually touched all the blocks.
2449	BLOCKS_PER_REQ = num_tokens / block_size
2450	assert (
2451	scheduler.kv_cache_manager.block_pool.get_num_free_blocks()
2452	== num_total_blocks - num_requests * BLOCKS_PER_REQ
2453	)
2454
2455
2456	def _step_until_done(
2457	scheduler: Scheduler,
2458	output: SchedulerOutput,
2459	model_runner_output: ModelRunnerOutput,
2460	):
2461	"""Loop over schedule(), update_from_output() until finished."""
2462	all_finished = False
2463	_ = scheduler.update_from_output(output, model_runner_output)
2464	while not all_finished:
2465	# Schedule + a few iterations until stopping.
2466	output = scheduler.schedule()
2467	assert len(scheduler.running)
2468	for _, num_scheduled_tokens in output.num_scheduled_tokens.items():
2469	# We should be in the decode phase now.
2470	assert num_scheduled_tokens == 1
2471	if scheduler.connector is not None:
2472	assert len(output.kv_connector_metadata.requests) == 0
2473	if scheduler.ec_connector is not None:
2474	assert len(output.ec_connector_metadata.mm_datas) == 0
2475	ecos = scheduler.update_from_output(output, model_runner_output)[0]
2476	all_done = True
2477	for eco in ecos.outputs:
2478	if eco.finish_reason is None:
2479	all_done = False
2480	all_finished = all_done
2481
2482
2483	def _num_waiting_requests(scheduler: Scheduler) -> int:
2484	return len(scheduler.waiting) + len(scheduler.kv_holding_waiting)
2485
2486
2487	def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
2488	"""Cycle requests through a KV transfer cycle."""
2489	# Requests should first transition to WAITING_FOR_REMOTE_KVS
2490	output = scheduler.schedule()
2491	assert _num_waiting_requests(scheduler) == len(req_ids)
2492	assert len(scheduler.running) == 0
2493	assert len(output.scheduled_new_reqs) == 0
2494	for req in scheduler.requests.values():
2495	assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
2496
2497	# No model execution yet
2498	EMPTY_OUTPUT = ModelRunnerOutput(
2499	req_ids=[],
2500	req_id_to_index={},
2501	sampled_token_ids=[],
2502	logprobs=None,
2503	prompt_logprobs_dict={},
2504	pooler_output=[],
2505	)
2506	initial_ecos = scheduler.update_from_output(output, EMPTY_OUTPUT)
2507
2508	# Simulate KV transfer completion using KVConnectorOutput.finished_recving
2509	output = scheduler.schedule()
2510	assert _num_waiting_requests(scheduler) == len(req_ids)
2511	assert len(scheduler.running) == 0
2512
2513	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2514	req_ids=[],
2515	req_id_to_index={},
2516	sampled_token_ids=[],
2517	logprobs=None,
2518	prompt_logprobs_dict={},
2519	pooler_output=[],
2520	kv_connector_output=KVConnectorOutput(finished_recving=req_ids),
2521	)
2522	scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
2523	for req_id in req_ids:
2524	assert req_id in scheduler.finished_recving_kv_req_ids
2525
2526	return initial_ecos
2527
2528
2529	@pytest.mark.parametrize(
2530	("load_modes", "expected_has_sync_loads"),
2531	[
2532	((True,), False),
2533	((False,), True),
2534	((True, False), True),
2535	],
2536	)
2537	@pytest.mark.skip_global_cleanup
2538	def test_has_sync_kv_loads(
2539	load_modes: tuple[bool, ...],
2540	expected_has_sync_loads: bool,
2541	tmp_path,
2542	):
2543	(tmp_path / "config.json").write_text(
2544	'{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
2545	)
2546	block_size = 16
2547	scheduler = create_scheduler(
2548	model=str(tmp_path),
2549	skip_tokenizer_init=True,
2550	use_kv_connector=mock_kv(matched_tokens=block_size, is_async=False),
2551	block_size=block_size,
2552	)
2553	requests = create_requests(
2554	num_requests=len(load_modes),
2555	num_tokens=block_size * 2,
2556	block_size=block_size,
2557	)
2558	for request in requests:
2559	scheduler.add_request(request)
2560
2561	scheduler.connector.get_num_new_matched_tokens = Mock(
2562	side_effect=[(block_size, is_async) for is_async in load_modes]
2563	)
2564
2565	output = scheduler.schedule()
2566
2567	assert output.has_sync_kv_loads is expected_has_sync_loads
2568
2569
2570	def test_kv_connector_honors_skip_reading_prefix_cache():
2571	"""A request that must score every prompt row takes no external hit."""
2572	BLOCK_SIZE = 16
2573	scheduler = create_scheduler(
2574	enable_prefix_caching=True,
2575	use_kv_connector=mock_kv(matched_tokens=BLOCK_SIZE * 2, is_async=False),
2576	block_size=BLOCK_SIZE,
2577	)
2578	plain, scoring = create_requests(num_requests=2, num_tokens=BLOCK_SIZE * 4)
2579	scoring.sampling_params.prompt_logprobs = 1
2580	scoring.skip_reading_prefix_cache = True
2581	scheduler.add_request(plain)
2582	scheduler.add_request(scoring)
2583
2584	output = scheduler.schedule()
2585	assert output.num_scheduled_tokens[plain.request_id] == BLOCK_SIZE * 2
2586	assert output.num_scheduled_tokens[scoring.request_id] == BLOCK_SIZE * 4
2587
2588
2589	@pytest.mark.parametrize("is_async", [False, True])
2590	def test_kv_connector_basic(is_async: bool):
2591	"""Test whether Scheduler with KVConnector schedules tokens, allocates
2592	memory, and cleans up requests as expected under normal operation.
2593	"""
2594	# Setup Scheduler.
2595	BLOCK_SIZE = 16
2596	NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
2597	scheduler = create_scheduler(
2598	enable_prefix_caching=True,
2599	use_kv_connector=mock_kv(
2600	matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async
2601	),
2602	block_size=BLOCK_SIZE,
2603	)
2604	NUM_TOTAL_BLOCKS = scheduler.kv_cache_manager.block_pool.get_num_free_blocks()
2605
2606	######################################################
2607	# FIRST SET OF REQUESTS - External Hit Only
2608	NUM_REQUESTS = 2
2609	NUM_TOKENS = NUM_MATCHED_NEW_TOKENS * 2
2610	MAX_TOKENS = 3
2611	requests = create_requests(
2612	num_requests=NUM_REQUESTS,
2613	num_tokens=NUM_TOKENS,
2614	max_tokens=MAX_TOKENS,
2615	block_size=BLOCK_SIZE,
2616	)
2617	req_ids = []
2618	req_to_index = {}
2619	for i, request in enumerate(requests):
2620	scheduler.add_request(request)
2621	req_ids.append(request.request_id)
2622	req_to_index[request.request_id] = i
2623
2624	if is_async:
2625	_step_until_kv_transfer_finished(scheduler, req_ids)
2626
2627	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2628	req_ids=req_ids,
2629	req_id_to_index=req_to_index,
2630	sampled_token_ids=[[1000]] * len(req_ids),
2631	logprobs=None,
2632	prompt_logprobs_dict={},
2633	pooler_output=[],
2634	)
2635
2636	# Ensure ScheduleOutput is correct.
2637	output = scheduler.schedule()
2638	_assert_right_scheduler_output(
2639	output=output,
2640	num_requests=NUM_REQUESTS,
2641	# Just the incremental tokens should be scheduled.
2642	expected_num_scheduled_tokens=NUM_TOKENS - NUM_MATCHED_NEW_TOKENS,
2643	)
2644
2645	# Ensure KVCacheManager is correct.
2646	_assert_right_kv_cache_manager(
2647	scheduler, requests, NUM_TOKENS, BLOCK_SIZE, NUM_REQUESTS, NUM_TOTAL_BLOCKS
2648	)
2649
2650	# Continue Generation until done.
2651	_step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
2652	_ = scheduler.schedule()
2653	# Confirm we clean up the memory properly.
2654	assert (
2655	scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_TOTAL_BLOCKS
2656	)
2657
2658	######################################################
2659	# SECOND SET OF REQUESTS - Local And External Hit
2660	NUM_TOKENS_PREFIX = NUM_TOKENS
2661	# We will get a local prefix cache hit for the first
2662	# NUM_TOKENS_PREFIX tokens since they are used above.
2663	NUM_TOKENS = NUM_TOKENS_PREFIX * 2
2664	requests = create_requests(
2665	num_requests=NUM_REQUESTS,
2666	num_tokens=NUM_TOKENS,
2667	max_tokens=MAX_TOKENS,
2668	block_size=BLOCK_SIZE,
2669	)
2670	req_ids = []
2671	req_to_index = {}
2672	for i, request in enumerate(requests):
2673	scheduler.add_request(request)
2674	req_ids.append(request.request_id)
2675	req_to_index[request.request_id] = i
2676
2677	if is_async:
2678	_step_until_kv_transfer_finished(scheduler, req_ids)
2679
2680	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2681	req_ids=req_ids,
2682	req_id_to_index=req_to_index,
2683	sampled_token_ids=[[1000]] * len(req_ids),
2684	logprobs=None,
2685	prompt_logprobs_dict={},
2686	pooler_output=[],
2687	)
2688
2689	# We should get a local cache hit of NUM_TOKENS_PREFIX and
2690	# a remote KV cache hit of NUM_MATCHED_NEW_TOKENS.
2691	output = scheduler.schedule()
2692	_assert_right_scheduler_output(
2693	output=output,
2694	num_requests=NUM_REQUESTS,
2695	# Just the incremental tokens after local + remote cache hit.
2696	expected_num_scheduled_tokens=(
2697	NUM_TOKENS - NUM_TOKENS_PREFIX - NUM_MATCHED_NEW_TOKENS
2698	),
2699	)
2700
2701	# Ensure KVCacheManager is correct.
2702	_assert_right_kv_cache_manager(
2703	scheduler, requests, NUM_TOKENS, BLOCK_SIZE, NUM_REQUESTS, NUM_TOTAL_BLOCKS
2704	)
2705
2706	# Continue Generation until done.
2707	_step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
2708	_ = scheduler.schedule()
2709	# Confirm we clean up the memory properly.
2710	assert (
2711	scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_TOTAL_BLOCKS
2712	)
2713
2714
2715	@pytest.mark.parametrize("is_async", [False, True])
2716	@pytest.mark.parametrize("local_cache_hits", [False, True])
2717	def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool):
2718	"""Verify connector prefix cache metrics are updated
2719	correctly when the scheduler processes requests with KV connector hits.
2720	"""
2721	BLOCK_SIZE = 16
2722	if local_cache_hits:
2723	NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2 # 32 tokens
2724	NUM_LOCAL_HITS = NUM_MATCHED_NEW_TOKENS * 2 # 64 tokens
2725	NUM_REQUESTS = 1
2726	NUM_TOKENS = NUM_LOCAL_HITS * 2 # 128 tokens
2727	else:
2728	NUM_MATCHED_NEW_TOKENS = 4
2729	NUM_LOCAL_HITS = 0
2730	NUM_REQUESTS = 2
2731	NUM_TOKENS = 8 # 8 tokens
2732
2733	# Setup Scheduler.
2734	scheduler = create_scheduler(
2735	enable_prefix_caching=local_cache_hits,
2736	use_kv_connector=mock_kv(
2737	matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async
2738	),
2739	block_size=BLOCK_SIZE,
2740	)
2741
2742	if local_cache_hits:
2743	# First, establish local cache by running a request to completion
2744	requests = create_requests(
2745	num_requests=1,
2746	num_tokens=NUM_LOCAL_HITS,
2747	max_tokens=2,
2748	block_size=BLOCK_SIZE,
2749	)
2750	req_ids = []
2751	req_to_index = {}
2752	for i, request in enumerate(requests):
2753	scheduler.add_request(request)
2754	req_ids.append(request.request_id)
2755	req_to_index[request.request_id] = i
2756
2757	if is_async:
2758	_step_until_kv_transfer_finished(scheduler, req_ids)
2759
2760	# Run first request to completion to establish local cache
2761	output = scheduler.schedule()
2762	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2763	req_ids=req_ids,
2764	req_id_to_index=req_to_index,
2765	sampled_token_ids=[[1000]] * len(req_ids),
2766	logprobs=None,
2767	prompt_logprobs_dict={},
2768	pooler_output=[],
2769	)
2770	_step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
2771	_ = scheduler.schedule()
2772
2773	# --- Prepare test requests ---
2774	MAX_TOKENS = 2
2775	requests = create_requests(
2776	num_requests=NUM_REQUESTS,
2777	num_tokens=NUM_TOKENS,
2778	max_tokens=MAX_TOKENS,
2779	block_size=BLOCK_SIZE,
2780	)
2781	req_ids = []
2782	req_to_index = {}
2783	for i, request in enumerate(requests):
2784	scheduler.add_request(request)
2785	req_ids.append(request.request_id)
2786	req_to_index[request.request_id] = i
2787
2788	initial_ecos = None
2789	if is_async:
2790	initial_ecos = _step_until_kv_transfer_finished(scheduler, req_ids)
2791
2792	# --- Trigger scheduling and simulate model output ---
2793	output = scheduler.schedule()
2794	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2795	req_ids=[r.request_id for r in requests],
2796	req_id_to_index={r.request_id: i for i, r in enumerate(requests)},
2797	sampled_token_ids=[[1000]] * NUM_REQUESTS,
2798	logprobs=None,
2799	prompt_logprobs_dict={},
2800	pooler_output=[],
2801	)
2802
2803	# Update scheduler stats
2804	ecos = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
2805
2806	# --- Assertions ---
2807	assert ecos is not None and len(ecos) > 0
2808	assert ecos[0].scheduler_stats is not None
2809
2810	if local_cache_hits:
2811	# For async, local cache stats come from the first step
2812	if initial_ecos:
2813	local_stats = initial_ecos[0].scheduler_stats.prefix_cache_stats
2814	else:
2815	local_stats = ecos[0].scheduler_stats.prefix_cache_stats
2816	assert local_stats is not None
2817	assert local_stats.queries == NUM_TOKENS * NUM_REQUESTS
2818	assert local_stats.hits == NUM_LOCAL_HITS * NUM_REQUESTS
2819
2820	if initial_ecos:
2821	external_stats = initial_ecos[0].scheduler_stats.connector_prefix_cache_stats
2822	else:
2823	external_stats = ecos[0].scheduler_stats.connector_prefix_cache_stats
2824	assert external_stats is not None
2825
2826	assert external_stats.queries == (NUM_TOKENS - NUM_LOCAL_HITS) * NUM_REQUESTS
2827	assert external_stats.hits == NUM_MATCHED_NEW_TOKENS * NUM_REQUESTS
2828	assert external_stats.requests == NUM_REQUESTS
2829	assert external_stats.preempted_requests == 0
2830
2831
2832	@pytest.mark.parametrize(
2833	"use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
2834	)
2835	def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):
2836	"""Test whether scheduler with KVConnector is able to handle
2837	unable to allocate (run out of blocks in allocate_slots().
2838	"""
2839	# Setup Scheduler With Mock External Cache Hit.
2840	BLOCK_SIZE = 4
2841	NUM_BLOCKS = 10
2842	NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
2843	scheduler = create_scheduler(
2844	enable_prefix_caching=True,
2845	use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=False),
2846	block_size=BLOCK_SIZE,
2847	num_blocks=NUM_BLOCKS,
2848	# encoder connector should not affect test results
2849	use_ec_connector=use_ec_connector,
2850	ec_role=ec_role,
2851	)
2852
2853	# Create two requests. The second request will not be able to
2854	# allocate slots because it will not have enough blocks.
2855	NUM_REQUESTS = 2
2856	NUM_TOKENS = (NUM_BLOCKS // 2 + 1) * BLOCK_SIZE
2857	MAX_TOKENS = 2
2858	requests = create_requests(
2859	num_requests=NUM_REQUESTS,
2860	num_tokens=NUM_TOKENS,
2861	max_tokens=MAX_TOKENS,
2862	block_size=BLOCK_SIZE,
2863	)
2864	req_ids = []
2865	req_to_index = {}
2866	for i, request in enumerate(requests):
2867	scheduler.add_request(request)
2868	req_ids.append(request.request_id)
2869	req_to_index[request.request_id] = i
2870
2871	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2872	req_ids=req_ids,
2873	req_id_to_index=req_to_index,
2874	sampled_token_ids=[[1000]] * len(req_ids),
2875	logprobs=None,
2876	prompt_logprobs_dict={},
2877	pooler_output=[],
2878	)
2879
2880	# Just one request should be running.
2881	output = scheduler.schedule()
2882	_assert_right_scheduler_output(
2883	output,
2884	num_requests=1,
2885	expected_num_scheduled_tokens=NUM_TOKENS - NUM_MATCHED_NEW_TOKENS,
2886	)
2887	assert len(scheduler.running) == 1
2888	assert len(scheduler.waiting) == 1
2889
2890	# All memory should be freed, with one request waiting.
2891	_step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
2892	assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
2893	assert len(scheduler.running) == 0
2894	assert len(scheduler.waiting) == 1
2895
2896	# Just one request should be running.
2897	output = scheduler.schedule()
2898	_assert_right_scheduler_output(
2899	output,
2900	num_requests=1,
2901	expected_num_scheduled_tokens=NUM_TOKENS - NUM_MATCHED_NEW_TOKENS,
2902	)
2903	assert len(scheduler.running) == 1
2904	assert len(scheduler.waiting) == 0
2905
2906	# All memory should be freed, with no requests waiting / running.
2907	_step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
2908	assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
2909	assert len(scheduler.running) == 0
2910	assert len(scheduler.waiting) == 0
2911
2912
2913	@pytest.mark.parametrize("use_v2_model_runner", [False, True])
2914	@pytest.mark.parametrize("is_async", [False, True])
2915	@pytest.mark.parametrize(
2916	"use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
2917	)
2918	def test_kv_connector_handles_preemption(
2919	is_async, use_ec_connector, ec_role, use_v2_model_runner
2920	):
2921	"""Test whether scheduler with KVConnector is able to handle
2922	unable to allocate (run out of blocks in allocate_slots().
2923	"""
2924	# Setup Scheduler With Mock External Cache Hit.
2925	BLOCK_SIZE = 2
2926	# NOTE: there is 1 null block, so this is 6 blocks.
2927	NUM_BLOCKS = 7
2928	NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE
2929	scheduler = create_scheduler(
2930	enable_prefix_caching=True,
2931	use_kv_connector=mock_kv(
2932	matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async
2933	),
2934	block_size=BLOCK_SIZE,
2935	num_blocks=NUM_BLOCKS,
2936	# encoder connector should not affect test results
2937	use_ec_connector=use_ec_connector,
2938	ec_role=ec_role,
2939	use_v2_model_runner=use_v2_model_runner,
2940	)
2941
2942	# Create two requests.
2943	# Both can be scheduled at first, but the second request
2944	# will be preempted and re-scheduled.
2945	NUM_REQUESTS = 2
2946	NUM_TOKENS = BLOCK_SIZE * 2 + 1
2947	MAX_TOKENS = BLOCK_SIZE * 2
2948	requests = create_requests(
2949	num_requests=NUM_REQUESTS,
2950	num_tokens=NUM_TOKENS,
2951	max_tokens=MAX_TOKENS,
2952	block_size=BLOCK_SIZE,
2953	)
2954	req_ids = []
2955	req_to_index = {}
2956	for i, request in enumerate(requests):
2957	scheduler.add_request(request)
2958	req_ids.append(request.request_id)
2959	req_to_index[request.request_id] = i
2960
2961	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
2962	req_ids=req_ids,
2963	req_id_to_index=req_to_index,
2964	sampled_token_ids=[[1000]] * len(req_ids),
2965	logprobs=None,
2966	prompt_logprobs_dict={},
2967	pooler_output=[],
2968	)
2969
2970	# All can be scheduled - 1st token.
2971	output = scheduler.schedule()
2972	if is_async:
2973	assert _num_waiting_requests(scheduler) == 2
2974	assert scheduler.running == []
2975	_step_until_kv_transfer_finished(scheduler, req_ids)
2976	output = scheduler.schedule()
2977
2978	_assert_right_scheduler_output(
2979	output,
2980	# 2 remote kv cache hits.
2981	num_requests=2,
2982	expected_num_scheduled_tokens=NUM_TOKENS - NUM_MATCHED_NEW_TOKENS,
2983	)
2984	assert len(scheduler.running) == 2
2985	_ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
2986
2987	# All can be scheduled - 2nd token.
2988	output = scheduler.schedule()
2989	_assert_right_scheduler_output(
2990	output,
2991	# no connector_metadata
2992	num_requests=0,
2993	expected_num_scheduled_tokens=1,
2994	)
2995	assert len(scheduler.running) == 2
2996	_ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
2997
2998	# This will generate a new block and cause a preemption - 3rd token.
2999	output = scheduler.schedule()
3000	_assert_right_scheduler_output(
3001	output,
3002	# no connector_metadata
3003	num_requests=0,
3004	expected_num_scheduled_tokens=1,
3005	)
3006	assert len(scheduler.running) == 1
3007	assert len(scheduler.waiting) == 1
3008	_ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
3009	assert len(scheduler.running) == 1
3010	assert len(scheduler.waiting) == 1
3011
3012	# Only 1 can be scheduled - 4th (and last token).
3013	output = scheduler.schedule()
3014	_assert_right_scheduler_output(
3015	output,
3016	# no connector_metadata
3017	num_requests=0,
3018	expected_num_scheduled_tokens=1,
3019	)
3020	assert len(scheduler.waiting) == 1
3021	assert len(scheduler.running) == 1
3022	_ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
3023	assert len(scheduler.running) == 0
3024	# All memory should be freed since nothing is running.
3025	assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
3026
3027	# Restarts the preempted request - generate 3rd token.
3028	# This will have a local and remote cache hit.
3029	output = scheduler.schedule()
3030	if is_async:
3031	waiting_req_ids = [
3032	req.request_id
3033	for req in scheduler.kv_holding_waiting
3034	if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
3035	]
3036	assert len(waiting_req_ids) == 1
3037	_step_until_kv_transfer_finished(scheduler, waiting_req_ids)
3038	output = scheduler.schedule()
3039
3040	_assert_right_scheduler_output(
3041	output,
3042	# 1 remote kv_cache hit!
3043	num_requests=1,
3044	# Only 1 block was preempted and there is a single
3045	# remote hit. So only single new token is scheduled.
3046	expected_num_scheduled_tokens=1,
3047	)
3048	assert len(scheduler.running) == 1
3049	assert len(scheduler.waiting) == 0
3050	if use_v2_model_runner:
3051	# V2 emits a resumed (previously preempted) request as a
3052	# NewRequestData rather than a cached request.
3053	assert output.scheduled_cached_reqs.num_reqs == 0
3054	assert len(output.scheduled_new_reqs) == 1
3055	else:
3056	assert output.scheduled_cached_reqs.num_reqs == 1
3057	assert output.scheduled_new_reqs == []
3058	_ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
3059	assert len(scheduler.running) == 1
3060	assert len(scheduler.waiting) == 0
3061
3062	# Only 1 can be scheduled - 4th (and last token).
3063	output = scheduler.schedule()
3064	_assert_right_scheduler_output(
3065	output,
3066	# no connector_metadata
3067	num_requests=0,
3068	expected_num_scheduled_tokens=1,
3069	)
3070	assert output.scheduled_cached_reqs.num_reqs == 1
3071	assert output.scheduled_new_reqs == []
3072	assert len(scheduler.running) == 1
3073	_ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
3074	assert len(scheduler.running) == 0
3075	# All memory should be freed since nothing is running.
3076	assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
3077
3078
3079	def make_output(scheduler: Scheduler):
3080	return ModelRunnerOutput(
3081	req_ids=[req.request_id for req in scheduler.running],
3082	req_id_to_index={req.request_id: i for i, req in enumerate(scheduler.running)},
3083	sampled_token_ids=[[1000]] * len(scheduler.running),
3084	logprobs=None,
3085	prompt_logprobs_dict={},
3086	pooler_output=[],
3087	)
3088
3089
3090	def assert_scheduler_empty(scheduler: Scheduler):
3091	"""Confirm the scheduler is "empty" - i.e. no leaks."""
3092	# Scheduler Metadata.
3093	assert len(scheduler.requests) == 0
3094	assert len(scheduler.waiting) == 0
3095	assert len(scheduler.running) == 0
3096	assert len(scheduler.finished_req_ids) == 0
3097
3098	# EncoderCacheManager.
3099	assert len(scheduler.encoder_cache_manager.freed) == 0
3100	assert len(scheduler.encoder_cache_manager.cached) == 0
3101
3102	# KVCache Manager.
3103	assert (
3104	len(
3105	scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks
3106	)
3107	== 0
3108	)
3109	assert (
3110	len(
3111	scheduler.kv_cache_manager.coordinator.single_type_managers[
3112	0
3113	].num_cached_block
3114	)
3115	== 0
3116	)
3117	num_free_blocks = (
3118	scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks
3119	)
3120	assert num_free_blocks == (scheduler.kv_cache_manager.block_pool.num_gpu_blocks - 1)
3121
3122	# NOTE(rob): just the ref count on blocks will be 0. The hash
3123	# value, etc will remain since we lazily evict for prefix cache.
3124	for block in scheduler.kv_cache_manager.block_pool.blocks:
3125	assert block.ref_cnt == 0
3126	# assert block._block_hash is None
3127	# assert (
3128	# len(scheduler.kv_cache_manager.block_pool.cached_block_hash_to_block
3129	# ) == 0)
3130
3131
3132	def test_memory_leak():
3133	"""Test that we do not have a memory leak."""
3134	scheduler = create_scheduler(enable_prefix_caching=True)
3135
3136	NUM_REQUESTS = 5
3137	NUM_TOKENS = 10
3138	MAX_TOKENS = 10
3139	requests = create_requests(
3140	num_requests=NUM_REQUESTS, num_tokens=NUM_TOKENS, max_tokens=MAX_TOKENS
3141	)
3142
3143	# Add each request.
3144	for request in requests:
3145	scheduler.add_request(request)
3146	scheduler_output = scheduler.schedule()
3147	model_runner_output = make_output(scheduler)
3148	scheduler.update_from_output(scheduler_output, model_runner_output)
3149
3150	# Iterate until done.
3151	while True:
3152	scheduler_output = scheduler.schedule()
3153	if len(scheduler.running) == 0:
3154	break
3155	model_runner_output = make_output(scheduler)
3156	scheduler.update_from_output(scheduler_output, model_runner_output)
3157
3158	# Confirm no memory leak.
3159	assert_scheduler_empty(scheduler)
3160
3161
3162	def create_scheduler_with_priority(
3163	model: str = "facebook/opt-125m",
3164	max_num_seqs: int = 16,
3165	max_num_batched_tokens: int = 8192,
3166	enable_prefix_caching: bool = False,
3167	long_prefill_token_threshold: int = 0,
3168	disable_chunked_mm_input: bool = False,
3169	use_kv_connector: bool = False,
3170	num_blocks: int = 10000,
3171	block_size: int = 16,
3172	max_model_len: int \| None = None,
3173	num_speculative_tokens: int \| None = None,
3174	use_ec_connector: bool = False,
3175	ec_role: str \| None = None,
3176	use_v2_model_runner: bool \| None = None,
3177	) -> Scheduler:
3178	"""Create scheduler with priority policy enabled.
3179
3180	Args:
3181	model: model under test
3182	max_num_seqs: max sequences to schedule
3183	max_num_batch_tokens: max num tokens to batch
3184	enable_prefix_caching: optionally force APC config
3185	(True/False) or use default
3186	(False)
3187
3188	Returns:
3189	{class}`Scheduler` instance with priority scheduling
3190
3191	"""
3192	model_config = ModelConfig(
3193	model=model,
3194	trust_remote_code=True,
3195	dtype="float16",
3196	seed=42,
3197	)
3198	if max_model_len is None:
3199	max_model_len = max_num_batched_tokens
3200	scheduler_config = SchedulerConfig(
3201	max_num_seqs=max_num_seqs,
3202	max_num_batched_tokens=max_num_batched_tokens,
3203	max_model_len=max_model_len,
3204	long_prefill_token_threshold=long_prefill_token_threshold,
3205	disable_chunked_mm_input=disable_chunked_mm_input,
3206	enable_chunked_prefill=True,
3207	is_encoder_decoder=model_config.is_encoder_decoder,
3208	policy="priority", # Enable priority scheduling
3209	# Ensure admission/preemption mechanics are deterministic
3210	watermark=0.0,
3211	)
3212	# Cache config, optionally force APC
3213	cache_config = CacheConfig(
3214	block_size=block_size,
3215	gpu_memory_utilization=0.9,
3216	cache_dtype="auto",
3217	enable_prefix_caching=enable_prefix_caching,
3218	)
3219	kv_transfer_config = (
3220	KVTransferConfig(
3221	kv_connector="ExampleConnector",
3222	kv_role="kv_both",
3223	kv_connector_extra_config={"shared_storage_path": "local_storage"},
3224	)
3225	if use_kv_connector
3226	else None
3227	)
3228
3229	speculative_config: SpeculativeConfig \| None = None
3230	if num_speculative_tokens is not None:
3231	speculative_config = SpeculativeConfig(
3232	model="ngram", num_speculative_tokens=num_speculative_tokens
3233	)
3234
3235	ec_transfer_config = (
3236	ECTransferConfig(
3237	ec_connector="ECExampleConnector",
3238	ec_role=ec_role,
3239	ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test"},
3240	)
3241	if use_ec_connector
3242	else None
3243	)
3244
3245	vllm_config = VllmConfig(
3246	scheduler_config=scheduler_config,
3247	model_config=model_config,
3248	cache_config=cache_config,
3249	kv_transfer_config=kv_transfer_config,
3250	speculative_config=speculative_config,
3251	ec_transfer_config=ec_transfer_config,
3252	)
3253	kv_cache_config = KVCacheConfig(
3254	num_blocks=num_blocks, # A large number of blocks to hold all requests
3255	kv_cache_tensors=[],
3256	kv_cache_groups=[
3257	KVCacheGroupSpec(
3258	["layer"],
3259	FullAttentionSpec(
3260	block_size=block_size,
3261	num_kv_heads=1,
3262	head_size=1,
3263	dtype=torch.float32,
3264	),
3265	)
3266	],
3267	)
3268	cache_config.num_gpu_blocks = num_blocks
3269	scheduler = Scheduler(
3270	vllm_config=vllm_config,
3271	kv_cache_config=kv_cache_config,
3272	log_stats=True,
3273	structured_output_manager=StructuredOutputManager(vllm_config),
3274	block_size=block_size,
3275	hash_block_size=block_size,
3276	)
3277	if use_v2_model_runner is None:
3278	use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER)
3279	scheduler.use_v2_model_runner = use_v2_model_runner
3280	return scheduler
3281
3282
3283	_none_hash_initialized = False
3284
3285
3286	def create_requests_with_priority(
3287	num_requests: int,
3288	priorities: list[int],
3289	arrival_times: list[float] \| None = None,
3290	num_tokens: int = 10,
3291	mm_hashes_list: list[list[str]] \| None = None,
3292	mm_positions: list[list[PlaceholderRange]] \| None = None,
3293	max_tokens: int = 16,
3294	stop_token_ids: list[int] \| None = None,
3295	prompt_logprobs: int \| None = None,
3296	starting_idx: int = 0,
3297	same_prompt: bool = False,
3298	block_size: int = 16,
3299	req_ids: list[str] \| None = None,
3300	):
3301	"""Create requests with specified priorities and arrival times."""
3302	assert len(priorities) == num_requests
3303	if arrival_times is not None:
3304	assert len(arrival_times) == num_requests
3305	else:
3306	arrival_times = [float(i) for i in range(num_requests)]
3307
3308	global _none_hash_initialized
3309	if not _none_hash_initialized:
3310	init_none_hash(sha256)
3311	_none_hash_initialized = True
3312
3313	block_hasher = get_request_block_hasher(block_size, sha256)
3314	sampling_params = SamplingParams(
3315	ignore_eos=False,
3316	max_tokens=max_tokens,
3317	stop_token_ids=stop_token_ids,
3318	prompt_logprobs=prompt_logprobs,
3319	)
3320	sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
3321	requests = []
3322
3323	if mm_hashes_list is not None:
3324	# NOTE: allow manual input; some mm items can have the same identifier
3325	# no. of mm_hashes and mm_positions for each request should be identical
3326	assert mm_positions is not None, (
3327	"mm_positions must be provided when mm_hashes_list is provided"
3328	)
3329	assert len(mm_hashes_list) == len(mm_positions) == num_requests
3330	assert [len(h) for h in mm_hashes_list] == [len(p) for p in mm_positions]
3331
3332	# Since same identifier would imply they are identical encoder output
3333	# Verify mm items with identical identifier are having mm_position.length
3334	seen_hashes: dict[str, int] = {}
3335
3336	if req_ids:
3337	assert len(req_ids) == num_requests
3338	else:
3339	req_ids = [f"{i + starting_idx}" for i in range(num_requests)]
3340
3341	for i in range(num_requests):
3342	mm_features = []
3343
3344	for j, position in enumerate(
3345	mm_positions[i] if mm_positions is not None else []
3346	):
3347	if mm_hashes_list is not None:
3348	identifier = mm_hashes_list[i][j]
3349
3350	# Verify if position length is identical
3351	position_length = position.length
3352	if identifier in seen_hashes:
3353	assert seen_hashes[identifier] == position_length, (
3354	f"mm_hash '{identifier}' has inconsistent position lengths: "
3355	f"previously {seen_hashes[identifier]}, now {position_length} "
3356	f"at request {i}, position {j}"
3357	)
3358	else:
3359	seen_hashes[identifier] = position_length
3360	else:
3361	# Unique dummy hash for each mm item
3362	identifier = f"hash{i}_{j}"
3363	mm_feature = MultiModalFeatureSpec(
3364	data=MultiModalKwargsItem.dummy(),
3365	mm_position=position,
3366	identifier=identifier,
3367	modality="image",
3368	)
3369	mm_features.append(mm_feature)
3370
3371	prompt_token_ids = (
3372	[starting_idx] * num_tokens
3373	if same_prompt
3374	else [i + starting_idx] * num_tokens
3375	)
3376	request = Request(
3377	request_id=req_ids[i],
3378	prompt_token_ids=prompt_token_ids,
3379	sampling_params=sampling_params,
3380	pooling_params=None,
3381	mm_features=mm_features if mm_features else None,
3382	arrival_time=arrival_times[i],
3383	priority=priorities[i],
3384	block_hasher=block_hasher,
3385	)
3386	requests.append(request)
3387	return requests
3388
3389
3390	def test_priority_scheduling_basic_ordering():
3391	"""Test that requests are scheduled in priority order
3392	(lower value = higher priority)."""
3393	scheduler = create_scheduler_with_priority()
3394
3395	# Create requests with different priorities
3396	# Priority 0 (highest), 1, 2 (lowest)
3397	priorities = [2, 0, 1] # Add in non-priority order
3398	arrival_times = [1.0, 2.0, 3.0] # All different arrival times
3399	requests = create_requests_with_priority(
3400	num_requests=3, priorities=priorities, arrival_times=arrival_times
3401	)
3402
3403	# Add requests in non-priority order
3404	for request in requests:
3405	scheduler.add_request(request)
3406
3407	# Schedule and verify priority order
3408	output = scheduler.schedule()
3409
3410	# Should schedule all requests since they fit in budget
3411	assert len(output.scheduled_new_reqs) == 3
3412
3413	# Verify they are scheduled in priority order:
3414	# req_1 (priority 0), req_2 (priority 1), req_0 (priority 2)
3415	scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
3416	assert scheduled_req_ids == ["1", "2", "0"]
3417
3418
3419	def test_priority_scheduling_arrival_time_tiebreaker():
3420	"""Test that arrival time is used
3421	as tiebreaker when priorities are equal."""
3422	scheduler = create_scheduler_with_priority()
3423
3424	# Create requests with same priority but different arrival times
3425	priorities = [1, 1, 1] # All same priority
3426	arrival_times = [3.0, 1.0, 2.0] # Different arrival times
3427	requests = create_requests_with_priority(
3428	num_requests=3, priorities=priorities, arrival_times=arrival_times
3429	)
3430
3431	# Add requests in non-arrival order
3432	for request in requests:
3433	scheduler.add_request(request)
3434
3435	# Schedule and verify arrival time order
3436	output = scheduler.schedule()
3437
3438	# Should schedule all requests since they fit in budget
3439	assert len(output.scheduled_new_reqs) == 3
3440
3441	# Verify they are scheduled in arrival time order:
3442	# req_1 (1.0), req_2 (2.0), req_0 (3.0)
3443	scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
3444	assert scheduled_req_ids == ["1", "2", "0"]
3445
3446
3447	def test_priority_scheduling_mixed_priority_and_arrival():
3448	"""Test priority scheduling with mixed priorities and arrival times."""
3449	scheduler = create_scheduler_with_priority()
3450
3451	# Create requests with mixed priorities and arrival times
3452	priorities = [2, 1, 1, 0] # Mixed priorities
3453	arrival_times = [1.0, 3.0, 2.0, 4.0] # Mixed arrival times
3454	requests = create_requests_with_priority(
3455	num_requests=4, priorities=priorities, arrival_times=arrival_times
3456	)
3457
3458	# Add requests
3459	for request in requests:
3460	scheduler.add_request(request)
3461
3462	# Schedule and verify order
3463	output = scheduler.schedule()
3464
3465	# Should schedule all requests since they fit in budget
3466	assert len(output.scheduled_new_reqs) == 4
3467
3468	# Expected order:
3469	# 1. req_3 (priority 0, arrival 4.0)
3470	# 2. req_2 (priority 1, arrival 2.0) - earlier arrival than req_1
3471	# 3. req_1 (priority 1, arrival 3.0)
3472	# 4. req_0 (priority 2, arrival 1.0)
3473	scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
3474	assert scheduled_req_ids == ["3", "2", "1", "0"]
3475
3476
3477	def test_priority_scheduling_preemption():
3478	"""Test that under KV block pressure the scheduler preempts the
3479	lowest-priority running request, not the highest-priority one.
3480
3481	A low-priority request starts running first. Then a high-priority
3482	request arrives and is admitted to running. When block pressure
3483	builds, the scheduler preempts the low-priority running request
3484	while keeping the high-priority one.
3485
3486	Block math
3487	----------
3488	block_size = 16, num_blocks = 6 (1 null → 5 usable).
3489
3490	Phase 1: lo1 (priority 5, 32 tokens) → 2 blocks. 3 free.
3491	Decode → lo1 has 33 tokens (needs 3rd block on next schedule).
3492	Phase 2: hi1 (priority 0, 32 tokens) arrives.
3493	schedule() allocates lo1's 3rd block (3 used) and admits
3494	hi1 (2 blocks) → 5 used, 0 free. Both running.
3495	Decode → lo1 34 tokens, hi1 33 tokens.
3496	Phase 3: schedule() → hi1 needs 3rd block, 0 free → preemption.
3497	lo1 (priority 5) is preempted, hi1 (priority 0) survives.
3498	"""
3499	block_size = 16
3500	num_blocks = 6 # 1 null block → 5 usable
3501	num_tokens = block_size * 2 # 32 tokens = exactly 2 blocks
3502
3503	scheduler = create_scheduler_with_priority(
3504	max_num_seqs=3,
3505	max_num_batched_tokens=200,
3506	num_blocks=num_blocks,
3507	block_size=block_size,
3508	)
3509
3510	# --- Phase 1: low-priority request starts running ---
3511	lo1 = create_requests_with_priority(
3512	num_requests=1,
3513	priorities=[5],
3514	arrival_times=[1.0],
3515	num_tokens=num_tokens,
3516	req_ids=["lo1"],
3517	)[0]
3518	scheduler.add_request(lo1)
3519
3520	output = scheduler.schedule()
3521	assert len(output.scheduled_new_reqs) == 1
3522
3523	# Decode: lo1 now has 33 tokens (crosses 32-token boundary).
3524	model_output = ModelRunnerOutput(
3525	req_ids=["lo1"],
3526	req_id_to_index={"lo1": 0},
3527	sampled_token_ids=[[100]],
3528	logprobs=None,
3529	prompt_logprobs_dict={},
3530	pooler_output=[],
3531	)
3532	scheduler.update_from_output(output, model_output)
3533
3534	# --- Phase 2: high-priority request arrives AFTER lo1 is running ---
3535	hi1 = create_requests_with_priority(
3536	num_requests=1,
3537	priorities=[0],
3538	arrival_times=[2.0],
3539	num_tokens=num_tokens,
3540	req_ids=["hi1"],
3541	)[0]
3542	scheduler.add_request(hi1)
3543
3544	# schedule(): lo1 gets its 3rd block (3 used), hi1 admitted (5 used,
3545	# 0 free). Both are now running.
3546	output = scheduler.schedule()
3547	assert any(r.req_id == "hi1" for r in output.scheduled_new_reqs)
3548	assert len(scheduler.running) == 2
3549
3550	# Decode: lo1 → 34 tokens, hi1 → 33 tokens.
3551	model_output = ModelRunnerOutput(
3552	req_ids=["lo1", "hi1"],
3553	req_id_to_index={"lo1": 0, "hi1": 1},
3554	sampled_token_ids=[[101], [100]],
3555	logprobs=None,
3556	prompt_logprobs_dict={},
3557	pooler_output=[],
3558	)
3559	scheduler.update_from_output(output, model_output)
3560
3561	# --- Phase 3: preemption with mixed-priority running requests ---
3562	# hi1 needs a 3rd block but 0 are free. The scheduler picks the
3563	# lowest-priority running request to preempt:
3564	# max(running, key=(priority, arrival_time)) → lo1 (5 > 0).
3565	output = scheduler.schedule()
3566
3567	lo1_req = scheduler.requests["lo1"]
3568	assert lo1_req.status == RequestStatus.PREEMPTED, (
3569	"Expected low-priority 'lo1' to be preempted"
3570	)
3571	assert any(req.request_id == "hi1" for req in scheduler.running), (
3572	"High-priority 'hi1' should still be running"
3573	)
3574
3575
3576	def test_priority_scheduling_no_preemption_when_space_available():
3577	"""Test that preemption doesn't happen
3578	when there's space for new requests."""
3579	scheduler = create_scheduler_with_priority(
3580	max_num_seqs=3, # Allow 3 concurrent requests
3581	max_num_batched_tokens=200, # Sufficient token budget
3582	)
3583
3584	# Add two low-priority running requests
3585	low_priority_requests = create_requests_with_priority(
3586	num_requests=2,
3587	priorities=[5, 5],
3588	arrival_times=[1.0, 2.0],
3589	num_tokens=30,
3590	req_ids=["lo1", "lo2"],
3591	)
3592
3593	for request in low_priority_requests:
3594	scheduler.add_request(request)
3595
3596	output = scheduler.schedule()
3597	model_output = ModelRunnerOutput(
3598	req_ids=[req.request_id for req in low_priority_requests],
3599	req_id_to_index={
3600	req.request_id: i for i, req in enumerate(low_priority_requests)
3601	},
3602	sampled_token_ids=[[100] for _ in low_priority_requests],
3603	logprobs=None,
3604	prompt_logprobs_dict={},
3605	pooler_output=[],
3606	)
3607	scheduler.update_from_output(output, model_output)
3608
3609	# Add high-priority request
3610	high_priority_request = create_requests_with_priority(
3611	num_requests=1,
3612	priorities=[0],
3613	arrival_times=[3.0],
3614	num_tokens=30,
3615	req_ids=["hi1"],
3616	)[0]
3617
3618	scheduler.add_request(high_priority_request)
3619
3620	# Schedule - should not preempt since there's space
3621	output = scheduler.schedule()
3622
3623	# Should schedule the new request without preemption
3624	assert len(output.scheduled_new_reqs) == 1
3625	assert len(scheduler.running) == 3 # All three requests running
3626	assert len(scheduler.waiting) == 0 # No requests waiting
3627
3628
3629	def test_priority_scheduling_preemption_victim_selection():
3630	"""Test that the correct victim is selected for
3631	preemption based on priority and arrival time."""
3632	# This test verifies the priority-based victim selection logic
3633	# by checking the waiting queue order after adding requests with different
3634	# priorities
3635	scheduler = create_scheduler_with_priority(
3636	max_num_seqs=1, # Force sequential processing to test priority order
3637	)
3638
3639	# Create requests with different priorities
3640	requests = create_requests_with_priority(
3641	num_requests=3,
3642	priorities=[3, 2, 0], # Different priorities: low, medium, high
3643	arrival_times=[1.0, 2.0, 3.0],
3644	num_tokens=10,
3645	)
3646
3647	# Add all requests
3648	for request in requests:
3649	scheduler.add_request(request)
3650
3651	# Schedule - should only schedule the highest priority request
3652	# (req_2, priority 0)
3653	output = scheduler.schedule()
3654	assert len(output.scheduled_new_reqs) == 1
3655	assert output.scheduled_new_reqs[0].req_id == "2" # Highest priority
3656
3657	# Verify the waiting queue has the remaining requests in priority order
3658	assert len(scheduler.waiting) == 2
3659
3660	# Extract waiting requests and verify priority order
3661	waiting_requests = list(scheduler.waiting)
3662
3663	waiting_priorities = [req.priority for req in waiting_requests]
3664	waiting_req_ids = [req.request_id for req in waiting_requests]
3665
3666	# Should be req_1 (priority 2) then req_0 (priority 3)
3667	assert waiting_priorities == [2, 3]
3668	assert waiting_req_ids == ["1", "0"]
3669
3670
3671	def test_priority_scheduling_equal_priority_preemption():
3672	"""Test arrival time tiebreaker when requests have equal priority."""
3673	# This test verifies that arrival time is used as a tiebreaker for equal
3674	# priorities
3675	scheduler = create_scheduler_with_priority(
3676	max_num_seqs=1, # Force sequential processing
3677	)
3678
3679	# Create requests with same priority but different arrival times
3680	requests = create_requests_with_priority(
3681	num_requests=3,
3682	priorities=[2, 2, 2], # Same priority
3683	arrival_times=[3.0, 1.0, 2.0], # Different arrival times
3684	num_tokens=10,
3685	)
3686
3687	# Add all requests
3688	for request in requests:
3689	scheduler.add_request(request)
3690
3691	# Schedule - should schedule the request with earliest arrival time
3692	output = scheduler.schedule()
3693	assert len(output.scheduled_new_reqs) == 1
3694	assert output.scheduled_new_reqs[0].req_id == "1" # Earliest arrival (1.0)
3695
3696	# Verify the waiting queue has remaining requests in arrival time order
3697	assert len(scheduler.waiting) == 2
3698
3699	# Extract waiting requests and verify arrival time order
3700	waiting_requests = list(scheduler.waiting)
3701
3702	waiting_arrival_times = [req.arrival_time for req in waiting_requests]
3703	waiting_req_ids = [req.request_id for req in waiting_requests]
3704
3705	# Should be req_2 (arrival 2.0) then req_0 (arrival 3.0)
3706	assert waiting_arrival_times == [2.0, 3.0]
3707	assert waiting_req_ids == ["2", "0"]
3708
3709
3710	def test_priority_scheduling_waiting_queue_order():
3711	"""Test that the waiting queue maintains priority order."""
3712	scheduler = create_scheduler_with_priority(
3713	max_num_seqs=1, # Only one request can run at a time
3714	)
3715
3716	# Create multiple requests with different priorities
3717	requests = create_requests_with_priority(
3718	num_requests=4,
3719	priorities=[3, 1, 2, 0], # Mixed priorities
3720	arrival_times=[1.0, 2.0, 3.0, 4.0],
3721	num_tokens=10,
3722	)
3723
3724	# Add all requests
3725	for request in requests:
3726	scheduler.add_request(request)
3727
3728	# Schedule - should only schedule the highest priority request
3729	# (req_3, priority 0)
3730	output = scheduler.schedule()
3731	assert len(output.scheduled_new_reqs) == 1
3732	assert output.scheduled_new_reqs[0].req_id == "3"
3733
3734	# Verify waiting queue has remaining requests in priority order
3735	assert len(scheduler.waiting) == 3
3736
3737	# Extract requests from waiting queue
3738	# (it's a heap, so we need to pop to see order)
3739	waiting_requests = list(scheduler.waiting)
3740
3741	waiting_priorities = [req.priority for req in waiting_requests]
3742	waiting_req_ids = [req.request_id for req in waiting_requests]
3743
3744	# Should be ordered by priority: req_1 (1), req_2 (2), req_0 (3)
3745	assert waiting_req_ids == ["1", "2", "0"]
3746	assert waiting_priorities == [1, 2, 3]
3747
3748
3749	def test_priority_scheduling_fcfs_fallback():
3750	"""Test that FCFS behavior is maintained when all
3751	requests have same priority."""
3752	scheduler = create_scheduler_with_priority()
3753
3754	# Create requests with same priority but different arrival times
3755	priorities = [1, 1, 1, 1] # All same priority
3756	arrival_times = [4.0, 1.0, 3.0, 2.0] # Different arrival times
3757	requests = create_requests_with_priority(
3758	num_requests=4, priorities=priorities, arrival_times=arrival_times
3759	)
3760
3761	# Add requests
3762	for request in requests:
3763	scheduler.add_request(request)
3764
3765	# Schedule
3766	output = scheduler.schedule()
3767
3768	# Should schedule all requests in arrival time order
3769	assert len(output.scheduled_new_reqs) == 4
3770	scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
3771
3772	# Expected order by arrival time:
3773	# req_1 (1.0), req_3 (2.0), req_2 (3.0), req_0 (4.0)
3774	assert scheduled_req_ids == ["1", "3", "2", "0"]
3775
3776
3777	def test_priority_scheduling_with_limited_slots():
3778	"""Test priority scheduling when max_num_seqs limits concurrent requests."""
3779	scheduler = create_scheduler_with_priority(
3780	max_num_seqs=2, # Only allow 2 concurrent requests
3781	max_num_batched_tokens=1000, # Plenty of token budget
3782	)
3783
3784	# Create requests with different priorities
3785	requests = create_requests_with_priority(
3786	num_requests=4,
3787	priorities=[3, 1, 2, 0], # Mixed priorities
3788	arrival_times=[1.0, 2.0, 3.0, 4.0],
3789	num_tokens=10,
3790	)
3791
3792	# Add all requests
3793	for request in requests:
3794	scheduler.add_request(request)
3795
3796	# Schedule - should only schedule the 2 highest priority requests
3797	output = scheduler.schedule()
3798	assert len(output.scheduled_new_reqs) == 2
3799
3800	# Should schedule req_3 (priority 0) and req_1 (priority 1)
3801	scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
3802	assert "3" in scheduled_req_ids # Priority 0
3803	assert "1" in scheduled_req_ids # Priority 1
3804
3805	# Remaining requests should be in waiting queue in priority order
3806	assert len(scheduler.waiting) == 2
3807
3808	# Extract waiting requests and verify order
3809	waiting_requests = list(scheduler.waiting)
3810	waiting_priorities = [req.priority for req in waiting_requests]
3811	waiting_req_ids = [req.request_id for req in waiting_requests]
3812
3813	# Should be req_2 (priority 2) then req_0 (priority 3)
3814	assert waiting_priorities == [2, 3]
3815	assert waiting_req_ids == ["2", "0"]
3816
3817
3818	def test_priority_scheduling_heap_property():
3819	"""Test that the waiting queue maintains heap
3820	property for priority scheduling."""
3821	scheduler = create_scheduler_with_priority(
3822	max_num_seqs=1, # Only one request can run at a time
3823	)
3824
3825	# Add requests in random priority order
3826	priorities = [5, 1, 8, 3, 2, 7, 4, 6]
3827	arrival_times = [float(i) for i in range(len(priorities))]
3828	requests = create_requests_with_priority(
3829	num_requests=len(priorities),
3830	priorities=priorities,
3831	arrival_times=arrival_times,
3832	num_tokens=10,
3833	)
3834
3835	# Add all requests
3836	for request in requests:
3837	scheduler.add_request(request)
3838
3839	# Schedule one request at a time and verify priority order
3840	scheduled_priorities = []
3841
3842	while scheduler.waiting:
3843	output = scheduler.schedule()
3844	if output.scheduled_new_reqs:
3845	req = output.scheduled_new_reqs[0]
3846	scheduled_priorities.append(requests[int(req.req_id)].priority)
3847
3848	# Simulate completion to make room for next request
3849	model_output = ModelRunnerOutput(
3850	req_ids=[req.req_id],
3851	req_id_to_index={req.req_id: 0},
3852	sampled_token_ids=[[100]],
3853	logprobs=None,
3854	prompt_logprobs_dict={},
3855	pooler_output=[],
3856	)
3857	scheduler.update_from_output(output, model_output)
3858
3859	# Finish the request to make room for the next one
3860	scheduler.finish_requests(req.req_id, RequestStatus.FINISHED_STOPPED)
3861
3862	# Verify requests were scheduled in priority order (lowest value first)
3863	expected_priorities = sorted(priorities)
3864	assert scheduled_priorities == expected_priorities
3865
3866
3867	def test_schedule_skip_tokenizer_init():
3868	scheduler = create_scheduler(skip_tokenizer_init=True)
3869	requests = create_requests(num_requests=5)
3870	for request in requests:
3871	scheduler.add_request(request)
3872	output = scheduler.schedule()
3873	assert len(output.scheduled_new_reqs) == len(requests)
3874
3875
3876	def test_schedule_skip_tokenizer_init_structured_output_request():
3877	scheduler = create_scheduler(skip_tokenizer_init=True)
3878	structured_outputs_params = StructuredOutputsParams(regex="[0-9]+")
3879	sampling_params = SamplingParams(
3880	ignore_eos=False,
3881	max_tokens=16,
3882	structured_outputs=structured_outputs_params,
3883	)
3884	sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
3885	request = Request(
3886	request_id="0",
3887	prompt_token_ids=[0, 1],
3888	mm_features=None,
3889	sampling_params=sampling_params,
3890	pooling_params=None,
3891	)
3892	scheduler.add_request(request)
3893	output = scheduler.schedule()
3894	assert len(output.scheduled_new_reqs) == 0
3895	assert len(scheduler.running) == 0
3896	assert len(scheduler.waiting) == 1
3897	assert len(scheduler.deferred_waiting) == 1
3898
3899
3900	@pytest.mark.parametrize("async_grammar", [True, False])
3901	def test_grammar_compile_error_finishes_only_request(async_grammar: bool):
3902	scheduler = create_scheduler()
3903	manager = scheduler.structured_output_manager
3904	manager.backend = Mock()
3905	manager.backend.compile_grammar.side_effect = RuntimeError(
3906	"forced FSM compilation error"
3907	)
3908	manager._use_async_grammar_compilation = async_grammar
3909
3910	sampling_params = SamplingParams(
3911	max_tokens=16,
3912	structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
3913	)
3914	sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
3915	request = Request(
3916	request_id="grammar-error",
3917	prompt_token_ids=[0, 1],
3918	sampling_params=sampling_params,
3919	pooling_params=None,
3920	)
3921
3922	manager.grammar_init(request)
3923	assert request.structured_output_request is not None
3924	grammar_future = request.structured_output_request._grammar
3925	assert isinstance(grammar_future, Future)
3926	assert isinstance(grammar_future.exception(timeout=5), RuntimeError)
3927
3928	scheduler.add_request(request)
3929	scheduler_output = scheduler.schedule()
3930	assert not scheduler_output.num_scheduled_tokens
3931
3932	engine_core_outputs = scheduler.update_from_output(
3933	scheduler_output,
3934	ModelRunnerOutput(req_ids=[], req_id_to_index={}),
3935	)
3936
3937	assert request.status == RequestStatus.FINISHED_ERROR
3938	assert request.request_id not in scheduler.requests
3939	output = engine_core_outputs[0].outputs[0]
3940	assert output.request_id == request.request_id
3941	assert output.finish_reason == FinishReason.ERROR
3942	assert output.stop_reason is None
3943
3944	healthy_request = create_requests(num_requests=1, req_ids=["healthy-request"])[0]
3945	scheduler.add_request(healthy_request)
3946	next_output = scheduler.schedule()
3947	assert [req.req_id for req in next_output.scheduled_new_reqs] == [
3948	healthy_request.request_id
3949	]
3950
3951
3952	def test_abort_request_when_structured_output_fsm_cannot_advance():
3953	scheduler = object.__new__(Scheduler)
3954	sampling_params = SamplingParams(ignore_eos=True, max_tokens=4)
3955	sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
3956
3957	request = Request(
3958	request_id="0",
3959	prompt_token_ids=[0, 1],
3960	mm_features=None,
3961	sampling_params=sampling_params,
3962	pooling_params=None,
3963	)
3964	request.structured_output_request = Mock()
3965	request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
3966	request.status = RequestStatus.RUNNING
3967	request.num_computed_tokens = request.num_tokens
3968
3969	scheduler.perf_metrics = None
3970	scheduler.connector = None
3971	scheduler.ec_connector = None
3972	scheduler.structured_output_manager = Mock()
3973	scheduler.structured_output_manager.accept_tokens.return_value = False
3974	scheduler.requests = {request.request_id: request}
3975	scheduler.running = [request]
3976	scheduler.waiting = Mock()
3977	scheduler.kv_cache_manager = Mock()
3978	scheduler.kv_cache_manager.take_events.return_value = None
3979	scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
3980	scheduler.kv_event_publisher = Mock()
3981	scheduler.finished_req_ids = set()
3982	scheduler.finished_req_ids_dict = None
3983	scheduler.aux_output_connector = None
3984	scheduler.grammar_compile_error_reqs = set()
3985	scheduler.encoder_cache_mismatch_reqs = set()
3986	scheduler.vllm_config = Mock()
3987	scheduler.return_sampling_mask = False
3988	scheduler.recompute_kv_load_failures = False
3989	scheduler.defer_block_free = False
3990	scheduler.make_stats = Mock(return_value=None)
3991	scheduler.max_model_len = 128
3992
3993	def free_request(req: Request, delay_free_blocks: bool = False):
3994	scheduler.finished_req_ids.add(req.request_id)
3995	scheduler.requests.pop(req.request_id, None)
3996	return None, None
3997
3998	scheduler._free_request = Mock(side_effect=free_request)
3999
4000	output = SchedulerOutput(
4001	scheduled_new_reqs=[],
4002	scheduled_cached_reqs=CachedRequestData.make_empty(),
4003	num_scheduled_tokens={request.request_id: 1},
4004	total_num_scheduled_tokens=1,
4005	scheduled_encoder_inputs={},
4006	scheduled_spec_decode_tokens={},
4007	num_common_prefix_blocks=[],
4008	finished_req_ids=set(),
4009	free_encoder_mm_hashes=[],
4010	)
4011
4012	model_runner_output = ModelRunnerOutput(
4013	req_ids=[request.request_id],
4014	req_id_to_index={request.request_id: 0},
4015	sampled_token_ids=[[123]],
4016	logprobs=None,
4017	prompt_logprobs_dict={},
4018	pooler_output=[],
4019	)
4020	engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
4021
4022	scheduler.structured_output_manager.accept_tokens.assert_called_once_with(
4023	request, [123]
4024	)
4025	assert request.resumable is False
4026	assert request.status == RequestStatus.FINISHED_ERROR
4027	assert request.request_id not in scheduler.requests
4028	assert not scheduler.running
4029	scheduler._free_request.assert_called_once_with(request)
4030	assert len(engine_core_outputs[0].outputs) == 1
4031	engine_core_output = engine_core_outputs[0].outputs[0]
4032	assert engine_core_output.request_id == request.request_id
4033	assert engine_core_output.new_token_ids == [123]
4034	assert engine_core_output.finish_reason == FinishReason.ERROR
4035
4036
4037	@pytest.mark.parametrize("use_v2_model_runner", [False, True])
4038	@pytest.mark.parametrize(
4039	"use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
4040	)
4041	def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
4042	use_ec_connector, ec_role, use_v2_model_runner
4043	):
4044	"""Test that priority scheduling preempts lower priority requests
4045	when out of KV cache space."""
4046	# Create scheduler with very limited memory to force preemption
4047	scheduler = create_scheduler_with_priority(
4048	max_num_seqs=2, # Allow multiple requests
4049	max_num_batched_tokens=200,
4050	num_blocks=5, # Can hold 64 tokens (first block is null)
4051	block_size=16, # Standard block size
4052	use_kv_connector=True,
4053	# encoder connector should not affect test results
4054	use_ec_connector=use_ec_connector,
4055	ec_role=ec_role,
4056	use_v2_model_runner=use_v2_model_runner,
4057	)
4058
4059	# Create a request and schedule it
4060	request_low = create_requests_with_priority(
4061	num_requests=1,
4062	priorities=[1],
4063	arrival_times=[0.0],
4064	num_tokens=30,
4065	starting_idx=0,
4066	)[0]
4067	scheduler.add_request(request_low)
4068	# 1st schedule
4069	output = scheduler.schedule()
4070	assert len(output.scheduled_new_reqs) == 1
4071	assert len(scheduler.waiting) == 0
4072	assert len(scheduler.running) == 1
4073
4074	# Simulate model execution - 1st decode
4075	model_output = ModelRunnerOutput(
4076	req_ids=[request_low.request_id],
4077	req_id_to_index={request_low.request_id: 0},
4078	sampled_token_ids=[[100]],
4079	# spec_token_ids=None,
4080	logprobs=None,
4081	prompt_logprobs_dict={},
4082	pooler_output=[],
4083	)
4084	scheduler.update_from_output(output, model_output)
4085
4086	# Create a high priority request and schedule it
4087	request_high = create_requests_with_priority(
4088	num_requests=1,
4089	priorities=[0],
4090	arrival_times=[1.0],
4091	num_tokens=32,
4092	starting_idx=1,
4093	)[0]
4094	scheduler.add_request(request_high)
4095	# 2nd schedule
4096	output = scheduler.schedule()
4097	# KV cache should be full at this point
4098	assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == 0
4099	assert len(output.scheduled_new_reqs) == 1
4100	assert output.scheduled_cached_reqs.num_reqs == 1
4101	assert len(scheduler.waiting) == 0
4102	assert len(scheduler.running) == 2
4103
4104	# Simulate model execution - 2nd decode
4105	requests = [request_low, request_high]
4106	model_output = ModelRunnerOutput(
4107	req_ids=[req.request_id for req in requests],
4108	req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
4109	sampled_token_ids=[[100] for _ in requests],
4110	# spec_token_ids=None,
4111	logprobs=None,
4112	prompt_logprobs_dict={},
4113	pooler_output=[],
4114	)
4115	scheduler.update_from_output(output, model_output)
4116
4117	# 3rd schedule - this should trigger preemption
4118	# req_low needs 32 tokens = 2 blocks
4119	# req_high needs 33 tokens = 3 blocks
4120	# so doesn't fit in 4 blocks.
4121	output = scheduler.schedule()
4122
4123	# Should have preempted req_low
4124	assert len(output.scheduled_new_reqs) == 0
4125	assert output.scheduled_cached_reqs.num_reqs == 1
4126	assert output.scheduled_cached_reqs.req_ids[0] == request_high.request_id
4127	assert scheduler.requests[request_low.request_id].status == RequestStatus.PREEMPTED
4128	assert len(scheduler.waiting) == 1
4129	assert len(scheduler.running) == 1
4130
4131	# Simulate model execution - 3rd decode
4132	model_output = ModelRunnerOutput(
4133	req_ids=[req.request_id for req in requests],
4134	req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
4135	sampled_token_ids=[[], [100]],
4136	# spec_token_ids=None,
4137	logprobs=None,
4138	prompt_logprobs_dict={},
4139	pooler_output=[],
4140	)
4141	# Finish the requests to make room for the preempted requests to resume
4142	scheduler.update_from_output(output, model_output)
4143	scheduler.finish_requests(request_high.request_id, RequestStatus.FINISHED_STOPPED)
4144
4145	# 4th Schedule - this should trigger the resumption
4146	output = scheduler.schedule()
4147	scheduled_cached_reqs = output.scheduled_cached_reqs
4148
4149	assert len(scheduler.waiting) == 0
4150	assert len(scheduler.running) == 1
4151
4152	if use_v2_model_runner:
4153	# V2 emits the resumed request as a NewRequestData, carrying its full
4154	# token ids in prefill_token_ids (instead of cached all_token_ids).
4155	assert scheduled_cached_reqs.num_reqs == 0
4156	assert len(output.scheduled_new_reqs) == 1
4157	new_req = output.scheduled_new_reqs[0]
4158	assert new_req.req_id == request_low.request_id
4159	# Resumed tokens include 30 prompt tokens and 2 decoded tokens.
4160	assert len(new_req.prefill_token_ids) == 32
4161	assert new_req.prefill_token_ids[31] == 100
4162	else:
4163	assert len(output.scheduled_new_reqs) == 0
4164	assert scheduled_cached_reqs.num_reqs == 1
4165	# Preempted request resumed in scheduled_cached_reqs
4166	assert len(scheduled_cached_reqs.resumed_req_ids) == 1
4167	assert len(scheduled_cached_reqs.all_token_ids) == 1
4168	assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
4169	assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids
4170	assert request_low.request_id in scheduled_cached_reqs.all_token_ids
4171	# Resumed tokens include 30 prompt tokens and 2 decoded tokens
4172	assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 32
4173	assert scheduled_cached_reqs.all_token_ids[request_low.request_id][31] == 100
4174
4175
4176	@pytest.mark.parametrize(
4177	("enable_chunked_prefill", "is_encoder_decoder", "expect_enabled"),
4178	[
4179	(True, False, True),
4180	(False, False, False),
4181	# Encoder-decoder models should always have it disabled
4182	(False, True, False),
4183	(True, True, False),
4184	],
4185	)
4186	def test_chunked_prefill_disabled_for_encoder_decoder(
4187	enable_chunked_prefill: bool, is_encoder_decoder: bool, expect_enabled: bool
4188	) -> None:
4189	"""Validate that chunked prefill is appropriately disabled for
4190	encoder-decoder models."""
4191	scheduler_config = SchedulerConfig(
4192	enable_chunked_prefill=enable_chunked_prefill,
4193	is_encoder_decoder=is_encoder_decoder,
4194	# Must <= max_num_batched_tokens if chunked prefill is disabled
4195	max_model_len=SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS,
4196	)
4197
4198	# `is_encoder_decoder` should only be used during construction
4199	# of the config, and otherwise stored in the model config.
4200	assert "is_encoder_decoder" not in vars(scheduler_config)
4201	assert "is_encoder_decoder" not in [
4202	f.name for f in dataclasses.fields(scheduler_config)
4203	]
4204	_validate_chunked_prefill_settings_for_encoder_decoder(
4205	scheduler_config, is_encoder_decoder, expect_enabled
4206	)
4207
4208	# Ensure it is retained in VllmConfig, even after its post-init.
4209	vllm_config = VllmConfig(scheduler_config=scheduler_config)
4210	_validate_chunked_prefill_settings_for_encoder_decoder(
4211	vllm_config.scheduler_config, is_encoder_decoder, expect_enabled
4212	)
4213
4214
4215	def _validate_chunked_prefill_settings_for_encoder_decoder(
4216	scheduler_config: SchedulerConfig, is_encoder_decoder: bool, expect_enabled: bool
4217	) -> None:
4218	"""Validate chunked prefill settings in the scheduler config for
4219	encoder-decoder models."""
4220	assert scheduler_config.enable_chunked_prefill is expect_enabled
4221	if is_encoder_decoder:
4222	# Encoder-decoder models should automatically disable chunked multimodal
4223	# inputs as well
4224	assert scheduler_config.disable_chunked_mm_input is not expect_enabled
4225	if is_encoder_decoder and not expect_enabled:
4226	assert scheduler_config.long_prefill_token_threshold == 0
4227
4228
4229	# ==============================================================================
4230	# EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests start
4231	# NOTE: In E->P->D disagg case, both KV and EC Connector works in P instance
4232	# Unless specify, the existence of KV Connector should not affect any test results
4233	# ==============================================================================
4234
4235
4236	def _assert_right_encoder_cache_allocated(
4237	scheduler: Scheduler,
4238	hashes_to_check: list[str] \| None = None,
4239	requests: list[Request] \| None = None,
4240	expected_total_allocated: int \| None = None,
4241	):
4242	"""Check whether encoder cache is allocated correctly."""
4243	encoder_cache_manager = scheduler.encoder_cache_manager
4244
4245	# Verify encoder cache manager exists
4246	assert encoder_cache_manager is not None, "Encoder cache manager should exist"
4247
4248	# Verify number of cache
4249	if expected_total_allocated is not None:
4250	assert len(encoder_cache_manager.cached) == expected_total_allocated
4251	if expected_total_allocated == 0:
4252	return
4253
4254	# Verify each request with MM data is in cache
4255	cached_hashes = set(encoder_cache_manager.cached.keys())
4256
4257	if hashes_to_check:
4258	missed_hashes = set(hashes_to_check) - cached_hashes
4259	assert not missed_hashes, (
4260	f"Miss hashes: {missed_hashes} "
4261	f"Existing encoder cache: {encoder_cache_manager.cached}"
4262	)
4263
4264	for req in requests if requests is not None else []:
4265	if req.mm_features:
4266	mm_hashes = [f.identifier for f in req.mm_features]
4267	req_hashes = set(mm_hashes) # unique hashes set
4268	missed_hashes = req_hashes - cached_hashes
4269	assert not missed_hashes, (
4270	f"Miss hashes in cache for request {req.request_id}: {missed_hashes} "
4271	f"Existing encoder cache: {encoder_cache_manager.cached}"
4272	)
4273
4274
4275	def _assert_right_ec_connector_metadata(
4276	output: SchedulerOutput,
4277	mm_features_list: list[MultiModalFeatureSpec],
4278	):
4279	"""Verify that ECConnector metadata EXACTLY matches the input MM data."""
4280	# Get the connector metadata
4281	metadata = output.ec_connector_metadata
4282
4283	# Create lookup dictionaries for efficient access
4284	metadata_dict = {mm_data.mm_hash: mm_data for mm_data in metadata.mm_datas}
4285
4286	# Check all required identifiers exist in metadata; and no extra
4287	# In ECExampleConnector format
4288	# NOTE: even having same identifier, the mm_features can be different
4289	# since their mm_position can be in different offsets, etc
4290	identifiers_dict = {f.identifier for f in mm_features_list}
4291	assert set(metadata_dict.keys()) == identifiers_dict
4292
4293	# Verify the info matches
4294	for i, mm_feature in enumerate(mm_features_list):
4295	identifier = mm_feature.identifier
4296	assert metadata_dict[identifier].mm_hash == identifier
4297	assert metadata_dict[identifier].num_token == mm_feature.mm_position.length
4298
4299
4300	def _assert_right_encoder_inputs(
4301	output: SchedulerOutput,
4302	check_exist: bool \| None = True,
4303	requests: list[Request] \| None = None,
4304	expected_encoder_inputs: list[list[int]] \| None = None,
4305	expected_total_reqs: int \| None = None,
4306	):
4307	"""Verify that requests/mm_hashes should (not) in scheduled encoder input
4308	If check_exist is False, this function returns True
4309	if requests are NOT in encoder inputs"""
4310	# Get the scheduled encoder inputs
4311	# NOTE: scheduled_encoder_inputs is a dictionary with request id as key
4312	scheduled_encoder_inputs = output.scheduled_encoder_inputs
4313
4314	# Check if scheduled_encoder_inputs is empty as expected
4315	if expected_total_reqs is not None:
4316	assert len(scheduled_encoder_inputs) == expected_total_reqs
4317	if expected_total_reqs == 0:
4318	return
4319
4320	# Number of expected encoder inputs should match number of requests
4321	if expected_encoder_inputs:
4322	assert check_exist and requests is not None # only support expect input exist
4323	assert len(requests) == len(expected_encoder_inputs)
4324
4325	# Check request (not) exist as expected
4326	for i, request in enumerate(requests if requests is not None else []):
4327	assert (request.request_id in scheduled_encoder_inputs) is check_exist, (
4328	f"Request {request.id} presence mismatch: expected {check_exist}, "
4329	f"got {request.id in scheduled_encoder_inputs}"
4330	)
4331	if expected_encoder_inputs:
4332	scheduled_encoder_input = scheduled_encoder_inputs[request.request_id]
4333	assert scheduled_encoder_input == expected_encoder_inputs[i]
4334
4335
4336	def test_scheduler_no_ec_connector_by_default():
4337	"""Test scheduler doesn't have EC connector by default."""
4338	scheduler = create_scheduler()
4339	assert scheduler.ec_connector is None
4340
4341
4342	def test_mamba_align_encoder_cache_cap_makes_progress():
4343	"""Two individually cacheable images must not deadlock Mamba alignment."""
4344	block_size = 768
4345	encoder_cache_size = 600
4346	scheduler = create_scheduler(
4347	max_num_batched_tokens=8192,
4348	block_size=block_size,
4349	enable_prefix_caching=True,
4350	)
4351	scheduler.need_mamba_block_aligned_split = True
4352	scheduler.max_num_encoder_input_tokens = encoder_cache_size
4353	scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=encoder_cache_size)
4354
4355	first_image_end = 510
4356	request = create_requests(
4357	num_requests=1,
4358	num_tokens=1010,
4359	mm_positions=[
4360	[
4361	PlaceholderRange(offset=0, length=500),
4362	PlaceholderRange(offset=first_image_end, length=500),
4363	]
4364	],
4365	max_tokens=1,
4366	block_size=block_size,
4367	req_ids=["req"],
4368	)[0]
4369	scheduler.add_request(request)
4370
4371	output = scheduler.schedule()
4372	assert output.num_scheduled_tokens[request.request_id] == first_image_end
4373	assert output.scheduled_encoder_inputs[request.request_id] == [0]
4374
4375	scheduler.update_from_output(
4376	output,
4377	ModelRunnerOutput(
4378	req_ids=[request.request_id],
4379	req_id_to_index={request.request_id: 0},
4380	sampled_token_ids=[[]],
4381	logprobs=None,
4382	prompt_logprobs_dict={},
4383	pooler_output=[],
4384	),
4385	)
4386
4387	output = scheduler.schedule()
4388	next_block_boundary = block_size
4389	assert output.num_scheduled_tokens[request.request_id] == (
4390	next_block_boundary - first_image_end
4391	)
4392	assert output.scheduled_encoder_inputs[request.request_id] == [1]
4393
4394
4395	def test_mamba_align_eagle_schedules_encoder_at_boundary():
4396	"""EAGLE lookahead at an aligned MM boundary requires encoder cache."""
4397	block_size = 512
4398	scheduler = create_scheduler(
4399	max_num_batched_tokens=700,
4400	max_model_len=2048,
4401	block_size=block_size,
4402	enable_prefix_caching=True,
4403	)
4404	scheduler.need_mamba_block_aligned_split = True
4405	scheduler.use_eagle = True
4406	scheduler.num_prefill_lookahead = 1
4407	scheduler.max_num_encoder_input_tokens = 2048
4408	scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=2048)
4409
4410	request = create_requests(
4411	num_requests=1,
4412	num_tokens=1200,
4413	mm_positions=[[PlaceholderRange(offset=block_size, length=100)]],
4414	max_tokens=1,
4415	block_size=block_size,
4416	req_ids=["req"],
4417	)[0]
4418	scheduler.add_request(request)
4419
4420	output = scheduler.schedule()
4421
4422	assert output.num_scheduled_tokens[request.request_id] == block_size
4423	assert output.scheduled_encoder_inputs[request.request_id] == [0]
4424
4425
4426	@pytest.mark.parametrize("use_kv_connector", [False, True])
4427	def test_ec_connector_text_only_request(use_kv_connector):
4428	"""Test text-only requests don't allocate encoder cache."""
4429	scheduler = create_scheduler(
4430	model="llava-hf/llava-1.5-7b-hf",
4431	use_kv_connector=use_kv_connector,
4432	use_ec_connector=True,
4433	ec_role="ec_consumer",
4434	)
4435
4436	NUM_PROMPT_TOKENS = 100
4437
4438	# Create text-only request (no mm_positions)
4439	requests = create_requests(
4440	num_requests=1,
4441	num_tokens=NUM_PROMPT_TOKENS,
4442	)
4443	assert not requests[0].mm_features # No MM data
4444
4445	scheduler.add_request(requests[0])
4446	output = scheduler.schedule()
4447
4448	# Should schedule
4449	assert len(output.scheduled_new_reqs) == 1
4450
4451	# Scheduled tokens should equal prompt tokens exactly
4452	scheduled = output.num_scheduled_tokens[requests[0].request_id]
4453	assert scheduled == NUM_PROMPT_TOKENS, (
4454	f"Text-only should schedule {NUM_PROMPT_TOKENS}, got {scheduled}"
4455	)
4456
4457	# Encoder cache should be empty
4458	_assert_right_encoder_cache_allocated(scheduler, expected_total_allocated=0)
4459
4460	# ECConnector should carry no metadata
4461	_assert_right_ec_connector_metadata(output, mm_features_list=[])
4462
4463	# Scheduled encoder input should be empty; no mm to compute
4464	_assert_right_encoder_inputs(output, expected_total_reqs=0)
4465
4466
4467	@pytest.mark.parametrize("use_kv_connector", [False, True])
4468	def test_ec_connector_cache_hit_external_load(use_kv_connector):
4469	"""Test ec_consumer loads from external cache when hit.
4470	A normal basic operation for EPD disaggrgation"""
4471	scheduler = create_scheduler(
4472	model="llava-hf/llava-1.5-7b-hf",
4473	enable_prefix_caching=True,
4474	# kv connector should not effect test results
4475	use_kv_connector=use_kv_connector,
4476	use_ec_connector=True,
4477	ec_role="ec_consumer",
4478	)
4479
4480	# Create MM request
4481	NUM_TOKENS = 200 # NOTE: includes mm tokens
4482	NUM_ENCODER_TOKENS = 100
4483	mm_hashes_list = [["hash_test1"]]
4484	mm_positions = [[PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS)]]
4485
4486	request = create_requests(
4487	num_requests=1,
4488	num_tokens=NUM_TOKENS,
4489	mm_hashes_list=mm_hashes_list,
4490	mm_positions=mm_positions,
4491	)[0]
4492
4493	# Mock cache hit - encoder cache has_exists externally
4494	scheduler.ec_connector.has_cache_item = Mock(return_value=True)
4495	scheduler.ec_connector.update_state_after_alloc = Mock(
4496	wraps=scheduler.ec_connector.update_state_after_alloc
4497	)
4498
4499	scheduler.add_request(request)
4500	output = scheduler.schedule()
4501	# Should schedule prompt tokens
4502	scheduled_tokens = output.num_scheduled_tokens[request.request_id]
4503	assert scheduled_tokens == NUM_TOKENS
4504
4505	# Should called update_state_after_alloc for external load
4506	scheduler.ec_connector.update_state_after_alloc.assert_called_with(request, 0)
4507
4508	# Encoder cache should contain mm items from request
4509	_assert_right_encoder_cache_allocated(scheduler, requests=[request])
4510
4511	# ECConnector should carry metadata of request
4512	_assert_right_ec_connector_metadata(output, mm_features_list=request.mm_features)
4513
4514	# Scheduled encoder input should be empty; no mm to compute
4515	_assert_right_encoder_inputs(output, expected_total_reqs=0)
4516
4517
4518	@pytest.mark.parametrize("use_kv_connector", [False, True])
4519	def test_ec_connector_cache_miss_computes_locally(use_kv_connector):
4520	"""Test consumer can compute encoder locally when cache miss (fallback)."""
4521	# encoder cache itself if it doesn't receive it from external storage
4522
4523	scheduler = create_scheduler(
4524	model="llava-hf/llava-1.5-7b-hf",
4525	enable_prefix_caching=True,
4526	use_kv_connector=use_kv_connector,
4527	use_ec_connector=True,
4528	ec_role="ec_consumer",
4529	)
4530
4531	# Verify consumer role
4532	assert scheduler.ec_connector is not None
4533	assert not scheduler.ec_connector.is_producer
4534
4535	# Create MM request
4536	request_mm_missed = create_requests(
4537	num_requests=1,
4538	num_tokens=200, # Total (including 100 MM)
4539	mm_positions=[[PlaceholderRange(offset=0, length=100)]], # 100 MM tokens
4540	)[0]
4541
4542	# Mock cache miss - encoder cache doesn't exist externally
4543	scheduler.ec_connector.has_cache_item = Mock(return_value=False)
4544
4545	scheduler.add_request(request_mm_missed)
4546	output = scheduler.schedule()
4547
4548	# SCHEDULER should decide to compute encoder locally (fallback)
4549	assert len(output.scheduled_new_reqs) == 1
4550
4551	# Should schedule full prompt tokens
4552	scheduled_tokens = output.num_scheduled_tokens[request_mm_missed.request_id]
4553	assert scheduled_tokens == 200, (
4554	f"Expected 200 tokens on cache miss, got {scheduled_tokens}"
4555	)
4556
4557	# Encoder cache should contain mm items from request
4558	_assert_right_encoder_cache_allocated(scheduler, requests=[request_mm_missed])
4559
4560	# ECConnector should carry no metadata (missed cache)
4561	_assert_right_ec_connector_metadata(output, mm_features_list=[])
4562
4563	# Scheduled encoder input contain mm for request_mm_missed
4564	_assert_right_encoder_inputs(
4565	output,
4566	requests=[request_mm_missed],
4567	expected_encoder_inputs=[[0]], # index 0 of the mm item
4568	expected_total_reqs=1,
4569	)
4570
4571	# Then MODEL_RUNNER will execute the encoder and cache the result
4572
4573
4574	@pytest.mark.parametrize("use_kv_connector", [False, True])
4575	def test_ec_connector_with_partial_cache_hit_multi_round(use_kv_connector):
4576	"""Test consumer with partial cache hit (local & connector) with 2 requests."""
4577	scheduler = create_scheduler(
4578	model="llava-hf/llava-1.5-7b-hf",
4579	enable_prefix_caching=True,
4580	use_kv_connector=use_kv_connector,
4581	use_ec_connector=True,
4582	ec_role="ec_consumer",
4583	)
4584
4585	# Create MM request
4586	NUM_TOKENS_1 = 300 # NOTE: includes mm tokens
4587	NUM_ENCODER_TOKENS_1 = 50
4588	mm_hashes_list_1 = [["hash1_A", "hash1_B", "hash1_A", "hash1_F"]]
4589	mm_positions_1 = [
4590	[
4591	PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS_1),
4592	PlaceholderRange(offset=100, length=NUM_ENCODER_TOKENS_1),
4593	PlaceholderRange(offset=200, length=NUM_ENCODER_TOKENS_1),
4594	PlaceholderRange(offset=250, length=NUM_ENCODER_TOKENS_1),
4595	]
4596	]
4597	has_cache_item_result_map_1 = {"hash1_A": False, "hash1_B": True, "hash1_F": True}
4598	# Create request with 4 MM items, with 2 identical items
4599	request1 = create_requests(
4600	num_requests=1,
4601	num_tokens=NUM_TOKENS_1,
4602	mm_hashes_list=mm_hashes_list_1,
4603	mm_positions=mm_positions_1,
4604	max_tokens=1, # For simplicity
4605	)[0]
4606
4607	# Mock partial cache hit: 1st and 3rd missing, 2nd and 4th exist
4608	scheduler.ec_connector.has_cache_item = Mock(
4609	side_effect=lambda hash_val: has_cache_item_result_map_1[hash_val]
4610	)
4611	scheduler.ec_connector.update_state_after_alloc = Mock(
4612	wraps=scheduler.ec_connector.update_state_after_alloc
4613	)
4614
4615	scheduler.add_request(request1)
4616	output = scheduler.schedule()
4617
4618	# Should schedule all tokens
4619	scheduled_tokens = output.num_scheduled_tokens[request1.request_id]
4620	assert scheduled_tokens == NUM_TOKENS_1
4621
4622	# Encoder cache should contain all mm items from request
4623	_assert_right_encoder_cache_allocated(scheduler, requests=[request1])
4624
4625	# Should have called update_state_after_alloc for external load
4626	scheduler.ec_connector.update_state_after_alloc.assert_called()
4627	scheduler.ec_connector.update_state_after_alloc.reset_mock()
4628
4629	# ECConnector should carry metadata for 2nd and 4th mm item
4630	_assert_right_ec_connector_metadata(
4631	output, mm_features_list=[request1.mm_features[1], request1.mm_features[3]]
4632	)
4633
4634	# Should schedule ONLY 1 encoder input (index 0), no repeat for identical items
4635	_assert_right_encoder_inputs(
4636	output,
4637	requests=[request1],
4638	expected_encoder_inputs=[[0]], # index 0 of the mm item ONLY
4639	expected_total_reqs=1,
4640	)
4641
4642	# Simulate model execution 1 step
4643	model_output = ModelRunnerOutput(
4644	req_ids=[request1.request_id],
4645	req_id_to_index={request1.request_id: 0},
4646	sampled_token_ids=[[100]],
4647	# spec_token_ids=None,
4648	logprobs=None,
4649	prompt_logprobs_dict={},
4650	pooler_output=[],
4651	)
4652	scheduler.update_from_output(output, model_output)
4653
4654	# request1 is finished after outputting 1 token
4655	# Finish request
4656	scheduler.finish_requests(request1.request_id, RequestStatus.FINISHED_LENGTH_CAPPED)
4657
4658	# Create another request with 4 MM items
4659	NUM_TOKENS_2 = 400
4660	NUM_ENCODER_TOKENS_2 = 50
4661	mm_hashes_list_2 = [["hash1_C", "hash1_D", "hash1_E", "hash1_A"]]
4662	mm_positions_2 = [
4663	[
4664	PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS_2),
4665	PlaceholderRange(offset=100, length=NUM_ENCODER_TOKENS_2),
4666	PlaceholderRange(offset=200, length=NUM_ENCODER_TOKENS_2),
4667	PlaceholderRange(offset=250, length=NUM_ENCODER_TOKENS_2),
4668	]
4669	]
4670	has_cache_item_result_map_2 = {
4671	"hash1_C": True,
4672	"hash1_D": False,
4673	"hash1_E": False,
4674	"hash1_A": True,
4675	}
4676	request2 = create_requests(
4677	num_requests=1,
4678	num_tokens=NUM_TOKENS_2,
4679	mm_hashes_list=mm_hashes_list_2,
4680	mm_positions=mm_positions_2,
4681	max_tokens=1, # For simplicity
4682	)[0]
4683
4684	# Mock partial cache hit: only hash1_A and hash1_C exist in connector
4685	scheduler.ec_connector.has_cache_item = Mock(
4686	side_effect=lambda hash_val: has_cache_item_result_map_2[hash_val]
4687	)
4688
4689	scheduler.add_request(request2)
4690	output = scheduler.schedule()
4691
4692	# Check
4693	# Should schedule all tokens
4694	scheduled_tokens = output.num_scheduled_tokens[request2.request_id]
4695	assert scheduled_tokens == 400
4696
4697	# Encoder cache should contain all mm items from request2
4698	_assert_right_encoder_cache_allocated(scheduler, requests=[request2])
4699
4700	# hash1_A should not be loaded from connector
4701	# since it's computed in last request & exist in local cache
4702	# Order of getting encoder cache should be: local cache -> connector-> compute
4703	# update_state_after_alloc is called for all paths:
4704	# index 0 (hash1_C): connector hit → queued for load
4705	# index 1 (hash1_D): cache miss → no-op inside connector
4706	# index 2 (hash1_E): cache miss → no-op inside connector
4707	scheduler.ec_connector.update_state_after_alloc.assert_any_call(request2, 0)
4708	scheduler.ec_connector.update_state_after_alloc.assert_any_call(request2, 1)
4709	scheduler.ec_connector.update_state_after_alloc.assert_any_call(request2, 2)
4710
4711	scheduler.ec_connector.update_state_after_alloc.reset_mock()
4712
4713	# ECConnector should carry metadata for hash1_C only (index 0)
4714	_assert_right_ec_connector_metadata(
4715	output, mm_features_list=[request2.mm_features[0]]
4716	)
4717
4718	# Should schedule 2 encoder input hash1_D and hash1_E (index 1, 2)
4719	_assert_right_encoder_inputs(
4720	output,
4721	requests=[request2],
4722	expected_encoder_inputs=[[1, 2]],
4723	expected_total_reqs=1,
4724	)
4725
4726
4727	@pytest.mark.parametrize("cache_exist", ["local", "connector_only", "no_where"])
4728	@pytest.mark.parametrize("use_kv_connector", [False, True])
4729	def test_ec_connector_schedule_multiple_requests(cache_exist, use_kv_connector):
4730	scheduler = create_scheduler(
4731	model="llava-hf/llava-1.5-7b-hf",
4732	max_num_seqs=10, # allow multiple requests
4733	max_num_batched_tokens=2048,
4734	enable_prefix_caching=True,
4735	use_kv_connector=use_kv_connector,
4736	use_ec_connector=True,
4737	ec_role="ec_consumer",
4738	)
4739	mm_hashes_list = [[f"hash_{i}"] for i in range(10)]
4740	mm_positions = [[PlaceholderRange(offset=i, length=100)] for i in range(10)]
4741	requests = create_requests(
4742	num_requests=10,
4743	num_tokens=200,
4744	mm_hashes_list=mm_hashes_list,
4745	mm_positions=mm_positions,
4746	)
4747	for request in requests:
4748	scheduler.add_request(request)
4749
4750	# Set up to test different encoder cache existence scenario after preemption
4751	# Order of getting encoder cache should be: local cache -> connector-> compute
4752	scheduler.ec_connector.update_state_after_alloc = Mock(
4753	wraps=scheduler.ec_connector.update_state_after_alloc
4754	)
4755
4756	if cache_exist == "local":
4757	# Allocate cache to cache manager manually to mimic
4758	for req in requests:
4759	scheduler.encoder_cache_manager.allocate(req, 0)
4760	else:
4761	# Make sure local encoder cache empty
4762	scheduler.encoder_cache_manager.cached = {}
4763
4764	if cache_exist == "connector_only":
4765	# Cache exist in ec_connector
4766	scheduler.ec_connector.has_cache_item = Mock(return_value=True)
4767	elif cache_exist == "no_where":
4768	scheduler.ec_connector.has_cache_item = Mock(return_value=False)
4769
4770	output = scheduler.schedule()
4771	assert len(output.scheduled_new_reqs) == len(requests)
4772	assert output.scheduled_cached_reqs.num_reqs == 0
4773	assert len(output.finished_req_ids) == 0
4774	for req_id, num_tokens in output.num_scheduled_tokens.items():
4775	assert num_tokens == len(requests[int(req_id)].prompt_token_ids)
4776
4777	## Encoder-cache-specific checks:
4778	# mm_hashes of requests exist in cache after scheduling for all scenario
4779	_assert_right_encoder_cache_allocated(scheduler, requests=requests)
4780
4781	if cache_exist == "connector_only":
4782	scheduler.ec_connector.update_state_after_alloc.assert_called_with(
4783	requests[-1], 0
4784	)
4785
4786	# Concat mm_features for the 10 requests together
4787	mm_features_list = [feature for req in requests for feature in req.mm_features]
4788
4789	# Check metadata should contain mm data for all 10 requests
4790	_assert_right_ec_connector_metadata(output, mm_features_list=mm_features_list)
4791	elif cache_exist == "local":
4792	# Local cache hit: items never reach update_state_after_alloc
4793	scheduler.ec_connector.update_state_after_alloc.assert_not_called()
4794	_assert_right_ec_connector_metadata(output, mm_features_list=[])
4795	else:
4796	# no_where: called from encoder_inputs_to_schedule but no-op
4797	# inside connector (has_cache_item returns False)
4798	assert cache_exist == "no_where"
4799	scheduler.ec_connector.update_state_after_alloc.assert_called()
4800	_assert_right_ec_connector_metadata(output, mm_features_list=[])
4801
4802	scheduler.ec_connector.update_state_after_alloc.reset_mock()
4803
4804	# Should only schedule encoder input when cache is not found anywhere
4805	if cache_exist == "no_where":
4806	_assert_right_encoder_inputs(
4807	output,
4808	requests=requests,
4809	expected_encoder_inputs=[[0] for _ in range(10)],
4810	expected_total_reqs=10,
4811	)
4812	else:
4813	_assert_right_encoder_inputs(output, expected_total_reqs=0)
4814
4815
4816	@pytest.mark.parametrize("use_kv_connector", [False, True])
4817	def test_ec_connector_unable_to_allocate(use_kv_connector):
4818	"""Test whether scheduler with ECConnector is able to handle
4819	unable to allocate (run out of blocks).
4820	"""
4821	# Setup Scheduler With Mock External Cache Hit.
4822	BLOCK_SIZE = 4
4823	NUM_BLOCKS = 10
4824	scheduler = create_scheduler(
4825	model="llava-hf/llava-1.5-7b-hf",
4826	enable_prefix_caching=True,
4827	use_kv_connector=use_kv_connector,
4828	block_size=BLOCK_SIZE,
4829	num_blocks=NUM_BLOCKS,
4830	use_ec_connector=True,
4831	ec_role="ec_consumer",
4832	)
4833
4834	# Mock ec_connector load external cache behavior
4835	scheduler.ec_connector.has_cache_item = Mock(return_value=True)
4836	scheduler.ec_connector.update_state_after_alloc = Mock(
4837	wraps=scheduler.ec_connector.update_state_after_alloc
4838	)
4839
4840	# Create two requests. The second request will not be able to
4841	# allocate slots because it will not have enough blocks.
4842	NUM_REQUESTS = 2
4843	NUM_TOKENS = (NUM_BLOCKS // 2 + 1) * BLOCK_SIZE
4844	MAX_TOKENS = 2
4845	requests = create_requests(
4846	num_requests=NUM_REQUESTS,
4847	num_tokens=NUM_TOKENS,
4848	mm_hashes_list=[["hash_1"], ["hash_2"]],
4849	mm_positions=[
4850	[PlaceholderRange(offset=1, length=10)] for _ in range(NUM_REQUESTS)
4851	],
4852	max_tokens=MAX_TOKENS,
4853	block_size=BLOCK_SIZE,
4854	)
4855	req_ids = []
4856	req_to_index = {}
4857	for i, request in enumerate(requests):
4858	scheduler.add_request(request)
4859	req_ids.append(request.request_id)
4860	req_to_index[request.request_id] = i
4861
4862	# Setup MODEL_RUNNER_OUTPUT to be run in _step_until_done later
4863	MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
4864	req_ids=req_ids,
4865	req_id_to_index=req_to_index,
4866	sampled_token_ids=[[1000]] * len(req_ids),
4867	logprobs=None,
4868	prompt_logprobs_dict={},
4869	pooler_output=[],
4870	)
4871
4872	# Just one request should be running.
4873	output = scheduler.schedule()
4874	scheduled_tokens = output.num_scheduled_tokens[scheduler.running[0].request_id]
4875	assert scheduled_tokens == NUM_TOKENS
4876	assert len(scheduler.running) == 1
4877	assert len(scheduler.waiting) == 1
4878
4879	# Should have called update_state_after_alloc for external load
4880	scheduler.ec_connector.update_state_after_alloc.assert_called_with(
4881	scheduler.running[0], 0
4882	)
4883	scheduler.ec_connector.update_state_after_alloc.reset_mock()
4884
4885	# All memory should be freed, with one request waiting.
4886	_step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
4887	assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
4888	assert len(scheduler.running) == 0
4889	assert len(scheduler.waiting) == 1
4890
4891	# Just one request should be running.
4892	output = scheduler.schedule()
4893	scheduled_tokens = output.num_scheduled_tokens[scheduler.running[0].request_id]
4894	assert scheduled_tokens == NUM_TOKENS
4895	assert len(scheduler.running) == 1
4896	assert len(scheduler.waiting) == 0
4897
4898	# update_state_after_alloc should be called for loading external cache
4899	scheduler.ec_connector.update_state_after_alloc.assert_called_with(
4900	scheduler.running[0], 0
4901	)
4902	scheduler.ec_connector.update_state_after_alloc.reset_mock()
4903
4904	# All memory should be freed, with no requests waiting / running.
4905	_step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
4906	assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == NUM_BLOCKS - 1
4907	assert len(scheduler.running) == 0
4908	assert len(scheduler.waiting) == 0
4909
4910
4911	@pytest.mark.parametrize("cache_exist", ["local", "connector_only", "no_where"])
4912	@pytest.mark.parametrize("use_kv_connector", [False, True])
4913	def test_priority_scheduling_ec_connector_preemption_and_resumption(
4914	cache_exist, use_kv_connector
4915	):
4916	"""Test that priority scheduling preempts lower priority requests
4917	when out of KV cache space."""
4918	# Create scheduler with very limited memory to force preemption
4919	scheduler = create_scheduler_with_priority(
4920	model="llava-hf/llava-1.5-7b-hf",
4921	enable_prefix_caching=True,
4922	max_num_seqs=2, # allow multiple requests
4923	# kv connector should not effect test results
4924	use_kv_connector=use_kv_connector,
4925	num_blocks=15, # can hold 244 tokens with 14 blocks (first block is null)
4926	block_size=16, # standard block size
4927	use_ec_connector=True,
4928	ec_role="ec_consumer",
4929	)
4930
4931	# Mock cache hit: Both cache exist in connector (at E->PD initially)
4932	scheduler.ec_connector.has_cache_item = Mock(return_value=True)
4933	scheduler.ec_connector.update_state_after_alloc = Mock(
4934	wraps=scheduler.ec_connector.update_state_after_alloc
4935	)
4936
4937	# Create a request and schedule it (and to be preempted)
4938	request_low = create_requests_with_priority(
4939	num_requests=1,
4940	priorities=[1],
4941	arrival_times=[0.0],
4942	num_tokens=94,
4943	mm_hashes_list=[["hash_low"]],
4944	# NOTE: this test only preempt the last block.
4945	# Setting mm_position at the last block can force to recompute encoding
4946	mm_positions=[[PlaceholderRange(offset=82, length=10)]],
4947	starting_idx=0,
4948	)[0]
4949	scheduler.add_request(request_low)
4950	# 1st schedule
4951	output = scheduler.schedule()
4952
4953	assert len(output.scheduled_new_reqs) == 1
4954	scheduled_tokens = output.num_scheduled_tokens[request_low.request_id]
4955	assert scheduled_tokens == 94
4956	assert len(scheduler.waiting) == 0
4957	assert len(scheduler.running) == 1
4958
4959	## Encoder-cache-specific checks:
4960	# Encoder cache should contain mm items from request
4961	_assert_right_encoder_cache_allocated(scheduler, requests=[request_low])
4962
4963	# Verify update_state_after_alloc called (external load)
4964	scheduler.ec_connector.update_state_after_alloc.assert_called_with(request_low, 0)
4965	scheduler.ec_connector.update_state_after_alloc.reset_mock()
4966
4967	# ECConnector should carry metadata of request
4968	_assert_right_ec_connector_metadata(
4969	output, mm_features_list=request_low.mm_features
4970	)
4971
4972	# Scheduled encoder input should be empty; no mm to compute
4973	_assert_right_encoder_inputs(output, expected_total_reqs=0)
4974
4975	# Simulate model execution - 1st decode
4976	model_output = ModelRunnerOutput(
4977	req_ids=[request_low.request_id],
4978	req_id_to_index={request_low.request_id: 0},
4979	sampled_token_ids=[[100]],
4980	# spec_token_ids=None,
4981	logprobs=None,
4982	prompt_logprobs_dict={},
4983	pooler_output=[],
4984	)
4985	scheduler.update_from_output(output, model_output)
4986
4987	# Create a high priority request and schedule it
4988	request_high = create_requests_with_priority(
4989	num_requests=1,
4990	priorities=[0],
4991	arrival_times=[1.0],
4992	num_tokens=128,
4993	mm_hashes_list=[["hash_high"]],
4994	mm_positions=[[PlaceholderRange(offset=1, length=10)]],
4995	max_tokens=2,
4996	starting_idx=1,
4997	)[0]
4998	scheduler.add_request(request_high)
4999	# 2nd schedule
5000	output = scheduler.schedule()

Showing first 5,000 of 7,350 lines. View raw