Harry MellorandClaude Opus 5
8c1557a79c
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor <hej.mellor@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
44 lines
1.3 KiB
Python
44 lines
1.3 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
|
|
|
|
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
|