@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ccee371e86
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Docs] Fix warnings in `mkdocs build` (continued) (#24092)

Signed-off-by: Zerohertz <ohg3417@gmail.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Hyogeun Oh (오효근) committed Sep 10, 2025 at 22:23 UTC ccee371e86f18643f251450397ca7e46a0356b5b
10 files changed +337 -342
vllm/model_executor/layers/fused_moe/layer.py
+1 -1
@@ -755,7 +755,7 @@ class FusedMoE(CustomOp):
755 intermediate_size: Intermediate size of the experts
756 params_dtype: Data type for the parameters.
757 reduce_results: Whether to all all_reduce on the output of the layer
758 - renomalize: Whether to renormalize the logits in the fused_moe kernel
758 + renormalize: Whether to renormalize the logits in the fused_moe kernel
759 quant_config: Quantization configure.
760 enable_eplb: Whether to enable expert parallelism load balancer.
761 """
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+2 -3
@@ -420,9 +420,8 @@ def shuffle_weights(
420
421 Args:
422 *tensors: Variable number of torch.Tensor objects.
423 - layout: A pair of integers specifying the
424 - block sizes used to divide the tensors during shuffling.
425 - Default is (16, 16).
423 + layout: A pair of integers specifying the block sizes used to divide
424 + the tensors during shuffling. Default is (16, 16).
425
426 Returns:
427 A Tuple of shuffled tensors.
vllm/model_executor/layers/fused_moe/routing_simulator.py
+5 -3
@@ -10,7 +10,7 @@ like uniform random routing.
10 """
11
12 from abc import ABC, abstractmethod
13 -from typing import Optional
13 +from typing import Any, Optional
14
15 import torch
16
@@ -50,7 +50,9 @@ class DistributionBasedRouting(RoutingStrategy):
50 distributions for testing different routing patterns.
51 """
52
53 - def __init__(self, distribution: str = "uniform", **distribution_params):
53 + def __init__(self,
54 + distribution: str = "uniform",
55 + **distribution_params: Any):
56 """
57 Initialize distribution-based routing.
58
@@ -244,7 +246,7 @@ class RoutingSimulator:
246 cls._routing_strategies[name] = strategy
247
248 @classmethod
247 - def get_available_strategies(cls):
249 + def get_available_strategies(cls) -> list[str]:
250 """
251 Get list of available routing strategy names.
252
vllm/model_executor/layers/quantization/bitblas.py
+5 -5
@@ -202,7 +202,7 @@ class BitBLASLinearMethod(LinearMethodBase):
202 output_size: int,
203 params_dtype: torch.dtype,
204 **extra_weight_attrs,
205 - ):
205 + ) -> None:
206 """Creates quantized weights for use in linear operations.
207
208 The function initializes and returns a dictionary containing quantized
@@ -211,7 +211,7 @@ class BitBLASLinearMethod(LinearMethodBase):
211
212 Args:
213 input_size_per_partition: The size of the input partition.
214 - output_size_per_partition: The size of the output partition.
214 + output_partition_sizes: List of output partition sizes.
215 input_size: The total size of the input (unused).
216 output_size: The total size of the output (unused).
217 params_dtype:
@@ -222,9 +222,9 @@ class BitBLASLinearMethod(LinearMethodBase):
222 scales ('scales'), and zeros ('zeros').
223
224 Raises:
225 - ValueError: If `params_dtype` is not `torch.float16` or if the
226 - input size per partition is not divisible by the group size in
227 - `quant_config`.
225 + ValueError: If `params_dtype` is not `torch.float16` or if the input
226 + size per partition is not divisible by the group size
227 + in `quant_config`.
228 """
229 del input_size, output_size # Unused arguments.
230 weight_loader = extra_weight_attrs["weight_loader"]
vllm/model_executor/layers/quantization/gptq_bitblas.py
+3 -3
@@ -265,9 +265,9 @@ class GPTQBitBLASLinearMethod(LinearMethodBase):
265 scales ('scales'), and zeros ('zeros').
266
267 Raises:
268 - ValueError: If `params_dtype` is not `torch.float16` or
269 - if the input size per partition is not divisible by the
270 - group size in `quant_config`.
268 + ValueError: If `params_dtype` is not `torch.float16` or if the input
269 + size per partition is not divisible by the group size
270 + in `quant_config`.
271 """
272 if params_dtype != torch.float16:
273 raise ValueError("Parameter data type must be torch.float16, "
vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py
+4 -4
@@ -46,11 +46,11 @@ def choose_mp_linear_kernel(
46 performance.
47
48 Args:
49 - config (MPLinearLayerConfig): Description of the linear layer to be
50 - implemented.
49 + config (MPLinearLayerConfig): Description of the linear layer to be
50 + implemented.
51 compute_capability (Optional[int], optional): The compute capability of
52 - the target device, if None uses `current_platform` to get the compute
53 - capability. Defaults to None.
52 + the target device, if None uses `current_platform` to get
53 + the compute capability. Defaults to None.
54
55 Raises:
56 ValueError: If no kernel can implement the given config.
vllm/model_executor/models/phi4mm_audio.py
+144 -123
@@ -7,7 +7,7 @@
7 #!/usr/bin/env python3
8 import abc
9 import math
10 -from typing import Literal, Optional
10 +from typing import Any, Literal, Optional, Union
11
12 import numpy as np
13 import torch
@@ -131,31 +131,31 @@ class ConformerEncoderLayer(nn.Module):
131
132 def __init__(
133 self,
134 - d_model=512,
135 - ext_pw_out_channel=0,
136 - depthwise_seperable_out_channel=256,
137 - depthwise_multiplier=1,
138 - n_head=4,
139 - d_ffn=2048,
140 - ext_pw_kernel_size=1,
141 - kernel_size=3,
142 - dropout_rate=0.1,
143 - causal=False,
144 - batch_norm=False,
145 - activation="relu",
146 - chunk_se=0,
147 - chunk_size=18,
148 - conv_activation="relu",
149 - conv_glu_type="sigmoid",
150 - bias_in_glu=True,
151 - linear_glu_in_convm=False,
152 - attention_inner_dim=-1,
153 - attention_glu_type="swish",
154 - activation_checkpointing="",
155 - export=False,
156 - use_pt_scaled_dot_product_attention=False,
134 + d_model: int = 512,
135 + ext_pw_out_channel: int = 0,
136 + depthwise_seperable_out_channel: int = 256,
137 + depthwise_multiplier: int = 1,
138 + n_head: int = 4,
139 + d_ffn: int = 2048,
140 + ext_pw_kernel_size: int = 1,
141 + kernel_size: int = 3,
142 + dropout_rate: float = 0.1,
143 + causal: bool = False,
144 + batch_norm: bool = False,
145 + activation: str = "relu",
146 + chunk_se: int = 0,
147 + chunk_size: int = 18,
148 + conv_activation: str = "relu",
149 + conv_glu_type: str = "sigmoid",
150 + bias_in_glu: bool = True,
151 + linear_glu_in_convm: bool = False,
152 + attention_inner_dim: int = -1,
153 + attention_glu_type: str = "swish",
154 + activation_checkpointing: str = "",
155 + export: bool = False,
156 + use_pt_scaled_dot_product_attention: bool = False,
157 attn_group_sizes: int = 1,
158 - ):
158 + ) -> None:
159 super().__init__()
160
161 self.feed_forward_in = FeedForward(
@@ -209,24 +209,21 @@ class ConformerEncoderLayer(nn.Module):
209
210 def forward(
211 self,
212 - x,
213 - pos_k,
214 - pos_v,
215 - mask,
212 + x: torch.Tensor,
213 + pos_k: torch.Tensor,
214 + pos_v: torch.Tensor,
215 + mask: torch.Tensor,
216 relative_attention_bias: Optional[Tensor] = None,
217 - ):
217 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
218 """ConformerEncoder forward.
219
220 Args:
221 - x: torch.Tensor
222 - input feature of shape (batch, max_time_in, size)
223 - pos_k: torch.Tensor
224 - positional key embedding.
225 - mask: torch.Tensor
226 - mask for x (batch, max_time_in)
227 - relative_attention_bias: Optional[torch.Tensor]
228 - bias added to attention logits w.r.t. relative positions
229 - (1, n_head, time1, time2)
221 + x: input feature of shape (batch, max_time_in, size)
222 + pos_k: positional key embedding.
223 + pos_v: positional value embedding.
224 + mask: mask for x (batch, max_time_in)
225 + relative_attention_bias: bias added to attention logits w.r.t.
226 + relative positions (1, n_head, time1, time2)
227 """
228 x = x + 0.5 * self.feed_forward_in(x)
229 norm_x = self.layer_norm_att(x)
@@ -323,25 +320,25 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
320
321 def __init__(
322 self,
326 - input_size,
327 - chunk_size,
328 - left_chunk,
329 - attention_dim=256,
330 - attention_heads=4,
331 - input_layer="nemo_conv",
332 - cnn_out=-1,
333 - cnn_layer_norm=False,
334 - time_reduction=4,
335 - dropout_rate=0.0,
336 - padding_idx=-1,
337 - relative_attention_bias_args=None,
338 - positional_dropout_rate=0.0,
339 - nemo_conv_settings=None,
323 + input_size: int,
324 + chunk_size: Union[int, list[int]],
325 + left_chunk: Union[int, list[int]],
326 + attention_dim: int = 256,
327 + attention_heads: int = 4,
328 + input_layer: str = "nemo_conv",
329 + cnn_out: int = -1,
330 + cnn_layer_norm: bool = False,
331 + time_reduction: int = 4,
332 + dropout_rate: float = 0.0,
333 + padding_idx: int = -1,
334 + relative_attention_bias_args: Optional[dict[str, Any]] = None,
335 + positional_dropout_rate: float = 0.0,
336 + nemo_conv_settings: Optional[dict[str, Any]] = None,
337 conv2d_extra_padding: Literal["feat", "feat_time", "none",
338 True] = "none",
342 - attention_group_size=1,
343 - encoder_embedding_config=None,
344 - ):
339 + attention_group_size: int = 1,
340 + encoder_embedding_config: Optional[dict[str, Any]] = None,
341 + ) -> None:
342 super().__init__()
343 self.input_size = input_size
344 self.input_layer = input_layer
@@ -399,7 +396,10 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
396 self.encoder_embedding = MeanVarianceNormLayer(
397 self.encoder_embedding_config["input_size"])
398
402 - def compute_lens_change(self, feature_lens):
399 + def compute_lens_change(
400 + self,
401 + feature_lens: Union[int,
402 + torch.Tensor]) -> Union[int, torch.Tensor]:
403 """feature_lens: int
404 return updated feature lens.
405
@@ -433,10 +433,14 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
433 return ceil_func(feature_lens / self.time_reduction)
434
435 @abc.abstractmethod
436 - def forward(self):
436 + def forward(self) -> Any:
437 """Abstract forward method implementation."""
438
439 - def _chunk_size_selection(self, chunk_size=None, left_chunk=None):
439 + def _chunk_size_selection(
440 + self,
441 + chunk_size: Optional[Union[int, list[int]]] = None,
442 + left_chunk: Optional[Union[int,
443 + list[int]]] = None) -> tuple[int, int]:
444 """If chunk size is a list, we will randomly select a chunk size."""
445
446 if chunk_size is None:
@@ -463,7 +467,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
467
468 return chunk_size_train_eff, left_chunk_train_eff
469
466 - def _get_embed_class(self, embed):
470 + def _get_embed_class(self, embed: nn.Module) -> nn.Module:
471 # pylint: disable=protected-access
472 is_embed_using_act_chkpt = isinstance(embed, CheckpointWrapper)
473 is_embed_fsdp_wrapped = isinstance(embed, FullyShardedDataParallel)
@@ -474,13 +478,17 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
478 embed_class = embed.module
479 return embed_class
480
477 - def _forward_embeddings_core(self, input_tensor, masks):
481 + def _forward_embeddings_core(
482 + self, input_tensor: torch.Tensor,
483 + masks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
484 embed_class = self._get_embed_class(self.embed)
485 assert isinstance(embed_class, NemoConvSubsampling)
486 input_tensor, masks = self.embed(input_tensor, masks)
487 return input_tensor, masks
488
483 - def _position_embedding(self, input_tensor):
489 + def _position_embedding(
490 + self, input_tensor: torch.Tensor
491 + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
492 pos_k = None
493 pos_v = None
494 if self.relative_attention_bias_layer is None:
@@ -488,7 +496,9 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
496 input_tensor) # default to add abs sinusoid embedding
497 return pos_k, pos_v
498
491 - def _streaming_mask(self, seq_len, batch_size, chunk_size, left_chunk):
499 + def _streaming_mask(self, seq_len: int, batch_size: int,
500 + chunk_size: Union[int, list[int]],
501 + left_chunk: Union[int, list[int]]) -> torch.Tensor:
502 chunk_size_train_eff, left_chunk_train_eff = \
503 self._chunk_size_selection(chunk_size, left_chunk)
504
@@ -502,11 +512,17 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
512 [batch_size, -1, -1]))
513 return enc_streaming_mask
514
505 - def forward_embeddings(self,
506 - xs_pad,
507 - masks,
508 - chunk_size_nc=None,
509 - left_chunk_nc=None):
515 + def forward_embeddings(
516 + self,
517 + xs_pad: torch.Tensor,
518 + masks: torch.Tensor,
519 + chunk_size_nc: Optional[Union[int, list[int]]] = None,
520 + left_chunk_nc: Optional[Union[int, list[int]]] = None
521 + ) -> Union[tuple[torch.Tensor, Optional[torch.Tensor],
522 + Optional[torch.Tensor], torch.Tensor, torch.Tensor],
523 + tuple[torch.Tensor, Optional[torch.Tensor],
524 + Optional[torch.Tensor], torch.Tensor, torch.Tensor,
525 + torch.Tensor]]:
526 """Forwarding the inputs through the top embedding layers
527
528 Args:
@@ -569,7 +585,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
585 return input_tensor, pos_k, pos_v, hs_mask, masks
586 return input_tensor, pos_k, pos_v, hs_mask, masks, hs_mask_nc
587
572 - def get_offset(self):
588 + def get_offset(self) -> int:
589 """Returns offset used when retaining inputs for decoding.
590
591 This is essentially, how many additional frames have to be added to
@@ -605,8 +621,6 @@ class ConformerEncoder(TransformerEncoderBase):
621 Some examples for the 2 cases:
622 left_chunk = 6
623 left_chunk = [12, 9, 6, 3]
608 - left_chunk: int
609 - number of chunks used for masking in streaming mode.
624 num_lang: int
625 This parameter is used to store the number of languages in the
626 lang_dict, only used for multiseed/multilingual models.
@@ -751,46 +765,46 @@ class ConformerEncoder(TransformerEncoderBase):
765
766 def __init__( # pylint: disable-all
767 self,
754 - input_size,
755 - chunk_size,
756 - left_chunk,
757 - num_lang=None,
758 - attention_dim=256,
759 - attention_heads=4,
760 - linear_units=2048,
761 - num_blocks=6,
762 - dropout_rate=0.1,
763 - input_layer="nemo_conv",
764 - causal=True,
765 - batch_norm=False,
766 - cnn_out=-1,
767 - cnn_layer_norm=False,
768 - ext_pw_out_channel=0,
769 - ext_pw_kernel_size=1,
770 - depthwise_seperable_out_channel=256,
771 - depthwise_multiplier=1,
772 - chunk_se=0,
773 - kernel_size=3,
774 - activation="relu",
775 - conv_activation="relu",
776 - conv_glu_type="sigmoid",
777 - bias_in_glu=True,
778 - linear_glu_in_convm=False,
779 - attention_glu_type="swish",
780 - export=False,
781 - extra_layer_output_idx=-1,
782 - extra_multi_layer_output_idxs=[], # noqa
783 - activation_checkpointing="",
784 - relative_attention_bias_args=None,
785 - time_reduction=4,
786 - use_pt_scaled_dot_product_attention=False,
787 - nemo_conv_settings=None,
768 + input_size: int,
769 + chunk_size: Union[int, list[int]],
770 + left_chunk: Union[int, list[int]],
771 + num_lang: Optional[int] = None,
772 + attention_dim: int = 256,
773 + attention_heads: int = 4,
774 + linear_units: int = 2048,
775 + num_blocks: int = 6,
776 + dropout_rate: float = 0.1,
777 + input_layer: str = "nemo_conv",
778 + causal: bool = True,
779 + batch_norm: bool = False,
780 + cnn_out: int = -1,
781 + cnn_layer_norm: bool = False,
782 + ext_pw_out_channel: int = 0,
783 + ext_pw_kernel_size: int = 1,
784 + depthwise_seperable_out_channel: int = 256,
785 + depthwise_multiplier: int = 1,
786 + chunk_se: int = 0,
787 + kernel_size: int = 3,
788 + activation: str = "relu",
789 + conv_activation: str = "relu",
790 + conv_glu_type: str = "sigmoid",
791 + bias_in_glu: bool = True,
792 + linear_glu_in_convm: bool = False,
793 + attention_glu_type: str = "swish",
794 + export: bool = False,
795 + extra_layer_output_idx: int = -1,
796 + extra_multi_layer_output_idxs: list[int] = [], # noqa
797 + activation_checkpointing: str = "",
798 + relative_attention_bias_args: Optional[dict[str, Any]] = None,
799 + time_reduction: int = 4,
800 + use_pt_scaled_dot_product_attention: bool = False,
801 + nemo_conv_settings: Optional[dict[str, Any]] = None,
802 conv2d_extra_padding: Literal["feat", "feat_time", "none",
803 True] = "none",
790 - replication_pad_for_subsample_embedding=False,
791 - attention_group_size=1,
792 - encoder_embedding_config=None,
793 - ):
804 + replication_pad_for_subsample_embedding: bool = False,
805 + attention_group_size: int = 1,
806 + encoder_embedding_config: Optional[dict[str, Any]] = None,
807 + ) -> None:
808 super().__init__(
809 input_size,
810 chunk_size,
@@ -852,11 +866,13 @@ class ConformerEncoder(TransformerEncoderBase):
866 # the device and the needed dtype:
867 self.register_buffer("dev_type", torch.zeros(()), persistent=False)
868
855 - def init_relative_attention_bias(self, input_tensor):
869 + def init_relative_attention_bias(
870 + self, input_tensor: torch.Tensor) -> Optional[torch.Tensor]:
871 if self.relative_attention_bias_layer:
872 return self.relative_attention_bias_layer(input_tensor)
873
859 - def calculate_hs_mask(self, xs_pad, device, mask):
874 + def calculate_hs_mask(self, xs_pad: torch.Tensor, device: torch.device,
875 + mask: Optional[torch.Tensor]) -> torch.Tensor:
876 max_audio_length = xs_pad.shape[1]
877 batch_size = xs_pad.shape[0]
878 enc_streaming_mask = self._streaming_mask(max_audio_length, batch_size,
@@ -877,7 +893,8 @@ class ConformerEncoder(TransformerEncoderBase):
893 return pad_mask
894
895 @torch.jit.ignore
880 - def forward(self, xs_pad, masks):
896 + def forward(self, xs_pad: torch.Tensor,
897 + masks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
898 """Conformer Forward function
899
900 Args:
@@ -997,7 +1014,12 @@ class WindowQformer(nn.Module):
1014 if normalize_before else None)
1015 self.window_size = window_size
1016
1000 - def forward(self, audio_embed, mask, embed_len=None):
1017 + def forward(
1018 + self,
1019 + audio_embed: torch.Tensor,
1020 + mask: Optional[torch.Tensor],
1021 + embed_len: Optional[int] = None
1022 + ) -> tuple[torch.Tensor, Optional[int]]:
1023 """forward decoder"""
1024 # audio_embed: N x T x D => N x D x T
1025
@@ -1042,7 +1064,7 @@ class WindowQformer(nn.Module):
1064 class AudioEmbedding(nn.Module):
1065 """Image embedding."""
1066
1045 - def __init__(self, config: PretrainedConfig, **kwargs) -> None:
1067 + def __init__(self, config: PretrainedConfig, **kwargs: Any) -> None:
1068 super().__init__()
1069 self.config = config
1070 # n_embed or hidden_size for text LM
@@ -1148,19 +1170,18 @@ class AudioEmbedding(nn.Module):
1170 self.input_embeds = None
1171 self.audio_embed_sizes = None
1172
1151 - def set_audio_embeds(self, input_embeds: torch.FloatTensor) -> None:
1173 + def set_audio_embeds(self, input_embeds: torch.Tensor) -> None:
1174 self.input_embeds = input_embeds
1175
1154 - def set_audio_embed_sizes(self,
1155 - audio_embed_sizes: torch.LongTensor) -> None:
1176 + def set_audio_embed_sizes(self, audio_embed_sizes: torch.Tensor) -> None:
1177 self.audio_embed_sizes = audio_embed_sizes
1178
1179 def get_audio_features(
1180 self,
1160 - input_embeds: torch.FloatTensor,
1161 - audio_attention_mask: torch.Tensor = None,
1181 + input_embeds: torch.Tensor,
1182 + audio_attention_mask: Optional[torch.Tensor] = None,
1183 audio_projection_mode: str = "speech",
1163 - ) -> torch.FloatTensor:
1184 + ) -> torch.Tensor:
1185 """
1186 arguments:
1187 input_embeds: audio features (B, T, D) B: num audios in a sequence
@@ -1214,10 +1235,10 @@ class AudioEmbedding(nn.Module):
1235
1236 def forward(
1237 self,
1217 - audio_features: torch.FloatTensor,
1218 - audio_attention_mask: torch.Tensor = None,
1238 + audio_features: torch.Tensor,
1239 + audio_attention_mask: Optional[torch.Tensor] = None,
1240 audio_projection_mode: str = "speech",
1220 - ) -> torch.FloatTensor:
1241 + ) -> torch.Tensor:
1242 """
1243 arguments:
1244 audio_features: audio features (T, D)
vllm/model_executor/models/phi4mm_utils.py
+164 -173
@@ -16,13 +16,13 @@ from torch import Tensor, nn
16 class BlockBase(nn.Module):
17 """Block abstract module"""
18
19 - def __init__(self, input_size, output_size):
19 + def __init__(self, input_size: int, output_size: int) -> None:
20 super().__init__()
21 self.input_size = input_size
22 self.output_size = output_size
23
24
25 -def get_activation(name="relu"):
25 +def get_activation(name: str = "relu") -> torch.nn.Module:
26 """Select an activation function by name
27
28 Args:
@@ -43,15 +43,18 @@ def get_activation(name="relu"):
43 return nn.Identity()
44
45
46 -def adaptive_enc_mask(x_len, chunk_start_idx, left_window=0, right_window=0):
46 +def adaptive_enc_mask(x_len: int,
47 + chunk_start_idx: list[int],
48 + left_window: int = 0,
49 + right_window: int = 0) -> torch.Tensor:
50 """
51 The function is very important for Transformer Transducer Streaming mode
52 Args:
50 - xs_len (int): sequence length
51 - chunk_start_idx (list): first idx of each chunk, such as [0,18,36,48].
53 + x_len: sequence length
54 + chunk_start_idx: first idx of each chunk, such as [0,18,36,48].
55 It also supports adaptive chunk size [0,10,15,45]
53 - left_window (int): how many left chunks can be seen
54 - right_window (int): how many right chunks can be seen. It is used for
56 + left_window: how many left chunks can be seen
57 + right_window: how many right chunks can be seen. It is used for
58 chunk overlap model.
59 Returns:
60 mask (torch.Tensor): a mask tensor for streaming model
@@ -172,13 +175,13 @@ class GLUPointWiseConv(nn.Module):
175
176 def __init__(
177 self,
175 - input_dim,
176 - output_dim,
177 - kernel_size,
178 - glu_type="sigmoid",
179 - bias_in_glu=True,
180 - causal=False,
181 - ):
178 + input_dim: int,
179 + output_dim: int,
180 + kernel_size: int,
181 + glu_type: str = "sigmoid",
182 + bias_in_glu: bool = True,
183 + causal: bool = False,
184 + ) -> None:
185 super().__init__()
186
187 self.glu_type = glu_type
@@ -216,11 +219,10 @@ class GLUPointWiseConv(nn.Module):
219 self.b1 = nn.Parameter(torch.zeros(1, output_dim, 1))
220 self.b2 = nn.Parameter(torch.zeros(1, output_dim, 1))
221
219 - def forward(self, x):
222 + def forward(self, x: Tensor) -> Tensor:
223 """
224 Args:
222 - x: torch.Tensor
223 - input tensor
225 + x: input tensor
226 """
227 # to be consistent with GLULinear, we assume the input always has the
228 # #channel (#dim) in the last dimension of the tensor, so need to
@@ -272,12 +274,12 @@ class DepthWiseSeperableConv1d(nn.Module):
274
275 def __init__(
276 self,
275 - input_dim,
276 - depthwise_seperable_out_channel,
277 - kernel_size,
278 - depthwise_multiplier,
279 - padding=0,
280 - ):
277 + input_dim: int,
278 + depthwise_seperable_out_channel: int,
279 + kernel_size: int,
280 + depthwise_multiplier: int,
281 + padding: int = 0,
282 + ) -> None:
283 super().__init__()
284
285 self.dw_conv = nn.Conv1d(
@@ -301,12 +303,11 @@ class DepthWiseSeperableConv1d(nn.Module):
303 self.pw_conv = nn.Identity()
304 self.depthwise_seperable_out_channel = depthwise_seperable_out_channel
305
304 - def forward(self, x):
306 + def forward(self, x: Tensor) -> Tensor:
307 """
308
309 Args:
308 - x: torch.Tensor
309 - input tensor
310 + x: input tensor
311 """
312 x = self.dw_conv(x)
313 if self.depthwise_seperable_out_channel != 0:
@@ -375,23 +376,23 @@ class ConvModule(nn.Module):
376
377 def __init__(
378 self,
378 - input_dim,
379 - ext_pw_out_channel,
380 - depthwise_seperable_out_channel,
381 - ext_pw_kernel_size,
382 - kernel_size,
383 - depthwise_multiplier,
384 - dropout_rate,
385 - causal=False,
386 - batch_norm=False,
387 - chunk_se=0,
388 - chunk_size=18,
389 - activation="relu",
390 - glu_type="sigmoid",
391 - bias_in_glu=True,
392 - linear_glu_in_convm=False,
393 - export=False,
394 - ):
379 + input_dim: int,
380 + ext_pw_out_channel: int,
381 + depthwise_seperable_out_channel: int,
382 + ext_pw_kernel_size: int,
383 + kernel_size: int,
384 + depthwise_multiplier: int,
385 + dropout_rate: float,
386 + causal: bool = False,
387 + batch_norm: bool = False,
388 + chunk_se: int = 0,
389 + chunk_size: int = 18,
390 + activation: str = "relu",
391 + glu_type: str = "sigmoid",
392 + bias_in_glu: bool = True,
393 + linear_glu_in_convm: bool = False,
394 + export: bool = False,
395 + ) -> None:
396 super().__init__()
397 self.layer_norm = nn.LayerNorm(input_dim)
398 self.input_dim = input_dim
@@ -437,7 +438,7 @@ class ConvModule(nn.Module):
438 self.ln2 = nn.Linear(input_dim * depthwise_multiplier,
439 input_dim)
440
440 - def _add_ext_pw_layer(self):
441 + def _add_ext_pw_layer(self) -> None:
442 """
443 This function is an extension of __init__ function
444 and dedicated to the convolution module creation
@@ -497,12 +498,11 @@ class ConvModule(nn.Module):
498 self.pw_conv_simplify_w = torch.nn.Parameter(torch.ones(3))
499 self.pw_conv_simplify_b = torch.nn.Parameter(torch.zeros(3))
500
500 - def forward(self, x):
501 + def forward(self, x: Tensor) -> Tensor:
502 """ConvModule Forward.
503
504 Args:
504 - x: torch.Tensor
505 - input tensor.
505 + x: input tensor.
506 """
507 x = self.layer_norm(x)
508
@@ -567,21 +567,20 @@ class GLULinear(nn.Module):
567
568 def __init__(
569 self,
570 - input_dim,
571 - output_dim,
572 - glu_type="sigmoid",
573 - bias_in_glu=True,
574 - ):
570 + input_dim: int,
571 + output_dim: int,
572 + glu_type: str = "sigmoid",
573 + bias_in_glu: bool = True,
574 + ) -> None:
575 super().__init__()
576 self.linear = nn.Linear(input_dim, output_dim * 2, bias_in_glu)
577 self.glu_act = GLU(-1, glu_type)
578
579 - def forward(self, x):
579 + def forward(self, x: Tensor) -> Tensor:
580 """GLULinear forward
581
582 Args:
583 - x: torch.Tensor
584 - inpute tensor.
583 + x: input tensor.
584 """
585 x = self.linear(x)
586 return self.glu_act(x)
@@ -609,12 +608,12 @@ class FeedForward(nn.Module):
608
609 def __init__(
610 self,
612 - d_model,
613 - d_inner,
614 - dropout_rate,
615 - activation="sigmoid",
616 - bias_in_glu=True,
617 - ):
611 + d_model: int,
612 + d_inner: int,
613 + dropout_rate: float,
614 + activation: str = "sigmoid",
615 + bias_in_glu: bool = True,
616 + ) -> None:
617 super().__init__()
618 self.d_model = d_model
619 self.d_inner = d_inner
@@ -628,12 +627,11 @@ class FeedForward(nn.Module):
627 nn.Dropout(dropout_rate),
628 )
629
631 - def forward(self, x):
630 + def forward(self, x: Tensor) -> Tensor:
631 """FeedForward forward function.
632
633 Args:
635 - x: torch.Tensor
636 - input tensor.
634 + x: input tensor.
635 """
636 out = self.net(self.layer_norm(x))
637
@@ -642,14 +640,14 @@ class FeedForward(nn.Module):
640
641 #### positional encoding starts here
642 def _pre_hook(
645 - state_dict,
646 - prefix,
647 - local_metadata,
648 - strict,
649 - missing_keys,
650 - unexpected_keys,
651 - error_msgs,
652 -):
643 + state_dict: dict,
644 + prefix: str,
645 + local_metadata: dict,
646 + strict: bool,
647 + missing_keys: list[str],
648 + unexpected_keys: list[str],
649 + error_msgs: list[str],
650 +) -> None:
651 """Perform pre-hook in load_state_dict for backward compatibility.
652
653 Note:
@@ -708,10 +706,10 @@ class T5RelativeAttentionLogitBias(nn.Module):
706 """
707
708 def __init__(self,
711 - num_heads,
712 - num_buckets=-1,
713 - max_distance=1000,
714 - symmetric=False):
709 + num_heads: int,
710 + num_buckets: int = -1,
711 + max_distance: int = 1000,
712 + symmetric: bool = False) -> None:
713 super().__init__()
714 self.num_heads = num_heads
715 self.num_buckets = num_buckets
@@ -727,7 +725,7 @@ class T5RelativeAttentionLogitBias(nn.Module):
725 self.num_buckets *= 2
726 self.bias_values = nn.Embedding(self.num_buckets, self.num_heads)
727
730 - def forward(self, x):
728 + def forward(self, x: Tensor) -> Tensor:
729 # instantiate bias compatible with shape of x
730 maxpos = x.size(1)
731 context_position = torch.arange(maxpos,
@@ -760,7 +758,7 @@ class T5RelativeAttentionLogitBias(nn.Module):
758
759 return t5_rel_att_bias
760
763 - def _bucket_relative_position(self, relative_position):
761 + def _bucket_relative_position(self, relative_position: Tensor) -> Tensor:
762 # this is a placeholder (isn't tested, likely buggy) using HuggingFace
763 # implem as a reference this also needs to be extended to support
764 # asymmetric +/- ve positions
@@ -810,7 +808,10 @@ class AbsolutePositionalEncoding(nn.Module):
808
809 """
810
813 - def __init__(self, d_model, dropout_rate, max_len=5000):
811 + def __init__(self,
812 + d_model: int,
813 + dropout_rate: float,
814 + max_len: int = 5000) -> None:
815 """Construct an PositionalEncoding object."""
816 super().__init__()
817 self.d_model = d_model
@@ -820,11 +821,11 @@ class AbsolutePositionalEncoding(nn.Module):
821 self.extend_pe(torch.tensor(0.0).expand(1, max_len))
822 self._register_load_state_dict_pre_hook(_pre_hook)
823
823 - def extend_pe(self, x):
824 + def extend_pe(self, x: torch.Tensor) -> None:
825 """Reset the positional encodings.
826
827 Args:
827 - x: torch.Tensor
828 + x: input tensor
829 """
830 if self.pe is not None and self.pe.size(1) >= x.size(1):
831 if self.pe.dtype != x.dtype or self.pe.device != x.device:
@@ -840,15 +841,14 @@ class AbsolutePositionalEncoding(nn.Module):
841 pe = pe.unsqueeze(0)
842 self.pe = pe.to(device=x.device, dtype=x.dtype)
843
843 - def forward(self, x: torch.Tensor):
844 + def forward(self, x: torch.Tensor) -> torch.Tensor:
845 """Add positional encoding.
846
847 Args:
847 - x: torch.Tensor
848 - Input tensor. shape is (batch, time, ...)
848 + x: Input tensor. shape is (batch, time, ...)
849
850 Returns:
851 - torch.Tensor: Encoded tensor. Its shape is (batch, time, ...)
851 + Encoded tensor. Its shape is (batch, time, ...)
852
853 """
854 self.extend_pe(x)
@@ -868,7 +868,7 @@ class MeanVarianceNormLayer(nn.Module):
868 layer input size.
869 """
870
871 - def __init__(self, input_size):
871 + def __init__(self, input_size: int) -> None:
872 super().__init__()
873 self.input_size = input_size
874 self.global_mean = nn.Parameter(torch.zeros(input_size))
@@ -878,8 +878,7 @@ class MeanVarianceNormLayer(nn.Module):
878 """MeanVarianceNormLayer Forward
879
880 Args:
881 - input_: torch.Tensor
882 - input tensor.
881 + input_: input tensor.
882 """
883 return (input_ - self.global_mean) * self.global_invstd
884
@@ -949,7 +948,10 @@ class CausalConv1D(nn.Conv1d):
948 dtype=dtype,
949 )
950
952 - def update_cache(self, x, cache=None):
951 + def update_cache(
952 + self,
953 + x: Tensor,
954 + cache: Optional[Tensor] = None) -> tuple[Tensor, Optional[Tensor]]:
955 if cache is None:
956 new_x = F.pad(x, pad=(self._left_padding, self._right_padding))
957 next_cache = cache
@@ -963,7 +965,11 @@ class CausalConv1D(nn.Conv1d):
965 next_cache = next_cache[:, :, -cache.size(-1):]
966 return new_x, next_cache
967
966 - def forward(self, x, cache=None):
968 + def forward(
969 + self,
970 + x: Tensor,
971 + cache: Optional[Tensor] = None
972 + ) -> Union[Tensor, tuple[Tensor, Optional[Tensor]]]:
973 x, cache = self.update_cache(x, cache=cache)
974 x = super().forward(x)
975 if cache is None:
@@ -1017,8 +1023,8 @@ class CausalConv2D(nn.Conv2d):
1023
1024 def forward(
1025 self,
1020 - x,
1021 - ):
1026 + x: Tensor,
1027 + ) -> Tensor:
1028 x = F.pad(
1029 x,
1030 pad=(self._left_padding, self._right_padding, 0, 0),
@@ -1062,16 +1068,16 @@ class NemoConvSubsampling(torch.nn.Module):
1068 """
1069
1070 def __init__(
1065 - self,
1066 - feat_in,
1067 - feat_out,
1068 - subsampling_factor=4,
1069 - subsampling="dw_striding",
1070 - conv_channels=256,
1071 - subsampling_conv_chunking_factor=1,
1072 - activation=nn.ReLU(), # noqa: B008
1073 - is_causal=False,
1074 - ):
1071 + self,
1072 + feat_in: int,
1073 + feat_out: int,
1074 + subsampling_factor: int = 4,
1075 + subsampling: str = "dw_striding",
1076 + conv_channels: int = 256,
1077 + subsampling_conv_chunking_factor: int = 1,
1078 + activation: torch.nn.Module = nn.ReLU(), # noqa: B008
1079 + is_causal: bool = False,
1080 + ) -> None:
1081 super().__init__()
1082 self._subsampling = subsampling
1083 self._conv_channels = conv_channels
@@ -1328,28 +1334,25 @@ class NemoConvSubsampling(torch.nn.Module):
1334
1335 self.conv = torch.nn.Sequential(*layers)
1336
1331 - def get_sampling_frames(self):
1337 + def get_sampling_frames(self) -> list[int]:
1338 return [1, self.subsampling_factor]
1339
1334 - def get_streaming_cache_size(self):
1340 + def get_streaming_cache_size(self) -> list[int]:
1341 return [0, self.subsampling_factor + 1]
1342
1337 - def forward(self, x, mask):
1343 + def forward(self, x: Tensor,
1344 + mask: Optional[Tensor]) -> tuple[Tensor, Optional[Tensor]]:
1345 """
1346 Forward method for NeMo subsampling.
1347
1348 Args:
1342 - x[Batch, Time, Filters]: torch.Tensor
1343 - input tensor
1344 - x_mask: torch.Tensor
1345 - input mask
1349 + x: input tensor
1350 + mask: input mask
1351
1352 Returns:
1348 - x: torch.Tensor
1349 - Resulting tensor from subsampling (B, T //
1353 + x: Resulting tensor from subsampling (B, T //
1354 time_reduction_factor, feat_out)
1351 - pad_mask: torch.Tensor
1352 - tensor of padded hidden state sequences (B, 1, T //
1355 + pad_mask: tensor of padded hidden state sequences (B, 1, T //
1356 time_reduction_factor)
1357 """
1358 x = x.unsqueeze(1) if self.conv2d_subsampling else x.transpose(1, 2)
@@ -1403,7 +1406,7 @@ class NemoConvSubsampling(torch.nn.Module):
1406 padding_length.size(0), -1) < padding_length.unsqueeze(1)
1407 return x, pad_mask.unsqueeze(1)
1408
1406 - def reset_parameters(self):
1409 + def reset_parameters(self) -> None:
1410 # initialize weights
1411 if self._subsampling == "dw_striding":
1412 with torch.no_grad():
@@ -1433,7 +1436,7 @@ class NemoConvSubsampling(torch.nn.Module):
1436 torch.nn.init.uniform_(self.out.weight, -fc_scale, fc_scale)
1437 torch.nn.init.uniform_(self.out.bias, -fc_scale, fc_scale)
1438
1436 - def conv_split_by_batch(self, x):
1439 + def conv_split_by_batch(self, x: Tensor) -> tuple[Tensor, bool]:
1440 """Tries to split input by batch, run conv and concat results"""
1441 b, _, _, _ = x.size()
1442 if b == 1: # can't split if batch size is 1
@@ -1460,7 +1463,7 @@ class NemoConvSubsampling(torch.nn.Module):
1463 True,
1464 )
1465
1463 - def conv_split_by_channel(self, x):
1466 + def conv_split_by_channel(self, x: Tensor) -> Tensor:
1467 """For dw convs, tries to split input by time, run conv and concat
1468 results"""
1469 x = self.conv[0](x) # full conv2D
@@ -1500,7 +1503,8 @@ class NemoConvSubsampling(torch.nn.Module):
1503 x = self.conv[i * 3 + 4](x) # activation
1504 return x
1505
1503 - def channel_chunked_conv(self, conv, chunk_size, x):
1506 + def channel_chunked_conv(self, conv: torch.nn.Module, chunk_size: int,
1507 + x: Tensor) -> Tensor:
1508 """Performs channel chunked convolution"""
1509
1510 ind = 0
@@ -1541,7 +1545,7 @@ class NemoConvSubsampling(torch.nn.Module):
1545 return torch.cat(out_chunks, 1)
1546
1547 def change_subsampling_conv_chunking_factor(
1544 - self, subsampling_conv_chunking_factor: int):
1548 + self, subsampling_conv_chunking_factor: int) -> None:
1549 if (subsampling_conv_chunking_factor != -1
1550 and subsampling_conv_chunking_factor != 1
1551 and subsampling_conv_chunking_factor % 2 != 0):
@@ -1552,12 +1556,12 @@ class NemoConvSubsampling(torch.nn.Module):
1556 self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
1557
1558
1555 -def calc_length(lengths,
1556 - all_paddings,
1557 - kernel_size,
1558 - stride,
1559 - ceil_mode,
1560 - repeat_num=1):
1559 +def calc_length(lengths: Tensor,
1560 + all_paddings: int,
1561 + kernel_size: int,
1562 + stride: int,
1563 + ceil_mode: bool,
1564 + repeat_num: int = 1) -> Tensor:
1565 """Calculates the output length of a Tensor passed through a convolution or
1566 max pooling layer"""
1567 add_pad: float = all_paddings - kernel_size
@@ -1573,11 +1577,11 @@ def calc_length(lengths,
1577 class AttModule(nn.Module):
1578 """Attention abstraction module"""
1579
1576 - def __init__(self):
1580 + def __init__(self) -> None:
1581 super().__init__()
1582 self.export_mode = False
1583
1580 - def set_export(self, mode=True):
1584 + def set_export(self, mode: bool = True) -> None:
1585 """set the export mode"""
1586 self.export_mode = mode
1587
@@ -1591,14 +1595,10 @@ class AttModule(nn.Module):
1595 """AttModule forward
1596
1597 Args:
1594 - x: torch.Tensor
1595 - input tensor.
1596 - memory: torch.Tensor, optional
1597 - memory tensor.
1598 - pos_emb: torch.Tensor, optional
1599 - positional encoder embedding.
1600 - att_mask: torch.Tensor, optional
1601 - attention mask tensor.
1598 + x: input tensor.
1599 + memory: memory tensor.
1600 + pos_emb: positional encoder embedding.
1601 + att_mask: attention mask tensor.
1602 """
1603 return x, memory, pos_emb, att_mask
1604
@@ -1606,15 +1606,15 @@ class AttModule(nn.Module):
1606 class AttBlock(BlockBase, AttModule):
1607 """Attention Block module to support both Attention and Block module."""
1608
1609 - def memory_dims(self, max_len=False):
1609 + def memory_dims(self, max_len: bool = False) -> tuple[int, int]:
1610 """memory dimensions"""
1611 return (1, self.input_size)
1612
1613
1614 def masked_softmax(
1615 - scores,
1615 + scores: Tensor,
1616 mask: Optional[Tensor],
1617 -):
1617 +) -> Tensor:
1618 if mask is not None:
1619 mask = mask.unsqueeze(1).eq(0) # (batch, 1, time1, time2)
1620 scores = scores.masked_fill(mask, -torch.inf)
@@ -1636,10 +1636,6 @@ class MultiHeadedAttention(nn.Module):
1636 input size features.
1637 dropout_rate: float
1638 dropout rate.
1639 - use_LN: bool
1640 - apply layer norm or not
1641 - dropout_at_output: bool
1642 - whether to apply dropout at output
1639 attention_inner_dim: int, optional
1640 the attention dimension used in the class,
1641 it can be different from the input dimension n_feat.
@@ -1666,16 +1662,16 @@ class MultiHeadedAttention(nn.Module):
1662
1663 def __init__(
1664 self,
1669 - n_head,
1670 - n_feat,
1671 - dropout_rate,
1672 - attention_inner_dim=-1,
1673 - glu_type="swish",
1674 - bias_in_glu=True,
1675 - use_pt_scaled_dot_product_attention=False,
1676 - n_value=-1,
1665 + n_head: int,
1666 + n_feat: int,
1667 + dropout_rate: float,
1668 + attention_inner_dim: int = -1,
1669 + glu_type: str = "swish",
1670 + bias_in_glu: bool = True,
1671 + use_pt_scaled_dot_product_attention: bool = False,
1672 + n_value: int = -1,
1673 group_size: int = 1,
1678 - ):
1674 + ) -> None:
1675 super().__init__()
1676 if n_value == -1:
1677 n_value = n_feat
@@ -1718,28 +1714,22 @@ class MultiHeadedAttention(nn.Module):
1714 query: Tensor,
1715 key: Tensor,
1716 value: Tensor,
1721 - pos_k: Tensor,
1722 - pos_v: Tensor,
1717 + pos_k: Optional[Tensor],
1718 + pos_v: Optional[Tensor],
1719 mask: Optional[Tensor],
1720 relative_attention_bias: Optional[Tensor] = None,
1725 - ):
1721 + ) -> Tensor:
1722 """Compute 'Scaled Dot Product Attention'.
1723
1724 Args:
1729 - query: torch.Tensor
1730 - query tensor (batch, time1, size)
1731 - key: torch.Tensor
1732 - key tensor (batch, time2, size)
1733 - value: torch.Tensor
1734 - value tensor (batch, time1, size)
1735 - pos_k: torch.Tensor
1736 - key tensor used for relative positional embedding.
1737 - pos_v: torch.Tensor
1738 - value tensor used for relative positional embedding.
1739 - mask: torch.Tensor
1740 - mask tensor (batch, time1, time2)
1741 - relative_attention_bias: torch.Tensor
1742 - bias added to attention logits w.r.t. relative positions
1725 + query: query tensor (batch, time1, size)
1726 + key: key tensor (batch, time2, size)
1727 + value: value tensor (batch, time1, size)
1728 + pos_k: key tensor used for relative positional embedding.
1729 + pos_v: value tensor used for relative positional embedding.
1730 + mask: mask tensor (batch, time1, time2)
1731 + relative_attention_bias: bias added to attention logits w.r.t.
1732 + relative positions
1733 (1, n_head, time1, time2)
1734 """
1735 n_batch = query.size(0)
@@ -1832,20 +1822,20 @@ class MultiSequential(torch.nn.Sequential):
1822 """Multi-input multi-output torch.nn.Sequential"""
1823
1824 @torch.jit.ignore
1835 - def forward(self, *args):
1825 + def forward(self, *args) -> tuple:
1826 """Forward method implementation."""
1827 for m in self:
1828 args = m(*args)
1829 return args
1830
1831
1842 -def get_offset(input_layer: str, time_reduction: int):
1832 +def get_offset(input_layer: str, time_reduction: int) -> int:
1833 """Get an offset. We will use the offset for determining #frames of a
1834 subsampled feature.
1835
1836 Args:
1847 - input_layer (str): Type of an input layer
1848 - time_reduction (int): time reduction factor for downsampling a feature
1837 + input_layer: Type of an input layer
1838 + time_reduction: time reduction factor for downsampling a feature
1839 Returns:
1840 int: offset
1841 """
@@ -1858,13 +1848,14 @@ def get_offset(input_layer: str, time_reduction: int):
1848 return 0
1849
1850
1861 -def unfold_tensor(xs_pad, max_seq_len):
1851 +def unfold_tensor(xs_pad: Tensor, max_seq_len: int) -> Tensor:
1852 """
1853 For a given tensor with shape of (N, T, D), if sequence length T is
1854 longer than max_seq_len, this function unfold it to a
1855 (NT', max_seq_len, D) where T' is T // max_seq_len.
1856 Args:
1867 - xs_pad: N, T, D
1857 + xs_pad: input tensor with shape (N, T, D)
1858 + max_seq_len: maximum sequence length
1859 """
1860 _, _, D = xs_pad.shape
1861 xs_pad = xs_pad.transpose(-1, -2) # convert to N, D, T
vllm/model_executor/models/qwen2_5_vl.py
+2 -14
@@ -1193,21 +1193,9 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1193 input_ids: Flattened (concatenated) input_ids corresponding to a
1194 batch.
1195 positions: Flattened (concatenated) position ids corresponding to a
1196 - batch.
1197 - **NOTE**: If mrope is enabled (default setting for Qwen2.5-VL
1198 - opensource models), the shape will be `(3, seq_len)`,
1196 + batch. **NOTE**: If mrope is enabled (default setting for
1197 + Qwen2.5-VL opensource models), the shape will be `(3, seq_len)`,
1198 otherwise it will be `(seq_len,).
1200 - pixel_values: Pixel values to be fed to a model.
1201 - `None` if no images are passed.
1202 - image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in LLM.
1203 - `None` if no images are passed.
1204 - pixel_values_videos: Pixel values of videos to be fed to a model.
1205 - `None` if no videos are passed.
1206 - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in LLM.
1207 - `None` if no videos are passed.
1208 - second_per_grid_ts: Tensor `(num_videos)` of video time interval (
1209 - in seconds) for each grid along the temporal dimension in the
1210 - 3D position IDs. `None` if no videos are passed.
1199 """
1200
1201 if intermediate_tensors is not None:
vllm/model_executor/models/zamba2.py
+7 -13
@@ -9,7 +9,7 @@ model alternates between state space model layers and attention-based layers.
9 """
10 from collections.abc import Iterable
11 from itertools import cycle
12 -from typing import Optional, Union
12 +from typing import Any, Optional, Union
13
14 import torch
15 from torch import nn
@@ -528,8 +528,6 @@ class Zamba2MambaDecoderLayer(nn.Module):
528 hidden_states: Input tensor [batch_size, seq_len, hidden_size]
529 mamba_cache_params: Parameters for Mamba's state caches
530 (one for conv, one for ssm)
531 - sequence_idx: Index tensor for identifying sequences in batch
532 - Required for proper chunked processing in prefill
531 transformer_hidden_states: Optional output from transformer path
532 Added to input if provided (used in hybrid architecture)
533 positions: Optional position IDs (unused in Mamba)
@@ -591,8 +589,6 @@ class Zamba2HybridLayer(nn.Module):
589
590 Args:
591 shared_transformer: Transformer decoder layer for attention pathway
594 - linear: Linear projection for transformer output before Mamba
595 - mamba: Mamba decoder layer for state space pathway
592 """
593 super().__init__()
594 self.block_idx = block_idx
@@ -630,8 +626,6 @@ class Zamba2HybridLayer(nn.Module):
626 positions: Position IDs for positional embeddings
627 mamba_cache_params: Parameters for Mamba's state caches
628 (one for conv, one for ssm)
633 - sequence_idx: Indices for identifying sequences in batch,
634 - required for proper chunked processing in prefill
629
630 Returns:
631 Output tensor combining transformer and Mamba representations
@@ -915,8 +909,8 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
909 prefix: Optional prefix for parameter names
910
911 Raises:
918 - AssertionError: If prefix caching is enabled (not supported by
919 - Mamba)
912 + AssertionError: If prefix caching is enabled
913 + (not supported by Mamba)
914 """
915 config = vllm_config.model_config.hf_config
916 cache_config = vllm_config.cache_config
@@ -971,7 +965,7 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
965 input_ids: torch.Tensor,
966 positions: torch.Tensor,
967 inputs_embeds: Optional[torch.Tensor] = None,
974 - **kwargs) -> torch.Tensor:
968 + **kwargs: Any) -> torch.Tensor:
969 """Forward pass through the model.
970
971 Args:
@@ -1012,9 +1006,9 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
1006
1007 return hidden_states
1008
1015 - def copy_inputs_before_cuda_graphs(self, input_buffers: dict[str,
1016 - torch.Tensor],
1017 - **kwargs) -> dict[str, torch.Tensor]:
1009 + def copy_inputs_before_cuda_graphs(
1010 + self, input_buffers: dict[str, torch.Tensor],
1011 + **kwargs: Any) -> dict[str, torch.Tensor]:
1012 """Copy inputs before CUDA graph capture.
1013
1014 Args: