4056c8ac1f
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Base class for attention-like layers."""
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
import torch
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.v1.attention.backend import AttentionBackend, AttentionImpl
|
|
from vllm.v1.kv_cache_interface import KVCacheSpec
|
|
|
|
|
|
class AttentionLayerBase(ABC):
|
|
"""Base class for attention-like layers (Attention, Mamba, etc.)
|
|
that support the v1 engine.
|
|
|
|
This provides a common interface for getting attention backends
|
|
from different layer types.
|
|
"""
|
|
|
|
impl: "AttentionImpl"
|
|
supports_dcp: bool = True
|
|
# Set by the speculator on its draft model's layers.
|
|
is_draft_layer: bool = False
|
|
|
|
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
|
|
"""Bind the allocated KV cache tensor to this layer.
|
|
|
|
The default stores the cache view as-is; subclasses (e.g. Mamba)
|
|
override this to unpack the raw buffer into per-state views.
|
|
"""
|
|
self.kv_cache = kv_cache
|
|
|
|
@abstractmethod
|
|
def get_attn_backend(self) -> type[AttentionBackend]:
|
|
"""Get the attention backend class for this layer."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
|
|
"""Get the KV cache spec for this layer.
|
|
May be None if the layer does not need KV cache.
|
|
"""
|
|
pass
|