@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8ad2076c4a

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

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Thomas Ortner committed Sep 4, 2026 at 18:26 UTC 8ad2076c4a0e2a7e5ccf5b11455f737430ad12db

14 files changed +522 -151

docs/models/supported_models.md

+33 -16

@@ -72,19 +72,26 @@ This means that, with the Transformers modeling backend for vLLM, new models can

73 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)).

75 -To make your model compatible with the Transformers modeling backend, it needs:

76 -

77 -1. `kwargs` passed down through all modules from `MyModel` to `MyAttention`.

78 - - If your model is encoder-only:

79 - 1. Add `is_causal = False` to `MyAttention`.

80 - - If your model is mixture-of-experts (MoE):

81 - 1. Your sparse MoE block must have an attribute called `experts`.

82 - 2. The class of `experts` (`MyExperts`) must either:

83 - - Inherit from `nn.ModuleList` (naive).

84 - - Or contain all 3D `nn.Parameters` (packed).

85 - 3. `MyExperts.forward` must accept `hidden_states`, `top_k_index`, `top_k_weights`.

86 -2. `MyAttention` must use `ALL_ATTENTION_FUNCTIONS` to call attention.

87 -3. `MyModel` must contain `_supports_attention_backend = True`.

75 +To make your model compatible with the Transformers modeling backend:

76 +

77 +1. `MyAttention` must use `ALL_ATTENTION_FUNCTIONS` to call attention.

78 + - It must make exactly one such call. vLLM attaches one `Attention` layer per `MyAttention` module.

79 + - `MyAttention` must contain a unique `layer_idx`. vLLM keys its KV cache using this index.

80 + - Pass `scaling=` to the interface if your scale is not `head_size**-0.5`. vLLM reads it from the call to the attention interface.

81 +2. If your model is encoder-only:

82 + 1. Add `is_causal = False` to `MyAttention`.

83 +3. If your model is mixture-of-experts (MoE):

84 + 1. Your sparse MoE block must have an attribute called `experts`.

85 + 2. The class of `experts` (`MyExperts`) must either:

86 + - Inherit from `nn.ModuleList` (naive).

87 + - Or contain all 3D `nn.Parameters` (packed).

88 + 3. `MyExperts.forward` must accept `hidden_states`, `top_k_index`, `top_k_weights`.

89 +

90 +!!! note

91 + `MyModel` no longer needs `_supports_attention_backend = True`, and `kwargs` no

92 + longer need to be passed down through every module from `MyModel` to

93 + `MyAttention`. vLLM reaches its attention layer through `MyAttention` itself, so

94 + all it asks is that `MyAttention` dispatches through the interface.

96 <details class="code">

97 <summary>modeling_my_model.py</summary>

@@ -97,14 +104,24 @@ from torch import nn

104 class MyAttention(nn.Module):

105 is_causal = False # Only do this for encoder-only models

106

107 + def __init__(self, config, layer_idx):

108 + ...

109 + self.config = config

110 + self.layer_idx = layer_idx

111 + self.scaling = self.head_dim**-0.5

112 + ...

113 +

114 def forward(self, hidden_states, **kwargs):

115 ...

102 - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

116 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(

117 + self.config._attn_implementation, eager_attention_forward

118 + )

119 attn_output, attn_weights = attention_interface(

120 self,

121 query_states,

122 key_states,

123 value_states,

124 + scaling=self.scaling,

125 **kwargs,

126 )

127 ...

@@ -127,7 +144,7 @@ class MySparseMoEBlock(nn.Module):

144 ...

145

146 class MyModel(PreTrainedModel):

130 - _supports_attention_backend = True

147 + ...

148 ```

149

150 </details>

@@ -135,7 +152,7 @@ class MyModel(PreTrainedModel):

152 Here is what happens in the background when this model is loaded:

153

154 1. The config is loaded.

138 -2. `MyModel` Python class is loaded from the `auto_map` in config, and we check that the model `is_backend_compatible()`.

155 +2. `MyModel` Python class is loaded from the `auto_map` in config, and we check that the model `_can_set_attn_implementation()`.

156 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.

157

158 That's it!

tests/models/transformers/fusers/test_linear.py

+41 -35

@@ -10,7 +10,7 @@ import torch

10 import torch.nn as nn

11 import torch.nn.functional as F

13 -from vllm.model_executor.models.transformers.fuser import get_fuser

13 +from vllm.model_executor.models.transformers.fuser import get_fuser, get_fusers

14 from vllm.model_executor.models.transformers.fusers import (

15 GLUFuser,

16 PackedQKVFuser,

@@ -307,7 +307,7 @@ class PerHeadSplitAttention(nn.Module):

307

308

309 class FakeSelfAttn(nn.Module):

310 - """Stand-in for the vLLM `Attention` looked up in `attention_instances`."""

310 + """Stand-in for the vLLM `Attention` attached to the dispatching module."""

311

312 def __init__(self):

313 super().__init__()

@@ -328,9 +328,9 @@ class FakeMQASelfAttn(FakeSelfAttn):

328

329 @pytest.fixture(autouse=True)

330 def _clear_fuser_cache():

331 - get_fuser.cache_clear()

331 + get_fusers.cache_clear()

332 yield

333 - get_fuser.cache_clear()

333 + get_fusers.cache_clear()

334

335

336 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)

392 def test_detects_and_rewrites_glu(mlp_cls, bias):

393 with torch.device("meta"):

394 meta = mlp_cls(bias=bias)

395 - fuser = get_fuser(meta)

395 + fuser = get_fuser(meta, GLUFuser)

396 assert isinstance(fuser, GLUFuser)

397 assert (

398 fuser.gate_name,

@@ -427,9 +427,9 @@ def test_glu_identifies_down_projection():

427 matches the column-parallel merged gate/up; `None` when there is no such

428 projection to force (fusion of gate/up still applies)."""

429 with torch.device("meta"):

430 - assert get_fuser(GLUMLP()).down_name == "down_proj"

431 - assert get_fuser(ReversedGLUMLP()).down_name == "down_proj"

432 - assert get_fuser(NoDownGLU()).down_name is None

430 + assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj"

431 + assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj"

432 + assert get_fuser(NoDownGLU(), GLUFuser).down_name is None

433

434

435 @pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])

@@ -439,7 +439,7 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):

439 pytest.skip("MHA q/k/v assignment is order-based by design")

440 with torch.device("meta"):

441 meta = attn_cls(kv_heads=kv_heads)

442 - fuser = get_fuser(meta)

442 + fuser = get_fuser(meta, QKVFuser)

443 assert isinstance(fuser, QKVFuser)

444 # q (sharded differently under TP) must be identified exactly; k/v may be

445 # swapped for non-canonical compute order, which is numerically consistent

@@ -467,27 +467,27 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):

467 for p in real.parameters():

468 nn.init.normal_(p, std=0.05)

469 x = torch.randn(1, 5, 32)

470 - attention_instances = {3: FakeSelfAttn()}

471 - expected, _ = real(x, attention_instances=attention_instances)

470 + real.attn = FakeSelfAttn()

471 + expected, _ = real(x)

472 fused = _apply_qkv_fuser_with_stubs(real, fuser)

473

474 # Fusion is in place: the module keeps its class and other attributes

475 assert fused is real and type(fused) is attn_cls

476 assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None

477 - out, _ = fused(x, attention_instances=attention_instances)

477 + out, _ = fused(x)

478 torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)

479

480

481 def test_qkv_identifies_output_projection():

482 with torch.device("meta"):

483 - assert get_fuser(FakeAttention()).o_name == "o_proj"

484 - assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"

485 - assert get_fuser(ExtraProjAttention()).o_name == "o_proj"

483 + assert get_fuser(FakeAttention(), QKVFuser).o_name == "o_proj"

484 + assert get_fuser(ReversedFakeAttention(), QKVFuser).o_name == "o_proj"

485 + assert get_fuser(ExtraProjAttention(), QKVFuser).o_name == "o_proj"

486 # Norm children (q_norm/k_norm) must not disturb o_proj identification.

487 - assert get_fuser(QKNormAttention()).o_name == "o_proj"

488 - assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"

487 + assert get_fuser(QKNormAttention(), QKVFuser).o_name == "o_proj"

488 + assert get_fuser(PerHeadQKNormAttention(), QKVFuser).o_name == "o_proj"

489 # A module between o_proj and the return is transparent.

490 - assert get_fuser(ResidDropoutAttention()).o_name == "o_proj"

490 + assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj"

491

492

493 @pytest.mark.parametrize("kv_heads", [1, 2])

@@ -498,7 +498,7 @@ def test_detects_and_rewrites_packed_qkv(kv_heads):

498 checkpoint weight as-is, and shards q by heads while replicating k/v."""

499 with torch.device("meta"):

500 meta = PackedQKVAttention(kv_heads=kv_heads)

501 - fuser = get_fuser(meta)

501 + fuser = get_fuser(meta, PackedQKVFuser)

502 assert isinstance(fuser, PackedQKVFuser)

503 assert (fuser.qkv_name, fuser.o_name) == ("c_attn", "c_proj")

504 assert (fuser.q_size, fuser.kv_size) == (32, 8 * kv_heads)

@@ -513,14 +513,14 @@ def test_detects_and_rewrites_packed_qkv(kv_heads):

513 for p in real.parameters():

514 nn.init.normal_(p, std=0.05)

515 x = torch.randn(1, 5, 32)

516 - attention_instances = {3: FakeMQASelfAttn()}

517 - expected, _ = real(x, attention_instances=attention_instances)

516 + real.attn = FakeMQASelfAttn()

517 + expected, _ = real(x)

518 fused = _apply_packed_qkv_fuser_with_stubs(real, fuser)

519

520 # Fusion is in place: the module keeps its class and other attributes

521 assert fused is real and type(fused) is PackedQKVAttention

522 assert fused.layer_idx == 3 and fused.is_causal

523 - out, _ = fused(x, attention_instances=attention_instances)

523 + out, _ = fused(x)

524 torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)

525

526

@@ -528,28 +528,28 @@ def test_per_head_split_is_not_packed_qkv():

528 """The split must consume the whole projection, else its sizes are head

529 widths and re-sharding by them would be wrong."""

530 with torch.device("meta"):

531 - assert get_fuser(PerHeadSplitAttention()) is None

531 + assert get_fuser(PerHeadSplitAttention(), PackedQKVFuser) is None

532

533

534 def test_fuser_is_cached_per_class_and_structure():

535 with torch.device("meta"):

536 - fuser_a = get_fuser(GLUMLP())

537 - fuser_b = get_fuser(GLUMLP())

536 + fuser_a = get_fuser(GLUMLP(), GLUFuser)

537 + fuser_b = get_fuser(GLUMLP(), GLUFuser)

538 assert fuser_a is fuser_b

539 - assert any(key[0] is GLUMLP for key in get_fuser.cache)

539 + assert any(key[0] is GLUMLP for key in get_fusers.cache)

540

541

542 @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])

543 def test_non_matching_modules_return_none(cls):

544 with torch.device("meta"):

545 module = cls()

546 - assert get_fuser(module) is None

546 + assert get_fuser(module, GLUFuser) is None

547

548

549 def test_untraceable_tail_still_fuses():

550 with torch.device("meta"):

551 meta = UntraceableTailGLUMLP()

552 - fuser = get_fuser(meta)

552 + fuser = get_fuser(meta, GLUFuser)

553 assert isinstance(fuser, GLUFuser)

554

555 # Numerics: the live tail must survive the rewrite

@@ -566,8 +566,8 @@ def test_weight_mappings_are_scoped_to_fused_prefixes():

566 from vllm.model_executor.models.utils import WeightsMapper

567

568 with torch.device("meta"):

569 - glu_fuser = get_fuser(GLUMLP())

570 - qkv_fuser = get_fuser(FakeAttention())

569 + glu_fuser = get_fuser(GLUMLP(), GLUFuser)

570 + qkv_fuser = get_fuser(FakeAttention(), QKVFuser)

571

572 mapper = WeightsMapper()

573 for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):

@@ -620,7 +620,7 @@ def test_weight_mappings_are_scoped_to_fused_prefixes():

620 def test_unfusable_modules_are_not_fused(cls, default_vllm_config):

621 with torch.device("meta"):

622 module = cls()

623 - fuser = get_fuser(module)

623 + fuser = get_fuser(module, GLUFuser)

624 # Either no pattern matches the class, or this instance fails validation

625 # (`recursive_replace` gates fusion and its weight mappings on `validate`)

626 assert fuser is None or not fuser.validate(module, default_vllm_config)

@@ -654,9 +654,15 @@ def _wider_model_config(head_dim: int) -> SimpleNamespace:

654

655

656 @pytest.mark.parametrize(

657 - "cls, fuser_module", [(FakeAttention, qkv), (PackedQKVAttention, packed_qkv)]

657 + "cls, fuser_module, fuser_cls",

658 + [

659 + (FakeAttention, qkv, QKVFuser),

660 + (PackedQKVAttention, packed_qkv, PackedQKVFuser),

661 + ],

662 )

659 -def test_head_counts_come_from_the_module_not_the_model(cls, fuser_module, monkeypatch):

663 +def test_head_counts_come_from_the_module_not_the_model(

664 + cls, fuser_module, fuser_cls, monkeypatch

665 +):

666 """A layer narrower than the model-wide head size must not be miscounted.

667

668 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

686 monkeypatch.setattr(

687 fuser_module, "replace_linear_class", lambda *a, **kw: nn.Identity()

688 )

683 - fuser = get_fuser(module)

689 + fuser = get_fuser(module, fuser_cls)

690 assert fuser is not None and fuser.validate(module, vllm_config)

691 fuser.update_attrs(module, "model.layers.0.self_attn", vllm_config)

692

@@ -700,5 +706,5 @@ def test_validate_accepts_a_layer_the_model_wide_head_size_would_reject():

706

707 # kv width is 24, not a multiple of the model-wide 16, but is of this

708 # layer's 8.

703 - fuser = get_fuser(module)

709 + fuser = get_fuser(module, QKVFuser)

710 assert fuser is not None and fuser.validate(module, vllm_config)

tests/models/transformers/fusers/test_mla.py

+1 -1

@@ -158,4 +158,4 @@ class GLU(nn.Module):

158

159 def test_non_mla_is_not_matched():

160 with torch.device("meta"):

161 - assert not isinstance(get_fuser(GLU()), MLAFuser)

161 + assert get_fuser(GLU(), MLAFuser) is None

tests/models/transformers/fusers/test_rms_norm.py

+11 -10

@@ -162,7 +162,7 @@ class UntraceableGatedRMSNorm(RMSNorm):

162 )

163 def test_detects_rms_norm_variants(cls, eps, zero_centered):

164 with torch.device("meta"):

165 - fuser = get_fuser(cls(16, eps=eps))

165 + fuser = get_fuser(cls(16, eps=eps), RMSNormFuser)

166 assert isinstance(fuser, RMSNormFuser)

167 assert fuser.zero_centered == zero_centered

168

@@ -171,7 +171,7 @@ def test_detects_rms_norm_variants(cls, eps, zero_centered):

171 def test_non_rms_norms_are_not_matched(cls):

172 with torch.device("meta"):

173 module = cls(16) if cls is nn.LayerNorm else cls()

174 - assert not isinstance(get_fuser(module), RMSNormFuser)

174 + assert get_fuser(module, RMSNormFuser) is None

175

176

177 @pytest.mark.parametrize(

@@ -179,7 +179,7 @@ def test_non_rms_norms_are_not_matched(cls):

179 )

180 def test_gated_rms_norm_is_not_fused(cls):

181 with torch.device("meta"):

182 - assert not isinstance(get_fuser(cls()), RMSNormFuser)

182 + assert get_fuser(cls(), RMSNormFuser) is None

183

184

185 @pytest.mark.parametrize(

@@ -196,7 +196,7 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c

196

197 with torch.device("meta"):

198 module = cls()

199 - fuser = get_fuser(module)

199 + fuser = get_fuser(module, RMSNormFuser)

200 built = fuser.fuse(module, "norm", default_vllm_config)

201 from vllm.model_executor.models.transformers.fusers.rms_norm import (

202 TPAwareNormMixin,

@@ -221,7 +221,7 @@ def test_weightless_norm_has_no_hidden_size(default_vllm_config):

221 matching the unfused norm and vLLM's native `RMSNorm(hidden_size=head_dim)`.

222 """

223 module = WeightlessRMSNorm(72)

224 - built = get_fuser(module).fuse(module, "norm", default_vllm_config)

224 + built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config)

225 assert built.hidden_size == 0

226

227 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):

236

237 with torch.device("meta"):

238 module = torch.nn.RMSNorm(16) # forward is a single `F.rms_norm` call

239 - fuser = get_fuser(module)

239 + fuser = get_fuser(module, RMSNormFuser)

240 assert isinstance(fuser, RMSNormFuser)

241 assert not fuser.zero_centered

242 vllm_config = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float32))

@@ -251,7 +251,8 @@ def test_eps_is_derived_per_instance(default_vllm_config):

251 with torch.device("meta"):

252 for eps in (1e-5, 1e-6):

253 module = RMSNorm(16, eps=eps)

254 - built = get_fuser(module).fuse(module, "norm", default_vllm_config)

254 + fuser = get_fuser(module, RMSNormFuser)

255 + built = fuser.fuse(module, "norm", default_vllm_config)

256 assert built.variance_epsilon == eps

257

258

@@ -262,7 +263,7 @@ def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config):

263 with torch.device("meta"):

264 for eps in (1e-5, 1e-6):

265 module = cls(16, eps=eps)

265 - fuser = get_fuser(module)

266 + fuser = get_fuser(module, RMSNormFuser)

267 assert fuser.eps_attr == cls.attr

268 built = fuser.fuse(module, "norm", default_vllm_config)

269 assert built.variance_epsilon == eps

@@ -274,7 +275,7 @@ def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplo

275 logger = "vllm.model_executor.models.transformers.fusers.rms_norm"

276 with caplog.at_level("DEBUG", logger=logger), torch.device("meta"):

277 module = LiteralEpsRMSNorm()

277 - fuser = get_fuser(module)

278 + fuser = get_fuser(module, RMSNormFuser)

279 built = fuser.fuse(module, "norm", default_vllm_config)

280 assert fuser.eps_attr is None

281 assert built.variance_epsilon == 1e-4

@@ -287,7 +288,7 @@ def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config):

288 with torch.device("meta"):

289 module = AmbiguousEpsRMSNorm(16, eps=1e-6)

290 before = dict(vars(module))

290 - fuser = get_fuser(module)

291 + fuser = get_fuser(module, RMSNormFuser)

292 assert fuser.eps_attr == "variance_epsilon"

293 assert vars(module) == before

294

tests/models/transformers/test_backend.py

+142

@@ -2,9 +2,11 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Test the functionality of the Transformers modeling backend."""

5 +import ast

6 import contextlib

7 import os

8 import tempfile

9 +from types import SimpleNamespace

10 from typing import Any

12 import pytest

@@ -15,6 +17,8 @@ from transformers import AutoConfig, AutoModel, PretrainedConfig

17 from vllm.config import ModelConfig, VllmConfig

18 from vllm.model_executor.models.interfaces import SupportsMultiModal

19 from vllm.model_executor.models.transformers.base import Base

20 +from vllm.model_executor.models.transformers.fusers import AttentionFuser

21 +from vllm.model_executor.models.transformers.fusers.attention import VLLM_ATTN_IMPL

22 from vllm.model_executor.models.transformers.multimodal import MultiModalMixin

23 from vllm.model_executor.models.utils import StageMissingLayer

@@ -583,3 +587,141 @@ def test_marking_skipped_without_tokenizer():

587 """

588 vision_tower = build_marked_model(0, skip_tokenizer_init=True).vision_tower

589 assert not isinstance(vision_tower, StageMissingLayer)

590 +

591 +

592 +NUM_LAYERS = 4

593 +

594 +

595 +def build_model(model_type: str, **overrides) -> nn.Module:

596 + """A tiny HF model of `NUM_LAYERS` layers, built on the meta device."""

597 + try:

598 + config = AutoConfig.for_model(

599 + model_type,

600 + num_hidden_layers=NUM_LAYERS,

601 + hidden_size=64,

602 + intermediate_size=128,

603 + num_attention_heads=4,

604 + num_key_value_heads=2,

605 + head_dim=16,

606 + vocab_size=VOCAB_SIZE,

607 + **overrides,

608 + )

609 + except ValueError:

610 + pytest.skip(f"The installed transformers has no {model_type!r} model")

611 + with torch.device("meta"):

612 + return AutoModel.from_config(config)

613 +

614 +

615 +ATTENTION_MODEL_TYPES = [

616 + # The default `head_size**-0.5` happens to be right here

617 + "llama",

618 + # Scales the query by a learnable per-dim weight, so it declares 1.0

619 + "gemma4_text",

620 + # MLA, and declares yarn's `mscale`

621 + "deepseek_v3",

622 +]

623 +

624 +

625 +@pytest.mark.parametrize("model_type", ATTENTION_MODEL_TYPES)

626 +def test_attention_dispatch_is_matched(model_type: str):

627 + """Exactly the decoder layers' attention modules match an `AttentionFuser`.

628 +

629 + That match is what `recursive_replace` records, and so what

630 + `create_attention_instances` attaches vLLM's attention layer to.

631 + """

632 + model = build_model(model_type)

633 + matched = {

634 + name

635 + for name, module in model.named_modules()

636 + if AttentionFuser.match(None, module) is not None

637 + }

638 + assert matched == {f"layers.{i}.self_attn" for i in range(NUM_LAYERS)}

639 +

640 +

641 +def test_attention_layer_index_is_the_modules_own():

642 + """The layer served comes from the module, not from its position in the stack.

643 +

644 + LongCat Flash gives each decoder layer two attention sublayers numbered

645 + `2i` and `2i + 1`, so `num_hidden_layers` is twice the length of the stack

646 + and the enclosing layer's position is not the index the KV cache is keyed by.

647 + A vision tower is excluded separately, by `validate`: only the text config is

648 + patched to dispatch to vLLM, so a tower's attention is left to Transformers.

649 + """

650 + model = build_model("longcat_flash", num_layers=2)

651 + text_config = model.config.get_text_config()

652 + assert len(model.layers) * 2 == text_config.num_hidden_layers

653 +

654 + attentions = [

655 + model.get_submodule(f"layers.{i}.self_attn.{j}") for i in (0, 1) for j in (0, 1)

656 + ]

657 + fusers = [AttentionFuser.match(None, attn) for attn in attentions]

658 + assert all(fuser is not None for fuser in fusers)

659 + assert [f.layer_index(a) for f, a in zip(fusers, attentions)] == [0, 1, 2, 3]

660 +

661 + # Not dispatching to vLLM yet, as a vision tower never would be. The vLLM

662 + # config is not consulted, only the one the module was built with.

663 + assert not any(f.validate(a, None) for f, a in zip(fusers, attentions))

664 + text_config._attn_implementation = VLLM_ATTN_IMPL

665 + assert all(f.validate(a, None) for f, a in zip(fusers, attentions))

666 +

667 +

668 +def test_attention_dispatch_is_required(monkeypatch: pytest.MonkeyPatch):

669 + """Models selected by the permissive registry fail clearly without dispatch."""

670 + monkeypatch.setattr(

671 + "vllm.model_executor.models.transformers.base.get_pp_indices",

672 + lambda *_: (0, 1),

673 + )

674 + model = SimpleNamespace(

675 + text_config=SimpleNamespace(num_hidden_layers=1),

676 + pp_group=SimpleNamespace(rank_in_group=0, world_size=1),

677 + attention_fusers={},

678 + _get_attn_cls=lambda: None,

679 + )

680 +

681 + with pytest.raises(

682 + ValueError,

683 + match="Layer 0 does not dispatch through the Transformers attention interface",

684 + ):

685 + Base.create_attention_instances(model)

686 +

687 +

688 +@pytest.mark.parametrize("model_type", ATTENTION_MODEL_TYPES)

689 +def test_attention_scale_is_the_declared_one(model_type: str):

690 + """The scale comes off the module, not from the `head_size**-0.5` default.

691 +

692 + Two of the three model types declare something the default would silently

693 + replace, so they are what makes this more than a tautology.

694 + """

695 + attention = build_model(model_type).layers[0].self_attn

696 + fuser = AttentionFuser.match(None, attention)

697 + assert fuser is not None

698 + assert fuser.scale(attention) == attention.scaling

699 +

700 +

701 +def test_attention_scale_is_the_argument_not_the_attribute():

702 + """OPT applies the scale to the query itself, then declares a literal 1.0.

703 +

704 + Its `self.scaling` is the `head_size**-0.5` it has already applied, so

705 + reading the attribute rather than the argument the module hands the

706 + interface would scale twice.

707 + """

708 + model = build_model("opt", ffn_dim=128, word_embed_proj_dim=64)

709 + attention = model.get_submodule("decoder.layers.0.self_attn")

710 + fuser = AttentionFuser.match(None, attention)

711 + assert fuser is not None

712 + assert fuser.scale(attention) == 1.0 != attention.scaling

713 +

714 +

715 +def test_attention_scale_rejects_unresolvable_expression():

716 + """Unresolvable declared scales must not silently use the Llama default."""

717 +

718 + class Attention(nn.Module):

719 + head_dim = 16

720 +

721 + fuser = AttentionFuser(

722 + source_cls="Attention",

723 + scale_expr=ast.parse("self.head_dim**-0.5 * self.factor", mode="eval").body,

724 + )

725 +

726 + with pytest.raises(ValueError, match="Cannot resolve attention scaling expression"):

727 + fuser.scale(Attention())

vllm/model_executor/models/registry.py

+4 -1

@@ -1253,7 +1253,10 @@ class _ModelRegistry:

1253 "'auto_map' (relevant if the model is custom)."

1254 )

1255

1256 - if not model_module.is_backend_compatible():

1256 + if not (

1257 + model_module.is_backend_compatible()

1258 + or model_module._can_set_attn_implementation()

1259 + ):

1260 if model_config.model_impl != "transformers":

1261 return None

1262

vllm/model_executor/models/transformers/__init__.py

+3 -15

@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING

21 import torch.nn.functional as F

22 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

24 -from vllm.model_executor.models.transformers.base import Base

24 +from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR, Base

25 from vllm.model_executor.models.transformers.causal import CausalMixin

26 from vllm.model_executor.models.transformers.legacy import LegacyMixin

27 from vllm.model_executor.models.transformers.moe import MoEMixin

@@ -40,8 +40,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY

40 if TYPE_CHECKING:

41 import torch

43 - from vllm.model_executor.layers.attention import Attention, MLAAttention

44 -

44 def vllm_attention_forward(

45 # Transformers args

@@ -50,15 +48,9 @@ def vllm_attention_forward(

48 key: "torch.Tensor",

49 value: "torch.Tensor",

50 attention_mask: "torch.Tensor",

53 - # Transformers kwargs

54 - scaling: float | None = None,

55 - # vLLM kwargs

56 - attention_instances: "dict[int, Attention] | None" = None,

51 **kwargs,

52 ):

59 - self_attn = attention_instances[module.layer_idx]

60 - if scaling is not None:

61 - self_attn.impl.scale = float(scaling)

53 + self_attn = getattr(module, VLLM_ATTN_ATTR)

54 hidden = query.shape[-2]

55 head_dim_qk = query.shape[-1]

56 head_dim_v = value.shape[-1]

@@ -85,13 +77,9 @@ def vllm_mla_attention_forward(

77 kv_c_normed: "torch.Tensor",

78 k_pe: "torch.Tensor",

79 attention_mask: "torch.Tensor",

88 - # Transformers kwargs

89 - scaling: float | None = None,

90 - # vLLM kwargs

91 - attention_instances: "dict[int, MLAAttention] | None" = None,

80 **kwargs,

81 ):

94 - self_attn = attention_instances[module.layer_idx]

82 + self_attn = getattr(module, VLLM_ATTN_ATTR)

83 # [batch=1, heads, num_tokens, qk_head_dim] -> [num_tokens, heads, qk_head_dim]

84 query = query.transpose(1, 2).flatten(0, 1)

85 num_tokens, num_heads = query.shape[:2]

vllm/model_executor/models/transformers/base.py

+78 -34

@@ -57,7 +57,11 @@ from vllm.model_executor.models.interfaces import (

57 )

58 from vllm.model_executor.models.interfaces_base import VllmModel

59 from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers

60 -from vllm.model_executor.models.transformers.fusers import MLAFuser

60 +from vllm.model_executor.models.transformers.fusers import AttentionFuser, MLAFuser

61 +from vllm.model_executor.models.transformers.fusers.attention import (

62 + VLLM_ATTN_IMPL,

63 + VLLM_MLA_ATTN_IMPL,

64 +)

65 from vllm.model_executor.models.transformers.utils import (

66 attrsetter,

67 can_enable_torch_compile,

@@ -73,7 +77,6 @@ from vllm.model_executor.models.utils import (

77 AutoWeightsLoader,

78 PPMissingLayer,

79 WeightsMapper,

76 - extract_layer_index,

80 make_empty_intermediate_tensors_factory,

81 maybe_prefix,

82 )

@@ -86,6 +89,11 @@ if TYPE_CHECKING:

90 logger = init_logger(__name__)

92 +VLLM_ATTN_ATTR = "attn"

93 +"""Attribute in which to store the vLLM `Attention` instance that serves a Transformers

94 +attention module. `"attn"` is chosen deliberately to match the mappings generated by

95 +`maybe_remap_kv_scale_name`."""

96 +

98 class PreTrainedModelClasses(NamedTuple):

99 decoder: type["PreTrainedModel"]

@@ -129,9 +137,13 @@ class Base(

137 self.packed_modules_mapping: dict[str, list[str]] = {}

138 """Fused module -> constituent projections, populated by `recursive_replace`

139 for the quantization machinery and loaders (e.g. bitsandbytes)."""

132 - self.fusers: dict[str, BaseFuser] = {}

133 - """Module qualname -> the fuser applied to it, populated

140 + self.fusers: dict[str, list[BaseFuser]] = {}

141 + """Module qualname -> the fusers applied to it, populated

142 by `recursive_replace` for `create_attention_instances`."""

143 + self.attention_fusers: dict[int, tuple[str, AttentionFuser]] = {}

144 + """`layer_idx` -> the qualname and fuser of the module computing that

145 + layer's attention, populated by `recursive_replace` for

146 + `create_attention_instances`."""

147

148 # Attrs for Eagle3 (see self.set_aux_hidden_state_layers)

149 self._target_class: type[nn.Module] = nn.Module

@@ -189,7 +201,7 @@ class Base(

201 - Sets the dtype to the default torch dtype set by vLLM because Transformers

202 uses the config dtype when creating the model

203 """

192 - self.text_config._attn_implementation = "vllm"

204 + self.text_config._attn_implementation = VLLM_ATTN_IMPL

205 self.config.dtype = torch.get_default_dtype()

206

207 @contextmanager

@@ -475,9 +487,22 @@ class Base(

487

488 vocab_embeddings = self._vocab_embeddings()

489

478 - def register_fusion(fuser: BaseFuser, prefix: str):

490 + def register_fusion(fuser: BaseFuser, prefix: str, module: nn.Module):

491 """Register a fused layer's mappings just before it is built."""

480 - self.fusers[prefix] = fuser

492 + self.fusers.setdefault(prefix, []).append(fuser)

493 +

494 + if (

495 + isinstance(fuser, AttentionFuser)

496 + and (index := fuser.layer_index(module)) is not None

497 + ):

498 + if index in self.attention_fusers:

499 + other, _ = self.attention_fusers[index]

500 + raise ValueError(

501 + f"Layer {index} dispatches attention from more than one "

502 + f"module ({other} and {prefix}), which one vLLM attention "

503 + "layer cannot serve."

504 + )

505 + self.attention_fusers[index] = (prefix, fuser)

506

507 orig_to_new_stacked = fuser.orig_to_new_stacked(prefix)

508 self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)

@@ -528,10 +553,11 @@ class Base(

553 new_module = replace_embedding_class(

554 child_module, self.quant_config, prefix=qual_name

555 )

531 - elif (fuser := fusers[child_module]) is not None:

532 - register_fusion(fuser, qual_name)

533 - new_module = fuser.fuse(child_module, qual_name, self.vllm_config)

534 - logger.info_once(fuser.info(child_name))

556 + elif child_module_fusers := fusers[child_module]:

557 + for fuser in child_module_fusers:

558 + register_fusion(fuser, qual_name, child_module)

559 + new_module = fuser.fuse(new_module, qual_name, self.vllm_config)

560 + logger.info_once(fuser.info(child_name))

561 _recursive_replace(new_module, prefix=qual_name)

562 elif not isinstance(child_module, MoERunner):

563 # MoERunner can contain aliases of shared experts and gates,

@@ -548,20 +574,14 @@ class Base(

574 """

575 Create `Attention` instances to inform KV cache allocation.

576 """

551 - mla_fusers = {}

577 attention_instances = {}

578 text_config = self.text_config

579 attn_cls = self._get_attn_cls()

580

581 # kv_lora_rank indicates that this is an MLA model

582 if getattr(text_config, "kv_lora_rank", None) is not None:

558 - mla_fusers = {

559 - extract_layer_index(prefix): (prefix, fuser)

560 - for prefix, fuser in self.fusers.items()

561 - if isinstance(fuser, MLAFuser)

562 - }

583 if attn_cls is MLAAttention:

564 - text_config._attn_implementation = "vllm_mla"

584 + text_config._attn_implementation = VLLM_MLA_ATTN_IMPL

585 else:

586 # MLA model not using MLAAttention: recompute head_size for full attn

587 qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", 0)

@@ -570,12 +590,23 @@ class Base(

590 self.model_config.model_arch_config.head_size = qk_head_dim

591

592 logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)

593 + layer_types = getattr(text_config, "layer_types", None)

594

595 pp_rank = self.pp_group.rank_in_group

596 pp_size = self.pp_group.world_size

597 start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)

598

599 for i in range(start, end):

600 + if i not in self.attention_fusers:

601 + in_range = layer_types and i < len(layer_types)

602 + layer = f"{i} ({layer_types[i]})" if in_range else str(i)

603 + raise ValueError(

604 + f"Layer {layer} does not dispatch through the Transformers "

605 + "attention interface and vLLM has no other way to handle it."

606 + )

607 + prefix, attn_fuser = self.attention_fusers[i]

608 + attn_module = self.get_submodule(prefix)

609 +

610 # `[i]` is the whole-model config unless the checkpoint is

611 # heterogeneous, in which case it is this layer's own geometry.

612 arch_config = self.model_config.model_arch_config[i]

@@ -583,8 +614,10 @@ class Base(

614 self.parallel_config, arch_config

615 )

616 head_size = arch_config.head_size

586 - # Default to Llama scale, maybe updated in vllm_attention_forward

587 - scale = head_size**-0.5

617 + scale = attn_fuser.scale(attn_module)

618 + # Default to Llama scale if AttentionFuser couldn't identify it

619 + if scale is None:

620 + scale = head_size**-0.5

621 num_kv_heads = self.model_config.get_num_kv_heads(

622 self.parallel_config, arch_config

623 )

@@ -598,17 +631,22 @@ class Base(

631 )

632

633 if attn_cls is MLAAttention:

601 - prefix, fuser = mla_fusers[i]

602 - mla_module = self.get_submodule(prefix)

634 + mla_fuser = next(

635 + (f for f in self.fusers[prefix] if isinstance(f, MLAFuser)), None

636 + )

637 + if mla_fuser is None:

638 + raise ValueError(

639 + f"Layer {i} ({prefix}) did not fuse as MLA, but the rest of "

640 + "the model did, so its latent KV cache cannot be allocated."

641 + )

642 dims = get_mla_dims(self.model_config)

643 kwargs.update(

605 - scale=mla_module.scaling,

644 qk_nope_head_dim=dims.qk_nope_head_dim,

645 qk_rope_head_dim=dims.qk_rope_head_dim,

646 v_head_dim=dims.v_head_dim,

647 q_lora_rank=dims.q_lora_rank,

648 kv_lora_rank=dims.kv_lora_rank,

611 - kv_b_proj=mla_module.get_submodule(fuser.kv_b_proj_name),

649 + kv_b_proj=attn_module.get_submodule(mla_fuser.kv_b_proj_name),

650 )

651 else:

652 kwargs.update(

@@ -618,17 +656,20 @@ class Base(

656 )

657

658 # Handle interleaved sliding window attention

621 - if (

622 - hasattr(text_config, "layer_types")

623 - and text_config.layer_types[i] == "sliding_attention"

624 - ):

659 + if layer_types and layer_types[i] == "sliding_attention":

660 kwargs["per_layer_sliding_window"] = text_config.sliding_window

661

662 attn_instance = attn_cls(**kwargs)

628 - if attn_cls is MLAAttention:

629 - # Attach MLA attn_instance to mla_module so it appears in

630 - # model.named_modules() and runs its process_weights_after_loading

631 - mla_module._vllm_mla_attn = attn_instance

663 +

664 + # Attach the instance to the module it serves, which is the only way it

665 + # is in the module tree at all: a plain dict is invisible to

666 + # `nn.Module.__setattr__`. Being in the tree runs its

667 + # `process_weights_after_loading` and moves its KV-scale buffers with

668 + # `model.to(device)`, and lets the attention interface read it off the

669 + # module rather than index by `module.layer_idx`, an `int` that bakes

670 + # layer identity into the traced graph and so costs one compiled

671 + # artifact per layer.

672 + setattr(attn_module, VLLM_ATTN_ATTR, attn_instance)

673 attention_instances[i] = attn_instance

674 return attention_instances

675

@@ -645,7 +686,11 @@ class Base(

686 return EncoderOnlyAttention

687 if self.model_config.use_mla:

688 self.check_version("5.15.0.dev0", "optimized MLA support")

648 - if any(isinstance(fuser, MLAFuser) for fuser in self.fusers.values()):

689 + if any(

690 + isinstance(fuser, MLAFuser)

691 + for fusers in self.fusers.values()

692 + for fuser in fusers

693 + ):

694 return MLAAttention

695 logger.warning_once(

696 "This model uses MLA but `MLAFuser` failed to match and/or fuse any "

@@ -731,7 +776,6 @@ class Base(

776 inputs_embeds=inputs_embeds,

777 use_cache=False,

778 position_ids=positions,

734 - attention_instances=self.attention_instances,

779 return_dict=False,

780 **self._output_aux_hidden_states_kwargs,

781 **kwargs,

vllm/model_executor/models/transformers/fuser.py

+69 -35

@@ -2,20 +2,22 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Fuser detection for the Transformers modeling backend.

5 -`get_fuser` traces a module class once (see `fx_utils`) and matches it against

6 -each concrete fuser in `fusers`; `Fusers` caches the result per class for a

7 -whole model. `base.recursive_replace` then applies the matched fuser per

5 +`get_fusers` traces a module class once (see `fx_utils`) and matches it against each

6 +concrete fuser in `fusers`; `Fusers` caches the result per class for a whole model.

7 +At most one match may give the module a different forward, and any number of fusers that

8 +leave the forward alone apply alongside. `base.recursive_replace` then applies them per

9 instance. RMSNorm-shaped modules the tracer cannot match are warned about.

10 """

12 from collections import UserDict

12 -from typing import TYPE_CHECKING

13 +from typing import TYPE_CHECKING, TypeVar

15 from cachetools import cached

16 from torch import nn

18 from vllm.logger import init_logger

19 from vllm.model_executor.models.transformers.fusers import (

20 + AttentionFuser,

21 BaseFuser,

22 GLUFuser,

23 MLAFuser,

@@ -31,58 +33,90 @@ if TYPE_CHECKING:

34 logger = init_logger(__name__)

36 +F = TypeVar("F", bound=BaseFuser)

37 +

39 def key(module: nn.Module) -> tuple:

36 - """Cache key for `get_fuser`. Considers module type and its immediate children."""

40 + """Cache key for `get_fusers`. Considers module type and its immediate children."""

41 return (type(module), tuple(name for name, _ in module.named_children()))

44 +FUSERS: tuple[type[BaseFuser], ...] = (

45 + MLAFuser,

46 + GLUFuser,

47 + QKVFuser,

48 + PackedQKVFuser,

49 + RMSNormFuser,

50 + AttentionFuser,

51 +)

52 +"""Every fuser, in match order: those that redefine the forward first, then those

53 +that leave it alone. A new fuser is added here."""

54 +

55 +

56 @cached(cache={}, key=key)

41 -def get_fuser(module: nn.Module) -> BaseFuser | None:

42 - """The fuser for `module`'s class and shape (cached), or `None` if no match."""

57 +def get_fusers(module: nn.Module) -> list[BaseFuser]:

58 + """Every fuser that applies to `module`'s class and shape (cached).

59 +

60 + The one that redefines the forward comes first, so applying them in order

61 + fuses before the rest read or adjust the result."""

62 # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a

63 # leaf module (raw tensor math, no submodules). Nothing else can match, and

64 # tracing is skipped for it.

65 n_linear = sum(isinstance(c, nn.Linear) for c in module.children())

66 is_leaf = next(module.children(), None) is None

48 - if n_linear < 2 and not is_leaf:

49 - return None

50 - if (graph := trace(module)) is None:

51 - return None

52 - for fuser_cls in (MLAFuser, GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser):

53 - if (fuser := fuser_cls.match(graph, module)) is not None:

54 - if isinstance(fuser, RewriteFuser):

55 - try:

56 - fuser.update_forward(module)

57 - except Exception as exc:

58 - logger.debug(

59 - "Attempted to fuse %s using %s but failed "

60 - "to update its forward method: %s",

61 - type(module),

62 - fuser_cls.__name__,

63 - exc,

64 - )

65 - continue

66 - return fuser

67 + graph = trace(module) if n_linear >= 2 or is_leaf else None

68 +

69 + fusers: list[BaseFuser] = []

70 + for fuser_cls in FUSERS:

71 + if fuser_cls.redefines_forward and (

72 + # A forward can only be rewritten from a trace of it, and only once,

73 + # because every rewrite starts from the original source.

74 + graph is None or any(fuser.redefines_forward for fuser in fusers)

75 + ):

76 + continue

77 + if (fuser := fuser_cls.match(graph, module)) is None:

78 + continue

79 + if isinstance(fuser, RewriteFuser):

80 + try:

81 + fuser.update_forward(module)

82 + except Exception as exc:

83 + logger.debug(

84 + "Attempted to fuse %s using %s but failed "

85 + "to update its forward method: %s",

86 + type(module),

87 + fuser_cls.__name__,

88 + exc,

89 + )

90 + continue

91 + fusers.append(fuser)

92 +

93 # A norm we could not match structurally is left unfused; flag likely misses.

68 - if module.__class__.__name__.endswith("RMSNorm"):

94 + if (

95 + module.__class__.__name__.endswith("RMSNorm")

96 + and graph is not None

97 + and not any(isinstance(fuser, RMSNormFuser) for fuser in fusers)

98 + ):

99 logger.warning_once(

100 "%s looks like an RMSNorm but its computation did not match the "

101 "expected pattern, so it was left unfused.",

102 module.__class__.__name__,

103 )

74 - return None

104 + return fusers

105 +

106 +

107 +def get_fuser(module: nn.Module, fuser_cls: type[F]) -> "F | None":

108 + """The first `fuser_cls` that applies to `module`, if one does."""

109 + return next((f for f in get_fusers(module) if isinstance(f, fuser_cls)), None)

110

111

112 class Fusers(UserDict):

78 - """Mapping from module class and shape to fuser, for all fusable modules."""

113 + """Mapping from module class and shape to fusers, for all fusable modules."""

114

115 def __init__(self, model: nn.Module, vllm_config: "VllmConfig"):

116 self.vllm_config = vllm_config

82 - super().__init__({key(m): get_fuser(m) for m in model.modules()})

117 + super().__init__({key(m): get_fusers(m) for m in model.modules()})

118

84 - def __getitem__(self, m: nn.Module) -> BaseFuser | None:

85 - fuser = self.data.get(key(m))

86 - if fuser is not None and fuser.validate(m, self.vllm_config):

87 - return fuser

88 - return None

119 + def __getitem__(self, m: nn.Module) -> list[BaseFuser]:

120 + """The fusers this instance can take, in the order to apply them."""

121 + fusers = self.data.get(key(m), ())

122 + return [fuser for fuser in fusers if fuser.validate(m, self.vllm_config)]

vllm/model_executor/models/transformers/fusers/__init__.py

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Concrete fusers for the Transformers modeling backend."""

5 +from vllm.model_executor.models.transformers.fusers.attention import AttentionFuser

6 from vllm.model_executor.models.transformers.fusers.base import (

7 BaseFuser,

8 RewriteFuser,

@@ -15,6 +16,7 @@ from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser

16 from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser

18 __all__ = [

19 + "AttentionFuser",

20 "BaseFuser",

21 "RewriteFuser",

22 "StackedFuser",

vllm/model_executor/models/transformers/fusers/attention.py new

+129

@@ -0,0 +1,129 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Attention fuser: the module that dispatches to the attention interface."""

4 +

5 +import ast

6 +import inspect

7 +import textwrap

8 +from dataclasses import dataclass

9 +from functools import cache

10 +from typing import TYPE_CHECKING, ClassVar

11 +

12 +from torch import fx, nn

13 +

14 +from vllm.model_executor.models.transformers.fusers.base import BaseFuser

15 +

16 +if TYPE_CHECKING:

17 + from vllm.config import VllmConfig

18 +

19 +VLLM_ATTN_IMPL = "vllm"

20 +VLLM_MLA_ATTN_IMPL = "vllm_mla"

21 +

22 +

23 +def _is_interface_lookup(node: ast.expr | None) -> bool:

24 + """Whether `node` reads an entry out of `ALL_ATTENTION_FUNCTIONS`."""

25 + # ALL_ATTENTION_FUNCTIONS.get_interface(...) or ALL_ATTENTION_FUNCTIONS[...]

26 + if isinstance(node, ast.Call):

27 + node = node.func

28 + if not isinstance(node, (ast.Attribute, ast.Subscript)):

29 + return False

30 + return (

31 + isinstance(node.value, ast.Name) and node.value.id == "ALL_ATTENTION_FUNCTIONS"

32 + )

33 +

34 +

35 +@cache

36 +def interface_call(cls: type[nn.Module]) -> ast.Call | None:

37 + """The attention interface call in `cls.forward`, if it makes exactly one."""

38 + try:

39 + source = inspect.getsource(inspect.unwrap(cls.forward))

40 + tree = ast.parse(textwrap.dedent(source))

41 + except (AttributeError, OSError, SyntaxError, TypeError):

42 + return None

43 +

44 + # Find the names of the local variables that read from the interface lookup

45 + names = set()

46 + for node in ast.walk(tree):

47 + if not isinstance(node, (ast.Assign, ast.AnnAssign)):

48 + continue

49 + if not _is_interface_lookup(node.value):

50 + continue

51 + targets = node.targets if isinstance(node, ast.Assign) else [node.target]

52 + names.update(t.id for t in targets if isinstance(t, ast.Name))

53 + # Find the calls to those names, and return the one if there is exactly one

54 + calls = [

55 + node

56 + for node in ast.walk(tree)

57 + if isinstance(node, ast.Call)

58 + and isinstance(node.func, ast.Name)

59 + and node.func.id in names

60 + ]

61 + return calls[0] if len(calls) == 1 else None

62 +

63 +

64 +def _resolve(node: ast.expr, module: nn.Module) -> object:

65 + """The value of `node` on `module`, for literals and `self.<attr>`."""

66 + if isinstance(node, ast.Constant):

67 + return node.value

68 + if (

69 + isinstance(node, ast.Attribute)

70 + and isinstance(node.value, ast.Name)

71 + and node.value.id == "self"

72 + ):

73 + return getattr(module, node.attr, None)

74 + return None

75 +

76 +

77 +@dataclass

78 +class AttentionFuser(BaseFuser):

79 + """A module that dispatches through the Transformers attention interface."""

80 +

81 + redefines_forward: ClassVar[bool] = False

82 +

83 + source_cls: str

84 + """Class of the HF module that dispatches (for logging)."""

85 + scale_expr: ast.expr | None

86 + """Source of the `scaling=` the module hands the interface, if it hands one."""

87 +

88 + def info(self, name: str) -> str:

89 + return f"Found: {name} ({self.source_cls}) -> attention interface"

90 +

91 + @classmethod

92 + def match(

93 + cls, graph: fx.Graph | None, module: nn.Module

94 + ) -> "AttentionFuser | None":

95 + if (call := interface_call(type(module))) is None:

96 + return None

97 + scaling = [kw.value for kw in call.keywords if kw.arg == "scaling"]

98 + scale_expr = scaling[0] if len(scaling) == 1 else None

99 + return cls(source_cls=type(module).__name__, scale_expr=scale_expr)

100 +

101 + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:

102 + """Whether `module` will actually dispatch to vLLM."""

103 + config = getattr(module, "config", None)

104 + # Only patched in the text config, this excludes attention based mm encoders

105 + vllm_attn_impls = {VLLM_ATTN_IMPL, VLLM_MLA_ATTN_IMPL}

106 + return getattr(config, "_attn_implementation", None) in vllm_attn_impls

107 +

108 + def fuse(

109 + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"

110 + ) -> nn.Module:

111 + return module

112 +

113 + def layer_index(self, module: nn.Module) -> int | None:

114 + """The layer `module` computes attention for, if it declares one."""

115 + layer_idx = getattr(module, "layer_idx", None)

116 + return layer_idx if isinstance(layer_idx, int) else None

117 +

118 + def scale(self, module: nn.Module) -> float | None:

119 + """The softmax scale `module` passes to the interface, or `None`."""

120 + if self.scale_expr is None:

121 + return None

122 + scale = _resolve(self.scale_expr, module)

123 + if not isinstance(scale, (int, float)) or isinstance(scale, bool):

124 + expression = ast.unparse(self.scale_expr)

125 + raise ValueError(

126 + f"Cannot resolve attention scaling expression {expression!r} in "

127 + f"{type(module).__name__}."

128 + )

129 + return float(scale)

vllm/model_executor/models/transformers/fusers/base.py

+5 -1

@@ -20,11 +20,15 @@ if TYPE_CHECKING:

20 class BaseFuser(ABC):

21 """A detected fusion and how to apply it.

23 - `match` analyses the module *class* once (cached, see `get_fuser`); `fuse`

23 + `match` analyses the module *class* once (cached, see `get_fusers`); `fuse`

24 then applies the fusion to an instance in `recursive_replace`, returning the

25 module to install in its place.

26 """

28 + redefines_forward: ClassVar[bool] = True

29 + """Whether `fuse` gives the module a different forward,

30 + by rewriting its source or by returning a different module in its place."""

31 +

32 @abstractmethod

33 def info(self, name: str) -> str:

34 """A human-readable description of the fusion at `name`, for logging."""

vllm/model_executor/models/transformers/fusers/mla.py

+1 -1

@@ -59,7 +59,7 @@ def _is_rms_norm(module: nn.Module) -> bool:

59 # Go via `get_fuser` because it caches

60 from vllm.model_executor.models.transformers.fuser import get_fuser

62 - return isinstance(get_fuser(module), RMSNormFuser)

62 + return get_fuser(module, RMSNormFuser) is not None

65 def _top_level_index(funcdef: ast.FunctionDef, node: ast.AST) -> int:

vllm/model_executor/models/transformers/moe.py

+3 -2

@@ -36,6 +36,7 @@ from vllm.model_executor.layers.fused_moe import (

36 )

37 from vllm.model_executor.models.interfaces import MixtureOfExperts

38 from vllm.model_executor.models.transformers.fuser import get_fuser

39 +from vllm.model_executor.models.transformers.fusers.glu import GLUFuser

40 from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser

41 from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix

42 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op

@@ -319,8 +320,8 @@ class MoEMixin(MixtureOfExperts):

320 # Store shared experts for later down projection adjustment

321 if shared_experts is not None:

322 hf_shared = shared_experts.shared_experts

322 - glu_fuser = get_fuser(hf_shared)

323 - down_name = getattr(glu_fuser, "down_name", None)

323 + glu_fuser = get_fuser(hf_shared, GLUFuser)

324 + down_name = glu_fuser and glu_fuser.down_name

325 if down_name is not None:

326 shared_down_projs.append((hf_shared, down_name))

327 # Prefer config, otherwise read it from fuser.