92 lines
2.5 KiB
Python
92 lines
2.5 KiB
Python
"""Attention backend registry with DP-aware backend selection.
|
|
|
|
Ported from xLLM upstream commit 78aa2a85 (PR #2258).
|
|
Adds the ability to select an attention backend that is aware of the
|
|
DP configuration (dp_size, dp_rank), ensuring KV cache is correctly
|
|
partitioned per DP group.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
|
|
@runtime_checkable
|
|
class AttentionBackend(Protocol):
|
|
"""Protocol for attention backends used by the Python model executor."""
|
|
|
|
def prepare(self, metadata: Any, graph_mode: bool = False) -> None:
|
|
...
|
|
|
|
def bind_kv_caches(self, layer_caches: list) -> None:
|
|
...
|
|
|
|
|
|
@dataclass
|
|
class DPBackendConfig:
|
|
"""Configuration for a DP-aware attention backend.
|
|
|
|
Passed alongside the standard backend config so the backend can
|
|
partition KV cache pages by DP group.
|
|
"""
|
|
|
|
dp_size: int = 1
|
|
dp_rank: int = 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Backend registry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_BACKEND_REGISTRY: dict[str, type] = {}
|
|
|
|
|
|
def register_backend(name: str, cls: type) -> None:
|
|
"""Register an attention backend class under ``name``."""
|
|
_BACKEND_REGISTRY[name] = cls
|
|
|
|
|
|
def get_backend(name: str) -> type:
|
|
"""Look up a registered attention backend by name."""
|
|
if name not in _BACKEND_REGISTRY:
|
|
available = ", ".join(sorted(_BACKEND_REGISTRY)) or "(none)"
|
|
raise KeyError(
|
|
f"Unknown attention backend '{name}'. Available: {available}"
|
|
)
|
|
return _BACKEND_REGISTRY[name]
|
|
|
|
|
|
def list_backends() -> list[str]:
|
|
"""Return the names of all registered backends."""
|
|
return sorted(_BACKEND_REGISTRY)
|
|
|
|
|
|
def create_attention_backend(
|
|
name: str,
|
|
*,
|
|
num_heads: int,
|
|
num_kv_heads: int,
|
|
head_dim: int,
|
|
scale: float,
|
|
dp_config: DPBackendConfig | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Instantiate a registered attention backend with DP config.
|
|
|
|
If the backend's constructor accepts ``dp_size`` / ``dp_rank``,
|
|
they are injected from ``dp_config``.
|
|
"""
|
|
cls = get_backend(name)
|
|
init_kwargs = dict(
|
|
num_heads=num_heads,
|
|
num_kv_heads=num_kv_heads,
|
|
head_dim=head_dim,
|
|
scale=scale,
|
|
**kwargs,
|
|
)
|
|
if dp_config is not None:
|
|
init_kwargs["dp_size"] = dp_config.dp_size
|
|
init_kwargs["dp_rank"] = dp_config.dp_rank
|
|
return cls(**init_kwargs)
|