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 14, 05:30 UTC
[deepseek] kernel block size for UniformTypeKVCacheSpecs (#26559)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Oct 10, 2025 at 16:40 UTC
6f0f570c436cb0d3da14ce9c68425430e8b4070b
2 files changed
+19
-6
vllm/v1/attention/backends/mla/indexer.py
+10
-2
@@ -1,11 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
-from typing import ClassVar, Optional
4
+from typing import ClassVar, Optional, Union
5
6
import torch
7
8
-from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
8
+from vllm.attention.backends.abstract import (
9
+ AttentionBackend,
10
+ AttentionMetadata,
11
+ MultipleOf,
12
+)
13
from vllm.config import VllmConfig
14
from vllm.logger import init_logger
15
from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata
@@ -47,6 +51,10 @@ class DeepseekV32IndexerBackend(AttentionBackend):
51
def get_kv_cache_stride_order() -> tuple[int, ...]:
52
return (0, 1, 2)
53
54
+ @classmethod
55
+ def get_supported_kernel_block_size(cls) -> list[Union[int, MultipleOf]]:
56
+ return [64]
57
+
58
59
@dataclass
60
class DeepseekV32IndexerPrefillChunkMetadata:
vllm/v1/worker/gpu_model_runner.py
+9
-4
@@ -4242,9 +4242,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4242
for kv_cache_group_id, kv_cache_group in enumerate(
4243
kv_cache_config.kv_cache_groups
4244
):
4245
- if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):
4245
+ kv_cache_spec = kv_cache_group.kv_cache_spec
4246
+ if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
4247
+ # All layers in the UniformTypeKVCacheSpecs have the same type,
4248
+ # Pick an arbitrary one to dispatch.
4249
+ kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values()))
4250
+ if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
4251
continue
4247
- elif isinstance(kv_cache_group.kv_cache_spec, AttentionSpec):
4252
+ elif isinstance(kv_cache_spec, AttentionSpec):
4253
# This is an attention backend that supports virtual
4254
# block splitting. Get the supported block sizes from
4255
# all backends in the group.
@@ -4254,10 +4259,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4259
kv_manager_block_size, attn_groups
4260
)
4261
kernel_block_sizes.append(selected_kernel_size)
4257
- elif isinstance(kv_cache_group.kv_cache_spec, MambaSpec):
4262
+ elif isinstance(kv_cache_spec, MambaSpec):
4263
# This is likely Mamba or other non-attention cache,
4264
# no splitting.
4260
- kernel_block_sizes.append(kv_cache_group.kv_cache_spec.block_size)
4265
+ kernel_block_sizes.append(kv_cache_spec.block_size)
4266
else:
4267
raise NotImplementedError(
4268
f"unknown kv cache spec {kv_cache_group.kv_cache_spec}"