[Transformers backend] Find attention with a fuser and attach vLLM's layer to it

docs/models/supported_models.md

@@ -72,19 +72,26 @@ This means that, with the Transformers modeling backend for vLLM, new models can
 
 This section details the necessary modifications to make to a Transformers compatible custom model that make it compatible with the Transformers modeling backend for vLLM. (We assume that a Transformers compatible custom model has already been created, see [Transformers - Customizing models](https://huggingface.co/docs/transformers/en/custom_models)).
 
-To make your model compatible with the Transformers modeling backend, it needs:
-
-1. `kwargs` passed down through all modules from `MyModel` to `MyAttention`.
-    - If your model is encoder-only:
-        1. Add `is_causal = False` to `MyAttention`.
-    - If your model is mixture-of-experts (MoE):
-        1. Your sparse MoE block must have an attribute called `experts`.
-        2. The class of `experts` (`MyExperts`) must either:
-            - Inherit from `nn.ModuleList` (naive).
-            - Or contain all 3D `nn.Parameters` (packed).
-        3. `MyExperts.forward` must accept `hidden_states`, `top_k_index`, `top_k_weights`.
-2. `MyAttention` must use `ALL_ATTENTION_FUNCTIONS` to call attention.
-3. `MyModel` must contain `_supports_attention_backend = True`.
+To make your model compatible with the Transformers modeling backend:
+
+1. `MyAttention` must use `ALL_ATTENTION_FUNCTIONS` to call attention.
+    - It must make exactly one such call. vLLM attaches one `Attention` layer per `MyAttention` module.
+    - `MyAttention` must contain a unique `layer_idx`. vLLM keys its KV cache using this index.
+    - Pass `scaling=` to the interface if your scale is not `head_size**-0.5`. vLLM reads it from the call to the attention interface.
+2. If your model is encoder-only:
+    1. Add `is_causal = False` to `MyAttention`.
+3. If your model is mixture-of-experts (MoE):
+    1. Your sparse MoE block must have an attribute called `experts`.
+    2. The class of `experts` (`MyExperts`) must either:
+        - Inherit from `nn.ModuleList` (naive).
+        - Or contain all 3D `nn.Parameters` (packed).
+    3. `MyExperts.forward` must accept `hidden_states`, `top_k_index`, `top_k_weights`.
+
+!!! note
+    `MyModel` no longer needs `_supports_attention_backend = True`, and `kwargs` no
+    longer need to be passed down through every module from `MyModel` to
+    `MyAttention`. vLLM reaches its attention layer through `MyAttention` itself, so
+    all it asks is that `MyAttention` dispatches through the interface.
 
 <details class="code">
 <summary>modeling_my_model.py</summary>
@@ -97,14 +104,24 @@ from torch import nn
 class MyAttention(nn.Module):
     is_causal = False  # Only do this for encoder-only models
 
+    def __init__(self, config, layer_idx):
+        ...
+        self.config = config
+        self.layer_idx = layer_idx
+        self.scaling = self.head_dim**-0.5
+        ...
+
     def forward(self, hidden_states, **kwargs):
         ...
-        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
         attn_output, attn_weights = attention_interface(
             self,
             query_states,
             key_states,
             value_states,
+            scaling=self.scaling,
             **kwargs,
         )
         ...
@@ -127,15 +144,15 @@ class MySparseMoEBlock(nn.Module):
         ...
 
 class MyModel(PreTrainedModel):
-    _supports_attention_backend = True
+    ...
 ```
 
 </details>
 
 Here is what happens in the background when this model is loaded:
 
 1. The config is loaded.
-2. `MyModel` Python class is loaded from the `auto_map` in config, and we check that the model `is_backend_compatible()`.
+2. `MyModel` Python class is loaded from the `auto_map` in config, and we check that the model `_can_set_attn_implementation()`.
 3. `MyModel` is loaded into one of the Transformers modeling backend classes in [vllm/model_executor/models/transformers](../../vllm/model_executor/models/transformers) which sets `self.config._attn_implementation = "vllm"` so that vLLM's attention layer is used.
 
 That's it!

tests/models/transformers/fusers/test_linear.py

@@ -10,7 +10,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from vllm.model_executor.models.transformers.fuser import get_fuser
+from vllm.model_executor.models.transformers.fuser import get_fuser, get_fusers
 from vllm.model_executor.models.transformers.fusers import (
     GLUFuser,
     PackedQKVFuser,
@@ -307,7 +307,7 @@ def forward(self, hidden_states):
 
 
 class FakeSelfAttn(nn.Module):
-    """Stand-in for the vLLM `Attention` looked up in `attention_instances`."""
+    """Stand-in for the vLLM `Attention` attached to the dispatching module."""
 
     def __init__(self):
         super().__init__()
@@ -328,9 +328,9 @@ def forward(self, q, k, v):
 
 @pytest.fixture(autouse=True)
 def _clear_fuser_cache():
-    get_fuser.cache_clear()
+    get_fusers.cache_clear()
     yield
-    get_fuser.cache_clear()
+    get_fusers.cache_clear()
 
 
 def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
@@ -392,7 +392,7 @@ def _apply_packed_qkv_fuser_with_stubs(module: nn.Module, fuser: PackedQKVFuser)
 def test_detects_and_rewrites_glu(mlp_cls, bias):
     with torch.device("meta"):
         meta = mlp_cls(bias=bias)
-    fuser = get_fuser(meta)
+    fuser = get_fuser(meta, GLUFuser)
     assert isinstance(fuser, GLUFuser)
     assert (
         fuser.gate_name,
@@ -427,9 +427,9 @@ def test_glu_identifies_down_projection():
     matches the column-parallel merged gate/up; `None` when there is no such
     projection to force (fusion of gate/up still applies)."""
     with torch.device("meta"):
-        assert get_fuser(GLUMLP()).down_name == "down_proj"
-        assert get_fuser(ReversedGLUMLP()).down_name == "down_proj"
-        assert get_fuser(NoDownGLU()).down_name is None
+        assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj"
+        assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj"
+        assert get_fuser(NoDownGLU(), GLUFuser).down_name is None
 
 
 @pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])
@@ -439,7 +439,7 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
         pytest.skip("MHA q/k/v assignment is order-based by design")
     with torch.device("meta"):
         meta = attn_cls(kv_heads=kv_heads)
-    fuser = get_fuser(meta)
+    fuser = get_fuser(meta, QKVFuser)
     assert isinstance(fuser, QKVFuser)
     # q (sharded differently under TP) must be identified exactly; k/v may be
     # swapped for non-canonical compute order, which is numerically consistent
@@ -467,27 +467,27 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
     for p in real.parameters():
         nn.init.normal_(p, std=0.05)
     x = torch.randn(1, 5, 32)
-    attention_instances = {3: FakeSelfAttn()}
-    expected, _ = real(x, attention_instances=attention_instances)
+    real.attn = FakeSelfAttn()
+    expected, _ = real(x)
     fused = _apply_qkv_fuser_with_stubs(real, fuser)
 
     # Fusion is in place: the module keeps its class and other attributes
     assert fused is real and type(fused) is attn_cls
     assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None
-    out, _ = fused(x, attention_instances=attention_instances)
+    out, _ = fused(x)
     torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
 
 
 def test_qkv_identifies_output_projection():
     with torch.device("meta"):
-        assert get_fuser(FakeAttention()).o_name == "o_proj"
-        assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"
-        assert get_fuser(ExtraProjAttention()).o_name == "o_proj"
+        assert get_fuser(FakeAttention(), QKVFuser).o_name == "o_proj"
+        assert get_fuser(ReversedFakeAttention(), QKVFuser).o_name == "o_proj"
+        assert get_fuser(ExtraProjAttention(), QKVFuser).o_name == "o_proj"
         # Norm children (q_norm/k_norm) must not disturb o_proj identification.
-        assert get_fuser(QKNormAttention()).o_name == "o_proj"
-        assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"
+        assert get_fuser(QKNormAttention(), QKVFuser).o_name == "o_proj"
+        assert get_fuser(PerHeadQKNormAttention(), QKVFuser).o_name == "o_proj"
         # A module between o_proj and the return is transparent.
-        assert get_fuser(ResidDropoutAttention()).o_name == "o_proj"
+        assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj"
 
 
 @pytest.mark.parametrize("kv_heads", [1, 2])
@@ -498,7 +498,7 @@ def test_detects_and_rewrites_packed_qkv(kv_heads):
     checkpoint weight as-is, and shards q by heads while replicating k/v."""
     with torch.device("meta"):
         meta = PackedQKVAttention(kv_heads=kv_heads)
-    fuser = get_fuser(meta)
+    fuser = get_fuser(meta, PackedQKVFuser)
     assert isinstance(fuser, PackedQKVFuser)
     assert (fuser.qkv_name, fuser.o_name) == ("c_attn", "c_proj")
     assert (fuser.q_size, fuser.kv_size) == (32, 8 * kv_heads)
@@ -513,43 +513,43 @@ def test_detects_and_rewrites_packed_qkv(kv_heads):
     for p in real.parameters():
         nn.init.normal_(p, std=0.05)
     x = torch.randn(1, 5, 32)
-    attention_instances = {3: FakeMQASelfAttn()}
-    expected, _ = real(x, attention_instances=attention_instances)
+    real.attn = FakeMQASelfAttn()
+    expected, _ = real(x)
     fused = _apply_packed_qkv_fuser_with_stubs(real, fuser)
 
     # Fusion is in place: the module keeps its class and other attributes
     assert fused is real and type(fused) is PackedQKVAttention
     assert fused.layer_idx == 3 and fused.is_causal
-    out, _ = fused(x, attention_instances=attention_instances)
+    out, _ = fused(x)
     torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
 
 
 def test_per_head_split_is_not_packed_qkv():
     """The split must consume the whole projection, else its sizes are head
     widths and re-sharding by them would be wrong."""
     with torch.device("meta"):
-        assert get_fuser(PerHeadSplitAttention()) is None
+        assert get_fuser(PerHeadSplitAttention(), PackedQKVFuser) is None
 
 
 def test_fuser_is_cached_per_class_and_structure():
     with torch.device("meta"):
-        fuser_a = get_fuser(GLUMLP())
-        fuser_b = get_fuser(GLUMLP())
+        fuser_a = get_fuser(GLUMLP(), GLUFuser)
+        fuser_b = get_fuser(GLUMLP(), GLUFuser)
     assert fuser_a is fuser_b
-    assert any(key[0] is GLUMLP for key in get_fuser.cache)
+    assert any(key[0] is GLUMLP for key in get_fusers.cache)
 
 
 @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])
 def test_non_matching_modules_return_none(cls):
     with torch.device("meta"):
         module = cls()
-    assert get_fuser(module) is None
+    assert get_fuser(module, GLUFuser) is None
 
 
 def test_untraceable_tail_still_fuses():
     with torch.device("meta"):
         meta = UntraceableTailGLUMLP()
-    fuser = get_fuser(meta)
+    fuser = get_fuser(meta, GLUFuser)
     assert isinstance(fuser, GLUFuser)
 
     # Numerics: the live tail must survive the rewrite
@@ -566,8 +566,8 @@ def test_weight_mappings_are_scoped_to_fused_prefixes():
     from vllm.model_executor.models.utils import WeightsMapper
 
     with torch.device("meta"):
-        glu_fuser = get_fuser(GLUMLP())
-        qkv_fuser = get_fuser(FakeAttention())
+        glu_fuser = get_fuser(GLUMLP(), GLUFuser)
+        qkv_fuser = get_fuser(FakeAttention(), QKVFuser)
 
     mapper = WeightsMapper()
     for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
@@ -620,7 +620,7 @@ def test_weight_mappings_are_scoped_to_fused_prefixes():
 def test_unfusable_modules_are_not_fused(cls, default_vllm_config):
     with torch.device("meta"):
         module = cls()
-    fuser = get_fuser(module)
+    fuser = get_fuser(module, GLUFuser)
     # Either no pattern matches the class, or this instance fails validation
     # (`recursive_replace` gates fusion and its weight mappings on `validate`)
     assert fuser is None or not fuser.validate(module, default_vllm_config)
@@ -654,9 +654,15 @@ def _wider_model_config(head_dim: int) -> SimpleNamespace:
 
 
 @pytest.mark.parametrize(
-    "cls, fuser_module", [(FakeAttention, qkv), (PackedQKVAttention, packed_qkv)]
+    "cls, fuser_module, fuser_cls",
+    [
+        (FakeAttention, qkv, QKVFuser),
+        (PackedQKVAttention, packed_qkv, PackedQKVFuser),
+    ],
 )
-def test_head_counts_come_from_the_module_not_the_model(cls, fuser_module, monkeypatch):
+def test_head_counts_come_from_the_module_not_the_model(
+    cls, fuser_module, fuser_cls, monkeypatch
+):
     """A layer narrower than the model-wide head size must not be miscounted.
 
     On a heterogeneous checkpoint (Gemma 4) the model-wide head size is the
@@ -680,7 +686,7 @@ def test_head_counts_come_from_the_module_not_the_model(cls, fuser_module, monke
     monkeypatch.setattr(
         fuser_module, "replace_linear_class", lambda *a, **kw: nn.Identity()
     )
-    fuser = get_fuser(module)
+    fuser = get_fuser(module, fuser_cls)
     assert fuser is not None and fuser.validate(module, vllm_config)
     fuser.update_attrs(module, "model.layers.0.self_attn", vllm_config)
 
@@ -700,5 +706,5 @@ def test_validate_accepts_a_layer_the_model_wide_head_size_would_reject():
 
     # kv width is 24, not a multiple of the model-wide 16, but is of this
     # layer's 8.
-    fuser = get_fuser(module)
+    fuser = get_fuser(module, QKVFuser)
     assert fuser is not None and fuser.validate(module, vllm_config)

tests/models/transformers/fusers/test_mla.py

@@ -158,4 +158,4 @@ def forward(self, x):
 
 def test_non_mla_is_not_matched():
     with torch.device("meta"):
-        assert not isinstance(get_fuser(GLU()), MLAFuser)
+        assert get_fuser(GLU(), MLAFuser) is None

tests/models/transformers/fusers/test_rms_norm.py

@@ -162,7 +162,7 @@ def forward(self, x, gate=None):
 )
 def test_detects_rms_norm_variants(cls, eps, zero_centered):
     with torch.device("meta"):
-        fuser = get_fuser(cls(16, eps=eps))
+        fuser = get_fuser(cls(16, eps=eps), RMSNormFuser)
     assert isinstance(fuser, RMSNormFuser)
     assert fuser.zero_centered == zero_centered
 
@@ -171,15 +171,15 @@ def test_detects_rms_norm_variants(cls, eps, zero_centered):
 def test_non_rms_norms_are_not_matched(cls):
     with torch.device("meta"):
         module = cls(16) if cls is nn.LayerNorm else cls()
-    assert not isinstance(get_fuser(module), RMSNormFuser)
+    assert get_fuser(module, RMSNormFuser) is None
 
 
 @pytest.mark.parametrize(
     "cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm]
 )
 def test_gated_rms_norm_is_not_fused(cls):
     with torch.device("meta"):
-        assert not isinstance(get_fuser(cls()), RMSNormFuser)
+        assert get_fuser(cls(), RMSNormFuser) is None
 
 
 @pytest.mark.parametrize(
@@ -196,7 +196,7 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c
 
     with torch.device("meta"):
         module = cls()
-        fuser = get_fuser(module)
+        fuser = get_fuser(module, RMSNormFuser)
         built = fuser.fuse(module, "norm", default_vllm_config)
     from vllm.model_executor.models.transformers.fusers.rms_norm import (
         TPAwareNormMixin,
@@ -221,7 +221,7 @@ def test_weightless_norm_has_no_hidden_size(default_vllm_config):
     matching the unfused norm and vLLM's native `RMSNorm(hidden_size=head_dim)`.
     """
     module = WeightlessRMSNorm(72)
-    built = get_fuser(module).fuse(module, "norm", default_vllm_config)
+    built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config)
     assert built.hidden_size == 0
 
     built.tp_size = 2  # emulate TP=2 without a real process group
@@ -236,7 +236,7 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config):
 
     with torch.device("meta"):
         module = torch.nn.RMSNorm(16)  # forward is a single `F.rms_norm` call
-        fuser = get_fuser(module)
+        fuser = get_fuser(module, RMSNormFuser)
         assert isinstance(fuser, RMSNormFuser)
         assert not fuser.zero_centered
         vllm_config = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float32))
@@ -251,7 +251,8 @@ def test_eps_is_derived_per_instance(default_vllm_config):
     with torch.device("meta"):
         for eps in (1e-5, 1e-6):
             module = RMSNorm(16, eps=eps)
-            built = get_fuser(module).fuse(module, "norm", default_vllm_config)
+            fuser = get_fuser(module, RMSNormFuser)
+            built = fuser.fuse(module, "norm", default_vllm_config)
             assert built.variance_epsilon == eps
 
 
@@ -262,7 +263,7 @@ def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config):
     with torch.device("meta"):
         for eps in (1e-5, 1e-6):
             module = cls(16, eps=eps)
-            fuser = get_fuser(module)
+            fuser = get_fuser(module, RMSNormFuser)
             assert fuser.eps_attr == cls.attr
             built = fuser.fuse(module, "norm", default_vllm_config)
             assert built.variance_epsilon == eps
@@ -274,7 +275,7 @@ def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplo
     logger = "vllm.model_executor.models.transformers.fusers.rms_norm"
     with caplog.at_level("DEBUG", logger=logger), torch.device("meta"):
         module = LiteralEpsRMSNorm()
-        fuser = get_fuser(module)
+        fuser = get_fuser(module, RMSNormFuser)
         built = fuser.fuse(module, "norm", default_vllm_config)
     assert fuser.eps_attr is None
     assert built.variance_epsilon == 1e-4
@@ -287,7 +288,7 @@ def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config):
     with torch.device("meta"):
         module = AmbiguousEpsRMSNorm(16, eps=1e-6)
         before = dict(vars(module))
-        fuser = get_fuser(module)
+        fuser = get_fuser(module, RMSNormFuser)
         assert fuser.eps_attr == "variance_epsilon"
         assert vars(module) == before