iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Attention]: Use KVCacheSpec for AttentionMetadataBuilder type hints (#50148)

Browse Source 
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
This commit is contained in:
Martin Hickey
2026-07-31 11:20:18 +00:00
committed by GitHub
parent 92643d68f5
commit b2fb83e7ff
14 changed files with 34 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/attention/chunked_local_attention.py
+1 -2
View File
@@ -20,7 +20,6 @@ from vllm.v1.attention.backends.utils import (
)
from vllm.v1.attention.selector import get_attn_backend
from vllm.v1.kv_cache_interface import (
AttentionSpec,
ChunkedLocalAttentionSpec,
KVCacheSpec,
get_kv_quant_mode,
@@ -42,7 +41,7 @@ def create_chunked_local_attention_backend(
def get_cudagraph_support(
cls: type["AttentionMetadataBuilder"],
vllm_config: VllmConfig,
kv_cache_spec: AttentionSpec,
kv_cache_spec: KVCacheSpec,
) -> AttentionCGSupport:
# Explicit override in case the underlying builder specialized this getter.
# @override omitted only because of mypy limitation due to type variable.
vllm/model_executor/layers/attention/mla_attention.py
+2
View File
@@ -1722,6 +1722,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
understand this class
"""
kv_cache_spec: AttentionSpec
# Defines the level of query length support for this backend.
# - SINGLE_ONLY: Only single-token queries (no spec decode support)
# - UNIFORM: Supports uniform multi-token queries (spec decode with uniform lengths)
vllm/v1/attention/backend.py
+3 -3
View File
@@ -25,7 +25,7 @@ if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
from vllm.platforms.interface import DeviceCapability
from vllm.v1.attention.backends.utils import KVCacheLayoutType
from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode
from vllm.v1.kv_cache_interface import KVCacheSpec, KVQuantMode
from vllm.v1.kv_cache_interface import get_kv_quant_mode
@@ -635,7 +635,7 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
@abstractmethod
def __init__(
self,
kv_cache_spec: "AttentionSpec",
kv_cache_spec: "KVCacheSpec",
layer_names: list[str],
vllm_config: "VllmConfig",
device: torch.device,
@@ -649,7 +649,7 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
def get_cudagraph_support(
cls: type["AttentionMetadataBuilder"],
vllm_config: "VllmConfig",
kv_cache_spec: "AttentionSpec",
kv_cache_spec: "KVCacheSpec",
) -> AttentionCGSupport:
"""Get the cudagraph support level of this builder class."""
return cls._cudagraph_support
vllm/v1/attention/backends/flash_attn.py
+2 -2
View File
@@ -58,7 +58,7 @@ from vllm.v1.attention.backend import (
CommonAttentionMetadata,
)
from vllm.v1.attention.backends.utils import get_kv_cache_layout
from vllm.v1.kv_cache_interface import AttentionSpec
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
from vllm.v1.worker.cp_utils import (
run_split_fa2_dcp_context_attention,
should_skip_dcp_context_attention,
@@ -360,7 +360,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
def get_cudagraph_support(
cls,
vllm_config: "VllmConfig",
kv_cache_spec: "AttentionSpec",
kv_cache_spec: "KVCacheSpec",
) -> AttentionCGSupport:
return cls._cudagraph_support
vllm/v1/attention/backends/flashinfer.py
+3 -1
View File
@@ -76,6 +76,7 @@ from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheSpec,
KVQuantMode,
UniformTypeKVCacheSpecs,
)
@@ -617,6 +618,7 @@ class FlashInferMetadata:
class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
kv_cache_spec: AttentionSpec
reorder_batch_threshold: int = 1
def __init__(
@@ -893,7 +895,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
def get_cudagraph_support(
cls: type["FlashInferMetadataBuilder"],
vllm_config: VllmConfig,
kv_cache_spec: AttentionSpec,
kv_cache_spec: KVCacheSpec,
) -> AttentionCGSupport:
"""Get the cudagraph support level for FlashInfer attention.
vllm/v1/attention/backends/gdn_attn.py
+3 -3
View File
@@ -21,7 +21,7 @@ from vllm.v1.attention.backends.utils import (
mamba_get_block_table_tensor,
split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
from vllm.v1.kv_cache_interface import MambaSpec
class GDNAttentionBackend(AttentionBackend):
@@ -80,18 +80,18 @@ class GDNAttentionMetadata:
class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]):
kv_cache_spec: MambaSpec
_cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
reorder_batch_threshold: int = 1
def __init__(
self,
kv_cache_spec: AttentionSpec,
kv_cache_spec: MambaSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
):
assert isinstance(kv_cache_spec, MambaSpec)
self.vllm_config = vllm_config
self.compilation_config = vllm_config.compilation_config
self.speculative_config = vllm_config.speculative_config
vllm/v1/attention/backends/hpc_attn.py
+2 -2
View File
@@ -33,7 +33,7 @@ from vllm.v1.attention.backends.utils import (
KVCacheLayoutType,
split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import AttentionSpec
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
logger = init_logger(__name__)
@@ -158,7 +158,7 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
def get_cudagraph_support(
cls: type["HpcAttnMetadataBuilder"],
vllm_config: VllmConfig,
kv_cache_spec: AttentionSpec,
kv_cache_spec: KVCacheSpec,
) -> AttentionCGSupport:
spec_config = vllm_config.speculative_config
if (
vllm/v1/attention/backends/linear_attn.py
+5 -5
View File
@@ -16,7 +16,7 @@ from vllm.v1.attention.backends.utils import (
mamba_get_block_table_tensor,
split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
class LinearAttentionBackend(AttentionBackend):
@@ -46,19 +46,19 @@ class LinearAttentionMetadata:
class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMetadata]):
kv_cache_spec: MambaSpec
reorder_batch_threshold: int = 1
_cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
def __init__(
self,
kv_cache_spec: AttentionSpec,
kv_cache_spec: MambaSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
):
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
assert isinstance(kv_cache_spec, MambaSpec)
def build(
self,
@@ -120,13 +120,13 @@ class BailingLinearAttentionMetadataBuilder(LinearAttentionMetadataBuilder):
def get_cudagraph_support(
cls,
vllm_config: VllmConfig,
kv_cache_spec: AttentionSpec,
kv_cache_spec: KVCacheSpec,
) -> AttentionCGSupport:
return AttentionCGSupport.UNIFORM_BATCH
def __init__(
self,
kv_cache_spec: AttentionSpec,
kv_cache_spec: MambaSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
vllm/v1/attention/backends/mamba2_attn.py
+2 -2
View File
@@ -16,7 +16,7 @@ from vllm.v1.attention.backends.mamba_attn import (
BaseMambaAttentionMetadata,
BaseMambaAttentionMetadataBuilder,
)
from vllm.v1.kv_cache_interface import AttentionSpec
from vllm.v1.kv_cache_interface import MambaSpec
def compute_varlen_chunk_metadata(
@@ -118,7 +118,7 @@ class Mamba2AttentionMetadataBuilder(
def __init__(
self,
kv_cache_spec: AttentionSpec,
kv_cache_spec: MambaSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
vllm/v1/attention/backends/mamba_attn.py
+3 -3
View File
@@ -21,7 +21,7 @@ from vllm.v1.attention.backends.utils import (
mamba_get_block_table_tensor,
split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
from vllm.v1.kv_cache_interface import MambaSpec
M = TypeVar("M", bound="BaseMambaAttentionMetadata")
@@ -83,6 +83,7 @@ class BaseMambaAttentionMetadata:
class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
kv_cache_spec: MambaSpec
metadata_cls: type[M]
reorder_batch_threshold: int = 1
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
@@ -92,7 +93,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
def __init__(
self,
kv_cache_spec: AttentionSpec,
kv_cache_spec: MambaSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
@@ -107,7 +108,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
self.use_replayssm = vllm_config.cache_config.use_replayssm
self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len
assert isinstance(kv_cache_spec, MambaSpec)
scheduler_config = vllm_config.scheduler_config
self.decode_cudagraph_max_bs: int = scheduler_config.max_num_seqs
if self.compilation_config.max_cudagraph_capture_size is not None:
vllm/v1/attention/backends/mla/indexer.py
+2 -2
View File
@@ -36,7 +36,7 @@ from vllm.v1.attention.backends.utils import (
get_dcp_local_seq_lens,
split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
from vllm.v1.worker.cp_utils import get_kv_cache_shard_count
logger = init_logger(__name__)
@@ -463,7 +463,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
def get_cudagraph_support(
cls,
vllm_config: VllmConfig,
kv_cache_spec: AttentionSpec,
kv_cache_spec: KVCacheSpec,
) -> AttentionCGSupport:
return AttentionCGSupport.UNIFORM_BATCH
vllm/v1/attention/backends/turboquant_attn.py
+2
View File
@@ -53,6 +53,7 @@ from vllm.v1.attention.ops.triton_turboquant_decode import (
triton_turboquant_decode_attention,
)
from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store
from vllm.v1.kv_cache_interface import AttentionSpec
from vllm.v1.worker.workspace import (
current_workspace_manager,
is_workspace_manager_initialized,
@@ -198,6 +199,7 @@ class TurboQuantMetadata(AttentionMetadata):
class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]):
"""Builds TurboQuantMetadata from scheduler output."""
kv_cache_spec: AttentionSpec
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
vllm/v1/worker/gpu/attn_utils.py
+1 -1
View File
@@ -169,7 +169,7 @@ def init_attn_backend(
# Check cudagraph support for the attention backend
cg_support = builder.get_cudagraph_support(
vllm_config,
cast(AttentionSpec, group.kv_cache_spec),
group.kv_cache_spec,
)
if cg_support.value < min_cg_support.value:
min_cg_support = cg_support
vllm/v1/worker/gpu/model_states/encoder_only.py
+3 -7
View File
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any, cast
from typing import Any
import torch
import torch.nn as nn
@@ -15,11 +15,7 @@ from vllm.v1.attention.backend import (
CommonAttentionMetadata,
)
from vllm.v1.core.sched.output import NewRequestData
from vllm.v1.kv_cache_interface import (
AttentionSpec,
EncoderOnlyAttentionSpec,
KVCacheConfig,
)
from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, KVCacheConfig
from vllm.v1.worker.gpu.input_batch import InputBatch
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
@@ -152,7 +148,7 @@ class EncoderOnlyModelState(DefaultModelState):
for group in self.encoder_attn_groups:
builder = group.get_metadata_builder(0)
cg_support = builder.get_cudagraph_support(
self.vllm_config, cast(AttentionSpec, group.kv_cache_spec)
self.vllm_config, group.kv_cache_spec
)
if cg_support.value < support.value:
support = cg_support