| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | import vllm.envs as envs |
| 7 | from tests.v1.kv_connector.unit.utils import MockKVConfig |
| 8 | from vllm.config import ( |
| 9 | CacheConfig, |
| 10 | DiffusionConfig, |
| 11 | ECTransferConfig, |
| 12 | KVTransferConfig, |
| 13 | ModelConfig, |
| 14 | MultiModalConfig, |
| 15 | ObservabilityConfig, |
| 16 | ParallelConfig, |
| 17 | SchedulerConfig, |
| 18 | SpeculativeConfig, |
| 19 | VllmConfig, |
| 20 | ) |
| 21 | from vllm.config.scheduler import SchedulerPolicy |
| 22 | from vllm.multimodal.inputs import ( |
| 23 | MultiModalFeatureSpec, |
| 24 | MultiModalKwargsItem, |
| 25 | PlaceholderRange, |
| 26 | ) |
| 27 | from vllm.sampling_params import SamplingParams |
| 28 | from vllm.utils.hashing import sha256 |
| 29 | from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash |
| 30 | from vllm.v1.core.sched.async_scheduler import AsyncScheduler |
| 31 | from vllm.v1.core.sched.scheduler import Scheduler |
| 32 | from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs |
| 33 | from vllm.v1.kv_cache_interface import ( |
| 34 | FullAttentionSpec, |
| 35 | KVCacheConfig, |
| 36 | KVCacheGroupSpec, |
| 37 | KVCacheSpec, |
| 38 | ) |
| 39 | from vllm.v1.request import Request |
| 40 | from vllm.v1.structured_output import StructuredOutputManager |
| 41 | |
| 42 | EOS_TOKEN_ID = 50256 |
| 43 | |
| 44 | |
| 45 | def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int = 0): |
| 46 | return MockKVConfig( |
| 47 | matched_tokens=matched_tokens, |
| 48 | is_async=is_async, |
| 49 | num_defers_before_matching=num_defers_before_matching, |
| 50 | ) |
| 51 | |
| 52 | |
| 53 | def create_scheduler( |
| 54 | model: str = "facebook/opt-125m", |
| 55 | max_num_seqs: int = 16, |
| 56 | max_num_active_seqs: int | None = None, |
| 57 | max_num_batched_tokens: int = 8192, |
| 58 | enable_chunked_prefill: bool = True, |
| 59 | enable_prefix_caching: bool = False, |
| 60 | long_prefill_token_threshold: int = 0, |
| 61 | long_prefill_token_threshold_adaptive: bool = False, |
| 62 | disable_chunked_mm_input: bool = False, |
| 63 | mm_encoder_only: bool = False, |
| 64 | use_kv_connector: None | bool | str | MockKVConfig = None, |
| 65 | kv_role: str = "kv_both", |
| 66 | num_blocks: int = 10000, |
| 67 | block_size: int = 16, |
| 68 | max_model_len: int | None = None, |
| 69 | num_speculative_tokens: int | None = None, |
| 70 | speculative_method: str | None = None, |
| 71 | parallel_drafting: bool = False, |
| 72 | skip_tokenizer_init: bool = False, |
| 73 | async_scheduling: bool = False, |
| 74 | pipeline_parallel_size: int = 1, |
| 75 | data_parallel_size: int = 1, |
| 76 | num_speculative_tokens_per_batch_size: list[tuple[int, int, int]] | None = None, |
| 77 | use_ec_connector: bool = False, |
| 78 | ec_role: str | None = None, |
| 79 | use_v2_model_runner: bool | None = None, |
| 80 | kv_cache_spec: KVCacheSpec | None = None, |
| 81 | per_request_spec_decode_metrics: str = "none", |
| 82 | scheduling_policy: SchedulerPolicy = "fcfs", |
| 83 | diffusion_canvas_length: int | None = None, |
| 84 | scheduler_cls: type[Scheduler] | None = None, |
| 85 | ) -> Scheduler | AsyncScheduler: |
| 86 | """Create scheduler under test. |
| 87 | |
| 88 | Args: |
| 89 | model: model under test |
| 90 | max_num_seqs: max sequences to schedule |
| 91 | max_num_batch_tokens: max num tokens to batch |
| 92 | enable_prefix_caching: optionally force APC config |
| 93 | (True/False) or use default |
| 94 | (False) |
| 95 | long_prefill_token_threshold: cap on prefill chunk size |
| 96 | long_prefill_token_threshold_adaptive: floor the cap at a |
| 97 | fair share of the token budget |
| 98 | |
| 99 | Returns: |
| 100 | {class}`Scheduler` instance |
| 101 | |
| 102 | """ |
| 103 | model_config = ModelConfig( |
| 104 | model=model, |
| 105 | trust_remote_code=True, |
| 106 | dtype="float16", |
| 107 | seed=42, |
| 108 | skip_tokenizer_init=skip_tokenizer_init, |
| 109 | # The scheduler reads model_config.max_model_len, not the |
| 110 | # SchedulerConfig one, so both must agree. |
| 111 | max_model_len=max_model_len, |
| 112 | ) |
| 113 | if mm_encoder_only: |
| 114 | model_config.multimodal_config.mm_encoder_only = True |
| 115 | if use_ec_connector and ec_role == "ec_producer": |
| 116 | model_config.multimodal_config = MultiModalConfig() |
| 117 | if max_model_len is None: |
| 118 | max_model_len = max_num_batched_tokens |
| 119 | scheduler_config = SchedulerConfig( |
| 120 | max_num_seqs=max_num_seqs, |
| 121 | max_num_active_seqs=max_num_active_seqs, |
| 122 | max_num_batched_tokens=max_num_batched_tokens, |
| 123 | max_model_len=max_model_len, |
| 124 | long_prefill_token_threshold=long_prefill_token_threshold, |
| 125 | long_prefill_token_threshold_adaptive=(long_prefill_token_threshold_adaptive), |
| 126 | disable_chunked_mm_input=disable_chunked_mm_input, |
| 127 | enable_chunked_prefill=enable_chunked_prefill, |
| 128 | async_scheduling=async_scheduling, |
| 129 | is_encoder_decoder=model_config.is_encoder_decoder, |
| 130 | # Ensure admission/preemption mechanics are deterministic |
| 131 | watermark=0.0, |
| 132 | policy=scheduling_policy, |
| 133 | ) |
| 134 | # Cache config, optionally force APC |
| 135 | cache_config = CacheConfig( |
| 136 | block_size=block_size, |
| 137 | gpu_memory_utilization=0.9, |
| 138 | cache_dtype="auto", |
| 139 | enable_prefix_caching=enable_prefix_caching, |
| 140 | ) |
| 141 | kv_transfer_config = None |
| 142 | if isinstance(use_kv_connector, MockKVConfig): |
| 143 | kv_transfer_config = KVTransferConfig( |
| 144 | kv_connector="MockKVConnector", |
| 145 | kv_role=kv_role, |
| 146 | kv_connector_extra_config={ |
| 147 | "matched_tokens": use_kv_connector.matched_tokens, |
| 148 | "is_async": use_kv_connector.is_async, |
| 149 | "num_defers_before_matching": ( |
| 150 | use_kv_connector.num_defers_before_matching |
| 151 | ), |
| 152 | }, |
| 153 | ) |
| 154 | elif isinstance(use_kv_connector, str): |
| 155 | kv_transfer_config = KVTransferConfig( |
| 156 | kv_connector=use_kv_connector, |
| 157 | kv_role=kv_role, |
| 158 | ) |
| 159 | elif use_kv_connector: |
| 160 | kv_transfer_config = KVTransferConfig( |
| 161 | kv_connector="ExampleConnector", |
| 162 | kv_role=kv_role, |
| 163 | kv_connector_extra_config={"shared_storage_path": "local_storage"}, |
| 164 | ) |
| 165 | |
| 166 | speculative_config: SpeculativeConfig | None = None |
| 167 | if num_speculative_tokens is not None: |
| 168 | spec_kwargs: dict = dict( |
| 169 | model="ngram", num_speculative_tokens=num_speculative_tokens |
| 170 | ) |
| 171 | if num_speculative_tokens_per_batch_size is not None: |
| 172 | spec_kwargs["num_speculative_tokens_per_batch_size"] = ( |
| 173 | num_speculative_tokens_per_batch_size |
| 174 | ) |
| 175 | if speculative_method is not None: |
| 176 | spec_kwargs["method"] = speculative_method |
| 177 | spec_kwargs["prompt_lookup_max"] = num_speculative_tokens |
| 178 | spec_kwargs["prompt_lookup_min"] = 1 |
| 179 | speculative_config = SpeculativeConfig(**spec_kwargs) |
| 180 | speculative_config.parallel_drafting = parallel_drafting |
| 181 | |
| 182 | ec_transfer_config = ( |
| 183 | ECTransferConfig( |
| 184 | ec_connector="ECExampleConnector", |
| 185 | ec_role=ec_role, |
| 186 | ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test"}, |
| 187 | ) |
| 188 | if use_ec_connector |
| 189 | else None |
| 190 | ) |
| 191 | |
| 192 | diffusion_config: DiffusionConfig | None = None |
| 193 | if diffusion_canvas_length is not None: |
| 194 | # A diffusion checkpoint declares its canvas in the HF config. |
| 195 | model_config.hf_config.canvas_length = diffusion_canvas_length |
| 196 | diffusion_config = DiffusionConfig(canvas_length=diffusion_canvas_length) |
| 197 | |
| 198 | vllm_config = VllmConfig( |
| 199 | scheduler_config=scheduler_config, |
| 200 | model_config=model_config, |
| 201 | cache_config=cache_config, |
| 202 | parallel_config=ParallelConfig( |
| 203 | pipeline_parallel_size=pipeline_parallel_size, |
| 204 | data_parallel_size=data_parallel_size, |
| 205 | ), |
| 206 | kv_transfer_config=kv_transfer_config, |
| 207 | speculative_config=speculative_config, |
| 208 | diffusion_config=diffusion_config, |
| 209 | ec_transfer_config=ec_transfer_config, |
| 210 | observability_config=ObservabilityConfig( |
| 211 | per_request_spec_decode_metrics=per_request_spec_decode_metrics, |
| 212 | ), |
| 213 | ) |
| 214 | if kv_cache_spec is None: |
| 215 | kv_cache_spec = FullAttentionSpec( |
| 216 | block_size=block_size, |
| 217 | num_kv_heads=1, |
| 218 | head_size=1, |
| 219 | dtype=torch.float32, |
| 220 | ) |
| 221 | kv_cache_config = KVCacheConfig( |
| 222 | num_blocks=num_blocks, # A large number of blocks to hold all requests |
| 223 | kv_cache_tensors=[], |
| 224 | kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], |
| 225 | ) |
| 226 | cache_config.num_gpu_blocks = num_blocks |
| 227 | register_all_kvcache_specs(vllm_config) |
| 228 | if scheduler_cls is None: |
| 229 | scheduler_cls = AsyncScheduler if async_scheduling else Scheduler |
| 230 | scheduler = scheduler_cls( |
| 231 | vllm_config=vllm_config, |
| 232 | kv_cache_config=kv_cache_config, |
| 233 | block_size=block_size, |
| 234 | log_stats=True, |
| 235 | structured_output_manager=StructuredOutputManager(vllm_config), |
| 236 | ) |
| 237 | if use_v2_model_runner is None: |
| 238 | use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER) |
| 239 | scheduler.use_v2_model_runner = use_v2_model_runner |
| 240 | return scheduler |
| 241 | |
| 242 | |
| 243 | _none_hash_initialized = False |
| 244 | |
| 245 | |
| 246 | def create_requests( |
| 247 | num_requests: int, |
| 248 | num_tokens: int = 10, |
| 249 | mm_hashes_list: list[list[str]] | None = None, |
| 250 | mm_positions: list[list[PlaceholderRange]] | None = None, |
| 251 | ignore_eos: bool = False, |
| 252 | max_tokens: int = 16, |
| 253 | stop_token_ids: list[int] | None = None, |
| 254 | prompt_logprobs: int | None = None, |
| 255 | same_prompt: bool = False, |
| 256 | block_size: int = 16, |
| 257 | req_ids: list[str] | None = None, |
| 258 | ) -> list[Request]: |
| 259 | global _none_hash_initialized |
| 260 | if not _none_hash_initialized: |
| 261 | init_none_hash(sha256) |
| 262 | _none_hash_initialized = True |
| 263 | |
| 264 | block_hasher = get_request_block_hasher(block_size, sha256) |
| 265 | sampling_params = SamplingParams( |
| 266 | ignore_eos=ignore_eos, |
| 267 | max_tokens=max_tokens, |
| 268 | stop_token_ids=stop_token_ids, |
| 269 | prompt_logprobs=prompt_logprobs, |
| 270 | ) |
| 271 | sampling_params.update_from_generation_config({}, EOS_TOKEN_ID) |
| 272 | requests = [] |
| 273 | |
| 274 | if mm_hashes_list is not None: |
| 275 | # NOTE: allow manual input; some mm items can have the same identifier |
| 276 | # no. of mm_hashes and mm_positions for each request should be identical |
| 277 | assert mm_positions is not None, ( |
| 278 | "mm_positions must be provided when mm_hashes_list is provided" |
| 279 | ) |
| 280 | assert len(mm_hashes_list) == len(mm_positions) == num_requests |
| 281 | assert [len(h) for h in mm_hashes_list] == [len(p) for p in mm_positions] |
| 282 | |
| 283 | # Since same identifier would imply they are identical encoder output |
| 284 | # Verify mm items with identical identifier are having mm_position.length |
| 285 | seen_hashes: dict[str, int] = {} |
| 286 | |
| 287 | if req_ids: |
| 288 | assert len(req_ids) == num_requests |
| 289 | else: |
| 290 | req_ids = [f"{i}" for i in range(num_requests)] |
| 291 | |
| 292 | for i in range(num_requests): |
| 293 | mm_features = [] |
| 294 | |
| 295 | for j, position in enumerate( |
| 296 | mm_positions[i] if mm_positions is not None else [] |
| 297 | ): |
| 298 | if mm_hashes_list is not None: |
| 299 | identifier = mm_hashes_list[i][j] |
| 300 | |
| 301 | # Verify if position length is identical |
| 302 | position_length = position.length |
| 303 | if identifier in seen_hashes: |
| 304 | assert seen_hashes[identifier] == position_length, ( |
| 305 | f"mm_hash '{identifier}' has inconsistent position lengths: " |
| 306 | f"previously {seen_hashes[identifier]}, now {position_length} " |
| 307 | f"at request {i}, position {j}" |
| 308 | ) |
| 309 | else: |
| 310 | seen_hashes[identifier] = position_length |
| 311 | else: |
| 312 | # Unique dummy hash for each mm item |
| 313 | identifier = f"hash{i}_{j}" |
| 314 | mm_feature = MultiModalFeatureSpec( |
| 315 | data=MultiModalKwargsItem.dummy(), |
| 316 | mm_position=position, |
| 317 | identifier=identifier, |
| 318 | modality="image", |
| 319 | ) |
| 320 | mm_features.append(mm_feature) |
| 321 | |
| 322 | prompt_token_ids = [0] * num_tokens if same_prompt else [i] * num_tokens |
| 323 | request = Request( |
| 324 | request_id=req_ids[i], |
| 325 | prompt_token_ids=prompt_token_ids, |
| 326 | sampling_params=sampling_params, |
| 327 | pooling_params=None, |
| 328 | mm_features=mm_features if mm_features else None, |
| 329 | block_hasher=block_hasher, |
| 330 | ) |
| 331 | requests.append(request) |
| 332 | return requests |