utils.py 14 KB
Newer Older
Robert Shaw's avatar
Robert Shaw committed
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import tempfile
from collections import defaultdict
5
from collections.abc import Callable
6
7
from dataclasses import dataclass
from itertools import chain, count
8
from typing import Any, Literal
Robert Shaw's avatar
Robert Shaw committed
9
10
11
12

import torch

from vllm import SamplingParams
13
from vllm.config import (
14
    AttentionConfig,
15
16
17
18
19
20
21
22
    CacheConfig,
    DeviceConfig,
    KVTransferConfig,
    ModelConfig,
    SchedulerConfig,
    VllmConfig,
)
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
23
24
25
26
27
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
    KVConnectorBase_V1,
    KVConnectorMetadata,
    KVConnectorRole,
)
28
29
from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (  # noqa
    ExampleConnector,
30
)
31
from vllm.utils.hashing import sha256
32
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
33
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
34
from vllm.v1.core.sched.scheduler import Scheduler, SchedulerOutput
35
36
37
38
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
39
    SlidingWindowSpec,
40
)
41
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
Robert Shaw's avatar
Robert Shaw committed
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
from vllm.v1.request import Request
from vllm.v1.structured_output import StructuredOutputManager

EOS_TOKEN_ID = 50256


def assert_scheduler_empty(scheduler: Scheduler):
    """Confirm the scheduler is "empty" - i.e. no leaks."""
    # Scheduler Metadata.
    assert len(scheduler.requests) == 0
    assert len(scheduler.waiting) == 0
    assert len(scheduler.running) == 0
    assert len(scheduler.finished_req_ids) == 0
    assert len(scheduler.finished_recving_kv_req_ids) == 0

    # EncoderCacheManager.
    assert len(scheduler.encoder_cache_manager.freed) == 0
    assert len(scheduler.encoder_cache_manager.cached) == 0

    # KVCache Manager.
62
63
64
65
66
67
68
69
70
71
72
73
74
75
    assert (
        len(
            scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks
        )
        == 0
    )
    assert (
        len(
            scheduler.kv_cache_manager.coordinator.single_type_managers[
                0
            ].num_cached_block
        )
        == 0
    )
Robert Shaw's avatar
Robert Shaw committed
76
    num_free_blocks = (
77
78
79
        scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks
    )
    assert num_free_blocks == (scheduler.kv_cache_manager.block_pool.num_gpu_blocks - 1)
Robert Shaw's avatar
Robert Shaw committed
80
81
82
83
84
85
86
87
88
89
90
91

    # NOTE(rob): just the ref count on blocks will be 0. The hash
    # value, etc will remain since we lazily evict for prefix cache.
    for block in scheduler.kv_cache_manager.block_pool.blocks:
        assert block.ref_cnt == 0


def create_vllm_config(
    model: str = "facebook/opt-125m",
    max_num_seqs: int = 16,
    max_num_batched_tokens: int = 64,
    block_size: int = 16,
92
93
    max_model_len: int = 10000,
    enable_chunked_prefill: bool = True,
94
    enable_permute_local_kv: bool = False,
95
96
97
98
    kv_connector_extra_config: dict[str, Any] | None = None,
    dtype: str = "float16",
    cache_dtype: str = "auto",
    hf_overrides: dict[str, Any] | None = None,
99
    attention_backend: str | None = None,
100
    kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
Robert Shaw's avatar
Robert Shaw committed
101
102
103
104
105
) -> VllmConfig:
    """Initialize VllmConfig For Testing."""
    model_config = ModelConfig(
        model=model,
        trust_remote_code=True,
106
        dtype=dtype,
Robert Shaw's avatar
Robert Shaw committed
107
        seed=42,
108
        hf_overrides=hf_overrides or {},
Robert Shaw's avatar
Robert Shaw committed
109
    )
110
111
112
113
114
115
116
    scheduler_config = SchedulerConfig(
        max_num_seqs=max_num_seqs,
        max_num_batched_tokens=max_num_batched_tokens,
        max_model_len=max_model_len,
        enable_chunked_prefill=enable_chunked_prefill,
        is_encoder_decoder=model_config.is_encoder_decoder,
    )
Robert Shaw's avatar
Robert Shaw committed
117
118
119
120
121
    # Cache config, optionally force APC
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
        swap_space=0,
122
        cache_dtype=cache_dtype,
Robert Shaw's avatar
Robert Shaw committed
123
124
125
126
127
        enable_prefix_caching=True,
    )
    kv_transfer_config = KVTransferConfig(
        kv_connector="NixlConnector",
        kv_role="kv_both",
128
        enable_permute_local_kv=enable_permute_local_kv,
129
        kv_connector_extra_config=kv_connector_extra_config or {},
130
        kv_load_failure_policy=kv_load_failure_policy,
Robert Shaw's avatar
Robert Shaw committed
131
    )
132
    attention_config = AttentionConfig(backend=attention_backend)
133
134
135
136
137
138
    return VllmConfig(
        scheduler_config=scheduler_config,
        model_config=model_config,
        cache_config=cache_config,
        kv_transfer_config=kv_transfer_config,
        device_config=DeviceConfig("cpu"),
139
        attention_config=attention_config,
140
    )
Robert Shaw's avatar
Robert Shaw committed
141
142
143
144
145


def create_scheduler(
    vllm_config: VllmConfig,
    num_blocks: int = 10000,
146
    kv_cache_config: KVCacheConfig | None = None,
Robert Shaw's avatar
Robert Shaw committed
147
148
149
) -> Scheduler:
    """Initialize Scheduler For Testing."""
    block_size = vllm_config.cache_config.block_size
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
    if kv_cache_config is None:
        kv_cache_config = KVCacheConfig(
            num_blocks=num_blocks,  # A large number of blocks to hold all requests
            kv_cache_tensors=[],
            kv_cache_groups=[
                KVCacheGroupSpec(
                    ["layer"],
                    FullAttentionSpec(
                        block_size=block_size,
                        num_kv_heads=1,
                        head_size=1,
                        dtype=torch.float32,
                    ),
                )
            ],
        )
Robert Shaw's avatar
Robert Shaw committed
166
167
168
169
170
171
    vllm_config.cache_config.num_gpu_blocks = num_blocks
    return Scheduler(
        vllm_config=vllm_config,
        kv_cache_config=kv_cache_config,
        log_stats=True,
        structured_output_manager=StructuredOutputManager(vllm_config),
172
        block_size=block_size,
Robert Shaw's avatar
Robert Shaw committed
173
174
175
    )


176
_request_count = count(1)
177
178
179
_none_hash_initialized = False


180
def create_request(
181
    request_id: int | None = None,
182
183
184
185
186
187
188
189
190
    num_tokens: int = 10,
    common_prefix_len=0,
    max_tokens: int = 16,
    do_remote_decode: bool = False,
    do_remote_prefill: bool = False,
    num_remote_blocks: int = 3,
    block_size: int = 16,
    hash_fn: Callable = sha256,
) -> Request:
Robert Shaw's avatar
Robert Shaw committed
191
    """Make dummy request for testing."""
192
193
194
195
196
    assert num_tokens >= common_prefix_len >= 0

    if request_id is None:
        request_id = next(_request_count)

197
198
    global _none_hash_initialized
    if not _none_hash_initialized:
199
        init_none_hash(hash_fn)
200
        _none_hash_initialized = True
Robert Shaw's avatar
Robert Shaw committed
201

202
    kv_transfer_params: dict[str, Any] | None = None
203

Robert Shaw's avatar
Robert Shaw committed
204
205
    if do_remote_decode:
        assert not do_remote_prefill
206
        kv_transfer_params = dict(do_remote_prefill=False, do_remote_decode=True)
Robert Shaw's avatar
Robert Shaw committed
207
    elif do_remote_prefill:
208
209
210
211
        kv_transfer_params = dict(
            do_remote_prefill=True,
            do_remote_decode=False,
            remote_engine_id="my-engine-id",
212
            remote_request_id=f"prefill-{request_id}",
213
214
215
216
            remote_block_ids=list(range(num_remote_blocks)),
            remote_host="my-host",
            remote_port=1234,
        )
Robert Shaw's avatar
Robert Shaw committed
217
218
219

    max_tokens = 1 if do_remote_decode else max_tokens
    sampling_params = SamplingParams(max_tokens=max_tokens)
220
    sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
Robert Shaw's avatar
Robert Shaw committed
221

222
223
224
    common_prefix = [1] * common_prefix_len if common_prefix_len > 0 else []
    suffix = [i * request_id for i in range(num_tokens - common_prefix_len)]
    prompt_token_ids = common_prefix + suffix
Robert Shaw's avatar
Robert Shaw committed
225
226
227
228
229

    req = Request(
        request_id=f"id-{request_id}",
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
230
        pooling_params=None,
231
        mm_features=None,
232
        block_hasher=get_request_block_hasher(block_size, hash_fn),
Robert Shaw's avatar
Robert Shaw committed
233
234
235
236
237
238
239
    )
    req.kv_transfer_params = kv_transfer_params
    return req


def create_model_runner_output(
    reqs: list[Request],
240
241
242
    finished_sending: set[str] | None = None,
    finished_recving: set[str] | None = None,
    invalid_block_ids: set[int] | None = None,
Robert Shaw's avatar
Robert Shaw committed
243
    use_eos: bool = False,
244
    token_id: int = 0,
Robert Shaw's avatar
Robert Shaw committed
245
246
247
248
249
250
251
252
) -> ModelRunnerOutput:
    """Make dummy model runner output for testing."""

    # Make request data.
    req_ids = [req.request_id for req in reqs]
    req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)}

    # Make sampled tokens.
253
    sampled_token = EOS_TOKEN_ID if use_eos else token_id
254
    sampled_token_ids = [[sampled_token] for _ in req_ids]
Robert Shaw's avatar
Robert Shaw committed
255

256
257
258
259
260
261
262
263
    kv_connector_output = (
        None
        if (
            finished_sending is None
            and finished_recving is None
            and invalid_block_ids is None
        )
        else KVConnectorOutput(
264
265
            finished_sending=finished_sending,
            finished_recving=finished_recving,
266
            invalid_block_ids=invalid_block_ids or set(),
267
        )
268
    )
269

Robert Shaw's avatar
Robert Shaw committed
270
271
272
273
274
275
276
    # Make output data structure.
    return ModelRunnerOutput(
        req_ids=req_ids,
        req_id_to_index=req_id_to_index,
        sampled_token_ids=sampled_token_ids,
        logprobs=None,
        prompt_logprobs_dict={},
277
        pooler_output=None,
278
        kv_connector_output=kv_connector_output,
Robert Shaw's avatar
Robert Shaw committed
279
    )
280
281


282
class TestExampleConnector(ExampleConnector):
283
    def __init__(self, config: VllmConfig, role, kv_cache_config):
284
        self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
285
        self._connector = ExampleConnector(config, role)
286
287
        self.call_record: dict[str, int] = defaultdict(int)
        # Use a unique temp file per connector
288
289
290
291
        self._event_file = (
            tempfile.gettempdir()
            + f"/connector_{self.name}-{self.role.name}_events.log"
        )
292
293
294
295
296
        # Start with an empty file
        with open(self._event_file, "w") as _:
            pass

    def __getattribute__(self, name):
297
298
299
300
301
302
303
304
305
306
        if name in (
            "_connector",
            "call_record",
            "name",
            "_event_file",
            "__class__",
            "__dict__",
            "__getattribute__",
            "__init__",
        ):  # avoid recursion
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
            return object.__getattribute__(self, name)
        if not hasattr(self._connector, name):
            return object.__getattribute__(self, name)
        attr = getattr(self._connector, name)

        # Intercept calls to the connector interface and write an event
        # for each one to a file, which can be read back in the main test proc.
        if callable(attr):

            def wrapper(*args, **kwargs):
                self.call_record[name] += 1

                # Include args that we're interested in
                to_log = [name]
                for arg in args:
                    if isinstance(arg, int):
                        to_log.append(str(arg))
                    elif isinstance(arg, KVCacheBlocks):
325
                        to_log.append(f"num_blocks={[len(b) for b in arg.blocks]}")
326
327
328
329

                # Log the event as a line to the file
                try:
                    with open(self._event_file, "a") as f:
330
                        f.write(" ".join(to_log) + "\n")
331
                except Exception as e:
332
                    print(f"[ERROR] Could not log event {name} for {self.name}: {e}")
333
334
335
336
337
338
                return attr(*args, **kwargs)

            return wrapper
        return attr


339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
@dataclass(frozen=True)
class MockKVConfig:
    matched_tokens: int = 0
    is_async: bool = False


class MockKVConnectorMetadata(KVConnectorMetadata):
    def __init__(self):
        # Scheduler tests check metadata.requests
        self.requests: list = []


class MockKVConnector(KVConnectorBase_V1):
    """Mock KV connector for scheduler tests, supporting both sync and async mode."""

    def __init__(
        self,
        vllm_config: VllmConfig,
        role: KVConnectorRole,
        kv_cache_config: KVCacheConfig | None = None,
    ):
        super().__init__(vllm_config, role, kv_cache_config)
        extra_config = self._kv_transfer_config.kv_connector_extra_config
        self.config = MockKVConfig(
            matched_tokens=extra_config["matched_tokens"],
            is_async=extra_config["is_async"],
        )

    def get_num_new_matched_tokens(
        self,
        request: Request,
        num_computed_tokens: int,
    ) -> tuple[int | None, bool]:
        return (self.config.matched_tokens, self.config.is_async)

    def update_state_after_alloc(
        self,
        request: Request,
        blocks: KVCacheBlocks,
        num_external_tokens: int,
    ):
        pass

    def build_connector_meta(
        self, scheduler_output: SchedulerOutput
    ) -> KVConnectorMetadata:
        metadata = MockKVConnectorMetadata()
        cached_reqs = scheduler_output.scheduled_cached_reqs
        for req_id in chain(
            (req.req_id for req in scheduler_output.scheduled_new_reqs),
            (
                req_id
                for req_id in cached_reqs.req_ids
                if req_id in cached_reqs.resumed_req_ids
            ),
        ):
            metadata.requests.append({"req_id": req_id})
        return metadata

    def start_load_kv(self, kv_caches, finished_req_ids):
        pass

    def wait_for_layer_load(self, layer_name):
        pass

    def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
        pass

    def wait_for_save(self):
        pass


411
KVConnectorFactory.register_connector(
412
    "TestExampleConnector", __name__, TestExampleConnector.__name__
413
)
414
415
416
417

KVConnectorFactory.register_connector(
    "MockKVConnector", __name__, MockKVConnector.__name__
)
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452


def make_kv_cache_config(
    block_size: int,
    hma_enabled: bool = False,
    sw_size: int = 128,
    num_blocks: int = 100,
) -> KVCacheConfig:
    kv_cache_groups = [
        KVCacheGroupSpec(
            ["layer0", "layer2"],
            FullAttentionSpec(
                block_size=block_size,
                num_kv_heads=4,
                head_size=16,
                dtype=torch.float16,
            ),
        )
    ]
    if hma_enabled:
        kv_cache_groups.append(
            KVCacheGroupSpec(
                ["layer1", "layer3"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=4,
                    head_size=16,
                    dtype=torch.float16,
                    sliding_window=sw_size,
                ),
            )
        )
    return KVCacheConfig(
        num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
    )