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

[Mypy] Fix typing for R/S models (#54157)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

Taneem Ibrahim committed Sep 10, 2026 at 08:43 UTC f9083cb83af44debe290f3753ca2effc292ea1c1

15 files changed +146 -85

tools/pre_commit/mypy.py

-2

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

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

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

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

119 - r"vllm/model_executor/models/[rR]",

120 - r"vllm/model_executor/models/[sS]",

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

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

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

vllm/model_executor/models/radio.py

+26 -8

@@ -137,8 +137,9 @@ class ViTPatchGenerator(nn.Module):

137 if isinstance(max_input_dims, int):

138 max_input_dims = (max_input_dims, max_input_dims)

139

140 - max_input_dims = tuple(

141 - int(math.ceil(d / patch_size) * patch_size) for d in max_input_dims

140 + max_input_dims = (

141 + int(math.ceil(max_input_dims[0] / patch_size) * patch_size),

142 + int(math.ceil(max_input_dims[1] / patch_size) * patch_size),

143 )

144

145 self.cpe_mode = max_input_dims != input_dims

@@ -154,7 +155,10 @@ class ViTPatchGenerator(nn.Module):

155

156 self.num_rows = max_input_dims[0] // patch_size

157 self.num_cols = max_input_dims[1] // patch_size

157 - self.input_dims = tuple(d // patch_size for d in input_dims)

158 + self.input_dims = (

159 + input_dims[0] // patch_size,

160 + input_dims[1] // patch_size,

161 + )

162 self.num_patches = self.num_rows * self.num_cols

163 self.max_input_dims = max_input_dims

164

@@ -369,7 +373,10 @@ class ViTPatchGenerator(nn.Module):

373 if input_size is None:

374 input_dims = self.input_dims

375 else:

372 - input_dims = tuple(d // self.patch_size for d in input_size)

376 + input_dims = (

377 + input_size[0] // self.patch_size,

378 + input_size[1] // self.patch_size,

379 + )

380

381 pos_embed = self._get_pos_embeddings(batch_size, input_dims)

382

@@ -538,8 +545,9 @@ class RadioInternVisionModel(nn.Module):

545

546 self.config = config

547 self.img_size, self.grid_size, self.num_patches = self._init_img_size(

541 - to_2tuple(config.patch_size), config.image_size

548 + config.patch_size, config.image_size

549 )

550 + assert self.img_size is not None

551 max_img_size = int(

552 round(config.cpe_max_size / config.patch_size) * config.patch_size

553 )

@@ -565,11 +573,21 @@ class RadioInternVisionModel(nn.Module):

573 prefix=f"{prefix}.encoder",

574 )

575

568 - def _init_img_size(self, patch_size, img_size: int | tuple[int, int]):

576 + def _init_img_size(

577 + self,

578 + patch_size: input_dim_t,

579 + img_size: input_dim_t | None,

580 + ) -> tuple[tuple[int, int] | None, tuple[int, int] | None, int | None]:

581 if img_size is None:

582 return None, None, None

571 - img_size = to_2tuple(img_size)

572 - grid_size = tuple([s // p for s, p in zip(img_size, patch_size)])

583 + if isinstance(img_size, int):

584 + img_size = (img_size, img_size)

585 + if isinstance(patch_size, int):

586 + patch_size = (patch_size, patch_size)

587 + grid_size = (

588 + img_size[0] // patch_size[0],

589 + img_size[1] // patch_size[1],

590 + )

591 num_patches = grid_size[0] * grid_size[1]

592 return img_size, grid_size, num_patches

593

vllm/model_executor/models/registry.py

+5 -2

@@ -1264,9 +1264,11 @@ class _ModelRegistry:

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

1265 )

1266

1267 + assert issubclass(model_module, transformers.PreTrainedModel)

1268 + transformers_model_cls: type[transformers.PreTrainedModel] = model_module

1269 if not (

1268 - model_module.is_backend_compatible()

1269 - or model_module._can_set_attn_implementation()

1270 + transformers_model_cls.is_backend_compatible()

1271 + or transformers_model_cls._can_set_attn_implementation()

1272 ):

1273 if model_config.model_impl != "transformers":

1274 return None

@@ -1323,6 +1325,7 @@ class _ModelRegistry:

1325 return (model_info, arch)

1326 elif model_config.model_impl == "terratorch":

1327 model_info = self._try_inspect_model_cls("Terratorch")

1328 + assert model_info is not None

1329 return (model_info, "Terratorch")

1330

1331 # Fallback to transformers impl (after resolving convert_type)

vllm/model_executor/models/roberta.py

+6 -3

@@ -139,11 +139,14 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

139 self, vllm_config: VllmConfig, prefix: str = ""

140 ) -> BertModel | BertWithRope:

141 hf_config = vllm_config.model_config.hf_config

142 - kwargs = dict(vllm_config=vllm_config, prefix=prefix)

142 if getattr(hf_config, "position_embedding_type", "absolute") == "absolute":

144 - return BertModel(**kwargs, embedding_class=RobertaEmbedding)

143 + return BertModel(

144 + vllm_config=vllm_config,

145 + prefix=prefix,

146 + embedding_class=RobertaEmbedding,

147 + )

148 else:

146 - return JinaRobertaModel(**kwargs)

149 + return JinaRobertaModel(vllm_config=vllm_config, prefix=prefix)

150

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

152 weights_list = list(weights)

vllm/model_executor/models/sarvam.py

+4 -6

@@ -325,6 +325,7 @@ class SarvamMLAMoE(nn.Module):

325

326 self.score_function = getattr(config, "score_function", "sigmoid")

327 self.num_shared_experts = getattr(config, "num_shared_experts", 1)

328 + self.shared_experts: SarvamMLAMLP | None

329 if self.num_shared_experts > 0:

330 if hasattr(config, "moe_shared_expert_intermediate_size"):

331 shared_int = config.moe_shared_expert_intermediate_size

@@ -590,6 +591,9 @@ class SarvamMLAModel(nn.Module, EagleModelMixin):

591

592

593 class SarvamMixtureOfExperts(MixtureOfExperts):

594 + moe_layers: list[MoERunner]

595 + moe_mlp_layers: list[SarvamMLAMoE]

596 +

597 def extract_moe_parameters(self, example_moe: SarvamMLAMoE | None) -> None:

598 if example_moe is None:

599 raise RuntimeError("No SarvamMLAMoE layer found in model.layers.")

@@ -626,12 +630,6 @@ class SarvamMixtureOfExperts(MixtureOfExperts):

630 if hasattr(fused, "update_expert_map"):

631 fused.update_expert_map()

632

629 - def set_eplb_state(self, eplb_state) -> None:

630 - self.eplb_state = eplb_state

631 - for moe in self.moe_layers:

632 - if hasattr(moe, "set_eplb_state"):

633 - moe.set_eplb_state(eplb_state)

634 -

633

634 class SarvamMLAForCausalLM(

635 nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3, SarvamMixtureOfExperts

vllm/model_executor/models/siglip.py

+6 -4

@@ -127,13 +127,15 @@ class SiglipProcessingInfo(BaseProcessingInfo):

127

128 pooler_config = self.ctx.model_config.pooler_config

129 assert pooler_config is not None

130 + pooling_type = pooler_config.seq_pooling_type

131 + assert pooling_type is not None

132

133 return get_num_selected_vision_tokens(

134 vision_encoder_info.get_num_image_tokens(

135 image_width=image_width,

136 image_height=image_height,

137 ),

136 - _get_vision_feature_select_strategy(pooler_config.seq_pooling_type),

138 + _get_vision_feature_select_strategy(pooling_type),

139 )

140

141 def get_image_size_with_most_features(self) -> ImageSize:

@@ -1021,9 +1023,9 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

1023 feature_select_strategy: VisionFeatureSelectStrategy | None = None,

1024 ) -> torch.Tensor:

1025 if feature_select_strategy is None:

1024 - feature_select_strategy = _get_vision_feature_select_strategy(

1025 - self.pooler_config.seq_pooling_type

1026 - )

1026 + pooling_type = self.pooler_config.seq_pooling_type

1027 + assert pooling_type is not None

1028 + feature_select_strategy = _get_vision_feature_select_strategy(pooling_type)

1029

1030 pooled_output = self.vision_model(

1031 pixel_values=pixel_values,

vllm/model_executor/models/siglip2navit.py

@@ -248,6 +248,7 @@ class Siglip2Attention(nn.Module):

248 values = values.view(seq_length, self.num_heads_per_partition, self.head_dim)

249

250 if self.use_rope:

251 + assert position_embeddings is not None

252 cos, sin = position_embeddings

253 queries, keys = apply_rotary_pos_emb(

254 queries.unsqueeze(0),

vllm/model_executor/models/skyworkr1v.py

+14 -7

@@ -29,7 +29,12 @@ from vllm.transformers_utils.processors.internvl import (

29 )

30 from vllm.utils.tensor_schema import TensorSchema, TensorShape

32 -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

32 +from .interfaces import (

33 + MultiModalEmbeddings,

34 + SupportsMultiModal,

35 + SupportsPP,

36 + supports_pp,

37 +)

38 from .internvl import (

39 BaseInternVLDummyInputsBuilder,

40 BaseInternVLMultiModalProcessor,

@@ -156,7 +161,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

161

162 config = vllm_config.model_config.hf_config

163 quant_config = vllm_config.quant_config

159 - multimodal_config = vllm_config.model_config.multimodal_config

164 + multimodal_config = vllm_config.model_config.get_multimodal_config()

165

166 self.config = config

167 self.multimodal_config = multimodal_config

@@ -182,20 +187,22 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

187 )

188

189 with self._mark_language_model(vllm_config):

185 - self.language_model = init_vllm_registered_model(

190 + language_model = init_vllm_registered_model(

191 vllm_config=vllm_config,

192 hf_config=config.text_config,

193 prefix=maybe_prefix(prefix, "language_model"),

194 )

195 + self.language_model = language_model

196 + assert supports_pp(language_model)

197

191 - self.img_context_token_id = None

198 + self.img_context_token_id: int | None = None

199 self.visual_token_mask = None

200 self.make_empty_intermediate_tensors = (

194 - self.language_model.make_empty_intermediate_tensors

201 + language_model.make_empty_intermediate_tensors

202 )

203

204 def _patch_quant_config(

198 - self, config: PretrainedConfig, quant_config: QuantizationConfig

205 + self, config: PretrainedConfig, quant_config: QuantizationConfig | None

206 ):

207 # the awq models from OpenGVLab missing `modules_to_not_convert`

208 # patch the quant_config to add `modules_to_not_convert` back

@@ -232,7 +239,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

239 def _init_mlp1(

240 self,

241 config: PretrainedConfig,

235 - quant_config: QuantizationConfig,

242 + quant_config: QuantizationConfig | None,

243 prefix: str = "",

244 ) -> nn.Module:

245 vit_hidden_size = config.vision_config.hidden_size

vllm/model_executor/models/solar.py

@@ -318,9 +318,11 @@ class SolarModel(nn.Module):

318 for i in range(self.start_layer, self.end_layer):

319 if i in self.config.bskcn_1:

320 bskcn_h_1 = hidden_states.clone()

321 + assert residual is not None

322 bskcn_r_1 = residual.clone()

323 if i in self.config.bskcn_2:

324 bskcn_h_2 = hidden_states.clone()

325 + assert residual is not None

326 bskcn_r_2 = residual.clone()

327 if i in self.config.bskcn_3:

328 hidden_states = bskcn_h_1 * bskcn_tv + hidden_states * (1 - bskcn_tv)

vllm/model_executor/models/step3_text.py

+7 -7

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

10 from torch import nn

12 from vllm.compilation.decorators import support_torch_compile

13 -from vllm.config import CacheConfig, ModelConfig, VllmConfig

13 +from vllm.config import CacheConfig, VllmConfig

14 from vllm.distributed import (

15 get_pp_group,

16 get_tensor_model_parallel_world_size,

@@ -55,7 +55,7 @@ logger = init_logger(__name__)

55 class FusedMoEBlock(nn.Module):

56 def __init__(

57 self,

58 - config: ModelConfig,

58 + config: Step3TextConfig,

59 quant_config: QuantizationConfig | None = None,

60 prefix: str = "",

61 ):

@@ -384,7 +384,7 @@ class Step3TextModel(nn.Module):

384 return hidden_states

385

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

387 - qkv_params_mapping = [

387 + qkv_params_mapping: list[tuple[str, str, float, float]] = [

388 # (param_name, shard_name, relative_start_idx, relative_end_idx)

389 (

390 ".qkv_proj",

@@ -410,7 +410,7 @@ class Step3TextModel(nn.Module):

410 / (self.config.share_q_dim + self.config.head_dim * 2),

411 ),

412 ]

413 - stacked_params_mapping = [

413 + stacked_params_mapping: list[tuple[str, str, int]] = [

414 # (param_name, shard_name, shard_id)

415 (".gate_up_proj", ".gate_proj", 0),

416 (".gate_up_proj", ".up_proj", 1),

@@ -421,7 +421,7 @@ class Step3TextModel(nn.Module):

421 "base_layer." if any(".base_layer." in name for name in params_dict) else ""

422 )

423

424 - expert_params_mapping = [

424 + expert_params_mapping: list[tuple[str, str, str]] = [

425 (

426 f".moe.experts.routed_experts.{base_layer}w13_weight",

427 ".moe.gate_proj.weight",

@@ -460,7 +460,7 @@ class Step3TextModel(nn.Module):

460 break

461 else:

462 for mapping in expert_params_mapping:

463 - param_name, weight_name, shard_id = mapping

463 + param_name, weight_name, expert_shard_id = mapping

464 if weight_name not in name:

465 continue

466 name = name.replace(weight_name, param_name)

@@ -480,7 +480,7 @@ class Step3TextModel(nn.Module):

480 param,

481 loaded_weight_expert,

482 name,

483 - shard_id=shard_id,

483 + shard_id=expert_shard_id,

484 expert_id=expert_id,

485 )

486 loaded_params.add(name)

vllm/model_executor/models/step3_vl.py

+19 -7

@@ -51,6 +51,7 @@ from .interfaces import (

51 SupportsEncoderCudaGraph,

52 SupportsMultiModal,

53 SupportsPP,

54 + supports_pp,

55 )

56 from .utils import (

57 AutoWeightsLoader,

@@ -105,10 +106,10 @@ class Step3VLProcessingInfo(BaseProcessingInfo):

106

107 return Step3VLImageProcessor(**kwargs)

108

108 - def get_hf_processor(self) -> Step3VLProcessor:

109 + def get_hf_processor(self, **kwargs: object) -> Step3VLProcessor:

110 return Step3VLProcessor(

111 tokenizer=self.get_tokenizer(),

111 - image_processor=self.get_image_processor(),

112 + image_processor=self.get_image_processor(**kwargs),

113 )

114

115 def get_supported_mm_limits(self) -> Mapping[str, int | None]:

@@ -169,8 +170,11 @@ class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo])

170

171 def get_replacement_step1o(item_idx: int):

172 out_item = out_mm_kwargs["image"][item_idx]

172 - num_patches = int(out_item["num_patches"].data)

173 + num_patches_data = out_item["num_patches"].data

174 patch_newline_mask = out_item["patch_newline_mask"].data

175 + assert isinstance(num_patches_data, torch.Tensor)

176 + assert isinstance(patch_newline_mask, torch.Tensor)

177 + num_patches = int(num_patches_data.item())

178 image_repl_ids = hf_processor.get_image_repl_feature_ids(

179 1, num_patches, patch_newline_mask.tolist()

180 )

@@ -514,7 +518,7 @@ class Step3VLForConditionalGeneration(

518 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

519 super().__init__()

520 config = vllm_config.model_config.hf_config

517 - multimodal_config = vllm_config.model_config.multimodal_config

521 + multimodal_config = vllm_config.model_config.get_multimodal_config()

522

523 self.config = config

524 self.model_config = vllm_config.model_config

@@ -560,14 +564,16 @@ class Step3VLForConditionalGeneration(

564 )

565

566 with self._mark_language_model(vllm_config):

563 - self.language_model = init_vllm_registered_model(

567 + language_model = init_vllm_registered_model(

568 vllm_config=vllm_config,

569 hf_config=config.text_config,

570 prefix=maybe_prefix(prefix, "language_model"),

571 )

572 + self.language_model = language_model

573 + assert supports_pp(language_model)

574

575 self.make_empty_intermediate_tensors = (

570 - self.language_model.make_empty_intermediate_tensors

576 + language_model.make_empty_intermediate_tensors

577 )

578

579 @property

@@ -626,6 +632,9 @@ class Step3VLForConditionalGeneration(

632 return None

633

634 if pixel_values is not None and patch_pixel_values is not None:

635 + assert isinstance(pixel_values, torch.Tensor)

636 + assert isinstance(patch_pixel_values, torch.Tensor)

637 + assert isinstance(num_patches, torch.Tensor)

638 return Step3VLImagePixelInputs(

639 type="pixel_values",

640 pixel_values=pixel_values.to(self.dtype),

@@ -634,6 +643,7 @@ class Step3VLForConditionalGeneration(

643 )

644

645 if image_embeds is not None:

646 + assert isinstance(image_embeds, torch.Tensor)

647 return Step3VLImageEmbeddingInputs(

648 type="image_embeds",

649 data=image_embeds.to(self.dtype),

@@ -657,7 +667,7 @@ class Step3VLForConditionalGeneration(

667

668 def _process_image_input(

669 self, image_input: Step3VLImageInputs

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

670 + ) -> list[torch.Tensor]:

671 if image_input["type"] == "image_embeds":

672 image_features = image_input["data"]

673 return [

@@ -686,6 +696,7 @@ class Step3VLForConditionalGeneration(

696 for i, num_patch in enumerate(num_patches_list):

697 cur_feature = []

698 if num_patch > 0:

699 + assert patch_image_features is not None

700 patch_slice = patch_image_features[

701 cur_patch_idx : cur_patch_idx + num_patch

702 ]

@@ -765,6 +776,7 @@ class Step3VLForConditionalGeneration(

776 from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

777

778 num_patches = mm_kwargs.get("num_patches")

779 + assert isinstance(num_patches, torch.Tensor)

780

781 img_grid = (

782 self.config.vision_config.image_size // self.config.vision_config.patch_size

vllm/model_executor/models/step3p5.py

+39 -29

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

11 from torch.nn.parameter import Parameter

13 from vllm.compilation.decorators import support_torch_compile

14 -from vllm.config import CacheConfig, ModelConfig, VllmConfig

14 +from vllm.config import CacheConfig, VllmConfig

15 from vllm.distributed import (

16 get_dp_group,

17 get_ep_group,

@@ -47,6 +47,11 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader

47 from vllm.sequence import IntermediateTensors

48 from vllm.v1.attention.backend import AttentionType

50 +if typing.TYPE_CHECKING:

51 + from vllm.transformers_utils.configs.step3p5 import Step3p5Config

52 +else:

53 + from transformers import PretrainedConfig as Step3p5Config

54 +

55 from .interfaces import MixtureOfExperts, SupportsPP

56 from .utils import (

57 AutoWeightsLoader,

@@ -79,7 +84,7 @@ class FP32ReplicatedLinear(ReplicatedLinear):

84 class Step3p5MLP(nn.Module):

85 def __init__(

86 self,

82 - config: ModelConfig,

87 + config: Step3p5Config,

88 hidden_size: int,

89 intermediate_size: int,

90 hidden_act: str,

@@ -147,9 +152,9 @@ class Step3p5Attention(nn.Module):

152 # Step3p5 specific args

153 sliding_window: int | None = None,

154 use_head_wise_attn_gate: bool = False,

150 - layer_types: list = None,

151 - use_rope_layers: list = None,

152 - yarn_only_types: list = None,

155 + layer_types: list[str] | None = None,

156 + use_rope_layers: list[bool] | None = None,

157 + yarn_only_types: list[str] | None = None,

158 swa_num_attention_heads: int | None = None,

159 partial_rotary_factor: float = 1.0,

160 ):

@@ -162,8 +167,10 @@ class Step3p5Attention(nn.Module):

167 enable_sliding_window = layer_types[self.layer_idx] == "sliding_attention"

168 else:

169 enable_sliding_window = self.layer_idx % 2 == 0

165 - if yarn_only_types and layer_types[self.layer_idx] not in yarn_only_types:

166 - rope_scaling = None

170 + if yarn_only_types:

171 + assert layer_types is not None

172 + if layer_types[self.layer_idx] not in yarn_only_types:

173 + rope_scaling = None

174

175 if sliding_window is not None and enable_sliding_window:

176 sliding_window = sliding_window

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

311

312 self.ep_size = get_ep_group().device_group.size()

313 config = vllm_config.model_config.hf_config

314 + assert isinstance(config, Step3p5Config)

315 quant_config = vllm_config.quant_config

316 parallel_config = vllm_config.parallel_config

317

@@ -349,6 +357,7 @@ class FusedMoEBlock(nn.Module):

357 else None

358 )

359 if swiglu_limit not in (None, 0):

360 + assert swiglu_limit is not None

361 swiglu_limit = float(swiglu_limit)

362 assert swiglu_limit == 7.0, (

363 "Swiglu limit in fused moe block only support 7.0 now."

@@ -408,6 +417,7 @@ class Step3p5DecoderLayer(nn.Module):

417 ) -> None:

418 super().__init__()

419 config = vllm_config.model_config.hf_config

420 + assert isinstance(config, Step3p5Config)

421 self.hidden_size = config.hidden_size

422 layer_idx = extract_layer_index(prefix)

423 self.layer_idx = layer_idx

@@ -419,20 +429,18 @@ class Step3p5DecoderLayer(nn.Module):

429 num_attention_heads = None

430 num_attention_groups = None

431 head_dim = None

432 + attention_other_setting = config.attention_other_setting

433 + layer_types = config.layer_types

434 if (

423 - getattr(config, "attention_other_setting", None)

424 - and getattr(config, "layer_types", [])

425 - and config.layer_types[layer_idx]

426 - == config.attention_other_setting["attention_type"]

435 + attention_other_setting

436 + and layer_types

437 + and layer_types[layer_idx] == attention_other_setting["attention_type"]

438 ):

428 - num_attention_heads = config.attention_other_setting[

429 - "num_attention_heads"

430 - ]

431 - num_attention_groups = config.attention_other_setting[

432 - "num_attention_groups"

433 - ]

434 - head_dim = config.attention_other_setting["head_dim"]

439 + num_attention_heads = attention_other_setting["num_attention_heads"]

440 + num_attention_groups = attention_other_setting["num_attention_groups"]

441 + head_dim = attention_other_setting["head_dim"]

442 partial_rotary_factors = getattr(config, "partial_rotary_factors", [])

443 + assert config.max_position_embeddings is not None

444 self.self_attn = Step3p5Attention(

445 hidden_size=self.hidden_size,

446 num_heads=num_attention_heads

@@ -539,6 +547,7 @@ class Step3p5Model(nn.Module):

547

548 self.vllm_config = vllm_config

549 config = vllm_config.model_config.hf_config

550 + assert isinstance(config, Step3p5Config)

551 self.vocab_size = config.vocab_size

552 self.config = config

553

@@ -605,8 +614,8 @@ class Step3p5Model(nn.Module):

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

615 config = self.config

616 assert config.num_attention_groups > 1, "Only support GQA"

608 - qkv_params_mapping = []

609 - stacked_params_mapping = [

617 + qkv_params_mapping: list[tuple[str, str, float, float]] = []

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

619 # (param_name, shard_name, shard_id)

620 ("qkv_proj", "q_proj", "q"),

621 ("qkv_proj", "k_proj", "k"),

@@ -622,7 +631,7 @@ class Step3p5Model(nn.Module):

631 )

632

633 # Old packed 3D format: .moe.gate_proj.weight [num_experts, out, in]

625 - expert_params_mapping = [

634 + expert_params_mapping: list[tuple[str, str, str]] = [

635 (

636 f".moe.experts.routed_experts.{base_layer}w13_weight",

637 ".moe.gate_proj.weight",

@@ -758,16 +767,16 @@ class Step3p5Model(nn.Module):

767 # Not an expert proj — use default loader

768 # (e.g. share_expert weights if they matched)

769 param = params_dict[local_name]

761 - weight_loader = getattr(

770 + fallback_loader = getattr(

771 param,

772 "weight_loader",

773 default_weight_loader,

774 )

766 - weight_loader(param, loaded_weight)

775 + fallback_loader(param, loaded_weight)

776 loaded_params.add(local_name)

777 continue

778

770 - for param_name, weight_name, shard_id in stacked_params_mapping:

779 + for param_name, weight_name, stacked_shard_id in stacked_params_mapping:

780 if weight_name not in local_name:

781 continue

782 if any(

@@ -782,11 +791,11 @@ class Step3p5Model(nn.Module):

791 continue

792 param = params_dict[replaced_name]

793 weight_loader = param.weight_loader

785 - weight_loader(param, loaded_weight, shard_id)

794 + weight_loader(param, loaded_weight, stacked_shard_id)

795 loaded_params.add(replaced_name)

796 break

797 else:

789 - for param_name, weight_name, shard_id in expert_params_mapping:

798 + for param_name, weight_name, expert_shard_id in expert_params_mapping:

799 if weight_name not in local_name:

800 continue

801 replaced_name = local_name.replace(weight_name, param_name)

@@ -823,7 +832,7 @@ class Step3p5Model(nn.Module):

832 param,

833 loaded_weight_expert,

834 replaced_name,

826 - shard_id=shard_id,

835 + shard_id=expert_shard_id,

836 expert_id=expert_id,

837 )

838 loaded_params.add(replaced_name)

@@ -861,10 +870,10 @@ class Step3p5Model(nn.Module):

870 if local_name not in params_dict:

871 continue

872 param = params_dict[local_name]

864 - weight_loader = getattr(

873 + fallback_loader = getattr(

874 param, "weight_loader", default_weight_loader

875 )

867 - weight_loader(param, loaded_weight)

876 + fallback_loader(param, loaded_weight)

877 loaded_params.add(local_name)

878 return loaded_params

879

@@ -918,6 +927,7 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

927 self.moe_layers.append(layer.moe.experts)

928

929 assert len(self.moe_layers) > 0, "No MoE layers found in the model."

930 + assert example_layer is not None

931 self.num_moe_layers = len(self.moe_layers)

932 self.num_expert_groups = 1

933 self.num_shared_experts = 0

vllm/model_executor/models/step3p5_mtp.py

+2 -1

@@ -305,7 +305,7 @@ class Step3p5MTP(nn.Module):

305 )

306 return loaded_params

307

308 - def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:

308 + def _rewrite_spec_layer_name(self, spec_layer: int | None, name: str) -> str:

309 """

310 Rewrite the weight name to match the format of the original model.

311 Add .mtp_block for modules in transformer layer block for spec layer

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

324 break

325 if not spec_layer_weight:

326 # treat rest weights as weights for transformer layer block

327 + assert spec_layer is not None

328 name = name.replace(

329 f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."

330 )

vllm/model_executor/models/step3p7.py

+6 -3

@@ -9,6 +9,7 @@ from vllm.logger import init_logger

9 from vllm.model_executor.layers.activation import get_act_fn

10 from vllm.model_executor.layers.linear import ColumnParallelLinear

12 +from .interfaces import supports_pp

13 from .step3_vl import Step3VLForConditionalGeneration

14 from .step_vl import PerceptionEncoder

15 from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix

@@ -41,7 +42,7 @@ class Step3p7ForConditionalGeneration(Step3VLForConditionalGeneration):

42 super(Step3VLForConditionalGeneration, self).__init__()

44 config = vllm_config.model_config.hf_config

44 - multimodal_config = vllm_config.model_config.multimodal_config

45 + multimodal_config = vllm_config.model_config.get_multimodal_config()

46 quant_config = vllm_config.quant_config

48 self.config = config

@@ -66,14 +67,16 @@ class Step3p7ForConditionalGeneration(Step3VLForConditionalGeneration):

67 )

69 with self._mark_language_model(vllm_config):

69 - self.language_model = init_vllm_registered_model(

70 + language_model = init_vllm_registered_model(

71 vllm_config=vllm_config,

72 hf_config=config.text_config,

73 prefix=maybe_prefix(prefix, "language_model"),

74 )

75 + self.language_model = language_model

76 + assert supports_pp(language_model)

78 self.make_empty_intermediate_tensors = (

76 - self.language_model.make_empty_intermediate_tensors

79 + language_model.make_empty_intermediate_tensors

80 )

82 def _get_vision_model_output(

vllm/model_executor/models/step_vl.py

+9 -6

@@ -22,6 +22,7 @@ from vllm.model_executor.layers.linear import (

22 )

23 from vllm.model_executor.layers.quantization import QuantizationConfig

25 +from .interfaces import supports_pp

26 from .step3_vl import Step3VLForConditionalGeneration

27 from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix

28 from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

@@ -148,7 +149,7 @@ class PerceptionEncoderMLP(nn.Module):

149 self,

150 input_dim: int,

151 hidden_dim: int,

151 - act_layer: Callable[[], nn.Module],

152 + act_layer: Callable[[torch.Tensor], torch.Tensor],

153 quant_config: QuantizationConfig | None = None,

154 prefix: str = "",

155 ):

@@ -257,7 +258,7 @@ class PerceptionEncoderVisionBlock(nn.Module):

258 max_grid_height: int,

259 max_grid_width: int,

260 mlp_ratio: float = 4.0,

260 - ls_init_value: float = None,

261 + ls_init_value: float | None = None,

262 act_layer: Callable = nn.GELU,

263 norm_layer: Callable = nn.LayerNorm,

264 use_cls_token: bool = False,

@@ -310,7 +311,7 @@ class PerceptionEncoderVisionTransformer(nn.Module):

311 max_grid_height: int,

312 max_grid_width: int,

313 mlp_ratio: float = 4.0,

313 - ls_init_value: float = None,

314 + ls_init_value: float | None = None,

315 act_layer: Callable = nn.GELU,

316 norm_layer: Callable = nn.LayerNorm,

317 use_cls_token: bool = False,

@@ -496,7 +497,7 @@ class StepVLForConditionalGeneration(Step3VLForConditionalGeneration):

497 super(Step3VLForConditionalGeneration, self).__init__()

498

499 config = vllm_config.model_config.hf_config

499 - multimodal_config = vllm_config.model_config.multimodal_config

500 + multimodal_config = vllm_config.model_config.get_multimodal_config()

501 quant_config = vllm_config.quant_config

502

503 self.config = config

@@ -522,14 +523,16 @@ class StepVLForConditionalGeneration(Step3VLForConditionalGeneration):

523 )

524

525 with self._mark_language_model(vllm_config):

525 - self.language_model = init_vllm_registered_model(

526 + language_model = init_vllm_registered_model(

527 vllm_config=vllm_config,

528 hf_config=config.text_config,

529 prefix=maybe_prefix(prefix, "language_model"),

530 )

531 + self.language_model = language_model

532 + assert supports_pp(language_model)

533

534 self.make_empty_intermediate_tensors = (

532 - self.language_model.make_empty_intermediate_tensors

535 + language_model.make_empty_intermediate_tensors

536 )

537

538 def _get_vision_model_output(