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

[Mypy Fix] Mypy fix for "vllm/model_executor/models/[aA][bB]" (#48977)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Aug 6, 2026 at 15:30 UTC adc3e03517d2e7333a3bb2083bb4d394a2986876

23 files changed +214 -167

tools/pre_commit/mypy.py

-2

@@ -113,8 +113,6 @@ SEPARATE_GROUPS = [

113

114 # TODO(woosuk): Include the code from Megatron and HuggingFace.

115 EXCLUDE = [

116 - r"vllm/model_executor/models/[aA]",

117 - r"vllm/model_executor/models/[bB]",

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

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

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

vllm/model_executor/layers/mamba/mamba_utils.py

+1 -1

@@ -144,7 +144,7 @@ class MambaStateShapeCalculator:

144 num_heads: int,

145 tp_size: int,

146 head_dim: int,

147 - ) -> tuple[tuple[int, int, int], ...]:

147 + ) -> tuple[tuple[int, int, int]]:

148 state_shape = (num_heads // tp_size, head_dim, head_dim)

149 return (state_shape,)

150

vllm/model_executor/models/AXK1.py

+28 -10

@@ -110,8 +110,12 @@ class AXK1MoE(nn.Module):

110 self.ep_group = get_ep_group().device_group

111 self.ep_rank = get_ep_group().rank_in_group

112 self.ep_size = self.ep_group.size()

113 + assert config.n_routed_experts is not None

114 + assert config.num_experts_per_tok is not None

115 + assert config.scoring_func is not None

116 + assert config.hidden_act is not None

117 self.n_routed_experts: int = config.n_routed_experts

114 - self.n_shared_experts: int = config.n_shared_experts

118 + self.n_shared_experts: int | None = config.n_shared_experts

119

120 self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe

121

@@ -245,7 +249,7 @@ class AXK1Attention(nn.Module):

249 qk_nope_head_dim: int,

250 qk_rope_head_dim: int,

251 v_head_dim: int,

248 - q_lora_rank: int,

252 + q_lora_rank: int | None,

253 kv_lora_rank: int,

254 max_position_embeddings: int = 8192,

255 cache_config: CacheConfig | None = None,

@@ -282,7 +286,7 @@ class AXK1Attention(nn.Module):

286 )

287 self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)

288 self.q_b_proj = ColumnParallelLinear(

285 - q_lora_rank,

289 + self.q_lora_rank,

290 self.num_heads * self.qk_head_dim,

291 bias=False,

292 quant_config=quant_config,

@@ -320,6 +324,7 @@ class AXK1Attention(nn.Module):

324 quant_config=quant_config,

325 prefix=f"{prefix}.o_proj",

326 )

327 + assert config.rope_parameters is not None

328 if config.rope_parameters["rope_type"] != "default":

329 config.rope_parameters["rope_type"] = (

330 "deepseek_yarn"

@@ -492,6 +497,7 @@ class AXK1MLAAttention(nn.Module):

497 prefix=f"{prefix}.o_proj",

498 )

499

500 + assert config.rope_parameters is not None

501 if config.rope_parameters["rope_type"] != "default":

502 config.rope_parameters["rope_type"] = (

503 "deepseek_yarn"

@@ -573,6 +579,13 @@ class AXK1DecoderLayer(nn.Module):

579 parallel_config = vllm_config.parallel_config

580 self.config = config

581

582 + assert config.max_position_embeddings is not None

583 + assert config.hidden_act is not None

584 + assert config.routed_scaling_factor is not None

585 + assert config.qk_nope_head_dim is not None

586 + assert config.qk_rope_head_dim is not None

587 + assert config.v_head_dim is not None

588 + assert config.kv_lora_rank is not None

589 self.hidden_size = config.hidden_size

590 max_position_embeddings = config.max_position_embeddings

591 # DecoderLayers are created with `make_layers` which passes the prefix

@@ -588,6 +601,7 @@ class AXK1DecoderLayer(nn.Module):

601 use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim))

602 self.use_mha = use_mha

603

604 + attn_cls: type[nn.Module]

605 if use_mha:

606 attn_cls = DeepseekAttention

607 elif model_config.use_mla:

@@ -635,6 +649,7 @@ class AXK1DecoderLayer(nn.Module):

649 self.routed_scaling_factor = config.routed_scaling_factor

650

651 def _is_layer_sparse(self) -> bool:

652 + assert self.config.moe_layer_freq is not None

653 return (

654 self.config.n_routed_experts is not None

655 and self.layer_idx >= self.config.first_k_dense_replace

@@ -789,7 +804,7 @@ class AXK1Model(nn.Module):

804 rocm_aiter_moe_shared_expert_enabled = (

805 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()

806 )

792 - stacked_params_mapping = [

807 + stacked_params_mapping: list[tuple[str, str, int | str]] = [

808 # (param_name, shard_name, shard_id)

809 ("gate_up_proj", "gate_proj", 0),

810 ("gate_up_proj", "up_proj", 1),

@@ -810,6 +825,7 @@ class AXK1Model(nn.Module):

825

826 # Params for weights, fp8 weight scales, fp8 activation scales

827 # (param_name, weight_name, expert_id, shard_id)

828 + assert self.config.n_routed_experts is not None

829 expert_params_mapping = fused_moe_make_expert_params_mapping(

830 self,

831 ckpt_gate_proj_name="gate_proj",

@@ -817,7 +833,7 @@ class AXK1Model(nn.Module):

833 ckpt_up_proj_name="up_proj",

834 num_experts=self.config.n_routed_experts

835 + (

820 - self.config.n_shared_experts

836 + (self.config.n_shared_experts or 0)

837 if rocm_aiter_moe_shared_expert_enabled

838 else 0

839 ),

@@ -926,7 +942,7 @@ class AXK1Model(nn.Module):

942 # param and delegate to its expert-aware weight_loader

943 # with expert_id.

944 for mapping in expert_params_mapping:

929 - param_name, weight_name, expert_id, shard_id = mapping

945 + param_name, weight_name, expert_id, expert_shard_id = mapping

946 if weight_name not in chunk_name:

947 continue

948

@@ -952,7 +968,7 @@ class AXK1Model(nn.Module):

968 param,

969 weight_to_load,

970 name_mapped,

955 - shard_id=shard_id,

971 + shard_id=expert_shard_id,

972 expert_id=expert_id,

973 return_success=True,

974 )

@@ -974,9 +990,10 @@ class AXK1Model(nn.Module):

990 continue

991

992 # Remapping the name of FP8 kv-scale.

977 - name = maybe_remap_kv_scale_name(name, params_dict)

978 - if name is None:

993 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

994 + if remapped_name is None:

995 continue

996 + name = remapped_name

997

998 if is_pp_missing_parameter(name, self):

999 continue

@@ -1014,7 +1031,7 @@ class AXK1MixtureOfExperts(MixtureOfExperts):

1031 self.num_physical_experts = example_moe.n_physical_experts

1032 self.num_local_physical_experts = example_moe.n_local_physical_experts

1033 self.num_routed_experts = example_moe.n_routed_experts

1017 - self.num_shared_experts = example_moe.n_shared_experts

1034 + self.num_shared_experts = example_moe.n_shared_experts or 0

1035 self.num_redundant_experts = example_moe.n_redundant_experts

1036

1037 def update_physical_experts_metadata(

@@ -1132,6 +1149,7 @@ class AXK1ForCausalLM(

1149 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

1150 # Params for weights, fp8 weight scales, fp8 activation scales

1151 # (param_name, weight_name, expert_id, shard_id)

1152 + assert self.config.n_routed_experts is not None

1153 return fused_moe_make_expert_params_mapping(

1154 self,

1155 ckpt_gate_proj_name="gate_proj",

vllm/model_executor/models/adapters.py

+34 -34

@@ -4,7 +4,7 @@

4 import itertools

5 from collections.abc import Iterable

6 from contextlib import contextmanager

7 -from typing import TYPE_CHECKING, Any, TypeVar, cast

7 +from typing import TYPE_CHECKING, Any, TypeVar

9 import torch

10 import torch.nn as nn

@@ -25,7 +25,7 @@ if TYPE_CHECKING:

25 from vllm.config import ModelConfig, VllmConfig

26 from vllm.model_executor.layers.pooler import Pooler

28 -_T = TypeVar("_T", bound=type[nn.Module])

28 +_T = TypeVar("_T", bound=nn.Module)

30 logger = init_logger(__name__)

@@ -45,7 +45,7 @@ def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None:

45 )

47 if dense_modules is None:

48 - return

48 + return None

50 try:

51 layers = []

@@ -126,14 +126,14 @@ def _get_pooling_model_name(orig_model_name: str, pooling_suffix: str) -> str:

126 return model_name + pooling_suffix

127

128

129 -def _create_pooling_model_cls(orig_cls: _T) -> _T:

129 +def _create_pooling_model_cls(orig_cls: type[_T]) -> type[_T]:

130 # Lazy import

131 from vllm.model_executor.layers.logits_processor import LogitsProcessor

132 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

133

134 from .utils import AutoWeightsLoader, StageMissingLayer, no_init_weights

135

136 - class ModelForPooling(orig_cls, VllmModelForPooling):

136 + class ModelForPooling(orig_cls, VllmModelForPooling): # type: ignore[valid-type,misc]

137 is_pooling_model = True

138

139 def __init__(

@@ -148,16 +148,19 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:

148 lambda mod: StageMissingLayer("output", mod),

149 targets=(LogitsProcessor, ParallelLMHead),

150 ):

151 - super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)

151 + super().__init__( # type: ignore[safe-super]

152 + vllm_config=vllm_config, prefix=prefix, **kwargs

153 + )

154

155 # Used by SEQ_CLS_LOAD_METHODS

156 self.vllm_config = vllm_config

157

158 # If the model already defines a pooler instance, don't overwrite it

159 pooler = getattr(self, "pooler", None)

158 - if not pooler and supports_multimodal(self):

160 + multimodal_model: object = self

161 + if not pooler and supports_multimodal(multimodal_model):

162 # Try to get the pooler from the LM backbone

160 - language_model = self.get_language_model()

163 + language_model = multimodal_model.get_language_model()

164 if hasattr(language_model, "pooler"):

165 pooler = language_model.pooler

166

@@ -173,7 +176,9 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:

176 ) -> "Pooler":

177 raise NotImplementedError

178

176 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

179 + def _load_pooling_model_weights(

180 + self, weights: Iterable[tuple[str, torch.Tensor]]

181 + ):

182 params_dict = dict(self.named_parameters())

183 # Use model-owned remapping only for the probe membership check

184 # (e.g. Gemma4 model.language_model.* -> language_model.model.*)

@@ -236,10 +241,13 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:

241 load_weights = getattr(super(), "load_weights", default_load_weights)

242 return load_weights(mapped_weights)

243

239 - return ModelForPooling # type: ignore

244 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

245 + return self._load_pooling_model_weights(weights)

246 +

247 + return ModelForPooling

248

249

242 -def as_embedding_model(cls: _T) -> _T:

250 +def as_embedding_model(cls: type[_T]) -> type[_T]:

251 """

252 Subclass an existing vLLM model to support embeddings.

253

@@ -257,7 +265,7 @@ def as_embedding_model(cls: _T) -> _T:

265 # Lazy import

266 from vllm.model_executor.layers.pooler import DispatchPooler

267

260 - class ModelForEmbedding(_create_pooling_model_cls(cls)):

268 + class ModelForEmbedding(_create_pooling_model_cls(cls)): # type: ignore[misc]

269 def _init_pooler(

270 self,

271 vllm_config: "VllmConfig",

@@ -270,7 +278,7 @@ def as_embedding_model(cls: _T) -> _T:

278

279 ModelForEmbedding.__name__ = _get_pooling_model_name(cls.__name__, "ForEmbedding")

280

273 - return ModelForEmbedding # type: ignore

281 + return ModelForEmbedding

282

283

284 def _resolve_num_labels(hf_config: Any, text_config: Any) -> int:

@@ -296,7 +304,7 @@ def _resolve_num_labels(hf_config: Any, text_config: Any) -> int:

304 return text_config.num_labels

305

306

299 -def as_seq_cls_model(cls: _T) -> _T:

307 +def as_seq_cls_model(cls: type[_T]) -> type[_T]:

308 """

309 Subclass an existing vLLM model to support classify and score tasks.

310

@@ -320,7 +328,8 @@ def as_seq_cls_model(cls: _T) -> _T:

328 from .utils import maybe_prefix

329

330 class ModelForSequenceClassification(

323 - _create_pooling_model_cls(cls), SupportsCrossEncoding

331 + _create_pooling_model_cls(cls), # type: ignore[misc]

332 + SupportsCrossEncoding,

333 ):

334 def _init_pooler(

335 self,

@@ -401,7 +410,7 @@ def as_seq_cls_model(cls: _T) -> _T:

410 cls.__name__, "ForSequenceClassification"

411 )

412

404 - return ModelForSequenceClassification # type: ignore

413 + return ModelForSequenceClassification

414

415

416 class SequenceClassificationConfig(VerifyAndUpdateConfig):

@@ -435,14 +444,15 @@ class SequenceClassificationConfig(VerifyAndUpdateConfig):

444 text_config.use_sep_token = use_sep_token

445

446

438 -def _get_language_model_for_seq_cls(model) -> nn.Module:

447 +def _get_language_model_for_seq_cls(model: nn.Module) -> nn.Module:

448 """

449 Get the language model component for sequence classification conversion.

450 For VLMs, returns the inner language model. For standard LLMs, returns model itself.

451 """

443 - if supports_multimodal(model):

452 + multimodal_model: object = model

453 + if supports_multimodal(multimodal_model):

454 try:

445 - lm = model.get_language_model()

455 + lm = multimodal_model.get_language_model()

456 if lm is not model:

457 return lm

458 except Exception:

@@ -516,12 +526,11 @@ def load_weights_using_from_2_way_softmax(

526 hf_config = model.config

527 text_config = hf_config.get_text_config()

528

519 - tokens = getattr(

529 + tokens: list[str] = getattr(

530 hf_config,

531 "classifier_from_token",

532 getattr(text_config, "classifier_from_token", []),

533 )

524 - tokens = cast(list[int], tokens)

534 assert len(tokens) == 2

535

536 language_model = _get_language_model_for_seq_cls(model)

@@ -545,12 +554,7 @@ def load_weights_using_from_2_way_softmax(

554 language_model.lm_head = language_model.lm_head.tie_weights(embed_tokens)

555

556 with _disable_seq_cls_loading_on_inner_model(language_model, is_vlm):

548 - # ModelForPooling is dynamically defined inside the _create_pooling_model_cls

549 - # function, so we need use this hacky method to obtain it.

550 - pooling_model_cls = next(

551 - x for x in type(model).__mro__ if x.__name__ == "ModelForPooling"

552 - )

553 - loaded_weights = pooling_model_cls.load_weights(model, weights)

557 + loaded_weights = model._load_pooling_model_weights(weights)

558

559 from vllm.tokenizers import get_tokenizer

560

@@ -594,8 +598,7 @@ def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Te

598 model_config = model.vllm_config.model_config

599 text_config = model.config.get_text_config()

600

597 - tokens = getattr(text_config, "classifier_from_token", [])

598 - tokens = cast(list[int], tokens)

601 + tokens: list[str] = getattr(text_config, "classifier_from_token", [])

602 assert len(tokens) > 0

603

604 language_model = _get_language_model_for_seq_cls(model)

@@ -619,11 +622,8 @@ def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Te

622 language_model.lm_head = language_model.lm_head.tie_weights(embed_tokens)

623

624 with _disable_seq_cls_loading_on_inner_model(language_model, is_vlm):

622 - pooling_model_cls = next(

623 - x for x in type(model).__mro__ if x.__name__ == "ModelForPooling"

624 - )

625 - # Skip ModelForSequenceClassification in MRO to avoid infinite recursion

626 - loaded_weights = pooling_model_cls.load_weights(model, weights)

625 + # Bypass ModelForSequenceClassification to avoid infinite recursion

626 + loaded_weights = model._load_pooling_model_weights(weights)

627

628 from vllm.tokenizers import get_tokenizer

629

vllm/model_executor/models/afmoe.py

+1 -6

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

47 PPMissingLayer,

48 WeightsMapper,

49 extract_layer_index,

50 - make_empty_intermediate_tensors_factory,

50 make_layers,

51 maybe_prefix,

52 )

@@ -232,7 +231,7 @@ class AfmoeAttention(nn.Module):

231

232 # Only create rotary embeddings for local attention

233 if self.is_local_attention:

235 - self.rotary_emb = get_rope(

234 + self.rotary_emb: nn.Module | None = get_rope(

235 self.head_dim,

236 max_position=max_position_embeddings,

237 rope_parameters=config.rope_parameters,

@@ -409,10 +408,6 @@ class AfmoeModel(nn.Module, EagleModelMixin):

408 else:

409 self.norm = PPMissingLayer()

410

412 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

413 - ["hidden_states", "residual"], config.hidden_size

414 - )

415 -

411 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

412 return self.embed_tokens(input_ids)

413

vllm/model_executor/models/aimv2.py

+3 -1

@@ -180,7 +180,9 @@ class AIMv2Transformer(nn.Module):

180 ]

181 )

182 if require_post_norm:

183 - self.post_trunk_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

183 + self.post_trunk_norm: RMSNorm | None = RMSNorm(

184 + config.hidden_size, eps=config.rms_norm_eps

185 + )

186 else:

187 self.post_trunk_norm = None

188

vllm/model_executor/models/arctic.py

+4 -4

@@ -501,7 +501,7 @@ class ArcticModel(nn.Module):

501 weight_loader(param, loaded_weight, shard_id)

502 break

503 else:

504 - for param_name, weight_name, shard_id in mlp_params_mapping:

504 + for param_name, weight_name, mlp_shard_id in mlp_params_mapping:

505 if weight_name not in name:

506 continue

507 name = name.replace(weight_name, param_name)

@@ -509,10 +509,10 @@ class ArcticModel(nn.Module):

509 continue

510 param = params_dict[name]

511 weight_loader = param.weight_loader

512 - weight_loader(param, loaded_weight, shard_id)

512 + weight_loader(param, loaded_weight, mlp_shard_id)

513 break

514 else:

515 - for param_name, weight_name, shard_id in expert_params_mapping:

515 + for param_name, weight_name, expert_id in expert_params_mapping:

516 if weight_name not in name:

517 continue

518 name = name.replace(weight_name, param_name)

@@ -521,7 +521,7 @@ class ArcticModel(nn.Module):

521 param = params_dict[name]

522 weight_loader = param.weight_loader

523 weight_loader(

524 - param, loaded_weight, weight_name, expert_id=shard_id

524 + param, loaded_weight, weight_name, expert_id=expert_id

525 )

526 break

527 else:

vllm/model_executor/models/aria.py

+4 -2

@@ -10,7 +10,7 @@ from transformers.models.aria.modeling_aria import AriaCrossAttention

10 from transformers.models.aria.processing_aria import AriaProcessor

12 from vllm.config import VllmConfig

13 -from vllm.config.multimodal import BaseDummyOptions

13 +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions

14 from vllm.inputs import MultiModalDataDict

15 from vllm.model_executor.layers.activation import get_act_fn

16 from vllm.model_executor.layers.fused_moe import (

@@ -319,7 +319,8 @@ class AriaDummyInputsBuilder(BaseDummyInputsBuilder[AriaProcessingInfo]):

319 num_images = mm_counts.get("image", 0)

320

321 processor = self.info.get_hf_processor()

322 - image_token: str = processor.tokenizer.image_token # type: ignore

322 + image_token = getattr(processor.tokenizer, "image_token", None)

323 + assert isinstance(image_token, str)

324

325 return image_token * num_images

326

@@ -335,6 +336,7 @@ class AriaDummyInputsBuilder(BaseDummyInputsBuilder[AriaProcessingInfo]):

336 num_images = mm_counts.get("image", 0)

337

338 image_overrides = mm_options.get("image")

339 + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)

340

341 return {

342 "image": self._get_dummy_images(

vllm/model_executor/models/audioflamingo3.py

+10 -7

@@ -17,7 +17,7 @@

17 # See the License for the specific language governing permissions and

18 # limitations under the License.

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

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

21 from typing import Annotated, Any, Literal, TypeAlias

23 import torch

@@ -30,7 +30,7 @@ from transformers.models.audioflamingo3 import (

30 from transformers.models.qwen2_audio import Qwen2AudioEncoder

32 from vllm.config import VllmConfig

33 -from vllm.config.multimodal import BaseDummyOptions

33 +from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions

34 from vllm.inputs import ModalityData, MultiModalDataDict

35 from vllm.model_executor.layers.activation import get_act_fn

36 from vllm.model_executor.models.module_mapping import MultiModelKeys

@@ -226,6 +226,7 @@ class AudioFlamingo3DummyInputsBuilder(

226 audio_len = int(hf_processor.max_audio_len * sampling_rate)

227 num_audios = mm_counts.get("audio", 0)

228 audio_overrides = mm_options.get("audio")

229 + assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)

230

231 return {

232 "audio": self._get_dummy_audios(

@@ -317,7 +318,7 @@ def _count_audio_tokens_from_mask(

318 if isinstance(chunk_counts, torch.Tensor):

319 counts = chunk_counts.tolist()

320 elif chunk_counts and isinstance(chunk_counts[0], torch.Tensor):

320 - counts = [count.item() for count in chunk_counts]

321 + counts = [int(count) for count in chunk_counts]

322 else:

323 counts = chunk_counts

324

@@ -375,7 +376,7 @@ class AudioFlamingo3MultiModalProcessor(

376 def _call_hf_processor(

377 self,

378 prompt: str,

378 - mm_data: dict[str, object],

379 + mm_data: Mapping[str, object],

380 mm_kwargs: Mapping[str, Any],

381 tok_kwargs: Mapping[str, object],

382 ) -> BatchFeature:

@@ -412,7 +413,8 @@ class AudioFlamingo3MultiModalProcessor(

413 chunk_counts = []

414 for audio in audio_list:

415 # audio is numpy array or list

415 - n_samples = len(audio) if isinstance(audio, list) else audio.shape[0]

416 + assert isinstance(audio, Sized)

417 + n_samples = len(audio)

418

419 n_win = max(1, (n_samples + window_size - 1) // window_size)

420 if n_win > max_windows:

@@ -457,6 +459,7 @@ class AudioFlamingo3MultiModalProcessor(

459 )

460 else:

461 audio_embeds = out_mm_data["audio_embeds"][item_idx]

462 + assert isinstance(audio_embeds, torch.Tensor)

463 num_features = audio_embeds.shape[0]

464

465 if num_features == 0:

@@ -556,8 +559,8 @@ class AudioFlamingo3ForConditionalGeneration(

559 def _process_audio_input(

560 self, audio_input: AudioFlamingo3Inputs

561 ) -> torch.Tensor | tuple[torch.Tensor, ...]:

559 - if audio_input["type"] == "audio_embeds":

560 - audio_embeds = audio_input["audio_embeds"]

562 + if audio_input.type == "audio_embeds":

563 + audio_embeds = audio_input.audio_embeds

564 return tuple(audio_embeds)

565

566 (

vllm/model_executor/models/bagel.py

+2 -1

@@ -14,7 +14,7 @@ import torch

14 import torch.nn as nn

16 from vllm.config import VllmConfig

17 -from vllm.config.multimodal import BaseDummyOptions

17 +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions

18 from vllm.inputs import MultiModalDataDict

19 from vllm.logger import init_logger

20 from vllm.model_executor.layers.activation import get_act_fn

@@ -258,6 +258,7 @@ class BagelDummyInputsBuilder(BaseDummyInputsBuilder[BagelProcessingInfo]):

258 # Use the configured image size

259 image_size = vit_config.image_size

260 image_overrides = mm_options.get("image")

261 + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)

262

263 return {

264 "image": self._get_dummy_images(

vllm/model_executor/models/bailing_moe.py

+2 -2

@@ -182,7 +182,7 @@ class BailingMLP(nn.Module):

182 intermediate_size: int,

183 config: PretrainedConfig,

184 quant_config: QuantizationConfig | None = None,

185 - reduce_results: bool | None = True,

185 + reduce_results: bool = True,

186 prefix: str = "",

187 ) -> None:

188 super().__init__()

@@ -279,7 +279,7 @@ class BailingMoE(nn.Module):

279 else:

280 intermediate_size = config.moe_intermediate_size

281 intermediate_size *= config.num_shared_experts

282 - self.shared_experts = BailingMLP(

282 + self.shared_experts: BailingMLP | None = BailingMLP(

283 intermediate_size=intermediate_size,

284 config=config,

285 quant_config=quant_config,

vllm/model_executor/models/bailing_moe_linear.py

+26 -15

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 import copy

4 -from collections.abc import Iterable

4 +from collections.abc import Callable, Iterable

6 import torch

7 import torch.nn as nn

@@ -133,19 +133,21 @@ class BailingMoeV25MLAAttention(nn.Module):

133

134 if self.q_lora_rank is not None:

135 # Use fused_qkv_a_proj when q_lora_rank is set

136 - self.fused_qkv_a_proj = MergedColumnParallelLinear(

137 - self.hidden_size,

138 - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

139 - bias=False,

140 - quant_config=quant_config,

141 - prefix=f"{prefix}.fused_qkv_a_proj",

142 - disable_tp=True,

136 + self.fused_qkv_a_proj: MergedColumnParallelLinear | None = (

137 + MergedColumnParallelLinear(

138 + self.hidden_size,

139 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

140 + bias=False,

141 + quant_config=quant_config,

142 + prefix=f"{prefix}.fused_qkv_a_proj",

143 + disable_tp=True,

144 + )

145 )

144 - self.q_a_layernorm = RMSNorm(

146 + self.q_a_layernorm: RMSNorm | None = RMSNorm(

147 self.q_lora_rank,

148 eps=config.rms_norm_eps,

149 )

148 - self.q_b_proj = ColumnParallelLinear(

150 + self.q_b_proj: ColumnParallelLinear | None = ColumnParallelLinear(

151 self.q_lora_rank,

152 self.num_heads * self.qk_head_dim,

153 bias=False,

@@ -311,6 +313,8 @@ class BailingMoeV25(nn.Module):

313 "score_function and correction_bias should be "

314 "(softmax, None) or (sigmoid, not None)"

315 )

316 + else:

317 + self.score_function = "softmax"

318

319 # Shared experts (using BailingMLP)

320 if self.num_shared_experts > 0:

@@ -319,7 +323,7 @@ class BailingMoeV25(nn.Module):

323 else:

324 intermediate_size = config.moe_intermediate_size

325 intermediate_size *= config.num_shared_experts

322 - self.shared_experts = BailingMLP(

326 + self.shared_experts: BailingMLP | None = BailingMLP(

327 intermediate_size=intermediate_size,

328 config=config,

329 quant_config=quant_config,

@@ -613,7 +617,9 @@ class BailingMoeV25Model(nn.Module):

617 return False

618

619 param = params_dict[name]

616 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

620 + weight_loader: Callable[..., None] = getattr(

621 + param, "weight_loader", default_weight_loader

622 + )

623

624 if shard_id is None:

625 weight_loader(param, tensor)

@@ -691,11 +697,16 @@ class BailingMoeV25Model(nn.Module):

697 continue

698

699 # Routed experts

694 - for param_name, weight_name, expert_id, shard_id in expert_mappings:

700 + for (

701 + param_name,

702 + weight_name,

703 + expert_id,

704 + expert_shard_id,

705 + ) in expert_mappings:

706 if weight_name not in norm_name:

707 continue

708 mapped = norm_name.replace(weight_name, param_name)

698 - if load_param(mapped, weight, (expert_id, shard_id)):

709 + if load_param(mapped, weight, (expert_id, expert_shard_id)):

710 break

711 continue

712

@@ -783,7 +794,7 @@ class BailingMoeV25ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

794 def get_mamba_state_shape_from_config(

795 cls,

796 vllm_config: VllmConfig,

786 - ) -> tuple[tuple[int, ...], ...]:

797 + ) -> tuple[tuple[int, int, int]]:

798 """Calculate shape for linear attention cache."""

799 config = vllm_config.model_config.hf_config

800 tp_size = vllm_config.parallel_config.tensor_parallel_size

vllm/model_executor/models/bailing_moe_mtp.py

+9 -6

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

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

3 """Inference-only Bailing MoE v2.5 MTP model."""

5 -from collections.abc import Iterable

5 +from collections.abc import Callable, Iterable

7 import torch

8 import torch.nn as nn

@@ -268,14 +268,17 @@ class BailingMoeV25MTPModel(nn.Module):

268 loaded_weight: torch.Tensor,

269 shard_id=None,

270 ) -> bool:

271 - name = maybe_remap_kv_scale_name(name, params_dict)

272 - if name is None:

271 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

272 + if remapped_name is None:

273 return False

274 + name = remapped_name

275 if name not in params_dict or is_pp_missing_parameter(name, self):

276 return False

277

278 param = params_dict[name]

278 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

279 + weight_loader: Callable[..., None] = getattr(

280 + param, "weight_loader", default_weight_loader

281 + )

282 if shard_id is None:

283 weight_loader(param, loaded_weight)

284 elif isinstance(shard_id, int):

@@ -351,14 +354,14 @@ class BailingMoeV25MTPModel(nn.Module):

354

355 if "mlp.experts" in name:

356 for mapping in expert_params_mapping:

354 - param_name, weight_name, expert_id, shard_id = mapping

357 + param_name, weight_name, expert_id, expert_shard_id = mapping

358 if weight_name not in name:

359 continue

360 mapped_name = name.replace(weight_name, param_name)

361 if load_param(

362 mapped_name,

363 loaded_weight,

361 - (expert_id, shard_id),

364 + (expert_id, expert_shard_id),

365 ):

366 loaded = True

367 break

vllm/model_executor/models/bailing_moe_v3.py

+21 -18

@@ -9,7 +9,7 @@ vLLM's parallel linear layers, MLA kernel, KDA kernel and fused MoE loader.

9 """

11 import copy

12 -from collections.abc import Iterable

12 +from collections.abc import Callable, Iterable

14 import torch

15 import torch.nn as nn

@@ -75,7 +75,6 @@ from vllm.third_party.flash_linear_attention.ops.kda import (

75 fused_recurrent_kda_fwd,

76 )

77 from vllm.utils.torch_utils import direct_register_custom_op

78 -from vllm.v1.attention.backend import AttentionMetadata

78 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

79 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum

@@ -136,7 +135,7 @@ def _is_kda_layer(

135

136 def _get_kda_state_shape_for_config(

137 vllm_config: VllmConfig,

139 -) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

138 +) -> tuple[tuple[int, int], tuple[int, int, int]]:

139 config = vllm_config.model_config.hf_config

140 num_spec = (

141 vllm_config.speculative_config.num_speculative_tokens

@@ -205,7 +204,7 @@ class BailingMoeV3MLP(nn.Module):

204 intermediate_size: int,

205 config: PretrainedConfig,

206 quant_config: QuantizationConfig | None = None,

208 - reduce_results: bool | None = True,

207 + reduce_results: bool = True,

208 prefix: str = "",

209 swiglu_limit: float | None = None,

210 ) -> None:

@@ -226,9 +225,7 @@ class BailingMoeV3MLP(nn.Module):

225 prefix=f"{prefix}.down_proj",

226 )

227 self.act_fn = (

229 - SwigluStepAndMul(limit=swiglu_limit)

230 - if swiglu_limit not in (None, 0)

231 - else SiluAndMul()

228 + SwigluStepAndMul(limit=swiglu_limit) if swiglu_limit else SiluAndMul()

229 )

230

231 def forward(self, x):

@@ -264,6 +261,8 @@ class BailingMoeV3MLAAttention(nn.Module):

261 self.num_local_heads = self.num_heads // tp_size

262 self.scaling = self.qk_head_dim**-0.5

263

264 + self.q_proj: ColumnParallelLinear | None

265 + self.kv_a_proj_with_mqa: ReplicatedLinear | None

266 if self.q_lora_rank is None:

267 self.q_proj = ColumnParallelLinear(

268 self.hidden_size,

@@ -400,7 +399,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

399

400 def get_state_dtype(

401 self,

403 - ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:

402 + ) -> tuple[torch.dtype, torch.dtype]:

403 if self.model_config is None or self.cache_config is None:

404 raise ValueError("model_config and cache_config must be set")

405 return MambaStateDtypeCalculator.kda_state_dtype(

@@ -409,7 +408,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

408

409 def get_state_shape(

410 self,

412 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

411 + ) -> tuple[tuple[int, int], tuple[int, int, int]]:

412 return MambaStateShapeCalculator.kda_state_shape(

413 self.tp_size,

414 self.num_heads,

@@ -579,12 +578,12 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

578 core_attn_out: torch.Tensor,

579 ) -> None:

580 forward_context = get_forward_context()

582 - attn_metadata: AttentionMetadata = forward_context.attn_metadata

583 - if attn_metadata is None:

581 + attn_metadata_map = forward_context.attn_metadata

582 + if attn_metadata_map is None:

583 return

584

586 - assert isinstance(attn_metadata, dict)

587 - attn_metadata = attn_metadata[self.prefix]

585 + assert isinstance(attn_metadata_map, dict)

586 + attn_metadata = attn_metadata_map[self.prefix]

587 assert isinstance(attn_metadata, GDNAttentionMetadata)

588 has_initial_state = attn_metadata.has_initial_state

589 spec_query_start_loc = attn_metadata.spec_query_start_loc

@@ -825,6 +824,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

824 assert q is not None and k is not None and v is not None

825 assert g1_non_spec is not None and beta_non_spec is not None

826 assert state_indices is not None

827 + assert has_initial_state is not None

828 zero_idx = state_indices[~has_initial_state]

829 recurrent_state[zero_idx] = 0

830 initial_state = recurrent_state_active[state_indices].contiguous()

@@ -846,6 +846,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

846 assert q is not None and k is not None and v is not None

847 assert g1_non_spec is not None and beta_non_spec is not None

848 assert state_indices is not None

849 + assert query_start_loc is not None

850 out, _ = fused_recurrent_kda(

851 q=q,

852 k=k,

@@ -1205,13 +1206,13 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1206 @classmethod

1207 def get_mamba_state_shape_from_config(

1208 cls, vllm_config: VllmConfig

1208 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

1209 + ) -> tuple[tuple[int, int], tuple[int, int, int]]:

1210 return _get_kda_state_shape_for_config(vllm_config)

1211

1212 @classmethod

1213 def get_mamba_state_dtype_from_config(

1214 cls, vllm_config: VllmConfig

1214 - ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:

1215 + ) -> tuple[torch.dtype, torch.dtype]:

1216 return MambaStateDtypeCalculator.kda_state_dtype(

1217 vllm_config.model_config.dtype,

1218 vllm_config.cache_config.mamba_cache_dtype,

@@ -1243,7 +1244,9 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1244 if name not in params_dict or is_pp_missing_parameter(name, self):

1245 return False

1246 param = params_dict[name]

1246 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

1247 + weight_loader: Callable[..., None] = getattr(

1248 + param, "weight_loader", default_weight_loader

1249 + )

1250 if shard_id is None:

1251 weight_loader(param, tensor)

1252 elif isinstance(shard_id, int):

@@ -1268,10 +1271,10 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1271 if name is None:

1272 continue

1273 loaded = False

1271 - for param_suf, weight_suf, shard_id in stacked_mappings:

1274 + for param_suf, weight_suf, stacked_shard_id in stacked_mappings:

1275 if weight_suf in name:

1276 mapped = name.replace(weight_suf, param_suf)

1274 - if load_param(mapped, weight, shard_id):

1277 + if load_param(mapped, weight, stacked_shard_id):

1278 loaded = True

1279 break

1280 if loaded:

vllm/model_executor/models/bailing_moe_v3_mtp.py

+11 -7

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

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

3 """Inference-only Bailing MoE v3 MTP model."""

5 -from collections.abc import Iterable

5 +from collections.abc import Callable, Iterable

7 import torch

8 import torch.nn as nn

@@ -233,11 +233,12 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):

233 self,

234 input_ids: torch.Tensor | None,

235 positions: torch.Tensor,

236 - hidden_states: torch.Tensor,

236 + hidden_states: torch.Tensor | None = None,

237 intermediate_tensors: IntermediateTensors | None = None,

238 inputs_embeds: torch.Tensor | None = None,

239 spec_step_idx: int = 0,

240 ) -> torch.Tensor:

241 + assert hidden_states is not None

242 del intermediate_tensors

243 return self.model(

244 input_ids,

@@ -280,14 +281,17 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):

281 loaded_weight: torch.Tensor,

282 shard_id=None,

283 ) -> bool:

283 - name = maybe_remap_kv_scale_name(name, params_dict)

284 - if name is None:

284 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

285 + if remapped_name is None:

286 return False

287 + name = remapped_name

288 if name not in params_dict or is_pp_missing_parameter(name, self):

289 return False

290

291 param = params_dict[name]

290 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

292 + weight_loader: Callable[..., None] = getattr(

293 + param, "weight_loader", default_weight_loader

294 + )

295 if shard_id is None:

296 weight_loader(param, loaded_weight)

297 elif isinstance(shard_id, int):

@@ -347,13 +351,13 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):

351 name = normalize_name(name)

352

353 loaded = False

350 - for param_name, weight_name, shard_id in stacked_params_mapping:

354 + for param_name, weight_name, stacked_shard_id in stacked_params_mapping:

355 if weight_name not in name:

356 continue

357 if "mlp.experts." in name and name not in params_dict:

358 continue

359 mapped_name = name.replace(weight_name, param_name)

356 - if load_param(mapped_name, loaded_weight, shard_id):

360 + if load_param(mapped_name, loaded_weight, stacked_shard_id):

361 loaded = True

362 break

363 if loaded:

vllm/model_executor/models/bee.py

+2 -1

@@ -8,7 +8,7 @@ import torch.nn as nn

8 from transformers.activations import GELUActivation

10 from vllm.config import VllmConfig

11 -from vllm.config.multimodal import BaseDummyOptions

11 +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions

12 from vllm.inputs import MultiModalDataDict

13 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -97,6 +97,7 @@ class BeeDummyInputsBuilder(LlavaDummyInputsBuilder[BeeProcessingInfo]):

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

99 image_overrides = mm_options.get("image")

100 + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)

101

102 return {

103 "image": self._get_dummy_images(

vllm/model_executor/models/bert.py

@@ -97,6 +97,7 @@ class BertPooler(SequencePooler):

97 def __init__(self, model_config: ModelConfig):

98 pooler_config = model_config.pooler_config

99 assert pooler_config is not None

100 + assert pooler_config.seq_pooling_type is not None

101

102 config: BertConfig = model_config.hf_config

103

vllm/model_executor/models/bert_with_rope.py

+18 -17

@@ -57,8 +57,8 @@ class BertWithRopeEmbedding(nn.Module):

57 config.vocab_size, config.hidden_size

58 )

59 if config.type_vocab_size > 0:

60 - self.token_type_embeddings = VocabParallelEmbedding(

61 - config.type_vocab_size, config.hidden_size

60 + self.token_type_embeddings: VocabParallelEmbedding | None = (

61 + VocabParallelEmbedding(config.type_vocab_size, config.hidden_size)

62 )

63 else:

64 self.token_type_embeddings = None

@@ -127,6 +127,7 @@ class BertWithRopeAttention(nn.Module):

127 prefix=f"{prefix}.qkv_proj",

128 )

129

130 + assert rotary_kwargs is not None

131 self.rotary_emb = get_rope(**rotary_kwargs)

132

133 self.attn = EncoderOnlyAttention(

@@ -466,7 +467,7 @@ class BertWithRope(nn.Module, SupportsQuant):

467 )

468

469 if add_pooling_layer:

469 - self.pooler = BertPooler(vllm_config.model_config)

470 + self.pooler: BertPooler | None = BertPooler(vllm_config.model_config)

471 else:

472 self.pooler = None

473

@@ -632,7 +633,7 @@ class JinaRobertaModel(BertWithRope):

633 scaling = self.config.lora_alpha / self.config.lora_rank

634 device = self.vllm_config.device_config.device

635

635 - weights = {name: weight for name, weight in weights}

636 + weights_dict = {name: weight for name, weight in weights}

637

638 o = ".original"

639 a = ".0.lora_A"

@@ -641,34 +642,34 @@ class JinaRobertaModel(BertWithRope):

642 # text-matching

643 i = -1

644

644 - for name in list(weights.keys()):

645 + for name in list(weights_dict.keys()):

646 if o in name:

646 - dtype = weights[name].dtype

647 - shape = weights[name].shape

647 + dtype = weights_dict[name].dtype

648 + shape = weights_dict[name].shape

649 weight_name = name[: -len(o)]

650

651 if "embeddings" in weight_name:

651 - B = weights[weight_name + a][i].to(device).float()

652 - A = weights[weight_name + b][i].to(device).float()

652 + B = weights_dict[weight_name + a][i].to(device).float()

653 + A = weights_dict[weight_name + b][i].to(device).float()

654 else:

654 - B = weights[weight_name + b][i].to(device).float()

655 - A = weights[weight_name + a][i].to(device).float()

655 + B = weights_dict[weight_name + b][i].to(device).float()

656 + A = weights_dict[weight_name + a][i].to(device).float()

657

658 weight = (

658 - weights[weight_name + o].to(device)

659 + weights_dict[weight_name + o].to(device)

660 + torch.matmul(B, A).view(shape) * scaling

661 )

662 weight = weight.cpu().to(dtype)

663

663 - weights[weight_name.replace(".parametrizations", "")] = weight

664 + weights_dict[weight_name.replace(".parametrizations", "")] = weight

665

666 del (

666 - weights[weight_name + o],

667 - weights[weight_name + a],

668 - weights[weight_name + b],

667 + weights_dict[weight_name + o],

668 + weights_dict[weight_name + a],

669 + weights_dict[weight_name + b],

670 )

671

671 - return [(name, weight) for name, weight in weights.items()]

672 + return [(name, weight) for name, weight in weights_dict.items()]

673

674 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

675 weights = self.jina_merge_lora_weights(weights)

vllm/model_executor/models/blip2.py

+4 -3

@@ -14,7 +14,7 @@ from transformers import (

14 )

16 from vllm.config import CacheConfig, VllmConfig

17 -from vllm.config.multimodal import BaseDummyOptions

17 +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions

18 from vllm.inputs import MultiModalDataDict

19 from vllm.model_executor.layers.activation import get_act_fn

20 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -454,6 +454,7 @@ class Blip2DummyInputsBuilder(BaseDummyInputsBuilder[Blip2ProcessingInfo]):

454 num_images = mm_counts.get("image", 0)

455

456 image_overrides = mm_options.get("image")

457 + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)

458

459 return {

460 "image": self._get_dummy_images(

@@ -620,8 +621,8 @@ class Blip2ForConditionalGeneration(

621 return self._image_pixels_to_features(self.vision_model, pixel_values)

622

623 def _process_image_input(self, image_input: Blip2ImageInputs) -> torch.Tensor:

623 - if image_input["type"] == "image_embeds":

624 - return image_input["data"]

624 + if image_input.type == "image_embeds":

625 + return image_input.data

626

627 image_features = self._process_image_pixels(image_input)

628

vllm/model_executor/models/idefics2_vision_model.py

+2 -1

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

19 """PyTorch Idefics2 model."""

21 from collections.abc import Iterable

22 +from typing import ClassVar

24 import torch

25 from torch import nn

@@ -352,7 +353,7 @@ class Idefics2Encoder(nn.Module):

353

354

355 class Idefics2VisionTransformer(nn.Module):

355 - hf_to_vllm_mapper = WeightsMapper(

356 + hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper(

357 orig_to_new_stacked={

358 ".q_proj": (".qkv_proj", "q"),

359 ".k_proj": (".qkv_proj", "k"),

vllm/model_executor/models/interfaces.py

+28 -25

@@ -71,6 +71,13 @@ The output embeddings must be one of the following formats:

71 - A single 3D tensor, with the batch dimension grouping the 2D tensors.

72 """

74 +MambaStateShapes: TypeAlias = (

75 + tuple[tuple[int, int]]

76 + | tuple[tuple[int, int, int]]

77 + | tuple[tuple[int, int], tuple[int, int]]

78 + | tuple[tuple[int, int], tuple[int, int, int]]

79 +)

80 +

82 @dataclass(frozen=True, slots=True)

83 class DiarizedTranscriptionSegment:

@@ -616,7 +623,7 @@ class SupportsLoRA(Protocol):

623 # The `embedding_module` and `embedding_padding_modules`

624 # are empty by default.

625 embedding_modules: ClassVar[dict[str, str]] = {}

619 - packed_modules_mapping: dict[str, list[str]] = {}

626 + packed_modules_mapping: ClassVar[dict[str, list[str]]] = {}

627 # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)

628 lora_skip_prefixes: ClassVar[list[str]] = []

629 lora_manager: "LoRAModelManager | None"

@@ -678,6 +685,15 @@ def _supports_lora(model: type[object] | object) -> bool:

685 return isinstance(model, SupportsLoRA)

686

687

688 +class _MakeEmptyIntermediateTensors(Protocol):

689 + def __call__(

690 + self,

691 + batch_size: int,

692 + dtype: torch.dtype,

693 + device: torch.device,

694 + ) -> "IntermediateTensors": ...

695 +

696 +

697 @runtime_checkable

698 class SupportsPP(Protocol):

699 """The interface required for all models that support pipeline parallel."""

@@ -691,14 +707,8 @@ class SupportsPP(Protocol):

707 MRO of your model class.

708 """

709

694 - def make_empty_intermediate_tensors(

695 - self,

696 - batch_size: int,

697 - dtype: torch.dtype,

698 - device: torch.device,

699 - ) -> "IntermediateTensors":

700 - """Called when PP rank > 0 for profiling purposes."""

701 - ...

710 + make_empty_intermediate_tensors: _MakeEmptyIntermediateTensors

711 + """Called when PP rank > 0 for profiling purposes."""

712

713 def forward(

714 self,

@@ -706,7 +716,7 @@ class SupportsPP(Protocol):

716 positions: Tensor,

717 *,

718 intermediate_tensors: "IntermediateTensors | None",

709 - ) -> "IntermediateTensors | None":

719 + ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]":

720 """

721 Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when

722 PP rank > 0.

@@ -723,12 +733,7 @@ class SupportsPP(Protocol):

733 class _SupportsPPType(Protocol):

734 supports_pp: Literal[True]

735

726 - def make_empty_intermediate_tensors(

727 - self,

728 - batch_size: int,

729 - dtype: torch.dtype,

730 - device: torch.device,

731 - ) -> "IntermediateTensors": ...

736 + make_empty_intermediate_tensors: _MakeEmptyIntermediateTensors

737

738 def forward(

739 self,

@@ -736,7 +741,7 @@ class _SupportsPPType(Protocol):

741 positions: Tensor,

742 *,

743 intermediate_tensors: "IntermediateTensors | None",

739 - ) -> "Tensor | IntermediateTensors": ...

744 + ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]": ...

745

746

747 @overload

@@ -785,7 +790,7 @@ def supports_pp(

790

791 def _supports_pp_attributes(model: type[object] | object) -> bool:

792 if isinstance(model, type):

788 - return isinstance(model, _SupportsPPType)

793 + return SupportsPP in model.__mro__ or isinstance(model, _SupportsPPType)

794

795 return isinstance(model, SupportsPP)

796

@@ -867,16 +872,14 @@ class IsHybrid(Protocol):

872 def get_mamba_state_shape_from_config(

873 cls,

874 vllm_config: "VllmConfig",

870 - ) -> tuple[tuple[int, int], tuple[int, int, int]]:

875 + ) -> MambaStateShapes:

876 """Calculate shapes for Mamba's convolutional and state caches.

877

878 Args:

879 vllm_config: vLLM config

880

881 Returns:

877 - Tuple containing:

878 - - conv_state_shape: Shape for convolutional state cache

879 - - temporal_state_shape: Shape for state space model cache

882 + Shapes for each state cache used by the model.

883 """

884 ...

885

@@ -1090,7 +1093,7 @@ class SupportsQuant:

1093 """The interface required for all models that support quantization."""

1094

1095 hf_to_vllm_mapper: ClassVar["WeightsMapper | None"] = None

1093 - packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None

1096 + packed_modules_mapping: ClassVar[dict[str, list[str]]]

1097 quant_config: QuantizationConfig | None = None

1098

1099 def __new__(cls, *args, **kwargs) -> Self:

@@ -1125,8 +1128,8 @@ class SupportsQuant:

1128 if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None:

1129 unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper()

1130 self.quant_config.apply_vllm_mapper(unstacked_mapper)

1128 - if self.packed_modules_mapping is not None:

1129 - self.quant_config.packed_modules_mapping.update(self.packed_modules_mapping)

1131 + if packed_modules_mapping := getattr(self, "packed_modules_mapping", None):

1132 + self.quant_config.packed_modules_mapping.update(packed_modules_mapping)

1133

1134

1135 @runtime_checkable

vllm/model_executor/models/llama.py

+2 -1

@@ -26,6 +26,7 @@

27 from collections.abc import Iterable

28 from itertools import islice

29 +from typing import ClassVar

31 import torch

32 from torch import nn

@@ -342,7 +343,7 @@ class LlamaDecoderLayer(nn.Module):

343 },

344 )

345 class LlamaModel(nn.Module, EagleModelMixin):

345 - hf_to_vllm_mapper = WeightsMapper(

346 + hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper(

347 orig_to_new_stacked={

348 # weight_name: (param_name, shard_id)

349 ".q_proj": (".qkv_proj", "q"),

vllm/model_executor/models/llava_next.py

+1 -3

@@ -204,9 +204,7 @@ class BaseLlavaNextMultiModalProcessor(BaseLlavaMultiModalProcessor[_I]):

204 raise NotImplementedError

205

206

207 -class LlavaNextMultiModalProcessor(

208 - BaseLlavaNextMultiModalProcessor[LlavaNextProcessingInfo]

209 -):

207 +class LlavaNextMultiModalProcessor(BaseLlavaNextMultiModalProcessor[_I]):

208 def _get_mm_fields_config(

209 self,

210 hf_inputs: BatchFeature,