| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """SWA bounded replay: after a prefix hit the scheduler recomputes the tail of |
| 4 | the hit to rebuild the non-cacheable sliding-window group, keeps the hit's |
| 5 | blocks, and hands the worker the replay start.""" |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from tests.v1.kv_connector.unit.utils import MockKVConfig, create_model_runner_output |
| 11 | from vllm.v1.core.sched.scheduler import Scheduler |
| 12 | from vllm.v1.kv_cache_interface import ( |
| 13 | FullAttentionSpec, |
| 14 | KVCacheConfig, |
| 15 | KVCacheGroupSpec, |
| 16 | SlidingWindowMLASpec, |
| 17 | ) |
| 18 | from vllm.v1.outputs import ModelRunnerOutput |
| 19 | from vllm.v1.request import RequestStatus |
| 20 | from vllm.v1.structured_output import StructuredOutputManager |
| 21 | |
| 22 | from .utils import create_requests, create_scheduler |
| 23 | |
| 24 | BLOCK_SIZE = 16 |
| 25 | WINDOW = 32 |
| 26 | NUM_PROMPT_TOKENS = 100 |
| 27 | # 100 tokens -> 6 full blocks cached -> 96-token hit. |
| 28 | HIT_TOKENS = NUM_PROMPT_TOKENS // BLOCK_SIZE * BLOCK_SIZE |
| 29 | FULL, SWA = 0, 1 |
| 30 | |
| 31 | |
| 32 | def _replay_scheduler( |
| 33 | *, |
| 34 | long_prefill_token_threshold: int = 0, |
| 35 | use_kv_connector: MockKVConfig | None = None, |
| 36 | ) -> Scheduler: |
| 37 | """A hybrid layout: one prefix-cacheable full-attention group and one |
| 38 | replayed sliding-window group.""" |
| 39 | base = create_scheduler( |
| 40 | block_size=BLOCK_SIZE, |
| 41 | enable_prefix_caching=True, |
| 42 | long_prefill_token_threshold=long_prefill_token_threshold, |
| 43 | use_kv_connector=use_kv_connector, |
| 44 | ) |
| 45 | vllm_config = base.vllm_config |
| 46 | kv_cache_config = KVCacheConfig( |
| 47 | num_blocks=10000, |
| 48 | kv_cache_tensors=[], |
| 49 | kv_cache_groups=[ |
| 50 | KVCacheGroupSpec( |
| 51 | ["full"], |
| 52 | FullAttentionSpec( |
| 53 | block_size=BLOCK_SIZE, |
| 54 | num_kv_heads=1, |
| 55 | head_size=1, |
| 56 | dtype=torch.float32, |
| 57 | ), |
| 58 | ), |
| 59 | KVCacheGroupSpec( |
| 60 | ["swa"], |
| 61 | SlidingWindowMLASpec( |
| 62 | block_size=BLOCK_SIZE, |
| 63 | num_kv_heads=1, |
| 64 | head_size=64, |
| 65 | dtype=torch.bfloat16, |
| 66 | sliding_window=WINDOW, |
| 67 | bounded_replay=True, |
| 68 | ), |
| 69 | ), |
| 70 | ], |
| 71 | ) |
| 72 | scheduler = Scheduler( |
| 73 | vllm_config=vllm_config, |
| 74 | kv_cache_config=kv_cache_config, |
| 75 | block_size=BLOCK_SIZE, |
| 76 | log_stats=True, |
| 77 | structured_output_manager=StructuredOutputManager(vllm_config), |
| 78 | ) |
| 79 | scheduler.use_v2_model_runner = True |
| 80 | assert scheduler.prefix_replay_tokens == WINDOW |
| 81 | return scheduler |
| 82 | |
| 83 | |
| 84 | def _prefill(scheduler: Scheduler, request) -> None: |
| 85 | scheduler.add_request(request) |
| 86 | out = scheduler.schedule() |
| 87 | assert out.num_scheduled_tokens[request.request_id] == request.num_prompt_tokens |
| 88 | scheduler.update_from_output(out, _step_output(out, [request])) |
| 89 | |
| 90 | |
| 91 | def _step_output(out, requests) -> ModelRunnerOutput: |
| 92 | """Model output for every scheduled request: a token once its prefill is |
| 93 | done (non-zero, so outputs never extend the all-zero shared prompt in the |
| 94 | cache), nothing while it is still prefilling. schedule() has already |
| 95 | advanced num_computed_tokens past this step's chunk.""" |
| 96 | scheduled = [r for r in requests if r.request_id in out.num_scheduled_tokens] |
| 97 | return ModelRunnerOutput( |
| 98 | req_ids=[r.request_id for r in scheduled], |
| 99 | req_id_to_index={r.request_id: i for i, r in enumerate(scheduled)}, |
| 100 | sampled_token_ids=[ |
| 101 | [1000] if r.num_computed_tokens >= r.num_prompt_tokens else [] |
| 102 | for r in scheduled |
| 103 | ], |
| 104 | logprobs=None, |
| 105 | prompt_logprobs_dict={}, |
| 106 | pooler_output=None, |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | def _new_req_data(out, request): |
| 111 | return next(r for r in out.scheduled_new_reqs if r.req_id == request.request_id) |
| 112 | |
| 113 | |
| 114 | def test_hit_replays_window_without_reallocating(): |
| 115 | scheduler = _replay_scheduler() |
| 116 | expected_replay = WINDOW |
| 117 | first, second = create_requests( |
| 118 | num_requests=2, |
| 119 | num_tokens=NUM_PROMPT_TOKENS, |
| 120 | same_prompt=True, |
| 121 | block_size=BLOCK_SIZE, |
| 122 | ) |
| 123 | _prefill(scheduler, first) |
| 124 | manager = scheduler.kv_cache_manager |
| 125 | |
| 126 | scheduler.add_request(second) |
| 127 | out = scheduler.schedule() |
| 128 | |
| 129 | replay_start = HIT_TOKENS - expected_replay |
| 130 | new_req = _new_req_data(out, second) |
| 131 | assert new_req.num_computed_tokens == replay_start |
| 132 | assert new_req.replay_start == replay_start |
| 133 | assert ( |
| 134 | out.num_scheduled_tokens[second.request_id] == NUM_PROMPT_TOKENS - replay_start |
| 135 | ) |
| 136 | |
| 137 | first_blocks = manager.get_blocks(first.request_id).blocks |
| 138 | second_blocks = manager.get_blocks(second.request_id).blocks |
| 139 | # The full-attention hit is shared (replayed tokens included); only the |
| 140 | # tail block is new. |
| 141 | num_hit_blocks = HIT_TOKENS // BLOCK_SIZE |
| 142 | assert second_blocks[FULL][:num_hit_blocks] == first_blocks[FULL][:num_hit_blocks] |
| 143 | assert all(b.ref_cnt == 2 for b in second_blocks[FULL][:num_hit_blocks]) |
| 144 | assert len(second_blocks[FULL]) == num_hit_blocks + 1 |
| 145 | # The window group never hits: blocks below the window are null, the |
| 146 | # window (replayed and new tokens) gets fresh blocks. |
| 147 | num_skipped_blocks = (HIT_TOKENS - WINDOW + 1) // BLOCK_SIZE |
| 148 | assert all(b.is_null for b in second_blocks[SWA][:num_skipped_blocks]) |
| 149 | num_window_blocks = len(second_blocks[SWA]) - num_skipped_blocks |
| 150 | assert num_window_blocks == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE) - num_skipped_blocks |
| 151 | assert not any(b.is_null for b in second_blocks[SWA][num_skipped_blocks:]) |
| 152 | |
| 153 | scheduler.update_from_output(out, _step_output(out, [first, second])) |
| 154 | assert second.num_computed_tokens == NUM_PROMPT_TOKENS |
| 155 | |
| 156 | |
| 157 | def test_replay_only_chunks_make_progress(): |
| 158 | """A chunk cap smaller than the replay window schedules replay-only |
| 159 | chunks; they advance the request instead of stalling it.""" |
| 160 | scheduler = _replay_scheduler(long_prefill_token_threshold=BLOCK_SIZE) |
| 161 | first, second = create_requests( |
| 162 | num_requests=2, |
| 163 | num_tokens=NUM_PROMPT_TOKENS, |
| 164 | same_prompt=True, |
| 165 | block_size=BLOCK_SIZE, |
| 166 | ) |
| 167 | scheduler.add_request(first) |
| 168 | while first.num_computed_tokens < NUM_PROMPT_TOKENS: |
| 169 | out = scheduler.schedule() |
| 170 | scheduler.update_from_output(out, _step_output(out, [first])) |
| 171 | |
| 172 | scheduler.add_request(second) |
| 173 | |
| 174 | def full_blocks(): |
| 175 | return scheduler.kv_cache_manager.get_blocks(second.request_id).blocks[FULL] |
| 176 | |
| 177 | scheduled: list[int] = [] |
| 178 | while second.num_computed_tokens < NUM_PROMPT_TOKENS: |
| 179 | out = scheduler.schedule() |
| 180 | if not scheduled: |
| 181 | new_req = _new_req_data(out, second) |
| 182 | assert new_req.num_computed_tokens == HIT_TOKENS - WINDOW |
| 183 | assert new_req.replay_start == HIT_TOKENS - WINDOW |
| 184 | # A replay-only chunk adopts the hit and allocates nothing new. |
| 185 | assert len(full_blocks()) == HIT_TOKENS // BLOCK_SIZE |
| 186 | scheduled.append(out.num_scheduled_tokens[second.request_id]) |
| 187 | scheduler.update_from_output(out, _step_output(out, [first, second])) |
| 188 | # 32 replayed + 4 new tokens in 16-token chunks: two replay-only chunks. |
| 189 | assert scheduled == [BLOCK_SIZE, BLOCK_SIZE, NUM_PROMPT_TOKENS - HIT_TOKENS] |
| 190 | assert len(full_blocks()) == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE) |
| 191 | assert second.status == RequestStatus.RUNNING |
| 192 | |
| 193 | |
| 194 | def test_async_remote_kv_hit_replays_after_load(): |
| 195 | """A KV-connector hit loaded asynchronously carries no window state |
| 196 | either; the replay is applied when the load completes.""" |
| 197 | matched = 64 |
| 198 | scheduler = _replay_scheduler( |
| 199 | use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True) |
| 200 | ) |
| 201 | request = create_requests( |
| 202 | num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE |
| 203 | )[0] |
| 204 | scheduler.add_request(request) |
| 205 | out = scheduler.schedule() |
| 206 | assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS |
| 207 | assert request.request_id not in out.num_scheduled_tokens |
| 208 | scheduler.update_from_output( |
| 209 | out, create_model_runner_output([], finished_recving={request.request_id}) |
| 210 | ) |
| 211 | |
| 212 | out = scheduler.schedule() |
| 213 | new_req = _new_req_data(out, request) |
| 214 | assert new_req.num_computed_tokens == matched - WINDOW |
| 215 | assert new_req.replay_start == matched - WINDOW |
| 216 | assert ( |
| 217 | out.num_scheduled_tokens[request.request_id] |
| 218 | == NUM_PROMPT_TOKENS - matched + WINDOW |
| 219 | ) |
| 220 | |
| 221 | |
| 222 | def test_remote_kv_hit_is_taken_in_whole_blocks(): |
| 223 | """A hit ending one token short of a block boundary would put the replay |
| 224 | window's first token in a block the sliding-window group retires, so a |
| 225 | connector hit is cut back to whole blocks.""" |
| 226 | matched = 4 * BLOCK_SIZE - 1 |
| 227 | scheduler = _replay_scheduler( |
| 228 | use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True) |
| 229 | ) |
| 230 | request = create_requests( |
| 231 | num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE |
| 232 | )[0] |
| 233 | scheduler.add_request(request) |
| 234 | out = scheduler.schedule() |
| 235 | scheduler.update_from_output( |
| 236 | out, create_model_runner_output([], finished_recving={request.request_id}) |
| 237 | ) |
| 238 | out = scheduler.schedule() |
| 239 | new_req = _new_req_data(out, request) |
| 240 | hit = 3 * BLOCK_SIZE |
| 241 | assert new_req.replay_start == hit - WINDOW |
| 242 | assert new_req.num_computed_tokens == hit - WINDOW |
| 243 | swa_manager = scheduler.kv_cache_manager.coordinator.single_type_managers[SWA] |
| 244 | swa_blocks = scheduler.kv_cache_manager.get_blocks(request.request_id).blocks[SWA] |
| 245 | assert swa_blocks[new_req.replay_start // BLOCK_SIZE] is not swa_manager._null_block |
| 246 | |
| 247 | |
| 248 | @pytest.mark.parametrize("via_connector", [False, True]) |
| 249 | def test_hit_no_longer_than_window_is_ignored(via_connector): |
| 250 | """Such a hit would be recomputed in full anyway. It is not adopted (an |
| 251 | async load of it is not even started), so a request has computed tokens |
| 252 | iff it replays.""" |
| 253 | if via_connector: |
| 254 | scheduler = _replay_scheduler( |
| 255 | use_kv_connector=MockKVConfig(matched_tokens=WINDOW, is_async=True) |
| 256 | ) |
| 257 | request = create_requests( |
| 258 | num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE |
| 259 | )[0] |
| 260 | else: |
| 261 | scheduler = _replay_scheduler() |
| 262 | first = create_requests( |
| 263 | num_requests=1, num_tokens=WINDOW, same_prompt=True, block_size=BLOCK_SIZE |
| 264 | )[0] |
| 265 | _prefill(scheduler, first) |
| 266 | request = create_requests( |
| 267 | num_requests=1, |
| 268 | num_tokens=NUM_PROMPT_TOKENS, |
| 269 | same_prompt=True, |
| 270 | block_size=BLOCK_SIZE, |
| 271 | req_ids=["long"], |
| 272 | )[0] |
| 273 | scheduler.add_request(request) |
| 274 | out = scheduler.schedule() |
| 275 | new_req = _new_req_data(out, request) |
| 276 | assert request.status == RequestStatus.RUNNING |
| 277 | assert new_req.num_computed_tokens == 0 |
| 278 | assert new_req.replay_start == 0 |
| 279 | assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS |