kv_cache_interface.py 14.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import copy
5
from dataclasses import dataclass, fields, replace
Chen Zhang's avatar
Chen Zhang committed
6
from math import prod
7
8

import torch
9
from typing_extensions import Self
10

11
from vllm.config import VllmConfig
12
from vllm.logger import init_logger
13
from vllm.utils.math_utils import cdiv
14
from vllm.utils.torch_utils import get_dtype_size
15
16
17
18

logger = init_logger(__name__)


19
@dataclass(frozen=True)
20
class KVCacheSpec:
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
    """
    A base class for specifying the KV cache format of one layer.
    """

    # number of tokens in a block
    block_size: int

    @property
    def page_size_bytes(self) -> int:
        """
        The size of a page with `block_size` tokens in bytes.

        Returns:
            The page size
        """
        raise NotImplementedError

38
    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
39
        """
40
        The maximum possible memory usage of this KV cache in bytes.
41
42

        Returns:
43
            The KV cache size in bytes
44
45
46
        """
        raise NotImplementedError

47
48
49
50
51
52
    def copy_with_new_block_size(self, block_size: int) -> Self:
        """
        Create a new KVCacheSpec from self but replacing the block size.
        """
        return replace(self, block_size=block_size)

53
54
55
56
57
    @classmethod
    def merge(cls, specs: list[Self]) -> Self:
        """
        Merge a list of KVCacheSpec objects into a single KVCacheSpec object.
        """
58
        assert all(spec == specs[0] for spec in specs[1:]), (
59
60
            "All layers in the same KV cache group must be the same."
        )
61
62
        return copy.deepcopy(specs[0])

63

64
@dataclass(frozen=True)
65
class AttentionSpec(KVCacheSpec):
66
67
68
69
70
71
    num_kv_heads: int
    head_size: int
    dtype: torch.dtype

    @property
    def page_size_bytes(self) -> int:
72
73
74
75
76
77
78
        return (
            2
            * self.block_size
            * self.num_kv_heads
            * self.head_size
            * get_dtype_size(self.dtype)
        )
79

80

81
@dataclass(frozen=True)
82
class FullAttentionSpec(AttentionSpec):
83
    """
84
85
86
87
    When hybrid allocator is disabled and the model contains both full
    attention layers and sliding window attention layers, sliding
    window attention are regarded as full attention in KV cache manager
    (blocks are allocated for all tokens), while computed as sliding window
88
89
    attention in model runner.
    In this case, we use FullAttentionSpec and record the sliding window size.
90
91
92
93
    """

    sliding_window: int | None = None
    """
94
95
    Default to None for not using sliding window attention.
    """
96
    attention_chunk_size: int | None = None
97
98
99

    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
        max_model_len = vllm_config.model_config.max_model_len
100
        dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
101
        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
102
103
        # Note(hc): each dcp rank only need save
        # (max_model_len//dcp_world_size) tokens locally.
104
105
        if dcp_world_size * pcp_world_size > 1:
            max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size)
106
107
        return cdiv(max_model_len, self.block_size) * self.page_size_bytes

108
    @classmethod
109
    def merge_window_sizes(cls, window_sizes: set[int]) -> int | None:
110
111
112
113
114
115
116
        if len(window_sizes) == 0:
            return None
        elif len(window_sizes) == 1:
            return window_sizes.pop()
        else:
            raise ValueError(
                "All attention layers in the same KV cache group must have the "
117
118
                "same window size."
            )
119

120
121
122
    @classmethod
    def merge(cls, specs: list[Self]) -> Self:
        """
123
        Merge a list of FullAttentionSpec objects into a single
124
125
        FullAttentionSpec object.
        """
126
        assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (
127
128
            "All attention layers in the same KV cache group must be FullAttentionSpec."
        )
129

130
131
132
133
134
135
136
137
        sliding_window = set(
            spec.sliding_window for spec in specs if spec.sliding_window is not None
        )
        attention_chunk_size = set(
            spec.attention_chunk_size
            for spec in specs
            if spec.attention_chunk_size is not None
        )
138
        assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), (
139
140
            "MLAAttentionSpec should be merged in MLAAttentionSpec.merge"
        )
141
142
143
144
145
146
147
148
149
150
151
152
        merged_spec = cls(
            block_size=specs[0].block_size,
            num_kv_heads=specs[0].num_kv_heads,
            head_size=specs[0].head_size,
            dtype=specs[0].dtype,
            sliding_window=cls.merge_window_sizes(sliding_window),
            attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
        )
        for spec in specs:
            for f in fields(AttentionSpec):
                assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
                    "All attention layers in the same KV cache group must have "
153
154
155
156
157
158
159
160
                    "the same attention spec."
                )
        assert (merged_spec.sliding_window is not None) + (
            merged_spec.attention_chunk_size is not None
        ) <= 1, (
            "Model with both sliding window layers and chunked local attention "
            "layers is not supported."
        )
161
162
        return merged_spec

163

164
165
166
@dataclass(frozen=True)
class MLAAttentionSpec(FullAttentionSpec):
    # TODO(Lucas/Chen): less hacky way to do this
167
    cache_dtype_str: str | None = None
168
169
170
171
172
173
174

    @property
    def page_size_bytes(self) -> int:
        if self.cache_dtype_str == "fp8_ds_mla":
            # See `vllm/v1/attention/backends/mla/flashmla_sparse.py`
            #  for details.
            return self.block_size * 656
175
176
177
178
179
180
        return (
            self.block_size
            * self.num_kv_heads
            * self.head_size
            * get_dtype_size(self.dtype)
        )
181
182
183
184

    @classmethod
    def merge(cls, specs: list[Self]) -> Self:
        assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), (
185
186
            "All attention layers in the same KV cache group must be MLAAttentionSpec."
        )
187
188
189
        cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
        assert len(cache_dtype_str_set) == 1, (
            "All attention layers in the same KV cache group must use the same "
190
191
            "quantization method."
        )
192
193
194
195
196
197
198
199
200
        return cls(
            block_size=specs[0].block_size,
            num_kv_heads=specs[0].num_kv_heads,
            head_size=specs[0].head_size,
            dtype=specs[0].dtype,
            cache_dtype_str=cache_dtype_str_set.pop(),
        )


201
@dataclass(frozen=True)
202
203
204
class ChunkedLocalAttentionSpec(AttentionSpec):
    attention_chunk_size: int

205
206
    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
        max_model_len = vllm_config.model_config.max_model_len
207
        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
208
209
210
211
212

        # During chunked prefill, we allocate KV cache for at most
        # `self.attention_chunk_size` computed tokens plus the newly scheduled
        # tokens. And we won't allocate KV cache for more than `max_model_len`
        # tokens.
213
214
215
        num_tokens = min(
            self.attention_chunk_size + max_num_batched_tokens, max_model_len
        )
216
217
218

        return cdiv(num_tokens, self.block_size) * self.page_size_bytes

219

220
@dataclass(frozen=True)
221
222
223
224
class SlidingWindowSpec(AttentionSpec):
    sliding_window: int

    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
225
        assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
226
            "DCP not support sliding window."
227
        )
228
        max_model_len = vllm_config.model_config.max_model_len
229
        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
230
231
232
233
234

        # During chunked prefill, we allocate KV cache for the last
        # `self.sliding_window-1` computed tokens plus the newly scheduled
        # tokens. And we won't allocate KV cache for more than `max_model_len`
        # tokens.
235
236
237
        num_tokens = min(
            self.sliding_window - 1 + max_num_batched_tokens, max_model_len
        )
238
239
240
241
242
243

        # +1 here because the sliding window may not start from the beginning
        # of the block. For example, if the block size is 4 and num_token
        # is 4, we need two blocks [XXCD] [EF] to store the sliding
        # window [CDEF] of 6 tokens.
        return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes
244
245


246
@dataclass(frozen=True)
Chen Zhang's avatar
Chen Zhang committed
247
248
class MambaSpec(KVCacheSpec):
    shapes: tuple[tuple[int, ...], ...]
249
    dtypes: tuple[torch.dtype]
250
    page_size_padded: int | None = None
251
    mamba_type: str = "mamba2"
252
    num_speculative_blocks: int = 0
Chen Zhang's avatar
Chen Zhang committed
253
254
255

    @property
    def page_size_bytes(self) -> int:
256
257
        page_size = sum(
            prod(shape) * get_dtype_size(dtype)
258
259
            for (shape, dtype) in zip(self.shapes, self.dtypes)
        )
260
261
262
263
        if self.page_size_padded is not None:
            assert self.page_size_padded >= page_size
            return self.page_size_padded
        return page_size
Chen Zhang's avatar
Chen Zhang committed
264
265

    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
266
267
        max_model_len = vllm_config.model_config.max_model_len
        return cdiv(max_model_len, self.block_size) * self.page_size_bytes
Chen Zhang's avatar
Chen Zhang committed
268
269


270
271
272
273
274
275
276
@dataclass(frozen=True)
class EncoderOnlyAttentionSpec(AttentionSpec):
    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
        # Encoder-only layers do not need KV cache
        return 0


277
278
279
280
281
282
283
284
285
@dataclass(frozen=True)
class CrossAttentionSpec(AttentionSpec):
    """
    KV cache spec for cross-attention layers in encoder-decoder models.
    """

    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
        # For cross-attention, we need to cache encoder states
        # Get encoder length (e.g., 1500 for Whisper).
286
        max_encoder_len = vllm_config.scheduler_config.max_num_encoder_input_tokens
287
288
289
        return cdiv(max_encoder_len, self.block_size) * self.page_size_bytes


290
291
292
293
294
295
296
297
@dataclass(frozen=True)
class UniformTypeKVCacheSpecs(KVCacheSpec):
    """
    A KV cache spec for multiple layers with the same type of attention. Here,
    same types means always need the same number of token slots. For example,
    sliding window attentions with different window sizes are not the same type
    and should not be merged into one UniformTypeKVCacheSpecs.
    """
298

299
300
301
302
    kv_cache_specs: dict[str, KVCacheSpec]

    @property
    def page_size_bytes(self) -> int:
303
        return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values())
304
305
306

    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
        max_num_pages = max(
307
308
309
            cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
            for spec in self.kv_cache_specs.values()
        )
310
311
312
313
314
315
316
317
318
319
320
321
        return max_num_pages * self.page_size_bytes

    @classmethod
    def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
        """
        Whether all layers have the same type of KV cache spec.
        """
        block_sizes = set(spec.block_size for spec in kv_cache_specs.values())
        if len(block_sizes) > 1:
            # Different block sizes, not uniform.
            return False
        one_spec = next(iter(kv_cache_specs.values()))
322
323
        if isinstance(one_spec, FullAttentionSpec):
            return all(
324
325
                isinstance(spec, FullAttentionSpec) for spec in kv_cache_specs.values()
            )
326
        elif isinstance(one_spec, CrossAttentionSpec):
327
            return all(
328
329
                isinstance(spec, CrossAttentionSpec) for spec in kv_cache_specs.values()
            )
330
331
332
333
        elif isinstance(one_spec, SlidingWindowSpec):
            return all(
                isinstance(spec, SlidingWindowSpec)
                and spec.sliding_window == one_spec.sliding_window
334
335
                for spec in kv_cache_specs.values()
            )
336
337
338
339
        elif isinstance(one_spec, ChunkedLocalAttentionSpec):
            return all(
                isinstance(spec, ChunkedLocalAttentionSpec)
                and spec.attention_chunk_size == one_spec.attention_chunk_size
340
341
                for spec in kv_cache_specs.values()
            )
342
343
        elif isinstance(one_spec, MambaSpec):
            return all(
344
345
346
347
                isinstance(spec, MambaSpec)
                and spec.num_speculative_blocks == one_spec.num_speculative_blocks
                for spec in kv_cache_specs.values()
            )
348
349
350
        else:
            # NOTE(Chen): Please add new branches for new KV cache spec types.
            raise NotImplementedError(
351
352
                f"Unsupported KV cache spec type: {type(one_spec)}"
            )
353
354

    @classmethod
355
    def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None:
356
357
358
359
360
361
362
363
364
365
366
        """
        Return a SameTypeKVCacheSpecs object if all layers have the same type
        of KV cache spec. Return None if not.
        """
        if cls.is_uniform_type(kv_cache_specs):
            block_size = next(iter(kv_cache_specs.values())).block_size
            return cls(block_size=block_size, kv_cache_specs=kv_cache_specs)
        else:
            return None


367
368
369
@dataclass
class KVCacheTensor:
    """
370
    A class for specifying how the workers should initialize the KV cache.
371
    """
372

373
374
    size: int  # size of the KV cache tensor in bytes
    shared_by: list[str]  # layer names that share the same KV cache tensor
375
376


377
378
379
380
381
382
@dataclass
class KVCacheGroupSpec:
    """
    Represents a group of model layers that share the same KV cache block table.
    These layers are regarded as one layer in the KV cache manager.
    """
383

384
385
386
387
388
389
    # The names of model layers in this group
    layer_names: list[str]
    # The KV cache spec of this manager layer
    kv_cache_spec: KVCacheSpec


390
391
392
393
394
@dataclass
class KVCacheConfig:
    """
    The KV cache configuration of a model.
    """
395

396
    num_blocks: int
397
    """The number of KV cache blocks"""
398
    kv_cache_tensors: list[KVCacheTensor]
399
400
    """How should model runner initialize the KV cache tensors for each layer"""
    kv_cache_groups: list[KVCacheGroupSpec]
401
    """
402
    The kv cache groups of the model.
403
404
405
406
    For models with only one type of attention, there is only one group that
    contains all layers.
    For models with multiple types of attention, there will be multiple groups,
    see `_get_kv_cache_config_uniform_page_size` for more details.
407
    """