@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e960ead0aa

[Mypy] Fix mypy typing for Transformers models (#54320)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Taneem Ibrahim committed Sep 17, 2026 at 11:11 UTC e960ead0aa70362a57f205be831083b68c7a5bdb

16 files changed +224 -145

tests/models/transformers/fusers/test_mla.py

+2 -4

@@ -53,7 +53,6 @@ def _match(q_lora_rank: int | None) -> MLAFuser | None:

53 def test_discovers_modules_without_q_lora():

54 fuser = _match(q_lora_rank=None)

55 assert isinstance(fuser, MLAFuser)

56 - assert not fuser.has_q_lora

56 assert fuser.q_proj_name == "q_proj"

57 assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa"

58 assert fuser.kv_a_layernorm_name == "kv_a_layernorm"

@@ -68,7 +67,6 @@ def test_discovers_modules_without_q_lora():

67 def test_discovers_modules_with_q_lora():

68 fuser = _match(q_lora_rank=64)

69 assert isinstance(fuser, MLAFuser)

71 - assert fuser.has_q_lora

70 assert fuser.q_a_proj_name == "q_a_proj"

71 assert fuser.q_a_layernorm_name == "q_a_layernorm"

72 assert fuser.q_b_proj_name == "q_b_proj"

@@ -107,7 +105,7 @@ def test_update_forward_rewrites_real_attention(q_lora_rank):

105 assert isinstance(fuser, MLAFuser)

106 fuser.update_forward(_attention(q_lora_rank))

107 names = set(fuser.fused_forward.__code__.co_names)

110 - if fuser.has_q_lora:

108 + if fuser.q_a_proj_name is not None:

109 assert _FUSED_QKV_A_PROJ in names

110 assert not {"q_a_proj", "kv_a_proj_with_mqa"} & names

111

@@ -158,7 +156,7 @@ def test_discovers_modules_under_arbitrary_names():

156 module = RenamedMLA()

157 fuser = MLAFuser.match(trace(module), module)

158 assert isinstance(fuser, MLAFuser)

161 - assert not fuser.has_q_lora

159 + assert fuser.q_a_proj_name is None

160 assert fuser.q_proj_name == "alpha"

161 assert fuser.kv_a_proj_name == "beta"

162 assert fuser.kv_a_layernorm_name == "gamma"

tools/pre_commit/mypy.py

-1

@@ -116,7 +116,6 @@ EXCLUDE = [

116 r"vllm/model_executor/models/[nN]",

117 r"vllm/model_executor/models/[oO]",

118 r"vllm/model_executor/models/[qQ]",

119 - r"vllm/model_executor/models/transformers",

119 r"vllm/model_executor/models/[uU]",

120 r"vllm/model_executor/models/[vV]",

121 r"vllm/model_executor/models/[wW]",

vllm/compilation/decorators.py

+8 -5

@@ -7,7 +7,7 @@ import inspect

7 import os

8 import sys

9 from collections.abc import Callable, Generator

10 -from typing import TYPE_CHECKING, Any, TypeVar, overload

10 +from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, overload

11 from unittest.mock import patch

13 import torch

@@ -32,6 +32,9 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer

33 from .monitor import monitor_profiling_run, monitor_torch_compile

35 +DynamicArgDims: TypeAlias = dict[str, int | list[int] | dict[int, str]]

36 +"""Argument name -> the dynamic dimension(s) of that argument."""

37 +

38 # shape_id parameter was added to mark_unbacked in PyTorch 2.11.0

39 _SUPPORTS_SHAPE_ID = is_torch_equal_or_newer("2.11.0")

@@ -89,7 +92,7 @@ def support_torch_compile(

92 @overload

93 def support_torch_compile(

94 *,

92 - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None,

95 + dynamic_arg_dims: DynamicArgDims | None,

96 ) -> Callable[[type[_T]], type[_T]]: ...

@@ -103,7 +106,7 @@ def support_torch_compile(

106 @overload

107 def support_torch_compile(

108 *,

106 - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None,

109 + dynamic_arg_dims: DynamicArgDims | None = None,

110 mark_unbacked_dims: dict[str, int | list[int]] | None = None,

111 enable_if: Callable[[VllmConfig], bool] | None = None,

112 is_encoder: bool = False,

@@ -117,7 +120,7 @@ def support_torch_compile(cls: type[_T]) -> type[_T]: ...

120 def support_torch_compile(

121 cls: type[_T] | None = None,

122 *,

120 - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None,

123 + dynamic_arg_dims: DynamicArgDims | None = None,

124 mark_unbacked_dims: dict[str, int | list[int]] | None = None,

125 enable_if: Callable[[VllmConfig], bool] | None = None,

126 is_encoder: bool = False,

@@ -328,7 +331,7 @@ def _try_load_aot_compiled_fn(

331

332 def _support_torch_compile(

333 cls: type[_T],

331 - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]],

334 + dynamic_arg_dims: DynamicArgDims,

335 mark_unbacked_dims: dict[str, int | list[int]] | None = None,

336 enable_if: Callable[[VllmConfig], bool] | None = None,

337 is_encoder: bool = False,

vllm/model_executor/models/transformers/__init__.py

+11 -13

@@ -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 VLLM_ATTN_ATTR, Base

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

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

@@ -126,10 +126,10 @@ ALL_ATTENTION_FUNCTIONS.register("vllm_mla", vllm_mla_attention_forward)

126

127

128 # Text only models

129 -class TransformersForCausalLM(CausalMixin, Base): ...

129 +class TransformersForCausalLM(CausalMixin): ...

130

131

132 -class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...

132 +class TransformersMoEForCausalLM(MoEMixin, CausalMixin): ...

133

134

135 # Multimodal models

@@ -138,7 +138,7 @@ class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...

138 info=MultiModalProcessingInfo,

139 dummy_inputs=MultiModalDummyInputsBuilder,

140 )

141 -class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...

141 +class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin): ...

142

143

144 @MULTIMODAL_REGISTRY.register_processor(

@@ -146,16 +146,14 @@ class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...

146 info=MultiModalProcessingInfo,

147 dummy_inputs=MultiModalDummyInputsBuilder,

148 )

149 -class TransformersMultiModalMoEForCausalLM(

150 - MoEMixin, MultiModalMixin, CausalMixin, Base

151 -): ...

149 +class TransformersMultiModalMoEForCausalLM(MoEMixin, MultiModalMixin, CausalMixin): ...

150

151

152 # Embedding models

155 -class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ...

153 +class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin): ...

154

155

158 -class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...

156 +class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin): ...

157

158

159 @MULTIMODAL_REGISTRY.register_processor(

@@ -163,17 +161,17 @@ class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...

161 info=MultiModalProcessingInfo,

162 dummy_inputs=MultiModalDummyInputsBuilder,

163 )

166 -class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ...

164 +class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin): ...

165

166

167 # Sequence classification models

168 class TransformersForSequenceClassification(

171 - SequenceClassificationMixin, LegacyMixin, Base

169 + SequenceClassificationMixin, LegacyMixin

170 ): ...

171

172

173 class TransformersMoEForSequenceClassification(

176 - SequenceClassificationMixin, MoEMixin, Base

174 + SequenceClassificationMixin, MoEMixin

175 ): ...

176

177

@@ -183,7 +181,7 @@ class TransformersMoEForSequenceClassification(

181 dummy_inputs=MultiModalDummyInputsBuilder,

182 )

183 class TransformersMultiModalForSequenceClassification(

186 - SequenceClassificationMixin, MultiModalMixin, Base

184 + SequenceClassificationMixin, MultiModalMixin

185 ): ...

186

187

vllm/model_executor/models/transformers/base.py

+29 -15

@@ -35,7 +35,7 @@ from transformers.conversion_mapping import (

35 get_model_conversion_mapping,

36 )

38 -from vllm.compilation.decorators import support_torch_compile

38 +from vllm.compilation.decorators import DynamicArgDims, support_torch_compile

39 from vllm.config.utils import getattr_iter

40 from vllm.distributed import get_pp_group, get_tp_group

41 from vllm.distributed.utils import get_pp_indices

@@ -76,6 +76,7 @@ from vllm.model_executor.models.transformers.utils import (

76 from vllm.model_executor.models.utils import (

77 AutoWeightsLoader,

78 PPMissingLayer,

79 + ShardId,

80 WeightsMapper,

81 make_empty_intermediate_tensors_factory,

82 maybe_prefix,

@@ -110,10 +111,13 @@ class Base(

111 SupportsEagle,

112 SupportsEagle3,

113 ):

113 - embedding_modules = ["embed_tokens"] # TODO transformers will have a util to get it

114 + hf_to_vllm_mapper: WeightsMapper

115 +

116 + # TODO transformers will have a util to get these

117 + embedding_modules = {"embed_tokens": "input_embeddings"}

118

119 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

116 - super().__init__()

120 + nn.Module.__init__(self)

121 logger.info("Using Transformers modeling backend.")

122

123 self.vllm_config = vllm_config

@@ -248,7 +252,7 @@ class Base(

252 def _decorate_cls_for_torch_compile(

253 self,

254 cls: type["PreTrainedModel"],

251 - dynamic_arg_dims: dict[str, int] | None,

255 + dynamic_arg_dims: DynamicArgDims | None,

256 enable_if: Callable[["VllmConfig"], bool],

257 is_encoder: bool,

258 ):

@@ -284,7 +288,7 @@ class Base(

288 self._decorate_cls_for_torch_compile(

289 cls=self._pre_trained_model_classes.decoder,

290 # Applied to a PreTrainedModel so the batch dimension will exist

287 - dynamic_arg_dims=dict[str, int](

291 + dynamic_arg_dims=DynamicArgDims(

292 input_ids=1, # shape: [1, seq_len]

293 inputs_embeds=1, # shape: [1, seq_len, hidden_size]

294 position_ids=-1, # shape: [1, seq_len] or [3, 1, seq_len] for mrope

@@ -303,9 +307,8 @@ class Base(

307 - Checkpoints saved with no base model prefix

308 - Any quantization config specific mappings

309 """

306 - self.hf_to_vllm_mapper = WeightsMapper()

307 - orig_to_new_renaming = self.hf_to_vllm_mapper.orig_to_new_renaming

308 - orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex

310 + orig_to_new_renaming: list[WeightRenaming] = []

311 + orig_to_new_regex: dict[re.Pattern, str | None] = {}

312

313 for mapping in get_model_conversion_mapping(self.model):

314 # Handle weights which have been renamed in Transformers

@@ -341,6 +344,11 @@ class Base(

344 nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)")

345 orig_to_new_regex[nested_lm_head_pattern] = r"\2"

346

347 + self.hf_to_vllm_mapper = WeightsMapper(

348 + orig_to_new_renaming=orig_to_new_renaming,

349 + orig_to_new_regex=orig_to_new_regex,

350 + )

351 +

352 # Apply mapping to quantization config if needed

353 self._maybe_apply_model_mapping()

354

@@ -371,8 +379,8 @@ class Base(

379 "modeling backend will infer the split from the layers of %s in order "

380 "of declaration and keep parameter-free modules on every rank. This "

381 "may fail if the model's structure is non-standard. %s",

374 - type(self.model),

375 - type(module),

382 + type(self.model).__name__,

383 + type(module).__name__,

384 tip,

385 )

386

@@ -470,7 +478,7 @@ class Base(

478 "backend will shard the model the best it can during graph fusion and "

479 "replicate the rest. This may be suboptimal or fail if the model does "

480 "not fuse cleanly. %s",

473 - type(self.model),

481 + type(self.model).__name__,

482 tip,

483 )

484

@@ -481,6 +489,8 @@ class Base(

489

490 vocab_embeddings = self._vocab_embeddings()

491

492 + orig_to_new_stacked: dict[str, tuple[str, ShardId]] = {}

493 +

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

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

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

@@ -498,8 +508,7 @@ class Base(

508 )

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

510

501 - orig_to_new_stacked = fuser.orig_to_new_stacked(prefix)

502 - self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)

511 + orig_to_new_stacked.update(fuser.orig_to_new_stacked(prefix))

512

513 packed_modules_mapping = fuser.packed_modules_mapping

514 self.packed_modules_mapping.update(packed_modules_mapping)

@@ -564,6 +573,8 @@ class Base(

573

574 _recursive_replace(self.model, prefix="model")

575

576 + self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked)

577 +

578 def _create_attention_instances(self):

579 """Create `Attention` instances to inform KV cache allocation."""

580 text_config = self.text_config

@@ -589,8 +600,11 @@ class Base(

600

601 for i in range(start, end):

602 if i not in self.attention_fusers:

592 - in_range = layer_types and i < len(layer_types)

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

603 + layer = (

604 + f"{i} ({layer_types[i]})"

605 + if layer_types and i < len(layer_types)

606 + else str(i)

607 + )

608 raise ValueError(

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

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

vllm/model_executor/models/transformers/causal.py

+3 -1

@@ -27,13 +27,15 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

27 from vllm.model_executor.models.interfaces_base import VllmModelForTextGeneration

28 from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix

30 +from .base import Base

31 +

32 if TYPE_CHECKING:

33 import torch

35 from vllm.config import VllmConfig

36 -class CausalMixin(VllmModelForTextGeneration):

38 +class CausalMixin(VllmModelForTextGeneration, Base):

39 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

40 # Skip VllmModelForTextGeneration.__init__ and call the next class in MRO

41 super(VllmModelForTextGeneration, self).__init__(

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

+7 -5

@@ -5,6 +5,7 @@

5 import ast

6 import inspect

7 import textwrap

8 +from collections.abc import Callable

9 from dataclasses import dataclass

10 from functools import cache

11 from typing import TYPE_CHECKING, ClassVar

@@ -36,16 +37,16 @@ def _is_interface_lookup(node: ast.expr | None) -> bool:

39 @cache

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

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

40 +def interface_call(forward: Callable) -> ast.Call | None:

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

42 try:

42 - source = inspect.getsource(inspect.unwrap(cls.forward))

43 + source = inspect.getsource(inspect.unwrap(forward))

44 tree = ast.parse(textwrap.dedent(source))

45 except (AttributeError, OSError, SyntaxError, TypeError):

46 return None

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

48 - names = set()

49 + names: set[str] = set()

50 for node in ast.walk(tree):

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

52 continue

@@ -97,7 +98,7 @@ class AttentionFuser(BaseFuser):

98 def match(

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

100 ) -> "AttentionFuser | None":

100 - if (call := interface_call(type(module))) is None:

101 + if (call := interface_call(type(module).forward)) is None:

102 return None

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

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

@@ -120,6 +121,7 @@ class AttentionFuser(BaseFuser):

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

122 ) -> nn.Module:

123 if (sinks := self.sinks(module)) is not None:

124 + assert isinstance(self.s_aux_expr, ast.Attribute)

125 size = sinks.numel() // vllm_config.parallel_config.tensor_parallel_size

126 device = vllm_config.device_config.device

127 data = torch.empty(size, dtype=sinks.dtype, device=device)

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

+5 -2

@@ -130,11 +130,14 @@ class StackedFuser(RewriteFuser):

130 """A fuser that merges sibling projections into one stacked linear and

131 rewrites the forward to call it."""

132

133 - merged_name: ClassVar[str]

134 - """Attribute name of the merged module created by `update_attrs`."""

133 merged_cls_name: ClassVar[str]

134 """Name of the vLLM class the merged projection becomes (for logging)."""

135

136 + @property

137 + @abstractmethod

138 + def merged_name(self) -> str:

139 + """Attribute name of the merged module created by `update_attrs`."""

140 +

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

142 sources = " + ".join(shard for shard, _ in self.shards)

143 return (

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

+18 -18

@@ -24,6 +24,7 @@ from vllm.model_executor.models.transformers.fx_utils import (

24 upstream_linear,

25 )

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

27 + Style,

28 log_replacement,

29 replace_linear_class,

30 )

@@ -128,30 +129,26 @@ class MLAFuser(StackedFuser):

129 merged_name: ClassVar[str] = "fused_qkv_a_proj"

130 merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear"

131

131 - @property

132 - def has_q_lora(self) -> bool:

133 - return self.q_a_proj_name is not None

134 -

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

133 info_str = (

134 f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)"

135 )

139 - if self.has_q_lora:

136 + if self.q_a_proj_name is not None:

137 info_str += "; " + super().info(name).removeprefix("Fused: ")

138 return info_str

139

140 @property

141 def shards(self) -> list[tuple[str, ShardId]]:

142 """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection."""

146 - if self.has_q_lora:

147 - return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)]

148 - return []

143 + if self.q_a_proj_name is None:

144 + return []

145 + return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)]

146

147 @property

148 def packed_modules_mapping(self) -> dict[str, list[str]]:

152 - if self.has_q_lora:

153 - return super().packed_modules_mapping

154 - return {}

149 + if self.q_a_proj_name is None:

150 + return {}

151 + return super().packed_modules_mapping

152

153 @classmethod

154 def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None":

@@ -234,10 +231,10 @@ class MLAFuser(StackedFuser):

231 Bypass the KV expansion method so the compressed latent reaches the `vllm_mla`

232 attention interface unexpanded."""

233 funcdef, fn = recover_forward(type(module))

237 - if self.has_q_lora:

234 + if (q_a_proj_name := self.q_a_proj_name) is not None:

235 # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.

236 # The fused call is inserted at the top-level statement preceding both.

240 - names = [self.q_a_proj_name, self.kv_a_proj_name]

237 + names = [q_a_proj_name, self.kv_a_proj_name]

238 calls = self._unguarded_calls(funcdef, names)

239 if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]):

240 raise ValueError("down-projections read different inputs")

@@ -258,7 +255,7 @@ class MLAFuser(StackedFuser):

255 def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"):

256 quant_config = vllm_config.quant_config

257

261 - def replace_linear_by_name(name: str, style: str):

258 + def replace_linear_by_name(name: str, style: Style):

259 linear = module.get_submodule(name)

260 replacement = replace_linear_class(

261 linear, style, quant_config, prefix=maybe_prefix(prefix, name)

@@ -266,8 +263,10 @@ class MLAFuser(StackedFuser):

263 setattr(module, name, replacement)

264 log_replacement(maybe_prefix(prefix, name), linear, replacement)

265

269 - if self.has_q_lora:

270 - q_a = module.get_submodule(self.q_a_proj_name)

266 + if (q_a_proj_name := self.q_a_proj_name) is not None:

267 + # `match` sets the q-LoRA names together, or none of them

268 + assert self.q_b_proj_name is not None

269 + q_a = module.get_submodule(q_a_proj_name)

270 kv_a = module.get_submodule(self.kv_a_proj_name)

271 merged = MergedColumnParallelLinear(

272 input_size=q_a.in_features,

@@ -280,7 +279,7 @@ class MLAFuser(StackedFuser):

279 )

280 logger.debug(

281 "%s: %s, %s: %s -> %s: %s",

283 - self.q_a_proj_name,

282 + q_a_proj_name,

283 q_a,

284 self.kv_a_proj_name,

285 kv_a,

@@ -289,10 +288,11 @@ class MLAFuser(StackedFuser):

288 )

289 setattr(module, self.merged_name, merged)

290 # The rewritten forward calls the merged projection instead.

292 - delattr(module, self.q_a_proj_name)

291 + delattr(module, q_a_proj_name)

292 delattr(module, self.kv_a_proj_name)

293 replace_linear_by_name(self.q_b_proj_name, "colwise")

294 else:

295 + assert self.q_proj_name is not None

296 replace_linear_by_name(self.kv_a_proj_name, "replicate")

297 replace_linear_by_name(self.q_proj_name, "colwise")

298

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

+6 -6

@@ -216,7 +216,7 @@ class RMSNormFuser(BaseFuser):

216 finally:

217 for name, value in candidates.items():

218 setattr(module, name, value)

219 - if (name := markers.get(marked)) is not None:

219 + if marked is not None and (name := markers.get(marked)) is not None:

220 return name, None

221 logger.debug_once(

222 "%s does not hold its eps (%s) in an attribute. Every instance in this "

@@ -251,18 +251,18 @@ class RMSNormFuser(BaseFuser):

251 ) -> nn.Module:

252 """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""

253 weight = getattr(module, "weight", None)

254 - has_weight = weight is not None

255 - hidden_size = weight.size(0) if has_weight else 0

254 + hidden_size = weight.size(0) if weight is not None else 0

255 eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps

256 if not isinstance(eps, (int, float)):

257 # If eps was not detected, match torch behaviour.

259 - dtype = weight.dtype if has_weight else vllm_config.model_config.dtype

258 + config_dtype = vllm_config.model_config.dtype

259 + dtype = weight.dtype if weight is not None else config_dtype

260 eps = torch.finfo(dtype).eps

261 if self.zero_centered:

262 return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)

263 return TPAwareRMSNorm(

264 hidden_size=hidden_size,

265 eps=eps,

266 - has_weight=has_weight,

267 - dtype=weight.dtype if has_weight else None,

266 + has_weight=weight is not None,

267 + dtype=weight.dtype if weight is not None else None,

268 )

vllm/model_executor/models/transformers/fx_utils.py

+14 -11

@@ -15,17 +15,19 @@ import operator

15 import textwrap

16 from collections.abc import Callable

17 from itertools import chain

18 +from typing import Any

19 from unittest import mock

21 import torch

22 from torch import fx, nn

23 from torch.nn import functional as F

24 +from typing_extensions import TypeIs

26 from vllm.logger import init_logger

28 logger = init_logger(__name__)

28 -_UNKNOWN = object()

30 +_UNKNOWN: Any = object()

31 """Sentinel meta value for proxies whose concrete value could not be inferred.

32 Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`)."""

@@ -34,7 +36,7 @@ _MODULE_CALL = nn.Module.__call__

36 `call_module` nodes; meta execution must call modules for real."""

37 -def is_leaf_call(node: object) -> bool:

39 +def is_leaf_call(node: object) -> TypeIs[fx.Node]:

40 """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface)."""

41 return isinstance(node, fx.Node) and node.meta.get("leaf_call", False)

@@ -60,7 +62,7 @@ class _MetaProxy(fx.Proxy):

62 executing each op on the meta values, so PyTorch's meta kernels are the

63 single source of shape inference — no per-op rules."""

63 - meta: object = _UNKNOWN

65 + meta: Any = _UNKNOWN

67 def __len__(self) -> int:

68 if self.meta is not _UNKNOWN:

@@ -149,6 +151,7 @@ class _AllLeafTracer(fx.Tracer):

151 if unknown:

152 return _UNKNOWN

153 if kind == "call_function":

154 + assert callable(target)

155 return target(*meta_args, **meta_kwargs)

156 if kind == "call_method":

157 receiver, *rest = meta_args

@@ -576,7 +579,7 @@ def _aliasing_reads(node: ast.expr) -> set[str]:

579 would taint nearly every later name.

580 """

581 names: set[str] = set()

579 - stack = [node]

582 + stack: list[ast.AST] = [node]

583 while stack:

584 current = stack.pop()

585 if isinstance(current, ast.Attribute) and current.attr in _METADATA_ATTRS:

@@ -646,7 +649,7 @@ def aliasing_names(

649 """

650 names = set(seed)

651 assigns = [

649 - node

652 + (node, node.value)

653 for node in ast.walk(funcdef)

654 if isinstance(node, (ast.Assign, ast.AnnAssign))

655 and node.value is not None

@@ -655,10 +658,10 @@ def aliasing_names(

658 changed = True

659 while changed:

660 changed = False

658 - for node in assigns:

659 - if not _aliasing_reads(node.value) & names:

661 + for node, value in assigns:

662 + if not _aliasing_reads(value) & names:

663 continue

661 - targets = getattr(node, "targets", None) or [node.target]

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

665 stores = {

666 child.id

667 for target in targets

@@ -770,7 +773,7 @@ def peel(node: object) -> object:

773 return node

774

775

773 -def is_fn(node: object, target: Callable) -> bool:

776 +def is_fn(node: object, target: object) -> TypeIs[fx.Node]:

777 """Is node `<target>()`."""

778 return (

779 isinstance(node, fx.Node)

@@ -779,14 +782,14 @@ def is_fn(node: object, target: Callable) -> bool:

782 )

783

784

782 -def is_method(node: object, name: str) -> bool:

785 +def is_method(node: object, name: str) -> TypeIs[fx.Node]:

786 """Is node `.<name>()`."""

787 return (

788 isinstance(node, fx.Node) and node.op == "call_method" and node.target == name

789 )

790

791

789 -def is_op(node: object, name: str) -> bool:

792 +def is_op(node: object, name: str) -> TypeIs[fx.Node]:

793 """Is node `<mod>.<name>()` for torch, F, operator, or Tensor."""

794 return any(

795 is_fn(node, getattr(module, name, None)) for module in (torch, F, operator)

vllm/model_executor/models/transformers/legacy.py

+18 -13

@@ -20,35 +20,38 @@ from typing import TYPE_CHECKING

21 import torch

23 +from vllm.model_executor.models.utils import WeightsMapper

24 from vllm.sequence import IntermediateTensors

26 +from .base import Base

27 +

28 if TYPE_CHECKING:

29 from vllm.config import VllmConfig

29 -class LegacyMixin:

32 +class LegacyMixin(Base):

33 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

34 super().__init__(vllm_config=vllm_config, prefix=prefix)

33 - # Drop unsupported/unwanted output embeddings layers.

34 - self.hf_to_vllm_mapper.orig_to_new_prefix.update(

35 - {

36 + self.hf_to_vllm_mapper |= WeightsMapper(

37 + # Drop unsupported/unwanted output embeddings layers.

38 + orig_to_new_prefix={

39 "model.lm_head.": None,

40 "model.predictions.": None,

41 "model.qa_outputs.": None,

42 "model.embeddings_project.": None,

43 "model.discriminator_predictions.": None,

41 - }

44 + },

45 + orig_to_new_substr={

46 + # Some encoder models have the position_ids buffer in the checkpoint.

47 + # vLLM always passes position_ids as an argument, so drop the buffer.

48 + "position_ids": None,

49 + # Some encoder models have the bias of the final classifier layer in

50 + # the checkpoint. vLLM does not use this bias, so drop it.

51 + "score.bias": None,

52 + },

53 )

44 - # Some encoder models have the position_ids buffer in the checkpoint. vLLM will

45 - # always pass position_ids as an argument, so we drop the buffer if it exists.

46 - self.hf_to_vllm_mapper.orig_to_new_substr["position_ids"] = None

47 -

48 - # Some encoder models have the bias of the final classifier layer in the

49 - # checkpoint. vLLM does not use this bias, so we drop it if it exists.

50 - self.hf_to_vllm_mapper.orig_to_new_substr["score.bias"] = None

51 -

55 # roberta-like models an extra padding in positions.

56 # FIXME(Isotr0py): This is quite hacky for roberta edge case,

57 # we should find a better way to handle this.

@@ -61,6 +64,7 @@ class LegacyMixin:

64 positions: torch.Tensor,

65 intermediate_tensors: IntermediateTensors | None = None,

66 inputs_embeds: torch.Tensor | None = None,

67 + **kwargs,

68 ) -> torch.Tensor | IntermediateTensors:

69 if self.is_roberta:

70 # RoBERTa positions start at padding_idx + 1.

@@ -72,4 +76,5 @@ class LegacyMixin:

76 positions=positions,

77 intermediate_tensors=intermediate_tensors,

78 inputs_embeds=inputs_embeds,

79 + **kwargs,

80 )

vllm/model_executor/models/transformers/moe.py

+31 -11

@@ -41,6 +41,7 @@ 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

44 +from .base import Base

45 from .utils import log_replacement, maybe_per_layer

47 if TYPE_CHECKING:

@@ -66,14 +67,15 @@ class TransformersMoERunner(MoERunner):

67 self._moe_state = moe_state

68 self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel

69 - def forward(

70 + def forward( # type: ignore[override]

71 self,

72 hidden_states: torch.Tensor,

73 topk_ids: torch.Tensor,

74 topk_weights: torch.Tensor,

75 **kwargs: Any,

76 ) -> torch.Tensor:

76 - """In Transformers `experts.forward` will have this signature.

77 + """In Transformers `experts.forward` will have this signature, which is why

78 + it deliberately does not match `MoERunner.forward`.

80 We discard any extra kwargs because we cannot use them here."""

81 # Note: we need to forward through a custom op so the topk_ids

@@ -126,10 +128,18 @@ direct_register_custom_op(

128

129 class TransformersRoutedExperts(RoutedExperts):

130 def get_expert_mapping(

129 - self, include_fused: bool = False

131 + self,

132 + ckpt_gate_proj_name: str | None = None,

133 + ckpt_down_proj_name: str | None = None,

134 + ckpt_up_proj_name: str | None = None,

135 + include_fused: bool = False,

136 ) -> list[tuple[str, str, int, str]]:

131 - common_names = ("gate_proj", "down_proj", "up_proj")

132 - common_map = super().get_expert_mapping(*common_names, include_fused)

137 + common_map = super().get_expert_mapping(

138 + ckpt_gate_proj_name or "gate_proj",

139 + ckpt_down_proj_name or "down_proj",

140 + ckpt_up_proj_name or "up_proj",

141 + include_fused,

142 + )

143 mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused)

144 if not include_fused:

145 return common_map + mixtral_map

@@ -138,7 +148,13 @@ class TransformersRoutedExperts(RoutedExperts):

148 return common_fused + mixtral_fused + common_unfused + mixtral_unfused

149

150

141 -class MoEMixin(MixtureOfExperts):

151 +class MoEMixin(MixtureOfExperts, Base):

152 + mlp_layers: list[nn.Module]

153 + """MoE blocks whose experts were replaced, for the `MixtureOfExperts` methods."""

154 +

155 + # Redeclared because `recursive_replace` assigns it below the first read

156 + num_local_physical_experts: int

157 +

158 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

159 self.check_version("5.0.0", "MoE models support")

160 # Skip MixtureOfExperts.__init__ and call the next class in MRO

@@ -197,7 +213,8 @@ class MoEMixin(MixtureOfExperts):

213

214 # Dtype the router computes in, if it is not the activation dtype.

215 config_router_dtype = getattr(text_config, "moe_router_dtype", None)

200 - config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)

216 + if config_router_dtype is not None:

217 + config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)

218

219 # Grouped topk routing kwargs

220 num_expert_group = getattr(text_config, "n_group", None)

@@ -222,7 +239,7 @@ class MoEMixin(MixtureOfExperts):

239 # MixtureOfExperts mixin settings

240 ep_size = get_ep_group().world_size

241

225 - self.mlp_layers = [] # Used for MixtureOfExperts methods

242 + self.mlp_layers = []

243 self.moe_layers = []

244 self.num_expert_groups = 1 if num_expert_group is None else num_expert_group

245 self.num_logical_experts = num_experts

@@ -321,9 +338,10 @@ class MoEMixin(MixtureOfExperts):

338 if shared_experts is not None:

339 hf_shared = shared_experts.shared_experts

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

341 + if glu_fuser is not None:

342 + down_name = glu_fuser.down_name

343 + if down_name is not None:

344 + shared_down_projs.append((hf_shared, down_name))

345 # Prefer config, otherwise read it from fuser.

346 router_dtype = config_router_dtype or fuser.router_dtype

347 gate = fuser.gate(moe_block, prefix, router_dtype)

@@ -374,7 +392,9 @@ class MoEMixin(MixtureOfExperts):

392 # Handle all gather in expert parallel

393 if topk_ids.size(0) != hidden_states.size(0):

394 dp_metadata = get_forward_context().dp_metadata

395 + assert dp_metadata is not None

396 sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

397 + assert sizes is not None

398 is_sp = moe_state.is_sequence_parallel

399 group = get_ep_group() if is_sp else get_dp_group()

400 assert sizes[group.rank_in_group] == topk_ids.shape[0]

vllm/model_executor/models/transformers/multimodal.py

+36 -18

@@ -17,17 +17,24 @@

17 """Transformers modeling backend mixin for multi-modal models."""

19 from collections import defaultdict

20 -from collections.abc import Mapping, Sequence

20 +from collections.abc import Iterable, Mapping, Sequence

21 from contextlib import ExitStack, contextmanager

22 from typing import TYPE_CHECKING, Any

24 import torch

25 import transformers

26 from packaging.version import Version

27 +from transformers.utils.generic import ModelOutput

29 from vllm.compilation.decorators import should_torch_compile_mm_encoder

30 +from vllm.config.multimodal import AudioDummyOptions

31 from vllm.config.utils import getattr_iter

30 -from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input

32 +from vllm.inputs import (

33 + MultiModalDataBuiltins,

34 + MultiModalDataDict,

35 + MultiModalInput,

36 + mm_input,

37 +)

38 from vllm.logger import init_logger

39 from vllm.model_executor.models.interfaces import (

40 MultiModalEmbeddings,

@@ -62,6 +69,8 @@ from vllm.sequence import IntermediateTensors

69 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

70 from vllm.utils.torch_utils import async_tensor_h2d

72 +from .base import Base

73 +

74 if TYPE_CHECKING:

75 from transformers import BatchFeature, PreTrainedModel

@@ -148,7 +157,7 @@ class MultiModalProcessingInfo(BaseProcessingInfo):

157 def get_max_image_tokens(self) -> int:

158 width, height = self.get_image_size_with_most_features()

159 processor = self.get_hf_processor()

151 - multimodal_config = self.ctx.model_config.multimodal_config

160 + multimodal_config = self.ctx.model_config.get_multimodal_config()

161 mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}

162 mm_tokens = processor._get_num_multimodal_tokens(

163 image_sizes=([height, width],), **mm_processor_kwargs

@@ -188,7 +197,7 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn

197 mm_counts: Mapping[str, int],

198 mm_options: Mapping[str, "BaseDummyOptions"],

199 ) -> MultiModalDataDict:

191 - data: MultiModalDataDict = {}

200 + data = MultiModalDataBuiltins()

201 if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)):

202 sampling_rate = self.info._get_audio_sampling_rate()

203 sub = self.info._get_audio_processor()

@@ -196,10 +205,14 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn

205 if chunk_length is None:

206 chunk_length = 30

207 audio_len = int(chunk_length * sampling_rate)

208 + audio_overrides = mm_options.get("audio")

209 + assert audio_overrides is None or isinstance(

210 + audio_overrides, AudioDummyOptions

211 + )

212 data["audio"] = self._get_dummy_audios(

213 length=audio_len,

214 num_audios=num_audios,

202 - overrides=mm_options.get("audio"),

215 + overrides=audio_overrides,

216 )

217 if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)):

218 target_width, target_height = self.info.get_image_size_with_most_features()

@@ -359,7 +372,9 @@ class _MultiModalProcessorBase(BaseMultiModalProcessor[MultiModalProcessingInfo]

372 itself is to process it by itself.

373 """

374 audios = mm_data.get("audio")

362 - if not audios or len({len(audio) for audio in audios}) == 1:

375 + if not isinstance(audios, Iterable) or not audios:

376 + return

377 + if len({len(audio) for audio in audios}) == 1:

378 return

379

380 alone = [

@@ -444,7 +459,8 @@ class LegacyMultiModalProcessor(_MultiModalProcessorBase):

459 token_id = getattr(processor, names[0], getattr_iter(config, names))

460 if token_id is None:

461 token = getattr(processor, f"{modality}_token", None)

447 - token_id = info.get_tokenizer().get_vocab().get(token)

462 + if token is not None:

463 + token_id = info.get_tokenizer().get_vocab().get(token)

464 if token_id is None:

465 raise ValueError(

466 f"Cannot find {modality}_token_id on processor or model config"

@@ -722,13 +738,15 @@ class OffsetsMultiModalProcessor(_MultiModalProcessorBase):

738 for modality, items in out_mm_kwargs.items():

739 # Popped so they are neither cached nor sent to the model; the updates

740 # they produce are cached alongside the item instead

725 - replacements = [

726 - PromptUpdateDetails.select_token_id(

727 - (ids := item.pop(f"{modality}_replacement_ids").data).tolist(),

728 - _get_embed_token_id(ids),

741 + replacements = []

742 + for item in items:

743 + ids = item.pop(f"{modality}_replacement_ids").data

744 + assert isinstance(ids, torch.Tensor)

745 + replacements.append(

746 + PromptUpdateDetails.select_token_id(

747 + ids.tolist(), _get_embed_token_id(ids)

748 + )

749 )

730 - for item in items

731 - ]

750 token = getattr(hf_processor, f"{modality}_token")

751 updates.append(

752 PromptReplacement(

@@ -777,7 +795,7 @@ class OffsetsMultiModalProcessor(_MultiModalProcessorBase):

795 ) -> list[int] | None:

796 """Ask the HF processor how many rows of image data each image produces."""

797 images = mm_data.get("images")

780 - if not images:

798 + if not isinstance(images, Iterable) or not images:

799 return None

800 try:

801 sizes = [(image.height, image.width) for image in images]

@@ -893,7 +911,7 @@ MultiModalProcessor = (

911 )

912

913

896 -class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):

914 +class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base):

915 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

916 # Skip SupportsMRoPE.__init__ and call the next class in MRO

917 super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)

@@ -925,7 +943,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):

943 if model_config.skip_tokenizer_init:

944 # Determining the supported modalities needs the HF processor, which in

945 # turn needs a tokenizer

928 - mm_config = model_config.multimodal_config

946 + mm_config = model_config.get_multimodal_config()

947 if mm_config.mm_encoder_only or any(

948 mm_config.get_limit_per_prompt(modality) == 0

949 for modality in encoder_classes

@@ -1029,7 +1047,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):

1047 # Exclude MultiModalMixin itself

1048 bases = [b for b in bases if b is not MultiModalMixin]

1049

1032 - class LanguageModel(*bases):

1050 + class LanguageModel(*bases): # type: ignore[misc]

1051 def __init__(self, multimodal_model):

1052 # Don't call super().__init__() to avoid re-initialization

1053 self.__dict__.update(multimodal_model.__dict__)

@@ -1203,7 +1221,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):

1221 # from `self.get_image_features`

1222 if isinstance(features, tuple):

1223 return features[0]

1206 - if isinstance(features, dict):

1224 + if isinstance(features, ModelOutput):

1225 return features.pooler_output

1226 return features

1227

vllm/model_executor/models/transformers/pooling.py

+24 -13

@@ -19,6 +19,7 @@

19 from typing import TYPE_CHECKING

21 import torch

22 +import torch.nn as nn

23 from transformers import AutoModelForSequenceClassification

25 from vllm.config.utils import getattr_iter

@@ -26,11 +27,25 @@ from vllm.model_executor.layers.pooler import DispatchPooler

27 from vllm.model_executor.models.interfaces import SupportsCrossEncoding

28 from vllm.model_executor.models.interfaces_base import VllmModelForPooling

30 +from .base import Base

31 +

32 if TYPE_CHECKING:

33 from vllm.config import VllmConfig

33 -class EmbeddingMixin(VllmModelForPooling):

36 +class ClassifierWithReshape(nn.Module):

37 + """Token extraction has already been applied in `pooler.pooling`.

38 +

39 + Add dim to match expected input shape of `classifier.forward`.

40 + """

41 +

42 + def forward(self, *args, **kwargs):

43 + if len(args) > 0:

44 + args = (args[0].unsqueeze(1), *args[1:])

45 + return super().forward(*args, **kwargs)

46 +

47 +

48 +class EmbeddingMixin(VllmModelForPooling, Base):

49 default_seq_pooling_type = "CLS"

51 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

@@ -45,7 +60,7 @@ class EmbeddingMixin(VllmModelForPooling):

60 self.pooler = DispatchPooler.for_embedding(pooler_config)

48 -class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):

63 +class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling, Base):

64 default_seq_pooling_type = "CLS"

66 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

@@ -83,17 +98,13 @@ class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):

98 )

99 self.init_parameters(self.classifier, dtype=self.model_config.head_dtype)

100

86 - class ClassifierWithReshape(self.classifier.__class__):

87 - """Token extraction has already been applied in `pooler.pooling`.

88 - Add dim to match expected input shape of `classifier.forward`.

89 - """

90 -

91 - def forward(self, *args, **kwargs):

92 - if len(args) > 0:

93 - args = (args[0].unsqueeze(1), *args[1:])

94 - return super().forward(*args, **kwargs)

95 -

96 - self.classifier.__class__ = ClassifierWithReshape

101 + # Order `ClassifierWithReshape` ahead of the classifier's own class so that

102 + # its `super().forward(...)` reaches the original implementation.

103 + self.classifier.__class__ = type(

104 + "ClassifierWithReshape",

105 + (ClassifierWithReshape, type(self.classifier)),

106 + {},

107 + )

108

109 self.pooler = DispatchPooler.for_seq_cls(

110 pooler_config,

vllm/model_executor/models/transformers/utils.py

+12 -9

@@ -18,11 +18,10 @@

19 from collections.abc import Callable, Iterator

20 from contextlib import contextmanager

21 -from functools import lru_cache

21 from itertools import chain

22 from operator import attrgetter

23 from pathlib import Path

25 -from typing import TYPE_CHECKING, Literal, TypeVar

24 +from typing import TYPE_CHECKING, Literal, TypeAlias, TypeVar

26 import torch

27 from torch import nn

@@ -70,7 +69,7 @@ def init_on_device_without_buffers(device: torch.device):

69 module._parameters[name].to(device), **kwargs

70 )

73 - tensor_constructors_to_patch = {}

72 + tensor_constructors_to_patch: dict[str, Callable] = {}

74 def patch_tensor_constructor(fn):

75 def wrapper(*args, **kwargs):

@@ -147,8 +146,8 @@ def replace_linear_class(

146 )

147

148

150 -TorchConv = nn.Conv2d | nn.Conv3d

151 -VllmConv = Conv2dLayer | Conv3dLayer

149 +TorchConv: TypeAlias = nn.Conv2d | nn.Conv3d

150 +VllmConv: TypeAlias = Conv2dLayer | Conv3dLayer

151

152

153 def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:

@@ -212,7 +211,9 @@ class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbeddin

211 `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding."""

212

213

215 -@lru_cache

214 +_rebased_embedding_classes: dict[type[nn.Embedding], type[VocabParallelEmbedding]] = {}

215 +

216 +

217 def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbedding]:

218 """Subclass `cls` so that `VocabParallelEmbedding` supersedes its `nn.Embedding`.

219

@@ -224,7 +225,10 @@ def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbe

225 The new class, to assign to `__class__` of an instance of `cls`.

226

227 """

227 - return type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {})

228 + if cls not in _rebased_embedding_classes:

229 + rebased_cls = type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {})

230 + _rebased_embedding_classes[cls] = rebased_cls

231 + return _rebased_embedding_classes[cls]

232

233

234 def replace_embedding_class(

@@ -278,10 +282,9 @@ def recursive_replace_linear(

282 qual_name = maybe_prefix(prefix, child_name)

283 # Replace modules as needed

284 if isinstance(child_module, nn.Linear):

281 - style = "replicate"

285 new_module = replace_linear_class(

286 child_module,

284 - style,

287 + "replicate",

288 quant_config,

289 prefix=qual_name,

290 )