AttentionBackend: Abstraktionsschicht für Attention-Backends

源码版本v0.25.1

Verantwortung

vLLM unterstützt mehr Attention-Implementierungen (attention), als man an einer Hand abzählen kann: FlashAttention 2/3/4, FlashInfer, Triton, MLA, Mamba, ROCm AITER, XPU … Jede Variante unterscheidet sich in KV-Cache (KV cache)-Layout, Kernel-Blockgröße und cudagraph-Support. Die AttentionBackend-Schicht existiert, um diese Unterschiede zu absorbieren, sodass die obere Schicht GPUModelRunner nur ein einheitliches forward aufruft, ohne wissen zu müssen, auf welchem Kernel gerade gelaufen wird.

Die Abstraktion hat zwei Hälften: AttentionBackend(backend.py:56) selbst ist eine Fabrik aus Klassenattributen und statischen Methoden, die beantwortet: «Wie heiße ich, welche impl-Klasse, welcher Metadata-Builder, welche Form soll der KV-Cache (KV cache) haben?»; AttentionMetadataBuilder(backend.py:600) ist dafür zuständig, in jedem Scheduler-Schritt (Scheduler) das CommonAttentionMetadata (die backend-übergreifend geteilten Per-Batch-Metadaten wie query_start_loc, seq_lens, block_table_tensor) in das AttentionMetadata zu übersetzen, das der Kernel des aktuellen Backends wirklich konsumiert. Darunter liegt AttentionImpl(backend.py:858), das Objekt, das tatsächlich die forward-Methode hält und pro Schicht vom Modellcode einmal aufgerufen wird.

Backends sind austauschbar: AttentionBackendEnum(registry.py:34) listet die «Klassenpfade» aller Backends als Enumeration auf; zur Laufzeit wählt get_attn_backend(selector.py:54) basierend auf head_size, dtype, kv_cache_dtype, MLA-Verwendung und Plattform-Rechenfähigkeit eines aus und lädt das entsprechende Modul lazily. Das Modell erhält bei der Initialisierung lediglich type[AttentionBackend], den Rest übernimmt diese Abstraktion.

Entwurfsmotivation

Hardware-übergreifende Einheit: Derselbe Modellcode läuft auf NVIDIA / AMD / Intel XPU / CPU; der Unterschied liegt nur in der Backend-Wahl, und der Selector (selector) bündelt diese Entscheidung an einer Stelle.
KV-Cache-Form entkoppelt: Jedes Backend bestimmt über get_kv_cache_shape das physische Layout des Block-Pools (FlashAttention: [num_blocks, 2, block_size, num_kv_heads, head_size], MLA: [num_blocks, block_size, head_size]). KVCacheManager sieht nur die Spec und interessiert sich nicht für das konkrete Backend.
cudagraph-Support abgestuft: AttentionCGSupport(backend.py:583) teilt in ALWAYS / UNIFORM_BATCH / UNIFORM_SINGLE_TOKEN_DECODE / NEVER ein, damit GPUModelRunner weiß, ob das aktuelle Backend in cudagraph gepackt werden kann.
argmax-Invarianz klassifiziert: AttentionImplBase(backend.py:769) unterscheidet Eigenschaften wie «kann softmax-lse zurückgeben», «lse ist ln oder log2», «unterstützt Prefill Context Parallelism»; der DCP-Fusion-Kernel wählt anhand dieser Flags seinen Zweig.
MLA geht separaten Weg: MLAAttentionImpl(backend.py:941) implementiert kein forward, sondern forward_mqa / forward_mha, weil der KV-Cache (KV cache) bei MLA das komprimierte Latent speichert, nicht das normale K/V.

Schlüsseldateien

AttentionBackend ABC:56 — Abstrakte Basisklasse, definiert statische Methoden wie get_name / get_impl_cls / get_builder_cls / get_kv_cache_shape.
vier Kern abstractmethods:74-97 — get_name, get_impl_cls, get_builder_cls, get_kv_cache_shape, die Unterklassen implementieren müssen.
CommonAttentionMetadata:395 — Backend-übergreifend geteilte Per-Batch-Metadaten; query_start_loc, seq_lens, block_table_tensor etc. leben hier.
AttentionMetadataBuilder ABC:600 — Die build-Methode übersetzt CommonAttentionMetadata in die backend-eigenen Metadaten.
AttentionImplBase:769 — Basisklasse mit dem echten forward; hält Fähigkeits-Flags wie can_return_lse_for_decode, lse_base_on_e, supports_pcp.
AttentionImpl:858 — ABC für Standard-Attention, definiert forward(layer, query, key, value, kv_cache, attn_metadata, output).
MLAAttentionImpl:941 — MLA-spezifische ABC; Schnittstelle ist forward_mqa / forward_mha, nicht das normale forward.
get_attn_backend:54 — Wählt Backend nach head_size/dtype/kv_cache_dtype/use_mla etc.; mit @cache.
AttentionBackendEnum:34 — Listet alle Backend-Klassenpfade als Enumeration; Standardwerte zeigen auf vllm.v1.attention.backends.*.
register_backend:233 — Überschreibt zur Laufzeit die Klasse, auf die ein Enum-Member zeigt; für Plugins.

Datenfluss

In einem Schritt (step) bereitet GPUModelRunner zuerst query/key/value und CommonAttentionMetadata für den gesamten Batch vor und ruft dann pro Schicht AttentionImpl.forward auf. Vor jeder Schicht hat AttentionMetadataBuilder.build des Backends die gemeinsamen Metadaten in die vom Kernel benötigte Form gebracht. Im Folgenden die Signatur von AttentionMetadataBuilder.build, die von allen Backends mit derselben Schnittstelle implementiert wird:

python

# vllm/v1/attention/backend.py L666-L678
@abstractmethod
def build(
    self,
    common_prefix_len: int,
    common_attn_metadata: CommonAttentionMetadata,
    fast_build: bool = False,
) -> M:
    """
    Central method that builds attention metadata.
    Some builders (MLA) require reorder_batch to be called prior to build.

    Args:
        common_prefix_len: The length of the common prefix of the batch.
        common_attn_metadata: The common attention metadata.

Welche Route das Backend nimmt, entscheidet get_attn_backend; es holt basierend auf der aktuellen vllm_config + Modell-dtype/head_size etc. den Backend-Klassenpfad aus AttentionBackendEnum und lädt ihn lazily:

python

# vllm/v1/attention/selector.py L75-L110
from vllm.config import get_current_vllm_config

vllm_config = get_current_vllm_config()

cache_config = vllm_config.cache_config
if cache_config is not None and cache_config.user_specified_block_size:
    block_size = cache_config.block_size
else:
    block_size = None

kv_transfer_config = vllm_config.kv_transfer_config
use_kv_connector = (
    kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance
)

attn_selector_config = AttentionSelectorConfig(
    head_size=head_size,
    dtype=dtype,
    kv_cache_dtype=cast(CacheDType | None, kv_cache_dtype),
    block_size=block_size,
    use_mla=use_mla,
    ...
)

return _cached_get_attn_backend(
    backend=vllm_config.attention_config.backend,
    attn_selector_config=attn_selector_config,
    num_heads=num_heads,
)

Das Modell erhält beim Initialisieren jedes Attention / MLAAttention diese Backend-Klasse, konstruiert via backend.get_impl_cls()(...) den impl und via backend.get_builder_cls()(...) den Metadata-Builder; diese beiden Objekte werden während der gesamten Inferenzrunde weiterverwendet.

Grenzen und Fehler

Nicht unterstützte Head-Größen: AttentionBackend.supports_head_size(backend.py:159) befragt das Backend zuerst; FlashAttention verlangt, dass head_size ein Vielfaches von 8 ist und 256 nicht überschreitet (FA4 erst bis 512), sonst wird direkt das Backend gewechselt.
Nicht unterstützte Datentypen / KV-Cache-dtype: supports_dtype / supports_kv_cache_dtype(backend.py:163-173) prüfen getrennt; FlashInfer unterstützt fp8/nvfp4, FlashAttention akzeptiert fp8 nur auf Hopper + FA3.
Blockgröße muss 16-ausgerichtet sein: get_supported_kernel_block_sizes von FlashAttention / Triton / MLA liefern jeweils MultipleOf(16); nicht ausgerichtete block_size wirft direkt ValueError(flash_attn.py:130-131).
MLA-Backend verwendet kein forward: MLAAttentionImpl(backend.py:941) deklariert nur forward_mqa / forward_mha; MLAAttention auf Modellebene ruft je nach Kontext die entsprechende Methode auf, falsche Verwendung von forward führt direkt zu AttributeError.
DCP-LSE-Basis darf nicht falsch sein: AttentionImplBase.lse_base_on_e(backend.py:795) unterscheidet ln vs. log2; eine falsche Wahl lässt den Softmax-Nenner beim schardübergreifenden Merge unbemerkt falsch werden.
Selector-Cache nimmt unveränderliche Konfiguration an: _cached_get_attn_backend(selector.py:114) nutzt @cache; Änderungen von dtype/head_size während des Laufens wählen das Backend nicht neu — ein Prozessneustart ist nötig.

Zusammenfassung

Die AttentionBackend-Abstraktionsschicht ist der Schlüssel, mit dem vLLM mehrere Hardware- und Kernel-Varianten parallel hält: Die Modellebene kennt nur die eine Schnittstelle AttentionImpl.forward; wie der Backend den KV-Cache (KV cache) anordnet, wie Metadaten erzeugt werden und ob FlashAttention oder Triton zum Einsatz kommen, wird vollständig vom Selector (selector) + Enum bei der Initialisierung festgelegt. Wie jedes Backend im Detail implementiert ist, siehe /attention/flash (FlashAttention / FlashInfer) und /attention/triton-mla (Triton / MLA). Wo die obere Schicht diese impls aufruft, siehe /worker/gpu-model-runner.

Siehe offizielle Dokumentation: vLLM 文档 · README.

AttentionBackend: Abstraktionsschicht für Attention-Backends ​

Verantwortung ​

Entwurfsmotivation ​

Schlüsseldateien ​

Datenfluss ​

Grenzen und Fehler ​

Zusammenfassung ​