6223dd8114 · tduncan/Nvidia-Optimization-Stack-VLLM

pyproject.toml

-1

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

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

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

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

83 -"vllm/model_executor/layers/**/*.py" = ["UP006", "UP035"]

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

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

85 "vllm/platforms/**/*.py" = ["UP006", "UP035"]

vllm/model_executor/layers/fused_moe/__init__.py

+3 -3

@@ -1,13 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 from contextlib import contextmanager

4 -from typing import Any, Dict, Optional

4 +from typing import Any, Optional

5

6 from vllm.model_executor.layers.fused_moe.layer import (

7 FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)

8 from vllm.triton_utils import HAS_TRITON

9

10 -_config: Optional[Dict[str, Any]] = None

10 +_config: Optional[dict[str, Any]] = None

11

12

13 @contextmanager

@@ -19,7 +19,7 @@ def override_config(config):

19 _config = old_config

20

21

22 -def get_config() -> Optional[Dict[str, Any]]:

22 +def get_config() -> Optional[dict[str, Any]]:

23 return _config

24

25

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 import importlib.util

3 -from typing import Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -61,7 +61,7 @@ def _moe_permute(

61 global_num_experts: int,

62 expert_map: Optional[torch.Tensor],

63 block_m: int,

64 -) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,

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

65 Optional[torch.Tensor]]:

66 """

67 Determine the sorted_token_ids, expert_ids for the given problem size.

vllm/model_executor/layers/fused_moe/fused_moe.py

+23 -23

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

3 import functools

4 import json

5 import os

6 -from typing import Any, Callable, Dict, List, Optional, Tuple

6 +from typing import Any, Callable, Optional

7

8 import torch

9

@@ -472,14 +472,14 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

472 num_tokens_post_padded: torch.Tensor,

473 mul_routed_weight: bool,

474 top_k: int,

475 - config: Dict[str, Any],

475 + config: dict[str, Any],

476 compute_type: tl.dtype,

477 use_fp8_w8a8: bool,

478 use_int8_w8a8: bool,

479 use_int8_w8a16: bool,

480 use_int4_w4a16: bool,

481 per_channel_quant: bool,

482 - block_shape: Optional[List[int]] = None) -> None:

482 + block_shape: Optional[list[int]] = None) -> None:

483 assert topk_weights is not None or not mul_routed_weight

484 assert topk_weights is None or topk_weights.stride(1) == 1

485 assert sorted_token_ids.stride(0) == 1

@@ -622,7 +622,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

622 def get_config_file_name(E: int,

623 N: int,

624 dtype: Optional[str],

625 - block_shape: Optional[List[int]] = None) -> str:

625 + block_shape: Optional[list[int]] = None) -> str:

626 device_name = current_platform.get_device_name().replace(" ", "_")

627 dtype_selector = "" if not dtype else f",dtype={dtype}"

628 block_shape_selector = ("" if not block_shape or not all(block_shape) else

@@ -638,7 +638,7 @@ def get_moe_configs(

638 dtype: Optional[str],

639 block_n: Optional[int] = None,

640 block_k: Optional[int] = None,

641 -) -> Optional[Dict[int, Any]]:

641 +) -> Optional[dict[int, Any]]:

642 """

643 Return optimized configurations for the fused MoE kernel.

644

@@ -670,7 +670,7 @@ def get_moe_configs(

670 return None

671

672

673 -def get_moe_wna16_block_config(config: Dict[str,

673 +def get_moe_wna16_block_config(config: dict[str,

674 int], use_moe_wna16_cuda: bool,

675 num_valid_tokens: int, size_k: int, size_n: int,

676 num_experts: int, group_size: int,

@@ -742,8 +742,8 @@ def get_default_config(

742 topk: int,

743 dtype: Optional[str],

744 is_marlin: bool,

745 - block_shape: Optional[List[int]] = None,

746 -) -> Dict[str, int]:

745 + block_shape: Optional[list[int]] = None,

746 +) -> dict[str, int]:

747 if dtype == "fp8_w8a8" and block_shape is not None:

748 # Block-wise quant: BLOCK_SIZE_N must be divisible by block_shape[0]

749 # BLOCK_SIZE_K must be divisible by block_shape[1]

@@ -795,13 +795,13 @@ def get_default_config(

795

796

797 def try_get_optimal_moe_config(

798 - w1_shape: Tuple[int, ...],

799 - w2_shape: Tuple[int, ...],

798 + w1_shape: tuple[int, ...],

799 + w2_shape: tuple[int, ...],

800 top_k: int,

801 dtype: Optional[str],

802 M: int,

803 is_marlin: bool = False,

804 - block_shape: Optional[List[int]] = None,

804 + block_shape: Optional[list[int]] = None,

805 ):

806 from vllm.model_executor.layers.fused_moe import get_config

807 override_config = get_config()

@@ -855,7 +855,7 @@ def fused_topk(

855 gating_output: torch.Tensor,

856 topk: int,

857 renormalize: bool,

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

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

859 assert hidden_states.shape[0] == gating_output.shape[0], (

860 "Number of tokens mismatch")

861

@@ -895,7 +895,7 @@ def grouped_topk(

895 topk_group: int = 0,

896 scoring_func: str = "softmax",

897 e_score_correction_bias: Optional[torch.Tensor] = None

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

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

899

900 assert hidden_states.shape[0] == gating_output.shape[0], (

901 "Number of tokens mismatch")

@@ -982,7 +982,7 @@ def inplace_fused_experts(hidden_states: torch.Tensor,

982 w2_zp: Optional[torch.Tensor] = None,

983 a1_scale: Optional[torch.Tensor] = None,

984 a2_scale: Optional[torch.Tensor] = None,

985 - block_shape: Optional[List[int]] = None) -> None:

985 + block_shape: Optional[list[int]] = None) -> None:

986 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,

987 activation, apply_router_weight_on_input, use_fp8_w8a8,

988 use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,

@@ -1012,7 +1012,7 @@ def inplace_fused_experts_fake(

1012 w2_zp: Optional[torch.Tensor] = None,

1013 a1_scale: Optional[torch.Tensor] = None,

1014 a2_scale: Optional[torch.Tensor] = None,

1015 - block_shape: Optional[List[int]] = None) -> None:

1015 + block_shape: Optional[list[int]] = None) -> None:

1016 pass

1017

1018

@@ -1046,7 +1046,7 @@ def outplace_fused_experts(

1046 w2_zp: Optional[torch.Tensor] = None,

1047 a1_scale: Optional[torch.Tensor] = None,

1048 a2_scale: Optional[torch.Tensor] = None,

1049 - block_shape: Optional[List[int]] = None) -> torch.Tensor:

1049 + block_shape: Optional[list[int]] = None) -> torch.Tensor:

1050 return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,

1051 False, activation, apply_router_weight_on_input,

1052 use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,

@@ -1076,7 +1076,7 @@ def outplace_fused_experts_fake(

1076 w2_zp: Optional[torch.Tensor] = None,

1077 a1_scale: Optional[torch.Tensor] = None,

1078 a2_scale: Optional[torch.Tensor] = None,

1079 - block_shape: Optional[List[int]] = None) -> torch.Tensor:

1079 + block_shape: Optional[list[int]] = None) -> torch.Tensor:

1080 return torch.empty_like(hidden_states)

1081

1082

@@ -1129,7 +1129,7 @@ def fused_experts(hidden_states: torch.Tensor,

1129 w2_zp: Optional[torch.Tensor] = None,

1130 a1_scale: Optional[torch.Tensor] = None,

1131 a2_scale: Optional[torch.Tensor] = None,

1132 - block_shape: Optional[List[int]] = None,

1132 + block_shape: Optional[list[int]] = None,

1133 allow_deep_gemm: bool = False) -> torch.Tensor:

1134 if (allow_deep_gemm and use_fp8_w8a8

1135 and _valid_deep_gemm(hidden_states, w1, w2, expert_map)):

@@ -1184,8 +1184,8 @@ def moe_kernel_prepare_input(

1184 use_int8_w8a16: bool,

1185 use_int4_w4a16: bool,

1186 per_channel_quant: bool,

1187 - block_shape: Optional[List[int]] = None,

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

1187 + block_shape: Optional[list[int]] = None,

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

1189 if use_fp8_w8a8:

1190 assert B_scale is not None

1191 if block_shape is None:

@@ -1248,7 +1248,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,

1248 w2_zp: Optional[torch.Tensor] = None,

1249 a1_scale: Optional[torch.Tensor] = None,

1250 a2_scale: Optional[torch.Tensor] = None,

1251 - block_shape: Optional[List[int]] = None):

1251 + block_shape: Optional[list[int]] = None):

1252 # Check constraints.

1253 if use_int4_w4a16:

1254 assert hidden_states.shape[1] // 2 == w1.shape[

@@ -1452,7 +1452,7 @@ def fused_moe(

1452 w2_zp: Optional[torch.Tensor] = None,

1453 a1_scale: Optional[torch.Tensor] = None,

1454 a2_scale: Optional[torch.Tensor] = None,

1455 - block_shape: Optional[List[int]] = None,

1455 + block_shape: Optional[list[int]] = None,

1456 ) -> torch.Tensor:

1457 """

1458 This function computes a Mixture of Experts (MoE) layer using two sets of

@@ -1497,7 +1497,7 @@ def fused_moe(

1497 a1.

1498 - a2_scale (Optional[torch.Tensor]): Optional scale to be used for

1499 a2.

1500 - - block_shape: (Optional[List[int]]): Optional block size for block-wise

1500 + - block_shape: (Optional[list[int]]): Optional block size for block-wise

1501 quantization.

1502

1503 Returns:

vllm/model_executor/layers/fused_moe/layer.py

+4 -4

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

2

3 from abc import abstractmethod

4 from enum import Enum

5 -from typing import Callable, List, Optional, Tuple

5 +from typing import Callable, Optional

6

7 import torch

8 import torch.nn.functional as F

@@ -326,7 +326,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

326

327 def determine_expert_map(

328 ep_size: int, ep_rank: int,

329 - global_num_experts: int) -> Tuple[int, Optional[torch.Tensor]]:

329 + global_num_experts: int) -> tuple[int, Optional[torch.Tensor]]:

330 """

331 Calculates how many experts should be assigned to each rank for EP and

332 creates a mapping from global to local expert index. Experts are

@@ -338,7 +338,7 @@ def determine_expert_map(

338 global_num_experts (int): The total number of experts in the model.

339

340 Returns:

341 - Tuple[int, Optional[torch.Tensor]]: A tuple containing:

341 + tuple[int, Optional[torch.Tensor]]: A tuple containing:

342 - local_num_experts (int): The number of experts assigned

343 to the current rank.

344 - expert_map (Optional[torch.Tensor]): A tensor of shape

@@ -909,7 +909,7 @@ class FusedMoE(torch.nn.Module):

909 def make_expert_params_mapping(

910 cls, ckpt_gate_proj_name: str, ckpt_down_proj_name: str,

911 ckpt_up_proj_name: str,

912 - num_experts: int) -> List[Tuple[str, str, int, str]]:

912 + num_experts: int) -> list[tuple[str, str, int, str]]:

913

914 return [

915 # (param_name, weight_name, expert_id, shard_id)

vllm/model_executor/layers/fused_moe/moe_align_block_size.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Optional, Tuple

2 +from typing import Optional

3

4 import torch

5

@@ -153,7 +153,7 @@ def moe_align_block_size(

153 num_experts: int,

154 expert_map: Optional[torch.Tensor] = None,

155 pad_sorted_ids: bool = False

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

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

157 """

158 Aligns the token distribution across experts to be compatible with block

159 size for matrix multiplication.

vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Optional, Tuple

2 +from typing import Optional

3

4 import torch

5

@@ -15,7 +15,7 @@ def moe_permute(

15 expert_map: Optional[torch.Tensor] = None,

16 align_block_size: Optional[int] = None,

17 fill_invalid_expert: int = -1

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

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

19 """

20 This function expands and permutes activation to gather uncontinuous tokens

21 for each expert.

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+7 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2 from functools import cache

3 -from typing import List, Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -97,7 +97,7 @@ def rocm_aiter_fmoe_fp8_blockscale_g1u1_impl(

97 w1_scale: torch.Tensor,

98 w2_scale: torch.Tensor,

99 a1_scale: torch.Tensor,

100 - block_shape: List[int],

100 + block_shape: list[int],

101 smooth_scale: Optional[torch.Tensor] = None) -> torch.Tensor:

102 from aiter import fmoe_fp8_blockscale_g1u1

103 from aiter.fused_moe_bf16_asm import moe_sorting_ck

@@ -142,7 +142,7 @@ def rocm_aiter_fmoe_fp8_blockscale_g1u1_fake(

142 w1_scale: torch.Tensor,

143 w2_scale: torch.Tensor,

144 a1_scale: torch.Tensor,

145 - block_shape: List[int],

145 + block_shape: list[int],

146 smooth_scale: Optional[torch.Tensor] = None) -> torch.Tensor:

147

148 return torch.empty_like(a1, dtype=hidden_states_dtype)

@@ -280,7 +280,7 @@ def rocm_aiter_fused_experts(hidden_states: torch.Tensor,

280 w2_zp: Optional[torch.Tensor] = None,

281 a1_scale: Optional[torch.Tensor] = None,

282 a2_scale: Optional[torch.Tensor] = None,

283 - block_shape: Optional[List[int]] = None,

283 + block_shape: Optional[list[int]] = None,

284 allow_deep_gemm: bool = False) -> torch.Tensor:

285

286 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

@@ -372,14 +372,14 @@ def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,

372 topk_indices: torch.Tensor,

373 token_expert_indices: torch.Tensor,

374 gating_output: torch.Tensor,

375 - renormalize: bool) -> Tuple[torch.Tensor, ...]:

375 + renormalize: bool) -> tuple[torch.Tensor, ...]:

376 torch.ops.vllm.rocm_aiter_topk_softmax(topk_weights, topk_indices,

377 token_expert_indices, gating_output,

378 renormalize)

379 return topk_weights, topk_indices

380

381

382 -def shuffle_weights(*tensors: torch.Tensor) -> Tuple[torch.Tensor, ...]:

382 +def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:

383 """

384 Applies shuffle_weight function from AITER to each

385 input tensor and returns them.

@@ -395,7 +395,7 @@ def shuffle_weights(*tensors: torch.Tensor) -> Tuple[torch.Tensor, ...]:

395

396

397 def expand_weights(*tensors: torch.Tensor,

398 - expansion_dims: list[int]) -> Tuple[torch.Tensor, ...]:

398 + expansion_dims: list[int]) -> tuple[torch.Tensor, ...]:

399 """

400 Expands the dimensions of input tensors.

401

vllm/model_executor/layers/fused_moe/utils.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2 from math import prod

3 -from typing import List, Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -10,7 +10,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

10 from vllm.utils import cdiv

11

12

13 -def _resize_cache(x: torch.Tensor, v: Tuple[int, ...]) -> torch.Tensor:

13 +def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:

14 """

15 Shrink the given tensor and apply the given view to it. This is

16 used to resize the intermediate fused_moe caches.

@@ -22,8 +22,8 @@ def _resize_cache(x: torch.Tensor, v: Tuple[int, ...]) -> torch.Tensor:

22 def _fp8_quantize(

23 A: torch.Tensor,

24 A_scale: Optional[torch.Tensor],

25 - block_shape: Optional[List[int]],

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

25 + block_shape: Optional[list[int]],

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

27 """

28 Perform fp8 quantization on the inputs. If a block_shape

29 is provided, the output will be blocked.

vllm/model_executor/layers/layernorm.py

+10 -10

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Custom normalization layers."""

3 -from typing import Optional, Tuple, Union

3 +from typing import Optional, Union

4

5 import torch

6 import torch.nn as nn

@@ -31,7 +31,7 @@ def rms_norm(x: torch.Tensor, weight: torch.Tensor,

31

32 def fused_add_rms_norm(

33 x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,

34 - variance_epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:

34 + variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:

35 from vllm import _custom_ops as ops

36 ops.fused_add_rms_norm(

37 x,

@@ -57,7 +57,7 @@ def rocm_aiter_rms_norm(x: torch.Tensor, weight: torch.Tensor,

57

58 def rocm_aiter_fused_add_rms_norm(

59 x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,

60 - variance_epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:

60 + variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:

61

62 import aiter as rocm_aiter

63

@@ -119,7 +119,7 @@ class RMSNorm(CustomOp):

119 self,

120 x: torch.Tensor,

121 residual: Optional[torch.Tensor] = None,

122 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

123 """PyTorch-native implementation equivalent to forward()."""

124 orig_dtype = x.dtype

125 x = x.to(torch.float32)

@@ -157,7 +157,7 @@ class RMSNorm(CustomOp):

157 self,

158 x: torch.Tensor,

159 residual: Optional[torch.Tensor] = None,

160 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

161 if self.variance_size_override is not None:

162 return self.forward_native(x, residual)

163

@@ -174,7 +174,7 @@ class RMSNorm(CustomOp):

174 self,

175 x: torch.Tensor,

176 residual: Optional[torch.Tensor] = None,

177 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

178 from vllm_hpu_extension.kernels import rms_norm

179 HPUFusedRMSNorm = rms_norm()

180 if HPUFusedRMSNorm is None:

@@ -194,7 +194,7 @@ class RMSNorm(CustomOp):

194 self,

195 x: torch.Tensor,

196 residual: Optional[torch.Tensor] = None,

197 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

198 if self.variance_size_override is not None:

199 return self.forward_native(x, residual)

200

@@ -244,7 +244,7 @@ class GemmaRMSNorm(CustomOp):

244 variance_epsilon: float,

245 x: torch.Tensor,

246 residual: Optional[torch.Tensor],

247 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

248 """PyTorch-native implementation equivalent to forward()."""

249 orig_dtype = x.dtype

250 if residual is not None:

@@ -267,7 +267,7 @@ class GemmaRMSNorm(CustomOp):

267 self,

268 x: torch.Tensor,

269 residual: Optional[torch.Tensor] = None,

270 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

271 """PyTorch-native implementation equivalent to forward()."""

272 return self.forward_static(self.weight.data, self.variance_epsilon, x,

273 residual)

@@ -276,7 +276,7 @@ class GemmaRMSNorm(CustomOp):

276 self,

277 x: torch.Tensor,

278 residual: Optional[torch.Tensor] = None,

279 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

280 if torch.compiler.is_compiling():

281 return self.forward_native(x, residual)

282

vllm/model_executor/layers/mamba/mamba_mixer2.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Optional, Union

4

5 import torch

6 from torch import nn

@@ -104,7 +104,7 @@ class Mixer2RMSNormGated(CustomOp):

104 self,

105 x: torch.Tensor,

106 gate: torch.Tensor,

107 - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

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

108

109 if self.tp_size > 1 or self.n_groups != 1:

110 return self.forward_native(x, gate)

@@ -136,7 +136,7 @@ def extra_groups_for_head_shards(ngroups: int, tp_size: int):

136

137

138 def mamba_v2_sharded_weight_loader(

139 - shard_spec: List[Tuple[int, int, float]],

139 + shard_spec: list[tuple[int, int, float]],

140 tp_size: int,

141 tp_rank: int,

142 ) -> LoaderFunction:

vllm/model_executor/layers/pooler.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 from enum import IntEnum

4 -from typing import List, Optional, Union

4 +from typing import Optional, Union

5

6 import torch

7 import torch.nn as nn

@@ -46,7 +46,7 @@ class SimplePooler(nn.Module):

46 normalize: bool,

47 softmax: bool,

48 step_tag_id: Optional[int] = None,

49 - returned_token_ids: Optional[List[int]] = None,

49 + returned_token_ids: Optional[list[int]] = None,

50 ) -> "SimplePooler":

51 if pooling_type == PoolingType.LAST:

52 assert step_tag_id is None and returned_token_ids is None

@@ -174,7 +174,7 @@ class StepPool(SimplePooler):

174 normalize: bool,

175 softmax: bool,

176 step_tag_id: Optional[int] = None,

177 - returned_token_ids: Optional[List[int]] = None,

177 + returned_token_ids: Optional[list[int]] = None,

178 ):

179 super().__init__(normalize=normalize, softmax=softmax)

180

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

259 normalize: bool,

260 softmax: bool,

261 step_tag_id: Optional[int] = None,

262 - returned_token_ids: Optional[List[int]] = None,

262 + returned_token_ids: Optional[list[int]] = None,

263 ) -> SimplePooler:

264 return SimplePooler.from_pooling_type(

265 pooling_type=PoolingType[pooler_config.pooling_type]

vllm/model_executor/layers/quantization/__init__.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Literal, Type, get_args

3 +from typing import Literal, get_args

4

5 from vllm.model_executor.layers.quantization.base_config import (

6 QuantizationConfig)

@@ -76,7 +76,7 @@ def register_quantization_config(quantization: str):

76 return _wrapper

77

78

79 -def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:

79 +def get_quantization_config(quantization: str) -> type[QuantizationConfig]:

80 if quantization not in QUANTIZATION_METHODS:

81 raise ValueError(f"Invalid quantization method: {quantization}")

82

@@ -110,7 +110,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:

110 from .torchao import TorchAOConfig

111 from .tpu_int8 import Int8TpuConfig

112

113 - method_to_config: dict[str, Type[QuantizationConfig]] = {

113 + method_to_config: dict[str, type[QuantizationConfig]] = {

114 "aqlm": AQLMConfig,

115 "awq": AWQConfig,

116 "deepspeedfp": DeepSpeedFPConfig,

vllm/model_executor/layers/quantization/aqlm.py

+7 -7

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

4 # and https://arxiv.org/pdf/2401.06118.pdf

5

6 import math

7 -from typing import Any, Dict, List, Optional

7 +from typing import Any, Optional

8

9 import torch

10 import torch.nn.functional as F

@@ -98,7 +98,7 @@ def generic_dequantize_gemm(

98 codebooks: torch.

99 Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]

100 scales: torch.Tensor, # [num_out_groups, 1, 1, 1]

101 - output_partition_sizes: List[int],

101 + output_partition_sizes: list[int],

102 bias: Optional[torch.Tensor],

103 ) -> torch.Tensor:

104 output_shape = input.shape[:-1] + (scales.shape[0], )

@@ -136,7 +136,7 @@ def optimized_dequantize_gemm(

136 codebooks: torch.

137 Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]

138 scales: torch.Tensor, # [num_out_groups, 1, 1, 1]

139 - output_partition_sizes: List[int],

139 + output_partition_sizes: list[int],

140 bias: Optional[torch.Tensor],

141 ) -> torch.Tensor:

142 weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)

@@ -191,7 +191,7 @@ class AQLMConfig(QuantizationConfig):

191 return "aqlm"

192

193 @classmethod

194 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

194 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

195 return [torch.half]

196

197 @classmethod

@@ -199,11 +199,11 @@ class AQLMConfig(QuantizationConfig):

199 return 60

200

201 @classmethod

202 - def get_config_filenames(cls) -> List[str]:

202 + def get_config_filenames(cls) -> list[str]:

203 return [] # no extra configs.

204

205 @classmethod

206 - def from_config(cls, config: Dict[str, Any]) -> "AQLMConfig":

206 + def from_config(cls, config: dict[str, Any]) -> "AQLMConfig":

207 in_group_size = cls.get_from_keys(config, ["in_group_size"])

208 nbits_per_codebook = cls.get_from_keys(config, ["nbits_per_codebook"])

209 num_code_books = cls.get_from_keys(config, ["num_codebooks"])

@@ -230,7 +230,7 @@ class AQLMLinearMethod(LinearMethodBase):

230

231 def create_weights(self, layer: torch.nn.Module,

232 input_size_per_partition: int,

233 - output_partition_sizes: List[int], input_size: int,

233 + output_partition_sizes: list[int], input_size: int,

234 output_size: int, params_dtype: torch.dtype,

235 **extra_weight_attrs):

236 del output_size # Unused.

vllm/model_executor/layers/quantization/awq.py

+7 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6

@@ -25,7 +25,7 @@ class AWQConfig(QuantizationConfig):

25 weight_bits: int,

26 group_size: int,

27 zero_point: bool,

28 - modules_to_not_convert: Optional[List[str]] = None,

28 + modules_to_not_convert: Optional[list[str]] = None,

29 ) -> None:

30 super().__init__()

31 self.weight_bits = weight_bits

@@ -48,7 +48,7 @@ class AWQConfig(QuantizationConfig):

48 def get_name(self) -> QuantizationMethods:

49 return "awq"

50

51 - def get_supported_act_dtypes(self) -> List[torch.dtype]:

51 + def get_supported_act_dtypes(self) -> list[torch.dtype]:

52 return [torch.half]

53

54 @classmethod

@@ -57,7 +57,7 @@ class AWQConfig(QuantizationConfig):

57 return 75

58

59 @staticmethod

60 - def get_config_filenames() -> List[str]:

60 + def get_config_filenames() -> list[str]:

61 return [

62 "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq

63 # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq

@@ -65,7 +65,7 @@ class AWQConfig(QuantizationConfig):

65 ]

66

67 @classmethod

68 - def from_config(cls, config: Dict[str, Any]) -> "AWQConfig":

68 + def from_config(cls, config: dict[str, Any]) -> "AWQConfig":

69 weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])

70 group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])

71 zero_point = cls.get_from_keys(config, ["zero_point"])

@@ -82,7 +82,7 @@ class AWQConfig(QuantizationConfig):

82 return None

83

84

85 -def is_layer_skipped_awq(prefix: str, modules_to_not_convert: List[str]):

85 +def is_layer_skipped_awq(prefix: str, modules_to_not_convert: list[str]):

86 return any(module_name in prefix for module_name in modules_to_not_convert)

87

88

@@ -98,7 +98,7 @@ class AWQLinearMethod(LinearMethodBase):

98

99 def create_weights(self, layer: torch.nn.Module,

100 input_size_per_partition: int,

101 - output_partition_sizes: List[int], input_size: int,

101 + output_partition_sizes: list[int], input_size: int,

102 output_size: int, params_dtype: torch.dtype,

103 **extra_weight_attrs):

104 if input_size_per_partition % self.quant_config.group_size != 0:

vllm/model_executor/layers/quantization/awq_marlin.py

+8 -8

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, List, Optional

3 +from typing import Any, Callable, Optional

4

5 import torch

6 from torch.nn import Parameter

@@ -46,8 +46,8 @@ class AWQMarlinConfig(QuantizationConfig):

46

47 def __init__(self, weight_bits: int, group_size: int, zero_point: bool,

48 lm_head_quantized: bool,

49 - modules_to_not_convert: Optional[List[str]],

50 - full_config: Dict[str, Any]) -> None:

49 + modules_to_not_convert: Optional[list[str]],

50 + full_config: dict[str, Any]) -> None:

51 super().__init__()

52 self.pack_factor = 32 // weight_bits # packed into int32

53 self.group_size = group_size

@@ -79,7 +79,7 @@ class AWQMarlinConfig(QuantizationConfig):

79 return "awq_marlin"

80

81 @classmethod

82 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

82 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

83 return [torch.half, torch.bfloat16]

84

85 @classmethod

@@ -87,11 +87,11 @@ class AWQMarlinConfig(QuantizationConfig):

87 return 80

88

89 @classmethod

90 - def get_config_filenames(cls) -> List[str]:

90 + def get_config_filenames(cls) -> list[str]:

91 return ["quantize_config.json"]

92

93 @classmethod

94 - def from_config(cls, config: Dict[str, Any]) -> "AWQMarlinConfig":

94 + def from_config(cls, config: dict[str, Any]) -> "AWQMarlinConfig":

95 weight_bits = cls.get_from_keys(config, ["bits"])

96 group_size = cls.get_from_keys(config, ["group_size"])

97 zero_point = cls.get_from_keys(config, ["zero_point"])

@@ -150,7 +150,7 @@ class AWQMarlinConfig(QuantizationConfig):

150 return None

151

152 @classmethod

153 - def is_awq_marlin_compatible(cls, quant_config: Dict[str, Any]):

153 + def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]):

154 # Extract data from quant config.

155 quant_method = quant_config.get("quant_method", "").lower()

156 num_bits = quant_config.get("bits")

@@ -189,7 +189,7 @@ class AWQMarlinLinearMethod(LinearMethodBase):

189 self,

190 layer: torch.nn.Module,

191 input_size_per_partition: int,

192 - output_partition_sizes: List[int],

192 + output_partition_sizes: list[int],

193 input_size: int,

194 output_size: int,

195 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/base_config.py

+8 -8

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

2

3 import inspect

4 from abc import ABC, abstractmethod

5 -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type

5 +from typing import TYPE_CHECKING, Any, Optional

6

7 import torch

8 from torch import nn

@@ -48,7 +48,7 @@ class QuantizeMethodBase(ABC):

48

49

50 def method_has_implemented_embedding(

51 - method_class: Type[QuantizeMethodBase]) -> bool:

51 + method_class: type[QuantizeMethodBase]) -> bool:

52 """

53 Not all quant methods have embedding implemented, so we need to check that

54 it exists for our given method. We check this by making sure the function

@@ -68,7 +68,7 @@ class QuantizationConfig(ABC):

68 def __init__(self):

69 super().__init__()

70 # mapping is updated by models as they initialize

71 - self.packed_modules_mapping: Dict[str, List[str]] = dict()

71 + self.packed_modules_mapping: dict[str, list[str]] = dict()

72

73 @abstractmethod

74 def get_name(self) -> QuantizationMethods:

@@ -76,7 +76,7 @@ class QuantizationConfig(ABC):

76 raise NotImplementedError

77

78 @abstractmethod

79 - def get_supported_act_dtypes(self) -> List[torch.dtype]:

79 + def get_supported_act_dtypes(self) -> list[torch.dtype]:

80 """List of supported activation dtypes."""

81 raise NotImplementedError

82

@@ -93,13 +93,13 @@ class QuantizationConfig(ABC):

93

94 @staticmethod

95 @abstractmethod

96 - def get_config_filenames() -> List[str]:

96 + def get_config_filenames() -> list[str]:

97 """List of filenames to search for in the model directory."""

98 raise NotImplementedError

99

100 @classmethod

101 @abstractmethod

102 - def from_config(cls, config: Dict[str, Any]) -> "QuantizationConfig":

102 + def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig":

103 """Create a config class from the model's quantization config."""

104 raise NotImplementedError

105

@@ -115,7 +115,7 @@ class QuantizationConfig(ABC):

115 return None

116

117 @staticmethod

118 - def get_from_keys(config: Dict[str, Any], keys: List[str]) -> Any:

118 + def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:

119 """Get a value from the model's quantization config."""

120 for key in keys:

121 if key in config:

@@ -124,7 +124,7 @@ class QuantizationConfig(ABC):

124 "quantization config.")

125

126 @staticmethod

127 - def get_from_keys_or(config: Dict[str, Any], keys: List[str],

127 + def get_from_keys_or(config: dict[str, Any], keys: list[str],

128 default: Any) -> Any:

129 """Get a optional value from the model's quantization config."""

130 try:

vllm/model_executor/layers/quantization/bitblas.py

+8 -8

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

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Any, Dict, List, Optional

2 +from typing import Any, Optional

3

4 import torch

5

@@ -105,7 +105,7 @@ class BitBLASConfig(QuantizationConfig):

105 return "bitblas"

106

107 @classmethod

108 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

108 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

109 return [torch.half, torch.bfloat16]

110

111 @classmethod

@@ -114,12 +114,12 @@ class BitBLASConfig(QuantizationConfig):

114 return 70

115

116 @classmethod

117 - def get_config_filenames(cls) -> List[str]:

117 + def get_config_filenames(cls) -> list[str]:

118 return ["quantize_config.json"]

119

120 @staticmethod

121 - def get_from_keys(config: Dict[str, Any],

122 - keys: List[str],

121 + def get_from_keys(config: dict[str, Any],

122 + keys: list[str],

123 default: Any = None) -> Any:

124 """Get a value from the model's quantization config."""

125 for key in keys:

@@ -128,7 +128,7 @@ class BitBLASConfig(QuantizationConfig):

128 return default

129

130 @classmethod

131 - def from_config(cls, config: Dict[str, Any]) -> "BitBLASConfig":

131 + def from_config(cls, config: dict[str, Any]) -> "BitBLASConfig":

132 weight_bits = cls.get_from_keys(config, ["bits"])

133 group_size = cls.get_from_keys(config, ["group_size"], -1)

134 desc_act = cls.get_from_keys(config, ["desc_act"], False)

@@ -193,7 +193,7 @@ class BitBLASLinearMethod(LinearMethodBase):

193 self,

194 layer: torch.nn.Module,

195 input_size_per_partition: int,

196 - output_partition_sizes: List[int],

196 + output_partition_sizes: list[int],

197 input_size: int,

198 output_size: int,

199 params_dtype: torch.dtype,

@@ -329,7 +329,7 @@ class BitBLASLinearMethod(LinearMethodBase):

329 self,

330 layer: torch.nn.Module,

331 input_size_per_partition: int,

332 - output_partition_sizes: List[int],

332 + output_partition_sizes: list[int],

333 input_size: int,

334 output_size: int,

335 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/bitsandbytes.py

+7 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6

@@ -29,7 +29,7 @@ class BitsAndBytesConfig(QuantizationConfig):

29 bnb_4bit_use_double_quant: bool = False,

30 llm_int8_enable_fp32_cpu_offload: bool = False,

31 llm_int8_has_fp16_weight: bool = False,

32 - llm_int8_skip_modules: Optional[List[str]] = None,

32 + llm_int8_skip_modules: Optional[list[str]] = None,

33 llm_int8_threshold: float = 6.0,

34 ) -> None:

35 super().__init__()

@@ -61,7 +61,7 @@ class BitsAndBytesConfig(QuantizationConfig):

61 return "bitsandbytes"

62

63 @classmethod

64 - def get_supported_act_dtypes(self) -> List[torch.dtype]:

64 + def get_supported_act_dtypes(self) -> list[torch.dtype]:

65 return [torch.float32, torch.float16, torch.bfloat16]

66

67 @classmethod

@@ -69,13 +69,13 @@ class BitsAndBytesConfig(QuantizationConfig):

69 return 70

70

71 @staticmethod

72 - def get_config_filenames() -> List[str]:

72 + def get_config_filenames() -> list[str]:

73 return [

74 "adapter_config.json",

75 ]

76

77 @classmethod

78 - def from_config(cls, config: Dict[str, Any]) -> "BitsAndBytesConfig":

78 + def from_config(cls, config: dict[str, Any]) -> "BitsAndBytesConfig":

79

80 def get_safe_value(config, keys, default_value=None):

81 try:

@@ -130,7 +130,7 @@ class BitsAndBytesConfig(QuantizationConfig):

130 return None

131

132

133 -def is_layer_skipped_bnb(prefix: str, llm_int8_skip_modules: List[str]):

133 +def is_layer_skipped_bnb(prefix: str, llm_int8_skip_modules: list[str]):

134 # Split the prefix into its dot-separated components

135 components = prefix.split('.')

136

@@ -169,7 +169,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):

169

170 def create_weights(self, layer: torch.nn.Module,

171 input_size_per_partition: int,

172 - output_partition_sizes: List[int], input_size: int,

172 + output_partition_sizes: list[int], input_size: int,

173 output_size: int, params_dtype: torch.dtype,

174 **extra_weight_attrs):

175 from bitsandbytes.nn import Int8Params

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

+19 -19

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 from contextlib import suppress

4 -from typing import Any, Dict, List, Literal, Optional, Tuple, cast

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

5

6 import torch

7 from compressed_tensors.config import (CompressionFormat,

@@ -38,20 +38,20 @@ logger = init_logger(__name__)

38 __all__ = ["CompressedTensorsLinearMethod"]

39

40 SPARSITY_CONFIG_NAME: Literal["sparsity_config"] = "sparsity_config"

41 -QUANTIZATION_SCHEME_MAP_TYPE = Dict[str, Optional[Dict[str, QuantizationArgs]]]

41 +QUANTIZATION_SCHEME_MAP_TYPE = dict[str, Optional[dict[str, QuantizationArgs]]]

42

43

44 class CompressedTensorsConfig(QuantizationConfig):

45

46 def __init__(

47 self,

48 - target_scheme_map: Dict[str, Any],

49 - ignore: List[str],

48 + target_scheme_map: dict[str, Any],

49 + ignore: list[str],

50 quant_format: str,

51 - sparsity_scheme_map: Dict[str, SparsityCompressionConfig],

52 - sparsity_ignore_list: List[str],

53 - kv_cache_scheme: Optional[Dict[str, Any]] = None,

54 - config: Optional[Dict[str, Any]] = None,

51 + sparsity_scheme_map: dict[str, SparsityCompressionConfig],

52 + sparsity_ignore_list: list[str],

53 + kv_cache_scheme: Optional[dict[str, Any]] = None,

54 + config: Optional[dict[str, Any]] = None,

55 ):

56 super().__init__()

57 self.ignore = ignore

@@ -66,7 +66,7 @@ class CompressedTensorsConfig(QuantizationConfig):

66 def get_linear_method(self) -> "CompressedTensorsLinearMethod":

67 return CompressedTensorsLinearMethod(self)

68

69 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

69 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

70 return [torch.float16, torch.bfloat16]

71

72 @classmethod

@@ -102,8 +102,8 @@ class CompressedTensorsConfig(QuantizationConfig):

102 return None

103

104 @classmethod

105 - def from_config(cls, config: Dict[str, Any]) -> "CompressedTensorsConfig":

106 - ignore: List[str] = cast(List[str], config.get("ignore", []))

105 + def from_config(cls, config: dict[str, Any]) -> "CompressedTensorsConfig":

106 + ignore: list[str] = cast(list[str], config.get("ignore", []))

107 quant_format = cast(str, config.get("format"))

108 target_scheme_map = cls._quantization_scheme_map_from_config(

109 config=config)

@@ -121,8 +121,8 @@ class CompressedTensorsConfig(QuantizationConfig):

121

122 @classmethod

123 def _parse_sparsity_config(

124 - cls, config: Dict[str, Any]

125 - ) -> Tuple[Dict[str, SparsityCompressionConfig], List[str]]:

124 + cls, config: dict[str, Any]

125 + ) -> tuple[dict[str, SparsityCompressionConfig], list[str]]:

126 """

127 :param config: The `quantization_config` dictionary from config.json

128 :return: A tuple with two elements

@@ -135,7 +135,7 @@ class CompressedTensorsConfig(QuantizationConfig):

135

136 sparsity_config = SparsityCompressionConfig.model_validate(

137 sparsity_config)

138 - sparse_scheme_map: Dict[str, SparsityCompressionConfig] = {

138 + sparse_scheme_map: dict[str, SparsityCompressionConfig] = {

139 target: sparsity_config

140 for target in sparsity_config.targets or list()

141 }

@@ -144,13 +144,13 @@ class CompressedTensorsConfig(QuantizationConfig):

144

145 @classmethod

146 def _quantization_scheme_map_from_config(

147 - cls, config: Dict[str, Any]) -> QUANTIZATION_SCHEME_MAP_TYPE:

147 + cls, config: dict[str, Any]) -> QUANTIZATION_SCHEME_MAP_TYPE:

148 """

149 :param config: The `quantization_config` dictionary from config.json

150 :return: A dictionary mapping target layer names to their corresponding

151 quantization_args for weights and input activations

152 """

153 - target_scheme_map: Dict[str, Any] = dict()

153 + target_scheme_map: dict[str, Any] = dict()

154 quant_format = cast(str, config.get("format"))

155

156 # The quant_config has multiple config_groups, each containing

@@ -188,7 +188,7 @@ class CompressedTensorsConfig(QuantizationConfig):

188 return target_scheme_map

189

190 @classmethod

191 - def get_config_filenames(cls) -> List[str]:

191 + def get_config_filenames(cls) -> list[str]:

192 return []

193

194 def _check_scheme_supported(self,

@@ -565,7 +565,7 @@ class CompressedTensorsLinearMethod(LinearMethodBase):

565

566 def create_weights(self, layer: torch.nn.Module,

567 input_size_per_partition: int,

568 - output_partition_sizes: List[int], input_size: int,

568 + output_partition_sizes: list[int], input_size: int,

569 output_size: int, params_dtype: torch.dtype,

570 **extra_weight_attrs):

571 """

@@ -611,7 +611,7 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):

611 super().__init__(quant_config)

612

613 @staticmethod

614 - def validate_kv_cache_scheme(kv_cache_scheme: Optional[Dict[str, Any]]):

614 + def validate_kv_cache_scheme(kv_cache_scheme: Optional[dict[str, Any]]):

615 """

616 Validator for the kv cache scheme. Useful for controlling the

617 kv cache quantization schemes, that are being supported in vLLM

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_24.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Callable, Optional

4

5 import torch

6 from compressed_tensors import CompressionFormat, ModelCompressor

@@ -31,7 +31,7 @@ class CompressedTensors24(CompressedTensorsScheme):

31 quantized: bool = False,

32 weight_quant: Optional[QuantizationArgs] = None,

33 input_quant: Optional[QuantizationArgs] = None,

34 - model_compression_config: Optional[Dict[str, Any]] = None,

34 + model_compression_config: Optional[dict[str, Any]] = None,

35 ):

36 self.quantized = quantized

37 self.weight_quant = weight_quant

@@ -53,7 +53,7 @@ class CompressedTensors24(CompressedTensorsScheme):

53 self,

54 layer: torch.nn.Module,

55 input_size: int,

56 - output_partition_sizes: List[int],

56 + output_partition_sizes: list[int],

57 input_size_per_partition: int,

58 params_dtype: torch.dtype,

59 weight_loader: Callable,

@@ -327,9 +327,9 @@ class CompressedTensors24(CompressedTensorsScheme):

327 )

328 return sparsity_compressor.decompress_weight(weight_data)

329

330 - split_weights: List[torch.Tensor] = []

331 - split_bitmask: List[torch.Tensor] = []

332 - split_shape: List[Tuple[int, int]] = []

330 + split_weights: list[torch.Tensor] = []

331 + split_bitmask: list[torch.Tensor] = []

332 + split_shape: list[tuple[int, int]] = []

333

334 if isinstance(layer, (QKVParallelLinear, MergedColumnParallelLinear)):

335 split_weights = torch.split(compressed, layer.logical_widths)

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_24.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional

3 +from typing import Callable, Optional

4

5 import torch

6 from torch.nn import Parameter

@@ -58,7 +58,7 @@ class CompressedTensorsW4A16Sparse24(CompressedTensorsScheme):

58 layer.meta = Parameter(layer.meta.data, requires_grad=False)

59

60 def create_weights(self, layer: torch.nn.Module, input_size: int,

61 - output_partition_sizes: List[int],

61 + output_partition_sizes: list[int],

62 input_size_per_partition: int,

63 params_dtype: torch.dtype, weight_loader: Callable,

64 **kwargs):

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Callable, List, Optional

2 +from typing import Callable, Optional

3

4 import torch

5 from torch.nn.parameter import Parameter

@@ -26,7 +26,7 @@ class CompressedTensorsW4A16Fp4(CompressedTensorsScheme):

26 return 80

27

28 def create_weights(self, layer: torch.nn.Module,

29 - output_partition_sizes: List[int],

29 + output_partition_sizes: list[int],

30 input_size_per_partition: int,

31 params_dtype: torch.dtype, weight_loader: Callable,

32 **kwargs):

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional

3 +from typing import Callable, Optional

4

5 import torch

6 from compressed_tensors.quantization import QuantizationStrategy

@@ -58,7 +58,7 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):

58 prepare_fp8_layer_for_marlin(layer)

59

60 def create_weights(self, layer: torch.nn.Module, input_size: int,

61 - output_partition_sizes: List[int],

61 + output_partition_sizes: list[int],

62 input_size_per_partition: int,

63 params_dtype: torch.dtype, weight_loader: Callable,

64 **kwargs):

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional

3 +from typing import Callable, Optional

4

5 import torch

6 from compressed_tensors.quantization import QuantizationStrategy

@@ -90,7 +90,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

90 layer.input_scale = None

91

92 def create_weights(self, layer: torch.nn.Module,

93 - output_partition_sizes: List[int],

93 + output_partition_sizes: list[int],

94 input_size_per_partition: int,

95 params_dtype: torch.dtype, weight_loader: Callable,

96 **kwargs):

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional, Set

3 +from typing import Callable, Optional

4

5 import torch

6 from compressed_tensors.quantization import QuantizationStrategy

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

19

20

21 class CompressedTensorsW8A8Int8(CompressedTensorsScheme):

22 - _kernel_backends_being_used: Set[str] = set()

22 + _kernel_backends_being_used: set[str] = set()

23

24 def __init__(self, strategy: str, is_static_input_scheme: bool,

25 input_symmetric: bool):

@@ -33,7 +33,7 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):

33 return 75

34

35 def create_weights(self, layer: torch.nn.Module,

36 - output_partition_sizes: List[int],

36 + output_partition_sizes: list[int],

37 input_size_per_partition: int,

38 params_dtype: torch.dtype, weight_loader: Callable,

39 **kwargs):

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional, Set

3 +from typing import Callable, Optional

4

5 import torch

6 from compressed_tensors.quantization import ActivationOrdering

@@ -35,7 +35,7 @@ WNA16_SUPPORTED_BITS = list(WNA16_SUPPORTED_TYPES_MAP.keys())

35

36

37 class CompressedTensorsWNA16(CompressedTensorsScheme):

38 - _kernel_backends_being_used: Set[str] = set()

38 + _kernel_backends_being_used: set[str] = set()

39

40 def __init__(self,

41 strategy: str,

@@ -70,7 +70,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):

70 return 80

71

72 def create_weights(self, layer: torch.nn.Module, output_size: int,

73 - input_size: int, output_partition_sizes: List[int],

73 + input_size: int, output_partition_sizes: list[int],

74 input_size_per_partition: int,

75 params_dtype: torch.dtype, weight_loader: Callable,

76 **kwargs):

vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Type

3 +from typing import Optional

4

5 import torch

6

@@ -126,7 +126,7 @@ def triton_scaled_mm(input: torch.Tensor,

126 weight: torch.Tensor,

127 scale_a: torch.Tensor,

128 scale_b: torch.Tensor,

129 - out_dtype: Type[torch.dtype],

129 + out_dtype: type[torch.dtype],

130 bias: Optional[torch.Tensor] = None,

131 block_size_m: int = 32,

132 block_size_n: int = 32,

vllm/model_executor/layers/quantization/compressed_tensors/utils.py

+6 -5

@@ -1,8 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 import re

4 +from collections.abc import Iterable, Mapping

5 from types import MappingProxyType

5 -from typing import Iterable, List, Mapping, Optional

6 +from typing import Optional

7

8 from compressed_tensors import CompressionFormat

9 from torch.nn import Module

@@ -20,7 +21,7 @@ def is_activation_quantization_format(format: str) -> bool:

21 def should_ignore_layer(

22 layer_name: Optional[str],

23 ignore: Iterable[str] = tuple(),

23 - fused_mapping: Mapping[str, List[str]] = MappingProxyType({})

24 + fused_mapping: Mapping[str, list[str]] = MappingProxyType({})

25 ) -> bool:

26 if layer_name is None:

27 return False

@@ -84,7 +85,7 @@ def find_matched_target(

85 layer_name: Optional[str],

86 module: Module,

87 targets: Iterable[str],

87 - fused_mapping: Mapping[str, List[str]] = MappingProxyType({})

88 + fused_mapping: Mapping[str, list[str]] = MappingProxyType({})

89 ) -> str:

90 """

91 Helper function to look up which "target" in the compressed-tensors

@@ -171,7 +172,7 @@ def _is_equal_or_regex_match(value: str,

172

173 def _match_fused_layer(

174 layer_name: str, target_layers: Iterable[str],

174 - fused_mapping: Mapping[str, List[str]]) -> Optional[str]:

175 + fused_mapping: Mapping[str, list[str]]) -> Optional[str]:

176 """

177 Match a fused layer name to its corresponding individual layer in

178 target_layers. Returns first value in fused_mapping which matches targets

@@ -201,7 +202,7 @@ def _match_fused_layer(

202 ]

203

204 # for each unfused component, find a match in targets

204 - unfused_matches: List[Optional[str]] = []

205 + unfused_matches: list[Optional[str]] = []

206 for unfused in unfused_paths:

207 for target in target_layers:

208 if _is_equal_or_regex_match(unfused, target):

vllm/model_executor/layers/quantization/deepspeedfp.py

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6 import torch.nn as nn

@@ -46,7 +46,7 @@ class DeepSpeedFPConfig(QuantizationConfig):

46 return "deepspeedfp"

47

48 @classmethod

49 - def from_config(cls, config: Dict[str, Any]) -> "DeepSpeedFPConfig":

49 + def from_config(cls, config: dict[str, Any]) -> "DeepSpeedFPConfig":

50 weight_bits = cls.get_from_keys(config, ["bits"])

51 group_size = cls.get_from_keys(config, ["group_size"])

52 return cls(weight_bits=weight_bits, group_size=group_size)

@@ -55,7 +55,7 @@ class DeepSpeedFPConfig(QuantizationConfig):

55 return DeepSpeedFPLinearMethod(self)

56

57 @classmethod

58 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

58 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

59 return [torch.half, torch.bfloat16]

60

61 @classmethod

@@ -64,7 +64,7 @@ class DeepSpeedFPConfig(QuantizationConfig):

64 return 60

65

66 @staticmethod

67 - def get_config_filenames() -> List[str]:

67 + def get_config_filenames() -> list[str]:

68 return [

69 "quant_config.json",

70 "quantize_config.json",

@@ -91,7 +91,7 @@ class DeepSpeedFPLinearMethod(LinearMethodBase):

91 def create_weights(self,

92 layer: torch.nn.Module,

93 input_size_per_partition: int,

94 - output_partition_sizes: List[int],

94 + output_partition_sizes: list[int],

95 input_size: int,

96 output_size: int,

97 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/experts_int8.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, List, Optional

3 +from typing import Any, Callable, Optional

4

5 import torch

6

@@ -25,7 +25,7 @@ class ExpertsInt8Config(QuantizationConfig):

25 return "experts_int8"

26

27 @classmethod

28 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

28 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

29 return [torch.bfloat16, torch.half]

30

31 @classmethod

@@ -33,11 +33,11 @@ class ExpertsInt8Config(QuantizationConfig):

33 return 80

34

35 @classmethod

36 - def get_config_filenames(cls) -> List[str]:

36 + def get_config_filenames(cls) -> list[str]:

37 return []

38

39 @classmethod

40 - def from_config(cls, config: Dict[str, Any]) -> "ExpertsInt8Config":

40 + def from_config(cls, config: dict[str, Any]) -> "ExpertsInt8Config":

41 return cls()

42

43 def get_quant_method(self, layer: torch.nn.Module,

vllm/model_executor/layers/quantization/fbgemm_fp8.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6 from torch.nn import Module

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

28 class FBGEMMFp8Config(QuantizationConfig):

29 """Config class for FBGEMM Fp8."""

30

31 - def __init__(self, ignore_list: List[str], input_scale_ub: float):

31 + def __init__(self, ignore_list: list[str], input_scale_ub: float):

32 super().__init__()

33 self.ignore_list = ignore_list if ignore_list else []

34 self.input_scale_ub = input_scale_ub

@@ -43,7 +43,7 @@ class FBGEMMFp8Config(QuantizationConfig):

43 return "fbgemm_fp8"

44

45 @classmethod

46 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

46 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

47 return [torch.bfloat16, torch.float16]

48

49 @classmethod

@@ -51,11 +51,11 @@ class FBGEMMFp8Config(QuantizationConfig):

51 return 80

52

53 @classmethod

54 - def get_config_filenames(cls) -> List[str]:

54 + def get_config_filenames(cls) -> list[str]:

55 return []

56

57 @classmethod

58 - def from_config(cls, config: Dict[str, Any]) -> "FBGEMMFp8Config":

58 + def from_config(cls, config: dict[str, Any]) -> "FBGEMMFp8Config":

59 ignore_list = cls.get_from_keys(config, ["modules_to_not_convert"])

60 input_scale_ub = cls.get_from_keys(config, ["activation_scale_ub"])

61 return cls(ignore_list=ignore_list, input_scale_ub=input_scale_ub)

@@ -82,7 +82,7 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):

82 self,

83 layer: torch.nn.Module,

84 input_size_per_partition: int,

85 - output_partition_sizes: List[int],

85 + output_partition_sizes: list[int],

86 input_size: int,

87 output_size: int,

88 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/fp8.py

+7 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 import importlib.util

4 -from typing import Any, Callable, Dict, List, Optional

4 +from typing import Any, Callable, Optional

5

6 import torch

7 import torch.nn.functional as F

@@ -57,8 +57,8 @@ class Fp8Config(QuantizationConfig):

57 self,

58 is_checkpoint_fp8_serialized: bool = False,

59 activation_scheme: str = "dynamic",

60 - ignored_layers: Optional[List[str]] = None,

61 - weight_block_size: Optional[List[int]] = None,

60 + ignored_layers: Optional[list[str]] = None,

61 + weight_block_size: Optional[list[int]] = None,

62 ) -> None:

63 super().__init__()

64 self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized

@@ -90,7 +90,7 @@ class Fp8Config(QuantizationConfig):

90 return "fp8"

91

92 @classmethod

93 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

93 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

94 return [torch.bfloat16, torch.half]

95

96 @classmethod

@@ -98,11 +98,11 @@ class Fp8Config(QuantizationConfig):

98 return 80

99

100 @classmethod

101 - def get_config_filenames(cls) -> List[str]:

101 + def get_config_filenames(cls) -> list[str]:

102 return []

103

104 @classmethod

105 - def from_config(cls, config: Dict[str, Any]) -> "Fp8Config":

105 + def from_config(cls, config: dict[str, Any]) -> "Fp8Config":

106 quant_method = cls.get_from_keys(config, ["quant_method"])

107 is_checkpoint_fp8_serialized = ("fp8" in quant_method)

108 activation_scheme = cls.get_from_keys(config, ["activation_scheme"])

@@ -191,7 +191,7 @@ class Fp8LinearMethod(LinearMethodBase):

191 self,

192 layer: torch.nn.Module,

193 input_size_per_partition: int,

194 - output_partition_sizes: List[int],

194 + output_partition_sizes: list[int],

195 input_size: int,

196 output_size: int,

197 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/gguf.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, List, Optional

3 +from typing import Any, Callable, Optional

4

5 import gguf

6 import torch

@@ -35,7 +35,7 @@ class GGUFConfig(QuantizationConfig):

35 def get_name(self) -> QuantizationMethods:

36 return "gguf"

37

38 - def get_supported_act_dtypes(self) -> List[torch.dtype]:

38 + def get_supported_act_dtypes(self) -> list[torch.dtype]:

39 return [torch.half, torch.bfloat16, torch.float32]

40

41 @classmethod

@@ -43,11 +43,11 @@ class GGUFConfig(QuantizationConfig):

43 return 60

44

45 @classmethod

46 - def get_config_filenames(cls) -> List[str]:

46 + def get_config_filenames(cls) -> list[str]:

47 return [] # no extra configs.

48

49 @classmethod

50 - def from_config(cls, config: Dict[str, Any]) -> "GGUFConfig":

50 + def from_config(cls, config: dict[str, Any]) -> "GGUFConfig":

51 return cls()

52

53 def get_quant_method(self, layer: torch.nn.Module,

@@ -215,7 +215,7 @@ class GGUFLinearMethod(LinearMethodBase):

215

216 def create_weights(self, layer: torch.nn.Module,

217 input_size_per_partition: int,

218 - output_partition_sizes: List[int], input_size: int,

218 + output_partition_sizes: list[int], input_size: int,

219 output_size: int, params_dtype: torch.dtype,

220 **extra_weight_attrs):

221 self.params_dtype = params_dtype

@@ -406,7 +406,7 @@ class GGUFEmbeddingMethod(GGUFLinearMethod):

406

407 class GGUFUninitializedParameter(UninitializedParameter):

408 cls_to_become = Parameter

409 - data_container: List[torch.Tensor]

409 + data_container: list[torch.Tensor]

410

411 def materialize_nested(self) -> Parameter:

412 dtype = {data.dtype for data in self.data_container}

vllm/model_executor/layers/quantization/gptq.py

+7 -7

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

3 import enum

4 from enum import Enum

5 from fractions import Fraction

6 -from typing import Any, Dict, List, Optional, Union

6 +from typing import Any, Optional, Union

7

8 import torch

9 from torch.nn.parameter import Parameter

@@ -34,11 +34,11 @@ class GPTQConfig(QuantizationConfig):

34 group_size: int,

35 desc_act: bool,

36 lm_head_quantized: bool,

37 - dynamic: Dict[str, Dict[str, Union[int, bool]]],

37 + dynamic: dict[str, dict[str, Union[int, bool]]],

38 ) -> None:

39 # GPTQModel use `dynamic` config property to allow per module

40 # quantization config so each module can be individually optimized.

41 - # Format is Dict[str, Dict] where key is a regex string that can

41 + # Format is dict[str, dict] where key is a regex string that can

42 # perform both positive ("+:" prefixed) or negative ("-:" prefixed)

43 # matching of a module.

44 # Default to positive match, override base quant config mode, if no

@@ -84,7 +84,7 @@ class GPTQConfig(QuantizationConfig):

84 return "gptq"

85

86 @classmethod

87 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

87 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

88 return [torch.half]

89

90 @classmethod

@@ -93,11 +93,11 @@ class GPTQConfig(QuantizationConfig):

93 return 60

94

95 @classmethod

96 - def get_config_filenames(cls) -> List[str]:

96 + def get_config_filenames(cls) -> list[str]:

97 return ["quantize_config.json"]

98

99 @classmethod

100 - def from_config(cls, config: Dict[str, Any]) -> "GPTQConfig":

100 + def from_config(cls, config: dict[str, Any]) -> "GPTQConfig":

101 dynamic = cls.get_from_keys_or(config, ["dynamic"], default={})

102 dynamic = {} if dynamic is None else dynamic

103

@@ -135,7 +135,7 @@ class GPTQLinearMethod(LinearMethodBase):

135 self,

136 layer: torch.nn.Module,

137 input_size_per_partition: int,

138 - output_partition_sizes: List[int],

138 + output_partition_sizes: list[int],

139 input_size: int,

140 output_size: int,

141 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/gptq_bitblas.py

+7 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Any, Dict, List, Optional, Set

2 +from typing import Any, Optional

3

4 import torch

5 from torch.nn.parameter import Parameter

@@ -129,7 +129,7 @@ class GPTQBitBLASConfig(QuantizationConfig):

129 return "gptq_bitblas"

130

131 @classmethod

132 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

132 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

133 return [torch.half, torch.bfloat16]

134

135 @classmethod

@@ -137,11 +137,11 @@ class GPTQBitBLASConfig(QuantizationConfig):

137 return 80

138

139 @classmethod

140 - def get_config_filenames(cls) -> List[str]:

140 + def get_config_filenames(cls) -> list[str]:

141 return ["quantize_config.json"]

142

143 @classmethod

144 - def from_config(cls, config: Dict[str, Any]) -> "GPTQBitBLASConfig":

144 + def from_config(cls, config: dict[str, Any]) -> "GPTQBitBLASConfig":

145 weight_bits = cls.get_from_keys(config, ["bits"])

146 group_size = cls.get_from_keys(config, ["group_size"])

147 desc_act = cls.get_from_keys(config, ["desc_act"])

@@ -185,7 +185,7 @@ class GPTQBitBLASConfig(QuantizationConfig):

185 return self.TORCH_BITBLAS_STORAGE_DTYPE

186

187 @classmethod

188 - def is_gptq_bitblas_compatible(cls, quant_config: Dict[str, Any]):

188 + def is_gptq_bitblas_compatible(cls, quant_config: dict[str, Any]):

189 # Extract data from quant config.

190 num_bits = quant_config.get("bits")

191 group_size = quant_config.get("group_size")

@@ -224,7 +224,7 @@ class GPTQBitBLASLinearMethod(LinearMethodBase):

224 """

225

226 kernel_type = BitBLASLinearKernel

227 - _kernel_backends_being_used: Set[str] = set()

227 + _kernel_backends_being_used: set[str] = set()

228

229 def __init__(self, quant_config: GPTQBitBLASConfig) -> None:

230 self.quant_config = quant_config

@@ -236,7 +236,7 @@ class GPTQBitBLASLinearMethod(LinearMethodBase):

236 self,

237 layer: torch.nn.Module,

238 input_size_per_partition: int,

239 - output_partition_sizes: List[int],

239 + output_partition_sizes: list[int],

240 input_size: int,

241 output_size: int,

242 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/gptq_marlin.py

+10 -10

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, List, Optional, Set, Union

3 +from typing import Any, Callable, Optional, Union

4

5 import torch

6

@@ -45,8 +45,8 @@ class GPTQMarlinConfig(QuantizationConfig):

45

46 def __init__(self, weight_bits: int, group_size: int, desc_act: bool,

47 is_sym: bool, lm_head_quantized: bool,

48 - dynamic: Dict[str, Dict[str, Union[int, bool]]],

49 - full_config: Dict[str, Any]) -> None:

48 + dynamic: dict[str, dict[str, Union[int, bool]]],

49 + full_config: dict[str, Any]) -> None:

50 super().__init__()

51 if desc_act and group_size == -1:

52 # In this case, act_order == True is the same as act_order == False

@@ -55,7 +55,7 @@ class GPTQMarlinConfig(QuantizationConfig):

55

56 # GPTQModel use `dynamic` config property to allow per module

57 # quantization config so each module can be individually optimized.

58 - # Format is Dict[str, Dict] where key is a regex string that can

58 + # Format is dict[str, dict] where key is a regex string that can

59 # perform both positive ("+:" prefixed) or negative ("-:" prefixed)

60 # matching of a module.

61 # Default to positive match, override base quant config mode, if no

@@ -105,7 +105,7 @@ class GPTQMarlinConfig(QuantizationConfig):

105 return "gptq_marlin"

106

107 @classmethod

108 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

108 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

109 return [torch.half, torch.bfloat16]

110

111 @classmethod

@@ -113,11 +113,11 @@ class GPTQMarlinConfig(QuantizationConfig):

113 return 80

114

115 @classmethod

116 - def get_config_filenames(cls) -> List[str]:

116 + def get_config_filenames(cls) -> list[str]:

117 return ["quantize_config.json"]

118

119 @classmethod

120 - def from_config(cls, config: Dict[str, Any]) -> "GPTQMarlinConfig":

120 + def from_config(cls, config: dict[str, Any]) -> "GPTQMarlinConfig":

121 dynamic = cls.get_from_keys_or(config, ["dynamic"], default={})

122 dynamic = {} if dynamic is None else dynamic

123

@@ -167,7 +167,7 @@ class GPTQMarlinConfig(QuantizationConfig):

167 GPTQMarlinLinearMethod)

168

169 @classmethod

170 - def is_gptq_marlin_compatible(cls, quant_config: Dict[str, Any]):

170 + def is_gptq_marlin_compatible(cls, quant_config: dict[str, Any]):

171 quant_method = quant_config.get("quant_method", "").lower()

172 num_bits = quant_config.get("bits")

173 group_size = quant_config.get("group_size")

@@ -199,7 +199,7 @@ class GPTQMarlinLinearMethod(LinearMethodBase):

199 quant_config: The GPTQ Marlin quantization config.

200 """

201

202 - _kernel_backends_being_used: Set[str] = set()

202 + _kernel_backends_being_used: set[str] = set()

203

204 def __init__(self, quant_config: GPTQMarlinConfig) -> None:

205 self.quant_config = quant_config

@@ -212,7 +212,7 @@ class GPTQMarlinLinearMethod(LinearMethodBase):

212 self,

213 layer: torch.nn.Module,

214 input_size_per_partition: int,

215 - output_partition_sizes: List[int],

215 + output_partition_sizes: list[int],

216 input_size: int,

217 output_size: int,

218 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/gptq_marlin_24.py

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6 from torch.nn.parameter import Parameter

@@ -90,7 +90,7 @@ class GPTQMarlin24Config(QuantizationConfig):

90 return "gptq_marlin_24"

91

92 @classmethod

93 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

93 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

94 return [torch.half]

95

96 @classmethod

@@ -99,11 +99,11 @@ class GPTQMarlin24Config(QuantizationConfig):

99 return 80

100

101 @classmethod

102 - def get_config_filenames(cls) -> List[str]:

102 + def get_config_filenames(cls) -> list[str]:

103 return ["quantize_config.json"]

104

105 @classmethod

106 - def from_config(cls, config: Dict[str, Any]) -> "GPTQMarlin24Config":

106 + def from_config(cls, config: dict[str, Any]) -> "GPTQMarlin24Config":

107 weight_bits = cls.get_from_keys(config, ["bits"])

108 group_size = cls.get_from_keys(config, ["group_size"])

109 return cls(weight_bits, group_size)

@@ -146,7 +146,7 @@ class GPTQMarlin24LinearMethod(LinearMethodBase):

146 self,

147 layer: torch.nn.Module,

148 input_size_per_partition: int,

149 - output_partition_sizes: List[int],

149 + output_partition_sizes: list[int],

150 input_size: int,

151 output_size: int,

152 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/hqq_marlin.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6

@@ -32,7 +32,7 @@ class HQQMarlinConfig(QuantizationConfig):

32 self,

33 weight_bits: int,

34 group_size: int,

35 - skip_modules: Optional[List[str]] = None,

35 + skip_modules: Optional[list[str]] = None,

36 ) -> None:

37 super().__init__()

38 assert group_size == 64, ("The only supported HQQ group size is "

@@ -55,7 +55,7 @@ class HQQMarlinConfig(QuantizationConfig):

55 return "hqq"

56

57 @classmethod

58 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

58 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

59 return [torch.half, torch.bfloat16]

60

61 @classmethod

@@ -63,11 +63,11 @@ class HQQMarlinConfig(QuantizationConfig):

63 return 80

64

65 @classmethod

66 - def get_config_filenames(cls) -> List[str]:

66 + def get_config_filenames(cls) -> list[str]:

67 return ["quantize_config.json"]

68

69 @classmethod

70 - def from_config(cls, config: Dict[str, Any]) -> "HQQMarlinConfig":

70 + def from_config(cls, config: dict[str, Any]) -> "HQQMarlinConfig":

71 wq_params = (config["quant_config"]["weight_quant_params"])

72 weight_bits = cls.get_from_keys(wq_params, ["nbits"])

73 group_size = cls.get_from_keys(wq_params, ["group_size"])

@@ -192,7 +192,7 @@ class HQQMarlinMethod(LinearMethodBase):

192 self,

193 layer: torch.nn.Module,

194 input_size_per_partition: int,

195 - output_partition_sizes: List[int],

195 + output_partition_sizes: list[int],

196 input_size: int,

197 output_size: int,

198 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/ipex_quant.py

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6

@@ -32,7 +32,7 @@ class IPEXConfig(QuantizationConfig):

32 method: str,

33 weight_bits: int,

34 group_size: int,

35 - modules_to_not_convert: Optional[List[str]] = None,

35 + modules_to_not_convert: Optional[list[str]] = None,

36 desc_act: Optional[bool] = None,

37 lm_head_quantized: Optional[bool] = None,

38 ) -> None:

@@ -63,7 +63,7 @@ class IPEXConfig(QuantizationConfig):

63 return "ipex"

64

65 @classmethod

66 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

66 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

67 return [torch.bfloat16, torch.float16]

68

69 @classmethod

@@ -71,14 +71,14 @@ class IPEXConfig(QuantizationConfig):

71 return -1

72

73 @staticmethod

74 - def get_config_filenames() -> List[str]:

74 + def get_config_filenames() -> list[str]:

75 return [

76 "quant_config.json",

77 "quantize_config.json",

78 ]

79

80 @classmethod

81 - def from_config(cls, config: Dict[str, Any]) -> "IPEXConfig":

81 + def from_config(cls, config: dict[str, Any]) -> "IPEXConfig":

82 method = cls.get_from_keys(config, ["quant_method"]).lower()

83 if method == "awq":

84 weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])

vllm/model_executor/layers/quantization/kernels/mixed_precision/MPLinearKernel.py

+5 -5

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

2

3 from abc import ABC, abstractmethod

4 from dataclasses import dataclass

5 -from typing import Callable, Optional, Tuple

5 +from typing import Callable, Optional

6

7 import torch

8

@@ -12,8 +12,8 @@ from vllm.scalar_type import ScalarType

12

13 @dataclass

14 class MPLinearLayerConfig:

15 - full_weight_shape: Tuple[int, int] # [in, out]

16 - partition_weight_shape: Tuple[int, int]

15 + full_weight_shape: tuple[int, int] # [in, out]

16 + partition_weight_shape: tuple[int, int]

17 weight_type: ScalarType

18 act_type: torch.dtype

19 group_size: int

@@ -31,7 +31,7 @@ class MPLinearKernel(ABC):

31 @classmethod

32 @abstractmethod

33 def can_implement(cls,

34 - c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:

34 + c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

35 raise NotImplementedError

36

37 def __init__(self,

@@ -75,7 +75,7 @@ class MPLinearKernel(ABC):

75 torch.nn.Parameter(new_param.data, requires_grad=False))

76

77 def _get_weight_params(

78 - self, layer: torch.nn.Module) -> Tuple[

78 + self, layer: torch.nn.Module) -> tuple[

79 torch.Tensor, # w_q

80 torch.Tensor, # w_s

81 Optional[torch.Tensor], # w_zp,

vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import List, Optional, Type

3 +from typing import Optional

4

5 import vllm.envs as envs

6 from vllm.model_executor.layers.quantization.kernels.mixed_precision.allspark import ( # noqa: E501

@@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.kernels.mixed_precision.MPLinearKer

18 from vllm.platforms import current_platform

19

20 # in priority/performance order (when available)

21 -_POSSIBLE_KERNELS: List[Type[MPLinearKernel]] = [

21 +_POSSIBLE_KERNELS: list[type[MPLinearKernel]] = [

22 MacheteLinearKernel,

23 AllSparkLinearKernel,

24 MarlinLinearKernel,

@@ -29,7 +29,7 @@ _POSSIBLE_KERNELS: List[Type[MPLinearKernel]] = [

29

30 def choose_mp_linear_kernel(

31 config: MPLinearLayerConfig,

32 - compute_capability: Optional[int] = None) -> Type[MPLinearKernel]:

32 + compute_capability: Optional[int] = None) -> type[MPLinearKernel]:

33 """

34 Choose an MPLinearKernel that can implement the given config for the given

35 compute capability. Attempts to choose the best kernel in terms of

@@ -46,7 +46,7 @@ def choose_mp_linear_kernel(

46 ValueError: If no kernel can implement the given config.

47

48 Returns:

49 - Type[MPLinearKernel]: Chosen kernel.

49 + type[MPLinearKernel]: Chosen kernel.

50 """

51 if compute_capability is None:

52 if current_platform is None:

vllm/model_executor/layers/quantization/kernels/mixed_precision/allspark.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -22,7 +22,7 @@ class AllSparkLinearKernel(MPLinearKernel):

22

23 @classmethod

24 def can_implement(cls,

25 - c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:

25 + c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

26 if c.has_g_idx:

27 return False, "Act reordering currently not supported by AllSpark"

28

vllm/model_executor/layers/quantization/kernels/mixed_precision/bitblas.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Optional

4

5 import torch

6

@@ -21,10 +21,10 @@ logger = init_logger(__name__)

21

22 class BitBLASLinearKernel(MPLinearKernel):

23

24 - OPT_FEATURES: List[int] = BITBLAS_OPTIMIZE_FEATURES

24 + OPT_FEATURES: list[int] = BITBLAS_OPTIMIZE_FEATURES

25 ENABLE_TUNING: bool = True

26 MATMUL_LAYOUT: str = "nt"

27 - BITBLAS_DTYPES: Dict[torch.dtype, str] = {

27 + BITBLAS_DTYPES: dict[torch.dtype, str] = {

28 torch.float32: "float32",

29 torch.float16: "float16",

30 torch.bfloat16: "bfloat16",

@@ -103,7 +103,7 @@ class BitBLASLinearKernel(MPLinearKernel):

103

104 @classmethod

105 def can_implement(cls,

106 - c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:

106 + c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

107

108 is_bitblas_installed = True

109

vllm/model_executor/layers/quantization/kernels/mixed_precision/exllama.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -25,7 +25,7 @@ class ExllamaLinearKernel(MPLinearKernel):

25

26 @classmethod

27 def can_implement(cls,

28 - c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:

28 + c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

29 if c.has_g_idx and\

30 c.partition_weight_shape[0] != c.full_weight_shape[0]:

31 return False, "Act reordering currently not supported by Exllama, "\

vllm/model_executor/layers/quantization/kernels/mixed_precision/machete.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 from functools import partial

4 -from typing import Optional, Tuple

4 +from typing import Optional

5

6 import torch

7

@@ -25,7 +25,7 @@ class MacheteLinearKernel(MPLinearKernel):

25

26 @classmethod

27 def can_implement(cls,

28 - c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:

28 + c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

29

30 if c.has_g_idx and\

31 c.partition_weight_shape[0] != c.full_weight_shape[0]:

vllm/model_executor/layers/quantization/kernels/mixed_precision/marlin.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -24,7 +24,7 @@ class MarlinLinearKernel(MPLinearKernel):

24

25 @classmethod

26 def can_implement(cls,

27 - c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:

27 + c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

28

29 quant_types = query_marlin_supported_quant_types(c.zero_points)

30 if c.weight_type not in quant_types:

vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py

+3 -3

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

2

3 from abc import ABC, abstractmethod

4 from dataclasses import dataclass

5 -from typing import Optional, Tuple

5 +from typing import Optional

6

7 import torch

8

@@ -24,7 +24,7 @@ class ScaledMMLinearKernel(ABC):

24 @classmethod

25 @abstractmethod

26 def can_implement(

27 - cls, c: ScaledMMLinearLayerConfig) -> Tuple[bool, Optional[str]]:

27 + cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

28 raise NotImplementedError

29

30 def __init__(self, c: ScaledMMLinearLayerConfig, w_q_param_name: str,

@@ -50,7 +50,7 @@ class ScaledMMLinearKernel(ABC):

50 raise NotImplementedError

51

52 def _get_weight_params(

53 - self, layer: torch.nn.Module) -> Tuple[

53 + self, layer: torch.nn.Module) -> tuple[

54 torch.Tensor, # weight

55 torch.Tensor, # weight_scale

56 Optional[torch.Tensor], # input_scale,

vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 import os

4 -from typing import Dict, List, Optional, Type

4 +from typing import Optional

5

6 from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (

7 AiterScaledMMLinearKernel)

@@ -16,7 +16,7 @@ from vllm.model_executor.layers.quantization.kernels.scaled_mm.xla import (

16 from vllm.platforms import PlatformEnum, current_platform

17

18 # in priority/performance order (when available)

19 -_POSSIBLE_KERNELS: Dict[PlatformEnum, List[Type[ScaledMMLinearKernel]]] = {

19 +_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[ScaledMMLinearKernel]]] = {

20 PlatformEnum.CPU: [CutlassScaledMMLinearKernel],

21 PlatformEnum.CUDA: [CutlassScaledMMLinearKernel],

22 PlatformEnum.ROCM: [AiterScaledMMLinearKernel, TritonScaledMMLinearKernel],

@@ -27,7 +27,7 @@ _POSSIBLE_KERNELS: Dict[PlatformEnum, List[Type[ScaledMMLinearKernel]]] = {

27 def choose_scaled_mm_linear_kernel(

28 config: ScaledMMLinearLayerConfig,

29 compute_capability: Optional[int] = None

30 -) -> Type[ScaledMMLinearKernel]:

30 +) -> type[ScaledMMLinearKernel]:

31 """

32 Choose an ScaledMMLinearKernel that can implement the given config for the

33 given compute capability. Attempts to choose the best kernel in terms of

@@ -44,7 +44,7 @@ def choose_scaled_mm_linear_kernel(

44 ValueError: If no kernel can implement the given config.

45

46 Returns:

47 - Type[ScaledMMLinearKernel]: Chosen kernel.

47 + type[ScaledMMLinearKernel]: Chosen kernel.

48 """

49

50 if compute_capability is None:

vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -20,7 +20,7 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):

20

21 @classmethod

22 def can_implement(

23 - cls, c: ScaledMMLinearLayerConfig) -> Tuple[bool, Optional[str]]:

23 + cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

24 if not current_platform.is_rocm():

25 return (

26 False,

vllm/model_executor/layers/quantization/kernels/scaled_mm/cutlass.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -22,7 +22,7 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):

22

23 @classmethod

24 def can_implement(

25 - cls, c: ScaledMMLinearLayerConfig) -> Tuple[bool, Optional[str]]:

25 + cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

26

27 if (not current_platform.is_cuda() and not current_platform.is_cpu()):

28 return False, "CutlassScaledMM requires running on CUDA or CPU."

vllm/model_executor/layers/quantization/kernels/scaled_mm/triton.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -18,7 +18,7 @@ class TritonScaledMMLinearKernel(CutlassScaledMMLinearKernel):

18

19 @classmethod

20 def can_implement(

21 - cls, c: ScaledMMLinearLayerConfig) -> Tuple[bool, Optional[str]]:

21 + cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

22 if current_platform.is_cpu():

23 return (

24 False,

vllm/model_executor/layers/quantization/kernels/scaled_mm/xla.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 import warnings

4 -from typing import Optional, Tuple

4 +from typing import Optional

5

6 import torch

7 from functorch.experimental.control_flow import cond # noqa: F401

@@ -25,7 +25,7 @@ class XLAScaledMMLinearKernel(ScaledMMLinearKernel):

25

26 @classmethod

27 def can_implement(

28 - cls, c: ScaledMMLinearLayerConfig) -> Tuple[bool, Optional[str]]:

28 + cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

29

30 if not current_platform.is_tpu():

31 return False, "ScaledMMXLA requires running on TPU."

vllm/model_executor/layers/quantization/marlin.py

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6 from torch.nn.parameter import Parameter

@@ -68,7 +68,7 @@ class MarlinConfig(QuantizationConfig):

68 return "marlin"

69

70 @classmethod

71 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

71 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

72 return [torch.half]

73

74 @classmethod

@@ -77,11 +77,11 @@ class MarlinConfig(QuantizationConfig):

77 return 80

78

79 @classmethod

80 - def get_config_filenames(cls) -> List[str]:

80 + def get_config_filenames(cls) -> list[str]:

81 return ["quantize_config.json"]

82

83 @classmethod

84 - def from_config(cls, config: Dict[str, Any]) -> "MarlinConfig":

84 + def from_config(cls, config: dict[str, Any]) -> "MarlinConfig":

85 group_size = cls.get_from_keys(config, ["group_size"])

86 lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],

87 default=False)

@@ -128,7 +128,7 @@ class MarlinLinearMethod(LinearMethodBase):

128 self,

129 layer: torch.nn.Module,

130 input_size_per_partition: int,

131 - output_partition_sizes: List[int],

131 + output_partition_sizes: list[int],

132 input_size: int,

133 output_size: int,

134 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/modelopt.py

+11 -11

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, List, Optional, Union

3 +from typing import Any, Callable, Optional, Union

4

5 import torch

6 from torch.nn import Module

@@ -53,7 +53,7 @@ class ModelOptFp8Config(QuantizationConfig):

53 return "modelopt"

54

55 @classmethod

56 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

56 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

57 return [torch.bfloat16, torch.half]

58

59 @classmethod

@@ -61,11 +61,11 @@ class ModelOptFp8Config(QuantizationConfig):

61 return 89

62

63 @classmethod

64 - def get_config_filenames(cls) -> List[str]:

64 + def get_config_filenames(cls) -> list[str]:

65 return ["hf_quant_config.json"]

66

67 @classmethod

68 - def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp8Config":

68 + def from_config(cls, config: dict[str, Any]) -> "ModelOptFp8Config":

69 quant_config = cls.get_from_keys(config, ["quantization"])

70 quant_method = quant_config["quant_algo"]

71 if quant_method not in QUANT_ALGOS:

@@ -107,7 +107,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):

107 self,

108 layer: torch.nn.Module,

109 input_size_per_partition: int,

110 - output_partition_sizes: List[int],

110 + output_partition_sizes: list[int],

111 input_size: int,

112 output_size: int,

113 params_dtype: torch.dtype,

@@ -177,7 +177,7 @@ class ModelOptNvFp4Config(QuantizationConfig):

177 self,

178 is_checkpoint_nvfp4_serialized: bool,

179 kv_cache_quant_algo: str,

180 - exclude_modules: List[str],

180 + exclude_modules: list[str],

181 group_size: int = 16,

182 ) -> None:

183 self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized

@@ -195,7 +195,7 @@ class ModelOptNvFp4Config(QuantizationConfig):

195 return "nvfp4"

196

197 @classmethod

198 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

198 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

199 return [torch.bfloat16, torch.half, torch.float8_e4m3fn]

200

201 @classmethod

@@ -203,11 +203,11 @@ class ModelOptNvFp4Config(QuantizationConfig):

203 return 80

204

205 @classmethod

206 - def get_config_filenames(cls) -> List[str]:

206 + def get_config_filenames(cls) -> list[str]:

207 return ["hf_quant_config.json"]

208

209 @classmethod

210 - def from_config(cls, config: Dict[str, Any]) -> "ModelOptNvFp4Config":

210 + def from_config(cls, config: dict[str, Any]) -> "ModelOptNvFp4Config":

211 quant_config = cls.get_from_keys(config, ["quantization"])

212 quant_method = quant_config["quant_algo"]

213 if quant_method not in QUANT_ALGOS:

@@ -227,7 +227,7 @@ class ModelOptNvFp4Config(QuantizationConfig):

227 return cls(is_checkpoint_nvfp4_serialized, kv_cache_quant_algo,

228 exclude_modules, group_size)

229

230 - def is_layer_excluded(self, prefix: str, exclude_modules: List):

230 + def is_layer_excluded(self, prefix: str, exclude_modules: list):

231 import re

232 for pattern in exclude_modules:

233 regex_str = pattern.replace('.', r'\.').replace('*', r'.*')

@@ -296,7 +296,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

296 self,

297 layer: torch.nn.Module,

298 input_size_per_partition: int,

299 - output_partition_sizes: List[int],

299 + output_partition_sizes: list[int],

300 input_size: int,

301 output_size: int,

302 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/moe_wna16.py

+8 -8

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, List, Optional

3 +from typing import Any, Callable, Optional

4

5 import torch

6

@@ -23,8 +23,8 @@ class MoeWNA16Config(QuantizationConfig):

23

24 def __init__(self, linear_quant_method: str, weight_bits: int,

25 group_size: int, has_zp: bool, lm_head_quantized: bool,

26 - modules_to_not_convert: Optional[List[str]],

27 - full_config: Dict[str, Any]) -> None:

26 + modules_to_not_convert: Optional[list[str]],

27 + full_config: dict[str, Any]) -> None:

28 super().__init__()

29 self.weight_bits = weight_bits

30 self.group_size = group_size

@@ -69,7 +69,7 @@ class MoeWNA16Config(QuantizationConfig):

69 return "moe_wna16"

70

71 @classmethod

72 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

72 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

73 return [torch.bfloat16, torch.half]

74

75 @classmethod

@@ -77,11 +77,11 @@ class MoeWNA16Config(QuantizationConfig):

77 return 70

78

79 @classmethod

80 - def get_config_filenames(cls) -> List[str]:

80 + def get_config_filenames(cls) -> list[str]:

81 return ["quantize_config.json"]

82

83 @classmethod

84 - def from_config(cls, config: Dict[str, Any]) -> "MoeWNA16Config":

84 + def from_config(cls, config: dict[str, Any]) -> "MoeWNA16Config":

85 linear_quant_method = cls.get_from_keys(config, ["quant_method"])

86 weight_bits = cls.get_from_keys(config, ["bits"])

87 group_size = cls.get_from_keys(config, ["group_size"])

@@ -109,7 +109,7 @@ class MoeWNA16Config(QuantizationConfig):

109 return None

110

111 @classmethod

112 - def is_moe_wna16_compatible(cls, quant_config: Dict[str, Any]):

112 + def is_moe_wna16_compatible(cls, quant_config: dict[str, Any]):

113 # Extract data from quant config.

114 quant_method = quant_config.get("quant_method", "").lower()

115 num_bits = quant_config.get("bits")

@@ -163,7 +163,7 @@ class MoeWNA16Config(QuantizationConfig):

163 return None

164

165

166 -def is_layer_skipped_quant(prefix: str, modules_to_not_convert: List[str]):

166 +def is_layer_skipped_quant(prefix: str, modules_to_not_convert: list[str]):

167 return any(module_name in prefix for module_name in modules_to_not_convert)

168

169

vllm/model_executor/layers/quantization/neuron_quant.py

+4 -4

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

2

3 import os

4 from importlib.util import find_spec

5 -from typing import Any, Dict, List, Optional

5 +from typing import Any, Optional

6

7 from torch.nn import Module

8

@@ -34,7 +34,7 @@ class NeuronQuantConfig(QuantizationConfig):

34 def get_name(self) -> QuantizationMethods:

35 return "neuron_quant"

36

37 - def get_supported_act_dtypes(self) -> List[str]:

37 + def get_supported_act_dtypes(self) -> list[str]:

38 return SUPPORTED_QUANT_DTYPE_LIST

39

40 @classmethod

@@ -43,11 +43,11 @@ class NeuronQuantConfig(QuantizationConfig):

43 "This function should not be called with Neuron Backend")

44

45 @staticmethod

46 - def get_config_filenames() -> List[str]:

46 + def get_config_filenames() -> list[str]:

47 return []

48

49 @classmethod

50 - def from_config(cls, config: Dict[str, Any]) -> "NeuronQuantConfig":

50 + def from_config(cls, config: dict[str, Any]) -> "NeuronQuantConfig":

51 quantize_method = cls.get_from_keys(config, ["quantize_method"])

52 dequant_dtype = cls.get_from_keys(config, ["dequant_dtype"])

53 return cls(dequant_dtype=dequant_dtype,

vllm/model_executor/layers/quantization/ptpc_fp8.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6 from torch.nn.parameter import Parameter

@@ -32,7 +32,7 @@ class PTPCFp8Config(Fp8Config):

32 def __init__(

33 self,

34 activation_scheme: str = "dynamic",

35 - ignored_layers: Optional[List[str]] = None,

35 + ignored_layers: Optional[list[str]] = None,

36 ) -> None:

37 if not current_platform.is_rocm():

38 raise ValueError(

@@ -55,7 +55,7 @@ class PTPCFp8Config(Fp8Config):

55 return "ptpc_fp8"

56

57 @classmethod

58 - def from_config(cls, config: Dict[str, Any]) -> "PTPCFp8Config":

58 + def from_config(cls, config: dict[str, Any]) -> "PTPCFp8Config":

59 activation_scheme = cls.get_from_keys(config, ["activation_scheme"])

60 ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)

61 return cls(activation_scheme=activation_scheme,

vllm/model_executor/layers/quantization/qqq.py

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6 from torch.nn.parameter import Parameter

@@ -89,7 +89,7 @@ class QQQConfig(QuantizationConfig):

89 return "qqq"

90

91 @classmethod

92 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

92 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

93 return [torch.half]

94

95 @classmethod

@@ -97,7 +97,7 @@ class QQQConfig(QuantizationConfig):

97 return 80

98

99 @classmethod

100 - def get_config_filenames(cls) -> List[str]:

100 + def get_config_filenames(cls) -> list[str]:

101 """List of filenames to search for in the model directory."""

102 return [

103 "quant_config.json",

@@ -105,7 +105,7 @@ class QQQConfig(QuantizationConfig):

105 ]

106

107 @classmethod

108 - def from_config(cls, config: Dict[str, Any]) -> "QQQConfig":

108 + def from_config(cls, config: dict[str, Any]) -> "QQQConfig":

109 weight_bits = cls.get_from_keys(config, ["wbits"])

110 group_size = cls.get_from_keys(config, ["group_size"])

111 return cls(weight_bits, group_size)

@@ -131,7 +131,7 @@ class QQQLinearMethod(LinearMethodBase):

131 self,

132 layer: torch.nn.Module,

133 input_size_per_partition: int,

134 - output_partition_sizes: List[int],

134 + output_partition_sizes: list[int],

135 input_size: int,

136 output_size: int,

137 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/quark/quark.py

+27 -27

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 import fnmatch

4 -from typing import Any, Dict, List, Optional, cast

4 +from typing import Any, Optional, cast

5

6 import torch

7

@@ -29,9 +29,9 @@ logger = init_logger(__name__)

29 class QuarkConfig(QuantizationConfig):

30

31 def __init__(self,

32 - quant_config: Dict[str, Any],

33 - kv_cache_group: Optional[List[str]] = None,

34 - kv_cache_config: Optional[Dict[str, Any]] = None,

32 + quant_config: dict[str, Any],

33 + kv_cache_group: Optional[list[str]] = None,

34 + kv_cache_config: Optional[dict[str, Any]] = None,

35 pack_method: str = "reorder"):

36 super().__init__()

37 if kv_cache_group is None:

@@ -44,7 +44,7 @@ class QuarkConfig(QuantizationConfig):

44 def get_linear_method(self) -> "QuarkLinearMethod":

45 return QuarkLinearMethod(self)

46

47 - def get_supported_act_dtypes(cls) -> List[torch.dtype]:

47 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:

48 return [torch.float16, torch.bfloat16]

49

50 @classmethod

@@ -59,7 +59,7 @@ class QuarkConfig(QuantizationConfig):

59 from vllm.attention.layer import Attention # Avoid circular import

60

61 # Check if the layer is skipped for quantization.

62 - exclude_layers = cast(List[str], self.quant_config.get("exclude"))

62 + exclude_layers = cast(list[str], self.quant_config.get("exclude"))

63 if should_ignore_layer(prefix,

64 ignore=exclude_layers,

65 fused_mapping=self.packed_modules_mapping):

@@ -78,12 +78,12 @@ class QuarkConfig(QuantizationConfig):

78 return None

79

80 @classmethod

81 - def from_config(cls, config: Dict[str, Any]) -> "QuarkConfig":

81 + def from_config(cls, config: dict[str, Any]) -> "QuarkConfig":

82 export_config = config.get("export")

83 if export_config is None:

84 raise ValueError("The export key should be included in "

85 "the configurations of Quark quantized model")

86 - kv_cache_group = cast(List[str], export_config.get("kv_cache_group"))

86 + kv_cache_group = cast(list[str], export_config.get("kv_cache_group"))

87 pack_method = cast(str, export_config.get("pack_method"))

88

89 # In the export model of quark, the quantization configuration

@@ -95,7 +95,7 @@ class QuarkConfig(QuantizationConfig):

95 kv_cache_config = None

96 else:

97 kv_cache_set = set(kv_cache_group)

98 - layer_quant_config = cast(Dict[str, Any],

98 + layer_quant_config = cast(dict[str, Any],

99 config.get("layer_quant_config"))

100 layer_quant_names = list(layer_quant_config.keys())

101 layer_quant_set = set(layer_quant_names)

@@ -108,7 +108,7 @@ class QuarkConfig(QuantizationConfig):

108 "configuration.")

109

110 q_configs = [

111 - cast(Dict[str, Any], layer_quant_config.get(name))

111 + cast(dict[str, Any], layer_quant_config.get(name))

112 for name in kv_cache_group

113 ]

114 if not all(

@@ -131,7 +131,7 @@ class QuarkConfig(QuantizationConfig):

131

132 # In case q_proj output is also quantized, remove the configuration

133 # to keep qkv consistency.

134 - q_proj_q_config = cast(Dict[str, Any],

134 + q_proj_q_config = cast(dict[str, Any],

135 layer_quant_config.get("*q_proj"))

136 if q_proj_q_config is not None:

137 q_proj_q_config["output_tensors"] = None

@@ -142,7 +142,7 @@ class QuarkConfig(QuantizationConfig):

142 pack_method=pack_method)

143

144 @classmethod

145 - def get_config_filenames(cls) -> List[str]:

145 + def get_config_filenames(cls) -> list[str]:

146 return []

147

148 def _check_scheme_supported(self,

@@ -162,8 +162,8 @@ class QuarkConfig(QuantizationConfig):

162 else:

163 return False

164

165 - def _is_fp8_w8a8(self, weight_quant: Optional[Dict[str, Any]],

166 - input_quant: Optional[Dict[str, Any]]) -> bool:

165 + def _is_fp8_w8a8(self, weight_quant: Optional[dict[str, Any]],

166 + input_quant: Optional[dict[str, Any]]) -> bool:

167 # Confirm weights and input quantized.

168 if weight_quant is None or input_quant is None:

169 return False

@@ -187,8 +187,8 @@ class QuarkConfig(QuantizationConfig):

187 is_per_tensor_activation = (input_quant.get("qscheme") == "per_tensor")

188 return is_per_tensor_activation

189

190 - def _is_static_tensor_w8a8(self, weight_quant: Optional[Dict[str, Any]],

191 - input_quant: Optional[Dict[str, Any]]) -> bool:

190 + def _is_static_tensor_w8a8(self, weight_quant: Optional[dict[str, Any]],

191 + input_quant: Optional[dict[str, Any]]) -> bool:

192 # Confirm weights and input quantized.

193 if weight_quant is None or input_quant is None:

194 return False

@@ -209,8 +209,8 @@ class QuarkConfig(QuantizationConfig):

209 # Only symmetric weight quantization supported.

210 return is_int8_dtype and is_tensor and is_weight_symmetric and is_static

211

212 - def _is_mx_fp4(self, weight_quant: Optional[Dict[str, Any]],

213 - input_quant: Optional[Dict[str, Any]]) -> bool:

212 + def _is_mx_fp4(self, weight_quant: Optional[dict[str, Any]],

213 + input_quant: Optional[dict[str, Any]]) -> bool:

214 # Confirm weights and input quantized.

215 if weight_quant is None or input_quant is None:

216 logger.debug("Quark model is not in MX-FP4 format: "

@@ -258,7 +258,7 @@ class QuarkConfig(QuantizationConfig):

258 return True

259

260 def _find_matched_config(self, layer_name: str,

261 - module: torch.nn.Module) -> Dict[str, Any]:

261 + module: torch.nn.Module) -> dict[str, Any]:

262

263 proj_name = layer_name.split(".")[-1]

264 if proj_name in self.packed_modules_mapping:

@@ -283,29 +283,29 @@ class QuarkConfig(QuantizationConfig):

283 return shard_configs[0]

284 else:

285 layer_quant_config = cast(

286 - Dict[str, Any], self.quant_config.get("layer_quant_config"))

286 + dict[str, Any], self.quant_config.get("layer_quant_config"))

287 for name_pattern in layer_quant_config:

288 if fnmatch.fnmatch(layer_name, name_pattern):

289 return layer_quant_config[name_pattern]

290

291 layer_type = cast(str, type(module))

292 layer_type_quant_config = cast(

293 - Dict[str, Any],

293 + dict[str, Any],

294 self.quant_config.get("layer_type_quant_config"))

295 if layer_type in layer_type_quant_config:

296 return layer_type_quant_config[layer_type]

297

298 global_quant_config = cast(

299 - Dict[str, Any], self.quant_config.get("global_quant_config"))

299 + dict[str, Any], self.quant_config.get("global_quant_config"))

300 return global_quant_config

301

302 - def _get_scheme_from_config(self, config: Dict[str, Any]) -> "QuarkScheme":

302 + def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkScheme":

303 if config.get("output_tensors") or config.get("bias"):

304 raise NotImplementedError(

305 "Currently, Quark models with output_tensors "

306 "and bias quantized are not supported")

307 - weight_config = cast(Dict[str, Any], config.get("weight"))

308 - input_config = cast(Dict[str, Any], config.get("input_tensors"))

307 + weight_config = cast(dict[str, Any], config.get("weight"))

308 + input_config = cast(dict[str, Any], config.get("input_tensors"))

309

310 if self._is_fp8_w8a8(weight_config, input_config):

311 is_fp8_w8a8_supported = self._check_scheme_supported(

@@ -373,7 +373,7 @@ class QuarkLinearMethod(LinearMethodBase):

373

374 def create_weights(self, layer: torch.nn.Module,

375 input_size_per_partition: int,

376 - output_partition_sizes: List[int], input_size: int,

376 + output_partition_sizes: list[int], input_size: int,

377 output_size: int, params_dtype: torch.dtype,

378 **extra_weight_attrs):

379 """

@@ -417,7 +417,7 @@ class QuarkKVCacheMethod(BaseKVCacheMethod):

417 super().__init__(quant_config)

418

419 @staticmethod

420 - def validate_kv_cache_config(kv_cache_config: Optional[Dict[str, Any]]):

420 + def validate_kv_cache_config(kv_cache_config: Optional[dict[str, Any]]):

421 """

422 Validator for the kv cache configuration. Useful for controlling the

423 kv cache quantization schemes, that are being supported in vLLM

vllm/model_executor/layers/quantization/quark/quark_moe.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, Optional

3 +from typing import Any, Callable, Optional

4

5 import torch

6

@@ -45,7 +45,7 @@ class QuarkMoEMethod(FusedMoEMethodBase):

45

46 class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

47

48 - def __init__(self, weight_config: Dict[str, Any], input_config: Dict[str,

48 + def __init__(self, weight_config: dict[str, Any], input_config: dict[str,

49 Any]):

50 self.weight_quant = weight_config

51 self.input_quant = input_config

vllm/model_executor/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Any, Callable, Dict, List, Optional

3 +from typing import Any, Callable, Optional

4

5 import torch

6 import torch.nn.functional as F

@@ -18,8 +18,8 @@ __all__ = ["QuarkW4A4MXFP4"]

18

19 class QuarkW4A4MXFP4(QuarkScheme):

20

21 - def __init__(self, weight_quant_spec: Dict[str, Any],

22 - input_quant_spec: Dict[str, Any]):

21 + def __init__(self, weight_quant_spec: dict[str, Any],

22 + input_quant_spec: dict[str, Any]):

23 self.out_dtype = torch.get_default_dtype()

24 self.qscheme = "per_group"

25 self.weight_quant_spec = weight_quant_spec

@@ -74,7 +74,7 @@ class QuarkW4A4MXFP4(QuarkScheme):

74 torch.cuda.empty_cache()

75

76 def create_weights(self, layer: torch.nn.Module,

77 - output_partition_sizes: List[int],

77 + output_partition_sizes: list[int],

78 input_size_per_partition: int,

79 params_dtype: torch.dtype, weight_loader: Callable,

80 **kwargs):

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional

3 +from typing import Callable, Optional

4

5 import torch

6 from torch.nn import Parameter

@@ -88,7 +88,7 @@ class QuarkW8A8Fp8(QuarkScheme):

88 layer.input_scale = None

89

90 def create_weights(self, layer: torch.nn.Module,

91 - output_partition_sizes: List[int],

91 + output_partition_sizes: list[int],

92 input_size_per_partition: int,

93 params_dtype: torch.dtype, weight_loader: Callable,

94 **kwargs):

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional, Set

3 +from typing import Callable, Optional

4

5 import torch

6

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

17

18

19 class QuarkW8A8Int8(QuarkScheme):

20 - _kernel_backends_being_used: Set[str] = set()

20 + _kernel_backends_being_used: set[str] = set()

21

22 def __init__(self, qscheme: str, is_static_input_scheme: Optional[bool],

23 input_symmetric: Optional[bool]):

@@ -31,7 +31,7 @@ class QuarkW8A8Int8(QuarkScheme):

31 return 75

32

33 def create_weights(self, layer: torch.nn.Module,

34 - output_partition_sizes: List[int],

34 + output_partition_sizes: list[int],

35 input_size_per_partition: int,

36 params_dtype: torch.dtype, weight_loader: Callable,

37 **kwargs):

vllm/model_executor/layers/quantization/quark/utils.py

+3 -2

@@ -1,8 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 import re

4 +from collections.abc import Iterable, Mapping

5 from types import MappingProxyType

5 -from typing import Any, Iterable, List, Mapping, Optional

6 +from typing import Any, Optional

7

8

9 def deep_compare(dict1: Any, dict2: Any) -> bool:

@@ -21,7 +22,7 @@ def deep_compare(dict1: Any, dict2: Any) -> bool:

22 def should_ignore_layer(

23 layer_name: Optional[str],

24 ignore: Iterable[str],

24 - fused_mapping: Mapping[str, List[str]] = MappingProxyType({})

25 + fused_mapping: Mapping[str, list[str]] = MappingProxyType({})

26 ) -> bool:

27 if layer_name is None:

28 return False

vllm/model_executor/layers/quantization/schema.py

+2 -2

@@ -12,7 +12,7 @@ possible on ROCm), the model can be optionally augmented with KV cache

12 scaling factors.

13 """

14

15 -from typing import Dict, Optional

15 +from typing import Optional

16

17 from pydantic import BaseModel, ConfigDict, ValidationInfo, model_validator

18

@@ -23,7 +23,7 @@ class KVCacheQuantSchema(BaseModel):

23 # layer indices to their per-tensor KV cache scaling factor.

24 # TODO: Consider pulling this and its validation methods out into its

25 # own schema class (tricky as its members are variable)

26 - scaling_factor: Dict[int, Dict[int, float]]

26 + scaling_factor: dict[int, dict[int, float]]

27

28 @model_validator(mode="after")

29 def check_is_fp8(self) -> "KVCacheQuantSchema":

vllm/model_executor/layers/quantization/torchao.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Any, Dict, List, Optional

2 +from typing import Any, Optional

3

4 import torch

5 import torch.nn.functional as F

@@ -24,7 +24,7 @@ class TorchAOConfig(QuantizationConfig):

24 def get_name(self) -> QuantizationMethods:

25 return "torchao"

26

27 - def get_supported_act_dtypes(self) -> List[torch.dtype]:

27 + def get_supported_act_dtypes(self) -> list[torch.dtype]:

28 return [torch.float32, torch.float16, torch.bfloat16]

29

30 @classmethod

@@ -32,11 +32,11 @@ class TorchAOConfig(QuantizationConfig):

32 return 75

33

34 @staticmethod

35 - def get_config_filenames() -> List[str]:

35 + def get_config_filenames() -> list[str]:

36 return ["config.json"]

37

38 @classmethod

39 - def from_config(cls, config: Dict[str, Any]) -> "TorchAOConfig":

39 + def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig":

40 """Create the quant config from an hf model config"""

41 try:

42 from torchao.core.config import config_from_dict

@@ -60,7 +60,7 @@ class TorchAOConfig(QuantizationConfig):

60 return TorchAOLinearMethod(self)

61 return None

62

63 - def get_scaled_act_names(self) -> List[str]:

63 + def get_scaled_act_names(self) -> list[str]:

64 return []

65

66

@@ -97,7 +97,7 @@ class TorchAOLinearMethod(LinearMethodBase):

97 self,

98 layer: torch.nn.Module,

99 input_size_per_partition: int,

100 - output_partition_sizes: List[int],

100 + output_partition_sizes: list[int],

101 input_size: int,

102 output_size: int,

103 params_dtype: torch.dtype,

vllm/model_executor/layers/quantization/tpu_int8.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

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

3 +from typing import Any, Optional

4

5 import torch

6 from torch.nn import Module

@@ -31,7 +31,7 @@ class Int8TpuConfig(QuantizationConfig):

31 def get_name(self) -> QuantizationMethods:

32 return "tpu_int8"

33

34 - def get_supported_act_dtypes(self) -> List[torch.dtype]:

34 + def get_supported_act_dtypes(self) -> list[torch.dtype]:

35 return [torch.float16, torch.bfloat16]

36

37 @classmethod

@@ -40,11 +40,11 @@ class Int8TpuConfig(QuantizationConfig):

40 "This function should not be called with TPU Backend")

41

42 @staticmethod

43 - def get_config_filenames() -> List[str]:

43 + def get_config_filenames() -> list[str]:

44 return []

45

46 @classmethod

47 - def from_config(cls, config: Dict[str, Any]) -> "Int8TpuConfig":

47 + def from_config(cls, config: dict[str, Any]) -> "Int8TpuConfig":

48 activation_scheme = cls.get_from_keys(config, ["activation_scheme"])

49 return cls(activation_scheme=activation_scheme)

50

@@ -62,7 +62,7 @@ class TPUInt8LinearMethod(LinearMethodBase):

62 self.quant_config = quant_config

63

64 def create_weights(self, layer: Module, input_size_per_partition: int,

65 - output_partition_sizes: List[int], input_size: int,

65 + output_partition_sizes: list[int], input_size: int,

66 output_size: int, params_dtype: torch.dtype,

67 **extra_weight_attrs):

68

@@ -77,7 +77,7 @@ class TPUInt8LinearMethod(LinearMethodBase):

77 layer.register_parameter("weight", weight)

78

79 def _quantize_weight(

80 - self, weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

80 + self, weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

81 weight_dtype = weight.dtype

82 weight = weight.cpu().to(torch.float32)

83 n_bit = 8

vllm/model_executor/layers/quantization/utils/bitblas_utils.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Optional, Tuple

2 +from typing import Optional

3

4 import torch

5

@@ -51,7 +51,7 @@ def _check_bitblas_supported(

51 quant_type: ScalarType,

52 group_size: Optional[int],

53 has_zp: bool,

54 - device_capability: Optional[int] = None) -> Tuple[bool, Optional[str]]:

54 + device_capability: Optional[int] = None) -> tuple[bool, Optional[str]]:

55

56 if device_capability is None:

57 capability_tuple = current_platform.get_device_capability()

@@ -133,7 +133,7 @@ def verify_bitblas_supports_shape(output_size_per_partition: int,

133 def check_bitblas_supports_shape(output_size_per_partition: int,

134 input_size_per_partition: int,

135 input_size: int, group_size: int) \

136 - -> Tuple[bool, Optional[str]]:

136 + -> tuple[bool, Optional[str]]:

137 try:

138 verify_bitblas_supports_shape(output_size_per_partition,

139 input_size_per_partition, input_size,

@@ -166,7 +166,7 @@ def bitblas_make_empty_zp(device: torch.device) -> torch.Tensor:

166

167

168 def bitblas_sort_g_idx(

169 - g_idx: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

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

170 g_idx_sort_indices = torch.argsort(g_idx).to(torch.int)

171 return g_idx[g_idx_sort_indices], g_idx_sort_indices

172

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+9 -9

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

4 import functools

5 import json

6 import os

7 -from typing import Any, Dict, List, Optional, Tuple, Union

7 +from typing import Any, Optional, Union

8

9 import torch

10

@@ -32,7 +32,7 @@ def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:

32 def apply_w8a8_block_fp8_linear(

33 input: torch.Tensor,

34 weight: torch.Tensor,

35 - block_size: List[int],

35 + block_size: list[int],

36 weight_scale: torch.Tensor,

37 input_scale: Optional[torch.Tensor] = None,

38 bias: Optional[torch.Tensor] = None,

@@ -95,7 +95,7 @@ def apply_w8a8_block_fp8_linear(

95 def apply_w8a8_block_fp8_linear_fake(

96 input: torch.Tensor,

97 weight: torch.Tensor,

98 - block_size: List[int],

98 + block_size: list[int],

99 weight_scale: torch.Tensor,

100 input_scale: Optional[torch.Tensor] = None,

101 ) -> torch.Tensor:

@@ -114,7 +114,7 @@ direct_register_custom_op(

114 def input_to_float8(

115 x: torch.Tensor,

116 dtype: Optional[torch.dtype] = None

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

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

118 """This function quantizes input values to float8 values "

119 "with tensor-wise quantization."""

120 dtype = current_platform.fp8_dtype() if dtype is None else dtype

@@ -129,7 +129,7 @@ def input_to_float8(

129 def block_quant_to_tensor_quant(

130 x_q_block: torch.Tensor,

131 x_s: torch.Tensor,

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

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

133 """This function converts block-wise quantization to tensor-wise

134 quantization. The inputs are block-wise quantization tensor `x_q_block`,

135 block-wise quantization scale and the block size.

@@ -247,7 +247,7 @@ def per_token_group_quant_fp8(

247 eps: float = 1e-10,

248 dtype: Optional[torch.dtype] = None,

249 column_major_scales: bool = False,

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

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

251 """Function to perform per-token-group quantization on an input tensor `x`.

252 It converts the tensor values into signed float8 values and returns the

253 quantized tensor along with the scaling factor used for quantization.

@@ -258,7 +258,7 @@ def per_token_group_quant_fp8(

258 dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn`

259 is supported for now.

260 Returns:

261 - Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the

261 + tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the

262 scaling factor for quantization.

263 """

264 dtype = current_platform.fp8_dtype() if dtype is None else dtype

@@ -412,7 +412,7 @@ def _w8a8_block_fp8_matmul(

412

413 @functools.lru_cache

414 def get_w8a8_block_fp8_configs(N: int, K: int, block_n: int,

415 - block_k: int) -> Optional[Dict[int, Any]]:

415 + block_k: int) -> Optional[dict[int, Any]]:

416 """

417 Return optimized configurations for the w8a8 block fp8 kernel.

418 The return value will be a dictionary that maps an irregular grid of

@@ -452,7 +452,7 @@ def w8a8_block_fp8_matmul(

452 B: torch.Tensor,

453 As: torch.Tensor,

454 Bs: torch.Tensor,

455 - block_size: List[int],

455 + block_size: list[int],

456 output_dtype: torch.dtype = torch.float16,

457 ) -> torch.Tensor:

458 """This function performs matrix multiplication with block-wise

vllm/model_executor/layers/quantization/utils/gptq_utils.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 import re

3 from copy import deepcopy

4 -from typing import Dict, Optional, Union

4 +from typing import Optional, Union

5

6 import torch

7

@@ -52,7 +52,7 @@ def get_dynamic_override(

52 layer_name: str,

53 key: Optional[str] = None,

54 default_value: Union[int, bool,

55 - None] = None) -> Union[Dict, int, bool, None]:

55 + None] = None) -> Union[dict, int, bool, None]:

56 for pattern, pattern_dict in config.dynamic.items():

57 # Negative match: matched modules are excluded from quantized init

58 if pattern.startswith("-:"):

vllm/model_executor/layers/quantization/utils/int8_utils.py

+8 -8

@@ -5,7 +5,7 @@ import functools

5 import json

6 import logging

7 import os

8 -from typing import Any, Dict, List, Optional, Tuple

8 +from typing import Any, Optional

9

10 import torch

11

@@ -18,7 +18,7 @@ logger = logging.getLogger(__name__)

18 def apply_w8a8_block_int8_linear(

19 input: torch.Tensor,

20 weight: torch.Tensor,

21 - block_size: List[int],

21 + block_size: list[int],

22 weight_scale: torch.Tensor,

23 input_scale: Optional[torch.Tensor] = None,

24 bias: Optional[torch.Tensor] = None,

@@ -43,7 +43,7 @@ def apply_w8a8_block_int8_linear(

43

44 def input_to_int8(

45 x: torch.Tensor,

46 - dtype: torch.dtype = torch.int8) -> Tuple[torch.Tensor, torch.Tensor]:

46 + dtype: torch.dtype = torch.int8) -> tuple[torch.Tensor, torch.Tensor]:

47 """This function quantizes input values to int8 values with

48 tensor-wise quantization."""

49 iinfo = torch.iinfo(dtype)

@@ -58,7 +58,7 @@ def input_to_int8(

58 def block_dequant(

59 x_q_block: torch.Tensor,

60 x_s: torch.Tensor,

61 - block_size: List[int],

61 + block_size: list[int],

62 ) -> torch.Tensor:

63 """This function conducts block-wise dequantization.

64 The inputs are block-wise quantization tensor `x_q_block`,

@@ -211,7 +211,7 @@ def per_token_group_quant_int8(

211 group_size: int,

212 eps: float = 1e-10,

213 dtype: torch.dtype = torch.int8,

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

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

215 """Function to perform per-token-group quantization on an input tensor `x`.

216

217 It converts the tensor values into signed int8 values and returns the

@@ -225,7 +225,7 @@ def per_token_group_quant_int8(

225 is supported for now.

226

227 Returns:

228 - Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the

228 + tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the

229 scaling factor for quantization.

230 """

231 assert (x.shape[-1] % group_size == 0

@@ -358,7 +358,7 @@ def _w8a8_block_int8_matmul(

358

359 @functools.lru_cache

360 def get_w8a8_block_int8_configs(N: int, K: int, block_n: int,

361 - block_k: int) -> Optional[Dict[int, Any]]:

361 + block_k: int) -> Optional[dict[int, Any]]:

362 """

363 Return optimized configurations for the w8a8 block fp8 kernel.

364

@@ -399,7 +399,7 @@ def w8a8_block_int8_matmul(

399 B: torch.Tensor,

400 As: torch.Tensor,

401 Bs: torch.Tensor,

402 - block_size: List[int],

402 + block_size: list[int],

403 output_dtype: torch.dtype = torch.float16,

404 ) -> torch.Tensor:

405 """This function performs matrix multiplication with block-wise

vllm/model_executor/layers/quantization/utils/machete_utils.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import List, Optional, Tuple

3 +from typing import Optional

4

5 import torch

6

@@ -10,19 +10,19 @@ MACHETE_SUPPORTED_GROUP_SIZES = [-1, 128]

10 MACHETE_PREPACKED_BLOCK_SHAPE = [64, 128]

11

12

13 -def query_machete_supported_quant_types(zero_points: bool) -> List[ScalarType]:

13 +def query_machete_supported_quant_types(zero_points: bool) -> list[ScalarType]:

14 if zero_points:

15 return [scalar_types.uint4, scalar_types.uint8]

16 else:

17 return [scalar_types.uint4b8, scalar_types.uint8b128]

18

19

20 -def query_machete_supported_act_types(zero_points: bool) -> List[ScalarType]:

20 +def query_machete_supported_act_types(zero_points: bool) -> list[ScalarType]:

21 return [torch.float16, torch.bfloat16]

22

23

24 def check_machete_supports_shape(in_features: int, out_featrues: int) \

25 - -> Tuple[bool, Optional[str]]:

25 + -> tuple[bool, Optional[str]]:

26 if in_features % MACHETE_PREPACKED_BLOCK_SHAPE[0] != 0:

27 return False, "Input features size must be divisible by "\

28 f"{MACHETE_PREPACKED_BLOCK_SHAPE[0]}"

vllm/model_executor/layers/quantization/utils/marlin_utils.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import List, Optional, Tuple

3 +from typing import Optional

4

5 import numpy

6 import torch

@@ -70,7 +70,7 @@ def _check_marlin_supported(

70 quant_type: ScalarType,

71 group_size: Optional[int],

72 has_zp: bool,

73 - device_capability: Optional[int] = None) -> Tuple[bool, Optional[str]]:

73 + device_capability: Optional[int] = None) -> tuple[bool, Optional[str]]:

74

75 if device_capability is None:

76 capability_tuple = current_platform.get_device_capability()

@@ -143,7 +143,7 @@ def verify_marlin_supports_shape(output_size_per_partition: int,

143 def check_marlin_supports_shape(output_size_per_partition: int,

144 input_size_per_partition: int,

145 input_size: int, group_size: int) \

146 - -> Tuple[bool, Optional[str]]:

146 + -> tuple[bool, Optional[str]]:

147 try:

148 verify_marlin_supports_shape(output_size_per_partition,

149 input_size_per_partition, input_size,

@@ -231,16 +231,16 @@ def marlin_make_empty_zp(device: torch.device) -> torch.Tensor:

231

232

233 def marlin_sort_g_idx(

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

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

235 g_idx_sort_indices = torch.argsort(g_idx).to(torch.int)

236 return g_idx[g_idx_sort_indices], g_idx_sort_indices

237

238

239 def get_scale_perms():

240 - scale_perm: List[int] = []

240 + scale_perm: list[int] = []

241 for i in range(8):

242 scale_perm.extend([i + 8 * j for j in range(8)])

243 - scale_perm_single: List[int] = []

243 + scale_perm_single: list[int] = []

244 for i in range(4):

245 scale_perm_single.extend(

246 [2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])

vllm/model_executor/layers/quantization/utils/marlin_utils_test.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Utility functions used for tests and benchmarks"""

3

4 -from typing import List, Optional

4 +from typing import Optional

5

6 import numpy as np

7 import torch

@@ -64,9 +64,9 @@ def marlin_weights(q_w, size_k, size_n, num_bits, perm):

64

65

66 def get_weight_perm(num_bits: int):

67 - perm_list: List[int] = []

67 + perm_list: list[int] = []

68 for i in range(32):

69 - perm1: List[int] = []

69 + perm1: list[int] = []

70 col = i // 4

71 for block in [0, 1]:

72 for row in [

vllm/model_executor/layers/quantization/utils/marlin_utils_test_24.py

+4 -5

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

2 """Utility functions used for tests and benchmarks"""

3

4 import random

5 -from typing import List

5

6 import numpy

7 import torch

@@ -373,19 +372,19 @@ def compress_quantized_24_weight(q_24, size_k, size_n, wtype: ScalarType):

372

373

374 def get_scale_perms_24():

376 - scale_perm: List[int] = []

375 + scale_perm: list[int] = []

376 for i in range(8):

377 scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]])

379 - scale_perm_single: List[int] = []

378 + scale_perm_single: list[int] = []

379 for i in range(8):

380 scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]])

381 return scale_perm, scale_perm_single

382

383

384 def get_weight_perm_24(num_bits: int):

386 - perm_list: List[int] = []

385 + perm_list: list[int] = []

386 for i in range(32):

388 - perm1: List[int] = []

387 + perm1: list[int] = []

388 col = i // 4

389 col_o = col // 2

390 for block in [0, 1]:

vllm/model_executor/layers/quantization/utils/marlin_utils_test_qqq.py

+4 -6

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import List

4 -

3 import numpy

4 import torch

5

@@ -34,10 +32,10 @@ def marlin_qqq_weights(q_w, size_k, size_n, num_bits, perm, group_size):

32

33

34 def get_qqq_scale_perms():

37 - scale_perm: List[int] = []

35 + scale_perm: list[int] = []

36 for i in range(8):

37 scale_perm.extend([i + 8 * j for j in range(8)])

40 - scale_perm_single: List[int] = []

38 + scale_perm_single: list[int] = []

39 for i in range(4):

40 scale_perm_single.extend(

41 [2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])

@@ -46,9 +44,9 @@ def get_qqq_scale_perms():

44

45 # NOTE(HandH1998): QQQ employs different perms for per-group and per-channel weight quantization. # noqa: E501

46 def get_qqq_weight_perm(num_bits: int, quant_type: str):

49 - perm_list: List[int] = []

47 + perm_list: list[int] = []

48 for i in range(32):

51 - perm1: List[int] = []

49 + perm1: list[int] = []

50 col = i // 4

51 for block in [0, 1]:

52 for row in [

vllm/model_executor/layers/quantization/utils/mxfp4_utils.py

+1 -2

@@ -1,5 +1,4 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Tuple

2

3 import torch

4

@@ -9,7 +8,7 @@ OCP_MX_BLOCK_SIZE = 32

8 def per_token_group_quant_mxfp4(x: torch.Tensor,

9 block_k: int,

10 scale_calculation_mode: str = "even"

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

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

12 try:

13 from quark.torch.kernel.hw_emulation.hw_emulation_interface import (

14 fake_quantize_fp4_fp6_per_group_with_scale)

vllm/model_executor/layers/quantization/utils/quant_utils.py

+9 -8

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

1 # SPDX-License-Identifier: Apache-2.0

2 """This file is used for /tests and /benchmarks"""

3 +from collections.abc import Mapping

4 from types import MappingProxyType

4 -from typing import List, Mapping, Optional, Tuple

5 +from typing import Optional

6

7 import numpy

8 import torch

@@ -15,7 +16,7 @@ SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]

16

17

18 # Normalize the group_shape to the full extent for any dims that are -1

18 -def _normalize_quant_group_shape(x: torch.Tensor, group_shape: Tuple[int,

19 +def _normalize_quant_group_shape(x: torch.Tensor, group_shape: tuple[int,

20 int]):

21 # -1 means full extent

22 return (group_shape[0] if group_shape[0] > 0 else x.shape[-2],

@@ -56,9 +57,9 @@ def group_broadcast(t, shape):

57 # (i.e. per-token-per-group)

58 def scaled_quantize(

59 x: torch.Tensor,

59 - group_shape: Tuple[int, int],

60 + group_shape: tuple[int, int],

61 quant_dtype: torch.dtype,

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

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

63 group_shape = _normalize_quant_group_shape(x, group_shape)

64 assert quant_dtype.is_floating_point, \

65 "currently `scaled_quantize` only supports floating point dtypes " \

@@ -97,9 +98,9 @@ def scaled_quantize(

98 def scaled_dequantize(

99 x_q: torch.Tensor,

100 x_s: torch.Tensor,

100 - group_shape: Optional[Tuple[int, int]] = None,

101 + group_shape: Optional[tuple[int, int]] = None,

102 out_dtype: torch.dtype = torch.float32,

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

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

104 if group_shape is not None:

105 group_shape = _normalize_quant_group_shape(x_q, group_shape)

106

@@ -173,8 +174,8 @@ def unpack_quantized_values_into_int32(w_q: torch.Tensor,

174

175 def is_layer_skipped(

176 prefix: str,

176 - ignored_layers: List[str],

177 - fused_mapping: Mapping[str, List[str]] = MappingProxyType({})

177 + ignored_layers: list[str],

178 + fused_mapping: Mapping[str, list[str]] = MappingProxyType({})

179 ) -> bool:

180 # prefix: model.layers.0.self_attn.q_proj

181 # proj_name: q_proj

vllm/model_executor/layers/quantization/utils/w8a8_utils.py

+9 -9

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Callable, List, Optional, Tuple, Union

3 +from typing import Callable, Optional, Union

4

5 import torch

6

@@ -81,7 +81,7 @@ def all_close_1d(x: torch.Tensor) -> bool:

81

82 def convert_to_channelwise(

83 weight_scale: torch.Tensor,

84 - logical_widths: List[int]) -> Tuple[torch.Tensor, torch.Tensor]:

84 + logical_widths: list[int]) -> tuple[torch.Tensor, torch.Tensor]:

85 # Create channelwise buffer

86 weight_scale_channel = torch.empty((sum(logical_widths), 1),

87 dtype=torch.float32,

@@ -99,7 +99,7 @@ def convert_to_channelwise(

99

100 def requantize_with_max_scale(

101 weight: torch.Tensor, weight_scale: torch.Tensor,

102 - logical_widths: List[int]) -> Tuple[torch.Tensor, torch.Tensor]:

102 + logical_widths: list[int]) -> tuple[torch.Tensor, torch.Tensor]:

103 # Max scale to be used for requanitzation.

104 max_w_scale = weight_scale.max()

105

@@ -136,7 +136,7 @@ def maybe_create_device_identity():

136 def cutlass_w8a8_scaled_mm(*, qinput: torch.Tensor, weight: torch.Tensor,

137 out_dtype: torch.dtype, scale_a: torch.Tensor,

138 scale_b: torch.Tensor, bias: torch.Tensor,

139 - output_shape: List, **kwargs) -> torch.Tensor:

139 + output_shape: list, **kwargs) -> torch.Tensor:

140

141 # Fused GEMM_DQ

142 output = ops.cutlass_scaled_mm(qinput,

@@ -154,7 +154,7 @@ def rocm_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,

154 scale_a: torch.Tensor,

155 scale_b: torch.Tensor, bias: torch.Tensor,

156 input_2d: torch.Tensor,

157 - output_shape: List) -> torch.Tensor:

157 + output_shape: list) -> torch.Tensor:

158 from vllm.platforms.rocm import on_mi250_mi300

159 if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi250_mi300(

160 ) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0:

@@ -177,7 +177,7 @@ def torch_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,

177 scale_a: torch.Tensor,

178 scale_b: torch.Tensor, bias: torch.Tensor,

179 input_2d: torch.Tensor,

180 - output_shape: List) -> torch.Tensor:

180 + output_shape: list) -> torch.Tensor:

181 output = torch._scaled_mm(qinput,

182 weight,

183 out_dtype=out_dtype,

@@ -198,7 +198,7 @@ def torch_per_token_w8a8_scaled_mm(*, qinput: torch.Tensor,

198 scale_a: torch.Tensor,

199 scale_b: torch.Tensor, bias: torch.Tensor,

200 input_2d: torch.Tensor,

201 - output_shape: List) -> torch.Tensor:

201 + output_shape: list) -> torch.Tensor:

202 # Note: Callers of this function should check USE_ROWWISE_TORCH_SCALED_MM

203 # when using it.

204 # For now it has only been validated on ROCm platform.

@@ -228,7 +228,7 @@ def torch_channelwise_w8a8_scaled_mm(*, qinput: torch.Tensor,

228 scale_a: torch.Tensor,

229 scale_b: torch.Tensor, bias: torch.Tensor,

230 input_2d: torch.Tensor,

231 - output_shape: List,

231 + output_shape: list,

232 **kwargs) -> torch.Tensor:

233 # Use unfused DQ due to limitations with scaled_mm

234

@@ -384,7 +384,7 @@ def normalize_e4m3fn_to_e4m3fnuz(

384 weight: torch.Tensor,

385 weight_scale: torch.Tensor,

386 input_scale: Optional[torch.Tensor] = None

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

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

388 assert weight.dtype == torch.float8_e4m3fn

389 # The bits pattern 10000000(-128) represents zero in e4m3fn

390 # but NaN in e4m3fnuz. So here we set it to 0.

vllm/model_executor/layers/rejection_sampler.py

+8 -8

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

2

3 from functools import cached_property

4 from importlib.util import find_spec

5 -from typing import Dict, Optional, Tuple

5 +from typing import Optional

6

7 import torch

8 import torch.jit

@@ -65,7 +65,7 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):

65 bonus_token_ids: torch.Tensor,

66 draft_probs: torch.Tensor,

67 draft_token_ids: torch.Tensor,

68 - seeded_seqs: Optional[Dict[int, torch.Generator]] = None,

68 + seeded_seqs: Optional[dict[int, torch.Generator]] = None,

69 ) -> torch.Tensor:

70 """Sample token ids using rejection sampling. This accepts or rejects

71 tokens proposed by the draft model using the probability of each token

@@ -161,8 +161,8 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):

161 target_probs: torch.Tensor, # [batch_size, k, vocab_size]

162 draft_probs: torch.Tensor, # [batch_size, k, vocab_size]

163 draft_token_ids: torch.Tensor, # [batch_size, k]

164 - seeded_seqs: Optional[Dict[int, torch.Generator]],

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

164 + seeded_seqs: Optional[dict[int, torch.Generator]],

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

166 """Perform modified rejection sampling on each sequence.

167

168 Returns:

@@ -194,7 +194,7 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):

194 return accepted, recovered_token_ids

195

196 def _create_uniform_samples(self,

197 - seeded_seqs: Optional[Dict[int,

197 + seeded_seqs: Optional[dict[int,

198 torch.Generator]],

199 batch_size: int, k: int,

200 device: torch.device) -> torch.Tensor:

@@ -210,7 +210,7 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):

210 a seed.

211

212 Args:

213 - seeded_seqs : Optional[Dict[int, torch.Generator]]

213 + seeded_seqs : Optional[dict[int, torch.Generator]]

214 A dictionary mapping indices in the batch to

215 `torch.Generator` objects. If `None`, all samples are

216 generated without a seed.

@@ -255,7 +255,7 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):

255 target_probs: torch.Tensor, # [batch_size, k, vocab_size]

256 draft_probs: torch.Tensor, # [batch_size, k, vocab_size]

257 draft_token_ids: torch.Tensor, # [batch_size, k]

258 - seeded_seqs: Optional[Dict[int, torch.Generator]],

258 + seeded_seqs: Optional[dict[int, torch.Generator]],

259 ) -> torch.Tensor:

260 r"""Create bool matrix over the proposed draft tokens. If

261 True, then a token can be accepted, else it should be

@@ -379,7 +379,7 @@ def _multinomial(

379 probs: torch.Tensor,

380 num_samples: int,

381 k: int,

382 - seeded_seqs: Dict[int, torch.Generator],

382 + seeded_seqs: dict[int, torch.Generator],

383 ) -> torch.Tensor:

384

385 if num_samples > 1:

vllm/model_executor/layers/resampler.py

+5 -5

@@ -33,7 +33,7 @@ Example models: Qwen (Qwen-VL), MiniCPM-V 2.0

33 """

34 import math

35 from functools import partial

36 -from typing import Callable, Optional, Tuple, Union

36 +from typing import Callable, Optional, Union

37

38 import numpy as np

39 import torch

@@ -69,7 +69,7 @@ def get_abs_pos(abs_pos: torch.Tensor, tgt_size: Union[torch.Tensor,

69 # https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20

70 def get_1d_sincos_pos_embed_from_grid(

71 embed_dim: int, pos: np.ndarray,

72 - version: Tuple[int, int] = (2, 0)) -> torch.Tensor:

72 + version: tuple[int, int] = (2, 0)) -> torch.Tensor:

73 """

74 embed_dim: output dimension for each position

75 pos: a list of positions to be encoded: size (M,) / (H, W)

@@ -96,7 +96,7 @@ def get_1d_sincos_pos_embed_from_grid(

96

97 def get_2d_sincos_pos_embed_from_grid(

98 embed_dim: int, grid: np.ndarray,

99 - version: Tuple[int, int] = (2, 0)) -> torch.Tensor:

99 + version: tuple[int, int] = (2, 0)) -> torch.Tensor:

100 assert embed_dim % 2 == 0

101

102 # use half of dimensions to encode grid_h

@@ -114,9 +114,9 @@ def get_2d_sincos_pos_embed_from_grid(

114

115 def get_2d_sincos_pos_embed(

116 embed_dim: int,

117 - grid_size: Union[int, Tuple[int, int]],

117 + grid_size: Union[int, tuple[int, int]],

118 cls_token: bool = False,

119 - version: Tuple[int, int] = (2, 0),

119 + version: tuple[int, int] = (2, 0),

120 ) -> torch.Tensor:

121 """

122 grid_size: int of the grid height and width

vllm/model_executor/layers/rotary_embedding.py

+53 -53

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

23 # limitations under the License.

24 """Rotary Positional Embeddings."""

25 import math

26 -from typing import Any, Dict, List, Optional, Tuple, Union

26 +from typing import Any, Optional, Union

27

28 import torch

29 import torch.nn as nn

@@ -140,7 +140,7 @@ class RotaryEmbedding(CustomOp):

140 query: torch.Tensor,

141 key: Optional[torch.Tensor] = None,

142 offsets: Optional[torch.Tensor] = None,

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

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

144 """A PyTorch-native implementation of forward()."""

145 if offsets is not None:

146 positions = positions + offsets

@@ -174,7 +174,7 @@ class RotaryEmbedding(CustomOp):

174 query: torch.Tensor,

175 key: Optional[torch.Tensor] = None,

176 offsets: Optional[torch.Tensor] = None,

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

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

178 from vllm import _custom_ops as ops

179

180 # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)

@@ -202,7 +202,7 @@ class RotaryEmbedding(CustomOp):

202 query: torch.Tensor,

203 key: Optional[torch.Tensor] = None,

204 offsets: Optional[torch.Tensor] = None,

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

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

206 from vllm._ipex_ops import ipex_ops as ops

207

208 self.cos_sin_cache = self.cos_sin_cache.to(positions.device,

@@ -232,7 +232,7 @@ class RotaryEmbedding(CustomOp):

232 query: torch.Tensor,

233 key: Optional[torch.Tensor] = None,

234 offsets: Optional[torch.Tensor] = None,

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

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

236 from habana_frameworks.torch.hpex.kernels import (

237 RotaryPosEmbeddingMode, apply_rotary_pos_emb)

238 if offsets is not None:

@@ -290,7 +290,7 @@ class RotaryEmbedding(CustomOp):

290 query: torch.Tensor,

291 key: Optional[torch.Tensor] = None,

292 offsets: Optional[torch.Tensor] = None,

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

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

294

295 def _apply_rotary_emb_neuron(

296 x: torch.Tensor,

@@ -406,23 +406,23 @@ class LinearScalingRotaryEmbedding(RotaryEmbedding):

406 max_position_embeddings: int,

407 base: int,

408 is_neox_style: bool,

409 - scaling_factors: Union[List[float], float],

409 + scaling_factors: Union[list[float], float],

410 dtype: torch.dtype,

411 ) -> None:

412 if isinstance(scaling_factors, float):

413 scaling_factors = [scaling_factors]

414 - self.scaling_factors: List[float] = scaling_factors # noqa

414 + self.scaling_factors: list[float] = scaling_factors # noqa

415 super().__init__(head_size, rotary_dim, max_position_embeddings, base,

416 is_neox_style, dtype)

417 # Lazy initialized.

418 - self._scaling_factor_to_offset: Dict[float, int]

418 + self._scaling_factor_to_offset: dict[float, int]

419

420 def _compute_cos_sin_cache(self) -> torch.Tensor:

421 inv_freq = self._compute_inv_freq(self.base)

422 - cache_list: List[torch.Tensor] = []

422 + cache_list: list[torch.Tensor] = []

423 # offsets to the next cache in a tensor.

424 # Each offset corresponds to the same index in scaling_factors.

425 - offsets: List[int] = []

425 + offsets: list[int] = []

426 for scaling_factor in self.scaling_factors:

427 # NOTE(woosuk): self.max_position_embeddings is the original

428 # maximum length before applying the rope scaling.

@@ -452,7 +452,7 @@ class LinearScalingRotaryEmbedding(RotaryEmbedding):

452 return torch.cat(cache_list, dim=0)

453

454 @property

455 - def scaling_factor_to_offset(self) -> Dict[float, int]:

455 + def scaling_factor_to_offset(self) -> dict[float, int]:

456 return self._scaling_factor_to_offset

457

458

@@ -512,7 +512,7 @@ def _yarn_find_correction_range(

512 high_rot: int,

513 dim: int,

514 base: float = 10000,

515 - max_position_embeddings: int = 2048) -> Tuple[int, int]:

515 + max_position_embeddings: int = 2048) -> tuple[int, int]:

516 low = math.floor(

517 _yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))

518 high = math.ceil(

@@ -613,8 +613,8 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):

613 base: int,

614 is_neox_style: bool,

615 dtype: torch.dtype,

616 - short_factor: List[float],

617 - long_factor: List[float],

616 + short_factor: list[float],

617 + long_factor: list[float],

618 short_mscale: Optional[float] = None,

619 long_mscale: Optional[float] = None,

620 ):

@@ -662,7 +662,7 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):

662 long_short_cache,

663 persistent=False)

664

665 - def _compute_inv_freq(self, rescale_factors: List[float]) -> torch.Tensor:

665 + def _compute_inv_freq(self, rescale_factors: list[float]) -> torch.Tensor:

666 rescale_factors = torch.tensor(rescale_factors, dtype=torch.float32)

667 inv_freq = 1.0 / (rescale_factors * (self.base**(torch.arange(

668 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim)))

@@ -671,7 +671,7 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):

671 def _compute_cos_sin_cache(

672 self,

673 max_position_embeddings: int,

674 - rescale_factors: List[float],

674 + rescale_factors: list[float],

675 mscale: float,

676 ) -> torch.Tensor:

677 inv_freq = self._compute_inv_freq(rescale_factors)

@@ -688,7 +688,7 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):

688 query: torch.Tensor,

689 key: Optional[torch.Tensor] = None,

690 offsets: Optional[torch.Tensor] = None,

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

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

692 assert key is not None

693 query = query.view(*query.shape[:-1], -1, self.head_size)

694 key = key.view(*key.shape[:-1], -1, self.head_size)

@@ -799,7 +799,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):

799 query: torch.Tensor,

800 key: Optional[torch.Tensor] = None,

801 offsets: Optional[torch.Tensor] = None,

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

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

803 """PyTorch-native implementation equivalent to forward()."""

804 assert key is not None

805 query_rot = query[..., :self.rotary_dim]

@@ -930,7 +930,7 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):

930 self,

931 query: torch.Tensor,

932 key: Optional[torch.Tensor] = None,

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

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

934 assert key is not None

935 self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)

936 query_ = torch.view_as_complex(query.float().reshape(

@@ -958,7 +958,7 @@ class MRotaryEmbedding(RotaryEmbedding):

958 base: int,

959 is_neox_style: bool,

960 dtype: torch.dtype,

961 - mrope_section: Optional[List[int]] = None,

961 + mrope_section: Optional[list[int]] = None,

962 ) -> None:

963 # In Qwen2.5-VL, the maximum index value is related to the duration of

964 # the input video. We enlarge max_position_embeddings to 4 times to get

@@ -976,7 +976,7 @@ class MRotaryEmbedding(RotaryEmbedding):

976 positions: torch.Tensor,

977 query: torch.Tensor,

978 key: Optional[torch.Tensor] = None,

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

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

980 """PyTorch-native implementation equivalent to forward().

981

982 Args:

@@ -1024,16 +1024,16 @@ class MRotaryEmbedding(RotaryEmbedding):

1024 @classmethod

1025 def get_input_positions(

1026 cls,

1027 - input_tokens: List[int],

1027 + input_tokens: list[int],

1028 hf_config: PretrainedConfig,

1029 - image_grid_thw: Optional[Union[List[List[int]], torch.Tensor]],

1030 - video_grid_thw: Optional[Union[List[List[int]], torch.Tensor]],

1031 - second_per_grid_ts: Optional[List[float]],

1029 + image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

1030 + video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

1031 + second_per_grid_ts: Optional[list[float]],

1032 context_len: int = 0,

1033 seq_len: Optional[int] = None,

1034 audio_feature_lengths: Optional[torch.Tensor] = None,

1035 use_audio_in_video: bool = False,

1036 - ) -> Tuple[List[List[int]], int]:

1036 + ) -> tuple[list[list[int]], int]:

1037 """Get mrope input positions and delta value."""

1038

1039 image_grid_thw = [] if image_grid_thw is None else image_grid_thw

@@ -1059,16 +1059,16 @@ class MRotaryEmbedding(RotaryEmbedding):

1059 @classmethod

1060 def get_input_positions_tensor(

1061 cls,

1062 - input_tokens: List[int],

1062 + input_tokens: list[int],

1063 hf_config: PretrainedConfig,

1064 - image_grid_thw: Union[List[List[int]], torch.Tensor],

1065 - video_grid_thw: Union[List[List[int]], torch.Tensor],

1066 - second_per_grid_ts: List[float],

1064 + image_grid_thw: Union[list[list[int]], torch.Tensor],

1065 + video_grid_thw: Union[list[list[int]], torch.Tensor],

1066 + second_per_grid_ts: list[float],

1067 context_len: int = 0,

1068 seq_len: Optional[int] = None,

1069 audio_feature_lengths: Optional[torch.Tensor] = None,

1070 use_audio_in_video: bool = False,

1071 - ) -> Tuple[torch.Tensor, int]:

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

1072 from vllm.transformers_utils.config import thinker_uses_mrope

1073 if thinker_uses_mrope(hf_config):

1074 return cls._omni_get_input_positions_tensor(

@@ -1096,14 +1096,14 @@ class MRotaryEmbedding(RotaryEmbedding):

1096 @classmethod

1097 def _vl_get_input_positions_tensor(

1098 cls,

1099 - input_tokens: List[int],

1099 + input_tokens: list[int],

1100 hf_config: PretrainedConfig,

1101 - image_grid_thw: Union[List[List[int]], torch.Tensor],

1102 - video_grid_thw: Union[List[List[int]], torch.Tensor],

1103 - second_per_grid_ts: List[float],

1101 + image_grid_thw: Union[list[list[int]], torch.Tensor],

1102 + video_grid_thw: Union[list[list[int]], torch.Tensor],

1103 + second_per_grid_ts: list[float],

1104 context_len: int = 0,

1105 seq_len: Optional[int] = None,

1106 - ) -> Tuple[torch.Tensor, int]:

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

1107 """Get mrope input positions and delta value."""

1108

1109 image_token_id = hf_config.image_token_id

@@ -1195,16 +1195,16 @@ class MRotaryEmbedding(RotaryEmbedding):

1195 @classmethod

1196 def _omni_get_input_positions_tensor(

1197 cls,

1198 - input_tokens: List[int],

1198 + input_tokens: list[int],

1199 hf_config: PretrainedConfig,

1200 - image_grid_thw: Union[List[List[int]], torch.Tensor],

1201 - video_grid_thw: Union[List[List[int]], torch.Tensor],

1202 - second_per_grid_ts: Optional[List[float]] = None,

1200 + image_grid_thw: Union[list[list[int]], torch.Tensor],

1201 + video_grid_thw: Union[list[list[int]], torch.Tensor],

1202 + second_per_grid_ts: Optional[list[float]] = None,

1203 context_len: int = 0,

1204 seq_len: Optional[int] = None,

1205 audio_feature_lengths: Optional[torch.Tensor] = None,

1206 use_audio_in_video: bool = False,

1207 - ) -> Tuple[torch.Tensor, int]:

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

1208 """Get mrope input positions and delta value (Qwen2.5-Omni version).

1209

1210 Differences from MRotaryEmbedding:

@@ -1329,7 +1329,7 @@ class MRotaryEmbedding(RotaryEmbedding):

1329 place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2

1330 pure_audio_len = place_num - 2

1331 added_audio_len = 0

1332 - audio_llm_pos_ids_list: List[torch.Tensor] = []

1332 + audio_llm_pos_ids_list: list[torch.Tensor] = []

1333 for t_chunk in t_index_split_chunk:

1334 vision_ntoken_per_chunk = len(

1335 t_chunk) * grid_h * grid_w // (spatial_merge_size**2)

@@ -1382,7 +1382,7 @@ class MRotaryEmbedding(RotaryEmbedding):

1382 start_idx: int,

1383 vision_idx: int,

1384 spatial_merge_size: int,

1385 - t_index: List[int],

1385 + t_index: list[int],

1386 grid_hs: torch.Tensor,

1387 grid_ws: torch.Tensor,

1388 ) -> torch.Tensor:

@@ -1402,8 +1402,8 @@ class MRotaryEmbedding(RotaryEmbedding):

1402

1403 @staticmethod

1404 def _split_list_into_ranges(lst: torch.Tensor,

1405 - interval: int) -> List[List[int]]:

1406 - ranges: List[List[int]] = [[]

1405 + interval: int) -> list[list[int]]:

1406 + ranges: list[list[int]] = [[]

1407 for _ in range((max(lst) // interval) + 1)]

1408 for num in lst:

1409 index = num // interval

@@ -1415,7 +1415,7 @@ class MRotaryEmbedding(RotaryEmbedding):

1415 mrope_position_delta: int,

1416 context_len: int,

1417 seq_len: int,

1418 - ) -> List[List[int]]:

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

1419 return [

1420 list(

1421 range(context_len + mrope_position_delta,

@@ -1438,9 +1438,9 @@ class MRotaryEmbedding(RotaryEmbedding):

1438 cls,

1439 thinker_config: PretrainedConfig,

1440 audio_len: int,

1441 - video_grid_thw: Union[List[int], torch.Tensor],

1441 + video_grid_thw: Union[list[int], torch.Tensor],

1442 video_second_per_grid_t: float,

1443 - ) -> List[int]:

1443 + ) -> list[int]:

1444 """Get video prompt updates when `use_audio_in_video` is True.

1445

1446 In this case, audio and vision update ids will be split into

@@ -1593,7 +1593,7 @@ class DualChunkRotaryEmbedding(CustomOp):

1593 query: torch.Tensor,

1594 key: torch.Tensor,

1595 offsets: Optional[torch.Tensor] = None,

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

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

1597 query = query.view(*query.shape[:-1], -1, self.head_size)

1598 key = key.view(*key.shape[:-1], -1, self.head_size)

1599 query_rot = query[..., :self.rotary_dim]

@@ -1664,7 +1664,7 @@ class DualChunkRotaryEmbedding(CustomOp):

1664 return s

1665

1666

1667 -_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {}

1667 +_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}

1668

1669

1670 def get_rope(

@@ -1673,10 +1673,10 @@ def get_rope(

1673 max_position: int,

1674 base: int,

1675 is_neox_style: bool = True,

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

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

1677 dtype: Optional[torch.dtype] = None,

1678 partial_rotary_factor: float = 1.0,

1679 - dual_chunk_attention_config: Optional[Dict[str, Any]] = None,

1679 + dual_chunk_attention_config: Optional[dict[str, Any]] = None,

1680 ) -> RotaryEmbedding:

1681 if dtype is None:

1682 dtype = torch.get_default_dtype()

vllm/model_executor/layers/sampler.py

+28 -27

@@ -2,10 +2,11 @@

2 """A layer that samples the next tokens from the model's outputs."""

3 import itertools

4 import warnings

5 +from collections.abc import Iterator

6 from dataclasses import dataclass

7 from importlib.util import find_spec

8 from math import inf

8 -from typing import Dict, Iterator, List, Optional, Tuple, Union

9 +from typing import Optional, Union

10

11 import msgspec

12 import torch

@@ -42,14 +43,14 @@ def get_sampler() -> torch.nn.Module:

43

44

45 # (num_token_ids, num_parent_ids) per sequence group.

45 -SampleResultType = List[Tuple[List[int], List[int]]]

46 +SampleResultType = list[tuple[list[int], list[int]]]

47

48 # Types of temporary data structures used for

49 # computing sample_result

49 -SampleMetadataType = Dict[SamplingType, Tuple[List[int],

50 - List[SequenceGroupToSample]]]

51 -MultinomialSamplesType = Dict[SamplingType, torch.Tensor]

52 -SampleResultsDictType = Dict[int, Tuple[List[int], List[int]]]

50 +SampleMetadataType = dict[SamplingType, tuple[list[int],

51 + list[SequenceGroupToSample]]]

52 +MultinomialSamplesType = dict[SamplingType, torch.Tensor]

53 +SampleResultsDictType = dict[int, tuple[list[int], list[int]]]

54

55

56 # Encapsulates temporary data structures for computing

@@ -76,7 +77,7 @@ class SampleResultArgsType:

77 MaybeDeferredSampleResultType = Union[SampleResultType, SampleResultArgsType]

78

79 # Abbreviation of the _sample() return type

79 -SampleReturnType = Tuple[MaybeDeferredSampleResultType, Optional[torch.Tensor]]

80 +SampleReturnType = tuple[MaybeDeferredSampleResultType, Optional[torch.Tensor]]

81

82

83 class SamplerOutput(

@@ -90,7 +91,7 @@ class SamplerOutput(

91 also has optional fields for device tensors.

92 """

93

93 - outputs: List[CompletionSequenceGroupOutput]

94 + outputs: list[CompletionSequenceGroupOutput]

95

96 # On-device tensor containing probabilities of each token.

97 sampled_token_probs: Optional[torch.Tensor] = None

@@ -350,7 +351,7 @@ def _apply_min_tokens_penalty(

351 have not been generated yet

352 """

353 # list of indices in logits that will be set to -inf

353 - logits_to_penalize: List[Tuple[int, int]] = []

354 + logits_to_penalize: list[tuple[int, int]] = []

355 logits_applied = 0

356 for seq_group in sampling_metadata.seq_groups:

357 seq_ids = seq_group.seq_ids

@@ -366,7 +367,7 @@ def _apply_min_tokens_penalty(

367 min_tokens = sampling_params.min_tokens

368 token_ids_to_penalize = sampling_params.all_stop_token_ids

369 if min_tokens > 0 and token_ids_to_penalize:

369 - seqs_to_penalize: List[int] = []

370 + seqs_to_penalize: list[int] = []

371 for j, seq_id in enumerate(seq_ids):

372 seq_data = seq_group.seq_data[seq_id]

373 if len(seq_data.output_token_ids_array) < min_tokens:

@@ -436,7 +437,7 @@ def _apply_min_p(

437

438

439 def _greedy_sample(

439 - selected_seq_groups: List[SequenceGroupToSample],

440 + selected_seq_groups: list[SequenceGroupToSample],

441 samples: torch.Tensor,

442 ) -> SampleResultType:

443 """Run greedy sampling on a given samples.

@@ -471,7 +472,7 @@ def _greedy_sample(

472

473

474 def _random_sample(

474 - selected_seq_groups: List[SequenceGroupToSample],

475 + selected_seq_groups: list[SequenceGroupToSample],

476 random_samples: torch.Tensor,

477 ) -> SampleResultType:

478 """Run random sampling on a given samples.

@@ -522,7 +523,7 @@ def _random_sample(

523 def _multinomial(

524 probs: torch.Tensor,

525 num_samples: int,

525 - seq_groups: Optional[List[SequenceGroupToSample]] = None,

526 + seq_groups: Optional[list[SequenceGroupToSample]] = None,

527 ) -> torch.Tensor:

528 if num_samples > 1:

529 probs = probs.repeat_interleave(num_samples, dim=0)

@@ -543,7 +544,7 @@ def _multinomial(

544

545 def _top_k_top_p_multinomial_with_flashinfer(

546 probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor,

546 - num_samples: int, seq_groups: Optional[List[SequenceGroupToSample]]):

547 + num_samples: int, seq_groups: Optional[list[SequenceGroupToSample]]):

548 max_top_k_round = 32

549 if num_samples > 1:

550 probs = probs.repeat_interleave(num_samples, dim=0)

@@ -648,7 +649,7 @@ def _sample_with_torch(

649 tensors required for Pythonization

650 '''

651

651 - categorized_seq_group_ids: Dict[SamplingType, List[int]] = {

652 + categorized_seq_group_ids: dict[SamplingType, list[int]] = {

653 t: []

654 for t in SamplingType

655 }

@@ -812,7 +813,7 @@ def get_logprobs(

813 logprobs: torch.Tensor,

814 sampling_metadata: SamplingMetadata,

815 sample_results: SampleResultType,

815 -) -> Tuple[List[Optional[PromptLogprobs]], List[SampleLogprobs]]:

816 +) -> tuple[list[Optional[PromptLogprobs]], list[SampleLogprobs]]:

817 """Return sample logprobs and prompt logprobs.

818

819 The logic consists of 3 parts.

@@ -841,9 +842,9 @@ def get_logprobs(

842 """

843 # The index of query token to calculate logprobs. It includes both

844 # prompt and sample logprob indices.

844 - query_indices: List[int] = []

845 + query_indices: list[int] = []

846 # The next token ids to get the logprob value from.

846 - next_token_ids: List[int] = []

847 + next_token_ids: list[int] = []

848 # The largest requested number of logprobs. We find logprobs as many as the

849 # largest num logprobs in this API. If every logprobs is None, it will be

850 # set to -1.

@@ -925,8 +926,8 @@ def get_logprobs(

926 ranks = ranks.to('cpu')

927

928 # Find prompt/sample logprobs.

928 - prompt_logprobs_per_seq_group: List[Optional[PromptLogprobs]] = []

929 - sample_logprobs_per_seq_group: List[SampleLogprobs] = []

929 + prompt_logprobs_per_seq_group: list[Optional[PromptLogprobs]] = []

930 + sample_logprobs_per_seq_group: list[SampleLogprobs] = []

931 top_logprob_idx = 0

932 selected_logprobs_idx = 0

933

@@ -977,7 +978,7 @@ def _get_prompt_logprob_if_needed(

978 for idx, token_id in enumerate(next_prompt_tokens):

979 # Calculate the prompt logprob of the real prompt tokens.

980 # {token_id: (logprob, rank_from_vocab)}

980 - prompt_logprobs_dict: Dict[int, Tuple[float, int]] = {

981 + prompt_logprobs_dict: dict[int, tuple[float, int]] = {

982 token_id: (selected_logprob_items[idx], rank_items[idx])

983 }

984

@@ -1009,7 +1010,7 @@ def _get_prompt_logprob_if_needed(

1010

1011 def _get_sampled_logprob_if_needed(

1012 seq_group: SequenceGroupToSample,

1012 - sample_result: Tuple[List[int], List[int]],

1013 + sample_result: tuple[list[int], list[int]],

1014 selected_logprobs: torch.Tensor,

1015 ranks: torch.Tensor,

1016 top_token_ids: torch.Tensor,

@@ -1130,9 +1131,9 @@ def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Tensor,

1131 def _build_sampler_output(

1132 maybe_deferred_sample_results: MaybeDeferredSampleResultType,

1133 sampling_metadata: SamplingMetadata,

1133 - prompt_logprobs: Optional[List[Optional[PromptLogprobs]]],

1134 - sample_logprobs: Optional[List[SampleLogprobs]],

1135 - on_device_tensors: Optional[Tuple[torch.Tensor, torch.Tensor,

1134 + prompt_logprobs: Optional[list[Optional[PromptLogprobs]]],

1135 + sample_logprobs: Optional[list[SampleLogprobs]],

1136 + on_device_tensors: Optional[tuple[torch.Tensor, torch.Tensor,

1137 torch.Tensor]],

1138 skip_sampler_cpu_output: bool = False,

1139 ) -> SamplerOutput:

@@ -1144,7 +1145,7 @@ def _build_sampler_output(

1145 allows post-processing without copies to CPU/serialization, e.g. in

1146 speculative decoding rejection sampling.

1147 """

1147 - sampler_output: List[CompletionSequenceGroupOutput] = []

1148 + sampler_output: list[CompletionSequenceGroupOutput] = []

1149

1150 if skip_sampler_cpu_output:

1151 assert isinstance(maybe_deferred_sample_results, SampleResultArgsType)

@@ -1166,7 +1167,7 @@ def _build_sampler_output(

1167 prompt_logprobs, sample_logprobs):

1168 seq_ids = seq_group.seq_ids

1169 next_token_ids, parent_ids = sample_result

1169 - seq_outputs: List[SequenceOutput] = []

1170 + seq_outputs: list[SequenceOutput] = []

1171 for parent_id, next_token_id, logprobs in zip(

1172 parent_ids, next_token_ids, group_sample_logprobs):

1173 seq_outputs.append(

vllm/model_executor/layers/spec_decode_base_sampler.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 from abc import abstractmethod

4 -from typing import Dict, Optional, Union

4 +from typing import Optional, Union

5

6 import torch

7 import torch.jit

@@ -253,6 +253,6 @@ class SpecDecodeStochasticBaseSampler(SpecDecodeBaseSampler):

253 bonus_token_ids: torch.Tensor,

254 draft_probs: torch.Tensor,

255 draft_token_ids: torch.Tensor,

256 - seeded_seqs: Optional[Dict[int, torch.Generator]] = None,

256 + seeded_seqs: Optional[dict[int, torch.Generator]] = None,

257 ) -> torch.Tensor:

258 raise NotImplementedError

vllm/model_executor/layers/utils.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Utility methods for model layers."""

3 -from typing import Callable, Optional, Tuple

3 +from typing import Callable, Optional

4

5 import torch

6

@@ -13,7 +13,7 @@ def get_token_bin_counts_and_mask(

13 tokens: torch.Tensor,

14 vocab_size: int,

15 num_seqs: int,

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

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

17 # Compute the bin counts for the tokens.

18 # vocab_size + 1 for padding.

19 bin_counts = torch.zeros((num_seqs, vocab_size + 1),

vllm/model_executor/layers/vocab_parallel_embedding.py

+8 -7

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

1 # SPDX-License-Identifier: Apache-2.0

2

3 +from collections.abc import Sequence

4 from dataclasses import dataclass

4 -from typing import List, Optional, Sequence, Tuple

5 +from typing import Optional

6

7 import torch

8 import torch.nn.functional as F

@@ -25,7 +26,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):

26

27 def create_weights(self, layer: torch.nn.Module,

28 input_size_per_partition: int,

28 - output_partition_sizes: List[int], input_size: int,

29 + output_partition_sizes: list[int], input_size: int,

30 output_size: int, params_dtype: torch.dtype,

31 **extra_weight_attrs):

32 """Create weights for embedding layer."""

@@ -141,7 +142,7 @@ def get_masked_input_and_mask(

142 input_: torch.Tensor, org_vocab_start_index: int,

143 org_vocab_end_index: int, num_org_vocab_padding: int,

144 added_vocab_start_index: int,

144 - added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]:

145 + added_vocab_end_index: int) -> tuple[torch.Tensor, torch.Tensor]:

146 # torch.compile will fuse all of the pointwise ops below

147 # into a single kernel, making it very fast

148 org_vocab_mask = (input_ >= org_vocab_start_index) & (

@@ -298,7 +299,7 @@ class VocabParallelEmbedding(torch.nn.Module):

299 org_vocab_start_index, org_vocab_end_index,

300 added_vocab_start_index, added_vocab_end_index)

301

301 - def get_sharded_to_full_mapping(self) -> Optional[List[int]]:

302 + def get_sharded_to_full_mapping(self) -> Optional[list[int]]:

303 """Get a mapping that can be used to reindex the gathered

304 logits for sampling.

305

@@ -312,9 +313,9 @@ class VocabParallelEmbedding(torch.nn.Module):

313 if self.tp_size < 2:

314 return None

315

315 - base_embeddings: List[int] = []

316 - added_embeddings: List[int] = []

317 - padding: List[int] = []

316 + base_embeddings: list[int] = []

317 + added_embeddings: list[int] = []

318 + padding: list[int] = []

319 for tp_rank in range(self.tp_size):

320 shard_indices = self._get_indices(self.num_embeddings_padded,

321 self.org_vocab_size_padded,

Update deprecated type hinting in `model_executor/layers` (#18056)