A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Bugfix][V1] Handle MLA in kv_cache_interface (#14462)
Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>
Tyler Michael Smith committed
Mar 8, 2025 at 01:18 UTC
333681408feabb97193880303b23f6571ba39045
3 files changed
+15
-10
vllm/v1/kv_cache_interface.py
+8
-5
@@ -23,9 +23,9 @@ class KVCacheSpecBase:
23
def type_id(self) -> str:
24
"""
25
The type identifier of this KV cache.
26
- Return different strings for layers with different KV cache type (e.g.,
27
- different number of tokens like full attention vs sliding window
28
- attention, different KV cache size per token like layers with different
26
+ Return different strings for layers with different KV cache type (e.g.,
27
+ different number of tokens like full attention vs sliding window
28
+ attention, different KV cache size per token like layers with different
29
number of heads)
30
31
Returns:
@@ -59,6 +59,7 @@ class FullAttentionSpec(KVCacheSpecBase):
59
num_kv_heads: int
60
head_size: int
61
dtype: torch.dtype
62
+ use_mla: bool
63
64
@property
65
def type_id(self) -> str:
@@ -66,7 +67,9 @@ class FullAttentionSpec(KVCacheSpecBase):
67
68
@property
69
def page_size_bytes(self) -> int:
69
- return 2 * self.block_size * self.num_kv_heads * self.head_size \
70
+ # For MLA we only store a single latent vector
71
+ coef = 1 if self.use_mla else 2
72
+ return coef * self.block_size * self.num_kv_heads * self.head_size \
73
* get_dtype_size(self.dtype)
74
75
def bytes_for_tokens(self, num_tokens: int) -> int:
@@ -104,7 +107,7 @@ class KVCacheConfig:
107
2. (not implemented yet) A model with the same number of full attention
108
layers and sliding window attention layers: two groups, one for full
109
attention layers and one for sliding window attention layers.
107
- 3. (not implemented yet) A model with 2 full attention layers and 4 sliding
110
+ 3. (not implemented yet) A model with 2 full attention layers and 4 sliding
111
window attention layers: three groups, (full * 2), (sw * 2), (sw * 2).
112
"""
113
groups: list[list[str]]
vllm/v1/worker/gpu_model_runner.py
+3
-2
@@ -1460,13 +1460,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1460
1461
forward_ctx = self.vllm_config.compilation_config.static_forward_context
1462
block_size = self.vllm_config.cache_config.block_size
1463
+ use_mla = self.vllm_config.model_config.use_mla
1464
kv_cache_spec: KVCacheSpec = {}
1465
for layer_name, attn_module in forward_ctx.items():
1466
if isinstance(attn_module, FusedMoE):
1467
continue
1468
1469
# TODO: Support other attention modules, e.g., sliding window,
1469
- # cross-attention, MLA.
1470
+ # cross-attention
1471
assert isinstance(attn_module, Attention)
1472
if attn_module.attn_type == AttentionType.DECODER:
1473
kv_cache_spec[layer_name] = FullAttentionSpec(
@@ -1474,7 +1475,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1475
num_kv_heads=attn_module.num_kv_heads,
1476
head_size=attn_module.head_size,
1477
dtype=attn_module.dtype,
1477
- )
1478
+ use_mla=use_mla)
1479
elif attn_module.attn_type in (AttentionType.ENCODER,
1480
AttentionType.ENCODER_ONLY):
1481
# encoder-only attention does not need KV cache.
vllm/v1/worker/tpu_model_runner.py
+4
-3
@@ -303,10 +303,10 @@ class TPUModelRunner:
303
304
def get_kv_cache_spec(self) -> KVCacheSpec:
305
"""
306
- Generates the KVCacheSpec by parsing the kv cache format from each
306
+ Generates the KVCacheSpec by parsing the kv cache format from each
307
Attention module in the static forward context.
308
Returns:
309
- KVCacheSpec: A dictionary mapping layer names to their KV cache
309
+ KVCacheSpec: A dictionary mapping layer names to their KV cache
310
format. Layers that do not need KV cache are not included.
311
"""
312
@@ -323,6 +323,7 @@ class TPUModelRunner:
323
num_kv_heads=attn_module.num_kv_heads,
324
head_size=attn_module.head_size,
325
dtype=attn_module.dtype,
326
+ use_mla=False,
327
)
328
elif attn_module.attn_type in (AttentionType.ENCODER,
329
AttentionType.ENCODER_ONLY):
@@ -764,7 +765,7 @@ class TPUModelRunner:
765
"""
766
Initialize KV cache based on `kv_cache_config`.
767
Args:
767
- kv_cache_config: Configuration for the KV cache, including the KV
768
+ kv_cache_config: Configuration for the KV cache, including the KV
769
cache size of each layer
770
"""
771
if len(kv_cache_config.groups) > 1: