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

[Mypy] Fix mypy typing for L models (#54177)

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

Taneem Ibrahim committed Sep 4, 2026 at 06:50 UTC c615b1fd67f5b149c75aea592de8e563cb0c2da9

19 files changed +255 -89

tools/pre_commit/mypy.py

-1

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

113

114 EXCLUDE = [

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

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

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

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

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

vllm/model_executor/models/laguna.py

@@ -317,6 +317,7 @@ class LagunaAttention(nn.Module):

317 # config.gating may be:

318 # - True / "per-element": one gate per (head, head_dim) channel

319 # - "per-head": one gate per head, broadcast across head_dim

320 + self.g_proj: ColumnParallelLinear | None

321 if self.gating:

322 # v5 LagunaConfig uses ``gating=True`` for per-head; older configs

323 # used ``"per-head"``. Accept both. ``"per-element"`` (or legacy

vllm/model_executor/models/laguna_dflash.py

+11 -15

@@ -69,7 +69,9 @@ class DFlashLagunaModel(DFlashQwen3Model, EagleModelMixin):

69 prefix: str = "",

70 ) -> None:

71 nn.Module.__init__(self)

72 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

72 + speculative_config = vllm_config.speculative_config

73 + assert speculative_config is not None

74 + self.config = speculative_config.draft_model_config.hf_config

75 self.vocab_size = self.config.vocab_size

76 self.quant_config = get_draft_quant_config(vllm_config)

@@ -216,24 +218,16 @@ class DFlashLagunaModel(DFlashQwen3Model, EagleModelMixin):

218 return all_k_normed

219

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

221 + if self.quant_config is not None:

222 + weights = self.quant_config.get_cache_scale_mapper().apply(weights)

223 params_dict = dict(self.named_parameters())

224 loaded_params: set[str] = set()

225 for name, loaded_weight in weights:

222 - if self.quant_config is not None and (

223 - scale_name := self.quant_config.get_cache_scale(name)

224 - ):

225 - param = params_dict[scale_name]

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

227 - loaded_weight = (

228 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]

229 - )

230 - weight_loader(param, loaded_weight)

231 - loaded_params.add(scale_name)

232 - continue

226 if "scale" in name:

234 - name = maybe_remap_kv_scale_name(name, params_dict)

235 - if name is None:

227 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

228 + if remapped_name is None:

229 continue

230 + name = remapped_name

231 param = params_dict[name]

232 weight_loader = getattr(param, "weight_loader", default_weight_loader)

233 weight_loader(param, loaded_weight)

@@ -244,7 +238,9 @@ class DFlashLagunaModel(DFlashQwen3Model, EagleModelMixin):

238 class DFlashLagunaForCausalLM(nn.Module, SupportsEagle3):

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

240 nn.Module.__init__(self)

247 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

241 + speculative_config = vllm_config.speculative_config

242 + assert speculative_config is not None

243 + self.config = speculative_config.draft_model_config.hf_config

244 if getattr(self.config, "draft_vocab_size", None) is None:

245 raise ValueError("Laguna DFlash config requires `draft_vocab_size`.")

246 self.has_own_embed_tokens = False

vllm/model_executor/models/lfm2_vl.py

+21

@@ -3,6 +3,7 @@

4 import itertools

5 import math

6 +import typing

7 from collections.abc import Iterable, Mapping, Sequence

8 from typing import Annotated, Any, Literal

@@ -17,6 +18,7 @@ from transformers.models.lfm2_vl.image_processing_lfm2_vl_fast import (

18 find_closest_aspect_ratio,

19 round_by_factor,

20 )

21 +from typing_extensions import Buffer

23 from vllm.config import VllmConfig

24 from vllm.config.multimodal import BaseDummyOptions

@@ -229,6 +231,13 @@ class Lfm2VLProcessingInfo(BaseProcessingInfo):

231 "max_image_tokens", image_processor.max_image_tokens

232 )

233 tile_size = mm_kwargs.get("tile_size", image_processor.tile_size)

234 + assert isinstance(downsample_factor, int)

235 + assert isinstance(encoder_patch_size, int)

236 + assert isinstance(max_pixels_tolerance, int | float)

237 + assert isinstance(min_tiles, int)

238 + assert isinstance(max_tiles, int)

239 + assert isinstance(max_image_tokens, int)

240 + assert isinstance(tile_size, int)

241

242 do_image_splitting = not min_tiles == max_tiles == 1

243 is_image_large = self._is_image_too_large(

@@ -337,6 +346,9 @@ class Lfm2VLProcessingInfo(BaseProcessingInfo):

346 "encoder_patch_size", image_processor.encoder_patch_size

347 )

348 tile_size = mm_kwargs.get("tile_size", image_processor.tile_size)

349 + assert isinstance(downsample_factor, int)

350 + assert isinstance(encoder_patch_size, int)

351 + assert isinstance(tile_size, int)

352

353 thumbnail_height_patches = int(spatial_shapes[-1][0].item())

354 thumbnail_width_patches = int(spatial_shapes[-1][1].item())

@@ -672,6 +684,7 @@ class Lfm2VLForConditionalGeneration(

684 super().__init__()

685 config: Lfm2VlConfig = vllm_config.model_config.hf_config

686 multimodal_config = vllm_config.model_config.multimodal_config

687 + assert multimodal_config is not None

688 vision_config = config.vision_config

689 quant_config = vllm_config.quant_config

690

@@ -915,6 +928,14 @@ class Lfm2VLForConditionalGeneration(

928 "min_image_tokens",

929 getattr(self.config, "min_image_tokens", None) or 64,

930 )

931 + if not isinstance(

932 + value,

933 + (str, Buffer, typing.SupportsInt, typing.SupportsIndex),

934 + ):

935 + raise TypeError(

936 + "int() argument must be a string, a bytes-like object "

937 + f"or a real number, not '{type(value).__name__}'"

938 + )

939 return max(1, int(value))

940

941 def _get_lfm2vl_item_tile_slices(

vllm/model_executor/models/lightonocr.py

+1 -1

@@ -49,7 +49,7 @@ class LightOnOCRProcessingInfo(Mistral3ProcessingInfo):

49 return PixtralHFEncoderInfo(self.get_hf_config())

52 -class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessingInfo]):

52 +class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]):

53 def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:

54 return self.dummy_inputs.get_dummy_text(mm_counts)

vllm/model_executor/models/llama.py

+15 -2

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

27 from collections.abc import Iterable

28 from itertools import islice

29 +from typing import TYPE_CHECKING

31 import torch

32 from torch import nn

@@ -544,13 +545,25 @@ class LlamaForCausalLM(

545 return loader.load_weights(weights)

546

547

547 -class LlamaBidirectionalForSequenceClassification(as_seq_cls_model(LlamaForCausalLM)):

548 +if TYPE_CHECKING:

549 + _LlamaBidirectionalForSequenceClassificationBase = LlamaForCausalLM

550 + _LlamaBidirectionalModelBase = LlamaForCausalLM

551 +else:

552 + _LlamaBidirectionalForSequenceClassificationBase = as_seq_cls_model(

553 + LlamaForCausalLM

554 + )

555 + _LlamaBidirectionalModelBase = as_embedding_model(LlamaForCausalLM)

556 +

557 +

558 +class LlamaBidirectionalForSequenceClassification(

559 + _LlamaBidirectionalForSequenceClassificationBase

560 +):

561 # This class sets the correct attention type and pooling type

562 # through LlamaBidirectionalConfig.

563 pass

564

565

553 -class LlamaBidirectionalModel(as_embedding_model(LlamaForCausalLM)):

566 +class LlamaBidirectionalModel(_LlamaBidirectionalModelBase):

567 # This class sets the correct attention type and pooling type

568 # through LlamaBidirectionalConfig.

569 pass

vllm/model_executor/models/llama4.py

+5 -3

@@ -18,7 +18,7 @@

18 # limitations under the License.

19 """Inference-only LLaMA model compatible with HuggingFace weights."""

21 -from collections.abc import Iterable

21 +from collections.abc import Callable, Iterable

23 import torch

24 from torch import nn

@@ -576,6 +576,7 @@ class Llama4Model(LlamaModel):

576 params_dict = dict(self.named_parameters())

577 # The module parameters that have been loaded.

578 loaded_params: set[str] = set()

579 + weight_loader: Callable[..., object]

580

581 # Iterate over all the weights and load them into module parameters.

582 for name, loaded_weight in weights:

@@ -612,9 +613,10 @@ class Llama4Model(LlamaModel):

613 # quant config's `get_cache_scale_mapper` does not cover

614 # (idempotent for names already renamed by the mapper).

615 if name.endswith("scale"):

615 - name = maybe_remap_kv_scale_name(name, params_dict)

616 - if name is None:

616 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

617 + if remapped_name is None:

618 continue

619 + name = remapped_name

620

621 # Load the weight into the module parameter with corresponding

622 # shard id and exit the for loop and the else block.

vllm/model_executor/models/llama4_eagle.py

+22 -4

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

17 # limitations under the License.

19 from collections.abc import Iterable

20 +from typing import TYPE_CHECKING

22 import torch

23 import torch.nn as nn

@@ -46,6 +47,19 @@ from .utils import (

47 logger = init_logger(__name__)

50 +if TYPE_CHECKING:

51 +

52 + class _EagleLlama4ForCausalLMBase(nn.Module):

53 + def set_moe_parameters(self) -> None: ...

54 +

55 + def permute_qk_weight_for_rotary(

56 + self, name: str, loaded_weight: torch.Tensor

57 + ) -> tuple[str, torch.Tensor]: ...

58 +

59 +else:

60 + _EagleLlama4ForCausalLMBase = Llama4ForCausalLM

61 +

62 +

63 @support_torch_compile

64 class LlamaModel(nn.Module):

65 hf_to_vllm_mapper = WeightsMapper(

@@ -68,7 +82,9 @@ class LlamaModel(nn.Module):

82 quant_config: QuantizationConfig | None = None,

83 ) -> None:

84 super().__init__()

71 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

85 + speculative_config = vllm_config.speculative_config

86 + assert speculative_config is not None

87 + self.config = speculative_config.draft_model_config.hf_config

88 self.validate_and_update_config(start_layer_id, quant_config)

89 self.vocab_size = self.config.vocab_size

90 self.embed_tokens = VocabParallelEmbedding(

@@ -152,16 +168,18 @@ class LlamaModel(nn.Module):

168 }

169

170

155 -class EagleLlama4ForCausalLM(Llama4ForCausalLM, SupportsMultiModalEmbeddings):

171 +class EagleLlama4ForCausalLM(_EagleLlama4ForCausalLMBase, SupportsMultiModalEmbeddings):

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

173 nn.Module.__init__(self)

158 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

174 + speculative_config = vllm_config.speculative_config

175 + assert speculative_config is not None

176 + self.config = speculative_config.draft_model_config.hf_config

177 target_layer_num = vllm_config.model_config.get_num_layers(

178 vllm_config.parallel_config

179 )

180 # draft model quantization config may differ from target model

181 quant_config = VllmConfig.get_quantization_config(

164 - vllm_config.speculative_config.draft_model_config, vllm_config.load_config

182 + speculative_config.draft_model_config, vllm_config.load_config

183 )

184 self.model = LlamaModel(

185 vllm_config=vllm_config,

vllm/model_executor/models/llama_eagle.py

+22 -5

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

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

4 from collections.abc import Iterable

5 +from typing import TYPE_CHECKING

7 import torch

8 import torch.nn as nn

@@ -14,7 +15,10 @@ from vllm.model_executor.layers.linear import ReplicatedLinear

15 from vllm.model_executor.layers.logits_processor import LogitsProcessor

16 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

17 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

17 -from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM

18 +from vllm.model_executor.models.llama import (

19 + LlamaDecoderLayer as BaseLlamaDecoderLayer,

20 +)

21 +from vllm.model_executor.models.llama import LlamaForCausalLM

23 from .utils import (

24 AutoWeightsLoader,

@@ -27,7 +31,16 @@ from .utils import (

31 logger = init_logger(__name__)

30 -class LlamaDecoderLayer(LlamaDecoderLayer):

34 +if TYPE_CHECKING:

35 +

36 + class _EagleLlamaForCausalLMBase(nn.Module):

37 + pass

38 +

39 +else:

40 + _EagleLlamaForCausalLMBase = LlamaForCausalLM

41 +

42 +

43 +class LlamaDecoderLayer(BaseLlamaDecoderLayer):

44 def __init__(

45 self,

46 vllm_config: VllmConfig,

@@ -69,7 +82,9 @@ class LlamaModel(nn.Module):

82 start_layer_id: int = 0,

83 ) -> None:

84 super().__init__()

72 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

85 + speculative_config = vllm_config.speculative_config

86 + assert speculative_config is not None

87 + self.config = speculative_config.draft_model_config.hf_config

88 self.vocab_size = self.config.vocab_size

90 # Get drafter's quantization config

@@ -128,10 +143,12 @@ class LlamaModel(nn.Module):

143 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

144

145

131 -class EagleLlamaForCausalLM(LlamaForCausalLM):

146 +class EagleLlamaForCausalLM(_EagleLlamaForCausalLMBase):

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

148 nn.Module.__init__(self)

134 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

149 + speculative_config = vllm_config.speculative_config

150 + assert speculative_config is not None

151 + self.config = speculative_config.draft_model_config.hf_config

152 # Ensure draft_vocab_size is set

153 # default to the base vocab size when absent

154 if getattr(self.config, "draft_vocab_size", None) is None:

vllm/model_executor/models/llama_eagle3.py

+44 -7

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

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

4 from collections.abc import Iterable

5 +from typing import TYPE_CHECKING

7 import torch

8 import torch.nn as nn

@@ -18,7 +19,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

19 ParallelLMHead,

20 VocabParallelEmbedding,

21 )

21 -from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM

22 +from vllm.model_executor.models.llama import (

23 + LlamaAttention,

24 + LlamaForCausalLM,

25 + LlamaMLP,

26 +)

27 +from vllm.model_executor.models.llama import (

28 + LlamaDecoderLayer as BaseLlamaDecoderLayer,

29 +)

30 from vllm.multimodal.inputs import NestedTensors

32 from .utils import (

@@ -32,7 +40,31 @@ from .utils import (

40 logger = init_logger(__name__)

35 -class LlamaDecoderLayer(LlamaDecoderLayer):

43 +if TYPE_CHECKING:

44 +

45 + class _Eagle3LlamaDecoderLayerBase(nn.Module):

46 + hidden_size: int

47 + self_attn: LlamaAttention

48 + mlp: LlamaMLP

49 + input_layernorm: RMSNorm

50 + post_attention_layernorm: RMSNorm

51 +

52 + def __init__(

53 + self,

54 + vllm_config: VllmConfig,

55 + prefix: str = "",

56 + config: LlamaConfig | None = None,

57 + ) -> None: ...

58 +

59 + class _Eagle3LlamaForCausalLMBase(nn.Module):

60 + pass

61 +

62 +else:

63 + _Eagle3LlamaDecoderLayerBase = BaseLlamaDecoderLayer

64 + _Eagle3LlamaForCausalLMBase = LlamaForCausalLM

65 +

66 +

67 +class LlamaDecoderLayer(_Eagle3LlamaDecoderLayerBase):

68 def __init__(

69 self,

70 vllm_config: VllmConfig,

@@ -136,7 +168,9 @@ class LlamaModel(nn.Module):

168 prefix: str = "",

169 ) -> None:

170 super().__init__()

139 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

171 + speculative_config = vllm_config.speculative_config

172 + assert speculative_config is not None

173 + self.config = speculative_config.draft_model_config.hf_config

174 self.vocab_size = self.config.vocab_size

175

176 # Get drafter's quantization config

@@ -187,7 +221,7 @@ class LlamaModel(nn.Module):

221 self.fc_input_size = target_hidden_size * self.num_aux_hidden_states

222

223 if self.norm_before_fc:

190 - self.input_norm = RMSNorm(

224 + self.input_norm: RMSNorm | None = RMSNorm(

225 self.fc_input_size,

226 eps=self.config.rms_norm_eps,

227 )

@@ -269,10 +303,12 @@ class LlamaModel(nn.Module):

303 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

304

305

272 -class Eagle3LlamaForCausalLM(LlamaForCausalLM):

306 +class Eagle3LlamaForCausalLM(_Eagle3LlamaForCausalLMBase):

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

308 nn.Module.__init__(self)

275 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

309 + speculative_config = vllm_config.speculative_config

310 + assert speculative_config is not None

311 + self.config = speculative_config.draft_model_config.hf_config

312 # Ensure draft_vocab_size is set

313 # default to the base vocab size when absent

314 if getattr(self.config, "draft_vocab_size", None) is None:

@@ -306,7 +342,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

342 requires_grad=False,

343 )

344

309 - self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting

345 + self.use_parallel_drafting = speculative_config.parallel_drafting

346

347 if self.use_parallel_drafting:

348 self.register_buffer(

@@ -365,6 +401,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

401 # combine multiple auxiliary hidden states returned by eagle3

402

403 if self.model.norm_before_fc:

404 + assert self.model.input_norm is not None

405 hidden_states = self.model.input_norm(hidden_states)

406

407 # `norm_before_fc` adds a single RMSNorm before the FC layer, whereas `fc_norm`

vllm/model_executor/models/llava.py

+9 -7

@@ -3,7 +3,7 @@

4 from abc import abstractmethod

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Annotated, Final, Literal, Protocol, TypeAlias, TypeVar

6 +from typing import Annotated, Literal, Protocol, TypeAlias, TypeVar

8 import torch

9 import torch.nn as nn

@@ -159,14 +159,14 @@ class LlavaMultiModalProjector(nn.Module):

159

160

161 class LlavaLikeConfig(Protocol):

162 - vision_config: Final[PretrainedConfig]

163 - image_token_index: Final[int]

164 - vision_feature_select_strategy: Final[str]

165 - vision_feature_layer: Final[int | list[int]]

162 + vision_config: PretrainedConfig

163 + image_token_index: int

164 + vision_feature_select_strategy: str

165 + vision_feature_layer: int | list[int]

166

167

168 class LlavaLikeProcessor(Protocol):

169 - image_token: Final[str]

169 + image_token: str

170

171

172 class BaseLlavaProcessingInfo(BaseProcessingInfo):

@@ -286,6 +286,7 @@ class BaseLlavaMultiModalProcessor(BaseMultiModalProcessor[_I]):

286 if isinstance(images, ImageEmbeddingItems):

287 num_image_tokens = images.get_feature_size(item_idx)

288 else:

289 + assert isinstance(images, ImageProcessorItems)

290 image_size = images.get_image_size(item_idx)

291 num_image_tokens = self.info.get_num_image_tokens(

292 image_width=image_size.width,

@@ -372,6 +373,7 @@ class PixtralHFMultiModalProcessor(BaseMultiModalProcessor[PixtralHFProcessingIn

373 image_end_id = vocab[processor.image_end_token]

374

375 assert isinstance(hf_config.vision_config, PixtralVisionConfig)

376 + assert isinstance(hf_config, LlavaConfig)

377 encoder_info = PixtralHFEncoderInfo(hf_config)

378

379 def get_replacement(item_idx: int):

@@ -642,7 +644,7 @@ class LlavaForConditionalGeneration(

644 self,

645 image_input: LlavaImageInputs,

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

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

647 + if isinstance(image_input, LlavaImageEmbeddingInputs):

648 return image_input["data"]

649

650 image_features = self._process_image_pixels(image_input)

vllm/model_executor/models/llava_next.py

+3 -3

@@ -3,7 +3,7 @@

4 from abc import abstractmethod

5 from collections.abc import Iterable, Mapping

6 -from typing import Annotated, Final, Literal, Protocol, TypeAlias, TypeVar

6 +from typing import Annotated, Literal, Protocol, TypeAlias, TypeVar

8 import torch

9 import torch.nn as nn

@@ -88,7 +88,7 @@ LlavaNextImageInputs: TypeAlias = (

90 class LlavaNextLikeConfig(LlavaLikeConfig, Protocol):

91 - image_grid_pinpoints: Final[list[list[int]]]

91 + image_grid_pinpoints: list[list[int]]

94 class LlavaNextProcessingInfo(BaseLlavaProcessingInfo):

@@ -483,7 +483,7 @@ class LlavaNextForConditionalGeneration(

483 self,

484 image_input: LlavaNextImageInputs,

485 ) -> torch.Tensor | list[torch.Tensor]:

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

486 + if isinstance(image_input, LlavaNextImageEmbeddingInputs):

487 return image_input["data"]

488

489 patch_embeddings = self._process_image_pixels(image_input)

vllm/model_executor/models/llava_next_video.py

+3 -1

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

10 from transformers import BatchFeature, LlavaNextVideoConfig, LlavaNextVideoProcessor

12 from vllm.config import VllmConfig

13 -from vllm.config.multimodal import BaseDummyOptions

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

14 from vllm.inputs import MultiModalDataDict

15 from vllm.model_executor.layers.activation import get_act_fn

16 from vllm.model_executor.models.clip import CLIPVisionModel

@@ -181,6 +181,7 @@ class LlavaNextVideoDummyInputsBuilder(

181 )

182

183 video_overrides = mm_options.get("video")

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

185

186 return {

187 "video": self._get_dummy_videos(

@@ -220,6 +221,7 @@ class LlavaNextVideoMultiModalProcessor(

221 if isinstance(videos, VideoEmbeddingItems):

222 num_video_tokens = videos.get_feature_size(item_idx)

223 else:

224 + assert isinstance(videos, VideoProcessorItems)

225 image_size = videos.get_frame_size(item_idx)

226 num_video_tokens = self.info.get_num_video_tokens(

227 image_width=image_size.width,

vllm/model_executor/models/llava_onevision.py

+21 -9

@@ -3,7 +3,7 @@

4 import math

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Annotated, Final, Literal, Protocol, TypeAlias

6 +from typing import Annotated, Literal, Protocol, TypeAlias, TypedDict

8 import torch

9 import torch.nn as nn

@@ -14,7 +14,7 @@ from transformers.models.llava_onevision.modeling_llava_onevision import (

14 )

16 from vllm.config import VllmConfig

17 -from vllm.config.multimodal import BaseDummyOptions

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

18 from vllm.inputs import MultiModalDataDict

19 from vllm.model_executor.layers.activation import get_act_fn

20 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -123,7 +123,12 @@ LlavaOnevisionMultiInputs: TypeAlias = (

123

124

125 class LlavaOnevisionLikeConfig(LlavaNextLikeConfig, Protocol):

126 - video_token_index: Final[int]

126 + video_token_index: int

127 +

128 +

129 +class LlavaOnevisionInputsByModality(TypedDict, total=False):

130 + image: LlavaOnevisionImageInputs | None

131 + video: LlavaOnevisionVideoPixelInputs | None

132

133

134 class LlavaOnevisionProcessingInfo(LlavaNextProcessingInfo):

@@ -288,6 +293,7 @@ class LlavaOnevisionDummyInputsBuilder(

293

294 image_overrides = mm_options.get("image")

295 video_overrides = mm_options.get("video")

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

297

298 return {

299 "image": self._get_dummy_images(

@@ -417,6 +423,7 @@ class LlavaOnevisionMultiModalProcessor(

423 if isinstance(videos, VideoEmbeddingItems):

424 num_video_tokens = videos.get_feature_size(item_idx)

425 else:

426 + assert isinstance(videos, VideoProcessorItems)

427 image_size = videos.get_frame_size(item_idx)

428 num_video_tokens = self.info.get_num_video_tokens(

429 image_width=image_size.width,

@@ -570,8 +577,10 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

577 },

578 )

579

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

574 - mm_input_by_modality = {}

580 + def _parse_and_validate_multimodal_inputs(

581 + self, **kwargs: object

582 + ) -> LlavaOnevisionInputsByModality:

583 + mm_input_by_modality = LlavaOnevisionInputsByModality()

584

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

586 # from the order of kwargs.

@@ -749,7 +758,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

758 self,

759 image_input: LlavaOnevisionImageInputs,

760 ) -> torch.Tensor | list[torch.Tensor]:

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

761 + if isinstance(image_input, LlavaOnevisionImageEmbeddingInputs):

762 return image_input["data"]

763

764 patch_embeddings = self._process_image_pixels(image_input)

@@ -860,12 +869,15 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

869 # NOTE: It is important to iterate over the keys in this dictionary

870 # to preserve the order of the modalities.

871 for modality in mm_input_by_modality:

863 - multimodal_input = mm_input_by_modality[modality]

872 if modality == "image":

865 - image_embeddings = self._process_image_input(multimodal_input)

873 + image_input = mm_input_by_modality["image"]

874 + assert image_input is not None

875 + image_embeddings = self._process_image_input(image_input)

876 multimodal_embeddings += tuple(image_embeddings)

877 if modality == "video":

868 - video_embeddings = self._process_video_pixels(multimodal_input)

878 + video_input = mm_input_by_modality["video"]

879 + assert video_input is not None

880 + video_embeddings = self._process_video_pixels(video_input)

881 multimodal_embeddings += tuple(video_embeddings)

882

883 return multimodal_embeddings

vllm/model_executor/models/llava_onevision2.py

+55 -18

@@ -24,12 +24,15 @@ from __future__ import annotations

25 import hashlib

26 import importlib

27 -from collections.abc import Callable, Iterable, Mapping, Sequence

27 +from collections.abc import Callable, Iterable, Mapping, Sequence, Sized

28 from functools import lru_cache

29 from typing import (

30 Annotated,

31 Any,

32 Literal,

33 + SupportsIndex,

34 + SupportsInt,

35 + TypedDict,

36 )

38 import numpy as np

@@ -45,6 +48,7 @@ from transformers.dynamic_module_utils import (

48 )

49 from transformers.models.qwen2_vl import Qwen2VLImageProcessor

50 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize

51 +from typing_extensions import Buffer

53 from vllm.compilation.decorators import (

54 should_torch_compile_mm_encoder,

@@ -730,6 +734,11 @@ LlavaOnevision2ImageInputs = (

734 )

735

736

737 +class LlavaOnevision2InputsByModality(TypedDict, total=False):

738 + images: LlavaOnevision2ImageInputs | None

739 + videos: LlavaOnevision2VideoPixelInputs | None

740 +

741 +

742 class LlavaOnevision2VisionRotaryEmbedding(nn.Module):

743 """3-D rotary frequency constructor with 4:6:6 (T:H:W) split.

744

@@ -1468,7 +1477,7 @@ class LlavaOnevision2DummyInputsBuilder(

1477 n_img = mm_counts.get("image", 0)

1478 n_vid = mm_counts.get("video", 0)

1479 w, h = self.info.get_image_size_with_most_features()

1471 - out: MultiModalDataDict = {}

1480 + out: dict[str, object] = {}

1481 if n_img:

1482 out["image"] = self._get_dummy_images(width=w, height=h, num_images=n_img)

1483 if n_vid:

@@ -1586,6 +1595,8 @@ class LlavaOnevision2MultiModalProcessor(

1595 # array, or tensor, and ``and <array>`` would raise on the ambiguous

1596 # truth value of a multi-element array.

1597 _videos = mm_data.get("videos")

1598 + if _videos is not None and not isinstance(_videos, Sized):

1599 + raise TypeError(f"object of type '{type(_videos).__name__}' has no len()")

1600 videos_present = _videos is not None and len(_videos) > 0

1601

1602 codec_video_paths = (

@@ -1666,15 +1677,24 @@ class LlavaOnevision2MultiModalProcessor(

1677 # provided by the backend's ``compute_frames_index_to_sample``; SSRF /

1678 # local-file gating is enforced by the connector before decoding.

1679 if videos_present:

1669 - timestamp_decimals = int(

1670 - hf_processor_mm_kwargs.get(

1671 - "timestamp_decimals", _DEFAULT_TIMESTAMP_DECIMALS

1672 - )

1680 + timestamp_decimals = hf_processor_mm_kwargs.get(

1681 + "timestamp_decimals", _DEFAULT_TIMESTAMP_DECIMALS

1682 )

1683 + if not isinstance(

1684 + timestamp_decimals, (str, Buffer, SupportsInt, SupportsIndex)

1685 + ):

1686 + raise TypeError(

1687 + "int() argument must be a string, a bytes-like object "

1688 + f"or a real number, not '{type(timestamp_decimals).__name__}'"

1689 + )

1690 + timestamp_decimals = int(timestamp_decimals)

1691

1692 per_video_frames: list[list[Image.Image]] = []

1693 per_video_timestamps: list[list[float]] = []

1677 - for item in mm_data["videos"]:

1694 + videos = mm_data["videos"]

1695 + if not isinstance(videos, Iterable):

1696 + raise TypeError(f"Unsupported videos input: {type(videos)}")

1697 + for item in videos:

1698 pil_frames, timestamps = _frame_video_to_pil_and_timestamps(item)

1699 per_video_frames.append(pil_frames)

1700 per_video_timestamps.append(timestamps)

@@ -1902,18 +1922,21 @@ class LlavaOnevision2MultiModalProcessor(

1922 def get_image_replacement(item_idx: int):

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

1924 grid_thw = out_item["image_grid_thw"].data

1925 + assert isinstance(grid_thw, torch.Tensor)

1926 n = int(grid_thw.prod(-1).sum()) // merge_length

1927 return [image_pad_id] * n

1928

1929 def get_video_replacement(item_idx: int):

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

1931 grid_thw = out_item["video_grid_thw"].data

1932 + assert isinstance(grid_thw, torch.Tensor)

1933 is_codec_field = out_item.get("video_is_codec")

1912 - is_codec = (

1913 - bool(int(is_codec_field.data.item()))

1914 - if is_codec_field is not None

1915 - else False

1916 - )

1934 + if is_codec_field is not None:

1935 + is_codec_data = is_codec_field.data

1936 + assert isinstance(is_codec_data, torch.Tensor)

1937 + is_codec = bool(int(is_codec_data.item()))

1938 + else:

1939 + is_codec = False

1940 tokens: list[int] = []

1941

1942 if is_codec:

@@ -1924,6 +1947,8 @@ class LlavaOnevision2MultiModalProcessor(

1947 # run, emit ``<sec seconds><|vision_start|><pad*N><|vision_end|>\n``.

1948 patch_positions = out_item["patch_positions_videos"].data

1949 fps_t = out_item["codec_fps"].data

1950 + assert isinstance(patch_positions, torch.Tensor)

1951 + assert isinstance(fps_t, torch.Tensor)

1952 fps = float(int(fps_t.item())) / 1_000_000_000.0

1953 runs = _codec_timestamp_runs(

1954 patch_positions, fps, image_processor.merge_size

@@ -1938,6 +1963,7 @@ class LlavaOnevision2MultiModalProcessor(

1963 tokens.extend(newline_ids)

1964 else:

1965 timestamps = out_item["frame_timestamps"].data

1966 + assert isinstance(timestamps, torch.Tensor)

1967 T_total = int(grid_thw.shape[0])

1968 for t in range(T_total):

1969 sec = float(timestamps[t].item())

@@ -2121,7 +2147,7 @@ class LlavaOnevision2ForConditionalGeneration(

2147 grid_thw = image_input["image_grid_thw"]

2148 assert grid_thw.ndim == 2

2149

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

2150 + if isinstance(image_input, LlavaOnevision2ImageEmbeddingInputs):

2151 image_embeds = image_input["image_embeds"]

2152 else:

2153 image_embeds = self.visual(

@@ -2166,8 +2192,13 @@ class LlavaOnevision2ForConditionalGeneration(

2192 )

2193 if isinstance(video_num_frames, list):

2194 video_num_frames = torch.cat([v.flatten() for v in video_num_frames])

2169 - else:

2195 + elif isinstance(video_num_frames, torch.Tensor):

2196 video_num_frames = video_num_frames.flatten()

2197 + else:

2198 + raise TypeError(

2199 + "video_num_frames must be a tensor or a list of tensors, "

2200 + f"got {type(video_num_frames)}"

2201 + )

2202 return LlavaOnevision2VideoPixelInputs(

2203 type="pixel_values_videos",

2204 pixel_values_videos=pixel_values_videos,

@@ -2202,8 +2233,10 @@ class LlavaOnevision2ForConditionalGeneration(

2233 cursor += n

2234 return video_embeds.split(sizes)

2235

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

2206 - modalities = {}

2236 + def _parse_and_validate_multimodal_inputs(

2237 + self, **kwargs: object

2238 + ) -> LlavaOnevision2InputsByModality:

2239 + modalities = LlavaOnevision2InputsByModality()

2240 for key in kwargs:

2241 if key in ("pixel_values", "image_embeds") and "images" not in modalities:

2242 modalities["images"] = self._parse_and_validate_image_input(**kwargs)

@@ -2221,9 +2254,13 @@ class LlavaOnevision2ForConditionalGeneration(

2254 multimodal_embeddings: tuple[torch.Tensor, ...] = ()

2255 for modality in modalities:

2256 if modality == "images":

2224 - multimodal_embeddings += self._process_image_input(modalities["images"])

2257 + image_input = modalities["images"]

2258 + assert image_input is not None

2259 + multimodal_embeddings += self._process_image_input(image_input)

2260 elif modality == "videos":

2226 - multimodal_embeddings += self._process_video_input(modalities["videos"])

2261 + video_input = modalities["videos"]

2262 + assert video_input is not None

2263 + multimodal_embeddings += self._process_video_input(video_input)

2264 return multimodal_embeddings

2265

2266 def get_input_embeddings(

vllm/model_executor/models/longcat_flash.py

+12 -7

@@ -83,6 +83,8 @@ logger = init_logger(__name__)

83 class FlashConfig(PretrainedConfig):

84 """Flash model configuration."""

86 + moe_intermediate_size: int

87 +

88 model_type = "longcat_flash"

89 keys_to_ignore_at_inference = ["past_key_values"]

@@ -604,7 +606,7 @@ class FlashModel(nn.Module):

606 else:

607 is_expert_weight = False

608 for mapping in expert_params_mapping:

607 - param_name, weight_name, expert_id, shard_id = mapping

609 + param_name, weight_name, expert_id, expert_shard_id = mapping

610 if weight_name not in name:

611 continue

612 is_expert_weight = True

@@ -627,7 +629,7 @@ class FlashModel(nn.Module):

629 param,

630 loaded_weight,

631 name_mapped,

630 - shard_id=shard_id,

632 + shard_id=expert_shard_id,

633 expert_id=expert_id,

634 return_success=True,

635 )

@@ -664,11 +666,14 @@ class FlashModel(nn.Module):

666 if isinstance(self.layers[layer_id], PPMissingLayer):

667 continue

668 self_attn = self.layers[layer_id].self_attn[i]

667 - if hasattr(

668 - self.quant_config, "weight_block_size"

669 - ) and self_attn.kv_b_proj.weight.dtype in (

670 - torch.float8_e4m3fn,

671 - torch.float8_e4m3fnuz,

669 + if (

670 + self.quant_config is not None

671 + and hasattr(self.quant_config, "weight_block_size")

672 + and self_attn.kv_b_proj.weight.dtype

673 + in (

674 + torch.float8_e4m3fn,

675 + torch.float8_e4m3fnuz,

676 + )

677 ):

678 weight_block_size = self.quant_config.weight_block_size

679 if weight_block_size is not None:

vllm/model_executor/models/longcat_flash_mtp.py

+8 -5

@@ -266,11 +266,14 @@ class LongCatFlashMTP(nn.Module):

266 loaded_params.add(name)

267 spec_layer_id = self.config.num_hidden_layers * 2

268 self_attn = self.model.layers[str(spec_layer_id)].mtp_block.self_attn

269 - if hasattr(

270 - self.quant_config, "weight_block_size"

271 - ) and self_attn.kv_b_proj.weight.dtype in (

272 - torch.float8_e4m3fn,

273 - torch.float8_e4m3fnuz,

269 + if (

270 + self.quant_config is not None

271 + and hasattr(self.quant_config, "weight_block_size")

272 + and self_attn.kv_b_proj.weight.dtype

273 + in (

274 + torch.float8_e4m3fn,

275 + torch.float8_e4m3fnuz,

276 + )

277 ):

278 weight_block_size = self.quant_config.weight_block_size

279 if weight_block_size is not None:

vllm/model_executor/models/longcat_flash_ngram.py

@@ -178,6 +178,7 @@ class FlashNgramModel(FlashModel):

178 if num_layers is not None and hf.num_hidden_layers != num_layers:

179 hf.num_hidden_layers = num_layers

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

181 + self.ngram_embeddings: NgramEmbedding | None

182 if get_pp_group().is_first_rank and uses_ngram_embedding(self.config):

183 self.ngram_embeddings = NgramEmbedding(self.config, self.embed_tokens)

184 else:

vllm/model_executor/models/mistral3.py

+1 -1

@@ -217,7 +217,7 @@ class Mistral3ProcessingInfo(BaseProcessingInfo):

217 def get_vision_encoder_info(

218 self,

219 mm_processor_kwargs: Mapping[str, object] | None = None,

220 - ) -> Mistral3HFEncoderInfo:

220 + ) -> PixtralHFEncoderInfo:

221 processor = self.get_hf_processor()

222 size = processor.image_processor.size

223 merged_kwargs = self.ctx.get_merged_mm_kwargs(mm_processor_kwargs or {})