@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70021a0081

[Mypy Fix] Mypy fix for "vllm/model_executor/models/[gG]" (#53616)

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

Wentao Ye committed Aug 25, 2026 at 12:47 UTC 70021a0081cc10d8862f7458c36353b72ea9b5f4

26 files changed +350 -123

tools/pre_commit/mypy.py

-1

@@ -112,7 +112,6 @@ SEPARATE_GROUPS = [

112 ]

113

114 EXCLUDE = [

115 - r"vllm/model_executor/models/[gG]",

115 r"vllm/model_executor/models/[hH]",

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

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

vllm/model_executor/models/deepseek_v2.py

+2 -2

@@ -457,7 +457,7 @@ class DeepseekV2Attention(nn.Module):

457 qk_nope_head_dim: int,

458 qk_rope_head_dim: int,

459 v_head_dim: int,

460 - q_lora_rank: int,

460 + q_lora_rank: int | None,

461 kv_lora_rank: int,

462 max_position_embeddings: int = 8192,

463 cache_config: CacheConfig | None = None,

@@ -495,7 +495,7 @@ class DeepseekV2Attention(nn.Module):

495 )

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

497 self.q_b_proj = ColumnParallelLinear(

498 - q_lora_rank,

498 + self.q_lora_rank,

499 self.num_heads * self.qk_head_dim,

500 bias=False,

501 quant_config=quant_config,

vllm/model_executor/models/gemma.py

@@ -311,6 +311,7 @@ class GemmaModel(nn.Module):

311 hidden_states = self.embed_input_ids(input_ids)

312 residual = None

313 else:

314 + assert intermediate_tensors is not None

315 hidden_states = intermediate_tensors["hidden_states"]

316 residual = intermediate_tensors["residual"]

317 for layer in islice(self.layers, self.start_layer, self.end_layer):

vllm/model_executor/models/gemma3_mm.py

@@ -852,7 +852,9 @@ class Gemma3ForConditionalGeneration(

852 def encoder_eager_forward(

853 self,

854 mm_kwargs: dict[str, Any],

855 + path: str = "default",

856 ) -> torch.Tensor:

857 image_input = self._parse_and_validate_image_input(**mm_kwargs)

858 + assert image_input is not None

859 results = self._process_image_input(image_input)

860 return torch.cat(results, dim=0)

vllm/model_executor/models/gemma3n.py

+1 -1

@@ -77,7 +77,7 @@ class Gemma3nAltUp(nn.Module):

77 altup_num_inputs: int,

78 altup_coef_clip: float,

79 altup_active_idx: int,

80 - quant_config: QuantizationConfig,

80 + quant_config: QuantizationConfig | None,

81 prefix: str,

82 ):

83 super().__init__()

vllm/model_executor/models/gemma3n_mm.py

+31 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import inspect

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

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

5 from typing import Annotated, Any, Literal

7 import torch

@@ -18,7 +18,7 @@ from transformers.models.gemma3n import (

18 from transformers.models.siglip import SiglipImageProcessorFast

20 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

21 -from vllm.config.multimodal import BaseDummyOptions

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

22 from vllm.config.speech_to_text import SpeechToTextParams

23 from vllm.inputs import MultiModalDataDict, PromptType, TextPrompt

24 from vllm.logger import init_logger

@@ -115,8 +115,10 @@ def batch_audio_features(

115 The `(bn, s_max, f)` features and their `(bn, s_max)` mask.

116 """

117 if isinstance(input_features, torch.Tensor):

118 + assert isinstance(input_features_mask, torch.Tensor)

119 return input_features.squeeze(1), input_features_mask.squeeze(1)

120

121 + assert isinstance(input_features_mask, list)

122 max_len = max(features.shape[0] for features in input_features)

123 batched_features = input_features[0].new_zeros(

124 (len(input_features), max_len, input_features[0].shape[-1])

@@ -248,6 +250,7 @@ class Gemma3nDummyInputsBuilder(BaseDummyInputsBuilder[Gemma3nProcessingInfo]):

250

251 image_overrides = mm_options.get("image")

252 audio_overrides = mm_options.get("audio")

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

254

255 return {

256 "image": self._get_dummy_images(

@@ -293,10 +296,14 @@ class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo])

296 # output (and thus its cache entry) is independent of the other

297 # items in the batch. `batch_audio_features` re-pads the

298 # per-item features so audio_tower can still run batched.

296 - num_frames = [

297 - self._get_num_audio_frames_alone(len(audio))

298 - for audio in mm_data["audio"]

299 - ]

299 + audios = mm_data["audio"]

300 + if not isinstance(audios, Sequence):

301 + raise TypeError("audio data must be a sequence")

302 + num_frames = []

303 + for audio in audios:

304 + if not isinstance(audio, Sized):

305 + raise TypeError("each audio item must have a length")

306 + num_frames.append(self._get_num_audio_frames_alone(len(audio)))

307 processed_data["input_features_padded"] = [

308 features[:n]

309 for features, n in zip(processed_data.pop("input_features"), num_frames)

@@ -685,7 +692,9 @@ class Gemma3nForConditionalGeneration(

692 )

693

694 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

688 - mm_input_by_modality = {}

695 + mm_input_by_modality: dict[

696 + str, Gemma3nImageInputs | Gemma3nAudioInputs | None

697 + ] = {}

698

699 # Preserve the order of modalities if there are multiple of them

700 # from the order of kwargs.

@@ -852,7 +861,13 @@ class Gemma3nForConditionalGeneration(

861 # select a chunk of pre-allocated PLEs. During normal execution,

862 # `embed_input_ids` is called before forward, hence this slice

863 # will contain PLEs computed from the actual input_ids.

855 - per_layer_inputs = self.per_layer_embeddings[: inputs_embeds.shape[0]]

864 + if inputs_embeds is not None:

865 + num_tokens = inputs_embeds.shape[0]

866 + elif intermediate_tensors is not None:

867 + num_tokens = intermediate_tensors["hidden_states"].shape[0]

868 + else:

869 + raise ValueError("inputs_embeds is required on the first PP rank")

870 + per_layer_inputs = self.per_layer_embeddings[:num_tokens]

871

872 hidden_states = self.language_model.model(

873 input_ids,

@@ -913,9 +928,15 @@ class Gemma3nForConditionalGeneration(

928 prompt += " this audio"

929

930 # We assume the language is a valid ISO 639-1 code.

916 - full_lang_name = cls.supported_languages.get(language, "")

931 + full_lang_name = (

932 + cls.supported_languages.get(language, "") if language is not None else ""

933 + )

934 # Translation only for now

918 - full_lang_name_to = cls.supported_languages.get(to_language, "")

935 + full_lang_name_to = (

936 + cls.supported_languages.get(to_language, "")

937 + if to_language is not None

938 + else ""

939 + )

940

941 if task_type == "transcribe" and full_lang_name:

942 prompt += f" into {full_lang_name}"

vllm/model_executor/models/gemma4.py

+31 -7

@@ -352,15 +352,19 @@ class Gemma4MoE(nn.Module):

352 )

353

354 # MoERunner experts with custom Gemma4 routing

355 + intermediate_size = getattr(

356 + config,

357 + "moe_intermediate_size",

358 + getattr(config, "expert_intermediate_size", None),

359 + )

360 + if intermediate_size is None:

361 + raise ValueError("Gemma4 MoE requires an expert intermediate size")

362 +

363 self.experts = FusedMoEFactory(

364 num_experts=config.num_experts,

365 top_k=config.top_k_experts,

366 hidden_size=config.hidden_size,

359 - intermediate_size=getattr(

360 - config,

361 - "moe_intermediate_size",

362 - getattr(config, "expert_intermediate_size", None),

363 - ),

367 + intermediate_size=intermediate_size,

368 renormalize=True,

369 quant_config=quant_config,

370 prefix=f"{prefix}.experts",

@@ -624,6 +628,11 @@ class Gemma4DecoderLayer(nn.Module):

628 self.enable_moe_block = getattr(config, "enable_moe_block", False) or getattr(

629 config, "use_second_mlp_block", False

630 )

631 + self.router: Gemma4Router | None

632 + self.moe: Gemma4MoE | None

633 + self.post_feedforward_layernorm_1: RMSNorm | None

634 + self.post_feedforward_layernorm_2: RMSNorm | None

635 + self.pre_feedforward_layernorm_2: RMSNorm | None

636 if self.enable_moe_block:

637 self.router = Gemma4Router(

638 config,

@@ -652,6 +661,9 @@ class Gemma4DecoderLayer(nn.Module):

661 self.pre_feedforward_layernorm_2 = None

662

663 # Per-Layer Embedding (PLE) components — present in each decoder layer

664 + self.per_layer_input_gate: ReplicatedLinear | None

665 + self.per_layer_projection: ReplicatedLinear | None

666 + self.post_per_layer_input_norm: RMSNorm | None

667 if (

668 self.hidden_size_per_layer_input is not None

669 and self.hidden_size_per_layer_input > 0

@@ -716,6 +728,11 @@ class Gemma4DecoderLayer(nn.Module):

728 hidden_states = self.mlp(hidden_states)

729

730 if self.enable_moe_block:

731 + assert self.post_feedforward_layernorm_1 is not None

732 + assert self.pre_feedforward_layernorm_2 is not None

733 + assert self.router is not None

734 + assert self.moe is not None

735 + assert self.post_feedforward_layernorm_2 is not None

736 hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states)

737

738 hidden_states_2 = self.pre_feedforward_layernorm_2(residual)

@@ -731,6 +748,8 @@ class Gemma4DecoderLayer(nn.Module):

748

749 # Apply PLE (Per-Layer Embedding) if configured

750 if per_layer_input is not None and self.per_layer_input_gate is not None:

751 + assert self.per_layer_projection is not None

752 + assert self.post_per_layer_input_norm is not None

753 gate = self.per_layer_input_gate(hidden_states)

754 gate = torch.nn.functional.gelu(gate, approximate="tanh")

755 gated_per_layer = gate * per_layer_input

@@ -865,6 +884,7 @@ class Gemma4SelfDecoderLayers(nn.Module):

884 """

885 if self.per_layer_model_projection is None:

886 return None

887 + assert self.per_layer_projection_norm is not None

888 per_layer_projection = self.per_layer_model_projection(inputs_embeds)

889 per_layer_projection = per_layer_projection * self.per_layer_projection_scale

890 per_layer_projection = per_layer_projection.reshape(

@@ -971,6 +991,9 @@ class Gemma4Model(nn.Module, EagleModelMixin):

991 )

992

993 # Per-Layer Embedding (PLE) components

994 + self.embed_tokens_per_layer: VocabParallelEmbedding | None

995 + self.per_layer_model_projection: ColumnParallelLinear | None

996 + self.per_layer_projection_norm: RMSNorm | None

997 if (

998 self.hidden_size_per_layer_input is not None

999 and self.hidden_size_per_layer_input > 0

@@ -1472,9 +1495,10 @@ class Gemma4Model(nn.Module, EagleModelMixin):

1495 else:

1496 if name.endswith(".bias") and name not in params_dict:

1497 continue

1475 - name = maybe_remap_kv_scale_name(name, params_dict)

1476 - if name is None:

1498 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

1499 + if remapped_name is None:

1500 continue

1501 + name = remapped_name

1502 if is_pp_missing_parameter(name, self):

1503 continue

1504 # Skip if name doesn't exist in params_dict (e.g., individual

vllm/model_executor/models/gemma4_dspark.py

+5 -1

@@ -89,7 +89,11 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):

89 ) -> tuple[torch.Tensor, torch.Tensor]:

90 k, _ = self.k_proj(hidden_states)

91 k_normed = self.k_norm(k.unflatten(-1, (self.num_kv_heads, self.head_dim)))

92 - v_src = k if self.use_k_eq_v else self.v_proj(hidden_states)[0]

92 + if self.use_k_eq_v:

93 + v_src = k

94 + else:

95 + assert self.v_proj is not None

96 + v_src = self.v_proj(hidden_states)[0]

97 v_normed = self.v_norm(v_src.unflatten(-1, (self.num_kv_heads, self.head_dim)))

98 return k_normed.flatten(-2, -1), v_normed.flatten(-2, -1)

vllm/model_executor/models/gemma4_mm.py

+58 -14

@@ -35,7 +35,11 @@ from transformers.models.gemma4.configuration_gemma4 import (

36 from vllm.config import VllmConfig

37 from vllm.config.model import get_served_model_name

38 -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions

38 +from vllm.config.multimodal import (

39 + AudioDummyOptions,

40 + BaseDummyOptions,

41 + VideoDummyOptions,

42 +)

43 from vllm.inputs import MultiModalDataDict

44 from vllm.logger import init_logger

45 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -189,6 +193,7 @@ class Gemma4VideoInputs(TensorSchema):

193 torch.Tensor,

194 TensorShape("bn", "np", 2),

195 ]

196 + video_frame_counts: Annotated[torch.Tensor, TensorShape("nv")]

197

198

199 # ---------------------------------------------------------------------------

@@ -490,8 +495,10 @@ class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]):

495 image_overrides = mm_options.get("image") if mm_options else None

496 audio_overrides = mm_options.get("audio") if mm_options else None

497 video_overrides = mm_options.get("video") if mm_options else None

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

499 + assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)

500

494 - data: MultiModalDataDict = {

501 + data: dict[str, Any] = {

502 "image": self._get_dummy_images(

503 width=img_width,

504 height=img_height,

@@ -592,6 +599,7 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):

599 # same vision tower, matching transformers processing_gemma4.py.

600 video_outputs: dict[str, Any] = {}

601 if videos := mm_data.pop("videos", []):

602 + assert isinstance(videos, list)

603 processor = self.info.get_hf_processor()

604

605 all_video_pixel_values: list[torch.Tensor] = []

@@ -858,8 +866,29 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):

866

867 def get_replacement_video(item_idx: int):

868 out_item = out_mm_kwargs["video"][item_idx]

861 - timestamps = out_item["video_timestamps"].data

862 - num_soft = out_item["video_num_soft_tokens"].data

869 + timestamps_data = out_item["video_timestamps"].data

870 + if isinstance(timestamps_data, torch.Tensor):

871 + timestamps_data = timestamps_data.tolist()

872 + if not isinstance(timestamps_data, list):

873 + raise TypeError("video timestamps must be a list")

874 + timestamps: list[float] = []

875 + for timestamp in timestamps_data:

876 + if not isinstance(timestamp, (int, float)):

877 + raise TypeError("each video timestamp must be numeric")

878 + timestamps.append(float(timestamp))

879 +

880 + num_soft_data = out_item["video_num_soft_tokens"].data

881 + if isinstance(num_soft_data, torch.Tensor):

882 + num_soft_data = num_soft_data.tolist()

883 + if not isinstance(num_soft_data, list):

884 + raise TypeError("video soft-token counts must be a list")

885 + num_soft: list[int] = []

886 + for count in num_soft_data:

887 + if not isinstance(count, int):

888 + raise TypeError(

889 + "each video soft-token count must be an integer"

890 + )

891 + num_soft.append(count)

892 return self.info.get_video_repl(

893 timestamps=timestamps,

894 num_soft_tokens_per_frame=num_soft,

@@ -1038,6 +1067,7 @@ class Gemma4ForConditionalGeneration(

1067 # (Marlin, FP8, …) require dimensions divisible by 64, which

1068 # the vision tower (intermediate_size=4304) does not satisfy.

1069 # TODO(mgoin): remove this by fixing kernel padding.

1070 + tower_quant: QuantizationConfig | None

1071 if quant_config and quant_config.get_name() in [

1072 "bitsandbytes",

1073 "torchao",

@@ -1068,6 +1098,7 @@ class Gemma4ForConditionalGeneration(

1098 )

1099

1100 # ---- Audio tower (variants with audio_config) ----

1101 + self.embed_audio: Gemma4MultimodalEmbedder | None

1102 if config.audio_config is not None:

1103 with self._mark_tower_model(vllm_config, "audio"):

1104 self.audio_tower = AutoModel.from_config(config=config.audio_config)

@@ -1177,22 +1208,24 @@ class Gemma4ForConditionalGeneration(

1208

1209 def _parse_and_validate_video_input(

1210 self, **kwargs: object

1180 - ) -> dict[str, torch.Tensor] | None:

1211 + ) -> Gemma4VideoInputs | None:

1212 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1213 pixel_position_ids_videos = kwargs.pop("pixel_position_ids_videos", None)

1214 video_frame_counts = kwargs.pop("video_frame_counts", None)

1215 if pixel_values_videos is None:

1216 return None

1186 - return {

1187 - "pixel_values_videos": pixel_values_videos,

1188 - "pixel_position_ids_videos": pixel_position_ids_videos,

1189 - "video_frame_counts": video_frame_counts,

1190 - }

1217 + return Gemma4VideoInputs(

1218 + pixel_values_videos=pixel_values_videos,

1219 + pixel_position_ids_videos=pixel_position_ids_videos,

1220 + video_frame_counts=video_frame_counts,

1221 + )

1222

1223 def _parse_and_validate_multimodal_inputs(

1224 self, **kwargs: object

1225 ) -> dict[str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None]:

1195 - mm_input_by_modality = {}

1226 + mm_input_by_modality: dict[

1227 + str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None

1228 + ] = {}

1229 for input_key in list(kwargs):

1230 if (

1231 input_key in ("pixel_values", "image_embeds")

@@ -1419,7 +1452,7 @@ class Gemma4ForConditionalGeneration(

1452

1453 def _process_video_input(

1454 self,

1422 - video_input: dict[str, torch.Tensor],

1455 + video_input: Gemma4VideoInputs,

1456 ) -> list[torch.Tensor]:

1457 """Batch-encode video frames through the vision tower.

1458

@@ -1543,6 +1576,8 @@ class Gemma4ForConditionalGeneration(

1576 audio_input["input_features_mask"],

1577 )

1578

1579 + assert self.audio_tower is not None

1580 + assert self.embed_audio is not None

1581 # Run audio tower — mask convention: True=valid, False=padding.

1582 audio_outputs = self.audio_tower(input_features, input_features_mask)

1583 if isinstance(audio_outputs, tuple):

@@ -1576,14 +1611,17 @@ class Gemma4ForConditionalGeneration(

1611 if multimodal_input is None:

1612 continue

1613 if modality == "image":

1614 + assert isinstance(multimodal_input, Gemma4ImageInputs)

1615 multimodal_embeddings.extend(

1616 self._process_image_input(multimodal_input)

1617 )

1618 elif modality == "video":

1619 + assert isinstance(multimodal_input, Gemma4VideoInputs)

1620 multimodal_embeddings.extend(

1621 self._process_video_input(multimodal_input)

1622 )

1623 elif modality == "audio":

1624 + assert isinstance(multimodal_input, Gemma4AudioInputs)

1625 multimodal_embeddings.extend(

1626 self._process_audio_input(multimodal_input)

1627 )

@@ -2204,13 +2242,15 @@ class Gemma4ForConditionalGeneration(

2242 mm_kwargs: MultiModalKwargsItem | None,

2243 num_mm_embeds: int,

2244 ) -> tuple[int, int | None]:

2245 + tower_tokens: int | None = None

2246 + connector_tokens: int | None = None

2247 if modality in ("image", "video"):

2248 vision_config = self.config.vision_config

2249 pooling_k2 = vision_config.pooling_kernel_size**2

2250

2251 if modality == "image":

2252 pixel_values_key = "pixel_values"

2213 - max_soft_tokens = vision_config.default_output_length

2253 + max_soft_tokens = int(vision_config.default_output_length)

2254 mm_processor_kwargs = getattr(

2255 getattr(self, "multimodal_config", None),

2256 "mm_processor_kwargs",

@@ -2241,7 +2281,7 @@ class Gemma4ForConditionalGeneration(

2281 * pooling_k2

2282 )

2283

2244 - if modality == "audio":

2284 + elif modality == "audio":

2285 tower_tokens = num_mm_embeds

2286 connector_tokens = num_mm_embeds

2287

@@ -2255,6 +2295,10 @@ class Gemma4ForConditionalGeneration(

2295 tower_tokens = audio_tokens

2296 connector_tokens = audio_tokens

2297

2298 + else:

2299 + raise ValueError(f"Unsupported modality: {modality}")

2300 +

2301 + assert tower_tokens is not None

2302 return tower_tokens, connector_tokens

2303

2304 @classmethod

vllm/model_executor/models/gemma4_mtp.py

+10 -3

@@ -345,7 +345,9 @@ class Gemma4MultiTokenPredictor(nn.Module):

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

346 super().__init__()

347

348 - config = vllm_config.speculative_config.draft_model_config.hf_config

348 + speculative_config = vllm_config.speculative_config

349 + assert speculative_config is not None

350 + config = speculative_config.draft_model_config.hf_config

351 text_config = _get_text_config(config)

352 quant_config = get_draft_quant_config(vllm_config)

353 self.config = text_config

@@ -467,7 +469,9 @@ class Gemma4MTP(nn.Module):

469

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

471 super().__init__()

470 - config = vllm_config.speculative_config.draft_model_config.hf_config

472 + speculative_config = vllm_config.speculative_config

473 + assert speculative_config is not None

474 + config = speculative_config.draft_model_config.hf_config

475 text_config = _get_text_config(config)

476 self.quant_config = get_draft_quant_config(vllm_config)

477 self.config = config

@@ -496,6 +500,7 @@ class Gemma4MTP(nn.Module):

500 soft_cap=getattr(text_config, "final_logit_softcapping", None),

501 )

502

503 + self.masked_embedding: Gemma4MTPMaskedEmbedder | None

504 if getattr(config, "use_ordered_embeddings", False):

505 num_centroids = getattr(config, "num_centroids", 2048)

506 top_k = getattr(config, "centroid_intermediate_top_k", 32)

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

521 else:

522 self.masked_embedding = None

523

519 - draft_cfg = vllm_config.speculative_config.draft_model_config

524 + draft_cfg = speculative_config.draft_model_config

525 gen_cfg = draft_cfg.try_get_generation_config()

526 self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None

527

@@ -545,6 +550,7 @@ class Gemma4MTP(nn.Module):

550 def _get_full_lm_head_weight(self) -> torch.Tensor:

551 if self._stable_full_lm_head_weight is not None:

552 return self._stable_full_lm_head_weight

553 + assert self.masked_embedding is not None

554 lm_head_weight = self.lm_head.weight

555 tp_size = get_tensor_model_parallel_world_size()

556 if tp_size > 1:

@@ -579,6 +585,7 @@ class Gemma4MTP(nn.Module):

585 hidden_states: torch.Tensor,

586 ) -> torch.Tensor:

587 """Sparse argmax via centroids masking. Returns token IDs directly."""

588 + assert self.masked_embedding is not None

589 return self.masked_embedding.get_top_tokens(

590 hidden_states,

591 self._get_full_lm_head_weight(),

vllm/model_executor/models/gemma4_unified.py

+12 -16

@@ -44,6 +44,7 @@ from vllm.model_executor.models.gemma4_mm import (

44 Gemma4MultimodalEmbedder,

45 Gemma4MultiModalProcessor,

46 Gemma4ProcessingInfo,

47 + Gemma4VideoInputs,

48 _get_max_soft_tokens,

49 )

50 from vllm.model_executor.models.module_mapping import MultiModelKeys

@@ -265,22 +266,16 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):

266 self.audio_tower = None

267

268 # ---- Encoder-free vision embedder ----

268 - self.vision_embedder = (

269 - Gemma4UnifiedVisionEmbedder(

270 - config.vision_config,

271 - quant_config=quant_config,

272 - prefix=maybe_prefix(prefix, "vision_embedder"),

273 - )

274 - if config.vision_config is not None

275 - else None

269 + if config.vision_config is None:

270 + raise ValueError("Gemma4 Unified requires a vision configuration")

271 + self.vision_embedder = Gemma4UnifiedVisionEmbedder(

272 + config.vision_config,

273 + quant_config=quant_config,

274 + prefix=maybe_prefix(prefix, "vision_embedder"),

275 )

277 - self.embed_vision = (

278 - Gemma4MultimodalEmbedder(

279 - config.vision_config,

280 - config.text_config,

281 - )

282 - if config.vision_config is not None

283 - else None

276 + self.embed_vision = Gemma4MultimodalEmbedder(

277 + config.vision_config,

278 + config.text_config,

279 )

280

281 # ---- Encoder-free audio embedder ----

@@ -382,7 +377,7 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):

377

378 def _process_video_input(

379 self,

385 - video_input: dict[str, torch.Tensor],

380 + video_input: Gemma4VideoInputs,

381 ) -> list[torch.Tensor]:

382 """Project video frames to LM space, one frame at a time.

383

@@ -429,6 +424,7 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):

424 audio_input["input_features_mask"],

425 )

426

427 + assert self.embed_audio is not None

428 target_dtype = self.embed_audio.embedding_projection.weight.dtype

429 audio_features = self.embed_audio(input_features.to(target_dtype))

430 per_audio: list[torch.Tensor] = []

vllm/model_executor/models/glm4_1v.py

+43 -15

@@ -1393,9 +1393,15 @@ class Glm4vProcessingInfo(BaseProcessingInfo):

1393 return list(timestamps_list)

1394

1395 def _get_video_frame_embed_token_id(self, hf_processor: object) -> int:

1396 - if isinstance(hf_processor, Glm4vProcessor) or TRANSFORMERS_WITH_GA:

1397 - return hf_processor.image_token_id

1398 - return hf_processor.video_token_id

1396 + attr = (

1397 + "image_token_id"

1398 + if isinstance(hf_processor, Glm4vProcessor) or TRANSFORMERS_WITH_GA

1399 + else "video_token_id"

1400 + )

1401 + token_id = getattr(hf_processor, attr, None)

1402 + if not isinstance(token_id, int):

1403 + raise ValueError(f"Processor has no valid {attr}")

1404 + return token_id

1405

1406 def _construct_video_placeholder(

1407 self,

@@ -1479,6 +1485,7 @@ class Glm4vDummyInputsBuilder(BaseDummyInputsBuilder[Glm4vProcessingInfo]):

1485

1486 image_overrides = mm_options.get("image")

1487 video_overrides = mm_options.get("video")

1488 + assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)

1489

1490 return {

1491 "image": self._get_dummy_images(

@@ -1519,7 +1526,7 @@ class Glm4vDummyInputsBuilder(BaseDummyInputsBuilder[Glm4vProcessingInfo]):

1526 )

1527 videos = [v.copy() for v in videos]

1528

1522 - video_items = []

1529 + video_items: list[VideoItem] = []

1530 for video in videos:

1531 video_num_frames = video.shape[0]

1532 video_metadata = {

@@ -1611,16 +1618,14 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]):

1618 processed_data.update(passthrough_data)

1619 return processed_data

1620

1614 - if (

1615 - "videos" in mm_data

1616 - and isinstance(mm_data["videos"], list)

1617 - and len(mm_data["videos"]) > 0

1618 - ):

1621 + videos = mm_data.get("videos")

1622 + if isinstance(videos, list) and videos:

1623 video_grid_thw_lst = []

1624 pixel_values_videos_lst = []

1625 frame_embed_token_id = self.info._get_video_frame_embed_token_id(processor)

1626 swap_video_frame_tokens = frame_embed_token_id == processor.image_token_id

1623 - for item in mm_data.pop("videos", []):

1627 + mm_data.pop("videos")

1628 + for item in videos:

1629 video_array, metadata = item

1630

1631 # don't update mm_kwargs inplace

@@ -1794,6 +1799,7 @@ class Glm4vForConditionalGeneration(

1799 self.config = config

1800 self.model_config = vllm_config.model_config

1801 self.multimodal_config = multimodal_config

1802 + assert multimodal_config is not None

1803 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1804 self.is_multimodal_pruning_enabled = (

1805 multimodal_config.is_multimodal_pruning_enabled()

@@ -1850,6 +1856,7 @@ class Glm4vForConditionalGeneration(

1856 image_embeds=image_embeds,

1857 image_grid_thw=image_grid_thw,

1858 )

1859 + raise AssertionError("image input must contain pixels or embeddings")

1860

1861 def _parse_and_validate_video_input(

1862 self, **kwargs: object

@@ -1874,6 +1881,7 @@ class Glm4vForConditionalGeneration(

1881 video_embeds=video_embeds,

1882 video_grid_thw=video_grid_thw,

1883 )

1884 + raise AssertionError("video input must contain pixels or embeddings")

1885

1886 def _process_image_input(

1887 self, image_input: Glm4vImageInputs

@@ -1996,12 +2004,12 @@ class Glm4vForConditionalGeneration(

2004 def _get_grid_thw_by_modality(

2005 self,

2006 mm_kwargs: dict[str, Any],

1999 - ) -> list[tuple[int, int, int]]:

2007 + ) -> list[list[int]]:

2008 grid_thw_key = f"{self.get_input_modality(mm_kwargs)}_grid_thw"

2009 grid_thw = mm_kwargs[grid_thw_key]

2010 if not isinstance(grid_thw, list):

2011 grid_thw = grid_thw.tolist()

2004 - return grid_thw

2012 + return [list(row) for row in grid_thw]

2013

2014 def get_encoder_cudagraph_item_specs(

2015 self,

@@ -2188,7 +2196,7 @@ class Glm4vForConditionalGeneration(

2196 return self.visual(pixel_values, grid_thw)

2197

2198 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

2191 - mm_input_by_modality = {}

2199 + mm_input_by_modality: dict[str, Glm4vImageInputs | Glm4vVideoInputs | None] = {}

2200

2201 # Preserve the order of modalities if there are multiple of them

2202 # from the order of kwargs.

@@ -2223,9 +2231,17 @@ class Glm4vForConditionalGeneration(

2231 for modality in mm_input_by_modality:

2232 multimodal_input = mm_input_by_modality[modality]

2233 if modality == "image":

2234 + assert isinstance(

2235 + multimodal_input,

2236 + (Glm4vImagePixelInputs, Glm4vImageEmbeddingInputs),

2237 + )

2238 image_embeddings = self._process_image_input(multimodal_input)

2239 multimodal_embeddings += tuple(image_embeddings)

2240 if modality == "video":

2241 + assert isinstance(

2242 + multimodal_input,

2243 + (Glm4vVideoPixelInputs, Glm4vVideoEmbeddingInputs),

2244 + )

2245 video_embeddings = self._process_video_input(multimodal_input)

2246 multimodal_embeddings += tuple(video_embeddings)

2247 return multimodal_embeddings

@@ -2238,14 +2254,26 @@ class Glm4vForConditionalGeneration(

2254 for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):

2255 embed_ranges = mm_feature.mm_position.extract_embeds_range()

2256 if mm_feature.modality == "image":

2241 - t, h, w = mm_feature.data["image_grid_thw"].data.tolist()

2257 + feature_data = mm_feature.data

2258 + assert feature_data is not None

2259 + grid_item = feature_data.get("image_grid_thw")

2260 + assert grid_item is not None

2261 + grid_data = grid_item.data

2262 + assert isinstance(grid_data, torch.Tensor)

2263 + t, h, w = grid_data.tolist()

2264 assert t == 1, f"Image must have 1 frame, got {t}"

2265 assert len(embed_ranges) == 1

2266 offset, end = embed_ranges[0]

2267 assert end - offset + 1 == h * w // spatial_merge_size**2

2268 yield offset, t, h // spatial_merge_size, w // spatial_merge_size

2269 elif mm_feature.modality == "video":

2248 - t, h, w = mm_feature.data["video_grid_thw"].data.tolist()

2270 + feature_data = mm_feature.data

2271 + assert feature_data is not None

2272 + grid_item = feature_data.get("video_grid_thw")

2273 + assert grid_item is not None

2274 + grid_data = grid_item.data

2275 + assert isinstance(grid_data, torch.Tensor)

2276 + t, h, w = grid_data.tolist()

2277 llm_grid_h = h // spatial_merge_size

2278 llm_grid_w = w // spatial_merge_size

2279 num_tokens_per_frame = llm_grid_h * llm_grid_w

vllm/model_executor/models/glm4_moe.py

@@ -502,6 +502,8 @@ class Glm4MoeModel(nn.Module):

502

503

504 class Glm4MixtureOfExperts(MixtureOfExperts):

505 + moe_mlp_layers: list[Glm4MoE]

506 +

507 def extract_moe_parameters(self, example_moe: Glm4MoE | None) -> None:

508 if example_moe is None:

509 raise RuntimeError("No Glm4MoE layer found in model.layers.")

vllm/model_executor/models/glm4_moe_lite.py

+10 -10

@@ -74,7 +74,6 @@ from .utils import (

74 PPMissingLayer,

75 get_spec_layer_idx_from_weight_name,

76 is_pp_missing_parameter,

77 - make_empty_intermediate_tensors_factory,

77 make_layers,

78 maybe_prefix,

79 )

@@ -132,6 +131,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module):

131 v_head_dim = getattr(config, "v_head_dim", 0)

132 kv_lora_rank = getattr(config, "kv_lora_rank", 0)

133

134 + attn_cls: Callable[..., Glm4MoeLiteMLAAttention | Glm4MoeLiteAttention]

135 if model_config.use_mla:

136 attn_cls = Glm4MoeLiteMLAAttention

137 else:

@@ -145,7 +145,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module):

145 qk_nope_head_dim=qk_nope_head_dim,

146 qk_rope_head_dim=qk_rope_head_dim,

147 v_head_dim=v_head_dim,

148 - q_lora_rank=config.q_lora_rank if hasattr(config, "q_lora_rank") else None,

148 + q_lora_rank=getattr(config, "q_lora_rank", None),

149 kv_lora_rank=kv_lora_rank,

150 max_position_embeddings=max_position_embeddings,

151 cache_config=cache_config,

@@ -266,9 +266,6 @@ class Glm4MoeLiteModel(nn.Module):

266 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

267 else:

268 self.norm = PPMissingLayer()

269 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

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

271 - )

269

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

271 return self.embed_tokens(input_ids)

@@ -454,7 +451,9 @@ class Glm4MoeLiteModel(nn.Module):

451 # param and delegate to its expert-aware weight_loader

452 # with expert_id.

453 for mapping in expert_params_mapping:

457 - param_name, weight_name, expert_id, shard_id = mapping

454 + param_name, weight_name, mapping_expert_id, mapping_shard_id = (

455 + mapping

456 + )

457 if weight_name not in chunk_name:

458 continue

459

@@ -480,8 +479,8 @@ class Glm4MoeLiteModel(nn.Module):

479 param,

480 weight_to_load,

481 name_mapped,

483 - shard_id=shard_id,

484 - expert_id=expert_id,

482 + shard_id=mapping_shard_id,

483 + expert_id=mapping_expert_id,

484 return_success=True,

485 )

486 if success:

@@ -502,9 +501,10 @@ class Glm4MoeLiteModel(nn.Module):

501 continue

502

503 # Remapping the name of FP8 kv-scale.

505 - name = maybe_remap_kv_scale_name(name, params_dict)

506 - if name is None:

504 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

505 + if remapped_name is None:

506 continue

507 + name = remapped_name

508

509 if is_pp_missing_parameter(name, self):

510 continue

vllm/model_executor/models/glm4_moe_lite_mtp.py

+12 -7

@@ -85,7 +85,9 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module):

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

86 super().__init__()

88 - config = vllm_config.speculative_config.draft_model_config.hf_config

88 + speculative_config = vllm_config.speculative_config

89 + assert speculative_config is not None

90 + config = speculative_config.draft_model_config.hf_config

91 self.config = config

92 quant_config = vllm_config.quant_config

@@ -235,7 +237,7 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):

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

238 return self.model.embed_input_ids(input_ids)

239

238 - def forward(

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

241 self,

242 input_ids: torch.Tensor | None,

243 positions: torch.Tensor,

@@ -370,7 +372,9 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):

372 # with expert_id.

373 is_expert_weight = False

374 for mapping in expert_params_mapping:

373 - param_name, weight_name, expert_id, shard_id = mapping

375 + param_name, weight_name, mapping_expert_id, mapping_shard_id = (

376 + mapping

377 + )

378 if weight_name not in chunk_name:

379 continue

380

@@ -393,8 +397,8 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):

397 param,

398 weight_to_load,

399 name_mapped,

396 - shard_id=shard_id,

397 - expert_id=expert_id,

400 + shard_id=mapping_shard_id,

401 + expert_id=mapping_expert_id,

402 return_success=True,

403 )

404 if success:

@@ -414,9 +418,10 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):

418 if name.endswith(".bias") and name not in params_dict:

419 continue

420

417 - name = maybe_remap_kv_scale_name(name, params_dict)

418 - if name is None:

421 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

422 + if remapped_name is None:

423 continue

424 + name = remapped_name

425

426 # According to DeepSeek-V3 Technical Report, MTP modules

427 # shares embedding layer. We only load the first weights.

vllm/model_executor/models/glm4_moe_mtp.py

@@ -93,6 +93,7 @@ class Glm4MoeMultiTokenPredictorLayer(nn.Module):

93 self.shared_head = SharedHead(

94 config=config, prefix=prefix, quant_config=quant_config

95 )

96 + assert parallel_config is not None

97 self.enable_eplb = parallel_config.enable_eplb

98 self.mtp_block = Glm4MoeDecoderLayer(

99 config=config,

vllm/model_executor/models/glm4v.py

+19 -2

@@ -582,7 +582,13 @@ class GLM4VForCausalLM(

582 for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):

583 offset = mm_feature.mm_position.offset

584 if mm_feature.modality == "image":

585 - t, h, w = mm_feature.data["image_grid_thw"].data.tolist()

585 + feature_data = mm_feature.data

586 + assert feature_data is not None

587 + grid_item = feature_data.get("image_grid_thw")

588 + assert grid_item is not None

589 + grid_data = grid_item.data

590 + assert isinstance(grid_data, torch.Tensor)

591 + t, h, w = grid_data.tolist()

592 assert t == 1, f"Image must have 1 frame, got {t}"

593 yield offset, t, h // spatial_merge_size, w // spatial_merge_size

594 else:

@@ -625,7 +631,18 @@ class GLM4VForCausalLM(

631 mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()

632 return torch.from_numpy(llm_positions), mrope_position_delta

633

628 - embed_input_ids = SupportsMultiModal.embed_input_ids

634 + def embed_input_ids(

635 + self,

636 + input_ids: torch.Tensor,

637 + multimodal_embeddings: MultiModalEmbeddings | None = None,

638 + is_multimodal: torch.Tensor | None = None,

639 + ) -> torch.Tensor:

640 + return SupportsMultiModal.embed_input_ids(

641 + self,

642 + input_ids,

643 + multimodal_embeddings,

644 + is_multimodal=is_multimodal,

645 + )

646

647 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:

648 image_input = self._parse_and_validate_image_input(**kwargs)

vllm/model_executor/models/glm_ocr.py

@@ -326,7 +326,9 @@ class GlmOcrVisionTransformer(Glm4vVisionTransformer):

326 self,

327 x: torch.Tensor,

328 grid_thw: torch.Tensor | list[list[int]],

329 + encoder_metadata: dict[str, torch.Tensor] | None = None,

330 ) -> torch.Tensor:

331 + del encoder_metadata

332 if isinstance(grid_thw, list):

333 grid_thw = torch.tensor(grid_thw, dtype=torch.int32)

334

vllm/model_executor/models/glm_ocr_mtp.py

+7 -4

@@ -58,7 +58,9 @@ class GlmOcrMultiTokenPredictorLayer(nn.Module):

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

59 nn.Module.__init__(self)

61 - config = vllm_config.speculative_config.draft_model_config.hf_config.text_config

61 + speculative_config = vllm_config.speculative_config

62 + assert speculative_config is not None

63 + config = speculative_config.draft_model_config.hf_config.text_config

64 self.config = config

65 quant_config = vllm_config.quant_config

@@ -144,7 +146,7 @@ class GlmOcrMTP(nn.Module, SupportsPP):

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

147 return self.model.embed_input_ids(input_ids)

148

147 - def forward(

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

150 self,

151 input_ids: torch.Tensor,

152 positions: torch.Tensor,

@@ -195,9 +197,10 @@ class GlmOcrMTP(nn.Module, SupportsPP):

197

198 if "scale" in name or "zero_point" in name:

199 # Remapping the name of FP8 kv-scale or zero point.

198 - name = maybe_remap_kv_scale_name(name, params_dict)

199 - if name is None:

200 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

201 + if remapped_name is None:

202 continue

203 + name = remapped_name

204

205 for param_name, weight_name, shard_id in stacked_params_mapping:

206 if weight_name not in name:

vllm/model_executor/models/glmasr.py

+15 -4

@@ -11,7 +11,7 @@ from transformers.models.glmasr import GlmAsrConfig, GlmAsrProcessor

11 from transformers.models.whisper import WhisperFeatureExtractor

13 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

14 -from vllm.config.multimodal import BaseDummyOptions

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

15 from vllm.config.speech_to_text import SpeechToTextParams

16 from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size

17 from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TokensPrompt

@@ -710,6 +710,7 @@ class GlmAsrDummyInputsBuilder(BaseDummyInputsBuilder[GlmAsrProcessingInfo]):

710 sampling_rate = feature_extractor.sampling_rate

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

712 audio_overrides = mm_options.get("audio")

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

714

715 max_audio_len = getattr(

716 self.info.get_hf_processor(), "max_audio_len", DEFAULT_MAX_AUDIO_LEN_S

@@ -793,8 +794,13 @@ class GlmAsrMultiModalProcessor(BaseMultiModalProcessor["GlmAsrProcessingInfo"])

794 )

795 processed_data["feature_attention_mask"] = mask

796

796 - audio = mm_data.get("audio", [])

797 - audio_list = [audio] if audio and not isinstance(audio, list) else audio

797 + audio = mm_data.get("audio")

798 + if audio is None:

799 + audio_list: list[Any] = []

800 + elif isinstance(audio, list):

801 + audio_list = audio

802 + else:

803 + audio_list = [audio]

804 if audio_list:

805 processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

806

@@ -878,6 +884,7 @@ class GlmAsrMultiModalProcessor(BaseMultiModalProcessor["GlmAsrProcessingInfo"])

884 audio_embeds = out_mm_data.get("audio_embeds")

885 if audio_embeds is not None:

886 embed = audio_embeds[item_idx]

887 + assert isinstance(embed, torch.Tensor)

888 num_features = embed.shape[0]

889 else:

890 raise ValueError(

@@ -1118,7 +1125,11 @@ class GlmAsrForConditionalGeneration(

1125 audio_token = cls._get_audio_token(model_config)

1126

1127 if task_type == "translate":

1121 - full_lang_name_to = cls.supported_languages.get(to_language, to_language)

1128 + full_lang_name_to = (

1129 + cls.supported_languages.get(to_language, to_language)

1130 + if to_language is not None

1131 + else ""

1132 + )

1133 user_content = f"{audio_token}translate the speech to {full_lang_name_to}"

1134 elif task_type == "transcribe":

1135 user_content = (

vllm/model_executor/models/gpt_j.py

@@ -227,6 +227,7 @@ class GPTJModel(nn.Module):

227 else:

228 hidden_states = self.embed_input_ids(input_ids)

229 else:

230 + assert intermediate_tensors is not None

231 hidden_states = intermediate_tensors["hidden_states"]

232 for layer in islice(self.h, self.start_layer, self.end_layer):

233 hidden_states = layer(position_ids, hidden_states)

vllm/model_executor/models/gpt_neox.py

@@ -244,6 +244,7 @@ class GPTNeoXModel(nn.Module):

244 else:

245 hidden_states = self.embed_input_ids(input_ids)

246 else:

247 + assert intermediate_tensors is not None

248 hidden_states = intermediate_tensors["hidden_states"]

249 for layer in islice(self.layers, self.start_layer, self.end_layer):

250 hidden_states = layer(position_ids, hidden_states)

vllm/model_executor/models/gpt_oss.py

+62 -21

@@ -79,6 +79,13 @@ _GPT_OSS_STREAMED_EXPERT_SUFFIX_TO_SHARD = {

79 }

82 +def _get_weight_loader(param: torch.Tensor) -> Callable[..., object]:

83 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

84 + if not callable(weight_loader):

85 + raise TypeError("weight_loader must be callable")

86 + return weight_loader

87 +

88 +

89 class OAIAttention(nn.Module):

90 # Override to switch RoPE convention. gpt-oss uses NeoX (chunk halves);

91 # privacy-filter and similar derivatives use GPT-J (interleaved pairs).

@@ -273,6 +280,28 @@ class GptOssRoutedExperts(RoutedExperts):

280

281 self._copy_to_expert(expert_data, loaded_weight)

282

283 + @typing.overload

284 + def weight_loader(

285 + self,

286 + param: torch.nn.Parameter,

287 + loaded_weight: torch.Tensor,

288 + weight_name: str,

289 + shard_id: str,

290 + expert_id: int,

291 + return_success: typing.Literal[False],

292 + ) -> None: ...

293 +

294 + @typing.overload

295 + def weight_loader(

296 + self,

297 + param: torch.nn.Parameter,

298 + loaded_weight: torch.Tensor,

299 + weight_name: str,

300 + shard_id: str,

301 + expert_id: int,

302 + return_success: typing.Literal[True],

303 + ) -> bool: ...

304 +

305 def weight_loader(

306 self,

307 param: torch.nn.Parameter,

@@ -283,14 +312,24 @@ class GptOssRoutedExperts(RoutedExperts):

312 return_success: bool = False,

313 ) -> bool | None:

314 if shard_id not in ("gpt_oss_w13", "gpt_oss_w2"):

286 - return super().weight_loader(

315 + if return_success:

316 + return super().weight_loader(

317 + param,

318 + loaded_weight,

319 + weight_name,

320 + shard_id,

321 + expert_id,

322 + True,

323 + )

324 + super().weight_loader(

325 param,

326 loaded_weight,

327 weight_name,

328 shard_id,

329 expert_id,

292 - return_success,

330 + False,

331 )

332 + return None

333

334 expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)

335 if expert_id == -1:

@@ -518,7 +557,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

557 heads_per_rank: int,

558 head_start: int,

559 weights: Iterable[tuple[str, torch.Tensor]],

521 - stacked_params_mapping: list[tuple[str, ...]],

560 + stacked_params_mapping: list[tuple[str, str, str]],

561 ) -> set[str]:

562 params_dict = dict(self.named_parameters())

563 loaded_params: set[str] = set()

@@ -564,7 +603,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

603 narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end, ...]

604

605 param = params_dict[name]

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

606 + weight_loader = _get_weight_loader(param)

607 weight_loader(

608 param,

609 narrow_weight,

@@ -586,7 +625,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

625 ]

626

627 param = params_dict[name]

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

628 + weight_loader = _get_weight_loader(param)

629 weight_loader(

630 param,

631 narrow_weight,

@@ -612,7 +651,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

651 narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end, ...]

652

653 param = params_dict[name]

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

654 + weight_loader = _get_weight_loader(param)

655 weight_loader(

656 param,

657 narrow_weight,

@@ -635,7 +674,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

674 narrow_weight = weight[..., tp_rank_start // 2 : tp_rank_end // 2]

675

676 param = params_dict[name]

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

677 + weight_loader = _get_weight_loader(param)

678 weight_loader(

679 param,

680 narrow_weight,

@@ -654,7 +693,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

693 narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end]

694

695 param = params_dict[name]

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

696 + weight_loader = _get_weight_loader(param)

697 weight_loader(

698 param,

699 narrow_weight,

@@ -667,7 +706,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

706 elif ".w2_bias" in name:

707 # Handle MLP down projection bias

708 param = params_dict[name]

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

709 + weight_loader = _get_weight_loader(param)

710 if use_ep:

711 weight = weight[ep_rank_start:ep_rank_end, ...]

712 else:

@@ -691,7 +730,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

730 continue

731 name = name.replace(weight_name, param_name)

732 param = params_dict[name]

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

733 + weight_loader = _get_weight_loader(param)

734 if weight_loader == default_weight_loader:

735 weight_loader(param, weight)

736 else:

@@ -702,7 +741,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

741 if name not in params_dict:

742 continue

743 param = params_dict[name]

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

744 + weight_loader = _get_weight_loader(param)

745 weight_loader(param, weight)

746 loaded_params.add(name)

747 return loaded_params

@@ -714,7 +753,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

753 heads_per_rank: int,

754 head_start: int,

755 weights: Iterable[tuple[str, torch.Tensor]],

717 - stacked_params_mapping: list[tuple[str, ...]],

756 + stacked_params_mapping: list[tuple[str, str, str]],

757 ) -> set[str]:

758 params_dict = dict(self.named_parameters())

759 loaded_params: set[str] = set()

@@ -820,6 +859,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

859 and expert_id is not None

860 ):

861 assert loaded_weight.numel() == 1

862 + assert fused_name is not None

863 expert_data = params_dict[fused_name].data[expert_id]

864 expert_data.copy_(loaded_weight)

865 loaded_params.add(fused_name)

@@ -901,6 +941,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

941 # fused gate_up_proj fused on disk, we cannot use the existing

942 # weight loaders without added complexity, so just do the

943 # direct load here.

944 + assert fused_name is not None

945 param = params_dict[fused_name]

946 expert_data = param.data[expert_id]

947 dim1 = sliced_weight.shape[0]

@@ -1062,9 +1103,10 @@ class GptOssModel(nn.Module, EagleModelMixin):

1103

1104 if name.endswith("scale"):

1105 # Remapping the name of FP8 kv-scale.

1065 - name = maybe_remap_kv_scale_name(name, params_dict)

1066 - if name is None:

1106 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

1107 + if remapped_name is None:

1108 continue

1109 + name = remapped_name

1110

1111 param = params_dict[name]

1112 weight_loader = param.weight_loader

@@ -1084,6 +1126,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

1126 if weight_name not in name:

1127 continue

1128

1129 + assert fused_name is not None

1130 param = params_dict[fused_name]

1131 # We should ask the weight loader to return success or not

1132 # here since otherwise we may skip experts with other

@@ -1112,9 +1155,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

1155 if name not in params_dict:

1156 continue

1157 param = params_dict[name]

1115 - weight_loader = getattr(

1116 - param, "weight_loader", default_weight_loader

1117 - )

1158 + weight_loader = _get_weight_loader(param)

1159 weight_loader(param, loaded_weight)

1160

1161 loaded_params.add(name)

@@ -1127,7 +1168,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

1168 heads_per_rank: int,

1169 head_start: int,

1170 weights: Iterable[tuple[str, torch.Tensor]],

1130 - stacked_params_mapping: list[tuple[str, ...]],

1171 + stacked_params_mapping: list[tuple[str, str, str]],

1172 ) -> set[str]:

1173 params_dict = dict(self.named_parameters())

1174 loaded_params: set[str] = set()

@@ -1224,7 +1265,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

1265 continue

1266 name = name.replace(weight_name, param_name)

1267 param = params_dict[name]

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

1268 + weight_loader = _get_weight_loader(param)

1269 if weight_loader == default_weight_loader:

1270 weight_loader(param, weight)

1271 else:

@@ -1235,7 +1276,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

1276 if name not in params_dict:

1277 continue

1278 param = params_dict[name]

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

1279 + weight_loader = _get_weight_loader(param)

1280 weight_loader(param, weight)

1281 loaded_params.add(name)

1282 return loaded_params

@@ -1367,7 +1408,7 @@ class GptOssModel(nn.Module, EagleModelMixin):

1408 class GptOssForCausalLM(

1409 nn.Module, SupportsPP, SupportsEagle, SupportsEagle3, SupportsLoRA

1410 ):

1370 - is_3d_moe_weight: bool = True

1411 + is_3d_moe_weight: typing.ClassVar[bool] = True

1412 packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}

1413

1414 hf_to_vllm_mapper = WeightsMapper(

vllm/model_executor/models/granite4_vision.py

@@ -130,6 +130,7 @@ class WindowQFormerDownsampler(nn.Module):

130

131 self.dropout = nn.Dropout(config.projector_dropout)

132

133 + self.downsampler: SpatialOffsetDownsampler | InterpolateDownsampler

134 if spatial_offset is not None:

135 self.downsampler = SpatialOffsetDownsampler(config, offset=spatial_offset)

136 else:

@@ -843,6 +844,7 @@ class Granite4VisionForConditionalGeneration(

844 # (possibly a chunk slice from the framework's encoder cache).

845 # Concatenate along token dim → (total_mm_tokens, lm_h * num_levels).

846 N, lm_h = inputs_embeds.shape

847 + assert multimodal_embeddings is not None

848 all_packed = torch.cat(

849 [t.to(dtype=inputs_embeds.dtype) for t in multimodal_embeddings],

850 dim=0,

vllm/model_executor/models/granite_speech.py

+13 -3

@@ -34,7 +34,7 @@ from torch import nn

34 from transformers import BatchFeature, PretrainedConfig

36 from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig

37 -from vllm.config.multimodal import BaseDummyOptions

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

38 from vllm.config.speech_to_text import SpeechToTextParams

39 from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt

40 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear

@@ -166,7 +166,9 @@ class GraniteSpeechMultiModalProcessor(

166 def get_replacement(item_idx: int):

167 audios = mm_items.get_items("audio", AudioProcessorItems)

168 audio = audios.get(item_idx)

169 - audio_length = audio.shape[-1]

169 + if audio is None:

170 + raise ValueError(f"Missing audio item {item_idx}")

171 + audio_length = len(audio) if isinstance(audio, list) else audio.shape[-1]

172 num_projector_features = feature_extractor._get_num_audio_features(

173 [audio_length]

174 )[0]

@@ -225,6 +227,7 @@ class GraniteSpeechDummyInputsBuilder(

227 ) -> MultiModalDataDict:

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

229 audio_overrides = mm_options.get("audio")

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

231

232 return {

233 "audio": self._get_dummy_audios(

@@ -659,6 +662,9 @@ class GraniteSpeechForConditionalGeneration(

662 if input_features is None:

663 return None

664

665 + if not isinstance(audio_embed_sizes, torch.Tensor):

666 + raise ValueError("audio_embed_sizes must be a tensor")

667 +

668 # If we have a batch of variable feature length audio clips, we need

669 # to mask the features; usually we would get an input_features_mask

670 # from the processor, but we handle rebuilding it here since

@@ -879,7 +885,11 @@ class GraniteSpeechForConditionalGeneration(

885 audio_tok = cls.get_placeholder_str("audio", 0)

886

887 if task_type == "translate":

882 - full_lang_name_to = cls.supported_languages.get(to_language, to_language)

888 + full_lang_name_to = (

889 + cls.supported_languages.get(to_language, to_language)

890 + if to_language is not None

891 + else ""

892 + )

893 user_prompt = f"{audio_tok}translate the speech to {full_lang_name_to}" # noqa: E501

894 elif task_type == "transcribe":

895 user_prompt = (

vllm/model_executor/models/granitemoehybrid.py

+7 -2

@@ -270,6 +270,7 @@ class GraniteMoeHybridAttention(nn.Module):

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

271 )

272

273 + self.rotary_emb: nn.Module | None

274 if config.position_embedding_type == "rope":

275 self.rotary_emb = get_rope(

276 self.head_dim,

@@ -454,7 +455,9 @@ class GraniteMoeHybridModel(nn.Module):

455 # Skip layers on other devices.

456 if not is_pp_missing_parameter(n, self):

457 param = params_dict[n]

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

458 + weight_loader = getattr(param, "weight_loader", None)

459 + if not callable(weight_loader):

460 + raise TypeError(f"{n} requires a shard-aware weight loader")

461 weight_loader(param, p, shard_id)

462 loaded_params.add(n)

463

@@ -462,7 +465,9 @@ class GraniteMoeHybridModel(nn.Module):

465 if n not in params_dict:

466 return

467 param = params_dict[n]

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

468 + weight_loader = getattr(param, "weight_loader", None)

469 + if not callable(weight_loader):

470 + raise TypeError(f"{n} requires an expert-aware weight loader")

471 weight_loader(param, p, name, shard_id=shard_id, expert_id=expert_id)

472 loaded_params.add(n)

473