26d0419309 · tduncan/Nvidia-Optimization-Stack-VLLM

pyproject.toml

-1

@@ -77,7 +77,6 @@ exclude = [

77 "vllm/engine/**/*.py" = ["UP006", "UP035"]

78 "vllm/executor/**/*.py" = ["UP006", "UP035"]

79 "vllm/model_executor/model_loader/**/*.py" = ["UP006", "UP035"]

80 -"vllm/model_executor/models/**/*.py" = ["UP006", "UP035"]

80 "vllm/prompt_adapter/**/*.py" = ["UP006", "UP035"]

81 "vllm/spec_decode/**/*.py" = ["UP006", "UP035"]

82 "vllm/worker/**/*.py" = ["UP006", "UP035"]

vllm/model_executor/models/arctic.py

+7 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Inference-only Snowflake Arctic model."""

3 -from typing import Iterable, List, Optional, Set, Tuple, Union

3 +from collections.abc import Iterable

4 +from typing import Optional, Union

5

6 import torch

7 from torch import nn

@@ -458,8 +459,8 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):

459 sampling_metadata)

460 return logits

461

461 - def load_weights(self, weights: Iterable[Tuple[str,

462 - torch.Tensor]]) -> Set[str]:

462 + def load_weights(self, weights: Iterable[tuple[str,

463 + torch.Tensor]]) -> set[str]:

464 stacked_params_mapping = [

465 # (param_name, shard_name, shard_id)

466 ("qkv_proj", "q_proj", "q"),

@@ -467,8 +468,8 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):

468 ("qkv_proj", "v_proj", "v"),

469 ]

470

470 - mlp_params_mapping: List[Tuple[str, str, int]] = []

471 - expert_params_mapping: List[Tuple[str, str, int]] = []

471 + mlp_params_mapping: list[tuple[str, str, int]] = []

472 + expert_params_mapping: list[tuple[str, str, int]] = []

473 num_layers = self.config.num_hidden_layers

474

475 for layer in range(num_layers):

@@ -497,7 +498,7 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):

498 ("ws", f"experts.{expert_id}.w3.weight", expert_id))

499

500 params_dict = dict(self.named_parameters())

500 - loaded_params: Set[str] = set()

501 + loaded_params: set[str] = set()

502

503 logger.info(

504 "It will take ~10 minutes loading from the 16-bit weights. "

vllm/model_executor/models/aria.py

+10 -10

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 from collections.abc import Iterable, Mapping, Sequence

3 -from typing import List, Optional, Set, Tuple, TypedDict, Union

3 +from typing import Optional, TypedDict, Union

4

5 import torch

6 import torch.nn as nn

@@ -66,8 +66,8 @@ class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant):

66 # Identity layer

67 self.post_layernorm = nn.Identity()

68

69 - def load_weights(self, weights: Iterable[Tuple[str,

70 - torch.Tensor]]) -> Set[str]:

69 + def load_weights(self, weights: Iterable[tuple[str,

70 + torch.Tensor]]) -> set[str]:

71 stacked_params_mapping = [

72 # (param_name, shard_name, shard_id)

73 ("qkv_proj", "q_proj", "q"),

@@ -75,7 +75,7 @@ class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant):

75 ("qkv_proj", "v_proj", "v"),

76 ]

77 params_dict = dict(self.named_parameters())

78 - loaded_params: Set[str] = set()

78 + loaded_params: set[str] = set()

79 for name, loaded_weight in weights:

80

81 # NOTE: post_layernorm is not used in Aria

@@ -326,8 +326,8 @@ class AriaTextModel(LlamaModel, SupportsQuant):

326

327 # Adapted from LlamaModel.load_weights with the modification of adding

328 # the expert weights mapping to `stacked_params_mapping`

329 - def load_weights(self, weights: Iterable[Tuple[str,

330 - torch.Tensor]]) -> Set[str]:

329 + def load_weights(self, weights: Iterable[tuple[str,

330 + torch.Tensor]]) -> set[str]:

331 stacked_params_mapping = [

332 # (param_name, shard_name, shard_id)

333 (".qkv_proj", ".q_proj", "q"),

@@ -339,7 +339,7 @@ class AriaTextModel(LlamaModel, SupportsQuant):

339 ("experts.w2_weight", "experts.fc2.weight", 'w2'),

340 ]

341 params_dict = dict(self.named_parameters())

342 - loaded_params: Set[str] = set()

342 + loaded_params: set[str] = set()

343 for name, loaded_weight in weights:

344 if "rotary_emb.inv_freq" in name:

345 continue

@@ -528,7 +528,7 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

528 self.vocab_size, logit_scale)

529

530 def _validate_image_sizes(

531 - self, images: List[torch.Tensor]) -> List[torch.Tensor]:

531 + self, images: list[torch.Tensor]) -> list[torch.Tensor]:

532 if not all(img.shape == images[0].shape for img in images):

533 raise ValueError("All images must be the same size")

534 return images

@@ -578,7 +578,7 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

578

579 def _process_image_input(

580 self, image_input: AriaImagePixelInputs

581 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

582 assert self.vision_tower is not None

583

584 pixel_values = image_input['pixel_values']

@@ -651,6 +651,6 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

651 sampling_metadata)

652 return logits

653

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

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

655 loader = AutoWeightsLoader(self)

656 loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/aya_vision.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0 Adapted from

2 # https://github.com/huggingface/transformers/tree/main/src/transformers/models/aya_vision

3 -from typing import (Iterable, Literal, Mapping, Optional, Sequence, Set, Tuple,

4 - TypedDict, Union, cast)

3 +from collections.abc import Iterable, Mapping, Sequence

4 +from typing import Literal, Optional, TypedDict, Union, cast

5

6 import torch

7 from torch import nn

@@ -315,8 +315,8 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,

315 def dtype(self):

316 return next(self.parameters()).dtype

317

318 - def load_weights(self, weights: Iterable[Tuple[str,

319 - torch.Tensor]]) -> Set[str]:

318 + def load_weights(self, weights: Iterable[tuple[str,

319 + torch.Tensor]]) -> set[str]:

320 loader = AutoWeightsLoader(self)

321 return loader.load_weights(weights)

322

vllm/model_executor/models/baichuan.py

+8 -7

@@ -20,7 +20,8 @@

20 # limitations under the License.

21 """Inference-only BaiChuan model compatible with HuggingFace weights."""

22 import math

23 -from typing import Iterable, Optional, Set, Tuple, Union

23 +from collections.abc import Iterable

24 +from typing import Optional, Union

25

26 import torch

27 from torch import nn

@@ -230,7 +231,7 @@ class BaiChuanDecoderLayer(nn.Module):

231 positions: torch.Tensor,

232 hidden_states: torch.Tensor,

233 residual: Optional[torch.Tensor],

233 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

235 # Self Attention

236 if residual is None:

237 residual = hidden_states

@@ -320,15 +321,15 @@ class BaiChuanModel(nn.Module):

321 hidden_states, _ = self.norm(hidden_states, residual)

322 return hidden_states

323

323 - def load_weights(self, weights: Iterable[Tuple[str,

324 - torch.Tensor]]) -> Set[str]:

324 + def load_weights(self, weights: Iterable[tuple[str,

325 + torch.Tensor]]) -> set[str]:

326 stacked_params_mapping = [

327 # (param_name, shard_name, shard_id)

328 ("gate_up_proj", "gate_proj", 0),

329 ("gate_up_proj", "up_proj", 1),

330 ]

331 params_dict = dict(self.named_parameters())

331 - loaded_params: Set[str] = set()

332 + loaded_params: set[str] = set()

333 for name, loaded_weight in weights:

334 if "rotary_emb.inv_freq" in name:

335 continue

@@ -421,8 +422,8 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP,

422 sampling_metadata)

423 return logits

424

424 - def load_weights(self, weights: Iterable[Tuple[str,

425 - torch.Tensor]]) -> Set[str]:

425 + def load_weights(self, weights: Iterable[tuple[str,

426 + torch.Tensor]]) -> set[str]:

427 loader = AutoWeightsLoader(self)

428 return loader.load_weights(weights)

429

vllm/model_executor/models/bamba.py

+8 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Inference-only Bamba model."""

3 # Added by the IBM Team, 2024

4 -from typing import Iterable, Optional, Set, Tuple

4 +from collections.abc import Iterable

5 +from typing import Optional

6

7 import torch

8 from torch import nn

@@ -355,8 +356,8 @@ class BambaModel(nn.Module):

356 hidden_states, _ = self.final_layernorm(hidden_states, residual)

357 return hidden_states

358

358 - def load_weights(self, weights: Iterable[Tuple[str,

359 - torch.Tensor]]) -> Set[str]:

359 + def load_weights(self, weights: Iterable[tuple[str,

360 + torch.Tensor]]) -> set[str]:

361 stacked_params_mapping = [

362 # (param_name, shard_name, shard_id)

363 ("qkv_proj", "q_proj", "q"),

@@ -367,7 +368,7 @@ class BambaModel(nn.Module):

368 ]

369

370 params_dict = dict(self.named_parameters())

370 - loaded_params: Set[str] = set()

371 + loaded_params: set[str] = set()

372 for name, loaded_weight in weights:

373 if "rotary_emb.inv_freq" in name:

374 continue

@@ -495,7 +496,7 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

496 return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

497

498 def _get_mamba_cache_shape(

498 - self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

499 + self) -> tuple[tuple[int, int], tuple[int, int]]:

500 world_size = get_tensor_model_parallel_world_size()

501 hidden_size = self.config.hidden_size

502

@@ -535,7 +536,7 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

536 sampling_metadata)

537 return logits

538

538 - def load_weights(self, weights: Iterable[Tuple[str,

539 - torch.Tensor]]) -> Set[str]:

539 + def load_weights(self, weights: Iterable[tuple[str,

540 + torch.Tensor]]) -> set[str]:

541 loader = AutoWeightsLoader(self)

542 return loader.load_weights(weights)

vllm/model_executor/models/bart.py

+4 -3

@@ -19,7 +19,8 @@

19 # limitations under the License.

20 """PyTorch BART model."""

21 import math

22 -from typing import Iterable, Optional, Tuple

22 +from collections.abc import Iterable

23 +from typing import Optional

24

25 import torch

26 from torch import nn

@@ -859,14 +860,14 @@ class BartForConditionalGeneration(nn.Module, SupportsV0Only, SupportsQuant):

860 def _rename_stacked_param(

861 self,

862 name: str,

862 - ) -> Tuple[str, Optional[str]]:

863 + ) -> tuple[str, Optional[str]]:

864 for key, mapping in self.stacked_params_mapping.items():

865 if key in name:

866 name = name.replace(key, mapping["param_name"])

867 return name, mapping["shard_id"]

868 return name, None

869

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

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

871

872 model_params_dict = dict(self.model.named_parameters())

873 top_params_dict = dict(self.named_parameters())

vllm/model_executor/models/bert.py

+7 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, Optional, Set, Tuple

3 +from collections.abc import Iterable

4 +from typing import Optional

5

6 import torch

7 from torch import nn

@@ -349,8 +350,8 @@ class BertModel(nn.Module, SupportsQuant):

350 token_type_ids=token_type_ids)

351 return self.encoder(hidden_states)

352

352 - def load_weights(self, weights: Iterable[Tuple[str,

353 - torch.Tensor]]) -> Set[str]:

353 + def load_weights(self, weights: Iterable[tuple[str,

354 + torch.Tensor]]) -> set[str]:

355 stacked_params_mapping = [

356 # (param_name, shard_name, shard_id)

357 ("qkv_proj", "query", "q"),

@@ -359,7 +360,7 @@ class BertModel(nn.Module, SupportsQuant):

360 ]

361

362 params_dict = dict(self.named_parameters())

362 - loaded_params: Set[str] = set()

363 + loaded_params: set[str] = set()

364 for name, loaded_weight in weights:

365 if self.pooler is None and "pooler" in name:

366 continue

@@ -424,7 +425,7 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):

425 ) -> Optional[PoolerOutput]:

426 return self._pooler(hidden_states, pooling_metadata)

427

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

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

429 weights = self.hf_to_vllm_mapper.apply(weights)

430 weights = ((name, data) for name, data in weights

431 if not name.startswith("lm_head."))

@@ -472,7 +473,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,

473 self._pooler = CrossEncodingPooler(config, self.classifier,

474 self.bert.pooler)

475

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

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

477

478 self_weights = []

479

vllm/model_executor/models/bert_with_rope.py

+13 -12

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Iterable, Optional, Set, Tuple

2 +from collections.abc import Iterable

3 +from typing import Optional

4

5 import torch

6 from torch import nn

@@ -208,7 +209,7 @@ class NomicRouter(nn.Module):

209

210 def forward(

211 self, x: torch.Tensor

211 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:

212 + ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:

213 weights = self.layer(x.view(-1, x.shape[-1]))[0].softmax(

214 dim=-1, dtype=torch.float32)

215 top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)

@@ -428,8 +429,8 @@ class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):

429 token_type_ids=token_type_ids)

430 return self.encoder(positions, hidden_states)

431

431 - def load_weights(self, weights: Iterable[Tuple[str,

432 - torch.Tensor]]) -> Set[str]:

432 + def load_weights(self, weights: Iterable[tuple[str,

433 + torch.Tensor]]) -> set[str]:

434 weights = self.hf_to_vllm_mapper.apply(weights)

435

436 if self.config.hidden_act in ["silu", "geglu"]:

@@ -442,7 +443,7 @@ class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):

443 stacked_params_mapping = []

444

445 params_dict = dict(self.named_parameters())

445 - loaded_params: Set[str] = set()

446 + loaded_params: set[str] = set()

447 for name, loaded_weight in weights:

448 if "pooler" in name:

449 continue

@@ -567,7 +568,7 @@ class GteNewModel(BertWithRope):

568 }

569 return config

570

570 - def split_up_gate_proj(self, weights: Iterable[Tuple[str, torch.Tensor]]):

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

572 n = "mlp.up_gate_proj"

573 for name, weight in weights:

574 if n in name:

@@ -578,14 +579,14 @@ class GteNewModel(BertWithRope):

579 yield name, weight

580

581 def ignore_unnecessary_layers(self,

581 - weights: Iterable[Tuple[str, torch.Tensor]]):

582 + weights: Iterable[tuple[str, torch.Tensor]]):

583 for name, weight in weights:

584 if name.startswith("classifier"):

585 continue

586 yield name, weight

587

587 - def load_weights(self, weights: Iterable[Tuple[str,

588 - torch.Tensor]]) -> Set[str]:

588 + def load_weights(self, weights: Iterable[tuple[str,

589 + torch.Tensor]]) -> set[str]:

590 weights = self.ignore_unnecessary_layers(weights)

591 weights = self.split_up_gate_proj(weights)

592 return super().load_weights(weights)

@@ -664,7 +665,7 @@ class JinaRobertaModel(BertWithRope):

665 token_type_ids=token_type_ids)

666

667 @torch.inference_mode()

667 - def jina_merge_lora_weights(self, weights: Iterable[Tuple[str,

668 + def jina_merge_lora_weights(self, weights: Iterable[tuple[str,

669 torch.Tensor]]):

670 # use for jina-embeddings-v3

671 # Merge Lora weights into a single weight tensor.

@@ -707,7 +708,7 @@ class JinaRobertaModel(BertWithRope):

708

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

710

710 - def load_weights(self, weights: Iterable[Tuple[str,

711 - torch.Tensor]]) -> Set[str]:

711 + def load_weights(self, weights: Iterable[tuple[str,

712 + torch.Tensor]]) -> set[str]:

713 weights = self.jina_merge_lora_weights(weights)

714 return super().load_weights(weights)

vllm/model_executor/models/blip.py

+5 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Minimal implementation of BlipVisionModel intended to be only used

3 within a vision language model."""

4 -from typing import Iterable, Optional, Set, Tuple, Union

4 +from collections.abc import Iterable

5 +from typing import Optional, Union

6

7 import torch

8 import torch.nn as nn

@@ -296,8 +297,8 @@ class BlipVisionModel(nn.Module, SupportsQuant):

297

298 return self.post_layernorm(hidden_states)

299

299 - def load_weights(self, weights: Iterable[Tuple[str,

300 - torch.Tensor]]) -> Set[str]:

300 + def load_weights(self, weights: Iterable[tuple[str,

301 + torch.Tensor]]) -> set[str]:

302 stacked_params_mapping = [

303 # (param_name, shard_name, shard_id)

304 ("qkv_proj", "q_proj", "q"),

@@ -305,7 +306,7 @@ class BlipVisionModel(nn.Module, SupportsQuant):

306 ("qkv_proj", "v_proj", "v"),

307 ]

308 params_dict = dict(self.named_parameters())

308 - loaded_params: Set[str] = set()

309 + loaded_params: set[str] = set()

310 layer_count = len(self.encoder.layers)

311

312 for name, loaded_weight in weights:

vllm/model_executor/models/blip2.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 from collections.abc import Iterable, Mapping, Sequence

4 -from typing import Literal, Optional, Set, Tuple, TypedDict, Union

4 +from typing import Literal, Optional, TypedDict, Union

5

6 import torch

7 import torch.nn as nn

@@ -186,7 +186,7 @@ class Blip2QFormerAttention(nn.Module):

186 self,

187 hidden_states: torch.Tensor,

188 encoder_hidden_states: Optional[torch.FloatTensor] = None,

189 - ) -> Tuple[torch.Tensor]:

189 + ) -> tuple[torch.Tensor]:

190 self_output = self.attention(

191 hidden_states,

192 encoder_hidden_states=encoder_hidden_states,

@@ -712,7 +712,7 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

712 return self.language_model.compute_logits(hidden_states,

713 sampling_metadata)

714

715 - def load_weights(self, weights: Iterable[Tuple[str,

716 - torch.Tensor]]) -> Set[str]:

715 + def load_weights(self, weights: Iterable[tuple[str,

716 + torch.Tensor]]) -> set[str]:

717 loader = AutoWeightsLoader(self)

718 return loader.load_weights(weights)

vllm/model_executor/models/bloom.py

+5 -4

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

18 # limitations under the License.

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

20 import math

21 -from typing import Iterable, Optional, Set, Tuple, Union

21 +from collections.abc import Iterable

22 +from typing import Optional, Union

23

24 import torch

25 from torch import nn

@@ -322,10 +323,10 @@ class BloomForCausalLM(nn.Module, SupportsPP, SupportsV0Only, SupportsQuant):

323 sampling_metadata)

324 return logits

325

325 - def load_weights(self, weights: Iterable[Tuple[str,

326 - torch.Tensor]]) -> Set[str]:

326 + def load_weights(self, weights: Iterable[tuple[str,

327 + torch.Tensor]]) -> set[str]:

328 params_dict = dict(self.named_parameters(remove_duplicate=False))

328 - loaded_params: Set[str] = set()

329 + loaded_params: set[str] = set()

330 for name, loaded_weight in weights:

331 if name == "lm_head.weight":

332 continue

vllm/model_executor/models/chameleon.py

+10 -10

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

2

3 from collections.abc import Iterable, Mapping, Sequence

4 from functools import cached_property

5 -from typing import Any, Dict, Literal, Optional, Set, Tuple, TypedDict, Union

5 +from typing import Any, Literal, Optional, TypedDict, Union

6

7 import torch

8 import torch.nn as nn

@@ -229,7 +229,7 @@ class ChameleonAttention(nn.Module):

229 num_heads: int,

230 num_kv_heads: int,

231 rope_theta: float = 10000,

232 - rope_scaling: Optional[Dict[str, Any]] = None,

232 + rope_scaling: Optional[dict[str, Any]] = None,

233 max_position_embeddings: int = 4096,

234 quant_config: Optional[QuantizationConfig] = None,

235 bias: bool = False,

@@ -292,7 +292,7 @@ class ChameleonAttention(nn.Module):

292 prefix=f"{prefix}.attn")

293

294 def _apply_qk_norm(self, q: torch.Tensor,

295 - k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

295 + k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

296 # reshape for layernorm

297 q = q.reshape(-1, self.num_heads, self.head_dim)

298 k = k.reshape(-1, self.num_kv_heads, self.head_dim)

@@ -367,7 +367,7 @@ class ChameleonDecoderLayer(nn.Module):

367 positions: torch.Tensor,

368 hidden_states: torch.Tensor,

369 residual: Optional[torch.Tensor],

370 - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:

370 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

371

372 if residual is None:

373 residual = hidden_states

@@ -438,7 +438,7 @@ class ChameleonSwinDecoderLayer(nn.Module):

438 positions: torch.Tensor,

439 hidden_states: torch.Tensor,

440 residual: Optional[torch.Tensor],

441 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

442

443 residual = hidden_states

444 hidden_states = self.self_attn(

@@ -773,7 +773,7 @@ class ChameleonVQVAE(nn.Module):

773

774 def encode(

775 self, pixel_values: torch.Tensor

776 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

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

777 hidden_states = self.encoder(pixel_values)

778 hidden_states = self.quant_conv(hidden_states)

779 quant, emb_loss, indices = self.quantize(hidden_states)

@@ -786,7 +786,7 @@ class ChameleonImageVocabularyMapping:

786 A class for mapping discrete image tokens from VQGAN to BPE tokens.

787 """

788

789 - def __init__(self, vocab_map: Dict[str, int]):

789 + def __init__(self, vocab_map: dict[str, int]):

790 self.vocab_map = vocab_map

791 self.image_token_id = vocab_map.get("<image>")

792

@@ -1052,8 +1052,8 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,

1052

1053 return logits

1054

1055 - def load_weights(self, weights: Iterable[Tuple[str,

1056 - torch.Tensor]]) -> Set[str]:

1055 + def load_weights(self, weights: Iterable[tuple[str,

1056 + torch.Tensor]]) -> set[str]:

1057 stacked_params_mapping = [

1058 # (param_name, shard_name, shard_id)

1059 (".qkv_proj", ".q_proj", "q"),

@@ -1063,7 +1063,7 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,

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

1064 ]

1065 params_dict = dict(self.named_parameters())

1066 - loaded_params: Set[str] = set()

1066 + loaded_params: set[str] = set()

1067 for name, loaded_weight in weights:

1068 if "rotary_emb.inv_freq" in name:

1069 continue

vllm/model_executor/models/chatglm.py

+6 -5

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

3 # https://github.com/THUDM/ChatGLM2-6B

4 """Inference-only ChatGLM model compatible with THUDM weights."""

5 import json

6 -from typing import Iterable, Optional, Set, Tuple, Union

6 +from collections.abc import Iterable

7 +from typing import Optional, Union

8

9 import torch

10 from torch import nn

@@ -358,15 +359,15 @@ class ChatGLMModel(nn.Module, SupportsQuant):

359

360 return hidden_states

361

361 - def load_weights(self, weights: Iterable[Tuple[str,

362 - torch.Tensor]]) -> Set[str]:

362 + def load_weights(self, weights: Iterable[tuple[str,

363 + torch.Tensor]]) -> set[str]:

364 stacked_params_mapping = [

365 # (param_name, shard_name, shard_id)

366 ("linear_proj.merged_proj", "linear_proj.gate_proj", 0),

367 ("linear_proj.merged_proj", "linear_proj.dense_h_to_4h", 1),

368 ]

369 params_dict = dict(self.named_parameters())

369 - loaded_params: Set[str] = set()

370 + loaded_params: set[str] = set()

371

372 for name, loaded_weight in weights:

373 for (param_name, weight_name, shard_id) in stacked_params_mapping:

@@ -440,7 +441,7 @@ class ChatGLMBaseModel(nn.Module):

441 sampling_metadata)

442 return logits

443

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

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

445 loader = AutoWeightsLoader(self)

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

447

vllm/model_executor/models/clip.py

+5 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Minimal implementation of CLIPVisionModel intended to be only used

3 within a vision language model."""

4 -from typing import Iterable, Optional, Set, Tuple, Union

4 +from collections.abc import Iterable

5 +from typing import Optional, Union

6

7 import torch

8 import torch.nn as nn

@@ -368,8 +369,8 @@ class CLIPVisionModel(nn.Module, SupportsQuant):

369

370 # (TODO) Add prefix argument for filtering out weights to be loaded

371 # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986

371 - def load_weights(self, weights: Iterable[Tuple[str,

372 - torch.Tensor]]) -> Set[str]:

372 + def load_weights(self, weights: Iterable[tuple[str,

373 + torch.Tensor]]) -> set[str]:

374 stacked_params_mapping = [

375 # (param_name, shard_name, shard_id)

376 ("qkv_proj", "q_proj", "q"),

@@ -377,7 +378,7 @@ class CLIPVisionModel(nn.Module, SupportsQuant):

378 ("qkv_proj", "v_proj", "v"),

379 ]

380 params_dict = dict(self.named_parameters())

380 - loaded_params: Set[str] = set()

381 + loaded_params: set[str] = set()

382 layer_count = len(self.vision_model.encoder.layers)

383

384 for name, loaded_weight in weights:

vllm/model_executor/models/commandr.py

+6 -5

@@ -21,7 +21,8 @@

21

22 # This file is based on the LLama model definition file in transformers

23 """PyTorch Cohere model."""

24 -from typing import Iterable, Optional, Set, Tuple, Union

24 +from collections.abc import Iterable

25 +from typing import Optional, Union

26

27 import torch

28 from torch import nn

@@ -259,7 +260,7 @@ class CohereDecoderLayer(nn.Module):

260 positions: torch.Tensor,

261 hidden_states: torch.Tensor,

262 residual: Optional[torch.Tensor],

262 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

264 # Self Attention

265 residual = hidden_states

266 hidden_states, residual = self.input_layernorm(hidden_states, residual)

@@ -404,8 +405,8 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

405

406 return logits

407

407 - def load_weights(self, weights: Iterable[Tuple[str,

408 - torch.Tensor]]) -> Set[str]:

408 + def load_weights(self, weights: Iterable[tuple[str,

409 + torch.Tensor]]) -> set[str]:

410 stacked_params_mapping = [

411 # (param_name, shard_name, shard_id)

412 ("qkv_proj", "q_proj", "q"),

@@ -415,7 +416,7 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

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

417 ]

418 params_dict = dict(self.named_parameters())

418 - loaded_params: Set[str] = set()

419 + loaded_params: set[str] = set()

420 for name, loaded_weight in weights:

421

422 # Skip loading rotary embeddings since vLLM has its own

vllm/model_executor/models/constant_size_cache.py

+6 -6

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 from abc import ABC, abstractmethod

3 -from typing import Any, Dict, List, Tuple

3 +from typing import Any

4

5 import torch

6

@@ -16,7 +16,7 @@ class ConstantSizeCache(ABC):

16 def __init__(self, max_batch_size: int):

17 # Maps between the request id and a dict that maps between the seq_id

18 # and its index inside the cache

19 - self.cache_indices_mapping: Dict[str, Dict[int, int]] = {}

19 + self.cache_indices_mapping: dict[str, dict[int, int]] = {}

20 self.free_cache_indices = list(range(max_batch_size))

21

22 @property

@@ -30,7 +30,7 @@ class ConstantSizeCache(ABC):

30 """Copy cache data from one index to another"""

31 pass

32

33 - def current_run_tensors(self, **kwargs) -> Tuple:

33 + def current_run_tensors(self, **kwargs) -> tuple:

34 """

35 Return the tensors for the current run's conv and ssm state.

36 """

@@ -117,8 +117,8 @@ class ConstantSizeCache(ABC):

117 return self.cache_indices_mapping[cur_rid][seq_id]

118

119 def _prepare_current_run_cache(

120 - self, request_ids_to_seq_ids: Dict[str, list[int]],

121 - finished_requests_ids: List[str]) -> List[int]:

120 + self, request_ids_to_seq_ids: dict[str, list[int]],

121 + finished_requests_ids: list[str]) -> list[int]:

122 return [

123 self._assign_seq_id_to_cache_index(req_id, seq_id,

124 finished_requests_ids)

@@ -127,7 +127,7 @@ class ConstantSizeCache(ABC):

127 ]

128

129 def _release_finished_requests(self,

130 - finished_seq_groups_req_ids: List[str]):

130 + finished_seq_groups_req_ids: list[str]):

131 for req_id in finished_seq_groups_req_ids:

132 if req_id in self.cache_indices_mapping:

133 for seq_id in self.cache_indices_mapping[req_id]:

vllm/model_executor/models/dbrx.py

+5 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, Optional, Set, Tuple, Union

3 +from collections.abc import Iterable

4 +from typing import Optional, Union

5

6 import torch

7 import torch.nn as nn

@@ -414,14 +415,14 @@ class DbrxForCausalLM(nn.Module, SupportsPP):

415 sampling_metadata)

416 return logits

417

417 - def load_weights(self, weights: Iterable[Tuple[str,

418 - torch.Tensor]]) -> Set[str]:

418 + def load_weights(self, weights: Iterable[tuple[str,

419 + torch.Tensor]]) -> set[str]:

420 expert_params_mapping = [(

421 "w13" if weight_name in ["w1", "v1"] else "w2",

422 f"mlp.{weight_name}",

423 ) for weight_name in ["w1", "v1", "w2"]]

424 params_dict = dict(self.named_parameters(remove_duplicate=False))

424 - loaded_params: Set[str] = set()

425 + loaded_params: set[str] = set()

426

427 for name, loaded_weight in weights:

428 if (self.quant_config is not None and

vllm/model_executor/models/deepseek.py

+8 -7

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only Deepseek model."""

25 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Any, Optional, Union

27

28 import torch

29 from torch import nn

@@ -184,7 +185,7 @@ class DeepseekAttention(nn.Module):

185 num_heads: int,

186 num_kv_heads: int,

187 rope_theta: float = 10000,

187 - rope_scaling: Optional[Dict[str, Any]] = None,

188 + rope_scaling: Optional[dict[str, Any]] = None,

189 max_position_embeddings: int = 8192,

190 cache_config: Optional[CacheConfig] = None,

191 quant_config: Optional[QuantizationConfig] = None,

@@ -385,8 +386,8 @@ class DeepseekModel(nn.Module):

386 hidden_states, _ = self.norm(hidden_states, residual)

387 return hidden_states

388

388 - def load_weights(self, weights: Iterable[Tuple[str,

389 - torch.Tensor]]) -> Set[str]:

389 + def load_weights(self, weights: Iterable[tuple[str,

390 + torch.Tensor]]) -> set[str]:

391 stacked_params_mapping = [

392 # (param_name, shard_name, shard_id)

393 ("qkv_proj", "q_proj", "q"),

@@ -397,7 +398,7 @@ class DeepseekModel(nn.Module):

398 ]

399

400 params_dict = dict(self.named_parameters())

400 - loaded_params: Set[str] = set()

401 + loaded_params: set[str] = set()

402 for name, loaded_weight in weights:

403 if "rotary_emb.inv_freq" in name:

404 continue

@@ -478,7 +479,7 @@ class DeepseekForCausalLM(nn.Module, SupportsPP):

479 sampling_metadata)

480 return logits

481

481 - def load_weights(self, weights: Iterable[Tuple[str,

482 - torch.Tensor]]) -> Set[str]:

482 + def load_weights(self, weights: Iterable[tuple[str,

483 + torch.Tensor]]) -> set[str]:

484 loader = AutoWeightsLoader(self)

485 return loader.load_weights(weights)

\ No newline at end of file

vllm/model_executor/models/deepseek_mtp.py

+5 -4

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Iterable, Optional, Set, Tuple

2 +from collections.abc import Iterable

3 +from typing import Optional

4

5 import torch

6 import torch.nn as nn

@@ -176,8 +177,8 @@ class DeepSeekMTP(nn.Module):

177 return self.model.compute_logits(hidden_states, sampling_metadata,

178 spec_step_idx)

179

179 - def load_weights(self, weights: Iterable[Tuple[str,

180 - torch.Tensor]]) -> Set[str]:

180 + def load_weights(self, weights: Iterable[tuple[str,

181 + torch.Tensor]]) -> set[str]:

182 stacked_params_mapping = [

183 ("gate_up_proj", "gate_proj", 0),

184 ("gate_up_proj", "up_proj", 1),

@@ -190,7 +191,7 @@ class DeepSeekMTP(nn.Module):

191 num_experts=self.config.n_routed_experts)

192

193 params_dict = dict(self.named_parameters())

193 - loaded_params: Set[str] = set()

194 + loaded_params: set[str] = set()

195 for name, loaded_weight in weights:

196 if "rotary_emb.inv_freq" in name:

197 continue

vllm/model_executor/models/deepseek_v2.py

+7 -6

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only DeepseekV2/DeepseekV3 model."""

25 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Any, Optional, Union

27

28 import torch

29 from torch import nn

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

201 q_lora_rank: int,

202 kv_lora_rank: int,

203 rope_theta: float = 10000,

203 - rope_scaling: Optional[Dict[str, Any]] = None,

204 + rope_scaling: Optional[dict[str, Any]] = None,

205 max_position_embeddings: int = 8192,

206 cache_config: Optional[CacheConfig] = None,

207 quant_config: Optional[QuantizationConfig] = None,

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

353 q_lora_rank: Optional[int],

354 kv_lora_rank: int,

355 rope_theta: float = 10000,

355 - rope_scaling: Optional[Dict[str, Any]] = None,

356 + rope_scaling: Optional[dict[str, Any]] = None,

357 max_position_embeddings: int = 8192,

358 cache_config: Optional[CacheConfig] = None,

359 quant_config: Optional[QuantizationConfig] = None,

@@ -736,8 +737,8 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):

737 device=device),

738 })

739

739 - def load_weights(self, weights: Iterable[Tuple[str,

740 - torch.Tensor]]) -> Set[str]:

740 + def load_weights(self, weights: Iterable[tuple[str,

741 + torch.Tensor]]) -> set[str]:

742 stacked_params_mapping = [

743 # (param_name, shard_name, shard_id)

744 ("gate_up_proj", "gate_proj", 0),

@@ -753,7 +754,7 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):

754 num_experts=self.config.n_routed_experts)

755

756 params_dict = dict(self.named_parameters())

756 - loaded_params: Set[str] = set()

757 + loaded_params: set[str] = set()

758 for name, loaded_weight in weights:

759 if "rotary_emb.inv_freq" in name:

760 continue

vllm/model_executor/models/deepseek_vl2.py

+9 -9

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

4 """Inference-only Deepseek-VL2 model compatible with HuggingFace weights."""

5 import math

6 from collections.abc import Iterable, Mapping, Sequence

7 -from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union

7 +from typing import Literal, Optional, TypedDict, Union

8

9 import torch

10 import torch.nn as nn

@@ -45,7 +45,7 @@ _IMAGE_TOKEN = "<image>"

45

46 class DeepseekVL2ImagePixelInputs(TypedDict):

47 type: Literal["pixel_values"]

48 - data: Union[torch.Tensor, List[torch.Tensor]]

48 + data: Union[torch.Tensor, list[torch.Tensor]]

49 """

50 Shape: `(batch_size * num_images, num_channels, height, width)`

51 """

@@ -57,7 +57,7 @@ class DeepseekVL2ImagePixelInputs(TypedDict):

57

58 class DeepseekVL2VImageEmbeddingInputs(TypedDict):

59 type: Literal["image_embeds"]

60 - data: Union[torch.Tensor, List[torch.Tensor]]

60 + data: Union[torch.Tensor, list[torch.Tensor]]

61 """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`

62

63 `hidden_size` must match the hidden size of language model backbone.

@@ -394,8 +394,8 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

394 return model

395

396 def _validate_pixel_values(

397 - self, data: Union[torch.Tensor, List[torch.Tensor]]

398 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

397 + self, data: Union[torch.Tensor, list[torch.Tensor]]

398 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

399

400 h = w = self.vision_config.image_size

401 expected_dims = (3, h, w)

@@ -415,8 +415,8 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

415 return data

416

417 def _validate_images_spatial_crop(

418 - self, data: Union[torch.Tensor, List[torch.Tensor]]

419 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

418 + self, data: Union[torch.Tensor, list[torch.Tensor]]

419 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

420 expected_dims = 2

421

422 def _validate_shape(d: torch.Tensor):

@@ -640,8 +640,8 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

640 return self.language_model.compute_logits(hidden_states,

641 sampling_metadata)

642

643 - def load_weights(self, weights: Iterable[Tuple[str,

644 - torch.Tensor]]) -> Set[str]:

643 + def load_weights(self, weights: Iterable[tuple[str,

644 + torch.Tensor]]) -> set[str]:

645

646 loader = AutoWeightsLoader(self)

647 autoloaded_weights = loader.load_weights(weights,

vllm/model_executor/models/eagle.py

+3 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, Optional, Tuple

3 +from collections.abc import Iterable

4 +from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -183,7 +184,7 @@ class EAGLE(nn.Module):

184

185 return logits

186

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

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

188 # This implementation is incompitable with https://huggingface.co/yuhuili/EAGLE-LLaMA3-Instruct-8B

189 # due to missing lm_head weights and its config being that of a

190 # Llama model. Here's a compatible version with the same weights:

vllm/model_executor/models/exaone.py

+10 -9

@@ -24,7 +24,8 @@

24 # limitations under the License.

25 """Inference-only Exaone model compatible with HuggingFace weights."""

26

27 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

27 +from collections.abc import Iterable

28 +from typing import Any, Optional, Union

29

30 import torch

31 from torch import nn

@@ -102,7 +103,7 @@ class ExaoneAttention(nn.Module):

103 num_heads: int,

104 num_kv_heads: int,

105 rope_theta: float = 10000,

105 - rope_scaling: Optional[Dict[str, Any]] = None,

106 + rope_scaling: Optional[dict[str, Any]] = None,

107 max_position_embeddings: int = 8192,

108 quant_config: Optional[QuantizationConfig] = None,

109 bias: bool = False,

@@ -196,7 +197,7 @@ class ExaoneBlockAttention(nn.Module):

197 num_heads: int,

198 num_kv_heads: int,

199 rope_theta: float = 10000,

199 - rope_scaling: Optional[Dict[str, Any]] = None,

200 + rope_scaling: Optional[dict[str, Any]] = None,

201 max_position_embeddings: int = 8192,

202 quant_config: Optional[QuantizationConfig] = None,

203 bias: bool = False,

@@ -282,7 +283,7 @@ class ExaoneDecoderLayer(nn.Module):

283 positions: torch.Tensor,

284 hidden_states: torch.Tensor,

285 residual: Optional[torch.Tensor],

285 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

287 # Self Attention

288 if residual is None:

289 residual = hidden_states

@@ -384,8 +385,8 @@ class ExaoneModel(nn.Module):

385 hidden_states, _ = self.ln_f(hidden_states, residual)

386 return hidden_states

387

387 - def load_weights(self, weights: Iterable[Tuple[str,

388 - torch.Tensor]]) -> Set[str]:

388 + def load_weights(self, weights: Iterable[tuple[str,

389 + torch.Tensor]]) -> set[str]:

390 stacked_params_mapping = [

391 # (param_name, shard_name, shard_id)

392 (".qkv_proj", ".q_proj", "q"),

@@ -395,7 +396,7 @@ class ExaoneModel(nn.Module):

396 (".gate_up_proj", ".c_fc_1", 1),

397 ]

398 params_dict = dict(self.named_parameters())

398 - loaded_params: Set[str] = set()

399 + loaded_params: set[str] = set()

400 for name, loaded_weight in weights:

401 if "rotary_emb.inv_freq" in name:

402 continue

@@ -535,8 +536,8 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

536 sampling_metadata)

537 return logits

538

538 - def load_weights(self, weights: Iterable[Tuple[str,

539 - torch.Tensor]]) -> Set[str]:

539 + def load_weights(self, weights: Iterable[tuple[str,

540 + torch.Tensor]]) -> set[str]:

541 loader = AutoWeightsLoader(

542 self,

543 # With tie_word_embeddings, we can skip lm_head.weight

vllm/model_executor/models/fairseq2_llama.py

+4 -4

@@ -16,7 +16,7 @@

16 # limitations under the License.

17 """Llama model for fairseq2 weights."""

18

19 -from typing import Iterable, Set, Tuple

19 +from collections.abc import Iterable

20

21 import torch

22 from torch.nn import Parameter

@@ -44,8 +44,8 @@ class Fairseq2LlamaForCausalLM(LlamaForCausalLM):

44 f"model.{self.tp_rank}.pt",

45 ]

46

47 - def load_weights(self, weights: Iterable[Tuple[str,

48 - torch.Tensor]]) -> Set[str]:

47 + def load_weights(self, weights: Iterable[tuple[str,

48 + torch.Tensor]]) -> set[str]:

49 # fairseq2's serialization adds a wrapper to usual .pt state_dict's:

50 # { "model_key": my_model_name, "my_model_name": state_dict }

51 # which we first need to unpack

@@ -102,7 +102,7 @@ class Fairseq2LlamaForCausalLM(LlamaForCausalLM):

102 name: str,

103 loaded_weight: torch.Tensor,

104 params: dict[str, Parameter],

105 - ) -> Tuple[str, torch.Tensor]:

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

106 """Reshape fairseq2's weights."""

107

108 def permute(w: torch.Tensor, n_heads: int) -> torch.Tensor:

vllm/model_executor/models/falcon.py

+7 -6

@@ -20,7 +20,8 @@

20 """PyTorch Falcon model."""

21

22 import math

23 -from typing import Iterable, Optional, Set, Tuple, Union

23 +from collections.abc import Iterable

24 +from typing import Optional, Union

25

26 import torch

27 from torch import nn

@@ -394,8 +395,8 @@ class FalconModel(nn.Module):

395 hidden_states = self.ln_f(hidden_states)

396 return hidden_states

397

397 - def load_weights(self, weights: Iterable[Tuple[str,

398 - torch.Tensor]]) -> Set[str]:

398 + def load_weights(self, weights: Iterable[tuple[str,

399 + torch.Tensor]]) -> set[str]:

400 total_num_heads = self.config.num_attention_heads

401 if self.config.new_decoder_architecture:

402 total_num_kv_heads = self.config.num_kv_heads

@@ -405,7 +406,7 @@ class FalconModel(nn.Module):

406 total_num_kv_heads = total_num_heads

407 num_query_heads_per_kv_head = total_num_heads // total_num_kv_heads

408 params_dict = dict(self.named_parameters(remove_duplicate=False))

408 - loaded_params: Set[str] = set()

409 + loaded_params: set[str] = set()

410 for name, loaded_weight in weights:

411 # Skip loading extra bias for GPTQ models.

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

@@ -498,8 +499,8 @@ class FalconForCausalLM(nn.Module, SupportsPP):

499 sampling_metadata)

500 return logits

501

501 - def load_weights(self, weights: Iterable[Tuple[str,

502 - torch.Tensor]]) -> Set[str]:

502 + def load_weights(self, weights: Iterable[tuple[str,

503 + torch.Tensor]]) -> set[str]:

504 loader = AutoWeightsLoader(

505 self,

506 skip_prefixes=(["lm_head."]

vllm/model_executor/models/florence2.py

+12 -12

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

3 import math

4 from collections import OrderedDict

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union

6 +from typing import Literal, Optional, TypedDict, Union

7

8 import torch

9 import torch.nn as nn

@@ -713,8 +713,8 @@ class Florence2LanguageForConditionalGeneration(nn.Module, SupportsV0Only):

713 sampling_metadata)

714 return logits

715

716 - def load_weights(self, weights: Iterable[Tuple[str,

717 - torch.Tensor]]) -> Set[str]:

716 + def load_weights(self, weights: Iterable[tuple[str,

717 + torch.Tensor]]) -> set[str]:

718 stacked_params_mapping = [

719 # (param_name, shard_name, shard_id)

720 ("qkv_proj", "q_proj", "q"),

@@ -723,7 +723,7 @@ class Florence2LanguageForConditionalGeneration(nn.Module, SupportsV0Only):

723 ]

724

725 params_dict = dict(self.named_parameters())

726 - loaded_params: Set[str] = set()

726 + loaded_params: set[str] = set()

727 for name, loaded_weight in weights:

728 for (param_name, weight_name, shard_id) in stacked_params_mapping:

729 if weight_name not in name:

@@ -922,8 +922,8 @@ class Florence2ForConditionalGeneration(nn.Module, SupportsMultiModal,

922 'Florence2 only supports COSINE as temporal embedding.')

923

924 def _validate_pixel_values(

925 - self, data: Union[torch.Tensor, List[torch.Tensor]]

926 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

925 + self, data: Union[torch.Tensor, list[torch.Tensor]]

926 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

927

928 size = self.processor_config["size"]

929 h, w = size["height"], size["width"]

@@ -944,12 +944,12 @@ class Florence2ForConditionalGeneration(nn.Module, SupportsMultiModal,

944 return data

945

946 def _parse_and_validate_image_input(self, **kwargs: object):

947 - pixel_values: Optional[Union[List[List[torch.Tensor]],

948 - List[torch.Tensor],

947 + pixel_values: Optional[Union[list[list[torch.Tensor]],

948 + list[torch.Tensor],

949 torch.Tensor]] = kwargs.pop(

950 "pixel_values", None)

951 - image_embeds: Optional[Union[List[List[torch.Tensor]],

952 - List[torch.Tensor],

951 + image_embeds: Optional[Union[list[list[torch.Tensor]],

952 + list[torch.Tensor],

953 torch.Tensor]] = kwargs.pop(

954 "image_embeds", None)

955

@@ -1096,7 +1096,7 @@ class Florence2ForConditionalGeneration(nn.Module, SupportsMultiModal,

1096 return self.language_model.compute_logits(hidden_states,

1097 sampling_metadata)

1098

1099 - def load_weights(self, weights: Iterable[Tuple[str,

1100 - torch.Tensor]]) -> Set[str]:

1099 + def load_weights(self, weights: Iterable[tuple[str,

1100 + torch.Tensor]]) -> set[str]:

1101 loader = AutoWeightsLoader(self)

1102 return loader.load_weights(weights)

vllm/model_executor/models/fuyu.py

+3 -3

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

18 """ PyTorch Fuyu model."""

19 import math

20 from collections.abc import Iterable, Mapping, Sequence

21 -from typing import Literal, Optional, Set, Tuple, TypedDict

21 +from typing import Literal, Optional, TypedDict

22

23 import torch

24 import torch.nn as nn

@@ -382,7 +382,7 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

382 self.language_model.lm_head, hidden_states, sampling_metadata)

383 return logits

384

385 - def load_weights(self, weights: Iterable[Tuple[str,

386 - torch.Tensor]]) -> Set[str]:

385 + def load_weights(self, weights: Iterable[tuple[str,

386 + torch.Tensor]]) -> set[str]:

387 loader = AutoWeightsLoader(self)

388 return loader.load_weights(weights)

vllm/model_executor/models/gemma.py

+8 -7

@@ -15,8 +15,9 @@

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

16 # limitations under the License.

17 """Inference-only Gemma model compatible with HuggingFace weights."""

18 +from collections.abc import Iterable

19 from functools import cache

19 -from typing import Iterable, Optional, Set, Tuple, Union

20 +from typing import Optional, Union

21

22 import torch

23 from torch import nn

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

232 positions: torch.Tensor,

233 hidden_states: torch.Tensor,

234 residual: Optional[torch.Tensor],

234 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

236 # Self Attention

237 if residual is None:

238 residual = hidden_states

@@ -318,8 +319,8 @@ class GemmaModel(nn.Module):

319 hidden_states, _ = self.norm(hidden_states, residual)

320 return hidden_states

321

321 - def load_weights(self, weights: Iterable[Tuple[str,

322 - torch.Tensor]]) -> Set[str]:

322 + def load_weights(self, weights: Iterable[tuple[str,

323 + torch.Tensor]]) -> set[str]:

324 stacked_params_mapping = [

325 # (param_name, shard_name, shard_id)

326 ("qkv_proj", "q_proj", "q"),

@@ -329,7 +330,7 @@ class GemmaModel(nn.Module):

330 ("gate_up_proj", "up_proj", 1),

331 ]

332 params_dict = dict(self.named_parameters())

332 - loaded_params: Set[str] = set()

333 + loaded_params: set[str] = set()

334 for name, loaded_weight in weights:

335 for (param_name, shard_name, shard_id) in stacked_params_mapping:

336 if shard_name not in name:

@@ -413,8 +414,8 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

414 sampling_metadata)

415 return logits

416

416 - def load_weights(self, weights: Iterable[Tuple[str,

417 - torch.Tensor]]) -> Set[str]:

417 + def load_weights(self, weights: Iterable[tuple[str,

418 + torch.Tensor]]) -> set[str]:

419 loader = AutoWeightsLoader(

420 self,

421 skip_prefixes=(["lm_head."]

vllm/model_executor/models/gemma2.py

+8 -7

@@ -15,7 +15,8 @@

15 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

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

17 # limitations under the License.

18 -from typing import Iterable, Optional, Set, Tuple, Union

18 +from collections.abc import Iterable

19 +from typing import Optional, Union

20

21 import torch

22 from torch import nn

@@ -218,7 +219,7 @@ class Gemma2DecoderLayer(nn.Module):

219 positions: torch.Tensor,

220 hidden_states: torch.Tensor,

221 residual: Optional[torch.Tensor],

221 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

223 if residual is None:

224 residual = hidden_states

225 hidden_states = self.input_layernorm(hidden_states)

@@ -305,8 +306,8 @@ class Gemma2Model(nn.Module):

306 hidden_states, _ = self.norm(hidden_states, residual)

307 return hidden_states

308

308 - def load_weights(self, weights: Iterable[Tuple[str,

309 - torch.Tensor]]) -> Set[str]:

309 + def load_weights(self, weights: Iterable[tuple[str,

310 + torch.Tensor]]) -> set[str]:

311 stacked_params_mapping = [

312 # (param_name, shard_name, shard_id)

313 ("qkv_proj", "q_proj", "q"),

@@ -316,7 +317,7 @@ class Gemma2Model(nn.Module):

317 ("gate_up_proj", "up_proj", 1),

318 ]

319 params_dict = dict(self.named_parameters())

319 - loaded_params: Set[str] = set()

320 + loaded_params: set[str] = set()

321 for name, loaded_weight in weights:

322 if (self.quant_config is not None and

323 (scale_name := self.quant_config.get_cache_scale(name))):

@@ -413,8 +414,8 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

414 sampling_metadata)

415 return logits

416

416 - def load_weights(self, weights: Iterable[Tuple[str,

417 - torch.Tensor]]) -> Set[str]:

417 + def load_weights(self, weights: Iterable[tuple[str,

418 + torch.Tensor]]) -> set[str]:

419 loader = AutoWeightsLoader(

420 self,

421 skip_prefixes=(["lm_head."]

vllm/model_executor/models/gemma3.py

+8 -7

@@ -14,7 +14,8 @@

14 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

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

16 # limitations under the License.

17 -from typing import Iterable, Optional, Set, Tuple, Union

17 +from collections.abc import Iterable

18 +from typing import Optional, Union

19

20 import torch

21 import torch.nn.functional as F

@@ -320,7 +321,7 @@ class Gemma3DecoderLayer(nn.Module):

321 hidden_states: torch.Tensor,

322 residual: Optional[torch.Tensor],

323 **kwargs,

323 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

325 if residual is None:

326 residual = hidden_states

327 hidden_states = self.input_layernorm(hidden_states)

@@ -412,8 +413,8 @@ class Gemma3Model(nn.Module):

413 hidden_states, _ = self.norm(hidden_states, residual)

414 return hidden_states

415

415 - def load_weights(self, weights: Iterable[Tuple[str,

416 - torch.Tensor]]) -> Set[str]:

416 + def load_weights(self, weights: Iterable[tuple[str,

417 + torch.Tensor]]) -> set[str]:

418 stacked_params_mapping = [

419 # (param_name, shard_name, shard_id)

420 ("qkv_proj", "q_proj", "q"),

@@ -423,7 +424,7 @@ class Gemma3Model(nn.Module):

424 ("gate_up_proj", "up_proj", 1),

425 ]

426 params_dict = dict(self.named_parameters())

426 - loaded_params: Set[str] = set()

427 + loaded_params: set[str] = set()

428 for name, loaded_weight in weights:

429 if (self.quant_config is not None and

430 (scale_name := self.quant_config.get_cache_scale(name))):

@@ -521,8 +522,8 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

522 sampling_metadata)

523 return logits

524

524 - def load_weights(self, weights: Iterable[Tuple[str,

525 - torch.Tensor]]) -> Set[str]:

525 + def load_weights(self, weights: Iterable[tuple[str,

526 + torch.Tensor]]) -> set[str]:

527 loader = AutoWeightsLoader(

528 self,

529 skip_prefixes=(["lm_head."]

vllm/model_executor/models/gemma3_mm.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2 import math

3 from collections.abc import Iterable, Mapping, Sequence

4 -from typing import Any, Literal, Optional, Set, Tuple, TypedDict

4 +from typing import Any, Literal, Optional, TypedDict

5

6 import torch

7 from torch import nn

@@ -701,8 +701,8 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,

701 return self.language_model.compute_logits(hidden_states,

702 sampling_metadata)

703

704 - def load_weights(self, weights: Iterable[Tuple[str,

705 - torch.Tensor]]) -> Set[str]:

704 + def load_weights(self, weights: Iterable[tuple[str,

705 + torch.Tensor]]) -> set[str]:

706 loader = AutoWeightsLoader(self)

707 return loader.load_weights(weights)

708

vllm/model_executor/models/glm4.py

+6 -5

@@ -21,7 +21,8 @@

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

22 # limitations under the License.

23 """Inference-only GLM-4-0414 model compatible with HuggingFace weights."""

24 -from typing import Iterable, Optional, Set, Tuple, Union

24 +from collections.abc import Iterable

25 +from typing import Optional, Union

26

27 import torch

28 from torch import nn

@@ -60,7 +61,7 @@ class Glm4Attention(nn.Module):

61 rope_theta: float = 10000,

62 cache_config: Optional[CacheConfig] = None,

63 quant_config: Optional[QuantizationConfig] = None,

63 - rope_scaling: Optional[Tuple] = None,

64 + rope_scaling: Optional[tuple] = None,

65 prefix: str = "",

66 attn_type: str = AttentionType.DECODER) -> None:

67 super().__init__()

@@ -183,7 +184,7 @@ class Glm4DecoderLayer(nn.Module):

184 positions: torch.Tensor,

185 hidden_states: torch.Tensor,

186 residual: Optional[torch.Tensor],

186 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

188 # Self Attention

189 if residual is None:

190 residual = hidden_states

@@ -293,8 +294,8 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

294 sampling_metadata)

295 return logits

296

296 - def load_weights(self, weights: Iterable[Tuple[str,

297 - torch.Tensor]]) -> Set[str]:

297 + def load_weights(self, weights: Iterable[tuple[str,

298 + torch.Tensor]]) -> set[str]:

299 loader = AutoWeightsLoader(

300 self,

301 skip_prefixes=(["lm_head."]

vllm/model_executor/models/gpt2.py

+5 -4

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

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

19 # limitations under the License.

20 """Inference-only GPT-2 model compatible with HuggingFace weights."""

21 -from typing import Iterable, Optional, Set, Tuple, Union

21 +from collections.abc import Iterable

22 +from typing import Optional, Union

23

24 import torch

25 from torch import nn

@@ -280,10 +281,10 @@ class GPT2LMHeadModel(nn.Module, SupportsPP):

281 sampling_metadata)

282 return logits

283

283 - def load_weights(self, weights: Iterable[Tuple[str,

284 - torch.Tensor]]) -> Set[str]:

284 + def load_weights(self, weights: Iterable[tuple[str,

285 + torch.Tensor]]) -> set[str]:

286 params_dict = dict(self.named_parameters(remove_duplicate=False))

286 - loaded_params: Set[str] = set()

287 + loaded_params: set[str] = set()

288 for name, loaded_weight in weights:

289 if ".attn.bias" in name or ".attn.masked_bias" in name:

290 # Skip attention mask.

vllm/model_executor/models/gpt_bigcode.py

+7 -6

@@ -19,7 +19,8 @@

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

20 # limitations under the License.

21 """Inference-only GPTBigCode model compatible with HuggingFace weights."""

22 -from typing import Iterable, Optional, Set, Tuple, Union

22 +from collections.abc import Iterable

23 +from typing import Optional, Union

24

25 import torch

26 from torch import nn

@@ -243,10 +244,10 @@ class GPTBigCodeModel(nn.Module):

244 hidden_states = self.ln_f(hidden_states)

245 return hidden_states

246

246 - def load_weights(self, weights: Iterable[Tuple[str,

247 - torch.Tensor]]) -> Set[str]:

247 + def load_weights(self, weights: Iterable[tuple[str,

248 + torch.Tensor]]) -> set[str]:

249 params_dict = dict(self.named_parameters(remove_duplicate=False))

249 - loaded_params: Set[str] = set()

250 + loaded_params: set[str] = set()

251 for name, loaded_weight in weights:

252 if ".attn.bias" in name:

253 # Skip attention mask.

@@ -327,8 +328,8 @@ class GPTBigCodeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

328 sampling_metadata)

329 return logits

330

330 - def load_weights(self, weights: Iterable[Tuple[str,

331 - torch.Tensor]]) -> Set[str]:

331 + def load_weights(self, weights: Iterable[tuple[str,

332 + torch.Tensor]]) -> set[str]:

333 loader = AutoWeightsLoader(

334 self,

335 skip_prefixes=(["lm_head."]),

vllm/model_executor/models/gpt_j.py

+7 -6

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

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

18 # limitations under the License.

19 """Inference-only GPT-J model compatible with HuggingFace weights."""

20 -from typing import Iterable, Optional, Set, Tuple, Union

20 +from collections.abc import Iterable

21 +from typing import Optional, Union

22

23 import torch

24 from torch import nn

@@ -228,8 +229,8 @@ class GPTJModel(nn.Module):

229 hidden_states = self.ln_f(hidden_states)

230 return hidden_states

231

231 - def load_weights(self, weights: Iterable[Tuple[str,

232 - torch.Tensor]]) -> Set[str]:

232 + def load_weights(self, weights: Iterable[tuple[str,

233 + torch.Tensor]]) -> set[str]:

234 stacked_params_mapping = [

235 # (param_name, shard_name, shard_id)

236 ("qkv_proj", "q_proj", "q"),

@@ -239,7 +240,7 @@ class GPTJModel(nn.Module):

240 ("gate_up_proj", "up_proj", 1),

241 ]

242 params_dict = dict(self.named_parameters())

242 - loaded_params: Set[str] = set()

243 + loaded_params: set[str] = set()

244 for name, loaded_weight in weights:

245 if "attn.bias" in name or "attn.masked_bias" in name:

246 continue

@@ -331,7 +332,7 @@ class GPTJForCausalLM(nn.Module, SupportsPP):

332 sampling_metadata, self.lm_head.bias)

333 return logits

334

334 - def load_weights(self, weights: Iterable[Tuple[str,

335 - torch.Tensor]]) -> Set[str]:

335 + def load_weights(self, weights: Iterable[tuple[str,

336 + torch.Tensor]]) -> set[str]:

337 loader = AutoWeightsLoader(self)

338 return loader.load_weights(weights)

\ No newline at end of file

vllm/model_executor/models/gpt_neox.py

+7 -6

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

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

18 # limitations under the License.

19 """Inference-only GPT-NeoX model compatible with HuggingFace weights."""

20 -from typing import Iterable, Optional, Set, Tuple, Union

20 +from collections.abc import Iterable

21 +from typing import Optional, Union

22

23 import torch

24 from torch import nn

@@ -240,10 +241,10 @@ class GPTNeoXModel(nn.Module):

241 hidden_states = self.final_layer_norm(hidden_states)

242 return hidden_states

243

243 - def load_weights(self, weights: Iterable[Tuple[str,

244 - torch.Tensor]]) -> Set[str]:

244 + def load_weights(self, weights: Iterable[tuple[str,

245 + torch.Tensor]]) -> set[str]:

246 params_dict = dict(self.named_parameters())

246 - loaded_params: Set[str] = set()

247 + loaded_params: set[str] = set()

248 for name, loaded_weight in weights:

249 if ("attention.bias" in name or "attention.masked_bias" in name

250 or "rotary_emb.inv_freq" in name):

@@ -324,7 +325,7 @@ class GPTNeoXForCausalLM(nn.Module, SupportsPP):

325 sampling_metadata)

326 return logits

327

327 - def load_weights(self, weights: Iterable[Tuple[str,

328 - torch.Tensor]]) -> Set[str]:

328 + def load_weights(self, weights: Iterable[tuple[str,

329 + torch.Tensor]]) -> set[str]:

330 loader = AutoWeightsLoader(self)

331 return loader.load_weights(weights)

vllm/model_executor/models/granite.py

+9 -8

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only IBM Granite model compatible with HuggingFace weights."""

25 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Any, Optional, Union

27

28 import torch

29 from torch import nn

@@ -97,7 +98,7 @@ class GraniteAttention(nn.Module):

98 num_heads: int,

99 num_kv_heads: int,

100 rope_theta: float = 10000,

100 - rope_scaling: Optional[Dict[str, Any]] = None,

101 + rope_scaling: Optional[dict[str, Any]] = None,

102 max_position_embeddings: int = 8192,

103 quant_config: Optional[QuantizationConfig] = None,

104 bias: bool = False,

@@ -230,7 +231,7 @@ class GraniteDecoderLayer(nn.Module):

231 self,

232 positions: torch.Tensor,

233 hidden_states: torch.Tensor,

233 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

235 # Self Attention

236 residual = hidden_states

237 hidden_states = self.input_layernorm(hidden_states)

@@ -321,8 +322,8 @@ class GraniteModel(nn.Module):

322 hidden_states = self.norm(hidden_states)

323 return hidden_states

324

324 - def load_weights(self, weights: Iterable[Tuple[str,

325 - torch.Tensor]]) -> Set[str]:

325 + def load_weights(self, weights: Iterable[tuple[str,

326 + torch.Tensor]]) -> set[str]:

327 stacked_params_mapping = [

328 # (param_name, shard_name, shard_id)

329 (".qkv_proj", ".q_proj", "q"),

@@ -332,7 +333,7 @@ class GraniteModel(nn.Module):

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

334 ]

335 params_dict = dict(self.named_parameters())

335 - loaded_params: Set[str] = set()

336 + loaded_params: set[str] = set()

337 for name, loaded_weight in weights:

338 if (self.quant_config is not None and

339 (scale_name := self.quant_config.get_cache_scale(name))):

@@ -475,8 +476,8 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

476 device=device),

477 })

478

478 - def load_weights(self, weights: Iterable[Tuple[str,

479 - torch.Tensor]]) -> Set[str]:

479 + def load_weights(self, weights: Iterable[tuple[str,

480 + torch.Tensor]]) -> set[str]:

481 skip_prefixes = [

482 "rotary_emb.inv_freq",

483 # Models trained using ColossalAI may include these tensors in

vllm/model_executor/models/granite_speech.py

+4 -3

@@ -23,7 +23,8 @@

23 # limitations under the License.

24 """Inference-only IBM Granite speeech model."""

25 import math

26 -from typing import Iterable, Mapping, Optional, Set, Tuple, TypedDict, Union

26 +from collections.abc import Iterable, Mapping

27 +from typing import Optional, TypedDict, Union

28

29 import torch

30 import torch.nn.functional as F

@@ -763,8 +764,8 @@ class GraniteSpeechForConditionalGeneration(

764

765 def load_weights(

766 self,

766 - weights: Iterable[Tuple[str, torch.Tensor]],

767 - ) -> Set[str]:

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

768 + ) -> set[str]:

769 loader = AutoWeightsLoader(self)

770 return loader.load_weights(weights)

771

vllm/model_executor/models/granitemoe.py

+6 -5

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only GraniteMoe model."""

25 -from typing import Iterable, Optional, Set, Tuple

25 +from collections.abc import Iterable

26 +from typing import Optional

27

28 import torch

29 from torch import nn

@@ -305,8 +306,8 @@ class GraniteMoeModel(nn.Module):

306 hidden_states = self.norm(hidden_states)

307 return hidden_states

308

308 - def load_weights(self, weights: Iterable[Tuple[str,

309 - torch.Tensor]]) -> Set[str]:

309 + def load_weights(self, weights: Iterable[tuple[str,

310 + torch.Tensor]]) -> set[str]:

311 new_weights = {}

312 for n, p in weights:

313 if n.endswith('.block_sparse_moe.input_linear.weight'):

@@ -425,8 +426,8 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

426 device=device),

427 })

428

428 - def load_weights(self, weights: Iterable[Tuple[str,

429 - torch.Tensor]]) -> Set[str]:

429 + def load_weights(self, weights: Iterable[tuple[str,

430 + torch.Tensor]]) -> set[str]:

431 loader = AutoWeightsLoader(

432 self,

433 skip_prefixes=(["lm_head."]

vllm/model_executor/models/granitemoehybrid.py

+8 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Inference-only GraniteMoeHybrid model."""

3 # Added by the IBM Team, 2025

4 -from typing import Iterable, Optional, Set, Tuple

4 +from collections.abc import Iterable

5 +from typing import Optional

6

7 import torch

8 from torch import nn

@@ -381,10 +382,10 @@ class GraniteMoeHybridModel(nn.Module):

382 hidden_states = self.norm(hidden_states)

383 return hidden_states

384

384 - def load_weights(self, weights: Iterable[Tuple[str,

385 - torch.Tensor]]) -> Set[str]:

385 + def load_weights(self, weights: Iterable[tuple[str,

386 + torch.Tensor]]) -> set[str]:

387 params_dict = dict(self.named_parameters())

387 - loaded_params: Set[str] = set()

388 + loaded_params: set[str] = set()

389

390 def _load(n, p):

391 param = params_dict[n]

@@ -538,7 +539,7 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,

539 return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

540

541 def _get_mamba_cache_shape(

541 - self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

542 + self) -> tuple[tuple[int, int], tuple[int, int]]:

543 world_size = get_tensor_model_parallel_world_size()

544 hidden_size = self.config.hidden_size

545

@@ -578,7 +579,7 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,

579 sampling_metadata)

580 return logits

581

581 - def load_weights(self, weights: Iterable[Tuple[str,

582 - torch.Tensor]]) -> Set[str]:

582 + def load_weights(self, weights: Iterable[tuple[str,

583 + torch.Tensor]]) -> set[str]:

584 loader = AutoWeightsLoader(self)

585 return loader.load_weights(weights)

vllm/model_executor/models/granitemoeshared.py

+6 -5

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

4 The architecture is the same as granitemoe but with the addition of shared

5 experts.

6 """

7 -from typing import Iterable, Optional, Set, Tuple

7 +from collections.abc import Iterable

8 +from typing import Optional

9

10 import torch

11 from torch import nn

@@ -208,8 +209,8 @@ class GraniteMoeSharedModel(nn.Module):

209 hidden_states = self.norm(hidden_states)

210 return hidden_states

211

211 - def load_weights(self, weights: Iterable[Tuple[str,

212 - torch.Tensor]]) -> Set[str]:

212 + def load_weights(self, weights: Iterable[tuple[str,

213 + torch.Tensor]]) -> set[str]:

214 new_weights = {}

215 for n, p in weights:

216 if n.endswith('.block_sparse_moe.input_linear.weight'):

@@ -329,8 +330,8 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

330 device=device),

331 })

332

332 - def load_weights(self, weights: Iterable[Tuple[str,

333 - torch.Tensor]]) -> Set[str]:

333 + def load_weights(self, weights: Iterable[tuple[str,

334 + torch.Tensor]]) -> set[str]:

335 loader = AutoWeightsLoader(

336 self,

337 skip_prefixes=(["lm_head."]

vllm/model_executor/models/grok1.py

+10 -9

@@ -21,7 +21,8 @@

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

22 # limitations under the License.

23 """Inference-only Grok1 model."""

24 -from typing import Iterable, List, Optional, Set, Tuple, Union

24 +from collections.abc import Iterable

25 +from typing import Optional, Union

26

27 import torch

28 import torch.nn.functional as F

@@ -263,7 +264,7 @@ class Grok1DecoderLayer(nn.Module):

264 kv_cache: torch.Tensor,

265 attn_metadata: AttentionMetadata,

266 residual: Optional[torch.Tensor],

266 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

268 # Self Attention

269 if residual is None:

270 residual = hidden_states

@@ -340,7 +341,7 @@ class Grok1Model(nn.Module):

341 self,

342 input_ids: torch.Tensor,

343 positions: torch.Tensor,

343 - kv_caches: List[torch.Tensor],

344 + kv_caches: list[torch.Tensor],

345 attn_metadata: AttentionMetadata,

346 intermediate_tensors: Optional[IntermediateTensors],

347 inputs_embeds: Optional[torch.Tensor] = None,

@@ -371,8 +372,8 @@ class Grok1Model(nn.Module):

372 hidden_states, _ = self.norm(hidden_states, residual)

373 return hidden_states

374

374 - def load_weights(self, weights: Iterable[Tuple[str,

375 - torch.Tensor]]) -> Set[str]:

375 + def load_weights(self, weights: Iterable[tuple[str,

376 + torch.Tensor]]) -> set[str]:

377 stacked_params_mapping = [

378 # (param_name, shard_name, shard_id)

379 ("qkv_proj", "q_proj", "q"),

@@ -390,7 +391,7 @@ class Grok1Model(nn.Module):

391 num_experts=num_experts)

392

393 params_dict = dict(self.named_parameters())

393 - loaded_params: Set[str] = set()

394 + loaded_params: set[str] = set()

395

396 for name, loaded_weight in weights:

397 if (self.quant_config is not None and

@@ -528,7 +529,7 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

529 self,

530 input_ids: torch.Tensor,

531 positions: torch.Tensor,

531 - kv_caches: List[torch.Tensor],

532 + kv_caches: list[torch.Tensor],

533 attn_metadata: AttentionMetadata,

534 intermediate_tensors: Optional[IntermediateTensors] = None,

535 inputs_embeds: Optional[torch.Tensor] = None,

@@ -547,8 +548,8 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

548 sampling_metadata)

549 return logits

550

550 - def load_weights(self, weights: Iterable[Tuple[str,

551 - torch.Tensor]]) -> Set[str]:

551 + def load_weights(self, weights: Iterable[tuple[str,

552 + torch.Tensor]]) -> set[str]:

553 skip_prefixes = ["rotary_emb.inv_freq"]

554 # Skip lm_head when tie_word_embeddings is True

555 if self.config.tie_word_embeddings:

vllm/model_executor/models/idefics2_vision_model.py

+5 -4

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

17 # limitations under the License.

18 """PyTorch Idefics2 model."""

19

20 -from typing import Iterable, Optional, Set, Tuple

20 +from collections.abc import Iterable

21 +from typing import Optional

22

23 import torch

24 from torch import nn

@@ -342,8 +343,8 @@ class Idefics2VisionTransformer(nn.Module):

343 last_hidden_state = self.post_layernorm(encoder_outputs)

344 return last_hidden_state

345

345 - def load_weights(self, weights: Iterable[Tuple[str,

346 - torch.Tensor]]) -> Set[str]:

346 + def load_weights(self, weights: Iterable[tuple[str,

347 + torch.Tensor]]) -> set[str]:

348 stacked_params_mapping = [

349 # (param_name, shard_name, shard_id)

350 ("qkv_proj", "q_proj", "q"),

@@ -351,7 +352,7 @@ class Idefics2VisionTransformer(nn.Module):

352 ("qkv_proj", "v_proj", "v"),

353 ]

354 params_dict = dict(self.named_parameters())

354 - loaded_params: Set[str] = set()

355 + loaded_params: set[str] = set()

356 layer_count = len(self.encoder.layers)

357

358 for name, loaded_weight in weights:

vllm/model_executor/models/idefics3.py

+4 -4

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

17

18 import math

19 from collections.abc import Iterable, Mapping, Sequence

20 -from typing import Dict, Literal, Optional, Set, Tuple, TypedDict, Union

20 +from typing import Literal, Optional, TypedDict, Union

21

22 import torch

23 from torch import nn

@@ -85,7 +85,7 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):

85 def get_hf_processor(

86 self,

87 *,

88 - size: Optional[Dict[str, int]] = None,

88 + size: Optional[dict[str, int]] = None,

89 **kwargs: object,

90 ) -> Idefics3Processor:

91 if size is not None:

@@ -752,8 +752,8 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,

752 sampling_metadata)

753 return logits

754

755 - def load_weights(self, weights: Iterable[Tuple[str,

756 - torch.Tensor]]) -> Set[str]:

755 + def load_weights(self, weights: Iterable[tuple[str,

756 + torch.Tensor]]) -> set[str]:

757 loader = AutoWeightsLoader(self)

758 return loader.load_weights(weights)

759

vllm/model_executor/models/interfaces.py

+44 -44

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import (TYPE_CHECKING, ClassVar, Dict, List, Literal, Optional,

4 - Protocol, Type, Union, overload, runtime_checkable)

3 +from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,

4 + Union, overload, runtime_checkable)

5

6 import torch

7 from torch import Tensor

@@ -102,7 +102,7 @@ class _SupportsMultiModalType(Protocol):

102

103 @overload

104 def supports_multimodal(

105 - model: Type[object]) -> TypeIs[Type[SupportsMultiModal]]:

105 + model: type[object]) -> TypeIs[type[SupportsMultiModal]]:

106 ...

107

108

@@ -112,8 +112,8 @@ def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]:

112

113

114 def supports_multimodal(

115 - model: Union[Type[object], object],

116 -) -> Union[TypeIs[Type[SupportsMultiModal]], TypeIs[SupportsMultiModal]]:

115 + model: Union[type[object], object],

116 +) -> Union[TypeIs[type[SupportsMultiModal]], TypeIs[SupportsMultiModal]]:

117 if isinstance(model, type):

118 return isinstance(model, _SupportsMultiModalType)

119

@@ -134,9 +134,9 @@ class SupportsLoRA(Protocol):

134 """

135 # The `embedding_module` and `embedding_padding_modules`

136 # are empty by default.

137 - embedding_modules: ClassVar[Dict[str, str]] = {}

138 - embedding_padding_modules: ClassVar[List[str]] = []

139 - packed_modules_mapping: ClassVar[Dict[str, List[str]]] = {}

137 + embedding_modules: ClassVar[dict[str, str]] = {}

138 + embedding_padding_modules: ClassVar[list[str]] = []

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

140

141

142 # We can't use runtime_checkable with ClassVar for issubclass checks

@@ -145,13 +145,13 @@ class SupportsLoRA(Protocol):

145 class _SupportsLoRAType(Protocol):

146 supports_lora: Literal[True]

147

148 - packed_modules_mapping: Dict[str, List[str]]

149 - embedding_modules: Dict[str, str]

150 - embedding_padding_modules: List[str]

148 + packed_modules_mapping: dict[str, list[str]]

149 + embedding_modules: dict[str, str]

150 + embedding_padding_modules: list[str]

151

152

153 @overload

154 -def supports_lora(model: Type[object]) -> TypeIs[Type[SupportsLoRA]]:

154 +def supports_lora(model: type[object]) -> TypeIs[type[SupportsLoRA]]:

155 ...

156

157

@@ -161,8 +161,8 @@ def supports_lora(model: object) -> TypeIs[SupportsLoRA]:

161

162

163 def supports_lora(

164 - model: Union[Type[object], object],

165 -) -> Union[TypeIs[Type[SupportsLoRA]], TypeIs[SupportsLoRA]]:

164 + model: Union[type[object], object],

165 +) -> Union[TypeIs[type[SupportsLoRA]], TypeIs[SupportsLoRA]]:

166 result = _supports_lora(model)

167

168 if not result:

@@ -191,7 +191,7 @@ def supports_lora(

191 return result

192

193

194 -def _supports_lora(model: Union[Type[object], object]) -> bool:

194 +def _supports_lora(model: Union[type[object], object]) -> bool:

195 if isinstance(model, type):

196 return isinstance(model, _SupportsLoRAType)

197

@@ -256,7 +256,7 @@ class _SupportsPPType(Protocol):

256

257

258 @overload

259 -def supports_pp(model: Type[object]) -> TypeIs[Type[SupportsPP]]:

259 +def supports_pp(model: type[object]) -> TypeIs[type[SupportsPP]]:

260 ...

261

262

@@ -266,8 +266,8 @@ def supports_pp(model: object) -> TypeIs[SupportsPP]:

266

267

268 def supports_pp(

269 - model: Union[Type[object], object],

270 -) -> Union[bool, TypeIs[Type[SupportsPP]], TypeIs[SupportsPP]]:

269 + model: Union[type[object], object],

270 +) -> Union[bool, TypeIs[type[SupportsPP]], TypeIs[SupportsPP]]:

271 supports_attributes = _supports_pp_attributes(model)

272 supports_inspect = _supports_pp_inspect(model)

273

@@ -298,14 +298,14 @@ def supports_pp(

298 return supports_attributes and supports_inspect

299

300

301 -def _supports_pp_attributes(model: Union[Type[object], object]) -> bool:

301 +def _supports_pp_attributes(model: Union[type[object], object]) -> bool:

302 if isinstance(model, type):

303 return isinstance(model, _SupportsPPType)

304

305 return isinstance(model, SupportsPP)

306

307

308 -def _supports_pp_inspect(model: Union[Type[object], object]) -> bool:

308 +def _supports_pp_inspect(model: Union[type[object], object]) -> bool:

309 model_forward = getattr(model, "forward", None)

310 if not callable(model_forward):

311 return False

@@ -336,13 +336,13 @@ def has_inner_state(model: object) -> TypeIs[HasInnerState]:

336

337

338 @overload

339 -def has_inner_state(model: Type[object]) -> TypeIs[Type[HasInnerState]]:

339 +def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]:

340 ...

341

342

343 def has_inner_state(

344 - model: Union[Type[object], object]

345 -) -> Union[TypeIs[Type[HasInnerState]], TypeIs[HasInnerState]]:

344 + model: Union[type[object], object]

345 +) -> Union[TypeIs[type[HasInnerState]], TypeIs[HasInnerState]]:

346 if isinstance(model, type):

347 return isinstance(model, _HasInnerStateType)

348

@@ -373,13 +373,13 @@ def is_attention_free(model: object) -> TypeIs[IsAttentionFree]:

373

374

375 @overload

376 -def is_attention_free(model: Type[object]) -> TypeIs[Type[IsAttentionFree]]:

376 +def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]:

377 ...

378

379

380 def is_attention_free(

381 - model: Union[Type[object], object]

382 -) -> Union[TypeIs[Type[IsAttentionFree]], TypeIs[IsAttentionFree]]:

381 + model: Union[type[object], object]

382 +) -> Union[TypeIs[type[IsAttentionFree]], TypeIs[IsAttentionFree]]:

383 if isinstance(model, type):

384 return isinstance(model, _IsAttentionFreeType)

385

@@ -410,13 +410,13 @@ def is_hybrid(model: object) -> TypeIs[IsHybrid]:

410

411

412 @overload

413 -def is_hybrid(model: Type[object]) -> TypeIs[Type[IsHybrid]]:

413 +def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]:

414 ...

415

416

417 def is_hybrid(

418 - model: Union[Type[object], object]

419 -) -> Union[TypeIs[Type[IsHybrid]], TypeIs[IsHybrid]]:

418 + model: Union[type[object], object]

419 +) -> Union[TypeIs[type[IsHybrid]], TypeIs[IsHybrid]]:

420 if isinstance(model, type):

421 return isinstance(model, _IsHybridType)

422

@@ -439,13 +439,13 @@ def has_noops(model: object) -> TypeIs[HasNoOps]:

439

440

441 @overload

442 -def has_noops(model: Type[object]) -> TypeIs[Type[HasNoOps]]:

442 +def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]:

443 ...

444

445

446 def has_noops(

447 - model: Union[Type[object], object]

448 -) -> Union[TypeIs[Type[HasNoOps]], TypeIs[HasNoOps]]:

447 + model: Union[type[object], object]

448 +) -> Union[TypeIs[type[HasNoOps]], TypeIs[HasNoOps]]:

449 if isinstance(model, type):

450 return isinstance(model, _HasNoOpsType)

451

@@ -461,7 +461,7 @@ class SupportsCrossEncoding(Protocol):

461

462 @overload

463 def supports_cross_encoding(

464 - model: Type[object]) -> TypeIs[Type[SupportsCrossEncoding]]:

464 + model: type[object]) -> TypeIs[type[SupportsCrossEncoding]]:

465 ...

466

467

@@ -471,8 +471,8 @@ def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]:

471

472

473 def _supports_cross_encoding(

474 - model: Union[Type[object], object],

475 -) -> Union[TypeIs[Type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:

474 + model: Union[type[object], object],

475 +) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:

476

477 if isinstance(model, type):

478 return isinstance(model, SupportsCrossEncoding)

@@ -481,15 +481,15 @@ def _supports_cross_encoding(

481

482

483 def supports_cross_encoding(

484 - model: Union[Type[object], object],

485 -) -> Union[TypeIs[Type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:

484 + model: Union[type[object], object],

485 +) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:

486 return is_pooling_model(model) and _supports_cross_encoding(model)

487

488

489 class SupportsQuant:

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

491

492 - packed_modules_mapping: ClassVar[Dict[str, List[str]]] = {}

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

493 quant_config: Optional[QuantizationConfig] = None

494

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

@@ -525,7 +525,7 @@ class SupportsTranscription(Protocol):

525

526 @overload

527 def supports_transcription(

528 - model: Type[object]) -> TypeIs[Type[SupportsTranscription]]:

528 + model: type[object]) -> TypeIs[type[SupportsTranscription]]:

529 ...

530

531

@@ -535,8 +535,8 @@ def supports_transcription(model: object) -> TypeIs[SupportsTranscription]:

535

536

537 def supports_transcription(

538 - model: Union[Type[object], object],

539 -) -> Union[TypeIs[Type[SupportsTranscription]], TypeIs[SupportsTranscription]]:

538 + model: Union[type[object], object],

539 +) -> Union[TypeIs[type[SupportsTranscription]], TypeIs[SupportsTranscription]]:

540 if isinstance(model, type):

541 return isinstance(model, SupportsTranscription)

542

@@ -551,7 +551,7 @@ class SupportsV0Only(Protocol):

551

552

553 @overload

554 -def supports_v0_only(model: Type[object]) -> TypeIs[Type[SupportsV0Only]]:

554 +def supports_v0_only(model: type[object]) -> TypeIs[type[SupportsV0Only]]:

555 ...

556

557

@@ -561,8 +561,8 @@ def supports_v0_only(model: object) -> TypeIs[SupportsV0Only]:

561

562

563 def supports_v0_only(

564 - model: Union[Type[object], object],

565 -) -> Union[TypeIs[Type[SupportsV0Only]], TypeIs[SupportsV0Only]]:

564 + model: Union[type[object], object],

565 +) -> Union[TypeIs[type[SupportsV0Only]], TypeIs[SupportsV0Only]]:

566 if isinstance(model, type):

567 return isinstance(model, SupportsV0Only)

568

vllm/model_executor/models/interfaces_base.py

+13 -13

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import (TYPE_CHECKING, Optional, Protocol, Type, Union, overload,

3 +from typing import (TYPE_CHECKING, Optional, Protocol, Union, overload,

4 runtime_checkable)

5

6 import torch

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

20

21 # The type of hidden states

22 # Currently, T = torch.Tensor for all models except for Medusa

23 -# which has T = List[torch.Tensor]

23 +# which has T = list[torch.Tensor]

24 T = TypeVar("T", default=torch.Tensor)

25 T_co = TypeVar("T_co", default=torch.Tensor, covariant=True)

26

@@ -48,12 +48,12 @@ class VllmModel(Protocol[T_co]):

48 ...

49

50

51 -def _check_vllm_model_init(model: Union[Type[object], object]) -> bool:

51 +def _check_vllm_model_init(model: Union[type[object], object]) -> bool:

52 model_init = model.__init__

53 return supports_kw(model_init, "vllm_config")

54

55

56 -def _check_vllm_model_forward(model: Union[Type[object], object]) -> bool:

56 +def _check_vllm_model_forward(model: Union[type[object], object]) -> bool:

57 model_forward = getattr(model, "forward", None)

58 if not callable(model_forward):

59 return False

@@ -75,7 +75,7 @@ def _check_vllm_model_forward(model: Union[Type[object], object]) -> bool:

75

76

77 @overload

78 -def is_vllm_model(model: Type[object]) -> TypeIs[Type[VllmModel]]:

78 +def is_vllm_model(model: type[object]) -> TypeIs[type[VllmModel]]:

79 ...

80

81

@@ -85,8 +85,8 @@ def is_vllm_model(model: object) -> TypeIs[VllmModel]:

85

86

87 def is_vllm_model(

88 - model: Union[Type[object], object],

89 -) -> Union[TypeIs[Type[VllmModel]], TypeIs[VllmModel]]:

88 + model: Union[type[object], object],

89 +) -> Union[TypeIs[type[VllmModel]], TypeIs[VllmModel]]:

90 return _check_vllm_model_init(model) and _check_vllm_model_forward(model)

91

92

@@ -105,7 +105,7 @@ class VllmModelForTextGeneration(VllmModel[T], Protocol[T]):

105

106 @overload

107 def is_text_generation_model(

108 - model: Type[object]) -> TypeIs[Type[VllmModelForTextGeneration]]:

108 + model: type[object]) -> TypeIs[type[VllmModelForTextGeneration]]:

109 ...

110

111

@@ -116,8 +116,8 @@ def is_text_generation_model(

116

117

118 def is_text_generation_model(

119 - model: Union[Type[object], object],

120 -) -> Union[TypeIs[Type[VllmModelForTextGeneration]],

119 + model: Union[type[object], object],

120 +) -> Union[TypeIs[type[VllmModelForTextGeneration]],

121 TypeIs[VllmModelForTextGeneration]]:

122 if not is_vllm_model(model):

123 return False

@@ -142,7 +142,7 @@ class VllmModelForPooling(VllmModel[T], Protocol[T]):

142

143

144 @overload

145 -def is_pooling_model(model: Type[object]) -> TypeIs[Type[VllmModelForPooling]]:

145 +def is_pooling_model(model: type[object]) -> TypeIs[type[VllmModelForPooling]]:

146 ...

147

148

@@ -152,8 +152,8 @@ def is_pooling_model(model: object) -> TypeIs[VllmModelForPooling]:

152

153

154 def is_pooling_model(

155 - model: Union[Type[object], object],

156 -) -> Union[TypeIs[Type[VllmModelForPooling]], TypeIs[VllmModelForPooling]]:

155 + model: Union[type[object], object],

156 +) -> Union[TypeIs[type[VllmModelForPooling]], TypeIs[VllmModelForPooling]]:

157 if not is_vllm_model(model):

158 return False

159

vllm/model_executor/models/intern_vit.py

+5 -4

@@ -6,8 +6,9 @@

7 # Licensed under The MIT License [see LICENSE for details]

8 # --------------------------------------------------------

9 +from collections.abc import Iterable

10 from functools import partial

10 -from typing import Iterable, Optional, Set, Tuple

11 +from typing import Optional

12

13 import torch

14 import torch.nn as nn

@@ -461,10 +462,10 @@ class InternVisionModel(nn.Module):

462

463 return encoder_outputs

464

464 - def load_weights(self, weights: Iterable[Tuple[str,

465 - torch.Tensor]]) -> Set[str]:

465 + def load_weights(self, weights: Iterable[tuple[str,

466 + torch.Tensor]]) -> set[str]:

467 params_dict = dict(self.named_parameters())

467 - loaded_params: Set[str] = set()

468 + loaded_params: set[str] = set()

469 for name, loaded_weight in weights:

470 param = params_dict[name]

471 weight_loader = getattr(param, "weight_loader",

vllm/model_executor/models/internlm2.py

+10 -9

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 +from collections.abc import Iterable

4 from functools import partial

4 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union

5 +from typing import Any, Optional, Union

6

7 import torch

8 from torch import nn

@@ -81,7 +82,7 @@ class InternLM2Attention(nn.Module):

82 num_heads: int,

83 num_kv_heads: int,

84 rope_theta: float = 10000,

84 - rope_scaling: Optional[Dict[str, Any]] = None,

85 + rope_scaling: Optional[dict[str, Any]] = None,

86 max_position_embeddings: int = 8192,

87 cache_config: Optional[CacheConfig] = None,

88 quant_config: Optional[QuantizationConfig] = None,

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

226 positions: torch.Tensor,

227 hidden_states: torch.Tensor,

228 residual: Optional[torch.Tensor],

228 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

230 # Self Attention

231 if residual is None:

232 residual = hidden_states

@@ -252,7 +253,7 @@ class InternLM2Model(nn.Module):

253 *,

254 vllm_config: VllmConfig,

255 prefix: str = "",

255 - layer_type: Type[InternLMDecoderLayer] = InternLMDecoderLayer):

256 + layer_type: type[InternLMDecoderLayer] = InternLMDecoderLayer):

257 super().__init__()

258

259 config = vllm_config.model_config.hf_config

@@ -316,7 +317,7 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

317 *,

318 vllm_config: VllmConfig,

319 prefix: str = "",

319 - model_type: Type[InternLM2Model] = InternLM2Model):

320 + model_type: type[InternLM2Model] = InternLM2Model):

321 super().__init__()

322 config = vllm_config.model_config.hf_config

323 quant_config = vllm_config.quant_config

@@ -361,15 +362,15 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

362 sampling_metadata)

363 return logits

364

364 - def load_weights(self, weights: Iterable[Tuple[str,

365 - torch.Tensor]]) -> Set[str]:

365 + def load_weights(self, weights: Iterable[tuple[str,

366 + torch.Tensor]]) -> set[str]:

367 stacked_params_mapping = [

368 # (param_name, shard_name, shard_id)

369 ("gate_up_proj", "w1", 0),

370 ("gate_up_proj", "w3", 1),

371 ]

372 params_dict = dict(self.named_parameters())

372 - loaded_params: Set[str] = set()

373 + loaded_params: set[str] = set()

374 for name, loaded_weight in weights:

375 if "rotary_emb.inv_freq" in name:

376 continue

@@ -407,7 +408,7 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):

408 *,

409 vllm_config: VllmConfig,

410 prefix: str = "",

410 - model_type: Type[InternLM2Model] = InternLM2Model,

411 + model_type: type[InternLM2Model] = InternLM2Model,

412 ):

413 super().__init__(vllm_config=vllm_config,

414 prefix=prefix,

vllm/model_executor/models/internlm2_ve.py

+2 -2

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Tuple, Union

3 +from typing import Optional, Union

4

5 import torch

6 from torch import nn

@@ -66,7 +66,7 @@ class InternLM2VEDecoderLayer(nn.Module):

66 hidden_states: torch.Tensor,

67 residual: Optional[torch.Tensor],

68 visual_token_mask: Optional[torch.Tensor] = None,

69 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

70 # Self Attention

71 if residual is None:

72 residual = hidden_states

vllm/model_executor/models/internvl.py

+3 -3

@@ -8,7 +8,7 @@

8 # --------------------------------------------------------

9 from abc import ABC, abstractmethod

10 from collections.abc import Iterable, Mapping, Sequence

11 -from typing import Literal, Optional, Set, Tuple, TypedDict, TypeVar, Union

11 +from typing import Literal, Optional, TypedDict, TypeVar, Union

12

13 import torch

14 import torch.nn as nn

@@ -932,8 +932,8 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):

932 return self.language_model.compute_logits(hidden_states,

933 sampling_metadata)

934

935 - def load_weights(self, weights: Iterable[Tuple[str,

936 - torch.Tensor]]) -> Set[str]:

935 + def load_weights(self, weights: Iterable[tuple[str,

936 + torch.Tensor]]) -> set[str]:

937 # unused modules appear in OpenGVLab/InternVideo2_5_Chat_8B

938 skip_prefixes = [

939 "action_embed", "temporal_embed", "track_embed",

vllm/model_executor/models/jais.py

+5 -4

@@ -21,7 +21,8 @@

21 """Inference-only Jais model compatible with HuggingFace weights."""

22

23 import math

24 -from typing import Iterable, Optional, Set, Tuple, Union

24 +from collections.abc import Iterable

25 +from typing import Optional, Union

26

27 import torch

28 from torch import nn

@@ -333,10 +334,10 @@ class JAISLMHeadModel(nn.Module, SupportsPP):

334 sampling_metadata)

335 return logits

336

336 - def load_weights(self, weights: Iterable[Tuple[str,

337 - torch.Tensor]]) -> Set[str]:

337 + def load_weights(self, weights: Iterable[tuple[str,

338 + torch.Tensor]]) -> set[str]:

339 params_dict = dict(self.named_parameters(remove_duplicate=False))

339 - loaded_params: Set[str] = set()

340 + loaded_params: set[str] = set()

341 for name, loaded_weight in weights:

342 if "lm_head.weight" in name:

343 # GPT-2 ties the weights of the embedding layer and the final

vllm/model_executor/models/jamba.py

+7 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Inference-only Jamba model."""

3 -from typing import Iterable, Optional, Set, Tuple

3 +from collections.abc import Iterable

4 +from typing import Optional

5

6 import torch

7 from torch import nn

@@ -442,7 +443,7 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

443 return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

444

445 def _get_mamba_cache_shape(

445 - self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

446 + self) -> tuple[tuple[int, int], tuple[int, int]]:

447 world_size = get_tensor_model_parallel_world_size()

448 hidden_size = self.config.hidden_size

449 conv_state_shape = (

@@ -464,8 +465,8 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

465 sampling_metadata)

466 return logits

467

467 - def load_weights(self, weights: Iterable[Tuple[str,

468 - torch.Tensor]]) -> Set[str]:

468 + def load_weights(self, weights: Iterable[tuple[str,

469 + torch.Tensor]]) -> set[str]:

470 stacked_params_mapping = [

471 # (param_name, shard_name, shard_id)

472 ("qkv_proj", "q_proj", "q"),

@@ -482,7 +483,7 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

483 num_experts=self.config.num_experts)

484

485 params_dict = dict(self.named_parameters())

485 - loaded_params: Set[str] = set()

486 + loaded_params: set[str] = set()

487 for name, loaded_weight in weights:

488 if "rotary_emb.inv_freq" in name:

489 continue

@@ -583,7 +584,7 @@ class JambaForSequenceClassification(JambaForCausalLM):

584 logits = self.score(hidden_states)

585 return self._pooler(logits, pooling_metadata)

586

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

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

588 # TODO: The reward weights themselves have float32 accuracy data, we

589 # would like to load them in fp32 to get that extra precision.

590 super().load_weights(weights)

vllm/model_executor/models/kimi_vl.py

+4 -5

@@ -43,10 +43,9 @@

43

44 import copy

45 import math

46 -from collections.abc import Mapping

46 +from collections.abc import Iterable, Mapping, Sequence

47 from dataclasses import dataclass

48 -from typing import (Any, Iterable, List, Literal, Optional, Sequence, Tuple,

49 - TypedDict, Union)

48 +from typing import Any, Literal, Optional, TypedDict, Union

49

50 import torch

51 from torch import nn

@@ -120,7 +119,7 @@ class KimiVLMultiModalProjector(nn.Module):

119

120 class KimiVLImagePixelInputs(TypedDict):

121 type: Literal["pixel_values"]

123 - pixel_values: Union[torch.Tensor, List[torch.Tensor]]

122 + pixel_values: Union[torch.Tensor, list[torch.Tensor]]

123 """

124 Shape:`(num_patches, num_channels, patch_size, patch_size)`

125 """

@@ -447,7 +446,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal):

446 sampling_metadata, **kwargs)

447 return logits

448

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

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

450 config = self.config.text_config

451 _KEYS_TO_MODIFY_MAPPING = {

452 "language_model.lm_head": "lm_head",

vllm/model_executor/models/llama.py

+10 -9

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

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

25 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Any, Optional, Union

27

28 import torch

29 from torch import nn

@@ -103,7 +104,7 @@ class LlamaAttention(nn.Module):

104 num_heads: int,

105 num_kv_heads: int,

106 rope_theta: float = 10000,

106 - rope_scaling: Optional[Dict[str, Any]] = None,

107 + rope_scaling: Optional[dict[str, Any]] = None,

108 max_position_embeddings: int = 8192,

109 quant_config: Optional[QuantizationConfig] = None,

110 bias: bool = False,

@@ -285,7 +286,7 @@ class LlamaDecoderLayer(nn.Module):

286 positions: torch.Tensor,

287 hidden_states: torch.Tensor,

288 residual: Optional[torch.Tensor],

288 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

290 # Self Attention

291 if residual is None:

292 residual = hidden_states

@@ -394,8 +395,8 @@ class LlamaModel(nn.Module):

395 return hidden_states, aux_hidden_states

396 return hidden_states

397

397 - def load_weights(self, weights: Iterable[Tuple[str,

398 - torch.Tensor]]) -> Set[str]:

398 + def load_weights(self, weights: Iterable[tuple[str,

399 + torch.Tensor]]) -> set[str]:

400 stacked_params_mapping = [

401 # (param_name, shard_name, shard_id)

402 (".qkv_proj", ".q_proj", "q"),

@@ -405,7 +406,7 @@ class LlamaModel(nn.Module):

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

407 ]

408 params_dict = dict(self.named_parameters())

408 - loaded_params: Set[str] = set()

409 + loaded_params: set[str] = set()

410 for name, loaded_weight in weights:

411 if "rotary_emb.inv_freq" in name:

412 continue

@@ -582,8 +583,8 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

583 sampling_metadata)

584 return logits

585

585 - def load_weights(self, weights: Iterable[Tuple[str,

586 - torch.Tensor]]) -> Set[str]:

586 + def load_weights(self, weights: Iterable[tuple[str,

587 + torch.Tensor]]) -> set[str]:

588 loader = AutoWeightsLoader(

589 self,

590 skip_prefixes=(["lm_head."]

@@ -599,7 +600,7 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

600 self,

601 name: str,

602 loaded_weight: torch.Tensor,

602 - ) -> Tuple[str, torch.Tensor]:

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

604

605 def permute(w: torch.Tensor, n_heads: int):

606 attn_in = self.config.head_dim * n_heads

vllm/model_executor/models/llama4.py

+14 -13

@@ -16,7 +16,8 @@

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

17 # limitations under the License.

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

19 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

19 +from collections.abc import Iterable

20 +from typing import Any, Optional

21

22 import torch

23 from torch import nn

@@ -48,7 +49,7 @@ class Llama4MoE(nn.Module):

49 gating_output: torch.Tensor,

50 topk: int,

51 renormalize: bool,

51 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

53 router_scores, router_indices = fast_topk(gating_output, topk, dim=-1)

54 # psuedo-standard is that the router scores are floats

55 router_scores = torch.sigmoid(router_scores.float())

@@ -115,7 +116,7 @@ class Llama4Attention(nn.Module):

116 num_heads: int,

117 num_kv_heads: int,

118 rope_theta: float = 10000,

118 - rope_scaling: Optional[Dict[str, Any]] = None,

119 + rope_scaling: Optional[dict[str, Any]] = None,

120 max_position_embeddings: int = 8192,

121 quant_config: Optional[QuantizationConfig] = None,

122 bias: bool = False,

@@ -300,7 +301,7 @@ class Llama4DecoderLayer(nn.Module):

301 positions: torch.Tensor,

302 hidden_states: torch.Tensor,

303 residual: Optional[torch.Tensor],

303 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

305 # Self Attention

306 if residual is None:

307 residual = hidden_states

@@ -335,9 +336,9 @@ class Llama4Model(LlamaModel):

336 self,

337 name: str,

338 loaded_weight: torch.Tensor,

338 - params_dict: Dict[str, nn.Parameter],

339 - loaded_params: Set[str],

340 - expert_params_mapping: List[Tuple[str, str, int, str]],

339 + params_dict: dict[str, nn.Parameter],

340 + loaded_params: set[str],

341 + expert_params_mapping: list[tuple[str, str, int, str]],

342 fused: bool = True,

343 ) -> bool:

344 expert_param_loaded = False

@@ -390,8 +391,8 @@ class Llama4Model(LlamaModel):

391 expert_param_loaded = True

392 return expert_param_loaded

393

393 - def load_weights(self, weights: Iterable[Tuple[str,

394 - torch.Tensor]]) -> Set[str]:

394 + def load_weights(self, weights: Iterable[tuple[str,

395 + torch.Tensor]]) -> set[str]:

396 stacked_params_mapping = [

397 # (param_name, shard_name, shard_id)

398 (".qkv_proj", ".q_proj", "q"),

@@ -412,7 +413,7 @@ class Llama4Model(LlamaModel):

413 ckpt_up_proj_name="gate_up_proj",

414 num_experts=1)

415 params_dict = dict(self.named_parameters())

415 - loaded_params: Set[str] = set()

416 + loaded_params: set[str] = set()

417 for name, loaded_weight in weights:

418 if "experts.gate_up_proj" in name or "experts.down_proj" in name:

419 fused_experts_params = True

@@ -489,8 +490,8 @@ class Llama4ForCausalLM(LlamaForCausalLM):

490 prefix=prefix,

491 layer_type=layer_type)

492

492 - def load_weights(self, weights: Iterable[Tuple[str,

493 - torch.Tensor]]) -> Set[str]:

493 + def load_weights(self, weights: Iterable[tuple[str,

494 + torch.Tensor]]) -> set[str]:

495 loader = AutoWeightsLoader(

496 self,

497 skip_prefixes=(["lm_head."]

@@ -506,7 +507,7 @@ class Llama4ForCausalLM(LlamaForCausalLM):

507 self,

508 name: str,

509 loaded_weight: torch.Tensor,

509 - ) -> Tuple[str, torch.Tensor]:

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

511

512 def permute(w: torch.Tensor, n_heads: int):

513 attn_in = self.config.head_dim * n_heads

vllm/model_executor/models/llama_eagle.py

+5 -5

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, Set, Tuple

3 +from collections.abc import Iterable

4

5 import torch

6 import torch.nn as nn

@@ -92,8 +92,8 @@ class LlamaModel(nn.Module):

92 hidden_states = hidden_states + residual

93 return hidden_states, hidden_states

94

95 - def load_weights(self, weights: Iterable[Tuple[str,

96 - torch.Tensor]]) -> Set[str]:

95 + def load_weights(self, weights: Iterable[tuple[str,

96 + torch.Tensor]]) -> set[str]:

97 stacked_params_mapping = [

98 # (param_name, shard_name, shard_id)

99 (".qkv_proj", ".q_proj", "q"),

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

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

104 ]

105 params_dict = dict(self.named_parameters())

106 - loaded_params: Set[str] = set()

106 + loaded_params: set[str] = set()

107 for name, loaded_weight in weights:

108 for param_name, weight_name, shard_id in stacked_params_mapping:

109 if weight_name not in name:

@@ -150,7 +150,7 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):

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

151 return self.model(input_ids, positions, hidden_states)

152

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

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

154 loader = AutoWeightsLoader(

155 self,

156 skip_prefixes=None,

vllm/model_executor/models/llama_eagle3.py

+7 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, Optional, Set, Tuple

3 +from collections.abc import Iterable

4 +from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -56,7 +57,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):

57 embeds: torch.Tensor,

58 hidden_states: torch.Tensor,

59 residual: Optional[torch.Tensor],

59 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

61

62 residual = hidden_states

63 embeds = self.input_layernorm(embeds)

@@ -140,8 +141,8 @@ class LlamaModel(nn.Module):

141 hidden_states, hidden_prenorm = self.norm(hidden_states, residual)

142 return hidden_states, hidden_prenorm

143

143 - def load_weights(self, weights: Iterable[Tuple[str,

144 - torch.Tensor]]) -> Set[str]:

144 + def load_weights(self, weights: Iterable[tuple[str,

145 + torch.Tensor]]) -> set[str]:

146 stacked_params_mapping = [

147 # (param_name, shard_name, shard_id)

148 (".qkv_proj", ".q_proj", "q"),

@@ -151,7 +152,7 @@ class LlamaModel(nn.Module):

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

153 ]

154 params_dict = dict(self.named_parameters())

154 - loaded_params: Set[str] = set()

155 + loaded_params: set[str] = set()

156 for name, loaded_weight in weights:

157 if 'midlayer.' in name:

158 name = name.replace('midlayer.', 'layers.0.')

@@ -228,7 +229,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

229 # combine multiple auxiliary hidden states returned by eagle3

230 return self.model.fc(hidden_states)

231

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

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

233 loader = AutoWeightsLoader(

234 self,

235 skip_prefixes=None,

vllm/model_executor/models/llava.py

+4 -4

@@ -2,8 +2,8 @@

2

3 from abc import abstractmethod

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import (Final, Literal, Optional, Protocol, Set, Tuple, TypedDict,

6 - TypeVar, Union, cast)

5 +from typing import (Final, Literal, Optional, Protocol, TypedDict, TypeVar,

6 + Union, cast)

7

8 import torch

9 import torch.nn as nn

@@ -751,8 +751,8 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

751 return self.language_model.compute_logits(hidden_states,

752 sampling_metadata)

753

754 - def load_weights(self, weights: Iterable[Tuple[str,

755 - torch.Tensor]]) -> Set[str]:

754 + def load_weights(self, weights: Iterable[tuple[str,

755 + torch.Tensor]]) -> set[str]:

756 loader = AutoWeightsLoader(self)

757 return loader.load_weights(weights)

758

vllm/model_executor/models/llava_next.py

+8 -7

@@ -1,8 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 from abc import abstractmethod

4 -from typing import (Final, Iterable, List, Literal, Mapping, Optional,

5 - Protocol, Set, Tuple, TypedDict, TypeVar, Union)

4 +from collections.abc import Iterable, Mapping

5 +from typing import (Final, Literal, Optional, Protocol, TypedDict, TypeVar,

6 + Union)

7

8 import torch

9 import torch.nn as nn

@@ -266,8 +267,8 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

267 return data

268

269 def _validate_pixel_values(

269 - self, data: Union[torch.Tensor, List[torch.Tensor]]

270 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

270 + self, data: Union[torch.Tensor, list[torch.Tensor]]

271 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

272

273 h = w = self.config.vision_config.image_size

274 expected_dims = (3, h, w)

@@ -450,7 +451,7 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

451 def _process_image_input(

452 self,

453 image_input: LlavaNextImageInputs,

453 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

454 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

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

456 return [image_input["data"]]

457

@@ -577,7 +578,7 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

578 return self.language_model.compute_logits(hidden_states,

579 sampling_metadata)

580

580 - def load_weights(self, weights: Iterable[Tuple[str,

581 - torch.Tensor]]) -> Set[str]:

581 + def load_weights(self, weights: Iterable[tuple[str,

582 + torch.Tensor]]) -> set[str]:

583 loader = AutoWeightsLoader(self)

584 return loader.load_weights(weights)

vllm/model_executor/models/llava_next_video.py

+7 -7

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

2

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union

5 +from typing import Literal, Optional, TypedDict, Union

6

7 import torch

8 import torch.nn as nn

@@ -35,7 +35,7 @@ from .vision import get_vision_encoder_info

35

36 class LlavaNextVideoPixelInputs(TypedDict):

37 type: Literal["pixel_values_videos"]

38 - data: Union[torch.Tensor, List[torch.Tensor]]

38 + data: Union[torch.Tensor, list[torch.Tensor]]

39 """

40 Shape: `(batch_size, num_frames, num_channels, height, width)`

41

@@ -300,8 +300,8 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

300 self.language_model.model.make_empty_intermediate_tensors)

301

302 def _validate_video_pixel_values(

303 - self, data: Union[torch.Tensor, List[torch.Tensor]]

304 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

303 + self, data: Union[torch.Tensor, list[torch.Tensor]]

304 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

305

306 h = w = self.config.vision_config.image_size

307 expected_dims = (3, h, w)

@@ -326,7 +326,7 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

326 A legal video input should have the following dimensions:

327 {

328 "pixel_values_videos" :

329 - List[b, Tensor(nb_frames, nb_channels, height, width)]

329 + list[b, Tensor(nb_frames, nb_channels, height, width)]

330 }

331 """

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

@@ -460,8 +460,8 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

460 return self.language_model.compute_logits(hidden_states,

461 sampling_metadata)

462

463 - def load_weights(self, weights: Iterable[Tuple[str,

464 - torch.Tensor]]) -> Set[str]:

463 + def load_weights(self, weights: Iterable[tuple[str,

464 + torch.Tensor]]) -> set[str]:

465 loader = AutoWeightsLoader(

466 self,

467 # This model doesn't support images for now

vllm/model_executor/models/llava_onevision.py

+10 -11

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

2

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import (Final, List, Literal, Optional, Protocol, Set, Tuple,

6 - TypedDict, Union)

5 +from typing import Final, Literal, Optional, Protocol, TypedDict, Union

6

7 import torch

8 import torch.nn as nn

@@ -471,8 +470,8 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

470 return data

471

472 def _validate_image_pixel_values(

474 - self, data: Union[torch.Tensor, List[torch.Tensor]]

475 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

473 + self, data: Union[torch.Tensor, list[torch.Tensor]]

474 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

475

476 h = w = self.config.vision_config.image_size

477 expected_dims = (3, h, w)

@@ -530,8 +529,8 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

529 raise AssertionError("This line should be unreachable.")

530

531 def _validate_video_pixel_values(

533 - self, data: Union[torch.Tensor, List[torch.Tensor]]

534 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

532 + self, data: Union[torch.Tensor, list[torch.Tensor]]

533 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

534

535 h = w = self.config.vision_config.image_size

536 expected_dims = (3, h, w)

@@ -557,7 +556,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

556 A legal video input should have the following dimensions:

557 {

558 "pixel_values_videos" :

560 - List[b, Tensor(nb_frames, nb_channels, height, width)]

559 + list[b, Tensor(nb_frames, nb_channels, height, width)]

560 }

561 """

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

@@ -706,7 +705,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

705 def _process_image_pixels(

706 self,

707 inputs: LlavaOnevisionImagePixelInputs,

709 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

708 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

709 assert self.vision_tower is not None

710

711 pixel_values = inputs["pixel_values"]

@@ -735,7 +734,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

734 def _process_image_input(

735 self,

736 image_input: LlavaOnevisionImageInputs,

738 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

737 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

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

739 return [image_input["data"]]

740

@@ -948,7 +947,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

947 return self.language_model.compute_logits(hidden_states,

948 sampling_metadata)

949

951 - def load_weights(self, weights: Iterable[Tuple[str,

952 - torch.Tensor]]) -> Set[str]:

950 + def load_weights(self, weights: Iterable[tuple[str,

951 + torch.Tensor]]) -> set[str]:

952 loader = AutoWeightsLoader(self)

953 return loader.load_weights(weights)

vllm/model_executor/models/mamba.py

+9 -8

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

1 # SPDX-License-Identifier: Apache-2.0

2 """PyTorch MAMBA model."""

3 -from typing import Iterable, Optional, Set, Tuple

3 +from collections.abc import Iterable

4 +from typing import Optional

5

6 import torch

7 from torch import nn

@@ -30,7 +31,7 @@ from .utils import (AutoWeightsLoader, is_pp_missing_parameter,

31 make_empty_intermediate_tensors_factory, make_layers,

32 maybe_prefix)

33

33 -KVCache = Tuple[torch.Tensor, torch.Tensor]

34 +KVCache = tuple[torch.Tensor, torch.Tensor]

35

36

37 class MambaDecoderLayer(nn.Module):

@@ -153,10 +154,10 @@ class MambaModel(nn.Module):

154

155 return hidden_states

156

156 - def load_weights(self, weights: Iterable[Tuple[str,

157 - torch.Tensor]]) -> Set[str]:

157 + def load_weights(self, weights: Iterable[tuple[str,

158 + torch.Tensor]]) -> set[str]:

159 params_dict = dict(self.named_parameters())

159 - loaded_params: Set[str] = set()

160 + loaded_params: set[str] = set()

161 for name, loaded_weight in weights:

162 if "A_log" in name:

163 name = name.replace("A_log", "A")

@@ -247,7 +248,7 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP,

248 return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

249

250 def _get_mamba_cache_shape(

250 - self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

251 + self) -> tuple[tuple[int, int], tuple[int, int]]:

252 world_size = get_tensor_model_parallel_world_size()

253 conv_state_shape = (

254 self.config.intermediate_size // world_size,

@@ -265,7 +266,7 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP,

266 sampling_metadata)

267 return logits

268

268 - def load_weights(self, weights: Iterable[Tuple[str,

269 - torch.Tensor]]) -> Set[str]:

269 + def load_weights(self, weights: Iterable[tuple[str,

270 + torch.Tensor]]) -> set[str]:

271 loader = AutoWeightsLoader(self)

272 return loader.load_weights(weights)

vllm/model_executor/models/mamba2.py

+7 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2 """PyTorch MAMBA2 model."""

3 -from typing import Iterable, Optional, Set, Tuple

3 +from collections.abc import Iterable

4 +from typing import Optional

5

6 import torch

7 from torch import nn

@@ -35,7 +36,7 @@ from .utils import (is_pp_missing_parameter,

36 make_empty_intermediate_tensors_factory, make_layers,

37 maybe_prefix)

38

38 -KVCache = Tuple[torch.Tensor, torch.Tensor]

39 +KVCache = tuple[torch.Tensor, torch.Tensor]

40

41

42 class Mamba2DecoderLayer(nn.Module):

@@ -241,7 +242,7 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree,

242 return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

243

244 def _get_mamba_cache_shape(

244 - self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

245 + self) -> tuple[tuple[int, int], tuple[int, int]]:

246 world_size = get_tensor_model_parallel_world_size()

247

248 conv_state_shape, temporal_state_shape = None, None

@@ -279,10 +280,10 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree,

280 sampling_metadata)

281 return logits

282

282 - def load_weights(self, weights: Iterable[Tuple[str,

283 - torch.Tensor]]) -> Set[str]:

283 + def load_weights(self, weights: Iterable[tuple[str,

284 + torch.Tensor]]) -> set[str]:

285 params_dict = dict(self.named_parameters())

285 - loaded_params: Set[str] = set()

286 + loaded_params: set[str] = set()

287 for name, loaded_weight in weights:

288 if "A_log" in name:

289 name = name.replace("A_log", "A")

vllm/model_executor/models/mamba_cache.py

+2 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 from dataclasses import dataclass

4 -from typing import Tuple

4

5 import torch

6

@@ -25,8 +24,8 @@ class MambaCacheParams:

24 class MambaCacheManager(ConstantSizeCache):

25

26 def __init__(self, vllm_config: VllmConfig, dtype: torch.dtype,

28 - num_mamba_layers: int, conv_state_shape: Tuple[int, int],

29 - temporal_state_shape: Tuple[int, int]):

27 + num_mamba_layers: int, conv_state_shape: tuple[int, int],

28 + temporal_state_shape: tuple[int, int]):

29

30 # Determine max batch size to set size of MambaCache

31 max_batch_size = vllm_config.scheduler_config.max_num_seqs

vllm/model_executor/models/medusa.py

+13 -12

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, List, Optional, Set, Tuple

3 +from collections.abc import Iterable

4 +from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -96,13 +97,13 @@ class Medusa(nn.Module):

97 # checkpoint file has token_map tensor.

98 self.token_map = None

99

99 - def forward(self, hidden_states: torch.Tensor) -> List[torch.Tensor]:

100 + def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]:

101 return [block(hidden_states) for block in self.blocks]

102

103 def compute_logits(

103 - self, hidden_states: List[torch.Tensor],

104 - sampling_metadata: SamplingMetadata) -> List[torch.Tensor]:

105 - logits_lst: List[torch.Tensor] = []

104 + self, hidden_states: list[torch.Tensor],

105 + sampling_metadata: SamplingMetadata) -> list[torch.Tensor]:

106 + logits_lst: list[torch.Tensor] = []

107

108 for hs, lm_head in zip(hidden_states, self.lm_heads):

109 _logits = self.logits_processor(lm_head, hs, sampling_metadata)

@@ -127,9 +128,9 @@ class Medusa(nn.Module):

128

129 def sample(

130 self,

130 - logits: List[torch.Tensor],

131 + logits: list[torch.Tensor],

132 sampling_metadata: SamplingMetadata,

132 - ) -> List[SamplerOutput]:

133 + ) -> list[SamplerOutput]:

134 logits = torch.stack(logits, dim=0).float()

135 logprobs = torch.log_softmax(logits, dim=-1)

136 token_ids = logits.argmax(-1) # support only top-1 for now

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

145 token_prob_list.append(probs[:, seq_group.sample_indices])

146 token_logprob_list.append(logprobs[:, seq_group.sample_indices])

147

147 - outputs: List[Optional[SamplerOutput]] = []

148 + outputs: list[Optional[SamplerOutput]] = []

149 for idx in range(len(sampling_metadata.seq_groups)):

150 outputs.append(

151 SamplerOutput(

@@ -160,7 +161,7 @@ class Medusa(nn.Module):

161 self,

162 previous_hidden_states: torch.Tensor,

163 sampling_metadata: SamplingMetadata,

163 - ) -> List[SamplerOutput]:

164 + ) -> list[SamplerOutput]:

165 return self.sample(

166 logits=self.compute_logits(

167 hidden_states=self.forward(previous_hidden_states),

@@ -169,10 +170,10 @@ class Medusa(nn.Module):

170 sampling_metadata=sampling_metadata,

171 )

172

172 - def load_weights(self, weights: Iterable[Tuple[str,

173 - torch.Tensor]]) -> Set[str]:

173 + def load_weights(self, weights: Iterable[tuple[str,

174 + torch.Tensor]]) -> set[str]:

175 params_dict = dict(self.named_parameters())

175 - loaded_params: Set[str] = set()

176 + loaded_params: set[str] = set()

177

178 weights_map = {}

179

vllm/model_executor/models/mimo.py

+5 -4

@@ -24,7 +24,8 @@

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

25 # limitations under the License.

26 """Inference-only MiMo model compatible with HuggingFace weights."""

27 -from typing import Iterable, Optional, Set, Tuple, Union

27 +from collections.abc import Iterable

28 +from typing import Optional, Union

29

30 import torch

31 import torch.nn as nn

@@ -87,8 +88,8 @@ class MiMoModel(Qwen2Model):

88 hidden_states = hidden_states + residual

89 return hidden_states

90

90 - def load_weights(self, weights: Iterable[Tuple[str,

91 - torch.Tensor]]) -> Set[str]:

91 + def load_weights(self, weights: Iterable[tuple[str,

92 + torch.Tensor]]) -> set[str]:

93 stacked_params_mapping = [

94 ("qkv_proj", "q_proj", "q"),

95 ("qkv_proj", "k_proj", "k"),

@@ -97,7 +98,7 @@ class MiMoModel(Qwen2Model):

98 ("gate_up_proj", "up_proj", 1),

99 ]

100 params_dict = dict(self.named_parameters(remove_duplicate=False))

100 - loaded_params: Set[str] = set()

101 + loaded_params: set[str] = set()

102 for name, loaded_weight in weights:

103 if "mtp_layers" in name:

104 continue

vllm/model_executor/models/mimo_mtp.py

+5 -4

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

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

19 # limitations under the License.

20 """Inference-only MiMo-MTP model."""

21 -from typing import Iterable, Optional, Set, Tuple

21 +from collections.abc import Iterable

22 +from typing import Optional

23

24 import torch

25 import torch.nn as nn

@@ -193,8 +194,8 @@ class MiMoMTP(nn.Module):

194 next_tokens = self.sampler(logits, sampling_metadata)

195 return next_tokens

196

196 - def load_weights(self, weights: Iterable[Tuple[str,

197 - torch.Tensor]]) -> Set[str]:

197 + def load_weights(self, weights: Iterable[tuple[str,

198 + torch.Tensor]]) -> set[str]:

199 stacked_params_mapping = [

200 ("qkv_proj", "q_proj", "q"),

201 ("qkv_proj", "k_proj", "k"),

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

205 ]

206

207 params_dict = dict(self.named_parameters())

207 - loaded_params: Set[str] = set()

208 + loaded_params: set[str] = set()

209 for name, loaded_weight in weights:

210

211 if "rotary_emb.inv_freq" in name:

vllm/model_executor/models/minicpm.py

+9 -8

@@ -23,7 +23,8 @@

23 # limitations under the License.

24 """Inference-only MiniCPM model compatible with HuggingFace weights."""

25 import math

26 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

26 +from collections.abc import Iterable

27 +from typing import Any, Optional, Union

28

29 import torch

30 from torch import nn

@@ -190,7 +191,7 @@ class MiniCPMAttention(nn.Module):

191 num_heads: int,

192 num_kv_heads: int,

193 rope_theta: float = 10000,

193 - rope_scaling: Optional[Dict[str, Any]] = None,

194 + rope_scaling: Optional[dict[str, Any]] = None,

195 max_position_embeddings: int = 8192,

196 cache_config: Optional[CacheConfig] = None,

197 quant_config: Optional[QuantizationConfig] = None,

@@ -329,7 +330,7 @@ class MiniCPMDecoderLayer(nn.Module):

330 positions: torch.Tensor,

331 hidden_states: torch.Tensor,

332 residual: Optional[torch.Tensor],

332 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

334 # Self Attention

335 residual = hidden_states

336 hidden_states = self.input_layernorm(hidden_states)

@@ -428,8 +429,8 @@ class MiniCPMModel(nn.Module):

429 hidden_states = self.norm(hidden_states)

430 return hidden_states

431

431 - def load_weights(self, weights: Iterable[Tuple[str,

432 - torch.Tensor]]) -> Set[str]:

432 + def load_weights(self, weights: Iterable[tuple[str,

433 + torch.Tensor]]) -> set[str]:

434 stacked_params_mapping = [

435 # (param_name, shard_name, shard_id)

436 ("qkv_proj", "q_proj", "q"),

@@ -446,7 +447,7 @@ class MiniCPMModel(nn.Module):

447 for weight_name in ["w1", "w2", "w3"]

448 ]

449 params_dict = dict(self.named_parameters())

449 - loaded_params: Set[str] = set()

450 + loaded_params: set[str] = set()

451 for name, loaded_weight in weights:

452 if "rotary_emb.inv_freq" in name:

453 continue

@@ -582,8 +583,8 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

583 sampling_metadata)

584 return logits

585

585 - def load_weights(self, weights: Iterable[Tuple[str,

586 - torch.Tensor]]) -> Set[str]:

586 + def load_weights(self, weights: Iterable[tuple[str,

587 + torch.Tensor]]) -> set[str]:

588 loader = AutoWeightsLoader(

589 self,

590 skip_prefixes=(["lm_head."]

vllm/model_executor/models/minicpm3.py

+2 -2

@@ -23,7 +23,7 @@

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

24 # limitations under the License.

25 """Inference-only MiniCPM3 model compatible with HuggingFace weights."""

26 -from typing import Any, Dict, Optional

26 +from typing import Any, Optional

27

28 import torch

29 from torch import nn

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

58 q_lora_rank: int,

59 kv_lora_rank: int,

60 rope_theta: float = 10000,

61 - rope_scaling: Optional[Dict[str, Any]] = None,

61 + rope_scaling: Optional[dict[str, Any]] = None,

62 max_position_embeddings: int = 8192,

63 cache_config: Optional[CacheConfig] = None,

64 quant_config: Optional[QuantizationConfig] = None,

vllm/model_executor/models/minicpmo.py

+3 -4

@@ -23,8 +23,7 @@

23 # limitations under the License.

24 """Inference-only MiniCPM-O model compatible with HuggingFace weights."""

25 from collections.abc import Iterable, Mapping, Sequence

26 -from typing import (Any, Callable, Literal, Optional, Set, Tuple, TypedDict,

27 - Union)

26 +from typing import Any, Callable, Literal, Optional, TypedDict, Union

27

28 import torch

29 from torch import nn

@@ -559,8 +558,8 @@ class MiniCPMO(MiniCPMV2_6):

558 self.audio_encoder_layer = -1

559 return model

560

562 - def load_weights(self, weights: Iterable[Tuple[str,

563 - torch.Tensor]]) -> Set[str]:

561 + def load_weights(self, weights: Iterable[tuple[str,

562 + torch.Tensor]]) -> set[str]:

563 loader = AutoWeightsLoader(self, skip_prefixes=["tts"])

564 return loader.load_weights(weights)

565

vllm/model_executor/models/minicpmv.py

+6 -7

@@ -26,8 +26,7 @@ import math

26 from collections import defaultdict

27 from collections.abc import Iterable, Mapping, Sequence

28 from functools import partial

29 -from typing import (Any, Callable, Literal, Optional, Set, Tuple, TypedDict,

30 - Union)

29 +from typing import Any, Callable, Literal, Optional, TypedDict, Union

30

31 import numpy as np

32 import torch

@@ -118,7 +117,7 @@ class Resampler2_5(BaseResampler):

117 num_heads: int,

118 kv_dim: Optional[int] = None,

119 norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,

121 - max_size: Tuple[int, int] = (70, 70),

120 + max_size: tuple[int, int] = (70, 70),

121 quant_config: Optional[QuantizationConfig] = None,

122 prefix: str = "") -> None:

123 super().__init__(num_queries,

@@ -133,7 +132,7 @@ class Resampler2_5(BaseResampler):

132 self._set_2d_pos_cache(self.max_size)

133

134 def _set_2d_pos_cache(self,

136 - max_size: Tuple[int, int],

135 + max_size: tuple[int, int],

136 device: torch.types.Device = "cpu") -> None:

137 pos_embed_arr = get_2d_sincos_pos_embed(self.embed_dim,

138 max_size,

@@ -203,7 +202,7 @@ class Resampler2_5(BaseResampler):

202 return x

203

204

206 -def get_version_by_config(config: PretrainedConfig) -> Tuple[int, ...]:

205 +def get_version_by_config(config: PretrainedConfig) -> tuple[int, ...]:

206 version_float = getattr(config, "version", None)

207

208 # The old configs do not include version number

@@ -938,8 +937,8 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

937 ) -> Optional[torch.Tensor]:

938 return self.llm.compute_logits(hidden_states, sampling_metadata)

939

941 - def load_weights(self, weights: Iterable[Tuple[str,

942 - torch.Tensor]]) -> Set[str]:

940 + def load_weights(self, weights: Iterable[tuple[str,

941 + torch.Tensor]]) -> set[str]:

942 loader = AutoWeightsLoader(self)

943 return loader.load_weights(weights)

944

vllm/model_executor/models/minimax_text_01.py

+9 -8

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

3 import copy

4 import math

5 import re

6 -from typing import Dict, Iterable, List, Optional, Set, Tuple, Union

6 +from collections.abc import Iterable

7 +from typing import Optional, Union

8

9 import torch

10 import torch.distributed

@@ -127,7 +128,7 @@ class MiniMaxText01RMSNormTP(CustomOp):

128 self,

129 x: torch.Tensor,

130 residual: Optional[torch.Tensor] = None,

130 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

131 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

132 assert residual is None, "RMSNorm does not support residual connection."

133 return self._forward(x)

134

@@ -178,7 +179,7 @@ class MiniMaxText01RotaryEmbedding(CustomOp):

179 positions: torch.Tensor,

180 query: torch.Tensor,

181 key: torch.Tensor,

181 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

183 from vllm import _custom_ops as ops

184 self.cos_sin_cache = self.cos_sin_cache.to(positions.device)

185 query_cast = query.to(self.cache_dtype)

@@ -708,11 +709,11 @@ class MiniMaxText01DecoderLayer(nn.Module):

709 def forward(self,

710 hidden_states: torch.Tensor,

711 positions: torch.Tensor,

711 - kv_caches: Union[List[Dict], Optional[torch.Tensor]],

712 + kv_caches: Union[list[dict], Optional[torch.Tensor]],

713 attn_metadata: AttentionMetadata,

714 residual: Optional[torch.Tensor],

715 is_warmup: bool = False,

715 - **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:

716 + **kwargs) -> tuple[torch.Tensor, torch.Tensor]:

717

718 forward_context = get_forward_context()

719 attn_metadata = forward_context.attn_metadata

@@ -1072,10 +1073,10 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid,

1073 device=device),

1074 })

1075

1075 - def load_weights(self, weights: Iterable[Tuple[str,

1076 - torch.Tensor]]) -> Set[str]:

1076 + def load_weights(self, weights: Iterable[tuple[str,

1077 + torch.Tensor]]) -> set[str]:

1078 params_dict = dict(self.named_parameters())

1078 - loaded_params: Set[str] = set()

1079 + loaded_params: set[str] = set()

1080

1081 def which_layer(name: str) -> int:

1082 if "layers" in name:

vllm/model_executor/models/minimax_vl_01.py

+3 -3

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 from collections.abc import Iterable, Mapping

3 -from typing import Literal, Optional, Set, Tuple, TypedDict, Union, cast

3 +from typing import Literal, Optional, TypedDict, Union, cast

4

5 import torch

6 import torch.nn as nn

@@ -357,7 +357,7 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal,

357 return self.language_model.compute_logits(hidden_states,

358 sampling_metadata)

359

360 - def load_weights(self, weights: Iterable[Tuple[str,

361 - torch.Tensor]]) -> Set[str]:

360 + def load_weights(self, weights: Iterable[tuple[str,

361 + torch.Tensor]]) -> set[str]:

362 loader = AutoWeightsLoader(self)

363 return loader.load_weights(weights)

vllm/model_executor/models/mistral3.py

+4 -4

@@ -2,8 +2,8 @@

2

3 from abc import abstractmethod

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import (Final, Literal, Optional, Protocol, Set, Tuple, TypedDict,

6 - TypeVar, Union)

5 +from typing import (Final, Literal, Optional, Protocol, TypedDict, TypeVar,

6 + Union)

7

8 import torch

9 import torch.nn as nn

@@ -589,8 +589,8 @@ class Mistral3ForConditionalGeneration(nn.Module, SupportsLoRA,

589 return self.language_model.compute_logits(hidden_states,

590 sampling_metadata)

591

592 - def load_weights(self, weights: Iterable[Tuple[str,

593 - torch.Tensor]]) -> Set[str]:

592 + def load_weights(self, weights: Iterable[tuple[str,

593 + torch.Tensor]]) -> set[str]:

594 loader = AutoWeightsLoader(self)

595 return loader.load_weights(weights)

596

vllm/model_executor/models/mixtral.py

+7 -6

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only Mixtral model."""

25 -from typing import Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -314,8 +315,8 @@ class MixtralModel(nn.Module):

315 hidden_states, _ = self.norm(hidden_states, residual)

316 return hidden_states

317

317 - def load_weights(self, weights: Iterable[Tuple[str,

318 - torch.Tensor]]) -> Set[str]:

318 + def load_weights(self, weights: Iterable[tuple[str,

319 + torch.Tensor]]) -> set[str]:

320 stacked_params_mapping = [

321 # (param_name, shard_name, shard_id)

322 ("qkv_proj", "q_proj", "q"),

@@ -332,7 +333,7 @@ class MixtralModel(nn.Module):

333 num_experts=self.config.num_local_experts)

334

335 params_dict = dict(self.named_parameters())

335 - loaded_params: Set[str] = set()

336 + loaded_params: set[str] = set()

337 for name, loaded_weight in weights:

338 if (self.quant_config is not None and

339 (scale_name := self.quant_config.get_cache_scale(name))):

@@ -479,7 +480,7 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

480 sampling_metadata)

481 return logits

482

482 - def load_weights(self, weights: Iterable[Tuple[str,

483 - torch.Tensor]]) -> Set[str]:

483 + def load_weights(self, weights: Iterable[tuple[str,

484 + torch.Tensor]]) -> set[str]:

485 loader = AutoWeightsLoader(self, skip_prefixes=["rotary_emb.inv_freq"])

486 return loader.load_weights(weights)

vllm/model_executor/models/mixtral_quant.py

+5 -4

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only Mixtral model."""

25 -from typing import Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Optional, Union

27

28 import numpy as np

29 import torch

@@ -397,8 +398,8 @@ class MixtralForCausalLM(nn.Module, SupportsPP):

398 sampling_metadata)

399 return logits

400

400 - def load_weights(self, weights: Iterable[Tuple[str,

401 - torch.Tensor]]) -> Set[str]:

401 + def load_weights(self, weights: Iterable[tuple[str,

402 + torch.Tensor]]) -> set[str]:

403 stacked_params_mapping = [

404 # (param_name, shard_name, shard_id)

405 ("qkv_proj", "q_proj", "q"),

@@ -407,7 +408,7 @@ class MixtralForCausalLM(nn.Module, SupportsPP):

408 ]

409

410 params_dict = dict(self.named_parameters())

410 - loaded_params: Set[str] = set()

411 + loaded_params: set[str] = set()

412 for name, loaded_weight in weights:

413 if "rotary_emb.inv_freq" in name:

414 continue

vllm/model_executor/models/mllama.py

+40 -40

@@ -16,7 +16,7 @@

16 """PyTorch Mllama model."""

17 import math

18 from collections.abc import Iterable, Mapping, Sequence

19 -from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union

19 +from typing import Literal, Optional, TypedDict, Union

20

21 import numpy as np

22 import torch

@@ -224,7 +224,7 @@ class MllamaMultiModalProcessor(EncDecMultiModalProcessor[MllamaProcessingInfo]

224

225 return mm_inputs

226

227 - def _get_num_image_in_last_group(self, prompt_token_ids: List[int]) -> int:

227 + def _get_num_image_in_last_group(self, prompt_token_ids: list[int]) -> int:

228 num_images = 0

229 for token_id in prompt_token_ids[::-1]:

230 if token_id == self.info.get_hf_config().image_token_index:

@@ -370,8 +370,8 @@ class ColumnParallelConv2dPatch(torch.nn.Module):

370 self,

371 in_channels: int,

372 out_channels: int,

373 - kernel_size: Union[int, Tuple[int, int]],

374 - stride: Union[int, Tuple[int, int]],

373 + kernel_size: Union[int, tuple[int, int]],

374 + stride: Union[int, tuple[int, int]],

375 bias: bool = False,

376 ) -> None:

377 super().__init__()

@@ -603,7 +603,7 @@ class MllamaVisionEncoder(nn.Module):

603 self,

604 hidden_states: torch.Tensor,

605 attention_mask: Optional[torch.Tensor] = None,

606 - ) -> Union[Tuple, BaseModelOutput]:

606 + ) -> Union[BaseModelOutput]:

607 encoder_states = ()

608

609 for i, encoder_layer in enumerate(self.layers):

@@ -878,7 +878,7 @@ class MllamaTextCrossAttention(nn.Module):

878 self,

879 hidden_states: torch.Tensor,

880 attention_mask: Optional[torch.Tensor],

881 - kv_range_for_decode: Optional[List[Tuple[int, int]]],

881 + kv_range_for_decode: Optional[list[tuple[int, int]]],

882 cross_attention_states: Optional[torch.Tensor],

883 ) -> torch.Tensor:

884 q, k, v = self.qkv_proj(hidden_states, cross_attention_states)

@@ -905,7 +905,7 @@ class MllamaTextCrossAttention(nn.Module):

905 k: torch.Tensor,

906 v: torch.Tensor,

907 attention_mask: torch.Tensor,

908 - kv_range_for_decode: List[Tuple[int, int]],

908 + kv_range_for_decode: list[tuple[int, int]],

909 ) -> torch.Tensor:

910 kv_cache = self.attn.kv_cache[self.pipeline_parallel_rank]

911 attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

@@ -1019,7 +1019,7 @@ class MllamaCrossAttentionDecoderLayer(torch.nn.Module):

1019 hidden_states: torch.Tensor,

1020 cross_attention_states: torch.Tensor,

1021 cross_attention_mask: torch.Tensor,

1022 - kv_range_for_decode: Optional[List[Tuple[int, int]]],

1022 + kv_range_for_decode: Optional[list[tuple[int, int]]],

1023 full_text_row_masked_out_mask: torch.Tensor,

1024 ) -> torch.Tensor:

1025 residual = hidden_states

@@ -1089,8 +1089,8 @@ class MllamaTextModel(nn.Module):

1089 positions: Optional[torch.LongTensor],

1090 cross_attention_states: Optional[torch.LongTensor],

1091 cross_attention_mask: Optional[torch.LongTensor],

1092 - kv_range_for_decode: Optional[List[Tuple[int, int]]],

1093 - full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor,

1092 + kv_range_for_decode: Optional[list[tuple[int, int]]],

1093 + full_text_row_masked_out_mask: Optional[tuple[torch.Tensor,

1094 torch.Tensor]],

1095 skip_cross_attention: bool,

1096 ) -> torch.Tensor:

@@ -1150,8 +1150,8 @@ class MllamaForCausalLM(nn.Module):

1150 positions: Optional[torch.LongTensor],

1151 cross_attention_states: Optional[torch.LongTensor],

1152 cross_attention_mask: Optional[torch.LongTensor],

1153 - kv_range_for_decode: Optional[List[Tuple[int, int]]],

1154 - full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor,

1153 + kv_range_for_decode: Optional[list[tuple[int, int]]],

1154 + full_text_row_masked_out_mask: Optional[tuple[torch.Tensor,

1155 torch.Tensor]],

1156 skip_cross_attention: bool,

1157 ) -> torch.Tensor:

@@ -1221,7 +1221,7 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1221 return logits

1222

1223 def unpack_data(self,

1224 - image_data: Union[List[torch.Tensor], torch.Tensor],

1224 + image_data: Union[list[torch.Tensor], torch.Tensor],

1225 padding_value=0) -> torch.Tensor:

1226 if isinstance(image_data, torch.Tensor):

1227 # torch.Tensor

@@ -1230,7 +1230,7 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1230 assert isinstance(

1231 image_data[0],

1232 torch.Tensor), "Image data is not properly batched."

1233 - # List[torch.Tensor]

1233 + # list[torch.Tensor]

1234 bsz = len(image_data)

1235 max_length = max(t.size(0) for t in image_data)

1236 trailing_dims = image_data[0].shape[1:]

@@ -1248,24 +1248,24 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1248 def _parse_and_validate_image_input(self, **kwargs: object):

1249 # tensor with the same shape will be batched together by

1250 # MultiModalKwargs.batch, so pixel_values here can be:

1251 - # - List[torch.Tensor]:

1251 + # - list[torch.Tensor]:

1252 # with shape (num_image, num_tiles, 3, image_res, image_res)

1253 # - torch.Tensor:

1254 # with shape (bs, num_image, num_tiles, 3, image_res, image_res)

1255 - pixel_values: Optional[Union[List[List[torch.Tensor]],

1256 - List[torch.Tensor],

1255 + pixel_values: Optional[Union[list[list[torch.Tensor]],

1256 + list[torch.Tensor],

1257 torch.Tensor]] = kwargs.pop(

1258 "pixel_values", None)

1259 - image_embeds: Optional[Union[List[List[torch.Tensor]],

1260 - List[torch.Tensor],

1259 + image_embeds: Optional[Union[list[list[torch.Tensor]],

1260 + list[torch.Tensor],

1261 torch.Tensor]] = kwargs.pop(

1262 "image_embeds", None)

1263 - aspect_ratio_ids: Optional[Union[List[List[torch.Tensor]],

1264 - List[torch.Tensor],

1263 + aspect_ratio_ids: Optional[Union[list[list[torch.Tensor]],

1264 + list[torch.Tensor],

1265 torch.Tensor]] = kwargs.pop(

1266 "aspect_ratio_ids", None)

1267 - aspect_ratio_mask: Optional[Union[List[List[torch.Tensor]],

1268 - List[torch.Tensor],

1267 + aspect_ratio_mask: Optional[Union[list[list[torch.Tensor]],

1268 + list[torch.Tensor],

1269 torch.Tensor]] = kwargs.pop(

1270 "aspect_ratio_mask", None)

1271

@@ -1293,10 +1293,10 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1293

1294 def _get_and_validate_encoder_lens(

1295 self,

1296 - encoder_seq_lens: List[int],

1297 - num_tiles: List[List[int]],

1296 + encoder_seq_lens: list[int],

1297 + num_tiles: list[list[int]],

1298 num_tokens_per_tile: int,

1299 - ) -> List[int]:

1299 + ) -> list[int]:

1300 # Get the actual number of encoder tokens for each sample.

1301 # Because attn_metadata.encoder_seq_lens only counts the last

1302 # group of images for each sample, which is used to cheat the

@@ -1318,7 +1318,7 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1318

1319 def flat_encoder_result(self, cross_attention_states: torch.Tensor,

1320 attn_metadata: AttentionMetadata,

1321 - actual_encoder_seq_lens: List[int]):

1321 + actual_encoder_seq_lens: list[int]):

1322

1323 cross_attention_states_flat = torch.zeros(

1324 sum(actual_encoder_seq_lens),

@@ -1342,8 +1342,8 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1342 self,

1343 image_inputs: MllamaImagePixelInputs,

1344 attn_metadata: AttentionMetadata,

1345 - actual_encoder_seq_lens: List[int],

1346 - ) -> Tuple[torch.Tensor]:

1345 + actual_encoder_seq_lens: list[int],

1346 + ) -> tuple[torch.Tensor]:

1347 # NOTE: llama's reference implementation runs vision model on CPU

1348 pixel_values = image_inputs['data']

1349 aspect_ratio_ids = image_inputs['aspect_ratio_ids']

@@ -1367,10 +1367,10 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1367 self,

1368 input_ids: torch.Tensor,

1369 attn_metadata: AttentionMetadata,

1370 - num_tiles: List[List[int]],

1370 + num_tiles: list[list[int]],

1371 num_tokens_per_tile: int,

1372 dtype: torch.dtype,

1373 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

1374 token_ids = input_ids.tolist()

1375 start = 0

1376 batch_token_ids = []

@@ -1422,7 +1422,7 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1422 input_ids: torch.Tensor,

1423 positions: torch.Tensor,

1424 **kwargs: object,

1425 - ) -> Union[Tuple, CausalLMOutputWithPast]:

1425 + ) -> Union[CausalLMOutputWithPast]:

1426 attn_metadata = get_forward_context().attn_metadata

1427 if attn_metadata.num_prefill_tokens > 0 and \

1428 attn_metadata.num_decode_tokens > 0:

@@ -1476,8 +1476,8 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1476

1477 return outputs

1478

1479 - def load_weights(self, weights: Iterable[Tuple[str,

1480 - torch.Tensor]]) -> Set[str]:

1479 + def load_weights(self, weights: Iterable[tuple[str,

1480 + torch.Tensor]]) -> set[str]:

1481 stacked_params_mapping = [

1482 # (param_name, shard_name, shard_id)

1483 (".qkv_proj", ".q_proj", "q"),

@@ -1487,7 +1487,7 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

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

1488 ]

1489 params_dict = dict(self.named_parameters())

1490 - updated_params: Set[str] = set()

1490 + updated_params: set[str] = set()

1491 for name, loaded_weight in weights:

1492 if 'patch_embedding.weight' in name:

1493 name = name.replace('patch_embedding.weight',

@@ -1538,7 +1538,7 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1538 tower_model="vision_model")

1539

1540

1541 -def skip_attention_mask(sparse_mask: List[List[int]]) -> bool:

1541 +def skip_attention_mask(sparse_mask: list[list[int]]) -> bool:

1542 for mask in sparse_mask:

1543 # Skip text-only samples.

1544 if len(mask) == 0:

@@ -1556,10 +1556,10 @@ def skip_attention_mask(sparse_mask: List[List[int]]) -> bool:

1556

1557

1558 def convert_sparse_cross_attention_mask_to_dense(

1559 - sparse_mask: List[List[List[int]]],

1560 - num_tiles: List[List[int]],

1561 - lengths: List[int],

1562 -) -> Tuple[np.ndarray, List[Tuple[int, int]]]:

1559 + sparse_mask: list[list[list[int]]],

1560 + num_tiles: list[list[int]],

1561 + lengths: list[int],

1562 +) -> tuple[np.ndarray, list[tuple[int, int]]]:

1563 total_length = sum(lengths)

1564 total_tiles = sum([sum(tiles) for tiles in num_tiles])

1565 dense_mask = np.zeros(shape=(total_length, total_tiles), dtype=np.int64)

vllm/model_executor/models/mllama4.py

+9 -9

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

18 import math

19 from collections.abc import Iterable, Mapping

20 from itertools import tee

21 -from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union

21 +from typing import Literal, Optional, TypedDict, Union

22

23 import torch

24 from torch import nn

@@ -582,7 +582,7 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]

582 mm_items: MultiModalDataItems,

583 hf_processor_mm_kwargs: Mapping[str, object],

584 out_mm_kwargs: MultiModalKwargs,

585 - ) -> List[PromptUpdate]:

585 + ) -> list[PromptUpdate]:

586 assert (

587 mm_items.get_count("image", strict=False) == 0

588 or "aspect_ratios" in out_mm_kwargs

@@ -778,26 +778,26 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

778

779 def separate_weights(

780 self,

781 - weights: Iterable[Tuple[str, torch.Tensor]],

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

782 prefix: str,

783 - ) -> Tuple[Iterable[Tuple[str, torch.Tensor]], Iterable[Tuple[

783 + ) -> tuple[Iterable[tuple[str, torch.Tensor]], Iterable[tuple[

784 str, torch.Tensor]]]:

785 weights1, weights2 = tee(weights, 2)

786

787 - def get_prefix_weights() -> Iterable[Tuple[str, torch.Tensor]]:

787 + def get_prefix_weights() -> Iterable[tuple[str, torch.Tensor]]:

788 for name, data in weights1:

789 if name.startswith(prefix):

790 yield (name, data)

791

792 - def get_other_weights() -> Iterable[Tuple[str, torch.Tensor]]:

792 + def get_other_weights() -> Iterable[tuple[str, torch.Tensor]]:

793 for name, data in weights2:

794 if not name.startswith(prefix):

795 yield (name, data)

796

797 return get_prefix_weights(), get_other_weights()

798

799 - def load_weights(self, weights: Iterable[Tuple[str,

800 - torch.Tensor]]) -> Set[str]:

799 + def load_weights(self, weights: Iterable[tuple[str,

800 + torch.Tensor]]) -> set[str]:

801

802 stacked_params_mapping = [

803 # (param_name, shard_name, shard_id)

@@ -806,7 +806,7 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

806 (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),

807 ]

808 params_dict = dict(self.named_parameters())

809 - updated_params: Set[str] = set()

809 + updated_params: set[str] = set()

810

811 # language_model is an Llama4ForCausalLM instance. We load it's

812 # using llama4's load_weights routine.

vllm/model_executor/models/mlp_speculator.py

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 import math

4 -from typing import Iterable, List, Set, Tuple

4 +from collections.abc import Iterable

5

6 import torch

7 import torch.nn as nn

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

148 previous_hidden_states: torch.Tensor,

149 num_predict_tokens: int,

150 sampling_metadata: SamplingMetadata,

151 - ) -> List[SamplerOutput]:

151 + ) -> list[SamplerOutput]:

152 if num_predict_tokens > self.max_speculative_tokens:

153 raise ValueError(f"Max speculative tokens for model is "

154 f"{self.max_speculative_tokens}, but "

@@ -190,10 +190,10 @@ class MLPSpeculator(nn.Module):

190

191 return next_tokens

192

193 - def load_weights(self, weights: Iterable[Tuple[str,

194 - torch.Tensor]]) -> Set[str]:

193 + def load_weights(self, weights: Iterable[tuple[str,

194 + torch.Tensor]]) -> set[str]:

195 params_dict = dict(self.named_parameters())

196 - loaded_params: Set[str] = set()

196 + loaded_params: set[str] = set()

197 for name, loaded_weight in weights:

198 name = name.replace("speculator.", "")

199 param = params_dict.get(name)

vllm/model_executor/models/modernbert.py

+6 -5

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Iterable, Optional, Set, Tuple

2 +from collections.abc import Iterable

3 +from typing import Optional

4

5 import torch

6 from torch import nn

@@ -212,11 +213,11 @@ class ModernBertModel(nn.Module):

213 eps=config.norm_eps,

214 bias=config.norm_bias)

215

215 - def load_weights(self, weights: Iterable[Tuple[str,

216 - torch.Tensor]]) -> Set[str]:

216 + def load_weights(self, weights: Iterable[tuple[str,

217 + torch.Tensor]]) -> set[str]:

218 weights = self.hf_to_vllm_mapper.apply(weights)

219 params_dict = dict(self.named_parameters())

219 - loaded_params: Set[str] = set()

220 + loaded_params: set[str] = set()

221 for name, loaded_weight in weights:

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

223 continue

@@ -280,7 +281,7 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):

281 self._pooler = CrossEncodingPooler(config, self.classifier,

282 ModernBertPooler(config))

283

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

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

285

286 self_weights = []

287

vllm/model_executor/models/module_mapping.py

+9 -9

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

4 # https://github.com/modelscope/ms-swift/blob/v2.4.2/swift/utils/module_mapping.py

5

6 from dataclasses import dataclass, field

7 -from typing import List, Union

7 +from typing import Union

8

9

10 @dataclass

@@ -46,17 +46,17 @@ class ModelKeys:

46

47 @dataclass

48 class MultiModelKeys(ModelKeys):

49 - language_model: List[str] = field(default_factory=list)

50 - connector: List[str] = field(default_factory=list)

49 + language_model: list[str] = field(default_factory=list)

50 + connector: list[str] = field(default_factory=list)

51 # vision tower and audio tower

52 - tower_model: List[str] = field(default_factory=list)

53 - generator: List[str] = field(default_factory=list)

52 + tower_model: list[str] = field(default_factory=list)

53 + generator: list[str] = field(default_factory=list)

54

55 @staticmethod

56 - def from_string_field(language_model: Union[str, List[str]] = None,

57 - connector: Union[str, List[str]] = None,

58 - tower_model: Union[str, List[str]] = None,

59 - generator: Union[str, List[str]] = None,

56 + def from_string_field(language_model: Union[str, list[str]] = None,

57 + connector: Union[str, list[str]] = None,

58 + tower_model: Union[str, list[str]] = None,

59 + generator: Union[str, list[str]] = None,

60 **kwargs) -> 'MultiModelKeys':

61

62 def to_list(value):

vllm/model_executor/models/molmo.py

+16 -16

@@ -4,7 +4,7 @@ import math

4 from collections.abc import Iterable, Mapping, Sequence

5 from dataclasses import dataclass

6 from functools import cached_property, partial

7 -from typing import List, Optional, Set, Tuple, TypedDict, Union

7 +from typing import Optional, TypedDict, Union

8

9 import numpy as np

10 import torch

@@ -90,7 +90,7 @@ class MolmoImageInputs(TypedDict):

90

91 @dataclass

92 class VisionBackboneConfig:

93 - image_default_input_size: Tuple[int, int] = (336, 336)

93 + image_default_input_size: tuple[int, int] = (336, 336)

94 image_patch_size: int = 14

95 image_pos_patch_size: int = 14

96 image_emb_dim: int = 1024

@@ -267,7 +267,7 @@ class BlockCollection(nn.Module):

267 for _ in range(config.image_num_layers)

268 ])

269

270 - def forward(self, x: torch.Tensor) -> List[torch.Tensor]:

270 + def forward(self, x: torch.Tensor) -> list[torch.Tensor]:

271 hidden_states = []

272 for r in self.resblocks:

273 x = r(x)

@@ -334,7 +334,7 @@ class VisionTransformer(nn.Module):

334

335 def forward(self,

336 x: torch.Tensor,

337 - patch_num: Optional[int] = None) -> List[torch.Tensor]:

337 + patch_num: Optional[int] = None) -> list[torch.Tensor]:

338 """

339 : param x: (batch_size, num_patch, n_pixels)

340 """

@@ -434,7 +434,7 @@ class MolmoAttention(nn.Module):

434 )

435

436 def _apply_qk_norm(self, q: torch.Tensor,

437 - k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

437 + k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

438 if self.tp_size > 1:

439 q = tensor_model_parallel_all_gather(q.contiguous())

440 k = tensor_model_parallel_all_gather(k.contiguous())

@@ -570,7 +570,7 @@ class MolmoDecoderLayer(nn.Module):

570 positions: torch.Tensor,

571 hidden_states: torch.Tensor,

572 residual: Optional[torch.Tensor],

573 - ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:

573 + ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:

574 # Self Attention

575 if residual is None:

576 residual = hidden_states

@@ -596,7 +596,7 @@ class MolmoDecoderNormAfterLayer(MolmoDecoderLayer):

596 positions: torch.Tensor,

597 hidden_states: torch.Tensor,

598 residual: Optional[torch.Tensor],

599 - ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:

599 + ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:

600 # Self Attention

601 residual = hidden_states

602 hidden_states = self.self_attn(

@@ -740,15 +740,15 @@ class MolmoVisionBackbone(nn.Module, SupportsQuant):

740 # image_features: (batch_size, num_image, num_patch, d_model)

741 return image_features

742

743 - def load_weights(self, weights: Iterable[Tuple[str,

744 - torch.Tensor]]) -> Set[str]:

743 + def load_weights(self, weights: Iterable[tuple[str,

744 + torch.Tensor]]) -> set[str]:

745 stacked_params_mapping = [

746 # (param_name, shard_name, shard_id)

747 ("merged_linear", "gate_proj", 0),

748 ("merged_linear", "up_proj", 1),

749 ]

750 params_dict = dict(self.named_parameters())

751 - loaded_params: Set[str] = set()

751 + loaded_params: set[str] = set()

752

753 for name, loaded_weight in weights:

754 for (param_name, weight_name, shard_id) in stacked_params_mapping:

@@ -855,10 +855,10 @@ class MolmoModel(nn.Module, SupportsQuant):

855 hidden_states = self.norm(hidden_states)

856 return hidden_states

857

858 - def load_weights(self, weights: Iterable[Tuple[str,

859 - torch.Tensor]]) -> Set[str]:

858 + def load_weights(self, weights: Iterable[tuple[str,

859 + torch.Tensor]]) -> set[str]:

860 params_dict = dict(self.named_parameters())

861 - loaded_params: Set[str] = set()

861 + loaded_params: set[str] = set()

862

863 for name, loaded_weight in weights:

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

@@ -1530,7 +1530,7 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1530 sampling_metadata)

1531 return logits

1532

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

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

1534

1535 loader = AutoWeightsLoader(self)

1536 weights = _get_weights_with_merged_embedding(weights)

@@ -1548,8 +1548,8 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1548

1549

1550 def _get_weights_with_merged_embedding(

1551 - weights: Iterable[Tuple[str, torch.Tensor]]

1552 -) -> Iterable[Tuple[str, torch.Tensor]]:

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

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

1553 embedding_weights = {}

1554 for name, weight in weights:

1555 if "wte.embedding" in name:

vllm/model_executor/models/moonvit.py

+4 -3

@@ -42,9 +42,10 @@

42 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

43 # SOFTWARE.

44 import math

45 +from collections.abc import Sequence

46 from copy import deepcopy

47 from functools import cached_property

47 -from typing import List, Optional, Sequence, Tuple, Union

48 +from typing import Optional, Union

49

50 import torch

51 import torch.nn as nn

@@ -222,7 +223,7 @@ class MoonVisionPatchEmbed(nn.Module):

223 self,

224 out_dim: int,

225 in_dim: int = 3,

225 - patch_size: Union[int, Tuple[int, int]] = (14, 14),

226 + patch_size: Union[int, tuple[int, int]] = (14, 14),

227 pos_emb_height: int = 14,

228 pos_emb_width: int = 14,

229 ):

@@ -526,7 +527,7 @@ def patch_merger(

527 x: torch.Tensor,

528 grid_hw: torch.Tensor,

529 merge_kernel_size: list[int, int] = (2, 2),

529 -) -> List[torch.Tensor]:

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

531 d_model = x.size(-1)

532

533 outputs = []

vllm/model_executor/models/mpt.py

+7 -6

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

2

3 # Adapted from https://huggingface.co/mosaicml/mpt-7b/tree/main

4 import math

5 -from typing import Iterable, Optional, Set, Tuple, Union

5 +from collections.abc import Iterable

6 +from typing import Optional, Union

7

8 import torch

9 import torch.nn as nn

@@ -265,10 +266,10 @@ class MPTModel(nn.Module):

266 hidden_states = self.norm_f(hidden_states)

267 return hidden_states

268

268 - def load_weights(self, weights: Iterable[Tuple[str,

269 - torch.Tensor]]) -> Set[str]:

269 + def load_weights(self, weights: Iterable[tuple[str,

270 + torch.Tensor]]) -> set[str]:

271 params_dict = dict(self.named_parameters(remove_duplicate=False))

271 - loaded_params: Set[str] = set()

272 + loaded_params: set[str] = set()

273 for name, loaded_weight in weights:

274 # Skip loading extra bias for GPTQ models.

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

@@ -323,7 +324,7 @@ class MPTForCausalLM(nn.Module, SupportsPP):

324 sampling_metadata)

325 return logits

326

326 - def load_weights(self, weights: Iterable[Tuple[str,

327 - torch.Tensor]]) -> Set[str]:

327 + def load_weights(self, weights: Iterable[tuple[str,

328 + torch.Tensor]]) -> set[str]:

329 loader = AutoWeightsLoader(self)

330 return loader.load_weights(weights)

vllm/model_executor/models/nemotron.py

+8 -7

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only Nemotron model compatible with HuggingFace weights."""

25 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Any, Optional, Union

27

28 import torch

29 from torch import nn

@@ -69,7 +70,7 @@ def _cast_if_autocast_enabled(*args):

70 class NemotronLayerNorm1P(nn.LayerNorm):

71

72 def __init__(self,

72 - normalized_shape: Union[int, List[int], torch.Size],

73 + normalized_shape: Union[int, list[int], torch.Size],

74 eps: float = 1e-5,

75 elementwise_affine: bool = True,

76 bias: bool = True,

@@ -133,7 +134,7 @@ class NemotronAttention(nn.Module):

134 num_heads: int,

135 num_kv_heads: int,

136 rope_theta: float = 10000,

136 - rope_scaling: Optional[Dict[str, Any]] = None,

137 + rope_scaling: Optional[dict[str, Any]] = None,

138 max_position_embeddings: int = 8192,

139 quant_config: Optional[QuantizationConfig] = None,

140 bias: bool = False,

@@ -267,7 +268,7 @@ class NemotronDecoderLayer(nn.Module):

268 positions: torch.Tensor,

269 hidden_states: torch.Tensor,

270 residual: Optional[torch.Tensor],

270 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

272 # Self Attention

273 if residual is None:

274 residual = hidden_states

@@ -441,8 +442,8 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

442 sampling_metadata)

443 return logits

444

444 - def load_weights(self, weights: Iterable[Tuple[str,

445 - torch.Tensor]]) -> Set[str]:

445 + def load_weights(self, weights: Iterable[tuple[str,

446 + torch.Tensor]]) -> set[str]:

447 stacked_params_mapping = [

448 # (param_name, shard_name, shard_id)

449 (".qkv_proj", ".q_proj", "q"),

@@ -450,7 +451,7 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

451 (".qkv_proj", ".v_proj", "v"),

452 ]

453 params_dict = dict(self.named_parameters())

453 - loaded_params: Set[str] = set()

454 + loaded_params: set[str] = set()

455 for name, loaded_weight in weights:

456 if "rotary_emb.inv_freq" in name:

457 continue

vllm/model_executor/models/nemotron_nas.py

+9 -8

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only deci model compatible with HuggingFace weights."""

25 -from typing import Iterable, Optional, Set, Tuple, Type, Union

25 +from collections.abc import Iterable

26 +from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -135,7 +136,7 @@ class DeciLMDecoderLayer(nn.Module):

136 positions: torch.Tensor,

137 hidden_states: torch.Tensor,

138 residual: Optional[torch.Tensor],

138 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

140 # Self Attention

141

142 if self._is_no_op_attention:

@@ -168,7 +169,7 @@ class DeciModel(nn.Module):

169 *,

170 vllm_config: VllmConfig,

171 prefix: str = "",

171 - layer_type: Type[DeciLMDecoderLayer] = DeciLMDecoderLayer,

172 + layer_type: type[DeciLMDecoderLayer] = DeciLMDecoderLayer,

173 ):

174 super().__init__()

175

@@ -260,8 +261,8 @@ class DeciModel(nn.Module):

261 hidden_states, _ = self.norm(hidden_states, residual)

262 return hidden_states

263

263 - def load_weights(self, weights: Iterable[Tuple[str,

264 - torch.Tensor]]) -> Set[str]:

264 + def load_weights(self, weights: Iterable[tuple[str,

265 + torch.Tensor]]) -> set[str]:

266 stacked_params_mapping = [

267 # (param_name, shard_name, shard_id)

268 (".qkv_proj", ".q_proj", "q"),

@@ -271,7 +272,7 @@ class DeciModel(nn.Module):

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

273 ]

274 params_dict = dict(self.named_parameters())

274 - loaded_params: Set[str] = set()

275 + loaded_params: set[str] = set()

276 for name, loaded_weight in weights:

277 if "rotary_emb.inv_freq" in name:

278 continue

@@ -428,8 +429,8 @@ class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):

429 sampling_metadata)

430 return logits

431

431 - def load_weights(self, weights: Iterable[Tuple[str,

432 - torch.Tensor]]) -> Set[str]:

432 + def load_weights(self, weights: Iterable[tuple[str,

433 + torch.Tensor]]) -> set[str]:

434 loader = AutoWeightsLoader(

435 self,

436 skip_prefixes=(["lm_head."]

vllm/model_executor/models/olmo.py

+6 -5

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only OLMo model compatible with HuggingFace weights."""

25 -from typing import Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -209,7 +210,7 @@ class OlmoDecoderLayer(nn.Module):

210 self,

211 positions: torch.Tensor,

212 hidden_states: torch.Tensor,

212 - ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:

213 + ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:

214 # Attention block.

215 residual = hidden_states

216 hidden_states = self.input_layernorm(hidden_states)

@@ -338,8 +339,8 @@ class OlmoForCausalLM(nn.Module, SupportsPP):

339 sampling_metadata)

340 return logits

341

341 - def load_weights(self, weights: Iterable[Tuple[str,

342 - torch.Tensor]]) -> Set[str]:

342 + def load_weights(self, weights: Iterable[tuple[str,

343 + torch.Tensor]]) -> set[str]:

344 stacked_params_mapping = [

345 # (param_name, shard_name, shard_id)

346 ("qkv_proj", "q_proj", "q"),

@@ -349,7 +350,7 @@ class OlmoForCausalLM(nn.Module, SupportsPP):

350 ("gate_up_proj", "up_proj", 1),

351 ]

352 params_dict = dict(self.named_parameters(remove_duplicate=False))

352 - loaded_params: Set[str] = set()

353 + loaded_params: set[str] = set()

354 for name, loaded_weight in weights:

355 if "rotary_emb.inv_freq" in name:

356 continue

vllm/model_executor/models/olmo2.py

+4 -3

@@ -23,8 +23,9 @@

23 # limitations under the License.

24 """Inference-only OLMo2 model compatible with HuggingFace weights."""

25

26 +from collections.abc import Iterable

27 from functools import partial

27 -from typing import Iterable, Optional, Tuple, Union

28 +from typing import Optional, Union

29

30 import torch

31 from torch import nn

@@ -135,7 +136,7 @@ class Olmo2Attention(nn.Module):

136 )

137

138 def _apply_qk_norm(self, q: torch.Tensor,

138 - k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

139 + k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

140 if self.tp_size > 1:

141 q = tensor_model_parallel_all_gather(q.contiguous())

142 k = tensor_model_parallel_all_gather(k.contiguous())

@@ -365,7 +366,7 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP):

366 sampling_metadata)

367 return logits

368

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

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

370 stacked_params_mapping = [

371 # (param_name, shard_name, shard_id)

372 ("qkv_proj", "q_proj", "q"),

vllm/model_executor/models/olmoe.py

+8 -7

@@ -12,7 +12,8 @@

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

13 # limitations under the License.

14 """Inference-only OLMoE model compatible with HuggingFace weights."""

15 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

15 +from collections.abc import Iterable

16 +from typing import Any, Optional, Union

17

18 import torch

19 from torch import nn

@@ -102,7 +103,7 @@ class OlmoeAttention(nn.Module):

103 num_heads: int,

104 num_kv_heads: int,

105 rope_theta: float = 10000,

105 - rope_scaling: Optional[Dict[str, Any]] = None,

106 + rope_scaling: Optional[dict[str, Any]] = None,

107 max_position_embeddings: int = 4096,

108 cache_config: Optional[CacheConfig] = None,

109 quant_config: Optional[QuantizationConfig] = None,

@@ -307,8 +308,8 @@ class OlmoeModel(nn.Module):

308 hidden_states, _ = self.norm(hidden_states, residual)

309 return hidden_states

310

310 - def load_weights(self, weights: Iterable[Tuple[str,

311 - torch.Tensor]]) -> Set[str]:

311 + def load_weights(self, weights: Iterable[tuple[str,

312 + torch.Tensor]]) -> set[str]:

313 stacked_params_mapping = [

314 # (param_name, shard_name, shard_id)

315 ("qkv_proj", "q_proj", "q"),

@@ -327,7 +328,7 @@ class OlmoeModel(nn.Module):

328 num_experts=self.config.num_experts)

329

330 params_dict = dict(self.named_parameters())

330 - loaded_params: Set[str] = set()

331 + loaded_params: set[str] = set()

332 for name, loaded_weight in weights:

333 for (param_name, weight_name, shard_id) in stacked_params_mapping:

334 # Skip non-stacked layers and experts (experts handled below).

@@ -439,8 +440,8 @@ class OlmoeForCausalLM(nn.Module, SupportsPP):

440 sampling_metadata)

441 return logits

442

442 - def load_weights(self, weights: Iterable[Tuple[str,

443 - torch.Tensor]]) -> Set[str]:

443 + def load_weights(self, weights: Iterable[tuple[str,

444 + torch.Tensor]]) -> set[str]:

445 loader = AutoWeightsLoader(

446 self,

447 skip_prefixes=["rotary_emb.inv_freq"],

vllm/model_executor/models/opt.py

+7 -6

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

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

19 # limitations under the License.

20 """Inference-only OPT model compatible with HuggingFace weights."""

21 -from typing import Iterable, Optional, Set, Tuple, Union

21 +from collections.abc import Iterable

22 +from typing import Optional, Union

23

24 import torch

25 from torch import nn

@@ -312,8 +313,8 @@ class OPTModel(nn.Module):

313 intermediate_tensors,

314 inputs_embeds=inputs_embeds)

315

315 - def load_weights(self, weights: Iterable[Tuple[str,

316 - torch.Tensor]]) -> Set[str]:

316 + def load_weights(self, weights: Iterable[tuple[str,

317 + torch.Tensor]]) -> set[str]:

318 stacked_params_mapping = [

319 # (param_name, shard_name, shard_id)

320 ("qkv_proj", "q_proj", "q"),

@@ -321,7 +322,7 @@ class OPTModel(nn.Module):

322 ("qkv_proj", "v_proj", "v"),

323 ]

324 params_dict = dict(self.named_parameters(remove_duplicate=False))

324 - loaded_params: Set[str] = set()

325 + loaded_params: set[str] = set()

326 for name, loaded_weight in weights:

327 for (param_name, weight_name, shard_id) in stacked_params_mapping:

328 if weight_name not in name:

@@ -400,8 +401,8 @@ class OPTForCausalLM(nn.Module, SupportsPP):

401 sampling_metadata)

402 return logits

403

403 - def load_weights(self, weights: Iterable[Tuple[str,

404 - torch.Tensor]]) -> Set[str]:

404 + def load_weights(self, weights: Iterable[tuple[str,

405 + torch.Tensor]]) -> set[str]:

406 loader = AutoWeightsLoader(

407 self,

408 skip_prefixes=(["lm_head.weight"]

vllm/model_executor/models/orion.py

+9 -8

@@ -5,7 +5,8 @@

5 # Copyright (c) OrionStar Inc.

6 # LICENSE: https://huggingface.co/OrionStarAI/Orion-14B-Base/blob/main/LICENSE

7 """Inference-only Orion-14B model compatible with HuggingFace weights."""

8 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

8 +from collections.abc import Iterable

9 +from typing import Any, Optional, Union

10

11 import torch

12 from torch import nn

@@ -72,7 +73,7 @@ class OrionAttention(nn.Module):

73 num_heads: int,

74 num_kv_heads: int,

75 rope_theta: float = 10000,

75 - rope_scaling: Optional[Dict[str, Any]] = None,

76 + rope_scaling: Optional[dict[str, Any]] = None,

77 max_position_embeddings: int = 8192,

78 cache_config: Optional[CacheConfig] = None,

79 quant_config: Optional[QuantizationConfig] = None,

@@ -186,7 +187,7 @@ class OrionDecoderLayer(nn.Module):

187 self,

188 positions: torch.Tensor,

189 hidden_states: torch.Tensor,

189 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

191 # Self Attention

192 residual = hidden_states

193 hidden_states = self.input_layernorm(hidden_states)

@@ -259,8 +260,8 @@ class OrionModel(nn.Module):

260 hidden_states = self.norm(hidden_states)

261 return hidden_states

262

262 - def load_weights(self, weights: Iterable[Tuple[str,

263 - torch.Tensor]]) -> Set[str]:

263 + def load_weights(self, weights: Iterable[tuple[str,

264 + torch.Tensor]]) -> set[str]:

265 stacked_params_mapping = [

266 # (param_name, shard_name, shard_id)

267 ("qkv_proj", "q_proj", "q"),

@@ -270,7 +271,7 @@ class OrionModel(nn.Module):

271 ("gate_up_proj", "up_proj", 1),

272 ]

273 params_dict = dict(self.named_parameters())

273 - loaded_params: Set[str] = set()

274 + loaded_params: set[str] = set()

275 for name, loaded_weight in weights:

276 for (param_name, weight_name, shard_id) in stacked_params_mapping:

277 if weight_name not in name:

@@ -341,8 +342,8 @@ class OrionForCausalLM(nn.Module, SupportsPP):

342 sampling_metadata)

343 return logits

344

344 - def load_weights(self, weights: Iterable[Tuple[str,

345 - torch.Tensor]]) -> Set[str]:

345 + def load_weights(self, weights: Iterable[tuple[str,

346 + torch.Tensor]]) -> set[str]:

347 loader = AutoWeightsLoader(

348 self,

349 skip_prefixes=([

vllm/model_executor/models/ovis.py

+5 -5

@@ -17,8 +17,8 @@

17 # limitations under the License.

18 """ PyTorch Ovis model."""

19 import math

20 -from typing import (Iterable, List, Literal, Mapping, Optional, Set, Tuple,

21 - TypedDict, Union)

20 +from collections.abc import Iterable, Mapping

21 +from typing import Literal, Optional, TypedDict, Union

22

23 import torch

24 import torch.nn as nn

@@ -211,7 +211,7 @@ class OvisImagePatchInputs(TypedDict):

211 `(batch_size * (num_patches + 1))`

212 """

213

214 - patches_per_image: List[int]

214 + patches_per_image: list[int]

215 """

216 List of number of total patches for each image in the batch.

217 This is used to restore the first two dimensions of `flat_data`.

@@ -545,8 +545,8 @@ class Ovis(nn.Module, SupportsMultiModal):

545 logits = self.llm.compute_logits(hidden_states, sampling_metadata)

546 return logits

547

548 - def load_weights(self, weights: Iterable[Tuple[str,

549 - torch.Tensor]]) -> Set[str]:

548 + def load_weights(self, weights: Iterable[tuple[str,

549 + torch.Tensor]]) -> set[str]:

550 loader = AutoWeightsLoader(self)

551 return loader.load_weights(weights)

552

vllm/model_executor/models/paligemma.py

+3 -3

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 from collections.abc import Iterable, Mapping, Sequence

3 -from typing import Literal, Optional, Set, Tuple, TypedDict, Union

3 +from typing import Literal, Optional, TypedDict, Union

4

5 import torch

6 from torch import nn

@@ -391,7 +391,7 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,

391 return self.language_model.compute_logits(hidden_states,

392 sampling_metadata)

393

394 - def load_weights(self, weights: Iterable[Tuple[str,

395 - torch.Tensor]]) -> Set[str]:

394 + def load_weights(self, weights: Iterable[tuple[str,

395 + torch.Tensor]]) -> set[str]:

396 loader = AutoWeightsLoader(self)

397 return loader.load_weights(weights)

vllm/model_executor/models/persimmon.py

+7 -6

@@ -21,7 +21,8 @@

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

22 # limitations under the License.

23 """Inference-only persimmon model compatible with HuggingFace weights."""

24 -from typing import Iterable, Optional, Set, Tuple, Union

24 +from collections.abc import Iterable

25 +from typing import Optional, Union

26

27 import torch

28 from torch import nn

@@ -260,10 +261,10 @@ class PersimmonModel(nn.Module):

261 hidden_states = self.final_layernorm(hidden_states)

262 return hidden_states

263

263 - def load_weights(self, weights: Iterable[Tuple[str,

264 - torch.Tensor]]) -> Set[str]:

264 + def load_weights(self, weights: Iterable[tuple[str,

265 + torch.Tensor]]) -> set[str]:

266 params_dict = dict(self.named_parameters(remove_duplicate=False))

266 - loaded_params: Set[str] = set()

267 + loaded_params: set[str] = set()

268 for name, loaded_weight in weights:

269 if is_pp_missing_parameter(name, self):

270 continue

@@ -336,7 +337,7 @@ class PersimmonForCausalLM(nn.Module, SupportsPP):

337 sampling_metadata)

338 return logits

339

339 - def load_weights(self, weights: Iterable[Tuple[str,

340 - torch.Tensor]]) -> Set[str]:

340 + def load_weights(self, weights: Iterable[tuple[str,

341 + torch.Tensor]]) -> set[str]:

342 loader = AutoWeightsLoader(self)

343 return loader.load_weights(weights)

vllm/model_executor/models/phi.py

+7 -6

@@ -36,7 +36,8 @@

36 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE

37 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

38 """Inference-only Phi-1.5 model compatible with HuggingFace weights."""

39 -from typing import Iterable, Optional, Set, Tuple, Union

39 +from collections.abc import Iterable

40 +from typing import Optional, Union

41

42 import torch

43 from torch import nn

@@ -248,8 +249,8 @@ class PhiModel(nn.Module):

249

250 return hidden_states

251

251 - def load_weights(self, weights: Iterable[Tuple[str,

252 - torch.Tensor]]) -> Set[str]:

252 + def load_weights(self, weights: Iterable[tuple[str,

253 + torch.Tensor]]) -> set[str]:

254 stacked_params_mapping = [

255 # (param_name, shard_name, shard_id)

256 ("qkv_proj", "q_proj", "q"),

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

258 ("qkv_proj", "v_proj", "v")

259 ]

260 params_dict = dict(self.named_parameters())

260 - loaded_params: Set[str] = set()

261 + loaded_params: set[str] = set()

262

263 for name, loaded_weight in weights:

264 if "rotary_emb.inv_freq" in name:

@@ -348,7 +349,7 @@ class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

349 sampling_metadata, self.lm_head.bias)

350 return logits

351

351 - def load_weights(self, weights: Iterable[Tuple[str,

352 - torch.Tensor]]) -> Set[str]:

352 + def load_weights(self, weights: Iterable[tuple[str,

353 + torch.Tensor]]) -> set[str]:

354 loader = AutoWeightsLoader(self)

355 return loader.load_weights(weights)

vllm/model_executor/models/phi3_small.py

+9 -8

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 import math

4 -from typing import Iterable, Optional, Set, Tuple, Union

4 +from collections.abc import Iterable

5 +from typing import Optional, Union

6

7 import torch

8 from torch import nn

@@ -230,8 +231,8 @@ class Phi3SmallSelfAttention(nn.Module):

231 self,

232 positions: torch.Tensor,

233 hidden_states: torch.Tensor,

233 - ) -> Tuple[torch.Tensor, Optional[torch.Tensor],

234 - Optional[Tuple[torch.Tensor]]]:

234 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],

235 + Optional[tuple[torch.Tensor]]]:

236 qkv, _ = self.query_key_value(hidden_states)

237

238 qkv = qkv.view(qkv.shape[:-1] +

@@ -352,10 +353,10 @@ class Phi3SmallModel(nn.Module):

353 hidden_states = self.final_layernorm(hidden_states)

354 return hidden_states

355

355 - def load_weights(self, weights: Iterable[Tuple[str,

356 - torch.Tensor]]) -> Set[str]:

356 + def load_weights(self, weights: Iterable[tuple[str,

357 + torch.Tensor]]) -> set[str]:

358 params_dict = dict(self.named_parameters())

358 - loaded_params: Set[str] = set()

359 + loaded_params: set[str] = set()

360 for name, loaded_weight in weights:

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

362 continue

@@ -454,8 +455,8 @@ class Phi3SmallForCausalLM(nn.Module, SupportsPP):

455 output_hidden_states = output_hidden_states

456 return output_hidden_states

457

457 - def load_weights(self, weights: Iterable[Tuple[str,

458 - torch.Tensor]]) -> Set[str]:

458 + def load_weights(self, weights: Iterable[tuple[str,

459 + torch.Tensor]]) -> set[str]:

460 loader = AutoWeightsLoader(

461 self,

462 skip_prefixes=(["lm_head.weight"]

vllm/model_executor/models/phi3v.py

+7 -7

@@ -16,7 +16,7 @@

16 # limitations under the License.

17 import re

18 from collections.abc import Iterable, Mapping, Sequence

19 -from typing import Any, List, Literal, Optional, Set, Tuple, TypedDict, Union

19 +from typing import Any, Literal, Optional, TypedDict, Union

20

21 import torch

22 import torch.nn as nn

@@ -94,7 +94,7 @@ def _init_img_processor(hf_config: PretrainedConfig,

94

95 class Phi3VImagePixelInputs(TypedDict):

96 type: Literal["pixel_values"]

97 - data: Union[torch.Tensor, List[torch.Tensor]]

97 + data: Union[torch.Tensor, list[torch.Tensor]]

98 """

99 Shape:

100 `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`

@@ -113,7 +113,7 @@ class Phi3VImagePixelInputs(TypedDict):

113

114 class Phi3VImageEmbeddingInputs(TypedDict):

115 type: Literal["image_embeds"]

116 - data: Union[torch.Tensor, List[torch.Tensor]]

116 + data: Union[torch.Tensor, list[torch.Tensor]]

117 """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`

118

119 `hidden_size` must match the hidden size of language model backbone.

@@ -571,8 +571,8 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

571 return data

572

573 def _validate_pixel_values(

574 - self, data: Union[torch.Tensor, List[torch.Tensor]]

575 - ) -> Union[torch.Tensor, List[torch.Tensor]]:

574 + self, data: Union[torch.Tensor, list[torch.Tensor]]

575 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

576

577 h = w = CLIP_VIT_LARGE_PATCH14_336_CONFIG.image_size

578 expected_dims = (3, h, w)

@@ -707,8 +707,8 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

707 return self.language_model.compute_logits(hidden_states,

708 sampling_metadata)

709

710 - def load_weights(self, weights: Iterable[Tuple[str,

711 - torch.Tensor]]) -> Set[str]:

710 + def load_weights(self, weights: Iterable[tuple[str,

711 + torch.Tensor]]) -> set[str]:

712

713 loader = AutoWeightsLoader(self)

714 autoloaded_weights = loader.load_weights(weights,

vllm/model_executor/models/phi4mm.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2 import math

3 from collections.abc import Iterable, Mapping, Sequence

4 -from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, Union

4 +from typing import Any, Literal, Optional, TypedDict, Union

5

6 import numpy as np

7 import torch

@@ -392,7 +392,7 @@ class Phi4MMImageEncoder(nn.Module):

392

393 class Phi4MMImagePixelInputs(TypedDict):

394 type: Literal["pixel_values"]

395 - data: Union[torch.Tensor, List[torch.Tensor]]

395 + data: Union[torch.Tensor, list[torch.Tensor]]

396 """

397 Shape:

398 `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`

@@ -417,7 +417,7 @@ class Phi4MMImagePixelInputs(TypedDict):

417

418 class Phi4MMImageEmbeddingInputs(TypedDict):

419 type: Literal["image_embeds"]

420 - data: Union[torch.Tensor, List[torch.Tensor]]

420 + data: Union[torch.Tensor, list[torch.Tensor]]

421 """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`

422

423 `hidden_size` must match the hidden size of language model backbone.

@@ -426,7 +426,7 @@ class Phi4MMImageEmbeddingInputs(TypedDict):

426

427 class Phi4MMAudioFeatureInputs(TypedDict):

428 type: Literal["audio_features"]

429 - data: Union[torch.Tensor, List[torch.Tensor]]

429 + data: Union[torch.Tensor, list[torch.Tensor]]

430 """Shape: `(batch_size * num_audios, 80, M)"""

431

432

@@ -1031,7 +1031,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1031 return audio_embeds

1032

1033 def _parse_and_validate_image_input(self,

1034 - **kwargs: object) -> Optional[Dict]:

1034 + **kwargs: object) -> Optional[dict]:

1035 input_image_embeds: NestedTensors = kwargs.get("input_image_embeds")

1036 if input_image_embeds is None:

1037 return None

@@ -1238,7 +1238,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1238 sampling_metadata)

1239 return logits

1240

1241 - def load_weights(self, weights: Iterable[Tuple[str,

1241 + def load_weights(self, weights: Iterable[tuple[str,

1242 torch.Tensor]]) -> None:

1243 weights = ((name, data) for name, data in weights

1244 if "lora" not in name)

vllm/model_executor/models/phi4mm_audio.py

+2 -2

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

6 #!/usr/bin/env python3

7 import abc

8 import math

9 -from typing import List, Literal, Optional

9 +from typing import Literal, Optional

10

11 import numpy as np

12 import torch

@@ -746,7 +746,7 @@ class ConformerEncoder(TransformerEncoderBase):

746 attention_group_size = attenion_heads = Multi-Query Attention

747 """

748

749 - extra_multi_layer_output_idxs: List[int]

749 + extra_multi_layer_output_idxs: list[int]

750

751 def __init__( # pylint: disable-all

752 self,

vllm/model_executor/models/phi4mm_utils.py

+2 -2

@@ -5,7 +5,7 @@

5 # but implemented by the Phi-Speech team

6 #!/usr/bin/env python3

7 import math

8 -from typing import Optional, Tuple, Union

8 +from typing import Optional, Union

9

10 import torch

11 import torch.nn.functional as F

@@ -1586,7 +1586,7 @@ class AttModule(nn.Module):

1586 memory: Optional[Tensor] = None,

1587 pos_emb: Optional[Tensor] = None,

1588 att_mask: Optional[Tensor] = None,

1589 - ) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:

1589 + ) -> tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:

1590 """AttModule forward

1591

1592 Args:

vllm/model_executor/models/phimoe.py

+7 -6

@@ -22,7 +22,8 @@

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

23 # limitations under the License.

24 """Inference-only PhiMoE model."""

25 -from typing import Iterable, Optional, Set, Tuple, Union

25 +from collections.abc import Iterable

26 +from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -505,8 +506,8 @@ class PhiMoEModel(nn.Module):

506 hidden_states = self.norm(hidden_states)

507 return hidden_states

508

508 - def load_weights(self, weights: Iterable[Tuple[str,

509 - torch.Tensor]]) -> Set[str]:

509 + def load_weights(self, weights: Iterable[tuple[str,

510 + torch.Tensor]]) -> set[str]:

511 stacked_params_mapping = [

512 # (param_name, shard_name, shard_id)

513 ("qkv_proj", "q_proj", "q"),

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

522 num_experts=self.config.num_local_experts)

523

524 params_dict = dict(self.named_parameters())

524 - loaded_params: Set[str] = set()

525 + loaded_params: set[str] = set()

526 for name, loaded_weight in weights:

527 if (self.quant_config is not None and

528 (scale_name := self.quant_config.get_cache_scale(name))):

@@ -657,8 +658,8 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

658 sampling_metadata)

659 return logits

660

660 - def load_weights(self, weights: Iterable[Tuple[str,

661 - torch.Tensor]]) -> Set[str]:

661 + def load_weights(self, weights: Iterable[tuple[str,

662 + torch.Tensor]]) -> set[str]:

663 loader = AutoWeightsLoader(

664 self,

665 skip_prefixes=(["rotary_emb.inv_freq"]),

vllm/model_executor/models/pixtral.py

+13 -13

@@ -4,7 +4,7 @@ import math

4 from collections.abc import Iterable, Mapping, Sequence

5 from dataclasses import dataclass, fields

6 from functools import cached_property

7 -from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union

7 +from typing import Literal, Optional, TypedDict, Union

8

9 import torch

10 import torch.nn as nn

@@ -438,18 +438,18 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

438 return self.language_model.compute_logits(hidden_states,

439 sampling_metadata)

440

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

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

442

443 - def is_vision_encoder_weights(weight: Tuple[str, torch.Tensor]):

443 + def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]):

444 return weight[0].startswith("vision_encoder")

445

446 - def is_vision_lang_adapter_weights(weight: Tuple[str, torch.Tensor]):

446 + def is_vision_lang_adapter_weights(weight: tuple[str, torch.Tensor]):

447 return weight[0].startswith("vision_language_adapter")

448

449 - def is_patch_merger(weight: Tuple[str, torch.Tensor]):

449 + def is_patch_merger(weight: tuple[str, torch.Tensor]):

450 return weight[0].startswith("patch_merger")

451

452 - def is_pre_mm_projector_norm(weight: Tuple[str, torch.Tensor]):

452 + def is_pre_mm_projector_norm(weight: tuple[str, torch.Tensor]):

453 return weight[0].startswith("pre_mm_projector_norm")

454

455 # Get references to parameters for direct loading

@@ -566,7 +566,7 @@ def apply_rotary_emb_vit(

566 xq: torch.Tensor,

567 xk: torch.Tensor,

568 freqs_cis: torch.Tensor,

569 -) -> Tuple[torch.Tensor, torch.Tensor]:

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

570 xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))

571 xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

572 assert freqs_cis.dtype == torch.complex64

@@ -671,7 +671,7 @@ class Transformer(nn.Module):

671 return x

672

673

674 -def position_meshgrid(patch_embeds_list: List[torch.Tensor], ) -> torch.Tensor:

674 +def position_meshgrid(patch_embeds_list: list[torch.Tensor], ) -> torch.Tensor:

675 positions = torch.cat([

676 torch.stack(

677 torch.meshgrid(

@@ -733,7 +733,7 @@ class VisionTransformer(nn.Module):

733

734 def forward(

735 self,

736 - images: List[torch.Tensor],

736 + images: list[torch.Tensor],

737 ) -> torch.Tensor:

738 """

739 Args:

@@ -1023,7 +1023,7 @@ class PixtralHFAttention(nn.Module):

1023 hidden_states: torch.Tensor,

1024 attention_mask: torch.Tensor,

1025 position_embeddings: torch.Tensor,

1026 - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:

1026 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

1027 batch, patches, _ = hidden_states.size()

1028

1029 qkv_states, _ = self.qkv_proj(hidden_states)

@@ -1249,8 +1249,8 @@ class PixtralHFVisionModel(nn.Module):

1249

1250 # (TODO) Add prefix argument for filtering out weights to be loaded

1251 # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986

1252 - def load_weights(self, weights: Iterable[Tuple[str,

1253 - torch.Tensor]]) -> Set[str]:

1252 + def load_weights(self, weights: Iterable[tuple[str,

1253 + torch.Tensor]]) -> set[str]:

1254 stacked_params_mapping = [

1255 # (param_name, shard_name, shard_id)

1256 (".qkv_proj", ".q_proj", "q"),

@@ -1260,7 +1260,7 @@ class PixtralHFVisionModel(nn.Module):

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

1261 ]

1262 params_dict = dict(self.named_parameters())

1263 - loaded_params: Set[str] = set()

1263 + loaded_params: set[str] = set()

1264 layer_count = len(self.transformer.layers)

1265

1266 for name, loaded_weight in weights:

vllm/model_executor/models/plamo2.py

+4 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Inference-only PLaMo2 model."""

3 import math

4 -from typing import Iterable, Optional, Tuple

4 +from collections.abc import Iterable

5 +from typing import Optional

6

7 import torch

8 from torch import nn

@@ -659,7 +660,7 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,

660 return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

661

662 def _get_mamba_cache_shape(

662 - self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

663 + self) -> tuple[tuple[int, int], tuple[int, int]]:

664 world_size = get_tensor_model_parallel_world_size()

665 hidden_size = (self.config.mamba_num_heads *

666 self.config.hidden_size_per_head)

@@ -682,7 +683,7 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,

683 sampling_metadata)

684 return logits

685

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

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

687 params_dict = dict(self.named_parameters())

688 for name, loaded_weight in weights:

689

vllm/model_executor/models/prithvi_geospatial_mae.py

+4 -4

@@ -16,7 +16,7 @@

16 # limitations under the License.

17 """Inference-only IBM/NASA Prithvi Geospatial model."""

18 from collections.abc import Iterable, Mapping, Sequence

19 -from typing import Optional, Set, Tuple, Union

19 +from typing import Optional, Union

20

21 import torch

22 import torch.nn as nn

@@ -154,7 +154,7 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,

154 "by PrithviGeospatialMAE.")

155

156 def _parse_and_validate_multimodal_data(

157 - self, **kwargs) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:

157 + self, **kwargs) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

158

159 pixel_values = kwargs.pop("pixel_values", None)

160 if not isinstance(pixel_values, torch.Tensor):

@@ -195,8 +195,8 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,

195 ) -> Optional[PoolerOutput]:

196 return PoolerOutput([PoolingSequenceGroupOutput(hidden_states)])

197

198 - def load_weights(self, weights: Iterable[Tuple[str,

199 - torch.Tensor]]) -> Set[str]:

198 + def load_weights(self, weights: Iterable[tuple[str,

199 + torch.Tensor]]) -> set[str]:

200 params_list = []

201 model_buffers = dict(self.named_buffers())

202 loaded_buffers = []

vllm/model_executor/models/qwen.py

+7 -6

@@ -6,7 +6,8 @@

6 # LICENSE: https://huggingface.co/Qwen/Qwen-7B/blob/main/LICENSE

7 """Inference-only QWen model compatible with HuggingFace weights."""

8 import json

9 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

9 +from collections.abc import Iterable

10 +from typing import Any, Optional, Union

11

12 import torch

13 from torch import nn

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

77 num_heads: int,

78 max_position_embeddings: int,

79 rope_theta: float = 10000,

79 - rope_scaling: Optional[Dict[str, Any]] = None,

80 + rope_scaling: Optional[dict[str, Any]] = None,

81 cache_config: Optional[CacheConfig] = None,

82 quant_config: Optional[QuantizationConfig] = None,

83 prefix: str = "",

@@ -166,7 +167,7 @@ class QWenBlock(nn.Module):

167 positions: torch.Tensor,

168 hidden_states: torch.Tensor,

169 residual: Optional[torch.Tensor],

169 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

171 # Self Attention

172 if residual is None:

173 residual = hidden_states

@@ -284,15 +285,15 @@ class QWenBaseModel(nn.Module):

285 sampling_metadata)

286 return logits

287

287 - def load_weights(self, weights: Iterable[Tuple[str,

288 - torch.Tensor]]) -> Set[str]:

288 + def load_weights(self, weights: Iterable[tuple[str,

289 + torch.Tensor]]) -> set[str]:

290 stacked_params_mapping = [

291 # (param_name, shard_name, shard_id)

292 ("gate_up_proj", "w2", 0),

293 ("gate_up_proj", "w1", 1),

294 ]

295 params_dict = dict(self.named_parameters())

295 - loaded_params: Set[str] = set()

296 + loaded_params: set[str] = set()

297 for name, loaded_weight in weights:

298 if "rotary_emb.inv_freq" in name:

299 continue

vllm/model_executor/models/qwen2.py

+10 -9

@@ -23,7 +23,8 @@

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

24 # limitations under the License.

25 """Inference-only Qwen2 model compatible with HuggingFace weights."""

26 -from typing import Any, Iterable, Optional, Set, Tuple, Union

26 +from collections.abc import Iterable

27 +from typing import Any, Optional, Union

28

29 import torch

30 from torch import nn

@@ -108,7 +109,7 @@ class Qwen2Attention(nn.Module):

109 rope_theta: float = 10000,

110 cache_config: Optional[CacheConfig] = None,

111 quant_config: Optional[QuantizationConfig] = None,

111 - rope_scaling: Optional[Tuple] = None,

112 + rope_scaling: Optional[tuple] = None,

113 prefix: str = "",

114 attn_type: str = AttentionType.DECODER,

115 dual_chunk_attention_config: Optional[dict[str, Any]] = None,

@@ -245,7 +246,7 @@ class Qwen2DecoderLayer(nn.Module):

246 positions: torch.Tensor,

247 hidden_states: torch.Tensor,

248 residual: Optional[torch.Tensor],

248 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

250 # Self Attention

251 if residual is None:

252 residual = hidden_states

@@ -367,8 +368,8 @@ class Qwen2Model(nn.Module):

368 hidden_states, _ = self.norm(hidden_states, residual)

369 return hidden_states

370

370 - def load_weights(self, weights: Iterable[Tuple[str,

371 - torch.Tensor]]) -> Set[str]:

371 + def load_weights(self, weights: Iterable[tuple[str,

372 + torch.Tensor]]) -> set[str]:

373 stacked_params_mapping = [

374 # (param_name, shard_name, shard_id)

375 ("qkv_proj", "q_proj", "q"),

@@ -378,7 +379,7 @@ class Qwen2Model(nn.Module):

379 ("gate_up_proj", "up_proj", 1),

380 ]

381 params_dict = dict(self.named_parameters(remove_duplicate=False))

381 - loaded_params: Set[str] = set()

382 + loaded_params: set[str] = set()

383 for name, loaded_weight in weights:

384 if "rotary_emb.inv_freq" in name:

385 continue

@@ -490,8 +491,8 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

491 sampling_metadata)

492 return logits

493

493 - def load_weights(self, weights: Iterable[Tuple[str,

494 - torch.Tensor]]) -> Set[str]:

494 + def load_weights(self, weights: Iterable[tuple[str,

495 + torch.Tensor]]) -> set[str]:

496 loader = AutoWeightsLoader(

497 self,

498 skip_prefixes=(["lm_head."]

@@ -559,7 +560,7 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):

560 ) -> Optional[PoolerOutput]:

561 return self._pooler(hidden_states, pooling_metadata)

562

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

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

564 weights = self.hf_to_vllm_mapper.apply(weights)

565 weights = ((name, data) for name, data in weights

566 if not name.startswith("lm_head."))

vllm/model_executor/models/qwen2_5_omni_thinker.py

+10 -10

@@ -21,10 +21,10 @@

21 # limitations under the License.

22 """Inference-only Qwen2.5-Omni model (thinker part)."""

23

24 +from collections.abc import Iterable, Mapping, Sequence

25 from copy import copy

26 from functools import partial

26 -from typing import (Any, Dict, Iterable, List, Mapping, Optional, Sequence,

27 - Set, Tuple, Union)

27 +from typing import Any, Optional, Union

28

29 import torch

30 import torch.nn as nn

@@ -138,7 +138,7 @@ class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo,

138 min_pixels: Optional[int] = None,

139 max_pixels: Optional[int] = None,

140 size: Optional[dict[str, int]] = None,

141 - fps: Optional[Union[float, List[float]]] = None,

141 + fps: Optional[Union[float, list[float]]] = None,

142 **kwargs: object,

143 ) -> Qwen2_5OmniProcessor:

144 if fps is not None:

@@ -550,7 +550,7 @@ class Qwen2_5OmniConditionalGenerationMixin:

550

551 def _parse_and_validate_image_input(

552 self,

553 - **kwargs: Dict[str, Any],

553 + **kwargs: dict[str, Any],

554 ) -> Optional[Qwen2_5_VLImageInputs]:

555 pixel_values = kwargs.pop("pixel_values", None)

556 image_embeds = kwargs.pop("image_embeds", None)

@@ -589,7 +589,7 @@ class Qwen2_5OmniConditionalGenerationMixin:

589

590 def _parse_and_validate_video_input(

591 self,

592 - **kwargs: Dict[str, Any],

592 + **kwargs: dict[str, Any],

593 ) -> Optional[Qwen2_5_VLVideoInputs]:

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

595 video_embeds = kwargs.pop("video_embeds", None)

@@ -627,7 +627,7 @@ class Qwen2_5OmniConditionalGenerationMixin:

627 def _process_audio_input(

628 self,

629 audio_input: Qwen2AudioInputs,

630 - audio_hashes: List[str] = None,

630 + audio_hashes: list[str] = None,

631 cached_audio_features: torch.Tensor = None,

632 ) -> torch.Tensor:

633

@@ -676,7 +676,7 @@ class Qwen2_5OmniConditionalGenerationMixin:

676 def _process_video_input(

677 self,

678 video_input: Qwen2_5_VLVideoInputs,

679 - video_hashes: List[str] = None,

679 + video_hashes: list[str] = None,

680 cached_video_embeds: torch.Tensor = None) -> torch.Tensor:

681 if video_input["type"] == "video_embeds":

682 return video_input["video_embeds"].type(self.visual.dtype)

@@ -825,7 +825,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

825 if audio_input is None and image_input is None and video_input is None:

826 return None

827

828 - multimodal_embeddings: List[Tuple[NestedTensors, str]] = []

828 + multimodal_embeddings: list[tuple[NestedTensors, str]] = []

829

830 if audio_input is not None:

831 audio_embeds = self._process_audio_input(audio_input)

@@ -891,8 +891,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

891 return self.language_model.compute_logits(hidden_states,

892 sampling_metadata)

893

894 - def load_weights(self, weights: Iterable[Tuple[str,

895 - torch.Tensor]]) -> Set[str]:

894 + def load_weights(self, weights: Iterable[tuple[str,

895 + torch.Tensor]]) -> set[str]:

896 loader = AutoWeightsLoader(

897 self,

898 skip_prefixes=["talker.", "token2wav."],

vllm/model_executor/models/qwen2_5_vl.py

+10 -10

@@ -24,9 +24,9 @@

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

25 # limitations under the License.

26 """Inference-only Qwen2.5-VL model compatible with HuggingFace weights."""

27 +from collections.abc import Iterable, Mapping

28 from functools import partial

28 -from typing import (Callable, Iterable, List, Literal, Mapping, Optional, Set,

29 - Tuple, TypedDict, Union)

29 +from typing import Callable, Literal, Optional, TypedDict, Union

30

31 import torch

32 import torch.nn as nn

@@ -91,7 +91,7 @@ class Qwen2_5_VLImageEmbeddingInputs(TypedDict):

91 type: Literal["image_embeds"]

92 image_embeds: torch.Tensor

93 """Supported types:

94 - - List[`torch.Tensor`]: A list of tensors holding all images' features.

94 + - list[`torch.Tensor`]: A list of tensors holding all images' features.

95 Each tensor holds an image's features.

96 - `torch.Tensor`: A tensor holding all images' features

97 (concatenation of all images' feature tensors).

@@ -137,7 +137,7 @@ class Qwen2_5_VLVideoEmbeddingInputs(TypedDict):

137 type: Literal["video_embeds"]

138 video_embeds: torch.Tensor

139 """Supported types:

140 - - List[`torch.Tensor`]: A list of tensors holding all videos' features.

140 + - list[`torch.Tensor`]: A list of tensors holding all videos' features.

141 Each tensor holds an video's features.

142 - `torch.Tensor`: A tensor holding all videos' features

143 (concatenation of all videos' feature tensors).

@@ -709,8 +709,8 @@ class Qwen2_5_VisionTransformer(nn.Module):

709 hidden_states = hidden_states[reverse_indices, :]

710 return hidden_states

711

712 - def load_weights(self, weights: Iterable[Tuple[str,

713 - torch.Tensor]]) -> Set[str]:

712 + def load_weights(self, weights: Iterable[tuple[str,

713 + torch.Tensor]]) -> set[str]:

714 stacked_params_mapping = [

715 # (param_name, shard_name, shard_id)

716 ("attn.qkv.", "attn.q.", "q"),

@@ -718,7 +718,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

718 ("attn.qkv.", "attn.v.", "v"),

719 ]

720 params_dict = dict(self.named_parameters(remove_duplicate=False))

721 - loaded_params: Set[str] = set()

721 + loaded_params: set[str] = set()

722

723 for name, loaded_weight in weights:

724 for (param_name, weight_name, shard_id) in stacked_params_mapping:

@@ -750,7 +750,7 @@ class Qwen2_5_VLProcessingInfo(Qwen2VLProcessingInfo):

750 min_pixels: Optional[int] = None,

751 max_pixels: Optional[int] = None,

752 size: Optional[dict[str, int]] = None,

753 - fps: Optional[Union[float, List[float]]] = None,

753 + fps: Optional[Union[float, list[float]]] = None,

754 **kwargs: object,

755 ) -> Qwen2_5_VLProcessor:

756 if fps is not None:

@@ -1116,8 +1116,8 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1116 return self.language_model.compute_logits(hidden_states,

1117 sampling_metadata)

1118

1119 - def load_weights(self, weights: Iterable[Tuple[str,

1120 - torch.Tensor]]) -> Set[str]:

1119 + def load_weights(self, weights: Iterable[tuple[str,

1120 + torch.Tensor]]) -> set[str]:

1121

1122 loader = AutoWeightsLoader(self)

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

vllm/model_executor/models/qwen2_audio.py

+3 -3

@@ -22,7 +22,7 @@

22 # limitations under the License.

23 """Inference-only Qwen2-Audio model compatible with HuggingFace weights."""

24 from collections.abc import Iterable, Mapping, Sequence

25 -from typing import Any, Optional, Set, Tuple, TypedDict, Union

25 +from typing import Any, Optional, TypedDict, Union

26

27 import torch

28 import torch.nn as nn

@@ -403,7 +403,7 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,

403 return self.language_model.compute_logits(hidden_states,

404 sampling_metadata)

405

406 - def load_weights(self, weights: Iterable[Tuple[str,

407 - torch.Tensor]]) -> Set[str]:

406 + def load_weights(self, weights: Iterable[tuple[str,

407 + torch.Tensor]]) -> set[str]:

408 loader = AutoWeightsLoader(self)

409 return loader.load_weights(weights)

vllm/model_executor/models/qwen2_moe.py

+9 -8

@@ -23,7 +23,8 @@

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

24 # limitations under the License.

25 """Inference-only Qwen2MoE model compatible with HuggingFace weights."""

26 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

26 +from collections.abc import Iterable

27 +from typing import Any, Optional, Union

28

29 import torch

30 import torch.nn.functional as F

@@ -169,12 +170,12 @@ class Qwen2MoeAttention(nn.Module):

170 num_heads: int,

171 num_kv_heads: int,

172 rope_theta: float = 10000,

172 - rope_scaling: Optional[Dict[str, Any]] = None,

173 + rope_scaling: Optional[dict[str, Any]] = None,

174 max_position_embeddings: int = 8192,

175 cache_config: Optional[CacheConfig] = None,

176 quant_config: Optional[QuantizationConfig] = None,

177 prefix: str = "",

177 - dual_chunk_attention_config: Optional[Dict[str, Any]] = None,

178 + dual_chunk_attention_config: Optional[dict[str, Any]] = None,

179 ) -> None:

180 super().__init__()

181 self.hidden_size = hidden_size

@@ -389,8 +390,8 @@ class Qwen2MoeModel(nn.Module):

390 hidden_states, _ = self.norm(hidden_states, residual)

391 return hidden_states

392

392 - def load_weights(self, weights: Iterable[Tuple[str,

393 - torch.Tensor]]) -> Set[str]:

393 + def load_weights(self, weights: Iterable[tuple[str,

394 + torch.Tensor]]) -> set[str]:

395 stacked_params_mapping = [

396 # (param_name, shard_name, shard_id)

397 ("qkv_proj", "q_proj", "q"),

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

410 num_experts=self.config.num_experts)

411

412 params_dict = dict(self.named_parameters())

412 - loaded_params: Set[str] = set()

413 + loaded_params: set[str] = set()

414 for name, loaded_weight in weights:

415 for (param_name, weight_name, shard_id) in stacked_params_mapping:

416 # Skip non-stacked layers and experts (experts handled below).

@@ -532,8 +533,8 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP):

533 sampling_metadata)

534 return logits

535

535 - def load_weights(self, weights: Iterable[Tuple[str,

536 - torch.Tensor]]) -> Set[str]:

536 + def load_weights(self, weights: Iterable[tuple[str,

537 + torch.Tensor]]) -> set[str]:

538 loader = AutoWeightsLoader(

539 self,

540 skip_prefixes=(["rotary_emb.inv_freq"]),

vllm/model_executor/models/qwen2_rm.py

+4 -3

@@ -5,7 +5,8 @@

7 """Inference-only Qwen2-RM model compatible with HuggingFace weights."""

8 -from typing import Iterable, Optional, Set, Tuple, Union

8 +from collections.abc import Iterable

9 +from typing import Optional, Union

10

11 import torch

12 from torch import nn

@@ -95,8 +96,8 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP,

96 ) -> Optional[PoolerOutput]:

97 return self._pooler(hidden_states, pooling_metadata)

98

98 - def load_weights(self, weights: Iterable[Tuple[str,

99 - torch.Tensor]]) -> Set[str]:

99 + def load_weights(self, weights: Iterable[tuple[str,

100 + torch.Tensor]]) -> set[str]:

101 loader = AutoWeightsLoader(self,

102 ignore_unexpected_prefixes=["lm_head."])

103 return loader.load_weights(weights)

vllm/model_executor/models/qwen2_vl.py

+8 -9

@@ -25,8 +25,7 @@

25 """Inference-only Qwen2-VL model compatible with HuggingFace weights."""

26 from collections.abc import Iterable, Mapping, Sequence

27 from functools import partial

28 -from typing import (Any, Callable, Literal, Optional, Set, Tuple, TypedDict,

29 - Union)

28 +from typing import Any, Callable, Literal, Optional, TypedDict, Union

29

30 import torch

31 import torch.nn as nn

@@ -102,7 +101,7 @@ class Qwen2VLImageEmbeddingInputs(TypedDict):

101 type: Literal["image_embeds"]

102 image_embeds: torch.Tensor

103 """Supported types:

105 - - List[`torch.Tensor`]: A list of tensors holding all images' features.

104 + - list[`torch.Tensor`]: A list of tensors holding all images' features.

105 Each tensor holds an image's features.

106 - `torch.Tensor`: A tensor holding all images' features

107 (concatenation of all images' feature tensors).

@@ -142,7 +141,7 @@ class Qwen2VLVideoEmbeddingInputs(TypedDict):

141 type: Literal["video_embeds"]

142 video_embeds: torch.Tensor

143 """Supported types:

145 - - List[`torch.Tensor`]: A list of tensors holding all videos' features.

144 + - list[`torch.Tensor`]: A list of tensors holding all videos' features.

145 Each tensor holds an video's features.

146 - `torch.Tensor`: A tensor holding all videos' features

147 (concatenation of all videos' feature tensors).

@@ -662,8 +661,8 @@ class Qwen2VisionTransformer(nn.Module):

661

662 return x

663

665 - def load_weights(self, weights: Iterable[Tuple[str,

666 - torch.Tensor]]) -> Set[str]:

664 + def load_weights(self, weights: Iterable[tuple[str,

665 + torch.Tensor]]) -> set[str]:

666 stacked_params_mapping = [

667 # (param_name, shard_name, shard_id)

668 ("qkv_proj", "q_proj", "q"),

@@ -671,7 +670,7 @@ class Qwen2VisionTransformer(nn.Module):

670 ("qkv_proj", "v_proj", "v"),

671 ]

672 params_dict = dict(self.named_parameters(remove_duplicate=False))

674 - loaded_params: Set[str] = set()

673 + loaded_params: set[str] = set()

674

675 for name, loaded_weight in weights:

676 for (param_name, weight_name, shard_id) in stacked_params_mapping:

@@ -1394,8 +1393,8 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1393 return self.language_model.compute_logits(hidden_states,

1394 sampling_metadata)

1395

1397 - def load_weights(self, weights: Iterable[Tuple[str,

1398 - torch.Tensor]]) -> Set[str]:

1396 + def load_weights(self, weights: Iterable[tuple[str,

1397 + torch.Tensor]]) -> set[str]:

1398

1399 loader = AutoWeightsLoader(self)

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

vllm/model_executor/models/qwen3.py

+6 -5

@@ -21,7 +21,8 @@

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

22 # limitations under the License.

23 """Inference-only Qwen3 model compatible with HuggingFace weights."""

24 -from typing import Iterable, Optional, Set, Tuple, Union

24 +from collections.abc import Iterable

25 +from typing import Optional, Union

26

27 import torch

28 from torch import nn

@@ -63,7 +64,7 @@ class Qwen3Attention(nn.Module):

64 rope_theta: float = 10000,

65 cache_config: Optional[CacheConfig] = None,

66 quant_config: Optional[QuantizationConfig] = None,

66 - rope_scaling: Optional[Tuple] = None,

67 + rope_scaling: Optional[tuple] = None,

68 prefix: str = "",

69 attn_type: str = AttentionType.DECODER) -> None:

70 super().__init__()

@@ -201,7 +202,7 @@ class Qwen3DecoderLayer(nn.Module):

202 positions: torch.Tensor,

203 hidden_states: torch.Tensor,

204 residual: Optional[torch.Tensor],

204 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

206 # Self Attention

207 if residual is None:

208 residual = hidden_states

@@ -309,8 +310,8 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

310 sampling_metadata)

311 return logits

312

312 - def load_weights(self, weights: Iterable[Tuple[str,

313 - torch.Tensor]]) -> Set[str]:

313 + def load_weights(self, weights: Iterable[tuple[str,

314 + torch.Tensor]]) -> set[str]:

315 loader = AutoWeightsLoader(

316 self,

317 skip_prefixes=(["lm_head."]

vllm/model_executor/models/qwen3_moe.py

+8 -7

@@ -21,7 +21,8 @@

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

22 # limitations under the License.

23 """Inference-only Qwen3MoE model compatible with HuggingFace weights."""

24 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

24 +from collections.abc import Iterable

25 +from typing import Any, Optional, Union

26

27 import torch

28 from torch import nn

@@ -149,7 +150,7 @@ class Qwen3MoeAttention(nn.Module):

150 num_heads: int,

151 num_kv_heads: int,

152 rope_theta: float = 10000,

152 - rope_scaling: Optional[Dict[str, Any]] = None,

153 + rope_scaling: Optional[dict[str, Any]] = None,

154 max_position_embeddings: int = 8192,

155 head_dim: Optional[int] = None,

156 rms_norm_eps: float = 1e-06,

@@ -373,8 +374,8 @@ class Qwen3MoeModel(nn.Module):

374 hidden_states, _ = self.norm(hidden_states, residual)

375 return hidden_states

376

376 - def load_weights(self, weights: Iterable[Tuple[str,

377 - torch.Tensor]]) -> Set[str]:

377 + def load_weights(self, weights: Iterable[tuple[str,

378 + torch.Tensor]]) -> set[str]:

379 stacked_params_mapping = [

380 # (param_name, shard_name, shard_id)

381 ("qkv_proj", "q_proj", "q"),

@@ -393,7 +394,7 @@ class Qwen3MoeModel(nn.Module):

394 num_experts=self.config.num_experts)

395

396 params_dict = dict(self.named_parameters())

396 - loaded_params: Set[str] = set()

397 + loaded_params: set[str] = set()

398 for name, loaded_weight in weights:

399 for (param_name, weight_name, shard_id) in stacked_params_mapping:

400 # Skip non-stacked layers and experts (experts handled below).

@@ -527,8 +528,8 @@ class Qwen3MoeForCausalLM(nn.Module, SupportsPP):

528 sampling_metadata)

529 return logits

530

530 - def load_weights(self, weights: Iterable[Tuple[str,

531 - torch.Tensor]]) -> Set[str]:

531 + def load_weights(self, weights: Iterable[tuple[str,

532 + torch.Tensor]]) -> set[str]:

533 loader = AutoWeightsLoader(

534 self,

535 skip_prefixes=(["rotary_emb.inv_freq"]),

vllm/model_executor/models/qwen_vl.py

+4 -5

@@ -9,10 +9,9 @@ import copy

9 import math

10 import re

11 import unicodedata

12 -from collections.abc import Collection, Mapping, Sequence

13 -from collections.abc import Set as AbstractSet

12 +from collections.abc import Collection, Mapping, Sequence, Set

13 from functools import lru_cache, partial

15 -from typing import Callable, List, Literal, Optional, TypedDict, Union

14 +from typing import Callable, Literal, Optional, TypedDict, Union

15

16 import torch

17 from torch import nn

@@ -395,7 +394,7 @@ def _get_tokenizer_without_image_pad(

394 def tokenize(

395 self,

396 text: str,

398 - allowed_special: Union[AbstractSet[str], str] = "all",

397 + allowed_special: Union[Set[str], str] = "all",

398 disallowed_special: Union[Collection[str], str] = (),

399 **kwargs,

400 ) -> list[Union[bytes, str]]:

@@ -411,7 +410,7 @@ def _get_tokenizer_without_image_pad(

410

411 def _decode(

412 self,

414 - token_ids: Union[int, List[int]],

413 + token_ids: Union[int, list[int]],

414 skip_special_tokens: bool = False,

415 errors: Optional[str] = None,

416 **kwargs,

vllm/model_executor/models/registry.py

+31 -31

@@ -10,10 +10,10 @@ import subprocess

10 import sys

11 import tempfile

12 from abc import ABC, abstractmethod

13 +from collections.abc import Set

14 from dataclasses import dataclass, field

15 from functools import lru_cache

15 -from typing import (AbstractSet, Callable, Dict, List, Optional, Tuple, Type,

16 - TypeVar, Union)

16 +from typing import Callable, Optional, TypeVar, Union

17

18 import cloudpickle

19 import torch.nn as nn

@@ -266,7 +266,7 @@ class _ModelInfo:

266 supports_v0_only: bool

267

268 @staticmethod

269 - def from_model_cls(model: Type[nn.Module]) -> "_ModelInfo":

269 + def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":

270 return _ModelInfo(

271 architecture=model.__name__,

272 is_text_generation_model=is_text_generation_model(model),

@@ -290,7 +290,7 @@ class _BaseRegisteredModel(ABC):

290 raise NotImplementedError

291

292 @abstractmethod

293 - def load_model_cls(self) -> Type[nn.Module]:

293 + def load_model_cls(self) -> type[nn.Module]:

294 raise NotImplementedError

295

296

@@ -301,10 +301,10 @@ class _RegisteredModel(_BaseRegisteredModel):

301 """

302

303 interfaces: _ModelInfo

304 - model_cls: Type[nn.Module]

304 + model_cls: type[nn.Module]

305

306 @staticmethod

307 - def from_model_cls(model_cls: Type[nn.Module]):

307 + def from_model_cls(model_cls: type[nn.Module]):

308 return _RegisteredModel(

309 interfaces=_ModelInfo.from_model_cls(model_cls),

310 model_cls=model_cls,

@@ -313,7 +313,7 @@ class _RegisteredModel(_BaseRegisteredModel):

313 def inspect_model_cls(self) -> _ModelInfo:

314 return self.interfaces

315

316 - def load_model_cls(self) -> Type[nn.Module]:

316 + def load_model_cls(self) -> type[nn.Module]:

317 return self.model_cls

318

319

@@ -330,7 +330,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel):

330 return _run_in_subprocess(

331 lambda: _ModelInfo.from_model_cls(self.load_model_cls()))

332

333 - def load_model_cls(self) -> Type[nn.Module]:

333 + def load_model_cls(self) -> type[nn.Module]:

334 mod = importlib.import_module(self.module_name)

335 return getattr(mod, self.class_name)

336

@@ -339,7 +339,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel):

339 def _try_load_model_cls(

340 model_arch: str,

341 model: _BaseRegisteredModel,

342 -) -> Optional[Type[nn.Module]]:

342 +) -> Optional[type[nn.Module]]:

343 from vllm.platforms import current_platform

344 current_platform.verify_model_arch(model_arch)

345 try:

@@ -366,15 +366,15 @@ def _try_inspect_model_cls(

366 @dataclass

367 class _ModelRegistry:

368 # Keyed by model_arch

369 - models: Dict[str, _BaseRegisteredModel] = field(default_factory=dict)

369 + models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)

370

371 - def get_supported_archs(self) -> AbstractSet[str]:

371 + def get_supported_archs(self) -> Set[str]:

372 return self.models.keys()

373

374 def register_model(

375 self,

376 model_arch: str,

377 - model_cls: Union[Type[nn.Module], str],

377 + model_cls: Union[type[nn.Module], str],

378 ) -> None:

379 """

380 Register an external model to be used in vLLM.

@@ -413,7 +413,7 @@ class _ModelRegistry:

413

414 self.models[model_arch] = model

415

416 - def _raise_for_unsupported(self, architectures: List[str]):

416 + def _raise_for_unsupported(self, architectures: list[str]):

417 all_supported_archs = self.get_supported_archs()

418

419 if any(arch in all_supported_archs for arch in architectures):

@@ -426,7 +426,7 @@ class _ModelRegistry:

426 f"Supported architectures: {all_supported_archs}")

427

428 def _try_load_model_cls(self,

429 - model_arch: str) -> Optional[Type[nn.Module]]:

429 + model_arch: str) -> Optional[type[nn.Module]]:

430 if model_arch not in self.models:

431 return None

432

@@ -440,8 +440,8 @@ class _ModelRegistry:

440

441 def _normalize_archs(

442 self,

443 - architectures: Union[str, List[str]],

444 - ) -> List[str]:

443 + architectures: Union[str, list[str]],

444 + ) -> list[str]:

445 if isinstance(architectures, str):

446 architectures = [architectures]

447 if not architectures:

@@ -458,8 +458,8 @@ class _ModelRegistry:

458

459 def inspect_model_cls(

460 self,

461 - architectures: Union[str, List[str]],

462 - ) -> Tuple[_ModelInfo, str]:

461 + architectures: Union[str, list[str]],

462 + ) -> tuple[_ModelInfo, str]:

463 architectures = self._normalize_archs(architectures)

464

465 for arch in architectures:

@@ -471,8 +471,8 @@ class _ModelRegistry:

471

472 def resolve_model_cls(

473 self,

474 - architectures: Union[str, List[str]],

475 - ) -> Tuple[Type[nn.Module], str]:

474 + architectures: Union[str, list[str]],

475 + ) -> tuple[type[nn.Module], str]:

476 architectures = self._normalize_archs(architectures)

477

478 for arch in architectures:

@@ -484,77 +484,77 @@ class _ModelRegistry:

484

485 def is_text_generation_model(

486 self,

487 - architectures: Union[str, List[str]],

487 + architectures: Union[str, list[str]],

488 ) -> bool:

489 model_cls, _ = self.inspect_model_cls(architectures)

490 return model_cls.is_text_generation_model

491

492 def is_pooling_model(

493 self,

494 - architectures: Union[str, List[str]],

494 + architectures: Union[str, list[str]],

495 ) -> bool:

496 model_cls, _ = self.inspect_model_cls(architectures)

497 return model_cls.is_pooling_model

498

499 def is_cross_encoder_model(

500 self,

501 - architectures: Union[str, List[str]],

501 + architectures: Union[str, list[str]],

502 ) -> bool:

503 model_cls, _ = self.inspect_model_cls(architectures)

504 return model_cls.supports_cross_encoding

505

506 def is_multimodal_model(

507 self,

508 - architectures: Union[str, List[str]],

508 + architectures: Union[str, list[str]],

509 ) -> bool:

510 model_cls, _ = self.inspect_model_cls(architectures)

511 return model_cls.supports_multimodal

512

513 def is_pp_supported_model(

514 self,

515 - architectures: Union[str, List[str]],

515 + architectures: Union[str, list[str]],

516 ) -> bool:

517 model_cls, _ = self.inspect_model_cls(architectures)

518 return model_cls.supports_pp

519

520 def model_has_inner_state(

521 self,

522 - architectures: Union[str, List[str]],

522 + architectures: Union[str, list[str]],

523 ) -> bool:

524 model_cls, _ = self.inspect_model_cls(architectures)

525 return model_cls.has_inner_state

526

527 def is_attention_free_model(

528 self,

529 - architectures: Union[str, List[str]],

529 + architectures: Union[str, list[str]],

530 ) -> bool:

531 model_cls, _ = self.inspect_model_cls(architectures)

532 return model_cls.is_attention_free

533

534 def is_hybrid_model(

535 self,

536 - architectures: Union[str, List[str]],

536 + architectures: Union[str, list[str]],

537 ) -> bool:

538 model_cls, _ = self.inspect_model_cls(architectures)

539 return model_cls.is_hybrid

540

541 def is_noops_model(

542 self,

543 - architectures: Union[str, List[str]],

543 + architectures: Union[str, list[str]],

544 ) -> bool:

545 model_cls, _ = self.inspect_model_cls(architectures)

546 return model_cls.has_noops

547

548 def is_transcription_model(

549 self,

550 - architectures: Union[str, List[str]],

550 + architectures: Union[str, list[str]],

551 ) -> bool:

552 model_cls, _ = self.inspect_model_cls(architectures)

553 return model_cls.supports_transcription

554

555 def is_v1_compatible(

556 self,

557 - architectures: Union[str, List[str]],

557 + architectures: Union[str, list[str]],

558 ) -> bool:

559 model_cls, _ = self.inspect_model_cls(architectures)

560 return not model_cls.supports_v0_only

vllm/model_executor/models/roberta.py

+6 -5

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 import itertools

4 -from typing import Iterable, Optional, Tuple, Union

4 +from collections.abc import Iterable

5 +from typing import Optional, Union

6

7 import torch

8 from torch import nn

@@ -135,7 +136,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

136 prefix=prefix,

137 embedding_class=RobertaEmbedding)

138

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

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

140 weights = self.hf_to_vllm_mapper.apply(weights)

141 # Separate weights in "roberta"-prefixed and all else (not in memory).

142 # For use with models like FacebookAI/roberta-base.

@@ -187,7 +188,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,

188 self.classifier = RobertaClassificationHead(config)

189 self._pooler = CrossEncodingPooler(config, self.classifier)

190

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

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

192 bert_weights, task_weights = roberta_task_weights_filter(weights)

193 bert_weights = self.jina_to_vllm_mapper.apply(bert_weights)

194

@@ -249,8 +250,8 @@ def create_position_ids_from_input_ids(input_ids,

250

251

252 def roberta_task_weights_filter(

252 - all_weights: Iterable[Tuple[str, torch.Tensor]]

253 -) -> Tuple[Iterable[Tuple[str, torch.Tensor]], Iterable[Tuple[str,

253 + all_weights: Iterable[tuple[str, torch.Tensor]]

254 +) -> tuple[Iterable[tuple[str, torch.Tensor]], Iterable[tuple[str,

255 torch.Tensor]]]:

256 """

257 Separate task-specific weights that are applied on top

vllm/model_executor/models/siglip.py

+6 -5

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

3 within a vision language model."""

4

5 import math

6 -from typing import Iterable, Optional, Set, Tuple, Union

6 +from collections.abc import Iterable

7 +from typing import Optional, Union

8

9 import torch

10 from torch import nn

@@ -265,7 +266,7 @@ class SiglipEncoderLayer(nn.Module):

266 def forward(

267 self,

268 hidden_states: torch.Tensor,

268 - ) -> Tuple[torch.Tensor, None]:

269 + ) -> tuple[torch.Tensor, None]:

270 residual = hidden_states

271

272 hidden_states = self.layer_norm1(hidden_states)

@@ -480,8 +481,8 @@ class SiglipVisionModel(nn.Module):

481 feature_sample_layers=feature_sample_layers,

482 )

483

483 - def load_weights(self, weights: Iterable[Tuple[str,

484 - torch.Tensor]]) -> Set[str]:

484 + def load_weights(self, weights: Iterable[tuple[str,

485 + torch.Tensor]]) -> set[str]:

486 stacked_params_mapping = [

487 # (param_name, shard_name, shard_id)

488 ("qkv_proj", "q_proj", "q"),

@@ -489,7 +490,7 @@ class SiglipVisionModel(nn.Module):

490 ("qkv_proj", "v_proj", "v"),

491 ]

492 params_dict = dict(self.named_parameters())

492 - loaded_params: Set[str] = set()

493 + loaded_params: set[str] = set()

494 layer_count = len(self.vision_model.encoder.layers)

495

496 for name, loaded_weight in weights:

vllm/model_executor/models/skyworkr1v.py

+3 -3

@@ -8,7 +8,7 @@

8 # --------------------------------------------------------

9 from abc import ABC, abstractmethod

10 from collections.abc import Iterable, Mapping, Sequence

11 -from typing import Literal, Optional, Set, Tuple, TypedDict, TypeVar, Union

11 +from typing import Literal, Optional, TypedDict, TypeVar, Union

12

13 import torch

14 import torch.nn as nn

@@ -937,8 +937,8 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

937 return self.language_model.compute_logits(hidden_states,

938 sampling_metadata)

939

940 - def load_weights(self, weights: Iterable[Tuple[str,

941 - torch.Tensor]]) -> Set[str]:

940 + def load_weights(self, weights: Iterable[tuple[str,

941 + torch.Tensor]]) -> set[str]:

942 skip_prefixes = [

943 "action_embed", "temporal_embed", "track_embed",

944 "track_embed_decoder", "box_token", "cg_criterion", "cg_model",

vllm/model_executor/models/smolvlm.py

+2 -2

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Dict, Optional

3 +from typing import Optional

4

5 from transformers import SmolVLMProcessor

6

@@ -21,7 +21,7 @@ class SmolVLMProcessingInfo(Idefics3ProcessingInfo):

21 def get_hf_processor(

22 self,

23 *,

24 - max_image_size: Optional[Dict[str, int]] = None,

24 + max_image_size: Optional[dict[str, int]] = None,

25 **kwargs: object,

26 ) -> SmolVLMProcessor:

27 if max_image_size is not None:

vllm/model_executor/models/solar.py

+7 -6

@@ -23,7 +23,8 @@

23 # limitations under the License.

24 """Inference-only Solar model compatible with HuggingFace weights."""

25

26 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

26 +from collections.abc import Iterable

27 +from typing import Any, Optional, Union

28

29 import torch

30 from torch import nn

@@ -101,7 +102,7 @@ class SolarAttention(nn.Module):

102 num_heads: int,

103 num_kv_heads: int,

104 rope_theta: float = 10000,

104 - rope_scaling: Optional[Dict[str, Any]] = None,

105 + rope_scaling: Optional[dict[str, Any]] = None,

106 max_position_embeddings: int = 8192,

107 quant_config: Optional[QuantizationConfig] = None,

108 bias: bool = False,

@@ -236,7 +237,7 @@ class SolarDecoderLayer(nn.Module):

237 positions: torch.Tensor,

238 hidden_states: torch.Tensor,

239 residual: Optional[torch.Tensor],

239 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

241 # Self Attention

242 if residual is None:

243 residual = hidden_states

@@ -437,8 +438,8 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

438 sampling_metadata)

439 return logits

440

440 - def load_weights(self, weights: Iterable[Tuple[str,

441 - torch.Tensor]]) -> Set[str]:

441 + def load_weights(self, weights: Iterable[tuple[str,

442 + torch.Tensor]]) -> set[str]:

443 stacked_params_mapping = [

444 # (param_name, shard_name, shard_id)

445 (".qkv_proj", ".q_proj", "q"),

@@ -448,7 +449,7 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

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

450 ]

451 params_dict = dict(self.named_parameters())

451 - loaded_params: Set[str] = set()

452 + loaded_params: set[str] = set()

453 for name, loaded_weight in weights:

454 if "rotary_emb.inv_freq" in name:

455 continue

vllm/model_executor/models/stablelm.py

+8 -7

@@ -20,7 +20,8 @@

20 # https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/config.json

21 """Inference-only StabeLM (https://github.com/Stability-AI/StableLM)

22 model compatible with HuggingFace weights."""

23 -from typing import Iterable, Optional, Set, Tuple, Union

23 +from collections.abc import Iterable

24 +from typing import Optional, Union

25

26 import torch

27 from torch import nn

@@ -180,7 +181,7 @@ class StablelmDecoderLayer(nn.Module):

181 self,

182 positions: torch.Tensor,

183 hidden_states: torch.Tensor,

183 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

185 # Self Attention

186 residual = hidden_states

187 hidden_states = self.input_layernorm(hidden_states)

@@ -252,8 +253,8 @@ class StableLMEpochModel(nn.Module):

253 hidden_states = self.norm(hidden_states)

254 return hidden_states

255

255 - def load_weights(self, weights: Iterable[Tuple[str,

256 - torch.Tensor]]) -> Set[str]:

256 + def load_weights(self, weights: Iterable[tuple[str,

257 + torch.Tensor]]) -> set[str]:

258 stacked_params_mapping = [

259 # (param_name, shard_name, shard_id)

260 ("qkv_proj", "q_proj", "q"),

@@ -263,7 +264,7 @@ class StableLMEpochModel(nn.Module):

264 ("gate_up_proj", "up_proj", 1),

265 ]

266 params_dict = dict(self.named_parameters())

266 - loaded_params: Set[str] = set()

267 + loaded_params: set[str] = set()

268 for name, loaded_weight in weights:

269 for (param_name, weight_name, shard_id) in stacked_params_mapping:

270 if weight_name not in name:

@@ -335,8 +336,8 @@ class StablelmForCausalLM(nn.Module, SupportsPP):

336 sampling_metadata)

337 return logits

338

338 - def load_weights(self, weights: Iterable[Tuple[str,

339 - torch.Tensor]]) -> Set[str]:

339 + def load_weights(self, weights: Iterable[tuple[str,

340 + torch.Tensor]]) -> set[str]:

341 loader = AutoWeightsLoader(

342 self,

343 # Models trained using ColossalAI may include these tensors in

vllm/model_executor/models/starcoder2.py

+7 -6

@@ -19,7 +19,8 @@

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

20 # limitations under the License.

21 """ PyTorch Starcoder2 model."""

22 -from typing import Iterable, Optional, Set, Tuple, Union

22 +from collections.abc import Iterable

23 +from typing import Optional, Union

24

25 import torch

26 from torch import nn

@@ -255,8 +256,8 @@ class Starcoder2Model(nn.Module):

256 hidden_states = self.norm(hidden_states)

257 return hidden_states

258

258 - def load_weights(self, weights: Iterable[Tuple[str,

259 - torch.Tensor]]) -> Set[str]:

259 + def load_weights(self, weights: Iterable[tuple[str,

260 + torch.Tensor]]) -> set[str]:

261 stacked_params_mapping = [

262 # (param_name, shard_name, shard_id)

263 ("qkv_proj", "q_proj", "q"),

@@ -265,7 +266,7 @@ class Starcoder2Model(nn.Module):

266 ]

267

268 params_dict = dict(self.named_parameters(remove_duplicate=False))

268 - loaded_params: Set[str] = set()

269 + loaded_params: set[str] = set()

270 for name, loaded_weight in weights:

271 for (param_name, weight_name, shard_id) in stacked_params_mapping:

272 if weight_name not in name:

@@ -342,8 +343,8 @@ class Starcoder2ForCausalLM(nn.Module, SupportsPP):

343 sampling_metadata)

344 return logits

345

345 - def load_weights(self, weights: Iterable[Tuple[str,

346 - torch.Tensor]]) -> Set[str]:

346 + def load_weights(self, weights: Iterable[tuple[str,

347 + torch.Tensor]]) -> set[str]:

348 loader = AutoWeightsLoader(

349 self,

350 # Models trained using ColossalAI may include these tensors in

vllm/model_executor/models/telechat2.py

+6 -6

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

19 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

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

21 # limitations under the License.

22 -from typing import Iterable, Set, Tuple

22 +from collections.abc import Iterable

23

24 import torch

25 import torch.nn as nn

@@ -50,14 +50,14 @@ class TeleChat2Model(LlamaModel):

50 layer.mlp.gate_up_proj.bias = None

51 layer.mlp.gate_up_proj.skip_bias_add = True

52

53 - def load_weights(self, weights: Iterable[Tuple[str,

54 - torch.Tensor]]) -> Set[str]:

53 + def load_weights(self, weights: Iterable[tuple[str,

54 + torch.Tensor]]) -> set[str]:

55 stacked_params_mapping = [

56 ('gate_up_proj', 'gate_proj', 0),

57 ('gate_up_proj', 'up_proj', 1),

58 ]

59 params_dict = dict(self.named_parameters())

60 - loaded_params: Set[str] = set()

60 + loaded_params: set[str] = set()

61 total_num_heads = self.config.n_head

62 head_dim = self.config.hidden_size // total_num_heads

63 for name, loaded_weight in weights:

@@ -128,8 +128,8 @@ class TeleChat2ForCausalLM(LlamaForCausalLM):

128 layer_type: type[nn.Module] = LlamaDecoderLayer):

129 return TeleChat2Model(vllm_config=vllm_config, prefix=prefix)

130

131 - def load_weights(self, weights: Iterable[Tuple[str,

132 - torch.Tensor]]) -> Set[str]:

131 + def load_weights(self, weights: Iterable[tuple[str,

132 + torch.Tensor]]) -> set[str]:

133

134 loader = AutoWeightsLoader(

135 self,

vllm/model_executor/models/transformers.py

+2 -1

@@ -15,7 +15,8 @@

15 # limitations under the License.

16 """Wrapper around `transformers` models"""

17 import re

18 -from typing import Iterable, Literal, Optional, Union

18 +from collections.abc import Iterable

19 +from typing import Literal, Optional, Union

20

21 import torch

22 from torch import nn

vllm/model_executor/models/ultravox.py

+3 -3

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

3 # Adapted from https://github.com/fixie-ai/ultravox/blob/ecd58c4041030bae2ad15aa6bcf04ab43199ea02/ultravox/model/ultravox_model.py

4 """PyTorch Ultravox model."""

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Any, Literal, Optional, Set, Tuple, TypedDict, Union

6 +from typing import Any, Literal, Optional, TypedDict, Union

7

8 import torch

9 from torch import nn

@@ -619,8 +619,8 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

619 return self.language_model.compute_logits(hidden_states,

620 sampling_metadata)

621

622 - def load_weights(self, weights: Iterable[Tuple[str,

623 - torch.Tensor]]) -> Set[str]:

622 + def load_weights(self, weights: Iterable[tuple[str,

623 + torch.Tensor]]) -> set[str]:

624

625 loader = AutoWeightsLoader(self,

626 ignore_unexpected_prefixes=["audio_tower."])

vllm/model_executor/models/utils.py

+25 -25

@@ -1,9 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 import itertools

4 +from collections.abc import Iterable, Mapping

5 from dataclasses import dataclass, field

5 -from typing import (Callable, Dict, Iterable, List, Literal, Mapping, Optional,

6 - Protocol, Set, Tuple, Union, overload)

6 +from typing import Callable, Literal, Optional, Protocol, Union, overload

7

8 import torch

9 import torch.nn as nn

@@ -58,8 +58,8 @@ class WeightsMapper:

58 return key

59

60 def apply(

61 - self, weights: Iterable[Tuple[str, torch.Tensor]]

62 - ) -> Iterable[Tuple[str, torch.Tensor]]:

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

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

63 return ((out_name, data) for name, data in weights

64 if (out_name := self._map_name(name)) is not None)

65

@@ -84,8 +84,8 @@ class AutoWeightsLoader:

84 self,

85 module: nn.Module,

86 *,

87 - skip_prefixes: Optional[List[str]] = None,

88 - ignore_unexpected_prefixes: Optional[List[str]] = None,

87 + skip_prefixes: Optional[list[str]] = None,

88 + ignore_unexpected_prefixes: Optional[list[str]] = None,

89 ) -> None:

90 super().__init__()

91

@@ -95,8 +95,8 @@ class AutoWeightsLoader:

95

96 def _groupby_prefix(

97 self,

98 - weights: Iterable[Tuple[str, torch.Tensor]],

99 - ) -> Iterable[Tuple[str, Iterable[Tuple[str, torch.Tensor]]]]:

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

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

100 weights_by_parts = ((weight_name.split(".", 1), weight_data)

101 for weight_name, weight_data in weights)

102

@@ -129,7 +129,7 @@ class AutoWeightsLoader:

129 self,

130 base_prefix: str,

131 param: nn.Parameter,

132 - weights: Iterable[Tuple[str, torch.Tensor]],

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

133 ) -> Iterable[str]:

134 for weight_name, weight_data in weights:

135 weight_qualname = self._get_qualname(base_prefix, weight_name)

@@ -159,7 +159,7 @@ class AutoWeightsLoader:

159 yield weight_qualname

160

161 def _add_loadable_non_param_tensors(self, module: nn.Module,

162 - child_params: Dict[str, torch.Tensor]):

162 + child_params: dict[str, torch.Tensor]):

163 """

164 Add tensor names that are not in the model params that may be in the

165 safetensors, e.g., batch normalization stats.

@@ -182,7 +182,7 @@ class AutoWeightsLoader:

182 self,

183 base_prefix: str,

184 module: nn.Module,

185 - weights: Iterable[Tuple[str, torch.Tensor]],

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

186 ) -> Iterable[str]:

187 if isinstance(module, PPMissingLayer):

188 return

@@ -251,10 +251,10 @@ class AutoWeightsLoader:

251

252 def load_weights(

253 self,

254 - weights: Iterable[Tuple[str, torch.Tensor]],

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

255 *,

256 mapper: Optional[WeightsMapper] = None,

257 - ) -> Set[str]:

257 + ) -> set[str]:

258 if mapper is not None:

259 weights = mapper.apply(weights)

260

@@ -292,13 +292,13 @@ def flatten_bn(x: torch.Tensor) -> torch.Tensor:

292

293

294 @overload

295 -def flatten_bn(x: List[torch.Tensor]) -> List[torch.Tensor]:

295 +def flatten_bn(x: list[torch.Tensor]) -> list[torch.Tensor]:

296 ...

297

298

299 @overload

300 def flatten_bn(

301 - x: Union[List[torch.Tensor], torch.Tensor],

301 + x: Union[list[torch.Tensor], torch.Tensor],

302 *,

303 concat: Literal[True],

304 ) -> torch.Tensor:

@@ -307,18 +307,18 @@ def flatten_bn(

307

308 @overload

309 def flatten_bn(

310 - x: Union[List[torch.Tensor], torch.Tensor],

310 + x: Union[list[torch.Tensor], torch.Tensor],

311 *,

312 concat: bool = False,

313 -) -> Union[List[torch.Tensor], torch.Tensor]:

313 +) -> Union[list[torch.Tensor], torch.Tensor]:

314 ...

315

316

317 def flatten_bn(

318 - x: Union[List[torch.Tensor], torch.Tensor],

318 + x: Union[list[torch.Tensor], torch.Tensor],

319 *,

320 concat: bool = False,

321 -) -> Union[List[torch.Tensor], torch.Tensor]:

321 +) -> Union[list[torch.Tensor], torch.Tensor]:

322 """

323 Flatten the ``B`` and ``N`` dimensions of batched multimodal inputs.

324

@@ -442,7 +442,7 @@ def merge_multimodal_embeddings(

442 input_ids: torch.Tensor,

443 inputs_embeds: torch.Tensor,

444 multimodal_embeddings: NestedTensors,

445 - placeholder_token_id: Union[int, List[int]],

445 + placeholder_token_id: Union[int, list[int]],

446 ) -> torch.Tensor:

447 """

448 Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the

@@ -596,7 +596,7 @@ def make_layers(

596 num_hidden_layers: int,

597 layer_fn: LayerFn,

598 prefix: str,

599 -) -> Tuple[int, int, torch.nn.ModuleList]:

599 +) -> tuple[int, int, torch.nn.ModuleList]:

600 """Make a list of layers with the given layer function, taking

601 pipeline parallelism into account.

602 """

@@ -614,10 +614,10 @@ def make_layers(

614

615

616 # NOTE: don't use lru_cache here because it can prevent garbage collection

617 -_model_to_pp_missing_layer_names: Dict[int, List[str]] = {}

617 +_model_to_pp_missing_layer_names: dict[int, list[str]] = {}

618

619

620 -def get_pp_missing_layer_names(model: torch.nn.Module) -> List[str]:

620 +def get_pp_missing_layer_names(model: torch.nn.Module) -> list[str]:

621 """Get the names of the missing layers in a pipeline parallel model."""

622 model_id = id(model)

623 if model_id in _model_to_pp_missing_layer_names:

@@ -645,7 +645,7 @@ def is_pp_missing_parameter(name: str, model: torch.nn.Module) -> bool:

645 for missing_layer_name in get_pp_missing_layer_names(model))

646

647

648 -def make_empty_intermediate_tensors_factory(keys: List[str], hidden_size: int):

648 +def make_empty_intermediate_tensors_factory(keys: list[str], hidden_size: int):

649

650 def make_empty_intermediate_tensors(

651 batch_size: int,

@@ -684,7 +684,7 @@ def extract_layer_index(layer_name: str) -> int:

684 - "model.encoder.layers.0.sub.1" -> ValueError

685 """

686 subnames = layer_name.split(".")

687 - int_vals: List[int] = []

687 + int_vals: list[int] = []

688 for subname in subnames:

689 try:

690 int_vals.append(int(subname))

vllm/model_executor/models/whisper.py

+11 -11

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

2

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import List, Optional, Set, Tuple, TypedDict, Union

5 +from typing import Optional, TypedDict, Union

6

7 import torch

8 from torch import nn

@@ -382,7 +382,7 @@ class WhisperEncoder(nn.Module):

382 self.embed_positions.weight.copy_(

383 sinusoids(*self.embed_positions.weight.shape))

384

385 - def forward(self, input_features: Union[torch.Tensor, List[torch.Tensor]]):

385 + def forward(self, input_features: Union[torch.Tensor, list[torch.Tensor]]):

386 hidden_states = []

387 for features in input_features:

388 embeds = nn.functional.gelu(self.conv1(features))

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

460

461 def forward(

462 self,

463 - input_features: Optional[Union[torch.Tensor, List[torch.Tensor]]],

463 + input_features: Optional[Union[torch.Tensor, list[torch.Tensor]]],

464 input_ids: Optional[torch.Tensor],

465 positions: torch.Tensor,

466 ) -> torch.Tensor:

@@ -474,14 +474,14 @@ class WhisperModel(nn.Module):

474

475 def get_encoder_outputs(

476 self,

477 - input_features: Optional[Union[torch.Tensor, List[torch.Tensor]]],

477 + input_features: Optional[Union[torch.Tensor, list[torch.Tensor]]],

478 ) -> Optional[torch.Tensor]:

479 if input_features is None:

480 return None

481 return self.encoder(input_features)

482

483 - def load_weights(self, weights: Iterable[Tuple[str,

484 - torch.Tensor]]) -> Set[str]:

483 + def load_weights(self, weights: Iterable[tuple[str,

484 + torch.Tensor]]) -> set[str]:

485 stacked_params_mapping = [

486 # (param_name, shard_name, shard_id)

487 (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),

@@ -491,7 +491,7 @@ class WhisperModel(nn.Module):

491 (".encoder_attn.kv_proj", ".encoder_attn.v_proj", "v"),

492 ]

493 params_dict = dict(self.named_parameters())

494 - loaded_params: Set[str] = set()

494 + loaded_params: set[str] = set()

495 for name, loaded_weight in weights:

496 for param_name, weight_name, shard_id in stacked_params_mapping:

497 if weight_name not in name:

@@ -722,8 +722,8 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,

722 sampling_metadata)

723 return logits

724

725 - def load_weights(self, weights: Iterable[Tuple[str,

726 - torch.Tensor]]) -> Set[str]:

725 + def load_weights(self, weights: Iterable[tuple[str,

726 + torch.Tensor]]) -> set[str]:

727 loader = AutoWeightsLoader(self, skip_prefixes=["proj_out."])

728

729 # add fake zeros bias for k_proj to state_dict

@@ -732,8 +732,8 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,

732

733

734 def _create_fake_bias_for_k_proj(

735 - weights: Iterable[Tuple[str, torch.Tensor]]

736 -) -> Iterable[Tuple[str, torch.Tensor]]:

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

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

737 """

738 Create full zeros bias for k_proj weight in self-attn and x-attn layers.

739 So that the bias for k_proj in qkv_proj can be initialized with zeros.

vllm/model_executor/models/zamba2.py

+13 -12

@@ -6,8 +6,9 @@ https://arxiv.org/abs/2411.15242, which combines Mamba and Transformer

6 architectures in a hybrid model optimized for efficient sequence modeling. The

7 model alternates between state space model layers and attention-based layers.

8 """

9 +from collections.abc import Iterable

10 from itertools import cycle

10 -from typing import Dict, Iterable, List, Optional, Set, Tuple, Union

11 +from typing import Optional, Union

12

13 import torch

14 from torch import nn

@@ -54,7 +55,7 @@ class Zamba2LoRA(nn.Module):

55 self,

56 input_dim: int,

57 rank: int,

57 - output_dim: Union[int, List[int]],

58 + output_dim: Union[int, list[int]],

59 quant_config: Optional[QuantizationConfig] = None,

60 ):

61 """Initialize the attention layer.

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

280 self,

281 config: Zamba2Config,

282 bare_block_idx: int,

282 - num_hybrid_layers: Dict[int, int],

283 + num_hybrid_layers: dict[int, int],

284 quant_config: Optional[QuantizationConfig] = None,

285 ) -> None:

286 """Initialize the MLP layer.

@@ -769,8 +770,8 @@ class Zamba2Model(nn.Module):

770 hidden_states = self.final_layernorm(hidden_states)

771 return hidden_states

772

772 - def load_weights(self, weights: Iterable[Tuple[str,

773 - torch.Tensor]]) -> Set[str]:

773 + def load_weights(self, weights: Iterable[tuple[str,

774 + torch.Tensor]]) -> set[str]:

775 stacked_params_mapping = [

776 # (param_name, shard_name, shard_id)

777 ("qkv_proj", "q_proj", "q"),

@@ -779,7 +780,7 @@ class Zamba2Model(nn.Module):

780 ]

781

782 params_dict = dict(self.named_parameters())

782 - loaded_params: Set[str] = set()

783 + loaded_params: set[str] = set()

784 for chkpt_weight_name, loaded_weight in weights:

785 for param_name, weight_name, shard_id in stacked_params_mapping:

786 if weight_name not in chkpt_weight_name:

@@ -914,9 +915,9 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

915

916 return hidden_states

917

917 - def copy_inputs_before_cuda_graphs(self, input_buffers: Dict[str,

918 + def copy_inputs_before_cuda_graphs(self, input_buffers: dict[str,

919 torch.Tensor],

919 - **kwargs) -> Dict[str, torch.Tensor]:

920 + **kwargs) -> dict[str, torch.Tensor]:

921 """Copy inputs before CUDA graph capture.

922

923 Args:

@@ -930,7 +931,7 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

931 input_buffers, **kwargs)

932

933 def get_seqlen_agnostic_capture_inputs(

933 - self, batch_size: int) -> Dict[str, torch.Tensor]:

934 + self, batch_size: int) -> dict[str, torch.Tensor]:

935 """Get inputs for sequence-length-agnostic graph capture.

936

937 Args:

@@ -941,7 +942,7 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

942 return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

943

944 def _get_mamba_cache_shape(

944 - self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

945 + self) -> tuple[tuple[int, int], tuple[int, int]]:

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

947

948 Returns:

@@ -1001,7 +1002,7 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

1002 sampling_metadata)

1003 return logits

1004

1004 - def load_weights(self, weights: Iterable[Tuple[str,

1005 - torch.Tensor]]) -> Set[str]:

1005 + def load_weights(self, weights: Iterable[tuple[str,

1006 + torch.Tensor]]) -> set[str]:

1007 loader = AutoWeightsLoader(self)

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

Update deprecated type hinting in `models` (#18132)