@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b26b70bec4

Fork 0

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

[Misc] Refactor `get_kv_cache_spec` into `AttentionLayerBase` (#26587)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Oct 18, 2025 at 15:51 UTC b26b70bec4faa3dae487db1f3eebaddd86c1d828

10 files changed +151 -118

vllm/attention/layer.py

+53 -5

@@ -16,6 +16,7 @@ from vllm.attention.backends.registry import _Backend, backend_name_to_enum

16 from vllm.attention.selector import get_attn_backend

17 from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target

18 from vllm.config import CacheConfig, get_current_vllm_config

19 +from vllm.config.vllm import VllmConfig

20 from vllm.distributed.kv_transfer import (

21 get_kv_transfer_group,

22 has_kv_transfer_group,

@@ -34,7 +35,16 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod

35 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

36 from vllm.model_executor.models.vision import get_vit_attn_backend

37 from vllm.platforms import current_platform

37 -from vllm.utils import direct_register_custom_op

38 +from vllm.utils import (

39 + direct_register_custom_op,

40 + kv_cache_dtype_str_to_dtype,

41 +)

42 +from vllm.v1.kv_cache_interface import (

43 + FullAttentionSpec,

44 + KVCacheSpec,

45 + MLAAttentionSpec,

46 + SlidingWindowSpec,

47 +)

49 FP8_DTYPE = current_platform.fp8_dtype()

50 logger = init_logger(__name__)

@@ -152,6 +162,7 @@ class Attention(nn.Module, AttentionLayerBase):

162 else:

163 sliding_window = None

164

165 + vllm_config = get_current_vllm_config()

166 if cache_config is not None:

167 kv_cache_dtype = cache_config.cache_dtype

168 block_size = cache_config.block_size

@@ -160,6 +171,9 @@ class Attention(nn.Module, AttentionLayerBase):

171 kv_cache_dtype = "auto"

172 block_size = 16

173 calculate_kv_scales = False

174 + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(

175 + kv_cache_dtype, vllm_config.model_config

176 + )

177 if num_kv_heads is None:

178 num_kv_heads = num_heads

179 assert num_heads % num_kv_heads == 0, (

@@ -256,7 +270,7 @@ class Attention(nn.Module, AttentionLayerBase):

270 self.use_direct_call = not current_platform.opaque_attention_op()

271

272 self.use_output = self.attn_backend.accept_output_buffer

259 - compilation_config = get_current_vllm_config().compilation_config

273 + compilation_config = vllm_config.compilation_config

274 if prefix in compilation_config.static_forward_context:

275 raise ValueError(f"Duplicate layer name: {prefix}")

276 compilation_config.static_forward_context[prefix] = self

@@ -276,9 +290,7 @@ class Attention(nn.Module, AttentionLayerBase):

290 # this variable will not be accessed if use_direct_call is True

291 self.kv_cache = [

292 torch.tensor([])

279 - for _ in range(

280 - get_current_vllm_config().parallel_config.pipeline_parallel_size

281 - )

293 + for _ in range(vllm_config.parallel_config.pipeline_parallel_size)

294 ]

295

296 # Initialize q/k/v range constants.

@@ -394,6 +406,30 @@ class Attention(nn.Module, AttentionLayerBase):

406 def get_attn_backend(self) -> type[AttentionBackend]:

407 return self.attn_backend

408

409 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

410 + # Block size may get updated after model loading, refresh it

411 + block_size = vllm_config.cache_config.block_size

412 + # Should not be called for enc-dec or encoder-only attention.

413 + assert self.attn_type == AttentionType.DECODER

414 + if self.sliding_window is not None:

415 + assert not vllm_config.model_config.use_mla, (

416 + "MLA is not supported for slidingwindow"

417 + )

418 + return SlidingWindowSpec(

419 + block_size=block_size,

420 + num_kv_heads=self.num_kv_heads,

421 + head_size=self.head_size,

422 + dtype=self.kv_cache_torch_dtype,

423 + sliding_window=self.sliding_window,

424 + )

425 + else:

426 + return FullAttentionSpec(

427 + block_size=block_size,

428 + num_kv_heads=self.num_kv_heads,

429 + head_size=self.head_size,

430 + dtype=self.kv_cache_torch_dtype,

431 + )

432 +

433

434 class MultiHeadAttention(nn.Module):

435 """Multi-headed attention without any cache, used for ViT."""

@@ -749,6 +785,18 @@ class MLAAttention(nn.Module, AttentionLayerBase):

785 def get_attn_backend(self) -> type[AttentionBackend]:

786 return self.attn_backend

787

788 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

789 + kv_cache_dtype = kv_cache_dtype_str_to_dtype(

790 + self.kv_cache_dtype, vllm_config.model_config

791 + )

792 + return MLAAttentionSpec(

793 + block_size=vllm_config.cache_config.block_size,

794 + num_kv_heads=1,

795 + head_size=self.head_size,

796 + dtype=kv_cache_dtype,

797 + cache_dtype_str=vllm_config.cache_config.cache_dtype,

798 + )

799 +

800

801 def wait_for_kv_layer_from_connector(layer_name: str):

802 if not has_kv_transfer_group() or not is_v1_kv_transfer_group():

vllm/attention/layers/chunked_local_attention.py

+13

@@ -9,6 +9,7 @@ from vllm import envs

9 from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata

10 from vllm.attention.selector import get_attn_backend

11 from vllm.config import CacheConfig

12 +from vllm.config.vllm import VllmConfig

13 from vllm.model_executor.layers.quantization import QuantizationConfig

14 from vllm.v1.attention.backends.utils import (

15 AttentionCGSupport,

@@ -16,6 +17,7 @@ from vllm.v1.attention.backends.utils import (

17 make_local_attention_virtual_batches,

18 subclass_attention_backend,

19 )

20 +from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, KVCacheSpec

22 from ..layer import Attention

@@ -67,6 +69,7 @@ class ChunkedLocalAttention(Attention):

69 kv_sharing_target_layer_name: str | None = None,

70 prefix: str = "",

71 ):

72 + self.attention_chunk_size = attention_chunk_size

73 dtype = torch.get_default_dtype()

74 if cache_config is not None:

75 kv_cache_dtype = cache_config.cache_dtype

@@ -99,3 +102,13 @@ class ChunkedLocalAttention(Attention):

102 kv_sharing_target_layer_name=kv_sharing_target_layer_name,

103 attn_backend=attn_backend,

104 )

105 +

106 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

107 + assert self.attention_chunk_size

108 + return ChunkedLocalAttentionSpec(

109 + block_size=vllm_config.cache_config.block_size,

110 + num_kv_heads=self.num_kv_heads,

111 + head_size=self.head_size,

112 + dtype=self.kv_cache_torch_dtype,

113 + attention_chunk_size=self.attention_chunk_size,

114 + )

vllm/attention/layers/cross_attention.py

+9 -1

@@ -21,7 +21,7 @@ from vllm.v1.attention.backends.utils import (

21 CommonAttentionMetadata,

22 subclass_attention_backend,

23 )

24 -from vllm.v1.kv_cache_interface import CrossAttentionSpec

24 +from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec

26 logger = init_logger(__name__)

@@ -174,3 +174,11 @@ class CrossAttention(Attention):

174 attn_type=AttentionType.ENCODER_DECODER,

175 **kwargs,

176 )

177 +

178 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

179 + return CrossAttentionSpec(

180 + block_size=vllm_config.cache_config.block_size,

181 + num_kv_heads=self.num_kv_heads,

182 + head_size=self.head_size,

183 + dtype=self.kv_cache_torch_dtype,

184 + )

vllm/attention/layers/encoder_only_attention.py

@@ -14,10 +14,12 @@ from vllm.attention.backends.abstract import (

14 from vllm.attention.layer import Attention

15 from vllm.attention.selector import get_attn_backend

16 from vllm.config import CacheConfig

17 +from vllm.config.vllm import VllmConfig

18 from vllm.v1.attention.backends.utils import (

19 CommonAttentionMetadata,

20 subclass_attention_backend,

21 )

22 +from vllm.v1.kv_cache_interface import KVCacheSpec

25 @functools.lru_cache

@@ -98,3 +100,7 @@ class EncoderOnlyAttention(Attention):

100 attn_type=AttentionType.ENCODER_ONLY,

101 **kwargs,

102 )

103 +

104 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

105 + # Does not need KV cache

106 + return None

vllm/model_executor/layers/attention_layer_base.py

+11

@@ -5,6 +5,9 @@

5 from abc import ABC, abstractmethod

6 from typing import TYPE_CHECKING

8 +from vllm.config import VllmConfig

9 +from vllm.v1.kv_cache_interface import KVCacheSpec

10 +

11 if TYPE_CHECKING:

12 from vllm.attention.backends.abstract import AttentionBackend

@@ -22,3 +25,11 @@ class AttentionLayerBase(ABC):

25 def get_attn_backend(self) -> type["AttentionBackend"]:

26 """Get the attention backend class for this layer."""

27 pass

28 +

29 + @abstractmethod

30 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:

31 + """

32 + Get the KV cache spec for this layer.

33 + May be None if the layer does not need KV cache.

34 + """

35 + pass

vllm/model_executor/layers/mamba/abstract.py

+29

@@ -6,7 +6,9 @@ from typing import TYPE_CHECKING

7 import torch

9 +from vllm.config import VllmConfig

10 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

11 +from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec

13 if TYPE_CHECKING:

14 from vllm.attention.backends.abstract import AttentionBackend

@@ -40,3 +42,30 @@ class MambaBase(AttentionLayerBase):

42 def get_attn_backend(self) -> type["AttentionBackend"]:

43 """Get the attention backend class for this Mamba layer."""

44 pass

45 +

46 + @abstractmethod

47 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:

48 + pass

49 +

50 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:

51 + if (

52 + vllm_config.speculative_config is not None

53 + and vllm_config.model_config.hf_config.model_type not in ["qwen3_next"]

54 + ):

55 + raise NotImplementedError(

56 + "Mamba with speculative decoding is not supported yet."

57 + )

58 + mamba_block_size = vllm_config.cache_config.mamba_block_size

59 + page_size_padded = vllm_config.cache_config.mamba_page_size_padded

60 + return MambaSpec(

61 + shapes=self.get_state_shape(),

62 + dtypes=self.get_state_dtype(),

63 + block_size=mamba_block_size,

64 + page_size_padded=page_size_padded,

65 + mamba_type=self.mamba_type,

66 + num_speculative_blocks=(

67 + vllm_config.speculative_config.num_speculative_tokens

68 + if vllm_config.speculative_config

69 + else 0

70 + ),

71 + )

vllm/model_executor/models/deepseek_v2.py

+1 -1

@@ -481,7 +481,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):

481 raise ValueError(f"Duplicate layer name: {prefix}")

482 compilation_config.static_forward_context[prefix] = self

483

484 - def get_kv_cache_spec(self) -> KVCacheSpec:

484 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

485 return MLAAttentionSpec( # Only has one vector instead of K + V

486 block_size=self.cache_config.block_size,

487 num_kv_heads=1,

vllm/utils/__init__.py

@@ -137,6 +137,15 @@ def set_default_torch_num_threads(num_threads: int):

137 torch.set_num_threads(old_num_threads)

138

139

140 +def kv_cache_dtype_str_to_dtype(

141 + kv_cache_dtype: str, model_config: ModelConfig

142 +) -> torch.dtype:

143 + if kv_cache_dtype == "auto":

144 + # Model config may not be specified for unit tests, default to float16

145 + return model_config.dtype if model_config else torch.half

146 + return STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]

147 +

148 +

149 T = TypeVar("T")

150 U = TypeVar("U")

151

vllm/v1/spec_decode/eagle.py

+1 -1

@@ -948,7 +948,7 @@ class EagleProposer:

948 indexer_layers[first_layer]

949 .get_attn_backend()

950 .get_builder_cls()(

951 - indexer_layers[first_layer].get_kv_cache_spec(),

951 + indexer_layers[first_layer].get_kv_cache_spec(self.vllm_config),

952 self.indexer_layer_names,

953 self.vllm_config,

954 self.device,

vllm/v1/worker/gpu_model_runner.py

+19 -110

@@ -19,8 +19,6 @@ from tqdm import tqdm

19 import vllm.envs as envs

20 from vllm.attention import Attention, AttentionType

21 from vllm.attention.backends.abstract import AttentionBackend, MultipleOf

22 -from vllm.attention.layer import MLAAttention

23 -from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention

22 from vllm.compilation.counter import compilation_counter

23 from vllm.compilation.cuda_graph import CUDAGraphWrapper

24 from vllm.compilation.monitor import set_cudagraph_capturing_enabled

@@ -44,10 +42,8 @@ from vllm.distributed.parallel_state import (

42 from vllm.forward_context import BatchDescriptor, set_forward_context

43 from vllm.logger import init_logger

44 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

47 -from vllm.model_executor.layers.mamba.abstract import MambaBase

45 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

46 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader

50 -from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache

47 from vllm.model_executor.models.interfaces import (

48 SupportsMultiModal,

49 is_mixture_of_experts,

@@ -73,11 +69,11 @@ from vllm.sampling_params import SamplingType

69 from vllm.sequence import IntermediateTensors

70 from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

71 from vllm.utils import (

76 - STR_DTYPE_TO_TORCH_DTYPE,

72 cdiv,

73 check_use_alibi,

74 get_dtype_size,

75 is_pin_memory_available,

76 + kv_cache_dtype_str_to_dtype,

77 length_from_prompt_token_ids_or_embeds,

78 round_up,

79 supports_dynamo,

@@ -106,7 +102,6 @@ from vllm.v1.kv_cache_interface import (

102 KVCacheGroupSpec,

103 KVCacheSpec,

104 MambaSpec,

109 - MLAAttentionSpec,

105 SlidingWindowSpec,

106 UniformTypeKVCacheSpecs,

107 )

@@ -239,10 +234,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

234 self.device = device

235 self.pin_memory = is_pin_memory_available()

236 self.dtype = self.model_config.dtype

242 - if cache_config.cache_dtype == "auto":

243 - self.kv_cache_dtype = self.dtype

244 - else:

245 - self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]

237 + self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(

238 + cache_config.cache_dtype, self.model_config

239 + )

240

241 self.is_pooling_model = model_config.runner_type == "pooling"

242 self.enable_prompt_embeds = model_config.enable_prompt_embeds

@@ -4577,109 +4571,24 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

4571 format. Layers that do not need KV cache are not included.

4572 """

4573

4580 - block_size = self.vllm_config.cache_config.block_size

4581 - use_mla = self.vllm_config.model_config.use_mla

4582 - cache_dtype_str = self.vllm_config.cache_config.cache_dtype

4574 kv_cache_spec: dict[str, KVCacheSpec] = {}

4575 attn_layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)

4576 for layer_name, attn_module in attn_layers.items():

4586 - if isinstance(attn_module, Attention):

4587 - if (

4588 - kv_tgt_layer := attn_module.kv_sharing_target_layer_name

4589 - ) is not None:

4590 - # The layer doesn't need its own KV cache and will use that of

4591 - # the target layer. We skip creating a KVCacheSpec for it, so

4592 - # that KV cache management logic will act as this layer does

4593 - # not exist, and doesn't allocate KV cache for the layer. This

4594 - # enables the memory saving of cross-layer kv sharing, allowing

4595 - # a given amount of memory to accommodate longer context lengths

4596 - # or enable more requests to be processed simultaneously.

4597 - self.shared_kv_cache_layers[layer_name] = kv_tgt_layer

4598 - continue

4599 -

4600 - # TODO(lucas): move the attention specs into the model layers like

4601 - # the attention backends

4602 - if attn_module.attn_type == AttentionType.DECODER:

4603 - if attn_module.sliding_window is not None:

4604 - assert not use_mla, "MLA is not supported for slidingwindow"

4605 - kv_cache_spec[layer_name] = SlidingWindowSpec(

4606 - block_size=block_size,

4607 - num_kv_heads=attn_module.num_kv_heads,

4608 - head_size=attn_module.head_size,

4609 - dtype=self.kv_cache_dtype,

4610 - sliding_window=attn_module.sliding_window,

4611 - )

4612 - elif self.attention_chunk_size is not None and isinstance(

4613 - attn_module, ChunkedLocalAttention

4614 - ):

4615 - kv_cache_spec[layer_name] = ChunkedLocalAttentionSpec(

4616 - block_size=block_size,

4617 - num_kv_heads=attn_module.num_kv_heads,

4618 - head_size=attn_module.head_size,

4619 - dtype=self.kv_cache_dtype,

4620 - attention_chunk_size=self.attention_chunk_size,

4621 - )

4622 - else:

4623 - kv_cache_spec[layer_name] = FullAttentionSpec(

4624 - block_size=block_size,

4625 - num_kv_heads=attn_module.num_kv_heads,

4626 - head_size=attn_module.head_size,

4627 - dtype=self.kv_cache_dtype,

4628 - )

4629 - elif attn_module.attn_type == AttentionType.ENCODER_DECODER:

4630 - kv_cache_spec[layer_name] = CrossAttentionSpec(

4631 - block_size=block_size,

4632 - num_kv_heads=attn_module.num_kv_heads,

4633 - head_size=attn_module.head_size,

4634 - dtype=self.kv_cache_dtype,

4635 - )

4636 - elif attn_module.attn_type in (

4637 - AttentionType.ENCODER,

4638 - AttentionType.ENCODER_ONLY,

4639 - ):

4640 - # encoder-only attention does not need KV cache.

4641 - continue

4642 - else:

4643 - raise ValueError(f"Unknown attention type: {attn_module.attn_type}")

4644 -

4645 - elif isinstance(attn_module, MLAAttention):

4646 - kv_cache_spec[layer_name] = MLAAttentionSpec(

4647 - block_size=block_size,

4648 - num_kv_heads=1,

4649 - head_size=attn_module.head_size,

4650 - dtype=self.kv_cache_dtype,

4651 - cache_dtype_str=cache_dtype_str,

4652 - )

4653 -

4654 - elif isinstance(attn_module, MambaBase):

4655 - if (

4656 - self.vllm_config.speculative_config is not None

4657 - and self.vllm_config.model_config.hf_config.model_type

4658 - not in ["qwen3_next"]

4659 - ):

4660 - raise NotImplementedError(

4661 - "Mamba with speculative decoding is not supported yet."

4662 - )

4663 - mamba_block_size = self.vllm_config.cache_config.mamba_block_size

4664 - page_size_padded = self.vllm_config.cache_config.mamba_page_size_padded

4665 - kv_cache_spec[layer_name] = MambaSpec(

4666 - shapes=attn_module.get_state_shape(),

4667 - dtypes=attn_module.get_state_dtype(),

4668 - block_size=mamba_block_size,

4669 - page_size_padded=page_size_padded,

4670 - mamba_type=attn_module.mamba_type,

4671 - num_speculative_blocks=(

4672 - self.speculative_config.num_speculative_tokens

4673 - if self.speculative_config

4674 - else 0

4675 - ),

4676 - )

4677 -

4678 - ds_indexer_layers = get_layers_from_vllm_config(

4679 - self.vllm_config, DeepseekV32IndexerCache

4680 - )

4681 - for layer_name, ds_indexer_module in ds_indexer_layers.items():

4682 - kv_cache_spec[layer_name] = ds_indexer_module.get_kv_cache_spec()

4577 + if isinstance(attn_module, Attention) and (

4578 + kv_tgt_layer := attn_module.kv_sharing_target_layer_name

4579 + ):

4580 + # The layer doesn't need its own KV cache and will use that of

4581 + # the target layer. We skip creating a KVCacheSpec for it, so

4582 + # that KV cache management logic will act as this layer does

4583 + # not exist, and doesn't allocate KV cache for the layer. This

4584 + # enables the memory saving of cross-layer kv sharing, allowing

4585 + # a given amount of memory to accommodate longer context lengths

4586 + # or enable more requests to be processed simultaneously.

4587 + self.shared_kv_cache_layers[layer_name] = kv_tgt_layer

4588 + continue

4589 + # Skip modules that don't need KV cache (eg encoder-only attention)

4590 + if spec := attn_module.get_kv_cache_spec(self.vllm_config):

4591 + kv_cache_spec[layer_name] = spec

4592

4593 return kv_cache_spec

4594