Update `Optional[x]` -> `x | None` and `Union[x, y]` to `x | y` (#26633)

benchmarks/backend_request_func.py

+13 -14

@@ -8,7 +8,6 @@ import sys

8 import time

9 import traceback

10 from dataclasses import dataclass, field

11 -from typing import Optional, Union

11

12 import aiohttp

13 import huggingface_hub.constants

@@ -28,13 +27,13 @@ class RequestFuncInput:

27 prompt_len: int

28 output_len: int

29 model: str

31 - model_name: Optional[str] = None

32 - logprobs: Optional[int] = None

33 - extra_body: Optional[dict] = None

34 - multi_modal_content: Optional[dict | list[dict]] = None

30 + model_name: str | None = None

31 + logprobs: int | None = None

32 + extra_body: dict | None = None

33 + multi_modal_content: dict | list[dict] | None = None

34 ignore_eos: bool = False

36 - language: Optional[str] = None

37 - request_id: Optional[str] = None

35 + language: str | None = None

36 + request_id: str | None = None

37

38

39 @dataclass

@@ -52,7 +51,7 @@ class RequestFuncOutput:

51

52 async def async_request_tgi(

53 request_func_input: RequestFuncInput,

55 - pbar: Optional[tqdm] = None,

54 + pbar: tqdm | None = None,

55 ) -> RequestFuncOutput:

56 api_url = request_func_input.api_url

57 assert api_url.endswith("generate_stream")

@@ -133,7 +132,7 @@ async def async_request_tgi(

132

133 async def async_request_trt_llm(

134 request_func_input: RequestFuncInput,

136 - pbar: Optional[tqdm] = None,

135 + pbar: tqdm | None = None,

136 ) -> RequestFuncOutput:

137 api_url = request_func_input.api_url

138 assert api_url.endswith("generate_stream")

@@ -204,7 +203,7 @@ async def async_request_trt_llm(

203

204 async def async_request_deepspeed_mii(

205 request_func_input: RequestFuncInput,

207 - pbar: Optional[tqdm] = None,

206 + pbar: tqdm | None = None,

207 ) -> RequestFuncOutput:

208 api_url = request_func_input.api_url

209 assert api_url.endswith(("completions", "profile")), (

@@ -267,7 +266,7 @@ async def async_request_deepspeed_mii(

266

267 async def async_request_openai_completions(

268 request_func_input: RequestFuncInput,

270 - pbar: Optional[tqdm] = None,

269 + pbar: tqdm | None = None,

270 ) -> RequestFuncOutput:

271 api_url = request_func_input.api_url

272 assert api_url.endswith(("completions", "profile")), (

@@ -367,7 +366,7 @@ async def async_request_openai_completions(

366

367 async def async_request_openai_chat_completions(

368 request_func_input: RequestFuncInput,

370 - pbar: Optional[tqdm] = None,

369 + pbar: tqdm | None = None,

370 ) -> RequestFuncOutput:

371 api_url = request_func_input.api_url

372 assert api_url.endswith(("chat/completions", "profile")), (

@@ -476,7 +475,7 @@ async def async_request_openai_chat_completions(

475

476 async def async_request_openai_audio(

477 request_func_input: RequestFuncInput,

479 - pbar: Optional[tqdm] = None,

478 + pbar: tqdm | None = None,

479 ) -> RequestFuncOutput:

480 # Lazy import without PlaceholderModule to avoid vllm dep.

481 import soundfile

@@ -610,7 +609,7 @@ def get_tokenizer(

609 tokenizer_mode: str = "auto",

610 trust_remote_code: bool = False,

611 **kwargs,

613 -) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:

612 +) -> PreTrainedTokenizer | PreTrainedTokenizerFast:

613 if pretrained_model_name_or_path is not None and not os.path.exists(

614 pretrained_model_name_or_path

615 ):

benchmarks/benchmark_prefix_caching.py

+2 -3

@@ -32,7 +32,6 @@ import dataclasses

32 import json

33 import random

34 import time

35 -from typing import Optional

35

36 from transformers import PreTrainedTokenizerBase

37

@@ -80,7 +79,7 @@ def sample_requests_from_dataset(

79 num_requests: int,

80 tokenizer: PreTrainedTokenizerBase,

81 input_length_range: tuple[int, int],

83 - fixed_output_len: Optional[int],

82 + fixed_output_len: int | None,

83 ) -> list[Request]:

84 if fixed_output_len is not None and fixed_output_len < 4:

85 raise ValueError("output_len too small")

@@ -128,7 +127,7 @@ def sample_requests_from_random(

127 num_requests: int,

128 tokenizer: PreTrainedTokenizerBase,

129 input_length_range: tuple[int, int],

131 - fixed_output_len: Optional[int],

130 + fixed_output_len: int | None,

131 prefix_len: int,

132 ) -> list[Request]:

133 requests = []

benchmarks/benchmark_prioritization.py

+1 -2

@@ -7,7 +7,6 @@ import dataclasses

7 import json

8 import random

9 import time

10 -from typing import Optional

10

11 from transformers import AutoTokenizer, PreTrainedTokenizerBase

12

@@ -24,7 +23,7 @@ def sample_requests(

23 dataset_path: str,

24 num_requests: int,

25 tokenizer: PreTrainedTokenizerBase,

27 - fixed_output_len: Optional[int],

26 + fixed_output_len: int | None,

27 ) -> list[tuple[str, int, int, int]]:

28 if fixed_output_len is not None and fixed_output_len < 4:

29 raise ValueError("output_len too small")

benchmarks/benchmark_serving_structured_output.py

+3 -4

@@ -32,7 +32,6 @@ import uuid

32 import warnings

33 from collections.abc import AsyncGenerator

34 from dataclasses import dataclass

35 -from typing import Optional

35

36 import datasets

37 import numpy as np

@@ -316,7 +315,7 @@ def calculate_metrics(

315 tokenizer: PreTrainedTokenizerBase,

316 selected_percentile_metrics: list[str],

317 selected_percentiles: list[float],

319 - goodput_config_dict: Optional[dict[str, float]] = None,

318 + goodput_config_dict: dict[str, float] | None = None,

319 ) -> tuple[BenchmarkMetrics, list[int]]:

320 actual_output_lens: list[int] = []

321 total_input = 0

@@ -436,9 +435,9 @@ async def benchmark(

435 selected_percentile_metrics: list[str],

436 selected_percentiles: list[str],

437 ignore_eos: bool,

439 - max_concurrency: Optional[int],

438 + max_concurrency: int | None,

439 structured_output_ratio: float,

441 - goodput_config_dict: Optional[dict[str, float]] = None,

440 + goodput_config_dict: dict[str, float] | None = None,

441 ):

442 if backend in ASYNC_REQUEST_FUNCS:

443 request_func = ASYNC_REQUEST_FUNCS[backend]

benchmarks/benchmark_utils.py

+8 -8

@@ -6,7 +6,7 @@ import math

6 import os

7 import time

8 from types import TracebackType

9 -from typing import Any, Optional, Union

9 +from typing import Any

10

11

12 def convert_to_pytorch_benchmark_format(

@@ -92,7 +92,7 @@ class TimeCollector:

92 def __init__(self, scale: int) -> None:

93 self.cnt: int = 0

94 self._sum: int = 0

95 - self._max: Optional[int] = None

95 + self._max: int | None = None

96 self.scale = scale

97 self.start_time: int = time.monotonic_ns()

98

@@ -104,13 +104,13 @@ class TimeCollector:

104 else:

105 self._max = max(self._max, v)

106

107 - def avg(self) -> Union[float, str]:

107 + def avg(self) -> float | str:

108 return self._sum * 1.0 / self.cnt / self.scale if self.cnt > 0 else "N/A"

109

110 - def max(self) -> Union[float, str]:

110 + def max(self) -> float | str:

111 return self._max / self.scale if self._max else "N/A"

112

113 - def dump_avg_max(self) -> list[Union[float, str]]:

113 + def dump_avg_max(self) -> list[float | str]:

114 return [self.avg(), self.max()]

115

116 def __enter__(self) -> None:

@@ -118,8 +118,8 @@ class TimeCollector:

118

119 def __exit__(

120 self,

121 - exc_type: Optional[type[BaseException]],

122 - exc_value: Optional[BaseException],

123 - exc_traceback: Optional[TracebackType],

121 + exc_type: type[BaseException] | None,

122 + exc_value: BaseException | None,

123 + exc_traceback: TracebackType | None,

124 ) -> None:

125 self.collect(time.monotonic_ns() - self.start_time)

benchmarks/cutlass_benchmarks/sparse_benchmarks.py

+1 -2

@@ -6,8 +6,7 @@ import copy

6 import itertools

7 import pickle as pkl

8 import time

9 -from collections.abc import Iterable

10 -from typing import Callable

9 +from collections.abc import Callable, Iterable

10

11 import torch

12 import torch.utils.benchmark as TBenchmark

benchmarks/cutlass_benchmarks/w8a8_benchmarks.py

+5 -6

@@ -6,8 +6,7 @@ import copy

6 import itertools

7 import pickle as pkl

8 import time

9 -from collections.abc import Iterable

10 -from typing import Callable, Optional

9 +from collections.abc import Callable, Iterable

10

11 import torch

12 import torch.utils.benchmark as TBenchmark

@@ -53,7 +52,7 @@ def bench_int8(

52 n: int,

53 label: str,

54 sub_label: str,

56 - bench_kernels: Optional[list[str]] = None,

55 + bench_kernels: list[str] | None = None,

56 ) -> Iterable[TMeasurement]:

57 """Benchmark INT8-based kernels."""

58 assert dtype == torch.int8

@@ -108,7 +107,7 @@ def bench_fp8(

107 n: int,

108 label: str,

109 sub_label: str,

111 - bench_kernels: Optional[list[str]] = None,

110 + bench_kernels: list[str] | None = None,

111 ) -> Iterable[TMeasurement]:

112 """Benchmark FP8-based kernels."""

113 assert dtype == torch.float8_e4m3fn

@@ -183,7 +182,7 @@ def bench(

182 n: int,

183 label: str,

184 sub_label: str,

186 - bench_kernels: Optional[list[str]] = None,

185 + bench_kernels: list[str] | None = None,

186 ) -> Iterable[TMeasurement]:

187 if dtype == torch.int8:

188 return bench_int8(dtype, m, k, n, label, sub_label, bench_kernels)

@@ -201,7 +200,7 @@ def print_timers(timers: Iterable[TMeasurement]):

200 def run(

201 dtype: torch.dtype,

202 MKNs: Iterable[tuple[int, int, int]],

204 - bench_kernels: Optional[list[str]] = None,

203 + bench_kernels: list[str] | None = None,

204 ) -> Iterable[TMeasurement]:

205 results = []

206 for m, k, n in MKNs:

benchmarks/fused_kernels/layernorm_rms_benchmarks.py

+4 -5

@@ -3,10 +3,9 @@

3

4 import pickle as pkl

5 import time

6 -from collections.abc import Iterable

6 +from collections.abc import Callable, Iterable

7 from dataclasses import dataclass

8 from itertools import product

9 -from typing import Callable, Optional

9

10 import torch

11 import torch.utils.benchmark as TBenchmark

@@ -51,7 +50,7 @@ def get_bench_params() -> list[bench_params_t]:

50 def unfused_int8_impl(

51 rms_norm_layer: RMSNorm,

52 x: torch.Tensor,

54 - residual: Optional[torch.Tensor],

53 + residual: torch.Tensor | None,

54 quant_dtype: torch.dtype,

55 ):

56 # Norm

@@ -68,7 +67,7 @@ def unfused_int8_impl(

67 def unfused_fp8_impl(

68 rms_norm_layer: RMSNorm,

69 x: torch.Tensor,

71 - residual: Optional[torch.Tensor],

70 + residual: torch.Tensor | None,

71 quant_dtype: torch.dtype,

72 ):

73 # Norm

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

84 def fused_impl(

85 rms_norm_layer: RMSNorm, # this stores the weights

86 x: torch.Tensor,

88 - residual: Optional[torch.Tensor],

87 + residual: torch.Tensor | None,

88 quant_dtype: torch.dtype,

89 ):

90 out, _ = ops.rms_norm_dynamic_per_token_quant(

benchmarks/kernels/bench_per_token_quant_fp8.py

+1 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import itertools

4 -from typing import Callable

4 +from collections.abc import Callable

5 from unittest.mock import patch

6

7 import pandas as pd

benchmarks/kernels/benchmark_device_communicators.py

+3 -3

@@ -22,8 +22,8 @@ Example:

22 import json

23 import os

24 import time

25 +from collections.abc import Callable

26 from contextlib import nullcontext

26 -from typing import Callable, Optional

27

28 import torch

29 import torch.distributed as dist

@@ -264,12 +264,12 @@ class CommunicatorBenchmark:

264 def benchmark_allreduce_single(

265 self,

266 sequence_length: int,

267 - allreduce_fn: Callable[[torch.Tensor], Optional[torch.Tensor]],

267 + allreduce_fn: Callable[[torch.Tensor], torch.Tensor | None],

268 should_use_fn: Callable[[torch.Tensor], bool],

269 context,

270 num_warmup: int,

271 num_trials: int,

272 - ) -> Optional[float]:

272 + ) -> float | None:

273 """Benchmark method with CUDA graph optimization."""

274 try:

275 # Create test tensor (2D: sequence_length x hidden_size)

benchmarks/kernels/benchmark_lora.py

+11 -10

@@ -6,11 +6,12 @@ import copy

6 import json

7 import pickle

8 import time

9 +from collections.abc import Callable

10 from dataclasses import dataclass

11 from enum import Enum, auto

12 from itertools import product

13 from pathlib import Path

13 -from typing import Any, Callable, Optional

14 +from typing import Any

15

16 import torch

17 import torch.utils.benchmark as TBenchmark

@@ -158,7 +159,7 @@ def ref_group_gemm(

159 seq_lens_cpu: torch.Tensor,

160 prompt_lora_mapping_cpu: torch.Tensor,

161 scaling: float,

161 - add_inputs: Optional[bool],

162 + add_inputs: bool | None,

163 ):

164 """

165 Torch group gemm reference implementation to test correctness of

@@ -316,8 +317,8 @@ class BenchmarkContext:

317 lora_rank: int

318 sort_by_lora_id: bool

319 dtype: torch.dtype

319 - seq_length: Optional[int] = None

320 - num_slices: Optional[int] = None # num_slices for slice based ops

320 + seq_length: int | None = None

321 + num_slices: int | None = None # num_slices for slice based ops

322

323 def with_seq_length(self, seq_length: int) -> "BenchmarkContext":

324 ctx = copy.copy(self)

@@ -561,7 +562,7 @@ class BenchmarkTensors:

562 }

563

564 def bench_fn_kwargs(

564 - self, op_type: OpType, add_inputs: Optional[bool] = None

565 + self, op_type: OpType, add_inputs: bool | None = None

566 ) -> dict[str, Any]:

567 if op_type.is_shrink_fn():

568 assert add_inputs is None

@@ -575,7 +576,7 @@ class BenchmarkTensors:

576 raise ValueError(f"Unrecognized optype {self}")

577

578 def test_correctness(

578 - self, op_type: OpType, expand_fn_add_inputs: Optional[bool]

579 + self, op_type: OpType, expand_fn_add_inputs: bool | None

580 ) -> bool:

581 """

582 Test correctness of op_type implementation against a grouped gemm

@@ -611,8 +612,8 @@ def bench_optype(

612 ctx: BenchmarkContext,

613 arg_pool_size: int,

614 op_type: OpType,

614 - cuda_graph_nops: Optional[int] = None,

615 - expand_fn_add_inputs: Optional[bool] = None,

615 + cuda_graph_nops: int | None = None,

616 + expand_fn_add_inputs: bool | None = None,

617 test_correctness: bool = False,

618 ) -> TMeasurement:

619 assert arg_pool_size >= 1

@@ -679,7 +680,7 @@ def bench_torch_mm(

680 ctx: BenchmarkContext,

681 arg_pool_size: int,

682 op_type: OpType,

682 - cuda_graph_nops: Optional[int] = None,

683 + cuda_graph_nops: int | None = None,

684 ) -> TMeasurement:

685 """

686 Benchmark basic torch.mm as a roofline.

@@ -744,7 +745,7 @@ def use_cuda_graph_recommendation() -> str:

745 """

746

747

747 -def print_timers(timers: list[TMeasurement], args: Optional[argparse.Namespace] = None):

748 +def print_timers(timers: list[TMeasurement], args: argparse.Namespace | None = None):

749 compare = TBenchmark.Compare(timers)

750 compare.print()

751

benchmarks/kernels/benchmark_machete.py

+15 -16

@@ -8,10 +8,9 @@ import math

8 import os

9 import pickle as pkl

10 import time

11 -from collections.abc import Iterable

11 +from collections.abc import Callable, Iterable

12 from dataclasses import dataclass

13 from itertools import product

14 -from typing import Callable, Optional

14

15 import pandas as pd

16 import torch

@@ -63,23 +62,23 @@ class BenchmarkTensors:

62 a: torch.Tensor

63

64 w_q: torch.Tensor

66 - group_size: Optional[int]

65 + group_size: int | None

66 wtype: ScalarType

67 w_g_s: torch.Tensor

69 - w_g_zp: Optional[torch.Tensor]

70 - w_ch_s: Optional[torch.Tensor]

71 - w_tok_s: Optional[torch.Tensor]

68 + w_g_zp: torch.Tensor | None

69 + w_ch_s: torch.Tensor | None

70 + w_tok_s: torch.Tensor | None

71

72

73 @dataclass

74 class TypeConfig:

75 act_type: torch.dtype

76 weight_type: ScalarType

78 - output_type: Optional[torch.dtype]

79 - group_scale_type: Optional[torch.dtype]

80 - group_zero_type: Optional[torch.dtype]

81 - channel_scale_type: Optional[torch.dtype]

82 - token_scale_type: Optional[torch.dtype]

77 + output_type: torch.dtype | None

78 + group_scale_type: torch.dtype | None

79 + group_zero_type: torch.dtype | None

80 + channel_scale_type: torch.dtype | None

81 + token_scale_type: torch.dtype | None

82

83

84 def rand_data(shape, dtype=torch.float16, scale=1):

@@ -93,8 +92,8 @@ def quantize_and_pack(

92 atype: torch.dtype,

93 w: torch.Tensor,

94 wtype: ScalarType,

96 - stype: Optional[torch.dtype],

97 - group_size: Optional[int],

95 + stype: torch.dtype | None,

96 + group_size: int | None,

97 zero_points: bool = False,

98 ):

99 assert wtype.is_integer(), "TODO: support floating point weights"

@@ -113,7 +112,7 @@ def quantize_and_pack(

112

113

114 def create_bench_tensors(

116 - shape: tuple[int, int, int], types: TypeConfig, group_size: Optional[int]

115 + shape: tuple[int, int, int], types: TypeConfig, group_size: int | None

116 ) -> list[BenchmarkTensors]:

117 m, n, k = shape

118

@@ -331,8 +330,8 @@ def bench_fns(label: str, sub_label: str, description: str, fns: list[Callable])

330 return res

331

332

334 -_SWEEP_SCHEDULES_RESULTS: Optional[pd.DataFrame] = None

335 -_SWEEP_SCHEDULES_RESULTS_CSV: Optional[str] = None

333 +_SWEEP_SCHEDULES_RESULTS: pd.DataFrame | None = None

334 +_SWEEP_SCHEDULES_RESULTS_CSV: str | None = None

335

336

337 def bench(

benchmarks/kernels/benchmark_paged_attention.py

+1 -2

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

3

4 import random

5 import time

6 -from typing import Optional

6

7 import torch

8

@@ -37,7 +36,7 @@ def main(

36 seed: int,

37 do_profile: bool,

38 device: str = "cuda",

40 - kv_cache_dtype: Optional[str] = None,

39 + kv_cache_dtype: str | None = None,

40 ) -> None:

41 current_platform.seed_everything(seed)

42

benchmarks/kernels/benchmark_per_token_group_quant.py

+1 -1

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

3

4 import argparse

5 import math

6 +from collections.abc import Callable

7 from contextlib import contextmanager

7 -from typing import Callable

8 from unittest.mock import patch

9

10 import torch

benchmarks/kernels/benchmark_reshape_and_cache.py

-2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import random

4 import time

5

benchmarks/kernels/benchmark_reshape_and_cache_flash.py

-2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import random

4 import time

5

benchmarks/kernels/benchmark_rmsnorm.py

+5 -6

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

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

3

4 import itertools

5 -from typing import Optional, Union

5

6 import torch

7 from flashinfer.norm import fused_add_rmsnorm, rmsnorm

@@ -21,8 +20,8 @@ class HuggingFaceRMSNorm(nn.Module):

20 def forward(

21 self,

22 x: torch.Tensor,

24 - residual: Optional[torch.Tensor] = None,

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

23 + residual: torch.Tensor | None = None,

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

25 orig_dtype = x.dtype

26 x = x.to(torch.float32)

27 if residual is not None:

@@ -41,7 +40,7 @@ class HuggingFaceRMSNorm(nn.Module):

40 def rmsnorm_naive(

41 x: torch.Tensor,

42 weight: torch.Tensor,

44 - residual: Optional[torch.Tensor] = None,

43 + residual: torch.Tensor | None = None,

44 eps: float = 1e-6,

45 ):

46 naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps)

@@ -65,7 +64,7 @@ def rmsnorm_naive(

64 def rmsnorm_flashinfer(

65 x: torch.Tensor,

66 weight: torch.Tensor,

68 - residual: Optional[torch.Tensor] = None,

67 + residual: torch.Tensor | None = None,

68 eps: float = 1e-6,

69 ):

70 orig_shape = x.shape

@@ -89,7 +88,7 @@ def rmsnorm_flashinfer(

88 def rmsnorm_vllm(

89 x: torch.Tensor,

90 weight: torch.Tensor,

92 - residual: Optional[torch.Tensor] = None,

91 + residual: torch.Tensor | None = None,

92 eps: float = 1e-6,

93 ):

94 orig_shape = x.shape

benchmarks/kernels/benchmark_rope.py

+1 -2

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

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

3

4 from itertools import accumulate

5 -from typing import Optional

5

6 import nvtx

7 import torch

@@ -18,7 +17,7 @@ def benchmark_rope_kernels_multi_lora(

17 seq_len: int,

18 num_heads: int,

19 head_size: int,

21 - rotary_dim: Optional[int],

20 + rotary_dim: int | None,

21 dtype: torch.dtype,

22 seed: int,

23 device: str,

benchmarks/kernels/benchmark_trtllm_decode_attention.py

+1 -4

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

4 import csv

5 import os

6 from datetime import datetime

7 -from typing import Optional

7

8 import flashinfer

9 import torch

@@ -28,9 +27,7 @@ def to_float8(x, dtype=torch.float8_e4m3fn):

27 @torch.no_grad()

28 def benchmark_decode(

29 dtype: torch.dtype,

31 - quant_dtypes: tuple[

32 - Optional[torch.dtype], Optional[torch.dtype], Optional[torch.dtype]

33 - ],

30 + quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None],

31 batch_size: int,

32 max_seq_len: int,

33 num_heads: tuple[int, int] = (64, 8),

benchmarks/kernels/benchmark_trtllm_prefill_attention.py

+1 -4

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

4 import csv

5 import os

6 from datetime import datetime

7 -from typing import Optional

7

8 import flashinfer

9 import torch

@@ -28,9 +27,7 @@ def to_float8(x, dtype=torch.float8_e4m3fn):

27 @torch.no_grad()

28 def benchmark_prefill(

29 dtype: torch.dtype,

31 - quant_dtypes: tuple[

32 - Optional[torch.dtype], Optional[torch.dtype], Optional[torch.dtype]

33 - ],

30 + quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None],

31 batch_size: int,

32 max_seq_len: int,

33 num_heads: tuple[int, int] = (64, 8),

benchmarks/kernels/utils.py

+3 -3

@@ -2,8 +2,8 @@

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

3

4 import dataclasses

5 -from collections.abc import Iterable

6 -from typing import Any, Callable, Optional

5 +from collections.abc import Callable, Iterable

6 +from typing import Any

7

8 import torch

9 import torch.utils.benchmark as TBenchmark

@@ -55,7 +55,7 @@ class Bench:

55

56 def __init__(

57 self,

58 - cuda_graph_params: Optional[CudaGraphBenchParams],

58 + cuda_graph_params: CudaGraphBenchParams | None,

59 label: str,

60 sub_label: str,

61 description: str,

benchmarks/multi_turn/bench_dataset.py

+11 -11

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

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

3 from abc import ABC, abstractmethod

4 from statistics import mean

5 -from typing import Any, NamedTuple, Optional, Union

5 +from typing import Any, NamedTuple

6

7 import numpy as np # type: ignore

8 import pandas as pd # type: ignore

@@ -35,8 +35,8 @@ class Distribution(ABC):

35 class UniformDistribution(Distribution):

36 def __init__(

37 self,

38 - min_val: Union[int, float],

39 - max_val: Union[int, float],

38 + min_val: int | float,

39 + max_val: int | float,

40 is_integer: bool = True,

41 ) -> None:

42 self.min_val = min_val

@@ -56,7 +56,7 @@ class UniformDistribution(Distribution):

56

57

58 class ConstantDistribution(Distribution):

59 - def __init__(self, value: Union[int, float]) -> None:

59 + def __init__(self, value: int | float) -> None:

60 self.value = value

61 self.max_val = value

62

@@ -68,7 +68,7 @@ class ConstantDistribution(Distribution):

68

69

70 class ZipfDistribution(Distribution):

71 - def __init__(self, alpha: float, max_val: Optional[int] = None) -> None:

71 + def __init__(self, alpha: float, max_val: int | None = None) -> None:

72 self.alpha = alpha

73 self.max_val = max_val

74

@@ -83,7 +83,7 @@ class ZipfDistribution(Distribution):

83

84

85 class PoissonDistribution(Distribution):

86 - def __init__(self, alpha: float, max_val: Optional[int] = None) -> None:

86 + def __init__(self, alpha: float, max_val: int | None = None) -> None:

87 self.alpha = alpha

88 self.max_val = max_val

89

@@ -100,11 +100,11 @@ class PoissonDistribution(Distribution):

100 class LognormalDistribution(Distribution):

101 def __init__(

102 self,

103 - mean: Optional[float] = None,

104 - sigma: Optional[float] = None,

105 - average: Optional[int] = None,

106 - median_ratio: Optional[float] = None,

107 - max_val: Optional[int] = None,

103 + mean: float | None = None,

104 + sigma: float | None = None,

105 + average: int | None = None,

106 + median_ratio: float | None = None,

107 + max_val: int | None = None,

108 ) -> None:

109 self.average = average

110 self.median_ratio = median_ratio

benchmarks/multi_turn/benchmark_serving_multi_turn.py

+5 -5

@@ -13,7 +13,7 @@ from datetime import datetime

13 from enum import Enum

14 from http import HTTPStatus

15 from statistics import mean

16 -from typing import NamedTuple, Union

16 +from typing import NamedTuple

17

18 import aiohttp # type: ignore

19 import numpy as np # type: ignore

@@ -169,7 +169,7 @@ class MovingAverage:

169 class DebugStats:

170 def __init__(self, logger: logging.Logger, window_size: int) -> None:

171 self.logger = logger

172 - self.metrics: dict[str, Union[MovingAverage, MetricStats]] = {

172 + self.metrics: dict[str, MovingAverage | MetricStats] = {

173 "moving_avg_ttft_ms": MovingAverage(window_size),

174 "moving_avg_tpot_ms": MovingAverage(window_size),

175 "ttft_ms": MetricStats(),

@@ -636,7 +636,7 @@ async def client_main(

636

637 if args.verbose:

638 curr_time_sec: float = time.perf_counter()

639 - time_since_last_turn: Union[str, float] = "N/A"

639 + time_since_last_turn: str | float = "N/A"

640 if conv_id in time_of_last_turn:

641 time_since_last_turn = round(

642 curr_time_sec - time_of_last_turn[conv_id], 3

@@ -928,13 +928,13 @@ async def main_mp(

928 f"{num_clients_finished} out of {bench_args.num_clients} clients finished, collected {len(client_metrics)} measurements, runtime {runtime_sec:.3f} sec{Color.RESET}" # noqa: E501

929 )

930

931 - rps: Union[str, float] = round(len(client_metrics) / runtime_sec, 3)

931 + rps: str | float = round(len(client_metrics) / runtime_sec, 3)

932 if len(client_metrics) < (5 * bench_args.num_clients):

933 # Do not estimate the RPS if the number of samples is very low

934 # (threshold can be tuned if needed)

935 rps = "N/A"

936

937 - runtime_left_sec: Union[str, float] = round(

937 + runtime_left_sec: str | float = round(

938 (runtime_sec / finished_convs) * (total_convs - finished_convs), 3

939 )

940 if percent < 0.05:

benchmarks/multi_turn/convert_sharegpt_to_openai.py

+9 -9

@@ -13,7 +13,7 @@ import argparse

13 import json

14 import random

15 from statistics import mean

16 -from typing import Any, Optional

16 +from typing import Any

17

18 import pandas as pd # type: ignore

19 import tqdm # type: ignore

@@ -25,7 +25,7 @@ def has_non_english_chars(text: str) -> bool:

25

26

27 def content_is_valid(

28 - content: str, min_content_len: Optional[int], max_content_len: Optional[int]

28 + content: str, min_content_len: int | None, max_content_len: int | None

29 ) -> bool:

30 if min_content_len and len(content) < min_content_len:

31 return False

@@ -37,7 +37,7 @@ def content_is_valid(

37

38

39 def print_stats(

40 - conversations: "list[dict[Any, Any]]", tokenizer: Optional[AutoTokenizer] = None

40 + conversations: "list[dict[Any, Any]]", tokenizer: AutoTokenizer | None = None

41 ) -> None:

42 # Collect statistics

43 stats = []

@@ -109,12 +109,12 @@ def convert_sharegpt_to_openai(

109 seed: int,

110 input_file: str,

111 output_file: str,

112 - max_items: Optional[int],

113 - min_content_len: Optional[int] = None,

114 - max_content_len: Optional[int] = None,

115 - min_turns: Optional[int] = None,

116 - max_turns: Optional[int] = None,

117 - model: Optional[str] = None,

112 + max_items: int | None,

113 + min_content_len: int | None = None,

114 + max_content_len: int | None = None,

115 + min_turns: int | None = None,

116 + max_turns: int | None = None,

117 + model: str | None = None,

118 ) -> None:

119 if min_turns and max_turns:

120 assert min_turns <= max_turns

csrc/cutlass_extensions/vllm_cutlass_library_extension.py

+6 -9

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

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

3

4 import enum

5 -from typing import Union

5

6 from cutlass_library import *

7

@@ -22,7 +21,7 @@ class MixedInputKernelScheduleType(enum.Enum):

21 TmaWarpSpecializedCooperative = enum_auto()

22

23

25 -VLLMDataTypeNames: dict[Union[VLLMDataType, DataType], str] = {

24 +VLLMDataTypeNames: dict[VLLMDataType | DataType, str] = {

25 **DataTypeNames, # type: ignore

26 **{

27 VLLMDataType.u4b8: "u4b8",

@@ -30,7 +29,7 @@ VLLMDataTypeNames: dict[Union[VLLMDataType, DataType], str] = {

29 },

30 }

31

33 -VLLMDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {

32 +VLLMDataTypeTag: dict[VLLMDataType | DataType, str] = {

33 **DataTypeTag, # type: ignore

34 **{

35 VLLMDataType.u4b8: "cutlass::vllm_uint4b8_t",

@@ -38,7 +37,7 @@ VLLMDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {

37 },

38 }

39

41 -VLLMDataTypeSize: dict[Union[VLLMDataType, DataType], int] = {

40 +VLLMDataTypeSize: dict[VLLMDataType | DataType, int] = {

41 **DataTypeSize, # type: ignore

42 **{

43 VLLMDataType.u4b8: 4,

@@ -46,7 +45,7 @@ VLLMDataTypeSize: dict[Union[VLLMDataType, DataType], int] = {

45 },

46 }

47

49 -VLLMDataTypeVLLMScalarTypeTag: dict[Union[VLLMDataType, DataType], str] = {

48 +VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = {

49 VLLMDataType.u4b8: "vllm::kU4B8",

50 VLLMDataType.u8b128: "vllm::kU8B128",

51 DataType.u4: "vllm::kU4",

@@ -57,7 +56,7 @@ VLLMDataTypeVLLMScalarTypeTag: dict[Union[VLLMDataType, DataType], str] = {

56 DataType.bf16: "vllm::kBfloat16",

57 }

58

60 -VLLMDataTypeTorchDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {

59 +VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = {

60 DataType.u8: "at::ScalarType::Byte",

61 DataType.s8: "at::ScalarType::Char",

62 DataType.e4m3: "at::ScalarType::Float8_e4m3fn",

@@ -67,9 +66,7 @@ VLLMDataTypeTorchDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {

66 DataType.f32: "at::ScalarType::Float",

67 }

68

70 -VLLMKernelScheduleTag: dict[

71 - Union[MixedInputKernelScheduleType, KernelScheduleType], str

72 -] = {

69 +VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = {

70 **KernelScheduleTag, # type: ignore

71 **{

72 MixedInputKernelScheduleType.TmaWarpSpecialized: "cutlass::gemm::KernelTmaWarpSpecialized", # noqa: E501

csrc/quantization/machete/generate.py

+2 -3

@@ -9,7 +9,6 @@ from collections.abc import Iterable

9 from copy import deepcopy

10 from dataclasses import dataclass, fields

11 from functools import reduce

12 -from typing import Optional, Union

12

13 import jinja2

14 from vllm_cutlass_library_extension import (

@@ -259,7 +258,7 @@ class ScheduleConfig:

258 @dataclass(frozen=True)

259 class TypeConfig:

260 a: DataType

262 - b: Union[DataType, VLLMDataType]

261 + b: DataType | VLLMDataType

262 b_group_scale: DataType

263 b_group_zeropoint: DataType

264 b_channel_scale: DataType

@@ -280,7 +279,7 @@ class PrepackTypeConfig:

279 class ImplConfig:

280 types: TypeConfig

281 schedules: list[ScheduleConfig]

283 - heuristic: list[tuple[Optional[str], ScheduleConfig]]

282 + heuristic: list[tuple[str | None, ScheduleConfig]]

283

284

285 def generate_sch_sig(schedule_config: ScheduleConfig) -> str:

docs/contributing/model/transcription.md

+7 -7

@@ -16,7 +16,7 @@ Declare supported languages and capabilities:

16

17 ??? code "supported_languages and supports_transcription_only"

18 ```python

19 - from typing import ClassVar, Mapping, Optional, Literal

19 + from typing import ClassVar, Mapping, Literal

20 import numpy as np

21 import torch

22 from torch import nn

@@ -81,10 +81,10 @@ Return a dict containing `multi_modal_data` with the audio, and either a `prompt

81 audio: np.ndarray,

82 stt_config: SpeechToTextConfig,

83 model_config: ModelConfig,

84 - language: Optional[str],

84 + language: str | None,

85 task_type: Literal["transcribe", "translate"],

86 request_prompt: str,

87 - to_language: Optional[str],

87 + to_language: str | None,

88 ) -> PromptType:

89 # Example with a free-form instruction prompt

90 task_word = "Transcribe" if task_type == "transcribe" else "Translate"

@@ -117,10 +117,10 @@ Return a dict with separate `encoder_prompt` and `decoder_prompt` entries:

117 audio: np.ndarray,

118 stt_config: SpeechToTextConfig,

119 model_config: ModelConfig,

120 - language: Optional[str],

120 + language: str | None,

121 task_type: Literal["transcribe", "translate"],

122 request_prompt: str,

123 - to_language: Optional[str],

123 + to_language: str | None,

124 ) -> PromptType:

125 if language is None:

126 raise ValueError("Language must be specified")

@@ -150,7 +150,7 @@ If your model requires a language and you want a default, override this method (

150 ??? code "validate_language()"

151 ```python

152 @classmethod

153 - def validate_language(cls, language: Optional[str]) -> Optional[str]:

153 + def validate_language(cls, language: str | None) -> str | None:

154 if language is None:

155 logger.warning(

156 "Defaulting to language='en'. If you wish to transcribe audio in a different language, pass the `language` field.")

@@ -175,7 +175,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:

175 audio_duration_s: float,

176 stt_config: SpeechToTextConfig,

177 model_config: ModelConfig,

178 - ) -> Optional[int]:

178 + ) -> int | None:

179 # Return None if unknown; otherwise return an estimate.

180 return int(audio_duration_s * stt_config.sample_rate // 320) # example

181 ```

docs/design/logits_processors.md

+3 -3

@@ -174,7 +174,7 @@ The previous sections alluded to the interfaces which vLLM logits processors mus

174 from collections.abc import Sequence

175 from dataclasses import dataclass

176 from enum import Enum, auto

177 - from typing import TYPE_CHECKING, Optional

177 + from typing import TYPE_CHECKING

178

179 import torch

180

@@ -244,7 +244,7 @@ The previous sections alluded to the interfaces which vLLM logits processors mus

244 @abstractmethod

245 def update_state(

246 self,

247 - batch_update: Optional["BatchUpdate"],

247 + batch_update: "BatchUpdate" | None,

248 ) -> None:

249 """Called when there are new output tokens, prior

250 to each forward pass.

@@ -274,7 +274,7 @@ A vLLM logits processor must subclass `LogitsProcessor` and define (at minimum)

274 * Return `True` if the logits processor is argmax invariant (never changes what is the highest-logit-value token ID for a given request), `False` if the logits processor may modify argmax

275 * `is_argmax_invariant()` is evaluated once at startup; if `True`, vLLM will skip applying this logits processor in a given step when all requests use greedy sampling

276

277 -* `update_state(self, batch_update: Optional["BatchUpdate"]) -> None`:

277 +* `update_state(self, batch_update: "BatchUpdate" | None) -> None`:

278 * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step

279 * Use the `BatchUpdate` members to update logits processor internal state

280 * **Note:** batch update data structure may be `None`, signaling no change to the batch constituents. In this case, the LogitsProcessor might still want to update its state based on the updated `output_token_ids` lists that it could have retained when they were added.

docs/features/custom_logitsprocs.md

+1 -2

@@ -93,7 +93,6 @@ The contrived example below implements a custom logits processor which consumes

93 ??? code "Example custom logits processor definition"

94

95 ``` python

96 - from typing import Optional

96 import torch

97 from vllm.config import VllmConfig

98 from vllm.sampling_params import SamplingParams

@@ -112,7 +111,7 @@ The contrived example below implements a custom logits processor which consumes

111 """Never impacts greedy sampling"""

112 return False

113

115 - def update_state(self, batch_update: Optional[BatchUpdate]):

114 + def update_state(self, batch_update: BatchUpdate | None):

115 if not batch_update:

116 return

117

examples/offline_inference/audio_language.py

+6 -6

@@ -10,7 +10,7 @@ on HuggingFace model repository.

10

11 import os

12 from dataclasses import asdict

13 -from typing import Any, NamedTuple, Optional

13 +from typing import Any, NamedTuple

14

15 from huggingface_hub import snapshot_download

16 from transformers import AutoTokenizer

@@ -30,11 +30,11 @@ question_per_audio_count = {

30

31 class ModelRequestData(NamedTuple):

32 engine_args: EngineArgs

33 - prompt: Optional[str] = None

34 - prompt_token_ids: Optional[dict[str, list[int]]] = None

35 - multi_modal_data: Optional[dict[str, Any]] = None

36 - stop_token_ids: Optional[list[int]] = None

37 - lora_requests: Optional[list[LoRARequest]] = None

33 + prompt: str | None = None

34 + prompt_token_ids: dict[str, list[int]] | None = None

35 + multi_modal_data: dict[str, Any] | None = None

36 + stop_token_ids: list[int] | None = None

37 + lora_requests: list[LoRARequest] | None = None

38

39

40 # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on

examples/offline_inference/kv_load_failure_recovery/rogue_shared_storage_connector.py

+2 -2

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

3 # ruff: noqa: E501

4 import logging

5 from dataclasses import dataclass, field

6 -from typing import TYPE_CHECKING, Optional

6 +from typing import TYPE_CHECKING

7

8 from vllm.config import VllmConfig

9 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

@@ -81,7 +81,7 @@ class RogueSharedStorageConnector(SharedStorageConnector):

81

82 def get_finished(

83 self, finished_req_ids: set[str]

84 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:

84 + ) -> tuple[set[str] | None, set[str] | None]:

85 if self._async_load:

86 meta = self._get_connector_metadata()

87 assert isinstance(meta, RogueSharedStorageConnectorMetadata)

examples/offline_inference/logits_processor/custom.py

+1 -3

@@ -33,8 +33,6 @@ Output: ' in the hands of the people.\n\nThe future of AI is in the'

33 ------------------------------------------------------------

34 """

35

36 -from typing import Optional

37 -

36 import torch

37

38 from vllm import LLM, SamplingParams

@@ -58,7 +56,7 @@ class DummyLogitsProcessor(LogitsProcessor):

56 def is_argmax_invariant(self) -> bool:

57 return False

58

61 - def update_state(self, batch_update: Optional[BatchUpdate]):

59 + def update_state(self, batch_update: BatchUpdate | None):

60 process_dict_updates(

61 self.req_info,

62 batch_update,

examples/offline_inference/logits_processor/custom_req.py

+3 -3

@@ -39,7 +39,7 @@ Output: ' in the hands of the people.\n\nThe future of AI is in the'

39 ------------------------------------------------------------

40 """

41

42 -from typing import Any, Optional

42 +from typing import Any

43

44 import torch

45

@@ -82,7 +82,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):

82 def new_req_logits_processor(

83 self,

84 params: SamplingParams,

85 - ) -> Optional[RequestLogitsProcessor]:

85 + ) -> RequestLogitsProcessor | None:

86 """This method returns a new request-level logits processor, customized

87 to the `target_token` value associated with a particular request.

88

@@ -96,7 +96,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):

96 Returns:

97 `Callable` request logits processor, or None

98 """

99 - target_token: Optional[Any] = params.extra_args and params.extra_args.get(

99 + target_token: Any | None = params.extra_args and params.extra_args.get(

100 "target_token"

101 )

102 if target_token is None:

examples/offline_inference/logits_processor/custom_req_init.py

+1 -3

@@ -41,8 +41,6 @@ which indicates that the logits processor is running. However, on a non-"cuda"

41 device, the first and third requests would not repeat the same token.

42 """

43

44 -from typing import Optional

45 -

44 import torch

45

46 from vllm import LLM, SamplingParams

@@ -91,7 +89,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):

89 def new_req_logits_processor(

90 self,

91 params: SamplingParams,

94 - ) -> Optional[RequestLogitsProcessor]:

92 + ) -> RequestLogitsProcessor | None:

93 """This method returns a new request-level logits processor, customized

94 to the `target_token` value associated with a particular request.

95

examples/offline_inference/lora_with_quantization_inference.py

+3 -4

@@ -8,7 +8,6 @@ Requires HuggingFace credentials for access.

8 """

9

10 import gc

11 -from typing import Optional

11

12 import torch

13 from huggingface_hub import snapshot_download

@@ -19,7 +18,7 @@ from vllm.lora.request import LoRARequest

18

19 def create_test_prompts(

20 lora_path: str,

22 -) -> list[tuple[str, SamplingParams, Optional[LoRARequest]]]:

21 +) -> list[tuple[str, SamplingParams, LoRARequest | None]]:

22 return [

23 # this is an example of using quantization without LoRA

24 (

@@ -56,7 +55,7 @@ def create_test_prompts(

55

56 def process_requests(

57 engine: LLMEngine,

59 - test_prompts: list[tuple[str, SamplingParams, Optional[LoRARequest]]],

58 + test_prompts: list[tuple[str, SamplingParams, LoRARequest | None]],

59 ):

60 """Continuously process a list of prompts and handle the outputs."""

61 request_id = 0

@@ -78,7 +77,7 @@ def process_requests(

77

78

79 def initialize_engine(

81 - model: str, quantization: str, lora_repo: Optional[str]

80 + model: str, quantization: str, lora_repo: str | None

81 ) -> LLMEngine:

82 """Initialize the LLMEngine."""

83

examples/offline_inference/multilora_inference.py

+2 -4

@@ -7,8 +7,6 @@ for offline inference.

7 Requires HuggingFace credentials for access to Llama2.

8 """

9

10 -from typing import Optional

11 -

10 from huggingface_hub import snapshot_download

11

12 from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams

@@ -17,7 +15,7 @@ from vllm.lora.request import LoRARequest

15

16 def create_test_prompts(

17 lora_path: str,

20 -) -> list[tuple[str, SamplingParams, Optional[LoRARequest]]]:

18 +) -> list[tuple[str, SamplingParams, LoRARequest | None]]:

19 """Create a list of test prompts with their sampling parameters.

20

21 2 requests for base model, 4 requests for the LoRA. We define 2

@@ -68,7 +66,7 @@ def create_test_prompts(

66

67 def process_requests(

68 engine: LLMEngine,

71 - test_prompts: list[tuple[str, SamplingParams, Optional[LoRARequest]]],

69 + test_prompts: list[tuple[str, SamplingParams, LoRARequest | None]],

70 ):

71 """Continuously process a list of prompts and handle the outputs."""

72 request_id = 0

examples/offline_inference/prithvi_geospatial_mae.py

+1 -2

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

3 import argparse

4 import datetime

5 import os

6 -from typing import Union

6

7 import albumentations

8 import numpy as np

@@ -160,7 +159,7 @@ def load_example(

159 file_paths: list[str],

160 mean: list[float] = None,

161 std: list[float] = None,

163 - indices: Union[list[int], None] = None,

162 + indices: list[int] | None = None,

163 ):

164 """Build an input example by loading images in *file_paths*.

165

examples/offline_inference/rlhf_utils.py

+4 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import gc

4 -from typing import Callable, Optional, TypedDict

4 +from collections.abc import Callable

5 +from typing import TypedDict

6

7 import torch

8 import zmq

@@ -71,7 +72,7 @@ class WorkerExtension:

72

73

74 def rebuild_ipc(

74 - handle: tuple[Callable, tuple], device_id: Optional[int] = None

75 + handle: tuple[Callable, tuple], device_id: int | None = None

76 ) -> torch.Tensor:

77 func, args = handle

78 list_args = list(args)

@@ -109,7 +110,7 @@ class ColocateWorkerExtension:

110 self._zmq_ctx = zmq.Context()

111 socket = self._zmq_ctx.socket(zmq.REP)

112 socket.connect(zmq_handles[self.report_device_id()])

112 - buffer: Optional[torch.Tensor] = None

113 + buffer: torch.Tensor | None = None

114 while True:

115 payload: tuple[Callable, tuple] | list[FlattenedTensorMetadata] | None = (

116 socket.recv_pyobj()

examples/offline_inference/vision_language.py

+3 -3

@@ -12,7 +12,7 @@ import os

12 import random

13 from contextlib import contextmanager

14 from dataclasses import asdict

15 -from typing import NamedTuple, Optional

15 +from typing import NamedTuple

16

17 from huggingface_hub import snapshot_download

18 from transformers import AutoTokenizer

@@ -28,8 +28,8 @@ from vllm.utils import FlexibleArgumentParser

28 class ModelRequestData(NamedTuple):

29 engine_args: EngineArgs

30 prompts: list[str]

31 - stop_token_ids: Optional[list[int]] = None

32 - lora_requests: Optional[list[LoRARequest]] = None

31 + stop_token_ids: list[int] | None = None

32 + lora_requests: list[LoRARequest] | None = None

33

34

35 # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on

examples/offline_inference/vision_language_multi_image.py

+6 -6

@@ -9,7 +9,7 @@ using the chat template defined by the model.

9 import os

10 from argparse import Namespace

11 from dataclasses import asdict

12 -from typing import NamedTuple, Optional

12 +from typing import NamedTuple

13

14 from huggingface_hub import snapshot_download

15 from PIL.Image import Image

@@ -41,9 +41,9 @@ class ModelRequestData(NamedTuple):

41 engine_args: EngineArgs

42 prompt: str

43 image_data: list[Image]

44 - stop_token_ids: Optional[list[int]] = None

45 - chat_template: Optional[str] = None

46 - lora_requests: Optional[list[LoRARequest]] = None

44 + stop_token_ids: list[int] | None = None

45 + chat_template: str | None = None

46 + lora_requests: list[LoRARequest] | None = None

47

48

49 # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on

@@ -1251,7 +1251,7 @@ model_example_map = {

1251 }

1252

1253

1254 -def run_generate(model, question: str, image_urls: list[str], seed: Optional[int]):

1254 +def run_generate(model, question: str, image_urls: list[str], seed: int | None):

1255 req_data = model_example_map[model](question, image_urls)

1256

1257 engine_args = asdict(req_data.engine_args) | {"seed": args.seed}

@@ -1277,7 +1277,7 @@ def run_generate(model, question: str, image_urls: list[str], seed: Optional[int

1277 print("-" * 50)

1278

1279

1280 -def run_chat(model: str, question: str, image_urls: list[str], seed: Optional[int]):

1280 +def run_chat(model: str, question: str, image_urls: list[str], seed: int | None):

1281 req_data = model_example_map[model](question, image_urls)

1282

1283 # Disable other modalities to save memory

examples/offline_inference/vision_language_pooling.py

+8 -8

@@ -11,7 +11,7 @@ on HuggingFace model repository.

11 from argparse import Namespace

12 from dataclasses import asdict

13 from pathlib import Path

14 -from typing import Literal, NamedTuple, Optional, TypedDict, Union, get_args

14 +from typing import Literal, NamedTuple, TypeAlias, TypedDict, get_args

15

16 from PIL.Image import Image

17

@@ -47,15 +47,15 @@ class TextImagesQuery(TypedDict):

47

48

49 QueryModality = Literal["text", "image", "text+image", "text+images"]

50 -Query = Union[TextQuery, ImageQuery, TextImageQuery, TextImagesQuery]

50 +Query: TypeAlias = TextQuery | ImageQuery | TextImageQuery | TextImagesQuery

51

52

53 class ModelRequestData(NamedTuple):

54 engine_args: EngineArgs

55 - prompt: Optional[str] = None

56 - image: Optional[Image] = None

57 - query: Optional[str] = None

58 - documents: Optional[ScoreMultiModalParam] = None

55 + prompt: str | None = None

56 + image: Image | None = None

57 + query: str | None = None

58 + documents: ScoreMultiModalParam | None = None

59

60

61 def run_clip(query: Query) -> ModelRequestData:

@@ -281,7 +281,7 @@ def get_query(modality: QueryModality):

281 raise ValueError(msg)

282

283

284 -def run_encode(model: str, modality: QueryModality, seed: Optional[int]):

284 +def run_encode(model: str, modality: QueryModality, seed: int | None):

285 query = get_query(modality)

286 req_data = model_example_map[model](query)

287

@@ -311,7 +311,7 @@ def run_encode(model: str, modality: QueryModality, seed: Optional[int]):

311 print("-" * 50)

312

313

314 -def run_score(model: str, modality: QueryModality, seed: Optional[int]):

314 +def run_score(model: str, modality: QueryModality, seed: int | None):

315 query = get_query(modality)

316 req_data = model_example_map[model](query)

317

examples/online_serving/disaggregated_serving/disagg_proxy_demo.py

+7 -10

@@ -23,7 +23,7 @@ import logging

23 import os

24 import sys

25 from abc import ABC, abstractmethod

26 -from typing import Callable, Optional

26 +from collections.abc import Callable

27

28 import aiohttp

29 import requests

@@ -49,12 +49,9 @@ class Proxy:

49 decode_instances: list[str],

50 model: str,

51 scheduling_policy: SchedulingPolicy,

52 - custom_create_completion: Optional[

53 - Callable[[Request], StreamingResponse]

54 - ] = None,

55 - custom_create_chat_completion: Optional[

56 - Callable[[Request], StreamingResponse]

57 - ] = None,

52 + custom_create_completion: Callable[[Request], StreamingResponse] | None = None,

53 + custom_create_chat_completion: Callable[[Request], StreamingResponse]

54 + | None = None,

55 ):

56 self.prefill_instances = prefill_instances

57 self.decode_instances = decode_instances

@@ -348,9 +345,9 @@ class ProxyServer:

345 def __init__(

346 self,

347 args: argparse.Namespace,

351 - scheduling_policy: Optional[SchedulingPolicy] = None,

352 - create_completion: Optional[Callable[[Request], StreamingResponse]] = None,

353 - create_chat_completion: Optional[Callable[[Request], StreamingResponse]] = None,

348 + scheduling_policy: SchedulingPolicy | None = None,

349 + create_completion: Callable[[Request], StreamingResponse] | None = None,

350 + create_chat_completion: Callable[[Request], StreamingResponse] | None = None,

351 ):

352 self.validate_parsed_serve_args(args)

353 self.port = args.port

examples/online_serving/kv_events_subscriber.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Any, Optional, Union

3 +from typing import Any

4

5 import msgspec

6 import zmq

@@ -25,16 +25,16 @@ class KVCacheEvent(

25

26 class BlockStored(KVCacheEvent):

27 block_hashes: list[ExternalBlockHash]

28 - parent_block_hash: Optional[ExternalBlockHash]

28 + parent_block_hash: ExternalBlockHash | None

29 token_ids: list[int]

30 block_size: int

31 - lora_id: Optional[int]

32 - medium: Optional[str]

31 + lora_id: int | None

32 + medium: str | None

33

34

35 class BlockRemoved(KVCacheEvent):

36 block_hashes: list[ExternalBlockHash]

37 - medium: Optional[str]

37 + medium: str | None

38

39

40 class AllBlocksCleared(KVCacheEvent):

@@ -42,7 +42,7 @@ class AllBlocksCleared(KVCacheEvent):

42

43

44 class KVEventBatch(EventBatch):

45 - events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]

45 + events: list[BlockStored | BlockRemoved | AllBlocksCleared]

46

47

48 def process_event(event_batch):

examples/online_serving/multi_instance_data_parallel.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import asyncio

4 -from typing import Optional

4

5 from vllm.engine.arg_utils import AsyncEngineArgs

6 from vllm.engine.async_llm_engine import AsyncLLMEngine

@@ -43,7 +42,7 @@ async def main():

42 )

43

44 prompt = "Who won the 2004 World Series?"

46 - final_output: Optional[RequestOutput] = None

45 + final_output: RequestOutput | None = None

46 async for output in engine_client.generate(

47 prompt=prompt,

48 sampling_params=sampling_params,

examples/online_serving/pooling/cohere_rerank_client.py

+1 -3

@@ -8,8 +8,6 @@ Note that `pip install cohere` is needed to run this example.

8 run: vllm serve BAAI/bge-reranker-base

9 """

10

11 -from typing import Union

12 -

11 import cohere

12 from cohere import Client, ClientV2

13

@@ -25,7 +23,7 @@ documents = [

23

24

25 def cohere_rerank(

28 - client: Union[Client, ClientV2], model: str, query: str, documents: list[str]

26 + client: Client | ClientV2, model: str, query: str, documents: list[str]

27 ) -> dict:

28 return client.rerank(model=model, query=query, documents=documents)

29

examples/online_serving/pooling/openai_chat_embedding_client_for_multimodal.py

+2 -2

@@ -9,7 +9,7 @@ Refer to each `run_*` function for the command to run the server for that model.

9 import argparse

10 import base64

11 import io

12 -from typing import Literal, Union

12 +from typing import Literal

13

14 from openai import OpenAI

15 from openai._types import NOT_GIVEN, NotGiven

@@ -29,7 +29,7 @@ def create_chat_embeddings(

29 *,

30 messages: list[ChatCompletionMessageParam],

31 model: str,

32 - encoding_format: Union[Literal["base64", "float"], NotGiven] = NOT_GIVEN,

32 + encoding_format: Literal["base64", "float"] | NotGiven = NOT_GIVEN,

33 ) -> CreateEmbeddingResponse:

34 """

35 Convenience function for accessing vLLM's Chat Embeddings API,

examples/online_serving/structured_outputs/structured_outputs.py

+2 -8

@@ -1,21 +1,15 @@

1 # ruff: noqa: E501

2 # SPDX-License-Identifier: Apache-2.0

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

4 -

5 -from __future__ import annotations

6 -

4 import argparse

5 import asyncio

6 import enum

7 import os

11 -from typing import TYPE_CHECKING, Any, Literal

8 +from typing import Any, Literal

9

10 import openai

11 import pydantic

15 -

16 -if TYPE_CHECKING:

17 - from openai.types.chat import ChatCompletionChunk

18 -

12 +from openai.types.chat import ChatCompletionChunk

13

14 ConstraintsFormat = Literal[

15 "choice",

pyproject.toml

-6

@@ -84,12 +84,6 @@ ignore = [

84 "B007",

85 # f-string format

86 "UP032",

87 - # Can remove once 3.10+ is the minimum Python version

88 - "UP007",

89 - "UP027",

90 - "UP035",

91 - "UP038",

92 - "UP045",

87 ]

88

89 [tool.ruff.format]

tests/benchmarks/test_random_dataset.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import random

4 -from typing import Any, NamedTuple, Optional, cast

4 +from typing import Any, NamedTuple, cast

5

6 import numpy as np

7 import pytest

@@ -185,8 +185,8 @@ def _collect_mm_samples(

185 output_len: int = 5,

186 base_items_per_request: int = 2,

187 num_mm_items_range_ratio: float = 0.0,

188 - limit_mm_per_prompt: Optional[dict[str, int]] = None,

189 - bucket_config: Optional[dict[tuple[int, int, int], float]] = None,

188 + limit_mm_per_prompt: dict[str, int] | None = None,

189 + bucket_config: dict[tuple[int, int, int], float] | None = None,

190 enable_multimodal_chat: bool = False,

191 ) -> list[SampleRequest]:

192 if limit_mm_per_prompt is None:

tests/ci_envs.py

+5 -4

@@ -5,13 +5,14 @@ These envs only work for a small part of the tests, fix what you need!

5 """

6

7 import os

8 -from typing import TYPE_CHECKING, Any, Callable, Optional

8 +from collections.abc import Callable

9 +from typing import TYPE_CHECKING, Any

10

11 if TYPE_CHECKING:

12 VLLM_CI_NO_SKIP: bool = False

12 - VLLM_CI_DTYPE: Optional[str] = None

13 - VLLM_CI_HEAD_DTYPE: Optional[str] = None

14 - VLLM_CI_HF_DTYPE: Optional[str] = None

13 + VLLM_CI_DTYPE: str | None = None

14 + VLLM_CI_HEAD_DTYPE: str | None = None

15 + VLLM_CI_HF_DTYPE: str | None = None

16

17 environment_variables: dict[str, Callable[[], Any]] = {

18 # A model family has many models with the same architecture.

tests/compile/backend.py

+2 -3

@@ -2,9 +2,8 @@

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

3

4 import weakref

5 -from collections.abc import Sequence

5 +from collections.abc import Callable, Sequence

6 from copy import deepcopy

7 -from typing import Callable, Union

7

8 from torch import fx

9 from torch._ops import OpOverload

@@ -44,7 +43,7 @@ class TestBackend:

43 Inductor config is default-initialized from VllmConfig.CompilationConfig.

44 """

45

47 - def __init__(self, *passes: Union[InductorPass, Callable[[fx.Graph], None]]):

46 + def __init__(self, *passes: InductorPass | Callable[[fx.Graph], None]):

47 self.custom_passes = list(passes)

48 compile_config = get_current_vllm_config().compilation_config

49 self.inductor_config = compile_config.inductor_compile_config

tests/compile/piecewise/test_toy_llama.py

+3 -3

@@ -10,7 +10,7 @@ initialized randomly with a fixed seed.

10 """

11

12 from dataclasses import dataclass

13 -from typing import Any, Optional

13 +from typing import Any

14

15 import pytest

16 import torch

@@ -162,7 +162,7 @@ class LlamaDecoderLayer(nn.Module):

162 self,

163 positions: torch.Tensor,

164 hidden_states: torch.Tensor,

165 - residual: Optional[torch.Tensor],

165 + residual: torch.Tensor | None,

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

167 """

168 For tractable computation:

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

217

218 def forward(

219 self,

220 - input_ids: Optional[torch.Tensor],

220 + input_ids: torch.Tensor | None,

221 positions: torch.Tensor,

222 ) -> torch.Tensor:

223 hidden_states = self.embedding_tokens(input_ids)

tests/compile/test_basic_correctness.py

-2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import dataclasses

4

5 import pytest

tests/compile/test_full_graph.py

+2 -4

@@ -1,11 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import logging

5 import tempfile

8 -from typing import Any, Union

6 +from typing import Any

7

8 import pytest

9 import torch

@@ -217,7 +215,7 @@ def test_inductor_graph_partition_attn_fusion(caplog_vllm):

215

216

217 def run_model(

220 - compile_config: Union[int, CompilationConfig],

218 + compile_config: int | CompilationConfig,

219 model: str,

220 model_kwargs: dict[str, Any],

221 ):

tests/compile/test_fusion_attn.py

+2 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import copy

4 -from typing import Optional

4

5 import pytest

6 import torch._dynamo

@@ -41,8 +40,8 @@ FP8_DTYPE = current_platform.fp8_dtype()

40 FP4_DTYPE = torch.uint8

41

42 # globals needed for string-import custom Dynamo backend field

44 -backend: Optional[TestBackend] = None

45 -backend_unfused: Optional[TestBackend] = None

43 +backend: TestBackend | None = None

44 +backend_unfused: TestBackend | None = None

45

46

47 class AttentionQuantPatternModel(torch.nn.Module):

tests/compile/test_wrapper.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6

@@ -10,7 +9,7 @@ from vllm.config import CompilationLevel

9

10

11 class MyMod(torch.nn.Module):

13 - def forward(self, x: torch.Tensor, cache: Optional[torch.Tensor] = None):

12 + def forward(self, x: torch.Tensor, cache: torch.Tensor | None = None):

13 if cache is not None:

14 return x + cache

15 return x * 2

@@ -24,11 +23,11 @@ class MyWrapper(TorchCompileWrapperWithCustomDispatcher):

23 compiled_callable, compilation_level=CompilationLevel.DYNAMO_ONCE

24 )

25

27 - def forward(self, x: torch.Tensor, cache: Optional[torch.Tensor] = None):

26 + def forward(self, x: torch.Tensor, cache: torch.Tensor | None = None):

27 # this is the function to be compiled

28 return self.model(x, cache)

29

31 - def __call__(self, x: torch.Tensor, cache: Optional[torch.Tensor] = None):

30 + def __call__(self, x: torch.Tensor, cache: torch.Tensor | None = None):

31 # let torch.compile compile twice

32 if len(self.compiled_codes) == 2:

33 dispatch_id = 0 if cache is None else 1

tests/conftest.py

+73 -75

@@ -21,7 +21,7 @@ import threading

21 from collections.abc import Generator

22 from contextlib import nullcontext

23 from enum import Enum

24 -from typing import Any, Callable, Optional, TypedDict, TypeVar, Union, cast

24 +from typing import Any, Callable, TypedDict, TypeVar, cast

25

26 import numpy as np

27 import pytest

@@ -68,7 +68,7 @@ _SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")

68

69 _M = TypeVar("_M")

70

71 -_PromptMultiModalInput = Union[list[_M], list[list[_M]]]

71 +_PromptMultiModalInput = list[_M] | list[list[_M]]

72

73 PromptImageInput = _PromptMultiModalInput[Image.Image]

74 PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]

@@ -267,7 +267,7 @@ class HfRunner:

267

268 return "cpu" if current_platform.is_cpu() else current_platform.device_type

269

270 - def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:

270 + def wrap_device(self, x: _T, device: str | None = None) -> _T:

271 if x is None or isinstance(x, (bool,)):

272 return x

273

@@ -287,14 +287,14 @@ class HfRunner:

287 model_name: str,

288 dtype: str = "auto",

289 *,

290 - model_kwargs: Optional[dict[str, Any]] = None,

290 + model_kwargs: dict[str, Any] | None = None,

291 trust_remote_code: bool = True,

292 is_sentence_transformer: bool = False,

293 is_cross_encoder: bool = False,

294 skip_tokenizer_init: bool = False,

295 auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,

296 # Set this to avoid hanging issue

297 - default_torch_num_threads: Optional[int] = None,

297 + default_torch_num_threads: int | None = None,

298 ) -> None:

299 init_ctx = (

300 nullcontext()

@@ -319,7 +319,7 @@ class HfRunner:

319 model_name: str,

320 dtype: str = "auto",

321 *,

322 - model_kwargs: Optional[dict[str, Any]] = None,

322 + model_kwargs: dict[str, Any] | None = None,

323 trust_remote_code: bool = True,

324 is_sentence_transformer: bool = False,

325 is_cross_encoder: bool = False,

@@ -406,11 +406,11 @@ class HfRunner:

406

407 def get_inputs(

408 self,

409 - prompts: Union[list[str], list[list[int]]],

410 - images: Optional[PromptImageInput] = None,

411 - videos: Optional[PromptVideoInput] = None,

412 - audios: Optional[PromptAudioInput] = None,

413 - ) -> list[Union[BatchFeature, BatchEncoding, dict[str, torch.Tensor]]]:

409 + prompts: list[str] | list[list[int]],

410 + images: PromptImageInput | None = None,

411 + videos: PromptVideoInput | None = None,

412 + audios: PromptAudioInput | None = None,

413 + ) -> list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]]:

414 if images is not None:

415 assert len(prompts) == len(images)

416

@@ -420,9 +420,7 @@ class HfRunner:

420 if audios is not None:

421 assert len(prompts) == len(audios)

422

423 - all_inputs: list[

424 - Union[BatchFeature, BatchEncoding, dict[str, torch.Tensor]]

425 - ] = []

423 + all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []

424 for i, prompt in enumerate(prompts):

425 if isinstance(prompt, str):

426 processor_kwargs: dict[str, Any] = {

@@ -494,10 +492,10 @@ class HfRunner:

492

493 def generate(

494 self,

497 - prompts: Union[list[str], list[list[int]]],

498 - images: Optional[PromptImageInput] = None,

499 - videos: Optional[PromptVideoInput] = None,

500 - audios: Optional[PromptAudioInput] = None,

495 + prompts: list[str] | list[list[int]],

496 + images: PromptImageInput | None = None,

497 + videos: PromptVideoInput | None = None,

498 + audios: PromptAudioInput | None = None,

499 **kwargs: Any,

500 ) -> list[tuple[list[list[int]], list[str]]]:

501 all_inputs = self.get_inputs(

@@ -522,11 +520,11 @@ class HfRunner:

520

521 def generate_greedy(

522 self,

525 - prompts: Union[list[str], list[list[int]]],

523 + prompts: list[str] | list[list[int]],

524 max_tokens: int,

527 - images: Optional[PromptImageInput] = None,

528 - videos: Optional[PromptVideoInput] = None,

529 - audios: Optional[PromptAudioInput] = None,

525 + images: PromptImageInput | None = None,

526 + videos: PromptVideoInput | None = None,

527 + audios: PromptAudioInput | None = None,

528 **kwargs: Any,

529 ) -> list[tuple[list[int], str]]:

530 outputs = self.generate(

@@ -546,9 +544,9 @@ class HfRunner:

544 prompts: list[str],

545 beam_width: int,

546 max_tokens: int,

549 - images: Optional[PromptImageInput] = None,

550 - videos: Optional[PromptVideoInput] = None,

551 - audios: Optional[PromptAudioInput] = None,

547 + images: PromptImageInput | None = None,

548 + videos: PromptVideoInput | None = None,

549 + audios: PromptAudioInput | None = None,

550 ) -> list[tuple[list[list[int]], list[str]]]:

551 outputs = self.generate(

552 prompts,

@@ -574,9 +572,9 @@ class HfRunner:

572 self,

573 prompts: list[str],

574 max_tokens: int,

577 - images: Optional[PromptImageInput] = None,

578 - videos: Optional[PromptVideoInput] = None,

579 - audios: Optional[PromptAudioInput] = None,

575 + images: PromptImageInput | None = None,

576 + videos: PromptVideoInput | None = None,

577 + audios: PromptAudioInput | None = None,

578 **kwargs: Any,

579 ) -> list[list[torch.Tensor]]:

580 all_inputs = self.get_inputs(

@@ -624,7 +622,7 @@ class HfRunner:

622 def _hidden_states_to_logprobs(

623 self,

624 hidden_states: tuple[tuple[torch.Tensor, ...], ...],

627 - num_logprobs: Optional[int],

625 + num_logprobs: int | None,

626 ) -> tuple[list[dict[int, float]], int]:

627 seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)

628 output_len = len(hidden_states)

@@ -652,10 +650,10 @@ class HfRunner:

650 self,

651 prompts: list[str],

652 max_tokens: int,

655 - num_logprobs: Optional[int],

656 - images: Optional[PromptImageInput] = None,

657 - audios: Optional[PromptAudioInput] = None,

658 - videos: Optional[PromptVideoInput] = None,

653 + num_logprobs: int | None,

654 + images: PromptImageInput | None = None,

655 + audios: PromptAudioInput | None = None,

656 + videos: PromptVideoInput | None = None,

657 **kwargs: Any,

658 ) -> list[TokensTextLogprobs]:

659 all_inputs = self.get_inputs(

@@ -734,20 +732,20 @@ class VllmRunner:

732 model_name: str,

733 runner: RunnerOption = "auto",

734 convert: ConvertOption = "auto",

737 - tokenizer_name: Optional[str] = None,

735 + tokenizer_name: str | None = None,

736 tokenizer_mode: str = "auto",

737 trust_remote_code: bool = True,

740 - seed: Optional[int] = 0,

741 - max_model_len: Optional[int] = 1024,

738 + seed: int | None = 0,

739 + max_model_len: int | None = 1024,

740 dtype: str = "auto",

741 disable_log_stats: bool = True,

742 tensor_parallel_size: int = 1,

743 block_size: int = 16 if not torch.xpu.is_available() else 64,

746 - enable_chunked_prefill: Optional[bool] = False,

744 + enable_chunked_prefill: bool | None = False,

745 swap_space: int = 4,

748 - enforce_eager: Optional[bool] = False,

746 + enforce_eager: bool | None = False,

747 # Set this to avoid hanging issue

750 - default_torch_num_threads: Optional[int] = None,

748 + default_torch_num_threads: int | None = None,

749 **kwargs,

750 ) -> None:

751 init_ctx = (

@@ -785,10 +783,10 @@ class VllmRunner:

783

784 def get_inputs(

785 self,

788 - prompts: Union[list[str], list[torch.Tensor], list[list[int]]],

789 - images: Optional[PromptImageInput] = None,

790 - videos: Optional[PromptVideoInput] = None,

791 - audios: Optional[PromptAudioInput] = None,

786 + prompts: list[str] | list[torch.Tensor] | list[list[int]],

787 + images: PromptImageInput | None = None,

788 + videos: PromptVideoInput | None = None,

789 + audios: PromptAudioInput | None = None,

790 ) -> list[dict[str, Any]]:

791 if any(

792 x is not None and len(x) != len(prompts) for x in [images, videos, audios]

@@ -824,11 +822,11 @@ class VllmRunner:

822

823 def generate(

824 self,

827 - prompts: Union[list[str], list[torch.Tensor], list[list[int]]],

825 + prompts: list[str] | list[torch.Tensor] | list[list[int]],

826 sampling_params: SamplingParams,

829 - images: Optional[PromptImageInput] = None,

830 - videos: Optional[PromptVideoInput] = None,

831 - audios: Optional[PromptAudioInput] = None,

827 + images: PromptImageInput | None = None,

828 + videos: PromptVideoInput | None = None,

829 + audios: PromptAudioInput | None = None,

830 **kwargs: Any,

831 ) -> list[tuple[list[list[int]], list[str]]]:

832 inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)

@@ -871,11 +869,11 @@ class VllmRunner:

869 self,

870 prompts: list[str],

871 sampling_params: SamplingParams,

874 - images: Optional[PromptImageInput] = None,

875 - audios: Optional[PromptAudioInput] = None,

876 - videos: Optional[PromptVideoInput] = None,

872 + images: PromptImageInput | None = None,

873 + audios: PromptAudioInput | None = None,

874 + videos: PromptVideoInput | None = None,

875 **kwargs: Any,

878 - ) -> Union[list[TokensTextLogprobs], list[TokensTextLogprobsPromptLogprobs]]:

876 + ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:

877 inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)

878

879 req_outputs = self.llm.generate(

@@ -894,11 +892,11 @@ class VllmRunner:

892

893 def generate_greedy(

894 self,

897 - prompts: Union[list[str], list[torch.Tensor], list[list[int]]],

895 + prompts: list[str] | list[torch.Tensor] | list[list[int]],

896 max_tokens: int,

899 - images: Optional[PromptImageInput] = None,

900 - videos: Optional[PromptVideoInput] = None,

901 - audios: Optional[PromptAudioInput] = None,

897 + images: PromptImageInput | None = None,

898 + videos: PromptVideoInput | None = None,

899 + audios: PromptAudioInput | None = None,

900 **kwargs: Any,

901 ) -> list[tuple[list[int], str]]:

902 greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)

@@ -916,15 +914,15 @@ class VllmRunner:

914 self,

915 prompts: list[str],

916 max_tokens: int,

919 - num_logprobs: Optional[int],

920 - num_prompt_logprobs: Optional[int] = None,

921 - images: Optional[PromptImageInput] = None,

922 - audios: Optional[PromptAudioInput] = None,

923 - videos: Optional[PromptVideoInput] = None,

924 - stop_token_ids: Optional[list[int]] = None,

925 - stop: Optional[list[str]] = None,

917 + num_logprobs: int | None,

918 + num_prompt_logprobs: int | None = None,

919 + images: PromptImageInput | None = None,

920 + audios: PromptAudioInput | None = None,

921 + videos: PromptVideoInput | None = None,

922 + stop_token_ids: list[int] | None = None,

923 + stop: list[str] | None = None,

924 **kwargs: Any,

927 - ) -> Union[list[TokensTextLogprobs], list[TokensTextLogprobsPromptLogprobs]]:

925 + ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:

926 greedy_logprobs_params = SamplingParams(

927 temperature=0.0,

928 max_tokens=max_tokens,

@@ -957,7 +955,7 @@ class VllmRunner:

955 perplexities = []

956 for output in outputs:

957 output = cast(TokensTextLogprobsPromptLogprobs, output)

960 - token_datas = cast(list[Optional[dict[int, Logprob]]], output[3])

958 + token_datas = cast(list[dict[int, Logprob] | None], output[3])

959 assert token_datas[0] is None

960 token_log_probs = []

961 for token_data in token_datas[1:]:

@@ -976,10 +974,10 @@ class VllmRunner:

974 prompts: list[str],

975 beam_width: int,

976 max_tokens: int,

979 - images: Optional[PromptImageInput] = None,

980 - videos: Optional[PromptVideoInput] = None,

981 - audios: Optional[PromptAudioInput] = None,

982 - concurrency_limit: Optional[int] = None,

977 + images: PromptImageInput | None = None,

978 + videos: PromptVideoInput | None = None,

979 + audios: PromptAudioInput | None = None,

980 + concurrency_limit: int | None = None,

981 ) -> list[tuple[list[list[int]], list[str]]]:

982 inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)

983

@@ -1002,9 +1000,9 @@ class VllmRunner:

1000 def embed(

1001 self,

1002 prompts: list[str],

1005 - images: Optional[PromptImageInput] = None,

1006 - videos: Optional[PromptVideoInput] = None,

1007 - audios: Optional[PromptAudioInput] = None,

1003 + images: PromptImageInput | None = None,

1004 + videos: PromptVideoInput | None = None,

1005 + audios: PromptAudioInput | None = None,

1006 *args,

1007 **kwargs,

1008 ) -> list[list[float]]:

@@ -1023,8 +1021,8 @@ class VllmRunner:

1021

1022 def score(

1023 self,

1026 - text_1: Union[str, list[str]],

1027 - text_2: Union[str, list[str]],

1024 + text_1: list[str] | str,

1025 + text_2: list[str] | str,

1026 *args,

1027 **kwargs,

1028 ) -> list[float]:

@@ -1226,8 +1224,8 @@ def _find_free_port() -> int:

1224 class LocalAssetServer:

1225 address: str

1226 port: int

1229 - server: Optional[http.server.ThreadingHTTPServer]

1230 - thread: Optional[threading.Thread]

1227 + server: http.server.ThreadingHTTPServer | None

1228 + thread: threading.Thread | None

1229

1230 def __init__(self, address: str = "127.0.0.1") -> None:

1231 self.address = address

tests/detokenizer/test_stop_strings.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Optional

4 +from typing import Any

5

6 import pytest

7

@@ -15,8 +15,8 @@ def _test_stopping(

15 llm: LLM,

16 expected_output: str,

17 expected_reason: Any,

18 - stop: Optional[list[str]] = None,

19 - stop_token_ids: Optional[list[int]] = None,

18 + stop: list[str] | None = None,

19 + stop_token_ids: list[int] | None = None,

20 include_in_output: bool = False,

21 ) -> None:

22 output = llm.generate(

tests/distributed/conftest.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import random

4 -from typing import Optional, Union

4

5 import msgspec

6 import msgspec.msgpack

@@ -78,8 +77,8 @@ class MockSubscriber:

77

78 def __init__(

79 self,

81 - pub_endpoints: Union[str, list[str]],

82 - replay_endpoints: Optional[Union[str, list[str]]] = None,

80 + pub_endpoints: str | list[str],

81 + replay_endpoints: str | list[str] | None = None,

82 topic: str = "",

83 decode_type=SampleBatch,

84 ):

@@ -111,7 +110,7 @@ class MockSubscriber:

110 self.last_seq = -1

111 self.decoder = msgspec.msgpack.Decoder(type=decode_type)

112

114 - def receive_one(self, timeout=1000) -> Union[tuple[int, SampleBatch], None]:

113 + def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None:

114 """Receive a single message with timeout"""

115 if not self.sub.poll(timeout):

116 return None

tests/distributed/test_comm_ops.py

+2 -3

@@ -5,9 +5,8 @@

5 Run `pytest tests/distributed/test_comm_ops.py`.

6 """

7

8 -from __future__ import annotations

9 -

10 -from typing import Any, Callable

8 +from collections.abc import Callable

9 +from typing import Any

10

11 import pytest

12 import ray

tests/distributed/test_context_parallel.py

+3 -3

@@ -11,7 +11,7 @@ WARNING: This test runs in both single-node (4 GPUs) and multi-node

11 import json

12 import os

13 from dataclasses import dataclass

14 -from typing import Literal, NamedTuple, Optional

14 +from typing import Literal, NamedTuple

15

16 import pytest

17

@@ -36,7 +36,7 @@ class ParallelSetup(NamedTuple):

36

37 class CPTestOptions(NamedTuple):

38 multi_node_only: bool

39 - load_format: Optional[str] = None

39 + load_format: str | None = None

40

41

42 @dataclass

@@ -54,7 +54,7 @@ class CPTestSettings:

54 dcp_base: int = 1,

55 multi_node_only: bool = False,

56 runner: RunnerOption = "auto",

57 - load_format: Optional[str] = None,

57 + load_format: str | None = None,

58 ):

59 parallel_setups = []

60 for eager_mode_val in [False]:

tests/distributed/test_expert_parallel.py

+10 -10

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

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

3

4 from dataclasses import dataclass

5 -from typing import Literal, NamedTuple, Optional

5 +from typing import Literal, NamedTuple

6

7 import pytest

8

@@ -22,9 +22,9 @@ class ParallelSetup(NamedTuple):

22

23 class EPTestOptions(NamedTuple):

24 trust_remote_code: bool

25 - tokenizer_mode: Optional[str]

26 - load_format: Optional[str] = None

27 - hf_overrides: Optional[str] = None

25 + tokenizer_mode: str | None

26 + load_format: str | None = None

27 + hf_overrides: str | None = None

28

29

30 @dataclass

@@ -40,9 +40,9 @@ class EPTestSettings:

40 tp_base: int = 2,

41 runner: RunnerOption = "auto",

42 trust_remote_code: bool = False,

43 - tokenizer_mode: Optional[str] = None,

44 - load_format: Optional[str] = None,

45 - hf_overrides: Optional[str] = None,

43 + tokenizer_mode: str | None = None,

44 + load_format: str | None = None,

45 + hf_overrides: str | None = None,

46 ):

47 return EPTestSettings(

48 parallel_setups=[

@@ -72,9 +72,9 @@ class EPTestSettings:

72 tp_base: int = 2,

73 runner: RunnerOption = "auto",

74 trust_remote_code: bool = False,

75 - tokenizer_mode: Optional[str] = None,

76 - load_format: Optional[str] = None,

77 - hf_overrides: Optional[str] = None,

75 + tokenizer_mode: str | None = None,

76 + load_format: str | None = None,

77 + hf_overrides: str | None = None,

78 ):

79 return EPTestSettings(

80 parallel_setups=[

tests/distributed/test_pipeline_parallel.py

+4 -4

@@ -11,7 +11,7 @@ WARNING: This test runs in both single-node (4 GPUs) and multi-node

11 import json

12 import os

13 from dataclasses import dataclass

14 -from typing import Literal, NamedTuple, Optional

14 +from typing import Literal, NamedTuple

15

16 import pytest

17

@@ -35,7 +35,7 @@ class ParallelSetup(NamedTuple):

35

36 class PPTestOptions(NamedTuple):

37 multi_node_only: bool

38 - load_format: Optional[str] = None

38 + load_format: str | None = None

39

40

41 @dataclass

@@ -52,7 +52,7 @@ class PPTestSettings:

52 pp_base: int = 2,

53 multi_node_only: bool = False,

54 runner: RunnerOption = "auto",

55 - load_format: Optional[str] = None,

55 + load_format: str | None = None,

56 ):

57 return PPTestSettings(

58 parallel_setups=[

@@ -76,7 +76,7 @@ class PPTestSettings:

76 pp_base: int = 2,

77 runner: RunnerOption = "auto",

78 multi_node_only: bool = False,

79 - load_format: Optional[str] = None,

79 + load_format: str | None = None,

80 ):

81 return PPTestSettings(

82 parallel_setups=[

tests/distributed/test_pp_cudagraph.py

+1 -7

@@ -1,16 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

5 -from typing import TYPE_CHECKING

6 -

3 import pytest

4 +from typing_extensions import LiteralString

5

6 from ..utils import compare_two_settings, create_new_process_for_each_test

7

11 -if TYPE_CHECKING:

12 - from typing_extensions import LiteralString

13 -

8

9 @pytest.mark.parametrize(

10 "PP_SIZE, MODEL_NAME",

tests/distributed/test_sequence_parallel.py

+5 -5

@@ -11,7 +11,7 @@ WARNING: This test runs in both single-node (4 GPUs) and multi-node

11 import json

12 import os

13 from dataclasses import dataclass

14 -from typing import Literal, NamedTuple, Optional

14 +from typing import Literal, NamedTuple

15

16 import pytest

17

@@ -36,7 +36,7 @@ class ParallelSetup(NamedTuple):

36

37 class SPTestOptions(NamedTuple):

38 multi_node_only: bool

39 - load_format: Optional[str] = None

39 + load_format: str | None = None

40

41

42 @dataclass

@@ -53,7 +53,7 @@ class SPTestSettings:

53 pp_base: int = 1,

54 multi_node_only: bool = False,

55 runner: RunnerOption = "auto",

56 - load_format: Optional[str] = None,

56 + load_format: str | None = None,

57 ):

58 parallel_setups = []

59 for eager_mode_val in [False, True]:

@@ -84,7 +84,7 @@ class SPTestSettings:

84 pp_base: int = 1,

85 runner: RunnerOption = "auto",

86 multi_node_only: bool = False,

87 - load_format: Optional[str] = None,

87 + load_format: str | None = None,

88 ):

89 parallel_setups = []

90 for eager_mode_val in [False, True]:

@@ -115,7 +115,7 @@ class SPTestSettings:

115 pp_base: int = 1,

116 runner: RunnerOption = "auto",

117 multi_node_only: bool = False,

118 - load_format: Optional[str] = None,

118 + load_format: str | None = None,

119 ):

120 parallel_setups = []

121 for fusion_val in [False, True]:

tests/engine/test_arg_utils.py

+8 -8

@@ -5,7 +5,7 @@ import json

5 from argparse import ArgumentError

6 from contextlib import nullcontext

7 from dataclasses import dataclass, field

8 -from typing import Annotated, Literal, Optional, Union

8 +from typing import Annotated, Literal

9

10 import pytest

11

@@ -115,9 +115,9 @@ class NestedConfig:

115 class DummyConfig:

116 regular_bool: bool = True

117 """Regular bool with default True"""

118 - optional_bool: Optional[bool] = None

118 + optional_bool: bool | None = None

119 """Optional bool with default None"""

120 - optional_literal: Optional[Literal["x", "y"]] = None

120 + optional_literal: Literal["x", "y"] | None = None

121 """Optional literal with default None"""

122 tuple_n: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3))

123 """Tuple with variable length"""

@@ -127,7 +127,7 @@ class DummyConfig:

127 """List with variable length"""

128 list_literal: list[Literal[1, 2]] = field(default_factory=list)

129 """List with literal choices"""

130 - list_union: list[Union[str, type[object]]] = field(default_factory=list)

130 + list_union: list[str | type[object]] = field(default_factory=list)

131 """List with union type"""

132 literal_literal: Literal[Literal[1], Literal[2]] = 1

133 """Literal of literals with default 1"""

@@ -152,11 +152,11 @@ def test_is_not_builtin(type_hint, expected):

152 ("type_hint", "expected"),

153 [

154 (Annotated[int, "annotation"], {int}),

155 - (Optional[int], {int, type(None)}),

156 - (Annotated[Optional[int], "annotation"], {int, type(None)}),

157 - (Optional[Annotated[int, "annotation"]], {int, type(None)}),

155 + (int | None, {int, type(None)}),

156 + (Annotated[int | None, "annotation"], {int, type(None)}),

157 + (Annotated[int, "annotation"] | None, {int, type(None)}),

158 ],

159 - ids=["Annotated", "Optional", "Annotated_Optional", "Optional_Annotated"],

159 + ids=["Annotated", "or_None", "Annotated_or_None", "or_None_Annotated"],

160 )

161 def test_get_type_hints(type_hint, expected):

162 assert get_type_hints(type_hint) == expected

tests/entrypoints/openai/test_async_tokenization.py

+1 -1

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

3

4 import asyncio

5 import random

6 -from typing import Callable

6 +from collections.abc import Callable

7

8 import openai

9 import pytest

tests/entrypoints/openai/test_chat.py

+1 -2

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

3

4 # imports for structured outputs tests

5 import json

6 -from typing import Optional

6

7 import jsonschema

8 import openai # use the official client for correctness check

@@ -176,7 +175,7 @@ async def test_too_many_chat_logprobs(client: openai.AsyncOpenAI, model_name: st

175 [(MODEL_NAME, 1), (MODEL_NAME, 0), (MODEL_NAME, -1), (MODEL_NAME, None)],

176 )

177 async def test_prompt_logprobs_chat(

179 - client: openai.AsyncOpenAI, model_name: str, prompt_logprobs: Optional[int]

178 + client: openai.AsyncOpenAI, model_name: str, prompt_logprobs: int | None

179 ):

180 params: dict = {

181 "messages": [

tests/entrypoints/openai/test_completion_with_function_calling.py

+1 -2

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

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

3

4 import datetime

5 -from typing import Union

5

6 import openai # use the official client for correctness check

7 import pytest

@@ -166,7 +165,7 @@ async def test_function_tool_use(

165 client: openai.AsyncOpenAI,

166 model_name: str,

167 stream: bool,

169 - tool_choice: Union[str, dict],

168 + tool_choice: str | dict,

169 enable_thinking: bool,

170 ):

171 if not stream:

tests/entrypoints/openai/test_lora_resolvers.py

+5 -6

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

4 from contextlib import suppress

5 from dataclasses import dataclass, field

6 from http import HTTPStatus

7 -from typing import Optional

7 from unittest.mock import AsyncMock, MagicMock

8

9 import pytest

@@ -38,13 +37,13 @@ class MockModelConfig:

37 trust_remote_code: bool = False

38 tokenizer_mode: str = "auto"

39 max_model_len: int = 100

41 - tokenizer_revision: Optional[str] = None

40 + tokenizer_revision: str | None = None

41 multimodal_config: MultiModalConfig = field(default_factory=MultiModalConfig)

42 hf_config: MockHFConfig = field(default_factory=MockHFConfig)

44 - logits_processor_pattern: Optional[str] = None

45 - diff_sampling_param: Optional[dict] = None

43 + logits_processor_pattern: str | None = None

44 + diff_sampling_param: dict | None = None

45 allowed_local_media_path: str = ""

47 - allowed_media_domains: Optional[list[str]] = None

46 + allowed_media_domains: list[str] | None = None

47 encoder_config = None

48 generation_config: str = "auto"

49 skip_tokenizer_init: bool = False

@@ -56,7 +55,7 @@ class MockModelConfig:

55 class MockLoRAResolver(LoRAResolver):

56 async def resolve_lora(

57 self, base_model_name: str, lora_name: str

59 - ) -> Optional[LoRARequest]:

58 + ) -> LoRARequest | None:

59 if lora_name == "test-lora":

60 return LoRARequest(

61 lora_name="test-lora",

tests/entrypoints/openai/test_serving_chat.py

+2 -7

@@ -1,16 +1,14 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from __future__ import annotations

5 -

3 import asyncio

4 from contextlib import suppress

5 from dataclasses import dataclass, field

9 -from typing import TYPE_CHECKING, Any

6 +from typing import Any

7 from unittest.mock import AsyncMock, MagicMock

8

9 import pytest

10 import pytest_asyncio

11 +from openai import OpenAI

12

13 from vllm.config.multimodal import MultiModalConfig

14 from vllm.entrypoints.openai.protocol import ChatCompletionRequest

@@ -21,9 +19,6 @@ from vllm.v1.engine.async_llm import AsyncLLM

19

20 from ...utils import RemoteOpenAIServer

21

24 -if TYPE_CHECKING:

25 - from openai import OpenAI

26 -

22 GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b"

23

24

tests/entrypoints/openai/tool_parsers/utils.py

+4 -5

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

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

3

4 from collections.abc import Iterable

5 -from typing import Union

5

6 from vllm.entrypoints.openai.protocol import (

7 ChatCompletionRequest,

@@ -84,10 +83,10 @@ class StreamingToolReconstructor:

83 def run_tool_extraction(

84 tool_parser: ToolParser,

85 model_output: str,

87 - request: Union[ChatCompletionRequest, None] = None,

86 + request: ChatCompletionRequest | None = None,

87 streaming: bool = False,

88 assert_one_tool_per_delta: bool = True,

90 -) -> tuple[Union[str, None], list[ToolCall]]:

89 +) -> tuple[str | None, list[ToolCall]]:

90 if streaming:

91 reconstructor = run_tool_extraction_streaming(

92 tool_parser,

@@ -105,7 +104,7 @@ def run_tool_extraction(

104 def run_tool_extraction_nonstreaming(

105 tool_parser: ToolParser,

106 model_output: str,

108 - request: Union[ChatCompletionRequest, None] = None,

107 + request: ChatCompletionRequest | None = None,

108 ) -> ExtractedToolCallInformation:

109 request = request or ChatCompletionRequest(messages=[], model="test-model")

110 return tool_parser.extract_tool_calls(model_output, request)

@@ -114,7 +113,7 @@ def run_tool_extraction_nonstreaming(

113 def run_tool_extraction_streaming(

114 tool_parser: ToolParser,

115 model_deltas: Iterable[str],

117 - request: Union[ChatCompletionRequest, None] = None,

116 + request: ChatCompletionRequest | None = None,

117 assert_one_tool_per_delta: bool = True,

118 ) -> StreamingToolReconstructor:

119 request = request or ChatCompletionRequest(messages=[], model="test-model")

tests/entrypoints/pooling/openai/test_embedding_dimensions.py

+2 -4

@@ -4,8 +4,6 @@

4 Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.

5 """

6

7 -from typing import Optional

8 -

7 import openai

8 import pytest

9

@@ -103,14 +101,14 @@ async def test_matryoshka(

101 run_embedding_correctness_test(hf_model, prompts, vllm_outputs, dimensions)

102

103 if model_info.is_matryoshka:

106 - valid_dimensions: list[Optional[int]] = [None]

104 + valid_dimensions: list[int | None] = [None]

105 if model_info.matryoshka_dimensions is not None:

106 valid_dimensions += model_info.matryoshka_dimensions[:2]

107

108 for dimensions in valid_dimensions:

109 await make_request_and_correctness_test(dimensions)

110

113 - invalid_dimensions: list[Optional[int]] = [-1]

111 + invalid_dimensions: list[int | None] = [-1]

112 if model_info.matryoshka_dimensions is not None:

113 assert 5 not in model_info.matryoshka_dimensions

114 invalid_dimensions.append(5)

tests/entrypoints/test_api_server_process_manager.py

+2 -3

@@ -5,7 +5,6 @@ import multiprocessing

5 import socket

6 import threading

7 import time

8 -from typing import Optional

8 from unittest.mock import patch

9

10 import pytest

@@ -105,7 +104,7 @@ def test_wait_for_completion_or_failure(api_server_args):

104 assert len(manager.processes) == 3

105

106 # Create a result capture for the thread

108 - result: dict[str, Optional[Exception]] = {"exception": None}

107 + result: dict[str, Exception | None] = {"exception": None}

108

109 def run_with_exception_capture():

110 try:

@@ -218,7 +217,7 @@ def test_external_process_monitoring(api_server_args):

217 assert len(manager.processes) == 3

218

219 # Create a result capture for the thread

221 - result: dict[str, Optional[Exception]] = {"exception": None}

220 + result: dict[str, Exception | None] = {"exception": None}

221

222 def run_with_exception_capture():

223 try:

tests/entrypoints/test_chat_utils.py

+7 -7

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

3

4 import warnings

5 from collections.abc import Mapping

6 -from typing import Literal, Optional

6 +from typing import Literal

7

8 import pytest

9 from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy

@@ -152,9 +152,9 @@ def audio_url():

152

153

154 def _assert_mm_data_is_image_input(

155 - mm_data: Optional[MultiModalDataDict],

155 + mm_data: MultiModalDataDict | None,

156 image_count: int,

157 - skipped_image_indices: Optional[list] = None,

157 + skipped_image_indices: list | None = None,

158 ) -> None:

159 assert mm_data is not None

160 assert set(mm_data.keys()) == {"image"}

@@ -169,9 +169,9 @@ def _assert_mm_data_is_image_input(

169

170

171 def _assert_mm_uuids(

172 - mm_uuids: Optional[MultiModalUUIDDict],

172 + mm_uuids: MultiModalUUIDDict | None,

173 media_count: int,

174 - expected_uuids: list[Optional[str]],

174 + expected_uuids: list[str | None],

175 modality: str = "image",

176 ) -> None:

177 if len(expected_uuids) > 0:

@@ -193,9 +193,9 @@ MultiModalDataCounts = Mapping[ModalityType, int]

193

194

195 def _assert_mm_data_inputs(

196 - mm_data: Optional[MultiModalDataDict],

196 + mm_data: MultiModalDataDict | None,

197 data_count: MultiModalDataCounts,

198 - skipped_media_indices: Optional[dict[str, list]] = None, # modality -> list[int]

198 + skipped_media_indices: dict[str, list] | None = None, # modality -> list[int]

199 ) -> None:

200 assert mm_data is not None

201 assert set(data_count.keys()) == (set(mm_data.keys()))

tests/entrypoints/test_renderer.py

+1 -2

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

3

4 import io

5 from dataclasses import dataclass

6 -from typing import Optional

6 from unittest.mock import AsyncMock, MagicMock

7

8 import pybase64

@@ -17,7 +16,7 @@ from vllm.inputs.data import is_embeds_prompt

16 @dataclass

17 class MockModelConfig:

18 max_model_len: int = 100

20 - encoder_config: Optional[dict] = None

19 + encoder_config: dict | None = None

20

21

22 class MockTokenizerResult:

tests/evals/gsm8k/gsm8k_eval.py

+6 -7

@@ -12,7 +12,6 @@ import json

12 import os

13 import time

14 from collections.abc import Generator

15 -from typing import Optional, Union

15

16 import aiohttp

17 import numpy as np

@@ -23,7 +22,7 @@ from tqdm.asyncio import tqdm

22 INVALID = -9999999

23

24

26 -def download_and_cache_file(url: str, filename: Optional[str] = None) -> str:

25 +def download_and_cache_file(url: str, filename: str | None = None) -> str:

26 """Download and cache a file from a URL."""

27 if filename is None:

28 filename = os.path.join("/tmp", url.split("/")[-1])

@@ -81,9 +80,9 @@ async def call_vllm_api(

80 prompt: str,

81 temperature: float,

82 max_tokens: int,

84 - stop: Optional[list[str]] = None,

85 - url: Optional[str] = None,

86 - seed: Optional[int] = None,

83 + stop: list[str] | None = None,

84 + url: str | None = None,

85 + seed: int | None = None,

86 ) -> str:

87 """Call vLLM's OpenAI-compatible completions endpoint."""

88 data = {

@@ -112,8 +111,8 @@ def evaluate_gsm8k(

111 host: str = "http://127.0.0.1",

112 port: int = 8000,

113 temperature: float = 0.0,

115 - seed: Optional[int] = 42,

116 -) -> dict[str, Union[float, int]]:

114 + seed: int | None = 42,

115 +) -> dict[str, float | int]:

116 """

117 Evaluate GSM8K accuracy using vLLM serve endpoint.

118

tests/kernels/attention/test_aiter_flash_attn.py

+5 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6 import torch

@@ -27,8 +26,8 @@ def ref_paged_attn(

26 kv_lens: list[int],

27 block_tables: torch.Tensor,

28 scale: float,

30 - sliding_window: Optional[int] = None,

31 - soft_cap: Optional[float] = None,

29 + sliding_window: int | None = None,

30 + soft_cap: float | None = None,

31 ) -> torch.Tensor:

32 num_seqs = len(query_lens)

33 block_tables = block_tables.cpu().numpy()

@@ -94,12 +93,12 @@ def test_varlen_with_paged_kv(

93 seq_lens: list[tuple[int, int]],

94 num_heads: tuple[int, int],

95 head_size: int,

97 - sliding_window: Optional[int],

96 + sliding_window: int | None,

97 dtype: torch.dtype,

98 block_size: int,

100 - soft_cap: Optional[float],

99 + soft_cap: float | None,

100 num_blocks: int,

102 - q_dtype: Optional[torch.dtype],

101 + q_dtype: torch.dtype | None,

102 ) -> None:

103 torch.set_default_device("cuda")

104 current_platform.seed_everything(0)

tests/kernels/attention/test_attention.py

+3 -4

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

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

3

4 import random

5 -from typing import Optional

5

6 import pytest

7 import torch

@@ -50,7 +49,7 @@ def ref_masked_attention(

49 key: torch.Tensor,

50 value: torch.Tensor,

51 scale: float,

53 - attn_mask: Optional[torch.Tensor] = None,

52 + attn_mask: torch.Tensor | None = None,

53 ) -> torch.Tensor:

54 attn_weights = scale * torch.einsum("qhd,khd->hqk", query, key).float()

55 if attn_mask is not None:

@@ -69,7 +68,7 @@ def ref_single_query_cached_kv_attention(

68 block_tables: torch.Tensor,

69 seq_lens: torch.Tensor,

70 scale: float,

72 - alibi_slopes: Optional[torch.Tensor],

71 + alibi_slopes: torch.Tensor | None,

72 ) -> None:

73 num_query_heads = query.shape[1]

74 num_kv_heads = value_cache.shape[1]

@@ -415,7 +414,7 @@ def ref_multi_query_kv_attention(

414 key: torch.Tensor,

415 value: torch.Tensor,

416 scale: float,

418 - alibi_bias: Optional[list[torch.Tensor]],

417 + alibi_bias: list[torch.Tensor] | None,

418 dtype: torch.dtype,

419 ) -> torch.Tensor:

420 num_seqs = len(cu_seq_lens) - 1

tests/kernels/attention/test_cascade_flash_attn.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6 import torch

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

84 head_size: int,

85 dtype: torch.dtype,

86 block_size: int,

88 - soft_cap: Optional[float],

87 + soft_cap: float | None,

88 num_blocks: int,

89 fa_version: int,

90 ) -> None:

tests/kernels/attention/test_cutlass_mla_decode.py

+1 -2

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

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

3 import math

4 import random

5 -from typing import Optional

5

6 import pytest

7 import torch

@@ -17,7 +16,7 @@ def cal_diff(

16 y: torch.Tensor,

17 name: str,

18 use_fp8: bool = False,

20 - diff_threshold: Optional[float] = None,

19 + diff_threshold: float | None = None,

20 ) -> None:

21 x, y = x.double(), y.double()

22 cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12)

tests/kernels/attention/test_flash_attn.py

+8 -9

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6 import torch

@@ -34,8 +33,8 @@ def ref_paged_attn(

33 kv_lens: list[int],

34 block_tables: torch.Tensor,

35 scale: float,

37 - sliding_window: Optional[int] = None,

38 - soft_cap: Optional[float] = None,

36 + sliding_window: int | None = None,

37 + soft_cap: float | None = None,

38 ) -> torch.Tensor:

39 num_seqs = len(query_lens)

40 block_tables = block_tables.cpu().numpy()

@@ -103,11 +102,11 @@ def test_flash_attn_with_paged_kv(

102 head_size: int,

103 dtype: torch.dtype,

104 block_size: int,

106 - soft_cap: Optional[float],

105 + soft_cap: float | None,

106 num_blocks: int,

108 - sliding_window: Optional[int],

107 + sliding_window: int | None,

108 fa_version: int,

110 - q_dtype: Optional[torch.dtype],

109 + q_dtype: torch.dtype | None,

110 ) -> None:

111 torch.set_default_device("cuda")

112 if not is_fa_version_supported(fa_version):

@@ -221,13 +220,13 @@ def test_varlen_with_paged_kv(

220 seq_lens: list[tuple[int, int]],

221 num_heads: tuple[int, int],

222 head_size: int,

224 - sliding_window: Optional[int],

223 + sliding_window: int | None,

224 dtype: torch.dtype,

225 block_size: int,

227 - soft_cap: Optional[float],

226 + soft_cap: float | None,

227 num_blocks: int,

228 fa_version: int,

230 - q_dtype: Optional[torch.dtype],

229 + q_dtype: torch.dtype | None,

230 ) -> None:

231 torch.set_default_device("cuda")

232 if not is_fa_version_supported(fa_version):

tests/kernels/attention/test_flashinfer.py

+8 -9

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import flashinfer

6 import pytest

@@ -26,8 +25,8 @@ def ref_paged_attn(

25 kv_lens: list[int],

26 block_tables: torch.Tensor,

27 scale: float,

29 - sliding_window: Optional[int] = None,

30 - soft_cap: Optional[float] = None,

28 + sliding_window: int | None = None,

29 + soft_cap: float | None = None,

30 ) -> torch.Tensor:

31 num_seqs = len(query_lens)

32 block_tables = block_tables.cpu().numpy()

@@ -90,8 +89,8 @@ def test_flashinfer_decode_with_paged_kv(

89 head_size: int,

90 dtype: torch.dtype,

91 block_size: int,

93 - soft_cap: Optional[float],

94 - sliding_window: Optional[int],

92 + soft_cap: float | None,

93 + sliding_window: int | None,

94 ) -> None:

95 torch.set_default_device("cuda")

96 current_platform.seed_everything(0)

@@ -185,8 +184,8 @@ def test_flashinfer_prefill_with_paged_kv(

184 head_size: int,

185 dtype: torch.dtype,

186 block_size: int,

188 - soft_cap: Optional[float],

189 - sliding_window: Optional[int],

187 + soft_cap: float | None,

188 + sliding_window: int | None,

189 ) -> None:

190 torch.set_default_device("cuda")

191 current_platform.seed_everything(0)

@@ -288,7 +287,7 @@ def test_flashinfer_prefill_with_paged_fp8_kv(

287 head_size: int,

288 dtype: torch.dtype,

289 block_size: int,

291 - soft_cap: Optional[float],

290 + soft_cap: float | None,

291 ) -> None:

292 pytest.skip("TODO: fix the accuracy issue")

293 torch.set_default_device("cuda")

@@ -398,7 +397,7 @@ def test_flashinfer_decode_with_paged_fp8_kv(

397 head_size: int,

398 dtype: torch.dtype,

399 block_size: int,

401 - soft_cap: Optional[float],

400 + soft_cap: float | None,

401 ) -> None:

402 # test doesn't work for num_heads = (16,16)

403 torch.set_default_device("cuda")

tests/kernels/attention/test_flashinfer_trtllm_attention.py

+4 -9

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import flashinfer

5 import pytest

@@ -68,9 +67,7 @@ NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.

67 @torch.inference_mode

68 def test_flashinfer_trtllm_decode_with_baseline(

69 dtype: torch.dtype,

71 - quant_dtypes: tuple[

72 - Optional[torch.dtype], Optional[torch.dtype], Optional[torch.dtype]

73 - ],

70 + quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None],

71 batch_size: int,

72 max_seq_lens: tuple[int, int],

73 num_heads: tuple[int, int],

@@ -78,7 +75,7 @@ def test_flashinfer_trtllm_decode_with_baseline(

75 kv_layout: str,

76 block_size: int,

77 window_left: int,

81 - soft_cap: Optional[float],

78 + soft_cap: float | None,

79 has_sinks: bool,

80 ) -> None:

81 torch.set_default_device("cuda")

@@ -267,9 +264,7 @@ def test_flashinfer_trtllm_decode_with_baseline(

264 @torch.inference_mode

265 def test_flashinfer_trtllm_prefill_with_baseline(

266 dtype: torch.dtype,

270 - quant_dtypes: tuple[

271 - Optional[torch.dtype], Optional[torch.dtype], Optional[torch.dtype]

272 - ],

267 + quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None],

268 batch_size: int,

269 max_seq_lens: tuple[int, int],

270 num_heads: tuple[int, int],

@@ -277,7 +272,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(

272 kv_layout: str,

273 block_size: int,

274 window_left: int,

280 - soft_cap: Optional[float],

275 + soft_cap: float | None,

276 has_sinks: bool,

277 ) -> None:

278 torch.set_default_device("cuda")

tests/kernels/attention/test_merge_attn_states.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import pytest

5 import torch

@@ -20,7 +19,7 @@ def merge_attn_states_torch(

19 prefix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]

20 suffix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]

21 suffix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]

23 - output_lse: Optional[torch.Tensor] = None, # [NUM_HEADS, NUM_TOKENS]

22 + output_lse: torch.Tensor | None = None, # [NUM_HEADS, NUM_TOKENS]

23 ):

24 p_lse = prefix_lse

25 s_lse = suffix_lse

tests/kernels/attention/test_triton_unified_attention.py

+5 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6 import torch

@@ -32,8 +31,8 @@ def ref_paged_attn(

31 kv_lens: list[int],

32 block_tables: torch.Tensor,

33 scale: float,

35 - sliding_window: Optional[int] = None,

36 - soft_cap: Optional[float] = None,

34 + sliding_window: int | None = None,

35 + soft_cap: float | None = None,

36 ) -> torch.Tensor:

37 num_seqs = len(query_lens)

38 block_tables = block_tables.cpu().numpy()

@@ -98,12 +97,12 @@ def test_triton_unified_attn(

97 seq_lens: list[tuple[int, int]],

98 num_heads: tuple[int, int],

99 head_size: int,

101 - sliding_window: Optional[int],

100 + sliding_window: int | None,

101 dtype: torch.dtype,

102 block_size: int,

104 - soft_cap: Optional[float],

103 + soft_cap: float | None,

104 num_blocks: int,

106 - q_dtype: Optional[torch.dtype],

105 + q_dtype: torch.dtype | None,

106 ) -> None:

107 torch.set_default_device("cuda")

108

tests/kernels/core/test_fused_quant_layernorm.py

+15 -16

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 import pytest

6 import torch

@@ -31,13 +30,13 @@ EPS = 1e-6

30 ## Helpers

31

32

34 -def as_float32_tensor(x: Union[float, torch.tensor]) -> torch.tensor:

33 +def as_float32_tensor(x: float | torch.Tensor) -> torch.Tensor:

34 return torch.as_tensor(x, dtype=torch.float32, device="cuda")

35

36

37 def ref_rms_norm(

39 - rms_norm_layer: RMSNorm, x: torch.Tensor, residual: Optional[torch.Tensor]

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

38 + rms_norm_layer: RMSNorm, x: torch.Tensor, residual: torch.Tensor | None

39 +) -> tuple[torch.Tensor, torch.Tensor | None]:

40 if residual is not None:

41 residual = residual.clone()

42 out, residual = rms_norm_layer.forward_native(x, residual)

@@ -51,9 +50,9 @@ def ref_dynamic_per_token_quant(

50 rms_norm_layer: RMSNorm,

51 x: torch.Tensor,

52 quant_dtype: torch.dtype,

54 - residual: Optional[torch.Tensor],

55 - scale_ub: Optional[torch.Tensor],

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

53 + residual: torch.Tensor | None,

54 + scale_ub: torch.Tensor | None,

55 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

56 if scale_ub is not None:

57 assert quant_dtype == torch.float8_e4m3fn

58

@@ -76,9 +75,9 @@ def ref_impl(

75 rms_norm_layer: RMSNorm,

76 x: torch.Tensor,

77 quant_dtype: torch.dtype,

79 - residual: Optional[torch.Tensor],

80 - scale_ub: Optional[torch.Tensor],

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

78 + residual: torch.Tensor | None,

79 + scale_ub: torch.Tensor | None,

80 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

81 return ref_dynamic_per_token_quant(

82 rms_norm_layer, x, quant_dtype, residual, scale_ub

83 )

@@ -88,9 +87,9 @@ def ops_dynamic_per_token_quant(

87 weight: torch.Tensor,

88 x: torch.Tensor,

89 quant_dtype: torch.dtype,

91 - residual: Optional[torch.Tensor],

92 - scale_ub: Optional[torch.Tensor],

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

90 + residual: torch.Tensor | None,

91 + scale_ub: torch.Tensor | None,

92 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

93 if residual is not None:

94 residual = residual.clone()

95 out, scales = ops.rms_norm_dynamic_per_token_quant(

@@ -103,9 +102,9 @@ def ops_impl(

102 weight: torch.Tensor,

103 x: torch.Tensor,

104 quant_dtype: torch.dtype,

106 - residual: Optional[torch.Tensor],

107 - scale_ub: Optional[torch.Tensor],

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

105 + residual: torch.Tensor | None,

106 + scale_ub: torch.Tensor | None,

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

108 return ops_dynamic_per_token_quant(weight, x, quant_dtype, residual, scale_ub)

109

110

tests/kernels/core/test_pos_encoding.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 +from collections.abc import Callable

5 from itertools import product

5 -from typing import Callable, Optional

6

7 import pytest

8 import torch

@@ -68,7 +68,7 @@ def test_rotary_embedding(

68 seq_len: int,

69 num_heads: int,

70 head_size: int,

71 - rotary_dim: Optional[int],

71 + rotary_dim: int | None,

72 dtype: torch.dtype,

73 seed: int,

74 device: str,

tests/kernels/core/test_rotary_embedding.py

+1 -3

@@ -4,8 +4,6 @@

4 Tests for miscellaneous utilities

5 """

6

7 -from typing import Optional

8 -

7 import pytest

8 import torch

9

@@ -17,7 +15,7 @@ def rotary_embedding_opcheck(

15 rot,

16 positions: torch.Tensor,

17 query: torch.Tensor,

20 - key: Optional[torch.Tensor] = None,

18 + key: torch.Tensor | None = None,

19 ):

20 cos_sin_cache = rot.cos_sin_cache.to(query.device, dtype=query.dtype)

21

tests/kernels/mamba/test_causal_conv1d.py

+10 -11

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6 import torch

@@ -19,11 +18,11 @@ from vllm.platforms import current_platform

18 def causal_conv1d_ref(

19 x: torch.Tensor,

20 weight: torch.Tensor,

22 - bias: Optional[torch.Tensor] = None,

23 - initial_states: Optional[torch.Tensor] = None,

21 + bias: torch.Tensor | None = None,

22 + initial_states: torch.Tensor | None = None,

23 return_final_states: bool = False,

25 - final_states_out: Optional[torch.Tensor] = None,

26 - activation: Optional[str] = "silu",

24 + final_states_out: torch.Tensor | None = None,

25 + activation: str | None = "silu",

26 ):

27 """

28 x: (batch, dim, seqlen)

@@ -117,12 +116,12 @@ def causal_conv1d_update_ref(

116 def causal_conv1d_opcheck_fn(

117 x: torch.Tensor,

118 weight: torch.Tensor,

120 - bias: Optional[torch.Tensor] = None,

121 - cu_seq_len: Optional[torch.Tensor] = None,

122 - cache_indices: Optional[torch.Tensor] = None,

123 - has_initial_state: Optional[torch.Tensor] = None,

124 - conv_states: Optional[torch.Tensor] = None,

125 - activation: Optional[str] = "silu",

119 + bias: torch.Tensor | None = None,

120 + cu_seq_len: torch.Tensor | None = None,

121 + cache_indices: torch.Tensor | None = None,

122 + has_initial_state: torch.Tensor | None = None,

123 + conv_states: torch.Tensor | None = None,

124 + activation: str | None = "silu",

125 pad_slot_id: int = PAD_SLOT_ID,

126 ):

127 """

tests/kernels/moe/modular_kernel_tools/common.py

+17 -17

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass

4 -from typing import Any, Optional, Union

4 +from typing import Any

5

6 import torch

7

@@ -35,7 +35,7 @@ from .mk_objects import (

35 from .parallel_utils import ProcessGroupInfo

36

37

38 -def _describe_tensor(t: Optional[torch.Tensor], name: str) -> str:

38 +def _describe_tensor(t: torch.Tensor | None, name: str) -> str:

39 if t is None:

40 return f"{name} : None"

41 else:

@@ -44,21 +44,21 @@ def _describe_tensor(t: Optional[torch.Tensor], name: str) -> str:

44

45 @dataclass

46 class Config:

47 - Ms: Union[list[int], int]

47 + Ms: list[int] | int

48 K: int

49 N: int

50 E: int

51 - topks: Union[list[int], int]

51 + topks: list[int] | int

52 dtype: torch.dtype

53 - quant_config: Optional[TestMoEQuantConfig]

53 + quant_config: TestMoEQuantConfig | None

54

55 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize

56 fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute

57

58 - fused_moe_chunk_size: Optional[int]

58 + fused_moe_chunk_size: int | None

59 world_size: int

60

61 - torch_trace_dir_path: Optional[str] = None

61 + torch_trace_dir_path: str | None = None

62

63 def __post_init__(self):

64 if self.quant_config is None:

@@ -93,7 +93,7 @@ class Config:

93 return self.Ms

94

95 @property

96 - def quant_dtype(self) -> Union[torch.dtype, str, None]:

96 + def quant_dtype(self) -> torch.dtype | str | None:

97 assert self.quant_config is not None

98 return self.quant_config.quant_dtype

99

@@ -112,7 +112,7 @@ class Config:

112 return self.quant_config.per_out_ch_quant

113

114 @property

115 - def quant_block_shape(self) -> Optional[list[int]]:

115 + def quant_block_shape(self) -> list[int] | None:

116 assert self.quant_config is not None

117 return self.quant_config.block_shape

118

@@ -209,7 +209,7 @@ class Config:

209 info = prepare_finalize_info(self.prepare_finalize_type)

210 return info.backend

211

212 - def is_valid(self) -> tuple[bool, Optional[str]]:

212 + def is_valid(self) -> tuple[bool, str | None]:

213 # Check prepare-finalize and fused-experts compatibility

214 if self.is_batched_prepare_finalize():

215 if not self.is_batched_fused_experts():

@@ -280,10 +280,10 @@ class Config:

280 class WeightTensors:

281 w1: torch.Tensor

282 w2: torch.Tensor

283 - w1_scale: Optional[torch.Tensor]

284 - w2_scale: Optional[torch.Tensor]

285 - w1_gs: Optional[torch.Tensor] = None

286 - w2_gs: Optional[torch.Tensor] = None

283 + w1_scale: torch.Tensor | None

284 + w2_scale: torch.Tensor | None

285 + w1_gs: torch.Tensor | None = None

286 + w2_gs: torch.Tensor | None = None

287

288 def describe(self):

289 s = ""

@@ -351,11 +351,11 @@ class WeightTensors:

351 @dataclass

352 class RankTensors:

353 hidden_states: torch.Tensor

354 - hidden_states_scale: Optional[torch.Tensor]

354 + hidden_states_scale: torch.Tensor | None

355

356 topk_weights: torch.Tensor

357 topk_ids: torch.Tensor

358 - expert_map: Optional[torch.Tensor]

358 + expert_map: torch.Tensor | None

359

360 def describe(self):

361 s = ""

@@ -370,7 +370,7 @@ class RankTensors:

370 @staticmethod

371 def make_hidden_states(

372 config: Config,

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

373 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

374 """

375 Return hidden_states

376 """

tests/kernels/moe/modular_kernel_tools/make_feature_matrix.py

+2 -3

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

4 import copy

5 from enum import Enum

6 from itertools import product

7 -from typing import Optional

7

8 import torch

9 from tqdm import tqdm

@@ -82,7 +81,7 @@ def make_feature_matrix(csv_file_path: str):

81 import pandas as pd

82

83 def add_to_results(

85 - config: Config, success: Result, results_df: Optional[pd.DataFrame] = None

84 + config: Config, success: Result, results_df: pd.DataFrame | None = None

85 ):

86 config_dict = asdict(config)

87 config_dict["prepare_finalize_type"] = config_dict[

@@ -121,7 +120,7 @@ def make_feature_matrix(csv_file_path: str):

120 product(Ms, Ks, Ns, Es, TOPKs, DTYPEs, PF_TYPES, FE_TYPES, Q_TYPES)

121 )

122

124 - results_df: Optional[pd.DataFrame] = None

123 + results_df: pd.DataFrame | None = None

124 for m, k, n, e, topks, dtype, pf_type, experts_type, quant_config in tqdm(

125 combinations

126 ):

tests/kernels/moe/modular_kernel_tools/mk_objects.py

+11 -12

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass

4 -from typing import Optional, Union

4

5 import torch

6

@@ -43,25 +42,25 @@ from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

42

43 @dataclass

44 class TestMoEQuantConfig:

46 - quant_dtype: Union[torch.dtype, str, None]

45 + quant_dtype: torch.dtype | str | None

46 per_out_ch_quant: bool

47 per_act_token_quant: bool

49 - block_shape: Optional[list[int]]

48 + block_shape: list[int] | None

49

50

51 @dataclass

52 class PrepareFinalizeInfo:

53 activation_format: mk.FusedMoEActivationFormat

55 - supported_dtypes: list[Union[torch.dtype, str]]

54 + supported_dtypes: list[torch.dtype | str]

55 blocked_quantization_support: bool

57 - backend: Optional[str]

56 + backend: str | None

57 supports_apply_weight_on_input: bool = True

58

59

60 @dataclass

61 class ExpertInfo:

62 activation_format: mk.FusedMoEActivationFormat

64 - supported_dtypes: list[Union[torch.dtype, str]]

63 + supported_dtypes: list[torch.dtype | str]

64 blocked_quantization_support: bool

65 supports_chunking: bool

66 supports_expert_map: bool

@@ -78,7 +77,7 @@ MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEPermuteExpertsUnpermute] = []

77

78 standard_format = mk.FusedMoEActivationFormat.Standard

79 batched_format = mk.FusedMoEActivationFormat.BatchedExperts

81 -common_float_types: list[Union[torch.dtype, str]] = [

80 +common_float_types: list[torch.dtype | str] = [

81 torch.float8_e4m3fn,

82 torch.bfloat16,

83 torch.float16,

@@ -92,9 +91,9 @@ fp8_types = [torch.float8_e4m3fn]

91 def register_prepare_and_finalize(

92 kind,

93 activation_format: mk.FusedMoEActivationFormat,

95 - supported_dtypes: list[Union[torch.dtype, str]],

94 + supported_dtypes: list[torch.dtype | str],

95 blocked_quantization_support: bool,

97 - backend: Optional[str],

96 + backend: str | None,

97 force_multigpu: bool = False,

98 supports_apply_weight_on_input: bool = True,

99 ):

@@ -121,7 +120,7 @@ def register_prepare_and_finalize(

120 def register_experts(

121 kind,

122 activation_format: mk.FusedMoEActivationFormat,

124 - supported_dtypes: list[Union[torch.dtype, str]],

123 + supported_dtypes: list[torch.dtype | str],

124 blocked_quantization_support: bool,

125 supports_chunking: bool,

126 supports_expert_map: bool,

@@ -340,7 +339,7 @@ if cutlass_fp4_supported():

339 supports_expert_map=False,

340 )

341

343 -MK_QUANT_CONFIGS: list[Optional[TestMoEQuantConfig]] = [

342 +MK_QUANT_CONFIGS: list[TestMoEQuantConfig | None] = [

343 None,

344 # per-channel / per-column weights and per-tensor activations

345 TestMoEQuantConfig(

@@ -395,7 +394,7 @@ if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():

394

395 def make_prepare_finalize(

396 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,

398 - backend: Optional[str],

397 + backend: str | None,

398 moe: FusedMoEConfig,

399 quant_config: FusedMoEQuantConfig,

400 ) -> mk.FusedMoEPrepareAndFinalize:

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

+6 -5

@@ -3,11 +3,12 @@

3 import dataclasses

4 import os

5 import traceback

6 -from typing import Any, Callable, Optional

6 +from collections.abc import Callable

7 +from typing import Any, Concatenate

8

9 import torch

10 from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]

10 -from typing_extensions import Concatenate, ParamSpec

11 +from typing_extensions import ParamSpec

12

13 from vllm.config import VllmConfig, set_current_vllm_config

14 from vllm.distributed import init_distributed_environment, initialize_model_parallel

@@ -58,9 +59,9 @@ def _worker_parallel_launch(

59 world_local_size: int,

60 node_rank: int,

61 init_method: str,

61 - worker: Callable[Concatenate[ProcessGroupInfo, Optional[VllmConfig], Any, P], None],

62 - vllm_config: Optional[VllmConfig],

63 - env_dict: Optional[dict],

62 + worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig | None, Any, P], None],

63 + vllm_config: VllmConfig | None,

64 + env_dict: dict | None,

65 *args: P.args,

66 **kwargs: P.kwargs,

67 ) -> None:

tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py

+2 -1

@@ -2,8 +2,9 @@

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

3

4 import copy

5 +from collections.abc import Callable

6 from itertools import product

6 -from typing import Any, Callable

7 +from typing import Any

8

9 import torch

10

tests/kernels/moe/parallel_utils.py

+11 -10

@@ -7,12 +7,13 @@ DeepEP test utilities

7 import dataclasses

8 import os

9 import traceback

10 -from typing import Callable, Optional

10 +from collections.abc import Callable

11 +from typing import Concatenate

12

13 import torch

14 from torch.distributed import ProcessGroup

15 from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]

15 -from typing_extensions import Concatenate, ParamSpec

16 +from typing_extensions import ParamSpec

17

18 from vllm.utils import get_open_port, has_deep_ep

19

@@ -126,8 +127,8 @@ def make_deepep_ht_a2a(

127 pgi: ProcessGroupInfo,

128 dp_size: int,

129 ht_args: DeepEPHTArgs,

129 - q_dtype: Optional[torch.dtype] = None,

130 - block_shape: Optional[list[int]] = None,

130 + q_dtype: torch.dtype | None = None,

131 + block_shape: list[int] | None = None,

132 ):

133 import deep_ep

134

@@ -153,8 +154,8 @@ def make_deepep_ll_a2a(

154 pg: ProcessGroup,

155 pgi: ProcessGroupInfo,

156 deepep_ll_args: DeepEPLLArgs,

156 - q_dtype: Optional[torch.dtype] = None,

157 - block_shape: Optional[list[int]] = None,

157 + q_dtype: torch.dtype | None = None,

158 + block_shape: list[int] | None = None,

159 ):

160 import deep_ep

161

@@ -185,10 +186,10 @@ def make_deepep_a2a(

186 pg: ProcessGroup,

187 pgi: ProcessGroupInfo,

188 dp_size: int,

188 - deepep_ht_args: Optional[DeepEPHTArgs],

189 - deepep_ll_args: Optional[DeepEPLLArgs],

190 - q_dtype: Optional[torch.dtype] = None,

191 - block_shape: Optional[list[int]] = None,

189 + deepep_ht_args: DeepEPHTArgs | None,

190 + deepep_ll_args: DeepEPLLArgs | None,

191 + q_dtype: torch.dtype | None = None,

192 + block_shape: list[int] | None = None,

193 ):

194 if deepep_ht_args is not None:

195 assert deepep_ll_args is None

tests/kernels/moe/test_batched_moe.py

+3 -4

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

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

3

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import pytest

7 import torch

@@ -55,7 +54,7 @@ vllm_config.scheduler_config.max_model_len = 8192

54 @dataclass

55 class BatchedMMConfig:

56 in_dtype: torch.dtype

58 - quant_dtype: Optional[torch.dtype]

57 + quant_dtype: torch.dtype | None

58 out_dtype: torch.dtype

59 num_experts: int

60 max_tokens_per_expert: int

@@ -115,7 +114,7 @@ def test_batched_mm(

114 K: int,

115 N: int,

116 dtype: torch.dtype,

118 - block_shape: Optional[list[int]],

117 + block_shape: list[int] | None,

118 per_act_token_quant: bool,

119 ):

120 current_platform.seed_everything(7)

@@ -242,7 +241,7 @@ def test_fused_moe_batched_experts(

241 topk: int,

242 dtype: torch.dtype,

243 per_act_token_quant: bool,

245 - block_shape: Optional[list[int]],

244 + block_shape: list[int] | None,

245 input_scales: bool,

246 ):

247 current_platform.seed_everything(7)

tests/kernels/moe/test_count_expert_num_tokens.py

+1 -2

@@ -5,7 +5,6 @@ Tests compute_expert_num_tokens kernels

5 """

6

7 import dataclasses

8 -from typing import Optional

8

9 import pytest

10 import torch

@@ -16,7 +15,7 @@ from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens

15 @dataclasses.dataclass

16 class TestTensors:

17 topk_ids: torch.Tensor

19 - expert_map: Optional[torch.Tensor] = None

18 + expert_map: torch.Tensor | None = None

19

20 def to_device(self, device: str):

21 self.topk_ids = self.topk_ids.to(device=device)

tests/kernels/moe/test_cutlass_moe.py

+11 -12

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

3 import copy

4 import dataclasses

5 from math import prod

6 -from typing import Optional

6

7 import pytest

8 import torch

@@ -85,16 +84,16 @@ class MOETensors:

84 @dataclasses.dataclass

85 class MOETensors8Bit(MOETensors):

86 # quantized

88 - a_q: Optional[torch.Tensor] = None # a -> a_q

89 - w1_q: Optional[torch.Tensor] = None # w1 -> w1_q

90 - w2_q: Optional[torch.Tensor] = None # w2 -> w2_q

91 - a_scale: Optional[torch.Tensor] = None

92 - w1_scale: Optional[torch.Tensor] = None

93 - w2_scale: Optional[torch.Tensor] = None

87 + a_q: torch.Tensor | None = None # a -> a_q

88 + w1_q: torch.Tensor | None = None # w1 -> w1_q

89 + w2_q: torch.Tensor | None = None # w2 -> w2_q

90 + a_scale: torch.Tensor | None = None

91 + w1_scale: torch.Tensor | None = None

92 + w2_scale: torch.Tensor | None = None

93 # dequantized

95 - a_d: Optional[torch.Tensor] = None # a -> a_q -> a_d

96 - w1_d: Optional[torch.Tensor] = None # w1 -> w1_q -> w1_d

97 - w2_d: Optional[torch.Tensor] = None # w2 -> w2_q -> w2_d

94 + a_d: torch.Tensor | None = None # a -> a_q -> a_d

95 + w1_d: torch.Tensor | None = None # w1 -> w1_q -> w1_d

96 + w2_d: torch.Tensor | None = None # w2 -> w2_q -> w2_d

97

98 @staticmethod

99 def make_moe_tensors_8bit(

@@ -209,7 +208,7 @@ def run_8_bit(

208 topk_ids: torch.Tensor,

209 per_act_token: bool,

210 per_out_ch: bool,

212 - num_local_experts: Optional[int] = None,

211 + num_local_experts: int | None = None,

212 ) -> torch.Tensor:

213 assert not any(

214 [

@@ -280,7 +279,7 @@ def test_cutlass_moe_8_bit_no_graph(

279 per_act_token: bool,

280 per_out_ch: bool,

281 monkeypatch,

283 - ep_size: Optional[int] = None,

282 + ep_size: int | None = None,

283 ):

284 current_platform.seed_everything(7)

285 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

tests/kernels/moe/test_deepep_deepgemm_moe.py

+6 -7

@@ -7,7 +7,6 @@ fp8 block-quantized case.

7 """

8

9 import dataclasses

10 -from typing import Optional

10

11 import pytest

12 import torch.distributed

@@ -92,13 +91,13 @@ class TestConfig:

91 block_size: list[int]

92 # configs for testing low-latency kernels

93 low_latency: bool

95 - use_fp8_dispatch: Optional[bool] = False

94 + use_fp8_dispatch: bool | None = False

95

96

97 @dataclasses.dataclass

98 class TestTensors:

99 rank_tokens: torch.Tensor # all ranks make this many tokens

101 - rank_token_scales: Optional[torch.Tensor]

100 + rank_token_scales: torch.Tensor | None

101 topk: torch.Tensor

102 topk_weights: torch.Tensor

103 config: TestConfig

@@ -143,7 +142,7 @@ def make_ll_modular_kernel(

142 max_tokens_per_rank: int,

143 dp_size: int,

144 hidden_size: int,

146 - q_dtype: Optional[torch.dtype],

145 + q_dtype: torch.dtype | None,

146 test_config: TestConfig,

147 quant_config: FusedMoEQuantConfig,

148 ) -> FusedMoEModularKernel:

@@ -179,7 +178,7 @@ def make_ht_modular_kernel(

178 pgi: ProcessGroupInfo,

179 dp_size: int,

180 num_local_experts: int,

182 - q_dtype: Optional[torch.dtype],

181 + q_dtype: torch.dtype | None,

182 test_config: TestConfig,

183 quant_config: FusedMoEQuantConfig,

184 ) -> FusedMoEModularKernel:

@@ -249,8 +248,8 @@ def deepep_deepgemm_moe_impl(

248 test_tensors: TestTensors,

249 w1: torch.Tensor,

250 w2: torch.Tensor,

252 - w1_scale: Optional[torch.Tensor],

253 - w2_scale: Optional[torch.Tensor],

251 + w1_scale: torch.Tensor | None,

252 + w2_scale: torch.Tensor | None,

253 ) -> torch.Tensor:

254 test_config = test_tensors.config

255 num_experts = test_config.num_experts

tests/kernels/moe/test_deepep_moe.py

+18 -21

@@ -5,7 +5,6 @@ Test deepep dispatch-combine logic

5 """

6

7 import dataclasses

8 -from typing import Optional, Union

8

9 import pytest

10 import torch.distributed

@@ -90,7 +89,7 @@ class TestConfig:

89 @dataclasses.dataclass

90 class TestTensors:

91 rank_tokens: torch.Tensor # all ranks make this many tokens

93 - rank_token_scales: Optional[torch.Tensor]

92 + rank_token_scales: torch.Tensor | None

93 topk: torch.Tensor

94 topk_weights: torch.Tensor

95 config: TestConfig

@@ -128,12 +127,12 @@ def make_modular_kernel(

127 dp_size: int,

128 num_experts: int,

129 num_local_experts: int,

131 - q_dtype: Optional[torch.dtype],

130 + q_dtype: torch.dtype | None,

131 use_fp8_dispatch: bool,

132 quant_config: FusedMoEQuantConfig,

133 ) -> FusedMoEModularKernel:

135 - ht_args: Optional[DeepEPHTArgs] = None

136 - ll_args: Optional[DeepEPLLArgs] = None

134 + ht_args: DeepEPHTArgs | None = None

135 + ll_args: DeepEPLLArgs | None = None

136

137 if low_latency_mode:

138 ll_args = DeepEPLLArgs(

@@ -148,16 +147,14 @@ def make_modular_kernel(

147 )

148 ht_args = DeepEPHTArgs(num_local_experts=num_local_experts)

149

151 - a2a: Union[DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize] = (

152 - make_deepep_a2a(

153 - pg=pg,

154 - pgi=pgi,

155 - dp_size=dp_size,

156 - q_dtype=q_dtype,

157 - block_shape=None,

158 - deepep_ht_args=ht_args,

159 - deepep_ll_args=ll_args,

160 - )

150 + a2a: DeepEPHTPrepareAndFinalize | DeepEPLLPrepareAndFinalize = make_deepep_a2a(

151 + pg=pg,

152 + pgi=pgi,

153 + dp_size=dp_size,

154 + q_dtype=q_dtype,

155 + block_shape=None,

156 + deepep_ht_args=ht_args,

157 + deepep_ll_args=ll_args,

158 )

159

160 num_dispatchers = pgi.world_size // dp_size

@@ -184,8 +181,8 @@ def deep_ep_moe_impl(

181 test_tensors: TestTensors,

182 w1: torch.Tensor,

183 w2: torch.Tensor,

187 - w1_scale: Optional[torch.Tensor],

188 - w2_scale: Optional[torch.Tensor],

184 + w1_scale: torch.Tensor | None,

185 + w2_scale: torch.Tensor | None,

186 num_experts: int,

187 use_fp8_dispatch: bool,

188 per_act_token_quant: bool,

@@ -281,8 +278,8 @@ def torch_moe_impl(

278 test_tensors: TestTensors,

279 w1: torch.Tensor,

280 w2: torch.Tensor,

284 - w1_scale: Optional[torch.Tensor],

285 - w2_scale: Optional[torch.Tensor],

281 + w1_scale: torch.Tensor | None,

282 + w2_scale: torch.Tensor | None,

283 using_fp8_dispatch: bool,

284 per_act_token_quant: bool,

285 ):

@@ -340,8 +337,8 @@ def _deep_ep_moe(

337 config: TestConfig,

338 w1: torch.Tensor,

339 w2: torch.Tensor,

343 - w1_scale: Optional[torch.Tensor],

344 - w2_scale: Optional[torch.Tensor],

340 + w1_scale: torch.Tensor | None,

341 + w2_scale: torch.Tensor | None,

342 use_fp8_dispatch: bool,

343 per_act_token_quant: bool,

344 ):

tests/kernels/moe/test_modular_kernel_combinations.py

+5 -5

@@ -5,7 +5,7 @@ import copy

5 import textwrap

6 import traceback

7 from itertools import product

8 -from typing import Any, Optional

8 +from typing import Any

9

10 import pytest

11 import torch

@@ -245,10 +245,10 @@ def test_modular_kernel_combinations_multigpu(

245 n: int,

246 e: int,

247 dtype: torch.dtype,

248 - quant_config: Optional[TestMoEQuantConfig],

248 + quant_config: TestMoEQuantConfig | None,

249 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,

250 fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,

251 - chunk_size: Optional[int],

251 + chunk_size: int | None,

252 world_size: int,

253 pytestconfig,

254 ):

@@ -287,10 +287,10 @@ def test_modular_kernel_combinations_singlegpu(

287 n: int,

288 e: int,

289 dtype: torch.dtype,

290 - quant_config: Optional[TestMoEQuantConfig],

290 + quant_config: TestMoEQuantConfig | None,

291 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,

292 fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,

293 - chunk_size: Optional[int],

293 + chunk_size: int | None,

294 world_size: int,

295 pytestconfig,

296 ):

tests/kernels/moe/test_moe.py

+4 -4

@@ -6,7 +6,7 @@ Run `pytest tests/kernels/test_moe.py`.

6 """

7

8 import functools

9 -from typing import Callable, Optional, Union

9 +from collections.abc import Callable

10

11 import pytest

12 import torch

@@ -80,7 +80,7 @@ vllm_config.scheduler_config.max_model_len = 8192

80

81

82 def run_moe_test(

83 - baseline: Union[Callable, torch.Tensor],

83 + baseline: Callable | torch.Tensor,

84 moe_fn: Callable,

85 a: torch.Tensor,

86 w1: torch.Tensor,

@@ -88,7 +88,7 @@ def run_moe_test(

88 score: torch.Tensor,

89 topk: int,

90 global_num_experts: int = -1,

91 - expert_map: Optional[torch.Tensor] = None,

91 + expert_map: torch.Tensor | None = None,

92 padding: bool = False,

93 use_compile: bool = False,

94 use_cudagraph: bool = False,

@@ -212,7 +212,7 @@ def test_fused_moe(

212 score: torch.Tensor,

213 topk: int,

214 global_num_experts: int = -1,

215 - expert_map: Optional[torch.Tensor] = None,

215 + expert_map: torch.Tensor | None = None,

216 ) -> torch.Tensor:

217 topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)

218 return m_fused_moe_fn(

tests/kernels/moe/test_moe_align_block_size.py

+1 -3

@@ -5,8 +5,6 @@

5 Run `pytest tests/kernels/moe/test_moe_align_block_size.py`.

6 """

7

8 -from typing import Optional

9 -

8 import pytest

9 import torch

10

@@ -94,7 +92,7 @@ def torch_moe_align_block_size(

92 topk_ids: torch.Tensor,

93 block_size: int,

94 num_experts: int,

97 - expert_map: Optional[torch.Tensor] = None,

95 + expert_map: torch.Tensor | None = None,

96 pad_sorted_ids: bool = False,

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

98 """

tests/kernels/moe/test_moe_permute_unpermute.py

+3 -5

@@ -5,8 +5,6 @@

5 Run `pytest tests/kernels/test_moe_permute_unpermute.py`.

6 """

7

8 -from typing import Optional

9 -

8 import numpy as np

9 import pytest

10 import torch

@@ -34,8 +32,8 @@ def torch_permute(

32 n_expert: int,

33 n_local_expert: int,

34 start_expert: int,

37 - expert_map: Optional[torch.Tensor] = None,

38 - align_block_size: Optional[int] = None,

35 + expert_map: torch.Tensor | None = None,

36 + align_block_size: int | None = None,

37 fill_invalid_expert: int = -1,

38 ) -> list[torch.Tensor]:

39 n_token, n_hidden = hidden_states.shape[0], hidden_states.shape[1]

@@ -210,7 +208,7 @@ def test_moe_permute_unpermute(

208 n_expert: int,

209 ep_size: int,

210 dtype: torch.dtype,

213 - align_block_size: Optional[int],

211 + align_block_size: int | None,

212 ):

213 if not moe_permute_unpermute_supported():

214 pytest.skip("moe_permute_unpermute is not supported on this platform.")

tests/kernels/moe/test_ocp_mx_moe.py

+6 -7

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

4 import importlib.metadata

5 from dataclasses import dataclass

6 from importlib.util import find_spec

7 -from typing import Optional

7

8 import pytest

9 import torch

@@ -103,7 +102,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):

102 assert output

103

104

106 -def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: Optional[float] = None):

105 +def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: float | None = None):

106 # Note we add an extra bias of 1 to the linear layer

107 x_glu, x_linear = torch.chunk(x, 2, dim=-1)

108 if limit is not None:

@@ -510,7 +509,7 @@ def test_trtllm_gen_mxfp4_fused_moe(

509 hidden_size: int,

510 alpha: float,

511 beta: float,

513 - limit: Optional[float],

512 + limit: float | None,

513 act_type: str,

514 transpose_optimized: bool,

515 ):

@@ -660,7 +659,7 @@ def test_flashinfer_cutlass_mxfp4_fused_moe(

659 hidden_size: int,

660 alpha: float,

661 beta: float,

663 - limit: Optional[float],

662 + limit: float | None,

663 ):

664 torch.manual_seed(42)

665 device = "cuda:0"

@@ -811,9 +810,9 @@ def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe(

810 num_tokens: int,

811 intermediate_size: int,

812 hidden_size: int,

814 - alpha: Optional[float],

815 - beta: Optional[float],

816 - limit: Optional[float],

813 + alpha: float | None,

814 + beta: float | None,

815 + limit: float | None,

816 ):

817 torch.manual_seed(42)

818 device = "cuda:0"

tests/kernels/moe/test_pplx_cutlass_moe.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6 import torch

@@ -73,7 +72,7 @@ def pplx_cutlass_moe(

72 out_dtype,

73 per_act_token: bool,

74 per_out_ch: bool,

76 - group_name: Optional[str],

75 + group_name: str | None,

76 ):

77 from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (

78 PplxPrepareAndFinalize,

tests/kernels/moe/test_pplx_moe.py

+29 -35

@@ -9,7 +9,7 @@ import copy

9 import itertools

10 import textwrap

11 import traceback

12 -from typing import Callable, Optional, Union

12 +from collections.abc import Callable

13

14 import pytest

15 import torch

@@ -89,7 +89,7 @@ def torch_prepare(

89 a: torch.Tensor,

90 topk_ids: torch.Tensor,

91 num_experts: int,

92 - max_num_tokens: Optional[int] = None,

92 + max_num_tokens: int | None = None,

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

94 assert topk_ids.dim() == 2

95 assert topk_ids.shape[0] == a.shape[0]

@@ -214,10 +214,10 @@ def create_pplx_prepare_finalize(

214 dp_size: int,

215 world_size: int,

216 in_dtype: torch.dtype,

217 - quant_dtype: Optional[torch.dtype],

218 - block_shape: Optional[list[int]],

217 + quant_dtype: torch.dtype | None,

218 + block_shape: list[int] | None,

219 per_act_token_quant: bool,

220 - group_name: Optional[str],

220 + group_name: str | None,

221 ):

222 from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (

223 PplxPrepareAndFinalize,

@@ -274,18 +274,14 @@ def chunk_by_rank(t: torch.Tensor, r: int, w: int) -> torch.Tensor:

274 return t[(r * chunk) : (r + 1) * chunk]

275

276

277 -def maybe_chunk_by_rank(

278 - t: Optional[torch.Tensor], r: int, w: int

279 -) -> Optional[torch.Tensor]:

277 +def maybe_chunk_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None:

278 if t is not None:

279 return chunk_by_rank(t, r, w)

280 else:

281 return t

282

283

286 -def chunk_scales_by_rank(

287 - t: Optional[torch.Tensor], r: int, w: int

288 -) -> Optional[torch.Tensor]:

284 +def chunk_scales_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None:

285 if t is not None and t.numel() > 1:

286 chunk = rank_chunk(t.shape[0], r, w)

287 return t[(r * chunk) : (r + 1) * chunk]

@@ -293,9 +289,7 @@ def chunk_scales_by_rank(

289 return t

290

291

296 -def chunk_scales(

297 - t: Optional[torch.Tensor], start: int, end: int

298 -) -> Optional[torch.Tensor]:

292 +def chunk_scales(t: torch.Tensor | None, start: int, end: int) -> torch.Tensor | None:

293 if t is not None and t.numel() > 1:

294 return t[start:end]

295 else:

@@ -313,10 +307,10 @@ def pplx_prepare_finalize(

307 topk_weight: torch.Tensor,

308 topk_ids: torch.Tensor,

309 num_experts: int,

316 - quant_dtype: Optional[torch.dtype],

317 - block_shape: Optional[list[int]],

310 + quant_dtype: torch.dtype | None,

311 + block_shape: list[int] | None,

312 per_act_token_quant: bool,

319 - group_name: Optional[str],

313 + group_name: str | None,

314 ) -> torch.Tensor:

315 assert torch.cuda.current_device() == pgi.local_rank

316

@@ -409,8 +403,8 @@ def _pplx_prepare_finalize(

403 score: torch.Tensor,

404 topk: torch.Tensor,

405 num_experts: int,

412 - quant_dtype: Optional[torch.dtype],

413 - block_shape: Optional[list[int]],

406 + quant_dtype: torch.dtype | None,

407 + block_shape: list[int] | None,

408 per_act_token_quant: bool,

409 use_internode: bool,

410 ):

@@ -479,7 +473,7 @@ def test_pplx_prepare_finalize_slow(

473 dtype: torch.dtype,

474 world_dp_size: tuple[int, int],

475 per_act_token_quant: bool,

482 - block_shape: Optional[list[int]],

476 + block_shape: list[int] | None,

477 use_internode: bool,

478 ):

479 if dtype == torch.float8_e4m3fn:

@@ -521,7 +515,7 @@ def test_pplx_prepare_finalize_slow(

515

516

517 def pplx_moe(

524 - group_name: Optional[str],

518 + group_name: str | None,

519 rank: int,

520 world_size: int,

521 dp_size: int,

@@ -530,17 +524,17 @@ def pplx_moe(

524 w2: torch.Tensor,

525 topk_weight: torch.Tensor,

526 topk_ids: torch.Tensor,

533 - w1_scale: Optional[torch.Tensor] = None,

534 - w2_scale: Optional[torch.Tensor] = None,

535 - a1_scale: Optional[torch.Tensor] = None,

536 - a2_scale: Optional[torch.Tensor] = None,

537 - quant_dtype: Optional[torch.dtype] = None,

527 + w1_scale: torch.Tensor | None = None,

528 + w2_scale: torch.Tensor | None = None,

529 + a1_scale: torch.Tensor | None = None,

530 + a2_scale: torch.Tensor | None = None,

531 + quant_dtype: torch.dtype | None = None,

532 per_act_token_quant=False,

539 - block_shape: Optional[list[int]] = None,

533 + block_shape: list[int] | None = None,

534 use_compile: bool = False,

535 use_cudagraphs: bool = True,

542 - shared_experts: Optional[torch.nn.Module] = None,

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

536 + shared_experts: torch.nn.Module | None = None,

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

538 num_tokens, hidden_dim = a.shape

539 num_experts = w1.shape[0]

540 topk = topk_ids.shape[1]

@@ -657,13 +651,13 @@ def _pplx_moe(

651 score: torch.Tensor,

652 topk: int,

653 num_experts: int,

660 - w1_s: Optional[torch.Tensor] = None,

661 - w2_s: Optional[torch.Tensor] = None,

662 - quant_dtype: Optional[torch.dtype] = None,

654 + w1_s: torch.Tensor | None = None,

655 + w2_s: torch.Tensor | None = None,

656 + quant_dtype: torch.dtype | None = None,

657 per_act_token_quant: bool = False,

664 - block_shape: Optional[list[int]] = None,

658 + block_shape: list[int] | None = None,

659 use_internode: bool = False,

666 - shared_experts: Optional[torch.nn.Module] = None,

660 + shared_experts: torch.nn.Module | None = None,

661 ):

662 try:

663 if use_internode:

@@ -812,7 +806,7 @@ def test_pplx_moe_slow(

806 dtype: torch.dtype,

807 world_dp_size: tuple[int, int],

808 per_act_token_quant: bool,

815 - block_shape: Optional[list[int]],

809 + block_shape: list[int] | None,

810 use_internode: bool,

811 ):

812 current_platform.seed_everything(7)

tests/kernels/moe/utils.py

+44 -45

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional, Union

3

4 import torch

5

@@ -27,13 +26,13 @@ def triton_moe(

26 w2: torch.Tensor,

27 topk_weight: torch.Tensor,

28 topk_ids: torch.Tensor,

30 - w1_scale: Optional[torch.Tensor] = None,

31 - w2_scale: Optional[torch.Tensor] = None,

32 - a1_scale: Optional[torch.Tensor] = None,

33 - a2_scale: Optional[torch.Tensor] = None,

34 - quant_dtype: Optional[torch.dtype] = None,

29 + w1_scale: torch.Tensor | None = None,

30 + w2_scale: torch.Tensor | None = None,

31 + a1_scale: torch.Tensor | None = None,

32 + a2_scale: torch.Tensor | None = None,

33 + quant_dtype: torch.dtype | None = None,

34 per_act_token_quant=False,

36 - block_shape: Optional[list[int]] = None,

35 + block_shape: list[int] | None = None,

36 ) -> torch.Tensor:

37 quant_config = FusedMoEQuantConfig.make(

38 quant_dtype,

@@ -54,13 +53,13 @@ def batched_moe(

53 w2: torch.Tensor,

54 topk_weight: torch.Tensor,

55 topk_ids: torch.Tensor,

57 - w1_scale: Optional[torch.Tensor] = None,

58 - w2_scale: Optional[torch.Tensor] = None,

59 - a1_scale: Optional[torch.Tensor] = None,

60 - a2_scale: Optional[torch.Tensor] = None,

61 - quant_dtype: Optional[torch.dtype] = None,

56 + w1_scale: torch.Tensor | None = None,

57 + w2_scale: torch.Tensor | None = None,

58 + a1_scale: torch.Tensor | None = None,

59 + a2_scale: torch.Tensor | None = None,

60 + quant_dtype: torch.dtype | None = None,

61 per_act_token_quant: bool = False,

63 - block_shape: Optional[list[int]] = None,

62 + block_shape: list[int] | None = None,

63 ) -> torch.Tensor:

64 max_num_tokens = round_up(a.shape[0], 64)

65

@@ -94,13 +93,13 @@ def naive_batched_moe(

93 w2: torch.Tensor,

94 topk_weight: torch.Tensor,

95 topk_ids: torch.Tensor,

97 - w1_scale: Optional[torch.Tensor] = None,

98 - w2_scale: Optional[torch.Tensor] = None,

99 - a1_scale: Optional[torch.Tensor] = None,

100 - a2_scale: Optional[torch.Tensor] = None,

101 - quant_dtype: Optional[torch.dtype] = None,

96 + w1_scale: torch.Tensor | None = None,

97 + w2_scale: torch.Tensor | None = None,

98 + a1_scale: torch.Tensor | None = None,

99 + a2_scale: torch.Tensor | None = None,

100 + quant_dtype: torch.dtype | None = None,

101 per_act_token_quant: bool = False,

103 - block_shape: Optional[list[int]] = None,

102 + block_shape: list[int] | None = None,

103 ) -> torch.Tensor:

104 max_num_tokens = round_up(a.shape[0], 64)

105

@@ -129,8 +128,8 @@ def naive_batched_moe(

128

129

130 def chunk_scales(

132 - scales: Optional[torch.Tensor], start: int, end: int

133 -) -> Optional[torch.Tensor]:

131 + scales: torch.Tensor | None, start: int, end: int

132 +) -> torch.Tensor | None:

133 if scales is not None:

134 if scales.numel() == 1:

135 return scales

@@ -144,10 +143,10 @@ def make_quantized_test_activations(

143 m: int,

144 k: int,

145 in_dtype: torch.dtype,

147 - quant_dtype: Optional[torch.dtype] = None,

148 - block_shape: Optional[list[int]] = None,

146 + quant_dtype: torch.dtype | None = None,

147 + block_shape: list[int] | None = None,

148 per_act_token_quant: bool = False,

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

149 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

150 a = torch.randn((E, m, k), device="cuda", dtype=in_dtype) / 10

151 a_q = a

152 a_scale = None

@@ -172,11 +171,11 @@ def make_quantized_test_activations(

171

172 def moe_quantize_weights(

173 w: torch.Tensor,

175 - w_s: Optional[torch.Tensor],

176 - quant_dtype: Union[torch.dtype, str, None],

174 + w_s: torch.Tensor | None,

175 + quant_dtype: torch.dtype | str | None,

176 per_token_quant: bool,

178 - block_shape: Optional[list[int]],

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

177 + block_shape: list[int] | None,

178 +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:

179 assert (

180 quant_dtype == torch.float8_e4m3fn

181 or quant_dtype == torch.int8

@@ -220,10 +219,10 @@ def make_test_weight(

219 rows: int,

220 cols: int,

221 in_dtype: torch.dtype = torch.bfloat16,

223 - quant_dtype: Union[torch.dtype, str, None] = None,

224 - block_shape: Optional[list[int]] = None,

222 + quant_dtype: torch.dtype | str | None = None,

223 + block_shape: list[int] | None = None,

224 per_out_ch_quant: bool = False,

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

225 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:

226 w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15

227 w_gs = None

228

@@ -262,12 +261,12 @@ def make_test_weights(

261 n: int,

262 k: int,

263 in_dtype: torch.dtype = torch.bfloat16,

265 - quant_dtype: Union[torch.dtype, str, None] = None,

266 - block_shape: Optional[list[int]] = None,

264 + quant_dtype: torch.dtype | str | None = None,

265 + block_shape: list[int] | None = None,

266 per_out_ch_quant: bool = False,

267 ) -> tuple[

269 - tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]],

270 - tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]],

268 + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None],

269 + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None],

270 ]:

271 return (

272 make_test_weight(

@@ -295,9 +294,9 @@ def make_test_quant_config(

294 n: int,

295 k: int,

296 in_dtype: torch.dtype,

298 - quant_dtype: Union[torch.dtype, str, None] = None,

297 + quant_dtype: torch.dtype | str | None = None,

298 per_act_token_quant: bool = False,

300 - block_shape: Optional[list[int]] = None,

299 + block_shape: list[int] | None = None,

300 ) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]:

301 (_, w1, w1_s, w1_gs), (_, w2, w2_s, w2_gs) = make_test_weights(

302 e,

@@ -310,8 +309,8 @@ def make_test_quant_config(

309 )

310

311 # Hacky/trivial scales for nvfp4.

313 - a1_gscale: Optional[torch.Tensor] = None

314 - a2_gscale: Optional[torch.Tensor] = None

312 + a1_gscale: torch.Tensor | None = None

313 + a2_gscale: torch.Tensor | None = None

314 if quant_dtype == "nvfp4":

315 a1_gscale = torch.ones((e,), device="cuda", dtype=torch.float32)

316 a2_gscale = torch.ones((e,), device="cuda", dtype=torch.float32)

@@ -348,9 +347,9 @@ def fused_moe(

347 score: torch.Tensor,

348 topk: int,

349 renormalize: bool = False,

351 - quant_config: Optional[FusedMoEQuantConfig] = None,

350 + quant_config: FusedMoEQuantConfig | None = None,

351 global_num_experts: int = -1,

353 - expert_map: Optional[torch.Tensor] = None,

352 + expert_map: torch.Tensor | None = None,

353 ) -> torch.Tensor:

354 topk_weights, topk_ids, _ = fused_topk(

355 hidden_states, score.float(), topk, renormalize

@@ -378,7 +377,7 @@ class BaselineMM(torch.nn.Module):

377 self.b = b.to(dtype=torch.float32)

378 self.out_dtype = out_dtype

379

381 - def forward(self, a: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

380 + def forward(self, a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:

381 return torch.mm(a.to(dtype=torch.float32), self.b).to(self.out_dtype), None

382

383

@@ -422,8 +421,8 @@ class RealMLP(torch.nn.Module):

421 quant_config=None,

422 reduce_results: bool = True,

423 prefix: str = "",

425 - w1_s: Optional[torch.Tensor] = None,

426 - w2_s: Optional[torch.Tensor] = None,

424 + w1_s: torch.Tensor | None = None,

425 + w2_s: torch.Tensor | None = None,

426 ) -> None:

427 from vllm.model_executor.layers.linear import (

428 MergedColumnParallelLinear,

@@ -481,7 +480,7 @@ def make_shared_experts(

480 N: int,

481 K: int,

482 in_dtype: torch.dtype = torch.bfloat16,

484 - quant_dtype: Union[torch.dtype, str, None] = None,

483 + quant_dtype: torch.dtype | str | None = None,

484 ) -> torch.nn.Module:

485 from vllm.model_executor.layers.quantization.fp8 import Fp8Config

486

tests/kernels/quant_utils.py

+14 -15

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 import torch

6

@@ -15,13 +14,13 @@ ROCM_FP8FNUZ_MAX = 224.0

14 FP8_DTYPE = current_platform.fp8_dtype()

15

16

18 -def as_float32_tensor(x: Union[float, torch.tensor]) -> torch.tensor:

17 +def as_float32_tensor(x: float | torch.Tensor) -> torch.Tensor:

18 return torch.as_tensor(x, dtype=torch.float32, device="cuda")

19

20

21 def ref_dynamic_per_token_quant(

23 - x: torch.tensor, quant_dtype: torch.dtype, scale_ub: Optional[torch.tensor] = None

24 -) -> tuple[torch.tensor, torch.tensor]:

22 + x: torch.Tensor, quant_dtype: torch.dtype, scale_ub: torch.Tensor | None = None

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

24 assert quant_dtype in [torch.int8, FP8_DTYPE]

25 if scale_ub is not None:

26 assert quant_dtype == FP8_DTYPE

@@ -76,8 +75,8 @@ def ref_dynamic_per_token_quant(

75 # ref_dynamic_per_token_quant, when we have a dynamic_per_tensor int8 quant

76 # kernel

77 def ref_dynamic_per_tensor_fp8_quant(

79 - x: torch.tensor,

80 -) -> tuple[torch.tensor, torch.tensor]:

78 + x: torch.Tensor,

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

80 fp8_traits = torch.finfo(FP8_DTYPE)

81 fp8_traits_max = (

82 ROCM_FP8FNUZ_MAX

@@ -250,10 +249,10 @@ def per_block_cast_to_int8(

249

250 def dequant(

251 t: torch.Tensor,

253 - scale: Optional[torch.Tensor],

254 - block_shape: Optional[list[int]],

252 + scale: torch.Tensor | None,

253 + block_shape: list[int] | None,

254 per_act_token_quant: bool,

256 - out_dtype: Optional[torch.dtype] = torch.float32,

255 + out_dtype: torch.dtype | None = torch.float32,

256 ) -> torch.Tensor:

257 if scale is not None:

258 f32 = torch.float32

@@ -267,10 +266,10 @@ def dequant(

266

267 def batched_dequant(

268 t: torch.Tensor,

270 - scale: Optional[torch.Tensor],

271 - block_shape: Optional[list[int]],

269 + scale: torch.Tensor | None,

270 + block_shape: list[int] | None,

271 per_act_token_quant: bool,

273 - out_dtype: Optional[torch.dtype] = torch.float32,

272 + out_dtype: torch.dtype | None = torch.float32,

273 ) -> torch.Tensor:

274 if scale is not None:

275 assert t.shape[0] == scale.shape[0]

@@ -289,9 +288,9 @@ def native_batched_masked_quant_matmul(

288 B: torch.Tensor,

289 C: torch.Tensor,

290 num_expert_tokens: torch.Tensor,

292 - A_scale: Optional[torch.Tensor] = None,

293 - B_scale: Optional[torch.Tensor] = None,

294 - block_shape: Optional[list[int]] = None,

291 + A_scale: torch.Tensor | None = None,

292 + B_scale: torch.Tensor | None = None,

293 + block_shape: list[int] | None = None,

294 per_act_token_quant: bool = False,

295 ) -> torch.Tensor:

296 num_expert_tokens_cpu = num_expert_tokens.clone()

tests/kernels/quantization/test_cutlass_w4a8.py

+10 -11

@@ -6,7 +6,6 @@ Run `pytest tests/kernels/quantization/test_cutlass_w4a8.py`.

6 """

7

8 from dataclasses import dataclass

9 -from typing import Optional

9

10 import pytest

11 import torch

@@ -60,10 +59,10 @@ SCHEDULES = [

59 class TypeConfig:

60 act_type: torch.dtype

61 weight_type: ScalarType

63 - output_type: Optional[torch.dtype]

64 - group_scale_type: Optional[torch.dtype]

65 - channel_scale_type: Optional[torch.dtype]

66 - token_scale_type: Optional[torch.dtype]

62 + output_type: torch.dtype | None

63 + group_scale_type: torch.dtype | None

64 + channel_scale_type: torch.dtype | None

65 + token_scale_type: torch.dtype | None

66

67

68 @dataclass

@@ -80,7 +79,7 @@ class Tensors:

79 # (Act Type, Weight Type, Output Type, Scale Type, ZeroPoints,

80 # Ch Scales Type, Tok Scales Type)

81 TestTypeTuple = tuple[

83 - list[torch.dtype], ScalarType, Optional[torch.dtype], Optional[torch.dtype], bool

82 + list[torch.dtype], ScalarType, torch.dtype | None, torch.dtype | None, bool

83 ]

84 TEST_TYPES = [

85 *(

@@ -116,8 +115,8 @@ def cutlass_quantize_and_pack(

115 atype: torch.dtype,

116 w: torch.Tensor,

117 wtype: ScalarType,

119 - stype: Optional[torch.dtype],

120 - group_size: Optional[int],

118 + stype: torch.dtype | None,

119 + group_size: int | None,

120 zero_points: bool = False,

121 ):

122 assert wtype.is_integer(), "TODO: support floating point weights"

@@ -143,7 +142,7 @@ def cutlass_quantize_and_pack(

142

143

144 def create_test_tensors(

146 - shape: tuple[int, int, int], types: TypeConfig, group_size: Optional[int]

145 + shape: tuple[int, int, int], types: TypeConfig, group_size: int | None

146 ) -> Tensors:

147 m, n, k = shape

148

@@ -185,8 +184,8 @@ def create_test_tensors(

184 def mm_test_helper(

185 types: TypeConfig,

186 tensors: Tensors,

188 - group_size: Optional[int] = None,

189 - schedule: Optional[str] = None,

187 + group_size: int | None = None,

188 + schedule: str | None = None,

189 ):

190 # CUTLASS upstream uses fp8 with fastaccum as reference

191 # https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu#L406

tests/kernels/quantization/test_machete_mm.py

+20 -21

@@ -7,7 +7,6 @@ Run `pytest tests/kernels/quantization/test_machete_mm.py`.

7

8 import math

9 from dataclasses import dataclass, fields

10 -from typing import Optional

10

11 import pytest

12 import torch

@@ -50,11 +49,11 @@ MNK_SHAPES = [

49 class TypeConfig:

50 act_type: torch.dtype

51 weight_type: ScalarType

53 - output_type: Optional[torch.dtype]

54 - group_scale_type: Optional[torch.dtype]

55 - group_zero_type: Optional[torch.dtype]

56 - channel_scale_type: Optional[torch.dtype]

57 - token_scale_type: Optional[torch.dtype]

52 + output_type: torch.dtype | None

53 + group_scale_type: torch.dtype | None

54 + group_zero_type: torch.dtype | None

55 + channel_scale_type: torch.dtype | None

56 + token_scale_type: torch.dtype | None

57

58

59 @dataclass

@@ -63,10 +62,10 @@ class Tensors:

62 a_ref: torch.Tensor

63 a: torch.Tensor

64 w_q: torch.Tensor

66 - w_g_s: Optional[torch.Tensor]

67 - w_g_zp: Optional[torch.Tensor]

68 - w_ch_s: Optional[torch.Tensor]

69 - w_tok_s: Optional[torch.Tensor]

65 + w_g_s: torch.Tensor | None

66 + w_g_zp: torch.Tensor | None

67 + w_ch_s: torch.Tensor | None

68 + w_tok_s: torch.Tensor | None

69

70

71 # (Act Type, Weight Type, Output Type, Scale Type, ZeroPoints,

@@ -74,7 +73,7 @@ class Tensors:

73 # NOTE: None "Scale Type" means the act type is floating point

74 # None "Output Type" means the output type is the same as the act type

75 TestTypeTuple = tuple[

77 - list[torch.dtype], ScalarType, Optional[torch.dtype], Optional[torch.dtype], bool

76 + list[torch.dtype], ScalarType, torch.dtype | None, torch.dtype | None, bool

77 ]

78 TEST_TYPES = [

79 # GPTQ style

@@ -139,11 +138,11 @@ def rand_data(shape, dtype=torch.float16, scale=1, offset=0):

138 return torch.randint(-8, 7, shape, dtype=dtype, device="cuda")

139

140

142 -def maybe_convert_zeropoints(zps: Optional[torch.Tensor], s: torch.Tensor):

141 +def maybe_convert_zeropoints(zps: torch.Tensor | None, s: torch.Tensor):

142 return zps if zps is None else -1 * s * (zps.to(s.dtype))

143

144

146 -def group_size_valid(shape: tuple[int, int, int], group_size: Optional[int]) -> bool:

145 +def group_size_valid(shape: tuple[int, int, int], group_size: int | None) -> bool:

146 return group_size is None or group_size == -1 or shape[2] % group_size == 0

147

148

@@ -151,8 +150,8 @@ def machete_quantize_and_pack(

150 atype: torch.dtype,

151 w: torch.Tensor,

152 wtype: ScalarType,

154 - stype: Optional[torch.dtype],

155 - group_size: Optional[int],

153 + stype: torch.dtype | None,

154 + group_size: int | None,

155 zero_points: bool = False,

156 ):

157 assert wtype.is_integer(), "TODO: support floating point weights"

@@ -178,8 +177,8 @@ def machete_quantize_and_pack(

177 def create_test_tensors(

178 shape: tuple[int, int, int],

179 types: TypeConfig,

181 - group_size: Optional[int],

182 - subset_stride_factor: Optional[int] = None,

180 + group_size: int | None,

181 + subset_stride_factor: int | None = None,

182 ) -> Tensors:

183 m, n, k = shape

184 factor = subset_stride_factor or 1

@@ -243,8 +242,8 @@ def create_test_tensors(

242 def machete_mm_test_helper(

243 types: TypeConfig,

244 tensors: Tensors,

246 - group_size: Optional[int] = None,

247 - schedule: Optional[str] = None,

245 + group_size: int | None = None,

246 + schedule: str | None = None,

247 ):

248 output_ref = torch.matmul(tensors.a_ref, tensors.w_ref)

249 output_ref_type = output_ref.dtype

@@ -294,7 +293,7 @@ def machete_mm_test_helper(

293 @pytest.mark.parametrize("shape", MNK_SHAPES, ids=lambda x: "x".join(str(v) for v in x))

294 @pytest.mark.parametrize("types", TEST_TYPES)

295 def test_machete_all_schedules(shape, types: TypeConfig):

297 - group_sizes: list[Optional[int]] = []

296 + group_sizes: list[int | None] = []

297 if types.group_scale_type is None:

298 group_sizes = [None]

299 else:

@@ -323,7 +322,7 @@ def test_machete_all_schedules(shape, types: TypeConfig):

322 @pytest.mark.parametrize("shape", MNK_SHAPES, ids=lambda x: "x".join(str(v) for v in x))

323 @pytest.mark.parametrize("types", TEST_TYPES)

324 def test_machete_heuristic(shape, types: TypeConfig):

326 - group_sizes: list[Optional[int]] = []

325 + group_sizes: list[int | None] = []

326 if types.group_scale_type is None:

327 group_sizes = [None]

328 else:

tests/kernels/quantization/test_triton_scaled_mm.py

+1 -2

@@ -6,7 +6,6 @@ Run `pytest tests/kernels/quantization/test_triton_scaled_mm.py`.

6 """

7

8 import importlib

9 -from typing import Optional

9

10 import pytest

11 import torch

@@ -27,7 +26,7 @@ def torch_scaled_mm(

26 scale_a: torch.Tensor,

27 scale_b: torch.Tensor,

28 out_dtype: type[torch.dtype],

30 - bias: Optional[torch.Tensor] = None,

29 + bias: torch.Tensor | None = None,

30 ) -> torch.Tensor:

31 out = torch.mm(a.to(torch.float32), b.to(torch.float32))

32 out = scale_a * out

tests/kernels/test_onednn.py

+2 -4

@@ -2,8 +2,6 @@

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

3 """Integration tests for FlexAttention backend vs default backend"""

4

5 -from typing import Optional

6 -

5 import pytest

6 import torch

7

@@ -38,8 +36,8 @@ def ref_int8_scaled_mm(

36 b: torch.Tensor,

37 scale_a: torch.Tensor,

38 scale_b: torch.Tensor,

41 - azp: Optional[torch.Tensor],

42 - bias: Optional[torch.Tensor],

39 + azp: torch.Tensor | None,

40 + bias: torch.Tensor | None,

41 output_type: torch.dtype,

42 ):

43 if azp is not None:

tests/kernels/utils.py

+54 -56

@@ -7,7 +7,7 @@ import random

7 import unittest

8 from collections.abc import Sequence

9 from numbers import Number

10 -from typing import Any, NamedTuple, Optional, Union

10 +from typing import Any, NamedTuple

11

12 import pytest

13 import torch

@@ -96,10 +96,10 @@ class PackedQKVInputs(NamedTuple):

96 query: torch.Tensor

97 key: torch.Tensor

98 value: torch.Tensor

99 - q_start_loc_list: Optional[list[int]]

100 - kv_start_loc_list: Optional[list[int]]

101 - q_seq_lens: Optional[list[int]]

102 - kv_seq_lens: Optional[list[int]]

99 + q_start_loc_list: list[int] | None

100 + kv_start_loc_list: list[int] | None

101 + q_seq_lens: list[int] | None

102 + kv_seq_lens: list[int] | None

103

104

105 class PackedQKVO(NamedTuple):

@@ -115,7 +115,7 @@ class PackedQKVO(NamedTuple):

115 x head_size) known-correct attention output

116 """

117

118 - packed_qkv: Optional[PackedQKVInputs]

118 + packed_qkv: PackedQKVInputs | None

119 ideal_output: torch.Tensor

120

121

@@ -149,12 +149,12 @@ class PhaseTestParameters(NamedTuple):

149 """

150

151 packed_qkvo: PackedQKVO

152 - kv_mmap: Optional[KVMemoryMap]

152 + kv_mmap: KVMemoryMap | None

153

154

155 def maybe_make_int_tensor(

156 - _list: Optional[list[int]],

157 - device: Union[torch.device, str],

156 + _list: list[int] | None,

157 + device: torch.device | str,

158 ) -> torch.Tensor:

159 """

160 Convert Python int list to a 1D int torch.Tensor on `device`

@@ -170,8 +170,8 @@ def maybe_make_int_tensor(

170

171

172 def maybe_make_long_tensor(

173 - _list: Optional[list[int]],

174 - device: Union[torch.device, str],

173 + _list: list[int] | None,

174 + device: torch.device | str,

175 ) -> torch.Tensor:

176 """

177 Convert Python int list to a 1D long torch.Tensor on `device`

@@ -186,7 +186,7 @@ def maybe_make_long_tensor(

186 )

187

188

189 -def maybe_max(_list: Optional[list]) -> Optional[Number]:

189 +def maybe_max(_list: list | None) -> Number | None:

190 """

191 Returns:

192

@@ -241,9 +241,9 @@ def ref_masked_attention(

241 key: torch.Tensor,

242 value: torch.Tensor,

243 scale: float,

244 - custom_mask: Optional[torch.Tensor] = None,

245 - q_seq_lens: Optional[list] = None,

246 - kv_seq_lens: Optional[list] = None,

244 + custom_mask: torch.Tensor | None = None,

245 + q_seq_lens: list | None = None,

246 + kv_seq_lens: list | None = None,

247 ) -> torch.Tensor:

248 """

249 "Golden" masked attention reference. Supports two types of masking:

@@ -302,11 +302,11 @@ def ref_masked_attention(

302 def make_qkv(

303 batch_size: int,

304 max_q_seq_len: int,

305 - max_kv_seq_len: Optional[int],

305 + max_kv_seq_len: int | None,

306 num_heads: int,

307 head_size: int,

308 - device: Union[torch.device, str],

309 - force_kv_seq_lens: Optional[list[int]] = None,

308 + device: torch.device | str,

309 + force_kv_seq_lens: list[int] | None = None,

310 attn_type: AttentionType = AttentionType.ENCODER_DECODER,

311 force_max_len: bool = False,

312 ) -> tuple[QKVInputs, QKVInputs, QKVInputs]:

@@ -436,7 +436,7 @@ def make_qkv(

436

437

438 def pack_tensor(

439 - unpacked_tensor: torch.Tensor, seq_lens: list[int], device: Union[torch.device, str]

439 + unpacked_tensor: torch.Tensor, seq_lens: list[int], device: torch.device | str

440 ) -> tuple[torch.Tensor, list[int]]:

441 """

442 Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an

@@ -470,7 +470,7 @@ def pack_tensor(

470 return packed_tensor, start_loc_list

471

472

473 -def pack_qkv(qkv: QKVInputs, device: Union[torch.device, str]) -> PackedQKVInputs:

473 +def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:

474 """

475 Individually pack each of Q, K and V, each with dimensions batch_size x

476 padded_seq_len x num_heads x head_size, into respective number_of_tokens x

@@ -594,19 +594,19 @@ def make_alibi_bias(

594

595

596 def _make_metadata_tensors(

597 - seq_lens: Optional[list[int]],

598 - context_lens: Optional[list[int]],

599 - encoder_seq_lens: Optional[list[int]],

600 - device: Union[torch.device, str],

597 + seq_lens: list[int] | None,

598 + context_lens: list[int] | None,

599 + encoder_seq_lens: list[int] | None,

600 + device: torch.device | str,

601 ) -> tuple[

602 torch.Tensor,

603 torch.Tensor,

604 Any,

605 Any,

606 - Optional[torch.Tensor],

606 + torch.Tensor | None,

607 torch.Tensor,

608 torch.Tensor,

609 - Optional[int],

609 + int | None,

610 ]:

611 """

612 Build scalar & tensor values required to build attention metadata structure.

@@ -678,7 +678,7 @@ def make_kv_cache(

678 num_heads: int,

679 head_size: int,

680 block_size: int,

681 - device: Union[torch.device, str],

681 + device: torch.device | str,

682 backend: str,

683 default_val: float = 0.0,

684 ) -> torch.Tensor:

@@ -726,18 +726,18 @@ def _num_tokens_to_min_blocks(num_tokens: int, block_size: int) -> int:

726 return (num_tokens + block_size) // block_size

727

728

729 -def make_empty_slot_mapping_tensor(device: Union[torch.device, str]):

729 +def make_empty_slot_mapping_tensor(device: torch.device | str):

730 return maybe_make_long_tensor([], device)

731

732

733 -def make_empty_block_tables_tensor(device: Union[torch.device, str]):

733 +def make_empty_block_tables_tensor(device: torch.device | str):

734 return torch.tensor([], device=device)

735

736

737 def split_slot_mapping(

738 slot_mapping_list: torch.Tensor,

739 seq_lens: list[int],

740 - device: Union[torch.device, str],

740 + device: torch.device | str,

741 ):

742 """

743 Split a slot mapping into valid prefill- and decode-phase slot mappings.

@@ -799,7 +799,7 @@ def split_slot_mapping(

799 def make_block_tables_slot_mapping(

800 block_size: int,

801 seq_lens: list[int],

802 - device: Union[torch.device, str],

802 + device: torch.device | str,

803 block_base_addr: int = 0,

804 ) -> tuple[torch.Tensor, list[int], int]:

805 """

@@ -880,11 +880,11 @@ def make_block_tables_slot_mapping(

880 def make_test_metadata(

881 attn_backend: _Backend,

882 is_prompt: bool,

883 - seq_lens: Optional[list[int]],

884 - decoder_test_params: Optional[PhaseTestParameters],

885 - device: Union[torch.device, str],

886 - encoder_test_params: Optional[PhaseTestParameters] = None,

887 - cross_test_params: Optional[PhaseTestParameters] = None,

883 + seq_lens: list[int] | None,

884 + decoder_test_params: PhaseTestParameters | None,

885 + device: torch.device | str,

886 + encoder_test_params: PhaseTestParameters | None = None,

887 + cross_test_params: PhaseTestParameters | None = None,

888 ) -> AttentionMetadata:

889 """

890 Construct fake attention metadata for a given test phase

@@ -1142,16 +1142,16 @@ def torch_experts(

1142 topk_weight: torch.Tensor,

1143 topk_ids: torch.Tensor,

1144 global_num_experts: int = -1,

1145 - b_bias1: Optional[torch.Tensor] = None,

1146 - b_bias2: Optional[torch.Tensor] = None,

1147 - expert_map: Optional[torch.Tensor] = None,

1148 - w1_scale: Optional[torch.Tensor] = None,

1149 - w2_scale: Optional[torch.Tensor] = None,

1150 - a1_scale: Optional[torch.Tensor] = None,

1151 - a2_scale: Optional[torch.Tensor] = None,

1152 - quant_dtype: Optional[torch.dtype] = None,

1145 + b_bias1: torch.Tensor | None = None,

1146 + b_bias2: torch.Tensor | None = None,

1147 + expert_map: torch.Tensor | None = None,

1148 + w1_scale: torch.Tensor | None = None,

1149 + w2_scale: torch.Tensor | None = None,

1150 + a1_scale: torch.Tensor | None = None,

1151 + a2_scale: torch.Tensor | None = None,

1152 + quant_dtype: torch.dtype | None = None,

1153 per_act_token_quant=False,

1154 - block_shape: Optional[list[int]] = None,

1154 + block_shape: list[int] | None = None,

1155 apply_router_weights_on_input: bool = False,

1156 ) -> torch.Tensor:

1157 assert (

@@ -1261,10 +1261,10 @@ def torch_moe(

1261 w2: torch.Tensor,

1262 score: torch.Tensor,

1263 topk: int,

1264 - b_bias1: Optional[torch.Tensor] = None,

1265 - b_bias2: Optional[torch.Tensor] = None,

1264 + b_bias1: torch.Tensor | None = None,

1265 + b_bias2: torch.Tensor | None = None,

1266 global_num_experts: int = -1,

1267 - expert_map: Optional[torch.Tensor] = None,

1267 + expert_map: torch.Tensor | None = None,

1268 ) -> torch.Tensor:

1269 score = torch.softmax(score, dim=-1, dtype=torch.float32)

1270 topk_weight, topk_ids = torch.topk(score, topk)

@@ -1298,15 +1298,13 @@ def torch_moe_single(a, w, score, topk):

1298 # A special version of op check that has a restricted default set of test_utils

1299 # and a patched version of allclose that supports fp8 types.

1300 def opcheck(

1301 - op: Union[

1302 - torch._ops.OpOverload,

1303 - torch._ops.OpOverloadPacket,

1304 - torch._library.custom_ops.CustomOpDef,

1305 - ],

1301 + op: torch._ops.OpOverload

1302 + | torch._ops.OpOverloadPacket

1303 + | torch._library.custom_ops.CustomOpDef,

1304 args: tuple[Any, ...],

1307 - kwargs: Optional[dict[str, Any]] = None,

1305 + kwargs: dict[str, Any] | None = None,

1306 *,

1309 - test_utils: Union[str, Sequence[str]] = ALL_OPCHECK_TEST_UTILS,

1307 + test_utils: str | Sequence[str] = ALL_OPCHECK_TEST_UTILS,

1308 raise_exception: bool = True,

1309 cond: bool = True,

1310 ) -> dict[str, str]:

@@ -1338,7 +1336,7 @@ def baseline_scaled_mm(

1336 scale_a: torch.Tensor,

1337 scale_b: torch.Tensor,

1338 out_dtype: type[torch.dtype],

1341 - bias: Optional[torch.Tensor] = None,

1339 + bias: torch.Tensor | None = None,

1340 ) -> torch.Tensor:

1341 # We treat N-dimensional group scaling as extended numpy-style broadcasting

1342 # in numpy simply stretches dimensions with an extent of 1 to match

tests/lora/test_layers.py

+3 -4

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

4 import random

5 from copy import deepcopy

6 from dataclasses import dataclass

7 -from typing import Optional

7 from unittest.mock import patch

8

9 import pytest

@@ -106,7 +105,7 @@ def skip_cuda_with_stage_false(request):

105

106 def get_random_id_to_index(

107 num_loras: int, num_slots: int, log: bool = True

109 -) -> list[Optional[int]]:

108 +) -> list[int | None]:

109 """Creates a random lora_id_to_index mapping.

110

111 Args:

@@ -122,7 +121,7 @@ def get_random_id_to_index(

121 "num_loras must be less than or equal to num_slots."

122 )

123

125 - slots: list[Optional[int]] = [None] * num_slots

124 + slots: list[int | None] = [None] * num_slots

125 random_slot_selections = (torch.randperm(num_slots)[:num_loras]).tolist()

126 for lora_id, slot_idx in enumerate(random_slot_selections, start=1):

127 slots[slot_idx] = lora_id

@@ -134,7 +133,7 @@ def get_random_id_to_index(

133

134

135 def populate_loras(

137 - id_to_index: list[Optional[int]],

136 + id_to_index: list[int | None],

137 layer: BaseLayerWithLoRA,

138 layer_weights: torch.Tensor,

139 generate_embeddings_tensor: int = 0,

tests/lora/test_llama_tp.py

+2 -5

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

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

3 import subprocess

4 import sys

5 -from typing import Union

5

6 import vllm

7 from vllm import LLM

@@ -27,7 +26,7 @@ def do_sample(

26 llm: vllm.LLM,

27 lora_path: str,

28 lora_id: int,

30 - tensorizer_config_dict: Union[dict, None] = None,

29 + tensorizer_config_dict: dict | None = None,

30 ) -> list[str]:

31 prompts = [

32 "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501

@@ -73,9 +72,7 @@ def do_sample(

72 return generated_texts

73

74

76 -def generate_and_test(

77 - llm, sql_lora_files, tensorizer_config_dict: Union[dict, None] = None

78 -):

75 +def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None = None):

76 print("lora adapter created")

77 print("lora 1")

78 assert (

tests/lora/test_qwen2vl.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass

4 -from typing import Optional

4

5 import pytest

6

@@ -20,7 +19,7 @@ class TestConfig:

19 max_loras: int = 2

20 max_lora_rank: int = 16

21 max_model_len: int = 4096

23 - mm_processor_kwargs: Optional[dict[str, int]] = None

22 + mm_processor_kwargs: dict[str, int] | None = None

23

24 def __post_init__(self):

25 if self.mm_processor_kwargs is None:

@@ -61,7 +60,7 @@ class Qwen2VLTester:

60 self,

61 images: list[ImageAsset],

62 expected_outputs: list[str],

64 - lora_id: Optional[int] = None,

63 + lora_id: int | None = None,

64 temperature: float = 0,

65 max_tokens: int = 5,

66 ):

@@ -92,7 +91,7 @@ class Qwen2VLTester:

91 self,

92 images: list[ImageAsset],

93 expected_outputs: list[list[str]],

95 - lora_id: Optional[int] = None,

94 + lora_id: int | None = None,

95 temperature: float = 0,

96 beam_width: int = 2,

97 max_tokens: int = 5,

tests/lora/test_resolver.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6

@@ -14,7 +13,7 @@ class DummyLoRAResolver(LoRAResolver):

13

14 async def resolve_lora(

15 self, base_model_name: str, lora_name: str

17 - ) -> Optional[LoRARequest]:

16 + ) -> LoRARequest | None:

17 if lora_name == "test_lora":

18 return LoRARequest(

19 lora_name=lora_name,

tests/lora/test_utils.py

+2 -2

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

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

3

4 from collections import OrderedDict

5 -from typing import NamedTuple, Optional

5 +from typing import NamedTuple

6 from unittest.mock import patch

7

8 import pytest

@@ -21,7 +21,7 @@ class LoRANameParserTestConfig(NamedTuple):

21 name: str

22 module_name: str

23 is_lora_a: bool

24 - weights_mapper: Optional[WeightsMapper] = None

24 + weights_mapper: WeightsMapper | None = None

25

26

27 def test_parse_fine_tuned_lora_name_valid():

tests/lora/utils.py

+2 -3

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

4 import json

5 import os

6 from dataclasses import dataclass

7 -from typing import Optional, Union

7

8 import torch

9 from safetensors.torch import save_file

@@ -81,7 +80,7 @@ class DummyLoRAManager:

80 module_name: str,

81 input_dim: int,

82 output_dims: list[int],

84 - noop_lora_index: Optional[list[int]] = None,

83 + noop_lora_index: list[int] | None = None,

84 rank: int = 8,

85 ):

86 base_loras: list[LoRALayerWeights] = []

@@ -113,7 +112,7 @@ def assert_close(a, b):

112 @dataclass

113 class PunicaTensors:

114 inputs_tensor: torch.Tensor

116 - lora_weights: Union[torch.Tensor, list[torch.Tensor]]

115 + lora_weights: torch.Tensor | list[torch.Tensor]

116 our_out_tensor: torch.Tensor

117 ref_out_tensor: torch.Tensor

118 b_seq_start_loc: torch.Tensor

tests/model_executor/model_loader/tensorizer_loader/conftest.py

+1 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Callable

3 +from collections.abc import Callable

4

5 import pytest

6

tests/model_executor/test_enabled_custom_ops.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import pytest

5 import torch

@@ -76,7 +75,7 @@ class Relu3(ReLUSquaredActivation):

75 ],

76 )

77 def test_enabled_ops(

79 - env: Optional[str],

78 + env: str | None,

79 torch_level: int,

80 use_inductor: bool,

81 ops_enabled: list[int],

tests/models/language/generation/test_common.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import pytest

5 import torch

@@ -138,7 +137,7 @@ def test_models(

137 example_prompts, max_tokens, num_logprobs

138 )

139

141 - prompt_embeds: Optional[list[torch.Tensor]] = [] if use_prompt_embeds else None

140 + prompt_embeds: list[torch.Tensor] | None = [] if use_prompt_embeds else None

141

142 prompt_token_ids = []

143 for prompt in example_prompts:

tests/models/language/generation/test_hybrid.py

+1 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable

4 +from collections.abc import Callable

5

6 import pytest

7

tests/models/language/generation_ppl_test/ppl_utils.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 # Adapted from https://huggingface.co/docs/transformers/perplexity

4 -from typing import Optional, cast

4 +from typing import cast

5

6 import pytest

7 import torch

@@ -85,7 +85,7 @@ def wikitext_ppl_test(

85 n_tokens = 0

86 for output in outputs:

87 output = cast(TokensTextLogprobsPromptLogprobs, output)

88 - token_datas = cast(list[Optional[dict[int, Logprob]]], output[3])

88 + token_datas = cast(list[dict[int, Logprob] | None], output[3])

89

90 assert token_datas[0] is None

91 token_log_probs = []

tests/models/language/pooling/embed_utils.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from collections.abc import Sequence

4 -from typing import Optional

4

5 import pytest

6

@@ -13,7 +12,7 @@ def run_embedding_correctness_test(

12 hf_model: "HfRunner",

13 inputs: list[str],

14 vllm_outputs: Sequence[list[float]],

16 - dimensions: Optional[int] = None,

15 + dimensions: int | None = None,

16 ):

17 hf_outputs = hf_model.encode(inputs)

18 if dimensions:

tests/models/language/pooling/test_embedding.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import pytest

5

@@ -66,7 +65,7 @@ def test_models(

65 pooling_type="MEAN", normalize=False

66 )

67

69 - max_model_len: Optional[int] = 512

68 + max_model_len: int | None = 512

69 if model in [

70 "sentence-transformers/all-MiniLM-L12-v2",

71 "sentence-transformers/stsb-roberta-base-v2",

tests/models/language/pooling/test_gritlm.py

-2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import numpy as np

4 import openai

5 import pytest

tests/models/language/pooling_mteb_test/mteb_utils.py

+3 -4

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

3

4 import tempfile

5 from collections.abc import Sequence

6 -from typing import Optional

6

7 import mteb

8 import numpy as np

@@ -51,7 +50,7 @@ class VllmMtebEncoder(mteb.Encoder):

50

51 def predict(

52 self,

54 - sentences: list[tuple[str, str, Optional[str]]], # query, corpus, prompt

53 + sentences: list[tuple[str, str, str | None]], # query, corpus, prompt

54 *args,

55 **kwargs,

56 ) -> np.ndarray:

@@ -100,7 +99,7 @@ class ScoreClientMtebEncoder(mteb.Encoder):

99

100 def predict(

101 self,

103 - sentences: list[tuple[str, str, Optional[str]]], # query, corpus, prompt

102 + sentences: list[tuple[str, str, str | None]], # query, corpus, prompt

103 *args,

104 **kwargs,

105 ) -> np.ndarray:

@@ -294,7 +293,7 @@ def mteb_test_rerank_models_hf(

293 original_predict = hf_model.predict

294

295 def _predict(

297 - sentences: list[tuple[str, str, Optional[str]]], # query, corpus, prompt

296 + sentences: list[tuple[str, str, str | None]], # query, corpus, prompt

297 *args,

298 **kwargs,

299 ):

tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Any, Optional

3 +from typing import Any

4

5 import numpy as np

6 import pytest

@@ -111,7 +111,7 @@ class GemmaMtebEncoder(VllmMtebEncoder):

111

112 def predict(

113 self,

114 - sentences: list[tuple[str, str, Optional[str]]], # query, corpus, prompt

114 + sentences: list[tuple[str, str, str | None]], # query, corpus, prompt

115 *args,

116 **kwargs,

117 ) -> np.ndarray:

tests/models/multimodal/generation/test_granite_speech.py

+3 -4

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

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

3

4 from collections.abc import Sequence

5 -from typing import Optional

5

6 import pytest

7 from transformers import AutoModelForSpeechSeq2Seq

@@ -18,8 +17,8 @@ HF_AUDIO_PROMPT = "<|start_of_role|>system<|end_of_role|>Knowledge Cutoff Date:

17

18

19 def vllm_to_hf_output(

21 - vllm_output: tuple[list[int], str, Optional[SampleLogprobs]],

22 -) -> tuple[list[int], str, Optional[SampleLogprobs]]:

20 + vllm_output: tuple[list[int], str, SampleLogprobs | None],

21 +) -> tuple[list[int], str, SampleLogprobs | None]:

22 """Sanitize hf output to be comparable with vllm output."""

23 output_ids, output_str, out_logprobs = vllm_output

24

@@ -46,7 +45,7 @@ def run_test(

45 max_tokens: int,

46 num_logprobs: int,

47 tensor_parallel_size: int,

49 - distributed_executor_backend: Optional[str] = None,

48 + distributed_executor_backend: str | None = None,

49 ):

50 """Inference result should be the same between hf and vllm.

51

tests/models/multimodal/generation/test_phi4_multimodal.py

+2 -3

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

3

4 import os

5 from collections.abc import Sequence

6 -from typing import Optional

6

7 import librosa

8 import pytest

@@ -57,7 +56,7 @@ if current_platform.is_rocm():

56 def run_test(

57 hf_runner: type[HfRunner],

58 vllm_runner: type[VllmRunner],

60 - inputs: Sequence[tuple[list[str], PromptImageInput, Optional[PromptAudioInput]]],

59 + inputs: Sequence[tuple[list[str], PromptImageInput, PromptAudioInput | None]],

60 model: str,

61 *,

62 max_model_len: int,

@@ -66,7 +65,7 @@ def run_test(

65 num_logprobs: int,

66 mm_limit: int,

67 tensor_parallel_size: int,

69 - distributed_executor_backend: Optional[str] = None,

68 + distributed_executor_backend: str | None = None,

69 ):

70 """Inference result should be the same between hf and vllm.

71

tests/models/multimodal/generation/test_phi4mm.py

+3 -4

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

3

4 import os

5 from collections.abc import Sequence

6 -from typing import Optional

6

7 import librosa

8 import pytest

@@ -48,7 +47,7 @@ models = [model_path]

47

48

49 def vllm_to_hf_output(

51 - vllm_output: tuple[list[int], str, Optional[SampleLogprobs]], model: str

50 + vllm_output: tuple[list[int], str, SampleLogprobs | None], model: str

51 ):

52 """Sanitize vllm output to be comparable with hf output."""

53 _, output_str, out_logprobs = vllm_output

@@ -79,7 +78,7 @@ if current_platform.is_rocm():

78 def run_test(

79 hf_runner: type[HfRunner],

80 vllm_runner: type[VllmRunner],

82 - inputs: Sequence[tuple[list[str], PromptImageInput, Optional[PromptAudioInput]]],

81 + inputs: Sequence[tuple[list[str], PromptImageInput, PromptAudioInput | None]],

82 model: str,

83 *,

84 max_model_len: int,

@@ -88,7 +87,7 @@ def run_test(

87 num_logprobs: int,

88 mm_limit: int,

89 tensor_parallel_size: int,

91 - distributed_executor_backend: Optional[str] = None,

90 + distributed_executor_backend: str | None = None,

91 ):

92 """Inference result should be the same between hf and vllm.

93

tests/models/multimodal/generation/test_pixtral.py

+2 -2

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

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

3 import json

4 from dataclasses import asdict

5 -from typing import TYPE_CHECKING, Any, Optional

5 +from typing import TYPE_CHECKING, Any

6

7 import pytest

8 from mistral_common.multimodal import download_image

@@ -117,7 +117,7 @@ FIXTURE_LOGPROBS_CHAT = {

117 MISTRAL_SMALL_3_1_ID: FIXTURES_PATH / "mistral_small_3_chat.json",

118 }

119

120 -OutputsLogprobs = list[tuple[list[int], str, Optional[SampleLogprobs]]]

120 +OutputsLogprobs = list[tuple[list[int], str, SampleLogprobs | None]]

121

122

123 # For the test author to store golden output in JSON

tests/models/multimodal/generation/test_qwen2_vl.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Optional, TypedDict, Union

4 +from typing import Any, TypedDict

5

6 import numpy.typing as npt

7 import pytest

@@ -83,7 +83,7 @@ class Qwen2VLPromptVideoEmbeddingInput(TypedDict):

83

84

85 def batch_make_image_embeddings(

86 - image_batches: list[Union[Image.Image, list[Image.Image]]],

86 + image_batches: list[Image.Image | list[Image.Image]],

87 processor,

88 llm: VllmRunner,

89 ) -> list[Qwen2VLPromptImageEmbeddingInput]:

@@ -272,7 +272,7 @@ def run_embedding_input_test(

272 num_logprobs: int,

273 mm_limit: int,

274 tensor_parallel_size: int,

275 - distributed_executor_backend: Optional[str] = None,

275 + distributed_executor_backend: str | None = None,

276 ):

277 """Inference result should be the same between

278 original image/video input and image/video embeddings input.

tests/models/multimodal/generation/test_whisper.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import pytest

5

@@ -92,7 +91,7 @@ def run_test(

91 model: str,

92 *,

93 tensor_parallel_size: int,

95 - distributed_executor_backend: Optional[str] = None,

94 + distributed_executor_backend: str | None = None,

95 ) -> None:

96 prompt_list = PROMPTS * 10

97 expected_list = EXPECTED[model] * 10

tests/models/multimodal/generation/vlm_utils/builders.py

+7 -10

@@ -2,9 +2,8 @@

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

3 """Helpers for building inputs that can be leveraged for different test types."""

4

5 -from collections.abc import Iterable

5 +from collections.abc import Callable, Iterable

6 from pathlib import PosixPath

7 -from typing import Callable, Optional, Union

7

8 import torch

9

@@ -47,9 +46,9 @@ def replace_test_placeholder(

46

47 def get_model_prompts(

48 base_prompts: Iterable[str],

50 - img_idx_to_prompt: Optional[Callable[[int], str]],

51 - video_idx_to_prompt: Optional[Callable[[int], str]],

52 - audio_idx_to_prompt: Optional[Callable[[int], str]],

49 + img_idx_to_prompt: Callable[[int], str] | None,

50 + video_idx_to_prompt: Callable[[int], str] | None,

51 + audio_idx_to_prompt: Callable[[int], str] | None,

52 prompt_formatter: Callable[[str], str],

53 ) -> list[str]:

54 """Given a model-agnostic base prompt and test configuration for a model(s)

@@ -93,7 +92,7 @@ def build_single_image_inputs_from_test_info(

92 test_info: VLMTestInfo,

93 image_assets: ImageTestAssets,

94 size_wrapper: ImageSizeWrapper,

96 - tmp_path: Optional[PosixPath] = None,

95 + tmp_path: PosixPath | None = None,

96 ) -> list[PromptWithMultiModalInput]:

97 if test_info.prompt_formatter is None:

98 raise ValueError("Prompt formatter must be set to build single image inputs")

@@ -147,7 +146,7 @@ def build_multi_image_inputs_from_test_info(

146 test_info: VLMTestInfo,

147 image_assets: ImageTestAssets,

148 size_wrapper: ImageSizeWrapper,

150 - tmp_path: Optional[PosixPath] = None,

149 + tmp_path: PosixPath | None = None,

150 ) -> list[PromptWithMultiModalInput]:

151 if test_info.prompt_formatter is None:

152 raise ValueError("Prompt formatter must be set to build multi image inputs")

@@ -266,9 +265,7 @@ def build_video_inputs_from_test_info(

265 ]

266

267

269 -def apply_image_size_scaling(

270 - image, size: Union[float, tuple[int, int]], size_type: SizeType

271 -):

268 +def apply_image_size_scaling(image, size: float | tuple[int, int], size_type: SizeType):

269 """Applies a size scaler to one image; this can be an image size factor,

270 which scales the image while maintaining the aspect ratio"""

271 # Special case for embeddings; if it's a tensor, it's only valid if we

tests/models/multimodal/generation/vlm_utils/core.py

+11 -10

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

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

3 """Core test implementation to be shared across modalities."""

4

5 -from typing import Any, Callable, Optional

5 +from collections.abc import Callable

6 +from typing import Any

7

8 import torch

9 from transformers.models.auto.auto_factory import _BaseAutoModelClass

@@ -27,21 +28,21 @@ def run_test(

28 enforce_eager: bool,

29 max_model_len: int,

30 max_num_seqs: int,

30 - hf_output_post_proc: Optional[Callable[[RunnerOutput, str], Any]],

31 - vllm_output_post_proc: Optional[Callable[[RunnerOutput, str], Any]],

31 + hf_output_post_proc: Callable[[RunnerOutput, str], Any] | None,

32 + vllm_output_post_proc: Callable[[RunnerOutput, str], Any] | None,

33 auto_cls: type[_BaseAutoModelClass],

34 use_tokenizer_eos: bool,

35 comparator: Callable[..., None],

35 - get_stop_token_ids: Optional[Callable[[AnyTokenizer], list[int]]],

36 - stop_str: Optional[list[str]],

36 + get_stop_token_ids: Callable[[AnyTokenizer], list[int]] | None,

37 + stop_str: list[str] | None,

38 limit_mm_per_prompt: dict[str, int],

38 - vllm_runner_kwargs: Optional[dict[str, Any]],

39 - hf_model_kwargs: Optional[dict[str, Any]],

40 - patch_hf_runner: Optional[Callable[[HfRunner], HfRunner]],

39 + vllm_runner_kwargs: dict[str, Any] | None,

40 + hf_model_kwargs: dict[str, Any] | None,

41 + patch_hf_runner: Callable[[HfRunner], HfRunner] | None,

42 runner: RunnerOption = "auto",

42 - distributed_executor_backend: Optional[str] = None,

43 + distributed_executor_backend: str | None = None,

44 tensor_parallel_size: int = 1,

44 - vllm_embeddings: Optional[torch.Tensor] = None,

45 + vllm_embeddings: torch.Tensor | None = None,

46 ):

47 """Modality agnostic test executor for comparing HF/vLLM outputs."""

48 # In the case of embeddings, vLLM takes separate input tensors

tests/models/multimodal/generation/vlm_utils/custom_inputs.py

+1 -1

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

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

3 """Custom input builders for edge-cases in different models."""

4

5 -from typing import Callable

5 +from collections.abc import Callable

6

7 from vllm.assets.image import ImageAsset

8 from vllm.multimodal.image import rescale_image_size

tests/models/multimodal/generation/vlm_utils/model_utils.py

+10 -11

@@ -7,7 +7,6 @@ typically specific to a small subset of models.

7

8 import types

9 from pathlib import PosixPath

10 -from typing import Optional, Union

10

11 import numpy as np

12 import numpy.typing as npt

@@ -58,7 +57,7 @@ def fuyu_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOutpu

57

58 def qwen_vllm_to_hf_output(

59 vllm_output: RunnerOutput, model: str

61 -) -> tuple[list[int], str, Optional[SampleLogprobs]]:

60 +) -> tuple[list[int], str, SampleLogprobs | None]:

61 """Sanitize vllm output [qwen models] to be comparable with hf output."""

62 output_ids, output_str, out_logprobs = vllm_output

63

@@ -69,7 +68,7 @@ def qwen_vllm_to_hf_output(

68

69 def qwen2_vllm_to_hf_output(

70 vllm_output: RunnerOutput, model: str

72 -) -> tuple[list[int], str, Optional[SampleLogprobs]]:

71 +) -> tuple[list[int], str, SampleLogprobs | None]:

72 """Sanitize vllm output [qwen2 models] to be comparable with hf output."""

73 output_ids, output_str, out_logprobs = vllm_output

74

@@ -80,7 +79,7 @@ def qwen2_vllm_to_hf_output(

79

80 def kimiv_vl_vllm_to_hf_output(

81 vllm_output: RunnerOutput, model: str

83 -) -> tuple[list[int], str, Optional[SampleLogprobs]]:

82 +) -> tuple[list[int], str, SampleLogprobs | None]:

83 """Sanitize vllm output [kimi_vl models] to be comparable with hf output."""

84 output_ids, output_str, out_logprobs = vllm_output

85

@@ -99,7 +98,7 @@ def llava_image_vllm_to_hf_output(

98

99 def llava_video_vllm_to_hf_output(

100 vllm_output: RunnerOutput, model: str

102 -) -> tuple[list[int], str, Optional[SampleLogprobs]]:

101 +) -> tuple[list[int], str, SampleLogprobs | None]:

102 config = AutoConfig.from_pretrained(model)

103 mm_token_id = config.video_token_index

104 return _llava_vllm_to_hf_output(vllm_output, model, mm_token_id)

@@ -263,7 +262,7 @@ def get_llava_embeddings(image_assets: ImageTestAssets):

262

263 ####### Prompt path encoders for models that need models on disk

264 def qwen_prompt_path_encoder(

266 - tmp_path: PosixPath, prompt: str, assets: Union[list[ImageAsset], ImageTestAssets]

265 + tmp_path: PosixPath, prompt: str, assets: list[ImageAsset] | ImageTestAssets

266 ) -> str:

267 """Given a temporary dir path, export one or more image assets into the

268 tempdir & replace its contents with the local path to the string so that

@@ -440,7 +439,7 @@ def h2ovl_patch_hf_runner(hf_model: HfRunner) -> HfRunner:

439 self.max_num = self.config.max_dynamic_patch

440 self.image_size = self.vision_config.image_size

441

443 - def __call__(self, text: str, images: Union[Image, list[Image]], **kwargs):

442 + def __call__(self, text: str, images: Image | list[Image], **kwargs):

443 from vllm.model_executor.models.h2ovl import (

444 IMG_CONTEXT,

445 IMG_END,

@@ -499,7 +498,7 @@ def skyworkr1v_patch_hf_runner(hf_model: HfRunner) -> HfRunner:

498 self.max_num = self.config.max_dynamic_patch

499 self.image_size = self.vision_config.image_size

500

502 - def __call__(self, text: str, images: Union[Image, list[Image]], **kwargs):

501 + def __call__(self, text: str, images: Image | list[Image], **kwargs):

502 from vllm.model_executor.models.skyworkr1v import (

503 IMG_CONTEXT,

504 IMG_END,

@@ -560,8 +559,8 @@ def internvl_patch_hf_runner(hf_model: HfRunner) -> HfRunner:

559 def __call__(

560 self,

561 text: str,

563 - images: Union[Image, list[Image]] = None,

564 - videos: Union[npt.NDArray, list[npt.NDArray]] = None,

562 + images: Image | list[Image] = None,

563 + videos: npt.NDArray | list[npt.NDArray] = None,

564 **kwargs,

565 ):

566 from vllm.model_executor.models.internvl import (

@@ -650,7 +649,7 @@ def _internvl_generate(

649 self,

650 pixel_values: torch.FloatTensor,

651 input_ids: torch.FloatTensor,

653 - attention_mask: Optional[torch.LongTensor] = None,

652 + attention_mask: torch.LongTensor | None = None,

653 **generate_kwargs,

654 ) -> torch.LongTensor:

655 """Generate method for InternVL2 model without fixed use_cache."""

tests/models/multimodal/generation/vlm_utils/types.py

+30 -30

@@ -2,10 +2,10 @@

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

3 """Types for writing multimodal model tests."""

4

5 -from collections.abc import Iterable

5 +from collections.abc import Callable, Iterable

6 from enum import Enum

7 from pathlib import PosixPath

8 -from typing import Any, Callable, NamedTuple, Optional, Union

8 +from typing import Any, NamedTuple

9

10 import torch

11 from pytest import MarkDecorator

@@ -52,16 +52,16 @@ VIDEO_BASE_PROMPT = f"{TEST_VIDEO_PLACEHOLDER}Why is this video funny?"

52

53 IMAGE_SIZE_FACTORS = [(), (1.0,), (1.0, 1.0, 1.0), (0.25, 0.5, 1.0)]

54 EMBEDDING_SIZE_FACTORS = [(), (1.0,), (1.0, 1.0, 1.0)]

55 -RunnerOutput = tuple[list[int], str, Optional[SampleLogprobs]]

55 +RunnerOutput = tuple[list[int], str, SampleLogprobs | None]

56

57

58 class PromptWithMultiModalInput(NamedTuple):

59 """Holds the multimodal input for a single test case."""

60

61 prompts: list[str]

62 - image_data: Optional[PromptImageInput] = None

63 - video_data: Optional[PromptVideoInput] = None

64 - audio_data: Optional[PromptAudioInput] = None

62 + image_data: PromptImageInput | None = None

63 + video_data: PromptVideoInput | None = None

64 + audio_data: PromptAudioInput | None = None

65

66

67 class VLMTestType(Enum):

@@ -87,17 +87,17 @@ class ImageSizeWrapper(NamedTuple):

87 type: SizeType

88 # A size factor is a wrapper of 0+ floats,

89 # while a fixed size contains an iterable of integer pairs

90 - data: Union[Iterable[float], Iterable[tuple[int, int]]]

90 + data: Iterable[float] | Iterable[tuple[int, int]]

91

92

93 class VLMTestInfo(NamedTuple):

94 """Holds the configuration for 1+ tests for one model architecture."""

95

96 models: list[str]

97 - test_type: Union[VLMTestType, Iterable[VLMTestType]]

97 + test_type: VLMTestType | Iterable[VLMTestType]

98

99 # Should be None only if this is a CUSTOM_INPUTS test

100 - prompt_formatter: Optional[Callable[[str], str]] = None

100 + prompt_formatter: Callable[[str], str] | None = None

101 img_idx_to_prompt: Callable[[int], str] = lambda idx: "<image>\n"

102 video_idx_to_prompt: Callable[[int], str] = lambda idx: "<video>\n"

103 audio_idx_to_prompt: Callable[[int], str] = lambda idx: "<audio>\n"

@@ -111,9 +111,9 @@ class VLMTestInfo(NamedTuple):

111

112 # Function for converting ImageAssets to image embeddings;

113 # We need to define this explicitly for embedding tests

114 - convert_assets_to_embeddings: Optional[

115 - Callable[[ImageTestAssets], list[torch.Tensor]]

116 - ] = None

114 + convert_assets_to_embeddings: (

115 + Callable[[ImageTestAssets], list[torch.Tensor]] | None

116 + ) = None

117

118 # Exposed options for vLLM runner; we change these in a several tests,

119 # but the defaults are derived from VllmRunner & the engine defaults

@@ -123,25 +123,25 @@ class VLMTestInfo(NamedTuple):

123 max_num_seqs: int = 256

124 runner: RunnerOption = "auto"

125 tensor_parallel_size: int = 1

126 - vllm_runner_kwargs: Optional[dict[str, Any]] = None

126 + vllm_runner_kwargs: dict[str, Any] | None = None

127

128 # Optional callable which gets a list of token IDs from the model tokenizer

129 - get_stop_token_ids: Optional[Callable[[AnyTokenizer], list[int]]] = None

129 + get_stop_token_ids: Callable[[AnyTokenizer], list[int]] | None = None

130 # Optional list of strings to stop generation, useful when stop tokens are

131 # not special tokens in the tokenizer

132 - stop_str: Optional[list[str]] = None

132 + stop_str: list[str] | None = None

133

134 # Exposed options for HF runner

135 - hf_model_kwargs: Optional[dict[str, Any]] = None

135 + hf_model_kwargs: dict[str, Any] | None = None

136 # Indicates we should explicitly pass the EOS from the tokenizer

137 use_tokenizer_eos: bool = False

138 auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM

139 - patch_hf_runner: Optional[Callable[[HfRunner], HfRunner]] = None

139 + patch_hf_runner: Callable[[HfRunner], HfRunner] | None = None

140

141 # Post processors that if defined, will run oun the outputs of the

142 # vLLM and HF runner, respectively (useful for sanitization, etc).

143 - vllm_output_post_proc: Optional[Callable[[RunnerOutput, str], Any]] = None

144 - hf_output_post_proc: Optional[Callable[[RunnerOutput, str], Any]] = None

143 + vllm_output_post_proc: Callable[[RunnerOutput, str], Any] | None = None

144 + hf_output_post_proc: Callable[[RunnerOutput, str], Any] | None = None

145

146 # Consumes the output of the callables above and checks if they're equal

147 comparator: Callable[..., None] = check_logprobs_close

@@ -152,7 +152,7 @@ class VLMTestInfo(NamedTuple):

152 max_tokens: int = 128

153 num_logprobs: int = 5

154 dtype: str = "auto"

155 - distributed_executor_backend: Optional[str] = None

155 + distributed_executor_backend: str | None = None

156 # Only expanded in video tests

157 num_video_frames: int = 16

158

@@ -162,19 +162,19 @@ class VLMTestInfo(NamedTuple):

162 # once per tests (much like concatenating and wrapping in one parametrize

163 # call)

164 image_size_factors: Iterable[Iterable[float]] = IMAGE_SIZE_FACTORS

165 - image_sizes: Optional[Iterable[Iterable[tuple[int, int]]]] = None

165 + image_sizes: Iterable[Iterable[tuple[int, int]]] | None = None

166

167 # Hack for updating a prompt to take into a local path; currently only used

168 # for Qwen-VL, which requires encoding the image path / url into the prompt

169 # for HF runner

170 - prompt_path_encoder: Optional[

171 - Callable[[PosixPath, str, Union[list[ImageAsset], ImageTestAssets]], str]

172 - ] = None # noqa: E501

170 + prompt_path_encoder: (

171 + Callable[[PosixPath, str, list[ImageAsset] | ImageTestAssets], str] | None

172 + ) = None # noqa: E501

173

174 # Allows configuring a test to run with custom inputs

175 - custom_test_opts: Optional[list[CustomTestOptions]] = None

175 + custom_test_opts: list[CustomTestOptions] | None = None

176

177 - marks: Optional[list[MarkDecorator]] = None

177 + marks: list[MarkDecorator] | None = None

178

179 def get_non_parametrized_runner_kwargs(self):

180 """Returns a dictionary of expandable kwargs for items that are used

@@ -207,10 +207,10 @@ class ExpandableVLMTestArgs(NamedTuple):

207 max_tokens: int

208 num_logprobs: int

209 dtype: str

210 - distributed_executor_backend: Optional[str]

210 + distributed_executor_backend: str | None

211 # Sizes are used for everything except for custom input tests

212 - size_wrapper: Optional[ImageSizeWrapper] = None

212 + size_wrapper: ImageSizeWrapper | None = None

213 # Video only

214 - num_video_frames: Optional[int] = None

214 + num_video_frames: int | None = None

215 # Custom inputs only

216 - custom_test_opts: Optional[CustomTestOptions] = None

216 + custom_test_opts: CustomTestOptions | None = None

tests/models/multimodal/pooling/test_dse_qwen2_vl.py

+1 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable

4 +from collections.abc import Callable

5

6 import pytest

7 import torch

tests/models/multimodal/pooling/test_jinavl_reranker.py

+2 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Union

3

4 import pytest

5 from transformers import AutoModel

@@ -32,7 +31,7 @@ def vllm_reranker(

31 def create_image_param(url: str) -> ChatCompletionContentPartImageParam:

32 return {"type": "image_url", "image_url": {"url": f"{url}"}}

33

35 - query: Union[list[str], ScoreMultiModalParam]

34 + query: list[str] | ScoreMultiModalParam

35 if query_type == "text":

36 query = query_strs

37 elif query_type == "image":

@@ -40,7 +39,7 @@ def vllm_reranker(

39 content=[create_image_param(url) for url in query_strs]

40 )

41

43 - documents: Union[list[str], ScoreMultiModalParam]

42 + documents: list[str] | ScoreMultiModalParam

43 if doc_type == "text":

44 documents = document_strs

45 elif doc_type == "image":

tests/models/multimodal/processing/test_common.py

+2 -3

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

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

3

4 from functools import partial

5 -from typing import Optional, Union

5

6 import numpy as np

7 import pytest

@@ -247,7 +246,7 @@ MM_DATA_PATCHES = {

246 def _test_processing_correctness_one(

247 model_config: ModelConfig,

248 tokenizer: AnyTokenizer,

250 - prompt: Union[str, list[int]],

249 + prompt: str | list[int],

250 mm_data: MultiModalDataDict,

251 baseline_processor: BaseMultiModalProcessor,

252 cached_processor: BaseMultiModalProcessor,

@@ -441,7 +440,7 @@ def _assert_inputs_equal(

440 a: MultiModalInputs,

441 b: MultiModalInputs,

442 *,

444 - ignore_mm_keys: Optional[set[str]] = None,

443 + ignore_mm_keys: set[str] | None = None,

444 msg: str = "",

445 ):

446 if ignore_mm_keys is None:

tests/models/multimodal/processing/test_h2ovl.py

+1 -2

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

3 """Tests for H2OVL's multimodal preprocessing kwargs."""

4

5 from collections.abc import Mapping

6 -from typing import Optional

6

7 import pytest

8 from PIL import Image

@@ -149,7 +148,7 @@ def test_processor_override(

148 size_factors: list[int],

149 min_dynamic_patch: int,

150 max_dynamic_patch: int,

152 - dynamic_image_size: Optional[bool],

151 + dynamic_image_size: bool | None,

152 kwargs_on_init: bool,

153 ):

154 mm_processor_kwargs = {

tests/models/multimodal/processing/test_internvl.py

+1 -2

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

3 """Tests for InternVL's multimodal preprocessing kwargs."""

4

5 from collections.abc import Mapping

6 -from typing import Optional

6

7 import pytest

8 from PIL import Image

@@ -103,7 +102,7 @@ def test_processor_override(

102 size_factors: list[int],

103 min_dynamic_patch: int,

104 max_dynamic_patch: int,

106 - dynamic_image_size: Optional[bool],

105 + dynamic_image_size: bool | None,

106 kwargs_on_init: bool,

107 ):

108 mm_processor_kwargs = {

tests/models/multimodal/processing/test_nemotron_vl.py

+1 -2

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

3 """Tests for Nemotron-Nano-VL's multimodal preprocessing kwargs."""

4

5 from collections.abc import Mapping

6 -from typing import Optional

6

7 import pytest

8 from PIL import Image

@@ -105,7 +104,7 @@ def test_processor_override(

104 size_factors: list[int],

105 min_dynamic_patch: int,

106 max_dynamic_patch: int,

108 - dynamic_image_size: Optional[bool],

107 + dynamic_image_size: bool | None,

108 kwargs_on_init: bool,

109 ):

110 mm_processor_kwargs = {

tests/models/multimodal/processing/test_tensor_schema.py

+8 -8

@@ -4,7 +4,7 @@ import tempfile

4 from collections.abc import Iterable

5 from contextlib import contextmanager

6 from functools import partial

7 -from typing import Any, Union

7 +from typing import Any, TypeAlias

8

9 import numpy as np

10 import pytest

@@ -55,15 +55,15 @@ REPO_ID_TO_SKIP = {

55 }

56

57 ImageInput = list[Image.Image]

58 -VideoInput = Union[

59 - list[Image.Image], list[np.ndarray], list[tuple[np.ndarray, dict[str, Any]]]

60 -]

58 +VideoInput: TypeAlias = (

59 + list[Image.Image] | list[np.ndarray] | list[tuple[np.ndarray, dict[str, Any]]]

60 +)

61 AudioInput = list[tuple[np.ndarray, int]]

62

63

64 def _resize_data(

65 - _data: Union[Image.Image, np.ndarray], size_factor: float

66 -) -> Union[Image.Image, np.ndarray]:

65 + _data: Image.Image | np.ndarray, size_factor: float

66 +) -> Image.Image | np.ndarray:

67 assert size_factor <= 1, "Size factor must be less than 1"

68 # Image input

69 if isinstance(_data, Image.Image):

@@ -88,8 +88,8 @@ def _resize_data(

88

89

90 def resize_mm_data(

91 - data: Union[ImageInput, VideoInput, AudioInput], size_factors: tuple[float, ...]

92 -) -> Union[ImageInput, VideoInput, AudioInput]:

91 + data: ImageInput | VideoInput | AudioInput, size_factors: tuple[float, ...]

92 +) -> ImageInput | VideoInput | AudioInput:

93 size_factors = size_factors[: len(data)]

94 if is_list_of(data, (Image.Image, np.ndarray, list)):

95 return [_resize_data(d, s) for d, s in zip(data, size_factors)]

tests/models/quantization/test_awq.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import pytest

6 import torch

@@ -30,7 +29,7 @@ def run_awq_test(

29 max_tokens: int,

30 num_logprobs: int,

31 tensor_parallel_size: int,

33 - distributed_executor_backend: Optional[str] = None,

32 + distributed_executor_backend: str | None = None,

33 ):

34 images = [asset.pil_image for asset in image_assets]

35

tests/models/registry.py

+10 -10

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

3

4 from collections.abc import Mapping, Set

5 from dataclasses import dataclass, field

6 -from typing import Any, Literal, Optional

6 +from typing import Any, Literal

7

8 import pytest

9 import torch

@@ -21,29 +21,29 @@ class _HfExamplesInfo:

21 extras: Mapping[str, str] = field(default_factory=dict)

22 """Extra models to use for testing this architecture."""

23

24 - tokenizer: Optional[str] = None

24 + tokenizer: str | None = None

25 """Set the tokenizer to load for this architecture."""

26

27 tokenizer_mode: TokenizerMode = "auto"

28 """Set the tokenizer type for this architecture."""

29

30 - speculative_model: Optional[str] = None

30 + speculative_model: str | None = None

31 """

32 The default model to use for testing this architecture, which is only used

33 for speculative decoding.

34 """

35

36 - min_transformers_version: Optional[str] = None

36 + min_transformers_version: str | None = None

37 """

38 The minimum version of HF Transformers that is required to run this model.

39 """

40

41 - max_transformers_version: Optional[str] = None

41 + max_transformers_version: str | None = None

42 """

43 The maximum version of HF Transformers that this model runs on.

44 """

45

46 - transformers_version_reason: Optional[str] = None

46 + transformers_version_reason: str | None = None

47 """

48 The reason for the minimum/maximum version requirement.

49 """

@@ -82,19 +82,19 @@ class _HfExamplesInfo:

82 hf_overrides: dict[str, Any] = field(default_factory=dict)

83 """The ``hf_overrides`` required to load the model."""

84

85 - max_model_len: Optional[int] = None

85 + max_model_len: int | None = None

86 """

87 The maximum model length to use for this model. Some models default to a

88 length that is too large to fit into memory in CI.

89 """

90

91 - revision: Optional[str] = None

91 + revision: str | None = None

92 """

93 The specific revision (commit hash, tag, or branch) to use for the model.

94 If not specified, the default revision will be used.

95 """

96

97 - max_num_seqs: Optional[int] = None

97 + max_num_seqs: int | None = None

98 """Maximum number of sequences to be processed in a single iteration."""

99

100 use_original_num_layers: bool = False

@@ -109,7 +109,7 @@ class _HfExamplesInfo:

109 on_fail: Literal["error", "skip", "return"],

110 check_min_version: bool = True,

111 check_max_version: bool = True,

112 - ) -> Optional[str]:

112 + ) -> str | None:

113 """

114 If the installed transformers version does not meet the requirements,

115 perform the given action.

tests/models/test_transformers.py

+4 -4

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

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

3 """Test the functionality of the Transformers backend."""

4

5 -from typing import Any, Optional, Union

5 +from typing import Any

6

7 import pytest

8

@@ -21,12 +21,12 @@ def get_model(arch: str) -> str:

21

22

23 def check_implementation(

24 - runner_ref: type[Union[HfRunner, VllmRunner]],

24 + runner_ref: type[HfRunner | VllmRunner],

25 runner_test: type[VllmRunner],

26 example_prompts: list[str],

27 model: str,

28 - kwargs_ref: Optional[dict[str, Any]] = None,

29 - kwargs_test: Optional[dict[str, Any]] = None,

28 + kwargs_ref: dict[str, Any] | None = None,

29 + kwargs_test: dict[str, Any] | None = None,

30 **kwargs,

31 ):

32 if kwargs_ref is None:

tests/models/utils.py

+19 -19

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

4 import warnings

5 from collections.abc import Sequence

6 from dataclasses import dataclass

7 -from typing import Any, Optional, Union

7 +from typing import Any

8

9 import torch

10 import torch.nn.functional as F

@@ -57,7 +57,7 @@ def check_outputs_equal(

57 #

58 # Assumes prompt logprobs were not requested.

59 TokensTextLogprobs = tuple[

60 - list[int], str, Optional[Union[list[dict[int, float]], SampleLogprobs]]

60 + list[int], str, list[dict[int, float]] | SampleLogprobs | None

61 ]

62

63 # Allow for tokens to be represented as str's rather than IDs;

@@ -68,7 +68,7 @@ TokensTextLogprobs = tuple[

68 #

69 # Assumes prompt logprobs were not requested.

70 TextTextLogprobs = tuple[

71 - list[str], str, Optional[Union[list[dict[str, float]], list[dict[str, Logprob]]]]

71 + list[str], str, list[dict[str, float]] | list[dict[str, Logprob]] | None

72 ]

73

74 # Representation of generated sequence as a tuple of

@@ -81,18 +81,18 @@ TextTextLogprobs = tuple[

81 TokensTextLogprobsPromptLogprobs = tuple[

82 list[int],

83 str,

84 - Optional[Union[list[dict[int, float]], SampleLogprobs]],

85 - Optional[Union[list[Optional[dict[int, float]]], PromptLogprobs]],

84 + list[dict[int, float]] | SampleLogprobs | None,

85 + list[dict[int, float] | None] | PromptLogprobs | None,

86 ]

87

88

89 def check_logprobs_close(

90 *,

91 outputs_0_lst: Sequence[

92 - Union[TokensTextLogprobs, TokensTextLogprobsPromptLogprobs, TextTextLogprobs]

92 + TokensTextLogprobs | TokensTextLogprobsPromptLogprobs | TextTextLogprobs

93 ],

94 outputs_1_lst: Sequence[

95 - Union[TokensTextLogprobs, TokensTextLogprobsPromptLogprobs, TextTextLogprobs]

95 + TokensTextLogprobs | TokensTextLogprobsPromptLogprobs | TextTextLogprobs

96 ],

97 name_0: str,

98 name_1: str,

@@ -273,9 +273,9 @@ def build_model_context(

273 model_id: str,

274 runner: RunnerOption = "auto",

275 dtype: ModelDType = "auto",

276 - model_config_kwargs: Optional[dict[str, Any]] = None,

277 - mm_processor_kwargs: Optional[dict[str, Any]] = None,

278 - limit_mm_per_prompt: Optional[dict[str, int]] = None,

276 + model_config_kwargs: dict[str, Any] | None = None,

277 + mm_processor_kwargs: dict[str, Any] | None = None,

278 + limit_mm_per_prompt: dict[str, int] | None = None,

279 mm_processor_cache_gb: int = 0,

280 ):

281 """Creates an InputProcessingContext for a given model.

@@ -369,18 +369,18 @@ class ModelInfo:

369 name: str

370 architecture: str = ""

371 dtype: str = "auto"

372 - max_model_len: Optional[int] = None

372 + max_model_len: int | None = None

373 hf_dtype: str = "float32"

374 - hf_overrides: Optional[dict[str, Any]] = None

374 + hf_overrides: dict[str, Any] | None = None

375 default_pooling_type: str = ""

376 enable_test: bool = True

377

378

379 @dataclass

380 class EmbedModelInfo(ModelInfo):

381 - mteb_score: Optional[float] = None

381 + mteb_score: float | None = None

382 is_matryoshka: bool = False

383 - matryoshka_dimensions: Optional[list[int]] = None

383 + matryoshka_dimensions: list[int] | None = None

384

385

386 @dataclass

@@ -395,7 +395,7 @@ class LASTPoolingEmbedModelInfo(EmbedModelInfo):

395

396 @dataclass

397 class RerankModelInfo(ModelInfo):

398 - mteb_score: Optional[float] = None

398 + mteb_score: float | None = None

399

400

401 @dataclass

@@ -411,14 +411,14 @@ class LASTPoolingRerankModelInfo(RerankModelInfo):

411 @dataclass

412 class GenerateModelInfo(ModelInfo):

413 hf_dtype: str = "auto"

414 - hf_ppl: Optional[float] = None

414 + hf_ppl: float | None = None

415

416

417 def dummy_hf_overrides(

418 hf_config: PretrainedConfig,

419 *,

420 model_arch: str = "",

421 - exist_overrides: Optional[dict[str, Any]] = None,

421 + exist_overrides: dict[str, Any] | None = None,

422 use_original_num_layers: bool = False,

423 ) -> PretrainedConfig:

424 """

@@ -507,8 +507,8 @@ def dummy_hf_overrides(

507

508 def check_transformers_version(

509 model: str,

510 - min_transformers_version: Optional[str] = None,

511 - max_transformers_version: Optional[str] = None,

510 + min_transformers_version: str | None = None,

511 + max_transformers_version: str | None = None,

512 ):

513 from .registry import _HfExamplesInfo

514

tests/multimodal/test_cache.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import numpy as np

5 import pytest

@@ -32,7 +31,7 @@ def _dummy_elem(

31 key: str,

32 size: int,

33 *,

35 - rng: Optional[np.random.RandomState] = None,

34 + rng: np.random.RandomState | None = None,

35 ):

36 if rng is None:

37 data = torch.empty((size,), dtype=torch.int8)

@@ -51,7 +50,7 @@ def _dummy_item(

50 modality: str,

51 size_by_key: dict[str, int],

52 *,

54 - rng: Optional[np.random.RandomState] = None,

53 + rng: np.random.RandomState | None = None,

54 ):

55 return MultiModalKwargsItem.from_elems(

56 [_dummy_elem(modality, key, size, rng=rng) for key, size in size_by_key.items()]

@@ -61,7 +60,7 @@ def _dummy_item(

60 def _dummy_items(

61 size_by_key_modality: dict[str, dict[str, int]],

62 *,

64 - rng: Optional[np.random.RandomState] = None,

63 + rng: np.random.RandomState | None = None,

64 ):

65 return MultiModalKwargsItems.from_seq(

66 [

tests/multimodal/test_processing.py

+2 -2

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

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

3

4 from contextlib import nullcontext

5 -from typing import Optional, cast

5 +from typing import cast

6

7 import numpy as np

8 import pytest

@@ -1003,7 +1003,7 @@ class DummyProcessor:

1003 self,

1004 a: int = 0,

1005 c: int = 0,

1006 - return_tensors: Optional[str] = None,

1006 + return_tensors: str | None = None,

1007 ) -> dict[str, int]:

1008 return dict(a=a, c=c)

1009

tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py

+4 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import base64

5 import datetime

6 import os

7 import tempfile

8 import urllib.request

9 from collections.abc import Sequence

12 -from typing import Any, Union

10 +from typing import Any

11

12 import albumentations

13 import numpy as np

@@ -160,11 +158,11 @@ def read_geotiff(

158

159

160 def load_image(

163 - data: Union[list[str]],

161 + data: list[str],

162 path_type: str,

163 mean: list[float] | None = None,

164 std: list[float] | None = None,

167 - indices: Union[list[int], None] | None = None,

165 + indices: list[int] | None | None = None,

166 ):

167 """Build an input example by loading images in *file_paths*.

168

@@ -280,7 +278,7 @@ class PrithviMultimodalDataProcessor(IOProcessor):

278 prompt: IOProcessorInput,

279 request_id: str | None = None,

280 **kwargs,

283 - ) -> Union[PromptType, Sequence[PromptType]]:

281 + ) -> PromptType | Sequence[PromptType]:

282 image_data = dict(prompt)

283

284 if request_id:

tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Literal, Optional, TypedDict, Union

4 +from typing import Any, Literal, TypedDict

5

6 import albumentations

7 from pydantic import BaseModel

@@ -38,7 +38,7 @@ class ImagePrompt(BaseModel):

38 """

39

40

41 -MultiModalPromptType = Union[ImagePrompt]

41 +MultiModalPromptType = ImagePrompt

42

43

44 class ImageRequestOutput(BaseModel):

@@ -54,4 +54,4 @@ class ImageRequestOutput(BaseModel):

54 type: Literal["path", "b64_json"]

55 format: str

56 data: str

57 - request_id: Optional[str] = None

57 + request_id: str | None = None

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py

+3 -4

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

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

3

4 from collections.abc import Iterable

5 -from typing import Optional, Union

5

6 import torch

7 import torch.nn as nn

@@ -44,9 +43,9 @@ class MyGemma2Embedding(nn.Module):

43 self,

44 input_ids: torch.Tensor,

45 positions: torch.Tensor,

47 - intermediate_tensors: Optional[IntermediateTensors] = None,

48 - inputs_embeds: Optional[torch.Tensor] = None,

49 - ) -> Union[torch.Tensor, IntermediateTensors]:

46 + intermediate_tensors: IntermediateTensors | None = None,

47 + inputs_embeds: torch.Tensor | None = None,

48 + ) -> torch.Tensor | IntermediateTensors:

49 hidden_states = self.model(

50 input_ids,

51 positions,

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_llava.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6

@@ -20,7 +19,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY

19 dummy_inputs=LlavaDummyInputsBuilder,

20 )

21 class MyLlava(LlavaForConditionalGeneration):

23 - def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:

22 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

23 # this dummy model always predicts the first token

24 logits = super().compute_logits(hidden_states)

25 if logits is not None:

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6

@@ -9,7 +8,7 @@ from vllm.model_executor.models.opt import OPTForCausalLM

8

9

10 class MyOPTForCausalLM(OPTForCausalLM):

12 - def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:

11 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

12 # this dummy model always predicts the first token

13 logits = super().compute_logits(hidden_states)

14 if logits is not None:

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py

+1 -3

@@ -1,10 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

6 -

7 -def dummy_platform_plugin() -> Optional[str]:

5 +def dummy_platform_plugin() -> str | None:

6 return "vllm_add_dummy_platform.dummy_platform.DummyPlatform"

7

8

tests/quantization/test_blackwell_moe.py

+1 -2

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

3

4 import json

5 import os

6 -from typing import Optional

6

7 import pytest

8

@@ -30,7 +29,7 @@ def set_test_environment():

29 dummy_hf_overrides = {"num_layers": 4, "num_hidden_layers": 4}

30

31

33 -def can_initialize(model: str, extra_args: Optional[list[str]] = None):

32 +def can_initialize(model: str, extra_args: list[str] | None = None):

33 # Server arguments

34 extra_args = extra_args if extra_args is not None else []

35 server_args = [

tests/quantization/test_compressed_tensors.py

+1 -3

@@ -5,8 +5,6 @@

5 Run `pytest tests/quantization/test_compressed_tensors.py`.

6 """

7

8 -from typing import Optional

9 -

8 import pytest

9 import torch

10 from compressed_tensors.quantization import QuantizationType

@@ -104,7 +102,7 @@ def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):

102 down_proj = layer.mlp.down_proj

103

104 # assert zp for symmetric and asymmetric cases

107 - def zp_valid(zp: Optional[torch.Tensor]):

105 + def zp_valid(zp: torch.Tensor | None):

106 if is_symmetric:

107 return zp is None

108

tests/quantization/test_quark.py

+2 -3

@@ -11,7 +11,6 @@ import importlib.metadata

11 import os

12 from dataclasses import dataclass

13 from importlib.util import find_spec

14 -from typing import Optional

14

15 import huggingface_hub

16 import lm_eval

@@ -156,8 +155,8 @@ class AccuracyTestConfig:

155 def get_model_args(

156 self,

157 tp_size: int,

159 - model_max_len: Optional[int] = None,

160 - kwargs: Optional[dict] = None,

158 + model_max_len: int | None = None,

159 + kwargs: dict | None = None,

160 ) -> dict:

161 if kwargs is None:

162 kwargs = {}

tests/quantization/test_register_quantization_config.py

+8 -8

@@ -7,7 +7,7 @@ See https://github.com/vllm-project/vllm/issues/11926 for more details.

7 Run `pytest tests/quantization/test_register_quantization_config.py`.

8 """

9

10 -from typing import Any, Optional

10 +from typing import Any

11

12 import pytest

13 import torch

@@ -37,10 +37,10 @@ class FakeQuantLinearMethod(UnquantizedLinearMethod):

37

38 def apply(

39 self,

40 - layer: "torch.nn.Module",

41 - x: "torch.Tensor",

42 - bias: Optional["torch.Tensor"] = None,

43 - ) -> "torch.Tensor":

40 + layer: torch.nn.Module,

41 + x: torch.Tensor,

42 + bias: torch.Tensor | None = None,

43 + ) -> torch.Tensor:

44 """Perform fake quantization before the linear layer."""

45

46 # Calculate the scales dynamically

@@ -72,7 +72,7 @@ class CustomQuantConfig(QuantizationConfig):

72 """Name of the quantization method."""

73 return "custom_quant"

74

75 - def get_supported_act_dtypes(self) -> list["torch.dtype"]:

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

76 """List of supported activation dtypes."""

77 return [torch.float16, torch.bfloat16]

78

@@ -92,8 +92,8 @@ class CustomQuantConfig(QuantizationConfig):

92 return CustomQuantConfig(num_bits=config.get("num_bits", 8))

93

94 def get_quant_method(

95 - self, layer: "torch.nn.Module", prefix: str

96 - ) -> Optional["FakeQuantLinearMethod"]:

95 + self, layer: torch.nn.Module, prefix: str

96 + ) -> FakeQuantLinearMethod | None:

97 """Get the quantize method to use for the quantized layer."""

98 if isinstance(layer, LinearBase):

99 return FakeQuantLinearMethod(num_bits=self.num_bits)

tests/reasoning/utils.py

+8 -9

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 from vllm.entrypoints.openai.protocol import ChatCompletionRequest, DeltaMessage

6 from vllm.reasoning import ReasoningParser

@@ -34,9 +33,9 @@ class StreamingReasoningReconstructor:

33 def run_reasoning_extraction(

34 reasoning_parser: ReasoningParser,

35 model_output: list[str],

37 - request: Union[ChatCompletionRequest, None] = None,

36 + request: ChatCompletionRequest | None = None,

37 streaming: bool = False,

39 -) -> tuple[Optional[str], Optional[str]]:

38 +) -> tuple[str | None, str | None]:

39 if streaming:

40 reconstructor = run_reasoning_extraction_streaming(

41 reasoning_parser,

@@ -57,9 +56,9 @@ def run_reasoning_extraction(

56 def run_reasoning_extraction_mistral(

57 reasoning_parser: ReasoningParser,

58 model_output: list[int],

60 - request: Union[ChatCompletionRequest, None] = None,

59 + request: ChatCompletionRequest | None = None,

60 streaming: bool = False,

62 -) -> tuple[Optional[str], Optional[str]]:

61 +) -> tuple[str | None, str | None]:

62 assert isinstance(reasoning_parser.model_tokenizer, MistralTokenizer), type(

63 reasoning_parser.model_tokenizer

64 )

@@ -86,8 +85,8 @@ def run_reasoning_extraction_mistral(

85 def run_reasoning_extraction_nonstreaming(

86 reasoning_parser: ReasoningParser,

87 model_output: list[str],

89 - request: Union[ChatCompletionRequest, None] = None,

90 -) -> tuple[Optional[str], Optional[str]]:

88 + request: ChatCompletionRequest | None = None,

89 +) -> tuple[str | None, str | None]:

90 request = request or ChatCompletionRequest(messages=[], model="test-model")

91 return reasoning_parser.extract_reasoning_content(

92 model_output="".join(model_output), request=request

@@ -97,7 +96,7 @@ def run_reasoning_extraction_nonstreaming(

96 def run_reasoning_extraction_streaming(

97 reasoning_parser: ReasoningParser,

98 model_deltas: list[str],

100 - request: Union[ChatCompletionRequest, None] = None,

99 + request: ChatCompletionRequest | None = None,

100 ) -> StreamingReasoningReconstructor:

101 request = request or ChatCompletionRequest(messages=[], model="test-model")

102 reconstructor = StreamingReasoningReconstructor()

@@ -129,7 +128,7 @@ def run_reasoning_extraction_streaming(

128 def run_reasoning_extraction_streaming_mistral(

129 reasoning_parser: ReasoningParser,

130 model_deltas: list[int],

132 - request: Union[ChatCompletionRequest, None] = None,

131 + request: ChatCompletionRequest | None = None,

132 ) -> StreamingReasoningReconstructor:

133 assert isinstance(reasoning_parser.model_tokenizer, MistralTokenizer), type(

134 reasoning_parser.model_tokenizer

tests/samplers/test_no_bad_words.py

+3 -5

@@ -6,8 +6,6 @@ Run `pytest tests/samplers/test_no_bad_words.py`.

6

7 """

8

9 -from typing import Optional

10 -

9 from transformers import AutoTokenizer

10

11 from vllm import LLM, SamplingParams

@@ -18,7 +16,7 @@ def _generate(

16 prompt: str,

17 num_prompt_tokens: int,

18 temperature: float = 0,

21 - bad_words: Optional[list[str]] = None,

19 + bad_words: list[str] | None = None,

20 ) -> list[int]:

21 sampling_params = SamplingParams(

22 temperature=temperature,

@@ -60,7 +58,7 @@ class TestOneTokenBadWord:

58 output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN])

59 assert self.target_token_id not in output_token_ids

60

63 - def _generate(self, llm: LLM, bad_words: Optional[list[str]] = None) -> list[int]:

61 + def _generate(self, llm: LLM, bad_words: list[str] | None = None) -> list[int]:

62 return _generate(

63 llm=llm,

64 prompt=self.PROMPT,

@@ -155,7 +153,7 @@ class TestTwoTokenBadWord:

153 self.neighbour_token_id2 in output_token_ids

154 )

155

158 - def _generate(self, llm: LLM, bad_words: Optional[list[str]] = None) -> list[int]:

156 + def _generate(self, llm: LLM, bad_words: list[str] | None = None) -> list[int]:

157 return _generate(

158 llm=llm,

159 prompt=self.PROMPT,

tests/tokenization/test_detokenize.py

+2 -2

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

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

3

4 from collections.abc import Generator

5 -from typing import Any, Optional

5 +from typing import Any

6

7 import pytest

8 from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast

@@ -52,7 +52,7 @@ def _run_incremental_decode(

52 skip_special_tokens: bool,

53 starting_index: int,

54 spaces_between_special_tokens: bool = True,

55 - fast: Optional[bool] = None,

55 + fast: bool | None = None,

56 ):

57 prompt_token_ids = all_input_ids[:starting_index]

58

tests/tokenization/test_tokenizer_registry.py

+8 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import TYPE_CHECKING, Any, Optional, Union

4 +from typing import TYPE_CHECKING, Any

5

6 from vllm.transformers_utils.tokenizer import get_tokenizer

7 from vllm.transformers_utils.tokenizer_base import TokenizerBase, TokenizerRegistry

@@ -61,11 +61,11 @@ class TestTokenizer(TokenizerBase):

61

62 def __call__(

63 self,

64 - text: Union[str, list[str], list[int]],

65 - text_pair: Optional[str] = None,

64 + text: str | list[str] | list[int],

65 + text_pair: str | None = None,

66 add_special_tokens: bool = False,

67 truncation: bool = False,

68 - max_length: Optional[int] = None,

68 + max_length: int | None = None,

69 ):

70 raise NotImplementedError()

71

@@ -79,17 +79,17 @@ class TestTokenizer(TokenizerBase):

79 self,

80 text: str,

81 truncation: bool = False,

82 - max_length: Optional[int] = None,

82 + max_length: int | None = None,

83 ) -> list[int]:

84 raise NotImplementedError()

85

86 - def encode(self, text: str, add_special_tokens: Optional[bool] = None) -> list[int]:

86 + def encode(self, text: str, add_special_tokens: bool | None = None) -> list[int]:

87 raise NotImplementedError()

88

89 def apply_chat_template(

90 self,

91 messages: list["ChatCompletionMessageParam"],

92 - tools: Optional[list[dict[str, Any]]] = None,

92 + tools: list[dict[str, Any]] | None = None,

93 **kwargs,

94 ) -> list[int]:

95 raise NotImplementedError()

@@ -97,9 +97,7 @@ class TestTokenizer(TokenizerBase):

97 def convert_tokens_to_string(self, tokens: list[str]) -> str:

98 raise NotImplementedError()

99

100 - def decode(

101 - self, ids: Union[list[int], int], skip_special_tokens: bool = True

102 - ) -> str:

100 + def decode(self, ids: list[int] | int, skip_special_tokens: bool = True) -> str:

101 raise NotImplementedError()

102

103 def convert_ids_to_tokens(

tests/tool_use/mistral/utils.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 from typing_extensions import TypedDict

6

@@ -9,9 +8,9 @@ from typing_extensions import TypedDict

8 class ServerConfig(TypedDict, total=False):

9 model: str

10 arguments: list[str]

12 - system_prompt: Optional[str]

13 - supports_parallel: Optional[bool]

14 - supports_rocm: Optional[bool]

11 + system_prompt: str | None

12 + supports_parallel: bool | None

13 + supports_rocm: bool | None

14

15

16 ARGS: list[str] = ["--max-model-len", "1024"]

tests/tool_use/test_jamba_tool_parser.py

+1 -2

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

3

4 import json

5 from collections.abc import Generator

6 -from typing import Optional

6

7 import partial_json_parser

8 import pytest

@@ -248,7 +247,7 @@ def test_extract_tool_calls_streaming(

247 function_names: list[str] = []

248 function_args_strs: list[str] = []

249 tool_call_idx: int = -1

251 - tool_call_ids: list[Optional[str]] = []

250 + tool_call_ids: list[str | None] = []

251

252 for delta_message in stream_delta_message_generator(

253 jamba_tool_parser, jamba_tokenizer, model_output

tests/tool_use/test_parallel_tool_calls.py

+1 -2

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

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

3

4 import json

5 -from typing import Optional

5

6 import openai

7 import pytest

@@ -80,7 +79,7 @@ async def test_parallel_tool_calls(

79 stream=True,

80 )

81

83 - role_name: Optional[str] = None

82 + role_name: str | None = None

83 finish_reason_count: int = 0

84

85 tool_call_names: list[str] = []

tests/tool_use/test_qwen3coder_tool_parser.py

+1 -2

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

3

4 import json

5 from collections.abc import Generator

6 -from typing import Optional

6

7 import pytest

8

@@ -107,7 +106,7 @@ def stream_delta_message_generator(

106 qwen3_tool_parser,

107 qwen3_tokenizer: AnyTokenizer,

108 model_output: str,

110 - request: Optional[ChatCompletionRequest] = None,

109 + request: ChatCompletionRequest | None = None,

110 ) -> Generator[DeltaMessage, None, None]:

111 all_token_ids = qwen3_tokenizer.encode(model_output, add_special_tokens=False)

112

tests/tool_use/test_seed_oss_tool_parser.py

+1 -2

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

4

5 import json

6 from collections.abc import Generator

7 -from typing import Optional

7

8 import pytest

9

@@ -259,7 +258,7 @@ def stream_delta_message_generator(

258 seed_oss_tool_parser: SeedOssToolParser,

259 seed_oss_tokenizer: AnyTokenizer,

260 model_output: str,

262 - request: Optional[ChatCompletionRequest] = None,

261 + request: ChatCompletionRequest | None = None,

262 ) -> Generator[DeltaMessage, None, None]:

263 all_token_ids = seed_oss_tokenizer.encode(model_output, add_special_tokens=False)

264

tests/tool_use/test_tool_calls.py

+3 -4

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

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

3

4 import json

5 -from typing import Optional

5

6 import openai

7 import pytest

@@ -58,10 +57,10 @@ async def test_tool_call_and_choice(client: openai.AsyncOpenAI):

57

58 assert stop_reason == "tool_calls"

59

61 - function_name: Optional[str] = None

60 + function_name: str | None = None

61 function_args_str: str = ""

63 - tool_call_id: Optional[str] = None

64 - role_name: Optional[str] = None

62 + tool_call_id: str | None = None

63 + role_name: str | None = None

64 finish_reason_count: int = 0

65

66 # make the same request, streaming

tests/tool_use/test_xlam_tool_parser.py

+1 -2

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

3

4 import json

5 from collections.abc import Generator

6 -from typing import Optional

6

7 import pytest

8

@@ -52,7 +51,7 @@ def stream_delta_message_generator(

51 xlam_tool_parser: xLAMToolParser,

52 xlam_tokenizer: AnyTokenizer,

53 model_output: str,

55 - request: Optional[ChatCompletionRequest] = None,

54 + request: ChatCompletionRequest | None = None,

55 ) -> Generator[DeltaMessage, None, None]:

56 all_token_ids = xlam_tokenizer.encode(model_output, add_special_tokens=False)

57

tests/tool_use/utils.py

+5 -5

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

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

3

4 from copy import deepcopy

5 -from typing import Any, Optional

5 +from typing import Any

6

7 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam

8 from typing_extensions import TypedDict

@@ -13,10 +13,10 @@ from tests.utils import VLLM_PATH

13 class ServerConfig(TypedDict, total=False):

14 model: str

15 arguments: list[str]

16 - system_prompt: Optional[str]

17 - supports_parallel: Optional[bool]

18 - supports_rocm: Optional[bool]

19 - extended: Optional[bool] # tests do not run in CI automatically

16 + system_prompt: str | None

17 + supports_parallel: bool | None

18 + supports_rocm: bool | None

19 + extended: bool | None # tests do not run in CI automatically

20

21

22 def patch_system_prompt(

tests/transformers_utils/test_config_parser_registry.py

+3 -4

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

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

3

4 from pathlib import Path

5 -from typing import Optional, Union

5

6 import pytest

7 from transformers import PretrainedConfig

@@ -15,10 +14,10 @@ from vllm.transformers_utils.config_parser_base import ConfigParserBase

14 class CustomConfigParser(ConfigParserBase):

15 def parse(

16 self,

18 - model: Union[str, Path],

17 + model: str | Path,

18 trust_remote_code: bool,

20 - revision: Optional[str] = None,

21 - code_revision: Optional[str] = None,

19 + revision: str | None = None,

20 + code_revision: str | None = None,

21 **kwargs,

22 ) -> tuple[dict, PretrainedConfig]:

23 raise NotImplementedError

tests/utils.py

+24 -23

@@ -15,10 +15,11 @@ import sys

15 import tempfile

16 import time

17 import warnings

18 +from collections.abc import Callable

19 from contextlib import ExitStack, contextmanager, suppress

20 from multiprocessing import Process

21 from pathlib import Path

21 -from typing import Any, Callable, Literal, Optional, Union

22 +from typing import Any, Literal

23 from unittest.mock import patch

24

25 import cloudpickle

@@ -94,7 +95,7 @@ class RemoteOpenAIServer:

95 DUMMY_API_KEY = "token-abc123" # vLLM's OpenAI server does not need API key

96

97 def _start_server(

97 - self, model: str, vllm_serve_args: list[str], env_dict: Optional[dict[str, str]]

98 + self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None

99 ) -> None:

100 """Subclasses override this method to customize server process launch"""

101 env = os.environ.copy()

@@ -117,11 +118,11 @@ class RemoteOpenAIServer:

118 model: str,

119 vllm_serve_args: list[str],

120 *,

120 - env_dict: Optional[dict[str, str]] = None,

121 - seed: Optional[int] = 0,

121 + env_dict: dict[str, str] | None = None,

122 + seed: int | None = 0,

123 auto_port: bool = True,

123 - max_wait_seconds: Optional[float] = None,

124 - override_hf_configs: Optional[dict[str, Any]] = None,

124 + max_wait_seconds: float | None = None,

125 + override_hf_configs: dict[str, Any] | None = None,

126 ) -> None:

127 if auto_port:

128 if "-p" in vllm_serve_args or "--port" in vllm_serve_args:

@@ -186,7 +187,7 @@ class RemoteOpenAIServer:

187 # force kill if needed

188 self.proc.kill()

189

189 - def _poll(self) -> Optional[int]:

190 + def _poll(self) -> int | None:

191 """Subclasses override this method to customize process polling"""

192 return self.proc.poll()

193

@@ -251,7 +252,7 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):

252 """Launch test server with custom child process"""

253

254 def _start_server(

254 - self, model: str, vllm_serve_args: list[str], env_dict: Optional[dict[str, str]]

255 + self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None

256 ) -> None:

257 self.proc: Process = Process(

258 target=self.child_process_fxn, args=(env_dict, model, vllm_serve_args)

@@ -262,12 +263,12 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):

263 self,

264 model: str,

265 vllm_serve_args: list[str],

265 - child_process_fxn: Callable[[Optional[dict[str, str]], str, list[str]], None],

266 + child_process_fxn: Callable[[dict[str, str] | None, str, list[str]], None],

267 *,

267 - env_dict: Optional[dict[str, str]] = None,

268 - seed: Optional[int] = 0,

268 + env_dict: dict[str, str] | None = None,

269 + seed: int | None = 0,

270 auto_port: bool = True,

270 - max_wait_seconds: Optional[float] = None,

271 + max_wait_seconds: float | None = None,

272 ) -> None:

273 """Store custom child process function then invoke superclass

274 constructor which will indirectly launch it."""

@@ -281,7 +282,7 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):

282 max_wait_seconds=max_wait_seconds,

283 )

284

284 - def _poll(self) -> Optional[int]:

285 + def _poll(self) -> int | None:

286 return self.proc.exitcode

287

288 def __exit__(self, exc_type, exc_value, traceback):

@@ -547,11 +548,11 @@ def compare_two_settings(

548 model: str,

549 arg1: list[str],

550 arg2: list[str],

550 - env1: Optional[dict[str, str]] = None,

551 - env2: Optional[dict[str, str]] = None,

551 + env1: dict[str, str] | None = None,

552 + env2: dict[str, str] | None = None,

553 *,

554 method: str = "generate",

554 - max_wait_seconds: Optional[float] = None,

555 + max_wait_seconds: float | None = None,

556 ) -> None:

557 """

558 Launch API server with two different sets of arguments/environments

@@ -577,10 +578,10 @@ def compare_two_settings(

578 def compare_all_settings(

579 model: str,

580 all_args: list[list[str]],

580 - all_envs: list[Optional[dict[str, str]]],

581 + all_envs: list[dict[str, str] | None],

582 *,

583 method: str = "generate",

583 - max_wait_seconds: Optional[float] = None,

584 + max_wait_seconds: float | None = None,

585 ) -> None:

586 """

587 Launch API server with several different sets of arguments/environments

@@ -785,8 +786,8 @@ def get_physical_device_indices(devices):

786 def wait_for_gpu_memory_to_clear(

787 *,

788 devices: list[int],

788 - threshold_bytes: Optional[int] = None,

789 - threshold_ratio: Optional[float] = None,

789 + threshold_bytes: int | None = None,

790 + threshold_ratio: float | None = None,

791 timeout_s: float = 120,

792 ) -> None:

793 assert threshold_bytes is not None or threshold_ratio is not None

@@ -1002,7 +1003,7 @@ def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]

1003

1004

1005 def create_new_process_for_each_test(

1005 - method: Optional[Literal["spawn", "fork"]] = None,

1006 + method: Literal["spawn", "fork"] | None = None,

1007 ) -> Callable[[Callable[_P, None]], Callable[_P, None]]:

1008 """Creates a decorator that runs each test function in a new process.

1009

@@ -1098,9 +1099,9 @@ async def completions_with_server_args(

1099 prompts: list[str],

1100 model_name: str,

1101 server_cli_args: list[str],

1101 - num_logprobs: Optional[int],

1102 + num_logprobs: int | None,

1103 max_wait_seconds: int = 240,

1103 - max_tokens: Union[int, list] = 5,

1104 + max_tokens: int | list = 5,

1105 ) -> list[Completion]:

1106 """Construct a remote OpenAI server, obtain an async client to the

1107 server & invoke the completions API to obtain completions.

tests/v1/attention/test_attention_backends.py

+2 -3

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

3 """Tests for v1 attention backends without GPUModelRunner dependency."""

4

5 from functools import partial

6 -from typing import Optional, Union

6

7 import pytest

8 import torch

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

201 key: torch.Tensor,

202 value: torch.Tensor,

203 kv_cache: torch.Tensor,

205 - sliding_window: Optional[int] = None,

204 + sliding_window: int | None = None,

205 ) -> torch.Tensor:

206 """Run attention computation using the specified backend's AttentionImpl."""

207

@@ -289,7 +288,7 @@ def run_attention_backend(

288 def _test_backend_correctness(

289 batch_spec: BatchSpec,

290 model: str,

292 - backend_to_test: list[Union[_Backend, str]],

291 + backend_to_test: list[_Backend | str],

292 mask_mod,

293 *,

294 block_size: int = 16,

tests/v1/attention/test_mla_backends.py

+2 -4

@@ -2,8 +2,6 @@

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

3 """Tests for v1 MLA backends without GPUModelRunner dependency."""

4

5 -from typing import Optional, Union

6 -

5 import pytest

6 import torch

7

@@ -81,8 +79,8 @@ def create_and_prepopulate_kv_cache(

79 num_blocks: int,

80 common_attn_metadata: CommonAttentionMetadata,

81 randomize_blocks: bool = True,

84 - kv_cache_dtype: Optional[str] = None,

85 - scale: Union[float, torch.Tensor] = 1.0,

82 + kv_cache_dtype: str | None = None,

83 + scale: float | torch.Tensor = 1.0,

84 ) -> torch.Tensor:

85 """Create and prepopulate an MLA KV cache with context data.

86

tests/v1/attention/utils.py

+3 -4

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

3 """Utility functions for attention-related v1 tests."""

4

5 from dataclasses import dataclass

6 -from typing import Optional, Union

6

7 import pytest

8 import torch

@@ -150,14 +149,14 @@ def create_vllm_config(

149 model_name: str = "meta-llama/Meta-Llama-3-8B",

150 tensor_parallel_size: int = 1,

151 max_model_len: int = 1024,

153 - dtype: Union[ModelDType, torch.dtype] = "auto",

152 + dtype: ModelDType | torch.dtype = "auto",

153 num_gpu_blocks: int = 1000,

154 block_size: int = 16,

155 max_num_seqs: int = 256,

156 max_num_batched_tokens: int = 8192,

157 enable_chunked_prefill: bool = True,

158 add_mock_model_methods: bool = True,

160 - hf_config_override: Optional[dict] = None,

159 + hf_config_override: dict | None = None,

160 ) -> VllmConfig:

161 """Create a VllmConfig for testing with reasonable defaults."""

162

@@ -252,7 +251,7 @@ class BackendConfig:

251 name: str

252 env_vars: dict

253 comp_config: dict # compilation config

255 - specific_gpu_arch: Optional[tuple] = None

254 + specific_gpu_arch: tuple | None = None

255

256

257 # Define all backend configurations of full cudagraph to be tested

tests/v1/core/test_kv_cache_utils.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import importlib

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import pytest

7 import torch

@@ -62,9 +62,9 @@ def make_request(

62 prompt_token_ids: list[int],

63 block_size: int = 3,

64 hash_fn: Callable = hash,

65 - mm_positions: Optional[list[PlaceholderRange]] = None,

66 - mm_hashes: Optional[list[str]] = None,

67 - cache_salt: Optional[str] = None,

65 + mm_positions: list[PlaceholderRange] | None = None,

66 + mm_hashes: list[str] | None = None,

67 + cache_salt: str | None = None,

68 ):

69 mm_features = []

70 if mm_positions is not None:

tests/v1/core/test_prefix_caching.py

+5 -5

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

3 """Compare the with and without prefix caching."""

4

5 import copy

6 -from typing import Callable, Optional

6 +from collections.abc import Callable

7

8 import pytest

9 import torch

@@ -55,10 +55,10 @@ def make_request(

55 prompt_token_ids: list[int],

56 block_size: int,

57 hash_fn: Callable,

58 - mm_positions: Optional[list[PlaceholderRange]] = None,

59 - mm_hashes: Optional[list[str]] = None,

60 - prompt_logprobs: Optional[int] = None,

61 - cache_salt: Optional[str] = None,

58 + mm_positions: list[PlaceholderRange] | None = None,

59 + mm_hashes: list[str] | None = None,

60 + prompt_logprobs: int | None = None,

61 + cache_salt: str | None = None,

62 ):

63 mm_features = []

64 if mm_positions is not None:

tests/v1/core/test_scheduler.py

+10 -13

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import dataclasses

4 -from typing import Optional

4 from unittest.mock import Mock

5

6 import pytest

@@ -78,9 +77,7 @@ def test_get_num_unfinished_requests():

77 (True, 5),

78 ],

79 )

81 -def test_schedule(

82 - enable_prefix_caching: Optional[bool], prompt_logprobs: Optional[int]

83 -):

80 +def test_schedule(enable_prefix_caching: bool | None, prompt_logprobs: int | None):

81 """Test scheduling.

82 Two cases: default APC/no prompt logprobs; APC=True + prompt logprobs

83 """

@@ -595,7 +592,7 @@ def test_check_stop_min_tokens():

592 ],

593 )

594 def test_schedule_concurrent_batches(

598 - enable_prefix_caching: Optional[bool], prompt_logprobs: Optional[int]

595 + enable_prefix_caching: bool | None, prompt_logprobs: int | None

596 ):

597 scheduler = create_scheduler(

598 max_num_batched_tokens=1024,

@@ -1323,14 +1320,14 @@ def create_scheduler_with_priority(

1320 model: str = "facebook/opt-125m",

1321 max_num_seqs: int = 16,

1322 max_num_batched_tokens: int = 8192,

1326 - enable_prefix_caching: Optional[bool] = None,

1323 + enable_prefix_caching: bool | None = None,

1324 long_prefill_token_threshold: int = 0,

1325 disable_chunked_mm_input: bool = False,

1326 use_kv_connector: bool = False,

1327 num_blocks: int = 10000,

1328 block_size: int = 16,

1332 - max_model_len: Optional[int] = None,

1333 - num_speculative_tokens: Optional[int] = None,

1329 + max_model_len: int | None = None,

1330 + num_speculative_tokens: int | None = None,

1331 ) -> Scheduler:

1332 """Create scheduler with priority policy enabled.

1333

@@ -1385,7 +1382,7 @@ def create_scheduler_with_priority(

1382 else None

1383 )

1384

1388 - speculative_config: Optional[SpeculativeConfig] = None

1385 + speculative_config: SpeculativeConfig | None = None

1386 if num_speculative_tokens is not None:

1387 speculative_config = SpeculativeConfig(

1388 model="ngram", num_speculative_tokens=num_speculative_tokens

@@ -1420,12 +1417,12 @@ def create_scheduler_with_priority(

1417 def create_requests_with_priority(

1418 num_requests: int,

1419 priorities: list[int],

1423 - arrival_times: Optional[list[float]] = None,

1420 + arrival_times: list[float] | None = None,

1421 num_tokens: int = 10,

1425 - mm_positions: Optional[list[list[PlaceholderRange]]] = None,

1422 + mm_positions: list[list[PlaceholderRange]] | None = None,

1423 max_tokens: int = 16,

1427 - stop_token_ids: Optional[list[int]] = None,

1428 - prompt_logprobs: Optional[int] = None,

1424 + stop_token_ids: list[int] | None = None,

1425 + prompt_logprobs: int | None = None,

1426 starting_idx: int = 0,

1427 ):

1428 """Create requests with specified priorities and arrival times."""

tests/v1/core/utils.py

+8 -9

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional, Union

3

4 import torch

5

@@ -37,17 +36,17 @@ def create_scheduler(

36 model: str = "facebook/opt-125m",

37 max_num_seqs: int = 16,

38 max_num_batched_tokens: int = 8192,

40 - enable_prefix_caching: Optional[bool] = None,

39 + enable_prefix_caching: bool | None = None,

40 long_prefill_token_threshold: int = 0,

41 disable_chunked_mm_input: bool = False,

42 use_kv_connector: bool = False,

43 num_blocks: int = 10000,

44 block_size: int = 16,

46 - max_model_len: Optional[int] = None,

47 - num_speculative_tokens: Optional[int] = None,

45 + max_model_len: int | None = None,

46 + num_speculative_tokens: int | None = None,

47 skip_tokenizer_init: bool = False,

48 async_scheduling: bool = False,

50 -) -> Union[Scheduler, AsyncScheduler]:

49 +) -> Scheduler | AsyncScheduler:

50 """Create scheduler under test.

51

52 Args:

@@ -102,7 +101,7 @@ def create_scheduler(

101 else None

102 )

103

105 - speculative_config: Optional[SpeculativeConfig] = None

104 + speculative_config: SpeculativeConfig | None = None

105 if num_speculative_tokens is not None:

106 speculative_config = SpeculativeConfig(

107 model="ngram", num_speculative_tokens=num_speculative_tokens

@@ -141,10 +140,10 @@ _none_hash_initialized = False

140 def create_requests(

141 num_requests: int,

142 num_tokens: int = 10,

144 - mm_positions: Optional[list[list[PlaceholderRange]]] = None,

143 + mm_positions: list[list[PlaceholderRange]] | None = None,

144 max_tokens: int = 16,

146 - stop_token_ids: Optional[list[int]] = None,

147 - prompt_logprobs: Optional[int] = None,

145 + stop_token_ids: list[int] | None = None,

146 + prompt_logprobs: int | None = None,

147 same_prompt: bool = False,

148 block_size: int = 16,

149 ) -> list[Request]:

tests/v1/distributed/test_async_llm_dp.py

+5 -6

@@ -5,7 +5,6 @@ import asyncio

5 import os

6 from contextlib import ExitStack

7 from dataclasses import dataclass

8 -from typing import Optional

8

9 import pytest

10

@@ -35,8 +34,8 @@ async def generate(

34 prompt: PromptType,

35 output_kind: RequestOutputKind,

36 max_tokens: int,

38 - prompt_logprobs: Optional[int] = None,

39 - data_parallel_rank: Optional[int] = None,

37 + prompt_logprobs: int | None = None,

38 + data_parallel_rank: int | None = None,

39 ) -> tuple[int, str]:

40 # Ensure generate doesn't complete too fast for cancellation test.

41 await asyncio.sleep(0.2)

@@ -91,9 +90,9 @@ async def test_load(

90

91 def record(

92 self,

94 - scheduler_stats: Optional[SchedulerStats],

95 - iteration_stats: Optional[IterationStats],

96 - mm_cache_stats: Optional[MultiModalCacheStats] = None,

93 + scheduler_stats: SchedulerStats | None,

94 + iteration_stats: IterationStats | None,

95 + mm_cache_stats: MultiModalCacheStats | None = None,

96 engine_idx: int = 0,

97 ):

98 if iteration_stats:

tests/v1/distributed/test_internal_lb_dp.py

+3 -3

@@ -5,7 +5,7 @@ import os

5 import threading

6 import time

7 import traceback

8 -from typing import Optional, cast

8 +from typing import cast

9

10 import openai # use the official client for correctness check

11 import pytest

@@ -46,7 +46,7 @@ class MultinodeInternalLBServerManager:

46 self.tp_size = tp_size

47 self.api_server_count = api_server_count

48 self.base_server_args = base_server_args

49 - self.servers: list[Optional[tuple[RemoteOpenAIServer, list[str]]]] = [None] * (

49 + self.servers: list[tuple[RemoteOpenAIServer, list[str]] | None] = [None] * (

50 dp_size // dp_per_node

51 )

52 self.server_threads: list[threading.Thread] = []

@@ -175,7 +175,7 @@ class APIOnlyServerManager:

175 self.tp_size = tp_size

176 self.api_server_count = api_server_count

177 self.base_server_args = base_server_args

178 - self.servers: list[Optional[tuple[RemoteOpenAIServer, list[str]]]] = [None] * 2

178 + self.servers: list[tuple[RemoteOpenAIServer, list[str]] | None] = [None] * 2

179 self.server_threads: list[threading.Thread] = []

180

181 def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:

tests/v1/e2e/test_min_tokens.py

+3 -5

@@ -13,8 +13,6 @@ Covers:

13 5) Multiple stop conditions

14 """

15

16 -from typing import Optional, Union

17 -

16 import pytest

17

18 from vllm import LLM, SamplingParams

@@ -33,9 +31,9 @@ class MinTokensTestCase:

31 name: str,

32 min_tokens: int,

33 max_tokens: int,

36 - stop: Optional[Union[str, list[str]]] = None,

37 - expected_min_len: Optional[int] = None,

38 - expected_exact_len: Optional[int] = None,

34 + stop: str | list[str] | None = None,

35 + expected_min_len: int | None = None,

36 + expected_exact_len: int | None = None,

37 ):

38 self.name = name

39 self.min_tokens = min_tokens

tests/v1/e2e/test_spec_decode.py

+2 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import random

6 -from typing import Any, Union

4 +from typing import Any

5

6 import pytest

7 import torch

@@ -34,7 +32,7 @@ def get_test_prompts(mm_enabled: bool):

32 for kind in random_prompt_type_choices:

33 word_choices = ["test", "temp", "hello", "where"]

34 word = random.choice(word_choices)

37 - prompt: Union[str, list[dict[str, Any]]] = ""

35 + prompt: str | list[dict[str, Any]] = ""

36 if kind == "repeat":

37 prompt = f"""

38 please repeat the word '{word}' 10 times.

tests/v1/engine/test_async_llm.py

+4 -5

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

3

4 import asyncio

5 from contextlib import ExitStack

6 -from typing import Optional

6 from unittest.mock import MagicMock

7

8 import pytest

@@ -53,8 +52,8 @@ async def generate(

52 output_kind: RequestOutputKind,

53 max_tokens: int,

54 n: int = 1,

56 - prompt_logprobs: Optional[int] = None,

57 - cancel_after: Optional[int] = None,

55 + prompt_logprobs: int | None = None,

56 + cancel_after: int | None = None,

57 ) -> tuple[int, str]:

58 # Ensure generate doesn't complete too fast for cancellation test.

59 await asyncio.sleep(0.2)

@@ -545,9 +544,9 @@ async def collect_outputs(

544 prompt: PromptType,

545 sampling_params: SamplingParams,

546 outputs_list: list[RequestOutput],

548 -) -> Optional[RequestOutput]:

547 +) -> RequestOutput | None:

548 """Helper to collect outputs and return the final one."""

550 - final_output: Optional[RequestOutput] = None

549 + final_output: RequestOutput | None = None

550 async for output in engine.generate(

551 request_id=request_id, prompt=prompt, sampling_params=sampling_params

552 ):

tests/v1/engine/test_engine_core_client.py

+5 -7

@@ -8,7 +8,7 @@ import time

8 import uuid

9 from dataclasses import dataclass

10 from threading import Thread

11 -from typing import Any, Optional, Union

11 +from typing import Any

12 from unittest.mock import MagicMock

13

14 import pytest

@@ -41,7 +41,7 @@ PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids

41

42

43 def make_request(

44 - params: SamplingParams, prompt_tokens_ids: Optional[list[int]] = None

44 + params: SamplingParams, prompt_tokens_ids: list[int] | None = None

45 ) -> EngineCoreRequest:

46 if not prompt_tokens_ids:

47 prompt_tokens_ids = PROMPT_TOKENS

@@ -113,9 +113,7 @@ async def loop_until_fully_done_async(client: EngineCoreClient, outputs: dict):

113

114

115 # Dummy utility function to monkey-patch into engine core.

116 -def echo(

117 - self, msg: str, err_msg: Optional[str] = None, sleep: Optional[float] = None

118 -) -> str:

116 +def echo(self, msg: str, err_msg: str | None = None, sleep: float | None = None) -> str:

117 print(f"echo util function called: {msg}, {err_msg}")

118 if sleep is not None:

119 time.sleep(sleep)

@@ -317,7 +315,7 @@ def echo_dc(

315 self,

316 msg: str,

317 return_list: bool = False,

320 -) -> Union[MyDataclass, list[MyDataclass]]:

318 +) -> MyDataclass | list[MyDataclass]:

319 print(f"echo dc util function called: {msg}")

320 val = None if msg is None else MyDataclass(msg)

321 # Return dataclass to verify support for returning custom types

@@ -330,7 +328,7 @@ def echo_dc_dict(

328 self,

329 msg: str,

330 return_dict: bool = False,

333 -) -> Union[MyDataclass, dict[str, MyDataclass]]:

331 +) -> MyDataclass | dict[str, MyDataclass]:

332 print(f"echo dc dict util function called: {msg}")

333 val = None if msg is None else MyDataclass(msg)

334 # Return dict of dataclasses to verify support for returning dicts

tests/v1/engine/test_llm_engine.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import random

4 from typing import TYPE_CHECKING

5

@@ -13,6 +11,8 @@ from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector

11

12 if TYPE_CHECKING:

13 from tests.conftest import VllmRunner

14 +else:

15 + VllmRunner = object

16

17 MODEL = "facebook/opt-125m"

18 DTYPE = "half"

tests/v1/engine/test_output_processor.py

+8 -9

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

3

4 import math

5 import time

6 -from typing import Optional

6

7 import pytest

8

@@ -118,13 +117,13 @@ def test_incremental_detokenization(

117

118 def _validate_logprobs(

119 gen_tokens: dict[str, list[int]],

121 - gen_logprobs: dict[str, Optional[SampleLogprobs]],

122 - gen_prompt_logprobs: dict[str, Optional[PromptLogprobs]],

120 + gen_logprobs: dict[str, SampleLogprobs | None],

121 + gen_prompt_logprobs: dict[str, PromptLogprobs | None],

122 gen_cumulative_logprob: dict[str, float],

123 dtv: DummyOutputProcessorTestVectors,

124 request_id_list: list[str],

126 - num_sample_logprobs: Optional[int],

127 - num_prompt_logprobs: Optional[int],

125 + num_sample_logprobs: int | None,

126 + num_prompt_logprobs: int | None,

127 ) -> None:

128 for req_idx, req_id in enumerate(request_id_list):

129 new_tokens = gen_tokens[req_id]

@@ -413,8 +412,8 @@ def _validate_logprobs(

412 @pytest.mark.parametrize("num_prompt_logprobs", [None, NUM_PROMPT_LOGPROBS_UNDER_TEST])

413 def test_logprobs_processor(

414 request_output_kind: RequestOutputKind,

416 - num_sample_logprobs: Optional[int],

417 - num_prompt_logprobs: Optional[int],

415 + num_sample_logprobs: int | None,

416 + num_prompt_logprobs: int | None,

417 dummy_test_vectors,

418 ):

419 output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=False)

@@ -530,7 +529,7 @@ def test_logprobs_processor(

529 )

530 def test_stop_token(

531 include_stop_str_in_output: bool,

533 - num_sample_logprobs: Optional[int],

532 + num_sample_logprobs: int | None,

533 stop_token_type: str,

534 ignore_eos: bool,

535 dummy_test_vectors,

@@ -696,7 +695,7 @@ def test_stop_token(

695 @pytest.mark.parametrize("num_sample_logprobs", [None, NUM_SAMPLE_LOGPROBS_UNDER_TEST])

696 def test_stop_string(

697 include_stop_str_in_output: bool,

699 - num_sample_logprobs: Optional[int],

698 + num_sample_logprobs: int | None,

699 dummy_test_vectors,

700 ):

701 output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=False)

tests/v1/engine/utils.py

+7 -8

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

3

4 import random

5 from dataclasses import dataclass

6 -from typing import Optional, Union

6 +from typing import TypeAlias

7

8 import torch

9 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast

@@ -12,7 +12,7 @@ from vllm.engine.arg_utils import EngineArgs

12 from vllm.v1.engine import EngineCoreOutput, FinishReason

13 from vllm.v1.outputs import LogprobsLists, LogprobsTensors

14

15 -GeneralTokenizerType = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]

15 +GeneralTokenizerType: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast

16

17 # Number of sample logprobs to request when testing sample logprobs

18 NUM_SAMPLE_LOGPROBS_UNDER_TEST = 5

@@ -332,16 +332,15 @@ class MockEngineCore:

332 # For each request, for each sampled token offset,

333 # a tuple of

334 # (list of topk token ids, list of sample logprob vals, rank)

335 - generated_logprobs_raw: Optional[

336 - list[list[tuple[list[int], list[float], int]]]

337 - ] = None,

335 + generated_logprobs_raw: list[list[tuple[list[int], list[float], int]]]

336 + | None = None,

337 # For each request, a tuple of

338 # (prompt logprob val matrix, prompt logprob tok id matrix);

339 # each matrix has dimensions

340 # (num prompt toks) x (num prompt logprobs+1)

342 - prompt_logprobs_raw: Optional[list[LogprobsTensors]] = None,

343 - eos_token_id: Optional[int] = None,

344 - stop_token_ids: Optional[list[int]] = None,

341 + prompt_logprobs_raw: list[LogprobsTensors] | None = None,

342 + eos_token_id: int | None = None,

343 + stop_token_ids: list[int] | None = None,

344 ignore_eos: bool = False,

345 ) -> None:

346 self.num_requests = len(tokens_list)

tests/v1/entrypoints/llm/test_struct_output_generate.py

+2 -2

@@ -2,8 +2,6 @@

2 # SPDX-License-Identifier: Apache-2.0

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

4

5 -from __future__ import annotations

6 -

5 import json

6 from dataclasses import fields

7 from enum import Enum

@@ -30,6 +28,8 @@ from vllm.sampling_params import (

28

29 if TYPE_CHECKING:

30 from vllm.config.model import TokenizerMode

31 +else:

32 + TokenizerMode = str

33

34 NGRAM_SPEC_CONFIG = {

35 "model": "[ngram]",

tests/v1/entrypoints/openai/test_completion.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import openai # use the official client for correctness check

6 import pytest

@@ -195,7 +194,7 @@ async def test_too_many_completion_logprobs(

194 [(MODEL_NAME, -1), (MODEL_NAME, 0), (MODEL_NAME, 1), (MODEL_NAME, None)],

195 )

196 async def test_prompt_logprobs_completion(

198 - client: openai.AsyncOpenAI, model_name: str, prompt_logprobs: Optional[int]

197 + client: openai.AsyncOpenAI, model_name: str, prompt_logprobs: int | None

198 ):

199 params: dict = {

200 "prompt": ["A robot may not injure another robot", "My name is"],

tests/v1/executor/test_executor.py

+6 -5

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

3

4 import asyncio

5 import os

6 -from typing import Any, Callable, Optional, Union

6 +from collections.abc import Callable

7 +from typing import Any

8

9 import pytest

10

@@ -20,12 +21,12 @@ class Mock: ...

21 class CustomMultiprocExecutor(MultiprocExecutor):

22 def collective_rpc(

23 self,

23 - method: Union[str, Callable],

24 - timeout: Optional[float] = None,

24 + method: str | Callable,

25 + timeout: float | None = None,

26 args: tuple = (),

26 - kwargs: Optional[dict] = None,

27 + kwargs: dict | None = None,

28 non_block: bool = False,

28 - unique_reply_rank: Optional[int] = None,

29 + unique_reply_rank: int | None = None,

30 ) -> list[Any]:

31 # Drop marker to show that this was run

32 with open(".marker", "w"):

tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py

+1 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable

4 +from collections.abc import Callable

5 from unittest.mock import Mock

6

7 import pytest

tests/v1/kv_connector/unit/test_nixl_connector.py

+2 -3

@@ -9,7 +9,6 @@ import textwrap

9 import time

10 import uuid

11 from collections import defaultdict

12 -from typing import Optional

12 from unittest.mock import patch

13

14 import pytest

@@ -154,7 +153,7 @@ class FakeNixlWrapper:

153 local_block_descs_ids: list[int],

154 remote_xfer_side_handle: int,

155 remote_block_descs_ids: list[int],

157 - notif_msg: Optional[bytes] = None,

156 + notif_msg: bytes | None = None,

157 ) -> int:

158 return uuid.uuid4().int

159

@@ -982,7 +981,7 @@ class FakePlatform(Platform):

981 return {"oot": ("oot",)}

982

983 @classmethod

985 - def get_nixl_memory_type(cls) -> Optional[str]:

984 + def get_nixl_memory_type(cls) -> str | None:

985 """

986 Returns the nixl memory type for the current platform.

987 """

tests/v1/kv_connector/unit/test_output_aggreagator.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from concurrent.futures import Future

4 -from typing import Optional

4

5 import pytest

6

@@ -14,9 +13,9 @@ pytestmark = pytest.mark.cpu_test

13 class DummyModelRunnerOutput(ModelRunnerOutput):

14 def __init__(

15 self,

17 - finished_sending: Optional[set[str]] = None,

18 - finished_recving: Optional[set[str]] = None,

19 - invalid_block_ids: Optional[set[int]] = None,

16 + finished_sending: set[str] | None = None,

17 + finished_recving: set[str] | None = None,

18 + invalid_block_ids: set[int] | None = None,

19 ):

20 self.kv_connector_output = KVConnectorOutput(

21 finished_sending=finished_sending,

tests/v1/kv_connector/unit/utils.py

+7 -6

@@ -2,8 +2,9 @@

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

3 import tempfile

4 from collections import defaultdict

5 +from collections.abc import Callable

6 from itertools import count

6 -from typing import Any, Callable, Optional

7 +from typing import Any

8

9 import torch

10

@@ -147,7 +148,7 @@ _none_hash_initialized = False

148

149

150 def create_request(

150 - request_id: Optional[int] = None,

151 + request_id: int | None = None,

152 num_tokens: int = 10,

153 common_prefix_len=0,

154 max_tokens: int = 16,

@@ -168,7 +169,7 @@ def create_request(

169 init_none_hash(hash_fn)

170 _none_hash_initialized = True

171

171 - kv_transfer_params: Optional[dict[str, Any]] = None

172 + kv_transfer_params: dict[str, Any] | None = None

173

174 if do_remote_decode:

175 assert not do_remote_prefill

@@ -205,9 +206,9 @@ def create_request(

206

207 def create_model_runner_output(

208 reqs: list[Request],

208 - finished_sending: Optional[set[str]] = None,

209 - finished_recving: Optional[set[str]] = None,

210 - invalid_block_ids: Optional[set[int]] = None,

209 + finished_sending: set[str] | None = None,

210 + finished_recving: set[str] | None = None,

211 + invalid_block_ids: set[int] | None = None,

212 use_eos: bool = False,

213 token_id: int = 0,

214 ) -> ModelRunnerOutput:

tests/v1/kv_offload/test_cpu_manager.py

+1 -2

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

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

3 from collections.abc import Iterable

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import numpy as np

7

@@ -29,7 +28,7 @@ def to_hashes(int_hashes: list[int]) -> list[BlockHash]:

28

29

30 def verify_store_output(

32 - prepare_store_output: Optional[PrepareStoreOutput],

31 + prepare_store_output: PrepareStoreOutput | None,

32 expected_prepare_store_output: ExpectedPrepareStoreOutput,

33 ):

34 assert prepare_store_output is not None

tests/v1/logits_processors/test_correctness.py

+4 -4

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

3

4 import random

5 from collections.abc import Callable

6 -from typing import NamedTuple, Optional, Union

6 +from typing import NamedTuple, TypeAlias

7

8 import numpy as np

9 import pytest

@@ -48,7 +48,7 @@ REQS_PER_LOGITPROC = 50

48 STR_NO_LOGITPROC = "none"

49

50 # LogitsProcessor subclass or "none"

51 -LogitprocType = Union[type[LogitsProcessor], str]

51 +LogitprocType: TypeAlias = type[LogitsProcessor] | str

52

53

54 class LogitsProcsRequestParams:

@@ -435,7 +435,7 @@ class LogitsprocTestHelpers(NamedTuple):

435 """Supports setting up and validating logitsprocs unit tests."""

436

437 eval_fxn: Callable

438 - gen_request_fxn: Optional[Callable] = None

438 + gen_request_fxn: Callable | None = None

439

440

441 logitsprocs_test_mapping = {

@@ -471,7 +471,7 @@ def _generate_fake_step_update(

471 workload_params: list[LogitsProcsRequestParams],

472 wdx: int,

473 batch_update_builder: BatchUpdateBuilder,

474 -) -> tuple[Optional[BatchUpdate], int, int]:

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

475 batch_size = len(persistent_batch)

476 workload_size = len(workload_params)

477 workload_reqs_remaining = workload_size - wdx

tests/v1/logits_processors/test_custom_offline.py

+2 -2

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

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

3 import random

4 import sys

5 -from typing import Any, Union

5 +from typing import Any

6

7 import pytest

8

@@ -159,7 +159,7 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource

159 _run_test({}, logitproc_loaded=True)

160 return

161

162 - kwargs: dict[str, list[Union[str, type[LogitsProcessor]]]] = {}

162 + kwargs: dict[str, list[str | type[LogitsProcessor]]] = {}

163 if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_FQCN:

164 # Scenario: load logitproc based on fully-qualified class name (FQCN)

165 # Inject dummy module which defines logitproc

tests/v1/logits_processors/test_custom_online.py

+3 -3

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

4 import os

5 import random

6 import sys

7 -from typing import Any, Optional

7 +from typing import Any

8

9 import openai

10 import pytest

@@ -25,7 +25,7 @@ from tests.v1.logits_processors.utils import entry_points as fake_entry_points

25

26

27 def _server_with_logitproc_entrypoint(

28 - env_dict: Optional[dict[str, str]],

28 + env_dict: dict[str, str] | None,

29 model: str,

30 vllm_serve_args: list[str],

31 ) -> None:

@@ -48,7 +48,7 @@ def _server_with_logitproc_entrypoint(

48

49

50 def _server_with_logitproc_module(

51 - env_dict: Optional[dict[str, str]],

51 + env_dict: dict[str, str] | None,

52 model: str,

53 vllm_serve_args: list[str],

54 ) -> None:

tests/v1/logits_processors/utils.py

+4 -4

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

3

4 import types

5 from enum import Enum, auto

6 -from typing import Any, Optional

6 +from typing import Any

7

8 import torch

9

@@ -61,7 +61,7 @@ class DummyLogitsProcessor(LogitsProcessor):

61 """Never impacts greedy sampling"""

62 return False

63

64 - def update_state(self, batch_update: Optional[BatchUpdate]):

64 + def update_state(self, batch_update: BatchUpdate | None):

65 process_dict_updates(

66 self.req_info,

67 batch_update,

@@ -145,7 +145,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):

145 def new_req_logits_processor(

146 self,

147 params: SamplingParams,

148 - ) -> Optional[RequestLogitsProcessor]:

148 + ) -> RequestLogitsProcessor | None:

149 """This method returns a new request-level logits processor, customized

150 to the `target_token` value associated with a particular request.

151

@@ -159,7 +159,7 @@ class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):

159 Returns:

160 `Callable` request logits processor, or None

161 """

162 - target_token: Optional[Any] = params.extra_args and params.extra_args.get(

162 + target_token: Any | None = params.extra_args and params.extra_args.get(

163 "target_token"

164 )

165 if target_token is None:

tests/v1/sample/test_rejection_sampler.py

+14 -14

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Any, Optional

3 +from typing import Any

4

5 import pytest

6 import torch

@@ -24,7 +24,7 @@ def rejection_sampler():

24 def create_logits_tensor(

25 output_token_ids: list[list[int]],

26 vocab_size: int = 100,

27 - token_idx_to_override: Optional[int] = None,

27 + token_idx_to_override: int | None = None,

28 ) -> torch.Tensor:

29 """Helper function to create logits tensor that

30 will produce desired token ids on argmax"""

@@ -43,18 +43,18 @@ def create_logits_tensor(

43

44 def create_sampling_metadata(

45 all_greedy: bool,

46 - output_token_ids: Optional[list[list[int]]] = None,

47 - prompt_token_ids: Optional[torch.Tensor] = None,

48 - spec_token_ids: Optional[torch.Tensor] = None,

49 - temperature: Optional[torch.Tensor] = None,

50 - top_k: Optional[torch.Tensor] = None,

51 - top_p: Optional[torch.Tensor] = None,

52 - generators: Optional[dict[int, Any]] = None,

53 - frequency_penalties: Optional[list[float]] = None,

54 - presence_penalties: Optional[list[float]] = None,

55 - repetition_penalties: Optional[list[float]] = None,

56 - bad_words_token_ids: Optional[dict[int, list[list[int]]]] = None,

57 - allowed_token_ids_mask: Optional[torch.Tensor] = None,

46 + output_token_ids: list[list[int]] | None = None,

47 + prompt_token_ids: torch.Tensor | None = None,

48 + spec_token_ids: torch.Tensor | None = None,

49 + temperature: torch.Tensor | None = None,

50 + top_k: torch.Tensor | None = None,

51 + top_p: torch.Tensor | None = None,

52 + generators: dict[int, Any] | None = None,

53 + frequency_penalties: list[float] | None = None,

54 + presence_penalties: list[float] | None = None,

55 + repetition_penalties: list[float] | None = None,

56 + bad_words_token_ids: dict[int, list[list[int]]] | None = None,

57 + allowed_token_ids_mask: torch.Tensor | None = None,

58 ) -> SamplingMetadata:

59 """Create a v1 sampling metadata object with all_greedy set

60 to the given value. Either all greedy or all random sampling

tests/v1/sample/utils.py

+4 -4

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

3

4 from collections.abc import Iterator

5 from enum import Enum

6 -from typing import NamedTuple, Optional

6 +from typing import NamedTuple

7

8 import regex as re

9 import torch

@@ -23,7 +23,7 @@ class BatchLogprobsComposition(Enum):

23 SAMPLE_PROMPT = 3

24

25

26 -BatchLogprobsSpecType = list[tuple[Optional[int], Optional[int]]]

26 +BatchLogprobsSpecType = list[tuple[int | None, int | None]]

27

28

29 def get_test_batch(

@@ -222,8 +222,8 @@ def create_allowed_token_ids(

222 vocab_size: int,

223 num_allowed_token_ids: int,

224 device: torch.device,

225 -) -> Optional[torch.Tensor]:

226 - mask: Optional[torch.Tensor] = None

225 +) -> torch.Tensor | None:

226 + mask: torch.Tensor | None = None

227 for i in range(batch_size):

228 if i % 2 == 1:

229 continue

tests/v1/spec_decode/test_eagle.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4 from unittest import mock

5

6 import pytest

@@ -39,7 +38,7 @@ eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"

38 def _create_proposer(

39 method: str,

40 num_speculative_tokens: int,

42 - speculative_token_tree: Optional[list[tuple[int, ...]]] = None,

41 + speculative_token_tree: list[tuple[int, ...]] | None = None,

42 ) -> EagleProposer:

43 model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)

44

tests/v1/spec_decode/test_tree_attention.py

+1 -2

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

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

3

4 import math

5 -from typing import Optional

5

6 import torch

7

@@ -37,7 +36,7 @@ def forward_attention(

36 slot_mapping: torch.Tensor,

37 seqlen_k: int,

38 backend: _Backend,

40 - spec_token_tree: Optional[str] = None,

39 + spec_token_tree: str | None = None,

40 num_spec_tokens: int = 0,

41 ) -> torch.Tensor:

42 batch_size, q_len, num_heads, dim_per_head = q.shape

tests/v1/test_serial_utils.py

+1 -2

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

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

3 from collections import UserDict

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import msgspec

7 import numpy as np

@@ -100,7 +99,7 @@ def test_encode_decode(monkeypatch: pytest.MonkeyPatch):

99

100

101 class MyRequest(msgspec.Struct):

103 - mm: Optional[list[MultiModalKwargsItems]]

102 + mm: list[MultiModalKwargsItems] | None

103

104

105 def test_multimodal_kwargs():

tests/v1/tpu/test_basic.py

+2 -2

@@ -5,8 +5,6 @@

5 Run `pytest tests/v1/tpu/test_basic.py`.

6 """

7

8 -from __future__ import annotations

9 -

8 from typing import TYPE_CHECKING

9

10 import pytest

@@ -16,6 +14,8 @@ from vllm.platforms import current_platform

14

15 if TYPE_CHECKING:

16 from tests.conftest import VllmRunner

17 +else:

18 + VllmRunner = object

19

20 MODELS = [

21 "Qwen/Qwen2.5-1.5B-Instruct",

tests/v1/tpu/test_perf.py

+2 -2

@@ -5,8 +5,6 @@

5 Run `pytest tests/v1/tpu/test_perf.py`.

6 """

7

8 -from __future__ import annotations

9 -

8 import time

9 from dataclasses import dataclass

10 from typing import TYPE_CHECKING

@@ -20,6 +18,8 @@ from vllm.transformers_utils.tokenizer import get_tokenizer

18

19 if TYPE_CHECKING:

20 from tests.conftest import VllmRunner

21 +else:

22 + VllmRunner = object

23

24

25 @dataclass

tests/v1/tracing/test_tracing.py

-2

@@ -2,8 +2,6 @@

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

3 # ruff: noqa

4 # type: ignore

5 -from __future__ import annotations

6 -

5 import threading

6 from collections.abc import Iterable

7 from concurrent import futures

tests/v1/worker/test_gpu_input_batch.py

+1 -2

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

3

4 import inspect

5 from collections.abc import Sequence

6 -from typing import Optional

6

7 import numpy as np

8 import pytest

@@ -271,7 +270,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

270 reqs, req_ids_retained, input_batch.req_id_to_index, device=torch.device(device)

271 )

272

274 - def same(t1: Optional[torch.Tensor], t2: Optional[torch.Tensor]) -> bool:

273 + def same(t1: torch.Tensor | None, t2: torch.Tensor | None) -> bool:

274 return (t1 is None and t2 is None) or (

275 t1 is not None and t2 is not None and torch.allclose(t1, t2)

276 )

tests/v1/worker/test_worker_memory_snapshot.py

+2 -3

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

4 import multiprocessing as mp

5 import os

6 import tempfile

7 -from multiprocessing import Queue

8 -from typing import Optional

7 +from multiprocessing.queues import Queue

8 from unittest.mock import patch

9

10 import pytest

@@ -16,7 +15,7 @@ from vllm.utils import MemorySnapshot

15 from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment

16

17 # Global queue to track operation order across processes

19 -_QUEUE: Optional[Queue] = None

18 +_QUEUE: Queue | None = None

19

20

21 def track_operation(operation: str, rank: int):

tests/vllm_test_utils/vllm_test_utils/blame.py

+1 -2

@@ -5,8 +5,7 @@ import contextlib

5 import dataclasses

6 import sys

7 import traceback

8 -from collections.abc import Generator

9 -from typing import Callable

8 +from collections.abc import Callable, Generator

9

10

11 @dataclasses.dataclass

tests/vllm_test_utils/vllm_test_utils/monitor.py

+2 -2

@@ -5,8 +5,8 @@ import contextlib

5 import dataclasses

6 import sys

7 import traceback

8 -from collections.abc import Generator

9 -from typing import Callable, Generic, TypeVar

8 +from collections.abc import Callable, Generator

9 +from typing import Generic, TypeVar

10

11 _T = TypeVar("_T")

12

tools/check_init_lazy_imports.py

-2

@@ -5,8 +5,6 @@ i.e: appears only within the ``if typing.TYPE_CHECKING:`` guard,

5 **except** for a short whitelist.

6 """

7

8 -from __future__ import annotations

9 -

8 import ast

9 import pathlib

10 import sys

tools/enforce_regex_import.py

-2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import subprocess

4 from pathlib import Path

5

tools/pre_commit/mypy.py

+2 -3

@@ -20,7 +20,6 @@ Args:

20

21 import subprocess

22 import sys

23 -from typing import Optional

23

24 import regex as re

25

@@ -96,8 +95,8 @@ def group_files(changed_files: list[str]) -> dict[str, list[str]]:

95

96 def mypy(

97 targets: list[str],

99 - python_version: Optional[str],

100 - follow_imports: Optional[str],

98 + python_version: str | None,

99 + follow_imports: str | None,

100 file_group: str,

101 ) -> int:

102 """

tools/profiler/visualize_layerwise_profile.py

+2 -2

@@ -7,7 +7,7 @@ import json

7 import math

8 import os

9 from pathlib import Path

10 -from typing import Any, Optional

10 +from typing import Any

11

12 import matplotlib.pyplot as plt

13 import pandas as pd

@@ -373,7 +373,7 @@ def plot_trace_df(

373 traces_df: pd.DataFrame,

374 plot_metric: str,

375 plot_title: str,

376 - output: Optional[Path] = None,

376 + output: Path | None = None,

377 ):

378 def get_phase_description(traces_df: pd.DataFrame, phase: str) -> str:

379 phase_df = traces_df.query(f'phase == "{phase}"')

vllm/_bc_linter.py

+2 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 # vllm/_bc_linter.py

4 -from __future__ import annotations

5 -

6 -from typing import Any, Callable, TypeVar, overload

4 +from collections.abc import Callable

5 +from typing import Any, TypeVar, overload

6

7 T = TypeVar("T")

8

vllm/_custom_ops.py

+103 -105

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import TYPE_CHECKING, Literal, Optional, Union

4 +from typing import TYPE_CHECKING, Literal

5

6 import torch

7

@@ -37,7 +37,7 @@ def paged_attention_v1(

37 seq_lens: torch.Tensor,

38 block_size: int,

39 max_seq_len: int,

40 - alibi_slopes: Optional[torch.Tensor],

40 + alibi_slopes: torch.Tensor | None,

41 kv_cache_dtype: str,

42 k_scale: torch.Tensor,

43 v_scale: torch.Tensor,

@@ -84,7 +84,7 @@ def paged_attention_v2(

84 seq_lens: torch.Tensor,

85 block_size: int,

86 max_seq_len: int,

87 - alibi_slopes: Optional[torch.Tensor],

87 + alibi_slopes: torch.Tensor | None,

88 kv_cache_dtype: str,

89 k_scale: torch.Tensor,

90 v_scale: torch.Tensor,

@@ -132,14 +132,14 @@ def paged_attention_rocm(

132 scale: float,

133 block_tables: torch.Tensor,

134 seq_lens: torch.Tensor,

135 - query_start_loc: Optional[torch.Tensor],

135 + query_start_loc: torch.Tensor | None,

136 block_size: int,

137 max_seq_len: int,

138 - alibi_slopes: Optional[torch.Tensor],

138 + alibi_slopes: torch.Tensor | None,

139 kv_cache_dtype: str,

140 k_scale: torch.Tensor,

141 v_scale: torch.Tensor,

142 - fp8_out_scale: Optional[torch.Tensor] = None,

142 + fp8_out_scale: torch.Tensor | None = None,

143 mfma_type: str = "fp8" if envs.VLLM_ROCM_FP8_MFMA_PAGE_ATTN else "f16",

144 ) -> None:

145 torch.ops._rocm_C.paged_attention(

@@ -186,7 +186,7 @@ def merge_attn_states(

186 prefix_lse: torch.Tensor,

187 suffix_output: torch.Tensor,

188 suffix_lse: torch.Tensor,

189 - output_lse: Optional[torch.Tensor] = None,

189 + output_lse: torch.Tensor | None = None,

190 ) -> None:

191 torch.ops._C.merge_attn_states(

192 output, output_lse, prefix_output, prefix_lse, suffix_output, suffix_lse

@@ -314,7 +314,7 @@ def convert_vertical_slash_indexes_mergehead(

314 def rotary_embedding(

315 positions: torch.Tensor,

316 query: torch.Tensor,

317 - key: Optional[torch.Tensor],

317 + key: torch.Tensor | None,

318 head_size: int,

319 cos_sin_cache: torch.Tensor,

320 is_neox: bool,

@@ -408,8 +408,8 @@ def rms_norm_dynamic_per_token_quant(

408 weight: torch.Tensor,

409 epsilon: float,

410 quant_dtype: torch.dtype,

411 - scale_ub: Optional[torch.Tensor] = None,

412 - residual: Optional[torch.Tensor] = None,

411 + scale_ub: torch.Tensor | None = None,

412 + residual: torch.Tensor | None = None,

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

414 output = torch.empty_like(input, dtype=quant_dtype)

415 scales = torch.empty(

@@ -527,14 +527,14 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):

527 @register_fake("_C::gptq_marlin_gemm")

528 def _gptq_marlin_gemm_fake(

529 a: torch.Tensor,

530 - c: Optional[torch.Tensor],

530 + c: torch.Tensor | None,

531 b_q_weight: torch.Tensor,

532 - b_bias: Optional[torch.Tensor],

532 + b_bias: torch.Tensor | None,

533 b_scales: torch.Tensor,

534 - global_scale: Optional[torch.Tensor],

535 - b_zeros: Optional[torch.Tensor],

536 - g_idx: Optional[torch.Tensor],

537 - perm: Optional[torch.Tensor],

534 + global_scale: torch.Tensor | None,

535 + b_zeros: torch.Tensor | None,

536 + g_idx: torch.Tensor | None,

537 + perm: torch.Tensor | None,

538 workspace: torch.Tensor,

539 b_q_type_id: int,

540 size_m: torch.SymInt,

@@ -582,13 +582,13 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):

582 # b_q Should be the tensor returned by machete_prepack_B

583 b_q: torch.Tensor,

584 b_type: ScalarType,

585 - out_type: Optional[torch.dtype] = None,

586 - b_group_scales: Optional[torch.Tensor] = None,

587 - b_group_zeros: Optional[torch.Tensor] = None,

588 - b_group_size: Optional[int] = None,

589 - b_channel_scales: Optional[torch.Tensor] = None,

590 - a_token_scales: Optional[torch.Tensor] = None,

591 - schedule: Optional[str] = None,

585 + out_type: torch.dtype | None = None,

586 + b_group_scales: torch.Tensor | None = None,

587 + b_group_zeros: torch.Tensor | None = None,

588 + b_group_size: int | None = None,

589 + b_channel_scales: torch.Tensor | None = None,

590 + a_token_scales: torch.Tensor | None = None,

591 + schedule: str | None = None,

592 ) -> torch.Tensor:

593 m = a.size(0)

594 n = b_q.size(1)

@@ -599,7 +599,7 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):

599 b_q_weight: torch.Tensor,

600 a_type: torch.dtype,

601 b_type: ScalarType,

602 - group_scales_type: Optional[torch.dtype],

602 + group_scales_type: torch.dtype | None,

603 ) -> torch.Tensor:

604 return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format)

605

@@ -612,8 +612,8 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):

612 b_group_size: int,

613 b_channel_scales: torch.Tensor,

614 a_token_scales: torch.Tensor,

615 - out_type: Optional[torch.dtype] = None,

616 - maybe_schedule: Optional[str] = None,

615 + out_type: torch.dtype | None = None,

616 + maybe_schedule: str | None = None,

617 ) -> torch.Tensor:

618 m = a.size(0)

619 n = b_q.size(1)

@@ -636,7 +636,7 @@ if hasattr(torch.ops._C, "allspark_w8a16_gemm"):

636 a: torch.Tensor,

637 b_qweight: torch.Tensor,

638 b_scales: torch.Tensor,

639 - b_qzeros: Optional[torch.Tensor],

639 + b_qzeros: torch.Tensor | None,

640 n: torch.SymInt,

641 group_size: torch.SymInt,

642 sm_count: torch.SymInt,

@@ -657,7 +657,7 @@ if hasattr(torch.ops._C, "ggml_dequantize"):

657 quant_type: int,

658 m: torch.SymInt,

659 n: torch.SymInt,

660 - dtype: Optional[torch.dtype] = None,

660 + dtype: torch.dtype | None = None,

661 ) -> torch.Tensor:

662 return torch.empty((m, n), dtype=torch.float16, device=W.device)

663

@@ -760,7 +760,7 @@ def cutlass_scaled_mm(

760 scale_a: torch.Tensor,

761 scale_b: torch.Tensor,

762 out_dtype: torch.dtype,

763 - bias: Optional[torch.Tensor] = None,

763 + bias: torch.Tensor | None = None,

764 ) -> torch.Tensor:

765 """

766 `cutlass_scaled_mm` implements a fused version of

@@ -812,8 +812,8 @@ def cutlass_scaled_mm_azp(

812 scale_b: torch.Tensor,

813 out_dtype: torch.dtype,

814 azp_adj: torch.Tensor,

815 - azp: Optional[torch.Tensor] = None,

816 - bias: Optional[torch.Tensor] = None,

815 + azp: torch.Tensor | None = None,

816 + bias: torch.Tensor | None = None,

817 ) -> torch.Tensor:

818 """

819 :param azp_adj: In the per-tensor case, this should include the azp.

@@ -890,7 +890,7 @@ def cutlass_scaled_sparse_mm(

890 scale_a: torch.Tensor,

891 scale_b: torch.Tensor,

892 out_dtype: torch.dtype,

893 - bias: Optional[torch.Tensor] = None,

893 + bias: torch.Tensor | None = None,

894 ) -> torch.Tensor:

895 """

896 Performs a scaled sparse matrix multiplication using Cutlass.

@@ -940,7 +940,7 @@ def get_cutlass_moe_mm_data(

940 num_experts: int,

941 n: int,

942 k: int,

943 - blockscale_offsets: Optional[torch.Tensor] = None,

943 + blockscale_offsets: torch.Tensor | None = None,

944 ):

945 """

946 Prepare data necessary to perform CUTLASS grouped matrix multiplications

@@ -986,7 +986,7 @@ def get_cutlass_moe_mm_problem_sizes(

986 num_experts: int,

987 n: int,

988 k: int,

989 - blockscale_offsets: Optional[torch.Tensor] = None,

989 + blockscale_offsets: torch.Tensor | None = None,

990 ):

991 """

992 Compute only the per-expert problem sizes needed by the two grouped matrix

@@ -1195,14 +1195,14 @@ def awq_marlin_moe_repack(

1195

1196 def gptq_marlin_gemm(

1197 a: torch.Tensor,

1198 - c: Optional[torch.Tensor],

1198 + c: torch.Tensor | None,

1199 b_q_weight: torch.Tensor,

1200 - b_bias: Optional[torch.Tensor],

1200 + b_bias: torch.Tensor | None,

1201 b_scales: torch.Tensor,

1202 - global_scale: Optional[torch.Tensor],

1203 - b_zeros: Optional[torch.Tensor],

1204 - g_idx: Optional[torch.Tensor],

1205 - perm: Optional[torch.Tensor],

1202 + global_scale: torch.Tensor | None,

1203 + b_zeros: torch.Tensor | None,

1204 + g_idx: torch.Tensor | None,

1205 + perm: torch.Tensor | None,

1206 workspace: torch.Tensor,

1207 b_q_type: ScalarType,

1208 size_m: int,

@@ -1239,11 +1239,11 @@ def gptq_marlin_gemm(

1239 def machete_supported_schedules(

1240 a_type: torch.dtype,

1241 b_type: ScalarType,

1242 - group_scales_type: Optional[torch.dtype],

1243 - group_zeros_type: Optional[torch.dtype] = None,

1244 - channel_scales_type: Optional[torch.dtype] = None,

1245 - token_scales_type: Optional[torch.dtype] = None,

1246 - out_type: Optional[torch.dtype] = None,

1242 + group_scales_type: torch.dtype | None,

1243 + group_zeros_type: torch.dtype | None = None,

1244 + channel_scales_type: torch.dtype | None = None,

1245 + token_scales_type: torch.dtype | None = None,

1246 + out_type: torch.dtype | None = None,

1247 ) -> list[str]:

1248 return torch.ops._C.machete_supported_schedules(

1249 a_type,

@@ -1261,13 +1261,13 @@ def machete_mm(

1261 # b_q Should be the tensor returned by machete_prepack_B

1262 b_q: torch.Tensor,

1263 b_type: ScalarType,

1264 - out_type: Optional[torch.dtype] = None,

1265 - b_group_scales: Optional[torch.Tensor] = None,

1266 - b_group_zeros: Optional[torch.Tensor] = None,

1267 - b_group_size: Optional[int] = None,

1268 - b_channel_scales: Optional[torch.Tensor] = None,

1269 - a_token_scales: Optional[torch.Tensor] = None,

1270 - schedule: Optional[str] = None,

1264 + out_type: torch.dtype | None = None,

1265 + b_group_scales: torch.Tensor | None = None,

1266 + b_group_zeros: torch.Tensor | None = None,

1267 + b_group_size: int | None = None,

1268 + b_channel_scales: torch.Tensor | None = None,

1269 + a_token_scales: torch.Tensor | None = None,

1270 + schedule: str | None = None,

1271 ) -> torch.Tensor:

1272 return torch.ops._C.machete_mm(

1273 a,

@@ -1287,7 +1287,7 @@ def machete_prepack_B(

1287 b_q_weight: torch.Tensor,

1288 a_type: torch.dtype,

1289 b_type: ScalarType,

1290 - group_scales_type: Optional[torch.dtype],

1290 + group_scales_type: torch.dtype | None,

1291 ) -> torch.Tensor:

1292 return torch.ops._C.machete_prepack_B(

1293 b_q_weight, a_type, b_type.id, group_scales_type

@@ -1303,8 +1303,8 @@ def cutlass_w4a8_mm(

1303 b_group_size: int,

1304 b_channel_scales: torch.Tensor,

1305 a_token_scales: torch.Tensor,

1306 - out_type: Optional[torch.dtype] = None,

1307 - maybe_schedule: Optional[str] = None,

1306 + out_type: torch.dtype | None = None,

1307 + maybe_schedule: str | None = None,

1308 ) -> torch.Tensor:

1309 return torch.ops._C.cutlass_w4a8_mm(

1310 a,

@@ -1458,11 +1458,11 @@ def scaled_fp4_experts_quant(

1458 # fp8

1459 def scaled_fp8_quant(

1460 input: torch.Tensor,

1461 - scale: Optional[torch.Tensor] = None,

1462 - num_token_padding: Optional[int] = None,

1463 - scale_ub: Optional[torch.Tensor] = None,

1461 + scale: torch.Tensor | None = None,

1462 + num_token_padding: int | None = None,

1463 + scale_ub: torch.Tensor | None = None,

1464 use_per_token_if_dynamic: bool = False,

1465 - output: Optional[torch.Tensor] = None,

1465 + output: torch.Tensor | None = None,

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

1467 """

1468 Quantize input tensor to FP8 and return quantized tensor and scale.

@@ -1489,7 +1489,7 @@ def scaled_fp8_quant(

1489 """

1490 # This code assumes batch_dim and num_tokens are flattened

1491 assert input.ndim == 2

1492 - shape: Union[tuple[int, int], torch.Size] = input.shape

1492 + shape: tuple[int, int] | torch.Size = input.shape

1493 # For ROCm on MI300, the output fp8 dtype is torch.float_e3m3fnuz

1494 out_dtype: torch.dtype = current_platform.fp8_dtype()

1495 if num_token_padding:

@@ -1520,7 +1520,7 @@ def scaled_fp8_quant(

1520 def allspark_repack_weight(

1521 qweight: torch.Tensor,

1522 scale: torch.Tensor,

1523 - zero_point: Optional[torch.Tensor] = None,

1523 + zero_point: torch.Tensor | None = None,

1524 has_zp: bool = False,

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

1526 """

@@ -1576,7 +1576,7 @@ def allspark_w8a16_gemm(

1576 a: torch.Tensor,

1577 b_qweight: torch.Tensor,

1578 b_scales: torch.Tensor,

1579 - b_qzeros: Optional[torch.Tensor],

1579 + b_qzeros: torch.Tensor | None,

1580 n: int,

1581 group_size: int,

1582 sm_count: int,

@@ -1603,10 +1603,10 @@ def allspark_w8a16_gemm(

1603 # int8

1604 def scaled_int8_quant(

1605 input: torch.Tensor,

1606 - scale: Optional[torch.Tensor] = None,

1607 - azp: Optional[torch.Tensor] = None,

1606 + scale: torch.Tensor | None = None,

1607 + azp: torch.Tensor | None = None,

1608 symmetric: bool = True,

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

1609 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

1610 """

1611 Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.

1612

@@ -1643,7 +1643,7 @@ def scaled_int8_quant(

1643

1644 # gguf

1645 def ggml_dequantize(

1646 - W: torch.Tensor, quant_type: int, m: int, n: int, dtype: Optional[torch.dtype]

1646 + W: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None

1647 ) -> torch.Tensor:

1648 return torch.ops._C.ggml_dequantize(W, quant_type, m, n, dtype)

1649

@@ -1713,13 +1713,13 @@ def selective_scan_fwd(

1713 A: torch.Tensor,

1714 B: torch.Tensor,

1715 C: torch.Tensor,

1716 - D_: Optional[torch.Tensor],

1717 - z_: Optional[torch.Tensor],

1718 - delta_bias_: Optional[torch.Tensor],

1716 + D_: torch.Tensor | None,

1717 + z_: torch.Tensor | None,

1718 + delta_bias_: torch.Tensor | None,

1719 delta_softplus: bool,

1720 - query_start_loc: Optional[torch.Tensor],

1721 - cache_indices: Optional[torch.Tensor],

1722 - has_initial_state: Optional[torch.Tensor],

1720 + query_start_loc: torch.Tensor | None,

1721 + cache_indices: torch.Tensor | None,

1722 + has_initial_state: torch.Tensor | None,

1723 ssm_states: torch.Tensor,

1724 pad_slot_id: int,

1725 ):

@@ -1794,8 +1794,8 @@ def moe_wna16_gemm(

1794 output: torch.Tensor,

1795 b_qweight: torch.Tensor,

1796 b_scales: torch.Tensor,

1797 - b_qzeros: Optional[torch.Tensor],

1798 - topk_weights: Optional[torch.Tensor],

1797 + b_qzeros: torch.Tensor | None,

1798 + topk_weights: torch.Tensor | None,

1799 sorted_token_ids: torch.Tensor,

1800 experts_ids: torch.Tensor,

1801 num_tokens_post_pad: torch.Tensor,

@@ -1864,14 +1864,14 @@ def grouped_topk(

1864

1865 def moe_wna16_marlin_gemm(

1866 input: torch.Tensor,

1867 - output: Optional[torch.Tensor],

1867 + output: torch.Tensor | None,

1868 b_qweight: torch.Tensor,

1869 - b_bias: Optional[torch.Tensor],

1869 + b_bias: torch.Tensor | None,

1870 b_scales: torch.Tensor,

1871 - global_scale: Optional[torch.Tensor],

1872 - b_qzeros: Optional[torch.Tensor],

1873 - g_idx: Optional[torch.Tensor],

1874 - perm: Optional[torch.Tensor],

1871 + global_scale: torch.Tensor | None,

1872 + b_qzeros: torch.Tensor | None,

1873 + g_idx: torch.Tensor | None,

1874 + perm: torch.Tensor | None,

1875 workspace: torch.Tensor,

1876 sorted_token_ids: torch.Tensor,

1877 expert_ids: torch.Tensor,

@@ -1950,12 +1950,12 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe")

1950 @register_fake("_moe_C::moe_wna16_marlin_gemm")

1951 def moe_wna16_marlin_gemm_fake(

1952 input: torch.Tensor,

1953 - output: Optional[torch.Tensor],

1953 + output: torch.Tensor | None,

1954 b_qweight: torch.Tensor,

1955 b_scales: torch.Tensor,

1956 - b_qzeros: Optional[torch.Tensor],

1957 - g_idx: Optional[torch.Tensor],

1958 - perm: Optional[torch.Tensor],

1956 + b_qzeros: torch.Tensor | None,

1957 + g_idx: torch.Tensor | None,

1958 + perm: torch.Tensor | None,

1959 workspace: torch.Tensor,

1960 sorted_token_ids: torch.Tensor,

1961 expert_ids: torch.Tensor,

@@ -2068,7 +2068,7 @@ def gather_and_maybe_dequant_cache(

2068 batch_size: int,

2069 kv_cache_dtype: str,

2070 scale: torch.Tensor,

2071 - seq_starts: Optional[torch.Tensor] = None,

2071 + seq_starts: torch.Tensor | None = None,

2072 ) -> None:

2073 torch.ops._C_cache_ops.gather_and_maybe_dequant_cache(

2074 src_cache,

@@ -2088,7 +2088,7 @@ def cp_gather_cache(

2088 block_table: torch.Tensor,

2089 cu_seq_lens: torch.Tensor,

2090 batch_size: int,

2091 - seq_starts: Optional[torch.Tensor] = None,

2091 + seq_starts: torch.Tensor | None = None,

2092 ) -> None:

2093 torch.ops._C_cache_ops.cp_gather_cache(

2094 src_cache, dst, block_table, cu_seq_lens, batch_size, seq_starts

@@ -2187,9 +2187,7 @@ def free_shared_buffer(ptr: int) -> None:

2187

2188

2189 # quick all reduce

2190 -def init_custom_qr(

2191 - rank: int, world_size: int, qr_max_size: Optional[int] = None

2192 -) -> int:

2190 +def init_custom_qr(rank: int, world_size: int, qr_max_size: int | None = None) -> int:

2191 return torch.ops._C_custom_ar.init_custom_qr(rank, world_size, qr_max_size)

2192

2193

@@ -2247,7 +2245,7 @@ def flash_mla_with_kvcache(

2245 head_dim_v: int,

2246 tile_scheduler_metadata: torch.Tensor,

2247 num_splits: torch.Tensor,

2250 - softmax_scale: Optional[float] = None,

2248 + softmax_scale: float | None = None,

2249 causal: bool = False,

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

2251 """

@@ -2324,7 +2322,7 @@ if hasattr(torch.ops._C, "weight_packed_linear"):

2322 def weight_packed_linear_fake(

2323 mat1: torch.Tensor,

2324 mat2: torch.Tensor,

2327 - bias: Optional[torch.Tensor],

2325 + bias: torch.Tensor | None,

2326 is_vnni: bool,

2327 ) -> torch.Tensor:

2328 return torch.empty(

@@ -2344,11 +2342,11 @@ if hasattr(torch.ops._C, "fused_experts_cpu"):

2342 inplace: bool,

2343 use_int8_w8a8: bool,

2344 use_fp8_w8a16: bool,

2347 - w1_scale: Optional[torch.Tensor],

2348 - w2_scale: Optional[torch.Tensor],

2349 - block_size: Optional[list[int]],

2350 - a1_scale: Optional[torch.Tensor],

2351 - a2_scale: Optional[torch.Tensor],

2345 + w1_scale: torch.Tensor | None,

2346 + w2_scale: torch.Tensor | None,

2347 + block_size: list[int] | None,

2348 + a1_scale: torch.Tensor | None,

2349 + a2_scale: torch.Tensor | None,

2350 is_vnni: bool,

2351 ) -> torch.Tensor:

2352 return torch.empty_like(hidden_states)

@@ -2361,7 +2359,7 @@ if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):

2359 mat1: torch.Tensor,

2360 mat2: torch.Tensor,

2361 scales2: torch.Tensor,

2364 - bias: Optional[torch.Tensor],

2362 + bias: torch.Tensor | None,

2363 out_dtype: torch.dtype,

2364 is_vnni: bool,

2365 ) -> torch.Tensor:

@@ -2372,7 +2370,7 @@ if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):

2370

2371 class CPUDNNLGEMMHandler:

2372 def __init__(self) -> None:

2375 - self.handler: Optional[int] = None

2373 + self.handler: int | None = None

2374 self.n = -1

2375 self.k = -1

2376

@@ -2403,7 +2401,7 @@ def create_onednn_mm(

2401 def onednn_mm(

2402 dnnl_handler: CPUDNNLGEMMHandler,

2403 x: torch.Tensor,

2406 - bias: Optional[torch.Tensor],

2404 + bias: torch.Tensor | None,

2405 ) -> torch.Tensor:

2406 output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype)

2407 torch.ops._C.onednn_mm(

@@ -2431,8 +2429,8 @@ def create_onednn_scaled_mm(

2429

2430 def onednn_scaled_int8_quant(

2431 input: torch.Tensor,

2434 - scale: Optional[torch.Tensor] = None,

2435 - azp: Optional[torch.Tensor] = None,

2432 + scale: torch.Tensor | None = None,

2433 + azp: torch.Tensor | None = None,

2434 symmetric: bool = True,

2435 ):

2436 """

@@ -2471,10 +2469,10 @@ def onednn_scaled_mm(

2469 dnnl_handler: CPUDNNLGEMMHandler,

2470 x: torch.Tensor,

2471 output: torch.Tensor,

2474 - input_scale: Optional[torch.Tensor],

2475 - input_zp: Optional[torch.Tensor],

2476 - input_zp_adj: Optional[torch.Tensor],

2477 - bias: Optional[torch.Tensor],

2472 + input_scale: torch.Tensor | None,

2473 + input_zp: torch.Tensor | None,

2474 + input_zp_adj: torch.Tensor | None,

2475 + bias: torch.Tensor | None,

2476 ) -> torch.Tensor:

2477 torch.ops._C.onednn_scaled_mm(

2478 output, x, input_scale, input_zp, input_zp_adj, bias, dnnl_handler.handler

vllm/_ipex_ops.py

+19 -20

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 import torch

6

@@ -65,7 +64,7 @@ class ipex_ops:

64 context_lens: torch.Tensor,

65 block_size: int,

66 max_context_len: int,

68 - alibi_slopes: Optional[torch.Tensor],

67 + alibi_slopes: torch.Tensor | None,

68 kv_cache_dtype: str,

69 k_scale: float,

70 v_scale: float,

@@ -107,7 +106,7 @@ class ipex_ops:

106 context_lens: torch.Tensor,

107 block_size: int,

108 max_context_len: int,

110 - alibi_slopes: Optional[torch.Tensor],

109 + alibi_slopes: torch.Tensor | None,

110 kv_cache_dtype: str,

111 k_scale: float,

112 v_scale: float,

@@ -174,7 +173,7 @@ class ipex_ops:

173 out: torch.Tensor,

174 seqlen_q: torch.Tensor,

175 seqlen_k: torch.Tensor,

177 - alibi_slopes: Optional[torch.Tensor],

176 + alibi_slopes: torch.Tensor | None,

177 max_seqlen_q: int,

178 max_seqlen_k: int,

179 pdropout: float,

@@ -254,8 +253,8 @@ class ipex_ops:

253 value_cache: torch.Tensor,

254 slot_mapping: torch.Tensor,

255 kv_cache_dtype: str,

257 - k_scale: Optional[torch.Tensor] = None,

258 - v_scale: Optional[torch.Tensor] = None,

256 + k_scale: torch.Tensor | None = None,

257 + v_scale: torch.Tensor | None = None,

258 k_scale_float: float = 1.0,

259 v_scale_float: float = 1.0,

260 ) -> None:

@@ -283,10 +282,10 @@ class ipex_ops:

282 softmax_scale: float,

283 causal: bool,

284 block_table: torch.Tensor,

286 - alibi_slopes: Optional[torch.Tensor],

287 - window_size: Optional[list[int]] = None,

288 - softcap: Optional[float] = 0.0,

289 - cu_seqlens_k: Optional[torch.Tensor] = None,

285 + alibi_slopes: torch.Tensor | None,

286 + window_size: list[int] | None = None,

287 + softcap: float | None = 0.0,

288 + cu_seqlens_k: torch.Tensor | None = None,

289 # The following parameters are not used in ipex kernel currently,

290 # we keep API compatible to CUDA's.

291 scheduler_metadata=None,

@@ -295,7 +294,7 @@ class ipex_ops:

294 k_descale=None,

295 v_descale=None,

296 num_splits=0,

298 - s_aux: Optional[torch.Tensor] = None,

297 + s_aux: torch.Tensor | None = None,

298 ):

299 if cu_seqlens_k is None:

300 # cu_seqlens_k is not used in ipex kernel.

@@ -344,10 +343,10 @@ class ipex_ops:

343 cache_seqlens: torch.Tensor,

344 qkv_dtype=torch.bfloat16,

345 headdim_v=None,

347 - cu_seqlens_q: Optional[torch.Tensor] = None,

348 - cu_seqlens_k_new: Optional[torch.Tensor] = None,

349 - cache_leftpad: Optional[torch.Tensor] = None,

350 - page_size: Optional[int] = None,

346 + cu_seqlens_q: torch.Tensor | None = None,

347 + cu_seqlens_k_new: torch.Tensor | None = None,

348 + cache_leftpad: torch.Tensor | None = None,

349 + page_size: int | None = None,

350 max_seqlen_k_new=0,

351 causal=False,

352 window_size=(-1, -1), # -1 means infinite context window

@@ -382,11 +381,11 @@ class ipex_ops:

381 @staticmethod

382 def scaled_fp8_quant(

383 input: torch.Tensor,

385 - scale: Optional[torch.Tensor] = None,

386 - num_token_padding: Optional[int] = None,

387 - scale_ub: Optional[torch.Tensor] = None,

384 + scale: torch.Tensor | None = None,

385 + num_token_padding: int | None = None,

386 + scale_ub: torch.Tensor | None = None,

387 use_per_token_if_dynamic: bool = False,

389 - output: Optional[torch.Tensor] = None,

388 + output: torch.Tensor | None = None,

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

390 """

391 Quantize input tensor to FP8 and return quantized tensor and scale.

@@ -414,7 +413,7 @@ class ipex_ops:

413 """

414 # This code assumes batch_dim and num_tokens are flattened

415 assert input.ndim == 2

417 - shape: Union[tuple[int, int], torch.Size] = input.shape

416 + shape: tuple[int, int] | torch.Size = input.shape

417 out_dtype: torch.dtype = current_platform.fp8_dtype()

418 if num_token_padding:

419 shape = (max(num_token_padding, input.shape[0]), shape[1])

vllm/assets/base.py

+1 -2

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

3

4 from functools import lru_cache

5 from pathlib import Path

6 -from typing import Optional

6

7 import vllm.envs as envs

8 from vllm.connections import global_http_connection

@@ -20,7 +19,7 @@ def get_cache_dir() -> Path:

19

20

21 @lru_cache

23 -def get_vllm_public_assets(filename: str, s3_prefix: Optional[str] = None) -> Path:

22 +def get_vllm_public_assets(filename: str, s3_prefix: str | None = None) -> Path:

23 """

24 Download an asset file from ``s3://vllm-public-assets``

25 and return the path to the downloaded file.

vllm/assets/video.py

+2 -2

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

3

4 from dataclasses import dataclass

5 from functools import lru_cache

6 -from typing import Any, ClassVar, Literal, Optional

6 +from typing import Any, ClassVar, Literal

7

8 import cv2

9 import numpy as np

@@ -137,7 +137,7 @@ class VideoAsset:

137 ret = video_get_metadata(self.video_path, self.num_frames)

138 return ret

139

140 - def get_audio(self, sampling_rate: Optional[float] = None) -> npt.NDArray:

140 + def get_audio(self, sampling_rate: float | None = None) -> npt.NDArray:

141 """

142 Read audio data from the video asset, used in Qwen2.5-Omni examples.

143

vllm/attention/backends/abstract.py

+20 -20

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

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

3

4 from abc import ABC, abstractmethod

5 -from typing import Generic, Optional, Protocol, TypeVar, Union

5 +from typing import Generic, Protocol, TypeVar

6

7 import torch

8

@@ -65,7 +65,7 @@ class AttentionBackend(ABC):

65 raise NotImplementedError

66

67 @classmethod

68 - def get_supported_kernel_block_size(cls) -> list[Union[int, MultipleOf]]:

68 + def get_supported_kernel_block_size(cls) -> list[int | MultipleOf]:

69 return cls.get_impl_cls().get_supported_kernel_block_size()

70

71 @classmethod

@@ -158,18 +158,18 @@ class AttentionImpl(ABC, Generic[T]):

158 num_heads: int,

159 head_size: int,

160 scale: float,

161 - num_kv_heads: Optional[int] = None,

162 - alibi_slopes: Optional[list[float]] = None,

163 - sliding_window: Optional[int] = None,

161 + num_kv_heads: int | None = None,

162 + alibi_slopes: list[float] | None = None,

163 + sliding_window: int | None = None,

164 kv_cache_dtype: str = "auto",

165 - logits_soft_cap: Optional[float] = None,

165 + logits_soft_cap: float | None = None,

166 attn_type: str = AttentionType.DECODER,

167 - kv_sharing_target_layer_name: Optional[str] = None,

167 + kv_sharing_target_layer_name: str | None = None,

168 ) -> None:

169 raise NotImplementedError

170

171 @staticmethod

172 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

172 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

173 # TODO: implement this function for all backends.

174 return [MultipleOf(1)]

175

@@ -182,9 +182,9 @@ class AttentionImpl(ABC, Generic[T]):

182 value: torch.Tensor,

183 kv_cache: torch.Tensor,

184 attn_metadata: T,

185 - output: Optional[torch.Tensor] = None,

186 - output_scale: Optional[torch.Tensor] = None,

187 - output_block_scale: Optional[torch.Tensor] = None,

185 + output: torch.Tensor | None = None,

186 + output_scale: torch.Tensor | None = None,

187 + output_block_scale: torch.Tensor | None = None,

188 ) -> torch.Tensor:

189 raise NotImplementedError

190

@@ -208,21 +208,21 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):

208 head_size: int,

209 scale: float,

210 num_kv_heads: int,

211 - alibi_slopes: Optional[list[float]],

212 - sliding_window: Optional[int],

211 + alibi_slopes: list[float] | None,

212 + sliding_window: int | None,

213 kv_cache_dtype: str,

214 - logits_soft_cap: Optional[float],

214 + logits_soft_cap: float | None,

215 attn_type: str,

216 - kv_sharing_target_layer_name: Optional[str],

216 + kv_sharing_target_layer_name: str | None,

217 # MLA Specific Arguments

218 - q_lora_rank: Optional[int],

218 + q_lora_rank: int | None,

219 kv_lora_rank: int,

220 qk_nope_head_dim: int,

221 qk_rope_head_dim: int,

222 qk_head_dim: int,

223 v_head_dim: int,

224 kv_b_proj: ColumnParallelLinear,

225 - indexer: Optional[object] = None,

225 + indexer: object | None = None,

226 ) -> None:

227 raise NotImplementedError

228

@@ -235,9 +235,9 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):

235 k_pe: torch.Tensor,

236 kv_cache: torch.Tensor,

237 attn_metadata: T,

238 - output: Optional[torch.Tensor] = None,

239 - output_scale: Optional[torch.Tensor] = None,

240 - output_block_scale: Optional[torch.Tensor] = None,

238 + output: torch.Tensor | None = None,

239 + output_scale: torch.Tensor | None = None,

240 + output_block_scale: torch.Tensor | None = None,

241 ) -> torch.Tensor:

242 raise NotImplementedError

243

vllm/attention/backends/registry.py

+2 -3

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

3 """Attention backend registry"""

4

5 import enum

6 -from typing import Optional

6

7 from vllm.utils import resolve_obj_by_qualname

8

@@ -53,7 +52,7 @@ BACKEND_MAP = {

52 }

53

54

56 -def register_attn_backend(backend: _Backend, class_path: Optional[str] = None):

55 +def register_attn_backend(backend: _Backend, class_path: str | None = None):

56 """

57 Decorator: register a custom attention backend into BACKEND_MAPPING.

58 - If class_path is provided, use it.

@@ -98,7 +97,7 @@ def backend_to_class(backend: _Backend) -> type:

97 return resolve_obj_by_qualname(backend_class_name)

98

99

101 -def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:

100 +def backend_name_to_enum(backend_name: str) -> _Backend | None:

101 """

102 Convert a string backend name to a _Backend enum value.

103

vllm/attention/backends/utils.py

+1 -2

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

3 """Attention backend utils"""

4

5 from dataclasses import dataclass

6 -from typing import Optional

6

7 from vllm.config import ModelConfig

8 from vllm.logger import init_logger

@@ -15,7 +14,7 @@ PAD_SLOT_ID = -1

14

15 @dataclass

16 class MLADims:

18 - q_lora_rank: Optional[int]

17 + q_lora_rank: int | None

18 kv_lora_rank: int

19 qk_nope_head_dim: int

20 qk_rope_head_dim: int

vllm/attention/layer.py

+27 -26

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

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

3 """Attention layer."""

4

5 -from typing import Callable, Optional, cast

5 +from collections.abc import Callable

6 +from typing import cast

7

8 import torch

9 import torch.nn as nn

@@ -128,16 +129,16 @@ class Attention(nn.Module, AttentionLayerBase):

129 num_heads: int,

130 head_size: int,

131 scale: float,

131 - num_kv_heads: Optional[int] = None,

132 - alibi_slopes: Optional[list[float]] = None,

133 - cache_config: Optional[CacheConfig] = None,

134 - quant_config: Optional[QuantizationConfig] = None,

135 - logits_soft_cap: Optional[float] = None,

136 - per_layer_sliding_window: Optional[int] = None,

132 + num_kv_heads: int | None = None,

133 + alibi_slopes: list[float] | None = None,

134 + cache_config: CacheConfig | None = None,

135 + quant_config: QuantizationConfig | None = None,

136 + logits_soft_cap: float | None = None,

137 + per_layer_sliding_window: int | None = None,

138 prefix: str = "",

139 attn_type: str = AttentionType.DECODER,

139 - kv_sharing_target_layer_name: Optional[str] = None,

140 - attn_backend: Optional[type[AttentionBackend]] = None,

140 + kv_sharing_target_layer_name: str | None = None,

141 + attn_backend: type[AttentionBackend] | None = None,

142 **extra_impl_args,

143 ) -> None:

144 """

@@ -191,7 +192,7 @@ class Attention(nn.Module, AttentionLayerBase):

192

193 # The output scale on host memory. This should be the input scale of

194 # the quant op after this attention layer.

194 - self._o_scale_float: Optional[float] = None

195 + self._o_scale_float: float | None = None

196

197 self.num_heads = num_heads

198 self.head_size = head_size

@@ -319,7 +320,7 @@ class Attention(nn.Module, AttentionLayerBase):

320 # For some alternate attention backends like MLA the attention output

321 # shape does not match the query shape, so we optionally let the model

322 # definition specify the output tensor shape.

322 - output_shape: Optional[torch.Size] = None,

323 + output_shape: torch.Size | None = None,

324 ) -> torch.Tensor:

325 """

326 The KV cache is stored inside this class and is accessed via

@@ -427,7 +428,7 @@ class MultiHeadAttention(nn.Module):

428 num_heads: int,

429 head_size: int,

430 scale: float,

430 - num_kv_heads: Optional[int] = None,

431 + num_kv_heads: int | None = None,

432 # This has no effect, it is only here to make it easier to swap

433 # between Attention and MultiHeadAttention

434 prefix: str = "",

@@ -582,14 +583,14 @@ class MLAAttention(nn.Module, AttentionLayerBase):

583 qk_nope_head_dim: int,

584 qk_rope_head_dim: int,

585 v_head_dim: int,

585 - q_lora_rank: Optional[int],

586 + q_lora_rank: int | None,

587 kv_lora_rank: int,

588 kv_b_proj: ColumnParallelLinear,

588 - cache_config: Optional[CacheConfig] = None,

589 - quant_config: Optional[QuantizationConfig] = None,

589 + cache_config: CacheConfig | None = None,

590 + quant_config: QuantizationConfig | None = None,

591 prefix: str = "",

592 use_sparse: bool = False,

592 - indexer: Optional[object] = None,

593 + indexer: object | None = None,

594 ):

595 super().__init__()

596 self.num_heads = num_heads

@@ -670,7 +671,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

671 self._q_scale_float = 1.0

672 self._k_scale_float = 1.0

673 self._v_scale_float = 1.0

673 - self._o_scale_float: Optional[float] = None

674 + self._o_scale_float: float | None = None

675

676 self.use_sparse = use_sparse

677

@@ -688,7 +689,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

689 q: torch.Tensor,

690 kv_c_normed: torch.Tensor,

691 k_pe: torch.Tensor,

691 - output_shape: Optional[torch.Size] = None,

692 + output_shape: torch.Size | None = None,

693 ) -> torch.Tensor:

694 if self.use_direct_call:

695 forward_context: ForwardContext = get_forward_context()

@@ -888,8 +889,8 @@ def unified_attention_with_output(

889 value: torch.Tensor,

890 output: torch.Tensor,

891 layer_name: str,

891 - output_scale: Optional[torch.Tensor] = None,

892 - output_block_scale: Optional[torch.Tensor] = None,

892 + output_scale: torch.Tensor | None = None,

893 + output_block_scale: torch.Tensor | None = None,

894 ) -> None:

895 wait_for_kv_layer_from_connector(layer_name)

896 forward_context: ForwardContext = get_forward_context()

@@ -919,8 +920,8 @@ def unified_attention_with_output_fake(

920 value: torch.Tensor,

921 output: torch.Tensor,

922 layer_name: str,

922 - output_scale: Optional[torch.Tensor] = None,

923 - output_block_scale: Optional[torch.Tensor] = None,

923 + output_scale: torch.Tensor | None = None,

924 + output_block_scale: torch.Tensor | None = None,

925 ) -> None:

926 return

927

@@ -978,8 +979,8 @@ def unified_mla_attention_with_output(

979 k_pe: torch.Tensor,

980 output: torch.Tensor,

981 layer_name: str,

981 - output_scale: Optional[torch.Tensor] = None,

982 - output_block_scale: Optional[torch.Tensor] = None,

982 + output_scale: torch.Tensor | None = None,

983 + output_block_scale: torch.Tensor | None = None,

984 ) -> None:

985 wait_for_kv_layer_from_connector(layer_name)

986 forward_context: ForwardContext = get_forward_context()

@@ -1009,8 +1010,8 @@ def unified_mla_attention_with_output_fake(

1010 k_pe: torch.Tensor,

1011 output: torch.Tensor,

1012 layer_name: str,

1012 - output_scale: Optional[torch.Tensor] = None,

1013 - output_block_scale: Optional[torch.Tensor] = None,

1013 + output_scale: torch.Tensor | None = None,

1014 + output_block_scale: torch.Tensor | None = None,

1015 ) -> None:

1016 return

1017

vllm/attention/layers/chunked_local_attention.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import functools

4 -from typing import ClassVar, Optional

4 +from typing import ClassVar

5

6 import torch

7

@@ -60,11 +60,11 @@ class ChunkedLocalAttention(Attention):

60 head_size: int,

61 scale: float,

62 attention_chunk_size: int,

63 - num_kv_heads: Optional[int] = None,

64 - alibi_slopes: Optional[list[float]] = None,

65 - cache_config: Optional[CacheConfig] = None,

66 - quant_config: Optional[QuantizationConfig] = None,

67 - kv_sharing_target_layer_name: Optional[str] = None,

63 + num_kv_heads: int | None = None,

64 + alibi_slopes: list[float] | None = None,

65 + cache_config: CacheConfig | None = None,

66 + quant_config: QuantizationConfig | None = None,

67 + kv_sharing_target_layer_name: str | None = None,

68 prefix: str = "",

69 ):

70 dtype = torch.get_default_dtype()

vllm/attention/layers/cross_attention.py

+2 -3

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

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

3 import functools

4 from copy import copy

5 -from typing import Optional

5

6 import numpy as np

7 import torch

@@ -138,8 +137,8 @@ class CrossAttention(Attention):

137 num_heads: int,

138 head_size: int,

139 scale: float,

141 - cache_config: Optional[CacheConfig] = None,

142 - attn_type: Optional[str] = None,

140 + cache_config: CacheConfig | None = None,

141 + attn_type: str | None = None,

142 **kwargs,

143 ):

144 dtype = torch.get_default_dtype()

vllm/attention/layers/encoder_only_attention.py

+2 -3

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

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

3 import functools

4 from copy import copy

5 -from typing import Optional

5

6 import torch

7

@@ -60,8 +59,8 @@ class EncoderOnlyAttention(Attention):

59 num_heads: int,

60 head_size: int,

61 scale: float,

63 - cache_config: Optional[CacheConfig] = None,

64 - attn_type: Optional[str] = None,

62 + cache_config: CacheConfig | None = None,

63 + attn_type: str | None = None,

64 **kwargs,

65 ):

66 dtype = torch.get_default_dtype()

vllm/attention/ops/flashmla.py

+9 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 # adapted from: https://github.com/deepseek-ai/FlashMLA/blob/main/flash_mla/flash_mla_interface.py

4 -from typing import Optional

4

5 import torch

6

@@ -31,7 +30,7 @@ else:

30 _flashmla_extension_C_AVAILABLE = False

31

32

34 -def _is_flashmla_available() -> tuple[bool, Optional[str]]:

33 +def _is_flashmla_available() -> tuple[bool, str | None]:

34 if not _flashmla_C_AVAILABLE:

35 return (

36 False,

@@ -49,7 +48,7 @@ def _is_flashmla_available() -> tuple[bool, Optional[str]]:

48 return True, None

49

50

52 -def is_flashmla_dense_supported() -> tuple[bool, Optional[str]]:

51 +def is_flashmla_dense_supported() -> tuple[bool, str | None]:

52 """

53 Return: is_supported_flag, unsupported_reason (optional).

54 """

@@ -61,7 +60,7 @@ def is_flashmla_dense_supported() -> tuple[bool, Optional[str]]:

60 return True, None

61

62

64 -def is_flashmla_sparse_supported() -> tuple[bool, Optional[str]]:

63 +def is_flashmla_sparse_supported() -> tuple[bool, str | None]:

64 """

65 Return: is_supported_flag, unsupported_reason (optional).

66 """

@@ -80,9 +79,9 @@ def get_mla_metadata(

79 cache_seqlens: torch.Tensor,

80 num_q_tokens_per_head_k: int,

81 num_heads_k: int,

83 - num_heads_q: Optional[int] = None,

82 + num_heads_q: int | None = None,

83 is_fp8_kvcache: bool = False,

85 - topk: Optional[int] = None,

84 + topk: int | None = None,

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

86 """

87 Arguments:

@@ -121,12 +120,12 @@ def flash_mla_with_kvcache(

120 head_dim_v: int,

121 tile_scheduler_metadata: torch.Tensor,

122 num_splits: torch.Tensor,

124 - softmax_scale: Optional[float] = None,

123 + softmax_scale: float | None = None,

124 causal: bool = False,

126 - descale_q: Optional[torch.Tensor] = None,

127 - descale_k: Optional[torch.Tensor] = None,

125 + descale_q: torch.Tensor | None = None,

126 + descale_k: torch.Tensor | None = None,

127 is_fp8_kvcache: bool = False,

129 - indices: Optional[torch.Tensor] = None,

128 + indices: torch.Tensor | None = None,

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

130 """

131 Arguments:

vllm/attention/ops/merge_attn_states.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -13,7 +12,7 @@ def merge_attn_states(

12 prefix_lse: torch.Tensor,

13 suffix_output: torch.Tensor,

14 suffix_lse: torch.Tensor,

16 - output_lse: Optional[torch.Tensor] = None,

15 + output_lse: torch.Tensor | None = None,

16 ) -> None:

17 # NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel

18 # is not support for FP8 dtype, fallback to use Triton kernel.

vllm/attention/ops/paged_attn.py

+5 -6

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

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

3

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import torch

7

@@ -27,7 +26,7 @@ class PagedAttentionMetadata:

26

27 # (batch_size,). The length of sequences (entire tokens seen so far) per

28 # sequence.

30 - seq_lens_tensor: Optional[torch.Tensor]

29 + seq_lens_tensor: torch.Tensor | None

30 # Maximum sequence length in the batch. 0 if it is prefill-only batch.

31 max_decode_seq_len: int

32 # (batch_size, max_blocks_per_seq).

@@ -36,7 +35,7 @@ class PagedAttentionMetadata:

35 # in the kv cache. Each block can contain up to block_size tokens.

36 # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph

37 # captured.

39 - block_tables: Optional[torch.Tensor]

38 + block_tables: torch.Tensor | None

39

40

41 class PagedAttention:

@@ -102,7 +101,7 @@ class PagedAttention:

101 kv_cache_dtype: str,

102 num_kv_heads: int,

103 scale: float,

105 - alibi_slopes: Optional[torch.Tensor],

104 + alibi_slopes: torch.Tensor | None,

105 k_scale: torch.Tensor,

106 v_scale: torch.Tensor,

107 tp_rank: int = 0,

@@ -211,8 +210,8 @@ class PagedAttention:

210 query_start_loc: torch.Tensor,

211 seq_lens_tensor: torch.Tensor,

212 max_query_len: int,

214 - alibi_slopes: Optional[torch.Tensor],

215 - sliding_window: Optional[int],

213 + alibi_slopes: torch.Tensor | None,

214 + sliding_window: int | None,

215 k_scale: torch.Tensor,

216 v_scale: torch.Tensor,

217 ) -> torch.Tensor:

vllm/attention/ops/rocm_aiter_mla.py

+9 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6

@@ -30,9 +29,9 @@ def aiter_mla_decode_fwd(

29 sm_scale: float,

30 qo_indptr: torch.Tensor,

31 max_seqlen_qo: int,

33 - kv_indptr: Optional[torch.Tensor] = None,

34 - kv_indices: Optional[torch.Tensor] = None,

35 - kv_last_page_lens: Optional[torch.Tensor] = None,

32 + kv_indptr: torch.Tensor | None = None,

33 + kv_indices: torch.Tensor | None = None,

34 + kv_last_page_lens: torch.Tensor | None = None,

35 logit_cap: float = 0.0,

36 ):

37 torch.ops.vllm.rocm_aiter_mla_decode_fwd(

@@ -55,9 +54,9 @@ def mla_decode_fwd_impl(

54 o: torch.Tensor,

55 qo_indptr: torch.Tensor,

56 max_seqlen_qo: int,

58 - kv_indptr: Optional[torch.Tensor] = None,

59 - kv_indices: Optional[torch.Tensor] = None,

60 - kv_last_page_lens: Optional[torch.Tensor] = None,

57 + kv_indptr: torch.Tensor | None = None,

58 + kv_indices: torch.Tensor | None = None,

59 + kv_last_page_lens: torch.Tensor | None = None,

60 sm_scale: float = 1.0,

61 logit_cap: float = 0.0,

62 ) -> None:

@@ -83,9 +82,9 @@ def mla_decode_fwd_fake(

82 o: torch.Tensor,

83 qo_indptr: torch.Tensor,

84 max_seqlen_qo: int,

86 - kv_indptr: Optional[torch.Tensor] = None,

87 - kv_indices: Optional[torch.Tensor] = None,

88 - kv_last_page_lens: Optional[torch.Tensor] = None,

85 + kv_indptr: torch.Tensor | None = None,

86 + kv_indices: torch.Tensor | None = None,

87 + kv_last_page_lens: torch.Tensor | None = None,

88 sm_scale: float = 1.0,

89 logit_cap: float = 0.0,

90 ) -> None:

vllm/attention/ops/rocm_aiter_paged_attn.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import aiter as rocm_aiter

5 import torch

@@ -62,7 +61,7 @@ class AITERPagedAttention(PagedAttention):

61 kv_cache_dtype: str,

62 num_kv_heads: int,

63 scale: float,

65 - alibi_slopes: Optional[torch.Tensor],

64 + alibi_slopes: torch.Tensor | None,

65 k_scale: torch.Tensor,

66 v_scale: torch.Tensor,

67 tp_rank: int = 0,

vllm/attention/ops/triton_merge_attn_states.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -15,7 +14,7 @@ def merge_attn_states(

14 prefix_lse: torch.Tensor,

15 suffix_output: torch.Tensor,

16 suffix_lse: torch.Tensor,

18 - output_lse: Optional[torch.Tensor] = None,

17 + output_lse: torch.Tensor | None = None,

18 ) -> None:

19 num_tokens = output.shape[0]

20 num_query_heads = output.shape[1]

vllm/attention/selector.py

+9 -10

@@ -6,7 +6,6 @@ from collections.abc import Generator

6 from contextlib import contextmanager

7 from dataclasses import dataclass

8 from functools import cache

9 -from typing import Optional, Union

9

10 import torch

11

@@ -19,7 +18,7 @@ from vllm.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname

18 logger = init_logger(__name__)

19

20

22 -def get_env_variable_attn_backend() -> Optional[_Backend]:

21 +def get_env_variable_attn_backend() -> _Backend | None:

22 """

23 Get the backend override specified by the vLLM attention

24 backend environment variable, if one is specified.

@@ -40,10 +39,10 @@ def get_env_variable_attn_backend() -> Optional[_Backend]:

39 #

40 # THIS SELECTION TAKES PRECEDENCE OVER THE

41 # VLLM_ATTENTION_BACKEND ENVIRONMENT VARIABLE

43 -forced_attn_backend: Optional[_Backend] = None

42 +forced_attn_backend: _Backend | None = None

43

44

46 -def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:

45 +def global_force_attn_backend(attn_backend: _Backend | None) -> None:

46 """

47 Force all attention operations to use a specified backend.

48

@@ -58,7 +57,7 @@ def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:

57 forced_attn_backend = attn_backend

58

59

61 -def get_global_forced_attn_backend() -> Optional[_Backend]:

60 +def get_global_forced_attn_backend() -> _Backend | None:

61 """

62 Get the currently-forced choice of attention backend,

63 or None if auto-selection is currently enabled.

@@ -77,7 +76,7 @@ class _IsSupported:

76

77

78 def is_attn_backend_supported(

80 - attn_backend: Union[str, type[AttentionBackend]],

79 + attn_backend: str | type[AttentionBackend],

80 head_size: int,

81 dtype: torch.dtype,

82 *,

@@ -127,7 +126,7 @@ def is_attn_backend_supported(

126 def get_attn_backend(

127 head_size: int,

128 dtype: torch.dtype,

130 - kv_cache_dtype: Optional[str],

129 + kv_cache_dtype: str | None,

130 block_size: int,

131 use_mla: bool = False,

132 has_sink: bool = False,

@@ -154,7 +153,7 @@ def get_attn_backend(

153 def _cached_get_attn_backend(

154 head_size: int,

155 dtype: torch.dtype,

157 - kv_cache_dtype: Optional[str],

156 + kv_cache_dtype: str | None,

157 block_size: int,

158 use_v1: bool = False,

159 use_mla: bool = False,

@@ -167,12 +166,12 @@ def _cached_get_attn_backend(

166 # THIS SELECTION OVERRIDES THE VLLM_ATTENTION_BACKEND

167 # ENVIRONMENT VARIABLE.

168 selected_backend = None

170 - backend_by_global_setting: Optional[_Backend] = get_global_forced_attn_backend()

169 + backend_by_global_setting: _Backend | None = get_global_forced_attn_backend()

170 if backend_by_global_setting is not None:

171 selected_backend = backend_by_global_setting

172 else:

173 # Check the environment variable and override if specified

175 - backend_by_env_var: Optional[str] = envs.VLLM_ATTENTION_BACKEND

174 + backend_by_env_var: str | None = envs.VLLM_ATTENTION_BACKEND

175 if backend_by_env_var is not None:

176 if backend_by_env_var.endswith("_VLLM_V1"):

177 logger.warning(

vllm/attention/utils/fa_utils.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 from vllm import envs

5 from vllm.logger import init_logger

@@ -21,7 +20,7 @@ elif current_platform.is_xpu():

20 get_scheduler_metadata = ops.get_scheduler_metadata

21

22

24 -def get_flash_attn_version(requires_alibi: bool = False) -> Optional[int]:

23 +def get_flash_attn_version(requires_alibi: bool = False) -> int | None:

24 # import here to avoid circular dependencies

25 from vllm.platforms import current_platform

26

vllm/beam_search.py

+8 -8

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

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

3

4 from dataclasses import dataclass

5 -from typing import TYPE_CHECKING, Any, Optional, Union

5 +from typing import TYPE_CHECKING, Any, Optional

6

7 from vllm.logprobs import Logprob

8 from vllm.lora.request import LoRARequest

@@ -22,13 +22,13 @@ class BeamSearchSequence:

22 # The tokens include the prompt.

23 tokens: list[int]

24 logprobs: list[dict[int, Logprob]]

25 - lora_request: Optional[LoRARequest] = None

25 + lora_request: LoRARequest | None = None

26 cum_logprob: float = 0.0

27 - text: Optional[str] = None

28 - finish_reason: Optional[str] = None

29 - stop_reason: Union[int, str, None] = None

27 + text: str | None = None

28 + finish_reason: str | None = None

29 + stop_reason: int | str | None = None

30 multi_modal_data: Optional["MultiModalDataDict"] = None

31 - mm_processor_kwargs: Optional[dict[str, Any]] = None

31 + mm_processor_kwargs: dict[str, Any] | None = None

32

33

34 @dataclass

@@ -45,8 +45,8 @@ class BeamSearchInstance:

45 def __init__(

46 self,

47 prompt_tokens: list[int],

48 - lora_request: Optional[LoRARequest] = None,

49 - logprobs: Optional[list[dict[int, Logprob]]] = None,

48 + lora_request: LoRARequest | None = None,

49 + logprobs: list[dict[int, Logprob]] | None = None,

50 **kwargs,

51 ):

52 self.beams: list[BeamSearchSequence] = [

vllm/benchmarks/datasets.py

+33 -33

@@ -21,13 +21,13 @@ import logging

21 import math

22 import random

23 from abc import ABC, abstractmethod

24 -from collections.abc import Iterator, Mapping

24 +from collections.abc import Callable, Iterator, Mapping

25 from contextlib import suppress

26 from copy import deepcopy

27 from dataclasses import dataclass

28 from functools import cache

29 from io import BytesIO

30 -from typing import Any, Callable, Optional, Union, cast

30 +from typing import Any, cast

31

32 import numpy as np

33 from PIL import Image

@@ -75,12 +75,12 @@ class SampleRequest:

75 Represents a single inference request for benchmarking.

76 """

77

78 - prompt: Union[str, list[str]]

78 + prompt: str | list[str]

79 prompt_len: int

80 expected_output_len: int

81 - multi_modal_data: Optional[Union[MultiModalDataDict, dict, list[dict]]] = None

82 - lora_request: Optional[LoRARequest] = None

83 - request_id: Optional[str] = None

81 + multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None

82 + lora_request: LoRARequest | None = None

83 + request_id: str | None = None

84

85

86 # -----------------------------------------------------------------------------

@@ -94,7 +94,7 @@ class BenchmarkDataset(ABC):

94

95 def __init__(

96 self,

97 - dataset_path: Optional[str] = None,

97 + dataset_path: str | None = None,

98 random_seed: int = DEFAULT_SEED,

99 disable_shuffle: bool = False,

100 **kwargs,

@@ -119,7 +119,7 @@ class BenchmarkDataset(ABC):

119 def apply_multimodal_chat_transformation(

120 self,

121 prompt: str,

122 - mm_content: Optional[Union[MultiModalDataDict, dict, list[dict]]] = None,

122 + mm_content: MultiModalDataDict | dict | list[dict] | None = None,

123 ) -> list[dict]:

124 """

125 Transform a prompt and optional multimodal content into a chat format.

@@ -154,9 +154,9 @@ class BenchmarkDataset(ABC):

154

155 def get_random_lora_request(

156 self,

157 - max_loras: Optional[int] = None,

158 - lora_path: Optional[str] = None,

159 - ) -> Optional[LoRARequest]:

157 + max_loras: int | None = None,

158 + lora_path: str | None = None,

159 + ) -> LoRARequest | None:

160 """

161 Optionally select a random LoRA request.

162

@@ -384,7 +384,7 @@ def gen_prompt_decode_to_target_len(

384 target_token_len: int,

385 max_retry: int = 10,

386 add_special_tokens: bool = False,

387 - rng: Optional[np.random.Generator] = None,

387 + rng: np.random.Generator | None = None,

388 ) -> tuple[str, list[int]]:

389 """

390 Ensure decoded-then-encoded prompt length matches the target token length.

@@ -1054,9 +1054,9 @@ class ShareGPTDataset(BenchmarkDataset):

1054 self,

1055 tokenizer: PreTrainedTokenizerBase,

1056 num_requests: int,

1057 - lora_path: Optional[str] = None,

1058 - max_loras: Optional[int] = None,

1059 - output_len: Optional[int] = None,

1057 + lora_path: str | None = None,

1058 + max_loras: int | None = None,

1059 + output_len: int | None = None,

1060 enable_multimodal_chat: bool = False,

1061 request_id_prefix: str = "",

1062 no_oversample: bool = False,

@@ -1766,9 +1766,9 @@ class CustomDataset(BenchmarkDataset):

1766 self,

1767 tokenizer: PreTrainedTokenizerBase,

1768 num_requests: int,

1769 - lora_path: Optional[str] = None,

1770 - max_loras: Optional[int] = None,

1771 - output_len: Optional[int] = None,

1769 + lora_path: str | None = None,

1770 + max_loras: int | None = None,

1771 + output_len: int | None = None,

1772 enable_multimodal_chat: bool = False,

1773 skip_chat_template: bool = False,

1774 request_id_prefix: str = "",

@@ -1997,8 +1997,8 @@ class BurstGPTDataset(BenchmarkDataset):

1997 self,

1998 tokenizer: PreTrainedTokenizerBase,

1999 num_requests: int,

2000 - max_loras: Optional[int] = None,

2001 - lora_path: Optional[str] = None,

2000 + max_loras: int | None = None,

2001 + lora_path: str | None = None,

2002 request_id_prefix: str = "",

2003 no_oversample: bool = False,

2004 **kwargs,

@@ -2034,15 +2034,15 @@ class BurstGPTDataset(BenchmarkDataset):

2034 class HuggingFaceDataset(BenchmarkDataset):

2035 """Base class for datasets hosted on HuggingFace."""

2036

2037 - SUPPORTED_DATASET_PATHS: Union[set[str], dict[str, Callable]] = set()

2037 + SUPPORTED_DATASET_PATHS: set[str] | dict[str, Callable] = set()

2038

2039 def __init__(

2040 self,

2041 dataset_path: str,

2042 dataset_split: str,

2043 no_stream: bool = False,

2044 - dataset_subset: Optional[str] = None,

2045 - hf_name: Optional[str] = None,

2044 + dataset_subset: str | None = None,

2045 + hf_name: str | None = None,

2046 **kwargs,

2047 ) -> None:

2048 super().__init__(dataset_path=dataset_path, **kwargs)

@@ -2083,7 +2083,7 @@ class ConversationDataset(HuggingFaceDataset):

2083 self,

2084 tokenizer: PreTrainedTokenizerBase,

2085 num_requests: int,

2086 - output_len: Optional[int] = None,

2086 + output_len: int | None = None,

2087 enable_multimodal_chat: bool = False,

2088 request_id_prefix: str = "",

2089 no_oversample: bool = False,

@@ -2152,7 +2152,7 @@ class VisionArenaDataset(HuggingFaceDataset):

2152 self,

2153 tokenizer: PreTrainedTokenizerBase,

2154 num_requests: int,

2155 - output_len: Optional[int] = None,

2155 + output_len: int | None = None,

2156 enable_multimodal_chat: bool = False,

2157 request_id_prefix: str = "",

2158 no_oversample: bool = False,

@@ -2206,7 +2206,7 @@ class MMVUDataset(HuggingFaceDataset):

2206 self,

2207 tokenizer: PreTrainedTokenizerBase,

2208 num_requests: int,

2209 - output_len: Optional[int] = None,

2209 + output_len: int | None = None,

2210 enable_multimodal_chat: bool = False,

2211 request_id_prefix: str = "",

2212 no_oversample: bool = False,

@@ -2267,7 +2267,7 @@ class InstructCoderDataset(HuggingFaceDataset):

2267 self,

2268 tokenizer: PreTrainedTokenizerBase,

2269 num_requests: int,

2270 - output_len: Optional[int] = None,

2270 + output_len: int | None = None,

2271 enable_multimodal_chat: bool = False,

2272 skip_chat_template: bool = False,

2273 request_id_prefix: str = "",

@@ -2331,7 +2331,7 @@ class MTBenchDataset(HuggingFaceDataset):

2331 self,

2332 tokenizer: PreTrainedTokenizerBase,

2333 num_requests: int,

2334 - output_len: Optional[int] = None,

2334 + output_len: int | None = None,

2335 enable_multimodal_chat: bool = False,

2336 skip_chat_template: bool = False,

2337 request_id_prefix: str = "",

@@ -2397,7 +2397,7 @@ class BlazeditDataset(HuggingFaceDataset):

2397 self,

2398 tokenizer: PreTrainedTokenizerBase,

2399 num_requests: int,

2400 - output_len: Optional[int] = None,

2400 + output_len: int | None = None,

2401 skip_chat_template: bool = False,

2402 request_id_prefix: str = "",

2403 no_oversample: bool = False,

@@ -2478,7 +2478,7 @@ class AIMODataset(HuggingFaceDataset):

2478 self,

2479 tokenizer: PreTrainedTokenizerBase,

2480 num_requests: int,

2481 - output_len: Optional[int] = None,

2481 + output_len: int | None = None,

2482 request_id_prefix: str = "",

2483 no_oversample: bool = False,

2484 **kwargs,

@@ -2660,7 +2660,7 @@ class ASRDataset(HuggingFaceDataset):

2660 self,

2661 tokenizer: PreTrainedTokenizerBase,

2662 num_requests: int,

2663 - output_len: Optional[int] = None,

2663 + output_len: int | None = None,

2664 request_id_prefix: str = "",

2665 no_oversample: bool = False,

2666 **kwargs,

@@ -2738,7 +2738,7 @@ class MLPerfDataset(HuggingFaceDataset):

2738 self,

2739 tokenizer: PreTrainedTokenizerBase,

2740 num_requests: int,

2741 - output_len: Optional[int] = None,

2741 + output_len: int | None = None,

2742 request_id_prefix: str = "",

2743 no_oversample: bool = False,

2744 **kwargs,

@@ -2902,7 +2902,7 @@ class MMStarDataset(HuggingFaceDataset):

2902 self,

2903 tokenizer: PreTrainedTokenizerBase,

2904 num_requests: int,

2905 - output_len: Optional[int] = None,

2905 + output_len: int | None = None,

2906 enable_multimodal_chat: bool = False,

2907 request_id_prefix: str = "",

2908 no_oversample: bool = False,

vllm/benchmarks/latency.py

+2 -2

@@ -7,7 +7,7 @@ import dataclasses

7 import json

8 import os

9 import time

10 -from typing import Any, Optional

10 +from typing import Any

11

12 import numpy as np

13 from tqdm import tqdm

@@ -127,7 +127,7 @@ def main(args: argparse.Namespace):

127 ),

128 )

129

130 - def run_to_completion(profile_dir: Optional[str] = None):

130 + def run_to_completion(profile_dir: str | None = None):

131 if profile_dir:

132 llm.start_profile()

133 llm_generate()

vllm/benchmarks/lib/endpoint_request_func.py

+20 -20

@@ -10,7 +10,7 @@ import time

10 import traceback

11 from collections.abc import Awaitable

12 from dataclasses import dataclass, field

13 -from typing import Any, Literal, Optional, Protocol, Union

13 +from typing import Any, Literal, Protocol

14

15 import aiohttp

16 import regex as re

@@ -69,14 +69,14 @@ class RequestFuncInput:

69 prompt_len: int

70 output_len: int

71 model: str

72 - model_name: Optional[str] = None

73 - logprobs: Optional[int] = None

74 - extra_headers: Optional[dict] = None

75 - extra_body: Optional[dict] = None

76 - multi_modal_content: Optional[Union[dict, list[dict]]] = None

72 + model_name: str | None = None

73 + logprobs: int | None = None

74 + extra_headers: dict | None = None

75 + extra_body: dict | None = None

76 + multi_modal_content: dict | list[dict] | None = None

77 ignore_eos: bool = False

78 - language: Optional[str] = None

79 - request_id: Optional[str] = None

78 + language: str | None = None

79 + request_id: str | None = None

80

81

82 @dataclass

@@ -100,14 +100,14 @@ class RequestFunc(Protocol):

100 self,

101 request_func_input: RequestFuncInput,

102 session: aiohttp.ClientSession,

103 - pbar: Optional[tqdm] = None,

103 + pbar: tqdm | None = None,

104 ) -> Awaitable[RequestFuncOutput]: ...

105

106

107 def _validate_api_url(

108 api_url: str,

109 api_name: str,

110 - expected_suffixes: Union[str, set[str]],

110 + expected_suffixes: str | set[str],

111 ) -> None:

112 if isinstance(expected_suffixes, str):

113 expected_suffixes = {expected_suffixes}

@@ -141,7 +141,7 @@ def _update_headers_common(

141 async def async_request_openai_completions(

142 request_func_input: RequestFuncInput,

143 session: aiohttp.ClientSession,

144 - pbar: Optional[tqdm] = None,

144 + pbar: tqdm | None = None,

145 ) -> RequestFuncOutput:

146 """The async request function for the OpenAI Completions API.

147

@@ -279,7 +279,7 @@ def _get_chat_content(

279 async def async_request_openai_chat_completions(

280 request_func_input: RequestFuncInput,

281 session: aiohttp.ClientSession,

282 - pbar: Optional[tqdm] = None,

282 + pbar: tqdm | None = None,

283 mm_position: Literal["first", "last"] = "last",

284 ) -> RequestFuncOutput:

285 api_url = request_func_input.api_url

@@ -376,7 +376,7 @@ async def async_request_openai_chat_completions(

376 async def async_request_openai_audio(

377 request_func_input: RequestFuncInput,

378 session: aiohttp.ClientSession,

379 - pbar: Optional[tqdm] = None,

379 + pbar: tqdm | None = None,

380 ) -> RequestFuncOutput:

381 # Lazy import without PlaceholderModule to avoid vllm dep.

382 import soundfile

@@ -489,7 +489,7 @@ async def _run_openai_embeddings(

489 api_url: str,

490 payload: dict[str, Any],

491 headers: dict[str, Any],

492 - pbar: Optional[tqdm] = None,

492 + pbar: tqdm | None = None,

493 ) -> RequestFuncOutput:

494 output = RequestFuncOutput()

495 st = time.perf_counter()

@@ -517,7 +517,7 @@ async def _run_openai_embeddings(

517 async def async_request_openai_embeddings(

518 request_func_input: RequestFuncInput,

519 session: aiohttp.ClientSession,

520 - pbar: Optional[tqdm] = None,

520 + pbar: tqdm | None = None,

521 ) -> RequestFuncOutput:

522 api_url = request_func_input.api_url

523 _validate_api_url(api_url, "OpenAI Embeddings API", "embeddings")

@@ -548,7 +548,7 @@ async def async_request_openai_embeddings(

548 async def async_request_openai_embeddings_chat(

549 request_func_input: RequestFuncInput,

550 session: aiohttp.ClientSession,

551 - pbar: Optional[tqdm] = None,

551 + pbar: tqdm | None = None,

552 mm_position: Literal["first", "last"] = "last",

553 ) -> RequestFuncOutput:

554 api_url = request_func_input.api_url

@@ -627,7 +627,7 @@ def _preprocess_vlm2vec(request_func_input: RequestFuncInput):

627 async def async_request_openai_embeddings_clip(

628 request_func_input: RequestFuncInput,

629 session: aiohttp.ClientSession,

630 - pbar: Optional[tqdm] = None,

630 + pbar: tqdm | None = None,

631 ) -> RequestFuncOutput:

632 _preprocess_clip(request_func_input)

633

@@ -641,7 +641,7 @@ async def async_request_openai_embeddings_clip(

641 async def async_request_openai_embeddings_vlm2vec(

642 request_func_input: RequestFuncInput,

643 session: aiohttp.ClientSession,

644 - pbar: Optional[tqdm] = None,

644 + pbar: tqdm | None = None,

645 ) -> RequestFuncOutput:

646 _preprocess_vlm2vec(request_func_input)

647

@@ -656,7 +656,7 @@ async def async_request_openai_embeddings_vlm2vec(

656 async def async_request_infinity_embeddings(

657 request_func_input: RequestFuncInput,

658 session: aiohttp.ClientSession,

659 - pbar: Optional[tqdm] = None,

659 + pbar: tqdm | None = None,

660 ) -> RequestFuncOutput:

661 api_url = request_func_input.api_url

662 _validate_api_url(api_url, "Infinity Embeddings API", "embeddings")

@@ -697,7 +697,7 @@ async def async_request_infinity_embeddings(

697 async def async_request_infinity_embeddings_clip(

698 request_func_input: RequestFuncInput,

699 session: aiohttp.ClientSession,

700 - pbar: Optional[tqdm] = None,

700 + pbar: tqdm | None = None,

701 ) -> RequestFuncOutput:

702 _preprocess_clip(request_func_input)

703

vllm/benchmarks/serve.py

+15 -15

@@ -31,7 +31,7 @@ from collections.abc import AsyncGenerator, Iterable

31 from dataclasses import dataclass

32 from datetime import datetime

33 from enum import Enum

34 -from typing import Any, Literal, Optional

34 +from typing import Any, Literal

35

36 import aiohttp

37 import numpy as np

@@ -107,9 +107,9 @@ class EmbedBenchmarkMetrics:

107

108

109 def _get_current_request_rate(

110 - ramp_up_strategy: Optional[Literal["linear", "exponential"]],

111 - ramp_up_start_rps: Optional[int],

112 - ramp_up_end_rps: Optional[int],

110 + ramp_up_strategy: Literal["linear", "exponential"] | None,

111 + ramp_up_start_rps: int | None,

112 + ramp_up_end_rps: int | None,

113 request_index: int,

114 total_requests: int,

115 request_rate: float,

@@ -135,9 +135,9 @@ async def get_request(

135 input_requests: list[SampleRequest],

136 request_rate: float,

137 burstiness: float = 1.0,

138 - ramp_up_strategy: Optional[Literal["linear", "exponential"]] = None,

139 - ramp_up_start_rps: Optional[int] = None,

140 - ramp_up_end_rps: Optional[int] = None,

138 + ramp_up_strategy: Literal["linear", "exponential"] | None = None,

139 + ramp_up_start_rps: int | None = None,

140 + ramp_up_end_rps: int | None = None,

141 ) -> AsyncGenerator[tuple[SampleRequest, float], None]:

142 """

143 Asynchronously generates requests at a specified rate

@@ -474,7 +474,7 @@ async def benchmark(

474 model_name: str,

475 tokenizer: PreTrainedTokenizerBase,

476 input_requests: list[SampleRequest],

477 - logprobs: Optional[int],

477 + logprobs: int | None,

478 request_rate: float,

479 burstiness: float,

480 disable_tqdm: bool,

@@ -483,13 +483,13 @@ async def benchmark(

483 selected_percentiles: list[float],

484 ignore_eos: bool,

485 goodput_config_dict: dict[str, float],

486 - max_concurrency: Optional[int],

487 - lora_modules: Optional[Iterable[str]],

488 - extra_headers: Optional[dict],

489 - extra_body: Optional[dict],

490 - ramp_up_strategy: Optional[Literal["linear", "exponential"]] = None,

491 - ramp_up_start_rps: Optional[int] = None,

492 - ramp_up_end_rps: Optional[int] = None,

486 + max_concurrency: int | None,

487 + lora_modules: Iterable[str] | None,

488 + extra_headers: dict | None,

489 + extra_body: dict | None,

490 + ramp_up_strategy: Literal["linear", "exponential"] | None = None,

491 + ramp_up_start_rps: int | None = None,

492 + ramp_up_end_rps: int | None = None,

493 ready_check_timeout_sec: int = 600,

494 ):

495 try:

vllm/benchmarks/throughput.py

+7 -7

@@ -9,7 +9,7 @@ import os

9 import random

10 import time

11 import warnings

12 -from typing import Any, Optional, Union

12 +from typing import Any

13

14 import torch

15 import uvloop

@@ -43,7 +43,7 @@ def run_vllm(

43 engine_args: EngineArgs,

44 do_profile: bool,

45 disable_detokenize: bool = False,

46 -) -> tuple[float, Optional[list[RequestOutput]]]:

46 +) -> tuple[float, list[RequestOutput] | None]:

47 from vllm import LLM, SamplingParams

48

49 llm = LLM(**dataclasses.asdict(engine_args))

@@ -56,7 +56,7 @@ def run_vllm(

56 " prompt_len and expected_output_len for all requests."

57 )

58 # Add the requests to the engine.

59 - prompts: list[Union[TextPrompt, TokensPrompt]] = []

59 + prompts: list[TextPrompt | TokensPrompt] = []

60 sampling_params: list[SamplingParams] = []

61 for request in requests:

62 prompt = (

@@ -79,7 +79,7 @@ def run_vllm(

79 detokenize=not disable_detokenize,

80 )

81 )

82 - lora_requests: Optional[list[LoRARequest]] = None

82 + lora_requests: list[LoRARequest] | None = None

83 if engine_args.enable_lora:

84 lora_requests = [request.lora_request for request in requests]

85

@@ -197,9 +197,9 @@ async def run_vllm_async(

197 )

198

199 # Add the requests to the engine.

200 - prompts: list[Union[TextPrompt, TokensPrompt]] = []

200 + prompts: list[TextPrompt | TokensPrompt] = []

201 sampling_params: list[SamplingParams] = []

202 - lora_requests: list[Optional[LoRARequest]] = []

202 + lora_requests: list[LoRARequest | None] = []

203 for request in requests:

204 prompt = (

205 TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"])

@@ -696,7 +696,7 @@ def main(args: argparse.Namespace):

696 )

697 requests = get_requests(args, tokenizer)

698 is_multi_modal = any(request.multi_modal_data is not None for request in requests)

699 - request_outputs: Optional[list[RequestOutput]] = None

699 + request_outputs: list[RequestOutput] | None = None

700 if args.backend == "vllm":

701 if args.async_engine:

702 elapsed_time = uvloop.run(

vllm/compilation/backends.py

+9 -8

@@ -7,9 +7,9 @@ import hashlib

7 import os

8 import pprint

9 import time

10 -from collections.abc import Sequence

10 +from collections.abc import Callable, Sequence

11 from contextlib import contextmanager

12 -from typing import Any, Callable, Optional

12 +from typing import Any

13

14 import torch

15 import torch.fx as fx

@@ -75,7 +75,7 @@ class CompilerManager:

75 """

76

77 def __init__(self, compilation_config: CompilationConfig):

78 - self.cache: dict[tuple[Optional[int], int, str], Any] = dict()

78 + self.cache: dict[tuple[int | None, int, str], Any] = dict()

79 self.is_cache_updated = False

80 self.compilation_config = compilation_config

81 self.compiler = make_compiler(compilation_config)

@@ -84,7 +84,7 @@ class CompilerManager:

84 return self.compiler.compute_hash(vllm_config)

85

86 @contextmanager

87 - def compile_context(self, runtime_shape: Optional[int] = None):

87 + def compile_context(self, runtime_shape: int | None = None):

88 """Provide compilation context for the duration of compilation to set

89 any torch global properties we want to scope to a single Inductor

90 compilation (e.g. partition rules, pass context)."""

@@ -145,8 +145,8 @@ class CompilerManager:

145 graph: fx.GraphModule,

146 example_inputs: list[Any],

147 graph_index: int,

148 - runtime_shape: Optional[int] = None,

149 - ) -> Optional[Callable]:

148 + runtime_shape: int | None = None,

149 + ) -> Callable | None:

150 if (runtime_shape, graph_index, self.compiler.name) not in self.cache:

151 return None

152 handle = self.cache[(runtime_shape, graph_index, self.compiler.name)]

@@ -178,7 +178,7 @@ class CompilerManager:

178 compilation_config: CompilationConfig,

179 graph_index: int = 0,

180 num_graphs: int = 1,

181 - runtime_shape: Optional[int] = None,

181 + runtime_shape: int | None = None,

182 ) -> Any:

183 if graph_index == 0:

184 # before compiling the first graph, record the start time

@@ -656,7 +656,8 @@ class VllmBackend:

656

657 graph_path = os.path.join(local_cache_dir, "computation_graph.py")

658 if not os.path.exists(graph_path):

659 - # code adapted from https://github.com/thuml/depyf/blob/dab831108a752d1facc00acdd6d4243891845c37/depyf/explain/patched_lazy_format_graph_code.py#L30 # noqa

659 + # code adapted from

660 + # https://github.com/thuml/depyf/blob/dab831108a752d1facc00acdd6d4243891845c37/depyf/explain/patched_lazy_format_graph_code.py#L30

661 # use `print_readable` because it can include submodules

662 src = (

663 "from __future__ import annotations\nimport torch\n"

vllm/compilation/base_static_graph.py

+2 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Callable, Protocol

4 +from collections.abc import Callable

5 +from typing import Any, Protocol

6

7 from vllm.config import CUDAGraphMode, VllmConfig

8

vllm/compilation/collective_fusion.py

+9 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from importlib.util import find_spec

4 -from typing import Optional

4

5 import torch

6 import torch._inductor.pattern_matcher as pm

@@ -432,7 +431,7 @@ class AsyncTPPass(VllmPatternMatcherPass):

431

432 self.dump_patterns(config, self.patterns)

433

435 - def is_applicable_for_shape(self, shape: Optional[int]) -> bool:

434 + def is_applicable_for_shape(self, shape: int | None) -> bool:

435 # only do replace for specific shapes

436 tp_size = get_tensor_model_parallel_world_size()

437 return shape is not None and shape % tp_size == 0

@@ -485,10 +484,10 @@ if flashinfer_comm is not None:

484 max_token_num: int,

485 pattern_code: int,

486 fuse_rms_quant: bool,

488 - norm_out: Optional[torch.Tensor] = None,

489 - quant_out: Optional[torch.Tensor] = None,

490 - scale_out: Optional[torch.Tensor] = None,

491 - scale_factor: Optional[torch.Tensor] = None,

487 + norm_out: torch.Tensor | None = None,

488 + quant_out: torch.Tensor | None = None,

489 + scale_out: torch.Tensor | None = None,

490 + scale_factor: torch.Tensor | None = None,

491 ) -> None:

492 num_tokens, hidden_size = allreduce_in.shape

493 element_size = allreduce_in.element_size()

@@ -589,10 +588,10 @@ if flashinfer_comm is not None:

588 max_token_num: int,

589 pattern_code: int,

590 fuse_rms_quant: bool,

592 - norm_out: Optional[torch.Tensor] = None,

593 - quant_out: Optional[torch.Tensor] = None,

594 - scale_out: Optional[torch.Tensor] = None,

595 - scale_factor: Optional[torch.Tensor] = None,

591 + norm_out: torch.Tensor | None = None,

592 + quant_out: torch.Tensor | None = None,

593 + scale_out: torch.Tensor | None = None,

594 + scale_factor: torch.Tensor | None = None,

595 ) -> None:

596 pass

597

vllm/compilation/compiler_interface.py

+17 -16

@@ -4,8 +4,9 @@ import contextlib

4 import copy

5 import hashlib

6 import os

7 +from collections.abc import Callable

8 from contextlib import ExitStack

8 -from typing import Any, Callable, Optional

9 +from typing import Any

10 from unittest.mock import patch

11

12 import torch

@@ -62,9 +63,9 @@ class CompilerInterface:

63 graph: fx.GraphModule,

64 example_inputs: list[Any],

65 compiler_config: dict[str, Any],

65 - runtime_shape: Optional[int] = None,

66 - key: Optional[str] = None,

67 - ) -> tuple[Optional[Callable], Optional[Any]]:

66 + runtime_shape: int | None = None,

67 + key: str | None = None,

68 + ) -> tuple[Callable | None, Any | None]:

69 """

70 Compile the graph with the given example inputs and compiler config,

71 with a runtime shape. If the `runtime_shape` is None, it means

@@ -97,7 +98,7 @@ class CompilerInterface:

98 graph: fx.GraphModule,

99 example_inputs: list[Any],

100 graph_index: int,

100 - runtime_shape: Optional[int] = None,

101 + runtime_shape: int | None = None,

102 ) -> Callable:

103 """

104 Load the compiled function from the handle.

@@ -191,9 +192,9 @@ class InductorStandaloneAdaptor(CompilerInterface):

192 graph: fx.GraphModule,

193 example_inputs: list[Any],

194 compiler_config: dict[str, Any],

194 - runtime_shape: Optional[int] = None,

195 - key: Optional[str] = None,

196 - ) -> tuple[Optional[Callable], Optional[Any]]:

195 + runtime_shape: int | None = None,

196 + key: str | None = None,

197 + ) -> tuple[Callable | None, Any | None]:

198 compilation_counter.num_inductor_compiles += 1

199 current_config = {}

200 if compiler_config is not None:

@@ -229,7 +230,7 @@ class InductorStandaloneAdaptor(CompilerInterface):

230 graph: fx.GraphModule,

231 example_inputs: list[Any],

232 graph_index: int,

232 - runtime_shape: Optional[int] = None,

233 + runtime_shape: int | None = None,

234 ) -> Callable:

235 assert isinstance(handle, tuple)

236 assert isinstance(handle[0], str)

@@ -293,9 +294,9 @@ class InductorAdaptor(CompilerInterface):

294 graph: fx.GraphModule,

295 example_inputs: list[Any],

296 compiler_config: dict[str, Any],

296 - runtime_shape: Optional[int] = None,

297 - key: Optional[str] = None,

298 - ) -> tuple[Optional[Callable], Optional[Any]]:

297 + runtime_shape: int | None = None,

298 + key: str | None = None,

299 + ) -> tuple[Callable | None, Any | None]:

300 compilation_counter.num_inductor_compiles += 1

301 from torch._inductor.compile_fx import compile_fx

302

@@ -492,7 +493,7 @@ class InductorAdaptor(CompilerInterface):

493 graph: fx.GraphModule,

494 example_inputs: list[Any],

495 graph_index: int,

495 - runtime_shape: Optional[int] = None,

496 + runtime_shape: int | None = None,

497 ) -> Callable:

498 assert isinstance(handle, tuple)

499 assert isinstance(handle[0], str)

@@ -610,9 +611,9 @@ class EagerAdaptor(CompilerInterface):

611 graph: fx.GraphModule,

612 example_inputs: list[Any],

613 compiler_config: dict[str, Any],

613 - runtime_shape: Optional[int] = None,

614 - key: Optional[str] = None,

615 - ) -> tuple[Optional[Callable], Optional[Any]]:

614 + runtime_shape: int | None = None,

615 + key: str | None = None,

616 + ) -> tuple[Callable | None, Any | None]:

617 compilation_counter.num_eager_compiles += 1

618 # we don't need to compile the graph, just return the graph itself.

619 # It does not support caching, return None for the handle.

vllm/compilation/cuda_graph.py

+6 -5

@@ -2,8 +2,9 @@

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

3

4 import dataclasses

5 +from collections.abc import Callable

6 from contextlib import ExitStack

6 -from typing import Any, Callable, Optional

7 +from typing import Any

8 from unittest.mock import patch

9

10 import torch

@@ -24,12 +25,12 @@ logger = init_logger(__name__)

25 @dataclasses.dataclass

26 class CUDAGraphEntry:

27 batch_descriptor: BatchDescriptor

27 - cudagraph: Optional[torch.cuda.CUDAGraph] = None

28 - output: Optional[Any] = None

28 + cudagraph: torch.cuda.CUDAGraph | None = None

29 + output: Any | None = None

30

31 # for cudagraph debugging, track the input addresses

32 # during capture, and check if they are the same during replay

32 - input_addresses: Optional[list[int]] = None

33 + input_addresses: list[int] | None = None

34

35

36 @dataclasses.dataclass

@@ -69,7 +70,7 @@ class CUDAGraphWrapper:

70 runnable: Callable,

71 vllm_config: VllmConfig,

72 runtime_mode: CUDAGraphMode,

72 - cudagraph_options: Optional[CUDAGraphOptions] = None,

73 + cudagraph_options: CUDAGraphOptions | None = None,

74 ):

75 self.runnable = runnable

76 self.vllm_config = vllm_config

vllm/compilation/decorators.py

+12 -11

@@ -6,7 +6,8 @@ import hashlib

6 import inspect

7 import os

8 import sys

9 -from typing import Callable, Optional, TypeVar, Union, overload

9 +from collections.abc import Callable

10 +from typing import TypeVar, overload

11 from unittest.mock import patch

12

13 import torch

@@ -61,14 +62,14 @@ def _should_ignore_torch_compile(cls) -> bool:

62 @overload

63 def support_torch_compile(

64 *,

64 - enable_if: Optional[Callable[[VllmConfig], bool]] = None,

65 + enable_if: Callable[[VllmConfig], bool] | None = None,

66 ) -> Callable[[_T], _T]: ...

67

68

69 @overload

70 def support_torch_compile(

71 *,

71 - dynamic_arg_dims: Optional[dict[str, Union[int, list[int]]]],

72 + dynamic_arg_dims: dict[str, int | list[int]] | None,

73 ) -> Callable[[_T], _T]: ...

74

75

@@ -77,11 +78,11 @@ def support_torch_compile(cls: _T) -> _T: ...

78

79

80 def support_torch_compile(

80 - cls: Optional[_T] = None,

81 + cls: _T | None = None,

82 *,

82 - dynamic_arg_dims: Optional[dict[str, Union[int, list[int]]]] = None,

83 - enable_if: Optional[Callable[[VllmConfig], bool]] = None,

84 -) -> Union[Callable[[_T], _T], _T]:

83 + dynamic_arg_dims: dict[str, int | list[int]] | None = None,

84 + enable_if: Callable[[VllmConfig], bool] | None = None,

85 +) -> Callable[[_T], _T] | _T:

86 """

87 A decorator to add support for compiling the forward method of a class.

88

@@ -147,9 +148,9 @@ def support_torch_compile(

148 for k, v in sig.parameters.items():

149 if v.annotation in [

150 torch.Tensor,

150 - Optional[torch.Tensor],

151 + torch.Tensor | None,

152 IntermediateTensors,

152 - Optional[IntermediateTensors],

153 + IntermediateTensors | None,

154 ]:

155 inferred_dynamic_arg_dims[k] = 0

156

@@ -209,8 +210,8 @@ def _verify_source_unchanged(source_info, vllm_config) -> None:

210

211 def _support_torch_compile(

212 cls: _T,

212 - dynamic_arg_dims: dict[str, Union[int, list[int]]],

213 - enable_if: Optional[Callable[[VllmConfig], bool]] = None,

213 + dynamic_arg_dims: dict[str, int | list[int]],

214 + enable_if: Callable[[VllmConfig], bool] | None = None,

215 ) -> _T:

216 """

217 A decorator to add support for compiling the forward method of a class.

vllm/compilation/fix_functionalization.py

+5 -6

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

3

4 import operator

5 from collections.abc import Iterable

6 -from typing import Optional, Union

6

7 import torch

8 from torch._higher_order_ops.auto_functionalize import auto_functionalized

@@ -150,7 +149,7 @@ class FixFunctionalizationPass(VllmInductorPass):

149 )

150 self.nodes_to_remove.clear()

151

153 - def _remove(self, node_or_nodes: Union[torch.fx.Node, Iterable[torch.fx.Node]]):

152 + def _remove(self, node_or_nodes: torch.fx.Node | Iterable[torch.fx.Node]):

153 """

154 Stage a node (or nodes) for removal at the end of the pass.

155 """

@@ -163,8 +162,8 @@ class FixFunctionalizationPass(VllmInductorPass):

162 self,

163 graph: torch.fx.Graph,

164 node: torch.fx.Node,

166 - mutated_args: dict[int, Union[torch.fx.Node, str]],

167 - args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None,

165 + mutated_args: dict[int, torch.fx.Node | str],

166 + args: tuple[torch.fx.Node | str, ...] | None = None,

167 ):

168 """

169 De-functionalize a node by replacing it with a call to the original.

@@ -176,7 +175,7 @@ class FixFunctionalizationPass(VllmInductorPass):

175 self._remove(node)

176

177 def replace_users_with_mutated_args(

179 - self, node: torch.fx.Node, mutated_args: dict[int, Union[torch.fx.Node, str]]

178 + self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]

179 ):

180 """

181 Replace all getitem users of the auto-functionalized node with the

@@ -207,7 +206,7 @@ class FixFunctionalizationPass(VllmInductorPass):

206 self,

207 graph: torch.fx.Graph,

208 node: torch.fx.Node,

210 - args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None,

209 + args: tuple[torch.fx.Node | str, ...] | None = None,

210 ):

211 """

212 Insert a new defunctionalized node into the graph before node.

vllm/compilation/fx_utils.py

+3 -6

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

3

4 import operator

5 from collections.abc import Iterable, Iterator

6 -from typing import Optional

6

7 from torch import fx

8 from torch._higher_order_ops.auto_functionalize import auto_functionalized

@@ -19,9 +18,7 @@ def is_auto_func(node: fx.Node, op: OpOverload) -> bool:

18

19

20 # Returns the first specified node with the given op (if it exists)

22 -def find_specified_fn_maybe(

23 - nodes: Iterable[fx.Node], op: OpOverload

24 -) -> Optional[fx.Node]:

21 +def find_specified_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None:

22 for node in nodes:

23 if node.target == op:

24 return node

@@ -36,7 +33,7 @@ def find_specified_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:

33

34

35 # Returns the first auto_functionalized node with the given op (if it exists)

39 -def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> Optional[fx.Node]:

36 +def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None:

37 for node in nodes:

38 if is_func(node, auto_functionalized) and node.args[0] == op: # noqa

39 return node

@@ -52,7 +49,7 @@ def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:

49

50 # Returns the getitem node that extracts the idx-th element from node

51 # (if it exists)

55 -def find_getitem_maybe(node: fx.Node, idx: int) -> Optional[fx.Node]:

52 +def find_getitem_maybe(node: fx.Node, idx: int) -> fx.Node | None:

53 for user in node.users:

54 if is_func(user, operator.getitem) and user.args[1] == idx:

55 return user

vllm/compilation/inductor_pass.py

+7 -8

@@ -6,8 +6,9 @@ import hashlib

6 import inspect

7 import json

8 import types

9 +from collections.abc import Callable

10 from contextlib import contextmanager

10 -from typing import Any, Callable, Optional, Union

11 +from typing import Any

12

13 import torch

14 from torch import fx

@@ -27,7 +28,7 @@ _pass_context = None

28

29

30 class PassContext:

30 - def __init__(self, runtime_shape: Optional[int]):

31 + def __init__(self, runtime_shape: int | None):

32 self.runtime_shape = runtime_shape

33

34

@@ -38,7 +39,7 @@ def get_pass_context() -> PassContext:

39

40

41 @contextmanager

41 -def pass_context(runtime_shape: Optional[int]):

42 +def pass_context(runtime_shape: int | None):

43 """A context manager that stores the current pass context,

44 usually it is a list of sizes to specialize.

45 """

@@ -67,7 +68,7 @@ class InductorPass(CustomGraphPass):

68 return InductorPass.hash_source(self)

69

70 @staticmethod

70 - def hash_source(*srcs: Union[str, Any]):

71 + def hash_source(*srcs: str | Any):

72 """

73 Utility method to hash the sources of functions or objects.

74 :param srcs: strings or objects to add to the hash.

@@ -95,7 +96,7 @@ class InductorPass(CustomGraphPass):

96 encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")

97 return hashlib.sha256(encoded).hexdigest()

98

98 - def is_applicable_for_shape(self, shape: Optional[int]):

99 + def is_applicable_for_shape(self, shape: int | None):

100 return True

101

102

@@ -105,9 +106,7 @@ class CallableInductorPass(InductorPass):

106 implementation of the UUID.

107 """

108

108 - def __init__(

109 - self, callable: Callable[[fx.Graph], None], uuid: Optional[Any] = None

110 - ):

109 + def __init__(self, callable: Callable[[fx.Graph], None], uuid: Any | None = None):

110 self.callable = callable

111 self._uuid = self.hash_source(callable) if uuid is None else uuid

112

vllm/compilation/noop_elimination.py

+2 -5

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

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

3

4 from collections.abc import Iterable

5 -from typing import Union

5

6 import torch.fx

7 from torch import SymInt

@@ -105,9 +104,7 @@ class NoOpEliminationPass(VllmInductorPass):

104 logger.debug("Removed %s no-op reshapes and slices", count)

105

106 # ---------------------- Shape comparison helpers ----------------------

108 - def dims_equivalent(

109 - self, dim: Union[int, SymInt], i_dim: Union[int, SymInt]

110 - ) -> bool:

107 + def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool:

108 """

109 This function checks if two dimensions are equivalent.

110 :param dim: The dimension arg to reshape/slice

@@ -127,7 +124,7 @@ class NoOpEliminationPass(VllmInductorPass):

124 return False

125

126 def all_dims_equivalent(

130 - self, dims: Iterable[Union[int, SymInt]], i_dims: Iterable[Union[int, SymInt]]

127 + self, dims: Iterable[int | SymInt], i_dims: Iterable[int | SymInt]

128 ) -> bool:

129 dims_ = list(dims)

130 i_dims_ = list(i_dims)

vllm/compilation/partition_rules.py

+2 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import contextlib

5 from typing import TYPE_CHECKING

6

@@ -16,7 +14,7 @@ if TYPE_CHECKING:

14 logger = init_logger(__name__)

15

16

19 -def resolve_defined_ops(op_names: list[str]) -> list[torch._ops.OpOverload]:

17 +def resolve_defined_ops(op_names: list[str]) -> list["torch._ops.OpOverload"]:

18 """Resolve operator names to OpOverload objects.

19

20 Skips operators that fail to resolve (e.g., operators not registered or

@@ -49,7 +47,7 @@ def resolve_defined_ops(op_names: list[str]) -> list[torch._ops.OpOverload]:

47

48

49 @contextlib.contextmanager

52 -def inductor_partition_rule_context(overloads: list[torch._ops.OpOverload]):

50 +def inductor_partition_rule_context(overloads: list["torch._ops.OpOverload"]):

51 """Context manager to temporarily register Inductor partition rules.

52

53 Registers custom partition rules for specified operators, forcing the

vllm/compilation/piecewise_backend.py

+2 -1

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

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

3

4 import dataclasses

5 -from typing import Any, Callable

5 +from collections.abc import Callable

6 +from typing import Any

7

8 import torch.fx as fx

9

vllm/compilation/sequence_parallelism.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5 import torch._inductor.pattern_matcher as pm

@@ -27,7 +26,7 @@ class _RMSNormAndQuantOpHelper:

26 epsilon: float,

27 dtype: torch.dtype,

28 device: str,

30 - quant_op: Optional[torch._ops.OpOverload] = None,

29 + quant_op: torch._ops.OpOverload | None = None,

30 **kwargs,

31 ):

32 self.epsilon = epsilon

@@ -110,7 +109,7 @@ class _SequenceParallelPatternHelper(_RMSNormAndQuantOpHelper):

109 epsilon: float,

110 dtype: torch.dtype,

111 device: str,

113 - quant_op: Optional[torch._ops.OpOverload] = None,

112 + quant_op: torch._ops.OpOverload | None = None,

113 **kwargs,

114 ):

115 super().__init__(epsilon, dtype, device, quant_op=quant_op, **kwargs)

@@ -483,7 +482,7 @@ class SequenceParallelismPass(VllmPatternMatcherPass):

482 ).register(self.patterns)

483 self.dump_patterns(config, self.patterns)

484

486 - def is_applicable_for_shape(self, shape: Optional[int]) -> bool:

485 + def is_applicable_for_shape(self, shape: int | None) -> bool:

486 tp_size = get_tensor_model_parallel_world_size()

487 return shape is not None and shape % tp_size == 0

488

vllm/compilation/torch25_custom_graph_pass.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from abc import ABC, abstractmethod

4 -from typing import Any, Optional

4 +from typing import Any

5

6 import torch

7

@@ -23,7 +23,7 @@ class Torch25CustomGraphPass(ABC): # noqa (redefinition)

23 """

24

25 @abstractmethod

26 - def uuid(self) -> Optional[Any]:

26 + def uuid(self) -> Any | None:

27 """

28 Return an ID to uniquely identify your custom pass implementation.

29 Return None to skip inductor code caching entirely.

vllm/compilation/vllm_inductor_pass.py

+2 -2

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

3 import functools

4 import operator

5 import time

6 -from typing import ClassVar, Optional

6 +from typing import ClassVar

7

8 import regex as re

9 import torch

@@ -24,7 +24,7 @@ class VllmInductorPass(InductorPass):

24 It provides timing, logging, and dumping utilities.

25 """

26

27 - dump_prefix: ClassVar[Optional[int]] = None

27 + dump_prefix: ClassVar[int | None] = None

28 """Keep track of pass index for debug dump ordering."""

29

30 def __init__(self, config: VllmConfig):

vllm/compilation/wrapper.py

+2 -2

@@ -4,9 +4,9 @@

4 import os

5 import sys

6 from abc import abstractmethod

7 +from collections.abc import Callable

8 from contextlib import contextmanager

9 from types import CodeType

9 -from typing import Callable, Optional

10

11 import torch

12

@@ -31,7 +31,7 @@ class TorchCompileWrapperWithCustomDispatcher:

31 """

32

33 def __init__(

34 - self, compiled_callable: Optional[Callable] = None, compilation_level: int = 0

34 + self, compiled_callable: Callable | None = None, compilation_level: int = 0

35 ):

36 vllm_config = get_current_vllm_config()

37 self.vllm_config = vllm_config

vllm/config/cache.py

+10 -10

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

3

4 import hashlib

5 from dataclasses import field

6 -from typing import TYPE_CHECKING, Any, Literal, Optional

6 +from typing import TYPE_CHECKING, Any, Literal

7

8 from pydantic import Field, SkipValidation, field_validator

9 from pydantic.dataclasses import dataclass

@@ -58,13 +58,13 @@ class CacheConfig:

58 is_attention_free: bool = False

59 """Whether the model is attention-free. This is primarily set in

60 `ModelConfig` and that value should be manually duplicated here."""

61 - num_gpu_blocks_override: Optional[int] = None

61 + num_gpu_blocks_override: int | None = None

62 """Number of GPU blocks to use. This overrides the profiled `num_gpu_blocks`

63 if specified. Does nothing if `None`. Used for testing preemption."""

64 - sliding_window: Optional[int] = None

64 + sliding_window: int | None = None

65 """Sliding window size for the KV cache. This is primarily set in

66 `ModelConfig` and that value should be manually duplicated here."""

67 - enable_prefix_caching: Optional[bool] = None

67 + enable_prefix_caching: bool | None = None

68 """Whether to enable prefix caching. Enabled by default for V1."""

69 prefix_caching_hash_algo: PrefixCachingHashAlgo = "sha256"

70 """Set the hash algorithm for prefix caching:\n

@@ -84,12 +84,12 @@ class CacheConfig:

84 """This enables dynamic calculation of `k_scale` and `v_scale` when

85 kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model

86 checkpoint if available. Otherwise, the scales will default to 1.0."""

87 - cpu_kvcache_space_bytes: Optional[int] = None

87 + cpu_kvcache_space_bytes: int | None = None

88 """(CPU backend only) CPU key-value cache space."""

89 - mamba_page_size_padded: Optional[int] = None

89 + mamba_page_size_padded: int | None = None

90 """ Optional override for mamba page size; used by hybrid mamba/attention

91 models to ensure exact alignment with attention page size."""

92 - mamba_block_size: Optional[int] = None

92 + mamba_block_size: int | None = None

93 """Size of a contiguous cache block in number of tokens for mamba cache."""

94 mamba_cache_dtype: MambaDType = "auto"

95 """The data type to use for the Mamba cache (both the conv as well as the

@@ -101,9 +101,9 @@ class CacheConfig:

101 for the ssm state will be determined by mamba_cache_dtype."""

102

103 # Will be set after profiling.

104 - num_gpu_blocks: Optional[int] = field(default=None, init=False)

104 + num_gpu_blocks: int | None = field(default=None, init=False)

105 """The number of blocks to allocate for GPU memory."""

106 - num_cpu_blocks: Optional[int] = field(default=None, init=False)

106 + num_cpu_blocks: int | None = field(default=None, init=False)

107 """The number of blocks to allocate for CPU memory."""

108

109 kv_sharing_fast_prefill: bool = False

@@ -116,7 +116,7 @@ class CacheConfig:

116 necessary for implementing this optimization in some models (e.g. Gemma3n)

117 """

118

119 - kv_cache_memory_bytes: Optional[int] = None

119 + kv_cache_memory_bytes: int | None = None

120 """Size of KV Cache per GPU in bytes. By default, this is set to None

121 and vllm can automatically infer the kv cache size based on

122 gpu_memory_utilization. However, users may want to manually specify

vllm/config/compilation.py

+10 -9

@@ -4,9 +4,10 @@

4 import enum

5 import hashlib

6 from collections import Counter

7 +from collections.abc import Callable

8 from dataclasses import asdict, field

9 from pathlib import Path

9 -from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Union

10 +from typing import TYPE_CHECKING, Any, ClassVar

11

12 from pydantic import TypeAdapter, field_validator

13 from pydantic.dataclasses import dataclass

@@ -168,7 +169,7 @@ class CompilationConfig:

169 """

170

171 # Top-level Compilation control

171 - level: Optional[int] = None

172 + level: int | None = None

173 """The level of compilation:

174

175 - None: If None, we will select the default compilation level.

@@ -177,7 +178,7 @@ class CompilationConfig:

178 - 1: dynamo as is.

179 - 2: dynamo once.

180 - 3: piecewise compilation."""

180 - debug_dump_path: Optional[Path] = None

181 + debug_dump_path: Path | None = None

182 """The path to dump the debug information."""

183 cache_dir: str = ""

184 """The directory to store the compiled graph, to accelerate Inductor

@@ -208,7 +209,7 @@ class CompilationConfig:

209 By default, all custom ops are enabled when running without Inductor and

210 disabled when running with Inductor: level>=PIECEWISE and use_inductor=True.

211 Inductor generates (fused) Triton kernels for disabled custom ops."""

211 - splitting_ops: Optional[list[str]] = None

212 + splitting_ops: list[str] | None = None

213 """A list of ops to exclude from cudagraphs, used in piecewise compilation.

214

215 The behavior depends on use_inductor_graph_partition:

@@ -238,7 +239,7 @@ class CompilationConfig:

239 are compiled using configurations in inductor_compile_config.

240

241 This setting is ignored if level<PIECEWISE."""

241 - compile_sizes: Optional[list[Union[int, str]]] = None

242 + compile_sizes: list[int | str] | None = None

243 """Sizes to compile for inductor. In addition

244 to integers, it also supports "cudagraph_capture_sizes" to

245 specify the sizes for cudagraph capture."""

@@ -253,7 +254,7 @@ class CompilationConfig:

254 constructor, e.g. `CompilationConfig(inductor_passes={"a": func})`."""

255

256 # CudaGraph compilation

256 - cudagraph_mode: Optional[CUDAGraphMode] = None

257 + cudagraph_mode: CUDAGraphMode | None = None

258 """

259 The mode of the cudagraph:

260

@@ -308,7 +309,7 @@ class CompilationConfig:

309 It means the first several runs will be treated as warmup runs.

310 Only after that, the execution will be recorded, and the recorded

311 cudagraph will be used for subsequent runs."""

311 - cudagraph_capture_sizes: Optional[list[int]] = None

312 + cudagraph_capture_sizes: list[int] | None = None

313 """Sizes to capture cudagraph.

314 - None (default): capture sizes are inferred from vllm config.

315 - list[int]: capture sizes are specified as given."""

@@ -320,7 +321,7 @@ class CompilationConfig:

321 internally managed buffer. Default is False.

322 Note that this flag is only effective when cudagraph_mode is PIECEWISE.

323 """

323 - full_cuda_graph: Optional[bool] = False

324 + full_cuda_graph: bool | None = False

325 """whether to use a full cuda graph for the entire forward pass rather than

326 splitting certain operations such as attention into subgraphs. Thus this

327 flag cannot be used together with splitting_ops. This may provide

@@ -544,7 +545,7 @@ class CompilationConfig:

545 "(where 'op' is the registered op name)"

546 )

547

547 - def init_backend(self, vllm_config: "VllmConfig") -> Union[str, Callable]:

548 + def init_backend(self, vllm_config: "VllmConfig") -> str | Callable:

549 if self.level == CompilationLevel.NO_COMPILATION:

550 raise ValueError("No compilation level is set.")

551

vllm/config/device.py

+2 -2

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

3

4 import hashlib

5 from dataclasses import field

6 -from typing import Any, Literal, Optional, Union

6 +from typing import Any, Literal

7

8 import torch

9 from pydantic import ConfigDict, SkipValidation

@@ -19,7 +19,7 @@ Device = Literal["auto", "cuda", "cpu", "tpu", "xpu"]

19 class DeviceConfig:

20 """Configuration for the device to use for vLLM execution."""

21

22 - device: SkipValidation[Optional[Union[Device, torch.device]]] = "auto"

22 + device: SkipValidation[Device | torch.device | None] = "auto"

23 """Device type for vLLM execution.

24 This parameter is deprecated and will be

25 removed in a future release.

vllm/config/kv_events.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 from pydantic.dataclasses import dataclass

6

@@ -26,7 +25,7 @@ class KVEventsConfig:

25 """The zmq endpoint to use for publishing kv events.

26 """

27

29 - replay_endpoint: Optional[str] = None

28 + replay_endpoint: str | None = None

29 """The zmq endpoint to use for replaying kv events.

30 """

31

vllm/config/kv_transfer.py

+7 -7

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

4 import hashlib

5 import uuid

6 from dataclasses import field

7 -from typing import Any, Literal, Optional, get_args

7 +from typing import Any, Literal, get_args

8

9 from pydantic.dataclasses import dataclass

10

@@ -20,14 +20,14 @@ KVRole = Literal[KVProducer, KVConsumer]

20 class KVTransferConfig:

21 """Configuration for distributed KV cache transfer."""

22

23 - kv_connector: Optional[str] = None

23 + kv_connector: str | None = None

24 """The KV connector for vLLM to transmit KV caches between vLLM instances.

25 """

26

27 - engine_id: Optional[str] = None

27 + engine_id: str | None = None

28 """The engine id for KV transfers."""

29

30 - kv_buffer_device: Optional[str] = "cuda"

30 + kv_buffer_device: str | None = "cuda"

31 """The device used by kv connector to buffer the KV cache. Choices are

32 'cuda' and 'cpu'."""

33

@@ -35,11 +35,11 @@ class KVTransferConfig:

35 """The buffer size for TorchDistributedConnector. Measured in number of

36 bytes. Recommended value: 1e9 (about 1GB)."""

37

38 - kv_role: Optional[KVRole] = None

38 + kv_role: KVRole | None = None

39 """Whether this vLLM instance produces, consumes KV cache, or both. Choices

40 are 'kv_producer', 'kv_consumer', and 'kv_both'."""

41

42 - kv_rank: Optional[int] = None

42 + kv_rank: int | None = None

43 """The rank of this vLLM instance in the KV cache transfer. Typical value:

44 0 for prefill instance, 1 for decode instance.

45 Currently only 1P1D is supported."""

@@ -57,7 +57,7 @@ class KVTransferConfig:

57 kv_connector_extra_config: dict[str, Any] = field(default_factory=dict)

58 """any extra config that the connector may need."""

59

60 - kv_connector_module_path: Optional[str] = None

60 + kv_connector_module_path: str | None = None

61 """The Python module path to dynamically load the KV connector from.

62 Only supported in V1."""

63

vllm/config/load.py

+9 -13

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

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

3

4 import hashlib

5 -from typing import TYPE_CHECKING, Any, Optional, Union

5 +from typing import TYPE_CHECKING, Any

6

7 from pydantic import Field, field_validator

8 from pydantic.dataclasses import dataclass

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

25 class LoadConfig:

26 """Configuration for loading the model weights."""

27

28 - load_format: Union[str, LoadFormats] = "auto"

28 + load_format: str | LoadFormats = "auto"

29 """The format of the model weights to load:\n

30 - "auto" will try to load the weights in the safetensors format and fall

31 back to the pytorch bin format if safetensors format is not available.\n

@@ -48,7 +48,7 @@ class LoadConfig:

48 - "mistral" will load weights from consolidated safetensors files used by

49 Mistral models.

50 - Other custom values can be supported via plugins."""

51 - download_dir: Optional[str] = None

51 + download_dir: str | None = None

52 """Directory to download and load the weights, default to the default

53 cache directory of Hugging Face."""

54 safetensors_load_strategy: str = "lazy"

@@ -64,23 +64,19 @@ class LoadConfig:

64 was quantized using torchao and saved using safetensors.

65 Needs torchao >= 0.14.0

66 """

67 - model_loader_extra_config: Union[dict, TensorizerConfig] = Field(

68 - default_factory=dict

69 - )

67 + model_loader_extra_config: dict | TensorizerConfig = Field(default_factory=dict)

68 """Extra config for model loader. This will be passed to the model loader

69 corresponding to the chosen load_format."""

72 - device: Optional[str] = None

70 + device: str | None = None

71 """Device to which model weights will be loaded, default to

72 device_config.device"""

75 - ignore_patterns: Union[list[str], str] = Field(

76 - default_factory=lambda: ["original/**/*"]

77 - )

73 + ignore_patterns: list[str] | str = Field(default_factory=lambda: ["original/**/*"])

74 """The list of patterns to ignore when loading the model. Default to

75 "original/**/*" to avoid repeated loading of llama's checkpoints."""

76 use_tqdm_on_load: bool = True

77 """Whether to enable tqdm for showing progress bar when loading model

78 weights."""

83 - pt_load_map_location: Union[str, dict[str, str]] = "cpu"

79 + pt_load_map_location: str | dict[str, str] = "cpu"

80 """

81 pt_load_map_location: the map location for loading pytorch checkpoint, to

82 support loading checkpoints can only be loaded on certain devices like

@@ -115,8 +111,8 @@ class LoadConfig:

111

112 @field_validator("ignore_patterns", mode="after")

113 def _validate_ignore_patterns(

118 - cls, ignore_patterns: Union[list[str], str]

119 - ) -> Union[list[str], str]:

114 + cls, ignore_patterns: list[str] | str

115 + ) -> list[str] | str:

116 if ignore_patterns != ["original/**/*"] and len(ignore_patterns) > 0:

117 logger.info(

118 "Ignoring the following patterns when downloading weights: %s",

vllm/config/lora.py

+4 -4

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

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

3

4 import hashlib

5 -from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, Union

5 +from typing import TYPE_CHECKING, Any, ClassVar, Literal

6

7 import torch

8 from pydantic import ConfigDict, Field, model_validator

@@ -42,10 +42,10 @@ class LoRAConfig:

42 parallelism. Enabling this will use the fully sharded layers. At high

43 sequence length, max rank or tensor parallel size, this is likely faster.

44 """

45 - max_cpu_loras: Optional[int] = None

45 + max_cpu_loras: int | None = None

46 """Maximum number of LoRAs to store in CPU memory. Must be >= than

47 `max_loras`."""

48 - lora_dtype: Union[torch.dtype, LoRADType] = "auto"

48 + lora_dtype: torch.dtype | LoRADType = "auto"

49 """Data type for LoRA. If auto, will default to base model dtype."""

50 lora_extra_vocab_size: LoRAExtraVocabSize = Field(

51 default=256,

@@ -60,7 +60,7 @@ class LoRAConfig:

60 lora_vocab_padding_size: ClassVar[int] = (

61 current_platform.get_lora_vocab_padding_size()

62 )

63 - default_mm_loras: Optional[dict[str, str]] = None

63 + default_mm_loras: dict[str, str] | None = None

64 """Dictionary mapping specific modalities to LoRA model paths; this field

65 is only applicable to multimodal models and should be leveraged when a

66 model always expects a LoRA to be active when a given modality is present.

vllm/config/model.py

+62 -72

@@ -4,18 +4,10 @@

4 import hashlib

5 import json

6 import warnings

7 +from collections.abc import Callable

8 from dataclasses import InitVar, field

9 from importlib.util import find_spec

9 -from typing import (

10 - TYPE_CHECKING,

11 - Any,

12 - Callable,

13 - Literal,

14 - Optional,

15 - Union,

16 - cast,

17 - get_args,

18 -)

10 +from typing import TYPE_CHECKING, Any, Literal, cast, get_args

11

12 import torch

13 from pydantic import ConfigDict, SkipValidation, field_validator, model_validator

@@ -89,7 +81,7 @@ ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]

81 LogprobsMode = Literal[

82 "raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs"

83 ]

92 -HfOverrides = Union[dict[str, Any], Callable[[PretrainedConfig], PretrainedConfig]]

84 +HfOverrides = dict[str, Any] | Callable[[PretrainedConfig], PretrainedConfig]

85 ModelImpl = Literal["auto", "vllm", "transformers", "terratorch"]

86

87 _RUNNER_TASKS: dict[RunnerType, list[TaskOption]] = {

@@ -121,7 +113,7 @@ class ModelConfig:

113 """Convert the model using adapters defined in

114 [vllm.model_executor.models.adapters][]. The most common use case is to

115 adapt a text generation model to be used for pooling tasks."""

124 - task: Optional[TaskOption] = None

116 + task: TaskOption | None = None

117 """[DEPRECATED] The task to use the model for. If the model supports more

118 than one model runner, this is used to select which model runner to run.

119

@@ -139,7 +131,7 @@ class ModelConfig:

131 trust_remote_code: bool = False

132 """Trust remote code (e.g., from HuggingFace) when downloading the model

133 and tokenizer."""

142 - dtype: Union[ModelDType, torch.dtype] = "auto"

134 + dtype: ModelDType | torch.dtype = "auto"

135 """Data type for model weights and activations:\n

136 - "auto" will use FP16 precision for FP32 and FP16 models, and BF16

137 precision for BF16 models.\n

@@ -148,33 +140,33 @@ class ModelConfig:

140 - "bfloat16" for a balance between precision and range.\n

141 - "float" is shorthand for FP32 precision.\n

142 - "float32" for FP32 precision."""

151 - seed: Optional[int] = None

143 + seed: int | None = None

144 """Random seed for reproducibility. Initialized to None in V0, but

145 initialized to 0 in V1."""

154 - hf_config_path: Optional[str] = None

146 + hf_config_path: str | None = None

147 """Name or path of the Hugging Face config to use. If unspecified, model

148 name or path will be used."""

149 allowed_local_media_path: str = ""

150 """Allowing API requests to read local images or videos from directories

151 specified by the server file system. This is a security risk. Should only

152 be enabled in trusted environments."""

161 - allowed_media_domains: Optional[list[str]] = None

153 + allowed_media_domains: list[str] | None = None

154 """If set, only media URLs that belong to this domain can be used for

155 multi-modal inputs. """

164 - revision: Optional[str] = None

156 + revision: str | None = None

157 """The specific model version to use. It can be a branch name, a tag name,

158 or a commit id. If unspecified, will use the default version."""

167 - code_revision: Optional[str] = None

159 + code_revision: str | None = None

160 """The specific revision to use for the model code on the Hugging Face Hub.

161 It can be a branch name, a tag name, or a commit id. If unspecified, will

162 use the default version."""

163 rope_scaling: dict[str, Any] = field(default_factory=dict)

164 """RoPE scaling configuration. For example,

165 `{"rope_type":"dynamic","factor":2.0}`."""

174 - rope_theta: Optional[float] = None

166 + rope_theta: float | None = None

167 """RoPE theta. Use with `rope_scaling`. In some cases, changing the RoPE

168 theta improves the performance of the scaled model."""

177 - tokenizer_revision: Optional[str] = None

169 + tokenizer_revision: str | None = None

170 """The specific revision to use for the tokenizer on the Hugging Face Hub.

171 It can be a branch name, a tag name, or a commit id. If unspecified, will

172 use the default version."""

@@ -187,9 +179,9 @@ class ModelConfig:

179 - 1k -> 1000\n

180 - 1K -> 1024\n

181 - 25.6k -> 25,600"""

190 - spec_target_max_model_len: Optional[int] = None

182 + spec_target_max_model_len: int | None = None

183 """Specify the maximum length for spec decoding draft models."""

192 - quantization: SkipValidation[Optional[QuantizationMethods]] = None

184 + quantization: SkipValidation[QuantizationMethods | None] = None

185 """Method used to quantize the weights. If `None`, we first check the

186 `quantization_config` attribute in the model config file. If that is

187 `None`, we assume the model weights are not quantized and use `dtype` to

@@ -230,7 +222,7 @@ class ModelConfig:

222 """If `True`, enables passing text embeddings as inputs via the

223 `prompt_embeds` key. Note that enabling this will double the time required

224 for graph compilation."""

233 - served_model_name: Optional[Union[str, list[str]]] = None

225 + served_model_name: str | list[str] | None = None

226 """The model name(s) used in the API. If multiple names are provided, the

227 server will respond to any of the provided names. The model name in the

228 model field of a response will be the first name in this list. If not

@@ -238,20 +230,20 @@ class ModelConfig:

230 that this name(s) will also be used in `model_name` tag content of

231 prometheus metrics, if multiple names provided, metrics tag will take the

232 first one."""

241 - config_format: Union[str, ConfigFormat] = "auto"

233 + config_format: str | ConfigFormat = "auto"

234 """The format of the model config to load:\n

235 - "auto" will try to load the config in hf format if available else it

236 will try to load in mistral format.\n

237 - "hf" will load the config in hf format.\n

238 - "mistral" will load the config in mistral format."""

247 - hf_token: Optional[Union[bool, str]] = None

239 + hf_token: bool | str | None = None

240 """The token to use as HTTP bearer authorization for remote files . If

241 `True`, will use the token generated when running `huggingface-cli login`

242 (stored in `~/.huggingface`)."""

243 hf_overrides: HfOverrides = field(default_factory=dict)

244 """If a dictionary, contains arguments to be forwarded to the Hugging Face

245 config. If a callable, it is called to update the HuggingFace config."""

254 - logits_processor_pattern: Optional[str] = None

246 + logits_processor_pattern: str | None = None

247 """Optional regex pattern specifying valid logits processor qualified names

248 that can be passed with the `logits_processors` extra completion argument.

249 Defaults to `None`, which allows no processors."""

@@ -269,7 +261,7 @@ class ModelConfig:

261 `--generation-config vllm`, only the override parameters are used."""

262 enable_sleep_mode: bool = False

263 """Enable sleep mode for the engine (only cuda platform is supported)."""

272 - model_impl: Union[str, ModelImpl] = "auto"

264 + model_impl: str | ModelImpl = "auto"

265 """Which implementation of the model to use:\n

266 - "auto" will try to use the vLLM implementation, if it exists, and fall

267 back to the Transformers implementation if no vLLM implementation is

@@ -278,36 +270,36 @@ class ModelConfig:

270 - "transformers" will use the Transformers model implementation.\n

271 - "terratorch" will use the TerraTorch model implementation.

272 """

281 - override_attention_dtype: Optional[str] = None

273 + override_attention_dtype: str | None = None

274 """Override dtype for attention"""

283 - logits_processors: Optional[list[Union[str, type[LogitsProcessor]]]] = None

275 + logits_processors: list[str | type[LogitsProcessor]] | None = None

276 """One or more logits processors' fully-qualified class names or class

277 definitions"""

286 - io_processor_plugin: Optional[str] = None

278 + io_processor_plugin: str | None = None

279 """IOProcessor plugin name to load at model startup"""

280

281 # Pooler config

290 - pooler_config: Optional[PoolerConfig] = None

282 + pooler_config: PoolerConfig | None = None

283 """Pooler config which controls the behaviour of output pooling in pooling

284 models."""

293 - override_pooler_config: Optional[Union[dict, PoolerConfig]] = None

285 + override_pooler_config: dict | PoolerConfig | None = None

286 """[DEPRECATED] Use `pooler_config` instead. This field will be removed in

287 v0.12.0 or v1.0.0, whichever is sooner."""

288

289 # Multimodal config and init vars

298 - multimodal_config: Optional[MultiModalConfig] = None

290 + multimodal_config: MultiModalConfig | None = None

291 """Configuration for multimodal model. If `None`, this will be inferred

292 from the architecture of `self.model`."""

301 - limit_mm_per_prompt: InitVar[Optional[dict[str, Union[int, dict[str, int]]]]] = None

302 - media_io_kwargs: InitVar[Optional[dict[str, dict[str, Any]]]] = None

303 - mm_processor_kwargs: InitVar[Optional[dict[str, Any]]] = None

304 - mm_processor_cache_gb: InitVar[Optional[float]] = None

305 - mm_processor_cache_type: InitVar[Optional[MMCacheType]] = None

306 - mm_shm_cache_max_object_size_mb: InitVar[Optional[int]] = None

307 - mm_encoder_tp_mode: InitVar[Optional[MMEncoderTPMode]] = None

308 - interleave_mm_strings: InitVar[Optional[bool]] = None

309 - skip_mm_profiling: InitVar[Optional[bool]] = None

310 - video_pruning_rate: InitVar[Optional[float]] = None

293 + limit_mm_per_prompt: InitVar[dict[str, int | dict[str, int]] | None] = None

294 + media_io_kwargs: InitVar[dict[str, dict[str, Any]] | None] = None

295 + mm_processor_kwargs: InitVar[dict[str, Any] | None] = None

296 + mm_processor_cache_gb: InitVar[float | None] = None

297 + mm_processor_cache_type: InitVar[MMCacheType | None] = None

298 + mm_shm_cache_max_object_size_mb: InitVar[int | None] = None

299 + mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None

300 + interleave_mm_strings: InitVar[bool | None] = None

301 + skip_mm_profiling: InitVar[bool | None] = None

302 + video_pruning_rate: InitVar[float | None] = None

303

304 def compute_hash(self) -> str:

305 """

@@ -369,7 +361,7 @@ class ModelConfig:

361

362 def _update_nested(

363 self,

372 - target: Union["PretrainedConfig", dict[str, Any]],

364 + target: PretrainedConfig | dict[str, Any],

365 updates: dict[str, Any],

366 ) -> None:

367 """Recursively updates a config or dict with nested updates."""

@@ -397,7 +389,7 @@ class ModelConfig:

389

390 def _apply_dict_overrides(

391 self,

400 - config: "PretrainedConfig",

392 + config: PretrainedConfig,

393 overrides: dict[str, Any],

394 ) -> None:

395 """Apply dict overrides, handling both nested configs and dict values."""

@@ -415,16 +407,16 @@ class ModelConfig:

407 def __post_init__(

408 self,

409 # Multimodal config init vars

418 - limit_mm_per_prompt: Optional[dict[str, int]],

419 - media_io_kwargs: Optional[dict[str, dict[str, Any]]],

420 - mm_processor_kwargs: Optional[dict[str, Any]],

421 - mm_processor_cache_gb: Optional[float],

422 - mm_processor_cache_type: Optional[MMCacheType],

423 - mm_shm_cache_max_object_size_mb: Optional[int],

424 - mm_encoder_tp_mode: Optional[MMEncoderTPMode],

425 - interleave_mm_strings: Optional[bool],

426 - skip_mm_profiling: Optional[bool],

427 - video_pruning_rate: Optional[float],

410 + limit_mm_per_prompt: dict[str, int] | None,

411 + media_io_kwargs: dict[str, dict[str, Any]] | None,

412 + mm_processor_kwargs: dict[str, Any] | None,

413 + mm_processor_cache_gb: float | None,

414 + mm_processor_cache_type: MMCacheType | None,

415 + mm_shm_cache_max_object_size_mb: int | None,

416 + mm_encoder_tp_mode: MMEncoderTPMode | None,

417 + interleave_mm_strings: bool | None,

418 + skip_mm_profiling: bool | None,

419 + video_pruning_rate: float | None,

420 ) -> None:

421 # Set the default seed to 0 in V1.

422 # NOTE(woosuk): In V0, we set the default seed to None because the

@@ -1209,7 +1201,7 @@ class ModelConfig:

1201 "Supported models implement the `SupportsPP` interface."

1202 )

1203

1212 - def get_sliding_window(self) -> Optional[int]:

1204 + def get_sliding_window(self) -> int | None:

1205 """Get the sliding window size from the HF text config if present."""

1206 return getattr(self.hf_text_config, "sliding_window", None)

1207

@@ -1479,7 +1471,7 @@ class ModelConfig:

1471 f"{block_type.value} layers"

1472 )

1473

1482 - def get_mamba_chunk_size(self) -> Optional[int]:

1474 + def get_mamba_chunk_size(self) -> int | None:

1475 """

1476 Returns the mamba chunk size if it exists

1477 """

@@ -1715,9 +1707,7 @@ class ModelConfig:

1707 return max_model_len

1708

1709

1718 -def get_served_model_name(

1719 - model: str, served_model_name: Optional[Union[str, list[str]]]

1720 -):

1710 +def get_served_model_name(model: str, served_model_name: str | list[str] | None):

1711 """

1712 If the input is a non-empty list, the first model_name in

1713 `served_model_name` is taken.

@@ -1761,9 +1751,9 @@ def iter_architecture_defaults():

1751 def try_match_architecture_defaults(

1752 architecture: str,

1753 *,

1764 - runner_type: Optional[RunnerType] = None,

1765 - convert_type: Optional[ConvertType] = None,

1766 -) -> Optional[tuple[str, tuple[RunnerType, ConvertType]]]:

1754 + runner_type: RunnerType | None = None,

1755 + convert_type: ConvertType | None = None,

1756 +) -> tuple[str, tuple[RunnerType, ConvertType]] | None:

1757 for suffix, (

1758 default_runner_type,

1759 default_convert_type,

@@ -1817,7 +1807,7 @@ def _find_dtype(

1807 model_id: str,

1808 config: PretrainedConfig,

1809 *,

1820 - revision: Optional[str],

1810 + revision: str | None,

1811 ):

1812 # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct

1813 # because config.torch_dtype can be None.

@@ -1902,10 +1892,10 @@ def _resolve_auto_dtype(

1892 def _get_and_verify_dtype(

1893 model_id: str,

1894 config: PretrainedConfig,

1905 - dtype: Union[str, torch.dtype],

1895 + dtype: str | torch.dtype,

1896 *,

1897 is_pooling_model: bool,

1908 - revision: Optional[str] = None,

1898 + revision: str | None = None,

1899 ) -> torch.dtype:

1900 config_dtype = _find_dtype(model_id, config, revision=revision)

1901 model_type = config.model_type

@@ -1947,7 +1937,7 @@ def _get_and_verify_dtype(

1937 def _get_head_dtype(

1938 config: PretrainedConfig, dtype: torch.dtype, runner_type: str

1939 ) -> torch.dtype:

1950 - head_dtype: Optional[Union[str, torch.dtype]] = getattr(config, "head_dtype", None)

1940 + head_dtype: str | torch.dtype | None = getattr(config, "head_dtype", None)

1941

1942 if head_dtype == "model":

1943 return dtype

@@ -1970,12 +1960,12 @@ def _get_head_dtype(

1960

1961 def _get_and_verify_max_len(

1962 hf_config: PretrainedConfig,

1973 - tokenizer_config: Optional[dict],

1974 - max_model_len: Optional[int],

1963 + tokenizer_config: dict | None,

1964 + max_model_len: int | None,

1965 disable_sliding_window: bool,

1976 - sliding_window: Optional[int],

1977 - spec_target_max_model_len: Optional[int] = None,

1978 - encoder_config: Optional[Any] = None,

1966 + sliding_window: int | None,

1967 + spec_target_max_model_len: int | None = None,

1968 + encoder_config: Any | None = None,

1969 ) -> int:

1970 """Get and verify the model's maximum length."""

1971 derived_max_model_len = float("inf")

vllm/config/multimodal.py

+14 -14

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

4 import hashlib

5 from collections.abc import Mapping

6 from dataclasses import field

7 -from typing import Any, Literal, Optional, Union

7 +from typing import Any, Literal, TypeAlias

8

9 from pydantic import ConfigDict, Field, field_validator

10 from pydantic.dataclasses import dataclass

@@ -23,31 +23,31 @@ class BaseDummyOptions:

23 class VideoDummyOptions(BaseDummyOptions):

24 """Options for generating dummy video data during profiling."""

25

26 - num_frames: Optional[int] = Field(None, gt=0)

27 - width: Optional[int] = Field(None, gt=0)

28 - height: Optional[int] = Field(None, gt=0)

26 + num_frames: int | None = Field(None, gt=0)

27 + width: int | None = Field(None, gt=0)

28 + height: int | None = Field(None, gt=0)

29

30

31 @dataclass(config=ConfigDict(extra="forbid"))

32 class ImageDummyOptions(BaseDummyOptions):

33 """Options for generating dummy image data during profiling."""

34

35 - width: Optional[int] = Field(None, gt=0)

36 - height: Optional[int] = Field(None, gt=0)

35 + width: int | None = Field(None, gt=0)

36 + height: int | None = Field(None, gt=0)

37

38

39 @dataclass(config=ConfigDict(extra="forbid"))

40 class AudioDummyOptions(BaseDummyOptions):

41 """Options for generating dummy audio data during profiling."""

42

43 - length: Optional[int] = Field(None, gt=0)

43 + length: int | None = Field(None, gt=0)

44

45

46 MMEncoderTPMode = Literal["weights", "data"]

47 MMCacheType = Literal["shm", "lru"]

48 -DummyOptions = Union[

49 - BaseDummyOptions, VideoDummyOptions, ImageDummyOptions, AudioDummyOptions

50 -]

48 +DummyOptions: TypeAlias = (

49 + BaseDummyOptions | VideoDummyOptions | ImageDummyOptions | AudioDummyOptions

50 +)

51

52

53 @config

@@ -75,7 +75,7 @@ class MultiModalConfig:

75 """Additional args passed to process media inputs, keyed by modalities.

76 For example, to set num_frames for video, set

77 `--media-io-kwargs '{"video": {"num_frames": 40} }'`"""

78 - mm_processor_kwargs: Optional[dict[str, object]] = None

78 + mm_processor_kwargs: dict[str, object] | None = None

79 """Arguments to be forwarded to the model's processor for multi-modal data,

80 e.g., image processor. Overrides for the multi-modal processor obtained

81 from `transformers.AutoProcessor.from_pretrained`.

@@ -123,7 +123,7 @@ class MultiModalConfig:

123 This reduces engine startup time but shifts the responsibility to users for

124 estimating the peak memory usage of the activation of multimodal encoder and

125 embedding cache."""

126 - video_pruning_rate: Optional[float] = None

126 + video_pruning_rate: float | None = None

127 """Sets pruning rate for video pruning via Efficient Video Sampling.

128 Value sits in range [0;1) and determines fraction of media tokens

129 from each video to be pruned.

@@ -132,7 +132,7 @@ class MultiModalConfig:

132 @field_validator("limit_per_prompt", mode="before")

133 @classmethod

134 def _validate_limit_per_prompt(

135 - cls, value: dict[str, Union[int, dict[str, int]]]

135 + cls, value: dict[str, int | dict[str, int]]

136 ) -> dict[str, DummyOptions]:

137 for k, v in value.items():

138 # Handle legacy format where only count is specified

@@ -179,7 +179,7 @@ class MultiModalConfig:

179 return 999

180 return limit_data.count

181

182 - def get_dummy_options(self, modality: str) -> Optional[BaseDummyOptions]:

182 + def get_dummy_options(self, modality: str) -> BaseDummyOptions | None:

183 """

184 Get the configurable dummy data options for a modality.

185 Returns None if no options are configured for this modality.

vllm/config/observability.py

+4 -4

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

3

4 import hashlib

5 from functools import cached_property

6 -from typing import Any, Literal, Optional, cast

6 +from typing import Any, Literal, cast

7

8 from pydantic.dataclasses import dataclass

9

@@ -18,7 +18,7 @@ DetailedTraceModules = Literal["model", "worker", "all"]

18 class ObservabilityConfig:

19 """Configuration for observability - metrics and tracing."""

20

21 - show_hidden_metrics_for_version: Optional[str] = None

21 + show_hidden_metrics_for_version: str | None = None

22 """Enable deprecated Prometheus metrics that have been hidden since the

23 specified version. For example, if a previously deprecated metric has been

24 hidden since the v0.7.0 release, you use

@@ -33,10 +33,10 @@ class ObservabilityConfig:

33 return False

34 return version._prev_minor_version_was(self.show_hidden_metrics_for_version)

35

36 - otlp_traces_endpoint: Optional[str] = None

36 + otlp_traces_endpoint: str | None = None

37 """Target URL to which OpenTelemetry traces will be sent."""

38

39 - collect_detailed_traces: Optional[list[DetailedTraceModules]] = None

39 + collect_detailed_traces: list[DetailedTraceModules] | None = None

40 """It makes sense to set this only if `--otlp-traces-endpoint` is set. If

41 set, it will collect detailed traces for the specified modules. This

42 involves use of possibly costly and or blocking operations and hence might

vllm/config/parallel.py

+13 -13

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

3

4 import hashlib

5 import os

6 -from typing import TYPE_CHECKING, Any, Literal, Optional, Union

6 +from typing import TYPE_CHECKING, Any, Literal

7

8 import torch

9 from pydantic import Field, model_validator

@@ -75,7 +75,7 @@ class ParallelConfig:

75 """Number of local data parallel groups."""

76 data_parallel_rank: int = 0

77 """Rank of the data parallel group."""

78 - data_parallel_rank_local: Optional[int] = None

78 + data_parallel_rank_local: int | None = None

79 """Local rank of the data parallel group,

80 set only in SPMD mode."""

81 data_parallel_master_ip: str = "127.0.0.1"

@@ -113,24 +113,24 @@ class ParallelConfig:

113 with 4 experts and 2 ranks, rank 0 will have experts [0, 2] and rank 1

114 will have experts [1, 3]. This strategy can help improve load balancing

115 for grouped expert models with no redundant experts."""

116 - num_redundant_experts: Optional[int] = None

116 + num_redundant_experts: int | None = None

117 """`num_redundant_experts` is deprecated and has been replaced with

118 `eplb_config.num_redundant_experts`. This will be removed in v0.12.0.

119 Please use `eplb_config.num_redundant_experts` instead."""

120 - eplb_window_size: Optional[int] = None

120 + eplb_window_size: int | None = None

121 """`eplb_window_size` is deprecated and has been replaced with

122 `eplb_config.window_size`. This will be removed in v0.12.0.

123 Please use `eplb_config.window_size` instead."""

124 - eplb_step_interval: Optional[int] = None

124 + eplb_step_interval: int | None = None

125 """`eplb_step_interval` is deprecated and has been replaced with

126 `eplb_config.step_interval`. This will be removed in v0.12.0.

127 Please use `eplb_config.step_interval` instead."""

128 - eplb_log_balancedness: Optional[bool] = None

128 + eplb_log_balancedness: bool | None = None

129 """`eplb_log_balancedness` is deprecated and has been replaced with

130 `eplb_config.log_balancedness`. This will be removed in v0.12.0.

131 Please use `eplb_config.log_balancedness` instead."""

132

133 - max_parallel_loading_workers: Optional[int] = None

133 + max_parallel_loading_workers: int | None = None

134 """Maximum number of parallel loading workers when loading model

135 sequentially in multiple batches. To avoid RAM OOM when using tensor

136 parallel and large models."""

@@ -159,15 +159,15 @@ class ParallelConfig:

159 ray_workers_use_nsight: bool = False

160 """Whether to profile Ray workers with nsight, see https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html#profiling-nsight-profiler."""

161

162 - ray_runtime_env: Optional[RuntimeEnv] = None

162 + ray_runtime_env: RuntimeEnv | None = None

163 """Ray runtime environment to pass to distributed workers."""

164

165 - placement_group: Optional[PlacementGroup] = None

165 + placement_group: PlacementGroup | None = None

166 """ray distributed model workers placement group."""

167

168 - distributed_executor_backend: Optional[

169 - Union[str, DistributedExecutorBackend, type[ExecutorBase]]

170 - ] = None

168 + distributed_executor_backend: (

169 + str | DistributedExecutorBackend | type[ExecutorBase] | None

170 + ) = None

171 """Backend to use for distributed model

172 workers, either "ray" or "mp" (multiprocessing). If the product

173 of pipeline_parallel_size and tensor_parallel_size is less than

@@ -306,7 +306,7 @@ class ParallelConfig:

306 )

307

308 max_retries = 5

309 - last_exc: Optional[Exception] = None

309 + last_exc: Exception | None = None

310 for _ in range(max_retries):

311 try:

312 # use gloo since the engine process might not have cuda device

vllm/config/pooler.py

+11 -11

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

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

3

4 import hashlib

5 -from typing import Any, Optional

5 +from typing import Any

6

7 from pydantic.dataclasses import dataclass

8

@@ -14,23 +14,23 @@ from vllm.config.utils import config

14 class PoolerConfig:

15 """Controls the behavior of output pooling in pooling models."""

16

17 - pooling_type: Optional[str] = None

17 + pooling_type: str | None = None

18 """

19 The pooling method of the pooling model. This should be a key in

20 [`vllm.model_executor.layers.pooler.PoolingType`][].

21 """

22

23 ## for embeddings models

24 - normalize: Optional[bool] = None

24 + normalize: bool | None = None

25 """

26 Whether to normalize the embeddings outputs. Defaults to True.

27 """

28 - dimensions: Optional[int] = None

28 + dimensions: int | None = None

29 """

30 Reduce the dimensions of embeddings if model

31 support matryoshka representation. Defaults to None.

32 """

33 - enable_chunked_processing: Optional[bool] = None

33 + enable_chunked_processing: bool | None = None

34 """

35 Whether to enable chunked processing for long inputs that exceed the model's

36 maximum position embeddings. When enabled, long inputs will be split into

@@ -38,7 +38,7 @@ class PoolerConfig:

38 This allows embedding models to handle arbitrarily long text without CUDA

39 errors. Defaults to False.

40 """

41 - max_embed_len: Optional[int] = None

41 + max_embed_len: int | None = None

42 """

43 Maximum input length allowed for embedding generation. When set, allows

44 inputs longer than max_embed_len to be accepted for embedding models.

@@ -48,29 +48,29 @@ class PoolerConfig:

48 """

49

50 ## for classification models

51 - activation: Optional[bool] = None

51 + activation: bool | None = None

52 """

53 Whether to apply activation function to the classification outputs.

54 Defaults to True.

55 """

56 - logit_bias: Optional[float] = None

56 + logit_bias: float | None = None

57 """

58 If provided, apply classification logit biases. Defaults to None.

59 """

60

61 ## for reward models

62 - softmax: Optional[bool] = None

62 + softmax: bool | None = None

63 """

64 Whether to apply softmax to the reward outputs.

65 Defaults to True.

66 """

67 - step_tag_id: Optional[int] = None

67 + step_tag_id: int | None = None

68 """

69 If set, only the score corresponding to the ``step_tag_id`` in the

70 generated sentence should be returned. Otherwise, the scores for all tokens

71 are returned.

72 """

73 - returned_token_ids: Optional[list[int]] = None

73 + returned_token_ids: list[int] | None = None

74 """

75 A list of indices for the vocabulary dimensions to be extracted,

76 such as the token IDs of ``good_token`` and ``bad_token`` in the

vllm/config/scheduler.py

+2 -2

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

3

4 import hashlib

5 from dataclasses import InitVar, field

6 -from typing import Any, Literal, Union

6 +from typing import Any, Literal

7

8 from pydantic import SkipValidation, model_validator

9 from pydantic.dataclasses import dataclass

@@ -133,7 +133,7 @@ class SchedulerConfig:

133

134 # scheduler class or path. "vllm.core.scheduler.Scheduler" (default)

135 # or "mod.custom_class".

136 - scheduler_cls: Union[str, type[object]] = "vllm.core.scheduler.Scheduler"

136 + scheduler_cls: str | type[object] = "vllm.core.scheduler.Scheduler"

137 """The scheduler class to use. "vllm.core.scheduler.Scheduler" is the

138 default scheduler. Can be a class directly or the path to a class of form

139 "mod.custom_class"."""

vllm/config/speculative.py

+15 -15

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

3

4 import ast

5 import hashlib

6 -from typing import TYPE_CHECKING, Any, Literal, Optional

6 +from typing import TYPE_CHECKING, Any, Literal

7

8 from pydantic import SkipValidation, model_validator

9 from pydantic.dataclasses import dataclass

@@ -59,16 +59,16 @@ MTP_MODEL_TYPES = (

59 class SpeculativeConfig:

60 """Configuration for speculative decoding."""

61

62 - enforce_eager: Optional[bool] = None

62 + enforce_eager: bool | None = None

63 """Override the default enforce_eager from model_config"""

64 # General speculative decoding control

65 num_speculative_tokens: SkipValidation[int] = None # type: ignore

66 """The number of speculative tokens, if provided. It will default to the

67 number in the draft model config if present, otherwise, it is required."""

68 - model: Optional[str] = None

68 + model: str | None = None

69 """The name of the draft model, eagle head, or additional weights, if

70 provided."""

71 - method: Optional[SpeculativeMethod] = None

71 + method: SpeculativeMethod | None = None

72 """The name of the speculative method to use. If users provide and set the

73 `model` param, the speculative method type will be detected automatically

74 if possible, if `model` param is not provided, the method name must be

@@ -76,7 +76,7 @@ class SpeculativeConfig:

76

77 If using `ngram` method, the related configuration `prompt_lookup_max` and

78 `prompt_lookup_min` should be considered."""

79 - draft_tensor_parallel_size: Optional[int] = None

79 + draft_tensor_parallel_size: int | None = None

80 """The degree of the tensor parallelism for the draft model. Can only be 1

81 or the same as the target model's tensor parallel size."""

82 disable_logprobs: bool = True

@@ -85,24 +85,24 @@ class SpeculativeConfig:

85 according to the log probability settings in SamplingParams."""

86

87 # Draft model configuration

88 - quantization: Optional[me_quant.QuantizationMethods] = None

88 + quantization: me_quant.QuantizationMethods | None = None

89 """Quantization method that was used to quantize the draft model weights.

90 If `None`, we assume the model weights are not quantized. Note that it only

91 takes effect when using the draft model-based speculative method."""

92 - max_model_len: Optional[int] = None

92 + max_model_len: int | None = None

93 """The maximum model length of the draft model. Used when testing the

94 ability to skip speculation for some sequences."""

95 - revision: Optional[str] = None

95 + revision: str | None = None

96 """The specific model version to use for the draft model. It can be a

97 branch name, a tag name, or a commit id. If unspecified, will use the

98 default version."""

99 - code_revision: Optional[str] = None

99 + code_revision: str | None = None

100 """The specific revision to use for the draft model code on Hugging Face

101 Hub. It can be a branch name, a tag name, or a commit id. If unspecified,

102 will use the default version."""

103

104 # Advanced control

105 - disable_by_batch_size: Optional[int] = None

105 + disable_by_batch_size: int | None = None

106 """Disable speculative decoding for new incoming requests when the number

107 of enqueued requests is larger than this value, if provided."""

108 disable_padded_drafter_batch: bool = False

@@ -112,14 +112,14 @@ class SpeculativeConfig:

112 only affects the EAGLE method of speculation."""

113

114 # Ngram proposer configuration

115 - prompt_lookup_max: Optional[int] = None

115 + prompt_lookup_max: int | None = None

116 """Maximum size of ngram token window when using Ngram proposer, required

117 when method is set to ngram."""

118 - prompt_lookup_min: Optional[int] = None

118 + prompt_lookup_min: int | None = None

119 """Minimum size of ngram token window when using Ngram proposer, if

120 provided. Defaults to 1."""

121

122 - speculative_token_tree: Optional[str] = None

122 + speculative_token_tree: str | None = None

123 """Specifies the tree structure for speculative token generation.

124 """

125 # required configuration params passed from engine

@@ -449,7 +449,7 @@ class SpeculativeConfig:

449

450 @staticmethod

451 def _maybe_override_draft_max_model_len(

452 - speculative_max_model_len: Optional[int],

452 + speculative_max_model_len: int | None,

453 draft_max_model_len: int,

454 target_max_model_len: int,

455 ) -> int:

@@ -488,7 +488,7 @@ class SpeculativeConfig:

488 @staticmethod

489 def _verify_and_get_draft_tp(

490 target_parallel_config: ParallelConfig,

491 - speculative_draft_tensor_parallel_size: Optional[int],

491 + speculative_draft_tensor_parallel_size: int | None,

492 draft_hf_config: PretrainedConfig,

493 ) -> int:

494 """

vllm/config/speech_to_text.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 from pydantic.dataclasses import dataclass

6

@@ -28,7 +27,7 @@ class SpeechToTextConfig:

27 splitting long audio. This helps maintain context across chunk boundaries

28 and improves transcription quality at split points."""

29

31 - min_energy_split_window_size: Optional[int] = 1600

30 + min_energy_split_window_size: int | None = 1600

31 """Window size in samples for finding low-energy (quiet) regions to split

32 audio chunks. The algorithm looks for the quietest moment within this

33 window to minimize cutting through speech. Default 1600 samples ≈ 100ms

vllm/config/vllm.py

+16 -16

@@ -10,7 +10,7 @@ from contextlib import contextmanager

10 from dataclasses import field, replace

11 from functools import lru_cache

12 from pathlib import Path

13 -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union

13 +from typing import TYPE_CHECKING, Any, TypeVar

14

15 import torch

16 from pydantic import ConfigDict

@@ -69,17 +69,17 @@ class VllmConfig:

69 """Device configuration."""

70 load_config: LoadConfig = field(default_factory=LoadConfig)

71 """Load configuration."""

72 - lora_config: Optional[LoRAConfig] = None

72 + lora_config: LoRAConfig | None = None

73 """LoRA configuration."""

74 - speculative_config: Optional[SpeculativeConfig] = None

74 + speculative_config: SpeculativeConfig | None = None

75 """Speculative decoding configuration."""

76 structured_outputs_config: StructuredOutputsConfig = field(

77 default_factory=StructuredOutputsConfig

78 )

79 """Structured outputs configuration."""

80 - observability_config: Optional[ObservabilityConfig] = None

80 + observability_config: ObservabilityConfig | None = None

81 """Observability configuration."""

82 - quant_config: Optional[QuantizationConfig] = None

82 + quant_config: QuantizationConfig | None = None

83 """Quantization configuration."""

84 compilation_config: CompilationConfig = field(default_factory=CompilationConfig)

85 """`torch.compile` and cudagraph capture configuration for the model.

@@ -96,14 +96,14 @@ class VllmConfig:

96 You can specify the full compilation config like so:

97 `{"level": 3, "cudagraph_capture_sizes": [1, 2, 4, 8]}`

98 """

99 - kv_transfer_config: Optional[KVTransferConfig] = None

99 + kv_transfer_config: KVTransferConfig | None = None

100 """The configurations for distributed KV cache transfer."""

101 - kv_events_config: Optional[KVEventsConfig] = None

101 + kv_events_config: KVEventsConfig | None = None

102 """The configurations for event publishing."""

103 # some opaque config, only used to provide additional information

104 # for the hash computation, mainly used for testing, debugging or out of

105 # tree config registration.

106 - additional_config: Union[dict, SupportsHash] = field(default_factory=dict)

106 + additional_config: dict | SupportsHash = field(default_factory=dict)

107 """Additional config for specified platform. Different platforms may

108 support different configs. Make sure the configs are valid for the platform

109 you are using. Contents must be hashable."""

@@ -212,7 +212,7 @@ class VllmConfig:

212 @staticmethod

213 def _get_quantization_config(

214 model_config: ModelConfig, load_config: LoadConfig

215 - ) -> Optional[QuantizationConfig]:

215 + ) -> QuantizationConfig | None:

216 """Get the quantization config."""

217 from vllm.platforms import current_platform

218

@@ -245,7 +245,7 @@ class VllmConfig:

245 @staticmethod

246 def get_quantization_config(

247 model_config: ModelConfig, load_config: LoadConfig

248 - ) -> Optional[QuantizationConfig]:

248 + ) -> QuantizationConfig | None:

249 import copy

250

251 # For some reason, the _ version of this modifies the model_config

@@ -257,7 +257,7 @@ class VllmConfig:

257 def with_hf_config(

258 self,

259 hf_config: PretrainedConfig,

260 - architectures: Optional[list[str]] = None,

260 + architectures: list[str] | None = None,

261 ) -> "VllmConfig":

262 if architectures is not None:

263 hf_config = copy.deepcopy(hf_config)

@@ -740,7 +740,7 @@ class VllmConfig:

740 f"Model: {self.model_config.model}"

741 )

742

743 - def compile_debug_dump_path(self) -> Optional[Path]:

743 + def compile_debug_dump_path(self) -> Path | None:

744 """Returns a rank-aware path for dumping

745 torch.compile debug information.

746 """

@@ -790,13 +790,13 @@ class VllmConfig:

790 )

791

792

793 -_current_vllm_config: Optional[VllmConfig] = None

794 -_current_prefix: Optional[str] = None

793 +_current_vllm_config: VllmConfig | None = None

794 +_current_prefix: str | None = None

795

796

797 @contextmanager

798 def set_current_vllm_config(

799 - vllm_config: VllmConfig, check_compile=False, prefix: Optional[str] = None

799 + vllm_config: VllmConfig, check_compile=False, prefix: str | None = None

800 ):

801 """

802 Temporarily set the current vLLM config.

@@ -866,7 +866,7 @@ T = TypeVar("T")

866 def get_layers_from_vllm_config(

867 vllm_config: VllmConfig,

868 layer_type: type[T],

869 - layer_names: Optional[list[str]] = None,

869 + layer_names: list[str] | None = None,

870 ) -> dict[str, T]:

871 """

872 Get layers from the vLLM config.

vllm/connections.py

+14 -15

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

3

4 from collections.abc import Mapping, MutableMapping

5 from pathlib import Path

6 -from typing import Optional

6 from urllib.parse import urlparse

7

8 import aiohttp

@@ -20,8 +19,8 @@ class HTTPConnection:

19

20 self.reuse_client = reuse_client

21

23 - self._sync_client: Optional[requests.Session] = None

24 - self._async_client: Optional[aiohttp.ClientSession] = None

22 + self._sync_client: requests.Session | None = None

23 + self._async_client: aiohttp.ClientSession | None = None

24

25 def get_sync_client(self) -> requests.Session:

26 if self._sync_client is None or not self.reuse_client:

@@ -53,8 +52,8 @@ class HTTPConnection:

52 url: str,

53 *,

54 stream: bool = False,

56 - timeout: Optional[float] = None,

57 - extra_headers: Optional[Mapping[str, str]] = None,

55 + timeout: float | None = None,

56 + extra_headers: Mapping[str, str] | None = None,

57 allow_redirects: bool = True,

58 ):

59 self._validate_http_url(url)

@@ -74,8 +73,8 @@ class HTTPConnection:

73 self,

74 url: str,

75 *,

77 - timeout: Optional[float] = None,

78 - extra_headers: Optional[Mapping[str, str]] = None,

76 + timeout: float | None = None,

77 + extra_headers: Mapping[str, str] | None = None,

78 allow_redirects: bool = True,

79 ):

80 self._validate_http_url(url)

@@ -91,7 +90,7 @@ class HTTPConnection:

90 )

91

92 def get_bytes(

94 - self, url: str, *, timeout: Optional[float] = None, allow_redirects: bool = True

93 + self, url: str, *, timeout: float | None = None, allow_redirects: bool = True

94 ) -> bytes:

95 with self.get_response(

96 url, timeout=timeout, allow_redirects=allow_redirects

@@ -104,7 +103,7 @@ class HTTPConnection:

103 self,

104 url: str,

105 *,

107 - timeout: Optional[float] = None,

106 + timeout: float | None = None,

107 allow_redirects: bool = True,

108 ) -> bytes:

109 async with await self.get_async_response(

@@ -114,7 +113,7 @@ class HTTPConnection:

113

114 return await r.read()

115

117 - def get_text(self, url: str, *, timeout: Optional[float] = None) -> str:

116 + def get_text(self, url: str, *, timeout: float | None = None) -> str:

117 with self.get_response(url, timeout=timeout) as r:

118 r.raise_for_status()

119

@@ -124,14 +123,14 @@ class HTTPConnection:

123 self,

124 url: str,

125 *,

127 - timeout: Optional[float] = None,

126 + timeout: float | None = None,

127 ) -> str:

128 async with await self.get_async_response(url, timeout=timeout) as r:

129 r.raise_for_status()

130

131 return await r.text()

132

134 - def get_json(self, url: str, *, timeout: Optional[float] = None) -> str:

133 + def get_json(self, url: str, *, timeout: float | None = None) -> str:

134 with self.get_response(url, timeout=timeout) as r:

135 r.raise_for_status()

136

@@ -141,7 +140,7 @@ class HTTPConnection:

140 self,

141 url: str,

142 *,

144 - timeout: Optional[float] = None,

143 + timeout: float | None = None,

144 ) -> str:

145 async with await self.get_async_response(url, timeout=timeout) as r:

146 r.raise_for_status()

@@ -153,7 +152,7 @@ class HTTPConnection:

152 url: str,

153 save_path: Path,

154 *,

156 - timeout: Optional[float] = None,

155 + timeout: float | None = None,

156 chunk_size: int = 128,

157 ) -> Path:

158 with self.get_response(url, timeout=timeout) as r:

@@ -170,7 +169,7 @@ class HTTPConnection:

169 url: str,

170 save_path: Path,

171 *,

173 - timeout: Optional[float] = None,

172 + timeout: float | None = None,

173 chunk_size: int = 128,

174 ) -> Path:

175 async with await self.get_async_response(url, timeout=timeout) as r:

vllm/device_allocator/cumem.py

+7 -6

@@ -11,8 +11,9 @@

11 import dataclasses

12 import gc

13 import os

14 +from collections.abc import Callable

15 from contextlib import contextmanager

15 -from typing import Any, Callable, Optional, Union

16 +from typing import Any

17

18 import torch

19

@@ -22,7 +23,7 @@ from vllm.utils import is_pin_memory_available

23 logger = init_logger(__name__)

24

25

25 -def find_loaded_library(lib_name) -> Optional[str]:

26 +def find_loaded_library(lib_name) -> str | None:

27 """

28 According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html,

29 the file `/proc/self/maps` contains the memory maps of the process, which includes the

@@ -78,7 +79,7 @@ HandleType = tuple[int, int, int, int]

79 class AllocationData:

80 handle: HandleType

81 tag: str

81 - cpu_backup_tensor: Optional[torch.Tensor] = None

82 + cpu_backup_tensor: torch.Tensor | None = None

83

84

85 def create_and_map(allocation_handle: HandleType) -> None:

@@ -197,7 +198,7 @@ class CuMemAllocator:

198 )

199 return data.handle

200

200 - def sleep(self, offload_tags: Optional[Union[tuple[str, ...], str]] = None) -> None:

201 + def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:

202 """

203 Put the allocator in sleep mode.

204 All data in the memory allocation with the specified tag will be

@@ -247,7 +248,7 @@ class CuMemAllocator:

248 gc.collect()

249 torch.cuda.empty_cache()

250

250 - def wake_up(self, tags: Optional[list[str]] = None) -> None:

251 + def wake_up(self, tags: list[str] | None = None) -> None:

252 """

253 Wake up the allocator from sleep mode.

254 All data that is previously offloaded will be loaded back to GPU

@@ -272,7 +273,7 @@ class CuMemAllocator:

273 data.cpu_backup_tensor = None

274

275 @contextmanager

275 - def use_memory_pool(self, tag: Optional[str] = None):

276 + def use_memory_pool(self, tag: str | None = None):

277 """

278 A context manager to use the memory pool.

279 All memory allocation created inside the context will be allocated

vllm/distributed/communication_op.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Optional, Union

4 +from typing import Any

5

6 import torch

7 import torch.distributed

@@ -30,13 +30,13 @@ def tensor_model_parallel_reduce_scatter(

30

31 def tensor_model_parallel_gather(

32 input_: torch.Tensor, dst: int = 0, dim: int = -1

33 -) -> Optional[torch.Tensor]:

33 +) -> torch.Tensor | None:

34 """Gather the input tensor across model parallel group."""

35 return get_tp_group().gather(input_, dst, dim)

36

37

38 def broadcast_tensor_dict(

39 - tensor_dict: Optional[dict[Any, Union[torch.Tensor, Any]]] = None, src: int = 0

39 + tensor_dict: dict[Any, torch.Tensor | Any] | None = None, src: int = 0

40 ):

41 if not torch.distributed.is_initialized():

42 return tensor_dict

vllm/distributed/device_communicators/all2all.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Any, Optional

3 +from typing import Any

4

5 import torch

6 import torch.distributed as dist

@@ -368,7 +368,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

368 return handle

369

370 # DeepEP LL uses RDMA so no SMs are used for communication

371 - def max_sms_used(self) -> Optional[int]:

371 + def max_sms_used(self) -> int | None:

372 return 0

373

374

vllm/distributed/device_communicators/all_reduce_utils.py

+4 -4

@@ -10,7 +10,7 @@ import sys

10 import tempfile

11 from collections.abc import Sequence

12 from itertools import product

13 -from typing import Any, Optional

13 +from typing import Any

14

15 import torch

16 import torch.distributed as dist

@@ -86,7 +86,7 @@ def producer(

86 producer_queue,

87 consumer_queue,

88 result_queue,

89 - cuda_visible_devices: Optional[str] = None,

89 + cuda_visible_devices: str | None = None,

90 ):

91 if cuda_visible_devices is not None:

92 update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})

@@ -120,7 +120,7 @@ def consumer(

120 producer_queue,

121 consumer_queue,

122 result_queue,

123 - cuda_visible_devices: Optional[str] = None,

123 + cuda_visible_devices: str | None = None,

124 ):

125 if cuda_visible_devices is not None:

126 update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})

@@ -253,7 +253,7 @@ def can_actually_p2p(

253 # e.g. used by different vllm engines. The device id in the cache file is a

254 # **local** device id, i.e. from 0 to num_dev-1, where num_dev is the number

255 # of visible devices in the vllm engine.

256 -_gpu_p2p_access_cache: Optional[dict[str, bool]] = None

256 +_gpu_p2p_access_cache: dict[str, bool] | None = None

257

258

259 def gpu_p2p_access_check(src: int, tgt: int) -> bool:

vllm/distributed/device_communicators/base_device_communicator.py

+11 -12

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import threading

4 -from typing import Optional, Union

4 from weakref import WeakValueDictionary

5

6 import torch

@@ -75,7 +74,7 @@ class All2AllManagerBase:

74 def set_num_sms(self, num_sms: int):

75 pass

76

78 - def max_sms_used(self) -> Optional[int]:

77 + def max_sms_used(self) -> int | None:

78 return None # None means it could use the whole GPU

79

80 def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False):

@@ -96,8 +95,8 @@ class DeviceCommunicatorBase:

95 def __init__(

96 self,

97 cpu_group: ProcessGroup,

99 - device: Optional[torch.device] = None,

100 - device_group: Optional[ProcessGroup] = None,

98 + device: torch.device | None = None,

99 + device_group: ProcessGroup | None = None,

100 unique_name: str = "",

101 ):

102 self.device = device or torch.device("cpu")

@@ -123,7 +122,7 @@ class DeviceCommunicatorBase:

122

123 self.is_ep_communicator = "ep" in unique_name

124 self.use_all2all = self.is_ep_communicator and use_ep

126 - self.all2all_manager: Optional[All2AllManagerBase] = None

125 + self.all2all_manager: All2AllManagerBase | None = None

126

127 def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:

128 dist.all_reduce(input_, group=self.device_group)

@@ -156,10 +155,10 @@ class DeviceCommunicatorBase:

155

156 def all_gatherv(

157 self,

159 - input_: Union[torch.Tensor, list[torch.Tensor]],

158 + input_: torch.Tensor | list[torch.Tensor],

159 dim: int = 0,

161 - sizes: Optional[list[int]] = None,

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

160 + sizes: list[int] | None = None,

161 + ) -> torch.Tensor | list[torch.Tensor]:

162 raise NotImplementedError

163

164 def reduce_scatter(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:

@@ -196,13 +195,13 @@ class DeviceCommunicatorBase:

195 return output_tensor.movedim(0, dim).contiguous()

196

197 def reduce_scatterv(

199 - self, input_: torch.Tensor, dim: int = -1, sizes: Optional[list[int]] = None

198 + self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None

199 ) -> torch.Tensor:

200 raise NotImplementedError

201

202 def gather(

203 self, input_: torch.Tensor, dst: int = 0, dim: int = -1

205 - ) -> Optional[torch.Tensor]:

204 + ) -> torch.Tensor | None:

205 """

206 NOTE: We assume that the input tensor is on the same device across

207 all the ranks.

@@ -231,7 +230,7 @@ class DeviceCommunicatorBase:

230 output_tensor = None

231 return output_tensor

232

234 - def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:

233 + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:

234 """Sends a tensor to the destination rank in a blocking way"""

235 """NOTE: `dst` is the local rank of the destination rank."""

236 if dst is None:

@@ -239,7 +238,7 @@ class DeviceCommunicatorBase:

238 torch.distributed.send(tensor, self.ranks[dst], self.device_group)

239

240 def recv(

242 - self, size: torch.Size, dtype: torch.dtype, src: Optional[int] = None

241 + self, size: torch.Size, dtype: torch.dtype, src: int | None = None

242 ) -> torch.Tensor:

243 """Receives a tensor from the source rank."""

244 """NOTE: `src` is the local rank of the source rank."""

vllm/distributed/device_communicators/cpu_communicator.py

+12 -12

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

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

3

4 import os

5 -from typing import Any, Optional, Union

5 +from typing import Any

6

7 import torch

8 from torch.distributed import ProcessGroup

@@ -18,8 +18,8 @@ class CpuCommunicator(DeviceCommunicatorBase):

18 def __init__(

19 self,

20 cpu_group: ProcessGroup,

21 - device: Optional[torch.device] = None,

22 - device_group: Optional[ProcessGroup] = None,

21 + device: torch.device | None = None,

22 + device_group: ProcessGroup | None = None,

23 unique_name: str = "",

24 ):

25 super().__init__(cpu_group, device, device_group, unique_name)

@@ -38,7 +38,7 @@ class CpuCommunicator(DeviceCommunicatorBase):

38

39 def gather(

40 self, input_: torch.Tensor, dst: int = 0, dim: int = -1

41 - ) -> Optional[torch.Tensor]:

41 + ) -> torch.Tensor | None:

42 """

43 NOTE: We assume that the input tensor is on the same device across

44 all the ranks.

@@ -99,7 +99,7 @@ class CpuCommunicator(DeviceCommunicatorBase):

99

100 def send_tensor_dict(

101 self,

102 - tensor_dict: dict[str, Union[torch.Tensor, Any]],

102 + tensor_dict: dict[str, torch.Tensor | Any],

103 dst: int,

104 ) -> None:

105 return self.dist_module.send_tensor_dict(tensor_dict, dst)

@@ -107,7 +107,7 @@ class CpuCommunicator(DeviceCommunicatorBase):

107 def recv_tensor_dict(

108 self,

109 src: int,

110 - ) -> dict[str, Union[torch.Tensor, Any]]:

110 + ) -> dict[str, torch.Tensor | Any]:

111 return self.dist_module.recv_tensor_dict(src)

112

113

@@ -140,16 +140,16 @@ class _CPUSHMDistributed:

140 return handle

141

142 def all_reduce(

143 - self, input: torch.Tensor, group: Optional[ProcessGroup] = None

143 + self, input: torch.Tensor, group: ProcessGroup | None = None

144 ) -> None:

145 torch.ops._C.shm_allreduce(self.handle, input)

146

147 def gather(

148 self,

149 input: torch.Tensor,

150 - gather_list: Optional[list[torch.Tensor]],

150 + gather_list: list[torch.Tensor] | None,

151 dst: int = -1,

152 - group: Optional[ProcessGroup] = None,

152 + group: ProcessGroup | None = None,

153 ) -> None:

154 # Note: different from the torch gather, here we use local dst rank.

155 torch.ops._C.shm_gather(

@@ -163,13 +163,13 @@ class _CPUSHMDistributed:

163 self,

164 output: torch.Tensor,

165 input: torch.Tensor,

166 - group: Optional[ProcessGroup] = None,

166 + group: ProcessGroup | None = None,

167 ) -> None:

168 torch.ops._C.shm_all_gather(self.handle, input, output)

169

170 def send_tensor_dict(

171 self,

172 - tensor_dict: dict[str, Union[torch.Tensor, Any]],

172 + tensor_dict: dict[str, torch.Tensor | Any],

173 dst: int,

174 ) -> None:

175 key_list = list(tensor_dict.keys())

@@ -191,7 +191,7 @@ class _CPUSHMDistributed:

191 def recv_tensor_dict(

192 self,

193 src: int,

194 - ) -> dict[str, Union[torch.Tensor, Any]]:

194 + ) -> dict[str, torch.Tensor | Any]:

195 tensor_list = torch.ops._C.shm_recv_tensor_list(self.handle, src)

196

197 value_list: list[torch.Tensor] = tensor_list[:-1]

vllm/distributed/device_communicators/cuda_communicator.py

+12 -13

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 import torch

6 from torch.distributed import ProcessGroup

@@ -26,8 +25,8 @@ class CudaCommunicator(DeviceCommunicatorBase):

25 def __init__(

26 self,

27 cpu_group: ProcessGroup,

29 - device: Optional[torch.device] = None,

30 - device_group: Optional[ProcessGroup] = None,

28 + device: torch.device | None = None,

29 + device_group: ProcessGroup | None = None,

30 unique_name: str = "",

31 ):

32 super().__init__(cpu_group, device, device_group, unique_name)

@@ -54,7 +53,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

53 )

54 from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator

55

57 - self.pynccl_comm: Optional[PyNcclCommunicator] = None

56 + self.pynccl_comm: PyNcclCommunicator | None = None

57 if self.world_size > 1:

58 self.pynccl_comm = PyNcclCommunicator(

59 group=self.cpu_group,

@@ -63,9 +62,9 @@ class CudaCommunicator(DeviceCommunicatorBase):

62 if is_symmetric_memory_enabled():

63 register_nccl_symmetric_ops(self.pynccl_comm)

64

66 - self.ca_comm: Optional[CustomAllreduce] = None

67 - self.qr_comm: Optional[QuickAllReduce] = None

68 - self.symm_mem_comm: Optional[SymmMemCommunicator] = None

65 + self.ca_comm: CustomAllreduce | None = None

66 + self.qr_comm: QuickAllReduce | None = None

67 + self.symm_mem_comm: SymmMemCommunicator | None = None

68 if use_torch_symm_mem and current_platform.is_cuda():

69 self.symm_mem_comm = SymmMemCommunicator(

70 group=self.cpu_group,

@@ -201,7 +200,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

200 return output.movedim(0, dim).contiguous()

201

202 def reduce_scatterv(

204 - self, input_: torch.Tensor, dim: int = -1, sizes: Optional[list[int]] = None

203 + self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None

204 ):

205 world_size = self.world_size

206 pynccl_comm = self.pynccl_comm

@@ -235,7 +234,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

234 # Reshape before returning

235 return output.movedim(0, dim).contiguous()

236

238 - def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:

237 + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:

238 """Sends a tensor to the destination rank in a blocking way"""

239 """NOTE: `dst` is the local rank of the destination rank."""

240 if dst is None:

@@ -248,7 +247,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

247 torch.distributed.send(tensor, self.ranks[dst], self.device_group)

248

249 def recv(

251 - self, size: torch.Size, dtype: torch.dtype, src: Optional[int] = None

250 + self, size: torch.Size, dtype: torch.dtype, src: int | None = None

251 ) -> torch.Tensor:

252 """Receives a tensor from the source rank."""

253 """NOTE: `src` is the local rank of the source rank."""

@@ -274,9 +273,9 @@ class CudaCommunicator(DeviceCommunicatorBase):

273

274 def all_gatherv(

275 self,

277 - input_: Union[torch.Tensor, list[torch.Tensor]],

276 + input_: torch.Tensor | list[torch.Tensor],

277 dim: int = 0,

279 - sizes: Optional[list[int]] = None,

278 + sizes: list[int] | None = None,

279 ):

280 if dim != 0:

281 raise NotImplementedError("only dim 0 all-gatherv is supported")

@@ -289,7 +288,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

288 if sizes is not None and all(s == sizes[0] for s in sizes):

289 sizes = None

290

292 - def _all_gather_single(input_: torch.Tensor, sizes: Optional[list[int]] = None):

291 + def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):

292 input_size = input_.size()

293 if sizes is not None:

294 assert len(sizes) == world_size

vllm/distributed/device_communicators/cuda_wrapper.py

+3 -3

@@ -7,7 +7,7 @@ convenient for use when we just need to call a few functions.

7

8 import ctypes

9 from dataclasses import dataclass

10 -from typing import Any, Optional

10 +from typing import Any

11

12 # this line makes it possible to directly load `libcudart.so` using `ctypes`

13 import torch # noqa

@@ -36,7 +36,7 @@ class Function:

36 argtypes: list[Any]

37

38

39 -def find_loaded_library(lib_name) -> Optional[str]:

39 +def find_loaded_library(lib_name) -> str | None:

40 """

41 According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html,

42 the file `/proc/self/maps` contains the memory maps of the process, which includes the

@@ -113,7 +113,7 @@ class CudaRTLibrary:

113 # to the corresponding dictionary

114 path_to_dict_mapping: dict[str, dict[str, Any]] = {}

115

116 - def __init__(self, so_file: Optional[str] = None):

116 + def __init__(self, so_file: str | None = None):

117 if so_file is None:

118 so_file = find_loaded_library("libcudart")

119 if so_file is None:

vllm/distributed/device_communicators/custom_all_reduce.py

+6 -7

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

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

3

4 from contextlib import contextmanager

5 -from typing import Optional, Union

5

6 import torch

7 import torch.distributed as dist

@@ -55,7 +54,7 @@ class CustomAllreduce:

54 def __init__(

55 self,

56 group: ProcessGroup,

58 - device: Union[int, str, torch.device],

57 + device: int | str | torch.device,

58 max_size=8192 * 1024,

59 symm_mem_enabled=False,

60 ) -> None:

@@ -260,7 +259,7 @@ class CustomAllreduce:

259 )

260 return out

261

263 - def custom_all_reduce(self, input: torch.Tensor) -> Optional[torch.Tensor]:

262 + def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor | None:

263 """The main allreduce API that provides support for cuda graph."""

264 # When custom allreduce is disabled, this will be None.

265 if self.disabled or not self.should_custom_ar(input):

@@ -292,8 +291,8 @@ class CustomAllreduce:

291 @staticmethod

292 def create_shared_buffer(

293 size_in_bytes: int,

295 - group: Optional[ProcessGroup] = None,

296 - uncached: Optional[bool] = False,

294 + group: ProcessGroup | None = None,

295 + uncached: bool | None = False,

296 ) -> list[int]:

297 pointer, handle = ops.allocate_shared_buffer_and_handle(size_in_bytes)

298

@@ -313,8 +312,8 @@ class CustomAllreduce:

312 @staticmethod

313 def free_shared_buffer(

314 pointers: list[int],

316 - group: Optional[ProcessGroup] = None,

317 - rank: Optional[int] = None,

315 + group: ProcessGroup | None = None,

316 + rank: int | None = None,

317 ) -> None:

318 if rank is None:

319 rank = dist.get_rank(group=group)

vllm/distributed/device_communicators/pynccl.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 # ===================== import region =====================

6 import torch

@@ -59,9 +58,9 @@ def register_nccl_symmetric_ops(pynccl_comm):

58 class PyNcclCommunicator:

59 def __init__(

60 self,

62 - group: Union[ProcessGroup, StatelessProcessGroup],

63 - device: Union[int, str, torch.device],

64 - library_path: Optional[str] = None,

61 + group: ProcessGroup | StatelessProcessGroup,

62 + device: int | str | torch.device,

63 + library_path: str | None = None,

64 ):

65 """

66 Args:

vllm/distributed/device_communicators/pynccl_allocator.py

+2 -2

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

3 import atexit

4 import contextlib

5 import tempfile

6 -from typing import Any, Optional

6 +from typing import Any

7

8 import torch

9 from packaging import version

@@ -141,7 +141,7 @@ class nccl_symm_mem_context:

141 or version.parse(torch.__version__) < version.parse("2.8.0.a0")

142 )

143 if self.disabled:

144 - self.pynccl_comm: Optional[PyNcclCommunicator] = None

144 + self.pynccl_comm: PyNcclCommunicator | None = None

145 self._mem_pool_ctx: contextlib.AbstractContextManager[Any] = (

146 contextlib.nullcontext()

147 )

vllm/distributed/device_communicators/pynccl_wrapper.py

+2 -2

@@ -25,7 +25,7 @@

25 import ctypes

26 import platform

27 from dataclasses import dataclass

28 -from typing import Any, Optional

28 +from typing import Any

29

30 import torch

31 from torch.distributed import ReduceOp

@@ -305,7 +305,7 @@ class NCCLLibrary:

305 # to the corresponding dictionary

306 path_to_dict_mapping: dict[str, dict[str, Any]] = {}

307

308 - def __init__(self, so_file: Optional[str] = None):

308 + def __init__(self, so_file: str | None = None):

309 so_file = so_file or find_nccl_library()

310

311 try:

vllm/distributed/device_communicators/quick_all_reduce.py

+1 -4

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

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

3

4 from enum import Enum

5 -from typing import Union

5

6 import torch

7 import torch.distributed as dist

@@ -58,9 +57,7 @@ class QuickAllReduce:

57 (torch.bfloat16, 8): [16 * MB, 2048 * MB, 2048 * MB, 2048 * MB],

58 }

59

61 - def __init__(

62 - self, group: ProcessGroup, device: Union[int, str, torch.device]

63 - ) -> None:

60 + def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> None:

61 """

62 Custom allreduce provides non-destructive acceleration and is

63 available for CUDA and ROCm MI300 series.

vllm/distributed/device_communicators/ray_communicator.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import uuid

4 -from typing import Any, Optional

4 +from typing import Any

5

6 import ray

7 import torch

@@ -27,15 +27,15 @@ class RayPPCommunicator(Communicator):

27 This class is not thread-safe.

28 """

29

30 - _comm: Optional[DeviceCommunicatorBase]

30 + _comm: DeviceCommunicatorBase | None

31

32 def __init__(

33 self,

34 world_size: int,

35 comm_id: Any,

36 - rank: Optional[int],

36 + rank: int | None,

37 actor_handles: list["ray.actor.ActorHandle"],

38 - cuda_stream: Optional[torch.cuda.Stream],

38 + cuda_stream: torch.cuda.Stream | None,

39 use_communication_streams: bool = False,

40 ):

41 """

@@ -56,7 +56,7 @@ class RayPPCommunicator(Communicator):

56 This is not supported.

57 """

58 self._world_size = world_size

59 - self._rank: Optional[int] = None

59 + self._rank: int | None = None

60 self._actor_handles = actor_handles

61 if use_communication_streams:

62 raise NotImplementedError("use_communication_streams is not supported")

@@ -143,7 +143,7 @@ class RayPPCommunicator(Communicator):

143 else:

144 raise ValueError(f"Actor {actor} not found in communicator group")

145

146 - def get_self_rank(self) -> Optional[int]:

146 + def get_self_rank(self) -> int | None:

147 """

148 Return this actor's rank.

149 """

vllm/distributed/device_communicators/shm_broadcast.py

+15 -15

@@ -7,7 +7,7 @@ from contextlib import contextmanager

7 from dataclasses import dataclass, field

8 from multiprocessing import shared_memory

9 from threading import Event

10 -from typing import Any, Optional, Union

10 +from typing import Any

11 from unittest.mock import patch

12

13 import torch

@@ -80,7 +80,7 @@ class ShmRingBuffer:

80 n_reader: int,

81 max_chunk_bytes: int,

82 max_chunks: int,

83 - name: Optional[str] = None,

83 + name: str | None = None,

84 ):

85 """

86 A shared memory ring buffer implementation for broadcast communication.

@@ -213,9 +213,9 @@ class ShmRingBuffer:

213 class Handle:

214 local_reader_ranks: list[int] = field(default_factory=list)

215

216 - buffer_handle: Optional[tuple[int, int, int, str]] = None

217 - local_subscribe_addr: Optional[str] = None

218 - remote_subscribe_addr: Optional[str] = None

216 + buffer_handle: tuple[int, int, int, str] | None = None

217 + local_subscribe_addr: str | None = None

218 + remote_subscribe_addr: str | None = None

219 remote_addr_ipv6: bool = False

220

221

@@ -224,10 +224,10 @@ class MessageQueue:

224 self,

225 n_reader, # number of all readers

226 n_local_reader, # number of local readers through shared memory

227 - local_reader_ranks: Optional[list[int]] = None,

227 + local_reader_ranks: list[int] | None = None,

228 max_chunk_bytes: int = 1024 * 1024 * 10,

229 max_chunks: int = 10,

230 - connect_ip: Optional[str] = None,

230 + connect_ip: str | None = None,

231 ):

232 if local_reader_ranks is None:

233 local_reader_ranks = list(range(n_local_reader))

@@ -384,7 +384,7 @@ class MessageQueue:

384 assert recv == b"READY"

385

386 @contextmanager

387 - def acquire_write(self, timeout: Optional[float] = None):

387 + def acquire_write(self, timeout: float | None = None):

388 assert self._is_writer, "Only writers can acquire write"

389 start_time = time.monotonic()

390 n_warning = 1

@@ -444,8 +444,8 @@ class MessageQueue:

444 @contextmanager

445 def acquire_read(

446 self,

447 - timeout: Optional[float] = None,

448 - cancel: Optional[Event] = None,

447 + timeout: float | None = None,

448 + cancel: Event | None = None,

449 indefinite: bool = False,

450 ):

451 assert self._is_local_reader, "Only readers can acquire read"

@@ -502,7 +502,7 @@ class MessageQueue:

502 self._read_spin_timer.record_activity()

503 break

504

505 - def enqueue(self, obj, timeout: Optional[float] = None):

505 + def enqueue(self, obj, timeout: float | None = None):

506 """Write to message queue with optional timeout (in seconds)"""

507 assert self._is_writer, "Only writers can enqueue"

508 serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)

@@ -520,8 +520,8 @@ class MessageQueue:

520

521 def dequeue(

522 self,

523 - timeout: Optional[float] = None,

524 - cancel: Optional[Event] = None,

523 + timeout: float | None = None,

524 + cancel: Event | None = None,

525 indefinite: bool = False,

526 ):

527 """Read from message queue with optional timeout (in seconds)"""

@@ -542,7 +542,7 @@ class MessageQueue:

542 return obj

543

544 @staticmethod

545 - def recv(socket: zmq.Socket, timeout: Optional[float]) -> Any:

545 + def recv(socket: zmq.Socket, timeout: float | None) -> Any:

546 timeout_ms = None if timeout is None else int(timeout * 1000)

547 if not socket.poll(timeout=timeout_ms):

548 raise TimeoutError

@@ -558,7 +558,7 @@ class MessageQueue:

558

559 @staticmethod

560 def create_from_process_group(

561 - pg: Union[ProcessGroup, StatelessProcessGroup],

561 + pg: ProcessGroup | StatelessProcessGroup,

562 max_chunk_bytes,

563 max_chunks,

564 writer_rank=0,

vllm/distributed/device_communicators/shm_object_storage.py

+8 -10

@@ -3,13 +3,13 @@

3

4 import pickle

5 from abc import ABC, abstractmethod

6 -from collections.abc import Iterable

6 +from collections.abc import Callable, Iterable

7 from contextlib import contextmanager

8 from dataclasses import dataclass

9 from itertools import chain

10 from multiprocessing import shared_memory

11 from multiprocessing.synchronize import Lock as LockType

12 -from typing import Any, Callable, Optional, Union

12 +from typing import Any

13 from unittest.mock import patch

14

15 import torch

@@ -109,7 +109,7 @@ class SingleWriterShmRingBuffer:

109 def __init__(

110 self,

111 data_buffer_size: int,

112 - name: Optional[str] = None,

112 + name: str | None = None,

113 create: bool = False,

114 ):

115 self.data_buffer_size = data_buffer_size

@@ -252,7 +252,7 @@ class SingleWriterShmRingBuffer:

252 def free_buf(

253 self,

254 is_free_fn: Callable[[int, memoryview], bool],

255 - nbytes: Optional[int] = None,

255 + nbytes: int | None = None,

256 ) -> Iterable[int]:

257 """

258 Free a buffer of the given size. This is a no-op in shared memory,

@@ -340,9 +340,7 @@ class MsgpackSerde(ObjectSerde):

340 self.mm_decoder = MsgpackDecoder(MultiModalKwargsItem)

341 self._mm_kwargs_item_cls = MultiModalKwargsItem

342

343 - def serialize(

344 - self, value: Any

345 - ) -> tuple[Union[bytes, list[bytes]], int, bytes, int]:

343 + def serialize(self, value: Any) -> tuple[bytes | list[bytes], int, bytes, int]:

344 len_arr = None

345 if isinstance(value, (torch.Tensor, self._mm_kwargs_item_cls)):

346 type_name = type(value).__name__

@@ -396,7 +394,7 @@ class ShmObjectStorageHandle:

394 n_readers: int

395 ring_buffer_handle: tuple[int, str]

396 serde_class: type[ObjectSerde]

399 - reader_lock: Optional[LockType]

397 + reader_lock: LockType | None

398

399

400 class SingleWriterShmObjectStorage:

@@ -444,7 +442,7 @@ class SingleWriterShmObjectStorage:

442 n_readers: int,

443 ring_buffer: SingleWriterShmRingBuffer,

444 serde_class: type[ObjectSerde] = MsgpackSerde,

447 - reader_lock: Optional[LockType] = None,

445 + reader_lock: LockType | None = None,

446 ):

447 """

448 Initialize the object storage.

@@ -492,7 +490,7 @@ class SingleWriterShmObjectStorage:

490

491 def copy_to_buffer(

492 self,

495 - data: Union[bytes, list[bytes]],

493 + data: bytes | list[bytes],

494 data_bytes: int,

495 metadata: bytes,

496 md_bytes: int,

vllm/distributed/device_communicators/symm_mem.py

+5 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional, Union

3

4 import torch

5 import torch.distributed as dist

@@ -31,10 +30,10 @@ class SymmMemCommunicator:

30 def __init__(

31 self,

32 group: ProcessGroup,

34 - device: Union[int, str, torch.device],

33 + device: int | str | torch.device,

34 # add options for testing

36 - force_multimem: Optional[bool] = None,

37 - max_size_override: Optional[int] = None,

35 + force_multimem: bool | None = None,

36 + max_size_override: int | None = None,

37 ):

38 self.disabled = True

39

@@ -108,8 +107,8 @@ class SymmMemCommunicator:

107 return inp_size < self.max_size

108

109 def all_reduce(

111 - self, inp: torch.Tensor, *, out: Optional[torch.Tensor] = None

112 - ) -> Optional[torch.Tensor]:

110 + self, inp: torch.Tensor, *, out: torch.Tensor | None = None

111 + ) -> torch.Tensor | None:

112 if not self.should_use_symm_mem(inp):

113 return None

114 if out is None:

vllm/distributed/device_communicators/tpu_communicator.py

+2 -3

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

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

3

4 import os

5 -from typing import Optional

5

6 import torch

7 from torch.distributed import ProcessGroup

@@ -39,8 +38,8 @@ class TpuCommunicator(DeviceCommunicatorBase):

38 def __init__(

39 self,

40 cpu_group: ProcessGroup,

42 - device: Optional[torch.device] = None,

43 - device_group: Optional[ProcessGroup] = None,

41 + device: torch.device | None = None,

42 + device_group: ProcessGroup | None = None,

43 unique_name: str = "",

44 ):

45 super().__init__(cpu_group, device, device_group, unique_name)

vllm/distributed/device_communicators/xpu_communicator.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6 import torch.distributed as dist

@@ -19,8 +18,8 @@ class XpuCommunicator(DeviceCommunicatorBase):

18 def __init__(

19 self,

20 cpu_group: ProcessGroup,

22 - device: Optional[torch.device] = None,

23 - device_group: Optional[ProcessGroup] = None,

21 + device: torch.device | None = None,

22 + device_group: ProcessGroup | None = None,

23 unique_name: str = "",

24 ):

25 super().__init__(cpu_group, device, device_group, unique_name)

@@ -45,7 +44,7 @@ class XpuCommunicator(DeviceCommunicatorBase):

44

45 def gather(

46 self, input_: torch.Tensor, dst: int = 0, dim: int = -1

48 - ) -> Optional[torch.Tensor]:

47 + ) -> torch.Tensor | None:

48 assert -input_.dim() <= dim < input_.dim(), (

49 f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"

50 )

vllm/distributed/eplb/eplb_state.py

+7 -8

@@ -29,7 +29,6 @@ physical experts.

29 import time

30 from collections.abc import Sequence

31 from dataclasses import dataclass

32 -from typing import Optional, Union

32

33 import torch

34 from torch.distributed import ProcessGroup, all_reduce

@@ -186,9 +185,9 @@ class EplbState:

185 model: MixtureOfExperts,

186 device: torch.device,

187 parallel_config: ParallelConfig,

189 - global_expert_load: Optional[torch.Tensor] = None,

190 - old_global_expert_indices: Optional[torch.Tensor] = None,

191 - rank_mapping: Optional[dict[int, int]] = None,

188 + global_expert_load: torch.Tensor | None = None,

189 + old_global_expert_indices: torch.Tensor | None = None,

190 + rank_mapping: dict[int, int] | None = None,

191 ) -> "EplbState":

192 """

193 Build the initial EPLB state.

@@ -439,9 +438,9 @@ class EplbState:

438 model: MixtureOfExperts,

439 is_profile: bool = False,

440 execute_shuffle: bool = True,

442 - global_expert_load: Optional[torch.Tensor] = None,

443 - rank_mapping: Optional[dict[int, int]] = None,

444 - ) -> Optional[torch.Tensor]:

441 + global_expert_load: torch.Tensor | None = None,

442 + rank_mapping: dict[int, int] | None = None,

443 + ) -> torch.Tensor | None:

444 """

445 Rearrange the experts according to the current load.

446 """

@@ -611,7 +610,7 @@ class EplbState:

610

611

612 def _node_count_with_rank_mapping(

614 - pg: Union[ProcessGroup, StatelessProcessGroup],

613 + pg: ProcessGroup | StatelessProcessGroup,

614 rank_mapping: dict[int, int],

615 ) -> int:

616 if isinstance(pg, ProcessGroup):

vllm/distributed/eplb/rebalance_execute.py

+1 -2

@@ -8,7 +8,6 @@ This involves the exchange of expert weights between GPUs.

8

9 from collections.abc import Iterable, MutableSequence, Sequence

10 from functools import partial

11 -from typing import Optional

11

12 import torch

13 from torch.distributed import (

@@ -253,7 +252,7 @@ def rearrange_expert_weights_inplace(

252 expert_weights: Sequence[Iterable[torch.Tensor]],

253 ep_group: ProcessGroup,

254 is_profile: bool = False,

256 - rank_mapping: Optional[dict[int, int]] = None,

255 + rank_mapping: dict[int, int] | None = None,

256 ) -> None:

257 """

258 Rearranges the expert weights in place according to the new expert indices.

vllm/distributed/kv_events.py

+15 -14

@@ -6,10 +6,11 @@ import threading

6 import time

7 from abc import ABC, abstractmethod

8 from collections import deque

9 +from collections.abc import Callable

10 from dataclasses import asdict

11 from itertools import count

12 from queue import Queue

12 -from typing import Any, Callable, Optional, Union

13 +from typing import Any

14

15 import msgspec

16 import zmq

@@ -29,7 +30,7 @@ class EventBatch(

30 ):

31 ts: float

32 events: list[Any]

32 - data_parallel_rank: Optional[int] = None

33 + data_parallel_rank: int | None = None

34

35

36 class KVCacheEvent(

@@ -47,16 +48,16 @@ MEDIUM_GPU = "GPU"

48

49 class BlockStored(KVCacheEvent):

50 block_hashes: list[ExternalBlockHash]

50 - parent_block_hash: Optional[ExternalBlockHash]

51 + parent_block_hash: ExternalBlockHash | None

52 token_ids: list[int]

53 block_size: int

53 - lora_id: Optional[int]

54 - medium: Optional[str]

54 + lora_id: int | None

55 + medium: str | None

56

57

58 class BlockRemoved(KVCacheEvent):

59 block_hashes: list[ExternalBlockHash]

59 - medium: Optional[str]

60 + medium: str | None

61

62

63 class AllBlocksCleared(KVCacheEvent):

@@ -64,7 +65,7 @@ class AllBlocksCleared(KVCacheEvent):

65

66

67 class KVEventBatch(EventBatch):

67 - events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]

68 + events: list[BlockStored | BlockRemoved | AllBlocksCleared]

69

70

71 class EventPublisher(ABC):

@@ -139,7 +140,7 @@ class ZmqEventPublisher(EventPublisher):

140 self,

141 data_parallel_rank: int,

142 endpoint: str = "tcp://*:5557",

142 - replay_endpoint: Optional[str] = None,

143 + replay_endpoint: str | None = None,

144 buffer_steps: int = 10_000,

145 hwm: int = 100_000,

146 max_queue_size: int = 100_000,

@@ -147,13 +148,13 @@ class ZmqEventPublisher(EventPublisher):

148 ) -> None:

149 # Storage

150 super().__init__(data_parallel_rank)

150 - self._event_queue = Queue[Optional[EventBatch]](maxsize=max_queue_size)

151 + self._event_queue = Queue[EventBatch | None](maxsize=max_queue_size)

152 self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps)

153

154 # ZMQ sockets

155 self._ctx = zmq.Context.instance()

155 - self._pub: Optional[zmq.Socket] = None

156 - self._replay: Optional[zmq.Socket] = None

156 + self._pub: zmq.Socket | None = None

157 + self._replay: zmq.Socket | None = None

158 self._dp_rank = data_parallel_rank

159

160 self._endpoint = self.offset_endpoint_port(endpoint, self._dp_rank)

@@ -303,8 +304,8 @@ class ZmqEventPublisher(EventPublisher):

304

305 @staticmethod

306 def offset_endpoint_port(

306 - endpoint: Optional[str], data_parallel_rank: int

307 - ) -> Optional[str]:

307 + endpoint: str | None, data_parallel_rank: int

308 + ) -> str | None:

309 """Helper function to offset the port in an endpoint by

310 the data parallel rank.

311

@@ -349,7 +350,7 @@ class EventPublisherFactory:

350

351 @classmethod

352 def create(

352 - cls, config: Optional[KVEventsConfig], data_parallel_rank: int = 0

353 + cls, config: KVEventsConfig | None, data_parallel_rank: int = 0

354 ) -> EventPublisher:

355 """Create publisher from a config mapping."""

356 if not config:

vllm/distributed/kv_transfer/kv_connector/factory.py

+2 -1

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

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

3

4 import importlib

5 -from typing import TYPE_CHECKING, Callable

5 +from collections.abc import Callable

6 +from typing import TYPE_CHECKING

7

8 import vllm.envs as envs

9 from vllm.distributed.kv_transfer.kv_connector.base import (

vllm/distributed/kv_transfer/kv_connector/utils.py

+5 -5

@@ -7,7 +7,7 @@ KV cache helper for store.

7 from collections import defaultdict

8 from collections.abc import Sequence

9 from concurrent.futures import CancelledError, Future

10 -from typing import Literal, Optional, Union, cast

10 +from typing import Literal, cast

11

12 import torch

13

@@ -136,7 +136,7 @@ class KVOutputAggregator:

136 # Aggregate kv_connector_output from all workers

137

138 def update_finished_set(

139 - req_ids: Optional[set[str]],

139 + req_ids: set[str] | None,

140 remaining_count_dict: dict[str, int],

141 finished_set: set[str],

142 ) -> None:

@@ -197,7 +197,7 @@ class KVOutputAggregator:

197 to the respective list of outputs."""

198 result_future: Future[ModelRunnerOutput] = Future()

199

200 - outputs: list[Optional[ModelRunnerOutput]] = [None] * len(output_futures)

200 + outputs: list[ModelRunnerOutput | None] = [None] * len(output_futures)

201

202 def make_callback(idx):

203 def callback(fut):

@@ -230,8 +230,8 @@ class KVOutputAggregator:

230 def _make_src_and_dst_indices(

231 src_block_ids: list[int],

232 dst_block_ids: list[int],

233 - src_device: Union[torch.device, str],

234 - dst_device: Union[torch.device, str],

233 + src_device: torch.device | str,

234 + dst_device: torch.device | str,

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

236 src_indices = torch.tensor(src_block_ids, device=src_device, dtype=torch.int64)

237 dst_indices = torch.tensor(dst_block_ids, device=dst_device, dtype=torch.int64)

vllm/distributed/kv_transfer/kv_connector/v1/base.py

+9 -9

@@ -37,8 +37,8 @@ The class provides the following primitives:

37

38 import enum

39 from abc import ABC, abstractmethod

40 -from collections.abc import Iterable

41 -from typing import TYPE_CHECKING, Any, Callable, Literal, Optional

40 +from collections.abc import Callable, Iterable

41 +from typing import TYPE_CHECKING, Any, Literal, Optional

42

43 import torch

44

@@ -93,7 +93,7 @@ class KVConnectorBase_V1(ABC):

93 "Initializing KVConnectorBase_V1. This API is experimental and "

94 "subject to change in the future as we iterate the design."

95 )

96 - self._connector_metadata: Optional[KVConnectorMetadata] = None

96 + self._connector_metadata: KVConnectorMetadata | None = None

97 self._vllm_config = vllm_config

98 self._role = role

99

@@ -222,7 +222,7 @@ class KVConnectorBase_V1(ABC):

222

223 def get_finished(

224 self, finished_req_ids: set[str]

225 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:

225 + ) -> tuple[set[str] | None, set[str] | None]:

226 """

227 Notifies worker-side connector ids of requests that have

228 finished generating tokens on the worker.

@@ -281,7 +281,7 @@ class KVConnectorBase_V1(ABC):

281 self,

282 request: "Request",

283 num_computed_tokens: int,

284 - ) -> tuple[Optional[int], bool]:

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

285 """

286 Get number of new tokens that can be loaded from the

287 external KV cache beyond the num_computed_tokens.

@@ -361,7 +361,7 @@ class KVConnectorBase_V1(ABC):

361 self,

362 request: "Request",

363 block_ids: list[int],

364 - ) -> tuple[bool, Optional[dict[str, Any]]]:

364 + ) -> tuple[bool, dict[str, Any] | None]:

365 """

366 Called exactly once when a request has finished, before its blocks are

367 freed.

@@ -388,7 +388,7 @@ class KVConnectorBase_V1(ABC):

388 return ()

389

390 @classmethod

391 - def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> Optional[str]:

391 + def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:

392 """

393 Get the required KV cache layout for this connector.

394 Args:

@@ -406,7 +406,7 @@ class KVConnectorBase_V1(ABC):

406 )

407 return None

408

409 - def get_finished_count(self) -> Optional[int]:

409 + def get_finished_count(self) -> int | None:

410 """

411 Get the count of requests expected to complete send/receive operations

412 via this connector.

@@ -419,7 +419,7 @@ class KVConnectorBase_V1(ABC):

419

420 @classmethod

421 def build_kv_connector_stats(

422 - cls, data: Optional[dict[str, Any]] = None

422 + cls, data: dict[str, Any] | None = None

423 ) -> Optional["KVConnectorStats"]:

424 """

425 KVConnectorStats resolution method. This method allows dynamically

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import TYPE_CHECKING, Any, Optional

3 +from typing import TYPE_CHECKING, Any

4

5 import torch

6 from lmcache.integration.vllm.vllm_v1_adapter import LMCacheConnectorV1Impl

@@ -96,7 +96,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):

96

97 def get_finished(

98 self, finished_req_ids: set[str]

99 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:

99 + ) -> tuple[set[str] | None, set[str] | None]:

100 """

101 Notifies worker-side connector ids of requests that have

102 finished generating tokens.

@@ -117,7 +117,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):

117 self,

118 request: "Request",

119 num_computed_tokens: int,

120 - ) -> tuple[Optional[int], bool]:

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

121 """

122 Get number of new tokens that can be loaded from the

123 external KV cache beyond the num_computed_tokens.

@@ -161,7 +161,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):

161 self,

162 request: "Request",

163 block_ids: list[int],

164 - ) -> tuple[bool, Optional[dict[str, Any]]]:

164 + ) -> tuple[bool, dict[str, Any] | None]:

165 """

166 Called when a request has finished, before its blocks are freed.

167

vllm/distributed/kv_transfer/kv_connector/v1/metrics.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass, field

4 -from typing import Any, Optional, Union

4 +from typing import Any

5

6 from vllm.config.kv_transfer import KVTransferConfig

7 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory

@@ -32,7 +32,7 @@ class KVConnectorStats:

32 """

33 raise NotImplementedError

34

35 - def reduce(self) -> dict[str, Union[int, float]]:

35 + def reduce(self) -> dict[str, int | float]:

36 """

37 Reduce the observations collected during a time interval to one or

38 more representative values (eg avg/median/sum of the series).

@@ -58,7 +58,7 @@ class KVConnectorLogging:

58 self.reset()

59

60 def reset(self):

61 - self.transfer_stats_accumulator: Optional[KVConnectorStats] = None

61 + self.transfer_stats_accumulator: KVConnectorStats | None = None

62

63 def observe(self, transfer_stats_data: dict[str, Any]):

64 # Should not be called when a KVConnector is not configured.

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

+11 -11

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

3 import copy

4 from collections.abc import Iterable

5 from dataclasses import dataclass

6 -from typing import TYPE_CHECKING, Any, Optional

6 +from typing import TYPE_CHECKING, Any

7

8 import torch

9

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

33 @dataclass

34 class MultiKVConnectorMetadata(KVConnectorMetadata):

35 metadata: tuple[KVConnectorMetadata, ...]

36 - extra_async_saves: Optional[dict[str, int]] = None

36 + extra_async_saves: dict[str, int] | None = None

37

38

39 @dataclass

@@ -130,7 +130,7 @@ class MultiConnector(KVConnectorBase_V1):

130 c.clear_connector_metadata()

131

132 def shutdown(self):

133 - exception: Optional[Exception] = None

133 + exception: Exception | None = None

134 for c in self._connectors:

135 try:

136 c.shutdown()

@@ -169,7 +169,7 @@ class MultiConnector(KVConnectorBase_V1):

169

170 def get_finished(

171 self, finished_req_ids: set[str]

172 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:

172 + ) -> tuple[set[str] | None, set[str] | None]:

173 finished_sending: set[str] = set()

174 finished_recving: set[str] = set()

175 for c in self._connectors:

@@ -207,7 +207,7 @@ class MultiConnector(KVConnectorBase_V1):

207 self,

208 request: "Request",

209 num_computed_tokens: int,

210 - ) -> tuple[Optional[int], bool]:

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

211 to_return = (0, False)

212 for i, c in enumerate(self._connectors):

213 toks, load_async = c.get_num_new_matched_tokens(

@@ -258,7 +258,7 @@ class MultiConnector(KVConnectorBase_V1):

258 self,

259 request: "Request",

260 blocks: list[int],

261 - ) -> tuple[bool, Optional[dict[str, Any]]]:

261 + ) -> tuple[bool, dict[str, Any] | None]:

262 async_saves = 0

263 kv_txfer_params = None

264 for c in self._connectors:

@@ -286,7 +286,7 @@ class MultiConnector(KVConnectorBase_V1):

286 yield from c.take_events()

287

288 @classmethod

289 - def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> Optional[str]:

289 + def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:

290 """

291 Get the required KV cache layout for this connector.

292 Args:

@@ -323,17 +323,17 @@ class MultiConnector(KVConnectorBase_V1):

323

324 @classmethod

325 def build_kv_connector_stats(

326 - cls, data: Optional[dict[str, Any]] = None

327 - ) -> Optional[KVConnectorStats]:

326 + cls, data: dict[str, Any] | None = None

327 + ) -> KVConnectorStats | None:

328 return (

329 MultiKVConnectorStats(data=data)

330 if data is not None

331 else MultiKVConnectorStats()

332 )

333

334 - def get_kv_connector_stats(self) -> Optional[MultiKVConnectorStats]:

334 + def get_kv_connector_stats(self) -> MultiKVConnectorStats | None:

335 # Group connector stats by connector type.

336 - stats_by_connector: Optional[MultiKVConnectorStats] = None

336 + stats_by_connector: MultiKVConnectorStats | None = None

337 for c in self._connectors:

338 stats = c.get_kv_connector_stats()

339 if stats is None:

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+16 -16

@@ -13,7 +13,7 @@ from collections import defaultdict

13 from collections.abc import Iterator

14 from concurrent.futures import Future, ThreadPoolExecutor

15 from dataclasses import dataclass

16 -from typing import TYPE_CHECKING, Any, Optional, Union

16 +from typing import TYPE_CHECKING, Any

17

18 import msgspec

19 import numpy as np

@@ -153,10 +153,10 @@ class NixlConnector(KVConnectorBase_V1):

153 self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id

154

155 if role == KVConnectorRole.SCHEDULER:

156 - self.connector_scheduler: Optional[NixlConnectorScheduler] = (

156 + self.connector_scheduler: NixlConnectorScheduler | None = (

157 NixlConnectorScheduler(vllm_config, self.engine_id)

158 )

159 - self.connector_worker: Optional[NixlConnectorWorker] = None

159 + self.connector_worker: NixlConnectorWorker | None = None

160 elif role == KVConnectorRole.WORKER:

161 self.connector_scheduler = None

162 self.connector_worker = NixlConnectorWorker(vllm_config, self.engine_id)

@@ -189,7 +189,7 @@ class NixlConnector(KVConnectorBase_V1):

189

190 def get_num_new_matched_tokens(

191 self, request: "Request", num_computed_tokens: int

192 - ) -> tuple[Optional[int], bool]:

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

193 assert self.connector_scheduler is not None

194 return self.connector_scheduler.get_num_new_matched_tokens(

195 request, num_computed_tokens

@@ -214,7 +214,7 @@ class NixlConnector(KVConnectorBase_V1):

214 self,

215 request: "Request",

216 block_ids: list[int],

217 - ) -> tuple[bool, Optional[dict[str, Any]]]:

217 + ) -> tuple[bool, dict[str, Any] | None]:

218 assert self.connector_scheduler is not None

219 return self.connector_scheduler.request_finished(request, block_ids)

220

@@ -234,14 +234,14 @@ class NixlConnector(KVConnectorBase_V1):

234 assert self.connector_worker is not None

235 return self.connector_worker.get_finished()

236

237 - def get_kv_connector_stats(self) -> Optional[KVConnectorStats]:

237 + def get_kv_connector_stats(self) -> KVConnectorStats | None:

238 assert self.connector_worker is not None

239 return self.connector_worker.get_kv_connector_stats()

240

241 @classmethod

242 def build_kv_connector_stats(

243 - cls, data: Optional[dict[str, Any]] = None

244 - ) -> Optional[KVConnectorStats]:

243 + cls, data: dict[str, Any] | None = None

244 + ) -> KVConnectorStats | None:

245 return (

246 NixlKVConnectorStats(data=data)

247 if data is not None

@@ -445,7 +445,7 @@ class NixlConnectorScheduler:

445 self,

446 request: "Request",

447 block_ids: list[int],

448 - ) -> tuple[bool, Optional[dict[str, Any]]]:

448 + ) -> tuple[bool, dict[str, Any] | None]:

449 """

450 Once a request is finished, determine whether request blocks

451 should be freed now or will be sent asynchronously and freed later.

@@ -584,7 +584,7 @@ class NixlConnectorWorker:

584 )

585

586 # Note: host xfer buffer ops when use_host_buffer is True

587 - self.copy_blocks: Optional[CopyBlocksOp] = None

587 + self.copy_blocks: CopyBlocksOp | None = None

588

589 # Map of engine_id -> kv_caches_base_addr. For TP case, each local

590 # rank will still only pull from a single remote TP worker.

@@ -615,7 +615,7 @@ class NixlConnectorWorker:

615 self._reqs_to_process: set[ReqId] = set()

616

617 # Background thread for handling new handshake requests.

618 - self._nixl_handshake_listener_t: Optional[threading.Thread] = None

618 + self._nixl_handshake_listener_t: threading.Thread | None = None

619 # Background thread for initializing new NIXL handshakes.

620 self._handshake_initiation_executor = ThreadPoolExecutor(

621 # NIXL is not guaranteed to be thread-safe, limit 1 worker.

@@ -635,7 +635,7 @@ class NixlConnectorWorker:

635 # TODO(mgoin): remove this once we have hybrid memory allocator

636 # Optimization for models with local attention (Llama 4)

637 # List of block window sizes for each layer for local attention

638 - self.block_window_per_layer: list[Optional[int]] = []

638 + self.block_window_per_layer: list[int | None] = []

639 self.use_mla = self.model_config.use_mla

640

641 backend = get_attn_backend(

@@ -1472,7 +1472,7 @@ class NixlConnectorWorker:

1472 self._recving_transfers[request_id].append((handle, time.perf_counter()))

1473

1474 def _get_block_descs_ids(

1475 - self, engine_id: str, block_ids: list[int], layer_idx: Optional[int] = None

1475 + self, engine_id: str, block_ids: list[int], layer_idx: int | None = None

1476 ) -> np.ndarray:

1477 """

1478 Get the descs ids for a set of block ids.

@@ -1518,7 +1518,7 @@ class NixlConnectorWorker:

1518 block_len = self.block_len_per_layer[layer_idx]

1519 return block_len

1520

1521 - def get_kv_connector_stats(self) -> Optional[KVConnectorStats]:

1521 + def get_kv_connector_stats(self) -> KVConnectorStats | None:

1522 """

1523 Get the KV transfer stats for the connector.

1524 """

@@ -1559,7 +1559,7 @@ def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:

1559 if socket_type not in (zmq.ROUTER, zmq.REQ):

1560 raise ValueError(f"Unexpected socket type: {socket_type}")

1561

1562 - ctx: Optional[zmq.Context] = None

1562 + ctx: zmq.Context | None = None

1563 try:

1564 ctx = zmq.Context() # type: ignore[attr-defined]

1565 yield make_zmq_socket(

@@ -1611,7 +1611,7 @@ class NixlKVConnectorStats(KVConnectorStats):

1611 accumulator.extend(v)

1612 return self

1613

1614 - def reduce(self) -> dict[str, Union[int, float]]:

1614 + def reduce(self) -> dict[str, int | float]:

1615 # Compute compact representative stats suitable for CLI logging

1616 if self.is_empty():

1617 return {

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

+6 -6

@@ -4,7 +4,7 @@ from collections import defaultdict

4 from collections.abc import Iterable, Iterator

5 from dataclasses import dataclass

6 from itertools import islice

7 -from typing import Any, Optional

7 +from typing import Any

8

9 import torch

10

@@ -46,8 +46,8 @@ class OffloadingConnector(KVConnectorBase_V1):

46

47 spec = OffloadingSpecFactory.create_spec(vllm_config)

48

49 - self.connector_scheduler: Optional[OffloadingConnectorScheduler] = None

50 - self.connector_worker: Optional[OffloadingConnectorWorker] = None

49 + self.connector_scheduler: OffloadingConnectorScheduler | None = None

50 + self.connector_worker: OffloadingConnectorWorker | None = None

51 if role == KVConnectorRole.SCHEDULER:

52 self.connector_scheduler = OffloadingConnectorScheduler(spec)

53 elif role == KVConnectorRole.WORKER:

@@ -113,7 +113,7 @@ class OffloadingConnector(KVConnectorBase_V1):

113 self,

114 request: "Request",

115 block_ids: list[int],

116 - ) -> tuple[bool, Optional[dict[str, Any]]]:

116 + ) -> tuple[bool, dict[str, Any] | None]:

117 assert self.connector_scheduler is not None

118 return self.connector_scheduler.request_finished(request, block_ids)

119

@@ -148,7 +148,7 @@ class OffloadingConnectorScheduler:

148 self,

149 req: Request,

150 start_idx: int = 0,

151 - end_idx: Optional[int] = None,

151 + end_idx: int | None = None,

152 ) -> Iterable[BlockHash]:

153 return islice(

154 req.block_hashes,

@@ -354,7 +354,7 @@ class OffloadingConnectorScheduler:

354 self,

355 request: Request,

356 block_ids: list[int],

357 - ) -> tuple[bool, Optional[dict[str, Any]]]:

357 + ) -> tuple[bool, dict[str, Any] | None]:

358 """

359 Called when a request has finished, before its blocks are freed.

360

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py

+3 -3

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

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

3

4 from dataclasses import dataclass

5 -from typing import TYPE_CHECKING, Any, Optional

5 +from typing import TYPE_CHECKING, Any

6

7 import regex as re

8 import torch

@@ -304,7 +304,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

304

305 def get_finished(

306 self, finished_req_ids: set[str], **kwargs: Any

307 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:

307 + ) -> tuple[set[str] | None, set[str] | None]:

308 """

309 Notifies worker-side connector ids of requests that have

310 finished generating tokens.

@@ -466,7 +466,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

466 self,

467 request: "Request",

468 block_ids: list[int],

469 - ) -> tuple[bool, Optional[dict[str, Any]]]:

469 + ) -> tuple[bool, dict[str, Any] | None]:

470 """

471 Called when a request has finished, before its blocks are freed.

472

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py

+6 -7

@@ -5,11 +5,10 @@ import logging

5 import os

6 import threading

7 import time

8 -import typing

8 from collections import deque

9 from contextlib import contextmanager

10 from dataclasses import dataclass

12 -from typing import Any, Optional

11 +from typing import Any

12

13 import msgpack

14 import torch

@@ -77,7 +76,7 @@ class P2pNcclEngine:

76 config: KVTransferConfig,

77 hostname: str = "",

78 port_offset: int = 0,

80 - library_path: Optional[str] = None,

79 + library_path: str | None = None,

80 ) -> None:

81 self.config = config

82 self.rank = port_offset

@@ -187,7 +186,7 @@ class P2pNcclEngine:

186 self.nccl_num_channels,

187 )

188

190 - def create_connect(self, remote_address: typing.Optional[str] = None):

189 + def create_connect(self, remote_address: str | None = None):

190 assert remote_address is not None

191 if remote_address not in self.socks:

192 sock = self.context.socket(zmq.DEALER)

@@ -224,7 +223,7 @@ class P2pNcclEngine:

223 self,

224 tensor_id: str,

225 tensor: torch.Tensor,

227 - remote_address: typing.Optional[str] = None,

226 + remote_address: str | None = None,

227 ) -> bool:

228 if remote_address is None:

229 with self.recv_store_cv:

@@ -296,7 +295,7 @@ class P2pNcclEngine:

295 def recv_tensor(

296 self,

297 tensor_id: str,

299 - remote_address: typing.Optional[str] = None,

298 + remote_address: str | None = None,

299 ) -> torch.Tensor:

300 if self.send_type == "PUT" or self.send_type == "PUT_ASYNC":

301 start_time = time.time()

@@ -527,7 +526,7 @@ class P2pNcclEngine:

526

527 def get_finished(

528 self, finished_req_ids: set[str], no_compile_layers

530 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:

529 + ) -> tuple[set[str] | None, set[str] | None]:

530 """

531 Notifies worker-side connector ids of requests that have

532 finished generating tokens.

vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py

+2 -2

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

3 import hashlib

4 import os

5 from dataclasses import dataclass, field

6 -from typing import TYPE_CHECKING, Any, Optional

6 +from typing import TYPE_CHECKING, Any

7

8 import safetensors

9 import torch

@@ -249,7 +249,7 @@ class SharedStorageConnector(KVConnectorBase_V1):

249 self,

250 request: "Request",

251 num_computed_tokens: int,

252 - ) -> tuple[Optional[int], bool]:

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

253 """

254 Get number of new tokens that can be loaded from the

255 external KV cache beyond the num_computed_tokens.

vllm/distributed/kv_transfer/kv_lookup_buffer/base.py

+4 -5

@@ -14,7 +14,6 @@ These classes above are abstracted behind class `KVCacheBufferBase`.

14 """

15

16 from abc import ABC, abstractmethod

17 -from typing import Optional

17

18 import torch

19

@@ -98,8 +97,8 @@ class KVLookupBufferBase(KVCacheBufferBase):

97

98 @abstractmethod

99 def drop_select(

101 - self, input_tokens: Optional[torch.Tensor], roi: Optional[torch.Tensor]

102 - ) -> list[Optional[torch.Tensor]]:

100 + self, input_tokens: torch.Tensor | None, roi: torch.Tensor | None

101 + ) -> list[torch.Tensor | None]:

102 """Select and *drop* KV cache entries from the lookup buffer.

103

104 The functionality is similar to the following python statements

@@ -143,7 +142,7 @@ class KVStoreBufferBase(KVCacheBufferBase):

142 def put(

143 self,

144 key: str,

146 - value: Optional[torch.Tensor],

145 + value: torch.Tensor | None,

146 ) -> None:

147 """Store a key-value pair in the buffer.

148

@@ -163,7 +162,7 @@ class KVStoreBufferBase(KVCacheBufferBase):

162 def get(

163 self,

164 key: str,

166 - ) -> Optional[torch.Tensor]:

165 + ) -> torch.Tensor | None:

166 """Retrieve a value from the buffer by key.

167

168 Args:

vllm/distributed/kv_transfer/kv_lookup_buffer/mooncake_store.py

+3 -4

@@ -10,7 +10,6 @@ from this remote lookup buffer.

10 import json

11 import os

12 from dataclasses import dataclass

13 -from typing import Optional

13

14 import torch

15 from safetensors.torch import load as safetensors_load

@@ -110,7 +109,7 @@ class MooncakeStore(KVStoreBufferBase):

109 def put(

110 self,

111 key: str,

113 - value: Optional[torch.Tensor],

112 + value: torch.Tensor | None,

113 ) -> None:

114 # A message queue needs to be introduced before making it asynchronous.

115 if value is not None:

@@ -119,7 +118,7 @@ class MooncakeStore(KVStoreBufferBase):

118 def get(

119 self,

120 key: str,

122 - ) -> Optional[torch.Tensor]:

121 + ) -> torch.Tensor | None:

122 # A message queue needs to be introduced before making it asynchronous.

123 value = self._get_impl(key)

124 return value

@@ -142,7 +141,7 @@ class MooncakeStore(KVStoreBufferBase):

141 def _get_impl(

142 self,

143 key: str,

145 - ) -> Optional[torch.Tensor]:

144 + ) -> torch.Tensor | None:

145 """Get KVCache from Mooncake Store"""

146 try:

147 data = self.store.get(key)

vllm/distributed/kv_transfer/kv_lookup_buffer/simple_buffer.py

+5 -6

@@ -13,7 +13,6 @@ Key Features:

13

14 import threading

15 from collections import deque

16 -from typing import Optional, Union

16

17 import torch

18

@@ -46,7 +45,7 @@ class SimpleBuffer(KVLookupBufferBase):

45 self.buffer_cv = threading.Condition()

46 self.signal_pipe = signal_pipe

47 self.data_pipe = data_pipe

49 - self.request_handling_thread: Optional[threading.Thread] = None

48 + self.request_handling_thread: threading.Thread | None = None

49

50 self.normal_signal = torch.tensor([0], device="cpu")

51 self.end_signal = None

@@ -81,14 +80,14 @@ class SimpleBuffer(KVLookupBufferBase):

80

81 return 0

82

84 - def _send_tensor_and_dec_size(self, tensor: Optional[torch.Tensor]) -> None:

83 + def _send_tensor_and_dec_size(self, tensor: torch.Tensor | None) -> None:

84 assert tensor is not None, "Use self.data_pipe.send(None) instead"

85 self.buffer_size -= tensor.element_size() * tensor.numel()

86 if tensor.dtype == torch.bool:

87 tensor = tensor.float()

88 self.data_pipe.send_tensor(tensor)

89

91 - def _get_element_size(self, data: Optional[Union[list, torch.Tensor]]):

90 + def _get_element_size(self, data: list | torch.Tensor | None):

91 if isinstance(data, torch.Tensor):

92 return data.element_size() * data.numel()

93 if not data:

@@ -184,8 +183,8 @@ class SimpleBuffer(KVLookupBufferBase):

183 logger.debug("Closing drop_select_handler")

184

185 def drop_select(

187 - self, input_tokens: Optional[torch.Tensor], roi: Optional[torch.Tensor]

188 - ) -> list[Optional[torch.Tensor]]:

186 + self, input_tokens: torch.Tensor | None, roi: torch.Tensor | None

187 + ) -> list[torch.Tensor | None]:

188 assert self.request_handling_thread is None, (

189 "drop_select should be called by the KV cache consumer "

190 "(e.g. the decode vLLM instance)"

vllm/distributed/kv_transfer/kv_pipe/base.py

+2 -3

@@ -12,7 +12,6 @@ you can bypass this interface and directly start from `kv_lookup_buffer`.

12 """

13

14 from abc import ABC, abstractmethod

15 -from typing import Optional

15

16 import torch

17

@@ -24,7 +23,7 @@ class KVPipeBase(ABC):

23 """

24

25 @abstractmethod

27 - def send_tensor(self, tensor: Optional[torch.Tensor]) -> None:

26 + def send_tensor(self, tensor: torch.Tensor | None) -> None:

27 """Send a tensor, or None, via the pipe.

28

29 Need to support sending None -- important for error handling.

@@ -42,7 +41,7 @@ class KVPipeBase(ABC):

41 raise NotImplementedError

42

43 @abstractmethod

45 - def recv_tensor(self) -> Optional[torch.Tensor]:

44 + def recv_tensor(self) -> torch.Tensor | None:

45 """Receive a tensor (can be None) from the pipeline.

46

47 Returns:

vllm/distributed/kv_transfer/kv_pipe/mooncake_pipe.py

+6 -7

@@ -6,7 +6,6 @@ import os

6 import struct

7 from concurrent.futures import ThreadPoolExecutor

8 from dataclasses import dataclass

9 -from typing import Optional, Union

9

10 import torch

11 import zmq

@@ -26,7 +25,7 @@ NONE_INT = -150886311

25 class MooncakeTransferEngineConfig:

26 prefill_url: str

27 decode_url: str

29 - metadata_backend: Union[str, None]

28 + metadata_backend: str | None

29 metadata_server: str

30 protocol: str

31 device_name: str

@@ -143,7 +142,7 @@ class MooncakeTransferEngine:

142 metadata_server: str,

143 protocol: str,

144 device_name: str,

146 - metadata_backend: Union[str, None],

145 + metadata_backend: str | None,

146 ) -> None:

147 """Initialize the mooncake instance."""

148 if metadata_backend is None:

@@ -231,7 +230,7 @@ class MooncakePipe(KVPipeBase):

230 """MooncakeTransferEngine based Pipe implementation."""

231

232 def __init__(

234 - self, local_rank: int, config: KVTransferConfig, device: Optional[str] = None

233 + self, local_rank: int, config: KVTransferConfig, device: str | None = None

234 ):

235 """Initialize the mooncake pipe and set related parameters."""

236 self.config = config

@@ -243,7 +242,7 @@ class MooncakePipe(KVPipeBase):

242 self.device = self._select_device(device)

243

244 self.transfer_engine = MooncakeTransferEngine(self.kv_rank, self.local_rank)

246 - self.transport_thread: Optional[ThreadPoolExecutor] = None

245 + self.transport_thread: ThreadPoolExecutor | None = None

246 self.none_tensor = torch.tensor([NONE_INT], device=self.device)

247

248 def _select_device(self, device: str) -> torch.device:

@@ -267,7 +266,7 @@ class MooncakePipe(KVPipeBase):

266 data = self.transfer_engine.recv_bytes()

267 return safetensors_load(data)["tensor"].to(self.device)

268

270 - def send_tensor(self, tensor: Optional[torch.Tensor]) -> None:

269 + def send_tensor(self, tensor: torch.Tensor | None) -> None:

270 """Send tensor to the target process."""

271 if self.transport_thread is None:

272 self.transport_thread = ThreadPoolExecutor(max_workers=1)

@@ -275,7 +274,7 @@ class MooncakePipe(KVPipeBase):

274 assert len(tensor.shape) > 0

275 self.transport_thread.submit(self._send_impl, tensor)

276

278 - def recv_tensor(self) -> Optional[torch.Tensor]:

277 + def recv_tensor(self) -> torch.Tensor | None:

278 """Receive tensor from other processes."""

279 if self.transport_thread is None:

280 self.transport_thread = ThreadPoolExecutor(max_workers=1)

vllm/distributed/kv_transfer/kv_pipe/pynccl_pipe.py

+10 -10

@@ -15,8 +15,8 @@ Key Features:

15

16 import threading

17 import time

18 +from collections.abc import Callable

19 from concurrent.futures import ThreadPoolExecutor

19 -from typing import Callable, Optional

20

21 import torch

22

@@ -35,7 +35,7 @@ class BrokenPipeException(Exception):

35 super().__init__(self.message)

36

37

38 -Metadata = dict[str, Optional[torch.Tensor]]

38 +Metadata = dict[str, torch.Tensor | None]

39

40

41 class PyNcclPipe(KVPipeBase):

@@ -47,7 +47,7 @@ class PyNcclPipe(KVPipeBase):

47 self,

48 local_rank: int,

49 config: KVTransferConfig,

50 - device: Optional[str] = None,

50 + device: str | None = None,

51 port_offset: int = 0,

52 ):

53 self.config = config

@@ -77,7 +77,7 @@ class PyNcclPipe(KVPipeBase):

77 self.target_rank_for_recv = (self.kv_rank - 1) % self.kv_parallel_size

78

79 # transportation-related variables

80 - self.transport_thread: Optional[ThreadPoolExecutor] = None

80 + self.transport_thread: ThreadPoolExecutor | None = None

81 self.buffer_size = 0

82 self.buffer_size_lock = threading.Lock()

83 self.buffer_size_thresh = self.config.kv_buffer_size

@@ -115,7 +115,7 @@ class PyNcclPipe(KVPipeBase):

115 else:

116 return torch.device("cpu")

117

118 - def _make_metadata(self, tensor: Optional[torch.Tensor]) -> Metadata:

118 + def _make_metadata(self, tensor: torch.Tensor | None) -> Metadata:

119 """

120 Create the metadata as a dictionary based on the input tensor.

121

@@ -167,7 +167,7 @@ class PyNcclPipe(KVPipeBase):

167 """

168 return self.group.recv_obj(self.target_rank_for_recv)

169

170 - def _send_impl(self, tensor: Optional[torch.Tensor]) -> None:

170 + def _send_impl(self, tensor: torch.Tensor | None) -> None:

171 """

172 The actual implementation of sending the tensor and its metadata to the

173 target rank.

@@ -181,7 +181,7 @@ class PyNcclPipe(KVPipeBase):

181 if tensor is not None:

182 self.device_send_func(tensor.to(self.device), self.target_rank_for_send)

183

184 - def _recv_impl(self) -> Optional[torch.Tensor]:

184 + def _recv_impl(self) -> torch.Tensor | None:

185 """

186 The actual implementation of receiving a tensor and its metadata from

187 the target rank.

@@ -198,7 +198,7 @@ class PyNcclPipe(KVPipeBase):

198 return buffer

199

200 def send_tensor_wrapper(

201 - self, tensor: Optional[torch.Tensor], tensor_size: int

201 + self, tensor: torch.Tensor | None, tensor_size: int

202 ) -> None:

203 """

204 Wrapper for _send_impl to handle exceptions and update buffer size.

@@ -228,7 +228,7 @@ class PyNcclPipe(KVPipeBase):

228 logger.debug("KV cache transfer pipe is full. Waiting...")

229 time.sleep(0.05)

230

231 - def send_tensor(self, tensor: Optional[torch.Tensor]) -> None:

231 + def send_tensor(self, tensor: torch.Tensor | None) -> None:

232 """

233 Sends a tensor and its metadata to the destination rank in a

234 non-blocking way.

@@ -251,7 +251,7 @@ class PyNcclPipe(KVPipeBase):

251

252 self.transport_thread.submit(self.send_tensor_wrapper, tensor, tensor_size)

253

254 - def recv_tensor(self) -> Optional[torch.Tensor]:

254 + def recv_tensor(self) -> torch.Tensor | None:

255 """

256 Receives a tensor and its metadata from the source rank. Blocking call.

257

vllm/distributed/kv_transfer/kv_transfer_state.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import TYPE_CHECKING, Optional

3 +from typing import TYPE_CHECKING

4

5 from vllm import envs

6 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType

@@ -13,7 +13,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (

13 if TYPE_CHECKING:

14 from vllm.config import VllmConfig

15

16 -_KV_CONNECTOR_AGENT: Optional[KVConnectorBaseType] = None

16 +_KV_CONNECTOR_AGENT: KVConnectorBaseType | None = None

17

18

19 def get_kv_transfer_group() -> KVConnectorBaseType:

@@ -27,7 +27,7 @@ def has_kv_transfer_group() -> bool:

27 return _KV_CONNECTOR_AGENT is not None

28

29

30 -def is_v1_kv_transfer_group(connector: Optional[KVConnectorBaseType] = None) -> bool:

30 +def is_v1_kv_transfer_group(connector: KVConnectorBaseType | None = None) -> bool:

31 """Check if the KV connector is the v1 connector.

32 If the argument is None, it will check the global KV connector

33

vllm/distributed/parallel_state.py

+43 -44

@@ -28,11 +28,12 @@ import gc

28 import pickle

29 import weakref

30 from collections import namedtuple

31 +from collections.abc import Callable

32 from contextlib import contextmanager, nullcontext

33 from dataclasses import dataclass

34 from datetime import timedelta

35 from multiprocessing import shared_memory

35 -from typing import Any, Callable, Optional, Union

36 +from typing import Any, Optional

37 from unittest.mock import patch

38

39 import torch

@@ -65,7 +66,7 @@ TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])

66

67

68 def _split_tensor_dict(

68 - tensor_dict: dict[str, Union[torch.Tensor, Any]],

69 + tensor_dict: dict[str, torch.Tensor | Any],

70 ) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:

71 """Split the tensor dictionary into two parts:

72 1. A list of (key, value) pairs. If the value is a tensor, it is replaced

@@ -300,17 +301,17 @@ class GroupCoordinator:

301 cpu_group: ProcessGroup # group for CPU communication

302 device_group: ProcessGroup # group for device communication

303 # device communicator (if use_device_communicator=True)

303 - device_communicator: Optional[DeviceCommunicatorBase]

304 - mq_broadcaster: Optional[Any] # shared memory broadcaster

304 + device_communicator: DeviceCommunicatorBase | None

305 + mq_broadcaster: Any | None # shared memory broadcaster

306

307 def __init__(

308 self,

309 group_ranks: list[list[int]],

310 local_rank: int,

310 - torch_distributed_backend: Union[str, Backend],

311 + torch_distributed_backend: str | Backend,

312 use_device_communicator: bool, # whether to use device communicator

313 use_message_queue_broadcaster: bool = False,

313 - group_name: Optional[str] = None,

314 + group_name: str | None = None,

315 ):

316 group_name = group_name or "anonymous"

317 self.unique_name = _get_unique_name(group_name)

@@ -368,7 +369,7 @@ class GroupCoordinator:

369

370 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue

371

371 - self.mq_broadcaster: Optional[MessageQueue] = None

372 + self.mq_broadcaster: MessageQueue | None = None

373 if use_message_queue_broadcaster and self.world_size > 1:

374 self.mq_broadcaster = MessageQueue.create_from_process_group(

375 self.cpu_group, 1 << 22, 6

@@ -419,9 +420,7 @@ class GroupCoordinator:

420 return self.ranks[(rank_in_group - 1) % world_size]

421

422 @contextmanager

422 - def graph_capture(

423 - self, graph_capture_context: Optional[GraphCaptureContext] = None

424 - ):

423 + def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None):

424 if graph_capture_context is None:

425 stream = torch.cuda.Stream()

426 graph_capture_context = GraphCaptureContext(stream)

@@ -502,9 +501,9 @@ class GroupCoordinator:

501

502 def all_gatherv(

503 self,

505 - input_: Union[torch.Tensor, list[torch.Tensor]],

504 + input_: torch.Tensor | list[torch.Tensor],

505 dim: int = 0,

507 - sizes: Optional[list[int]] = None,

506 + sizes: list[int] | None = None,

507 ):

508 if self.device_communicator is None:

509 raise ValueError("No device communicator found")

@@ -527,7 +526,7 @@ class GroupCoordinator:

526 return self._reduce_scatter_out_place(input_, dim)

527

528 def reduce_scatterv(

530 - self, input_: torch.Tensor, dim: int = -1, sizes: Optional[list[int]] = None

529 + self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None

530 ) -> torch.Tensor:

531 if self.device_communicator is None:

532 raise ValueError("No device communicator found")

@@ -540,7 +539,7 @@ class GroupCoordinator:

539

540 def gather(

541 self, input_: torch.Tensor, dst: int = 0, dim: int = -1

543 - ) -> Optional[torch.Tensor]:

542 + ) -> torch.Tensor | None:

543 """

544 NOTE: We assume that the input tensor is on the same device across

545 all the ranks.

@@ -569,7 +568,7 @@ class GroupCoordinator:

568 )

569 return input_

570

572 - def broadcast_object(self, obj: Optional[Any] = None, src: int = 0):

571 + def broadcast_object(self, obj: Any | None = None, src: int = 0):

572 """Broadcast the input object.

573 NOTE: `src` is the local rank of the source rank.

574 """

@@ -594,7 +593,7 @@ class GroupCoordinator:

593 return recv[0]

594

595 def broadcast_object_list(

597 - self, obj_list: list[Any], src: int = 0, group: Optional[ProcessGroup] = None

596 + self, obj_list: list[Any], src: int = 0, group: ProcessGroup | None = None

597 ):

598 """Broadcast the input object list.

599 NOTE: `src` is the local rank of the source rank.

@@ -675,11 +674,11 @@ class GroupCoordinator:

674

675 def broadcast_tensor_dict(

676 self,

678 - tensor_dict: Optional[dict[str, Union[torch.Tensor, Any]]] = None,

677 + tensor_dict: dict[str, torch.Tensor | Any] | None = None,

678 src: int = 0,

680 - group: Optional[ProcessGroup] = None,

681 - metadata_group: Optional[ProcessGroup] = None,

682 - ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

679 + group: ProcessGroup | None = None,

680 + metadata_group: ProcessGroup | None = None,

681 + ) -> dict[str, torch.Tensor | Any] | None:

682 """Broadcast the input tensor dictionary.

683 NOTE: `src` is the local rank of the source rank.

684 """

@@ -757,11 +756,11 @@ class GroupCoordinator:

756

757 def send_tensor_dict(

758 self,

760 - tensor_dict: dict[str, Union[torch.Tensor, Any]],

761 - dst: Optional[int] = None,

759 + tensor_dict: dict[str, torch.Tensor | Any],

760 + dst: int | None = None,

761 all_gather_group: Optional["GroupCoordinator"] = None,

763 - all_gather_tensors: Optional[dict[str, bool]] = None,

764 - ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

762 + all_gather_tensors: dict[str, bool] | None = None,

763 + ) -> dict[str, torch.Tensor | Any] | None:

764 """Send the input tensor dictionary.

765 NOTE: `dst` is the local rank of the source rank.

766

@@ -845,10 +844,10 @@ class GroupCoordinator:

844

845 def recv_tensor_dict(

846 self,

848 - src: Optional[int] = None,

847 + src: int | None = None,

848 all_gather_group: Optional["GroupCoordinator"] = None,

850 - all_gather_tensors: Optional[dict[str, bool]] = None,

851 - ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

849 + all_gather_tensors: dict[str, bool] | None = None,

850 + ) -> dict[str, torch.Tensor | Any] | None:

851 """Recv the input tensor dictionary.

852 NOTE: `src` is the local rank of the source rank.

853

@@ -943,7 +942,7 @@ class GroupCoordinator:

942 """

943 torch.distributed.barrier(group=self.cpu_group)

944

946 - def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:

945 + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:

946 """Sends a tensor to the destination rank in a blocking way"""

947 """NOTE: `dst` is the local rank of the destination rank."""

948 if self.device_communicator is None:

@@ -951,7 +950,7 @@ class GroupCoordinator:

950 self.device_communicator.send(tensor, dst)

951

952 def recv(

954 - self, size: torch.Size, dtype: torch.dtype, src: Optional[int] = None

953 + self, size: torch.Size, dtype: torch.dtype, src: int | None = None

954 ) -> torch.Tensor:

955 """Receives a tensor from the source rank."""

956 """NOTE: `src` is the local rank of the source rank."""

@@ -997,8 +996,8 @@ class GroupCoordinator:

996 return hidden_states

997

998

1000 -_WORLD: Optional[GroupCoordinator] = None

1001 -_NODE_COUNT: Optional[int] = None

999 +_WORLD: GroupCoordinator | None = None

1000 +_NODE_COUNT: int | None = None

1001

1002

1003 def get_world_group() -> GroupCoordinator:

@@ -1023,7 +1022,7 @@ def init_model_parallel_group(

1022 local_rank: int,

1023 backend: str,

1024 use_message_queue_broadcaster: bool = False,

1026 - group_name: Optional[str] = None,

1025 + group_name: str | None = None,

1026 ) -> GroupCoordinator:

1027 return GroupCoordinator(

1028 group_ranks=group_ranks,

@@ -1035,7 +1034,7 @@ def init_model_parallel_group(

1034 )

1035

1036

1038 -_TP: Optional[GroupCoordinator] = None

1037 +_TP: GroupCoordinator | None = None

1038

1039

1040 def get_tp_group() -> GroupCoordinator:

@@ -1052,7 +1051,7 @@ def get_tensor_model_parallel_group():

1051 return get_tp_group()

1052

1053

1055 -_DCP: Optional[GroupCoordinator] = None

1054 +_DCP: GroupCoordinator | None = None

1055

1056

1057 def get_dcp_group() -> GroupCoordinator:

@@ -1063,9 +1062,9 @@ def get_dcp_group() -> GroupCoordinator:

1062 # kept for backward compatibility

1063 get_context_model_parallel_group = get_dcp_group

1064

1066 -_PP: Optional[GroupCoordinator] = None

1065 +_PP: GroupCoordinator | None = None

1066

1068 -_DP: Optional[GroupCoordinator] = None

1067 +_DP: GroupCoordinator | None = None

1068

1069

1070 def get_dp_group() -> GroupCoordinator:

@@ -1073,7 +1072,7 @@ def get_dp_group() -> GroupCoordinator:

1072 return _DP

1073

1074

1076 -_EP: Optional[GroupCoordinator] = None

1075 +_EP: GroupCoordinator | None = None

1076

1077

1078 def get_ep_group() -> GroupCoordinator:

@@ -1131,7 +1130,7 @@ def init_distributed_environment(

1130 distributed_init_method: str = "env://",

1131 local_rank: int = -1,

1132 backend: str = "nccl",

1134 - timeout: Optional[timedelta] = None,

1133 + timeout: timedelta | None = None,

1134 ):

1135 logger.debug(

1136 "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",

@@ -1208,8 +1207,8 @@ def init_distributed_environment(

1207 def initialize_model_parallel(

1208 tensor_model_parallel_size: int = 1,

1209 pipeline_model_parallel_size: int = 1,

1211 - decode_context_model_parallel_size: Optional[int] = 1,

1212 - backend: Optional[str] = None,

1210 + decode_context_model_parallel_size: int | None = 1,

1211 + backend: str | None = None,

1212 ) -> None:

1213 """

1214 Initialize model parallel groups.

@@ -1338,8 +1337,8 @@ def initialize_model_parallel(

1337 def ensure_model_parallel_initialized(

1338 tensor_model_parallel_size: int,

1339 pipeline_model_parallel_size: int,

1341 - decode_context_model_parallel_size: Optional[int] = 1,

1342 - backend: Optional[str] = None,

1340 + decode_context_model_parallel_size: int | None = 1,

1341 + backend: str | None = None,

1342 ) -> None:

1343 """Helper to initialize model parallel groups if they are not initialized,

1344 or ensure tensor-parallel and pipeline-parallel sizes are equal to expected

@@ -1504,7 +1503,7 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):

1503

1504

1505 def in_the_same_node_as(

1507 - pg: Union[ProcessGroup, StatelessProcessGroup], source_rank: int = 0

1506 + pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0

1507 ) -> list[bool]:

1508 """

1509 This is a collective operation that returns if each rank is in the same node

@@ -1624,7 +1623,7 @@ def is_global_first_rank() -> bool:

1623 return True

1624

1625

1627 -def _node_count(pg: Union[ProcessGroup, StatelessProcessGroup]) -> int:

1626 +def _node_count(pg: ProcessGroup | StatelessProcessGroup) -> int:

1627 """

1628 Returns the total number of nodes in the process group.

1629

vllm/distributed/tpu_distributed_utils.py

+3 -3

@@ -30,9 +30,9 @@ class XlaQKVParallelLinear(nn.Module):

30 self.q_weight: Parameter

31 self.k_weight: Parameter

32 self.v_weight: Parameter

33 - self.q_bias: Optional[Parameter]

34 - self.k_bias: Optional[Parameter]

35 - self.v_bias: Optional[Parameter]

33 + self.q_bias: Parameter | None

34 + self.k_bias: Parameter | None

35 + self.v_bias: Parameter | None

36 self._load_weights_from_qkv_linear(qkv_linear)

37 if mesh is not None:

38 self._shard_weight(mesh)

vllm/distributed/utils.py

+3 -3

@@ -15,7 +15,7 @@ import uuid

15 from collections import deque

16 from collections.abc import Sequence

17 from datetime import timedelta

18 -from typing import Any, Optional

18 +from typing import Any

19

20 import torch

21 from torch.distributed import ProcessGroup, TCPStore

@@ -150,7 +150,7 @@ class StatelessProcessGroup:

150 store: torch._C._distributed_c10d.Store

151

152 # stores a reference to the socket so that the file descriptor stays alive

153 - socket: Optional[socket.socket]

153 + socket: socket.socket | None

154

155 data_expiration_seconds: int = 3600 # 1 hour

156

@@ -197,7 +197,7 @@ class StatelessProcessGroup:

197 self.recv_src_counter[src] += 1

198 return obj

199

200 - def broadcast_obj(self, obj: Optional[Any], src: int) -> Any:

200 + def broadcast_obj(self, obj: Any | None, src: int) -> Any:

201 """Broadcast an object from a source rank to all other ranks.

202 It does not clean up after all ranks have received the object.

203 Use it for limited times, e.g., for initialization.

vllm/engine/arg_utils.py

+71 -68

@@ -7,15 +7,16 @@ import dataclasses

7 import functools

8 import json

9 import sys

10 +from collections.abc import Callable

11 from dataclasses import MISSING, dataclass, fields, is_dataclass

12 from itertools import permutations

13 +from types import UnionType

14 from typing import (

15 TYPE_CHECKING,

16 Annotated,

17 Any,

16 - Callable,

18 Literal,

18 - Optional,

19 + TypeAlias,

20 TypeVar,

21 Union,

22 cast,

@@ -98,8 +99,8 @@ logger = init_logger(__name__)

99

100 # object is used to allow for special typing forms

101 T = TypeVar("T")

101 -TypeHint = Union[type[Any], object]

102 -TypeHintT = Union[type[T], object]

102 +TypeHint: TypeAlias = type[Any] | object

103 +TypeHintT: TypeAlias = type[T] | object

104

105

106 def parse_type(return_type: Callable[[str], T]) -> Callable[[str], T]:

@@ -114,8 +115,8 @@ def parse_type(return_type: Callable[[str], T]) -> Callable[[str], T]:

115 return _parse_type

116

117

117 -def optional_type(return_type: Callable[[str], T]) -> Callable[[str], Optional[T]]:

118 - def _optional_type(val: str) -> Optional[T]:

118 +def optional_type(return_type: Callable[[str], T]) -> Callable[[str], T | None]:

119 + def _optional_type(val: str) -> T | None:

120 if val == "" or val == "None":

121 return None

122 return parse_type(return_type)(val)

@@ -123,7 +124,7 @@ def optional_type(return_type: Callable[[str], T]) -> Callable[[str], Optional[T

124 return _optional_type

125

126

126 -def union_dict_and_str(val: str) -> Optional[Union[str, dict[str, str]]]:

127 +def union_dict_and_str(val: str) -> str | dict[str, str] | None:

128 if not re.match(r"(?s)^\s*{.*}\s*$", val):

129 return str(val)

130 return optional_type(json.loads)(val)

@@ -174,7 +175,8 @@ def get_type_hints(type_hint: TypeHint) -> set[TypeHint]:

175

176 if origin is Annotated:

177 type_hints.update(get_type_hints(args[0]))

177 - elif origin is Union:

178 + elif origin in {Union, UnionType}:

179 + # Union for Union[X, Y] and UnionType for X | Y

180 for arg in args:

181 type_hints.update(get_type_hints(arg))

182 else:

@@ -195,7 +197,7 @@ NEEDS_HELP = (

197

198

199 @functools.lru_cache(maxsize=30)

198 -def _compute_kwargs(cls: ConfigType) -> dict[str, Any]:

200 +def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:

201 # Save time only getting attr docs if we're generating help text

202 cls_docs = get_attr_docs(cls) if NEEDS_HELP else {}

203 kwargs = {}

@@ -262,7 +264,8 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, Any]:

264 type_hint = get_type(type_hints, list)

265 types = get_args(type_hint)

266 list_type = types[0]

265 - if get_origin(list_type) is Union:

267 + if get_origin(list_type) in {Union, UnionType}:

268 + # Union for Union[X, Y] and UnionType for X | Y

269 msg = "List type must contain str if it is a Union."

270 assert str in get_args(list_type), msg

271 list_type = str

@@ -310,7 +313,7 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, Any]:

313 return kwargs

314

315

313 -def get_kwargs(cls: ConfigType) -> dict[str, Any]:

316 +def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:

317 """Return argparse kwargs for the given Config dataclass.

318

319 If `--help` or `mkdocs` are not present in the command line command, the

@@ -328,43 +331,43 @@ class EngineArgs:

331 """Arguments for vLLM engine."""

332

333 model: str = ModelConfig.model

331 - served_model_name: Optional[Union[str, list[str]]] = ModelConfig.served_model_name

332 - tokenizer: Optional[str] = ModelConfig.tokenizer

333 - hf_config_path: Optional[str] = ModelConfig.hf_config_path

334 + served_model_name: str | list[str] | None = ModelConfig.served_model_name

335 + tokenizer: str | None = ModelConfig.tokenizer

336 + hf_config_path: str | None = ModelConfig.hf_config_path

337 runner: RunnerOption = ModelConfig.runner

338 convert: ConvertOption = ModelConfig.convert

336 - task: Optional[TaskOption] = ModelConfig.task

339 + task: TaskOption | None = ModelConfig.task

340 skip_tokenizer_init: bool = ModelConfig.skip_tokenizer_init

341 enable_prompt_embeds: bool = ModelConfig.enable_prompt_embeds

342 tokenizer_mode: TokenizerMode = ModelConfig.tokenizer_mode

343 trust_remote_code: bool = ModelConfig.trust_remote_code

344 allowed_local_media_path: str = ModelConfig.allowed_local_media_path

342 - allowed_media_domains: Optional[list[str]] = ModelConfig.allowed_media_domains

343 - download_dir: Optional[str] = LoadConfig.download_dir

345 + allowed_media_domains: list[str] | None = ModelConfig.allowed_media_domains

346 + download_dir: str | None = LoadConfig.download_dir

347 safetensors_load_strategy: str = LoadConfig.safetensors_load_strategy

345 - load_format: Union[str, LoadFormats] = LoadConfig.load_format

348 + load_format: str | LoadFormats = LoadConfig.load_format

349 config_format: str = ModelConfig.config_format

350 dtype: ModelDType = ModelConfig.dtype

351 kv_cache_dtype: CacheDType = CacheConfig.cache_dtype

349 - seed: Optional[int] = ModelConfig.seed

350 - max_model_len: Optional[int] = ModelConfig.max_model_len

352 + seed: int | None = ModelConfig.seed

353 + max_model_len: int | None = ModelConfig.max_model_len

354 cuda_graph_sizes: list[int] = get_field(SchedulerConfig, "cuda_graph_sizes")

355 # Note: Specifying a custom executor backend by passing a class

356 # is intended for expert use only. The API may change without

357 # notice.

355 - distributed_executor_backend: Optional[

356 - Union[str, DistributedExecutorBackend, type[ExecutorBase]]

357 - ] = ParallelConfig.distributed_executor_backend

358 + distributed_executor_backend: (

359 + str | DistributedExecutorBackend | type[ExecutorBase] | None

360 + ) = ParallelConfig.distributed_executor_backend

361 # number of P/D disaggregation (or other disaggregation) workers

362 pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size

363 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size

364 decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size

365 data_parallel_size: int = ParallelConfig.data_parallel_size

363 - data_parallel_rank: Optional[int] = None

364 - data_parallel_start_rank: Optional[int] = None

365 - data_parallel_size_local: Optional[int] = None

366 - data_parallel_address: Optional[str] = None

367 - data_parallel_rpc_port: Optional[int] = None

366 + data_parallel_rank: int | None = None

367 + data_parallel_start_rank: int | None = None

368 + data_parallel_size_local: int | None = None

369 + data_parallel_address: str | None = None

370 + data_parallel_rpc_port: int | None = None

371 data_parallel_hybrid_lb: bool = False

372 data_parallel_backend: str = ParallelConfig.data_parallel_backend

373 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel

@@ -385,11 +388,11 @@ class EngineArgs:

388 eplb_window_size: int = EPLBConfig.window_size

389 eplb_step_interval: int = EPLBConfig.step_interval

390 eplb_log_balancedness: bool = EPLBConfig.log_balancedness

388 - max_parallel_loading_workers: Optional[int] = (

391 + max_parallel_loading_workers: int | None = (

392 ParallelConfig.max_parallel_loading_workers

393 )

391 - block_size: Optional[BlockSize] = CacheConfig.block_size

392 - enable_prefix_caching: Optional[bool] = CacheConfig.enable_prefix_caching

394 + block_size: BlockSize | None = CacheConfig.block_size

395 + enable_prefix_caching: bool | None = CacheConfig.enable_prefix_caching

396 prefix_caching_hash_algo: PrefixCachingHashAlgo = (

397 CacheConfig.prefix_caching_hash_algo

398 )

@@ -398,62 +401,62 @@ class EngineArgs:

401 swap_space: float = CacheConfig.swap_space

402 cpu_offload_gb: float = CacheConfig.cpu_offload_gb

403 gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization

401 - kv_cache_memory_bytes: Optional[int] = CacheConfig.kv_cache_memory_bytes

402 - max_num_batched_tokens: Optional[int] = SchedulerConfig.max_num_batched_tokens

404 + kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes

405 + max_num_batched_tokens: int | None = SchedulerConfig.max_num_batched_tokens

406 max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills

407 max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills

408 long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold

406 - max_num_seqs: Optional[int] = SchedulerConfig.max_num_seqs

409 + max_num_seqs: int | None = SchedulerConfig.max_num_seqs

410 max_logprobs: int = ModelConfig.max_logprobs

411 logprobs_mode: LogprobsMode = ModelConfig.logprobs_mode

412 disable_log_stats: bool = False

410 - revision: Optional[str] = ModelConfig.revision

411 - code_revision: Optional[str] = ModelConfig.code_revision

413 + revision: str | None = ModelConfig.revision

414 + code_revision: str | None = ModelConfig.code_revision

415 rope_scaling: dict[str, Any] = get_field(ModelConfig, "rope_scaling")

413 - rope_theta: Optional[float] = ModelConfig.rope_theta

414 - hf_token: Optional[Union[bool, str]] = ModelConfig.hf_token

416 + rope_theta: float | None = ModelConfig.rope_theta

417 + hf_token: bool | str | None = ModelConfig.hf_token

418 hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides")

416 - tokenizer_revision: Optional[str] = ModelConfig.tokenizer_revision

417 - quantization: Optional[QuantizationMethods] = ModelConfig.quantization

419 + tokenizer_revision: str | None = ModelConfig.tokenizer_revision

420 + quantization: QuantizationMethods | None = ModelConfig.quantization

421 enforce_eager: bool = ModelConfig.enforce_eager

422 disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce

420 - limit_mm_per_prompt: dict[str, Union[int, dict[str, int]]] = get_field(

423 + limit_mm_per_prompt: dict[str, int | dict[str, int]] = get_field(

424 MultiModalConfig, "limit_per_prompt"

425 )

426 interleave_mm_strings: bool = MultiModalConfig.interleave_mm_strings

427 media_io_kwargs: dict[str, dict[str, Any]] = get_field(

428 MultiModalConfig, "media_io_kwargs"

429 )

427 - mm_processor_kwargs: Optional[dict[str, Any]] = MultiModalConfig.mm_processor_kwargs

430 + mm_processor_kwargs: dict[str, Any] | None = MultiModalConfig.mm_processor_kwargs

431 disable_mm_preprocessor_cache: bool = False # DEPRECATED

432 mm_processor_cache_gb: float = MultiModalConfig.mm_processor_cache_gb

430 - mm_processor_cache_type: Optional[MMCacheType] = (

433 + mm_processor_cache_type: MMCacheType | None = (

434 MultiModalConfig.mm_processor_cache_type

435 )

436 mm_shm_cache_max_object_size_mb: int = (

437 MultiModalConfig.mm_shm_cache_max_object_size_mb

438 )

439 mm_encoder_tp_mode: MMEncoderTPMode = MultiModalConfig.mm_encoder_tp_mode

437 - io_processor_plugin: Optional[str] = None

440 + io_processor_plugin: str | None = None

441 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling

442 video_pruning_rate: float = MultiModalConfig.video_pruning_rate

443 # LoRA fields

444 enable_lora: bool = False

445 max_loras: int = LoRAConfig.max_loras

446 max_lora_rank: int = LoRAConfig.max_lora_rank

444 - default_mm_loras: Optional[dict[str, str]] = LoRAConfig.default_mm_loras

447 + default_mm_loras: dict[str, str] | None = LoRAConfig.default_mm_loras

448 fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras

446 - max_cpu_loras: Optional[int] = LoRAConfig.max_cpu_loras

447 - lora_dtype: Optional[Union[str, torch.dtype]] = LoRAConfig.lora_dtype

449 + max_cpu_loras: int | None = LoRAConfig.max_cpu_loras

450 + lora_dtype: str | torch.dtype | None = LoRAConfig.lora_dtype

451 lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size

452

453 ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight

451 - num_gpu_blocks_override: Optional[int] = CacheConfig.num_gpu_blocks_override

454 + num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override

455 num_lookahead_slots: int = SchedulerConfig.num_lookahead_slots

456 model_loader_extra_config: dict = get_field(LoadConfig, "model_loader_extra_config")

454 - ignore_patterns: Union[str, list[str]] = get_field(LoadConfig, "ignore_patterns")

457 + ignore_patterns: str | list[str] = get_field(LoadConfig, "ignore_patterns")

458

456 - enable_chunked_prefill: Optional[bool] = SchedulerConfig.enable_chunked_prefill

459 + enable_chunked_prefill: bool | None = SchedulerConfig.enable_chunked_prefill

460 disable_chunked_mm_input: bool = SchedulerConfig.disable_chunked_mm_input

461

462 disable_hybrid_kv_cache_manager: bool = (

@@ -465,35 +468,35 @@ class EngineArgs:

468 )

469 reasoning_parser: str = StructuredOutputsConfig.reasoning_parser

470 # Deprecated guided decoding fields

468 - guided_decoding_backend: Optional[str] = None

469 - guided_decoding_disable_fallback: Optional[bool] = None

470 - guided_decoding_disable_any_whitespace: Optional[bool] = None

471 - guided_decoding_disable_additional_properties: Optional[bool] = None

471 + guided_decoding_backend: str | None = None

472 + guided_decoding_disable_fallback: bool | None = None

473 + guided_decoding_disable_any_whitespace: bool | None = None

474 + guided_decoding_disable_additional_properties: bool | None = None

475

473 - logits_processor_pattern: Optional[str] = ModelConfig.logits_processor_pattern

476 + logits_processor_pattern: str | None = ModelConfig.logits_processor_pattern

477

475 - speculative_config: Optional[dict[str, Any]] = None

478 + speculative_config: dict[str, Any] | None = None

479

477 - show_hidden_metrics_for_version: Optional[str] = (

480 + show_hidden_metrics_for_version: str | None = (

481 ObservabilityConfig.show_hidden_metrics_for_version

482 )

480 - otlp_traces_endpoint: Optional[str] = ObservabilityConfig.otlp_traces_endpoint

481 - collect_detailed_traces: Optional[list[DetailedTraceModules]] = (

483 + otlp_traces_endpoint: str | None = ObservabilityConfig.otlp_traces_endpoint

484 + collect_detailed_traces: list[DetailedTraceModules] | None = (

485 ObservabilityConfig.collect_detailed_traces

486 )

487 scheduling_policy: SchedulerPolicy = SchedulerConfig.policy

485 - scheduler_cls: Union[str, type[object]] = SchedulerConfig.scheduler_cls

488 + scheduler_cls: str | type[object] = SchedulerConfig.scheduler_cls

489

487 - pooler_config: Optional[PoolerConfig] = ModelConfig.pooler_config

488 - override_pooler_config: Optional[Union[dict, PoolerConfig]] = (

490 + pooler_config: PoolerConfig | None = ModelConfig.pooler_config

491 + override_pooler_config: dict | PoolerConfig | None = (

492 ModelConfig.override_pooler_config

493 )

494 compilation_config: CompilationConfig = get_field(VllmConfig, "compilation_config")

495 worker_cls: str = ParallelConfig.worker_cls

496 worker_extension_cls: str = ParallelConfig.worker_extension_cls

497

495 - kv_transfer_config: Optional[KVTransferConfig] = None

496 - kv_events_config: Optional[KVEventsConfig] = None

498 + kv_transfer_config: KVTransferConfig | None = None

499 + kv_events_config: KVEventsConfig | None = None

500

501 generation_config: str = ModelConfig.generation_config

502 enable_sleep_mode: bool = ModelConfig.enable_sleep_mode

@@ -515,7 +518,7 @@ class EngineArgs:

518 # DEPRECATED

519 enable_multimodal_encoder_data_parallel: bool = False

520

518 - logits_processors: Optional[list[Union[str, type[LogitsProcessor]]]] = (

521 + logits_processors: list[str | type[LogitsProcessor]] | None = (

522 ModelConfig.logits_processors

523 )

524 """Custom logitproc types"""

@@ -1187,7 +1190,7 @@ class EngineArgs:

1190 target_parallel_config: ParallelConfig,

1191 enable_chunked_prefill: bool,

1192 disable_log_stats: bool,

1190 - ) -> Optional["SpeculativeConfig"]:

1193 + ) -> SpeculativeConfig | None:

1194 """Initializes and returns a SpeculativeConfig object based on

1195 `speculative_config`.

1196

@@ -1214,7 +1217,7 @@ class EngineArgs:

1217

1218 def create_engine_config(

1219 self,

1217 - usage_context: Optional[UsageContext] = None,

1220 + usage_context: UsageContext | None = None,

1221 headless: bool = False,

1222 ) -> VllmConfig:

1223 """

@@ -1282,7 +1285,7 @@ class EngineArgs:

1285 self.enable_chunked_prefill = False

1286 assert self.enable_chunked_prefill is not None

1287

1285 - sliding_window: Optional[int] = None

1288 + sliding_window: int | None = None

1289 if not is_interleaved(model_config.hf_text_config):

1290 # Only set CacheConfig.sliding_window if the model is all sliding

1291 # window. Otherwise CacheConfig.sliding_window will override the

vllm/engine/metrics.py

+13 -13

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

3

4 import time

5 from collections import Counter as CollectionsCounter

6 -from typing import Optional, Union, cast

6 +from typing import cast

7

8 import numpy as np

9 import prometheus_client

@@ -304,7 +304,7 @@ class _RayGaugeWrapper:

304 self,

305 name: str,

306 documentation: str = "",

307 - labelnames: Optional[list[str]] = None,

307 + labelnames: list[str] | None = None,

308 multiprocess_mode: str = "",

309 ):

310 del multiprocess_mode

@@ -317,7 +317,7 @@ class _RayGaugeWrapper:

317 self._gauge.set_default_tags(labels)

318 return self

319

320 - def set(self, value: Union[int, float]):

320 + def set(self, value: int | float):

321 return self._gauge.set(value)

322

323 def set_to_current_time(self):

@@ -330,7 +330,7 @@ class _RayCounterWrapper:

330 prometheus_client.Counter"""

331

332 def __init__(

333 - self, name: str, documentation: str = "", labelnames: Optional[list[str]] = None

333 + self, name: str, documentation: str = "", labelnames: list[str] | None = None

334 ):

335 labelnames_tuple = tuple(labelnames) if labelnames else None

336 self._counter = ray_metrics.Counter(

@@ -341,7 +341,7 @@ class _RayCounterWrapper:

341 self._counter.set_default_tags(labels)

342 return self

343

344 - def inc(self, value: Union[int, float] = 1.0):

344 + def inc(self, value: int | float = 1.0):

345 if value == 0:

346 return

347 return self._counter.inc(value)

@@ -355,8 +355,8 @@ class _RayHistogramWrapper:

355 self,

356 name: str,

357 documentation: str = "",

358 - labelnames: Optional[list[str]] = None,

359 - buckets: Optional[list[float]] = None,

358 + labelnames: list[str] | None = None,

359 + buckets: list[float] | None = None,

360 ):

361 labelnames_tuple = tuple(labelnames) if labelnames else None

362 boundaries = buckets if buckets else []

@@ -371,7 +371,7 @@ class _RayHistogramWrapper:

371 self._histogram.set_default_tags(labels)

372 return self

373

374 - def observe(self, value: Union[int, float]):

374 + def observe(self, value: int | float):

375 return self._histogram.observe(value)

376

377

@@ -451,8 +451,8 @@ class LoggingStatLogger(StatLoggerBase):

451

452 def __init__(self, local_interval: float, vllm_config: VllmConfig) -> None:

453 super().__init__(local_interval, vllm_config)

454 - self.last_prompt_throughput: Optional[float] = None

455 - self.last_generation_throughput: Optional[float] = None

454 + self.last_prompt_throughput: float | None = None

455 + self.last_generation_throughput: float | None = None

456

457 def log(self, stats: Stats) -> None:

458 """Called by LLMEngine.

@@ -539,11 +539,11 @@ class PrometheusStatLogger(StatLoggerBase):

539 labelnames=list(labels.keys()), vllm_config=vllm_config

540 )

541

542 - def _log_gauge(self, gauge, data: Union[int, float]) -> None:

542 + def _log_gauge(self, gauge, data: int | float) -> None:

543 # Convenience function for logging to gauge.

544 gauge.labels(**self.labels).set(data)

545

546 - def _log_counter(self, counter, data: Union[int, float]) -> None:

546 + def _log_counter(self, counter, data: int | float) -> None:

547 # Convenience function for logging to counter.

548 # Prevent ValueError from negative increment

549 if data < 0:

@@ -558,7 +558,7 @@ class PrometheusStatLogger(StatLoggerBase):

558 for label, count in data.items():

559 counter.labels(**{**self.labels, label_key: label}).inc(count)

560

561 - def _log_histogram(self, histogram, data: Union[list[int], list[float]]) -> None:

561 + def _log_histogram(self, histogram, data: list[int] | list[float]) -> None:

562 # Convenience function for logging list to histogram.

563 for datum in data:

564 histogram.labels(**self.labels).observe(datum)

vllm/engine/protocol.py

+16 -16

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

3

4 from abc import ABC, abstractmethod

5 from collections.abc import AsyncGenerator, Iterable, Mapping

6 -from typing import Any, Optional, Union

6 +from typing import Any

7

8 from vllm.config import ModelConfig, VllmConfig

9 from vllm.inputs.data import PromptType

@@ -28,7 +28,7 @@ class EngineClient(ABC):

28 vllm_config: VllmConfig

29 model_config: ModelConfig

30 processor: Processor

31 - io_processor: Optional[IOProcessor]

31 + io_processor: IOProcessor | None

32

33 @property

34 @abstractmethod

@@ -49,16 +49,16 @@ class EngineClient(ABC):

49 @abstractmethod

50 def generate(

51 self,

52 - prompt: Union[EngineCoreRequest, PromptType],

52 + prompt: EngineCoreRequest | PromptType,

53 sampling_params: SamplingParams,

54 request_id: str,

55 *,

56 - prompt_text: Optional[str] = None,

57 - lora_request: Optional[LoRARequest] = None,

58 - tokenization_kwargs: Optional[dict[str, Any]] = None,

59 - trace_headers: Optional[Mapping[str, str]] = None,

56 + prompt_text: str | None = None,

57 + lora_request: LoRARequest | None = None,

58 + tokenization_kwargs: dict[str, Any] | None = None,

59 + trace_headers: Mapping[str, str] | None = None,

60 priority: int = 0,

61 - data_parallel_rank: Optional[int] = None,

61 + data_parallel_rank: int | None = None,

62 ) -> AsyncGenerator[RequestOutput, None]:

63 """Generate outputs for a request."""

64 ...

@@ -69,16 +69,16 @@ class EngineClient(ABC):

69 prompt: PromptType,

70 pooling_params: PoolingParams,

71 request_id: str,

72 - lora_request: Optional[LoRARequest] = None,

73 - trace_headers: Optional[Mapping[str, str]] = None,

72 + lora_request: LoRARequest | None = None,

73 + trace_headers: Mapping[str, str] | None = None,

74 priority: int = 0,

75 - tokenization_kwargs: Optional[dict[str, Any]] = None,

75 + tokenization_kwargs: dict[str, Any] | None = None,

76 ) -> AsyncGenerator[PoolingRequestOutput, None]:

77 """Generate outputs for a request from a pooling model."""

78 ...

79

80 @abstractmethod

81 - async def abort(self, request_id: Union[str, Iterable[str]]) -> None:

81 + async def abort(self, request_id: str | Iterable[str]) -> None:

82 """Abort a request.

83

84 Args:

@@ -119,7 +119,7 @@ class EngineClient(ABC):

119 ...

120

121 @abstractmethod

122 - async def reset_prefix_cache(self, device: Optional[Device] = None) -> None:

122 + async def reset_prefix_cache(self, device: Device | None = None) -> None:

123 """Reset the prefix cache"""

124 ...

125

@@ -129,7 +129,7 @@ class EngineClient(ABC):

129 ...

130

131 @abstractmethod

132 - async def wake_up(self, tags: Optional[list[str]] = None) -> None:

132 + async def wake_up(self, tags: list[str] | None = None) -> None:

133 """Wake up the engine"""

134 ...

135

@@ -152,9 +152,9 @@ class EngineClient(ABC):

152 async def collective_rpc(

153 self,

154 method: str,

155 - timeout: Optional[float] = None,

155 + timeout: float | None = None,

156 args: tuple = (),

157 - kwargs: Optional[dict] = None,

157 + kwargs: dict | None = None,

158 ):

159 """Perform a collective RPC call to the given path."""

160 raise NotImplementedError

vllm/entrypoints/api_server.py

+3 -3

@@ -13,7 +13,7 @@ import json

13 import ssl

14 from argparse import Namespace

15 from collections.abc import AsyncGenerator

16 -from typing import Any, Optional

16 +from typing import Any

17

18 from fastapi import FastAPI, Request

19 from fastapi.responses import JSONResponse, Response, StreamingResponse

@@ -101,7 +101,7 @@ def build_app(args: Namespace) -> FastAPI:

101

102 async def init_app(

103 args: Namespace,

104 - llm_engine: Optional[AsyncLLMEngine] = None,

104 + llm_engine: AsyncLLMEngine | None = None,

105 ) -> FastAPI:

106 app = build_app(args)

107

@@ -120,7 +120,7 @@ async def init_app(

120

121

122 async def run_server(

123 - args: Namespace, llm_engine: Optional[AsyncLLMEngine] = None, **uvicorn_kwargs: Any

123 + args: Namespace, llm_engine: AsyncLLMEngine | None = None, **uvicorn_kwargs: Any

124 ) -> None:

125 logger.info("vLLM API server version %s", VLLM_VERSION)

126 logger.info("args: %s", args)

vllm/entrypoints/chat_utils.py

+104 -107

@@ -5,10 +5,10 @@ import asyncio

5 import json

6 from abc import ABC, abstractmethod

7 from collections import Counter, defaultdict, deque

8 -from collections.abc import Awaitable, Iterable

8 +from collections.abc import Awaitable, Callable, Iterable

9 from functools import cached_property, lru_cache, partial

10 from pathlib import Path

11 -from typing import Any, Callable, Generic, Literal, Optional, TypeVar, Union, cast

11 +from typing import Any, Generic, Literal, TypeAlias, TypeVar, cast

12

13 import jinja2

14 import jinja2.ext

@@ -40,7 +40,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter

40 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast, ProcessorMixin

41

42 # pydantic needs the TypedDict from typing_extensions

43 -from typing_extensions import Required, TypeAlias, TypedDict

43 +from typing_extensions import Required, TypedDict

44

45 from vllm.config import ModelConfig

46 from vllm.logger import init_logger

@@ -76,7 +76,7 @@ class ChatCompletionContentPartAudioParam(TypedDict, total=False):

76

77

78 class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):

79 - image_embeds: Optional[Union[str, dict[str, str]]]

79 + image_embeds: str | dict[str, str] | None

80 """

81 The image embeddings. It can be either:

82 - A single base64 string.

@@ -84,7 +84,7 @@ class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):

84 """

85 type: Required[Literal["image_embeds"]]

86 """The type of the content part."""

87 - uuid: Optional[str]

87 + uuid: str | None

88 """

89 User-provided UUID of a media. User must guarantee that it is properly

90 generated and unique for different medias.

@@ -123,8 +123,8 @@ class CustomChatCompletionContentPILImageParam(TypedDict, total=False):

123 }

124 """

125

126 - image_pil: Optional[PILImage]

127 - uuid: Optional[str]

126 + image_pil: PILImage | None

127 + uuid: str | None

128 """

129 User-provided UUID of a media. User must guarantee that it is properly

130 generated and unique for different medias.

@@ -141,8 +141,8 @@ class CustomChatCompletionContentSimpleImageParam(TypedDict, total=False):

141 }

142 """

143

144 - image_url: Optional[str]

145 - uuid: Optional[str]

144 + image_url: str | None

145 + uuid: str | None

146 """

147 User-provided UUID of a media. User must guarantee that it is properly

148 generated and unique for different medias.

@@ -158,7 +158,7 @@ class CustomChatCompletionContentSimpleAudioParam(TypedDict, total=False):

158 }

159 """

160

161 - audio_url: Optional[str]

161 + audio_url: str | None

162

163

164 class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False):

@@ -170,8 +170,8 @@ class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False):

170 }

171 """

172

173 - video_url: Optional[str]

174 - uuid: Optional[str]

173 + video_url: str | None

174 + uuid: str | None

175 """

176 User-provided UUID of a media. User must guarantee that it is properly

177 generated and unique for different medias.

@@ -199,20 +199,20 @@ class CustomThinkCompletionContentParam(TypedDict, total=False):

199 """The thinking type."""

200

201

202 -ChatCompletionContentPartParam: TypeAlias = Union[

203 - OpenAIChatCompletionContentPartParam,

204 - ChatCompletionContentPartAudioParam,

205 - ChatCompletionContentPartInputAudioParam,

206 - ChatCompletionContentPartVideoParam,

207 - ChatCompletionContentPartRefusalParam,

208 - CustomChatCompletionContentPILImageParam,

209 - CustomChatCompletionContentSimpleImageParam,

210 - ChatCompletionContentPartImageEmbedsParam,

211 - CustomChatCompletionContentSimpleAudioParam,

212 - CustomChatCompletionContentSimpleVideoParam,

213 - str,

214 - CustomThinkCompletionContentParam,

215 -]

202 +ChatCompletionContentPartParam: TypeAlias = (

203 + OpenAIChatCompletionContentPartParam

204 + | ChatCompletionContentPartAudioParam

205 + | ChatCompletionContentPartInputAudioParam

206 + | ChatCompletionContentPartVideoParam

207 + | ChatCompletionContentPartRefusalParam

208 + | CustomChatCompletionContentPILImageParam

209 + | CustomChatCompletionContentSimpleImageParam

210 + | ChatCompletionContentPartImageEmbedsParam

211 + | CustomChatCompletionContentSimpleAudioParam

212 + | CustomChatCompletionContentSimpleVideoParam

213 + | str

214 + | CustomThinkCompletionContentParam

215 +)

216

217

218 class CustomChatCompletionMessageParam(TypedDict, total=False):

@@ -221,7 +221,7 @@ class CustomChatCompletionMessageParam(TypedDict, total=False):

221 role: Required[str]

222 """The role of the message's author."""

223

224 - content: Union[str, list[ChatCompletionContentPartParam]]

224 + content: str | list[ChatCompletionContentPartParam]

225 """The contents of the message."""

226

227 name: str

@@ -231,18 +231,18 @@ class CustomChatCompletionMessageParam(TypedDict, total=False):

231 same role.

232 """

233

234 - tool_call_id: Optional[str]

234 + tool_call_id: str | None

235 """Tool call that this message is responding to."""

236

237 - tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]

237 + tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None

238 """The tool calls generated by the model, such as function calls."""

239

240

241 -ChatCompletionMessageParam = Union[

242 - OpenAIChatCompletionMessageParam,

243 - CustomChatCompletionMessageParam,

244 - OpenAIHarmonyMessage,

245 -]

241 +ChatCompletionMessageParam: TypeAlias = (

242 + OpenAIChatCompletionMessageParam

243 + | CustomChatCompletionMessageParam

244 + | OpenAIHarmonyMessage

245 +)

246

247

248 # TODO: Make fields ReadOnly once mypy supports it

@@ -250,16 +250,16 @@ class ConversationMessage(TypedDict, total=False):

250 role: Required[str]

251 """The role of the message's author."""

252

253 - content: Union[Optional[str], list[dict[str, str]]]

253 + content: str | None | list[dict[str, str]]

254 """The contents of the message"""

255

256 - tool_call_id: Optional[str]

256 + tool_call_id: str | None

257 """Tool call that this message is responding to."""

258

259 - name: Optional[str]

259 + name: str | None

260 """The name of the function to call"""

261

262 - tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]

262 + tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None

263 """The tool calls generated by the model, such as function calls."""

264

265

@@ -294,7 +294,7 @@ def _is_attr_access(node: jinja2.nodes.Node, varname: str, key: str) -> bool:

294 def _is_var_or_elems_access(

295 node: jinja2.nodes.Node,

296 varname: str,

297 - key: Optional[str] = None,

297 + key: str | None = None,

298 ) -> bool:

299 if isinstance(node, jinja2.nodes.Filter):

300 return node.node is not None and _is_var_or_elems_access(

@@ -369,7 +369,7 @@ def _iter_nodes_assign_content_item(root: jinja2.nodes.Node):

369 break

370

371

372 -def _try_extract_ast(chat_template: str) -> Optional[jinja2.nodes.Template]:

372 +def _try_extract_ast(chat_template: str) -> jinja2.nodes.Template | None:

373 try:

374 jinja_compiled = hf_chat_utils._compile_jinja_template(chat_template)

375 return jinja_compiled.environment.parse(chat_template)

@@ -400,9 +400,9 @@ def _detect_content_format(

400

401

402 def resolve_mistral_chat_template(

403 - chat_template: Optional[str],

403 + chat_template: str | None,

404 **kwargs: Any,

405 -) -> Optional[str]:

405 +) -> str | None:

406 if chat_template is not None or kwargs.get("chat_template_kwargs") is not None:

407 raise ValueError(

408 "'chat_template' or 'chat_template_kwargs' cannot be overridden "

@@ -412,7 +412,7 @@ def resolve_mistral_chat_template(

412 return None

413

414

415 -_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], Optional[str]]()

415 +_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()

416 """

417 Used in `_try_get_processor_chat_template` to avoid calling

418 `cached_get_processor` again if the processor fails to be loaded.

@@ -422,9 +422,9 @@ This is needed because `lru_cache` does not cache when an exception happens.

422

423

424 def _try_get_processor_chat_template(

425 - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

425 + tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,

426 model_config: ModelConfig,

427 -) -> Optional[str]:

427 +) -> str | None:

428 cache_key = (tokenizer.name_or_path, model_config.trust_remote_code)

429 if cache_key in _PROCESSOR_CHAT_TEMPLATES:

430 return _PROCESSOR_CHAT_TEMPLATES[cache_key]

@@ -458,12 +458,12 @@ def _try_get_processor_chat_template(

458

459

460 def resolve_hf_chat_template(

461 - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

462 - chat_template: Optional[str],

463 - tools: Optional[list[dict[str, Any]]],

461 + tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,

462 + chat_template: str | None,

463 + tools: list[dict[str, Any]] | None,

464 *,

465 model_config: ModelConfig,

466 -) -> Optional[str]:

466 +) -> str | None:

467 # 1st priority: The given chat template

468 if chat_template is not None:

469 return chat_template

@@ -505,8 +505,8 @@ def resolve_hf_chat_template(

505

506

507 def _resolve_chat_template_content_format(

508 - chat_template: Optional[str],

509 - tools: Optional[list[dict[str, Any]]],

508 + chat_template: str | None,

509 + tools: list[dict[str, Any]] | None,

510 tokenizer: AnyTokenizer,

511 *,

512 model_config: ModelConfig,

@@ -538,7 +538,7 @@ def _resolve_chat_template_content_format(

538

539 @lru_cache

540 def _log_chat_template_content_format(

541 - chat_template: Optional[str],

541 + chat_template: str | None,

542 given_format: ChatTemplateContentFormatOption,

543 detected_format: ChatTemplateContentFormatOption,

544 ):

@@ -561,8 +561,8 @@ def _log_chat_template_content_format(

561

562

563 def resolve_chat_template_content_format(

564 - chat_template: Optional[str],

565 - tools: Optional[list[dict[str, Any]]],

564 + chat_template: str | None,

565 + tools: list[dict[str, Any]] | None,

566 given_format: ChatTemplateContentFormatOption,

567 tokenizer: AnyTokenizer,

568 *,

@@ -604,8 +604,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

604 self._model_config = model_config

605 self._tokenizer = tokenizer

606

607 - self._items_by_modality = defaultdict[str, list[Optional[_T]]](list)

608 - self._uuids_by_modality = defaultdict[str, list[Optional[str]]](list)

607 + self._items_by_modality = defaultdict[str, list[_T | None]](list)

608 + self._uuids_by_modality = defaultdict[str, list[str | None]](list)

609

610 @property

611 def model_config(self) -> ModelConfig:

@@ -637,9 +637,9 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

637 def add(

638 self,

639 modality: ModalityStr,

640 - item: Optional[_T],

641 - uuid: Optional[str] = None,

642 - ) -> Optional[str]:

640 + item: _T | None,

641 + uuid: str | None = None,

642 + ) -> str | None:

643 """

644 Add a multi-modal item to the current prompt and returns the

645 placeholder string to use, if any.

@@ -657,7 +657,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

657

658 return self.model_cls.get_placeholder_str(modality, num_items)

659

660 - def all_mm_uuids(self) -> Optional[MultiModalUUIDDict]:

660 + def all_mm_uuids(self) -> MultiModalUUIDDict | None:

661 if not self._items_by_modality:

662 return None

663 mm_uuids = {}

@@ -684,7 +684,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

684

685

686 class MultiModalItemTracker(BaseMultiModalItemTracker[object]):

687 - def all_mm_data(self) -> Optional[MultiModalDataDict]:

687 + def all_mm_data(self) -> MultiModalDataDict | None:

688 if not self._items_by_modality:

689 return None

690 mm_inputs = {}

@@ -710,7 +710,7 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[object]):

710

711

712 class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):

713 - async def all_mm_data(self) -> Optional[MultiModalDataDict]:

713 + async def all_mm_data(self) -> MultiModalDataDict | None:

714 if not self._items_by_modality:

715 return None

716 mm_inputs = {}

@@ -756,7 +756,7 @@ class BaseMultiModalContentParser(ABC):

756 # }

757 self._placeholder_storage: dict[str, list] = defaultdict(list)

758

759 - def _add_placeholder(self, modality: ModalityStr, placeholder: Optional[str]):

759 + def _add_placeholder(self, modality: ModalityStr, placeholder: str | None):

760 mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]

761 if placeholder:

762 self._placeholder_storage[mod_placeholder].append(placeholder)

@@ -765,35 +765,35 @@ class BaseMultiModalContentParser(ABC):

765 return dict(self._placeholder_storage)

766

767 @abstractmethod

768 - def parse_image(self, image_url: Optional[str], uuid: Optional[str] = None) -> None:

768 + def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:

769 raise NotImplementedError

770

771 @abstractmethod

772 def parse_image_embeds(

773 self,

774 - image_embeds: Union[str, dict[str, str], None],

775 - uuid: Optional[str] = None,

774 + image_embeds: str | dict[str, str] | None,

775 + uuid: str | None = None,

776 ) -> None:

777 raise NotImplementedError

778

779 @abstractmethod

780 def parse_image_pil(

781 - self, image_pil: Optional[Image.Image], uuid: Optional[str] = None

781 + self, image_pil: Image.Image | None, uuid: str | None = None

782 ) -> None:

783 raise NotImplementedError

784

785 @abstractmethod

786 - def parse_audio(self, audio_url: Optional[str], uuid: Optional[str] = None) -> None:

786 + def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:

787 raise NotImplementedError

788

789 @abstractmethod

790 def parse_input_audio(

791 - self, input_audio: Optional[InputAudio], uuid: Optional[str] = None

791 + self, input_audio: InputAudio | None, uuid: str | None = None

792 ) -> None:

793 raise NotImplementedError

794

795 @abstractmethod

796 - def parse_video(self, video_url: Optional[str], uuid: Optional[str] = None) -> None:

796 + def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:

797 raise NotImplementedError

798

799

@@ -810,7 +810,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

810 allowed_media_domains=tracker.allowed_media_domains,

811 )

812

813 - def parse_image(self, image_url: Optional[str], uuid: Optional[str] = None) -> None:

813 + def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:

814 image = self._connector.fetch_image(image_url) if image_url else None

815

816 placeholder = self._tracker.add("image", image, uuid)

@@ -818,8 +818,8 @@ class MultiModalContentParser(BaseMultiModalContentParser):

818

819 def parse_image_embeds(

820 self,

821 - image_embeds: Union[str, dict[str, str], None],

822 - uuid: Optional[str] = None,

821 + image_embeds: str | dict[str, str] | None,

822 + uuid: str | None = None,

823 ) -> None:

824 if isinstance(image_embeds, dict):

825 embeds = {

@@ -838,19 +838,19 @@ class MultiModalContentParser(BaseMultiModalContentParser):

838 self._add_placeholder("image", placeholder)

839

840 def parse_image_pil(

841 - self, image_pil: Optional[Image.Image], uuid: Optional[str] = None

841 + self, image_pil: Image.Image | None, uuid: str | None = None

842 ) -> None:

843 placeholder = self._tracker.add("image", image_pil, uuid)

844 self._add_placeholder("image", placeholder)

845

846 - def parse_audio(self, audio_url: Optional[str], uuid: Optional[str] = None) -> None:

846 + def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:

847 audio = self._connector.fetch_audio(audio_url) if audio_url else None

848

849 placeholder = self._tracker.add("audio", audio, uuid)

850 self._add_placeholder("audio", placeholder)

851

852 def parse_input_audio(

853 - self, input_audio: Optional[InputAudio], uuid: Optional[str] = None

853 + self, input_audio: InputAudio | None, uuid: str | None = None

854 ) -> None:

855 if input_audio:

856 audio_data = input_audio.get("data", "")

@@ -865,7 +865,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

865

866 return self.parse_audio(audio_url, uuid)

867

868 - def parse_video(self, video_url: Optional[str], uuid: Optional[str] = None) -> None:

868 + def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:

869 video = self._connector.fetch_video(video_url=video_url) if video_url else None

870

871 placeholder = self._tracker.add("video", video, uuid)

@@ -885,7 +885,7 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

885 allowed_media_domains=tracker.allowed_media_domains,

886 )

887

888 - def parse_image(self, image_url: Optional[str], uuid: Optional[str] = None) -> None:

888 + def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:

889 image_coro = self._connector.fetch_image_async(image_url) if image_url else None

890

891 placeholder = self._tracker.add("image", image_coro, uuid)

@@ -893,10 +893,10 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

893

894 def parse_image_embeds(

895 self,

896 - image_embeds: Union[str, dict[str, str], None],

897 - uuid: Optional[str] = None,

896 + image_embeds: str | dict[str, str] | None,

897 + uuid: str | None = None,

898 ) -> None:

899 - future: asyncio.Future[Union[str, dict[str, str], None]] = asyncio.Future()

899 + future: asyncio.Future[str | dict[str, str] | None] = asyncio.Future()

900

901 if isinstance(image_embeds, dict):

902 embeds = {

@@ -916,9 +916,9 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

916 self._add_placeholder("image", placeholder)

917

918 def parse_image_pil(

919 - self, image_pil: Optional[Image.Image], uuid: Optional[str] = None

919 + self, image_pil: Image.Image | None, uuid: str | None = None

920 ) -> None:

921 - future: asyncio.Future[Optional[Image.Image]] = asyncio.Future()

921 + future: asyncio.Future[Image.Image | None] = asyncio.Future()

922 if image_pil:

923 future.set_result(image_pil)

924 else:

@@ -927,14 +927,14 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

927 placeholder = self._tracker.add("image", future, uuid)

928 self._add_placeholder("image", placeholder)

929

930 - def parse_audio(self, audio_url: Optional[str], uuid: Optional[str] = None) -> None:

930 + def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:

931 audio_coro = self._connector.fetch_audio_async(audio_url) if audio_url else None

932

933 placeholder = self._tracker.add("audio", audio_coro, uuid)

934 self._add_placeholder("audio", placeholder)

935

936 def parse_input_audio(

937 - self, input_audio: Optional[InputAudio], uuid: Optional[str] = None

937 + self, input_audio: InputAudio | None, uuid: str | None = None

938 ) -> None:

939 if input_audio:

940 audio_data = input_audio.get("data", "")

@@ -949,7 +949,7 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

949

950 return self.parse_audio(audio_url, uuid)

951

952 - def parse_video(self, video_url: Optional[str], uuid: Optional[str] = None) -> None:

952 + def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:

953 video = (

954 self._connector.fetch_video_async(video_url=video_url)

955 if video_url

@@ -960,7 +960,7 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

960 self._add_placeholder("video", placeholder)

961

962

963 -def validate_chat_template(chat_template: Optional[Union[Path, str]]):

963 +def validate_chat_template(chat_template: Path | str | None):

964 """Raises if the provided chat template appears invalid."""

965 if chat_template is None:

966 return

@@ -984,10 +984,10 @@ def validate_chat_template(chat_template: Optional[Union[Path, str]]):

984

985

986 def _load_chat_template(

987 - chat_template: Optional[Union[Path, str]],

987 + chat_template: Path | str | None,

988 *,

989 is_literal: bool = False,

990 -) -> Optional[str]:

990 +) -> str | None:

991 if chat_template is None:

992 return None

993

@@ -1024,10 +1024,10 @@ _cached_load_chat_template = lru_cache(_load_chat_template)

1024

1025

1026 def load_chat_template(

1027 - chat_template: Optional[Union[Path, str]],

1027 + chat_template: Path | str | None,

1028 *,

1029 is_literal: bool = False,

1030 -) -> Optional[str]:

1030 +) -> str | None:

1031 return _cached_load_chat_template(chat_template, is_literal=is_literal)

1032

1033

@@ -1107,7 +1107,7 @@ _AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python

1107 _VideoParser = TypeAdapter(ChatCompletionContentPartVideoParam).validate_python

1108

1109 _ResponsesInputImageParser = TypeAdapter(ResponseInputImageParam).validate_python

1110 -_ContentPart: TypeAlias = Union[str, dict[str, str], InputAudio, PILImage]

1110 +_ContentPart: TypeAlias = str | dict[str, str] | InputAudio | PILImage

1111

1112 # Define a mapping from part types to their corresponding parsing functions.

1113 MM_PARSER_MAP: dict[

@@ -1264,7 +1264,7 @@ def _parse_chat_message_content_part(

1264 *,

1265 wrap_dicts: bool,

1266 interleave_strings: bool,

1267 -) -> Optional[_ContentPart]:

1267 +) -> _ContentPart | None:

1268 """Parses a single part of a conversation. If wrap_dicts is True,

1269 structured dictionary pieces for texts and images will be

1270 wrapped in dictionaries, i.e., {"type": "text", "text", ...} and

@@ -1310,10 +1310,7 @@ def _parse_chat_message_content_part(

1310 mm_parser.parse_image(str_content, uuid)

1311 modality = "image"

1312 elif part_type == "image_embeds":

1313 - if content is not None:

1314 - content = cast(Union[str, dict[str, str]], content)

1315 - else:

1316 - content = None

1313 + content = cast(str | dict[str, str], content) if content is not None else None

1314 mm_parser.parse_image_embeds(content, uuid)

1315 modality = "image"

1316 elif part_type == "audio_url":

@@ -1411,8 +1408,8 @@ def parse_chat_messages(

1408 content_format: _ChatTemplateContentFormat,

1409 ) -> tuple[

1410 list[ConversationMessage],

1414 - Optional[MultiModalDataDict],

1415 - Optional[MultiModalUUIDDict],

1411 + MultiModalDataDict | None,

1412 + MultiModalUUIDDict | None,

1413 ]:

1414 conversation: list[ConversationMessage] = []

1415 mm_tracker = MultiModalItemTracker(model_config, tokenizer)

@@ -1443,8 +1440,8 @@ def parse_chat_messages_futures(

1440 content_format: _ChatTemplateContentFormat,

1441 ) -> tuple[

1442 list[ConversationMessage],

1446 - Awaitable[Optional[MultiModalDataDict]],

1447 - Optional[MultiModalUUIDDict],

1443 + Awaitable[MultiModalDataDict | None],

1444 + MultiModalUUIDDict | None,

1445 ]:

1446 conversation: list[ConversationMessage] = []

1447 mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)

@@ -1498,7 +1495,7 @@ _cached_resolve_chat_template_kwargs = lru_cache(_resolve_chat_template_kwargs)

1495

1496

1497 def resolve_chat_template_kwargs(

1501 - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

1498 + tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,

1499 chat_template: str,

1500 chat_template_kwargs: dict[str, Any],

1501 ) -> dict[str, Any]:

@@ -1518,10 +1515,10 @@ def resolve_chat_template_kwargs(

1515

1516

1517 def apply_hf_chat_template(

1521 - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

1518 + tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,

1519 conversation: list[ConversationMessage],

1523 - chat_template: Optional[str],

1524 - tools: Optional[list[dict[str, Any]]],

1520 + chat_template: str | None,

1521 + tools: list[dict[str, Any]] | None,

1522 *,

1523 model_config: ModelConfig,

1524 tokenize: bool = False, # Different from HF's default

@@ -1569,8 +1566,8 @@ def apply_hf_chat_template(

1566 def apply_mistral_chat_template(

1567 tokenizer: MistralTokenizer,

1568 messages: list[ChatCompletionMessageParam],

1572 - chat_template: Optional[str],

1573 - tools: Optional[list[dict[str, Any]]],

1569 + chat_template: str | None,

1570 + tools: list[dict[str, Any]] | None,

1571 **kwargs: Any,

1572 ) -> list[int]:

1573 from mistral_common.exceptions import MistralCommonException

vllm/entrypoints/cli/benchmark/main.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import argparse

5 import typing

6

@@ -12,6 +10,8 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG

10

11 if typing.TYPE_CHECKING:

12 from vllm.utils import FlexibleArgumentParser

13 +else:

14 + FlexibleArgumentParser = argparse.ArgumentParser

15

16

17 class BenchmarkSubcommand(CLISubcommand):

vllm/entrypoints/cli/collect_env.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import argparse

5 import typing

6

@@ -11,6 +9,8 @@ from vllm.entrypoints.cli.types import CLISubcommand

9

10 if typing.TYPE_CHECKING:

11 from vllm.utils import FlexibleArgumentParser

12 +else:

13 + FlexibleArgumentParser = argparse.ArgumentParser

14

15

16 class CollectEnvSubcommand(CLISubcommand):

vllm/entrypoints/cli/main.py

-2

@@ -5,8 +5,6 @@

5 Note that all future modules must be lazily loaded within main

6 to avoid certain eager import breakage."""

7

8 -from __future__ import annotations

9 -

8 import importlib.metadata

9 import sys

10

vllm/entrypoints/cli/openai.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import argparse

5 import os

6 import signal

@@ -16,6 +14,8 @@ from vllm.entrypoints.cli.types import CLISubcommand

14

15 if TYPE_CHECKING:

16 from vllm.utils import FlexibleArgumentParser

17 +else:

18 + FlexibleArgumentParser = argparse.ArgumentParser

19

20

21 def _register_signal_handlers():

vllm/entrypoints/cli/run_batch.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import argparse

5 import asyncio

6 import importlib.metadata

@@ -14,6 +12,8 @@ from vllm.logger import init_logger

12

13 if typing.TYPE_CHECKING:

14 from vllm.utils import FlexibleArgumentParser

15 +else:

16 + FlexibleArgumentParser = argparse.ArgumentParser

17

18 logger = init_logger(__name__)

19

vllm/entrypoints/cli/serve.py

+1 -2

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

3

4 import argparse

5 import signal

6 -from typing import Optional

6

7 import uvloop

8

@@ -179,7 +178,7 @@ def run_multi_api_server(args: argparse.Namespace):

178 hybrid_dp_lb = parallel_config.data_parallel_hybrid_lb

179 assert external_dp_lb or hybrid_dp_lb or dp_rank == 0

180

182 - api_server_manager: Optional[APIServerProcessManager] = None

181 + api_server_manager: APIServerProcessManager | None = None

182

183 with launch_core_engines(

184 vllm_config, executor_class, log_stats, num_api_servers

vllm/entrypoints/cli/types.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import argparse

5 import typing

6

7 if typing.TYPE_CHECKING:

8 from vllm.utils import FlexibleArgumentParser

9 +else:

10 + FlexibleArgumentParser = argparse.ArgumentParser

11

12

13 class CLISubcommand:

vllm/entrypoints/context.py

+8 -8

@@ -6,7 +6,7 @@ import json

6 import logging

7 from abc import ABC, abstractmethod

8 from contextlib import AsyncExitStack

9 -from typing import TYPE_CHECKING, Optional, Union

9 +from typing import TYPE_CHECKING, Union

10

11 from openai.types.responses.tool import Mcp

12 from openai_harmony import Author, Message, Role, StreamState, TextContent

@@ -82,7 +82,7 @@ class ConversationContext(ABC):

82 @abstractmethod

83 async def init_tool_sessions(

84 self,

85 - tool_server: Optional[ToolServer],

85 + tool_server: ToolServer | None,

86 exit_stack: AsyncExitStack,

87 request_id: str,

88 mcp_tools: dict[str, Mcp],

@@ -122,7 +122,7 @@ class SimpleContext(ConversationContext):

122

123 async def init_tool_sessions(

124 self,

125 - tool_server: Optional[ToolServer],

125 + tool_server: ToolServer | None,

126 exit_stack: AsyncExitStack,

127 request_id: str,

128 mcp_tools: dict[str, Mcp],

@@ -140,9 +140,9 @@ class HarmonyContext(ConversationContext):

140 available_tools: list[str],

141 ):

142 self._messages = messages

143 - self.finish_reason: Optional[str] = None

143 + self.finish_reason: str | None = None

144 self.available_tools = available_tools

145 - self._tool_sessions: dict[str, Union[ClientSession, Tool]] = {}

145 + self._tool_sessions: dict[str, ClientSession | Tool] = {}

146 self.called_tools: set[str] = set()

147

148 self.parser = get_streamable_parser_for_assistant()

@@ -164,7 +164,7 @@ class HarmonyContext(ConversationContext):

164 if self.parser.current_channel in {"analysis", "commentary"}:

165 self.num_reasoning_tokens += 1

166

167 - def append_output(self, output: Union[RequestOutput, list[Message]]) -> None:

167 + def append_output(self, output: RequestOutput | list[Message]) -> None:

168 if isinstance(output, RequestOutput):

169 output_token_ids = output.outputs[0].token_ids

170 self.parser = get_streamable_parser_for_assistant()

@@ -358,7 +358,7 @@ class HarmonyContext(ConversationContext):

358

359 async def init_tool_sessions(

360 self,

361 - tool_server: Optional[ToolServer],

361 + tool_server: ToolServer | None,

362 exit_stack: AsyncExitStack,

363 request_id: str,

364 mcp_tools: dict[str, Mcp],

@@ -446,7 +446,7 @@ class StreamingHarmonyContext(HarmonyContext):

446 def messages(self) -> list:

447 return self._messages

448

449 - def append_output(self, output: Union[RequestOutput, list[Message]]) -> None:

449 + def append_output(self, output: RequestOutput | list[Message]) -> None:

450 if isinstance(output, RequestOutput):

451 # append_output is called for each output token in streaming case,

452 # so we only want to add the prompt tokens once for each message.

vllm/entrypoints/harmony_utils.py

+5 -7

@@ -1,12 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from __future__ import annotations

5 -

4 import datetime

5 import json

6 from collections.abc import Iterable, Sequence

9 -from typing import Literal, Union

7 +from typing import Literal

8

9 from openai.types.responses import (

10 ResponseFunctionToolCall,

@@ -122,7 +120,7 @@ def get_system_message(

120 return sys_msg

121

122

125 -def create_tool_definition(tool: Union[ChatCompletionToolsParam, Tool]):

123 +def create_tool_definition(tool: ChatCompletionToolsParam | Tool):

124 if isinstance(tool, ChatCompletionToolsParam):

125 return ToolDescription.new(

126 name=tool.function.name,

@@ -138,13 +136,13 @@ def create_tool_definition(tool: Union[ChatCompletionToolsParam, Tool]):

136

137 def get_developer_message(

138 instructions: str | None = None,

141 - tools: list[Union[Tool, ChatCompletionToolsParam]] | None = None,

139 + tools: list[Tool | ChatCompletionToolsParam] | None = None,

140 ) -> Message:

141 dev_msg_content = DeveloperContent.new()

142 if instructions is not None and not envs.VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS:

143 dev_msg_content = dev_msg_content.with_instructions(instructions)

144 if tools is not None:

147 - function_tools: list[Union[Tool, ChatCompletionToolsParam]] = []

145 + function_tools: list[Tool | ChatCompletionToolsParam] = []

146 for tool in tools:

147 if tool.type in (

148 "web_search_preview",

@@ -178,7 +176,7 @@ def get_user_message(content: str) -> Message:

176

177 def parse_response_input(

178 response_msg: ResponseInputOutputItem,

181 - prev_responses: list[Union[ResponseOutputItem, ResponseReasoningItem]],

179 + prev_responses: list[ResponseOutputItem | ResponseReasoningItem],

180 ) -> Message:

181 if not isinstance(response_msg, dict):

182 response_msg = response_msg.model_dump()

vllm/entrypoints/launcher.py

+2 -2

@@ -5,7 +5,7 @@ import asyncio

5 import signal

6 import socket

7 from http import HTTPStatus

8 -from typing import Any, Optional

8 +from typing import Any

9

10 import uvicorn

11 from fastapi import FastAPI, Request, Response

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

26

27 async def serve_http(

28 app: FastAPI,

29 - sock: Optional[socket.socket],

29 + sock: socket.socket | None,

30 enable_ssl_refresh: bool = False,

31 **uvicorn_kwargs: Any,

32 ):

vllm/entrypoints/llm.py

+113 -121

@@ -2,8 +2,8 @@

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

3

4 import itertools

5 -from collections.abc import Sequence

6 -from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast

5 +from collections.abc import Callable, Sequence

6 +from typing import TYPE_CHECKING, Any, cast

7

8 import cloudpickle

9 import torch.nn as nn

@@ -191,36 +191,34 @@ class LLM:

191 *,

192 runner: RunnerOption = "auto",

193 convert: ConvertOption = "auto",

194 - tokenizer: Optional[str] = None,

194 + tokenizer: str | None = None,

195 tokenizer_mode: TokenizerMode = "auto",

196 skip_tokenizer_init: bool = False,

197 trust_remote_code: bool = False,

198 allowed_local_media_path: str = "",

199 - allowed_media_domains: Optional[list[str]] = None,

199 + allowed_media_domains: list[str] | None = None,

200 tensor_parallel_size: int = 1,

201 dtype: ModelDType = "auto",

202 - quantization: Optional[QuantizationMethods] = None,

203 - revision: Optional[str] = None,

204 - tokenizer_revision: Optional[str] = None,

205 - seed: Optional[int] = None,

202 + quantization: QuantizationMethods | None = None,

203 + revision: str | None = None,

204 + tokenizer_revision: str | None = None,

205 + seed: int | None = None,

206 gpu_memory_utilization: float = 0.9,

207 swap_space: float = 4,

208 cpu_offload_gb: float = 0,

209 enforce_eager: bool = False,

210 disable_custom_all_reduce: bool = False,

211 - hf_token: Optional[Union[bool, str]] = None,

212 - hf_overrides: Optional[HfOverrides] = None,

213 - mm_processor_kwargs: Optional[dict[str, Any]] = None,

214 - pooler_config: Optional[PoolerConfig] = None,

215 - override_pooler_config: Optional[PoolerConfig] = None,

216 - structured_outputs_config: Optional[

217 - Union[dict[str, Any], StructuredOutputsConfig]

218 - ] = None,

219 - kv_cache_memory_bytes: Optional[int] = None,

220 - compilation_config: Optional[

221 - Union[int, dict[str, Any], CompilationConfig]

222 - ] = None,

223 - logits_processors: Optional[list[Union[str, type[LogitsProcessor]]]] = None,

211 + hf_token: bool | str | None = None,

212 + hf_overrides: HfOverrides | None = None,

213 + mm_processor_kwargs: dict[str, Any] | None = None,

214 + pooler_config: PoolerConfig | None = None,

215 + override_pooler_config: PoolerConfig | None = None,

216 + structured_outputs_config: dict[str, Any]

217 + | StructuredOutputsConfig

218 + | None = None,

219 + kv_cache_memory_bytes: int | None = None,

220 + compilation_config: int | dict[str, Any] | CompilationConfig | None = None,

221 + logits_processors: list[str | type[LogitsProcessor]] | None = None,

222 **kwargs: Any,

223 ) -> None:

224 """LLM constructor."""

@@ -331,7 +329,7 @@ class LLM:

329 self.engine_class = type(self.llm_engine)

330

331 self.request_counter = Counter()

334 - self.default_sampling_params: Union[dict[str, Any], None] = None

332 + self.default_sampling_params: dict[str, Any] | None = None

333

334 supported_tasks = self.llm_engine.get_supported_tasks()

335 logger.info("Supported tasks: %s", supported_tasks)

@@ -367,14 +365,12 @@ class LLM:

365

366 def generate(

367 self,

370 - prompts: Union[PromptType, Sequence[PromptType]],

371 - sampling_params: Optional[

372 - Union[SamplingParams, Sequence[SamplingParams]]

373 - ] = None,

368 + prompts: PromptType | Sequence[PromptType],

369 + sampling_params: SamplingParams | Sequence[SamplingParams] | None = None,

370 *,

375 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

376 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

377 - priority: Optional[list[int]] = None,

371 + use_tqdm: bool | Callable[..., tqdm] = True,

372 + lora_request: list[LoRARequest] | LoRARequest | None = None,

373 + priority: list[int] | None = None,

374 ) -> list[RequestOutput]:

375 """Generates the completions for the input prompts.

376

@@ -437,8 +433,8 @@ class LLM:

433

434 def _get_modality_specific_lora_reqs(

435 self,

440 - prompts: Union[PromptType, Sequence[PromptType]],

441 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]],

436 + prompts: PromptType | Sequence[PromptType],

437 + lora_request: list[LoRARequest] | LoRARequest | None,

438 ):

439 # Grab the lora config off the vllm config on the engine,

440 # since this is the same for both v0 & v1.

@@ -474,8 +470,8 @@ class LLM:

470 def _resolve_single_prompt_mm_lora(

471 self,

472 prompt: PromptType,

477 - lora_request: Optional[LoRARequest],

478 - default_mm_loras: Optional[dict[str, str]],

473 + lora_request: LoRARequest | None,

474 + default_mm_loras: dict[str, str] | None,

475 ):

476 if (

477 not default_mm_loras

@@ -526,10 +522,10 @@ class LLM:

522

523 def collective_rpc(

524 self,

529 - method: Union[str, Callable[..., _R]],

530 - timeout: Optional[float] = None,

525 + method: str | Callable[..., _R],

526 + timeout: float | None = None,

527 args: tuple = (),

532 - kwargs: Optional[dict[str, Any]] = None,

528 + kwargs: dict[str, Any] | None = None,

529 ) -> list[_R]:

530 """

531 Execute an RPC call on all workers.

@@ -571,9 +567,9 @@ class LLM:

567

568 def _get_beam_search_lora_requests(

569 self,

574 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]],

575 - prompts: list[Union[TokensPrompt, TextPrompt]],

576 - ) -> list[Optional[LoRARequest]]:

570 + lora_request: list[LoRARequest] | LoRARequest | None,

571 + prompts: list[TokensPrompt | TextPrompt],

572 + ) -> list[LoRARequest | None]:

573 """Get the optional lora request corresponding to each prompt."""

574 if isinstance(lora_request, Sequence) and len(lora_request) != len(prompts):

575 raise ValueError(

@@ -587,11 +583,11 @@ class LLM:

583

584 def beam_search(

585 self,

590 - prompts: list[Union[TokensPrompt, TextPrompt]],

586 + prompts: list[TokensPrompt | TextPrompt],

587 params: BeamSearchParams,

592 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

588 + lora_request: list[LoRARequest] | LoRARequest | None = None,

589 use_tqdm: bool = False,

594 - concurrency_limit: Optional[int] = None,

590 + concurrency_limit: int | None = None,

591 ) -> list[BeamSearchOutput]:

592 """

593 Generate sequences using beam search.

@@ -770,16 +766,15 @@ class LLM:

766

767 def preprocess_chat(

768 self,

773 - messages: Union[

774 - list[ChatCompletionMessageParam], list[list[ChatCompletionMessageParam]]

775 - ],

776 - chat_template: Optional[str] = None,

769 + messages: list[ChatCompletionMessageParam]

770 + | list[list[ChatCompletionMessageParam]],

771 + chat_template: str | None = None,

772 chat_template_content_format: ChatTemplateContentFormatOption = "auto",

773 add_generation_prompt: bool = True,

774 continue_final_message: bool = False,

780 - tools: Optional[list[dict[str, Any]]] = None,

781 - chat_template_kwargs: Optional[dict[str, Any]] = None,

782 - mm_processor_kwargs: Optional[dict[str, Any]] = None,

775 + tools: list[dict[str, Any]] | None = None,

776 + chat_template_kwargs: dict[str, Any] | None = None,

777 + mm_processor_kwargs: dict[str, Any] | None = None,

778 ) -> list[TokensPrompt]:

779 """

780 Generate prompt for a chat conversation. The pre-processed

@@ -868,19 +863,18 @@ class LLM:

863

864 def chat(

865 self,

871 - messages: Union[

872 - list[ChatCompletionMessageParam], list[list[ChatCompletionMessageParam]]

873 - ],

874 - sampling_params: Optional[Union[SamplingParams, list[SamplingParams]]] = None,

875 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

876 - lora_request: Optional[LoRARequest] = None,

877 - chat_template: Optional[str] = None,

866 + messages: list[ChatCompletionMessageParam]

867 + | list[list[ChatCompletionMessageParam]],

868 + sampling_params: SamplingParams | list[SamplingParams] | None = None,

869 + use_tqdm: bool | Callable[..., tqdm] = True,

870 + lora_request: LoRARequest | None = None,

871 + chat_template: str | None = None,

872 chat_template_content_format: ChatTemplateContentFormatOption = "auto",

873 add_generation_prompt: bool = True,

874 continue_final_message: bool = False,

881 - tools: Optional[list[dict[str, Any]]] = None,

882 - chat_template_kwargs: Optional[dict[str, Any]] = None,

883 - mm_processor_kwargs: Optional[dict[str, Any]] = None,

875 + tools: list[dict[str, Any]] | None = None,

876 + chat_template_kwargs: dict[str, Any] | None = None,

877 + mm_processor_kwargs: dict[str, Any] | None = None,

878 ) -> list[RequestOutput]:

879 """

880 Generate responses for a chat conversation.

@@ -953,14 +947,14 @@ class LLM:

947

948 def encode(

949 self,

956 - prompts: Union[PromptType, Sequence[PromptType], DataPrompt],

957 - pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,

950 + prompts: PromptType | Sequence[PromptType] | DataPrompt,

951 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,

952 *,

959 - truncate_prompt_tokens: Optional[int] = None,

960 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

961 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

953 + truncate_prompt_tokens: int | None = None,

954 + use_tqdm: bool | Callable[..., tqdm] = True,

955 + lora_request: list[LoRARequest] | LoRARequest | None = None,

956 pooling_task: PoolingTask = "encode",

963 - tokenization_kwargs: Optional[dict[str, Any]] = None,

957 + tokenization_kwargs: dict[str, Any] | None = None,

958 ) -> list[PoolingRequestOutput]:

959 """Apply pooling to the hidden states corresponding to the input

960 prompts.

@@ -1086,12 +1080,12 @@ class LLM:

1080

1081 def embed(

1082 self,

1089 - prompts: Union[PromptType, Sequence[PromptType]],

1083 + prompts: PromptType | Sequence[PromptType],

1084 *,

1091 - truncate_prompt_tokens: Optional[int] = None,

1092 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1093 - pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,

1094 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1085 + truncate_prompt_tokens: int | None = None,

1086 + use_tqdm: bool | Callable[..., tqdm] = True,

1087 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,

1088 + lora_request: list[LoRARequest] | LoRARequest | None = None,

1089 ) -> list[EmbeddingRequestOutput]:

1090 """

1091 Generate an embedding vector for each prompt.

@@ -1135,11 +1129,11 @@ class LLM:

1129

1130 def classify(

1131 self,

1138 - prompts: Union[PromptType, Sequence[PromptType]],

1132 + prompts: PromptType | Sequence[PromptType],

1133 *,

1140 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1141 - pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,

1142 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1134 + use_tqdm: bool | Callable[..., tqdm] = True,

1135 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,

1136 + lora_request: list[LoRARequest] | LoRARequest | None = None,

1137 ) -> list[ClassificationRequestOutput]:

1138 """

1139 Generate class logits for each prompt.

@@ -1181,13 +1175,13 @@ class LLM:

1175

1176 def reward(

1177 self,

1184 - prompts: Union[PromptType, Sequence[PromptType]],

1178 + prompts: PromptType | Sequence[PromptType],

1179 /,

1180 *,

1187 - truncate_prompt_tokens: Optional[int] = None,

1188 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1189 - pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,

1190 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1181 + truncate_prompt_tokens: int | None = None,

1182 + use_tqdm: bool | Callable[..., tqdm] = True,

1183 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,

1184 + lora_request: list[LoRARequest] | LoRARequest | None = None,

1185 ) -> list[PoolingRequestOutput]:

1186 """

1187 Generate rewards for each prompt.

@@ -1220,12 +1214,12 @@ class LLM:

1214 def _embedding_score(

1215 self,

1216 tokenizer: AnyTokenizer,

1223 - text_1: list[Union[str, TextPrompt, TokensPrompt]],

1224 - text_2: list[Union[str, TextPrompt, TokensPrompt]],

1225 - truncate_prompt_tokens: Optional[int] = None,

1226 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1227 - pooling_params: Optional[PoolingParams] = None,

1228 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1217 + text_1: list[str | TextPrompt | TokensPrompt],

1218 + text_2: list[str | TextPrompt | TokensPrompt],

1219 + truncate_prompt_tokens: int | None = None,

1220 + use_tqdm: bool | Callable[..., tqdm] = True,

1221 + pooling_params: PoolingParams | None = None,

1222 + lora_request: list[LoRARequest] | LoRARequest | None = None,

1223 ) -> list[ScoringRequestOutput]:

1224 encoded_output: list[PoolingRequestOutput] = self.encode(

1225 text_1 + text_2,

@@ -1252,12 +1246,12 @@ class LLM:

1246 def _cross_encoding_score(

1247 self,

1248 tokenizer: AnyTokenizer,

1255 - data_1: Union[list[str], list[ScoreContentPartParam]],

1256 - data_2: Union[list[str], list[ScoreContentPartParam]],

1257 - truncate_prompt_tokens: Optional[int] = None,

1258 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1259 - pooling_params: Optional[PoolingParams] = None,

1260 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1249 + data_1: list[str] | list[ScoreContentPartParam],

1250 + data_2: list[str] | list[ScoreContentPartParam],

1251 + truncate_prompt_tokens: int | None = None,

1252 + use_tqdm: bool | Callable[..., tqdm] = True,

1253 + pooling_params: PoolingParams | None = None,

1254 + lora_request: list[LoRARequest] | LoRARequest | None = None,

1255 ) -> list[ScoringRequestOutput]:

1256 model_config = self.model_config

1257

@@ -1316,14 +1310,14 @@ class LLM:

1310

1311 def score(

1312 self,

1319 - data_1: Union[SingletonPrompt, Sequence[SingletonPrompt], ScoreMultiModalParam],

1320 - data_2: Union[SingletonPrompt, Sequence[SingletonPrompt], ScoreMultiModalParam],

1313 + data_1: SingletonPrompt | Sequence[SingletonPrompt] | ScoreMultiModalParam,

1314 + data_2: SingletonPrompt | Sequence[SingletonPrompt] | ScoreMultiModalParam,

1315 /,

1316 *,

1323 - truncate_prompt_tokens: Optional[int] = None,

1324 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1325 - pooling_params: Optional[PoolingParams] = None,

1326 - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1317 + truncate_prompt_tokens: int | None = None,

1318 + use_tqdm: bool | Callable[..., tqdm] = True,

1319 + pooling_params: PoolingParams | None = None,

1320 + lora_request: list[LoRARequest] | LoRARequest | None = None,

1321 ) -> list[ScoringRequestOutput]:

1322 """Generate similarity scores for all pairs `<text,text_pair>` or

1323 `<multi-modal data, multi-modal data pair>`.

@@ -1391,9 +1385,9 @@ class LLM:

1385 if not model_config.is_multimodal_model:

1386

1387 def check_data_type(

1394 - data: Union[

1395 - SingletonPrompt, Sequence[SingletonPrompt], ScoreMultiModalParam

1396 - ],

1388 + data: SingletonPrompt

1389 + | Sequence[SingletonPrompt]

1390 + | ScoreMultiModalParam,

1391 ):

1392 if isinstance(data, dict) and "content" in data:

1393 raise ValueError(

@@ -1470,7 +1464,7 @@ class LLM:

1464 def stop_profile(self) -> None:

1465 self.llm_engine.stop_profile()

1466

1473 - def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

1467 + def reset_prefix_cache(self, device: Device | None = None) -> bool:

1468 return self.llm_engine.reset_prefix_cache(device)

1469

1470 def sleep(self, level: int = 1):

@@ -1495,7 +1489,7 @@ class LLM:

1489 self.reset_prefix_cache()

1490 self.llm_engine.sleep(level=level)

1491

1498 - def wake_up(self, tags: Optional[list[str]] = None):

1492 + def wake_up(self, tags: list[str] | None = None):

1493 """

1494 Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]

1495 method for more details.

@@ -1523,17 +1517,15 @@ class LLM:

1517

1518 def _validate_and_add_requests(

1519 self,

1526 - prompts: Union[PromptType, Sequence[PromptType], DataPrompt],

1527 - params: Union[

1528 - SamplingParams,

1529 - Sequence[SamplingParams],

1530 - PoolingParams,

1531 - Sequence[PoolingParams],

1532 - ],

1520 + prompts: PromptType | Sequence[PromptType] | DataPrompt,

1521 + params: SamplingParams

1522 + | Sequence[SamplingParams]

1523 + | PoolingParams

1524 + | Sequence[PoolingParams],

1525 *,

1534 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1535 - lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],

1536 - priority: Optional[list[int]] = None,

1526 + use_tqdm: bool | Callable[..., tqdm] = True,

1527 + lora_request: Sequence[LoRARequest] | LoRARequest | None,

1528 + priority: list[int] | None = None,

1529 ) -> None:

1530 if isinstance(prompts, (str, dict)):

1531 # Convert a single prompt to a list.

@@ -1575,8 +1567,8 @@ class LLM:

1567

1568 def _validate_mm_data_and_uuids(

1569 self,

1578 - multi_modal_data: Optional[Any], # MultiModalDataDict

1579 - multi_modal_uuids: Optional[Any], # MultiModalUUIDDict

1570 + multi_modal_data: Any | None, # MultiModalDataDict

1571 + multi_modal_uuids: Any | None, # MultiModalUUIDDict

1572 ):

1573 """

1574 Validate that if any multi-modal data is skipped (i.e. None),

@@ -1625,9 +1617,9 @@ class LLM:

1617 self,

1618 request_id: str,

1619 engine_prompt: PromptType,

1628 - params: Union[SamplingParams, PoolingParams],

1620 + params: SamplingParams | PoolingParams,

1621 *,

1630 - lora_request: Optional[LoRARequest],

1622 + lora_request: LoRARequest | None,

1623 priority: int,

1624 ) -> tuple[EngineCoreRequest, dict[str, Any]]:

1625 """Use the Processor to process inputs for LLMEngine."""

@@ -1651,8 +1643,8 @@ class LLM:

1643 def _add_request(

1644 self,

1645 prompt: PromptType,

1654 - params: Union[SamplingParams, PoolingParams],

1655 - lora_request: Optional[LoRARequest] = None,

1646 + params: SamplingParams | PoolingParams,

1647 + lora_request: LoRARequest | None = None,

1648 priority: int = 0,

1649 ) -> None:

1650 prompt_text, _, _ = get_prompt_components(prompt)

@@ -1677,8 +1669,8 @@ class LLM:

1669 )

1670

1671 def _run_engine(

1680 - self, *, use_tqdm: Union[bool, Callable[..., tqdm]] = True

1681 - ) -> list[Union[RequestOutput, PoolingRequestOutput]]:

1672 + self, *, use_tqdm: bool | Callable[..., tqdm] = True

1673 + ) -> list[RequestOutput | PoolingRequestOutput]:

1674 # Initialize tqdm.

1675 if use_tqdm:

1676 num_requests = self.llm_engine.get_num_unfinished_requests()

@@ -1691,7 +1683,7 @@ class LLM:

1683 )

1684

1685 # Run the engine.

1694 - outputs: list[Union[RequestOutput, PoolingRequestOutput]] = []

1686 + outputs: list[RequestOutput | PoolingRequestOutput] = []

1687 total_in_toks = 0

1688 total_out_toks = 0

1689 while self.llm_engine.has_unfinished_requests():

vllm/entrypoints/logger.py

+8 -9

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

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

3

4 from collections.abc import Sequence

5 -from typing import Optional, Union

5

6 import torch

7

@@ -15,17 +14,17 @@ logger = init_logger(__name__)

14

15

16 class RequestLogger:

18 - def __init__(self, *, max_log_len: Optional[int]) -> None:

17 + def __init__(self, *, max_log_len: int | None) -> None:

18 self.max_log_len = max_log_len

19

20 def log_inputs(

21 self,

22 request_id: str,

24 - prompt: Optional[str],

25 - prompt_token_ids: Optional[list[int]],

26 - prompt_embeds: Optional[torch.Tensor],

27 - params: Optional[Union[SamplingParams, PoolingParams, BeamSearchParams]],

28 - lora_request: Optional[LoRARequest],

23 + prompt: str | None,

24 + prompt_token_ids: list[int] | None,

25 + prompt_embeds: torch.Tensor | None,

26 + params: SamplingParams | PoolingParams | BeamSearchParams | None,

27 + lora_request: LoRARequest | None,

28 ) -> None:

29 max_log_len = self.max_log_len

30 if max_log_len is not None:

@@ -52,8 +51,8 @@ class RequestLogger:

51 self,

52 request_id: str,

53 outputs: str,

55 - output_token_ids: Optional[Sequence[int]],

56 - finish_reason: Optional[str] = None,

54 + output_token_ids: Sequence[int] | None,

55 + finish_reason: str | None = None,

56 is_streaming: bool = False,

57 delta: bool = False,

58 ) -> None:

vllm/entrypoints/openai/api_server.py

+20 -20

@@ -16,10 +16,10 @@ import socket

16 import tempfile

17 import uuid

18 from argparse import Namespace

19 -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable

19 +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable

20 from contextlib import asynccontextmanager

21 from http import HTTPStatus

22 -from typing import Annotated, Any, Callable, Literal, Optional

22 +from typing import Annotated, Any, Literal

23

24 import prometheus_client

25 import pydantic

@@ -166,8 +166,8 @@ async def build_async_engine_client(

166 args: Namespace,

167 *,

168 usage_context: UsageContext = UsageContext.OPENAI_API_SERVER,

169 - disable_frontend_multiprocessing: Optional[bool] = None,

170 - client_config: Optional[dict[str, Any]] = None,

169 + disable_frontend_multiprocessing: bool | None = None,

170 + client_config: dict[str, Any] | None = None,

171 ) -> AsyncIterator[EngineClient]:

172 if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":

173 # The executor is expected to be mp.

@@ -203,7 +203,7 @@ async def build_async_engine_client_from_engine_args(

203 *,

204 usage_context: UsageContext = UsageContext.OPENAI_API_SERVER,

205 disable_frontend_multiprocessing: bool = False,

206 - client_config: Optional[dict[str, Any]] = None,

206 + client_config: dict[str, Any] | None = None,

207 ) -> AsyncIterator[EngineClient]:

208 """

209 Create EngineClient, either:

@@ -227,7 +227,7 @@ async def build_async_engine_client_from_engine_args(

227

228 from vllm.v1.engine.async_llm import AsyncLLM

229

230 - async_llm: Optional[AsyncLLM] = None

230 + async_llm: AsyncLLM | None = None

231

232 # Don't mutate the input client_config

233 client_config = dict(client_config) if client_config else {}

@@ -308,35 +308,35 @@ def models(request: Request) -> OpenAIServingModels:

308 return request.app.state.openai_serving_models

309

310

311 -def responses(request: Request) -> Optional[OpenAIServingResponses]:

311 +def responses(request: Request) -> OpenAIServingResponses | None:

312 return request.app.state.openai_serving_responses

313

314

315 -def chat(request: Request) -> Optional[OpenAIServingChat]:

315 +def chat(request: Request) -> OpenAIServingChat | None:

316 return request.app.state.openai_serving_chat

317

318

319 -def completion(request: Request) -> Optional[OpenAIServingCompletion]:

319 +def completion(request: Request) -> OpenAIServingCompletion | None:

320 return request.app.state.openai_serving_completion

321

322

323 -def pooling(request: Request) -> Optional[OpenAIServingPooling]:

323 +def pooling(request: Request) -> OpenAIServingPooling | None:

324 return request.app.state.openai_serving_pooling

325

326

327 -def embedding(request: Request) -> Optional[OpenAIServingEmbedding]:

327 +def embedding(request: Request) -> OpenAIServingEmbedding | None:

328 return request.app.state.openai_serving_embedding

329

330

331 -def score(request: Request) -> Optional[ServingScores]:

331 +def score(request: Request) -> ServingScores | None:

332 return request.app.state.openai_serving_scores

333

334

335 -def classify(request: Request) -> Optional[ServingClassification]:

335 +def classify(request: Request) -> ServingClassification | None:

336 return request.app.state.openai_serving_classification

337

338

339 -def rerank(request: Request) -> Optional[ServingScores]:

339 +def rerank(request: Request) -> ServingScores | None:

340 return request.app.state.openai_serving_scores

341

342

@@ -542,8 +542,8 @@ async def create_responses(request: ResponsesRequest, raw_request: Request):

542 async def retrieve_responses(

543 response_id: str,

544 raw_request: Request,

545 - starting_after: Optional[int] = None,

546 - stream: Optional[bool] = False,

545 + starting_after: int | None = None,

546 + stream: bool | None = False,

547 ):

548 handler = responses(raw_request)

549 if handler is None:

@@ -1039,7 +1039,7 @@ if envs.VLLM_SERVER_DEV_MODE:

1039 # User-defined `method` is responsible for deserialization if needed.

1040 args: list[str] = body.get("args", [])

1041 kwargs: dict[str, str] = body.get("kwargs", {})

1042 - timeout: Optional[float] = body.get("timeout")

1042 + timeout: float | None = body.get("timeout")

1043 results = await engine_client(raw_request).collective_rpc(

1044 method=method, timeout=timeout, args=tuple(args), kwargs=kwargs

1045 )

@@ -1120,7 +1120,7 @@ async def is_scaling_elastic_ep(raw_request: Request):

1120 # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers

1121 # (requires typing_extensions >= 4.13)

1122 RequestType = Any

1123 -GetHandlerFn = Callable[[Request], Optional[OpenAIServing]]

1123 +GetHandlerFn = Callable[[Request], OpenAIServing | None]

1124 EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]

1125

1126 # NOTE: Items defined earlier take higher priority

@@ -1236,7 +1236,7 @@ if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING:

1236 return Response(status_code=200, content=response)

1237

1238

1239 -def load_log_config(log_config_file: Optional[str]) -> Optional[dict]:

1239 +def load_log_config(log_config_file: str | None) -> dict | None:

1240 if not log_config_file:

1241 return None

1242 try:

@@ -1655,7 +1655,7 @@ async def init_app_state(

1655 )

1656

1657 if args.tool_server == "demo":

1658 - tool_server: Optional[ToolServer] = DemoToolServer()

1658 + tool_server: ToolServer | None = DemoToolServer()

1659 assert isinstance(tool_server, DemoToolServer)

1660 await tool_server.init_and_validate()

1661 elif args.tool_server:

vllm/entrypoints/openai/cli_args.py

+16 -16

@@ -11,7 +11,7 @@ import json

11 import ssl

12 from collections.abc import Sequence

13 from dataclasses import field

14 -from typing import Literal, Optional, Union

14 +from typing import Literal

15

16 from pydantic.dataclasses import dataclass

17

@@ -39,8 +39,8 @@ class LoRAParserAction(argparse.Action):

39 self,

40 parser: argparse.ArgumentParser,

41 namespace: argparse.Namespace,

42 - values: Optional[Union[str, Sequence[str]]],

43 - option_string: Optional[str] = None,

42 + values: str | Sequence[str] | None,

43 + option_string: str | None = None,

44 ):

45 if values is None:

46 values = []

@@ -73,11 +73,11 @@ class LoRAParserAction(argparse.Action):

73 class FrontendArgs:

74 """Arguments for the OpenAI-compatible frontend server."""

75

76 - host: Optional[str] = None

76 + host: str | None = None

77 """Host name."""

78 port: int = 8000

79 """Port number."""

80 - uds: Optional[str] = None

80 + uds: str | None = None

81 """Unix domain socket path. If set, host and port arguments are ignored."""

82 uvicorn_log_level: Literal[

83 "debug", "info", "warning", "error", "critical", "trace"

@@ -93,15 +93,15 @@ class FrontendArgs:

93 """Allowed methods."""

94 allowed_headers: list[str] = field(default_factory=lambda: ["*"])

95 """Allowed headers."""

96 - api_key: Optional[list[str]] = None

96 + api_key: list[str] | None = None

97 """If provided, the server will require one of these keys to be presented in

98 the header."""

99 - lora_modules: Optional[list[LoRAModulePath]] = None

99 + lora_modules: list[LoRAModulePath] | None = None

100 """LoRA modules configurations in either 'name=path' format or JSON format

101 or JSON list format. Example (old format): `'name=path'` Example (new

102 format): `{\"name\": \"name\", \"path\": \"lora_path\",

103 \"base_model_name\": \"id\"}`"""

104 - chat_template: Optional[str] = None

104 + chat_template: str | None = None

105 """The file path to the chat template, or the template in single-line form

106 for the specified model."""

107 chat_template_content_format: ChatTemplateContentFormatOption = "auto"

@@ -116,17 +116,17 @@ class FrontendArgs:

116 or the ones from tokenizer."""

117 response_role: str = "assistant"

118 """The role name to return if `request.add_generation_prompt=true`."""

119 - ssl_keyfile: Optional[str] = None

119 + ssl_keyfile: str | None = None

120 """The file path to the SSL key file."""

121 - ssl_certfile: Optional[str] = None

121 + ssl_certfile: str | None = None

122 """The file path to the SSL cert file."""

123 - ssl_ca_certs: Optional[str] = None

123 + ssl_ca_certs: str | None = None

124 """The CA certificates file."""

125 enable_ssl_refresh: bool = False

126 """Refresh SSL Context when SSL certificate files change"""

127 ssl_cert_reqs: int = int(ssl.CERT_NONE)

128 """Whether client certificate is required (see stdlib ssl module's)."""

129 - root_path: Optional[str] = None

129 + root_path: str | None = None

130 """FastAPI root_path when app is behind a path based routing proxy."""

131 middleware: list[str] = field(default_factory=lambda: [])

132 """Additional ASGI middleware to apply to the app. We accept multiple

@@ -149,7 +149,7 @@ class FrontendArgs:

149 exclude_tools_when_tool_choice_none: bool = False

150 """If specified, exclude tool definitions in prompts when

151 tool_choice='none'."""

152 - tool_call_parser: Optional[str] = None

152 + tool_call_parser: str | None = None

153 """Select the tool call parser depending on the model that you're using.

154 This is used to parse the model-generated tool call into OpenAI API format.

155 Required for `--enable-auto-tool-choice`. You can choose any option from

@@ -158,13 +158,13 @@ class FrontendArgs:

158 """Special the tool parser plugin write to parse the model-generated tool

159 into OpenAI API format, the name register in this plugin can be used in

160 `--tool-call-parser`."""

161 - tool_server: Optional[str] = None

161 + tool_server: str | None = None

162 """Comma-separated list of host:port pairs (IPv4, IPv6, or hostname).

163 Examples: 127.0.0.1:8000, [::1]:8000, localhost:1234. Or `demo` for demo

164 purpose."""

165 - log_config_file: Optional[str] = envs.VLLM_LOGGING_CONFIG_PATH

165 + log_config_file: str | None = envs.VLLM_LOGGING_CONFIG_PATH

166 """Path to logging config JSON file for both vllm and uvicorn"""

167 - max_log_len: Optional[int] = None

167 + max_log_len: int | None = None

168 """Max number of prompt characters or prompt ID numbers being printed in

169 log. The default of None means unlimited."""

170 disable_fastapi_docs: bool = False

vllm/entrypoints/openai/logits_processors.py

+4 -5

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

3

4 from collections.abc import Iterable

5 from functools import lru_cache, partial

6 -from typing import Optional, Union

6

7 import torch

8

@@ -16,8 +15,8 @@ class AllowedTokenIdsLogitsProcessor:

15 specific set of token ids."""

16

17 def __init__(self, allowed_ids: Iterable[int]):

19 - self.allowed_ids: Optional[list[int]] = list(allowed_ids)

20 - self.mask: Optional[torch.Tensor] = None

18 + self.allowed_ids: list[int] | None = list(allowed_ids)

19 + self.mask: torch.Tensor | None = None

20

21 def __call__(self, token_ids: list[int], logits: torch.Tensor) -> torch.Tensor:

22 if self.mask is None:

@@ -53,8 +52,8 @@ def logit_bias_logits_processor(

52

53

54 def get_logits_processors(

56 - logit_bias: Optional[Union[dict[int, float], dict[str, float]]],

57 - allowed_token_ids: Optional[list[int]],

55 + logit_bias: dict[int, float] | dict[str, float] | None,

56 + allowed_token_ids: list[int] | None,

57 tokenizer: AnyTokenizer,

58 ) -> list[LogitsProcessor]:

59 logits_processors: list[LogitsProcessor] = []

vllm/entrypoints/openai/protocol.py

+339 -339

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

6 import json

7 import time

8 from http import HTTPStatus

9 -from typing import Annotated, Any, ClassVar, Generic, Literal, Optional, TypeVar, Union

9 +from typing import Annotated, Any, ClassVar, Generic, Literal, TypeAlias, TypeVar

10

11 import regex as re

12 import torch

@@ -54,6 +54,7 @@ try: # For older openai versions (< 1.100.0)

54 except ImportError: # For newer openai versions (>= 1.100.0)

55 from openai.types.responses import ResponseFormatTextConfig as ResponseTextConfig

56

57 +

58 from openai.types.responses.response import IncompleteDetails, ToolChoice

59 from openai.types.responses.tool import Tool

60 from openai.types.shared import Metadata, Reasoning

@@ -67,7 +68,6 @@ from pydantic import (

68 field_validator,

69 model_validator,

70 )

70 -from typing_extensions import TypeAlias

71

72 from vllm import envs

73 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam, make_tool_call_id

@@ -93,7 +93,7 @@ class OpenAIBaseModel(BaseModel):

93 model_config = ConfigDict(extra="allow")

94

95 # Cache class field names

96 - field_names: ClassVar[Optional[set[str]]] = None

96 + field_names: ClassVar[set[str] | None] = None

97

98 @model_validator(mode="wrap")

99 @classmethod

@@ -123,7 +123,7 @@ class OpenAIBaseModel(BaseModel):

123 class ErrorInfo(OpenAIBaseModel):

124 message: str

125 type: str

126 - param: Optional[str] = None

126 + param: str | None = None

127 code: int

128

129

@@ -142,7 +142,7 @@ class ModelPermission(OpenAIBaseModel):

142 allow_view: bool = True

143 allow_fine_tuning: bool = False

144 organization: str = "*"

145 - group: Optional[str] = None

145 + group: str | None = None

146 is_blocking: bool = False

147

148

@@ -151,9 +151,9 @@ class ModelCard(OpenAIBaseModel):

151 object: str = "model"

152 created: int = Field(default_factory=lambda: int(time.time()))

153 owned_by: str = "vllm"

154 - root: Optional[str] = None

155 - parent: Optional[str] = None

156 - max_model_len: Optional[int] = None

154 + root: str | None = None

155 + parent: str | None = None

156 + max_model_len: int | None = None

157 permission: list[ModelPermission] = Field(default_factory=list)

158

159

@@ -163,37 +163,35 @@ class ModelList(OpenAIBaseModel):

163

164

165 class PromptTokenUsageInfo(OpenAIBaseModel):

166 - cached_tokens: Optional[int] = None

166 + cached_tokens: int | None = None

167

168

169 class UsageInfo(OpenAIBaseModel):

170 prompt_tokens: int = 0

171 total_tokens: int = 0

172 - completion_tokens: Optional[int] = 0

173 - prompt_tokens_details: Optional[PromptTokenUsageInfo] = None

172 + completion_tokens: int | None = 0

173 + prompt_tokens_details: PromptTokenUsageInfo | None = None

174

175

176 class RequestResponseMetadata(BaseModel):

177 request_id: str

178 - final_usage_info: Optional[UsageInfo] = None

178 + final_usage_info: UsageInfo | None = None

179

180

181 class JsonSchemaResponseFormat(OpenAIBaseModel):

182 name: str

183 - description: Optional[str] = None

183 + description: str | None = None

184 # schema is the field in openai but that causes conflicts with pydantic so

185 # instead use json_schema with an alias

186 - json_schema: Optional[dict[str, Any]] = Field(default=None, alias="schema")

187 - strict: Optional[bool] = None

186 + json_schema: dict[str, Any] | None = Field(default=None, alias="schema")

187 + strict: bool | None = None

188

189

190 class StructuralTag(OpenAIBaseModel):

191 begin: str

192 # schema is the field, but that causes conflicts with pydantic so

193 # instead use structural_tag_schema with an alias

194 - structural_tag_schema: Optional[dict[str, Any]] = Field(

195 - default=None, alias="schema"

196 - )

194 + structural_tag_schema: dict[str, Any] | None = Field(default=None, alias="schema")

195 end: str

196

197

@@ -206,21 +204,21 @@ class StructuralTagResponseFormat(OpenAIBaseModel):

204 class ResponseFormat(OpenAIBaseModel):

205 # type must be "json_schema", "json_object", or "text"

206 type: Literal["text", "json_object", "json_schema"]

209 - json_schema: Optional[JsonSchemaResponseFormat] = None

207 + json_schema: JsonSchemaResponseFormat | None = None

208

209

212 -AnyResponseFormat = Union[ResponseFormat, StructuralTagResponseFormat]

210 +AnyResponseFormat: TypeAlias = ResponseFormat | StructuralTagResponseFormat

211

212

213 class StreamOptions(OpenAIBaseModel):

216 - include_usage: Optional[bool] = True

217 - continuous_usage_stats: Optional[bool] = False

214 + include_usage: bool | None = True

215 + continuous_usage_stats: bool | None = False

216

217

218 class FunctionDefinition(OpenAIBaseModel):

219 name: str

222 - description: Optional[str] = None

223 - parameters: Optional[dict[str, Any]] = None

220 + description: str | None = None

221 + parameters: dict[str, Any] | None = None

222

223

224 class ChatCompletionToolsParam(OpenAIBaseModel):

@@ -241,18 +239,18 @@ class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel):

239 # see https://github.com/pydantic/pydantic/issues/3125

240 class LogitsProcessorConstructor(BaseModel):

241 qualname: str

244 - args: Optional[list[Any]] = None

245 - kwargs: Optional[dict[str, Any]] = None

242 + args: list[Any] | None = None

243 + kwargs: dict[str, Any] | None = None

244

245 model_config = ConfigDict(extra="forbid")

246

247

250 -LogitsProcessors = list[Union[str, LogitsProcessorConstructor]]

248 +LogitsProcessors = list[str | LogitsProcessorConstructor]

249

250

251 def get_logits_processors(

254 - processors: Optional[LogitsProcessors], pattern: Optional[str]

255 -) -> Optional[list[Any]]:

252 + processors: LogitsProcessors | None, pattern: str | None

253 +) -> list[Any] | None:

254 if processors and pattern:

255 logits_processors = []

256 for processor in processors:

@@ -284,16 +282,16 @@ def get_logits_processors(

282 return None

283

284

287 -ResponseInputOutputItem: TypeAlias = Union[

288 - ResponseInputItemParam, ResponseReasoningItem, ResponseFunctionToolCall

289 -]

285 +ResponseInputOutputItem: TypeAlias = (

286 + ResponseInputItemParam | ResponseReasoningItem | ResponseFunctionToolCall

287 +)

288

289

290 class ResponsesRequest(OpenAIBaseModel):

291 # Ordered by official OpenAI API documentation

292 # https://platform.openai.com/docs/api-reference/responses/create

295 - background: Optional[bool] = False

296 - include: Optional[

293 + background: bool | None = False

294 + include: (

295 list[

296 Literal[

297 "code_interpreter_call.outputs",

@@ -304,28 +302,29 @@ class ResponsesRequest(OpenAIBaseModel):

302 "reasoning.encrypted_content",

303 ],

304 ]

307 - ] = None

308 - input: Union[str, list[ResponseInputOutputItem]]

309 - instructions: Optional[str] = None

310 - max_output_tokens: Optional[int] = None

311 - max_tool_calls: Optional[int] = None

312 - metadata: Optional[Metadata] = None

313 - model: Optional[str] = None

314 - parallel_tool_calls: Optional[bool] = True

315 - previous_response_id: Optional[str] = None

316 - prompt: Optional[ResponsePrompt] = None

317 - reasoning: Optional[Reasoning] = None

305 + | None

306 + ) = None

307 + input: str | list[ResponseInputOutputItem]

308 + instructions: str | None = None

309 + max_output_tokens: int | None = None

310 + max_tool_calls: int | None = None

311 + metadata: Metadata | None = None

312 + model: str | None = None

313 + parallel_tool_calls: bool | None = True

314 + previous_response_id: str | None = None

315 + prompt: ResponsePrompt | None = None

316 + reasoning: Reasoning | None = None

317 service_tier: Literal["auto", "default", "flex", "scale", "priority"] = "auto"

319 - store: Optional[bool] = True

320 - stream: Optional[bool] = False

321 - temperature: Optional[float] = None

322 - text: Optional[ResponseTextConfig] = None

318 + store: bool | None = True

319 + stream: bool | None = False

320 + temperature: float | None = None

321 + text: ResponseTextConfig | None = None

322 tool_choice: ToolChoice = "auto"

323 tools: list[Tool] = Field(default_factory=list)

325 - top_logprobs: Optional[int] = 0

326 - top_p: Optional[float] = None

327 - truncation: Optional[Literal["auto", "disabled"]] = "disabled"

328 - user: Optional[str] = None

324 + top_logprobs: int | None = 0

325 + top_p: float | None = None

326 + truncation: Literal["auto", "disabled"] | None = "disabled"

327 + user: str | None = None

328

329 # --8<-- [start:responses-extra-params]

330 request_id: str = Field(

@@ -336,7 +335,7 @@ class ResponsesRequest(OpenAIBaseModel):

335 "through out the inference process and return in response."

336 ),

337 )

339 - mm_processor_kwargs: Optional[dict[str, Any]] = Field(

338 + mm_processor_kwargs: dict[str, Any] | None = Field(

339 default=None,

340 description=("Additional kwargs to pass to the HF processor."),

341 )

@@ -348,7 +347,7 @@ class ResponsesRequest(OpenAIBaseModel):

347 "if the served model does not use priority scheduling."

348 ),

349 )

351 - cache_salt: Optional[str] = Field(

350 + cache_salt: str | None = Field(

351 default=None,

352 description=(

353 "If specified, the prefix cache will be salted with the provided "

@@ -378,7 +377,7 @@ class ResponsesRequest(OpenAIBaseModel):

377 def to_sampling_params(

378 self,

379 default_max_tokens: int,

381 - default_sampling_params: Optional[dict] = None,

380 + default_sampling_params: dict | None = None,

381 ) -> SamplingParams:

382 if self.max_output_tokens is None:

383 max_tokens = default_max_tokens

@@ -465,58 +464,57 @@ class ChatCompletionRequest(OpenAIBaseModel):

464 # Ordered by official OpenAI API documentation

465 # https://platform.openai.com/docs/api-reference/chat/create

466 messages: list[ChatCompletionMessageParam]

468 - model: Optional[str] = None

469 - frequency_penalty: Optional[float] = 0.0

470 - logit_bias: Optional[dict[str, float]] = None

471 - logprobs: Optional[bool] = False

472 - top_logprobs: Optional[int] = 0

473 - max_tokens: Optional[int] = Field(

467 + model: str | None = None

468 + frequency_penalty: float | None = 0.0

469 + logit_bias: dict[str, float] | None = None

470 + logprobs: bool | None = False

471 + top_logprobs: int | None = 0

472 + max_tokens: int | None = Field(

473 default=None,

474 deprecated="max_tokens is deprecated in favor of "

475 "the max_completion_tokens field",

476 )

478 - max_completion_tokens: Optional[int] = None

479 - n: Optional[int] = 1

480 - presence_penalty: Optional[float] = 0.0

481 - response_format: Optional[AnyResponseFormat] = None

482 - seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

483 - stop: Optional[Union[str, list[str]]] = []

484 - stream: Optional[bool] = False

485 - stream_options: Optional[StreamOptions] = None

486 - temperature: Optional[float] = None

487 - top_p: Optional[float] = None

488 - tools: Optional[list[ChatCompletionToolsParam]] = None

489 - tool_choice: Optional[

490 - Union[

491 - Literal["none"],

492 - Literal["auto"],

493 - Literal["required"],

494 - ChatCompletionNamedToolChoiceParam,

495 - ]

496 - ] = "none"

497 - reasoning_effort: Optional[Literal["low", "medium", "high"]] = None

477 + max_completion_tokens: int | None = None

478 + n: int | None = 1

479 + presence_penalty: float | None = 0.0

480 + response_format: AnyResponseFormat | None = None

481 + seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

482 + stop: str | list[str] | None = []

483 + stream: bool | None = False

484 + stream_options: StreamOptions | None = None

485 + temperature: float | None = None

486 + top_p: float | None = None

487 + tools: list[ChatCompletionToolsParam] | None = None

488 + tool_choice: (

489 + Literal["none"]

490 + | Literal["auto"]

491 + | Literal["required"]

492 + | ChatCompletionNamedToolChoiceParam

493 + | None

494 + ) = "none"

495 + reasoning_effort: Literal["low", "medium", "high"] | None = None

496 include_reasoning: bool = True

497

498 # NOTE this will be ignored by vLLM -- the model determines the behavior

501 - parallel_tool_calls: Optional[bool] = False

502 - user: Optional[str] = None

499 + parallel_tool_calls: bool | None = False

500 + user: str | None = None

501

502 # --8<-- [start:chat-completion-sampling-params]

505 - best_of: Optional[int] = None

503 + best_of: int | None = None

504 use_beam_search: bool = False

507 - top_k: Optional[int] = None

508 - min_p: Optional[float] = None

509 - repetition_penalty: Optional[float] = None

505 + top_k: int | None = None

506 + min_p: float | None = None

507 + repetition_penalty: float | None = None

508 length_penalty: float = 1.0

511 - stop_token_ids: Optional[list[int]] = []

509 + stop_token_ids: list[int] | None = []

510 include_stop_str_in_output: bool = False

511 ignore_eos: bool = False

512 min_tokens: int = 0

513 skip_special_tokens: bool = True

514 spaces_between_special_tokens: bool = True

517 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

518 - prompt_logprobs: Optional[int] = None

519 - allowed_token_ids: Optional[list[int]] = None

515 + truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None

516 + prompt_logprobs: int | None = None

517 + allowed_token_ids: list[int] | None = None

518 bad_words: list[str] = Field(default_factory=list)

519 # --8<-- [end:chat-completion-sampling-params]

520

@@ -556,7 +554,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

554 "default)."

555 ),

556 )

559 - documents: Optional[list[dict[str, str]]] = Field(

557 + documents: list[dict[str, str]] | None = Field(

558 default=None,

559 description=(

560 "A list of dicts representing documents that will be accessible to "

@@ -566,7 +564,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

564 '"title" and "text" keys.'

565 ),

566 )

569 - chat_template: Optional[str] = Field(

567 + chat_template: str | None = Field(

568 default=None,

569 description=(

570 "A Jinja template to use for this conversion. "

@@ -575,22 +573,22 @@ class ChatCompletionRequest(OpenAIBaseModel):

573 "does not define one."

574 ),

575 )

578 - chat_template_kwargs: Optional[dict[str, Any]] = Field(

576 + chat_template_kwargs: dict[str, Any] | None = Field(

577 default=None,

578 description=(

579 "Additional keyword args to pass to the template renderer. "

580 "Will be accessible by the chat template."

581 ),

582 )

585 - mm_processor_kwargs: Optional[dict[str, Any]] = Field(

583 + mm_processor_kwargs: dict[str, Any] | None = Field(

584 default=None,

585 description=("Additional kwargs to pass to the HF processor."),

586 )

589 - structured_outputs: Optional[StructuredOutputsParams] = Field(

587 + structured_outputs: StructuredOutputsParams | None = Field(

588 default=None,

589 description="Additional kwargs for structured outputs",

590 )

593 - guided_json: Optional[Union[str, dict, BaseModel]] = Field(

591 + guided_json: str | dict | BaseModel | None = Field(

592 default=None,

593 description=(

594 "`guided_json` is deprecated. "

@@ -598,7 +596,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

596 "Please pass `json` to `structured_outputs` instead."

597 ),

598 )

601 - guided_regex: Optional[str] = Field(

599 + guided_regex: str | None = Field(

600 default=None,

601 description=(

602 "`guided_regex` is deprecated. "

@@ -606,7 +604,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

604 "Please pass `regex` to `structured_outputs` instead."

605 ),

606 )

609 - guided_choice: Optional[list[str]] = Field(

607 + guided_choice: list[str] | None = Field(

608 default=None,

609 description=(

610 "`guided_choice` is deprecated. "

@@ -614,7 +612,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

612 "Please pass `choice` to `structured_outputs` instead."

613 ),

614 )

617 - guided_grammar: Optional[str] = Field(

615 + guided_grammar: str | None = Field(

616 default=None,

617 description=(

618 "`guided_grammar` is deprecated. "

@@ -622,7 +620,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

620 "Please pass `grammar` to `structured_outputs` instead."

621 ),

622 )

625 - structural_tag: Optional[str] = Field(

623 + structural_tag: str | None = Field(

624 default=None,

625 description=(

626 "`structural_tag` is deprecated. "

@@ -630,7 +628,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

628 "Please pass `structural_tag` to `structured_outputs` instead."

629 ),

630 )

633 - guided_decoding_backend: Optional[str] = Field(

631 + guided_decoding_backend: str | None = Field(

632 default=None,

633 description=(

634 "`guided_decoding_backend` is deprecated. "

@@ -638,7 +636,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

636 "Please remove it from your request."

637 ),

638 )

641 - guided_whitespace_pattern: Optional[str] = Field(

639 + guided_whitespace_pattern: str | None = Field(

640 default=None,

641 description=(

642 "`guided_whitespace_pattern` is deprecated. "

@@ -662,7 +660,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

660 "through out the inference process and return in response."

661 ),

662 )

665 - logits_processors: Optional[LogitsProcessors] = Field(

663 + logits_processors: LogitsProcessors | None = Field(

664 default=None,

665 description=(

666 "A list of either qualified names of logits processors, or "

@@ -675,7 +673,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

673 "{'param': 'value'}}."

674 ),

675 )

678 - return_tokens_as_token_ids: Optional[bool] = Field(

676 + return_tokens_as_token_ids: bool | None = Field(

677 default=None,

678 description=(

679 "If specified with 'logprobs', tokens are represented "

@@ -683,7 +681,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

681 "that are not JSON-encodable can be identified."

682 ),

683 )

686 - return_token_ids: Optional[bool] = Field(

684 + return_token_ids: bool | None = Field(

685 default=None,

686 description=(

687 "If specified, the result will include token IDs alongside the "

@@ -693,7 +691,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

691 "need to map generated text back to input tokens."

692 ),

693 )

696 - cache_salt: Optional[str] = Field(

694 + cache_salt: str | None = Field(

695 default=None,

696 description=(

697 "If specified, the prefix cache will be salted with the provided "

@@ -704,12 +702,12 @@ class ChatCompletionRequest(OpenAIBaseModel):

702 "to 256 bit). Not supported by vLLM engine V0."

703 ),

704 )

707 - kv_transfer_params: Optional[dict[str, Any]] = Field(

705 + kv_transfer_params: dict[str, Any] | None = Field(

706 default=None,

707 description="KVTransfer parameters used for disaggregated serving.",

708 )

709

712 - vllm_xargs: Optional[dict[str, Union[str, int, float]]] = Field(

710 + vllm_xargs: dict[str, str | int | float] | None = Field(

711 default=None,

712 description=(

713 "Additional request parameters with string or "

@@ -749,7 +747,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

747 def to_sampling_params(

748 self,

749 max_tokens: int,

752 - logits_processor_pattern: Optional[str],

750 + logits_processor_pattern: str | None,

751 default_sampling_params: dict,

752 ) -> SamplingParams:

753 # Default parameters

@@ -860,7 +858,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

858 extra_args=extra_args or None,

859 )

860

863 - def _get_json_schema_from_tool(self) -> Optional[Union[str, dict]]:

861 + def _get_json_schema_from_tool(self) -> str | dict | None:

862 # user has chosen to not use any tool

863 if self.tool_choice == "none" or self.tools is None:

864 return None

@@ -1098,44 +1096,44 @@ class ChatCompletionRequest(OpenAIBaseModel):

1096 class CompletionRequest(OpenAIBaseModel):

1097 # Ordered by official OpenAI API documentation

1098 # https://platform.openai.com/docs/api-reference/completions/create

1101 - model: Optional[str] = None

1102 - prompt: Optional[Union[list[int], list[list[int]], str, list[str]]] = None

1103 - best_of: Optional[int] = None

1104 - echo: Optional[bool] = False

1105 - frequency_penalty: Optional[float] = 0.0

1106 - logit_bias: Optional[dict[str, float]] = None

1107 - logprobs: Optional[int] = None

1108 - max_tokens: Optional[int] = 16

1099 + model: str | None = None

1100 + prompt: list[int] | list[list[int]] | str | list[str] | None = None

1101 + best_of: int | None = None

1102 + echo: bool | None = False

1103 + frequency_penalty: float | None = 0.0

1104 + logit_bias: dict[str, float] | None = None

1105 + logprobs: int | None = None

1106 + max_tokens: int | None = 16

1107 n: int = 1

1110 - presence_penalty: Optional[float] = 0.0

1111 - seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

1112 - stop: Optional[Union[str, list[str]]] = []

1113 - stream: Optional[bool] = False

1114 - stream_options: Optional[StreamOptions] = None

1115 - suffix: Optional[str] = None

1116 - temperature: Optional[float] = None

1117 - top_p: Optional[float] = None

1118 - user: Optional[str] = None

1108 + presence_penalty: float | None = 0.0

1109 + seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

1110 + stop: str | list[str] | None = []

1111 + stream: bool | None = False

1112 + stream_options: StreamOptions | None = None

1113 + suffix: str | None = None

1114 + temperature: float | None = None

1115 + top_p: float | None = None

1116 + user: str | None = None

1117

1118 # --8<-- [start:completion-sampling-params]

1119 use_beam_search: bool = False

1122 - top_k: Optional[int] = None

1123 - min_p: Optional[float] = None

1124 - repetition_penalty: Optional[float] = None

1120 + top_k: int | None = None

1121 + min_p: float | None = None

1122 + repetition_penalty: float | None = None

1123 length_penalty: float = 1.0

1126 - stop_token_ids: Optional[list[int]] = []

1124 + stop_token_ids: list[int] | None = []

1125 include_stop_str_in_output: bool = False

1126 ignore_eos: bool = False

1127 min_tokens: int = 0

1128 skip_special_tokens: bool = True

1129 spaces_between_special_tokens: bool = True

1132 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

1133 - allowed_token_ids: Optional[list[int]] = None

1134 - prompt_logprobs: Optional[int] = None

1130 + truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None

1131 + allowed_token_ids: list[int] | None = None

1132 + prompt_logprobs: int | None = None

1133 # --8<-- [end:completion-sampling-params]

1134

1135 # --8<-- [start:completion-extra-params]

1138 - prompt_embeds: Optional[Union[bytes, list[bytes]]] = None

1136 + prompt_embeds: bytes | list[bytes] | None = None

1137 add_special_tokens: bool = Field(

1138 default=True,

1139 description=(

@@ -1143,7 +1141,7 @@ class CompletionRequest(OpenAIBaseModel):

1141 "the prompt."

1142 ),

1143 )

1146 - response_format: Optional[AnyResponseFormat] = Field(

1144 + response_format: AnyResponseFormat | None = Field(

1145 default=None,

1146 description=(

1147 "Similar to chat completion, this parameter specifies the format "

@@ -1151,11 +1149,11 @@ class CompletionRequest(OpenAIBaseModel):

1149 ", {'type': 'structural_tag'}, or {'type': 'text' } is supported."

1150 ),

1151 )

1154 - structured_outputs: Optional[StructuredOutputsParams] = Field(

1152 + structured_outputs: StructuredOutputsParams | None = Field(

1153 default=None,

1154 description="Additional kwargs for structured outputs",

1155 )

1158 - guided_json: Optional[Union[str, dict, BaseModel]] = Field(

1156 + guided_json: str | dict | BaseModel | None = Field(

1157 default=None,

1158 description=(

1159 "`guided_json` is deprecated. "

@@ -1163,7 +1161,7 @@ class CompletionRequest(OpenAIBaseModel):

1161 "Please pass `json` to `structured_outputs` instead."

1162 ),

1163 )

1166 - guided_regex: Optional[str] = Field(

1164 + guided_regex: str | None = Field(

1165 default=None,

1166 description=(

1167 "`guided_regex` is deprecated. "

@@ -1171,7 +1169,7 @@ class CompletionRequest(OpenAIBaseModel):

1169 "Please pass `regex` to `structured_outputs` instead."

1170 ),

1171 )

1174 - guided_choice: Optional[list[str]] = Field(

1172 + guided_choice: list[str] | None = Field(

1173 default=None,

1174 description=(

1175 "`guided_choice` is deprecated. "

@@ -1179,7 +1177,7 @@ class CompletionRequest(OpenAIBaseModel):

1177 "Please pass `choice` to `structured_outputs` instead."

1178 ),

1179 )

1182 - guided_grammar: Optional[str] = Field(

1180 + guided_grammar: str | None = Field(

1181 default=None,

1182 description=(

1183 "`guided_grammar` is deprecated. "

@@ -1187,7 +1185,7 @@ class CompletionRequest(OpenAIBaseModel):

1185 "Please pass `grammar` to `structured_outputs` instead."

1186 ),

1187 )

1190 - guided_decoding_backend: Optional[str] = Field(

1188 + guided_decoding_backend: str | None = Field(

1189 default=None,

1190 description=(

1191 "`guided_decoding_backend` is deprecated. "

@@ -1195,7 +1193,7 @@ class CompletionRequest(OpenAIBaseModel):

1193 "Please remove it from your request."

1194 ),

1195 )

1198 - guided_whitespace_pattern: Optional[str] = Field(

1196 + guided_whitespace_pattern: str | None = Field(

1197 default=None,

1198 description=(

1199 "`guided_whitespace_pattern` is deprecated. "

@@ -1219,7 +1217,7 @@ class CompletionRequest(OpenAIBaseModel):

1217 "through out the inference process and return in response."

1218 ),

1219 )

1222 - logits_processors: Optional[LogitsProcessors] = Field(

1220 + logits_processors: LogitsProcessors | None = Field(

1221 default=None,

1222 description=(

1223 "A list of either qualified names of logits processors, or "

@@ -1233,7 +1231,7 @@ class CompletionRequest(OpenAIBaseModel):

1231 ),

1232 )

1233

1236 - return_tokens_as_token_ids: Optional[bool] = Field(

1234 + return_tokens_as_token_ids: bool | None = Field(

1235 default=None,

1236 description=(

1237 "If specified with 'logprobs', tokens are represented "

@@ -1241,7 +1239,7 @@ class CompletionRequest(OpenAIBaseModel):

1239 "that are not JSON-encodable can be identified."

1240 ),

1241 )

1244 - return_token_ids: Optional[bool] = Field(

1242 + return_token_ids: bool | None = Field(

1243 default=None,

1244 description=(

1245 "If specified, the result will include token IDs alongside the "

@@ -1252,7 +1250,7 @@ class CompletionRequest(OpenAIBaseModel):

1250 ),

1251 )

1252

1255 - cache_salt: Optional[str] = Field(

1253 + cache_salt: str | None = Field(

1254 default=None,

1255 description=(

1256 "If specified, the prefix cache will be salted with the provided "

@@ -1264,12 +1262,12 @@ class CompletionRequest(OpenAIBaseModel):

1262 ),

1263 )

1264

1267 - kv_transfer_params: Optional[dict[str, Any]] = Field(

1265 + kv_transfer_params: dict[str, Any] | None = Field(

1266 default=None,

1267 description="KVTransfer parameters used for disaggregated serving.",

1268 )

1269

1272 - vllm_xargs: Optional[dict[str, Union[str, int, float]]] = Field(

1270 + vllm_xargs: dict[str, str | int | float] | None = Field(

1271 default=None,

1272 description=(

1273 "Additional request parameters with string or "

@@ -1291,7 +1289,7 @@ class CompletionRequest(OpenAIBaseModel):

1289 def to_beam_search_params(

1290 self,

1291 max_tokens: int,

1294 - default_sampling_params: Optional[dict] = None,

1292 + default_sampling_params: dict | None = None,

1293 ) -> BeamSearchParams:

1294 if default_sampling_params is None:

1295 default_sampling_params = {}

@@ -1312,8 +1310,8 @@ class CompletionRequest(OpenAIBaseModel):

1310 def to_sampling_params(

1311 self,

1312 max_tokens: int,

1315 - logits_processor_pattern: Optional[str],

1316 - default_sampling_params: Optional[dict] = None,

1313 + logits_processor_pattern: str | None,

1314 + default_sampling_params: dict | None = None,

1315 ) -> SamplingParams:

1316 if default_sampling_params is None:

1317 default_sampling_params = {}

@@ -1488,12 +1486,12 @@ class CompletionRequest(OpenAIBaseModel):

1486 class EmbeddingCompletionRequest(OpenAIBaseModel):

1487 # Ordered by official OpenAI API documentation

1488 # https://platform.openai.com/docs/api-reference/embeddings

1491 - model: Optional[str] = None

1492 - input: Union[list[int], list[list[int]], str, list[str]]

1489 + model: str | None = None

1490 + input: list[int] | list[list[int]] | str | list[str]

1491 encoding_format: Literal["float", "base64"] = "float"

1494 - dimensions: Optional[int] = None

1495 - user: Optional[str] = None

1496 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

1492 + dimensions: int | None = None

1493 + user: str | None = None

1494 + truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None

1495

1496 # --8<-- [start:embedding-extra-params]

1497 add_special_tokens: bool = Field(

@@ -1519,7 +1517,7 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):

1517 "through out the inference process and return in response."

1518 ),

1519 )

1522 - normalize: Optional[bool] = None

1520 + normalize: bool | None = None

1521

1522 # --8<-- [end:embedding-extra-params]

1523

@@ -1532,13 +1530,13 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):

1530

1531

1532 class EmbeddingChatRequest(OpenAIBaseModel):

1535 - model: Optional[str] = None

1533 + model: str | None = None

1534 messages: list[ChatCompletionMessageParam]

1535

1536 encoding_format: Literal["float", "base64"] = "float"

1539 - dimensions: Optional[int] = None

1540 - user: Optional[str] = None

1541 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

1537 + dimensions: int | None = None

1538 + user: str | None = None

1539 + truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None

1540

1541 # --8<-- [start:chat-embedding-extra-params]

1542 add_generation_prompt: bool = Field(

@@ -1560,7 +1558,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):

1558 "default)."

1559 ),

1560 )

1563 - chat_template: Optional[str] = Field(

1561 + chat_template: str | None = Field(

1562 default=None,

1563 description=(

1564 "A Jinja template to use for this conversion. "

@@ -1569,14 +1567,14 @@ class EmbeddingChatRequest(OpenAIBaseModel):

1567 "does not define one."

1568 ),

1569 )

1572 - chat_template_kwargs: Optional[dict[str, Any]] = Field(

1570 + chat_template_kwargs: dict[str, Any] | None = Field(

1571 default=None,

1572 description=(

1573 "Additional keyword args to pass to the template renderer. "

1574 "Will be accessible by the chat template."

1575 ),

1576 )

1579 - mm_processor_kwargs: Optional[dict[str, Any]] = Field(

1577 + mm_processor_kwargs: dict[str, Any] | None = Field(

1578 default=None,

1579 description=("Additional kwargs to pass to the HF processor."),

1580 )

@@ -1596,7 +1594,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):

1594 "through out the inference process and return in response."

1595 ),

1596 )

1599 - normalize: Optional[bool] = None

1597 + normalize: bool | None = None

1598 # --8<-- [end:chat-embedding-extra-params]

1599

1600 @model_validator(mode="before")

@@ -1617,7 +1615,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):

1615 )

1616

1617

1620 -EmbeddingRequest = Union[EmbeddingCompletionRequest, EmbeddingChatRequest]

1618 +EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest

1619

1620 PoolingCompletionRequest = EmbeddingCompletionRequest

1621 PoolingChatRequest = EmbeddingChatRequest

@@ -1626,7 +1624,7 @@ T = TypeVar("T")

1624

1625

1626 class IOProcessorRequest(OpenAIBaseModel, Generic[T]):

1629 - model: Optional[str] = None

1627 + model: str | None = None

1628

1629 priority: int = Field(default=0)

1630 """

@@ -1646,7 +1644,7 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):

1644

1645

1646 class IOProcessorResponse(OpenAIBaseModel, Generic[T]):

1649 - request_id: Optional[str] = None

1647 + request_id: str | None = None

1648 """

1649 The request_id associated with this response

1650 """

@@ -1659,18 +1657,20 @@ class IOProcessorResponse(OpenAIBaseModel, Generic[T]):

1657 """

1658

1659

1662 -PoolingRequest = Union[PoolingCompletionRequest, PoolingChatRequest, IOProcessorRequest]

1660 +PoolingRequest: TypeAlias = (

1661 + PoolingCompletionRequest | PoolingChatRequest | IOProcessorRequest

1662 +)

1663

1664

1665 class ScoreRequest(OpenAIBaseModel):

1666 - model: Optional[str] = None

1667 - text_1: Union[list[str], str, ScoreMultiModalParam]

1668 - text_2: Union[list[str], str, ScoreMultiModalParam]

1669 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

1666 + model: str | None = None

1667 + text_1: list[str] | str | ScoreMultiModalParam

1668 + text_2: list[str] | str | ScoreMultiModalParam

1669 + truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None

1670

1671 # --8<-- [start:score-extra-params]

1672

1673 - mm_processor_kwargs: Optional[dict[str, Any]] = Field(

1673 + mm_processor_kwargs: dict[str, Any] | None = Field(

1674 default=None,

1675 description=("Additional kwargs to pass to the HF processor."),

1676 )

@@ -1684,7 +1684,7 @@ class ScoreRequest(OpenAIBaseModel):

1684 ),

1685 )

1686

1687 - activation: Optional[bool] = None

1687 + activation: bool | None = None

1688

1689 # --8<-- [end:score-extra-params]

1690

@@ -1696,15 +1696,15 @@ class ScoreRequest(OpenAIBaseModel):

1696

1697

1698 class RerankRequest(OpenAIBaseModel):

1699 - model: Optional[str] = None

1700 - query: Union[str, ScoreMultiModalParam]

1701 - documents: Union[list[str], ScoreMultiModalParam]

1699 + model: str | None = None

1700 + query: str | ScoreMultiModalParam

1701 + documents: list[str] | ScoreMultiModalParam

1702 top_n: int = Field(default_factory=lambda: 0)

1703 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

1703 + truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None

1704

1705 # --8<-- [start:rerank-extra-params]

1706

1707 - mm_processor_kwargs: Optional[dict[str, Any]] = Field(

1707 + mm_processor_kwargs: dict[str, Any] | None = Field(

1708 default=None,

1709 description=("Additional kwargs to pass to the HF processor."),

1710 )

@@ -1718,7 +1718,7 @@ class RerankRequest(OpenAIBaseModel):

1718 ),

1719 )

1720

1721 - activation: Optional[bool] = None

1721 + activation: bool | None = None

1722

1723 # --8<-- [end:rerank-extra-params]

1724

@@ -1730,8 +1730,8 @@ class RerankRequest(OpenAIBaseModel):

1730

1731

1732 class RerankDocument(BaseModel):

1733 - text: Optional[str] = None

1734 - multi_modal: Optional[ScoreContentPartParam] = None

1733 + text: str | None = None

1734 + multi_modal: ScoreContentPartParam | None = None

1735

1736

1737 class RerankResult(BaseModel):

@@ -1753,17 +1753,17 @@ class RerankResponse(OpenAIBaseModel):

1753

1754 class CompletionLogProbs(OpenAIBaseModel):

1755 text_offset: list[int] = Field(default_factory=list)

1756 - token_logprobs: list[Optional[float]] = Field(default_factory=list)

1756 + token_logprobs: list[float | None] = Field(default_factory=list)

1757 tokens: list[str] = Field(default_factory=list)

1758 - top_logprobs: list[Optional[dict[str, float]]] = Field(default_factory=list)

1758 + top_logprobs: list[dict[str, float] | None] = Field(default_factory=list)

1759

1760

1761 class CompletionResponseChoice(OpenAIBaseModel):

1762 index: int

1763 text: str

1764 - logprobs: Optional[CompletionLogProbs] = None

1765 - finish_reason: Optional[str] = None

1766 - stop_reason: Optional[Union[int, str]] = Field(

1764 + logprobs: CompletionLogProbs | None = None

1765 + finish_reason: str | None = None

1766 + stop_reason: int | str | None = Field(

1767 default=None,

1768 description=(

1769 "The stop string or token id that caused the completion "

@@ -1771,9 +1771,9 @@ class CompletionResponseChoice(OpenAIBaseModel):

1771 "including encountering the EOS token"

1772 ),

1773 )

1774 - token_ids: Optional[list[int]] = None # For response

1775 - prompt_logprobs: Optional[list[Optional[dict[int, Logprob]]]] = None

1776 - prompt_token_ids: Optional[list[int]] = None # For prompt

1774 + token_ids: list[int] | None = None # For response

1775 + prompt_logprobs: list[dict[int, Logprob] | None] | None = None

1776 + prompt_token_ids: list[int] | None = None # For prompt

1777

1778

1779 class CompletionResponse(OpenAIBaseModel):

@@ -1782,14 +1782,12 @@ class CompletionResponse(OpenAIBaseModel):

1782 created: int = Field(default_factory=lambda: int(time.time()))

1783 model: str

1784 choices: list[CompletionResponseChoice]

1785 - service_tier: Optional[Literal["auto", "default", "flex", "scale", "priority"]] = (

1786 - None

1787 - )

1788 - system_fingerprint: Optional[str] = None

1785 + service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None

1786 + system_fingerprint: str | None = None

1787 usage: UsageInfo

1788

1789 # vLLM-specific fields that are not in OpenAI spec

1792 - kv_transfer_params: Optional[dict[str, Any]] = Field(

1790 + kv_transfer_params: dict[str, Any] | None = Field(

1791 default=None, description="KVTransfer parameters."

1792 )

1793

@@ -1797,9 +1795,9 @@ class CompletionResponse(OpenAIBaseModel):

1795 class CompletionResponseStreamChoice(OpenAIBaseModel):

1796 index: int

1797 text: str

1800 - logprobs: Optional[CompletionLogProbs] = None

1801 - finish_reason: Optional[str] = None

1802 - stop_reason: Optional[Union[int, str]] = Field(

1798 + logprobs: CompletionLogProbs | None = None

1799 + finish_reason: str | None = None

1800 + stop_reason: int | str | None = Field(

1801 default=None,

1802 description=(

1803 "The stop string or token id that caused the completion "

@@ -1809,8 +1807,8 @@ class CompletionResponseStreamChoice(OpenAIBaseModel):

1807 )

1808 # not part of the OpenAI spec but for tracing the tokens

1809 # prompt tokens is put into choice to align with CompletionResponseChoice

1812 - prompt_token_ids: Optional[list[int]] = None

1813 - token_ids: Optional[list[int]] = None

1810 + prompt_token_ids: list[int] | None = None

1811 + token_ids: list[int] | None = None

1812

1813

1814 class CompletionStreamResponse(OpenAIBaseModel):

@@ -1819,13 +1817,13 @@ class CompletionStreamResponse(OpenAIBaseModel):

1817 created: int = Field(default_factory=lambda: int(time.time()))

1818 model: str

1819 choices: list[CompletionResponseStreamChoice]

1822 - usage: Optional[UsageInfo] = Field(default=None)

1820 + usage: UsageInfo | None = Field(default=None)

1821

1822

1823 class EmbeddingResponseData(OpenAIBaseModel):

1824 index: int

1825 object: str = "embedding"

1828 - embedding: Union[list[float], str]

1826 + embedding: list[float] | str

1827

1828

1829 class EmbeddingResponse(OpenAIBaseModel):

@@ -1840,7 +1838,7 @@ class EmbeddingResponse(OpenAIBaseModel):

1838 class PoolingResponseData(OpenAIBaseModel):

1839 index: int

1840 object: str = "pooling"

1843 - data: Union[list[list[float]], list[float], str]

1841 + data: list[list[float]] | list[float] | str

1842

1843

1844 class PoolingResponse(OpenAIBaseModel):

@@ -1868,10 +1866,10 @@ class ScoreResponse(OpenAIBaseModel):

1866

1867

1868 class ClassificationRequest(OpenAIBaseModel):

1871 - model: Optional[str] = None

1872 - input: Union[list[str], str]

1873 - truncate_prompt_tokens: Optional[int] = None

1874 - user: Optional[str] = None

1869 + model: str | None = None

1870 + input: list[str] | str

1871 + truncate_prompt_tokens: int | None = None

1872 + user: str | None = None

1873

1874 # --8<-- [start:classification-extra-params]

1875 priority: int = Field(

@@ -1883,7 +1881,7 @@ class ClassificationRequest(OpenAIBaseModel):

1881 ),

1882 )

1883

1886 - activation: Optional[bool] = None

1884 + activation: bool | None = None

1885

1886 # --8<-- [end:classification-extra-params]

1887

@@ -1896,7 +1894,7 @@ class ClassificationRequest(OpenAIBaseModel):

1894

1895 class ClassificationData(OpenAIBaseModel):

1896 index: int

1899 - label: Optional[str]

1897 + label: str | None

1898 probs: list[float]

1899 num_classes: int

1900

@@ -1922,16 +1920,16 @@ class ToolCall(OpenAIBaseModel):

1920

1921

1922 class DeltaFunctionCall(BaseModel):

1925 - name: Optional[str] = None

1926 - arguments: Optional[str] = None

1923 + name: str | None = None

1924 + arguments: str | None = None

1925

1926

1927 # a tool call delta where everything is optional

1928 class DeltaToolCall(OpenAIBaseModel):

1931 - id: Optional[str] = None

1932 - type: Optional[Literal["function"]] = None

1929 + id: str | None = None

1930 + type: Literal["function"] | None = None

1931 index: int

1934 - function: Optional[DeltaFunctionCall] = None

1932 + function: DeltaFunctionCall | None = None

1933

1934

1935 class ExtractedToolCallInformation(BaseModel):

@@ -1943,50 +1941,50 @@ class ExtractedToolCallInformation(BaseModel):

1941

1942 # content - per OpenAI spec, content AND tool calls can be returned rarely

1943 # But some models will do this intentionally

1946 - content: Optional[str] = None

1944 + content: str | None = None

1945

1946

1947 class ChatMessage(OpenAIBaseModel):

1948 role: str

1951 - content: Optional[str] = None

1952 - refusal: Optional[str] = None

1953 - annotations: Optional[OpenAIAnnotation] = None

1954 - audio: Optional[OpenAIChatCompletionAudio] = None

1955 - function_call: Optional[FunctionCall] = None

1949 + content: str | None = None

1950 + refusal: str | None = None

1951 + annotations: OpenAIAnnotation | None = None

1952 + audio: OpenAIChatCompletionAudio | None = None

1953 + function_call: FunctionCall | None = None

1954 tool_calls: list[ToolCall] = Field(default_factory=list)

1955

1956 # vLLM-specific fields that are not in OpenAI spec

1959 - reasoning_content: Optional[str] = None

1957 + reasoning_content: str | None = None

1958

1959

1960 class ChatCompletionLogProb(OpenAIBaseModel):

1961 token: str

1962 logprob: float = -9999.0

1965 - bytes: Optional[list[int]] = None

1963 + bytes: list[int] | None = None

1964

1965

1966 class ChatCompletionLogProbsContent(ChatCompletionLogProb):

1967 # Workaround: redefine fields name cache so that it's not

1968 # shared with the super class.

1971 - field_names: ClassVar[Optional[set[str]]] = None

1969 + field_names: ClassVar[set[str] | None] = None

1970 top_logprobs: list[ChatCompletionLogProb] = Field(default_factory=list)

1971

1972

1973 class ChatCompletionLogProbs(OpenAIBaseModel):

1976 - content: Optional[list[ChatCompletionLogProbsContent]] = None

1974 + content: list[ChatCompletionLogProbsContent] | None = None

1975

1976

1977 class ChatCompletionResponseChoice(OpenAIBaseModel):

1978 index: int

1979 message: ChatMessage

1982 - logprobs: Optional[ChatCompletionLogProbs] = None

1980 + logprobs: ChatCompletionLogProbs | None = None

1981 # per OpenAI spec this is the default

1984 - finish_reason: Optional[str] = "stop"

1982 + finish_reason: str | None = "stop"

1983 # not part of the OpenAI spec but included in vLLM for legacy reasons

1986 - stop_reason: Optional[Union[int, str]] = None

1984 + stop_reason: int | str | None = None

1985 # not part of the OpenAI spec but is useful for tracing the tokens

1986 # in agent scenarios

1989 - token_ids: Optional[list[int]] = None

1987 + token_ids: list[int] | None = None

1988

1989

1990 class ChatCompletionResponse(OpenAIBaseModel):

@@ -1995,35 +1993,33 @@ class ChatCompletionResponse(OpenAIBaseModel):

1993 created: int = Field(default_factory=lambda: int(time.time()))

1994 model: str

1995 choices: list[ChatCompletionResponseChoice]

1998 - service_tier: Optional[Literal["auto", "default", "flex", "scale", "priority"]] = (

1999 - None

2000 - )

2001 - system_fingerprint: Optional[str] = None

1996 + service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None

1997 + system_fingerprint: str | None = None

1998 usage: UsageInfo

1999

2000 # vLLM-specific fields that are not in OpenAI spec

2005 - prompt_logprobs: Optional[list[Optional[dict[int, Logprob]]]] = None

2006 - prompt_token_ids: Optional[list[int]] = None

2007 - kv_transfer_params: Optional[dict[str, Any]] = Field(

2001 + prompt_logprobs: list[dict[int, Logprob] | None] | None = None

2002 + prompt_token_ids: list[int] | None = None

2003 + kv_transfer_params: dict[str, Any] | None = Field(

2004 default=None, description="KVTransfer parameters."

2005 )

2006

2007

2008 class DeltaMessage(OpenAIBaseModel):

2013 - role: Optional[str] = None

2014 - content: Optional[str] = None

2015 - reasoning_content: Optional[str] = None

2009 + role: str | None = None

2010 + content: str | None = None

2011 + reasoning_content: str | None = None

2012 tool_calls: list[DeltaToolCall] = Field(default_factory=list)

2013

2014

2015 class ChatCompletionResponseStreamChoice(OpenAIBaseModel):

2016 index: int

2017 delta: DeltaMessage

2022 - logprobs: Optional[ChatCompletionLogProbs] = None

2023 - finish_reason: Optional[str] = None

2024 - stop_reason: Optional[Union[int, str]] = None

2018 + logprobs: ChatCompletionLogProbs | None = None

2019 + finish_reason: str | None = None

2020 + stop_reason: int | str | None = None

2021 # not part of the OpenAI spec but for tracing the tokens

2026 - token_ids: Optional[list[int]] = None

2022 + token_ids: list[int] | None = None

2023

2024

2025 class ChatCompletionStreamResponse(OpenAIBaseModel):

@@ -2032,15 +2028,15 @@ class ChatCompletionStreamResponse(OpenAIBaseModel):

2028 created: int = Field(default_factory=lambda: int(time.time()))

2029 model: str

2030 choices: list[ChatCompletionResponseStreamChoice]

2035 - usage: Optional[UsageInfo] = Field(default=None)

2031 + usage: UsageInfo | None = Field(default=None)

2032 # not part of the OpenAI spec but for tracing the tokens

2037 - prompt_token_ids: Optional[list[int]] = None

2033 + prompt_token_ids: list[int] | None = None

2034

2035

2036 class TranscriptionResponseStreamChoice(OpenAIBaseModel):

2037 delta: DeltaMessage

2042 - finish_reason: Optional[str] = None

2043 - stop_reason: Optional[Union[int, str]] = None

2038 + finish_reason: str | None = None

2039 + stop_reason: int | str | None = None

2040

2041

2042 class TranscriptionStreamResponse(OpenAIBaseModel):

@@ -2049,7 +2045,7 @@ class TranscriptionStreamResponse(OpenAIBaseModel):

2045 created: int = Field(default_factory=lambda: int(time.time()))

2046 model: str

2047 choices: list[TranscriptionResponseStreamChoice]

2052 - usage: Optional[UsageInfo] = Field(default=None)

2048 + usage: UsageInfo | None = Field(default=None)

2049

2050

2051 class InputTokensDetails(OpenAIBaseModel):

@@ -2073,9 +2069,9 @@ class ResponsesResponse(OpenAIBaseModel):

2069 id: str = Field(default_factory=lambda: f"resp_{random_uuid()}")

2070 created_at: int = Field(default_factory=lambda: int(time.time()))

2071 # error: Optional[ResponseError] = None

2076 - incomplete_details: Optional[IncompleteDetails] = None

2077 - instructions: Optional[str] = None

2078 - metadata: Optional[Metadata] = None

2072 + incomplete_details: IncompleteDetails | None = None

2073 + instructions: str | None = None

2074 + metadata: Metadata | None = None

2075 model: str

2076 object: Literal["response"] = "response"

2077 output: list[ResponseOutputItem]

@@ -2086,24 +2082,24 @@ class ResponsesResponse(OpenAIBaseModel):

2082 top_p: float

2083 background: bool

2084 max_output_tokens: int

2089 - max_tool_calls: Optional[int] = None

2090 - previous_response_id: Optional[str] = None

2091 - prompt: Optional[ResponsePrompt] = None

2092 - reasoning: Optional[Reasoning] = None

2085 + max_tool_calls: int | None = None

2086 + previous_response_id: str | None = None

2087 + prompt: ResponsePrompt | None = None

2088 + reasoning: Reasoning | None = None

2089 service_tier: Literal["auto", "default", "flex", "scale", "priority"]

2090 status: ResponseStatus

2095 - text: Optional[ResponseTextConfig] = None

2096 - top_logprobs: Optional[int] = None

2091 + text: ResponseTextConfig | None = None

2092 + top_logprobs: int | None = None

2093 truncation: Literal["auto", "disabled"]

2098 - usage: Optional[ResponseUsage] = None

2099 - user: Optional[str] = None

2094 + usage: ResponseUsage | None = None

2095 + user: str | None = None

2096

2097 # --8<-- [start:responses-extra-params]

2098 # These are populated when enable_response_messages is set to True

2099 # NOTE: custom serialization is needed

2100 # see serialize_input_messages and serialize_output_messages

2105 - input_messages: Optional[list[ChatCompletionMessageParam]] = None

2106 - output_messages: Optional[list[ChatCompletionMessageParam]] = None

2101 + input_messages: list[ChatCompletionMessageParam] | None = None

2102 + output_messages: list[ChatCompletionMessageParam] | None = None

2103 # --8<-- [end:responses-extra-params]

2104

2105 # NOTE: openAI harmony doesn't serialize TextContent properly,

@@ -2150,11 +2146,11 @@ class ResponsesResponse(OpenAIBaseModel):

2146 created_time: int,

2147 output: list[ResponseOutputItem],

2148 status: ResponseStatus,

2153 - usage: Optional[ResponseUsage] = None,

2154 - input_messages: Optional[list[ChatCompletionMessageParam]] = None,

2155 - output_messages: Optional[list[ChatCompletionMessageParam]] = None,

2149 + usage: ResponseUsage | None = None,

2150 + input_messages: list[ChatCompletionMessageParam] | None = None,

2151 + output_messages: list[ChatCompletionMessageParam] | None = None,

2152 ) -> "ResponsesResponse":

2157 - incomplete_details: Optional[IncompleteDetails] = None

2153 + incomplete_details: IncompleteDetails | None = None

2154 if status == "incomplete":

2155 incomplete_details = IncompleteDetails(reason="max_output_tokens")

2156 # TODO: implement the other reason for incomplete_details,

@@ -2249,31 +2245,31 @@ class ResponseInProgressEvent(OpenAIResponseInProgressEvent):

2245 response: ResponsesResponse # type: ignore[override]

2246

2247

2252 -StreamingResponsesResponse: TypeAlias = Union[

2253 - "ResponseCreatedEvent",

2254 - "ResponseInProgressEvent",

2255 - "ResponseCompletedEvent",

2256 - ResponseOutputItemAddedEvent,

2257 - ResponseOutputItemDoneEvent,

2258 - ResponseContentPartAddedEvent,

2259 - ResponseContentPartDoneEvent,

2260 - ResponseReasoningTextDeltaEvent,

2261 - ResponseReasoningTextDoneEvent,

2262 - ResponseReasoningPartAddedEvent,

2263 - ResponseReasoningPartDoneEvent,

2264 - ResponseCodeInterpreterCallInProgressEvent,

2265 - ResponseCodeInterpreterCallCodeDeltaEvent,

2266 - ResponseWebSearchCallInProgressEvent,

2267 - ResponseWebSearchCallSearchingEvent,

2268 - ResponseWebSearchCallCompletedEvent,

2269 - ResponseCodeInterpreterCallCodeDoneEvent,

2270 - ResponseCodeInterpreterCallInterpretingEvent,

2271 - ResponseCodeInterpreterCallCompletedEvent,

2272 -]

2248 +StreamingResponsesResponse: TypeAlias = (

2249 + ResponseCreatedEvent

2250 + | ResponseInProgressEvent

2251 + | ResponseCompletedEvent

2252 + | ResponseOutputItemAddedEvent

2253 + | ResponseOutputItemDoneEvent

2254 + | ResponseContentPartAddedEvent

2255 + | ResponseContentPartDoneEvent

2256 + | ResponseReasoningTextDeltaEvent

2257 + | ResponseReasoningTextDoneEvent

2258 + | ResponseReasoningPartAddedEvent

2259 + | ResponseReasoningPartDoneEvent

2260 + | ResponseCodeInterpreterCallInProgressEvent

2261 + | ResponseCodeInterpreterCallCodeDeltaEvent

2262 + | ResponseWebSearchCallInProgressEvent

2263 + | ResponseWebSearchCallSearchingEvent

2264 + | ResponseWebSearchCallCompletedEvent

2265 + | ResponseCodeInterpreterCallCodeDoneEvent

2266 + | ResponseCodeInterpreterCallInterpretingEvent

2267 + | ResponseCodeInterpreterCallCompletedEvent

2268 +)

2269

2274 -BatchRequestInputBody = Union[

2275 - ChatCompletionRequest, EmbeddingRequest, ScoreRequest, RerankRequest

2276 -]

2270 +BatchRequestInputBody: TypeAlias = (

2271 + ChatCompletionRequest | EmbeddingRequest | ScoreRequest | RerankRequest

2272 +)

2273

2274

2275 class BatchRequestInput(OpenAIBaseModel):

@@ -2322,9 +2318,13 @@ class BatchResponseData(OpenAIBaseModel):

2318 request_id: str

2319

2320 # The body of the response.

2325 - body: Optional[

2326 - Union[ChatCompletionResponse, EmbeddingResponse, ScoreResponse, RerankResponse]

2327 - ] = None

2321 + body: (

2322 + ChatCompletionResponse

2323 + | EmbeddingResponse

2324 + | ScoreResponse

2325 + | RerankResponse

2326 + | None

2327 + ) = None

2328

2329

2330 class BatchRequestOutput(OpenAIBaseModel):

@@ -2338,15 +2338,15 @@ class BatchRequestOutput(OpenAIBaseModel):

2338 # inputs.

2339 custom_id: str

2340

2341 - response: Optional[BatchResponseData]

2341 + response: BatchResponseData | None

2342

2343 # For requests that failed with a non-HTTP error, this will contain more

2344 # information on the cause of the failure.

2345 - error: Optional[Any]

2345 + error: Any | None

2346

2347

2348 class TokenizeCompletionRequest(OpenAIBaseModel):

2349 - model: Optional[str] = None

2349 + model: str | None = None

2350 prompt: str

2351

2352 add_special_tokens: bool = Field(

@@ -2356,7 +2356,7 @@ class TokenizeCompletionRequest(OpenAIBaseModel):

2356 "the prompt."

2357 ),

2358 )

2359 - return_token_strs: Optional[bool] = Field(

2359 + return_token_strs: bool | None = Field(

2360 default=False,

2361 description=(

2362 "If true, also return the token strings corresponding to the token ids."

@@ -2365,7 +2365,7 @@ class TokenizeCompletionRequest(OpenAIBaseModel):

2365

2366

2367 class TokenizeChatRequest(OpenAIBaseModel):

2368 - model: Optional[str] = None

2368 + model: str | None = None

2369 messages: list[ChatCompletionMessageParam]

2370

2371 add_generation_prompt: bool = Field(

@@ -2376,7 +2376,7 @@ class TokenizeChatRequest(OpenAIBaseModel):

2376 "model."

2377 ),

2378 )

2379 - return_token_strs: Optional[bool] = Field(

2379 + return_token_strs: bool | None = Field(

2380 default=False,

2381 description=(

2382 "If true, also return the token strings corresponding to the token ids."

@@ -2402,7 +2402,7 @@ class TokenizeChatRequest(OpenAIBaseModel):

2402 "default)."

2403 ),

2404 )

2405 - chat_template: Optional[str] = Field(

2405 + chat_template: str | None = Field(

2406 default=None,

2407 description=(

2408 "A Jinja template to use for this conversion. "

@@ -2411,18 +2411,18 @@ class TokenizeChatRequest(OpenAIBaseModel):

2411 "does not define one."

2412 ),

2413 )

2414 - chat_template_kwargs: Optional[dict[str, Any]] = Field(

2414 + chat_template_kwargs: dict[str, Any] | None = Field(

2415 default=None,

2416 description=(

2417 "Additional keyword args to pass to the template renderer. "

2418 "Will be accessible by the chat template."

2419 ),

2420 )

2421 - mm_processor_kwargs: Optional[dict[str, Any]] = Field(

2421 + mm_processor_kwargs: dict[str, Any] | None = Field(

2422 default=None,

2423 description=("Additional kwargs to pass to the HF processor."),

2424 )

2425 - tools: Optional[list[ChatCompletionToolsParam]] = Field(

2425 + tools: list[ChatCompletionToolsParam] | None = Field(

2426 default=None,

2427 description=("A list of tools the model may call."),

2428 )

@@ -2438,18 +2438,18 @@ class TokenizeChatRequest(OpenAIBaseModel):

2438 return data

2439

2440

2441 -TokenizeRequest = Union[TokenizeCompletionRequest, TokenizeChatRequest]

2441 +TokenizeRequest: TypeAlias = TokenizeCompletionRequest | TokenizeChatRequest

2442

2443

2444 class TokenizeResponse(OpenAIBaseModel):

2445 count: int

2446 max_model_len: int

2447 tokens: list[int]

2448 - token_strs: Optional[list[str]] = None

2448 + token_strs: list[str] | None = None

2449

2450

2451 class DetokenizeRequest(OpenAIBaseModel):

2452 - model: Optional[str] = None

2452 + model: str | None = None

2453 tokens: list[int]

2454

2455

@@ -2474,7 +2474,7 @@ class LoadLoRAAdapterRequest(BaseModel):

2474

2475 class UnloadLoRAAdapterRequest(BaseModel):

2476 lora_name: str

2477 - lora_int_id: Optional[int] = Field(default=None)

2477 + lora_int_id: int | None = Field(default=None)

2478

2479

2480 ## Protocols for Audio

@@ -2491,11 +2491,11 @@ class TranscriptionRequest(OpenAIBaseModel):

2491 formats: flac, mp3, mp4, mpeg, mpga, m4a, ogg, wav, or webm.

2492 """

2493

2494 - model: Optional[str] = None

2494 + model: str | None = None

2495 """ID of the model to use.

2496 """

2497

2498 - language: Optional[str] = None

2498 + language: str | None = None

2499 """The language of the input audio.

2500

2501 Supplying the input language in

@@ -2530,16 +2530,16 @@ class TranscriptionRequest(OpenAIBaseModel):

2530 timestamps incurs additional latency.

2531 """

2532

2533 - stream: Optional[bool] = False

2533 + stream: bool | None = False

2534 """When set, it will enable output to be streamed in a similar fashion

2535 as the Chat Completion endpoint.

2536 """

2537 # --8<-- [start:transcription-extra-params]

2538 # Flattened stream option to simplify form data.

2539 - stream_include_usage: Optional[bool] = False

2540 - stream_continuous_usage_stats: Optional[bool] = False

2539 + stream_include_usage: bool | None = False

2540 + stream_continuous_usage_stats: bool | None = False

2541

2542 - vllm_xargs: Optional[dict[str, Union[str, int, float]]] = Field(

2542 + vllm_xargs: dict[str, str | int | float] | None = Field(

2543 default=None,

2544 description=(

2545 "Additional request parameters with string or "

@@ -2548,7 +2548,7 @@ class TranscriptionRequest(OpenAIBaseModel):

2548 )

2549 # --8<-- [end:transcription-extra-params]

2550

2551 - to_language: Optional[str] = None

2551 + to_language: str | None = None

2552 """The language of the output audio we transcribe to.

2553

2554 Please note that this is not currently used by supported models at this

@@ -2565,29 +2565,29 @@ class TranscriptionRequest(OpenAIBaseModel):

2565 to automatically increase the temperature until certain thresholds are hit.

2566 """

2567

2568 - top_p: Optional[float] = None

2568 + top_p: float | None = None

2569 """Enables nucleus (top-p) sampling, where tokens are selected from the

2570 smallest possible set whose cumulative probability exceeds `p`.

2571 """

2572

2573 - top_k: Optional[int] = None

2573 + top_k: int | None = None

2574 """Limits sampling to the `k` most probable tokens at each step."""

2575

2576 - min_p: Optional[float] = None

2576 + min_p: float | None = None

2577 """Filters out tokens with a probability lower than `min_p`, ensuring a

2578 minimum likelihood threshold during sampling.

2579 """

2580

2581 - seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

2581 + seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

2582 """The seed to use for sampling."""

2583

2584 - frequency_penalty: Optional[float] = 0.0

2584 + frequency_penalty: float | None = 0.0

2585 """The frequency penalty to use for sampling."""

2586

2587 - repetition_penalty: Optional[float] = None

2587 + repetition_penalty: float | None = None

2588 """The repetition penalty to use for sampling."""

2589

2590 - presence_penalty: Optional[float] = 0.0

2590 + presence_penalty: float | None = 0.0

2591 """The presence penalty to use for sampling."""

2592 # --8<-- [end:transcription-sampling-params]

2593

@@ -2601,7 +2601,7 @@ class TranscriptionRequest(OpenAIBaseModel):

2601 }

2602

2603 def to_sampling_params(

2604 - self, default_max_tokens: int, default_sampling_params: Optional[dict] = None

2604 + self, default_max_tokens: int, default_sampling_params: dict | None = None

2605 ) -> SamplingParams:

2606 max_tokens = default_max_tokens

2607

@@ -2740,17 +2740,17 @@ class TranscriptionResponseVerbose(OpenAIBaseModel):

2740 text: str

2741 """The transcribed text."""

2742

2743 - segments: Optional[list[TranscriptionSegment]] = None

2743 + segments: list[TranscriptionSegment] | None = None

2744 """Segments of the transcribed text and their corresponding details."""

2745

2746 - words: Optional[list[TranscriptionWord]] = None

2746 + words: list[TranscriptionWord] | None = None

2747 """Extracted words and their corresponding timestamps."""

2748

2749

2750 class TranslationResponseStreamChoice(OpenAIBaseModel):

2751 delta: DeltaMessage

2752 - finish_reason: Optional[str] = None

2753 - stop_reason: Optional[Union[int, str]] = None

2752 + finish_reason: str | None = None

2753 + stop_reason: int | str | None = None

2754

2755

2756 class TranslationStreamResponse(OpenAIBaseModel):

@@ -2759,7 +2759,7 @@ class TranslationStreamResponse(OpenAIBaseModel):

2759 created: int = Field(default_factory=lambda: int(time.time()))

2760 model: str

2761 choices: list[TranslationResponseStreamChoice]

2762 - usage: Optional[UsageInfo] = Field(default=None)

2762 + usage: UsageInfo | None = Field(default=None)

2763

2764

2765 class TranslationRequest(OpenAIBaseModel):

@@ -2772,7 +2772,7 @@ class TranslationRequest(OpenAIBaseModel):

2772 formats: flac, mp3, mp4, mpeg, mpga, m4a, ogg, wav, or webm.

2773 """

2774

2775 - model: Optional[str] = None

2775 + model: str | None = None

2776 """ID of the model to use.

2777 """

2778

@@ -2792,7 +2792,7 @@ class TranslationRequest(OpenAIBaseModel):

2792

2793 # TODO support additional sampling parameters

2794 # --8<-- [start:translation-sampling-params]

2795 - seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

2795 + seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)

2796 """The seed to use for sampling."""

2797

2798 temperature: float = Field(default=0.0)

@@ -2806,7 +2806,7 @@ class TranslationRequest(OpenAIBaseModel):

2806 # --8<-- [end:translation-sampling-params]

2807

2808 # --8<-- [start:translation-extra-params]

2809 - language: Optional[str] = None

2809 + language: str | None = None

2810 """The language of the input audio we translate from.

2811

2812 Supplying the input language in

@@ -2814,7 +2814,7 @@ class TranslationRequest(OpenAIBaseModel):

2814 will improve accuracy.

2815 """

2816

2817 - to_language: Optional[str] = None

2817 + to_language: str | None = None

2818 """The language of the input audio we translate to.

2819

2820 Please note that this is not supported by all models, refer to the specific

@@ -2822,14 +2822,14 @@ class TranslationRequest(OpenAIBaseModel):

2822 For instance, Whisper only supports `to_language=en`.

2823 """

2824

2825 - stream: Optional[bool] = False

2825 + stream: bool | None = False

2826 """Custom field not present in the original OpenAI definition. When set,

2827 it will enable output to be streamed in a similar fashion as the Chat

2828 Completion endpoint.

2829 """

2830 # Flattened stream option to simplify form data.

2831 - stream_include_usage: Optional[bool] = False

2832 - stream_continuous_usage_stats: Optional[bool] = False

2831 + stream_include_usage: bool | None = False

2832 + stream_continuous_usage_stats: bool | None = False

2833 # --8<-- [end:translation-extra-params]

2834

2835 # Default sampling parameters for translation requests.

@@ -2838,7 +2838,7 @@ class TranslationRequest(OpenAIBaseModel):

2838 }

2839

2840 def to_sampling_params(

2841 - self, default_max_tokens: int, default_sampling_params: Optional[dict] = None

2841 + self, default_max_tokens: int, default_sampling_params: dict | None = None

2842 ) -> SamplingParams:

2843 max_tokens = default_max_tokens

2844

@@ -2939,8 +2939,8 @@ class TranslationResponseVerbose(OpenAIBaseModel):

2939 text: str

2940 """The translated text."""

2941

2942 - segments: Optional[list[TranslationSegment]] = None

2942 + segments: list[TranslationSegment] | None = None

2943 """Segments of the translated text and their corresponding details."""

2944

2945 - words: Optional[list[TranslationWord]] = None

2945 + words: list[TranslationWord] | None = None

2946 """Extracted words and their corresponding timestamps."""

vllm/entrypoints/openai/run_batch.py

+2 -3

@@ -4,10 +4,9 @@

4 import asyncio

5 import tempfile

6 from argparse import Namespace

7 -from collections.abc import Awaitable

7 +from collections.abc import Awaitable, Callable

8 from http import HTTPStatus

9 from io import StringIO

10 -from typing import Callable, Optional

10

11 import aiohttp

12 import torch

@@ -124,7 +123,7 @@ _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elap

123 class BatchProgressTracker:

124 def __init__(self):

125 self._total = 0

127 - self._pbar: Optional[tqdm] = None

126 + self._pbar: tqdm | None = None

127

128 def submitted(self):

129 self._total += 1

vllm/entrypoints/openai/serving_chat.py

+22 -22

@@ -6,7 +6,7 @@ import json

6 import time

7 from collections.abc import AsyncGenerator, AsyncIterator

8 from collections.abc import Sequence as GenericSequence

9 -from typing import Final, Optional, Union

9 +from typing import Final

10

11 import jinja2

12 import partial_json_parser

@@ -82,15 +82,15 @@ class OpenAIServingChat(OpenAIServing):

82 models: OpenAIServingModels,

83 response_role: str,

84 *,

85 - request_logger: Optional[RequestLogger],

86 - chat_template: Optional[str],

85 + request_logger: RequestLogger | None,

86 + chat_template: str | None,

87 chat_template_content_format: ChatTemplateContentFormatOption,

88 trust_request_chat_template: bool = False,

89 return_tokens_as_token_ids: bool = False,

90 reasoning_parser: str = "",

91 enable_auto_tools: bool = False,

92 exclude_tools_when_tool_choice_none: bool = False,

93 - tool_parser: Optional[str] = None,

93 + tool_parser: str | None = None,

94 enable_prompt_tokens_details: bool = False,

95 enable_force_include_usage: bool = False,

96 enable_log_outputs: bool = False,

@@ -159,8 +159,8 @@ class OpenAIServingChat(OpenAIServing):

159 async def create_chat_completion(

160 self,

161 request: ChatCompletionRequest,

162 - raw_request: Optional[Request] = None,

163 - ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, ErrorResponse]:

162 + raw_request: Request | None = None,

163 + ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:

164 """

165 Chat Completion API similar to OpenAI's API.

166

@@ -281,7 +281,7 @@ class OpenAIServingChat(OpenAIServing):

281 default_sampling_params=self.default_sampling_params,

282 )

283

284 - sampling_params: Union[SamplingParams, BeamSearchParams]

284 + sampling_params: SamplingParams | BeamSearchParams

285 if request.use_beam_search:

286 sampling_params = request.to_beam_search_params(

287 max_tokens, self.default_sampling_params

@@ -416,11 +416,11 @@ class OpenAIServingChat(OpenAIServing):

416 def extract_tool_call_required_streaming(

417 self,

418 previous_text: str,

419 - current_text: Optional[str],

419 + current_text: str | None,

420 delta_text: str,

421 function_name_returned: bool,

422 - tool_call_idx: Optional[int] = None,

423 - ) -> tuple[Optional[DeltaMessage], bool]:

422 + tool_call_idx: int | None = None,

423 + ) -> tuple[DeltaMessage | None, bool]:

424 if current_text is None or current_text == "":

425 # if the current text is empty, we cannot parse it

426 return None, function_name_returned

@@ -548,7 +548,7 @@ class OpenAIServingChat(OpenAIServing):

548 and self._should_stream_with_auto_tool_parsing(request)

549 )

550

551 - all_previous_token_ids: Optional[list[list[int]]]

551 + all_previous_token_ids: list[list[int]] | None

552 function_name_returned = [False] * num_choices

553 if self.tool_call_id_type == "kimi_k2":

554 history_tool_call_cnt = get_history_tool_calls_cnt(conversation)

@@ -583,7 +583,7 @@ class OpenAIServingChat(OpenAIServing):

583 # Prepare the tool parser if it's needed

584 try:

585 if tool_choice_auto and self.tool_parser:

586 - tool_parsers: list[Optional[ToolParser]] = [

586 + tool_parsers: list[ToolParser | None] = [

587 self.tool_parser(tokenizer)

588 ] * num_choices

589 else:

@@ -661,7 +661,7 @@ class OpenAIServingChat(OpenAIServing):

661 # Send response to echo the input portion of the

662 # last message

663 if request.echo:

664 - last_msg_content: Union[str, list[dict[str, str]]] = ""

664 + last_msg_content: str | list[dict[str, str]] = ""

665 if (

666 conversation

667 and "content" in conversation[-1]

@@ -734,7 +734,7 @@ class OpenAIServingChat(OpenAIServing):

734 # Chunked prefill case, don't return empty chunks

735 continue

736

737 - delta_message: Optional[DeltaMessage]

737 + delta_message: DeltaMessage | None

738

739 # just update previous_texts and previous_token_ids

740 if tool_choice_auto or self.reasoning_parser:

@@ -1260,9 +1260,9 @@ class OpenAIServingChat(OpenAIServing):

1260 conversation: list[ConversationMessage],

1261 tokenizer: AnyTokenizer,

1262 request_metadata: RequestResponseMetadata,

1263 - ) -> Union[ErrorResponse, ChatCompletionResponse]:

1263 + ) -> ErrorResponse | ChatCompletionResponse:

1264 created_time = int(time.time())

1265 - final_res: Optional[RequestOutput] = None

1265 + final_res: RequestOutput | None = None

1266

1267 try:

1268 async for res in result_generator:

@@ -1512,7 +1512,7 @@ class OpenAIServingChat(OpenAIServing):

1512 choices.append(choice_data)

1513

1514 if request.echo:

1515 - last_msg_content: Union[str, list[dict[str, str]]] = ""

1515 + last_msg_content: str | list[dict[str, str]] = ""

1516 if (

1517 conversation

1518 and "content" in conversation[-1]

@@ -1597,7 +1597,7 @@ class OpenAIServingChat(OpenAIServing):

1597 def _get_top_logprobs(

1598 self,

1599 logprobs: dict[int, Logprob],

1600 - top_logprobs: Optional[int],

1600 + top_logprobs: int | None,

1601 tokenizer: AnyTokenizer,

1602 should_return_as_token_id: bool,

1603 ) -> list[ChatCompletionLogProb]:

@@ -1621,10 +1621,10 @@ class OpenAIServingChat(OpenAIServing):

1621 def _create_chat_logprobs(

1622 self,

1623 token_ids: GenericSequence[int],

1624 - top_logprobs: GenericSequence[Optional[dict[int, Logprob]]],

1624 + top_logprobs: GenericSequence[dict[int, Logprob] | None],

1625 tokenizer: AnyTokenizer,

1626 - num_output_top_logprobs: Optional[int] = None,

1627 - return_as_token_id: Optional[bool] = None,

1626 + num_output_top_logprobs: int | None = None,

1627 + return_as_token_id: bool | None = None,

1628 ) -> ChatCompletionLogProbs:

1629 """Create OpenAI-style logprobs."""

1630 logprobs_content: list[ChatCompletionLogProbsContent] = []

@@ -1693,7 +1693,7 @@ class OpenAIServingChat(OpenAIServing):

1693

1694 def _should_check_for_unstreamed_tool_arg_tokens(

1695 self,

1696 - delta_message: Optional[DeltaMessage],

1696 + delta_message: DeltaMessage | None,

1697 output: CompletionOutput,

1698 ) -> bool:

1699 """

vllm/entrypoints/openai/serving_classification.py

+6 -6

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

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

3

4 from http import HTTPStatus

5 -from typing import Optional, Union, cast

5 +from typing import cast

6

7 import numpy as np

8 from fastapi import Request

@@ -36,7 +36,7 @@ class ClassificationMixin(OpenAIServing):

36 async def _preprocess(

37 self,

38 ctx: ServeContext,

39 - ) -> Optional[ErrorResponse]:

39 + ) -> ErrorResponse | None:

40 """

41 Process classification inputs: tokenize text, resolve adapters,

42 and prepare model-specific inputs.

@@ -70,7 +70,7 @@ class ClassificationMixin(OpenAIServing):

70 def _build_response(

71 self,

72 ctx: ServeContext,

73 - ) -> Union[ClassificationResponse, ErrorResponse]:

73 + ) -> ClassificationResponse | ErrorResponse:

74 """

75 Convert model outputs to a formatted classification response

76 with probabilities and labels.

@@ -129,7 +129,7 @@ class ServingClassification(ClassificationMixin):

129 engine_client: EngineClient,

130 models: OpenAIServingModels,

131 *,

132 - request_logger: Optional[RequestLogger],

132 + request_logger: RequestLogger | None,

133 log_error_stack: bool = False,

134 ) -> None:

135 super().__init__(

@@ -143,7 +143,7 @@ class ServingClassification(ClassificationMixin):

143 self,

144 request: ClassificationRequest,

145 raw_request: Request,

146 - ) -> Union[ClassificationResponse, ErrorResponse]:

146 + ) -> ClassificationResponse | ErrorResponse:

147 model_name = self.models.model_name()

148 request_id = f"{self.request_id_prefix}-{self._base_request_id(raw_request)}"

149

@@ -160,7 +160,7 @@ class ServingClassification(ClassificationMixin):

160 def _create_pooling_params(

161 self,

162 ctx: ClassificationServeContext,

163 - ) -> Union[PoolingParams, ErrorResponse]:

163 + ) -> PoolingParams | ErrorResponse:

164 pooling_params = super()._create_pooling_params(ctx)

165 if isinstance(pooling_params, ErrorResponse):

166 return pooling_params

vllm/entrypoints/openai/serving_completion.py

+16 -16

@@ -5,7 +5,7 @@ import asyncio

5 import time

6 from collections.abc import AsyncGenerator, AsyncIterator

7 from collections.abc import Sequence as GenericSequence

8 -from typing import Optional, Union, cast

8 +from typing import cast

9

10 import jinja2

11 from fastapi import Request

@@ -45,7 +45,7 @@ class OpenAIServingCompletion(OpenAIServing):

45 engine_client: EngineClient,

46 models: OpenAIServingModels,

47 *,

48 - request_logger: Optional[RequestLogger],

48 + request_logger: RequestLogger | None,

49 return_tokens_as_token_ids: bool = False,

50 enable_prompt_tokens_details: bool = False,

51 enable_force_include_usage: bool = False,

@@ -73,8 +73,8 @@ class OpenAIServingCompletion(OpenAIServing):

73 async def create_completion(

74 self,

75 request: CompletionRequest,

76 - raw_request: Optional[Request] = None,

77 - ) -> Union[AsyncGenerator[str, None], CompletionResponse, ErrorResponse]:

76 + raw_request: Request | None = None,

77 + ) -> AsyncGenerator[str, None] | CompletionResponse | ErrorResponse:

78 """Completion API similar to OpenAI's API.

79

80 See https://platform.openai.com/docs/api-reference/completions/create

@@ -166,7 +166,7 @@ class OpenAIServingCompletion(OpenAIServing):

166 default_sampling_params=self.default_sampling_params,

167 )

168

169 - sampling_params: Union[SamplingParams, BeamSearchParams]

169 + sampling_params: SamplingParams | BeamSearchParams

170 if request.use_beam_search:

171 sampling_params = request.to_beam_search_params(

172 max_tokens, self.default_sampling_params

@@ -196,7 +196,7 @@ class OpenAIServingCompletion(OpenAIServing):

196 # Mypy inconsistently requires this second cast in different

197 # environments. It shouldn't be necessary (redundant from above)

198 # but pre-commit in CI fails without it.

199 - engine_prompt = cast(Union[EmbedsPrompt, TokensPrompt], engine_prompt)

199 + engine_prompt = cast(EmbedsPrompt | TokensPrompt, engine_prompt)

200 if isinstance(sampling_params, BeamSearchParams):

201 generator = self.beam_search(

202 prompt=engine_prompt,

@@ -260,7 +260,7 @@ class OpenAIServingCompletion(OpenAIServing):

260 )

261

262 # Non-streaming response

263 - final_res_batch: list[Optional[RequestOutput]] = [None] * num_prompts

263 + final_res_batch: list[RequestOutput | None] = [None] * num_prompts

264 try:

265 async for i, res in result_generator:

266 final_res_batch[i] = res

@@ -312,7 +312,7 @@ class OpenAIServingCompletion(OpenAIServing):

312 async def completion_stream_generator(

313 self,

314 request: CompletionRequest,

315 - engine_prompts: list[Union[TokensPrompt, EmbedsPrompt]],

315 + engine_prompts: list[TokensPrompt | EmbedsPrompt],

316 result_generator: AsyncIterator[tuple[int, RequestOutput]],

317 request_id: str,

318 created_time: int,

@@ -362,7 +362,7 @@ class OpenAIServingCompletion(OpenAIServing):

362 num_prompt_tokens[prompt_idx] = len(prompt_token_ids)

363

364 delta_token_ids: GenericSequence[int]

365 - out_logprobs: Optional[GenericSequence[Optional[dict[int, Logprob]]]]

365 + out_logprobs: GenericSequence[dict[int, Logprob] | None] | None

366

367 for output in res.outputs:

368 i = output.index + prompt_idx * num_choices

@@ -370,7 +370,7 @@ class OpenAIServingCompletion(OpenAIServing):

370 # Useful when request.return_token_ids is True

371 # Returning prompt token IDs shares the same logic

372 # with the echo implementation.

373 - prompt_token_ids_to_return: Optional[list[int]] = None

373 + prompt_token_ids_to_return: list[int] | None = None

374

375 assert request.max_tokens is not None

376 if request.echo and not has_echoed[i]:

@@ -524,7 +524,7 @@ class OpenAIServingCompletion(OpenAIServing):

524 prompt_text = final_res.prompt

525

526 token_ids: GenericSequence[int]

527 - out_logprobs: Optional[GenericSequence[Optional[dict[int, Logprob]]]]

527 + out_logprobs: GenericSequence[dict[int, Logprob] | None] | None

528

529 for output in final_res.outputs:

530 assert request.max_tokens is not None

@@ -617,17 +617,17 @@ class OpenAIServingCompletion(OpenAIServing):

617 def _create_completion_logprobs(

618 self,

619 token_ids: GenericSequence[int],

620 - top_logprobs: GenericSequence[Optional[dict[int, Logprob]]],

620 + top_logprobs: GenericSequence[dict[int, Logprob] | None],

621 num_output_top_logprobs: int,

622 tokenizer: AnyTokenizer,

623 initial_text_offset: int = 0,

624 - return_as_token_id: Optional[bool] = None,

624 + return_as_token_id: bool | None = None,

625 ) -> CompletionLogProbs:

626 """Create logprobs for OpenAI Completion API."""

627 out_text_offset: list[int] = []

628 - out_token_logprobs: list[Optional[float]] = []

628 + out_token_logprobs: list[float | None] = []

629 out_tokens: list[str] = []

630 - out_top_logprobs: list[Optional[dict[str, float]]] = []

630 + out_top_logprobs: list[dict[str, float] | None] = []

631

632 last_token_len = 0

633

@@ -695,7 +695,7 @@ class OpenAIServingCompletion(OpenAIServing):

695 def _build_render_config(

696 self,

697 request: CompletionRequest,

698 - max_input_length: Optional[int] = None,

698 + max_input_length: int | None = None,

699 ) -> RenderConfig:

700 max_input_tokens_len = self.max_model_len - (request.max_tokens or 0)

701 return RenderConfig(

vllm/entrypoints/openai/serving_embedding.py

+17 -19

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

3

4 import base64

5 from collections.abc import AsyncGenerator, Mapping

6 -from typing import Any, Final, Literal, Optional, Union, cast

6 +from typing import Any, Final, Literal, cast

7

8 import numpy as np

9 import torch

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

48 def _get_embedding(

49 output: EmbeddingOutput,

50 encoding_format: Literal["float", "base64"],

51 -) -> Union[list[float], str]:

51 +) -> list[float] | str:

52 if encoding_format == "float":

53 return output.embedding

54 elif encoding_format == "base64":

@@ -80,7 +80,7 @@ class EmbeddingMixin(OpenAIServing):

80 async def _preprocess(

81 self,

82 ctx: ServeContext,

83 - ) -> Optional[ErrorResponse]:

83 + ) -> ErrorResponse | None:

84 ctx = cast(EmbeddingServeContext, ctx)

85 try:

86 ctx.lora_request = self._maybe_get_adapters(ctx.request)

@@ -130,7 +130,7 @@ class EmbeddingMixin(OpenAIServing):

130 def _build_response(

131 self,

132 ctx: ServeContext,

133 - ) -> Union[EmbeddingResponse, ErrorResponse]:

133 + ) -> EmbeddingResponse | ErrorResponse:

134 items: list[EmbeddingResponseData] = []

135 num_prompt_tokens = 0

136

@@ -314,9 +314,9 @@ class EmbeddingMixin(OpenAIServing):

314 ctx: EmbeddingServeContext,

315 engine_prompt: EngineTokensPrompt,

316 pooling_params: PoolingParams,

317 - trace_headers: Optional[Mapping[str, str]],

317 + trace_headers: Mapping[str, str] | None,

318 prompt_index: int,

319 - ) -> AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]:

319 + ) -> AsyncGenerator[RequestOutput | PoolingRequestOutput, None]:

320 """Create a generator for a single prompt using standard processing."""

321 request_id_item = f"{ctx.request_id}-{prompt_index}"

322

@@ -341,7 +341,7 @@ class EmbeddingMixin(OpenAIServing):

341 async def _prepare_generators(

342 self,

343 ctx: ServeContext,

344 - ) -> Optional[ErrorResponse]:

344 + ) -> ErrorResponse | None:

345 """Override to support chunked processing."""

346 ctx = cast(EmbeddingServeContext, ctx)

347

@@ -354,7 +354,7 @@ class EmbeddingMixin(OpenAIServing):

354

355 # Custom logic for chunked processing

356 generators: list[

357 - AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]

357 + AsyncGenerator[RequestOutput | PoolingRequestOutput, None]

358 ] = []

359

360 try:

@@ -413,7 +413,7 @@ class EmbeddingMixin(OpenAIServing):

413 async def _collect_batch(

414 self,

415 ctx: ServeContext,

416 - ) -> Optional[ErrorResponse]:

416 + ) -> ErrorResponse | None:

417 """Collect and aggregate batch results

418 with support for chunked processing.

419

@@ -522,9 +522,7 @@ class EmbeddingMixin(OpenAIServing):

522 )

523

524 # Finalize aggregated results

525 - final_res_batch: list[

526 - Union[PoolingRequestOutput, EmbeddingRequestOutput]

527 - ] = []

525 + final_res_batch: list[PoolingRequestOutput | EmbeddingRequestOutput] = []

526 num_prompts = len(ctx.engine_prompts)

527

528 for prompt_idx in range(num_prompts):

@@ -581,7 +579,7 @@ class EmbeddingMixin(OpenAIServing):

579 )

580

581 ctx.final_res_batch = cast(

584 - list[Union[RequestOutput, PoolingRequestOutput]], final_res_batch

582 + list[RequestOutput | PoolingRequestOutput], final_res_batch

583 )

584

585 return None

@@ -598,8 +596,8 @@ class OpenAIServingEmbedding(EmbeddingMixin):

596 engine_client: EngineClient,

597 models: OpenAIServingModels,

598 *,

601 - request_logger: Optional[RequestLogger],

602 - chat_template: Optional[str],

599 + request_logger: RequestLogger | None,

600 + chat_template: str | None,

601 chat_template_content_format: ChatTemplateContentFormatOption,

602 trust_request_chat_template: bool = False,

603 log_error_stack: bool = False,

@@ -618,8 +616,8 @@ class OpenAIServingEmbedding(EmbeddingMixin):

616 async def create_embedding(

617 self,

618 request: EmbeddingRequest,

621 - raw_request: Optional[Request] = None,

622 - ) -> Union[EmbeddingResponse, ErrorResponse]:

619 + raw_request: Request | None = None,

620 + ) -> EmbeddingResponse | ErrorResponse:

621 """

622 Embedding API similar to OpenAI's API.

623

@@ -647,7 +645,7 @@ class OpenAIServingEmbedding(EmbeddingMixin):

645 def _create_pooling_params(

646 self,

647 ctx: ServeContext[EmbeddingRequest],

650 - ) -> Union[PoolingParams, ErrorResponse]:

648 + ) -> PoolingParams | ErrorResponse:

649 pooling_params = super()._create_pooling_params(ctx)

650 if isinstance(pooling_params, ErrorResponse):

651 return pooling_params

@@ -662,7 +660,7 @@ class OpenAIServingEmbedding(EmbeddingMixin):

660 async def _preprocess(

661 self,

662 ctx: ServeContext,

665 - ) -> Optional[ErrorResponse]:

663 + ) -> ErrorResponse | None:

664 if isinstance(ctx.request, EmbeddingChatRequest):

665 error_check_ret = self._validate_chat_template(

666 request_chat_template=ctx.request.chat_template,

vllm/entrypoints/openai/serving_engine.py

+94 -96

@@ -5,10 +5,10 @@ import json

5 import sys

6 import time

7 import traceback

8 -from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence

8 +from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence

9 from concurrent.futures import ThreadPoolExecutor

10 from http import HTTPStatus

11 -from typing import Any, Callable, ClassVar, Generic, Optional, TypeVar, Union

11 +from typing import Any, ClassVar, Generic, TypeAlias, TypeVar

12

13 import torch

14 from fastapi import Request

@@ -102,38 +102,38 @@ from vllm.v1.engine import EngineCoreRequest

102

103 logger = init_logger(__name__)

104

105 -CompletionLikeRequest = Union[

106 - CompletionRequest,

107 - DetokenizeRequest,

108 - EmbeddingCompletionRequest,

109 - RerankRequest,

110 - ClassificationRequest,

111 - ScoreRequest,

112 - TokenizeCompletionRequest,

113 -]

114 -

115 -ChatLikeRequest = Union[

116 - ChatCompletionRequest, EmbeddingChatRequest, TokenizeChatRequest

117 -]

118 -SpeechToTextRequest = Union[TranscriptionRequest, TranslationRequest]

119 -AnyRequest = Union[

120 - CompletionLikeRequest,

121 - ChatLikeRequest,

122 - SpeechToTextRequest,

123 - ResponsesRequest,

124 - IOProcessorRequest,

125 -]

105 +CompletionLikeRequest: TypeAlias = (

106 + CompletionRequest

107 + | DetokenizeRequest

108 + | EmbeddingCompletionRequest

109 + | RerankRequest

110 + | ClassificationRequest

111 + | ScoreRequest

112 + | TokenizeCompletionRequest

113 +)

114

127 -AnyResponse = Union[

128 - CompletionResponse,

129 - ChatCompletionResponse,

130 - EmbeddingResponse,

131 - TranscriptionResponse,

132 - TokenizeResponse,

133 - PoolingResponse,

134 - ClassificationResponse,

135 - ScoreResponse,

136 -]

115 +ChatLikeRequest: TypeAlias = (

116 + ChatCompletionRequest | EmbeddingChatRequest | TokenizeChatRequest

117 +)

118 +SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest

119 +AnyRequest: TypeAlias = (

120 + CompletionLikeRequest

121 + | ChatLikeRequest

122 + | SpeechToTextRequest

123 + | ResponsesRequest

124 + | IOProcessorRequest

125 +)

126 +

127 +AnyResponse: TypeAlias = (

128 + CompletionResponse

129 + | ChatCompletionResponse

130 + | EmbeddingResponse

131 + | TranscriptionResponse

132 + | TokenizeResponse

133 + | PoolingResponse

134 + | ClassificationResponse

135 + | ScoreResponse

136 +)

137

138

139 class TextTokensPrompt(TypedDict):

@@ -145,7 +145,7 @@ class EmbedsPrompt(TypedDict):

145 prompt_embeds: torch.Tensor

146

147

148 -RequestPrompt = Union[list[int], str, TextTokensPrompt, EmbedsPrompt]

148 +RequestPrompt: TypeAlias = list[int] | str | TextTokensPrompt | EmbedsPrompt

149

150

151 def is_text_tokens_prompt(prompt: RequestPrompt) -> TypeIs[TextTokensPrompt]:

@@ -173,8 +173,8 @@ class RequestProcessingMixin(BaseModel):

173 handling prompt preparation and engine input.

174 """

175

176 - request_prompts: Optional[Sequence[RequestPrompt]] = []

177 - engine_prompts: Optional[list[EngineTokensPrompt]] = []

176 + request_prompts: Sequence[RequestPrompt] | None = []

177 + engine_prompts: list[EngineTokensPrompt] | None = []

178

179 model_config = ConfigDict(arbitrary_types_allowed=True)

180

@@ -185,10 +185,10 @@ class ResponseGenerationMixin(BaseModel):

185 managing result generators and final batch results.

186 """

187

188 - result_generator: Optional[

189 - AsyncGenerator[tuple[int, Union[RequestOutput, PoolingRequestOutput]], None]

190 - ] = None

191 - final_res_batch: list[Union[RequestOutput, PoolingRequestOutput]] = Field(

188 + result_generator: (

189 + AsyncGenerator[tuple[int, RequestOutput | PoolingRequestOutput], None] | None

190 + ) = None

191 + final_res_batch: list[RequestOutput | PoolingRequestOutput] = Field(

192 default_factory=list

193 )

194

@@ -203,14 +203,14 @@ class ServeContext(

203 ):

204 # Shared across all requests

205 request: RequestT

206 - raw_request: Optional[Request] = None

206 + raw_request: Request | None = None

207 model_name: str

208 request_id: str

209 created_time: int = Field(default_factory=lambda: int(time.time()))

210 - lora_request: Optional[LoRARequest] = None

210 + lora_request: LoRARequest | None = None

211

212 # Shared across most requests

213 - tokenizer: Optional[AnyTokenizer] = None

213 + tokenizer: AnyTokenizer | None = None

214

215 # `protected_namespaces` resolves Pydantic v2's warning

216 # on conflict with protected namespace "model_"

@@ -224,7 +224,7 @@ ClassificationServeContext = ServeContext[ClassificationRequest]

224

225

226 class EmbeddingServeContext(ServeContext[EmbeddingRequest]):

227 - chat_template: Optional[str] = None

227 + chat_template: str | None = None

228 chat_template_content_format: ChatTemplateContentFormatOption

229

230

@@ -247,7 +247,7 @@ class OpenAIServing:

247 engine_client: EngineClient,

248 models: OpenAIServingModels,

249 *,

250 - request_logger: Optional[RequestLogger],

250 + request_logger: RequestLogger | None,

251 return_tokens_as_token_ids: bool = False,

252 enable_force_include_usage: bool = False,

253 log_error_stack: bool = False,

@@ -276,8 +276,8 @@ class OpenAIServing:

276 self.max_model_len = self.model_config.max_model_len

277

278 def _get_tool_parser(

279 - self, tool_parser_name: Optional[str] = None, enable_auto_tools: bool = False

280 - ) -> Optional[Callable[[AnyTokenizer], ToolParser]]:

279 + self, tool_parser_name: str | None = None, enable_auto_tools: bool = False

280 + ) -> Callable[[AnyTokenizer], ToolParser] | None:

281 """Get the tool parser based on the name."""

282 parser = None

283 if not enable_auto_tools or tool_parser_name is None:

@@ -307,7 +307,7 @@ class OpenAIServing:

307 def _get_reasoning_parser(

308 self,

309 reasoning_parser_name: str,

310 - ) -> Optional[Callable[[AnyTokenizer], ReasoningParser]]:

310 + ) -> Callable[[AnyTokenizer], ReasoningParser] | None:

311 """Get the reasoning parser based on the name."""

312 parser = None

313 if not reasoning_parser_name:

@@ -328,7 +328,7 @@ class OpenAIServing:

328 prompt: PromptType,

329 request_id: str,

330 params: BeamSearchParams,

331 - lora_request: Optional[LoRARequest] = None,

331 + lora_request: LoRARequest | None = None,

332 ) -> AsyncGenerator[RequestOutput, None]:

333 beam_width = params.beam_width

334 max_tokens = params.max_tokens

@@ -364,9 +364,9 @@ class OpenAIServing:

364 # this happens again in generation, so the double expansion causes

365 # a mismatch.

366 # TODO - would be ideal to handle this more gracefully.

367 - prompt_text: Optional[str]

367 + prompt_text: str | None

368 prompt_token_ids: list[int]

369 - multi_modal_data: Optional[MultiModalDataDict]

369 + multi_modal_data: MultiModalDataDict | None

370 if isinstance(prompt, str):

371 prompt_text = prompt

372 prompt_token_ids = []

@@ -376,7 +376,7 @@ class OpenAIServing:

376 prompt_token_ids = prompt.get("prompt_token_ids", []) # type: ignore

377 multi_modal_data = prompt.get("multi_modal_data") # type: ignore

378

379 - mm_processor_kwargs: Optional[dict[str, Any]] = processed_inputs.get(

379 + mm_processor_kwargs: dict[str, Any] | None = processed_inputs.get(

380 "mm_processor_kwargs"

381 ) # type: ignore

382

@@ -507,7 +507,7 @@ class OpenAIServing:

507 prompt_logprobs=None,

508 )

509

510 - def _get_renderer(self, tokenizer: Optional[AnyTokenizer]) -> BaseRenderer:

510 + def _get_renderer(self, tokenizer: AnyTokenizer | None) -> BaseRenderer:

511 """

512 Get a Renderer instance with the provided tokenizer.

513 Uses shared async tokenizer pool for efficiency.

@@ -545,7 +545,7 @@ class OpenAIServing:

545 async def _preprocess(

546 self,

547 ctx: ServeContext,

548 - ) -> Optional[ErrorResponse]:

548 + ) -> ErrorResponse | None:

549 """

550 Default preprocessing hook. Subclasses may override

551 to prepare `ctx` (classification, embedding, etc.).

@@ -555,7 +555,7 @@ class OpenAIServing:

555 def _build_response(

556 self,

557 ctx: ServeContext,

558 - ) -> Union[AnyResponse, ErrorResponse]:

558 + ) -> AnyResponse | ErrorResponse:

559 """

560 Default response builder. Subclass may override this method

561 to return the appropriate response object.

@@ -565,8 +565,8 @@ class OpenAIServing:

565 async def handle(

566 self,

567 ctx: ServeContext,

568 - ) -> Union[AnyResponse, ErrorResponse]:

569 - generation: AsyncGenerator[Union[AnyResponse, ErrorResponse], None]

568 + ) -> AnyResponse | ErrorResponse:

569 + generation: AsyncGenerator[AnyResponse | ErrorResponse, None]

570 generation = self._pipeline(ctx)

571

572 async for response in generation:

@@ -577,7 +577,7 @@ class OpenAIServing:

577 async def _pipeline(

578 self,

579 ctx: ServeContext,

580 - ) -> AsyncGenerator[Union[AnyResponse, ErrorResponse], None]:

580 + ) -> AsyncGenerator[AnyResponse | ErrorResponse, None]:

581 """Execute the request processing pipeline yielding responses."""

582 if error := await self._check_model(ctx.request):

583 yield error

@@ -598,7 +598,7 @@ class OpenAIServing:

598

599 yield self._build_response(ctx)

600

601 - def _validate_request(self, ctx: ServeContext) -> Optional[ErrorResponse]:

601 + def _validate_request(self, ctx: ServeContext) -> ErrorResponse | None:

602 truncate_prompt_tokens = getattr(ctx.request, "truncate_prompt_tokens", None)

603

604 if (

@@ -615,7 +615,7 @@ class OpenAIServing:

615 def _create_pooling_params(

616 self,

617 ctx: ServeContext,

618 - ) -> Union[PoolingParams, ErrorResponse]:

618 + ) -> PoolingParams | ErrorResponse:

619 if not hasattr(ctx.request, "to_pooling_params"):

620 return self.create_error_response(

621 "Request type does not support pooling parameters"

@@ -626,10 +626,10 @@ class OpenAIServing:

626 async def _prepare_generators(

627 self,

628 ctx: ServeContext,

629 - ) -> Optional[ErrorResponse]:

629 + ) -> ErrorResponse | None:

630 """Schedule the request and get the result generator."""

631 generators: list[

632 - AsyncGenerator[Union[RequestOutput, PoolingRequestOutput], None]

632 + AsyncGenerator[RequestOutput | PoolingRequestOutput, None]

633 ] = []

634

635 try:

@@ -678,14 +678,14 @@ class OpenAIServing:

678 async def _collect_batch(

679 self,

680 ctx: ServeContext,

681 - ) -> Optional[ErrorResponse]:

681 + ) -> ErrorResponse | None:

682 """Collect batch results from the result generator."""

683 try:

684 if ctx.engine_prompts is None:

685 return self.create_error_response("Engine prompts not available")

686

687 num_prompts = len(ctx.engine_prompts)

688 - final_res_batch: list[Optional[Union[RequestOutput, PoolingRequestOutput]]]

688 + final_res_batch: list[RequestOutput | PoolingRequestOutput | None]

689 final_res_batch = [None] * num_prompts

690

691 if ctx.result_generator is None:

@@ -738,7 +738,7 @@ class OpenAIServing:

738 async def _check_model(

739 self,

740 request: AnyRequest,

741 - ) -> Optional[ErrorResponse]:

741 + ) -> ErrorResponse | None:

742 error_response = None

743

744 if self._is_model_supported(request.model):

@@ -764,9 +764,7 @@ class OpenAIServing:

764 status_code=HTTPStatus.NOT_FOUND,

765 )

766

767 - def _get_active_default_mm_loras(

768 - self, request: AnyRequest

769 - ) -> Optional[LoRARequest]:

767 + def _get_active_default_mm_loras(self, request: AnyRequest) -> LoRARequest | None:

768 """Determine if there are any active default multimodal loras."""

769 # TODO: Currently this is only enabled for chat completions

770 # to be better aligned with only being enabled for .generate

@@ -793,7 +791,7 @@ class OpenAIServing:

791 self,

792 request: AnyRequest,

793 supports_default_mm_loras: bool = False,

796 - ) -> Optional[LoRARequest]:

794 + ) -> LoRARequest | None:

795 if request.model in self.models.lora_requests:

796 return self.models.lora_requests[request.model]

797

@@ -877,7 +875,7 @@ class OpenAIServing:

875 self,

876 request: AnyRequest,

877 prompt_ids: list[int],

880 - tokenizer: Optional[AnyTokenizer],

878 + tokenizer: AnyTokenizer | None,

879 ) -> TextTokensPrompt:

880 truncate_prompt_tokens = getattr(request, "truncate_prompt_tokens", None)

881

@@ -972,7 +970,7 @@ class OpenAIServing:

970 self,

971 request: AnyRequest,

972 tokenizer: AnyTokenizer,

975 - prompt_input: Union[str, list[int]],

973 + prompt_input: str | list[int],

974 add_special_tokens: bool = True,

975 ) -> TextTokensPrompt:

976 """

@@ -991,7 +989,7 @@ class OpenAIServing:

989 self,

990 request: AnyRequest,

991 tokenizer: AnyTokenizer,

994 - prompt_inputs: Iterable[Union[str, list[int]]],

992 + prompt_inputs: Iterable[str | list[int]],

993 add_special_tokens: bool = True,

994 ) -> AsyncGenerator[TextTokensPrompt, None]:

995 """

@@ -1014,10 +1012,10 @@ class OpenAIServing:

1012

1013 def _validate_chat_template(

1014 self,

1017 - request_chat_template: Optional[str],

1018 - chat_template_kwargs: Optional[dict[str, Any]],

1015 + request_chat_template: str | None,

1016 + chat_template_kwargs: dict[str, Any] | None,

1017 trust_request_chat_template: bool,

1020 - ) -> Optional[ErrorResponse]:

1018 + ) -> ErrorResponse | None:

1019 if not trust_request_chat_template and (

1020 request_chat_template is not None

1021 or (

@@ -1034,17 +1032,17 @@ class OpenAIServing:

1032

1033 async def _preprocess_chat(

1034 self,

1037 - request: Union[ChatLikeRequest, ResponsesRequest],

1035 + request: ChatLikeRequest | ResponsesRequest,

1036 tokenizer: AnyTokenizer,

1037 messages: list[ChatCompletionMessageParam],

1040 - chat_template: Optional[str],

1038 + chat_template: str | None,

1039 chat_template_content_format: ChatTemplateContentFormatOption,

1040 add_generation_prompt: bool = True,

1041 continue_final_message: bool = False,

1044 - tool_dicts: Optional[list[dict[str, Any]]] = None,

1045 - documents: Optional[list[dict[str, str]]] = None,

1046 - chat_template_kwargs: Optional[dict[str, Any]] = None,

1047 - tool_parser: Optional[Callable[[AnyTokenizer], ToolParser]] = None,

1042 + tool_dicts: list[dict[str, Any]] | None = None,

1043 + documents: list[dict[str, str]] | None = None,

1044 + chat_template_kwargs: dict[str, Any] | None = None,

1045 + tool_parser: Callable[[AnyTokenizer], ToolParser] | None = None,

1046 add_special_tokens: bool = False,

1047 ) -> tuple[

1048 list[ConversationMessage],

@@ -1076,7 +1074,7 @@ class OpenAIServing:

1074 )

1075 _chat_template_kwargs.update(chat_template_kwargs or {})

1076

1079 - request_prompt: Union[str, list[int]]

1077 + request_prompt: str | list[int]

1078

1079 if tokenizer is None:

1080 request_prompt = "placeholder"

@@ -1158,10 +1156,10 @@ class OpenAIServing:

1156 self,

1157 request_id: str,

1158 engine_prompt: PromptType,

1161 - params: Union[SamplingParams, PoolingParams],

1159 + params: SamplingParams | PoolingParams,

1160 *,

1163 - lora_request: Optional[LoRARequest],

1164 - trace_headers: Optional[Mapping[str, str]],

1161 + lora_request: LoRARequest | None,

1162 + trace_headers: Mapping[str, str] | None,

1163 priority: int,

1164 ) -> tuple[EngineCoreRequest, dict[str, Any]]:

1165 """Use the Processor to process inputs for AsyncLLM."""

@@ -1188,7 +1186,7 @@ class OpenAIServing:

1186 engine_prompt: EngineTokensPrompt,

1187 sampling_params: SamplingParams,

1188 context: ConversationContext,

1191 - lora_request: Optional[LoRARequest] = None,

1189 + lora_request: LoRARequest | None = None,

1190 priority: int = 0,

1191 **kwargs,

1192 ):

@@ -1250,7 +1248,7 @@ class OpenAIServing:

1248

1249 def _get_prompt_components(

1250 self,

1253 - prompt: Union[RequestPrompt, PromptType],

1251 + prompt: RequestPrompt | PromptType,

1252 ) -> PromptComponents:

1253 if isinstance(prompt, list):

1254 return PromptComponents(token_ids=prompt)

@@ -1260,9 +1258,9 @@ class OpenAIServing:

1258 def _log_inputs(

1259 self,

1260 request_id: str,

1263 - inputs: Union[RequestPrompt, PromptType],

1264 - params: Optional[Union[SamplingParams, PoolingParams, BeamSearchParams]],

1265 - lora_request: Optional[LoRARequest],

1261 + inputs: RequestPrompt | PromptType,

1262 + params: SamplingParams | PoolingParams | BeamSearchParams | None,

1263 + lora_request: LoRARequest | None,

1264 ) -> None:

1265 if self.request_logger is None:

1266 return

@@ -1281,7 +1279,7 @@ class OpenAIServing:

1279 async def _get_trace_headers(

1280 self,

1281 headers: Headers,

1284 - ) -> Optional[Mapping[str, str]]:

1282 + ) -> Mapping[str, str] | None:

1283 is_tracing_enabled = await self.engine_client.is_tracing_enabled()

1284

1285 if is_tracing_enabled:

@@ -1294,8 +1292,8 @@ class OpenAIServing:

1292

1293 @staticmethod

1294 def _base_request_id(

1297 - raw_request: Optional[Request], default: Optional[str] = None

1298 - ) -> Optional[str]:

1295 + raw_request: Request | None, default: str | None = None

1296 + ) -> str | None:

1297 """Pulls the request id to use from a header, if provided"""

1298 default = default or random_uuid()

1299 if raw_request is None:

@@ -1317,15 +1315,15 @@ class OpenAIServing:

1315 return logprob.decoded_token

1316 return tokenizer.decode(token_id)

1317

1320 - def _is_model_supported(self, model_name: Optional[str]) -> bool:

1318 + def _is_model_supported(self, model_name: str | None) -> bool:

1319 if not model_name:

1320 return True

1321 return self.models.is_base_model(model_name)

1322

1323

1324 def clamp_prompt_logprobs(

1327 - prompt_logprobs: Union[PromptLogprobs, None],

1328 -) -> Union[PromptLogprobs, None]:

1325 + prompt_logprobs: PromptLogprobs | None,

1326 +) -> PromptLogprobs | None:

1327 if prompt_logprobs is None:

1328 return prompt_logprobs

1329

vllm/entrypoints/openai/serving_models.py

+9 -10

@@ -5,7 +5,6 @@ from asyncio import Lock

5 from collections import defaultdict

6 from dataclasses import dataclass

7 from http import HTTPStatus

8 -from typing import Optional, Union

8

9 from vllm.engine.protocol import EngineClient

10 from vllm.entrypoints.openai.protocol import (

@@ -35,7 +34,7 @@ class BaseModelPath:

34 class LoRAModulePath:

35 name: str

36 path: str

38 - base_model_name: Optional[str] = None

37 + base_model_name: str | None = None

38

39

40 class OpenAIServingModels:

@@ -52,7 +51,7 @@ class OpenAIServingModels:

51 engine_client: EngineClient,

52 base_model_paths: list[BaseModelPath],

53 *,

55 - lora_modules: Optional[list[LoRAModulePath]] = None,

54 + lora_modules: list[LoRAModulePath] | None = None,

55 ):

56 super().__init__()

57

@@ -93,7 +92,7 @@ class OpenAIServingModels:

92 def is_base_model(self, model_name) -> bool:

93 return any(model.name == model_name for model in self.base_model_paths)

94

96 - def model_name(self, lora_request: Optional[LoRARequest] = None) -> str:

95 + def model_name(self, lora_request: LoRARequest | None = None) -> str:

96 """Returns the appropriate model name depending on the availability

97 and support of the LoRA or base model.

98 Parameters:

@@ -132,8 +131,8 @@ class OpenAIServingModels:

131 return ModelList(data=model_cards)

132

133 async def load_lora_adapter(

135 - self, request: LoadLoRAAdapterRequest, base_model_name: Optional[str] = None

136 - ) -> Union[ErrorResponse, str]:

134 + self, request: LoadLoRAAdapterRequest, base_model_name: str | None = None

135 + ) -> ErrorResponse | str:

136 lora_name = request.lora_name

137

138 # Ensure atomicity based on the lora name

@@ -173,7 +172,7 @@ class OpenAIServingModels:

172

173 async def unload_lora_adapter(

174 self, request: UnloadLoRAAdapterRequest

176 - ) -> Union[ErrorResponse, str]:

175 + ) -> ErrorResponse | str:

176 lora_name = request.lora_name

177

178 # Ensure atomicity based on the lora name

@@ -189,7 +188,7 @@ class OpenAIServingModels:

188

189 async def _check_load_lora_adapter_request(

190 self, request: LoadLoRAAdapterRequest

192 - ) -> Optional[ErrorResponse]:

191 + ) -> ErrorResponse | None:

192 # Check if both 'lora_name' and 'lora_path' are provided

193 if not request.lora_name or not request.lora_path:

194 return create_error_response(

@@ -211,7 +210,7 @@ class OpenAIServingModels:

210

211 async def _check_unload_lora_adapter_request(

212 self, request: UnloadLoRAAdapterRequest

214 - ) -> Optional[ErrorResponse]:

213 + ) -> ErrorResponse | None:

214 # Check if 'lora_name' is not provided return an error

215 if not request.lora_name:

216 return create_error_response(

@@ -230,7 +229,7 @@ class OpenAIServingModels:

229

230 return None

231

233 - async def resolve_lora(self, lora_name: str) -> Union[LoRARequest, ErrorResponse]:

232 + async def resolve_lora(self, lora_name: str) -> LoRARequest | ErrorResponse:

233 """Attempt to resolve a LoRA adapter using available resolvers.

234

235 Args:

vllm/entrypoints/openai/serving_pooling.py

+7 -7

@@ -5,7 +5,7 @@ import asyncio

5 import base64

6 import time

7 from collections.abc import AsyncGenerator

8 -from typing import Final, Literal, Optional, Union, cast

8 +from typing import Final, Literal, cast

9

10 import jinja2

11 import numpy as np

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

41 def _get_data(

42 output: PoolingOutput,

43 encoding_format: Literal["float", "base64"],

44 -) -> Union[list[float], str]:

44 +) -> list[float] | str:

45 if encoding_format == "float":

46 return output.data.tolist()

47 elif encoding_format == "base64":

@@ -60,8 +60,8 @@ class OpenAIServingPooling(OpenAIServing):

60 engine_client: EngineClient,

61 models: OpenAIServingModels,

62 *,

63 - request_logger: Optional[RequestLogger],

64 - chat_template: Optional[str],

63 + request_logger: RequestLogger | None,

64 + chat_template: str | None,

65 chat_template_content_format: ChatTemplateContentFormatOption,

66 trust_request_chat_template: bool = False,

67 log_error_stack: bool = False,

@@ -80,8 +80,8 @@ class OpenAIServingPooling(OpenAIServing):

80 async def create_pooling(

81 self,

82 request: PoolingRequest,

83 - raw_request: Optional[Request] = None,

84 - ) -> Union[PoolingResponse, IOProcessorResponse, ErrorResponse]:

83 + raw_request: Request | None = None,

84 + ) -> PoolingResponse | IOProcessorResponse | ErrorResponse:

85 """

86 See https://platform.openai.com/docs/api-reference/embeddings/create

87 for the API specification. This API mimics the OpenAI Embedding API.

@@ -219,7 +219,7 @@ class OpenAIServingPooling(OpenAIServing):

219 num_prompts = len(engine_prompts)

220

221 # Non-streaming response

222 - final_res_batch: list[Optional[PoolingRequestOutput]]

222 + final_res_batch: list[PoolingRequestOutput | None]

223 final_res_batch = [None] * num_prompts

224 try:

225 async for i, res in result_generator:

vllm/entrypoints/openai/serving_responses.py

+36 -36

@@ -6,11 +6,11 @@ import json

6 import time

7 import uuid

8 from collections import deque

9 -from collections.abc import AsyncGenerator, AsyncIterator, Sequence

9 +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence

10 from contextlib import AsyncExitStack

11 from copy import copy

12 from http import HTTPStatus

13 -from typing import Callable, Final, Optional, Union

13 +from typing import Final

14

15 import jinja2

16 from fastapi import Request

@@ -109,14 +109,14 @@ class OpenAIServingResponses(OpenAIServing):

109 engine_client: EngineClient,

110 models: OpenAIServingModels,

111 *,

112 - request_logger: Optional[RequestLogger],

113 - chat_template: Optional[str],

112 + request_logger: RequestLogger | None,

113 + chat_template: str | None,

114 chat_template_content_format: ChatTemplateContentFormatOption,

115 return_tokens_as_token_ids: bool = False,

116 reasoning_parser: str = "",

117 enable_auto_tools: bool = False,

118 - tool_parser: Optional[str] = None,

119 - tool_server: Optional[ToolServer] = None,

118 + tool_parser: str | None = None,

119 + tool_server: ToolServer | None = None,

120 enable_prompt_tokens_details: bool = False,

121 enable_force_include_usage: bool = False,

122 enable_log_outputs: bool = False,

@@ -210,7 +210,7 @@ class OpenAIServingResponses(OpenAIServing):

210

211 def _validate_generator_input(

212 self, engine_prompt: EngineTokensPrompt

213 - ) -> Optional[ErrorResponse]:

213 + ) -> ErrorResponse | None:

214 """Add validations to the input to the generator here."""

215 if self.max_model_len <= len(engine_prompt["prompt_token_ids"]):

216 error_message = (

@@ -229,12 +229,12 @@ class OpenAIServingResponses(OpenAIServing):

229 async def create_responses(

230 self,

231 request: ResponsesRequest,

232 - raw_request: Optional[Request] = None,

233 - ) -> Union[

234 - AsyncGenerator[StreamingResponsesResponse, None],

235 - ResponsesResponse,

236 - ErrorResponse,

237 - ]:

232 + raw_request: Request | None = None,

233 + ) -> (

234 + AsyncGenerator[StreamingResponsesResponse, None]

235 + | ResponsesResponse

236 + | ErrorResponse

237 + ):

238 error_check_ret = await self._check_model(request)

239 if error_check_ret is not None:

240 logger.error("Error with model %s", error_check_ret)

@@ -460,7 +460,7 @@ class OpenAIServingResponses(OpenAIServing):

460 async def _make_request(

461 self,

462 request: ResponsesRequest,

463 - prev_response: Optional[ResponsesResponse],

463 + prev_response: ResponsesResponse | None,

464 tokenizer: AnyTokenizer,

465 ):

466 if len(request.tools) > 0:

@@ -481,7 +481,7 @@ class OpenAIServingResponses(OpenAIServing):

481 def _make_request_with_harmony(

482 self,

483 request: ResponsesRequest,

484 - prev_response: Optional[ResponsesResponse],

484 + prev_response: ResponsesResponse | None,

485 ):

486 if request.tool_choice != "auto":

487 raise NotImplementedError(

@@ -522,8 +522,8 @@ class OpenAIServingResponses(OpenAIServing):

522 model_name: str,

523 tokenizer: AnyTokenizer,

524 request_metadata: RequestResponseMetadata,

525 - created_time: Optional[int] = None,

526 - ) -> Union[ErrorResponse, ResponsesResponse]:

525 + created_time: int | None = None,

526 + ) -> ErrorResponse | ResponsesResponse:

527 if created_time is None:

528 created_time = int(time.time())

529

@@ -642,9 +642,9 @@ class OpenAIServingResponses(OpenAIServing):

642 def _create_response_logprobs(

643 self,

644 token_ids: Sequence[int],

645 - logprobs: Optional[SampleLogprobs],

645 + logprobs: SampleLogprobs | None,

646 tokenizer: AnyTokenizer,

647 - top_logprobs: Optional[int] = None,

647 + top_logprobs: int | None = None,

648 ) -> list[Logprob]:

649 assert logprobs is not None, "logprobs must be provided"

650 assert len(token_ids) == len(logprobs), (

@@ -676,9 +676,9 @@ class OpenAIServingResponses(OpenAIServing):

676 def _create_stream_response_logprobs(

677 self,

678 token_ids: Sequence[int],

679 - logprobs: Optional[SampleLogprobs],

679 + logprobs: SampleLogprobs | None,

680 tokenizer: AnyTokenizer,

681 - top_logprobs: Optional[int] = None,

681 + top_logprobs: int | None = None,

682 ) -> list[response_text_delta_event.Logprob]:

683 lgs = self._create_response_logprobs(

684 token_ids=token_ids,

@@ -793,7 +793,7 @@ class OpenAIServingResponses(OpenAIServing):

793 def _construct_input_messages(

794 self,

795 request: ResponsesRequest,

796 - prev_response: Optional[ResponsesResponse] = None,

796 + prev_response: ResponsesResponse | None = None,

797 ) -> list[ChatCompletionMessageParam]:

798 messages: list[ChatCompletionMessageParam] = []

799 if request.instructions:

@@ -833,7 +833,7 @@ class OpenAIServingResponses(OpenAIServing):

833 def _construct_input_messages_with_harmony(

834 self,

835 request: ResponsesRequest,

836 - prev_response: Optional[ResponsesResponse],

836 + prev_response: ResponsesResponse | None,

837 ) -> list[OpenAIHarmonyMessage]:

838 messages: list[OpenAIHarmonyMessage] = []

839 if prev_response is None:

@@ -986,7 +986,7 @@ class OpenAIServingResponses(OpenAIServing):

986 async def responses_background_stream_generator(

987 self,

988 response_id: str,

989 - starting_after: Optional[int] = None,

989 + starting_after: int | None = None,

990 ) -> AsyncGenerator[StreamingResponsesResponse, None]:

991 if response_id not in self.event_store:

992 raise ValueError(f"Unknown response_id: {response_id}")

@@ -1011,13 +1011,13 @@ class OpenAIServingResponses(OpenAIServing):

1011 async def retrieve_responses(

1012 self,

1013 response_id: str,

1014 - starting_after: Optional[int],

1015 - stream: Optional[bool],

1016 - ) -> Union[

1017 - ErrorResponse,

1018 - ResponsesResponse,

1019 - AsyncGenerator[StreamingResponsesResponse, None],

1020 - ]:

1014 + starting_after: int | None,

1015 + stream: bool | None,

1016 + ) -> (

1017 + ErrorResponse

1018 + | ResponsesResponse

1019 + | AsyncGenerator[StreamingResponsesResponse, None]

1020 + ):

1021 async with self.response_store_lock:

1022 response = self.response_store.get(response_id)

1023

@@ -1034,7 +1034,7 @@ class OpenAIServingResponses(OpenAIServing):

1034 async def cancel_responses(

1035 self,

1036 response_id: str,

1037 - ) -> Union[ErrorResponse, ResponsesResponse]:

1037 + ) -> ErrorResponse | ResponsesResponse:

1038 async with self.response_store_lock:

1039 response = self.response_store.get(response_id)

1040 if response is None:

@@ -1082,7 +1082,7 @@ class OpenAIServingResponses(OpenAIServing):

1082 self,

1083 request: ResponsesRequest,

1084 sampling_params: SamplingParams,

1085 - result_generator: AsyncIterator[Optional[ConversationContext]],

1085 + result_generator: AsyncIterator[ConversationContext | None],

1086 context: ConversationContext,

1087 model_name: str,

1088 tokenizer: AnyTokenizer,

@@ -1385,7 +1385,7 @@ class OpenAIServingResponses(OpenAIServing):

1385 self,

1386 request: ResponsesRequest,

1387 sampling_params: SamplingParams,

1388 - result_generator: AsyncIterator[Optional[ConversationContext]],

1388 + result_generator: AsyncIterator[ConversationContext | None],

1389 context: ConversationContext,

1390 model_name: str,

1391 tokenizer: AnyTokenizer,

@@ -1772,12 +1772,12 @@ class OpenAIServingResponses(OpenAIServing):

1772 self,

1773 request: ResponsesRequest,

1774 sampling_params: SamplingParams,

1775 - result_generator: AsyncIterator[Optional[ConversationContext]],

1775 + result_generator: AsyncIterator[ConversationContext | None],

1776 context: ConversationContext,

1777 model_name: str,

1778 tokenizer: AnyTokenizer,

1779 request_metadata: RequestResponseMetadata,

1780 - created_time: Optional[int] = None,

1780 + created_time: int | None = None,

1781 ) -> AsyncGenerator[StreamingResponsesResponse, None]:

1782 # TODO:

1783 # 1. Handle disconnect

vllm/entrypoints/openai/serving_score.py

+29 -29

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

3 import asyncio

4 import time

5 from collections.abc import AsyncGenerator, Mapping

6 -from typing import Any, Optional, Union

6 +from typing import Any

7

8 from fastapi import Request

9

@@ -48,7 +48,7 @@ class ServingScores(OpenAIServing):

48 engine_client: EngineClient,

49 models: OpenAIServingModels,

50 *,

51 - request_logger: Optional[RequestLogger],

51 + request_logger: RequestLogger | None,

52 log_error_stack: bool = False,

53 ) -> None:

54 super().__init__(

@@ -63,12 +63,12 @@ class ServingScores(OpenAIServing):

63 tokenizer: AnyTokenizer,

64 texts_1: list[str],

65 texts_2: list[str],

66 - request: Union[RerankRequest, ScoreRequest],

66 + request: RerankRequest | ScoreRequest,

67 request_id: str,

68 - tokenization_kwargs: Optional[dict[str, Any]] = None,

69 - lora_request: Optional[Union[LoRARequest, None]] = None,

70 - trace_headers: Optional[Mapping[str, str]] = None,

71 - ) -> Union[list[PoolingRequestOutput], ErrorResponse]:

68 + tokenization_kwargs: dict[str, Any] | None = None,

69 + lora_request: LoRARequest | None | None = None,

70 + trace_headers: Mapping[str, str] | None = None,

71 + ) -> list[PoolingRequestOutput] | ErrorResponse:

72 input_texts = texts_1 + texts_2

73

74 engine_prompts: list[TokensPrompt] = []

@@ -125,7 +125,7 @@ class ServingScores(OpenAIServing):

125 # Non-streaming response

126 final_res_batch: list[PoolingRequestOutput] = []

127

128 - embeddings: list[Optional[PoolingRequestOutput]] = [None] * len(engine_prompts)

128 + embeddings: list[PoolingRequestOutput | None] = [None] * len(engine_prompts)

129

130 async for i, res in result_generator:

131 embeddings[i] = res

@@ -152,11 +152,11 @@ class ServingScores(OpenAIServing):

152

153 def _preprocess_score(

154 self,

155 - request: Union[RerankRequest, ScoreRequest],

155 + request: RerankRequest | ScoreRequest,

156 tokenizer: AnyTokenizer,

157 tokenization_kwargs: dict[str, Any],

158 - data_1: Union[str, ScoreContentPartParam],

159 - data_2: Union[str, ScoreContentPartParam],

158 + data_1: str | ScoreContentPartParam,

159 + data_2: str | ScoreContentPartParam,

160 ) -> tuple[str, TokensPrompt]:

161 model_config = self.model_config

162

@@ -176,14 +176,14 @@ class ServingScores(OpenAIServing):

176 async def _cross_encoding_score(

177 self,

178 tokenizer: AnyTokenizer,

179 - data_1: Union[list[str], list[ScoreContentPartParam]],

180 - data_2: Union[list[str], list[ScoreContentPartParam]],

181 - request: Union[RerankRequest, ScoreRequest],

179 + data_1: list[str] | list[ScoreContentPartParam],

180 + data_2: list[str] | list[ScoreContentPartParam],

181 + request: RerankRequest | ScoreRequest,

182 request_id: str,

183 - tokenization_kwargs: Optional[dict[str, Any]] = None,

184 - lora_request: Optional[Union[LoRARequest, None]] = None,

185 - trace_headers: Optional[Mapping[str, str]] = None,

186 - ) -> Union[list[PoolingRequestOutput], ErrorResponse]:

183 + tokenization_kwargs: dict[str, Any] | None = None,

184 + lora_request: LoRARequest | None | None = None,

185 + trace_headers: Mapping[str, str] | None = None,

186 + ) -> list[PoolingRequestOutput] | ErrorResponse:

187 request_prompts: list[str] = []

188 engine_prompts: list[TokensPrompt] = []

189

@@ -259,7 +259,7 @@ class ServingScores(OpenAIServing):

259 result_generator = merge_async_iterators(*generators)

260

261 # Non-streaming response

262 - final_res_batch: list[Optional[PoolingRequestOutput]] = [None] * len(

262 + final_res_batch: list[PoolingRequestOutput | None] = [None] * len(

263 engine_prompts

264 )

265

@@ -270,12 +270,12 @@ class ServingScores(OpenAIServing):

270

271 async def _run_scoring(

272 self,

273 - data_1: Union[list[str], str, ScoreMultiModalParam],

274 - data_2: Union[list[str], str, ScoreMultiModalParam],

275 - request: Union[ScoreRequest, RerankRequest],

273 + data_1: list[str] | str | ScoreMultiModalParam,

274 + data_2: list[str] | str | ScoreMultiModalParam,

275 + request: ScoreRequest | RerankRequest,

276 request_id: str,

277 - raw_request: Optional[Request] = None,

278 - ) -> Union[list[PoolingRequestOutput], ErrorResponse]:

277 + raw_request: Request | None = None,

278 + ) -> list[PoolingRequestOutput] | ErrorResponse:

279 lora_request = self._maybe_get_adapters(request)

280

281 tokenizer = await self.engine_client.get_tokenizer()

@@ -339,8 +339,8 @@ class ServingScores(OpenAIServing):

339 async def create_score(

340 self,

341 request: ScoreRequest,

342 - raw_request: Optional[Request] = None,

343 - ) -> Union[ScoreResponse, ErrorResponse]:

342 + raw_request: Request | None = None,

343 + ) -> ScoreResponse | ErrorResponse:

344 """

345 Score API similar to Sentence Transformers cross encoder

346

@@ -377,8 +377,8 @@ class ServingScores(OpenAIServing):

377 return self.create_error_response(str(e))

378

379 async def do_rerank(

380 - self, request: RerankRequest, raw_request: Optional[Request] = None

381 - ) -> Union[RerankResponse, ErrorResponse]:

380 + self, request: RerankRequest, raw_request: Request | None = None

381 + ) -> RerankResponse | ErrorResponse:

382 """

383 Rerank API based on JinaAI's rerank API; implements the same

384 API interface. Designed for compatibility with off-the-shelf

@@ -468,7 +468,7 @@ class ServingScores(OpenAIServing):

468 final_res_batch: list[PoolingRequestOutput],

469 request_id: str,

470 model_name: str,

471 - documents: Union[list[str], ScoreMultiModalParam],

471 + documents: list[str] | ScoreMultiModalParam,

472 top_n: int,

473 ) -> RerankResponse:

474 """

vllm/entrypoints/openai/serving_tokenization.py

+7 -7

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass

4 -from typing import Any, Final, Optional, Union

4 +from typing import Any, Final

5

6 import jinja2

7 from fastapi import Request

@@ -33,8 +33,8 @@ class OpenAIServingTokenization(OpenAIServing):

33 engine_client: EngineClient,

34 models: OpenAIServingModels,

35 *,

36 - request_logger: Optional[RequestLogger],

37 - chat_template: Optional[str],

36 + request_logger: RequestLogger | None,

37 + chat_template: str | None,

38 chat_template_content_format: ChatTemplateContentFormatOption,

39 trust_request_chat_template: bool = False,

40 log_error_stack: bool = False,

@@ -54,7 +54,7 @@ class OpenAIServingTokenization(OpenAIServing):

54 self,

55 request: TokenizeRequest,

56 raw_request: Request,

57 - ) -> Union[TokenizeResponse, ErrorResponse]:

57 + ) -> TokenizeResponse | ErrorResponse:

58 error_check_ret = await self._check_model(request)

59 if error_check_ret is not None:

60 return error_check_ret

@@ -129,7 +129,7 @@ class OpenAIServingTokenization(OpenAIServing):

129 self,

130 request: DetokenizeRequest,

131 raw_request: Request,

132 - ) -> Union[DetokenizeResponse, ErrorResponse]:

132 + ) -> DetokenizeResponse | ErrorResponse:

133 error_check_ret = await self._check_model(request)

134 if error_check_ret is not None:

135 return error_check_ret

@@ -155,7 +155,7 @@ class OpenAIServingTokenization(OpenAIServing):

155

156 async def get_tokenizer_info(

157 self,

158 - ) -> Union[TokenizerInfoResponse, ErrorResponse]:

158 + ) -> TokenizerInfoResponse | ErrorResponse:

159 """Get comprehensive tokenizer information."""

160 try:

161 tokenizer = await self.engine_client.get_tokenizer()

@@ -171,7 +171,7 @@ class OpenAIServingTokenization(OpenAIServing):

171 @dataclass

172 class TokenizerInfo:

173 tokenizer: AnyTokenizer

174 - chat_template: Optional[str]

174 + chat_template: str | None

175

176 def to_dict(self) -> dict[str, Any]:

177 """Return the tokenizer configuration."""

vllm/entrypoints/openai/serving_transcription.py

+4 -5

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from collections.abc import AsyncGenerator

4 -from typing import Optional, Union

4

5 from fastapi import Request

6

@@ -35,7 +34,7 @@ class OpenAIServingTranscription(OpenAISpeechToText):

34 engine_client: EngineClient,

35 models: OpenAIServingModels,

36 *,

38 - request_logger: Optional[RequestLogger],

37 + request_logger: RequestLogger | None,

38 return_tokens_as_token_ids: bool = False,

39 log_error_stack: bool = False,

40 ):

@@ -50,7 +49,7 @@ class OpenAIServingTranscription(OpenAISpeechToText):

49

50 async def create_transcription(

51 self, audio_data: bytes, request: TranscriptionRequest, raw_request: Request

53 - ) -> Union[TranscriptionResponse, AsyncGenerator[str, None], ErrorResponse]:

52 + ) -> TranscriptionResponse | AsyncGenerator[str, None] | ErrorResponse:

53 """Transcription API similar to OpenAI's API.

54

55 See https://platform.openai.com/docs/api-reference/audio/createTranscription

@@ -94,7 +93,7 @@ class OpenAIServingTranslation(OpenAISpeechToText):

93 engine_client: EngineClient,

94 models: OpenAIServingModels,

95 *,

97 - request_logger: Optional[RequestLogger],

96 + request_logger: RequestLogger | None,

97 return_tokens_as_token_ids: bool = False,

98 log_error_stack: bool = False,

99 ):

@@ -109,7 +108,7 @@ class OpenAIServingTranslation(OpenAISpeechToText):

108

109 async def create_translation(

110 self, audio_data: bytes, request: TranslationRequest, raw_request: Request

112 - ) -> Union[TranslationResponse, AsyncGenerator[str, None], ErrorResponse]:

111 + ) -> TranslationResponse | AsyncGenerator[str, None] | ErrorResponse:

112 """Translation API similar to OpenAI's API.

113

114 See https://platform.openai.com/docs/api-reference/audio/createTranslation

vllm/entrypoints/openai/speech_to_text.py

+10 -15

@@ -4,9 +4,9 @@ import asyncio

4 import io

5 import math

6 import time

7 -from collections.abc import AsyncGenerator

7 +from collections.abc import AsyncGenerator, Callable

8 from functools import cached_property

9 -from typing import Callable, Literal, Optional, TypeVar, Union, cast

9 +from typing import Literal, TypeAlias, TypeVar, cast

10

11 import numpy as np

12 from fastapi import Request

@@ -39,7 +39,7 @@ try:

39 except ImportError:

40 librosa = PlaceholderModule("librosa") # type: ignore[assignment]

41

42 -SpeechToTextResponse = Union[TranscriptionResponse, TranslationResponse]

42 +SpeechToTextResponse: TypeAlias = TranscriptionResponse | TranslationResponse

43 T = TypeVar("T", bound=SpeechToTextResponse)

44

45 logger = init_logger(__name__)

@@ -54,7 +54,7 @@ class OpenAISpeechToText(OpenAIServing):

54 engine_client: EngineClient,

55 models: OpenAIServingModels,

56 *,

57 - request_logger: Optional[RequestLogger],

57 + request_logger: RequestLogger | None,

58 return_tokens_as_token_ids: bool = False,

59 task_type: Literal["transcribe", "translate"] = "transcribe",

60 log_error_stack: bool = False,

@@ -140,7 +140,7 @@ class OpenAISpeechToText(OpenAIServing):

140 raw_request: Request,

141 response_class: type[T],

142 stream_generator_method: Callable[..., AsyncGenerator[str, None]],

143 - ) -> Union[T, AsyncGenerator[str, None], ErrorResponse]:

143 + ) -> T | AsyncGenerator[str, None] | ErrorResponse:

144 """Base method for speech-to-text operations like transcription and

145 translation."""

146 error_check_ret = await self._check_model(request)

@@ -181,9 +181,7 @@ class OpenAISpeechToText(OpenAIServing):

181 logger.exception("Error in preprocessing prompt inputs")

182 return self.create_error_response(str(e))

183

184 - list_result_generator: Optional[list[AsyncGenerator[RequestOutput, None]]] = (

185 - None

186 - )

184 + list_result_generator: list[AsyncGenerator[RequestOutput, None]] | None = None

185 try:

186 # Unlike most decoder-only models, whisper generation length is not

187 # constrained by the size of the input audio, which is mapped to a

@@ -252,13 +250,10 @@ class OpenAISpeechToText(OpenAIServing):

250 request_metadata: RequestResponseMetadata,

251 audio_duration_s: float,

252 chunk_object_type: Literal["translation.chunk", "transcription.chunk"],

255 - response_stream_choice_class: Union[

256 - type[TranscriptionResponseStreamChoice],

257 - type[TranslationResponseStreamChoice],

258 - ],

259 - stream_response_class: Union[

260 - type[TranscriptionStreamResponse], type[TranslationStreamResponse]

261 - ],

253 + response_stream_choice_class: type[TranscriptionResponseStreamChoice]

254 + | type[TranslationResponseStreamChoice],

255 + stream_response_class: type[TranscriptionStreamResponse]

256 + | type[TranslationStreamResponse],

257 ) -> AsyncGenerator[str, None]:

258 created_time = int(time.time())

259 model_name = request.model

vllm/entrypoints/openai/tool_parsers/abstract_tool_parser.py

+6 -7

@@ -2,9 +2,8 @@

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

3

4 import os

5 -from collections.abc import Sequence

5 +from collections.abc import Callable, Sequence

6 from functools import cached_property

7 -from typing import Callable, Optional, Union

7

8 from vllm.entrypoints.openai.protocol import (

9 ChatCompletionRequest,

@@ -69,7 +68,7 @@ class ToolParser:

68 current_token_ids: Sequence[int],

69 delta_token_ids: Sequence[int],

70 request: ChatCompletionRequest,

72 - ) -> Union[DeltaMessage, None]:

71 + ) -> DeltaMessage | None:

72 """

73 Instance method that should be implemented for extracting tool calls

74 from an incomplete response; for use when handling tool calls and

@@ -101,7 +100,7 @@ class ToolParserManager:

100 def _register_module(

101 cls,

102 module: type,

104 - module_name: Optional[Union[str, list[str]]] = None,

103 + module_name: str | list[str] | None = None,

104 force: bool = True,

105 ) -> None:

106 if not issubclass(module, ToolParser):

@@ -123,10 +122,10 @@ class ToolParserManager:

122 @classmethod

123 def register_module(

124 cls,

126 - name: Optional[Union[str, list[str]]] = None,

125 + name: str | list[str] | None = None,

126 force: bool = True,

128 - module: Union[type, None] = None,

129 - ) -> Union[type, Callable]:

127 + module: type | None = None,

128 + ) -> type | Callable:

129 """

130 Register module with the given name or name list. it can be used as a

131 decoder(with module as None) or normal function(with module as not

vllm/entrypoints/openai/tool_parsers/deepseekv31_tool_parser.py

+2 -3

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

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

3

4 from collections.abc import Sequence

5 -from typing import Union

5

6 import regex as re

7

@@ -129,7 +128,7 @@ class DeepSeekV31ToolParser(ToolParser):

128 current_token_ids: Sequence[int],

129 delta_token_ids: Sequence[int],

130 request: ChatCompletionRequest,

132 - ) -> Union[DeltaMessage, None]:

131 + ) -> DeltaMessage | None:

132 logger.debug("delta_text: %s", delta_text)

133 logger.debug("delta_token_ids: %s", delta_token_ids)

134 # check to see if we should be streaming a tool call - is there a

@@ -272,7 +271,7 @@ class DeepSeekV31ToolParser(ToolParser):

271 if not self.current_tool_name_sent:

272 if current_tool_call is None:

273 return None

275 - function_name: Union[str, None] = current_tool_call.get("name")

274 + function_name: str | None = current_tool_call.get("name")

275 if function_name:

276 self.current_tool_name_sent = True

277 return DeltaMessage(

vllm/entrypoints/openai/tool_parsers/deepseekv3_tool_parser.py

+2 -3

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

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

3

4 from collections.abc import Sequence

5 -from typing import Union

5

6 import regex as re

7

@@ -129,7 +128,7 @@ class DeepSeekV3ToolParser(ToolParser):

128 current_token_ids: Sequence[int],

129 delta_token_ids: Sequence[int],

130 request: ChatCompletionRequest,

132 - ) -> Union[DeltaMessage, None]:

131 + ) -> DeltaMessage | None:

132 logger.debug("delta_text: %s", delta_text)

133 logger.debug("delta_token_ids: %s", delta_token_ids)

134 # check to see if we should be streaming a tool call - is there a

@@ -272,7 +271,7 @@ class DeepSeekV3ToolParser(ToolParser):

271 if not self.current_tool_name_sent:

272 if current_tool_call is None:

273 return None

275 - function_name: Union[str, None] = current_tool_call.get("name")

274 + function_name: str | None = current_tool_call.get("name")

275 if function_name:

276 self.current_tool_name_sent = True

277 return DeltaMessage(

vllm/entrypoints/openai/tool_parsers/glm4_moe_tool_parser.py

+3 -3

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

4 import ast

5 import json

6 from collections.abc import Sequence

7 -from typing import Any, Optional, Union

7 +from typing import Any

8

9 import regex as re

10

@@ -66,7 +66,7 @@ class Glm4MoeModelToolParser(ToolParser):

66 def _is_string_type(

67 tool_name: str,

68 arg_name: str,

69 - tools: Optional[list[ChatCompletionToolsParam]],

69 + tools: list[ChatCompletionToolsParam] | None,

70 ) -> bool:

71 if tools is None:

72 return False

@@ -144,7 +144,7 @@ class Glm4MoeModelToolParser(ToolParser):

144 current_token_ids: Sequence[int],

145 delta_token_ids: Sequence[int],

146 request: ChatCompletionRequest,

147 - ) -> Union[DeltaMessage, None]:

147 + ) -> DeltaMessage | None:

148 self._buffer += delta_text

149 cur_text = self._buffer

150 start_idx = cur_text.find(self.tool_call_start_token)

vllm/entrypoints/openai/tool_parsers/granite_20b_fc_tool_parser.py

+1 -2

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

4 import json

5 from collections.abc import Sequence

6 from json import JSONDecoder

7 -from typing import Union

7

8 import partial_json_parser

9 import regex as re

@@ -121,7 +120,7 @@ class Granite20bFCToolParser(ToolParser):

120 current_token_ids: Sequence[int],

121 delta_token_ids: Sequence[int],

122 request: ChatCompletionRequest,

124 - ) -> Union[DeltaMessage, None]:

123 + ) -> DeltaMessage | None:

124 if len(current_text) < len(self.bot_token) and self.bot_token.startswith(

125 current_text

126 ):

vllm/entrypoints/openai/tool_parsers/granite_tool_parser.py

+1 -2

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

3

4 import json

5 from collections.abc import Sequence

6 -from typing import Union

6

7 import partial_json_parser

8 from partial_json_parser.core.options import Allow

@@ -108,7 +107,7 @@ class GraniteToolParser(ToolParser):

107 current_token_ids: Sequence[int],

108 delta_token_ids: Sequence[int],

109 request: ChatCompletionRequest,

111 - ) -> Union[DeltaMessage, None]:

110 + ) -> DeltaMessage | None:

111 start_idx = consume_space(0, current_text)

112 if current_text[start_idx:].startswith(self.bot_token):

113 start_idx = consume_space(start_idx + len(self.bot_token), current_text)

vllm/entrypoints/openai/tool_parsers/hermes_tool_parser.py

+2 -3

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

3

4 import json

5 from collections.abc import Sequence

6 -from typing import Union

6

7 import partial_json_parser

8 import regex as re

@@ -181,7 +180,7 @@ class Hermes2ProToolParser(ToolParser):

180 current_token_ids: Sequence[int],

181 delta_token_ids: Sequence[int],

182 request: ChatCompletionRequest,

184 - ) -> Union[DeltaMessage, None]:

183 + ) -> DeltaMessage | None:

184 # 1. All tokens are parsed based on _text, not token_ids.

185 # 2. All incoming text data is processed by the tool_call_delta_buffer

186 # function for buffering before being used for parsing.

@@ -333,7 +332,7 @@ class Hermes2ProToolParser(ToolParser):

332 if not self.current_tool_name_sent:

333 if current_tool_call is None:

334 return None

336 - function_name: Union[str, None] = current_tool_call.get("name")

335 + function_name: str | None = current_tool_call.get("name")

336 if function_name:

337 self.current_tool_name_sent = True

338 return DeltaMessage(

vllm/entrypoints/openai/tool_parsers/hunyuan_a13b_tool_parser.py

+3 -3

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

4

5 import json

6 from collections.abc import Sequence

7 -from typing import Any, Optional, Union

7 +from typing import Any

8

9 import regex as re

10

@@ -73,7 +73,7 @@ class HunyuanA13BToolParser(ToolParser):

73

74 def preprocess_model_output(

75 self, model_output: str

76 - ) -> tuple[Optional[str], Optional[str]]:

76 + ) -> tuple[str | None, str | None]:

77 # find the location tool call

78 for match in self.answer_tool_calls_pattern.finditer(model_output):

79 start, end = match.span()

@@ -176,7 +176,7 @@ class HunyuanA13BToolParser(ToolParser):

176 current_token_ids: Sequence[int],

177 delta_token_ids: Sequence[int],

178 request: ChatCompletionRequest,

179 - ) -> Union[DeltaMessage, None]:

179 + ) -> DeltaMessage | None:

180 """

181 Extract tool calls for streaming mode.

182 """

vllm/entrypoints/openai/tool_parsers/internlm2_tool_parser.py

+1 -2

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

3

4 import json

5 from collections.abc import Sequence

6 -from typing import Union

6

7 import partial_json_parser

8 from partial_json_parser.core.options import Allow

@@ -59,7 +58,7 @@ class Internlm2ToolParser(ToolParser):

58 current_token_ids: Sequence[int],

59 delta_token_ids: Sequence[int],

60 request: ChatCompletionRequest,

62 - ) -> Union[DeltaMessage, None]:

61 + ) -> DeltaMessage | None:

62 if "<|action_start|>" not in current_text:

63 self.position = len(current_text)

64 return DeltaMessage(content=delta_text)

vllm/entrypoints/openai/tool_parsers/jamba_tool_parser.py

+2 -3

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

3

4 import json

5 from collections.abc import Sequence

6 -from typing import Union

6

7 import partial_json_parser

8 import regex as re

@@ -129,7 +128,7 @@ class JambaToolParser(ToolParser):

128 current_token_ids: Sequence[int],

129 delta_token_ids: Sequence[int],

130 request: ChatCompletionRequest,

132 - ) -> Union[DeltaMessage, None]:

131 + ) -> DeltaMessage | None:

132 # if the tool call token is not in the tokens generated so far, append

133 # output to contents since it's not a tool

134 if self.tool_calls_start_token not in current_text:

@@ -190,7 +189,7 @@ class JambaToolParser(ToolParser):

189 # auto-generated due to JSON completions, but wasn't

190 # streamed to the client yet.

191 if self.current_tool_id >= 0:

193 - diff: Union[str, None] = current_tool_call.get("arguments")

192 + diff: str | None = current_tool_call.get("arguments")

193

194 if diff:

195 diff = json.dumps(diff, ensure_ascii=False).replace(

vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py

+2 -3

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

3 # code modified from deepseekv3_tool_parser.py

4

5 from collections.abc import Sequence

6 -from typing import Union

6

7 import regex as re

8

@@ -131,7 +130,7 @@ class KimiK2ToolParser(ToolParser):

130 current_token_ids: Sequence[int],

131 delta_token_ids: Sequence[int],

132 request: ChatCompletionRequest,

134 - ) -> Union[DeltaMessage, None]:

133 + ) -> DeltaMessage | None:

134 logger.debug("delta_text: %s", delta_text)

135 logger.debug("delta_token_ids: %s", delta_token_ids)

136 # check to see if we should be streaming a tool call - is there a

@@ -278,7 +277,7 @@ class KimiK2ToolParser(ToolParser):

277 if not self.current_tool_name_sent:

278 if current_tool_call is None:

279 return None

281 - function_name: Union[str, None] = current_tool_call.get("name")

280 + function_name: str | None = current_tool_call.get("name")

281 tool_id = current_tool_call.get("id")

282 if function_name:

283 self.current_tool_name_sent = True

vllm/entrypoints/openai/tool_parsers/llama4_pythonic_tool_parser.py

+4 -4

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

3 import ast

4 import json

5 from collections.abc import Sequence

6 -from typing import Any, Union

6 +from typing import Any

7

8 import regex as re

9 from transformers import PreTrainedTokenizerBase

@@ -128,7 +128,7 @@ class Llama4PythonicToolParser(ToolParser):

128 current_token_ids: Sequence[int],

129 delta_token_ids: Sequence[int],

130 request: ChatCompletionRequest,

131 - ) -> Union[DeltaMessage, None]:

131 + ) -> DeltaMessage | None:

132 if not current_text.startswith("[") and not current_text.startswith(

133 "<|python_start|>"

134 ):

@@ -245,7 +245,7 @@ def _handle_single_tool(call: ast.Call) -> ToolCall:

245 )

246

247

248 -def _make_valid_python(text: str) -> Union[tuple[str, str], None]:

248 +def _make_valid_python(text: str) -> tuple[str, str] | None:

249 bracket_stack = []

250 for index, char in enumerate(text):

251 if char in {"[", "(", "{"}:

@@ -317,7 +317,7 @@ def _make_valid_python(text: str) -> Union[tuple[str, str], None]:

317

318 def _compute_tool_delta(

319 previously_sent_args: str, new_call: ToolCall, index: int, withheld_suffix: str

320 -) -> Union[DeltaToolCall, None]:

320 +) -> DeltaToolCall | None:

321 new_call_args = new_call.function.arguments

322 if withheld_suffix:

323 assert new_call_args.endswith(withheld_suffix)

vllm/entrypoints/openai/tool_parsers/llama_tool_parser.py

+1 -2

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

3

4 import json

5 from collections.abc import Sequence

6 -from typing import Union

6

7 import partial_json_parser

8 import regex as re

@@ -134,7 +133,7 @@ class Llama3JsonToolParser(ToolParser):

133 current_token_ids: Sequence[int],

134 delta_token_ids: Sequence[int],

135 request: ChatCompletionRequest,

137 - ) -> Union[DeltaMessage, None]:

136 + ) -> DeltaMessage | None:

137 if not (

138 current_text.startswith(self.bot_token) or current_text.startswith("{")

139 ):

vllm/entrypoints/openai/tool_parsers/minimax_tool_parser.py

+7 -7

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

3

4 import json

5 from collections.abc import Sequence

6 -from typing import Any, Optional, Union

6 +from typing import Any

7

8 import regex as re

9

@@ -509,7 +509,7 @@ class MinimaxToolParser(ToolParser):

509

510 def _get_current_tool_content(

511 self, text: str, tool_index: int

512 - ) -> tuple[Optional[str], Optional[str]]:

512 + ) -> tuple[str | None, str | None]:

513 """

514 Get the content of a specific tool by index.

515

@@ -545,7 +545,7 @@ class MinimaxToolParser(ToolParser):

545

546 def _handle_tool_name_streaming(

547 self, tool_content: str, tool_count: int

548 - ) -> Union[DeltaMessage, None]:

548 + ) -> DeltaMessage | None:

549 """

550 Handle streaming of tool names.

551

@@ -595,7 +595,7 @@ class MinimaxToolParser(ToolParser):

595

596 def _handle_tool_args_streaming(

597 self, tool_content: str, tool_count: int

598 - ) -> Union[DeltaMessage, None]:

598 + ) -> DeltaMessage | None:

599 """

600 Handle streaming of tool arguments.

601

@@ -702,7 +702,7 @@ class MinimaxToolParser(ToolParser):

702 current_token_ids: Sequence[int],

703 delta_token_ids: Sequence[int],

704 request: ChatCompletionRequest,

705 - ) -> Union[DeltaMessage, None]:

705 + ) -> DeltaMessage | None:

706 self._update_thinking_state(current_text)

707

708 if self.in_thinking_tag:

@@ -776,7 +776,7 @@ class MinimaxToolParser(ToolParser):

776 )

777 return None

778

779 - def _find_tool_start_outside_thinking(self, current_text: str) -> Optional[int]:

779 + def _find_tool_start_outside_thinking(self, current_text: str) -> int | None:

780 """

781 Find the start position of tool calls outside of thinking tags.

782

@@ -809,7 +809,7 @@ class MinimaxToolParser(ToolParser):

809

810 def _extract_content_before_tools(

811 self, current_text: str, delta_text: str, tool_start: int

812 - ) -> Optional[str]:

812 + ) -> str | None:

813 """

814 Extract content that appears before tool calls.

815

vllm/entrypoints/openai/tool_parsers/mistral_tool_parser.py

+2 -3

@@ -5,7 +5,6 @@ import json

5 from collections.abc import Sequence

6 from random import choices

7 from string import ascii_letters, digits

8 -from typing import Union

8

9 import partial_json_parser

10 import regex as re

@@ -194,7 +193,7 @@ class MistralToolParser(ToolParser):

193 current_token_ids: Sequence[int],

194 delta_token_ids: Sequence[int],

195 request: ChatCompletionRequest,

197 - ) -> Union[DeltaMessage, None]:

196 + ) -> DeltaMessage | None:

197 # if the tool call token is not in the tokens generated so far, append

198 # output to contents since it's not a tool

199 if self.bot_token not in current_text:

@@ -252,7 +251,7 @@ class MistralToolParser(ToolParser):

251 # auto-generated due to JSON completions, but wasn't

252 # streamed to the client yet.

253 if self.current_tool_id >= 0:

255 - diff: Union[str, None] = current_tool_call.get("arguments")

254 + diff: str | None = current_tool_call.get("arguments")

255

256 if diff:

257 diff = json.dumps(diff, ensure_ascii=False).replace(

vllm/entrypoints/openai/tool_parsers/openai_tool_parser.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from __future__ import annotations

4 -

3 import json

4 from collections.abc import Sequence

5 from typing import TYPE_CHECKING

@@ -22,13 +20,15 @@ from vllm.logger import init_logger

20

21 if TYPE_CHECKING:

22 from vllm.transformers_utils.tokenizer import AnyTokenizer

23 +else:

24 + AnyTokenizer = object

25

26 logger = init_logger(__name__)

27

28

29 @ToolParserManager.register_module("openai")

30 class OpenAIToolParser(ToolParser):

31 - def __init__(self, tokenizer: AnyTokenizer):

31 + def __init__(self, tokenizer: "AnyTokenizer"):

32 super().__init__(tokenizer)

33

34 def extract_tool_calls(

vllm/entrypoints/openai/tool_parsers/phi4mini_tool_parser.py

+2 -2

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

3

4 import json

5 from collections.abc import Sequence

6 -from typing import Any, Optional

6 +from typing import Any

7

8 import regex as re

9 from transformers import PreTrainedTokenizerBase

@@ -118,5 +118,5 @@ class Phi4MiniJsonToolParser(ToolParser):

118 current_token_ids: Sequence[int],

119 delta_token_ids: Sequence[int],

120 request: ChatCompletionRequest,

121 - ) -> Optional[DeltaMessage]:

121 + ) -> DeltaMessage | None:

122 return None

vllm/entrypoints/openai/tool_parsers/pythonic_tool_parser.py

+4 -4

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

4 import ast

5 import json

6 from collections.abc import Sequence

7 -from typing import Any, Union

7 +from typing import Any

8

9 import regex as re

10 from transformers import PreTrainedTokenizerBase

@@ -124,7 +124,7 @@ class PythonicToolParser(ToolParser):

124 current_token_ids: Sequence[int],

125 delta_token_ids: Sequence[int],

126 request: ChatCompletionRequest,

127 - ) -> Union[DeltaMessage, None]:

127 + ) -> DeltaMessage | None:

128 if not current_text.startswith("["):

129 return DeltaMessage(content=delta_text)

130

@@ -236,7 +236,7 @@ def _handle_single_tool(call: ast.Call) -> ToolCall:

236 )

237

238

239 -def _make_valid_python(text: str) -> Union[tuple[str, str], None]:

239 +def _make_valid_python(text: str) -> tuple[str, str] | None:

240 bracket_stack = []

241 for index, char in enumerate(text):

242 if char in {"[", "(", "{"}:

@@ -308,7 +308,7 @@ def _make_valid_python(text: str) -> Union[tuple[str, str], None]:

308

309 def _compute_tool_delta(

310 previously_sent_args: str, new_call: ToolCall, index: int, withheld_suffix: str

311 -) -> Union[DeltaToolCall, None]:

311 +) -> DeltaToolCall | None:

312 new_call_args = new_call.function.arguments

313 if withheld_suffix:

314 assert new_call_args.endswith(withheld_suffix)

vllm/entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py

+6 -6

@@ -4,7 +4,7 @@ import ast

4 import json

5 import uuid

6 from collections.abc import Sequence

7 -from typing import Any, Optional, Union

7 +from typing import Any

8

9 import regex as re

10

@@ -36,7 +36,7 @@ class Qwen3CoderToolParser(ToolParser):

36 self.current_tool_name_sent: bool = False

37 self.prev_tool_call_arr: list[dict] = []

38 # Override base class type - we use string IDs for tool calls

39 - self.current_tool_id: Optional[str] = None # type: ignore

39 + self.current_tool_id: str | None = None # type: ignore

40 self.streamed_args_for_tool: list[str] = []

41

42 # Sentinel tokens for streaming mode

@@ -110,7 +110,7 @@ class Qwen3CoderToolParser(ToolParser):

110 self.streaming_request = None

111

112 def _get_arguments_config(

113 - self, func_name: str, tools: Optional[list[ChatCompletionToolsParam]]

113 + self, func_name: str, tools: list[ChatCompletionToolsParam] | None

114 ) -> dict:

115 """Extract argument configuration for a function."""

116 if tools is None:

@@ -240,8 +240,8 @@ class Qwen3CoderToolParser(ToolParser):

240 return param_value

241

242 def _parse_xml_function_call(

243 - self, function_call_str: str, tools: Optional[list[ChatCompletionToolsParam]]

244 - ) -> Optional[ToolCall]:

243 + self, function_call_str: str, tools: list[ChatCompletionToolsParam] | None

244 + ) -> ToolCall | None:

245 # Extract function name

246 end_index = function_call_str.index(">")

247 function_name = function_call_str[:end_index]

@@ -349,7 +349,7 @@ class Qwen3CoderToolParser(ToolParser):

349 current_token_ids: Sequence[int],

350 delta_token_ids: Sequence[int],

351 request: ChatCompletionRequest,

352 - ) -> Union[DeltaMessage, None]:

352 + ) -> DeltaMessage | None:

353 # Store request for type conversion

354 if not previous_text:

355 self._reset_streaming_state()

vllm/entrypoints/openai/tool_parsers/qwen3xml_tool_parser.py

+8 -10

@@ -4,7 +4,7 @@ import ast

4 import json

5 import uuid

6 from collections.abc import Sequence

7 -from typing import Any, Optional, Union

7 +from typing import Any

8 from xml.parsers.expat import ParserCreate

9

10 import regex as re

@@ -39,7 +39,7 @@ class StreamingXMLToolCallParser:

39 self.reset_streaming_state()

40

41 # Tool configuration information

42 - self.tools: Union[list[ChatCompletionToolsParam], None] = None

42 + self.tools: list[ChatCompletionToolsParam] | None = None

43 self.tool_call_start_token: str = "<tool_call>"

44 self.tool_call_end_token: str = "</tool_call>"

45 self.function_start_token: str = "<function="

@@ -341,7 +341,7 @@ class StreamingXMLToolCallParser:

341 # Skip blank content

342 return not element

343

344 - def _find_next_complete_element(self, start_pos: int) -> tuple[Optional[str], int]:

344 + def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]:

345 """

346 Find next complete XML element from specified position

347

@@ -584,7 +584,7 @@ class StreamingXMLToolCallParser:

584 """Emit Delta response (streaming output)"""

585 self.deltas.append(delta)

586

587 - def _auto_close_open_parameter_if_needed(self, incoming_tag: Optional[str] = None):

587 + def _auto_close_open_parameter_if_needed(self, incoming_tag: str | None = None):

588 """Before starting to process new elements,

589 if there are unclosed tags from before,

590 automatically complete their endings to the parser.

@@ -953,7 +953,7 @@ class StreamingXMLToolCallParser:

953 self.parser.EndElementHandler = self._end_element

954 self.parser.CharacterDataHandler = self._char_data

955

956 - def set_tools(self, tools: Union[list[ChatCompletionToolsParam], None]):

956 + def set_tools(self, tools: list[ChatCompletionToolsParam] | None):

957 """Set tool configuration information"""

958 self.tools = tools

959

@@ -961,7 +961,7 @@ class StreamingXMLToolCallParser:

961 """Generate unique call ID"""

962 return f"call_{uuid.uuid4().hex[:24]}"

963

964 - def _extract_function_name(self, name: str, attrs: dict[str, str]) -> Optional[str]:

964 + def _extract_function_name(self, name: str, attrs: dict[str, str]) -> str | None:

965 """Extract function name from various formats"""

966 if attrs and "name" in attrs:

967 return attrs["name"]

@@ -973,9 +973,7 @@ class StreamingXMLToolCallParser:

973

974 return None

975

976 - def _extract_parameter_name(

977 - self, name: str, attrs: dict[str, str]

978 - ) -> Optional[str]:

976 + def _extract_parameter_name(self, name: str, attrs: dict[str, str]) -> str | None:

977 """Extract parameter name from various formats"""

978 if attrs and "name" in attrs:

979 return attrs["name"]

@@ -1218,7 +1216,7 @@ class Qwen3XMLToolParser(ToolParser):

1216 current_token_ids: Sequence[int],

1217 delta_token_ids: Sequence[int],

1218 request: ChatCompletionRequest,

1221 - ) -> Union[DeltaMessage, None]:

1219 + ) -> DeltaMessage | None:

1220 if not previous_text:

1221 self.parser.reset_streaming_state()

1222 if request:

vllm/entrypoints/openai/tool_parsers/seed_oss_tool_parser.py

+4 -4

@@ -7,7 +7,7 @@ import ast

7 import json

8 import uuid

9 from collections.abc import Sequence

10 -from typing import Any, Optional, Union

10 +from typing import Any

11

12 import regex as re

13

@@ -109,8 +109,8 @@ class SeedOssToolParser(ToolParser):

109 self.json_closed = False

110

111 def _parse_xml_function_call(

112 - self, function_call_str: str, tools: Optional[list[ChatCompletionToolsParam]]

113 - ) -> Optional[ToolCall]:

112 + self, function_call_str: str, tools: list[ChatCompletionToolsParam] | None

113 + ) -> ToolCall | None:

114 def get_arguments_config(func_name: str) -> dict:

115 if tools is None:

116 return {}

@@ -357,7 +357,7 @@ class SeedOssToolParser(ToolParser):

357 current_token_ids: Sequence[int],

358 delta_token_ids: Sequence[int],

359 request: ChatCompletionRequest,

360 - ) -> Union[DeltaMessage, None]:

360 + ) -> DeltaMessage | None:

361 # If no delta text, return None unless

362 # it's an EOS token after tool calls

363 if not delta_text:

vllm/entrypoints/openai/tool_parsers/step3_tool_parser.py

+3 -3

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

4 import contextlib

5 import json

6 from collections.abc import Sequence

7 -from typing import Any, Optional, Union

7 +from typing import Any

8

9 import regex as re

10

@@ -58,7 +58,7 @@ class Step3ToolParser(ToolParser):

58 @staticmethod

59 def _parse_steptml_invoke(

60 action_text: str,

61 - ) -> tuple[Optional[str], Optional[dict[str, str]]]:

61 + ) -> tuple[str | None, dict[str, str] | None]:

62 func_name_match = re.search(r'<steptml:invoke name="([^"]+)">', action_text)

63 if not func_name_match:

64 return None, None

@@ -117,7 +117,7 @@ class Step3ToolParser(ToolParser):

117 current_token_ids: Sequence[int],

118 delta_token_ids: Sequence[int],

119 request: ChatCompletionRequest,

120 - ) -> Union[DeltaMessage, None]:

120 + ) -> DeltaMessage | None:

121 # The main loop processes the stream from the last known position.

122 while True:

123 if self.position >= len(current_text):

vllm/entrypoints/renderer.py

+38 -43

@@ -5,7 +5,7 @@ import asyncio

5 import io

6 from abc import ABC, abstractmethod

7 from dataclasses import dataclass

8 -from typing import Annotated, Optional, Union

8 +from typing import Annotated

9

10 import pybase64

11 import torch

@@ -24,25 +24,25 @@ from vllm.utils import AsyncMicrobatchTokenizer

24 class RenderConfig:

25 """Configuration to control how prompts are prepared."""

26

27 - max_length: Optional[int] = None

27 + max_length: int | None = None

28 """Maximum allowable total input token length. If provided,

29 token inputs longer than this raise ``ValueError``."""

30

31 - truncate_prompt_tokens: Optional[int] = None

31 + truncate_prompt_tokens: int | None = None

32 """Number of tokens to keep. ``None`` means no truncation.

33 ``0`` yields an empty list (and skips embeds).

34 ``-1`` maps to ``model_config.max_model_len``."""

35

36 - add_special_tokens: Optional[bool] = True

36 + add_special_tokens: bool | None = True

37 """Whether to add model-specific special tokens during tokenization."""

38

39 - cache_salt: Optional[str] = None

39 + cache_salt: str | None = None

40 """String to disambiguate prefix cache entries."""

41

42 - needs_detokenization: Optional[bool] = False

42 + needs_detokenization: bool | None = False

43 """If True, detokenize IDs back to text for inclusion in outputs."""

44

45 - def verify_truncate_prompt_tokens(self, model_config: ModelConfig) -> Optional[int]:

45 + def verify_truncate_prompt_tokens(self, model_config: ModelConfig) -> int | None:

46 """Validate and normalize `truncate_prompt_tokens` parameter."""

47 truncate_prompt_tokens = self.truncate_prompt_tokens

48 if truncate_prompt_tokens is None:

@@ -85,7 +85,7 @@ class BaseRenderer(ABC):

85 def __init__(

86 self,

87 model_config: ModelConfig,

88 - tokenizer: Optional[AnyTokenizer] = None,

88 + tokenizer: AnyTokenizer | None = None,

89 ):

90 super().__init__()

91 self.model_config = model_config

@@ -95,7 +95,7 @@ class BaseRenderer(ABC):

95 async def render_prompt(

96 self,

97 *,

98 - prompt_or_prompts: Union[str, list[str], list[int], list[list[int]]],

98 + prompt_or_prompts: str | list[str] | list[int] | list[list[int]],

99 config: RenderConfig,

100 ) -> list[EngineTokensPrompt]:

101 """

@@ -126,12 +126,10 @@ class BaseRenderer(ABC):

126 async def render_prompt_and_embeds(

127 self,

128 *,

129 - prompt_or_prompts: Optional[

130 - Union[str, list[str], list[int], list[list[int]]]

131 - ] = None,

132 - prompt_embeds: Optional[Union[bytes, list[bytes]]] = None,

129 + prompt_or_prompts: str | list[str] | list[int] | list[list[int]] | None = None,

130 + prompt_embeds: bytes | list[bytes] | None = None,

131 config: RenderConfig,

134 - ) -> list[Union[EngineTokensPrompt, EngineEmbedsPrompt]]:

132 + ) -> list[EngineTokensPrompt | EngineEmbedsPrompt]:

133 """

134 Convert text/token and/or base64-encoded embeddings inputs into

135 engine-ready prompt objects using a unified RenderConfig.

@@ -161,9 +159,9 @@ class BaseRenderer(ABC):

159 @classmethod

160 def load_prompt_embeds(

161 cls,

164 - prompt_embeds: Union[bytes, list[bytes]],

165 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=0)]] = None,

166 - cache_salt: Optional[str] = None,

162 + prompt_embeds: bytes | list[bytes],

163 + truncate_prompt_tokens: Annotated[int, Field(ge=0)] | None = None,

164 + cache_salt: str | None = None,

165 ) -> list[EngineEmbedsPrompt]:

166 """Load and validate base64-encoded embeddings into prompt objects."""

167

@@ -199,19 +197,18 @@ class CompletionRenderer(BaseRenderer):

197 def __init__(

198 self,

199 model_config: ModelConfig,

202 - tokenizer: Optional[AnyTokenizer] = None,

203 - async_tokenizer_pool: Optional[

204 - dict[AnyTokenizer, AsyncMicrobatchTokenizer]

205 - ] = None,

200 + tokenizer: AnyTokenizer | None = None,

201 + async_tokenizer_pool: dict[AnyTokenizer, AsyncMicrobatchTokenizer]

202 + | None = None,

203 ):

204 super().__init__(model_config, tokenizer)

205 self.async_tokenizer_pool = async_tokenizer_pool

209 - self.async_tokenizer: Optional[AsyncMicrobatchTokenizer] = None

206 + self.async_tokenizer: AsyncMicrobatchTokenizer | None = None

207

208 async def render_prompt(

209 self,

210 *,

214 - prompt_or_prompts: Union[str, list[str], list[int], list[list[int]]],

211 + prompt_or_prompts: str | list[str] | list[int] | list[list[int]],

212 config: RenderConfig,

213 ) -> list[EngineTokensPrompt]:

214 """Implementation of prompt rendering for completion-style requests.

@@ -237,12 +234,10 @@ class CompletionRenderer(BaseRenderer):

234 async def render_prompt_and_embeds(

235 self,

236 *,

240 - prompt_or_prompts: Optional[

241 - Union[str, list[str], list[int], list[list[int]]]

242 - ] = None,

243 - prompt_embeds: Optional[Union[bytes, list[bytes]]] = None,

237 + prompt_or_prompts: str | list[str] | list[int] | list[list[int]] | None = None,

238 + prompt_embeds: bytes | list[bytes] | None = None,

239 config: RenderConfig,

245 - ) -> list[Union[EngineTokensPrompt, EngineEmbedsPrompt]]:

240 + ) -> list[EngineTokensPrompt | EngineEmbedsPrompt]:

241 """

242 Render text/token prompts and/or precomputed embedding prompts. At

243 least one of `prompt_or_prompts` or `prompt_embeds` must be provided.

@@ -251,7 +246,7 @@ class CompletionRenderer(BaseRenderer):

246 if truncate_prompt_tokens == 0:

247 return []

248

254 - rendered: list[Union[EngineTokensPrompt, EngineEmbedsPrompt]] = []

249 + rendered: list[EngineTokensPrompt | EngineEmbedsPrompt] = []

250

251 if prompt_embeds is not None:

252 rendered.extend(

@@ -271,7 +266,7 @@ class CompletionRenderer(BaseRenderer):

266 return rendered

267

268 def _maybe_apply_truncation(

274 - self, token_ids: list[int], truncate_prompt_tokens: Optional[int]

269 + self, token_ids: list[int], truncate_prompt_tokens: int | None

270 ) -> list[int]:

271 """Apply truncation to token sequence."""

272 if truncate_prompt_tokens is None:

@@ -283,9 +278,9 @@ class CompletionRenderer(BaseRenderer):

278

279 async def _create_prompt(

280 self,

286 - prompt_input: Union[EngineTextPrompt, EngineTokensPrompt],

281 + prompt_input: EngineTextPrompt | EngineTokensPrompt,

282 config: RenderConfig,

288 - truncate_prompt_tokens: Optional[int],

283 + truncate_prompt_tokens: int | None,

284 ) -> EngineTokensPrompt:

285 prompt, prompt_token_ids, _ = get_prompt_components(prompt_input)

286

@@ -315,10 +310,10 @@ class CompletionRenderer(BaseRenderer):

310 async def _create_prompt_from_text(

311 self,

312 text: str,

318 - max_length: Optional[int],

319 - truncate_prompt_tokens: Optional[int],

320 - add_special_tokens: Optional[bool],

321 - cache_salt: Optional[str],

313 + max_length: int | None,

314 + truncate_prompt_tokens: int | None,

315 + add_special_tokens: bool | None,

316 + cache_salt: str | None,

317 ) -> EngineTokensPrompt:

318 """Tokenize text input asynchronously."""

319 async_tokenizer = self._get_async_tokenizer()

@@ -348,10 +343,10 @@ class CompletionRenderer(BaseRenderer):

343 async def _create_prompt_from_token_ids(

344 self,

345 token_ids: list[int],

351 - max_length: Optional[int],

352 - truncate_prompt_tokens: Optional[int],

353 - cache_salt: Optional[str],

354 - needs_detokenization: Optional[bool] = False,

346 + max_length: int | None,

347 + truncate_prompt_tokens: int | None,

348 + cache_salt: str | None,

349 + needs_detokenization: bool | None = False,

350 ) -> EngineTokensPrompt:

351 """Optionally detokenize token IDs and build a tokens prompt."""

352 token_ids = self._maybe_apply_truncation(token_ids, truncate_prompt_tokens)

@@ -391,9 +386,9 @@ class CompletionRenderer(BaseRenderer):

386 def _create_tokens_prompt(

387 self,

388 token_ids: list[int],

394 - max_length: Optional[int] = None,

395 - cache_salt: Optional[str] = None,

396 - prompt: Optional[str] = None,

389 + max_length: int | None = None,

390 + cache_salt: str | None = None,

391 + prompt: str | None = None,

392 ) -> EngineTokensPrompt:

393 """Create validated EngineTokensPrompt."""

394 if max_length is not None and len(token_ids) > max_length:

vllm/entrypoints/score_utils.py

+17 -17

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Any, Optional, Union, cast

3 +from typing import Any, TypeAlias, cast

4

5 from torch.nn import CosineSimilarity

6 -from typing_extensions import Required, TypeAlias, TypedDict

6 +from typing_extensions import Required, TypedDict

7

8 from vllm.config import ModelConfig

9 from vllm.entrypoints.chat_utils import (

@@ -25,9 +25,9 @@ from vllm.transformers_utils.tokenizer import (

25 PreTrainedTokenizerFast,

26 )

27

28 -ScoreContentPartParam: TypeAlias = Union[

29 - ChatCompletionContentPartImageParam, ChatCompletionContentPartImageEmbedsParam

30 -]

28 +ScoreContentPartParam: TypeAlias = (

29 + ChatCompletionContentPartImageParam | ChatCompletionContentPartImageEmbedsParam

30 +)

31

32

33 class ScoreMultiModalParam(TypedDict, total=False):

@@ -45,12 +45,12 @@ class ScoreMultiModalParam(TypedDict, total=False):

45

46

47 def _cosine_similarity(

48 - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

48 + tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,

49 embed_1: list[PoolingRequestOutput],

50 embed_2: list[PoolingRequestOutput],

51 ) -> list[PoolingRequestOutput]:

52 scorer = CosineSimilarity(0)

53 - scores: Union[list[PoolingRequestOutput]] = []

53 + scores: list[PoolingRequestOutput] = []

54

55 for emb_1, emb_2 in zip(embed_1, embed_2):

56 pair_score = scorer(emb_1.outputs.data, emb_2.outputs.data)

@@ -74,8 +74,8 @@ def _cosine_similarity(

74

75

76 def _validate_score_input_lens(

77 - data_1: Union[list[str], list[ScoreContentPartParam]],

78 - data_2: Union[list[str], list[ScoreContentPartParam]],

77 + data_1: list[str] | list[ScoreContentPartParam],

78 + data_2: list[str] | list[ScoreContentPartParam],

79 ):

80 len_1 = len(data_1)

81 len_2 = len(data_2)

@@ -89,18 +89,18 @@ def _validate_score_input_lens(

89

90

91 def parse_score_data(

92 - data_1: Union[str, ScoreContentPartParam],

93 - data_2: Union[str, ScoreContentPartParam],

92 + data_1: str | ScoreContentPartParam,

93 + data_2: str | ScoreContentPartParam,

94 model_config: ModelConfig,

95 tokenizer: AnyTokenizer,

96 -) -> tuple[str, str, Optional[MultiModalDataDict]]:

96 +) -> tuple[str, str, MultiModalDataDict | None]:

97 mm_tracker = MultiModalItemTracker(model_config, tokenizer)

98

99 content_1 = _parse_score_content(data_1, mm_tracker)

100

101 content_2 = _parse_score_content(data_2, mm_tracker)

102

103 - def ensure_str(content: Optional[_ContentPart]) -> str:

103 + def ensure_str(content: _ContentPart | None) -> str:

104 if content is not None and isinstance(content, str):

105 return cast(str, content)

106 else:

@@ -113,9 +113,9 @@ def parse_score_data(

113

114

115 def _parse_score_content(

116 - data: Union[str, ScoreContentPartParam],

116 + data: str | ScoreContentPartParam,

117 mm_tracker: BaseMultiModalItemTracker,

118 -) -> Optional[_ContentPart]:

118 +) -> _ContentPart | None:

119 if isinstance(data, str):

120 data = ChatCompletionContentPartTextParam(type="text", text=data)

121

@@ -182,8 +182,8 @@ def get_score_prompt(

182 model_config: ModelConfig,

183 tokenizer: AnyTokenizer,

184 tokenization_kwargs: dict[str, Any],

185 - data_1: Union[str, ScoreContentPartParam],

186 - data_2: Union[str, ScoreContentPartParam],

185 + data_1: str | ScoreContentPartParam,

186 + data_2: str | ScoreContentPartParam,

187 ) -> tuple[str, TokensPrompt]:

188 prompt_1, prompt_2, mm_data = parse_score_data(

189 data_1,

vllm/entrypoints/ssl.py

+4 -4

@@ -2,8 +2,8 @@

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

3

4 import asyncio

5 +from collections.abc import Callable

6 from ssl import SSLContext

6 -from typing import Callable, Optional

7

8 from watchfiles import Change, awatch

9

@@ -20,9 +20,9 @@ class SSLCertRefresher:

20 def __init__(

21 self,

22 ssl_context: SSLContext,

23 - key_path: Optional[str] = None,

24 - cert_path: Optional[str] = None,

25 - ca_path: Optional[str] = None,

23 + key_path: str | None = None,

24 + cert_path: str | None = None,

25 + ca_path: str | None = None,

26 ) -> None:

27 self.ssl = ssl_context

28 self.key_path = key_path

vllm/entrypoints/tool_server.py

+6 -6

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

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

3 from abc import ABC, abstractmethod

4 from contextlib import AbstractAsyncContextManager, asynccontextmanager

5 -from typing import TYPE_CHECKING, Any, Optional

5 +from typing import TYPE_CHECKING, Any

6

7 from openai_harmony import ToolDescription, ToolNamespaceConfig

8

@@ -80,7 +80,7 @@ class ToolServer(ABC):

80 pass

81

82 @abstractmethod

83 - def get_tool_description(self, tool_name: str) -> Optional[ToolNamespaceConfig]:

83 + def get_tool_description(self, tool_name: str) -> ToolNamespaceConfig | None:

84 """

85 Return the tool description for the given tool name.

86 If the tool is not supported, return None.

@@ -89,7 +89,7 @@ class ToolServer(ABC):

89

90 @abstractmethod

91 def new_session(

92 - self, tool_name: str, session_id: str, headers: Optional[dict[str, str]] = None

92 + self, tool_name: str, session_id: str, headers: dict[str, str] | None = None

93 ) -> AbstractAsyncContextManager[Any]:

94 """

95 Create a session for the tool.

@@ -152,7 +152,7 @@ class MCPToolServer(ToolServer):

152

153 @asynccontextmanager

154 async def new_session(

155 - self, tool_name: str, session_id: str, headers: Optional[dict[str, str]] = None

155 + self, tool_name: str, session_id: str, headers: dict[str, str] | None = None

156 ):

157 from mcp import ClientSession

158 from mcp.client.sse import sse_client

@@ -190,7 +190,7 @@ class DemoToolServer(ToolServer):

190 def has_tool(self, tool_name: str) -> bool:

191 return tool_name in self.tools

192

193 - def get_tool_description(self, tool_name: str) -> Optional[ToolNamespaceConfig]:

193 + def get_tool_description(self, tool_name: str) -> ToolNamespaceConfig | None:

194 if tool_name not in self.tools:

195 return None

196 if tool_name == "browser":

@@ -202,7 +202,7 @@ class DemoToolServer(ToolServer):

202

203 @asynccontextmanager

204 async def new_session(

205 - self, tool_name: str, session_id: str, headers: Optional[dict[str, str]] = None

205 + self, tool_name: str, session_id: str, headers: dict[str, str] | None = None

206 ):

207 if tool_name not in self.tools:

208 raise KeyError(f"Tool '{tool_name}' is not supported")

vllm/entrypoints/utils.py

+6 -6

@@ -6,7 +6,7 @@ import dataclasses

6 import functools

7 import os

8 from argparse import Namespace

9 -from typing import Any, Optional, Union

9 +from typing import Any

10

11 from fastapi import Request

12 from fastapi.responses import JSONResponse, StreamingResponse

@@ -164,9 +164,9 @@ def cli_env_setup():

164

165 def _validate_truncation_size(

166 max_model_len: int,

167 - truncate_prompt_tokens: Optional[int],

168 - tokenization_kwargs: Optional[dict[str, Any]] = None,

169 -) -> Optional[int]:

167 + truncate_prompt_tokens: int | None,

168 + tokenization_kwargs: dict[str, Any] | None = None,

169 +) -> int | None:

170 if truncate_prompt_tokens is not None:

171 if truncate_prompt_tokens <= -1:

172 truncate_prompt_tokens = max_model_len

@@ -191,7 +191,7 @@ def _validate_truncation_size(

191

192 def get_max_tokens(

193 max_model_len: int,

194 - request: Union[ChatCompletionRequest, CompletionRequest],

194 + request: ChatCompletionRequest | CompletionRequest,

195 input_length: int,

196 default_sampling_params: dict,

197 ) -> int:

@@ -211,7 +211,7 @@ def get_max_tokens(

211 )

212

213

214 -def log_non_default_args(args: Union[Namespace, EngineArgs]):

214 +def log_non_default_args(args: Namespace | EngineArgs):

215 non_default_args = {}

216

217 # Handle Namespace

vllm/envs.py

+43 -42

@@ -6,27 +6,28 @@ import json

6 import os

7 import sys

8 import tempfile

9 -from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union

9 +from collections.abc import Callable

10 +from typing import TYPE_CHECKING, Any, Literal

11

12 if TYPE_CHECKING:

13 VLLM_HOST_IP: str = ""

13 - VLLM_PORT: Optional[int] = None

14 + VLLM_PORT: int | None = None

15 VLLM_RPC_BASE_PATH: str = tempfile.gettempdir()

16 VLLM_USE_MODELSCOPE: bool = False

17 VLLM_RINGBUFFER_WARNING_INTERVAL: int = 60

17 - VLLM_NCCL_SO_PATH: Optional[str] = None

18 - LD_LIBRARY_PATH: Optional[str] = None

18 + VLLM_NCCL_SO_PATH: str | None = None

19 + LD_LIBRARY_PATH: str | None = None

20 VLLM_USE_TRITON_FLASH_ATTN: bool = True

21 VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False

21 - VLLM_FLASH_ATTN_VERSION: Optional[int] = None

22 + VLLM_FLASH_ATTN_VERSION: int | None = None

23 LOCAL_RANK: int = 0

23 - CUDA_VISIBLE_DEVICES: Optional[str] = None

24 + CUDA_VISIBLE_DEVICES: str | None = None

25 VLLM_ENGINE_ITERATION_TIMEOUT_S: int = 60

25 - VLLM_API_KEY: Optional[str] = None

26 - S3_ACCESS_KEY_ID: Optional[str] = None

27 - S3_SECRET_ACCESS_KEY: Optional[str] = None

28 - S3_ENDPOINT_URL: Optional[str] = None

29 - VLLM_MODEL_REDIRECT_PATH: Optional[str] = None

26 + VLLM_API_KEY: str | None = None

27 + S3_ACCESS_KEY_ID: str | None = None

28 + S3_SECRET_ACCESS_KEY: str | None = None

29 + S3_ENDPOINT_URL: str | None = None

30 + VLLM_MODEL_REDIRECT_PATH: str | None = None

31 VLLM_CACHE_ROOT: str = os.path.expanduser("~/.cache/vllm")

32 VLLM_CONFIG_ROOT: str = os.path.expanduser("~/.config/vllm")

33 VLLM_USAGE_STATS_SERVER: str = "https://stats.vllm.ai"

@@ -38,16 +39,16 @@ if TYPE_CHECKING:

39 VLLM_LOGGING_LEVEL: str = "INFO"

40 VLLM_LOGGING_PREFIX: str = ""

41 VLLM_LOGGING_STREAM: str = "ext://sys.stdout"

41 - VLLM_LOGGING_CONFIG_PATH: Optional[str] = None

42 - VLLM_LOGITS_PROCESSOR_THREADS: Optional[int] = None

42 + VLLM_LOGGING_CONFIG_PATH: str | None = None

43 + VLLM_LOGITS_PROCESSOR_THREADS: int | None = None

44 VLLM_LOG_STATS_INTERVAL: float = 10.0

45 VLLM_TRACE_FUNCTION: int = 0

45 - VLLM_ATTENTION_BACKEND: Optional[str] = None

46 - VLLM_USE_FLASHINFER_SAMPLER: Optional[bool] = None

47 - VLLM_PP_LAYER_PARTITION: Optional[str] = None

48 - VLLM_CPU_KVCACHE_SPACE: Optional[int] = 0

46 + VLLM_ATTENTION_BACKEND: str | None = None

47 + VLLM_USE_FLASHINFER_SAMPLER: bool | None = None

48 + VLLM_PP_LAYER_PARTITION: str | None = None

49 + VLLM_CPU_KVCACHE_SPACE: int | None = 0

50 VLLM_CPU_OMP_THREADS_BIND: str = ""

50 - VLLM_CPU_NUM_OF_RESERVED_CPU: Optional[int] = None

51 + VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None

52 VLLM_CPU_MOE_PREPACK: bool = True

53 VLLM_CPU_SGL_KERNEL: bool = False

54 VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")

@@ -73,20 +74,20 @@ if TYPE_CHECKING:

74 VLLM_MM_INPUT_CACHE_GIB: int = 4

75 VLLM_TARGET_DEVICE: str = "cuda"

76 VLLM_MAIN_CUDA_VERSION: str = "12.8"

76 - MAX_JOBS: Optional[str] = None

77 - NVCC_THREADS: Optional[str] = None

77 + MAX_JOBS: str | None = None

78 + NVCC_THREADS: str | None = None

79 VLLM_USE_PRECOMPILED: bool = False

80 VLLM_DOCKER_BUILD_CONTEXT: bool = False

81 VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL: bool = False

82 VLLM_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False

82 - CMAKE_BUILD_TYPE: Optional[Literal["Debug", "Release", "RelWithDebInfo"]] = None

83 + CMAKE_BUILD_TYPE: Literal["Debug", "Release", "RelWithDebInfo"] | None = None

84 VERBOSE: bool = False

85 VLLM_ALLOW_LONG_MAX_MODEL_LEN: bool = False

86 VLLM_RPC_TIMEOUT: int = 10000 # ms

87 VLLM_HTTP_TIMEOUT_KEEP_ALIVE: int = 5 # seconds

87 - VLLM_PLUGINS: Optional[list[str]] = None

88 - VLLM_LORA_RESOLVER_CACHE_DIR: Optional[str] = None

89 - VLLM_TORCH_PROFILER_DIR: Optional[str] = None

88 + VLLM_PLUGINS: list[str] | None = None

89 + VLLM_LORA_RESOLVER_CACHE_DIR: str | None = None

90 + VLLM_TORCH_PROFILER_DIR: str | None = None

91 VLLM_TORCH_PROFILER_RECORD_SHAPES: bool = False

92 VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False

93 VLLM_USE_AOT_COMPILE: bool = False

@@ -126,7 +127,7 @@ if TYPE_CHECKING:

127 VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH: int = 32

128 VLLM_RAY_PER_WORKER_GPUS: float = 1.0

129 VLLM_RAY_BUNDLE_INDICES: str = ""

129 - VLLM_CUDART_SO_PATH: Optional[str] = None

130 + VLLM_CUDART_SO_PATH: str | None = None

131 VLLM_DP_RANK: int = 0

132 VLLM_DP_RANK_LOCAL: int = -1

133 VLLM_DP_SIZE: int = 1

@@ -137,11 +138,11 @@ if TYPE_CHECKING:

138 VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False

139 VLLM_RAY_DP_PACK_STRATEGY: str = "strict"

140 VLLM_MARLIN_USE_ATOMIC_ADD: bool = False

140 - VLLM_MXFP4_USE_MARLIN: Optional[bool] = None

141 + VLLM_MXFP4_USE_MARLIN: bool | None = None

142 VLLM_V0_USE_OUTLINES_CACHE: bool = False

143 VLLM_V1_USE_OUTLINES_CACHE: bool = False

144 VLLM_TPU_BUCKET_PADDING_GAP: int = 0

144 - VLLM_TPU_MOST_MODEL_LEN: Optional[int] = None

145 + VLLM_TPU_MOST_MODEL_LEN: int | None = None

146 VLLM_TPU_USING_PATHWAYS: bool = False

147 VLLM_USE_DEEP_GEMM: bool = True

148 VLLM_USE_DEEP_GEMM_E8M0: bool = True

@@ -169,21 +170,21 @@ if TYPE_CHECKING:

170 VLLM_SLEEP_WHEN_IDLE: bool = False

171 VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16

172 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300

172 - VLLM_KV_CACHE_LAYOUT: Optional[Literal["NHD", "HND"]] = None

173 + VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None

174 VLLM_COMPUTE_NANS_IN_LOGITS: bool = False

175 VLLM_USE_NVFP4_CT_EMULATIONS: bool = False

176 VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[

177 "FP", "INT8", "INT6", "INT4", "NONE"

178 ] = "NONE"

179 VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True

179 - VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: Optional[int] = None

180 + VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None

181 VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 480

182 VLLM_USE_CUDNN_PREFILL: bool = False

183 VLLM_ENABLE_CUDAGRAPH_GC: bool = False

184 VLLM_LOOPBACK_IP: str = ""

185 VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = False

186 VLLM_ENABLE_RESPONSES_API_STORE: bool = False

186 - VLLM_USE_TRTLLM_ATTENTION: Optional[str] = None

187 + VLLM_USE_TRTLLM_ATTENTION: str | None = None

188 VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION: bool = False

189 VLLM_HAS_FLASHINFER_CUBIN: bool = False

190 VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False

@@ -191,7 +192,7 @@ if TYPE_CHECKING:

192 VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False

193 VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False

194 VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True

194 - VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None

195 + VLLM_TUNED_CONFIG_FOLDER: str | None = None

196 VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False

197 VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False

198 VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False

@@ -201,12 +202,12 @@ if TYPE_CHECKING:

202 VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024

203 VLLM_DBO_COMM_SMS: int = 20

204 GPT_OSS_SYSTEM_TOOL_MCP_LABELS: list[str] = []

204 - VLLM_PATTERN_MATCH_DEBUG: Optional[str] = None

205 - VLLM_DEBUG_DUMP_PATH: Optional[str] = None

205 + VLLM_PATTERN_MATCH_DEBUG: str | None = None

206 + VLLM_DEBUG_DUMP_PATH: str | None = None

207 VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE: bool = True

208 VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING: bool = True

209 VLLM_USE_NCCL_SYMM_MEM: bool = False

209 - VLLM_NCCL_INCLUDE_PATH: Optional[str] = None

210 + VLLM_NCCL_INCLUDE_PATH: str | None = None

211 VLLM_USE_FBGEMM: bool = False

212 VLLM_GC_DEBUG: str = ""

213

@@ -225,13 +226,13 @@ def get_default_config_root():

226 )

227

228

228 -def maybe_convert_int(value: Optional[str]) -> Optional[int]:

229 +def maybe_convert_int(value: str | None) -> int | None:

230 if value is None:

231 return None

232 return int(value)

233

234

234 -def maybe_convert_bool(value: Optional[str]) -> Optional[bool]:

235 +def maybe_convert_bool(value: str | None) -> bool | None:

236 if value is None:

237 return None

238 return bool(int(value))

@@ -246,10 +247,10 @@ def use_aot_compile() -> bool:

247

248 def env_with_choices(

249 env_name: str,

249 - default: Optional[str],

250 - choices: Union[list[str], Callable[[], list[str]]],

250 + default: str | None,

251 + choices: list[str] | Callable[[], list[str]],

252 case_sensitive: bool = True,

252 -) -> Callable[[], Optional[str]]:

253 +) -> Callable[[], str | None]:

254 """

255 Create a lambda that validates environment variable against allowed choices

256

@@ -263,7 +264,7 @@ def env_with_choices(

264 Lambda function for environment_variables dict

265 """

266

266 - def _get_validated_env() -> Optional[str]:

267 + def _get_validated_env() -> str | None:

268 value = os.getenv(env_name)

269 if value is None:

270 return default

@@ -292,7 +293,7 @@ def env_with_choices(

293 def env_list_with_choices(

294 env_name: str,

295 default: list[str],

295 - choices: Union[list[str], Callable[[], list[str]]],

296 + choices: list[str] | Callable[[], list[str]],

297 case_sensitive: bool = True,

298 ) -> Callable[[], list[str]]:

299 """

@@ -344,7 +345,7 @@ def env_list_with_choices(

345 return _get_validated_env_list

346

347

347 -def get_vllm_port() -> Optional[int]:

348 +def get_vllm_port() -> int | None:

349 """Get the port from VLLM_PORT environment variable.

350

351 Returns:

vllm/executor/executor_base.py

+19 -19

@@ -4,9 +4,9 @@

4 import asyncio

5 import time

6 from abc import ABC, abstractmethod

7 -from collections.abc import Awaitable

7 +from collections.abc import Awaitable, Callable

8 from functools import cached_property

9 -from typing import Any, Callable, Optional, Union

9 +from typing import Any

10

11 from typing_extensions import TypeVar

12

@@ -63,10 +63,10 @@ class ExecutorBase(ABC):

63 @abstractmethod

64 def collective_rpc(

65 self,

66 - method: Union[str, Callable[[WorkerBase], _R]],

67 - timeout: Optional[float] = None,

66 + method: str | Callable[[WorkerBase], _R],

67 + timeout: float | None = None,

68 args: tuple = (),

69 - kwargs: Optional[dict[str, Any]] = None,

69 + kwargs: dict[str, Any] | None = None,

70 ) -> list[_R]:

71 """

72 Execute an RPC call on all workers.

@@ -143,7 +143,7 @@ class ExecutorBase(ABC):

143

144 def execute_model(

145 self, execute_model_req: ExecuteModelRequest

146 - ) -> Optional[list[Union[SamplerOutput, PoolerOutput]]]:

146 + ) -> list[SamplerOutput | PoolerOutput] | None:

147 output = self.collective_rpc("execute_model", args=(execute_model_req,))

148 return output[0]

149

@@ -192,7 +192,7 @@ class ExecutorBase(ABC):

192 "It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep

193 )

194

195 - def wake_up(self, tags: Optional[list[str]] = None):

195 + def wake_up(self, tags: list[str] | None = None):

196 if not self.is_sleeping:

197 logger.warning("Executor is not sleeping.")

198 return

@@ -222,8 +222,8 @@ class ExecutorBase(ABC):

222 def save_sharded_state(

223 self,

224 path: str,

225 - pattern: Optional[str] = None,

226 - max_size: Optional[int] = None,

225 + pattern: str | None = None,

226 + max_size: int | None = None,

227 ) -> None:

228 self.collective_rpc(

229 "save_sharded_state",

@@ -256,7 +256,7 @@ class ExecutorBase(ABC):

256 exception."""

257 self.check_health()

258

259 - def init_kv_output_aggregator(self, finished_count: Optional[int]) -> None:

259 + def init_kv_output_aggregator(self, finished_count: int | None) -> None:

260 """Init KVOutputAggregator"""

261 self.kv_output_aggregator = KVOutputAggregator(

262 finished_count or self.parallel_config.world_size

@@ -269,7 +269,7 @@ class DistributedExecutorBase(ExecutorBase):

269 def __init__(self, *args, **kwargs):

270 # This is non-None when the execute model loop is running

271 # in the parallel workers. It's a coroutine in the AsyncLLMEngine case.

272 - self.parallel_worker_tasks: Optional[Union[Any, Awaitable[Any]]] = None

272 + self.parallel_worker_tasks: Any | Awaitable[Any] | None = None

273

274 super().__init__(*args, **kwargs)

275

@@ -302,8 +302,8 @@ class DistributedExecutorBase(ExecutorBase):

302

303 @abstractmethod

304 def _driver_execute_model(

305 - self, execute_model_req: Optional[ExecuteModelRequest]

306 - ) -> Optional[list[SamplerOutput]]:

305 + self, execute_model_req: ExecuteModelRequest | None

306 + ) -> list[SamplerOutput] | None:

307 """Run execute_model in the driver worker.

308

309 Passing None will cause the driver to stop the model execution loop

@@ -314,20 +314,20 @@ class DistributedExecutorBase(ExecutorBase):

314

315 def collective_rpc(

316 self,

317 - method: Union[str, Callable],

318 - timeout: Optional[float] = None,

317 + method: str | Callable,

318 + timeout: float | None = None,

319 args: tuple = (),

320 - kwargs: Optional[dict[str, Any]] = None,

320 + kwargs: dict[str, Any] | None = None,

321 ) -> list[Any]:

322 return self._run_workers(method, *args, **(kwargs or {}))

323

324 @abstractmethod

325 def _run_workers(

326 self,

327 - method: Union[str, Callable],

327 + method: str | Callable,

328 *args,

329 async_run_tensor_parallel_workers_only: bool = False,

330 - max_concurrent_workers: Optional[int] = None,

330 + max_concurrent_workers: int | None = None,

331 **kwargs,

332 ) -> Any:

333 """Runs the given method on all workers.

@@ -374,7 +374,7 @@ class DistributedExecutorBase(ExecutorBase):

374 @abstractmethod

375 async def _driver_execute_model_async(

376 self,

377 - execute_model_req: Optional[ExecuteModelRequest] = None,

377 + execute_model_req: ExecuteModelRequest | None = None,

378 ) -> list[SamplerOutput]:

379 """Execute the model asynchronously in the driver worker.

380

vllm/executor/ray_distributed_executor.py

+11 -10

@@ -4,8 +4,9 @@

4 import asyncio

5 import os

6 from collections import defaultdict

7 +from collections.abc import Callable

8 from dataclasses import dataclass

8 -from typing import TYPE_CHECKING, Any, Callable, Optional, Union

9 +from typing import TYPE_CHECKING, Any

10

11 import cloudpickle

12 import msgspec

@@ -71,7 +72,7 @@ class RayDistributedExecutor(DistributedExecutorBase):

72 uses_ray: bool = True

73

74 def _init_executor(self) -> None:

74 - self.forward_dag: Optional[ray.dag.CompiledDAG] = None

75 + self.forward_dag: ray.dag.CompiledDAG | None = None

76 if envs.VLLM_USE_V1:

77 # V1 uses SPMD worker and compiled DAG

78 os.environ["VLLM_USE_RAY_SPMD_WORKER"] = "1"

@@ -114,10 +115,10 @@ class RayDistributedExecutor(DistributedExecutorBase):

115 self._init_workers_ray(placement_group)

116

117 self.input_encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)

117 - self.output_decoder = msgspec.msgpack.Decoder(Optional[list[SamplerOutput]])

118 + self.output_decoder = msgspec.msgpack.Decoder(list[SamplerOutput] | None)

119 self.use_v1 = envs.VLLM_USE_V1

120

120 - self.pp_locks: Optional[list[asyncio.Lock]] = None

121 + self.pp_locks: list[asyncio.Lock] | None = None

122 if not self.use_ray_compiled_dag:

123 self.driver_exec_method = make_async(self.driver_worker.execute_method)

124

@@ -162,7 +163,7 @@ class RayDistributedExecutor(DistributedExecutorBase):

163

164 # The driver dummy worker does not actually use any resources.

165 # It holds the resource for the driver worker.

165 - self.driver_dummy_worker: Optional[RayWorkerWrapper] = None

166 + self.driver_dummy_worker: RayWorkerWrapper | None = None

167 # The remaining workers are the actual ray actors.

168 self.workers: list[RayWorkerWrapper] = []

169

@@ -432,8 +433,8 @@ class RayDistributedExecutor(DistributedExecutorBase):

433 self.non_driver_workers.append(worker)

434

435 def _driver_execute_model(

435 - self, execute_model_req: Optional[ExecuteModelRequest]

436 - ) -> Optional[list[SamplerOutput]]:

436 + self, execute_model_req: ExecuteModelRequest | None

437 + ) -> list[SamplerOutput] | None:

438 """Run execute_model in the driver worker.

439

440 Passing None will cause the driver to stop the model execution

@@ -463,10 +464,10 @@ class RayDistributedExecutor(DistributedExecutorBase):

464

465 def _run_workers(

466 self,

466 - method: Union[str, Callable],

467 + method: str | Callable,

468 *args,

469 async_run_tensor_parallel_workers_only: bool = False,

469 - max_concurrent_workers: Optional[int] = None,

470 + max_concurrent_workers: int | None = None,

471 **kwargs,

472 ) -> Any:

473 """Runs the given method on all workers. Can be used in the following

@@ -683,7 +684,7 @@ class RayDistributedExecutor(DistributedExecutorBase):

684 return self.output_decoder.decode(output)

685

686 async def _driver_execute_model_async(

686 - self, execute_model_req: Optional[ExecuteModelRequest] = None

687 + self, execute_model_req: ExecuteModelRequest | None = None

688 ) -> list[SamplerOutput]:

689 assert not self.use_ray_spmd_worker, (

690 "driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1"

vllm/executor/ray_utils.py

+3 -3

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

4 import os

5 import time

6 from collections import defaultdict

7 -from typing import TYPE_CHECKING, Optional, Union

7 +from typing import TYPE_CHECKING, Union

8

9 import msgspec

10

@@ -72,7 +72,7 @@ try:

72

73 def execute_model_spmd(

74 self,

75 - req_or_tuple: Union[bytes, tuple[bytes, Optional[IntermediateTensors]]],

75 + req_or_tuple: bytes | tuple[bytes, IntermediateTensors | None],

76 ) -> bytes:

77 """Execute model in SPMD fashion: used only when SPMD worker and

78 compiled DAG are both enabled.

@@ -300,7 +300,7 @@ def _wait_until_pg_removed(current_placement_group: "PlacementGroup"):

300

301 def initialize_ray_cluster(

302 parallel_config: ParallelConfig,

303 - ray_address: Optional[str] = None,

303 + ray_address: str | None = None,

304 ):

305 """Initialize the distributed cluster with Ray.

306

vllm/executor/uniproc_executor.py

+6 -5

@@ -1,10 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 import os

4 +from collections.abc import Callable

5 from concurrent.futures import Future, ThreadPoolExecutor

6 from functools import cached_property

7 from multiprocessing import Lock

7 -from typing import Any, Callable, Optional, Union

8 +from typing import Any

9

10 import torch

11 import torch.distributed as dist

@@ -36,7 +37,7 @@ class UniProcExecutor(ExecutorBase):

37 shared_worker_lock=Lock(),

38 )

39

39 - self.async_output_thread: Optional[ThreadPoolExecutor] = None

40 + self.async_output_thread: ThreadPoolExecutor | None = None

41 if self.max_concurrent_batches > 1:

42 self.async_output_thread = ThreadPoolExecutor(

43 max_workers=1, thread_name_prefix="WorkerAsyncOutput"

@@ -60,10 +61,10 @@ class UniProcExecutor(ExecutorBase):

61

62 def collective_rpc(

63 self,

63 - method: Union[str, Callable],

64 - timeout: Optional[float] = None,

64 + method: str | Callable,

65 + timeout: float | None = None,

66 args: tuple = (),

66 - kwargs: Optional[dict] = None,

67 + kwargs: dict | None = None,

68 non_block: bool = False,

69 ) -> list[Any]:

70 if kwargs is None:

vllm/forward_context.py

+16 -16

@@ -5,7 +5,7 @@ import time

5 from collections import defaultdict

6 from contextlib import contextmanager

7 from dataclasses import dataclass

8 -from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union

8 +from typing import TYPE_CHECKING, Any, NamedTuple, Union

9

10 import torch

11

@@ -84,7 +84,7 @@ class DPMetadata:

84 num_tokens_across_dp_cpu: torch.Tensor

85

86 # NOTE: local_sizes should only be set by the chunked_sizes context manager

87 - local_sizes: Optional[list[int]] = None

87 + local_sizes: list[int] | None = None

88

89 @staticmethod

90 def make(

@@ -158,7 +158,7 @@ class DPMetadata:

158 finally:

159 self.local_sizes = None

160

161 - def get_chunk_sizes_across_dp_rank(self) -> Optional[list[int]]:

161 + def get_chunk_sizes_across_dp_rank(self) -> list[int] | None:

162 assert self.local_sizes is not None

163 return self.local_sizes

164

@@ -194,13 +194,13 @@ class ForwardContext:

194 # TODO: remove after making all virtual_engines share the same kv cache

195 virtual_engine: int # set dynamically for each forward pass

196 # set dynamically for each forward pass

197 - dp_metadata: Optional[DPMetadata] = None

197 + dp_metadata: DPMetadata | None = None

198 # determine the cudagraph style at runtime to be FULL, PIECEWISE, or NONE.

199 # by default NONE, no cudagraph is used.

200 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE

201 - batch_descriptor: Optional[BatchDescriptor] = None

201 + batch_descriptor: BatchDescriptor | None = None

202

203 - ubatch_slices: Optional[UBatchSlices] = None

203 + ubatch_slices: UBatchSlices | None = None

204

205 def __post_init__(self):

206 assert self.cudagraph_runtime_mode.valid_runtime_modes(), (

@@ -208,7 +208,7 @@ class ForwardContext:

208 )

209

210

211 -_forward_context: Optional[ForwardContext] = None

211 +_forward_context: ForwardContext | None = None

212

213

214 def get_forward_context() -> ForwardContext:

@@ -224,10 +224,10 @@ def create_forward_context(

224 attn_metadata: Any,

225 vllm_config: VllmConfig,

226 virtual_engine: int = 0,

227 - dp_metadata: Optional[DPMetadata] = None,

227 + dp_metadata: DPMetadata | None = None,

228 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

229 - batch_descriptor: Optional[BatchDescriptor] = None,

230 - ubatch_slices: Optional[UBatchSlices] = None,

229 + batch_descriptor: BatchDescriptor | None = None,

230 + ubatch_slices: UBatchSlices | None = None,

231 ):

232 return ForwardContext(

233 no_compile_layers=vllm_config.compilation_config.static_forward_context,

@@ -241,7 +241,7 @@ def create_forward_context(

241

242

243 @contextmanager

244 -def override_forward_context(forward_context: Optional[ForwardContext]):

244 +def override_forward_context(forward_context: ForwardContext | None):

245 """A context manager that overrides the current forward context.

246 This is used to override the forward context for a specific

247 forward pass.

@@ -260,11 +260,11 @@ def set_forward_context(

260 attn_metadata: Any,

261 vllm_config: VllmConfig,

262 virtual_engine: int = 0,

263 - num_tokens: Optional[int] = None,

264 - num_tokens_across_dp: Optional[torch.Tensor] = None,

263 + num_tokens: int | None = None,

264 + num_tokens_across_dp: torch.Tensor | None = None,

265 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

266 - batch_descriptor: Optional[BatchDescriptor] = None,

267 - ubatch_slices: Optional[UBatchSlices] = None,

266 + batch_descriptor: BatchDescriptor | None = None,

267 + ubatch_slices: UBatchSlices | None = None,

268 ):

269 """A context manager that stores the current forward context,

270 can be attention metadata, etc.

@@ -275,7 +275,7 @@ def set_forward_context(

275 if need_to_track_batchsize:

276 forward_start_time = time.perf_counter()

277

278 - dp_metadata: Optional[DPMetadata] = None

278 + dp_metadata: DPMetadata | None = None

279 if vllm_config.parallel_config.data_parallel_size > 1 and (

280 attn_metadata is not None or num_tokens is not None

281 ):

vllm/inputs/data.py

+26 -24

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from collections.abc import Iterable

4 -from typing import TYPE_CHECKING, Any, Generic, Literal, Optional, Union, cast

4 +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, cast

5

6 import torch

7 from typing_extensions import NotRequired, TypedDict, TypeIs, TypeVar

@@ -12,6 +12,10 @@ if TYPE_CHECKING:

12 MultiModalInputs,

13 MultiModalUUIDDict,

14 )

15 +else:

16 + MultiModalDataDict = object

17 + MultiModalInputs = object

18 + MultiModalUUIDDict = object

19

20

21 class TextPrompt(TypedDict):

@@ -20,13 +24,13 @@ class TextPrompt(TypedDict):

24 prompt: str

25 """The input text to be tokenized before passing to the model."""

26

23 - multi_modal_data: NotRequired[Optional["MultiModalDataDict"]]

27 + multi_modal_data: NotRequired[MultiModalDataDict | None]

28 """

29 Optional multi-modal data to pass to the model,

30 if the model supports it.

31 """

32

29 - mm_processor_kwargs: NotRequired[Optional[dict[str, Any]]]

33 + mm_processor_kwargs: NotRequired[dict[str, Any] | None]

34 """

35 Optional multi-modal processor kwargs to be forwarded to the

36 multimodal input mapper & processor. Note that if multiple modalities

@@ -34,7 +38,7 @@ class TextPrompt(TypedDict):

38 to pass the mm_processor_kwargs to each of them.

39 """

40

37 - multi_modal_uuids: NotRequired["MultiModalUUIDDict"]

41 + multi_modal_uuids: NotRequired[MultiModalUUIDDict]

42 """

43 Optional user-specified UUIDs for multimodal items, mapped by modality.

44 Lists must match the number of items per modality and may contain `None`.

@@ -61,13 +65,13 @@ class TokensPrompt(TypedDict):

65 token_type_ids: NotRequired[list[int]]

66 """A list of token type IDs to pass to the cross encoder model."""

67

64 - multi_modal_data: NotRequired[Optional["MultiModalDataDict"]]

68 + multi_modal_data: NotRequired[MultiModalDataDict | None]

69 """

70 Optional multi-modal data to pass to the model,

71 if the model supports it.

72 """

73

70 - mm_processor_kwargs: NotRequired[Optional[dict[str, Any]]]

74 + mm_processor_kwargs: NotRequired[dict[str, Any] | None]

75 """

76 Optional multi-modal processor kwargs to be forwarded to the

77 multimodal input mapper & processor. Note that if multiple modalities

@@ -75,7 +79,7 @@ class TokensPrompt(TypedDict):

79 to pass the mm_processor_kwargs to each of them.

80 """

81

78 - multi_modal_uuids: NotRequired["MultiModalUUIDDict"]

82 + multi_modal_uuids: NotRequired[MultiModalUUIDDict]

83 """

84 Optional user-specified UUIDs for multimodal items, mapped by modality.

85 Lists must match the number of items per modality and may contain `None`.

@@ -111,7 +115,7 @@ class DataPrompt(TypedDict):

115 """The input data format"""

116

117

114 -SingletonPrompt = Union[str, TextPrompt, TokensPrompt, EmbedsPrompt]

118 +SingletonPrompt: TypeAlias = str | TextPrompt | TokensPrompt | EmbedsPrompt

119 """

120 Set of possible schemas for a single prompt:

121

@@ -185,12 +189,12 @@ class ExplicitEncoderDecoderPrompt(TypedDict, Generic[_T1_co, _T2_co]):

189

190 encoder_prompt: _T1_co

191

188 - decoder_prompt: Optional[_T2_co]

192 + decoder_prompt: _T2_co | None

193

194 mm_processor_kwargs: NotRequired[dict[str, Any]]

195

196

193 -PromptType = Union[SingletonPrompt, ExplicitEncoderDecoderPrompt]

197 +PromptType: TypeAlias = SingletonPrompt | ExplicitEncoderDecoderPrompt

198 """

199 Set of possible schemas for an LLM input, including

200 both decoder-only and encoder/decoder input types:

@@ -220,7 +224,7 @@ class TokenInputs(TypedDict):

224

225 def token_inputs(

226 prompt_token_ids: list[int],

223 - cache_salt: Optional[str] = None,

227 + cache_salt: str | None = None,

228 ) -> TokenInputs:

229 """Construct [`TokenInputs`][vllm.inputs.data.TokenInputs] from optional

230 values."""

@@ -249,7 +253,7 @@ class EmbedsInputs(TypedDict):

253

254 def embeds_inputs(

255 prompt_embeds: torch.Tensor,

252 - cache_salt: Optional[str] = None,

256 + cache_salt: str | None = None,

257 ) -> EmbedsInputs:

258 """Construct [`EmbedsInputs`][vllm.inputs.data.EmbedsInputs] from optional

259 values."""

@@ -261,7 +265,7 @@ def embeds_inputs(

265 return inputs

266

267

264 -DecoderOnlyInputs = Union[TokenInputs, EmbedsInputs, "MultiModalInputs"]

268 +DecoderOnlyInputs: TypeAlias = TokenInputs | EmbedsInputs | MultiModalInputs

269 """

270 The inputs in [`LLMEngine`][vllm.engine.llm_engine.LLMEngine] before they are

271 passed to the model executor.

@@ -277,20 +281,20 @@ class EncoderDecoderInputs(TypedDict):

281 This specifies the required data for encoder-decoder models.

282 """

283

280 - encoder: Union[TokenInputs, "MultiModalInputs"]

284 + encoder: TokenInputs | MultiModalInputs

285 """The inputs for the encoder portion."""

286

283 - decoder: Union[TokenInputs, "MultiModalInputs"]

287 + decoder: TokenInputs | MultiModalInputs

288 """The inputs for the decoder portion."""

289

290

287 -SingletonInputs = Union[TokenInputs, EmbedsInputs, "MultiModalInputs"]

291 +SingletonInputs: TypeAlias = TokenInputs | EmbedsInputs | MultiModalInputs

292 """

293 A processed [`SingletonPrompt`][vllm.inputs.data.SingletonPrompt] which can be

294 passed to [`Sequence`][collections.abc.Sequence].

295 """

296

293 -ProcessorInputs = Union[DecoderOnlyInputs, EncoderDecoderInputs]

297 +ProcessorInputs: TypeAlias = DecoderOnlyInputs | EncoderDecoderInputs

298 """

299 The outputs from [`vllm.inputs.preprocess.InputPreprocessor`][].

300 """

@@ -301,8 +305,8 @@ _T2 = TypeVar("_T2", bound=SingletonPrompt, default=SingletonPrompt)

305

306 def build_explicit_enc_dec_prompt(

307 encoder_prompt: _T1,

304 - decoder_prompt: Optional[_T2],

305 - mm_processor_kwargs: Optional[dict[str, Any]] = None,

308 + decoder_prompt: _T2 | None,

309 + mm_processor_kwargs: dict[str, Any] | None = None,

310 ) -> ExplicitEncoderDecoderPrompt[_T1, _T2]:

311 if mm_processor_kwargs is None:

312 mm_processor_kwargs = {}

@@ -315,10 +319,8 @@ def build_explicit_enc_dec_prompt(

319

320 def zip_enc_dec_prompts(

321 enc_prompts: Iterable[_T1],

318 - dec_prompts: Iterable[Optional[_T2]],

319 - mm_processor_kwargs: Optional[

320 - Union[Iterable[dict[str, Any]], dict[str, Any]]

321 - ] = None,

322 + dec_prompts: Iterable[_T2 | None],

323 + mm_processor_kwargs: Iterable[dict[str, Any]] | dict[str, Any] | None = None,

324 ) -> list[ExplicitEncoderDecoderPrompt[_T1, _T2]]:

325 """

326 Zip encoder and decoder prompts together into a list of

@@ -350,7 +352,7 @@ def zip_enc_dec_prompts(

352

353 def to_enc_dec_tuple_list(

354 enc_dec_prompts: Iterable[ExplicitEncoderDecoderPrompt[_T1, _T2]],

353 -) -> list[tuple[_T1, Optional[_T2]]]:

355 +) -> list[tuple[_T1, _T2 | None]]:

356 return [

357 (enc_dec_prompt["encoder_prompt"], enc_dec_prompt["decoder_prompt"])

358 for enc_dec_prompt in enc_dec_prompts

vllm/inputs/parse.py

+10 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from collections.abc import Sequence

4 -from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, TypedDict, Union, cast

4 +from typing import TYPE_CHECKING, Literal, NamedTuple, TypeAlias, TypedDict, cast

5

6 from typing_extensions import TypeIs

7

@@ -23,8 +23,8 @@ if TYPE_CHECKING:

23

24

25 def parse_raw_prompts(

26 - prompt: Union[str, list[str], list[int], list[list[int]]],

27 -) -> Union[Sequence[TextPrompt], Sequence[TokensPrompt]]:

26 + prompt: str | list[str] | list[int] | list[list[int]],

27 +) -> Sequence[TextPrompt] | Sequence[TokensPrompt]:

28 if isinstance(prompt, str):

29 # case 1: a string

30 return [TextPrompt(prompt=prompt)]

@@ -76,9 +76,9 @@ class ParsedEmbedsPrompt(TypedDict):

76 content: EmbedsPrompt

77

78

79 -ParsedSingletonPrompt = Union[

80 - ParsedStrPrompt, ParsedTextPrompt, ParsedTokensPrompt, ParsedEmbedsPrompt

81 -]

79 +ParsedSingletonPrompt: TypeAlias = (

80 + ParsedStrPrompt | ParsedTextPrompt | ParsedTokensPrompt | ParsedEmbedsPrompt

81 +)

82

83

84 def parse_singleton_prompt(prompt: SingletonPrompt) -> ParsedSingletonPrompt:

@@ -106,7 +106,7 @@ def is_explicit_encoder_decoder_prompt(

106

107 def split_enc_dec_inputs(

108 inputs: ProcessorInputs,

109 -) -> tuple[Optional[SingletonInputs], SingletonInputs]:

109 +) -> tuple[SingletonInputs | None, SingletonInputs]:

110 if "encoder" in inputs and "decoder" in inputs:

111 # NOTE: This passes pyright but not mypy

112 return (

@@ -118,9 +118,9 @@ def split_enc_dec_inputs(

118

119

120 class PromptComponents(NamedTuple):

121 - text: Optional[str] = None

122 - token_ids: Optional[list[int]] = None

123 - embeds: Optional["torch.Tensor"] = None

121 + text: str | None = None

122 + token_ids: list[int] | None = None

123 + embeds: "torch.Tensor | None" = None

124

125

126 def get_prompt_components(prompt: PromptType) -> PromptComponents:

vllm/inputs/preprocess.py

+42 -44

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

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

3

4 from collections.abc import Mapping

5 -from typing import Any, Optional, Union, cast

5 +from typing import Any, cast

6

7 from typing_extensions import assert_never

8

@@ -46,9 +46,9 @@ class InputPreprocessor:

46 def __init__(

47 self,

48 model_config: ModelConfig,

49 - tokenizer: Optional[AnyTokenizer],

49 + tokenizer: AnyTokenizer | None,

50 mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,

51 - mm_processor_cache: Optional[BaseMultiModalProcessorCache] = None,

51 + mm_processor_cache: BaseMultiModalProcessorCache | None = None,

52 ) -> None:

53 super().__init__()

54

@@ -67,7 +67,7 @@ class InputPreprocessor:

67

68 return self.tokenizer

69

70 - def get_bos_token_id(self) -> Optional[int]:

70 + def get_bos_token_id(self) -> int | None:

71 if self.tokenizer is None:

72 logger.warning_once(

73 "Using None for BOS token id because tokenizer is not initialized"

@@ -76,7 +76,7 @@ class InputPreprocessor:

76

77 return self.tokenizer.bos_token_id

78

79 - def get_eos_token_id(self) -> Optional[int]:

79 + def get_eos_token_id(self) -> int | None:

80 if self.tokenizer is None:

81 logger.warning_once(

82 "Using None for EOS token id because tokenizer is not initialized"

@@ -85,7 +85,7 @@ class InputPreprocessor:

85

86 return self.tokenizer.eos_token_id

87

88 - def get_decoder_start_token_id(self) -> Optional[int]:

88 + def get_decoder_start_token_id(self) -> int | None:

89 """

90 Obtain the decoder start token id employed by an encoder/decoder

91 model. Returns None for non-encoder/decoder models or if the

@@ -157,7 +157,7 @@ class InputPreprocessor:

157

158 def _prepare_decoder_input_ids_for_generation(

159 self,

160 - decoder_input_ids: Optional[list[int]],

160 + decoder_input_ids: list[int] | None,

161 ) -> list[int]:

162 """

163 Prepares `decoder_input_ids` for generation with encoder-decoder models.

@@ -194,7 +194,7 @@ class InputPreprocessor:

194

195 def _get_tokenization_kw(

196 self,

197 - overrides: Optional[dict[str, Any]] = None,

197 + overrides: dict[str, Any] | None = None,

198 ) -> dict[str, Any]:

199 kwargs = dict[str, Any]()

200

@@ -212,7 +212,7 @@ class InputPreprocessor:

212 def _tokenize_prompt(

213 self,

214 prompt: str,

215 - tokenization_kwargs: Optional[dict[str, Any]] = None,

215 + tokenization_kwargs: dict[str, Any] | None = None,

216 ) -> list[int]:

217 """

218 Apply the model's tokenizer to a text prompt, returning the

@@ -251,12 +251,12 @@ class InputPreprocessor:

251

252 def _process_multimodal(

253 self,

254 - prompt: Union[str, list[int]],

254 + prompt: str | list[int],

255 mm_data: MultiModalDataDict,

256 - mm_processor_kwargs: Optional[Mapping[str, object]],

257 - tokenization_kwargs: Optional[dict[str, Any]] = None,

256 + mm_processor_kwargs: Mapping[str, object] | None,

257 + tokenization_kwargs: dict[str, Any] | None = None,

258 *,

259 - mm_uuids: Optional[MultiModalUUIDDict] = None,

259 + mm_uuids: MultiModalUUIDDict | None = None,

260 ) -> MultiModalInputs:

261 """

262 Apply the model's multi-modal processor to a multi-modal prompt,

@@ -320,7 +320,7 @@ class InputPreprocessor:

320 )

321

322 def _truncate_inputs(

323 - self, inputs: list[int], tokenization_kwargs: Optional[dict[str, Any]] = None

323 + self, inputs: list[int], tokenization_kwargs: dict[str, Any] | None = None

324 ) -> list[int]:

325 if (

326 not tokenization_kwargs

@@ -339,15 +339,15 @@ class InputPreprocessor:

339 def _process_tokens(

340 self,

341 parsed_content: TokensPrompt,

342 - tokenization_kwargs: Optional[dict[str, Any]] = None,

342 + tokenization_kwargs: dict[str, Any] | None = None,

343 *,

344 - mm_uuids: Optional[MultiModalUUIDDict] = None,

345 - ) -> Union[TokenInputs, MultiModalInputs]:

344 + mm_uuids: MultiModalUUIDDict | None = None,

345 + ) -> TokenInputs | MultiModalInputs:

346 prompt_token_ids = self._truncate_inputs(

347 parsed_content["prompt_token_ids"], tokenization_kwargs

348 )

349

350 - inputs: Union[TokenInputs, MultiModalInputs]

350 + inputs: TokenInputs | MultiModalInputs

351 if self.model_config.is_multimodal_model:

352 inputs = self._process_multimodal(

353 prompt_token_ids,

@@ -370,13 +370,13 @@ class InputPreprocessor:

370 def _process_text(

371 self,

372 parsed_content: TextPrompt,

373 - tokenization_kwargs: Optional[dict[str, Any]] = None,

373 + tokenization_kwargs: dict[str, Any] | None = None,

374 *,

375 - mm_uuids: Optional[MultiModalUUIDDict] = None,

376 - ) -> Union[TokenInputs, MultiModalInputs]:

375 + mm_uuids: MultiModalUUIDDict | None = None,

376 + ) -> TokenInputs | MultiModalInputs:

377 prompt_text = parsed_content["prompt"]

378

379 - inputs: Union[TokenInputs, MultiModalInputs]

379 + inputs: TokenInputs | MultiModalInputs

380 if self.model_config.is_multimodal_model:

381 inputs = self._process_multimodal(

382 prompt_text,

@@ -403,9 +403,9 @@ class InputPreprocessor:

403 def _prompt_to_llm_inputs(

404 self,

405 prompt: SingletonPrompt,

406 - tokenization_kwargs: Optional[dict[str, Any]] = None,

406 + tokenization_kwargs: dict[str, Any] | None = None,

407 *,

408 - mm_uuids: Optional[MultiModalUUIDDict] = None,

408 + mm_uuids: MultiModalUUIDDict | None = None,

409 ) -> SingletonInputs:

410 """

411 Extract the singleton inputs from a prompt.

@@ -445,7 +445,7 @@ class InputPreprocessor:

445 def _build_enc_dec_llm_inputs(

446 self,

447 encoder_inputs: SingletonInputs,

448 - decoder_inputs: Optional[SingletonInputs],

448 + decoder_inputs: SingletonInputs | None,

449 ) -> EncoderDecoderInputs:

450 if (

451 encoder_inputs["type"] == "embeds"

@@ -457,10 +457,8 @@ class InputPreprocessor:

457 )

458

459 # Needed for mypy

460 - encoder_inputs = cast(Union[TokenInputs, MultiModalInputs], encoder_inputs)

461 - decoder_inputs = cast(

462 - Optional[Union[TokenInputs, MultiModalInputs]], decoder_inputs

463 - )

460 + encoder_inputs = cast(TokenInputs | MultiModalInputs, encoder_inputs)

461 + decoder_inputs = cast(TokenInputs | MultiModalInputs | None, decoder_inputs)

462

463 if decoder_inputs is None:

464 if self.model_config.hf_config.model_type == "whisper":

@@ -491,8 +489,8 @@ class InputPreprocessor:

489

490 def _split_enc_dec_mm_inputs(

491 self,

494 - inputs: Union[SingletonInputs, MultiModalEncDecInputs],

495 - decoder_inputs_to_override: Optional[SingletonInputs] = None,

492 + inputs: SingletonInputs | MultiModalEncDecInputs,

493 + decoder_inputs_to_override: SingletonInputs | None = None,

494 ) -> tuple[SingletonInputs, SingletonInputs]:

495 """

496 For encoder/decoder models only:

@@ -509,11 +507,11 @@ class InputPreprocessor:

507

508 # Needed for mypy

509 inputs = cast(

512 - Union[TokenInputs, MultiModalInputs, MultiModalEncDecInputs],

510 + TokenInputs | MultiModalInputs | MultiModalEncDecInputs,

511 inputs,

512 )

513 decoder_inputs_to_override = cast(

516 - Optional[Union[TokenInputs, MultiModalInputs]],

514 + TokenInputs | MultiModalInputs | None,

515 decoder_inputs_to_override,

516 )

517

@@ -553,9 +551,9 @@ class InputPreprocessor:

551 def _process_encoder_decoder_prompt(

552 self,

553 prompt: PromptType,

556 - tokenization_kwargs: Optional[dict[str, Any]] = None,

554 + tokenization_kwargs: dict[str, Any] | None = None,

555 *,

558 - mm_uuids: Optional[MultiModalUUIDDict] = None,

556 + mm_uuids: MultiModalUUIDDict | None = None,

557 ) -> EncoderDecoderInputs:

558 """

559 For encoder/decoder models only:

@@ -591,7 +589,7 @@ class InputPreprocessor:

589 instance

590 """

591 encoder_inputs: SingletonInputs

594 - decoder_inputs: Optional[SingletonInputs]

592 + decoder_inputs: SingletonInputs | None

593

594 if is_explicit_encoder_decoder_prompt(prompt):

595 # `cast` is needed for mypy, but not pyright

@@ -633,7 +631,7 @@ class InputPreprocessor:

631 ) -> DecoderOnlyInputs:

632 if "prompt_token_ids" in prompt_inputs:

633 prompt_inputs = cast(

636 - Union[TokenInputs, MultiModalInputs], prompt_inputs

634 + TokenInputs | MultiModalInputs, prompt_inputs

635 ) # Needed for mypy

636

637 return prompt_inputs

@@ -641,9 +639,9 @@ class InputPreprocessor:

639 def _process_decoder_only_prompt(

640 self,

641 prompt: SingletonPrompt,

644 - tokenization_kwargs: Optional[dict[str, Any]] = None,

642 + tokenization_kwargs: dict[str, Any] | None = None,

643 *,

646 - mm_uuids: Optional[MultiModalUUIDDict] = None,

644 + mm_uuids: MultiModalUUIDDict | None = None,

645 ) -> DecoderOnlyInputs:

646 """

647 For decoder-only models:

@@ -670,9 +668,9 @@ class InputPreprocessor:

668 def _preprocess(

669 self,

670 prompt: PromptType,

673 - tokenization_kwargs: Optional[dict[str, Any]] = None,

671 + tokenization_kwargs: dict[str, Any] | None = None,

672 *,

675 - mm_uuids: Optional[MultiModalUUIDDict] = None,

673 + mm_uuids: MultiModalUUIDDict | None = None,

674 ) -> ProcessorInputs:

675 if self.model_config.is_encoder_decoder:

676 # Encoder-decoder model requires special mapping of

@@ -699,9 +697,9 @@ class InputPreprocessor:

697 def preprocess(

698 self,

699 prompt: PromptType,

702 - tokenization_kwargs: Optional[dict[str, Any]] = None,

700 + tokenization_kwargs: dict[str, Any] | None = None,

701 *,

704 - mm_uuids: Optional[MultiModalUUIDDict] = None,

702 + mm_uuids: MultiModalUUIDDict | None = None,

703 ) -> ProcessorInputs:

704 """Preprocess the input prompt."""

705 res = self._preprocess(

@@ -718,7 +716,7 @@ class InputPreprocessor:

716

717 return res

718

721 - def stat_mm_cache(self) -> Optional[MultiModalCacheStats]:

719 + def stat_mm_cache(self) -> MultiModalCacheStats | None:

720 mm_cache_stats = self.mm_cache_stats

721 if mm_cache_stats is None:

722 return None

vllm/logger.py

+2 -2

@@ -13,7 +13,7 @@ from logging import Logger

13 from logging.config import dictConfig

14 from os import path

15 from types import MethodType

16 -from typing import Any, Optional, cast

16 +from typing import Any, cast

17

18 import vllm.envs as envs

19

@@ -217,7 +217,7 @@ def _trace_calls(log_path, root_dir, frame, event, arg=None):

217 return partial(_trace_calls, log_path, root_dir)

218

219

220 -def enable_trace_function_call(log_file_path: str, root_dir: Optional[str] = None):

220 +def enable_trace_function_call(log_file_path: str, root_dir: str | None = None):

221 """

222 Enable tracing of every function call in code under `root_dir`.

223 This is useful for debugging hangs or crashes.

vllm/logging_utils/dump_input.py

+2 -3

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

4 import contextlib

5 import enum

6 import json

7 -from typing import Optional

7

8 import torch

9

@@ -57,7 +56,7 @@ def prepare_object_to_dump(obj) -> str:

56 def dump_engine_exception(

57 config: VllmConfig,

58 scheduler_output: SchedulerOutput,

60 - scheduler_stats: Optional[SchedulerStats],

59 + scheduler_stats: SchedulerStats | None,

60 ):

61 # NOTE: ensure we can log extra info without risking raises

62 # unexpected errors during logging

@@ -68,7 +67,7 @@ def dump_engine_exception(

67 def _dump_engine_exception(

68 config: VllmConfig,

69 scheduler_output: SchedulerOutput,

71 - scheduler_stats: Optional[SchedulerStats],

70 + scheduler_stats: SchedulerStats | None,

71 ):

72 logger.error(

73 "Dumping input data for V1 LLM engine (v%s) with config: %s, ",

vllm/logits_process.py

+6 -6

@@ -1,16 +1,16 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from collections.abc import Sequence

4 -from typing import Callable, Union

3 +from collections.abc import Callable, Sequence

4 +from typing import TypeAlias

5

6 import torch

7

8 from vllm.transformers_utils.tokenizer import AnyTokenizer

9

10 -LogitsProcessor = Union[

11 - Callable[[list[int], torch.Tensor], torch.Tensor],

12 - Callable[[list[int], list[int], torch.Tensor], torch.Tensor],

13 -]

10 +LogitsProcessor: TypeAlias = (

11 + Callable[[list[int], torch.Tensor], torch.Tensor]

12 + | Callable[[list[int], list[int], torch.Tensor], torch.Tensor]

13 +)

14 """LogitsProcessor is a function that takes a list

15 of previously generated tokens, the logits tensor

16 for the next token and, optionally, prompt tokens as a

vllm/logprobs.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass

4 -from typing import Optional

4

5

6 # We use dataclass for now because it is used for

@@ -18,12 +17,12 @@ class Logprob:

17 """

18

19 logprob: float

21 - rank: Optional[int] = None

22 - decoded_token: Optional[str] = None

20 + rank: int | None = None

21 + decoded_token: str | None = None

22

23

24 # {token_id -> logprob} per each sequence group. None if the corresponding

25 # sequence group doesn't require prompt logprob.

27 -PromptLogprobs = list[Optional[dict[int, Logprob]]]

26 +PromptLogprobs = list[dict[int, Logprob] | None]

27 # {token_id -> logprob} for each sequence group.

28 SampleLogprobs = list[dict[int, Logprob]]

vllm/lora/layers/base.py

+8 -8

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import TYPE_CHECKING, Optional, Union

4 +from typing import TYPE_CHECKING

5

6 import torch

7 import torch.nn as nn

@@ -15,14 +15,14 @@ if TYPE_CHECKING:

15

16 class BaseLayerWithLoRA(nn.Module):

17 def slice_lora_a(

18 - self, lora_a: Union[torch.Tensor, list[Union[torch.Tensor, None]]]

19 - ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:

18 + self, lora_a: torch.Tensor | list[torch.Tensor | None]

19 + ) -> torch.Tensor | list[torch.Tensor | None]:

20 """Slice lora a if splitting for tensor parallelism."""

21 ...

22

23 def slice_lora_b(

24 - self, lora_b: Union[torch.Tensor, list[Union[torch.Tensor, None]]]

25 - ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:

24 + self, lora_b: torch.Tensor | list[torch.Tensor | None]

25 + ) -> torch.Tensor | list[torch.Tensor | None]:

26 """Slice lora b if splitting with tensor parallelism."""

27 ...

28

@@ -30,7 +30,7 @@ class BaseLayerWithLoRA(nn.Module):

30 self,

31 max_loras: int,

32 lora_config: LoRAConfig,

33 - model_config: Optional[PretrainedConfig] = None,

33 + model_config: PretrainedConfig | None = None,

34 ) -> None:

35 """Initializes lora matrices."""

36 ...

@@ -44,7 +44,7 @@ class BaseLayerWithLoRA(nn.Module):

44 index: int,

45 lora_a: torch.Tensor,

46 lora_b: torch.Tensor,

47 - embeddings_tensor: Optional[torch.Tensor],

47 + embeddings_tensor: torch.Tensor | None,

48 ):

49 """Overwrites lora tensors at index."""

50 ...

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

61 source_layer: nn.Module,

62 lora_config: LoRAConfig,

63 packed_modules_list: list,

64 - model_config: Optional[PretrainedConfig],

64 + model_config: PretrainedConfig | None,

65 ) -> bool:

66 """Returns True if the layer can be replaced by this LoRA layer."""

67 raise NotImplementedError

vllm/lora/layers/base_linear.py

+5 -8

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6 from transformers import PretrainedConfig

@@ -37,7 +36,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

36 self,

37 max_loras: int,

38 lora_config: LoRAConfig,

40 - model_config: Optional[PretrainedConfig] = None,

39 + model_config: PretrainedConfig | None = None,

40 ) -> None:

41 self.lora_config = lora_config

42 #

@@ -97,7 +96,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

96 index: int,

97 lora_a: torch.Tensor,

98 lora_b: torch.Tensor,

100 - embeddings_tensor: Optional[torch.Tensor],

99 + embeddings_tensor: torch.Tensor | None,

100 ):

101 # Except for QKVParallelLinearWithLoRA and

102 # MergedColumnParallelLinearWithLoRA, all other linear LoRA layers

@@ -119,9 +118,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

118 lora_b, non_blocking=True

119 )

120

122 - def apply(

123 - self, x: torch.Tensor, bias: Optional[torch.Tensor] = None

124 - ) -> torch.Tensor:

121 + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

122 output = self.base_layer.quant_method.apply(self.base_layer, x, bias)

123

124 # In transformers backend, x and output have extra batch dimension like

@@ -131,7 +128,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

128 output = output.flatten(0, 1)

129 x = x.flatten(0, 1)

130

134 - lora_output: Optional[torch.Tensor] = self.punica_wrapper.add_lora_linear(

131 + lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_linear(

132 output, x, self.lora_a_stacked, self.lora_b_stacked, 1.0, self.output_slices

133 )

134 if not current_platform.can_update_inplace():

@@ -160,7 +157,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

157 raise ValueError(f"Unsupported base layer: {self.base_layer}")

158

159 @property

163 - def bias(self) -> Optional[torch.Tensor]:

160 + def bias(self) -> torch.Tensor | None:

161 if hasattr(self.base_layer, "bias"):

162 return self.base_layer.bias

163 else:

vllm/lora/layers/column_parallel_linear.py

+27 -36

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 import torch

6 import torch.nn as nn

@@ -46,7 +45,7 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):

45 device=x.device,

46 )

47

49 - shrunk_buffers: Optional[torch.Tensor] = layer.punica_wrapper.add_shrink(

48 + shrunk_buffers: torch.Tensor | None = layer.punica_wrapper.add_shrink(

49 buffers, x, layer.lora_a_stacked, 1.0

50 )

51

@@ -55,7 +54,7 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):

54

55 buffers = tensor_model_parallel_all_gather(buffers)

56

58 - lora_output: Optional[torch.Tensor] = layer.punica_wrapper.add_expand(

57 + lora_output: torch.Tensor | None = layer.punica_wrapper.add_expand(

58 output,

59 buffers,

60 layer.lora_b_stacked,

@@ -121,7 +120,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

120

121 def forward(

122 self, input_: torch.Tensor

124 - ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:

123 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:

124 """Forward of ColumnParallelLinear

125

126 Args:

@@ -154,7 +153,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

153 source_layer: nn.Module,

154 lora_config: LoRAConfig,

155 packed_modules_list: list,

157 - model_config: Optional[PretrainedConfig],

156 + model_config: PretrainedConfig | None,

157 ) -> bool:

158 return type(source_layer) is ColumnParallelLinear or (

159 type(source_layer) is MergedColumnParallelLinear

@@ -172,7 +171,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

171 """

172

173 def __init__(

175 - self, base_layer: Union[MergedColumnParallelLinear, QKVParallelLinear]

174 + self, base_layer: MergedColumnParallelLinear | QKVParallelLinear

175 ) -> None:

176 super().__init__(base_layer)

177 # There are two LoRA layers

@@ -189,7 +188,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

188 self,

189 max_loras: int,

190 lora_config: LoRAConfig,

192 - model_config: Optional[PretrainedConfig] = None,

191 + model_config: PretrainedConfig | None = None,

192 ) -> None:

193 """

194 The main reason for overriding this function is to enhance code

@@ -227,13 +226,13 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

226 )

227

228 def slice_lora_a(

230 - self, lora_a: list[Union[torch.Tensor, None]]

231 - ) -> list[Union[torch.Tensor, None]]:

229 + self, lora_a: list[torch.Tensor | None]

230 + ) -> list[torch.Tensor | None]:

231 return lora_a

232

233 def slice_lora_b(

235 - self, lora_b: list[Union[torch.Tensor, None]]

236 - ) -> list[Union[torch.Tensor, None]]:

234 + self, lora_b: list[torch.Tensor | None]

235 + ) -> list[torch.Tensor | None]:

236 sliced_lora_b = [None] * self.n_slices

237 for i, (shard_id, shard_size) in enumerate(

238 zip(self.output_ids, self.output_slices)

@@ -249,7 +248,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

248 index: int,

249 lora_a: torch.Tensor,

250 lora_b: torch.Tensor,

252 - embeddings_tensor: Optional[torch.Tensor],

251 + embeddings_tensor: torch.Tensor | None,

252 ):

253 self.reset_lora(index)

254

@@ -274,7 +273,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

273 source_layer: nn.Module,

274 lora_config: LoRAConfig,

275 packed_modules_list: list,

277 - model_config: Optional[PretrainedConfig],

276 + model_config: PretrainedConfig | None,

277 ) -> bool:

278 return (

279 type(source_layer) is MergedColumnParallelLinear

@@ -340,7 +339,7 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

339 source_layer: nn.Module,

340 lora_config: LoRAConfig,

341 packed_modules_list: list,

343 - model_config: Optional[PretrainedConfig],

342 + model_config: PretrainedConfig | None,

343 ) -> bool:

344 return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 1

345

@@ -383,7 +382,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):

382 self,

383 max_loras: int,

384 lora_config: LoRAConfig,

386 - model_config: Optional[PretrainedConfig] = None,

385 + model_config: PretrainedConfig | None = None,

386 ) -> None:

387 """

388 The main reason for overloading this function is to handle inconsistent

@@ -398,7 +397,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):

397 source_layer: nn.Module,

398 lora_config: LoRAConfig,

399 packed_modules_list: list,

401 - model_config: Optional[PretrainedConfig],

400 + model_config: PretrainedConfig | None,

401 ) -> bool:

402 return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 3

403

@@ -426,9 +425,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):

425 lora_a = lora_a[start_idx : start_idx + shard_size, :]

426 return lora_a

427

429 - def apply(

430 - self, x: torch.Tensor, bias: Optional[torch.Tensor] = None

431 - ) -> torch.Tensor:

428 + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

429 return _mcp_apply(x, bias, self)

430

431 @classmethod

@@ -438,7 +435,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):

435 source_layer: nn.Module,

436 lora_config: LoRAConfig,

437 packed_modules_list: list,

441 - model_config: Optional[PretrainedConfig],

438 + model_config: PretrainedConfig | None,

439 ) -> bool:

440 # specifying kwargs so they can be easily accessed in decorator

441 return super().can_replace_layer(

@@ -459,8 +456,8 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo

456 """

457

458 def slice_lora_a(

462 - self, lora_a: list[Union[torch.Tensor, None]]

463 - ) -> list[Union[torch.Tensor, None]]:

459 + self, lora_a: list[torch.Tensor | None]

460 + ) -> list[torch.Tensor | None]:

461 # NOTE: lora_a contains 2 subloras, and each sublora could be None.

462 output_shard_size = self.lora_a_stacked[0].shape[2]

463 output_start_idx = self.tp_rank * output_shard_size

@@ -474,9 +471,7 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo

471 ]

472 return lora_a

473

477 - def apply(

478 - self, x: torch.Tensor, bias: Optional[torch.Tensor] = None

479 - ) -> torch.Tensor:

474 + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

475 return _mcp_apply(x, bias, self)

476

477 @classmethod

@@ -486,7 +481,7 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo

481 source_layer: nn.Module,

482 lora_config: LoRAConfig,

483 packed_modules_list: list,

489 - model_config: Optional[PretrainedConfig],

484 + model_config: PretrainedConfig | None,

485 ) -> bool:

486 # specifying kwargs so they can be easily accessed in decorator

487 return super().can_replace_layer(

@@ -512,9 +507,7 @@ class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):

507 lora_a = lora_a[start_idx : start_idx + shard_size, :]

508 return lora_a

509

515 - def apply(

516 - self, x: torch.Tensor, bias: Optional[torch.Tensor] = None

517 - ) -> torch.Tensor:

510 + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

511 return _mcp_apply(x, bias, self)

512

513 @classmethod

@@ -524,7 +517,7 @@ class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):

517 source_layer: nn.Module,

518 lora_config: LoRAConfig,

519 packed_modules_list: list,

527 - model_config: Optional[PretrainedConfig],

520 + model_config: PretrainedConfig | None,

521 ) -> bool:

522 # specifying kwargs so they can be easily accessed in decorator

523 return super().can_replace_layer(

@@ -545,8 +538,8 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):

538 """

539

540 def slice_lora_a(

548 - self, lora_a: list[Union[torch.Tensor, None]]

549 - ) -> list[Union[torch.Tensor, None]]:

541 + self, lora_a: list[torch.Tensor | None]

542 + ) -> list[torch.Tensor | None]:

543 # NOTE: lora_a contains 3 subloras, and each sublora could be None.

544 shard_size = [self.lora_a_stacked[i].shape[2] for i in range(3)]

545 start_idx = [self.tp_rank * shard_size[i] for i in range(3)]

@@ -563,9 +556,7 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):

556 ]

557 return lora_a

558

566 - def apply(

567 - self, x: torch.Tensor, bias: Optional[torch.Tensor] = None

568 - ) -> torch.Tensor:

559 + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

560 return _mcp_apply(x, bias, self)

561

562 @classmethod

@@ -575,7 +566,7 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):

566 source_layer: nn.Module,

567 lora_config: LoRAConfig,

568 packed_modules_list: list,

578 - model_config: Optional[PretrainedConfig],

569 + model_config: PretrainedConfig | None,

570 ) -> bool:

571 # specifying kwargs so they can be easily accessed in decorator

572 return super().can_replace_layer(

vllm/lora/layers/logits_processor.py

+7 -8

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

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

3

4 import math

5 -from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -41,7 +40,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

40 hidden_size: int,

41 dtype: torch.dtype,

42 device: torch.device,

44 - sharded_to_full_mapping: Optional[list[int]],

43 + sharded_to_full_mapping: list[int] | None,

44 ) -> None:

45 super().__init__()

46 self.base_layer = base_layer

@@ -88,7 +87,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

87 self,

88 max_loras: int,

89 lora_config: LoRAConfig,

91 - model_config: Optional[PretrainedConfig] = None,

90 + model_config: PretrainedConfig | None = None,

91 ) -> None:

92 # TODO: Verify if this condition can be further relaxed

93 if 32000 < self.base_layer.vocab_size > 257024:

@@ -142,7 +141,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

141 index: int,

142 lora_a: torch.Tensor,

143 lora_b: torch.Tensor,

145 - embeddings_tensor: Optional[torch.Tensor],

144 + embeddings_tensor: torch.Tensor | None,

145 ):

146 self.reset_lora(index)

147 self.lora_a_stacked[index, 0, : lora_a.shape[0], : lora_a.shape[1]].copy_(

@@ -162,8 +161,8 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

161 self,

162 hidden_states: torch.Tensor,

163 lm_head: VocabParallelEmbedding,

165 - embedding_bias: Optional[torch.Tensor] = None,

166 - ) -> Optional[torch.Tensor]:

164 + embedding_bias: torch.Tensor | None = None,

165 + ) -> torch.Tensor | None:

166 # Get the logits for the next tokens.

167 logits = lm_head.quant_method.apply(lm_head, hidden_states)

168 if embedding_bias is not None:

@@ -227,7 +226,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

226 + lora_logits.shape[1],

227 ] = lora_logits

228

230 - lora_output: Optional[torch.Tensor] = self.punica_wrapper.add_lora_logits(

229 + lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits(

230 logits, hidden_states, self.lora_a_stacked, self.lora_b_stacked, 1.0

231 )

232

@@ -247,7 +246,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

246 source_layer: nn.Module,

247 lora_config: LoRAConfig,

248 packed_modules_list: list,

250 - model_config: Optional[PretrainedConfig],

249 + model_config: PretrainedConfig | None,

250 ) -> bool:

251 # Special handling for the LogitsProcessor.

252 return False

vllm/lora/layers/replicated_linear.py

+2 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 import torch

6 import torch.nn as nn

@@ -24,7 +23,7 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

23

24 def forward(

25 self, input_: torch.Tensor

27 - ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:

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

27 """Forward of ReplicatedLinearWithLoRA

28

29 Args:

@@ -54,6 +53,6 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

53 source_layer: nn.Module,

54 lora_config: LoRAConfig,

55 packed_modules_list: list,

57 - model_config: Optional[PretrainedConfig],

56 + model_config: PretrainedConfig | None,

57 ) -> bool:

58 return type(source_layer) is ReplicatedLinear

vllm/lora/layers/row_parallel_linear.py

+6 -9

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional, Union

4

5 import torch

6 import torch.nn as nn

@@ -41,7 +40,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

40

41 def forward(

42 self, input_: torch.Tensor

44 - ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:

43 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:

44 """Forward of RowParallelLinear

45

46 Args:

@@ -93,7 +92,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

92 source_layer: nn.Module,

93 lora_config: LoRAConfig,

94 packed_modules_list: list,

96 - model_config: Optional[PretrainedConfig],

95 + model_config: PretrainedConfig | None,

96 ) -> bool:

97 return type(source_layer) is RowParallelLinear

98

@@ -120,9 +119,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

119 lora_b = lora_b[start_idx:end_idx, :]

120 return lora_b

121

123 - def apply(

124 - self, x: torch.Tensor, bias: Optional[torch.Tensor] = None

125 - ) -> torch.Tensor:

122 + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:

123 output = self.base_layer.quant_method.apply(self.base_layer, x)

124

125 x = x.view(-1, x.shape[-1])

@@ -133,7 +130,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

130 device=x.device,

131 )

132

136 - shrunk_buffer: Optional[torch.Tensor] = self.punica_wrapper.add_shrink(

133 + shrunk_buffer: torch.Tensor | None = self.punica_wrapper.add_shrink(

134 buffer, x, self.lora_a_stacked, 1.0

135 )

136 if not current_platform.can_update_inplace():

@@ -150,7 +147,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

147 # NOTE offset are based on the rank.

148 shard_size = self.lora_b_stacked[0].shape[2]

149 offset_start = self.tp_rank * shard_size

153 - lora_output: Optional[torch.Tensor] = self.punica_wrapper.add_expand(

150 + lora_output: torch.Tensor | None = self.punica_wrapper.add_expand(

151 output,

152 buffer,

153 self.lora_b_stacked,

@@ -172,7 +169,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

169 source_layer: nn.Module,

170 lora_config: LoRAConfig,

171 packed_modules_list: list,

175 - model_config: Optional[PretrainedConfig],

172 + model_config: PretrainedConfig | None,

173 ) -> bool:

174 # specifying kwargs so they can be easily accessed in decorator

175 return super().can_replace_layer(

vllm/lora/layers/vocal_parallel_embedding.py

+6 -7

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6 import torch.nn as nn

@@ -19,14 +18,14 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):

18 def __init__(self, base_layer: VocabParallelEmbedding) -> None:

19 super().__init__()

20 self.base_layer = base_layer

22 - self.embeddings_slice: Optional[tuple[int, int]]

23 - self.embeddings_weights: Optional[torch.Tensor]

21 + self.embeddings_slice: tuple[int, int] | None

22 + self.embeddings_weights: torch.Tensor | None

23

24 def create_lora_weights(

25 self,

26 max_loras: int,

27 lora_config: LoRAConfig,

29 - model_config: Optional[PretrainedConfig] = None,

28 + model_config: PretrainedConfig | None = None,

29 ) -> None:

30 if self.base_layer.num_added_embeddings_per_partition > 0:

31 # We can start adding lora weights

@@ -90,7 +89,7 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):

89 index: int,

90 lora_a: torch.Tensor,

91 lora_b: torch.Tensor,

93 - embeddings_tensor: Optional[torch.Tensor],

92 + embeddings_tensor: torch.Tensor | None,

93 ):

94 self.reset_lora(index)

95 # NOTE self.lora_a_stacked is row-major, and lora_a is col-major,

@@ -143,7 +142,7 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):

142 -1,

143 )

144

146 - lora_output: Optional[torch.Tensor] = self.punica_wrapper.add_lora_embedding(

145 + lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_embedding(

146 full_output, full_lora_a_embeddings, self.lora_b_stacked, add_input=True

147 )

148

@@ -158,7 +157,7 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):

157 source_layer: nn.Module,

158 lora_config: LoRAConfig,

159 packed_modules_list: list,

161 - model_config: Optional[PretrainedConfig],

160 + model_config: PretrainedConfig | None,

161 ) -> bool:

162 return type(source_layer) is VocabParallelEmbedding

163

vllm/lora/lora_weights.py

+8 -8

@@ -21,8 +21,8 @@ class LoRALayerWeights:

21 lora_alpha: int,

22 lora_a: torch.Tensor,

23 lora_b: torch.Tensor,

24 - embeddings_tensor: Optional[torch.Tensor] = None,

25 - scaling: Optional[float] = None,

24 + embeddings_tensor: torch.Tensor | None = None,

25 + scaling: float | None = None,

26 ) -> None:

27 self.module_name = module_name

28 self.rank = rank

@@ -67,7 +67,7 @@ class LoRALayerWeights:

67 cls,

68 module_name: str,

69 peft_helper: PEFTHelper,

70 - embeddings_tensor: Optional[torch.Tensor] = None,

70 + embeddings_tensor: torch.Tensor | None = None,

71 ) -> "LoRALayerWeights":

72 # lora_a and lora_b are set to None for config-based construction

73 return cls(

@@ -89,7 +89,7 @@ class LoRALayerWeights:

89 rank: int,

90 dtype: torch.dtype,

91 device: torch.types.Device,

92 - embeddings_tensor_dim: Optional[int] = None,

92 + embeddings_tensor_dim: int | None = None,

93 ) -> "LoRALayerWeights":

94 pin_memory = str(device) == "cpu" and is_pin_memory_available()

95 lora_a = torch.zeros(

@@ -127,10 +127,10 @@ class PackedLoRALayerWeights(LoRALayerWeights):

127 self,

128 module_name: str,

129 rank: int,

130 - lora_alphas: list[Optional[int]],

131 - lora_a: list[Optional[torch.Tensor]],

132 - lora_b: list[Optional[torch.Tensor]],

133 - scaling: Optional[list[float]] = None,

130 + lora_alphas: list[int | None],

131 + lora_a: list[torch.Tensor | None],

132 + lora_b: list[torch.Tensor | None],

133 + scaling: list[float] | None = None,

134 ) -> None:

135 super().__init__(

136 module_name=module_name,

vllm/lora/models.py

+25 -24

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

3

4 import math

5 import os

6 -from typing import Callable, Optional, TypeVar, Union

6 +from collections.abc import Callable

7 +from typing import TypeVar

8

9 import regex as re

10 import safetensors.torch

@@ -44,7 +45,7 @@ class AdapterLRUCache(LRUCache[int, T]):

45 super().__init__(capacity)

46 self.deactivate_fn = deactivate_fn

47

47 - def _on_remove(self, key: int, value: Optional[T]):

48 + def _on_remove(self, key: int, value: T | None):

49 logger.debug("Removing adapter int id: %d", key)

50 self.deactivate_fn(key)

51 return super()._on_remove(key, value)

@@ -113,7 +114,7 @@ class LoRAModel:

114 else 0

115 )

116

116 - def get_lora(self, module_name: str) -> Optional[LoRALayerWeights]:

117 + def get_lora(self, module_name: str) -> LoRALayerWeights | None:

118 """Get LoRA for a given module by name"""

119 return self.loras.get(module_name, None)

120

@@ -128,12 +129,12 @@ class LoRAModel:

129 tensors: dict[str, torch.Tensor],

130 peft_helper: PEFTHelper,

131 device: str = "cuda",

131 - dtype: Optional[torch.dtype] = None,

132 - embeddings: Optional[dict[str, torch.Tensor]] = None,

133 - target_embedding_padding: Optional[int] = None,

134 - embedding_modules: Optional[dict[str, str]] = None,

135 - embedding_padding_modules: Optional[list[str]] = None,

136 - weights_mapper: Optional[WeightsMapper] = None,

132 + dtype: torch.dtype | None = None,

133 + embeddings: dict[str, torch.Tensor] | None = None,

134 + target_embedding_padding: int | None = None,

135 + embedding_modules: dict[str, str] | None = None,

136 + embedding_padding_modules: list[str] | None = None,

137 + weights_mapper: WeightsMapper | None = None,

138 ) -> "LoRAModel":

139 """Create a LoRAModel from a dictionary of tensors."""

140 pin_memory = str(device) == "cpu" and is_pin_memory_available()

@@ -191,14 +192,14 @@ class LoRAModel:

192 expected_lora_modules: list[str],

193 peft_helper: PEFTHelper,

194 *,

194 - lora_model_id: Optional[int] = None,

195 + lora_model_id: int | None = None,

196 device: str = "cuda",

196 - dtype: Optional[torch.dtype] = None,

197 - target_embedding_padding: Optional[int] = None,

198 - embedding_modules: Optional[dict[str, str]] = None,

199 - embedding_padding_modules: Optional[list[str]] = None,

200 - weights_mapper: Optional[WeightsMapper] = None,

201 - tensorizer_config_dict: Optional[dict] = None,

197 + dtype: torch.dtype | None = None,

198 + target_embedding_padding: int | None = None,

199 + embedding_modules: dict[str, str] | None = None,

200 + embedding_padding_modules: list[str] | None = None,

201 + weights_mapper: WeightsMapper | None = None,

202 + tensorizer_config_dict: dict | None = None,

203 ) -> "LoRAModel":

204 """Create a LoRAModel from a local checkpoint.

205

@@ -223,7 +224,7 @@ class LoRAModel:

224 )

225 new_embeddings_bin_file_path = os.path.join(lora_dir, "new_embeddings.bin")

226 tensors: dict[str, torch.Tensor] = {}

226 - unexpected_modules: list[Union[list[str], str]] = []

227 + unexpected_modules: list[list[str] | str] = []

228

229 def check_unexpected_modules(modules: dict):

230 for lora_module in modules.keys(): # noqa

@@ -357,7 +358,7 @@ class LoRAModelManager:

358 self.max_num_seqs = max_num_seqs

359 assert self.capacity >= self.lora_slots

360 self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8

360 - self.lora_index_to_id: list[Optional[int]] = [None] * self.lora_slots

361 + self.lora_index_to_id: list[int | None] = [None] * self.lora_slots

362 self.vocab_size = vocab_size

363 self.punica_wrapper = get_punica_wrapper(

364 max_num_batched_tokens,

@@ -383,7 +384,7 @@ class LoRAModelManager:

384 self.packed_modules: dict[str, list[str]] = {}

385 self.modules: dict[str, BaseLayerWithLoRA] = {}

386 # Dict instead of a set for compatibility with LRUCache.

386 - self._last_mapping: Optional[LoRAMapping] = None

387 + self._last_mapping: LoRAMapping | None = None

388 self._create_lora_modules()

389 self.model.lora_manager = self

390

@@ -555,7 +556,7 @@ class LoRAModelManager:

556 self,

557 lora_id: int,

558 rank: int,

558 - embedding_modules: Optional[dict[str, str]] = None,

559 + embedding_modules: dict[str, str] | None = None,

560 ) -> LoRAModel:

561 """Create zero-initialized LoRAModel for warmup."""

562 model = LoRAModel(lora_id, rank, {})

@@ -607,7 +608,7 @@ class LoRAModelManager:

608 else:

609 parts = module_name.split(".")

610 replacements = self.packed_modules_mapping[parts[-1]]

610 - subloras: list[Optional[LoRALayerWeights]] = []

611 + subloras: list[LoRALayerWeights | None] = []

612 for i, r in enumerate(replacements):

613 lora = LoRALayerWeights.create_dummy_lora_weights(

614 module_name + "." + r,

@@ -658,7 +659,7 @@ class LoRAModelManager:

659

660 def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None:

661 for module_name, new_module_names in self.packed_modules.items():

661 - replacement_loras: list[Optional[LoRALayerWeights]] = []

662 + replacement_loras: list[LoRALayerWeights | None] = []

663 replaced_module: set[str] = set()

664 has_replacement = False

665 for r in new_module_names:

@@ -687,7 +688,7 @@ class LoRAModelManager:

688

689 def _get_lora_layer_weights(

690 self, lora_model: LoRAModel, module_name: str

690 - ) -> Optional[LoRALayerWeights]:

691 + ) -> LoRALayerWeights | None:

692 org_module_name = module_name

693 if self.is_pooling_model and not lora_model.check_lora_name(module_name):

694 # If it's a pool model, and the layer name is not found,

@@ -732,7 +733,7 @@ class LoRAModelManager:

733 def list_adapters(self) -> dict[int, LoRAModel]:

734 return dict(self._registered_adapters)

735

735 - def get_adapter(self, adapter_id: int) -> Optional[LoRAModel]:

736 + def get_adapter(self, adapter_id: int) -> LoRAModel | None:

737 return self._registered_adapters.get(adapter_id)

738

739

vllm/lora/ops/triton_ops/lora_kernel_metadata.py

+1 -2

@@ -5,7 +5,6 @@ LoRA kernels metadata preparation utilities.

5 """

6

7 from dataclasses import dataclass

8 -from typing import Union

8

9 import torch

10

@@ -31,7 +30,7 @@ class LoRAKernelMeta:

30

31 @staticmethod

32 def make(

34 - max_loras: int, max_num_tokens: int, device: Union[torch.device, str]

33 + max_loras: int, max_num_tokens: int, device: torch.device | str

34 ) -> "LoRAKernelMeta":

35 token_lora_mapping = torch.empty(

36 max_num_tokens, dtype=torch.int32, device=device

vllm/lora/peft_helper.py

+6 -6

@@ -7,7 +7,7 @@ import json

7 import math

8 import os

9 from dataclasses import MISSING, dataclass, field, fields

10 -from typing import Literal, Optional, Union

10 +from typing import Literal

11

12 from vllm.config.lora import LoRAConfig

13 from vllm.logger import init_logger

@@ -27,17 +27,17 @@ class PEFTHelper:

27 # Required fields

28 r: int

29 lora_alpha: int

30 - target_modules: Union[list[str], str]

30 + target_modules: list[str] | str

31

32 bias: Literal["none"] = field(default="none")

33 - modules_to_save: Optional[list[str]] = field(default=None)

33 + modules_to_save: list[str] | None = field(default=None)

34 # True to use Rank-Stabilized LoRA (rsLoRA, see: https://arxiv.org/abs/2312.03732)

35 use_rslora: bool = field(default=False)

36 # True to use Weight-Decomposed Low-Rank Adaptation (DoRA, see: https://arxiv.org/abs/2402.09353)

37 use_dora: bool = field(default=False)

38 # Extra vllm field, start with 'vllm_' to avoid conflict

39 vllm_lora_scaling_factor: float = field(default=1.0)

40 - vllm_max_position_embeddings: Optional[int] = field(default=False)

40 + vllm_max_position_embeddings: int | None = field(default=False)

41

42 def _validate_features(self) -> list[str]:

43 """

@@ -81,8 +81,8 @@ class PEFTHelper:

81 def from_local_dir(

82 cls,

83 lora_path: str,

84 - max_position_embeddings: Optional[int],

85 - tensorizer_config_dict: Optional[dict] = None,

84 + max_position_embeddings: int | None,

85 + tensorizer_config_dict: dict | None = None,

86 ) -> "PEFTHelper":

87 lora_config_path = os.path.join(lora_path, "adapter_config.json")

88

vllm/lora/punica_wrapper/punica_base.py

+24 -24

@@ -8,7 +8,7 @@ https://arxiv.org/abs/2310.18547

8 """

9

10 from abc import ABC, abstractmethod

11 -from typing import TYPE_CHECKING, Optional, Union

11 +from typing import TYPE_CHECKING

12

13 import torch

14

@@ -28,7 +28,7 @@ class PunicaWrapperABC(ABC):

28 def update_metadata(

29 self,

30 mapping: "LoRAMapping",

31 - lora_index_to_id: list[Optional[int]],

31 + lora_index_to_id: list[int | None],

32 max_loras: int,

33 vocab_size: int,

34 extra_vocab_size: int,

@@ -42,12 +42,12 @@ class PunicaWrapperABC(ABC):

42 @abstractmethod

43 def add_shrink(

44 self,

45 - y: Union[tuple[torch.Tensor, ...], torch.Tensor],

45 + y: tuple[torch.Tensor, ...] | torch.Tensor,

46 x: torch.Tensor,

47 lora_a_stacked: tuple[torch.Tensor, ...],

48 scale: float,

49 **kwargs,

50 - ) -> Optional[torch.Tensor]:

50 + ) -> torch.Tensor | None:

51 """

52 Performs GEMM for multiple slices of lora_a.

53 """

@@ -58,13 +58,13 @@ class PunicaWrapperABC(ABC):

58 def add_expand(

59 self,

60 y: torch.Tensor,

61 - x: Union[tuple[torch.Tensor, ...], torch.Tensor],

61 + x: tuple[torch.Tensor, ...] | torch.Tensor,

62 lora_b_stacked: tuple[torch.Tensor, ...],

63 output_slices: tuple[int, ...],

64 offset_start: int = 0,

65 add_inputs=True,

66 **kwargs,

67 - ) -> Optional[torch.Tensor]:

67 + ) -> torch.Tensor | None:

68 """

69 Performs GEMM for multiple slices of lora_b.

70 """

@@ -78,7 +78,7 @@ class PunicaWrapperABC(ABC):

78 lora_b_stacked: torch.Tensor,

79 add_inputs: bool = True,

80 **kwargs,

81 - ) -> Optional[torch.Tensor]:

81 + ) -> torch.Tensor | None:

82 """

83 Applies lora specifically for VocabParallelEmbeddingWithLoRA,

84 and this layer only requires the expand operation.

@@ -95,9 +95,9 @@ class PunicaWrapperABC(ABC):

95 scale: float,

96 output_slices: tuple[int, ...],

97 *,

98 - buffer: Optional[tuple[torch.Tensor, ...]] = None,

98 + buffer: tuple[torch.Tensor, ...] | None = None,

99 **kwargs,

100 - ) -> Optional[torch.Tensor]:

100 + ) -> torch.Tensor | None:

101 """

102 Applicable to linear-related lora.

103 """

@@ -113,9 +113,9 @@ class PunicaWrapperABC(ABC):

113 lora_b_stacked: torch.Tensor,

114 scale,

115 *,

116 - buffer: Optional[torch.Tensor] = None,

116 + buffer: torch.Tensor | None = None,

117 **kwargs,

118 - ) -> Optional[torch.Tensor]:

118 + ) -> torch.Tensor | None:

119 """

120 Applies lora specifically for LogitsProcessorWithLoRA.

121 """

@@ -133,7 +133,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

133 self,

134 max_num_batched_tokens: int,

135 max_batches: int,

136 - device: Union[torch.device, str],

136 + device: torch.device | str,

137 **kwargs,

138 ):

139 self._token_lora_indices = torch.empty(

@@ -152,7 +152,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

152 # 4 is the number of indices tensors.

153 # base_indices, sampler_indices, sampler_indices_padded,

154 # embeddings_indices

155 - self.indices_len: list[Optional[int]] = [None] * 4

155 + self.indices_len: list[int | None] = [None] * 4

156 # these attributes are the information required for sgmv kernel

157 self._seq_start_locs = torch.empty(max_batches, dtype=torch.long, device=device)

158 self._seq_lengths = torch.empty(max_batches, dtype=torch.long, device=device)

@@ -169,7 +169,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

169 def _update_base_metadata(

170 self,

171 mapping: "LoRAMapping",

172 - lora_index_to_id: list[Optional[int]],

172 + lora_index_to_id: list[int | None],

173 max_loras: int,

174 vocab_size: int,

175 extra_vocab_size: int,

@@ -282,7 +282,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

282 def update_metadata(

283 self,

284 mapping: "LoRAMapping",

285 - lora_index_to_id: list[Optional[int]],

285 + lora_index_to_id: list[int | None],

286 max_loras: int,

287 vocab_size: int,

288 extra_vocab_size: int,

@@ -302,12 +302,12 @@ class PunicaWrapperBase(PunicaWrapperABC):

302 @abstractmethod

303 def add_shrink(

304 self,

305 - y: Union[tuple[torch.Tensor, ...], torch.Tensor],

305 + y: tuple[torch.Tensor, ...] | torch.Tensor,

306 x: torch.Tensor,

307 lora_a_stacked: tuple[torch.Tensor, ...],

308 scale: float,

309 **kwargs,

310 - ) -> Optional[torch.Tensor]:

310 + ) -> torch.Tensor | None:

311 """

312 Performs GEMM for multiple slices of lora_a.

313

@@ -329,13 +329,13 @@ class PunicaWrapperBase(PunicaWrapperABC):

329 def add_expand(

330 self,

331 y: torch.Tensor,

332 - x: Union[tuple[torch.Tensor, ...], torch.Tensor],

332 + x: tuple[torch.Tensor, ...] | torch.Tensor,

333 lora_b_stacked: tuple[torch.Tensor, ...],

334 output_slices: tuple[int, ...],

335 offset_start: int = 0,

336 add_inputs=True,

337 **kwargs,

338 - ) -> Optional[torch.Tensor]:

338 + ) -> torch.Tensor | None:

339 """

340 Performs GEMM for multiple slices of lora_b.

341

@@ -366,7 +366,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

366 lora_b_stacked: torch.Tensor,

367 add_inputs: bool = True,

368 **kwargs,

369 - ) -> Optional[torch.Tensor]:

369 + ) -> torch.Tensor | None:

370 """

371 Applies lora specifically for VocabParallelEmbeddingWithLoRA.

372 and this layer only requires the expand operation.

@@ -392,9 +392,9 @@ class PunicaWrapperBase(PunicaWrapperABC):

392 scale: float,

393 output_slices: tuple[int, ...],

394 *,

395 - buffer: Optional[tuple[torch.Tensor, ...]] = None,

395 + buffer: tuple[torch.Tensor, ...] | None = None,

396 **kwargs,

397 - ) -> Optional[torch.Tensor]:

397 + ) -> torch.Tensor | None:

398 """

399 Applicable to linear-related lora.

400

@@ -428,9 +428,9 @@ class PunicaWrapperBase(PunicaWrapperABC):

428 lora_b_stacked: torch.Tensor,

429 scale,

430 *,

431 - buffer: Optional[torch.Tensor] = None,

431 + buffer: torch.Tensor | None = None,

432 **kwargs,

433 - ) -> Optional[torch.Tensor]:

433 + ) -> torch.Tensor | None:

434 """

435 Applies lora specifically for LogitsProcessorWithLoRA.

436

vllm/lora/punica_wrapper/punica_cpu.py

+6 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional, Union

4 +from collections.abc import Callable

5

6 import torch

7

@@ -30,7 +30,7 @@ class PunicaWrapperCPU(PunicaWrapperBase):

30 self,

31 max_num_batched_tokens: int,

32 max_batches: int,

33 - device: Union[torch.device, str],

33 + device: torch.device | str,

34 **kwargs,

35 ):

36 PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)

@@ -165,7 +165,7 @@ class PunicaWrapperCPU(PunicaWrapperBase):

165

166 def add_shrink(

167 self,

168 - y: Union[tuple[torch.Tensor, ...], torch.Tensor],

168 + y: tuple[torch.Tensor, ...] | torch.Tensor,

169 x: torch.Tensor,

170 lora_a_stacked: tuple[torch.Tensor, ...],

171 scale: float,

@@ -197,7 +197,7 @@ class PunicaWrapperCPU(PunicaWrapperBase):

197 def add_expand(

198 self,

199 y: torch.Tensor,

200 - x: Union[tuple[torch.Tensor, ...], torch.Tensor],

200 + x: tuple[torch.Tensor, ...] | torch.Tensor,

201 lora_b_stacked: tuple[torch.Tensor, ...],

202 output_slices: tuple[int, ...],

203 offset_start: int = 0,

@@ -271,7 +271,7 @@ class PunicaWrapperCPU(PunicaWrapperBase):

271 scale: float,

272 output_slices: tuple[int, ...],

273 *,

274 - buffer: Optional[tuple[torch.Tensor, ...]] = None,

274 + buffer: tuple[torch.Tensor, ...] | None = None,

275 **kwargs,

276 ) -> None:

277 """

@@ -319,7 +319,7 @@ class PunicaWrapperCPU(PunicaWrapperBase):

319 lora_b_stacked: torch.Tensor,

320 scale,

321 *,

322 - buffer: Optional[torch.Tensor] = None,

322 + buffer: torch.Tensor | None = None,

323 **kwargs,

324 ) -> None:

325 """

vllm/lora/punica_wrapper/punica_gpu.py

+5 -5

@@ -7,7 +7,7 @@ Punica: Multi-Tenant LoRA Serving.

7 https://arxiv.org/abs/2310.18547

8 """

9

10 -from typing import Optional, Union, final

10 +from typing import final

11

12 import torch

13

@@ -32,7 +32,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

32 self,

33 max_num_batched_tokens: int,

34 max_batches: int,

35 - device: Union[torch.device, str],

35 + device: torch.device | str,

36 **kwargs,

37 ):

38 PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)

@@ -50,7 +50,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

50 def update_metadata(

51 self,

52 mapping: LoRAMapping,

53 - lora_index_to_id: list[Optional[int]],

53 + lora_index_to_id: list[int | None],

54 max_loras: int,

55 vocab_size: int,

56 extra_vocab_size: int,

@@ -179,7 +179,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

179 scale: float,

180 output_slices: tuple[int, ...],

181 *,

182 - buffer: Optional[torch.Tensor] = None,

182 + buffer: torch.Tensor | None = None,

183 **kwargs,

184 ) -> None:

185 """

@@ -238,7 +238,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

238 lora_b_stacked: torch.Tensor,

239 scale,

240 *,

241 - buffer: Optional[torch.Tensor] = None,

241 + buffer: torch.Tensor | None = None,

242 **kwargs,

243 ) -> None:

244 """

vllm/lora/punica_wrapper/punica_tpu.py

+8 -8

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

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

3

4 import math

5 -from typing import TYPE_CHECKING, Optional, Union

5 +from typing import TYPE_CHECKING

6

7 import torch

8 import torch.nn.functional as F

@@ -29,7 +29,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):

29 self,

30 max_num_batched_tokens: int,

31 max_batches: int,

32 - device: Union[torch.device, str],

32 + device: torch.device | str,

33 **kwargs,

34 ):

35 PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)

@@ -105,12 +105,12 @@ class PunicaWrapperTPU(PunicaWrapperBase):

105

106 def add_shrink(

107 self,

108 - y: Union[tuple[torch.Tensor, ...], torch.Tensor],

108 + y: tuple[torch.Tensor, ...] | torch.Tensor,

109 x: torch.Tensor,

110 lora_a_stacked: tuple[torch.Tensor, ...],

111 scale: float,

112 **kwargs,

113 - ) -> Optional[torch.Tensor]:

113 + ) -> torch.Tensor | None:

114 """

115 Performs GEMM for multiple slices of lora_a.

116

@@ -137,7 +137,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):

137 def add_expand(

138 self,

139 y: torch.Tensor,

140 - x: Union[tuple[torch.Tensor, ...], torch.Tensor],

140 + x: tuple[torch.Tensor, ...] | torch.Tensor,

141 lora_b_stacked: tuple[torch.Tensor, ...],

142 output_slices: tuple[int, ...],

143 offset_start: int = 0,

@@ -209,7 +209,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):

209 scale: float,

210 output_slices: tuple[int, ...],

211 *,

212 - buffer: Optional[tuple[torch.Tensor, ...]] = None,

212 + buffer: tuple[torch.Tensor, ...] | None = None,

213 **kwargs,

214 ) -> torch.Tensor:

215 """

@@ -257,7 +257,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):

257 lora_b_stacked: torch.Tensor,

258 scale,

259 *,

260 - buffer: Optional[torch.Tensor] = None,

260 + buffer: torch.Tensor | None = None,

261 **kwargs,

262 ) -> torch.Tensor:

263 """

@@ -289,7 +289,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):

289 def _update_base_metadata(

290 self,

291 mapping: "LoRAMapping",

292 - lora_index_to_id: list[Optional[int]],

292 + lora_index_to_id: list[int | None],

293 max_loras: int,

294 vocab_size: int,

295 extra_vocab_size: int,

vllm/lora/punica_wrapper/punica_xpu.py

+5 -5

@@ -7,7 +7,7 @@ Punica: Multi-Tenant LoRA Serving.

7 https://arxiv.org/abs/2310.18547

8 """

9

10 -from typing import Optional, Union, final

10 +from typing import final

11

12 import torch

13

@@ -29,7 +29,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

29 self,

30 max_num_batched_tokens: int,

31 max_batches: int,

32 - device: Union[torch.device, str],

32 + device: torch.device | str,

33 **kwargs,

34 ):

35 PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)

@@ -40,7 +40,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

40 def update_metadata(

41 self,

42 mapping: LoRAMapping,

43 - lora_index_to_id: list[Optional[int]],

43 + lora_index_to_id: list[int | None],

44 max_loras: int,

45 vocab_size: int,

46 extra_vocab_size: int,

@@ -180,7 +180,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

180 scale: float,

181 output_slices: tuple[int, ...],

182 *,

183 - buffer: Optional[torch.Tensor] = None,

183 + buffer: torch.Tensor | None = None,

184 **kwargs,

185 ) -> None:

186 """

@@ -247,7 +247,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

247 lora_b_stacked: torch.Tensor,

248 scale,

249 *,

250 - buffer: Optional[torch.Tensor] = None,

250 + buffer: torch.Tensor | None = None,

251 **kwargs,

252 ) -> None:

253 """

vllm/lora/punica_wrapper/utils.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import TYPE_CHECKING, Optional, Union

4 +from typing import TYPE_CHECKING

5

6 import torch

7

@@ -51,7 +51,7 @@ def compute_meta(

51 # TODO see if this can be vectorized

52 def convert_mapping(

53 mapping: "LoRAMapping",

54 - lora_index_to_id: list[Optional[int]],

54 + lora_index_to_id: list[int | None],

55 max_loras: int,

56 vocab_size: int,

57 extra_vocab_size: int,

@@ -104,7 +104,7 @@ def convert_mapping(

104 embedding_indices[i] = lora_idx if index_mapping_indices[i] > 0 else 0

105 lora_indices[i] = lora_idx

106

107 - indices_list: list[Union[list[int], torch.Tensor]] = [

107 + indices_list: list[list[int] | torch.Tensor] = [

108 index_mapping_indices,

109 lora_indices,

110 embedding_indices,

vllm/lora/request.py

+4 -5

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

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

3

4 import warnings

5 -from typing import Optional

5

6 import msgspec

7

@@ -27,10 +26,10 @@ class LoRARequest(

26 lora_name: str

27 lora_int_id: int

28 lora_path: str = ""

30 - lora_local_path: Optional[str] = msgspec.field(default=None)

31 - long_lora_max_len: Optional[int] = None

32 - base_model_name: Optional[str] = msgspec.field(default=None)

33 - tensorizer_config_dict: Optional[dict] = None

29 + lora_local_path: str | None = msgspec.field(default=None)

30 + long_lora_max_len: int | None = None

31 + base_model_name: str | None = msgspec.field(default=None)

32 + tensorizer_config_dict: dict | None = None

33

34 def __post_init__(self):

35 if self.lora_int_id < 1:

vllm/lora/resolver.py

+1 -2

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

4 from abc import ABC, abstractmethod

5 from collections.abc import Set

6 from dataclasses import dataclass, field

7 -from typing import Optional

7

8 from vllm.logger import init_logger

9 from vllm.lora.request import LoRARequest

@@ -24,7 +23,7 @@ class LoRAResolver(ABC):

23 @abstractmethod

24 async def resolve_lora(

25 self, base_model_name: str, lora_name: str

27 - ) -> Optional[LoRARequest]:

26 + ) -> LoRARequest | None:

27 """Abstract method to resolve and fetch a LoRA model adapter.

28

29 Implements logic to locate and download LoRA adapter based on the name.

vllm/lora/utils.py

+4 -4

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

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

3

4 import os

5 -from typing import TYPE_CHECKING, Optional, Union

5 +from typing import TYPE_CHECKING, Optional

6

7 import huggingface_hub

8 import regex as re

@@ -66,7 +66,7 @@ def from_layer(

66 max_loras: int,

67 lora_config: LoRAConfig,

68 packed_modules_list: list,

69 - model_config: Optional[PretrainedConfig] = None,

69 + model_config: PretrainedConfig | None = None,

70 ) -> nn.Module:

71 for lora_cls in _all_lora_classes:

72 # specifying kwargs so they can be easily accessed in decorator

@@ -87,7 +87,7 @@ def from_layer_logits_processor(

87 lm_head: "ParallelLMHead",

88 max_loras: int,

89 lora_config: LoRAConfig,

90 - model_config: Optional[PretrainedConfig] = None,

90 + model_config: PretrainedConfig | None = None,

91 ) -> LogitsProcessorWithLoRA:

92 ret = LogitsProcessorWithLoRA(

93 layer,

@@ -155,7 +155,7 @@ def parse_fine_tuned_lora_name(

155

156

157 def is_regex_target_modules(

158 - load_modules: Union[str, list[str]], expected_lora_modules: list[str]

158 + load_modules: str | list[str], expected_lora_modules: list[str]

159 ) -> bool:

160 """

161 PEFT supports passing `target_modules` in the form of regular expressions,

vllm/lora/worker_manager.py

+3 -3

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

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

3

4 from contextlib import contextmanager

5 -from typing import Any, Literal, Optional, Union

5 +from typing import Any, Literal

6

7 import torch

8

@@ -40,7 +40,7 @@ class WorkerLoRAManager:

40 self._lora_model_cls = lora_model_cls

41 self.embedding_modules = embedding_modules

42 self.embedding_padding_modules = embedding_padding_modules

43 - self._cached_dummy_lora: Union[None, Literal[False], LoRAModel] = False

43 + self._cached_dummy_lora: None | Literal[False] | LoRAModel = False

44 self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs

45 self.max_num_batched_tokens = (

46 vllm_config.scheduler_config.max_num_batched_tokens

@@ -166,7 +166,7 @@ class WorkerLoRAManager:

166 def pin_adapter(self, adapter_id: int) -> bool:

167 return self._adapter_manager.pin_adapter(adapter_id)

168

169 - def set_active_adapters(self, requests: set[Any], mapping: Optional[Any]) -> None:

169 + def set_active_adapters(self, requests: set[Any], mapping: Any | None) -> None:

170 self._apply_adapters(requests)

171 if mapping is not None:

172 self._adapter_manager.set_adapter_mapping(mapping)

vllm/model_executor/custom_op.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch.nn as nn

6

@@ -171,7 +170,7 @@ class CustomOp(nn.Module):

170 # or

171 # - @CustomOP.register_oot(name="UnquantizedFusedMoEMethod")

172 @classmethod

174 - def register_oot(cls, _decorated_op_cls=None, name: Optional[str] = None):

173 + def register_oot(cls, _decorated_op_cls=None, name: str | None = None):

174 def decorator(op_cls):

175 reg_name = name if name is not None else cls.__name__

176 assert reg_name not in cls.op_registry_oot, f"Duplicate op name: {reg_name}"

vllm/model_executor/layers/activation.py

+1 -2

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

3 """Custom activation functions."""

4

5 import math

6 -from typing import Optional

6

7 import torch

8 import torch.nn as nn

@@ -486,7 +485,7 @@ class ScaledActivation(nn.Module):

485 act_module: nn.Module,

486 intermediate_size: int,

487 input_is_parallel: bool = True,

489 - params_dtype: Optional[torch.dtype] = None,

488 + params_dtype: torch.dtype | None = None,

489 ):

490 super().__init__()

491 self.act = act_module

vllm/model_executor/layers/batch_invariant.py

+4 -6

@@ -4,7 +4,7 @@ import contextlib

4 import os

5 from collections import namedtuple

6 from collections.abc import Callable

7 -from typing import Any, Union

7 +from typing import Any

8

9 import torch

10

@@ -138,7 +138,7 @@ def matmul_kernel_persistent(

138

139

140 def matmul_persistent(

141 - a: torch.Tensor, b: torch.Tensor, bias: Union[torch.Tensor, None] = None

141 + a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None

142 ):

143 # Check constraints.

144 assert a.shape[1] == b.shape[0], "Incompatible dimensions"

@@ -375,7 +375,7 @@ def mean_dim(

375 input: torch.Tensor,

376 dim: int,

377 keepdim: bool = False,

378 - dtype: Union[torch.dtype, None] = None,

378 + dtype: torch.dtype | None = None,

379 ) -> torch.Tensor:

380 """

381 Triton implementation of torch.mean with single dimension reduction.

@@ -475,9 +475,7 @@ def _log_softmax_batch_invariant(input, dim, _half_to_float):

475 return log_softmax(input, dim=dim)

476

477

478 -def mean_batch_invariant(

479 - input, dim, keepdim=False, dtype: Union[torch.dtype, None] = None

480 -):

478 +def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None):

479 assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}"

480

481 result = input.to(torch.float32)

vllm/model_executor/layers/fla/ops/chunk.py

+3 -4

@@ -8,7 +8,6 @@

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9 # ruff: noqa: E501

10 import warnings

11 -from typing import Optional

11

12 import torch

13 from einops import rearrange

@@ -32,7 +31,7 @@ def chunk_gated_delta_rule_fwd(

31 scale: float,

32 initial_state: torch.Tensor,

33 output_final_state: bool,

35 - cu_seqlens: Optional[torch.LongTensor] = None,

34 + cu_seqlens: torch.LongTensor | None = None,

35 ):

36 g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)

37 # obtain WY representation. u is actually the new v.

@@ -86,7 +85,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):

85 scale: float,

86 initial_state: torch.Tensor,

87 output_final_state: bool,

89 - cu_seqlens: Optional[torch.LongTensor] = None,

88 + cu_seqlens: torch.LongTensor | None = None,

89 use_qk_l2norm_in_kernel: bool = False,

90 ):

91 if use_qk_l2norm_in_kernel:

@@ -119,7 +118,7 @@ def chunk_gated_delta_rule(

118 scale: float = None,

119 initial_state: torch.Tensor = None,

120 output_final_state: bool = False,

122 - cu_seqlens: Optional[torch.LongTensor] = None,

121 + cu_seqlens: torch.LongTensor | None = None,

122 head_first: bool = False,

123 use_qk_l2norm_in_kernel: bool = False,

124 ):

vllm/model_executor/layers/fla/ops/chunk_delta_h.py

+3 -4

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

7 # the following copyright notice:

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9 # ruff: noqa: E501

10 -from typing import Optional

10

11 import torch

12

@@ -257,12 +256,12 @@ def chunk_gated_delta_rule_fwd_h(

256 k: torch.Tensor,

257 w: torch.Tensor,

258 u: torch.Tensor,

260 - g: Optional[torch.Tensor] = None,

261 - initial_state: Optional[torch.Tensor] = None,

259 + g: torch.Tensor | None = None,

260 + initial_state: torch.Tensor | None = None,

261 output_final_state: bool = False,

262 chunk_size: int = 64, # SY: remove this argument and force chunk size 64?

263 save_new_value: bool = True,

265 - cu_seqlens: Optional[torch.LongTensor] = None,

264 + cu_seqlens: torch.LongTensor | None = None,

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

266 B, T, Hg, K, V = *k.shape, u.shape[-1]

267 H = u.shape[-2]

vllm/model_executor/layers/fla/ops/chunk_o.py

+3 -4

@@ -9,7 +9,6 @@

9

10 # ruff: noqa: E501

11

12 -from typing import Optional

12

13 import torch

14

@@ -144,9 +143,9 @@ def chunk_fwd_o(

143 k: torch.Tensor,

144 v: torch.Tensor,

145 h: torch.Tensor,

147 - g: Optional[torch.Tensor] = None, # cumsum of log decay

148 - scale: Optional[float] = None,

149 - cu_seqlens: Optional[torch.LongTensor] = None,

146 + g: torch.Tensor | None = None, # cumsum of log decay

147 + scale: float | None = None,

148 + cu_seqlens: torch.LongTensor | None = None,

149 chunk_size: int = 64,

150 ) -> torch.Tensor:

151 B, T, Hg, K, V = *q.shape, v.shape[-1]

vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py

+2 -3

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

7 # the following copyright notice:

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9 # ruff: noqa: E501

10 -from typing import Optional

10

11 import torch

12

@@ -104,8 +103,8 @@ def chunk_scaled_dot_kkt_fwd_kernel(

103 def chunk_scaled_dot_kkt_fwd(

104 k: torch.Tensor,

105 beta: torch.Tensor,

107 - g_cumsum: Optional[torch.Tensor] = None,

108 - cu_seqlens: Optional[torch.LongTensor] = None,

106 + g_cumsum: torch.Tensor | None = None,

107 + cu_seqlens: torch.LongTensor | None = None,

108 chunk_size: int = 64,

109 output_dtype: torch.dtype = torch.float32,

110 ) -> torch.Tensor:

vllm/model_executor/layers/fla/ops/cumsum.py

+6 -7

@@ -8,7 +8,6 @@

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9 # ruff: noqa: E501

10 import warnings

11 -from typing import Optional

11

12 import torch

13

@@ -163,9 +162,9 @@ def chunk_local_cumsum_scalar(

162 g: torch.Tensor,

163 chunk_size: int,

164 reverse: bool = False,

166 - cu_seqlens: Optional[torch.Tensor] = None,

165 + cu_seqlens: torch.Tensor | None = None,

166 head_first: bool = False,

168 - output_dtype: Optional[torch.dtype] = torch.float,

167 + output_dtype: torch.dtype | None = torch.float,

168 ) -> torch.Tensor:

169 if head_first:

170 B, H, T = g.shape

@@ -200,9 +199,9 @@ def chunk_local_cumsum_vector(

199 g: torch.Tensor,

200 chunk_size: int,

201 reverse: bool = False,

203 - cu_seqlens: Optional[torch.Tensor] = None,

202 + cu_seqlens: torch.Tensor | None = None,

203 head_first: bool = False,

205 - output_dtype: Optional[torch.dtype] = torch.float,

204 + output_dtype: torch.dtype | None = torch.float,

205 ) -> torch.Tensor:

206 if head_first:

207 B, H, T, S = g.shape

@@ -248,9 +247,9 @@ def chunk_local_cumsum(

247 g: torch.Tensor,

248 chunk_size: int,

249 reverse: bool = False,

251 - cu_seqlens: Optional[torch.Tensor] = None,

250 + cu_seqlens: torch.Tensor | None = None,

251 head_first: bool = False,

253 - output_dtype: Optional[torch.dtype] = torch.float,

252 + output_dtype: torch.dtype | None = torch.float,

253 **kwargs,

254 ) -> torch.Tensor:

255 if not head_first and g.shape[1] < g.shape[2]:

vllm/model_executor/layers/fla/ops/fused_recurrent.py

+9 -10

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

7 # the following copyright notice:

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9 # ruff: noqa: E501

10 -from typing import Optional

10

11 import torch

12

@@ -169,9 +168,9 @@ def fused_recurrent_gated_delta_rule_fwd(

168 scale: float,

169 initial_state: torch.Tensor,

170 inplace_final_state: bool = True,

172 - cu_seqlens: Optional[torch.LongTensor] = None,

173 - ssm_state_indices: Optional[torch.Tensor] = None,

174 - num_accepted_tokens: Optional[torch.Tensor] = None,

171 + cu_seqlens: torch.LongTensor | None = None,

172 + ssm_state_indices: torch.Tensor | None = None,

173 + num_accepted_tokens: torch.Tensor | None = None,

174 use_qk_l2norm_in_kernel: bool = False,

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

176 B, T, H, K, V = *k.shape, v.shape[-1]

@@ -248,9 +247,9 @@ class FusedRecurrentFunction(torch.autograd.Function):

247 scale: float,

248 initial_state: torch.Tensor,

249 inplace_final_state: bool = True,

251 - cu_seqlens: Optional[torch.LongTensor] = None,

252 - ssm_state_indices: Optional[torch.Tensor] = None,

253 - num_accepted_tokens: Optional[torch.Tensor] = None,

250 + cu_seqlens: torch.LongTensor | None = None,

251 + ssm_state_indices: torch.Tensor | None = None,

252 + num_accepted_tokens: torch.Tensor | None = None,

253 use_qk_l2norm_in_kernel: bool = False,

254 ):

255 o, final_state = fused_recurrent_gated_delta_rule_fwd(

@@ -280,9 +279,9 @@ def fused_recurrent_gated_delta_rule(

279 scale: float = None,

280 initial_state: torch.Tensor = None,

281 inplace_final_state: bool = True,

283 - cu_seqlens: Optional[torch.LongTensor] = None,

284 - ssm_state_indices: Optional[torch.Tensor] = None,

285 - num_accepted_tokens: Optional[torch.Tensor] = None,

282 + cu_seqlens: torch.LongTensor | None = None,

283 + ssm_state_indices: torch.Tensor | None = None,

284 + num_accepted_tokens: torch.Tensor | None = None,

285 use_qk_l2norm_in_kernel: bool = False,

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

287 r"""

vllm/model_executor/layers/fla/ops/l2norm.py

+1 -2

@@ -8,7 +8,6 @@

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9

10 import os

11 -from typing import Optional

11

12 import torch

13

@@ -90,7 +89,7 @@ def l2norm_fwd_kernel2(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr):

89

90

91 def l2norm_fwd(

93 - x: torch.Tensor, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None

92 + x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None

93 ):

94 x_shape_og = x.shape

95 x = x.view(-1, x.shape[-1])

vllm/model_executor/layers/fla/ops/layernorm_guard.py

+6 -7

@@ -14,7 +14,6 @@

14 # The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.

15

16 from functools import lru_cache

17 -from typing import Optional

17

18 import torch

19 import torch.nn as nn

@@ -324,10 +323,10 @@ class LayerNormGated(nn.Module):

323 self,

324 hidden_size,

325 eps: float = 1e-5,

327 - group_size: Optional[int] = None,

326 + group_size: int | None = None,

327 norm_before_gate: bool = True,

329 - device: Optional[torch.device] = None,

330 - dtype: Optional[torch.dtype] = None,

328 + device: torch.device | None = None,

329 + dtype: torch.dtype | None = None,

330 ):

331 """If group_size is not None, we do GroupNorm with each group having group_size elements.

332 group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group).

@@ -364,10 +363,10 @@ class RMSNormGated(nn.Module):

363 self,

364 hidden_size,

365 eps: float = 1e-5,

367 - group_size: Optional[int] = None,

366 + group_size: int | None = None,

367 norm_before_gate: bool = False,

369 - device: Optional[torch.device] = None,

370 - dtype: Optional[torch.dtype] = None,

368 + device: torch.device | None = None,

369 + dtype: torch.dtype | None = None,

370 ):

371 """If group_size is not None, we do GroupNorm with each group having group_size elements.

372 group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group).

vllm/model_executor/layers/fla/ops/solve_tril.py

+1 -2

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

7 # the following copyright notice:

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9 # ruff: noqa: E501

10 -from typing import Optional

10

11 import torch

12

@@ -407,7 +406,7 @@ def merge_16x16_to_64x64_inverse_kernel(

406 @input_guard

407 def solve_tril(

408 A: torch.Tensor,

410 - cu_seqlens: Optional[torch.Tensor] = None,

409 + cu_seqlens: torch.Tensor | None = None,

410 output_dtype: torch.dtype = torch.float,

411 ) -> torch.Tensor:

412 """

vllm/model_executor/layers/fla/ops/utils.py

+3 -2

@@ -11,8 +11,9 @@ import contextlib

11 import functools

12 import logging

13 import os

14 +from collections.abc import Callable

15 from enum import Enum

15 -from typing import Any, Callable, Literal, Optional

16 +from typing import Any, Literal

17

18 import torch

19

@@ -43,7 +44,7 @@ def tensor_cache(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]

44 A wrapped version of the input function with single-entry caching.

45 """

46

46 - cache_entries: tuple[Optional[tuple], Optional[dict], Any] = []

47 + cache_entries: tuple[tuple | None, dict | None, Any] = []

48 cache_size = 4

49

50 @functools.wraps(fn)

vllm/model_executor/layers/fla/ops/wy_fast.py

+1 -2

@@ -8,7 +8,6 @@

8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

9

10 # ruff: noqa: E501

11 -from typing import Optional

11

12 import torch

13

@@ -123,7 +122,7 @@ def recompute_w_u_fwd(

122 beta: torch.Tensor,

123 g_cumsum: torch.Tensor,

124 A: torch.Tensor,

126 - cu_seqlens: Optional[torch.LongTensor],

125 + cu_seqlens: torch.LongTensor | None,

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

127 B, T, Hg, K, V = *k.shape, v.shape[-1]

128 H = v.shape[-2]

vllm/model_executor/layers/fused_moe/__init__.py

+3 -3

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

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

3

4 from contextlib import contextmanager

5 -from typing import Any, Optional

5 +from typing import Any

6

7 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig

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

@@ -19,7 +19,7 @@ from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

19 from vllm.model_executor.layers.fused_moe.utils import activation_without_mul

20 from vllm.triton_utils import HAS_TRITON

21

22 -_config: Optional[dict[str, Any]] = None

22 +_config: dict[str, Any] | None = None

23

24

25 @contextmanager

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

31 _config = old_config

32

33

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

34 +def get_config() -> dict[str, Any] | None:

35 return _config

36

37

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+5 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -259,7 +258,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

258 topk: int,

259 global_num_experts: int,

260 local_num_experts: int,

262 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

261 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

262 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

263 # FIXME (varun): We should be able to dispatch only from the leader

264 # DP ranks in the case of TP > 1. At the moment, all the Ranks

@@ -282,12 +281,12 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

281 topk_ids: torch.Tensor,

282 activation: str,

283 global_num_experts: int,

285 - expert_map: Optional[torch.Tensor],

286 - a1q_scale: Optional[torch.Tensor],

287 - a2_scale: Optional[torch.Tensor],

284 + expert_map: torch.Tensor | None,

285 + a1q_scale: torch.Tensor | None,

286 + a2_scale: torch.Tensor | None,

287 workspace13: torch.Tensor,

288 workspace2: torch.Tensor,

290 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

289 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

290 apply_router_weight_on_input: bool,

291 ):

292 assert expert_tokens_meta is not None

vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py

+5 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -110,7 +109,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

109 topk: int,

110 global_num_experts: int,

111 local_num_experts: int,

113 - expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],

112 + expert_tokens_metadata: mk.ExpertTokensMetadata | None,

113 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

114 # Note: the deep gemm workspaces are strictly larger than the triton

115 # workspaces so we can be pessimistic here and allocate for DeepGemm

@@ -148,12 +147,12 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

147 topk_ids: torch.Tensor,

148 activation: str,

149 global_num_experts: int,

151 - expert_map: Optional[torch.Tensor],

152 - a1q_scale: Optional[torch.Tensor],

153 - a2_scale: Optional[torch.Tensor],

150 + expert_map: torch.Tensor | None,

151 + a1q_scale: torch.Tensor | None,

152 + a2_scale: torch.Tensor | None,

153 workspace13: torch.Tensor,

154 workspace2: torch.Tensor,

156 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

155 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

156 apply_router_weight_on_input: bool,

157 ):

158 experts = (

vllm/model_executor/layers/fused_moe/config.py

+64 -64

@@ -34,8 +34,8 @@ def _get_config_dtype_str(

34 use_fp8_w8a8: bool = False,

35 use_int8_w8a16: bool = False,

36 use_int4_w4a16: bool = False,

37 - ocp_mx_scheme: Optional[str] = None,

38 -) -> Optional[str]:

37 + ocp_mx_scheme: str | None = None,

38 +) -> str | None:

39 """

40 Return a string used to construct the filename that contains the

41 tuning info for a particular quantization scheme. See

@@ -60,16 +60,16 @@ def _get_config_dtype_str(

60

61

62 def _quant_flags_to_group_shape(

63 - quant_dtype: Union[torch.dtype, str, None],

63 + quant_dtype: torch.dtype | str | None,

64 per_act_token_quant: bool,

65 per_out_ch_quant: bool,

66 - block_shape: Optional[list[int]],

67 -) -> tuple[Optional[GroupShape], Optional[GroupShape]]:

66 + block_shape: list[int] | None,

67 +) -> tuple[GroupShape | None, GroupShape | None]:

68 """

69 Convert MoE quantization flags into more generic GroupShapes.

70 """

71 - a_shape: Optional[GroupShape]

72 - w_shape: Optional[GroupShape]

71 + a_shape: GroupShape | None

72 + w_shape: GroupShape | None

73 if block_shape is not None:

74 assert not per_act_token_quant

75 assert not per_out_ch_quant

@@ -100,7 +100,7 @@ class FusedMoEQuantDesc:

100 # The quantized type of this parameters. None means unquantized or

101 # already quantized.

102 # TODO (bnell): use scalar_type instead of Union.

103 - dtype: Union[torch.dtype, str, None] = None

103 + dtype: torch.dtype | str | None = None

104

105 # A field that describes the quantization group shape, from quant_utils.py.

106 # * (-1, -1) for per-tensor quantization

@@ -109,7 +109,7 @@ class FusedMoEQuantDesc:

109 # * (128, 128) for 128x128 deepseek style block quantization

110 # * (1, 128) for deepseek style activation quantization

111 # (i.e. per-token-per-group)

112 - shape: Optional[GroupShape] = None

112 + shape: GroupShape | None = None

113

114 # Quantization scales.

115 # TODO(bnell): maybe put PrecisionConfigs in subclass of QuantDesc?

@@ -117,13 +117,13 @@ class FusedMoEQuantDesc:

117

118 # Quantization alphas or gscales, used for nvfp4 types.

119 # TODO(bnell): put some of these in subclasses

120 - alpha_or_gscale: Optional[torch.Tensor] = None

120 + alpha_or_gscale: torch.Tensor | None = None

121

122 # Zero points for int4/int8 types

123 - zp: Optional[torch.Tensor] = None

123 + zp: torch.Tensor | None = None

124

125 # Biases for GPT triton MoE

126 - bias: Optional[torch.Tensor] = None

126 + bias: torch.Tensor | None = None

127

128

129 # TODO(bnell): have subclasses for specific moe methods?

@@ -179,7 +179,7 @@ class FusedMoEQuantConfig:

179 #

180

181 @property

182 - def quant_dtype(self) -> Union[torch.dtype, str, None]:

182 + def quant_dtype(self) -> torch.dtype | str | None:

183 return self._a1.dtype

184

185 @property

@@ -203,7 +203,7 @@ class FusedMoEQuantConfig:

203 return self._a1.shape == GroupShape.PER_TENSOR

204

205 @property

206 - def block_shape(self) -> Optional[list[int]]:

206 + def block_shape(self) -> list[int] | None:

207 if (

208 self._a1.shape is not None

209 and self._a1.shape != GroupShape.PER_TENSOR

@@ -218,34 +218,34 @@ class FusedMoEQuantConfig:

218 return self.block_shape is not None

219

220 @property

221 - def a1_scale(self) -> Optional[torch.Tensor]:

221 + def a1_scale(self) -> torch.Tensor | None:

222 assert self._a1.scale is None or isinstance(self._a1.scale, torch.Tensor)

223 return self._a1.scale

224

225 @property

226 - def a1_gscale(self) -> Optional[torch.Tensor]:

226 + def a1_gscale(self) -> torch.Tensor | None:

227 return self._a1.alpha_or_gscale

228

229 @property

230 - def a2_scale(self) -> Optional[torch.Tensor]:

230 + def a2_scale(self) -> torch.Tensor | None:

231 assert self._a2.scale is None or isinstance(self._a2.scale, torch.Tensor)

232 return self._a2.scale

233

234 @property

235 - def a2_gscale(self) -> Optional[torch.Tensor]:

235 + def a2_gscale(self) -> torch.Tensor | None:

236 return self._a2.alpha_or_gscale

237

238 @property

239 - def w1_scale(self) -> Optional[torch.Tensor]:

239 + def w1_scale(self) -> torch.Tensor | None:

240 assert self._w1.scale is None or isinstance(self._w1.scale, torch.Tensor)

241 return self._w1.scale

242

243 @property

244 - def w1_zp(self) -> Optional[torch.Tensor]:

244 + def w1_zp(self) -> torch.Tensor | None:

245 return self._w1.zp

246

247 @property

248 - def w1_bias(self) -> Optional[torch.Tensor]:

248 + def w1_bias(self) -> torch.Tensor | None:

249 return self._w1.bias

250

251 @property

@@ -254,20 +254,20 @@ class FusedMoEQuantConfig:

254 return self._w1.scale

255

256 @property

257 - def g1_alphas(self) -> Optional[torch.Tensor]:

257 + def g1_alphas(self) -> torch.Tensor | None:

258 return self._w1.alpha_or_gscale

259

260 @property

261 - def w2_scale(self) -> Optional[torch.Tensor]:

261 + def w2_scale(self) -> torch.Tensor | None:

262 assert self._w2.scale is None or isinstance(self._w2.scale, torch.Tensor)

263 return self._w2.scale

264

265 @property

266 - def w2_zp(self) -> Optional[torch.Tensor]:

266 + def w2_zp(self) -> torch.Tensor | None:

267 return self._w2.zp

268

269 @property

270 - def w2_bias(self) -> Optional[torch.Tensor]:

270 + def w2_bias(self) -> torch.Tensor | None:

271 return self._w2.bias

272

273 @property

@@ -276,7 +276,7 @@ class FusedMoEQuantConfig:

276 return self._w2.scale

277

278 @property

279 - def g2_alphas(self) -> Optional[torch.Tensor]:

279 + def g2_alphas(self) -> torch.Tensor | None:

280 return self._w2.alpha_or_gscale

281

282 @property

@@ -296,7 +296,7 @@ class FusedMoEQuantConfig:

296 return self._a1.dtype is None and self._w1.dtype == "int4"

297

298 @property

299 - def ocp_mx_scheme(self) -> Union[str, None]:

299 + def ocp_mx_scheme(self) -> str | None:

300 if not hasattr(self, "_ocp_mx_scheme"):

301 if (self._a1.dtype is not None and not isinstance(self._a1.dtype, str)) or (

302 self._w1.dtype is not None and not isinstance(self._w1.dtype, str)

@@ -322,7 +322,7 @@ class FusedMoEQuantConfig:

322 def use_nvfp4_w4a4(self) -> bool:

323 return self.quant_dtype == "nvfp4"

324

325 - def config_name(self, dtype: torch.dtype) -> Optional[str]:

325 + def config_name(self, dtype: torch.dtype) -> str | None:

326 """

327 Return a string used to construct the filename that contains the

328 tuning info for a particular quantization scheme. See

@@ -340,7 +340,7 @@ class FusedMoEQuantConfig:

340 self,

341 max_tokens: int,

342 hidden_dim: int,

343 - ) -> Optional[tuple[int, int]]:

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

344 """

345 Construct the proper activation scale shape for this

346 config.

@@ -363,7 +363,7 @@ class FusedMoEQuantConfig:

363 num_experts: int,

364 max_tokens: int,

365 hidden_dim: int,

366 - ) -> Optional[tuple[int, int, int]]:

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

367 """

368 Construct the proper activation batched scale shape for this

369 config, e.g. (num experts, *scale_shape).

@@ -377,23 +377,23 @@ class FusedMoEQuantConfig:

377

378 @staticmethod

379 def make(

380 - quant_dtype: Union[torch.dtype, str, None] = None,

380 + quant_dtype: torch.dtype | str | None = None,

381 per_act_token_quant: bool = False,

382 per_out_ch_quant: bool = False,

383 - block_shape: Optional[list[int]] = None,

383 + block_shape: list[int] | None = None,

384 w1_scale: Union[torch.Tensor, "PrecisionConfig", None] = None,

385 w2_scale: Union[torch.Tensor, "PrecisionConfig", None] = None,

386 - a1_scale: Optional[torch.Tensor] = None,

387 - a2_scale: Optional[torch.Tensor] = None,

388 - g1_alphas: Optional[torch.Tensor] = None,

389 - g2_alphas: Optional[torch.Tensor] = None,

390 - a1_gscale: Optional[torch.Tensor] = None,

391 - a2_gscale: Optional[torch.Tensor] = None,

392 - w1_bias: Optional[torch.Tensor] = None,

393 - w2_bias: Optional[torch.Tensor] = None,

394 - w1_zp: Optional[torch.Tensor] = None,

395 - w2_zp: Optional[torch.Tensor] = None,

396 - weight_dtype: Union[torch.dtype, str, None] = None,

386 + a1_scale: torch.Tensor | None = None,

387 + a2_scale: torch.Tensor | None = None,

388 + g1_alphas: torch.Tensor | None = None,

389 + g2_alphas: torch.Tensor | None = None,

390 + a1_gscale: torch.Tensor | None = None,

391 + a2_gscale: torch.Tensor | None = None,

392 + w1_bias: torch.Tensor | None = None,

393 + w2_bias: torch.Tensor | None = None,

394 + w1_zp: torch.Tensor | None = None,

395 + w2_zp: torch.Tensor | None = None,

396 + weight_dtype: torch.dtype | str | None = None,

397 ) -> "FusedMoEQuantConfig":

398 """

399 General builder function for a FusedMoEQuantConfig.

@@ -457,11 +457,11 @@ class FusedMoEQuantConfig:

457 def fp8_w8a8_moe_quant_config(

458 w1_scale: torch.Tensor,

459 w2_scale: torch.Tensor,

460 - a1_scale: Optional[torch.Tensor] = None,

461 - a2_scale: Optional[torch.Tensor] = None,

460 + a1_scale: torch.Tensor | None = None,

461 + a2_scale: torch.Tensor | None = None,

462 per_act_token_quant: bool = False,

463 per_out_ch_quant: bool = False,

464 - block_shape: Optional[list[int]] = None,

464 + block_shape: list[int] | None = None,

465 ) -> FusedMoEQuantConfig:

466 """

467 Construct a quant config for fp8 activations and fp8 weights.

@@ -481,8 +481,8 @@ def fp8_w8a8_moe_quant_config(

481 def int8_w8a8_moe_quant_config(

482 w1_scale: torch.Tensor,

483 w2_scale: torch.Tensor,

484 - a1_scale: Optional[torch.Tensor],

485 - a2_scale: Optional[torch.Tensor],

484 + a1_scale: torch.Tensor | None,

485 + a2_scale: torch.Tensor | None,

486 per_act_token_quant: bool = False,

487 ) -> FusedMoEQuantConfig:

488 """

@@ -503,8 +503,8 @@ def int8_w8a8_moe_quant_config(

503 def mxfp4_w4a16_moe_quant_config(

504 w1_scale: Union[torch.Tensor, "PrecisionConfig"],

505 w2_scale: Union[torch.Tensor, "PrecisionConfig"],

506 - w1_bias: Optional[torch.Tensor] = None,

507 - w2_bias: Optional[torch.Tensor] = None,

506 + w1_bias: torch.Tensor | None = None,

507 + w2_bias: torch.Tensor | None = None,

508 ) -> FusedMoEQuantConfig:

509 """

510 Construct a quant config for unquantized activations and mxfp4 weights.

@@ -521,12 +521,12 @@ def ocp_mx_moe_quant_config(

521 quant_dtype: str,

522 w1_scale: Union[torch.Tensor, "PrecisionConfig"],

523 w2_scale: Union[torch.Tensor, "PrecisionConfig"],

524 - weight_dtype: Optional[str] = None,

525 - a1_scale: Optional[torch.Tensor] = None,

526 - a2_scale: Optional[torch.Tensor] = None,

527 - w1_bias: Optional[torch.Tensor] = None,

528 - w2_bias: Optional[torch.Tensor] = None,

529 - block_shape: Optional[list[int]] = None,

524 + weight_dtype: str | None = None,

525 + a1_scale: torch.Tensor | None = None,

526 + a2_scale: torch.Tensor | None = None,

527 + w1_bias: torch.Tensor | None = None,

528 + w2_bias: torch.Tensor | None = None,

529 + block_shape: list[int] | None = None,

530 ) -> FusedMoEQuantConfig:

531 """

532 Construct a quant config for mxfp4 activations and mxfp4 weights.

@@ -575,9 +575,9 @@ def nvfp4_moe_quant_config(

575 def int4_w4a16_moe_quant_config(

576 w1_scale: torch.Tensor,

577 w2_scale: torch.Tensor,

578 - w1_zp: Optional[torch.Tensor],

579 - w2_zp: Optional[torch.Tensor],

580 - block_shape: Optional[list[int]] = None,

578 + w1_zp: torch.Tensor | None,

579 + w2_zp: torch.Tensor | None,

580 + block_shape: list[int] | None = None,

581 ) -> FusedMoEQuantConfig:

582 """

583 Construct a quant config for 16-bit float activations and int4 weights.

@@ -595,9 +595,9 @@ def int4_w4a16_moe_quant_config(

595 def int8_w8a16_moe_quant_config(

596 w1_scale: torch.Tensor,

597 w2_scale: torch.Tensor,

598 - w1_zp: Optional[torch.Tensor],

599 - w2_zp: Optional[torch.Tensor],

600 - block_shape: Optional[list[int]] = None,

598 + w1_zp: torch.Tensor | None,

599 + w2_zp: torch.Tensor | None,

600 + block_shape: list[int] | None = None,

601 ) -> FusedMoEQuantConfig:

602 """

603 Construct a quant config for 16-bit float activations and int8 weights.

@@ -613,8 +613,8 @@ def int8_w8a16_moe_quant_config(

613

614

615 def biased_moe_quant_config(

616 - w1_bias: Optional[torch.Tensor],

617 - w2_bias: Optional[torch.Tensor],

616 + w1_bias: torch.Tensor | None,

617 + w2_bias: torch.Tensor | None,

618 ) -> FusedMoEQuantConfig:

619 """

620 Construct a quant config for unquantized activations with biases.

vllm/model_executor/layers/fused_moe/cpu_fused_moe.py

+21 -21

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Callable, Optional

3 +from collections.abc import Callable

4

5 import torch

6 from torch.nn import functional as F

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

33 topk_group: int = 0,

34 scoring_func: str = "softmax",

35 routed_scaling_factor: float = 1.0,

36 - e_score_correction_bias: Optional[torch.Tensor] = None,

36 + e_score_correction_bias: torch.Tensor | None = None,

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

38 assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"

39

@@ -88,12 +88,12 @@ def select_experts(

88 top_k: int,

89 use_grouped_topk: bool,

90 renormalize: bool,

91 - topk_group: Optional[int] = None,

92 - num_expert_group: Optional[int] = None,

93 - custom_routing_function: Optional[Callable] = None,

91 + topk_group: int | None = None,

92 + num_expert_group: int | None = None,

93 + custom_routing_function: Callable | None = None,

94 scoring_func: str = "softmax",

95 routed_scaling_factor: float = 1.0,

96 - e_score_correction_bias: Optional[torch.Tensor] = None,

96 + e_score_correction_bias: torch.Tensor | None = None,

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

98 if use_grouped_topk:

99 assert topk_group is not None

@@ -147,14 +147,14 @@ class IPEXFusedMOE:

147 top_k: int,

148 router_logits: torch.Tensor,

149 renormalize: bool,

150 - topk_group: Optional[int] = None,

151 - num_expert_group: Optional[int] = None,

150 + topk_group: int | None = None,

151 + num_expert_group: int | None = None,

152 global_num_experts: int = -1,

153 - expert_map: Optional[torch.Tensor] = None,

154 - custom_routing_function: Optional[Callable] = None,

153 + expert_map: torch.Tensor | None = None,

154 + custom_routing_function: Callable | None = None,

155 scoring_func: str = "softmax",

156 routed_scaling_factor: float = 1.0,

157 - e_score_correction_bias: Optional[torch.Tensor] = None,

157 + e_score_correction_bias: torch.Tensor | None = None,

158 apply_router_weight_on_input: bool = False,

159 activation: str = "silu",

160 ) -> torch.Tensor:

@@ -189,14 +189,14 @@ class SGLFusedMOE:

189 top_k: int,

190 router_logits: torch.Tensor,

191 renormalize: bool,

192 - topk_group: Optional[int] = None,

193 - num_expert_group: Optional[int] = None,

192 + topk_group: int | None = None,

193 + num_expert_group: int | None = None,

194 global_num_experts: int = -1,

195 - expert_map: Optional[torch.Tensor] = None,

196 - custom_routing_function: Optional[Callable] = None,

195 + expert_map: torch.Tensor | None = None,

196 + custom_routing_function: Callable | None = None,

197 scoring_func: str = "softmax",

198 routed_scaling_factor: float = 1.0,

199 - e_score_correction_bias: Optional[torch.Tensor] = None,

199 + e_score_correction_bias: torch.Tensor | None = None,

200 apply_router_weight_on_input: bool = False,

201 activation: str = "silu",

202 ) -> torch.Tensor:

@@ -247,14 +247,14 @@ class CPUFusedMOE:

247 top_k: int,

248 router_logits: torch.Tensor,

249 renormalize: bool,

250 - topk_group: Optional[int] = None,

251 - num_expert_group: Optional[int] = None,

250 + topk_group: int | None = None,

251 + num_expert_group: int | None = None,

252 global_num_experts: int = -1,

253 - expert_map: Optional[torch.Tensor] = None,

254 - custom_routing_function: Optional[Callable] = None,

253 + expert_map: torch.Tensor | None = None,

254 + custom_routing_function: Callable | None = None,

255 scoring_func: str = "softmax",

256 routed_scaling_factor: float = 1.0,

257 - e_score_correction_bias: Optional[torch.Tensor] = None,

257 + e_score_correction_bias: torch.Tensor | None = None,

258 apply_router_weight_on_input: bool = False,

259 activation: str = "silu",

260 ) -> torch.Tensor:

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+27 -27

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

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

3 """CUTLASS based Fused MoE kernels."""

4

5 -from typing import Callable, Optional

5 +from collections.abc import Callable

6

7 import torch

8

@@ -35,23 +35,23 @@ def run_cutlass_moe_fp8(

35 topk_ids: torch.Tensor,

36 activation_callable: Callable,

37 global_num_experts: int,

38 - expert_map: Optional[torch.Tensor],

39 - w1_scale: Optional[torch.Tensor],

40 - w2_scale: Optional[torch.Tensor],

41 - a1q_scale: Optional[torch.Tensor],

42 - a2_scale: Optional[torch.Tensor],

38 + expert_map: torch.Tensor | None,

39 + w1_scale: torch.Tensor | None,

40 + w2_scale: torch.Tensor | None,

41 + a1q_scale: torch.Tensor | None,

42 + a2_scale: torch.Tensor | None,

43 ab_strides1: torch.Tensor,

44 ab_strides2: torch.Tensor,

45 c_strides1: torch.Tensor,

46 c_strides2: torch.Tensor,

47 workspace13: torch.Tensor,

48 workspace2: torch.Tensor,

49 - expert_num_tokens: Optional[torch.Tensor],

49 + expert_num_tokens: torch.Tensor | None,

50 out_dtype: torch.dtype,

51 per_act_token: bool,

52 per_out_ch: bool,

53 use_batched_format: bool,

54 - topk_weights: Optional[torch.Tensor],

54 + topk_weights: torch.Tensor | None,

55 ):

56 a1q = hidden_states

57

@@ -249,7 +249,7 @@ def run_cutlass_moe_fp8(

249 class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

250 def __init__(

251 self,

252 - out_dtype: Optional[torch.dtype],

252 + out_dtype: torch.dtype | None,

253 ab_strides1: torch.Tensor,

254 ab_strides2: torch.Tensor,

255 c_strides1: torch.Tensor,

@@ -278,12 +278,12 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

278 topk_ids: torch.Tensor,

279 activation: str,

280 global_num_experts: int,

281 - expert_map: Optional[torch.Tensor],

282 - a1q_scale: Optional[torch.Tensor],

283 - a2_scale: Optional[torch.Tensor],

281 + expert_map: torch.Tensor | None,

282 + a1q_scale: torch.Tensor | None,

283 + a2_scale: torch.Tensor | None,

284 workspace13: torch.Tensor,

285 workspace2: torch.Tensor,

286 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

286 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

287 apply_router_weight_on_input: bool,

288 ):

289 assert self.w1_zp is None, "w1_zp is not supported in CUTLASS MoE"

@@ -331,7 +331,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

331 class CutlassExpertsFp8(CutlassExpertsFp8Base):

332 def __init__(

333 self,

334 - out_dtype: Optional[torch.dtype],

334 + out_dtype: torch.dtype | None,

335 ab_strides1: torch.Tensor,

336 ab_strides2: torch.Tensor,

337 c_strides1: torch.Tensor,

@@ -377,7 +377,7 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):

377 topk: int,

378 global_num_experts: int,

379 local_num_experts: int,

380 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

380 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

381 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

382 workspace1 = (M * topk, max(N, K))

383 workspace2 = (M * topk, max(N // 2, K))

@@ -390,7 +390,7 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

390 self,

391 max_experts_per_worker: int,

392 num_dispatchers: int,

393 - out_dtype: Optional[torch.dtype],

393 + out_dtype: torch.dtype | None,

394 ab_strides1: torch.Tensor,

395 ab_strides2: torch.Tensor,

396 c_strides1: torch.Tensor,

@@ -435,7 +435,7 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

435 topk: int,

436 global_num_experts: int,

437 local_num_experts: int,

438 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

438 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

439 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

440 num_dp = self.num_dispatchers

441 assert num_dp is not None

@@ -457,7 +457,7 @@ def cutlass_moe_fp8(

457 c_strides2: torch.Tensor,

458 quant_config: FusedMoEQuantConfig,

459 activation: str = "silu",

460 - expert_map: Optional[torch.Tensor] = None,

460 + expert_map: torch.Tensor | None = None,

461 apply_router_weight_on_input: bool = False,

462 global_num_experts: int = -1,

463 ) -> torch.Tensor:

@@ -768,7 +768,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

768 topk: int,

769 global_num_experts: int,

770 local_num_experts: int,

771 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

771 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

772 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

773 workspace1: tuple[int, ...] = ()

774 workspace2: tuple[int, ...] = ()

@@ -793,12 +793,12 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

793 topk_ids: torch.Tensor,

794 activation: str,

795 global_num_experts: int,

796 - expert_map: Optional[torch.Tensor],

797 - a1q_scale: Optional[torch.Tensor], # unused

798 - a2_scale: Optional[torch.Tensor], # unused

799 - workspace13: Optional[torch.Tensor],

800 - workspace2: Optional[torch.Tensor],

801 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

796 + expert_map: torch.Tensor | None,

797 + a1q_scale: torch.Tensor | None, # unused

798 + a2_scale: torch.Tensor | None, # unused

799 + workspace13: torch.Tensor | None,

800 + workspace2: torch.Tensor | None,

801 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

802 apply_router_weight_on_input: bool,

803 ):

804 e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids)

@@ -839,7 +839,7 @@ def cutlass_moe_fp4(

839 n: int,

840 k: int,

841 e: int,

842 - expert_map: Optional[torch.Tensor] = None,

842 + expert_map: torch.Tensor | None = None,

843 apply_router_weight_on_input: bool = False,

844 ) -> torch.Tensor:

845 assert expert_map is None, (

@@ -896,7 +896,7 @@ def _valid_cutlass_block_scaled_grouped_gemm(

896 inplace: bool,

897 activation: str,

898 apply_router_weight_on_input: bool,

899 - expert_map: Optional[torch.Tensor],

899 + expert_map: torch.Tensor | None,

900 ) -> bool:

901 def _valid_cutlass_block_scaled_grouped_gemm_shape(N: int, K: int):

902 return N % 128 == 0 and K % 128 == 0

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+9 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5 from tqdm import tqdm

@@ -204,7 +203,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

203 topk: int,

204 global_num_experts: int,

205 local_num_experts: int,

207 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

206 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

207 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

208 assert self.block_shape is not None

209 block_m = self.block_shape[0]

@@ -228,12 +227,12 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

227 topk_ids: torch.Tensor,

228 activation: str,

229 global_num_experts: int,

231 - expert_map: Optional[torch.Tensor],

232 - a1q_scale: Optional[torch.Tensor],

233 - a2_scale: Optional[torch.Tensor],

230 + expert_map: torch.Tensor | None,

231 + a1q_scale: torch.Tensor | None,

232 + a2_scale: torch.Tensor | None,

233 workspace13: torch.Tensor,

234 workspace2: torch.Tensor,

236 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

235 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

236 apply_router_weight_on_input: bool,

237 ):

238 assert a1q_scale is not None

@@ -284,7 +283,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

283

284 self.activation(activation, act_out, mm1_out.view(-1, N))

285

287 - a2q_scale: Optional[torch.Tensor] = None

286 + a2q_scale: torch.Tensor | None = None

287 a2q, a2q_scale = per_token_group_quant_fp8(

288 act_out, self.block_shape[1], column_major_scales=True, out_q=quant_out

289 )

@@ -317,9 +316,9 @@ def deep_gemm_moe_fp8(

316 inplace: bool = False,

317 activation: str = "silu",

318 global_num_experts: int = -1,

320 - expert_map: Optional[torch.Tensor] = None,

321 - a1_scale: Optional[torch.Tensor] = None,

322 - a2_scale: Optional[torch.Tensor] = None,

319 + expert_map: torch.Tensor | None = None,

320 + a1_scale: torch.Tensor | None = None,

321 + a2_scale: torch.Tensor | None = None,

322 apply_router_weight_on_input=False,

323 ) -> torch.Tensor:

324 """

vllm/model_executor/layers/fused_moe/deep_gemm_utils.py

+7 -8

@@ -6,7 +6,6 @@ and updated to fit vllm needs and terminology.

6 """

7

8 import functools

9 -from typing import Optional

9

10 import torch

11

@@ -39,7 +38,7 @@ def compute_aligned_M(

38 num_topk: int,

39 local_num_experts: int,

40 alignment: int,

42 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

41 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

42 ):

43 if (expert_tokens_meta is not None) and (

44 expert_tokens_meta.expert_num_tokens_cpu is not None

@@ -175,7 +174,7 @@ def ep_scatter(

174 recv_x_scale: torch.Tensor,

175 recv_topk: torch.Tensor,

176 num_recv_tokens_per_expert: torch.Tensor,

178 - expert_map: Optional[torch.Tensor],

177 + expert_map: torch.Tensor | None,

178 expert_start_loc: torch.Tensor,

179 output_tensor: torch.Tensor,

180 output_tensor_scale: torch.Tensor,

@@ -305,7 +304,7 @@ def ep_gather(

304 recv_topk_ids: torch.Tensor,

305 recv_topk_weight: torch.Tensor,

306 input_index: torch.Tensor,

308 - expert_map: Optional[torch.Tensor],

307 + expert_map: torch.Tensor | None,

308 output_tensor: torch.Tensor,

309 ):

310 num_warps = 2

@@ -346,9 +345,9 @@ def deepgemm_moe_permute(

345 aq_scale: torch.Tensor,

346 topk_ids: torch.Tensor,

347 local_num_experts: int,

349 - expert_map: Optional[torch.Tensor],

350 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

351 - aq_out: Optional[torch.Tensor] = None,

348 + expert_map: torch.Tensor | None,

349 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

350 + aq_out: torch.Tensor | None = None,

351 ):

352 assert aq.ndim == 2

353 assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"

@@ -415,7 +414,7 @@ def deepgemm_unpermute_and_reduce(

414 topk_ids: torch.Tensor,

415 topk_weights: torch.Tensor,

416 inv_perm: torch.Tensor,

418 - expert_map: Optional[torch.Tensor],

417 + expert_map: torch.Tensor | None,

418 output: torch.Tensor,

419 ):

420 return ep_gather(

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

+14 -14

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Callable, Optional, Union

3 +from collections.abc import Callable

4

5 import deep_ep

6 import torch

@@ -77,18 +77,18 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

77 def activation_format(self) -> mk.FusedMoEActivationFormat:

78 return mk.FusedMoEActivationFormat.Standard

79

80 - def max_num_tokens_per_rank(self) -> Optional[int]:

80 + def max_num_tokens_per_rank(self) -> int | None:

81 return None

82

83 - def topk_indices_dtype(self) -> Optional[torch.dtype]:

83 + def topk_indices_dtype(self) -> torch.dtype | None:

84 return torch.int64

85

86 - def _get_dispatch_config(self) -> Optional[deep_ep.Config]:

86 + def _get_dispatch_config(self) -> deep_ep.Config | None:

87 if self.num_dispatchers_ not in self.available_rank_configs:

88 return None

89 return deep_ep.Buffer.get_dispatch_config(self.num_dispatchers_)

90

91 - def _get_combine_config(self) -> Optional[deep_ep.Config]:

91 + def _get_combine_config(self) -> deep_ep.Config | None:

92 if self.num_dispatchers_ not in self.available_rank_configs:

93 return None

94 return deep_ep.Buffer.get_combine_config(self.num_dispatchers_)

@@ -96,11 +96,11 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

96 def _do_dispatch(

97 self,

98 tokens: torch.Tensor,

99 - token_scales: Optional[torch.Tensor],

99 + token_scales: torch.Tensor | None,

100 rank_topk_ids: torch.Tensor,

101 rank_topk_weights: torch.Tensor,

102 num_experts: int,

103 - a1_scale: Optional[torch.Tensor],

103 + a1_scale: torch.Tensor | None,

104 quant_config: FusedMoEQuantConfig,

105 ) -> Callable:

106 has_scales = token_scales is not None

@@ -175,12 +175,12 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

175 self,

176 event: deep_ep.EventOverlap,

177 has_scales: bool,

178 - token_data: Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor],

179 - expert_topk_ids: Optional[torch.Tensor],

178 + token_data: tuple[torch.Tensor, torch.Tensor] | torch.Tensor,

179 + expert_topk_ids: torch.Tensor | None,

180 num_experts: int,

181 expert_num_tokens_per_expert_list: list[int],

182 - expert_topk_weights: Optional[torch.Tensor],

183 - a1_scale: Optional[torch.Tensor],

182 + expert_topk_weights: torch.Tensor | None,

183 + a1_scale: torch.Tensor | None,

184 quant_config: FusedMoEQuantConfig,

185 ) -> mk.PrepareResultType:

186 if event.event is not None:

@@ -249,7 +249,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

249 topk_weights: torch.Tensor,

250 topk_ids: torch.Tensor,

251 num_experts: int,

252 - expert_map: Optional[torch.Tensor],

252 + expert_map: torch.Tensor | None,

253 apply_router_weight_on_input: bool,

254 quant_config: FusedMoEQuantConfig,

255 ) -> mk.ReceiverType:

@@ -294,7 +294,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

294 topk_weights: torch.Tensor,

295 topk_ids: torch.Tensor,

296 num_experts: int,

297 - expert_map: Optional[torch.Tensor],

297 + expert_map: torch.Tensor | None,

298 apply_router_weight_on_input: bool,

299 quant_config: FusedMoEQuantConfig,

300 ) -> mk.PrepareResultType:

@@ -318,7 +318,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

318 apply_router_weight_on_input: bool,

319 weight_and_reduce_impl: mk.TopKWeightAndReduce,

320 do_async: bool,

321 - ) -> Optional[Callable]:

321 + ) -> Callable | None:

322 a2a_idx = dbo_current_ubatch_id()

323 handle = self.handles[a2a_idx]

324 assert handle is not None

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+10 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Callable, Optional, Union

3 +from collections.abc import Callable

4

5 import deep_ep

6 import torch

@@ -67,7 +67,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

67 # The dispatch function returns a handle that the combine function

68 # requires. We store the handle here so it is available to the

69 # combine function.

70 - self.handles: list[Optional[tuple]] = [None, None]

70 + self.handles: list[tuple | None] = [None, None]

71 self.num_dispatchers_ = num_dispatchers

72

73 def num_dispatchers(self) -> int:

@@ -80,18 +80,18 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

80 def activation_format(self) -> mk.FusedMoEActivationFormat:

81 return mk.FusedMoEActivationFormat.BatchedExperts

82

83 - def max_num_tokens_per_rank(self) -> Optional[int]:

83 + def max_num_tokens_per_rank(self) -> int | None:

84 return self.max_tokens_per_rank

85

86 - def topk_indices_dtype(self) -> Optional[torch.dtype]:

86 + def topk_indices_dtype(self) -> torch.dtype | None:

87 return torch.int64

88

89 def _do_quant(

90 self,

91 - x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

91 + x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

92 a1_dtype: torch.dtype,

93 quant_config: FusedMoEQuantConfig,

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

94 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

95 if self.use_fp8_dispatch:

96 block_k = (

97 quant_config.block_shape[1]

@@ -137,7 +137,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

137 topk_weights: torch.Tensor,

138 topk_ids: torch.Tensor,

139 num_experts: int,

140 - expert_map: Optional[torch.Tensor],

140 + expert_map: torch.Tensor | None,

141 apply_router_weight_on_input: bool,

142 quant_config: FusedMoEQuantConfig,

143 ) -> tuple[Callable, mk.ReceiverType]:

@@ -200,9 +200,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

200

201 def _receiver(

202 self,

203 - expert_x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

203 + expert_x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

204 expert_num_tokens: torch.Tensor,

205 - a1_scale: Optional[torch.Tensor],

205 + a1_scale: torch.Tensor | None,

206 a1_dtype: torch.dtype,

207 quant_config: FusedMoEQuantConfig,

208 ) -> mk.PrepareResultType:

@@ -220,7 +220,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

220 topk_weights: torch.Tensor,

221 topk_ids: torch.Tensor,

222 num_experts: int,

223 - expert_map: Optional[torch.Tensor],

223 + expert_map: torch.Tensor | None,

224 apply_router_weight_on_input: bool,

225 quant_config: FusedMoEQuantConfig,

226 ) -> mk.PrepareResultType:

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+10 -11

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -96,7 +95,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

95 topk: int,

96 global_num_experts: int,

97 local_num_experts: int,

99 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

98 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

99 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

100 # We use global_num_experts due to how moe_align_block_size handles

101 # expert_maps.

@@ -133,13 +132,13 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

132 topk_ids: torch.Tensor,

133 activation: str,

134 global_num_experts: int,

136 - expert_map: Optional[torch.Tensor],

137 - a1q_scale: Optional[torch.Tensor],

138 - a2_scale: Optional[torch.Tensor],

139 - workspace13: Optional[torch.Tensor],

140 - workspace2: Optional[torch.Tensor],

141 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

142 - apply_router_weight_on_input: Optional[bool],

135 + expert_map: torch.Tensor | None,

136 + a1q_scale: torch.Tensor | None,

137 + a2_scale: torch.Tensor | None,

138 + workspace13: torch.Tensor | None,

139 + workspace2: torch.Tensor | None,

140 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

141 + apply_router_weight_on_input: bool | None,

142 ):

143 assert activation == "silu", (

144 "Only activation silu is supported in FlashInferExperts"

@@ -207,7 +206,7 @@ def flashinfer_cutlass_moe_fp4(

206 inplace: bool = False,

207 activation: str = "silu",

208 global_num_experts: int = -1,

210 - expert_map: Optional[torch.Tensor] = None,

209 + expert_map: torch.Tensor | None = None,

210 apply_router_weight_on_input: bool = False,

211 ) -> torch.Tensor:

212 fused_experts = mk.FusedMoEModularKernel(

@@ -242,7 +241,7 @@ def flashinfer_cutlass_moe(

241 inplace: bool = False,

242 activation: str = "silu",

243 global_num_experts: int = -1,

245 - expert_map: Optional[torch.Tensor] = None,

244 + expert_map: torch.Tensor | None = None,

245 apply_router_weight_on_input: bool = False,

246 tp_rank: int = 0,

247 tp_size: int = 1,

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+4 -5

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -39,10 +38,10 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

38 def activation_format(self) -> mk.FusedMoEActivationFormat:

39 return mk.FusedMoEActivationFormat.Standard

40

42 - def max_num_tokens_per_rank(self) -> Optional[int]:

41 + def max_num_tokens_per_rank(self) -> int | None:

42 return None

43

45 - def topk_indices_dtype(self) -> Optional[torch.dtype]:

44 + def topk_indices_dtype(self) -> torch.dtype | None:

45 return None

46

47 def num_dispatchers(self) -> int:

@@ -89,7 +88,7 @@ class FlashInferAllToAllMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFina

88 topk_weights: torch.Tensor,

89 topk_ids: torch.Tensor,

90 num_experts: int,

92 - expert_map: Optional[torch.Tensor],

91 + expert_map: torch.Tensor | None,

92 apply_router_weight_on_input: bool,

93 quant_config: FusedMoEQuantConfig,

94 ) -> mk.PrepareResultType:

@@ -164,7 +163,7 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin

163 topk_weights: torch.Tensor,

164 topk_ids: torch.Tensor,

165 num_experts: int,

167 - expert_map: Optional[torch.Tensor],

166 + expert_map: torch.Tensor | None,

167 apply_router_weight_on_input: bool,

168 quant_config: FusedMoEQuantConfig,

169 ) -> mk.PrepareResultType:

vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py

+6 -7

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -105,7 +104,7 @@ direct_register_custom_op(

104

105 def flashinfer_fused_moe_per_tensor_scale_fp8(

106 routing_logits: torch.Tensor,

108 - routing_bias: Optional[torch.Tensor],

107 + routing_bias: torch.Tensor | None,

108 hidden_states: torch.Tensor,

109 input_scale: torch.Tensor,

110 gemm1_weights: torch.Tensor,

@@ -115,8 +114,8 @@ def flashinfer_fused_moe_per_tensor_scale_fp8(

114 output2_scales_scalar: torch.Tensor,

115 num_experts: int,

116 top_k: int,

118 - num_expert_group: Optional[int],

119 - topk_group: Optional[int],

117 + num_expert_group: int | None,

118 + topk_group: int | None,

119 intermediate_size: int,

120 local_expert_offset: int,

121 local_num_experts: int,

@@ -163,7 +162,7 @@ def flashinfer_fused_moe_per_tensor_scale_fp8(

162

163 def flashinfer_fused_moe_per_tensor_scale_fp8_fake(

164 routing_logits: torch.Tensor,

166 - routing_bias: Optional[torch.Tensor],

165 + routing_bias: torch.Tensor | None,

166 hidden_states: torch.Tensor,

167 input_scale: torch.Tensor,

168 gemm1_weights: torch.Tensor,

@@ -173,8 +172,8 @@ def flashinfer_fused_moe_per_tensor_scale_fp8_fake(

172 output2_scales_scalar: torch.Tensor,

173 num_experts: int,

174 top_k: int,

176 - num_expert_group: Optional[int],

177 - topk_group: Optional[int],

175 + num_expert_group: int | None,

176 + topk_group: int | None,

177 intermediate_size: int,

178 local_expert_offset: int,

179 local_num_experts: int,

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+20 -22

@@ -2,8 +2,6 @@

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

3 """Fused batched MoE kernel."""

4

5 -from typing import Optional

6 -

5 import torch

6

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

@@ -370,8 +368,8 @@ def invoke_moe_batched_triton_kernel(

368 expert_num_tokens: torch.Tensor, # [E]

369 compute_type: tl.dtype,

370 # Quantization data

373 - A_scale: Optional[torch.Tensor],

374 - B_scale: Optional[torch.Tensor],

371 + A_scale: torch.Tensor | None,

372 + B_scale: torch.Tensor | None,

373 B_zp: torch.Tensor,

374 # Quantization schemes

375 use_fp8_w8a8: bool,

@@ -379,7 +377,7 @@ def invoke_moe_batched_triton_kernel(

377 use_int4_w4a16: bool,

378 config: dict[str, int],

379 per_act_token_quant: bool,

382 - block_shape: Optional[list[int]] = None,

380 + block_shape: list[int] | None = None,

381 ):

382 assert not use_int4_w4a16

383 max_num_tokens = A.size(1)

@@ -500,10 +498,10 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

498 def activation_format(self) -> mk.FusedMoEActivationFormat:

499 return mk.FusedMoEActivationFormat.BatchedExperts

500

503 - def max_num_tokens_per_rank(self) -> Optional[int]:

501 + def max_num_tokens_per_rank(self) -> int | None:

502 return self.max_num_tokens

503

506 - def topk_indices_dtype(self) -> Optional[torch.dtype]:

504 + def topk_indices_dtype(self) -> torch.dtype | None:

505 return None

506

507 def num_dispatchers(self) -> int:

@@ -518,7 +516,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

516 topk_weights: torch.Tensor,

517 topk_ids: torch.Tensor,

518 num_experts: int,

521 - expert_map: Optional[torch.Tensor],

519 + expert_map: torch.Tensor | None,

520 apply_router_weight_on_input: bool,

521 quant_config: FusedMoEQuantConfig,

522 ) -> mk.PrepareResultType:

@@ -674,7 +672,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

672 topk: int,

673 global_num_experts: int,

674 local_num_experts: int,

677 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

675 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

676 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

677 num_dp = self.num_dispatchers

678 num_experts = local_num_experts

@@ -701,12 +699,12 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

699 topk_ids: torch.Tensor,

700 activation: str,

701 global_num_experts: int,

704 - expert_map: Optional[torch.Tensor],

705 - a1q_scale: Optional[torch.Tensor],

706 - a2_scale: Optional[torch.Tensor],

702 + expert_map: torch.Tensor | None,

703 + a1q_scale: torch.Tensor | None,

704 + a2_scale: torch.Tensor | None,

705 workspace13: torch.Tensor,

706 workspace2: torch.Tensor,

709 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

707 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

708 apply_router_weight_on_input: bool,

709 ):

710 assert hidden_states.dim() == 3

@@ -754,15 +752,15 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

752

753 def batched_moe_kernel_quantize_input(

754 A: torch.Tensor,

757 - A_scale: Optional[torch.Tensor],

755 + A_scale: torch.Tensor | None,

756 num_tokens: int,

757 E: int,

758 N: int,

759 expert_num_tokens: torch.Tensor,

762 - qtype: Optional[torch.dtype],

760 + qtype: torch.dtype | None,

761 per_act_token_quant: bool,

764 - block_shape: Optional[list[int]] = None,

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

762 + block_shape: list[int] | None = None,

763 +) -> tuple[torch.Tensor, torch.Tensor | None]:

764 if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing():

765 # Note: this does a bunch of extra work because expert_num_tokens is

766 # ignored but it does support torch.compile + cudagraphs.

@@ -868,7 +866,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

866 topk: int,

867 global_num_experts: int,

868 local_num_experts: int,

871 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

869 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

870 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

871 num_dp = self.num_dispatchers

872 num_experts = local_num_experts

@@ -888,12 +886,12 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

886 topk_ids: torch.Tensor,

887 activation: str,

888 global_num_experts: int,

891 - expert_map: Optional[torch.Tensor],

892 - a1q_scale: Optional[torch.Tensor],

893 - a2_scale: Optional[torch.Tensor],

889 + expert_map: torch.Tensor | None,

890 + a1q_scale: torch.Tensor | None,

891 + a2_scale: torch.Tensor | None,

892 workspace13: torch.Tensor,

893 workspace2: torch.Tensor,

896 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

894 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

895 apply_router_weight_on_input: bool,

896 ):

897 # Check constraints.

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+36 -38

@@ -2,8 +2,6 @@

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

3 """Fused MoE utilities for GPTQ."""

4

5 -from typing import Optional

6 -

5 import torch

6 from typing_extensions import override

7

@@ -28,31 +26,31 @@ def fused_marlin_moe(

26 hidden_states: torch.Tensor,

27 w1: torch.Tensor,

28 w2: torch.Tensor,

31 - bias1: Optional[torch.Tensor],

32 - bias2: Optional[torch.Tensor],

29 + bias1: torch.Tensor | None,

30 + bias2: torch.Tensor | None,

31 w1_scale: torch.Tensor,

32 w2_scale: torch.Tensor,

35 - gating_output: Optional[torch.Tensor],

33 + gating_output: torch.Tensor | None,

34 topk_weights: torch.Tensor,

35 topk_ids: torch.Tensor,

36 quant_type_id: int,

37 apply_router_weight_on_input: bool = False,

38 global_num_experts: int = -1,

41 - activation: Optional[str] = "silu",

42 - expert_map: Optional[torch.Tensor] = None,

43 - global_scale1: Optional[torch.Tensor] = None,

44 - global_scale2: Optional[torch.Tensor] = None,

45 - g_idx1: Optional[torch.Tensor] = None,

46 - g_idx2: Optional[torch.Tensor] = None,

47 - sort_indices1: Optional[torch.Tensor] = None,

48 - sort_indices2: Optional[torch.Tensor] = None,

49 - w1_zeros: Optional[torch.Tensor] = None,

50 - w2_zeros: Optional[torch.Tensor] = None,

51 - workspace: Optional[torch.Tensor] = None,

52 - intermediate_cache13: Optional[torch.Tensor] = None,

53 - intermediate_cache2: Optional[torch.Tensor] = None,

39 + activation: str | None = "silu",

40 + expert_map: torch.Tensor | None = None,

41 + global_scale1: torch.Tensor | None = None,

42 + global_scale2: torch.Tensor | None = None,

43 + g_idx1: torch.Tensor | None = None,

44 + g_idx2: torch.Tensor | None = None,

45 + sort_indices1: torch.Tensor | None = None,

46 + sort_indices2: torch.Tensor | None = None,

47 + w1_zeros: torch.Tensor | None = None,

48 + w2_zeros: torch.Tensor | None = None,

49 + workspace: torch.Tensor | None = None,

50 + intermediate_cache13: torch.Tensor | None = None,

51 + intermediate_cache2: torch.Tensor | None = None,

52 is_k_full: bool = True,

55 - output: Optional[torch.Tensor] = None,

53 + output: torch.Tensor | None = None,

54 inplace: bool = False,

55 ) -> torch.Tensor:

56 """

@@ -249,26 +247,26 @@ def fused_marlin_moe_fake(

247 w2: torch.Tensor,

248 w1_scale: torch.Tensor,

249 w2_scale: torch.Tensor,

252 - gating_output: Optional[torch.Tensor],

250 + gating_output: torch.Tensor | None,

251 topk_weights: torch.Tensor,

252 topk_ids: torch.Tensor,

253 quant_type_id: int,

254 apply_router_weight_on_input: bool = False,

255 global_num_experts: int = -1,

258 - global_scale1: Optional[torch.Tensor] = None,

259 - global_scale2: Optional[torch.Tensor] = None,

260 - expert_map: Optional[torch.Tensor] = None,

261 - g_idx1: Optional[torch.Tensor] = None,

262 - g_idx2: Optional[torch.Tensor] = None,

263 - sort_indices1: Optional[torch.Tensor] = None,

264 - sort_indices2: Optional[torch.Tensor] = None,

265 - w1_zeros: Optional[torch.Tensor] = None,

266 - w2_zeros: Optional[torch.Tensor] = None,

267 - workspace: Optional[torch.Tensor] = None,

268 - intermediate_cache13: Optional[torch.Tensor] = None,

269 - intermediate_cache2: Optional[torch.Tensor] = None,

256 + global_scale1: torch.Tensor | None = None,

257 + global_scale2: torch.Tensor | None = None,

258 + expert_map: torch.Tensor | None = None,

259 + g_idx1: torch.Tensor | None = None,

260 + g_idx2: torch.Tensor | None = None,

261 + sort_indices1: torch.Tensor | None = None,

262 + sort_indices2: torch.Tensor | None = None,

263 + w1_zeros: torch.Tensor | None = None,

264 + w2_zeros: torch.Tensor | None = None,

265 + workspace: torch.Tensor | None = None,

266 + intermediate_cache13: torch.Tensor | None = None,

267 + intermediate_cache2: torch.Tensor | None = None,

268 is_k_full: bool = True,

271 - output: Optional[torch.Tensor] = None,

269 + output: torch.Tensor | None = None,

270 inplace: bool = False,

271 ) -> torch.Tensor:

272 return torch.empty_like(hidden_states)

@@ -341,7 +339,7 @@ class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):

339 topk: int,

340 global_num_experts: int,

341 local_num_experts: int,

344 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

342 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

343 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

344 # Modular Kernel provisions output buffer from workspace1. However in

345 # the fused_marlin_moe() function, the final torch.sum(), is defined

@@ -374,12 +372,12 @@ class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):

372 topk_ids: torch.Tensor,

373 activation: str,

374 global_num_experts: int,

377 - expert_map: Optional[torch.Tensor],

378 - a1q_scale: Optional[torch.Tensor],

379 - a2_scale: Optional[torch.Tensor],

375 + expert_map: torch.Tensor | None,

376 + a1q_scale: torch.Tensor | None,

377 + a2_scale: torch.Tensor | None,

378 workspace13: torch.Tensor,

379 workspace2: torch.Tensor,

382 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

380 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

381 apply_router_weight_on_input: bool,

382 ):

383 assert self.w1_scale is not None

vllm/model_executor/layers/fused_moe/fused_moe.py

+85 -84

@@ -5,7 +5,8 @@

5 import functools

6 import json

7 import os

8 -from typing import Any, Callable, Optional, Union

8 +from collections.abc import Callable

9 +from typing import Any

10

11 import torch

12 import torch.nn.functional as F

@@ -539,10 +540,10 @@ def invoke_fused_moe_kernel(

540 A: torch.Tensor,

541 B: torch.Tensor,

542 C: torch.Tensor,

542 - A_scale: Optional[torch.Tensor],

543 - B_scale: Optional[torch.Tensor],

544 - B_zp: Optional[torch.Tensor],

545 - topk_weights: Optional[torch.Tensor],

543 + A_scale: torch.Tensor | None,

544 + B_scale: torch.Tensor | None,

545 + B_zp: torch.Tensor | None,

546 + topk_weights: torch.Tensor | None,

547 sorted_token_ids: torch.Tensor,

548 expert_ids: torch.Tensor,

549 num_tokens_post_padded: torch.Tensor,

@@ -555,8 +556,8 @@ def invoke_fused_moe_kernel(

556 use_int8_w8a16: bool,

557 use_int4_w4a16: bool,

558 per_channel_quant: bool,

558 - block_shape: Optional[list[int]] = None,

559 - B_bias: Optional[torch.Tensor] = None,

559 + block_shape: list[int] | None = None,

560 + B_bias: torch.Tensor | None = None,

561 ) -> None:

562 assert topk_weights is not None or not mul_routed_weight

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

@@ -808,7 +809,7 @@ def zero_experts_compute_triton(

809

810 # Adapted from: https://github.com/sgl-project/sglang/pull/2628

811 def get_config_file_name(

811 - E: int, N: int, dtype: Optional[str], block_shape: Optional[list[int]] = None

812 + E: int, N: int, dtype: str | None, block_shape: list[int] | None = None

813 ) -> str:

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

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

@@ -823,10 +824,10 @@ def get_config_file_name(

824 def get_moe_configs(

825 E: int,

826 N: int,

826 - dtype: Optional[str],

827 - block_n: Optional[int] = None,

828 - block_k: Optional[int] = None,

829 -) -> Optional[dict[int, Any]]:

827 + dtype: str | None,

828 + block_n: int | None = None,

829 + block_k: int | None = None,

830 +) -> dict[int, Any] | None:

831 """

832 Return optimized configurations for the fused MoE kernel.

833

@@ -965,8 +966,8 @@ def get_default_config(

966 N: int,

967 K: int,

968 topk: int,

968 - dtype: Optional[str],

969 - block_shape: Optional[list[int]] = None,

969 + dtype: str | None,

970 + block_shape: list[int] | None = None,

971 ) -> dict[str, int]:

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

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

@@ -1016,9 +1017,9 @@ def try_get_optimal_moe_config(

1017 w1_shape: tuple[int, ...],

1018 w2_shape: tuple[int, ...],

1019 top_k: int,

1019 - dtype: Optional[str],

1020 + dtype: str | None,

1021 M: int,

1021 - block_shape: Optional[list[int]] = None,

1022 + block_shape: list[int] | None = None,

1023 ) -> dict[str, int]:

1024 from vllm.model_executor.layers.fused_moe import get_config

1025

@@ -1076,7 +1077,7 @@ def fused_topk(

1077 gating_output: torch.Tensor,

1078 topk: int,

1079 renormalize: bool,

1079 - indices_type: Optional[torch.dtype] = None,

1080 + indices_type: torch.dtype | None = None,

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

1082 assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"

1083

@@ -1135,7 +1136,7 @@ def grouped_topk(

1136 topk_group: int = 0,

1137 scoring_func: str = "softmax",

1138 routed_scaling_factor: float = 1.0,

1138 - e_score_correction_bias: Optional[torch.Tensor] = None,

1139 + e_score_correction_bias: torch.Tensor | None = None,

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

1141 if (

1142 envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK

@@ -1211,7 +1212,7 @@ def eplb_map_to_physical_and_record(

1212 expert_load_view: torch.Tensor,

1213 logical_to_physical_map: torch.Tensor,

1214 logical_replica_count: torch.Tensor,

1214 - indices_type: Optional[torch.dtype] = None,

1215 + indices_type: torch.dtype | None = None,

1216 ) -> torch.Tensor:

1217 """

1218 Map the logical expert ids to physical expert ids

@@ -1326,19 +1327,19 @@ def inplace_fused_experts(

1327 use_int8_w8a8: bool = False,

1328 use_int8_w8a16: bool = False,

1329 use_int4_w4a16: bool = False,

1329 - ocp_mx_scheme: Optional[str] = None,

1330 + ocp_mx_scheme: str | None = None,

1331 per_channel_quant: bool = False,

1332 global_num_experts: int = -1,

1332 - expert_map: Optional[torch.Tensor] = None,

1333 - w1_scale: Optional[torch.Tensor] = None,

1334 - w2_scale: Optional[torch.Tensor] = None,

1335 - w1_zp: Optional[torch.Tensor] = None,

1336 - w2_zp: Optional[torch.Tensor] = None,

1337 - a1_scale: Optional[torch.Tensor] = None,

1338 - a2_scale: Optional[torch.Tensor] = None,

1339 - block_shape: Optional[list[int]] = None,

1340 - w1_bias: Optional[torch.Tensor] = None,

1341 - w2_bias: Optional[torch.Tensor] = None,

1333 + expert_map: torch.Tensor | None = None,

1334 + w1_scale: torch.Tensor | None = None,

1335 + w2_scale: torch.Tensor | None = None,

1336 + w1_zp: torch.Tensor | None = None,

1337 + w2_zp: torch.Tensor | None = None,

1338 + a1_scale: torch.Tensor | None = None,

1339 + a2_scale: torch.Tensor | None = None,

1340 + block_shape: list[int] | None = None,

1341 + w1_bias: torch.Tensor | None = None,

1342 + w2_bias: torch.Tensor | None = None,

1343 ) -> None:

1344 fused_experts_impl(

1345 hidden_states,

@@ -1381,19 +1382,19 @@ def inplace_fused_experts_fake(

1382 use_int8_w8a8: bool = False,

1383 use_int8_w8a16: bool = False,

1384 use_int4_w4a16: bool = False,

1384 - ocp_mx_scheme: Optional[str] = None,

1385 + ocp_mx_scheme: str | None = None,

1386 per_channel_quant: bool = False,

1387 global_num_experts: int = -1,

1387 - expert_map: Optional[torch.Tensor] = None,

1388 - w1_scale: Optional[torch.Tensor] = None,

1389 - w2_scale: Optional[torch.Tensor] = None,

1390 - w1_zp: Optional[torch.Tensor] = None,

1391 - w2_zp: Optional[torch.Tensor] = None,

1392 - a1_scale: Optional[torch.Tensor] = None,

1393 - a2_scale: Optional[torch.Tensor] = None,

1394 - block_shape: Optional[list[int]] = None,

1395 - w1_bias: Optional[torch.Tensor] = None,

1396 - w2_bias: Optional[torch.Tensor] = None,

1388 + expert_map: torch.Tensor | None = None,

1389 + w1_scale: torch.Tensor | None = None,

1390 + w2_scale: torch.Tensor | None = None,

1391 + w1_zp: torch.Tensor | None = None,

1392 + w2_zp: torch.Tensor | None = None,

1393 + a1_scale: torch.Tensor | None = None,

1394 + a2_scale: torch.Tensor | None = None,

1395 + block_shape: list[int] | None = None,

1396 + w1_bias: torch.Tensor | None = None,

1397 + w2_bias: torch.Tensor | None = None,

1398 ) -> None:

1399 pass

1400

@@ -1423,19 +1424,19 @@ def outplace_fused_experts(

1424 use_int8_w8a8: bool = False,

1425 use_int8_w8a16: bool = False,

1426 use_int4_w4a16: bool = False,

1426 - ocp_mx_scheme: Optional[str] = None,

1427 + ocp_mx_scheme: str | None = None,

1428 per_channel_quant: bool = False,

1429 global_num_experts: int = -1,

1429 - expert_map: Optional[torch.Tensor] = None,

1430 - w1_scale: Optional[torch.Tensor] = None,

1431 - w2_scale: Optional[torch.Tensor] = None,

1432 - w1_zp: Optional[torch.Tensor] = None,

1433 - w2_zp: Optional[torch.Tensor] = None,

1434 - a1_scale: Optional[torch.Tensor] = None,

1435 - a2_scale: Optional[torch.Tensor] = None,

1436 - block_shape: Optional[list[int]] = None,

1437 - w1_bias: Optional[torch.Tensor] = None,

1438 - w2_bias: Optional[torch.Tensor] = None,

1430 + expert_map: torch.Tensor | None = None,

1431 + w1_scale: torch.Tensor | None = None,

1432 + w2_scale: torch.Tensor | None = None,

1433 + w1_zp: torch.Tensor | None = None,

1434 + w2_zp: torch.Tensor | None = None,

1435 + a1_scale: torch.Tensor | None = None,

1436 + a2_scale: torch.Tensor | None = None,

1437 + block_shape: list[int] | None = None,

1438 + w1_bias: torch.Tensor | None = None,

1439 + w2_bias: torch.Tensor | None = None,

1440 ) -> torch.Tensor:

1441 return fused_experts_impl(

1442 hidden_states,

@@ -1477,19 +1478,19 @@ def outplace_fused_experts_fake(

1478 use_int8_w8a8: bool = False,

1479 use_int8_w8a16: bool = False,

1480 use_int4_w4a16: bool = False,

1480 - ocp_mx_scheme: Optional[str] = None,

1481 + ocp_mx_scheme: str | None = None,

1482 per_channel_quant: bool = False,

1483 global_num_experts: int = -1,

1483 - expert_map: Optional[torch.Tensor] = None,

1484 - w1_scale: Optional[torch.Tensor] = None,

1485 - w2_scale: Optional[torch.Tensor] = None,

1486 - w1_zp: Optional[torch.Tensor] = None,

1487 - w2_zp: Optional[torch.Tensor] = None,

1488 - a1_scale: Optional[torch.Tensor] = None,

1489 - a2_scale: Optional[torch.Tensor] = None,

1490 - block_shape: Optional[list[int]] = None,

1491 - w1_bias: Optional[torch.Tensor] = None,

1492 - w2_bias: Optional[torch.Tensor] = None,

1484 + expert_map: torch.Tensor | None = None,

1485 + w1_scale: torch.Tensor | None = None,

1486 + w2_scale: torch.Tensor | None = None,

1487 + w1_zp: torch.Tensor | None = None,

1488 + w2_zp: torch.Tensor | None = None,

1489 + a1_scale: torch.Tensor | None = None,

1490 + a2_scale: torch.Tensor | None = None,

1491 + block_shape: list[int] | None = None,

1492 + w1_bias: torch.Tensor | None = None,

1493 + w2_bias: torch.Tensor | None = None,

1494 ) -> torch.Tensor:

1495 return torch.empty_like(hidden_states)

1496

@@ -1534,8 +1535,8 @@ def fused_experts(

1535 activation: str = "silu",

1536 apply_router_weight_on_input: bool = False,

1537 global_num_experts: int = -1,

1537 - expert_map: Optional[torch.Tensor] = None,

1538 - quant_config: Optional[FusedMoEQuantConfig] = None,

1538 + expert_map: torch.Tensor | None = None,

1539 + quant_config: FusedMoEQuantConfig | None = None,

1540 allow_deep_gemm: bool = False,

1541 allow_cutlass_block_scaled_grouped_gemm: bool = False,

1542 ) -> torch.Tensor:

@@ -1625,8 +1626,8 @@ GELU_NO_MUL: str = activation_without_mul("gelu")

1626 def _get_config_quant_dtype(

1627 use_fp8_w8a8: bool,

1628 use_int8_w8a8: bool,

1628 - ocp_mx_scheme: Optional[str],

1629 -) -> Union[None, torch.dtype, str]:

1629 + ocp_mx_scheme: str | None,

1630 +) -> None | torch.dtype | str:

1631 """

1632 Get the quantization type based on the quantization strategy flags.

1633 We don't have a quant_config at this point so we need to work backwards.

@@ -1660,19 +1661,19 @@ def fused_experts_impl(

1661 use_int8_w8a8: bool = False,

1662 use_int8_w8a16: bool = False,

1663 use_int4_w4a16: bool = False,

1663 - ocp_mx_scheme: Optional[str] = None,

1664 + ocp_mx_scheme: str | None = None,

1665 per_channel_quant: bool = False,

1666 global_num_experts: int = -1,

1666 - expert_map: Optional[torch.Tensor] = None,

1667 - w1_scale: Optional[torch.Tensor] = None,

1668 - w2_scale: Optional[torch.Tensor] = None,

1669 - w1_zp: Optional[torch.Tensor] = None,

1670 - w2_zp: Optional[torch.Tensor] = None,

1671 - a1_scale: Optional[torch.Tensor] = None,

1672 - a2_scale: Optional[torch.Tensor] = None,

1673 - block_shape: Optional[list[int]] = None,

1674 - w1_bias: Optional[torch.Tensor] = None,

1675 - w2_bias: Optional[torch.Tensor] = None,

1667 + expert_map: torch.Tensor | None = None,

1668 + w1_scale: torch.Tensor | None = None,

1669 + w2_scale: torch.Tensor | None = None,

1670 + w1_zp: torch.Tensor | None = None,

1671 + w2_zp: torch.Tensor | None = None,

1672 + a1_scale: torch.Tensor | None = None,

1673 + a2_scale: torch.Tensor | None = None,

1674 + block_shape: list[int] | None = None,

1675 + w1_bias: torch.Tensor | None = None,

1676 + w2_bias: torch.Tensor | None = None,

1677 ) -> torch.Tensor:

1678 # Check constraints.

1679 if use_int4_w4a16:

@@ -1964,7 +1965,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1965 topk: int,

1966 global_num_experts: int,

1967 local_num_experts: int,

1967 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

1968 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

1969 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

1970 workspace1 = (M, topk, max(N // 2, K))

1971 workspace2 = (M, topk, max(N, K))

@@ -1981,12 +1982,12 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1982 topk_ids: torch.Tensor,

1983 activation: str,

1984 global_num_experts: int,

1984 - expert_map: Optional[torch.Tensor],

1985 - a1q_scale: Optional[torch.Tensor],

1986 - a2_scale: Optional[torch.Tensor],

1985 + expert_map: torch.Tensor | None,

1986 + a1q_scale: torch.Tensor | None,

1987 + a2_scale: torch.Tensor | None,

1988 workspace13: torch.Tensor,

1989 workspace2: torch.Tensor,

1989 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

1990 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

1991 apply_router_weight_on_input: bool,

1992 ):

1993 # Check constraints.

@@ -2074,7 +2075,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

2075 activation, intermediate_cache2, intermediate_cache1.view(-1, N)

2076 )

2077

2077 - a2q_scale: Optional[torch.Tensor] = None

2078 + a2q_scale: torch.Tensor | None = None

2079

2080 qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(

2081 intermediate_cache2,

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

+10 -11

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -80,10 +79,10 @@ def triton_kernel_moe_forward(

79 topk: int,

80 renormalize: bool,

81 activation: str = "silu",

83 - quant_config: Optional[FusedMoEQuantConfig] = None,

82 + quant_config: FusedMoEQuantConfig | None = None,

83 apply_router_weight_on_input: bool = False,

84 global_num_experts: int = -1,

86 - expert_map: Optional[torch.Tensor] = None,

85 + expert_map: torch.Tensor | None = None,

86 ) -> torch.Tensor:

87 routing_data, gather_idx, scatter_idx = routing(

88 gating_output, topk, sm_first=not renormalize

@@ -115,13 +114,13 @@ def triton_kernel_fused_experts(

114 gather_indx, # GatherIndx

115 scatter_indx, # ScatterIndx

116 activation: str = "silu",

118 - quant_config: Optional[FusedMoEQuantConfig] = None,

117 + quant_config: FusedMoEQuantConfig | None = None,

118 swiglu_alpha: float = 1.702,

119 swiglu_limit: float = 7.0,

120 apply_router_weight_on_input: bool = False,

121 global_num_experts: int = -1,

123 - expert_map: Optional[torch.Tensor] = None,

124 - a1q_scale: Optional[torch.Tensor] = None,

122 + expert_map: torch.Tensor | None = None,

123 + a1q_scale: torch.Tensor | None = None,

124 ) -> torch.Tensor:

125 if quant_config is None:

126 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

@@ -261,7 +260,7 @@ class OAITritonExperts(BaseOAITritonExperts):

260 topk: int,

261 global_num_experts: int,

262 local_num_experts: int,

264 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

263 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

264 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

265 # workspace are allocated inside the kernel

266 workspace1 = (M, K)

@@ -279,12 +278,12 @@ class OAITritonExperts(BaseOAITritonExperts):

278 topk_ids: torch.Tensor,

279 activation: str,

280 global_num_experts: int,

282 - expert_map: Optional[torch.Tensor],

283 - a1q_scale: Optional[torch.Tensor],

284 - a2_scale: Optional[torch.Tensor],

281 + expert_map: torch.Tensor | None,

282 + a1q_scale: torch.Tensor | None,

283 + a2_scale: torch.Tensor | None,

284 workspace13: torch.Tensor,

285 workspace2: torch.Tensor,

287 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

286 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

287 apply_router_weight_on_input: bool,

288 ):

289 if expert_map is not None:

vllm/model_executor/layers/fused_moe/layer.py

+108 -108

@@ -2,10 +2,10 @@

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

3

4 from abc import abstractmethod

5 -from collections.abc import Iterable

5 +from collections.abc import Callable, Iterable

6 from contextlib import nullcontext

7 from enum import Enum

8 -from typing import Callable, Literal, Optional, Union, get_args, overload

8 +from typing import Literal, get_args, overload

9

10 import torch

11 import torch.nn.functional as F

@@ -70,15 +70,15 @@ if current_platform.is_cuda_alike():

70 )

71 else:

72 fused_experts = None # type: ignore

73 - FusedMoEPermuteExpertsUnpermute = None # type: ignore

74 - FusedMoEPrepareAndFinalize = None # type: ignore

73 + FusedMoEPermuteExpertsUnpermute = object # type: ignore

74 + FusedMoEPrepareAndFinalize = object # type: ignore

75

76 def _eplb_map_to_physical_and_record(

77 topk_ids: torch.Tensor,

78 expert_load_view: torch.Tensor,

79 logical_to_physical_map: torch.Tensor,

80 logical_replica_count: torch.Tensor,

81 - indices_type: Optional[torch.dtype],

81 + indices_type: torch.dtype | None,

82 ) -> torch.Tensor:

83 # CPU fallback: no EPLB so just return as is

84 return topk_ids

@@ -110,8 +110,8 @@ class FusedMoEMethodBase(QuantizeMethodBase):

110 def __init__(self, moe: FusedMoEConfig):

111 super().__init__()

112 self.moe = moe

113 - self.moe_quant_config: Optional[FusedMoEQuantConfig] = None

114 - self.fused_experts: Optional[FusedMoEModularKernel] = None

113 + self.moe_quant_config: FusedMoEQuantConfig | None = None

114 + self.fused_experts: FusedMoEModularKernel | None = None

115 self.topk_indices_dtype = None

116

117 @abstractmethod

@@ -139,12 +139,12 @@ class FusedMoEMethodBase(QuantizeMethodBase):

139 @staticmethod

140 def _maybe_make_prepare_finalize(

141 moe: FusedMoEConfig,

142 - quant_config: Optional[FusedMoEQuantConfig],

143 - ) -> Optional[FusedMoEPrepareAndFinalize]:

142 + quant_config: FusedMoEQuantConfig | None,

143 + ) -> FusedMoEPrepareAndFinalize | None:

144 all2all_manager = get_ep_group().device_communicator.all2all_manager

145 assert all2all_manager is not None

146

147 - prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None

147 + prepare_finalize: FusedMoEPrepareAndFinalize | None = None

148

149 # TODO: could allow this now

150 assert not moe.use_flashinfer_cutlass_kernels, "Must be created in modelopt.py"

@@ -229,7 +229,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):

229

230 return prepare_finalize

231

232 - def maybe_make_prepare_finalize(self) -> Optional[FusedMoEPrepareAndFinalize]:

232 + def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:

233 if self.moe.moe_parallel_config.use_all2all_kernels:

234 return FusedMoEMethodBase._maybe_make_prepare_finalize(

235 self.moe, self.moe_quant_config

@@ -280,7 +280,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):

280 @abstractmethod

281 def get_fused_moe_quant_config(

282 self, layer: torch.nn.Module

283 - ) -> Optional[FusedMoEQuantConfig]:

283 + ) -> FusedMoEQuantConfig | None:

284 raise NotImplementedError

285

286 @property

@@ -296,21 +296,21 @@ class FusedMoEMethodBase(QuantizeMethodBase):

296 top_k: int,

297 renormalize: bool,

298 use_grouped_topk: bool = False,

299 - topk_group: Optional[int] = None,

300 - num_expert_group: Optional[int] = None,

299 + topk_group: int | None = None,

300 + num_expert_group: int | None = None,

301 global_num_experts: int = -1,

302 - expert_map: Optional[torch.Tensor] = None,

303 - custom_routing_function: Optional[Callable] = None,

302 + expert_map: torch.Tensor | None = None,

303 + custom_routing_function: Callable | None = None,

304 scoring_func: str = "softmax",

305 routed_scaling_factor: float = 1.0,

306 - e_score_correction_bias: Optional[torch.Tensor] = None,

306 + e_score_correction_bias: torch.Tensor | None = None,

307 apply_router_weight_on_input: bool = False,

308 activation: str = "silu",

309 enable_eplb: bool = False,

310 - expert_load_view: Optional[torch.Tensor] = None,

311 - logical_to_physical_map: Optional[torch.Tensor] = None,

312 - logical_replica_count: Optional[torch.Tensor] = None,

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

310 + expert_load_view: torch.Tensor | None = None,

311 + logical_to_physical_map: torch.Tensor | None = None,

312 + logical_replica_count: torch.Tensor | None = None,

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

314 raise NotImplementedError

315

316

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

368 )

369 self.flashinfer_cutlass_moe = None # type: ignore

370

371 - def maybe_make_prepare_finalize(self) -> Optional[FusedMoEPrepareAndFinalize]:

371 + def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:

372 if self.rocm_aiter_moe_enabled:

373 return None

374 else:

@@ -532,21 +532,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

532 top_k: int,

533 renormalize: bool,

534 use_grouped_topk: bool = False,

535 - topk_group: Optional[int] = None,

536 - num_expert_group: Optional[int] = None,

535 + topk_group: int | None = None,

536 + num_expert_group: int | None = None,

537 global_num_experts: int = -1,

538 - expert_map: Optional[torch.Tensor] = None,

539 - custom_routing_function: Optional[Callable] = None,

538 + expert_map: torch.Tensor | None = None,

539 + custom_routing_function: Callable | None = None,

540 scoring_func: str = "softmax",

541 routed_scaling_factor: float = 1.0,

542 - e_score_correction_bias: Optional[torch.Tensor] = None,

542 + e_score_correction_bias: torch.Tensor | None = None,

543 apply_router_weight_on_input: bool = False,

544 activation: str = "silu",

545 enable_eplb: bool = False,

546 - expert_load_view: Optional[torch.Tensor] = None,

547 - logical_to_physical_map: Optional[torch.Tensor] = None,

548 - logical_replica_count: Optional[torch.Tensor] = None,

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

546 + expert_load_view: torch.Tensor | None = None,

547 + logical_to_physical_map: torch.Tensor | None = None,

548 + logical_replica_count: torch.Tensor | None = None,

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

550 if enable_eplb:

551 assert expert_load_view is not None

552 assert logical_to_physical_map is not None

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

578

579 def get_fused_moe_quant_config(

580 self, layer: torch.nn.Module

581 - ) -> Optional[FusedMoEQuantConfig]:

581 + ) -> FusedMoEQuantConfig | None:

582 if self.moe.has_bias:

583 return biased_moe_quant_config(

584 layer.w13_bias,

@@ -595,21 +595,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

595 top_k: int,

596 router_logits: torch.Tensor,

597 renormalize: bool,

598 - topk_group: Optional[int] = None,

599 - num_expert_group: Optional[int] = None,

598 + topk_group: int | None = None,

599 + num_expert_group: int | None = None,

600 global_num_experts: int = -1,

601 - expert_map: Optional[torch.Tensor] = None,

602 - custom_routing_function: Optional[Callable] = None,

601 + expert_map: torch.Tensor | None = None,

602 + custom_routing_function: Callable | None = None,

603 scoring_func: str = "softmax",

604 routed_scaling_factor: float = 1.0,

605 - e_score_correction_bias: Optional[torch.Tensor] = None,

605 + e_score_correction_bias: torch.Tensor | None = None,

606 apply_router_weight_on_input: bool = False,

607 activation: str = "silu",

608 enable_eplb: bool = False,

609 - expert_load_view: Optional[torch.Tensor] = None,

610 - logical_to_physical_map: Optional[torch.Tensor] = None,

611 - logical_replica_count: Optional[torch.Tensor] = None,

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

609 + expert_load_view: torch.Tensor | None = None,

610 + logical_to_physical_map: torch.Tensor | None = None,

611 + logical_replica_count: torch.Tensor | None = None,

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

613 zero_expert_num = getattr(layer, "zero_expert_num", 0)

614 zero_expert_type = getattr(layer, "zero_expert_type", None)

615

@@ -705,21 +705,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

705 top_k: int,

706 router_logits: torch.Tensor,

707 renormalize: bool,

708 - topk_group: Optional[int] = None,

709 - num_expert_group: Optional[int] = None,

708 + topk_group: int | None = None,

709 + num_expert_group: int | None = None,

710 global_num_experts: int = -1,

711 - expert_map: Optional[torch.Tensor] = None,

712 - custom_routing_function: Optional[Callable] = None,

711 + expert_map: torch.Tensor | None = None,

712 + custom_routing_function: Callable | None = None,

713 scoring_func: str = "softmax",

714 routed_scaling_factor: float = 1.0,

715 - e_score_correction_bias: Optional[torch.Tensor] = None,

715 + e_score_correction_bias: torch.Tensor | None = None,

716 apply_router_weight_on_input: bool = False,

717 activation: str = "silu",

718 enable_eplb: bool = False,

719 - expert_load_view: Optional[torch.Tensor] = None,

720 - logical_to_physical_map: Optional[torch.Tensor] = None,

721 - logical_replica_count: Optional[torch.Tensor] = None,

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

719 + expert_load_view: torch.Tensor | None = None,

720 + logical_to_physical_map: torch.Tensor | None = None,

721 + logical_replica_count: torch.Tensor | None = None,

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

723 if (

724 enable_eplb is not False

725 or expert_load_view is not None

@@ -754,21 +754,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

754 top_k: int,

755 router_logits: torch.Tensor,

756 renormalize: bool,

757 - topk_group: Optional[int] = None,

758 - num_expert_group: Optional[int] = None,

757 + topk_group: int | None = None,

758 + num_expert_group: int | None = None,

759 global_num_experts: int = -1,

760 - expert_map: Optional[torch.Tensor] = None,

761 - custom_routing_function: Optional[Callable] = None,

760 + expert_map: torch.Tensor | None = None,

761 + custom_routing_function: Callable | None = None,

762 scoring_func: str = "softmax",

763 routed_scaling_factor: float = 1.0,

764 - e_score_correction_bias: Optional[torch.Tensor] = None,

764 + e_score_correction_bias: torch.Tensor | None = None,

765 apply_router_weight_on_input: bool = False,

766 activation: str = "silu",

767 enable_eplb: bool = False,

768 - expert_load_view: Optional[torch.Tensor] = None,

769 - logical_to_physical_map: Optional[torch.Tensor] = None,

770 - logical_replica_count: Optional[torch.Tensor] = None,

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

768 + expert_load_view: torch.Tensor | None = None,

769 + logical_to_physical_map: torch.Tensor | None = None,

770 + logical_replica_count: torch.Tensor | None = None,

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

772 if (

773 enable_eplb is not False

774 or expert_load_view is not None

@@ -795,21 +795,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

795 top_k: int,

796 router_logits: torch.Tensor,

797 renormalize: bool,

798 - topk_group: Optional[int] = None,

799 - num_expert_group: Optional[int] = None,

798 + topk_group: int | None = None,

799 + num_expert_group: int | None = None,

800 global_num_experts: int = -1,

801 - expert_map: Optional[torch.Tensor] = None,

802 - custom_routing_function: Optional[Callable] = None,

801 + expert_map: torch.Tensor | None = None,

802 + custom_routing_function: Callable | None = None,

803 scoring_func: str = "softmax",

804 routed_scaling_factor: float = 1.0,

805 - e_score_correction_bias: Optional[torch.Tensor] = None,

805 + e_score_correction_bias: torch.Tensor | None = None,

806 apply_router_weight_on_input: bool = False,

807 activation: str = "silu",

808 enable_eplb: bool = False,

809 - expert_load_view: Optional[torch.Tensor] = None,

810 - logical_to_physical_map: Optional[torch.Tensor] = None,

811 - logical_replica_count: Optional[torch.Tensor] = None,

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

809 + expert_load_view: torch.Tensor | None = None,

810 + logical_to_physical_map: torch.Tensor | None = None,

811 + logical_replica_count: torch.Tensor | None = None,

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

813 assert not use_grouped_topk

814 assert num_expert_group is None

815 assert topk_group is None

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

860 ep_rank: int,

861 global_num_experts: int,

862 expert_placement_strategy: ExpertPlacementStrategy = "linear",

863 -) -> tuple[int, Optional[torch.Tensor]]:

863 +) -> tuple[int, torch.Tensor | None]:

864 """

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

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

@@ -941,7 +941,7 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str:

941 def maybe_roundup_hidden_size(

942 hidden_size: int,

943 act_dtype: torch.dtype,

944 - quant_config: Optional[QuantizationConfig],

944 + quant_config: QuantizationConfig | None,

945 moe_parallel_config: FusedMoEParallelConfig,

946 ) -> int:

947 """

@@ -1016,30 +1016,30 @@ class FusedMoE(CustomOp):

1016 top_k: int,

1017 hidden_size: int,

1018 intermediate_size: int,

1019 - params_dtype: Optional[torch.dtype] = None,

1019 + params_dtype: torch.dtype | None = None,

1020 reduce_results: bool = False,

1021 renormalize: bool = True,

1022 use_grouped_topk: bool = False,

1023 - num_expert_group: Optional[int] = None,

1024 - topk_group: Optional[int] = None,

1025 - quant_config: Optional[QuantizationConfig] = None,

1026 - tp_size: Optional[int] = None,

1027 - ep_size: Optional[int] = None,

1028 - dp_size: Optional[int] = None,

1023 + num_expert_group: int | None = None,

1024 + topk_group: int | None = None,

1025 + quant_config: QuantizationConfig | None = None,

1026 + tp_size: int | None = None,

1027 + ep_size: int | None = None,

1028 + dp_size: int | None = None,

1029 prefix: str = "",

1030 - custom_routing_function: Optional[Callable] = None,

1030 + custom_routing_function: Callable | None = None,

1031 scoring_func: str = "softmax",

1032 routed_scaling_factor: float = 1.0,

1033 - e_score_correction_bias: Optional[torch.Tensor] = None,

1033 + e_score_correction_bias: torch.Tensor | None = None,

1034 apply_router_weight_on_input: bool = False,

1035 activation: str = "silu",

1036 enable_eplb: bool = False,

1037 num_redundant_experts: int = 0,

1038 has_bias: bool = False,

1039 is_sequence_parallel=False,

1040 - zero_expert_num: Optional[int] = 0,

1041 - zero_expert_type: Optional[str] = None,

1042 - expert_mapping: Optional[list[tuple[str, str, int, str]]] = None,

1040 + zero_expert_num: int | None = 0,

1041 + zero_expert_type: str | None = None,

1042 + expert_mapping: list[tuple[str, str, int, str]] | None = None,

1043 ):

1044 super().__init__()

1045 if params_dtype is None:

@@ -1092,9 +1092,9 @@ class FusedMoE(CustomOp):

1092 self.layer_name = prefix

1093

1094 self.enable_eplb = enable_eplb

1095 - self.expert_load_view: Optional[torch.Tensor] = None

1096 - self.logical_to_physical_map: Optional[torch.Tensor] = None

1097 - self.logical_replica_count: Optional[torch.Tensor] = None

1095 + self.expert_load_view: torch.Tensor | None = None

1096 + self.logical_to_physical_map: torch.Tensor | None = None

1097 + self.logical_replica_count: torch.Tensor | None = None

1098

1099 # Determine expert maps

1100 if self.use_ep:

@@ -1128,7 +1128,7 @@ class FusedMoE(CustomOp):

1128 )

1129 expert_placement_strategy = "linear"

1130

1131 - self.expert_map: Optional[torch.Tensor]

1131 + self.expert_map: torch.Tensor | None

1132 local_num_experts, expert_map = determine_expert_map(

1133 ep_size=self.ep_size,

1134 ep_rank=self.ep_rank,

@@ -1187,12 +1187,12 @@ class FusedMoE(CustomOp):

1187 has_bias=has_bias,

1188 )

1189 self.moe_config = moe

1190 - self.moe_quant_config: Optional[FusedMoEQuantConfig] = None

1190 + self.moe_quant_config: FusedMoEQuantConfig | None = None

1191 self.quant_config = quant_config

1192

1193 # Note: get_quant_method will look at the layer's local_num_experts

1194 # for heuristic purposes, so it must be initialized first.

1195 - quant_method: Optional[QuantizeMethodBase] = None

1195 + quant_method: QuantizeMethodBase | None = None

1196 quant_method = (

1197 UnquantizedFusedMoEMethod(moe)

1198 if quant_config is None

@@ -1238,8 +1238,8 @@ class FusedMoE(CustomOp):

1238 self.quant_method.create_weights(layer=self, **moe_quant_params)

1239

1240 # Chunked all2all staging tensor

1241 - self.batched_hidden_states: Optional[torch.Tensor] = None

1242 - self.batched_router_logits: Optional[torch.Tensor] = None

1241 + self.batched_hidden_states: torch.Tensor | None = None

1242 + self.batched_router_logits: torch.Tensor | None = None

1243

1244 if self.use_dp_chunking:

1245 states_shape: tuple[int, ...]

@@ -1262,7 +1262,7 @@ class FusedMoE(CustomOp):

1262 )

1263

1264 @property

1265 - def shared_experts(self) -> Optional[torch.nn.Module]:

1265 + def shared_experts(self) -> torch.nn.Module | None:

1266 return None

1267

1268 @property

@@ -1534,7 +1534,7 @@ class FusedMoE(CustomOp):

1534 shard_id: str,

1535 expert_id: int,

1536 return_success: bool = False,

1537 - ) -> Optional[bool]:

1537 + ) -> bool | None:

1538 if self.quant_config and self.quant_config.get_name() == "mxfp4":

1539 # (FIXME) for gpt-oss all experts are combined

1540 if "bias" in weight_name:

@@ -1851,21 +1851,21 @@ class FusedMoE(CustomOp):

1851 top_k: int,

1852 use_grouped_topk: bool,

1853 renormalize: bool,

1854 - topk_group: Optional[int] = None,

1855 - num_expert_group: Optional[int] = None,

1856 - custom_routing_function: Optional[Callable] = None,

1854 + topk_group: int | None = None,

1855 + num_expert_group: int | None = None,

1856 + custom_routing_function: Callable | None = None,

1857 scoring_func: str = "softmax",

1858 routed_scaling_factor: float = 1.0,

1859 - e_score_correction_bias: Optional[torch.Tensor] = None,

1860 - indices_type: Optional[torch.dtype] = None,

1859 + e_score_correction_bias: torch.Tensor | None = None,

1860 + indices_type: torch.dtype | None = None,

1861 enable_eplb: bool = False,

1862 - expert_map: Optional[torch.Tensor] = None,

1863 - expert_load_view: Optional[torch.Tensor] = None,

1864 - logical_to_physical_map: Optional[torch.Tensor] = None,

1865 - logical_replica_count: Optional[torch.Tensor] = None,

1866 - global_num_experts: Optional[int] = None,

1867 - zero_expert_num: Optional[int] = None,

1868 - zero_expert_type: Optional[str] = None,

1862 + expert_map: torch.Tensor | None = None,

1863 + expert_load_view: torch.Tensor | None = None,

1864 + logical_to_physical_map: torch.Tensor | None = None,

1865 + logical_replica_count: torch.Tensor | None = None,

1866 + global_num_experts: int | None = None,

1867 + zero_expert_num: int | None = None,

1868 + zero_expert_type: str | None = None,

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

1870 """

1871 Route the input hidden states to the top-k experts based on the

@@ -2006,7 +2006,7 @@ class FusedMoE(CustomOp):

2006 self,

2007 hidden_states: torch.Tensor,

2008 router_logits: torch.Tensor,

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

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

2010 og_hidden_states = hidden_states.shape[-1]

2011 if self.hidden_size != og_hidden_states:

2012 hidden_states = F.pad(

@@ -2047,14 +2047,14 @@ class FusedMoE(CustomOp):

2047 self,

2048 hidden_states: torch.Tensor,

2049 router_logits: torch.Tensor,

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

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

2051 return self.forward_native(hidden_states, router_logits)

2052

2053 def forward_impl_chunked(

2054 self,

2055 full_hidden_states: torch.Tensor,

2056 full_router_logits: torch.Tensor,

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

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

2058 assert self.batched_hidden_states is not None

2059 assert self.batched_router_logits is not None

2060 assert self.batched_hidden_states.dtype == full_hidden_states.dtype

@@ -2200,7 +2200,7 @@ class FusedMoE(CustomOp):

2200 self,

2201 hidden_states: torch.Tensor,

2202 router_logits: torch.Tensor,

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

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

2204 assert self.quant_method is not None

2205

2206 self.ensure_moe_quant_config()

vllm/model_executor/layers/fused_moe/modular_kernel.py

+50 -49

@@ -1,10 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 from abc import ABC, abstractmethod

4 +from collections.abc import Callable

5 from dataclasses import dataclass

6 from enum import Enum

7 from math import prod

7 -from typing import Callable, Optional, Union, final

8 +from typing import final

9

10 import torch

11

@@ -81,7 +82,7 @@ class ExpertTokensMetadata:

82 """

83

84 expert_num_tokens: torch.Tensor

84 - expert_num_tokens_cpu: Optional[torch.Tensor]

85 + expert_num_tokens_cpu: torch.Tensor | None

86

87 @staticmethod

88 def make_from_list(

@@ -104,7 +105,7 @@ class TopKWeightAndReduce(ABC):

105 @abstractmethod

106 def apply(

107 self,

107 - output: Optional[torch.Tensor],

108 + output: torch.Tensor | None,

109 fused_expert_output: torch.Tensor,

110 topk_weights: torch.Tensor,

111 topk_ids: torch.Tensor,

@@ -132,10 +133,10 @@ class TopKWeightAndReduce(ABC):

133 #

134 PrepareResultType = tuple[

135 torch.Tensor,

135 - Optional[torch.Tensor],

136 - Optional[ExpertTokensMetadata],

137 - Optional[torch.Tensor],

138 - Optional[torch.Tensor],

136 + torch.Tensor | None,

137 + ExpertTokensMetadata | None,

138 + torch.Tensor | None,

139 + torch.Tensor | None,

140 ]

141

142 ReceiverType = Callable[[], PrepareResultType]

@@ -155,7 +156,7 @@ class FusedMoEPrepareAndFinalize(ABC):

156 topk_weights: torch.Tensor,

157 topk_ids: torch.Tensor,

158 num_experts: int,

158 - expert_map: Optional[torch.Tensor],

159 + expert_map: torch.Tensor | None,

160 apply_router_weight_on_input: bool,

161 quant_config: FusedMoEQuantConfig,

162 ) -> PrepareResultType:

@@ -195,10 +196,10 @@ class FusedMoEPrepareAndFinalize(ABC):

196 topk_weights: torch.Tensor,

197 topk_ids: torch.Tensor,

198 num_experts: int,

198 - expert_map: Optional[torch.Tensor],

199 + expert_map: torch.Tensor | None,

200 apply_router_weight_on_input: bool,

201 quant_config: FusedMoEQuantConfig,

201 - ) -> Union[tuple[Callable, ReceiverType], ReceiverType]:

202 + ) -> tuple[Callable, ReceiverType] | ReceiverType:

203 """

204 Perform any quantization (and/or) dispatching needed for this kernel

205 but do not wait for results from other workers.

@@ -270,7 +271,7 @@ class FusedMoEPrepareAndFinalize(ABC):

271 topk_ids: torch.Tensor,

272 apply_router_weight_on_input: bool,

273 weight_and_reduce_impl: TopKWeightAndReduce,

273 - ) -> Union[tuple[Callable, Callable], Callable]:

274 + ) -> tuple[Callable, Callable] | Callable:

275 """

276 Perform any combine plus apply weights and perform a reduction on the

277 fused experts output but do not wait for results from other workers.

@@ -314,7 +315,7 @@ class FusedMoEPrepareAndFinalize(ABC):

315 raise NotImplementedError

316

317 @abstractmethod

317 - def topk_indices_dtype(self) -> Optional[torch.dtype]:

318 + def topk_indices_dtype(self) -> torch.dtype | None:

319 """

320 The PrepareFinalize All2All implementations generally constrain the

321 dtype of the topk_ids they support. This function returns the

@@ -324,7 +325,7 @@ class FusedMoEPrepareAndFinalize(ABC):

325 raise NotImplementedError

326

327 @abstractmethod

327 - def max_num_tokens_per_rank(self) -> Optional[int]:

328 + def max_num_tokens_per_rank(self) -> int | None:

329 """

330 Some PrepareFinalize All2All implementations are batched. Meaning,

331 they can process only as set of tokens at a time. This

@@ -423,11 +424,11 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

424 #

425

426 @property

426 - def quant_dtype(self) -> Optional[torch.dtype]:

427 + def quant_dtype(self) -> torch.dtype | None:

428 return self.quant_config.quant_dtype

429

430 @property

430 - def block_shape(self) -> Optional[list[int]]:

431 + def block_shape(self) -> list[int] | None:

432 return self.quant_config.block_shape

433

434 @property

@@ -439,51 +440,51 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

440 return self.quant_config.per_out_ch_quant

441

442 @property

442 - def a1_scale(self) -> Optional[torch.Tensor]:

443 + def a1_scale(self) -> torch.Tensor | None:

444 return self.quant_config.a1_scale

445

446 @property

446 - def a2_scale(self) -> Optional[torch.Tensor]:

447 + def a2_scale(self) -> torch.Tensor | None:

448 return self.quant_config.a2_scale

449

450 @property

450 - def a1_gscale(self) -> Optional[torch.Tensor]:

451 + def a1_gscale(self) -> torch.Tensor | None:

452 return self.quant_config.a1_gscale

453

454 @property

454 - def a2_gscale(self) -> Optional[torch.Tensor]:

455 + def a2_gscale(self) -> torch.Tensor | None:

456 return self.quant_config.a2_gscale

457

458 @property

458 - def w1_scale(self) -> Optional[torch.Tensor]:

459 + def w1_scale(self) -> torch.Tensor | None:

460 return self.quant_config.w1_scale

461

462 @property

462 - def w2_scale(self) -> Optional[torch.Tensor]:

463 + def w2_scale(self) -> torch.Tensor | None:

464 return self.quant_config.w2_scale

465

466 @property

466 - def w1_zp(self) -> Optional[torch.Tensor]:

467 + def w1_zp(self) -> torch.Tensor | None:

468 return self.quant_config.w1_zp

469

470 @property

470 - def w2_zp(self) -> Optional[torch.Tensor]:

471 + def w2_zp(self) -> torch.Tensor | None:

472 return self.quant_config.w2_zp

473

474 @property

474 - def w1_bias(self) -> Optional[torch.Tensor]:

475 + def w1_bias(self) -> torch.Tensor | None:

476 return self.quant_config.w1_bias

477

478 @property

478 - def w2_bias(self) -> Optional[torch.Tensor]:

479 + def w2_bias(self) -> torch.Tensor | None:

480 return self.quant_config.w2_bias

481

482 @property

482 - def g1_alphas(self) -> Optional[torch.Tensor]:

483 + def g1_alphas(self) -> torch.Tensor | None:

484 return self.quant_config.g1_alphas

485

486 @property

486 - def g2_alphas(self) -> Optional[torch.Tensor]:

487 + def g2_alphas(self) -> torch.Tensor | None:

488 return self.quant_config.g2_alphas

489

490 # TODO (bnell): make this return a CHUNK_SIZE or None instead?

@@ -517,7 +518,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

518 topk: int,

519 global_num_experts: int,

520 local_num_experts: int,

520 - expert_tokens_meta: Optional[ExpertTokensMetadata],

521 + expert_tokens_meta: ExpertTokensMetadata | None,

522 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

523 """

524 Compute the shapes for the temporary and final outputs of the two gemms

@@ -578,12 +579,12 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

579 topk_ids: torch.Tensor,

580 activation: str,

581 global_num_experts: int,

581 - expert_map: Optional[torch.Tensor],

582 - a1q_scale: Optional[torch.Tensor],

583 - a2_scale: Optional[torch.Tensor],

582 + expert_map: torch.Tensor | None,

583 + a1q_scale: torch.Tensor | None,

584 + a2_scale: torch.Tensor | None,

585 workspace13: torch.Tensor,

586 workspace2: torch.Tensor,

586 - expert_tokens_meta: Optional[ExpertTokensMetadata],

587 + expert_tokens_meta: ExpertTokensMetadata | None,

588 apply_router_weight_on_input: bool,

589 ) -> None:

590 """

@@ -625,8 +626,8 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

626

627

628 def _slice_scales(

628 - scales: Optional[torch.Tensor], start: int, end: int

629 -) -> Optional[torch.Tensor]:

629 + scales: torch.Tensor | None, start: int, end: int

630 +) -> torch.Tensor | None:

631 if scales is not None:

632 if scales.numel() == 1:

633 return scales

@@ -688,7 +689,7 @@ class FusedMoEModularKernel(torch.nn.Module):

689 self,

690 prepare_finalize: FusedMoEPrepareAndFinalize,

691 fused_experts: FusedMoEPermuteExpertsUnpermute,

691 - shared_experts: Optional[torch.nn.Module] = None,

692 + shared_experts: torch.nn.Module | None = None,

693 ):

694 super().__init__()

695 self.prepare_finalize = prepare_finalize

@@ -741,7 +742,7 @@ class FusedMoEModularKernel(torch.nn.Module):

742 top_k: int,

743 global_num_experts: int,

744 local_num_experts: int,

744 - expert_tokens_meta: Optional[ExpertTokensMetadata],

745 + expert_tokens_meta: ExpertTokensMetadata | None,

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

747 """

748 Allocate temporary and output buffers for the fused experts op.

@@ -825,11 +826,11 @@ class FusedMoEModularKernel(torch.nn.Module):

826 @staticmethod

827 def _slice_expert_tokens_metadata(

828 num_chunks: int,

828 - full_expert_tokens_meta: Optional[ExpertTokensMetadata],

829 + full_expert_tokens_meta: ExpertTokensMetadata | None,

830 chunk_topk_ids: torch.Tensor,

831 local_num_experts: int,

831 - expert_map: Optional[torch.Tensor],

832 - ) -> Optional[ExpertTokensMetadata]:

832 + expert_map: torch.Tensor | None,

833 + ) -> ExpertTokensMetadata | None:

834 if num_chunks == 1 or full_expert_tokens_meta is None:

835 return full_expert_tokens_meta

836

@@ -861,12 +862,12 @@ class FusedMoEModularKernel(torch.nn.Module):

862 topk_weights: torch.Tensor,

863 topk_ids: torch.Tensor,

864 global_num_experts: int,

864 - expert_map: Optional[torch.Tensor],

865 + expert_map: torch.Tensor | None,

866 apply_router_weight_on_input: bool,

867 ) -> tuple[

868 torch.Tensor,

868 - Optional[torch.Tensor],

869 - Optional[ExpertTokensMetadata],

869 + torch.Tensor | None,

870 + ExpertTokensMetadata | None,

871 torch.Tensor,

872 torch.Tensor,

873 ]:

@@ -945,7 +946,7 @@ class FusedMoEModularKernel(torch.nn.Module):

946 self,

947 in_dtype: torch.dtype,

948 a1q: torch.Tensor,

948 - a1q_scale: Optional[torch.Tensor],

949 + a1q_scale: torch.Tensor | None,

950 w1: torch.Tensor,

951 w2: torch.Tensor,

952 topk_weights: torch.Tensor,

@@ -953,9 +954,9 @@ class FusedMoEModularKernel(torch.nn.Module):

954 activation: str,

955 global_num_experts: int,

956 local_num_experts: int,

956 - expert_map: Optional[torch.Tensor],

957 + expert_map: torch.Tensor | None,

958 apply_router_weight_on_input: bool,

958 - expert_tokens_meta: Optional[ExpertTokensMetadata],

959 + expert_tokens_meta: ExpertTokensMetadata | None,

960 ) -> torch.Tensor:

961 _, M_full, N, K, top_k = self.fused_experts.moe_problem_size(

962 a1q, w1, w2, topk_ids

@@ -1042,12 +1043,12 @@ class FusedMoEModularKernel(torch.nn.Module):

1043 topk_weights: torch.Tensor,

1044 topk_ids: torch.Tensor,

1045 apply_router_weight_on_input: bool,

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

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

1047 """

1048 The _finalize method is a wrapper around self.prepare_finalize.finalize

1049 that handles DBO, async and shared expert overlap.

1050 """

1050 - shared_output: Optional[torch.Tensor] = None

1051 + shared_output: torch.Tensor | None = None

1052

1053 if not self.prepare_finalize.supports_async():

1054 assert not dbo_enabled()

@@ -1112,9 +1113,9 @@ class FusedMoEModularKernel(torch.nn.Module):

1113 inplace: bool = False,

1114 activation: str = "silu",

1115 global_num_experts: int = -1,

1115 - expert_map: Optional[torch.Tensor] = None,

1116 + expert_map: torch.Tensor | None = None,

1117 apply_router_weight_on_input: bool = False,

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

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

1119 """

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

1121 of weights, w1 and w2, and top-k gating mechanism.

vllm/model_executor/layers/fused_moe/moe_align_block_size.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -13,7 +12,7 @@ def moe_align_block_size(

12 topk_ids: torch.Tensor,

13 block_size: int,

14 num_experts: int,

16 - expert_map: Optional[torch.Tensor] = None,

15 + expert_map: torch.Tensor | None = None,

16 pad_sorted_ids: bool = False,

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

18 """

vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py

+11 -16

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -13,14 +12,12 @@ from vllm.model_executor.layers.fused_moe.utils import _fp8_perm

12

13 def _moe_permute(

14 curr_hidden_states: torch.Tensor,

16 - a1q_scale: Optional[torch.Tensor],

15 + a1q_scale: torch.Tensor | None,

16 curr_topk_ids: torch.Tensor,

17 global_num_experts: int,

19 - expert_map: Optional[torch.Tensor],

18 + expert_map: torch.Tensor | None,

19 block_m: int,

21 -) -> tuple[

22 - torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor

23 -]:

20 +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]:

21 """

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

23 Permute the hidden states and scales according to `sorted_token_ids`.

@@ -33,7 +30,7 @@ def _moe_permute(

30 curr_topk_ids, block_m, global_num_experts, expert_map, pad_sorted_ids=True

31 )

32

36 - inv_perm: Optional[torch.Tensor] = None

33 + inv_perm: torch.Tensor | None = None

34

35 num_tokens = top_k_num * tokens_in_chunk

36 expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)

@@ -53,7 +50,7 @@ def _moe_permute(

50 def _moe_unpermute_and_reduce(

51 out: torch.Tensor,

52 curr_hidden: torch.Tensor,

56 - inv_perm: Optional[torch.Tensor],

53 + inv_perm: torch.Tensor | None,

54 topk_weight: torch.Tensor,

55 apply_router_weight_on_input: bool,

56 ) -> None:

@@ -73,17 +70,15 @@ def _moe_unpermute_and_reduce(

70

71 def moe_permute(

72 hidden_states: torch.Tensor,

76 - a1q_scale: Optional[torch.Tensor],

73 + a1q_scale: torch.Tensor | None,

74 topk_ids: torch.Tensor,

75 n_expert: int,

76 n_local_expert: int = -1,

80 - expert_map: Optional[torch.Tensor] = None,

81 - align_block_size: Optional[int] = None,

77 + expert_map: torch.Tensor | None = None,

78 + align_block_size: int | None = None,

79 fill_invalid_expert: int = -1,

83 - permuted_hidden_states: Optional[torch.Tensor] = None,

84 -) -> tuple[

85 - torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor

86 -]:

80 + permuted_hidden_states: torch.Tensor | None = None,

81 +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]:

82 """

83 This function expands and permutes activation to gather uncontinuous tokens

84 for each expert.

@@ -198,7 +193,7 @@ def moe_unpermute(

193 permuted_hidden_states: torch.Tensor,

194 topk_weights: torch.Tensor,

195 inv_permuted_idx: torch.Tensor,

201 - expert_first_token_offset: Optional[torch.Tensor] = None,

196 + expert_first_token_offset: torch.Tensor | None = None,

197 ) -> None:

198 """

199 This function expands and permutes activation to gathering uncontinuous

vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py

+13 -13

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Callable, Optional, Union

3 +from collections.abc import Callable

4

5 import pplx_kernels as pplx

6 import torch

@@ -24,9 +24,9 @@ def pplx_hidden_dim_scale_bytes(

24 max_num_tokens: int,

25 hidden_dim: int,

26 in_dtype: torch.dtype,

27 - quant_dtype: Union[torch.dtype, str, None],

27 + quant_dtype: torch.dtype | str | None,

28 per_act_token_quant: bool,

29 - block_shape: Optional[list[int]],

29 + block_shape: list[int] | None,

30 ):

31 # All pplx byte sizes must be 16-byte aligned.

32 align = 16

@@ -82,10 +82,10 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

82 def activation_format(self) -> mk.FusedMoEActivationFormat:

83 return mk.FusedMoEActivationFormat.BatchedExperts

84

85 - def max_num_tokens_per_rank(self) -> Optional[int]:

85 + def max_num_tokens_per_rank(self) -> int | None:

86 return self.max_num_tokens

87

88 - def topk_indices_dtype(self) -> Optional[torch.dtype]:

88 + def topk_indices_dtype(self) -> torch.dtype | None:

89 return torch.uint32

90

91 def num_dispatchers(self) -> int:

@@ -103,7 +103,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

103 topk_weights: torch.Tensor,

104 topk_ids: torch.Tensor,

105 num_experts: int,

106 - expert_map: Optional[torch.Tensor],

106 + expert_map: torch.Tensor | None,

107 apply_router_weight_on_input: bool,

108 quant_config: FusedMoEQuantConfig,

109 ) -> tuple[Callable, mk.ReceiverType]:

@@ -148,7 +148,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

148 a1q, a1q_scale, quant_config.per_act_token_quant, quant_config.block_shape

149 )

150

151 - orig_a_scale_block_shape: Optional[int] = None

151 + orig_a_scale_block_shape: int | None = None

152

153 if a1q_scale is not None:

154 scalar_scales = a1q_scale.numel() == 1

@@ -184,7 +184,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

184 device=device,

185 )

186

187 - expert_x_scale: Optional[torch.Tensor] = None

187 + expert_x_scale: torch.Tensor | None = None

188 if a1q.dtype.itemsize == 1:

189 if quant_config.is_per_act_token:

190 # (M x 1) -> (E x M x K)

@@ -212,7 +212,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

212

213 # This argument is optional, defaults to indices.size(0)

214 # There's not much point setting this unless it is != indices.size(0)

215 - bound_m: Optional[torch.Tensor] = None

215 + bound_m: torch.Tensor | None = None

216

217 self.a2a.dispatch(

218 out_expert_num_tokens=expert_num_tokens,

@@ -252,8 +252,8 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

252 self,

253 expert_num_tokens: torch.Tensor,

254 expert_x: torch.Tensor,

255 - expert_x_scale: Optional[torch.Tensor],

256 - orig_a_scale_block_shape: Optional[int],

255 + expert_x_scale: torch.Tensor | None,

256 + orig_a_scale_block_shape: int | None,

257 ) -> mk.PrepareResultType:

258 if expert_x_scale is not None:

259 expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape]

@@ -271,7 +271,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

271 topk_weights: torch.Tensor,

272 topk_ids: torch.Tensor,

273 num_experts: int,

274 - expert_map: Optional[torch.Tensor],

274 + expert_map: torch.Tensor | None,

275 apply_router_weight_on_input: bool,

276 quant_config: FusedMoEQuantConfig,

277 ) -> mk.PrepareResultType:

@@ -302,7 +302,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

302

303 # This argument is optional

304 # There's not much point setting this unless it is != topk_ids.size(0)

305 - bound_m: Optional[torch.Tensor] = None

305 + bound_m: torch.Tensor | None = None

306

307 # TODO (bnell): fails in test_pplx_moe.py, figure out what's going on

308 # num_tokens = output.size(0) # M

vllm/model_executor/layers/fused_moe/prepare_finalize.py

+3 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -18,10 +17,10 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

17 def activation_format(self) -> mk.FusedMoEActivationFormat:

18 return mk.FusedMoEActivationFormat.Standard

19

21 - def max_num_tokens_per_rank(self) -> Optional[int]:

20 + def max_num_tokens_per_rank(self) -> int | None:

21 return None

22

24 - def topk_indices_dtype(self) -> Optional[torch.dtype]:

23 + def topk_indices_dtype(self) -> torch.dtype | None:

24 return None

25

26 def num_dispatchers(self) -> int:

@@ -36,7 +35,7 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

35 topk_weights: torch.Tensor,

36 topk_ids: torch.Tensor,

37 num_experts: int,

39 - expert_map: Optional[torch.Tensor],

38 + expert_map: torch.Tensor | None,

39 apply_router_weight_on_input: bool,

40 quant_config: FusedMoEQuantConfig,

41 ) -> mk.PrepareResultType:

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+25 -26

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

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

3 from enum import IntEnum

4 from functools import cache

5 -from typing import Optional

5

6 import torch

7

@@ -53,13 +52,13 @@ def rocm_aiter_asm_moe_tkw1_impl(

52 w2: torch.Tensor,

53 topk_weights: torch.Tensor,

54 topk_ids: torch.Tensor,

56 - fc1_scale: Optional[torch.Tensor] = None,

57 - fc2_scale: Optional[torch.Tensor] = None,

58 - fc1_smooth_scale: Optional[torch.Tensor] = None,

59 - fc2_smooth_scale: Optional[torch.Tensor] = None,

55 + fc1_scale: torch.Tensor | None = None,

56 + fc2_scale: torch.Tensor | None = None,

57 + fc1_smooth_scale: torch.Tensor | None = None,

58 + fc2_smooth_scale: torch.Tensor | None = None,

59 a16: bool = False,

61 - per_tensor_quant_scale: Optional[torch.Tensor] = None,

62 - expert_mask: Optional[torch.Tensor] = None,

60 + per_tensor_quant_scale: torch.Tensor | None = None,

61 + expert_mask: torch.Tensor | None = None,

62 activation_method: int = ActivationMethod.SILU.value,

63 ) -> torch.Tensor:

64 from aiter import ActivationType

@@ -90,13 +89,13 @@ def rocm_aiter_asm_moe_tkw1_fake(

89 w2: torch.Tensor,

90 topk_weights: torch.Tensor,

91 topk_ids: torch.Tensor,

93 - fc1_scale: Optional[torch.Tensor] = None,

94 - fc2_scale: Optional[torch.Tensor] = None,

95 - fc1_smooth_scale: Optional[torch.Tensor] = None,

96 - fc2_smooth_scale: Optional[torch.Tensor] = None,

92 + fc1_scale: torch.Tensor | None = None,

93 + fc2_scale: torch.Tensor | None = None,

94 + fc1_smooth_scale: torch.Tensor | None = None,

95 + fc2_smooth_scale: torch.Tensor | None = None,

96 a16: bool = False,

98 - per_tensor_quant_scale: Optional[torch.Tensor] = None,

99 - expert_mask: Optional[torch.Tensor] = None,

97 + per_tensor_quant_scale: torch.Tensor | None = None,

98 + expert_mask: torch.Tensor | None = None,

99 activation_method: int = ActivationMethod.SILU.value,

100 ) -> torch.Tensor:

101 return torch.empty_like(hidden_states)

@@ -206,14 +205,14 @@ def rocm_aiter_fused_moe_impl(

205 w2: torch.Tensor,

206 topk_weight: torch.Tensor,

207 topk_ids: torch.Tensor,

209 - expert_mask: Optional[torch.Tensor] = None,

208 + expert_mask: torch.Tensor | None = None,

209 activation_method: int = ActivationMethod.SILU.value,

210 quant_method: int = QuantMethod.NO.value,

211 doweight_stage1: bool = False,

213 - w1_scale: Optional[torch.Tensor] = None,

214 - w2_scale: Optional[torch.Tensor] = None,

215 - a1_scale: Optional[torch.Tensor] = None,

216 - a2_scale: Optional[torch.Tensor] = None,

212 + w1_scale: torch.Tensor | None = None,

213 + w2_scale: torch.Tensor | None = None,

214 + a1_scale: torch.Tensor | None = None,

215 + a2_scale: torch.Tensor | None = None,

216 ) -> torch.Tensor:

217 from aiter import ActivationType, QuantType

218 from aiter.fused_moe import fused_moe

@@ -244,14 +243,14 @@ def rocm_aiter_fused_moe_fake(

243 w2: torch.Tensor,

244 topk_weight: torch.Tensor,

245 topk_ids: torch.Tensor,

247 - expert_mask: Optional[torch.Tensor] = None,

246 + expert_mask: torch.Tensor | None = None,

247 activation_method: int = ActivationMethod.SILU.value,

248 quant_method: int = QuantMethod.NO.value,

249 doweight_stage1: bool = False,

251 - w1_scale: Optional[torch.Tensor] = None,

252 - w2_scale: Optional[torch.Tensor] = None,

253 - a1_scale: Optional[torch.Tensor] = None,

254 - a2_scale: Optional[torch.Tensor] = None,

250 + w1_scale: torch.Tensor | None = None,

251 + w2_scale: torch.Tensor | None = None,

252 + a1_scale: torch.Tensor | None = None,

253 + a2_scale: torch.Tensor | None = None,

254 ) -> torch.Tensor:

255 return torch.empty_like(hidden_states)

256

@@ -300,7 +299,7 @@ def rocm_aiter_grouped_topk(

299 topk_group: int = 0,

300 scoring_func: str = "softmax",

301 routed_scaling_factor: float = 1.0,

303 - e_score_correction_bias: Optional[torch.Tensor] = None,

302 + e_score_correction_bias: torch.Tensor | None = None,

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

304 token = hidden_states.shape[0]

305 device = hidden_states.device

@@ -342,8 +341,8 @@ def rocm_aiter_fused_experts(

341 topk_ids: torch.Tensor,

342 activation: str = "silu",

343 apply_router_weight_on_input: bool = False,

345 - expert_map: Optional[torch.Tensor] = None,

346 - quant_config: Optional[FusedMoEQuantConfig] = None,

344 + expert_map: torch.Tensor | None = None,

345 + quant_config: FusedMoEQuantConfig | None = None,

346 ) -> torch.Tensor:

347 if quant_config is None:

348 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

vllm/model_executor/layers/fused_moe/routing_simulator.py

+4 -4

@@ -10,7 +10,7 @@ like uniform random routing.

10 """

11

12 from abc import ABC, abstractmethod

13 -from typing import Any, Optional

13 +from typing import Any

14

15 import torch

16

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

24 hidden_states: torch.Tensor,

25 router_logits: torch.Tensor,

26 top_k: int,

27 - indices_type: Optional[torch.dtype] = None,

27 + indices_type: torch.dtype | None = None,

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

29 """

30 Route tokens to experts.

@@ -89,7 +89,7 @@ class DistributionBasedRouting(RoutingStrategy):

89 hidden_states: torch.Tensor,

90 router_logits: torch.Tensor,

91 top_k: int,

92 - indices_type: Optional[torch.dtype] = None,

92 + indices_type: torch.dtype | None = None,

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

94 """

95 Randomly select experts for each token using the specified distribution.

@@ -269,7 +269,7 @@ class RoutingSimulator:

269 router_logits: torch.Tensor,

270 strategy_name: str,

271 top_k: int,

272 - indices_type: Optional[torch.dtype] = None,

272 + indices_type: torch.dtype | None = None,

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

274 """

275 Simulate token-to-expert routing using the specified strategy.

vllm/model_executor/layers/fused_moe/shared_fused_moe.py

+2 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -18,7 +17,7 @@ class SharedFusedMoE(FusedMoE):

17

18 def __init__(

19 self,

21 - shared_experts: Optional[torch.nn.Module],

20 + shared_experts: torch.nn.Module | None,

21 use_overlapped: bool = True,

22 **kwargs,

23 ):

@@ -35,7 +34,7 @@ class SharedFusedMoE(FusedMoE):

34 )

35

36 @property

38 - def shared_experts(self) -> Optional[torch.nn.Module]:

37 + def shared_experts(self) -> torch.nn.Module | None:

38 return self._shared_experts if self.use_overlapped else None

39

40 def forward(

vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py

+4 -5

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6

@@ -29,7 +28,7 @@ class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce):

28

29 def apply(

30 self,

32 - output: Optional[torch.Tensor],

31 + output: torch.Tensor | None,

32 fused_expert_output: torch.Tensor,

33 topk_weights: torch.Tensor,

34 topk_ids: torch.Tensor,

@@ -52,7 +51,7 @@ class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):

51

52 def apply(

53 self,

55 - output: Optional[torch.Tensor],

54 + output: torch.Tensor | None,

55 fused_expert_output: torch.Tensor,

56 topk_weights: torch.Tensor,

57 topk_ids: torch.Tensor,

@@ -84,7 +83,7 @@ class TopKWeightAndReduceContiguous(mk.TopKWeightAndReduce):

83

84 def apply(

85 self,

87 - output: Optional[torch.Tensor],

86 + output: torch.Tensor | None,

87 fused_expert_output: torch.Tensor,

88 topk_weights: torch.Tensor,

89 topk_ids: torch.Tensor,

@@ -133,7 +132,7 @@ class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce):

132

133 def apply(

134 self,

136 - output: Optional[torch.Tensor],

135 + output: torch.Tensor | None,

136 fused_expert_output: torch.Tensor,

137 topk_weights: torch.Tensor,

138 topk_ids: torch.Tensor,

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+5 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -89,7 +88,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

88 topk: int,

89 global_num_experts: int,

90 local_num_experts: int,

92 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

91 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

92 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

93 # Note: the deep gemm workspaces are strictly larger than the triton

94 # workspaces so we can be pessimistic here and allocate for DeepGemm

@@ -128,12 +127,12 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

127 topk_ids: torch.Tensor,

128 activation: str,

129 global_num_experts: int,

131 - expert_map: Optional[torch.Tensor],

132 - a1q_scale: Optional[torch.Tensor],

133 - a2_scale: Optional[torch.Tensor],

130 + expert_map: torch.Tensor | None,

131 + a1q_scale: torch.Tensor | None,

132 + a2_scale: torch.Tensor | None,

133 workspace13: torch.Tensor,

134 workspace2: torch.Tensor,

136 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

135 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

136 apply_router_weight_on_input: bool,

137 ):

138 use_deep_gemm = self.allow_deep_gemm and (

vllm/model_executor/layers/fused_moe/trtllm_moe.py

+5 -6

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -58,7 +57,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

57 topk: int,

58 global_num_experts: int,

59 local_num_experts: int,

61 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

60 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

61 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

62 # The workspaces for this implementation are managed by flashinfer.

63 workspace1 = (0,)

@@ -100,12 +99,12 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

99 topk_ids: torch.Tensor,

100 activation: str,

101 global_num_experts: int,

103 - expert_map: Optional[torch.Tensor],

104 - a1q_scale: Optional[torch.Tensor],

105 - a2_scale: Optional[torch.Tensor],

102 + expert_map: torch.Tensor | None,

103 + a1q_scale: torch.Tensor | None,

104 + a2_scale: torch.Tensor | None,

105 workspace13: torch.Tensor,

106 workspace2: torch.Tensor,

108 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

107 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

108 apply_router_weight_on_input: bool,

109 ):

110 topk = topk_ids.size(-1)

vllm/model_executor/layers/fused_moe/utils.py

+23 -24

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from math import prod

4 -from typing import Optional, Union

4

5 import torch

6

@@ -60,7 +59,7 @@ def _count_expert_num_tokens(

59

60

61 def count_expert_num_tokens(

63 - topk_ids: torch.Tensor, num_local_experts: int, expert_map: Optional[torch.Tensor]

62 + topk_ids: torch.Tensor, num_local_experts: int, expert_map: torch.Tensor | None

63 ) -> torch.Tensor:

64 """

65 Count the number to tokens assigned to each expert.

@@ -112,7 +111,7 @@ def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:

111

112 def _nvfp4_quantize(

113 A: torch.Tensor,

115 - A_scale: Optional[torch.Tensor],

114 + A_scale: torch.Tensor | None,

115 is_sf_swizzled_layout: bool,

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

117 return flashinfer_fp4_quantize(

@@ -122,9 +121,9 @@ def _nvfp4_quantize(

121

122 def _fp8_quantize(

123 A: torch.Tensor,

125 - A_scale: Optional[torch.Tensor],

124 + A_scale: torch.Tensor | None,

125 per_act_token: bool,

127 - block_shape: Optional[list[int]] = None,

126 + block_shape: list[int] | None = None,

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

128 """

129 Perform fp8 quantization on the inputs. If a block_shape

@@ -148,9 +147,9 @@ def _fp8_quantize(

147

148 def _int8_quantize(

149 A: torch.Tensor,

151 - A_scale: Optional[torch.Tensor],

150 + A_scale: torch.Tensor | None,

151 per_act_token: bool,

153 - block_shape: Optional[list[int]] = None,

152 + block_shape: list[int] | None = None,

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

154 """

155 Perform int8 quantization on the inputs. If a block_shape

@@ -175,9 +174,9 @@ def _int8_quantize(

174

175 def _mxfp4_quantize(

176 A: torch.Tensor,

178 - A_scale: Optional[torch.Tensor],

177 + A_scale: torch.Tensor | None,

178 per_act_token_quant: bool,

180 - block_shape: Optional[list[int]] = None,

179 + block_shape: list[int] | None = None,

180 ) -> tuple[torch.Tensor, None]:

181 assert block_shape is None

182 # TODO: native mxfp4 is currently not integrated in vllm,

@@ -191,9 +190,9 @@ def _mxfp4_quantize(

190

191 def _mxfp8_e4m3_quantize(

192 A: torch.Tensor,

194 - A_scale: Optional[torch.Tensor],

193 + A_scale: torch.Tensor | None,

194 per_act_token_quant: bool,

196 - block_shape: Optional[list[int]] = None,

195 + block_shape: list[int] | None = None,

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

197 assert A_scale is None

198 assert not per_act_token_quant

@@ -203,9 +202,9 @@ def _mxfp8_e4m3_quantize(

202

203 def _mxfp6_e3m2_quantize(

204 A: torch.Tensor,

206 - A_scale: Optional[torch.Tensor],

205 + A_scale: torch.Tensor | None,

206 per_act_token_quant: bool,

208 - block_shape: Optional[list[int]] = None,

207 + block_shape: list[int] | None = None,

208 ) -> tuple[torch.Tensor, None]:

209 assert block_shape is None

210

@@ -220,9 +219,9 @@ def _mxfp6_e3m2_quantize(

219

220 def _mxfp6_e2m3_quantize(

221 A: torch.Tensor,

223 - A_scale: Optional[torch.Tensor],

222 + A_scale: torch.Tensor | None,

223 per_act_token_quant: bool,

225 - block_shape: Optional[list[int]] = None,

224 + block_shape: list[int] | None = None,

225 ) -> tuple[torch.Tensor, None]:

226 assert block_shape is None

227

@@ -237,12 +236,12 @@ def _mxfp6_e2m3_quantize(

236

237 def moe_kernel_quantize_input(

238 A: torch.Tensor,

240 - A_scale: Optional[torch.Tensor],

241 - quant_dtype: Union[None, torch.dtype, str],

239 + A_scale: torch.Tensor | None,

240 + quant_dtype: None | torch.dtype | str,

241 per_act_token_quant: bool,

243 - block_shape: Optional[list[int]] = None,

242 + block_shape: list[int] | None = None,

243 is_fp4_scale_swizzled: bool = True,

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

244 +) -> tuple[torch.Tensor, torch.Tensor | None]:

245 if quant_dtype == torch.float8_e4m3fn:

246 return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)

247 elif quant_dtype == torch.int8:

@@ -273,7 +272,7 @@ def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:

272 return m[idx, ...]

273

274

276 -def normalize_scales_shape(scales: Optional[torch.Tensor]) -> Optional[torch.Tensor]:

275 +def normalize_scales_shape(scales: torch.Tensor | None) -> torch.Tensor | None:

276 if scales is not None:

277 if scales.numel() == 1:

278 scales = scales.view(1, 1)

@@ -283,9 +282,9 @@ def normalize_scales_shape(scales: Optional[torch.Tensor]) -> Optional[torch.Ten

282

283

284 def normalize_batched_scales_shape(

286 - scales: Optional[torch.Tensor],

285 + scales: torch.Tensor | None,

286 num_experts: int,

288 -) -> Optional[torch.Tensor]:

287 +) -> torch.Tensor | None:

288 if scales is not None and scales.ndim < 3:

289 if scales.numel() == 1:

290 scales = scales.view(1)

@@ -300,9 +299,9 @@ def normalize_batched_scales_shape(

299

300 def _validate_scale_shape(

301 a: torch.Tensor,

303 - a_scale: Optional[torch.Tensor],

302 + a_scale: torch.Tensor | None,

303 per_act_token_quant: bool,

305 - block_shape: Optional[list[int]],

304 + block_shape: list[int] | None,

305 ) -> None:

306 if a_scale is None:

307 return

vllm/model_executor/layers/layernorm.py

+16 -18

@@ -2,8 +2,6 @@

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

3 """Custom normalization layers."""

4

5 -from typing import Optional, Union

6 -

5 import torch

6 import torch.nn as nn

7 import torch.nn.functional as F

@@ -159,9 +157,9 @@ class RMSNorm(CustomOp):

157 self,

158 hidden_size: int,

159 eps: float = 1e-6,

162 - var_hidden_size: Optional[int] = None,

160 + var_hidden_size: int | None = None,

161 has_weight: bool = True,

164 - dtype: Optional[torch.dtype] = None,

162 + dtype: torch.dtype | None = None,

163 ) -> None:

164 super().__init__()

165

@@ -190,8 +188,8 @@ class RMSNorm(CustomOp):

188 def forward_native(

189 self,

190 x: torch.Tensor,

193 - residual: Optional[torch.Tensor] = None,

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

191 + residual: torch.Tensor | None = None,

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

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

194 orig_dtype = x.dtype

195 x = x.to(torch.float32)

@@ -231,8 +229,8 @@ class RMSNorm(CustomOp):

229 def forward_cuda(

230 self,

231 x: torch.Tensor,

234 - residual: Optional[torch.Tensor] = None,

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

232 + residual: torch.Tensor | None = None,

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

234 if self.variance_size_override is not None:

235 return self.forward_native(x, residual)

236

@@ -247,8 +245,8 @@ class RMSNorm(CustomOp):

245 def forward_hip(

246 self,

247 x: torch.Tensor,

250 - residual: Optional[torch.Tensor] = None,

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

248 + residual: torch.Tensor | None = None,

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

250 if self.variance_size_override is not None:

251 return self.forward_native(x, residual)

252

@@ -263,8 +261,8 @@ class RMSNorm(CustomOp):

261 def forward_xpu(

262 self,

263 x: torch.Tensor,

266 - residual: Optional[torch.Tensor] = None,

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

264 + residual: torch.Tensor | None = None,

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

266 if self.variance_size_override is not None:

267 return self.forward_native(x, residual)

268

@@ -313,8 +311,8 @@ class GemmaRMSNorm(CustomOp):

311 weight: torch.Tensor,

312 variance_epsilon: float,

313 x: torch.Tensor,

316 - residual: Optional[torch.Tensor],

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

314 + residual: torch.Tensor | None,

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

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

317 orig_dtype = x.dtype

318 if residual is not None:

@@ -337,16 +335,16 @@ class GemmaRMSNorm(CustomOp):

335 def forward_native(

336 self,

337 x: torch.Tensor,

340 - residual: Optional[torch.Tensor] = None,

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

338 + residual: torch.Tensor | None = None,

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

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

341 return self.forward_static(self.weight.data, self.variance_epsilon, x, residual)

342

343 def forward_cuda(

344 self,

345 x: torch.Tensor,

348 - residual: Optional[torch.Tensor] = None,

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

346 + residual: torch.Tensor | None = None,

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

348 if torch.compiler.is_compiling():

349 return self.forward_native(x, residual)

350

vllm/model_executor/layers/lightning_attn.py

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5 from einops import rearrange

@@ -529,7 +528,7 @@ def lightning_attention(

528 v: torch.Tensor,

529 ed: torch.Tensor,

530 block_size: int = 256,

532 - kv_history: Optional[torch.Tensor] = None,

531 + kv_history: torch.Tensor | None = None,

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

533 """

534 Apply lightning attention algorithm

vllm/model_executor/layers/linear.py

+25 -25

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

3

4 import itertools

5 from abc import abstractmethod

6 -from typing import Any, Optional, Union

6 +from typing import Any

7

8 import torch

9 from torch.nn.parameter import Parameter, UninitializedParameter

@@ -187,7 +187,7 @@ class LinearMethodBase(QuantizeMethodBase):

187 self,

188 layer: torch.nn.Module,

189 x: torch.Tensor,

190 - bias: Optional[torch.Tensor] = None,

190 + bias: torch.Tensor | None = None,

191 ) -> torch.Tensor:

192 """Apply the weights in layer to the input tensor.

193 Expects create_weights to have been called before on the layer."""

@@ -252,7 +252,7 @@ class UnquantizedLinearMethod(LinearMethodBase):

252 self,

253 layer: torch.nn.Module,

254 x: torch.Tensor,

255 - bias: Optional[torch.Tensor] = None,

255 + bias: torch.Tensor | None = None,

256 ) -> torch.Tensor:

257 return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)

258

@@ -276,8 +276,8 @@ class LinearBase(CustomOp):

276 input_size: int,

277 output_size: int,

278 skip_bias_add: bool = False,

279 - params_dtype: Optional[torch.dtype] = None,

280 - quant_config: Optional[QuantizationConfig] = None,

279 + params_dtype: torch.dtype | None = None,

280 + quant_config: QuantizationConfig | None = None,

281 prefix: str = "",

282 *,

283 return_bias: bool = True,

@@ -295,7 +295,7 @@ class LinearBase(CustomOp):

295 self.quant_config = quant_config

296 self.prefix = prefix

297 if quant_config is None:

298 - self.quant_method: Optional[QuantizeMethodBase] = UnquantizedLinearMethod()

298 + self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod()

299 else:

300 self.quant_method = quant_config.get_quant_method(self, prefix=prefix)

301 self.return_bias = return_bias

@@ -333,8 +333,8 @@ class ReplicatedLinear(LinearBase):

333 output_size: int,

334 bias: bool = True,

335 skip_bias_add: bool = False,

336 - params_dtype: Optional[torch.dtype] = None,

337 - quant_config: Optional[QuantizationConfig] = None,

336 + params_dtype: torch.dtype | None = None,

337 + quant_config: QuantizationConfig | None = None,

338 prefix: str = "",

339 *,

340 return_bias: bool = True,

@@ -409,7 +409,7 @@ class ReplicatedLinear(LinearBase):

409 def forward(

410 self,

411 x: torch.Tensor,

412 - ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:

412 + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:

413 bias = self.bias if not self.skip_bias_add else None

414 assert self.quant_method is not None

415

@@ -461,9 +461,9 @@ class ColumnParallelLinear(LinearBase):

461 bias: bool = True,

462 gather_output: bool = False,

463 skip_bias_add: bool = False,

464 - params_dtype: Optional[torch.dtype] = None,

465 - quant_config: Optional[QuantizationConfig] = None,

466 - output_sizes: Optional[list[int]] = None,

464 + params_dtype: torch.dtype | None = None,

465 + quant_config: QuantizationConfig | None = None,

466 + output_sizes: list[int] | None = None,

467 prefix: str = "",

468 *,

469 return_bias: bool = True,

@@ -574,7 +574,7 @@ class ColumnParallelLinear(LinearBase):

574 def forward(

575 self,

576 input_,

577 - ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:

577 + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:

578 bias = self.bias if not self.skip_bias_add else None

579

580 # Matrix multiply.

@@ -633,8 +633,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

633 bias: bool = True,

634 gather_output: bool = False,

635 skip_bias_add: bool = False,

636 - params_dtype: Optional[torch.dtype] = None,

637 - quant_config: Optional[QuantizationConfig] = None,

636 + params_dtype: torch.dtype | None = None,

637 + quant_config: QuantizationConfig | None = None,

638 prefix: str = "",

639 *,

640 return_bias: bool = True,

@@ -662,7 +662,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

662 self,

663 param: Parameter,

664 loaded_weight: torch.Tensor,

665 - loaded_shard_id: Optional[int] = None,

665 + loaded_shard_id: int | None = None,

666 ):

667 # Special case for GGUF

668 # initialize GGUF param after we know the quantize type

@@ -838,7 +838,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

838 self,

839 param: BasevLLMParameter,

840 loaded_weight: torch.Tensor,

841 - loaded_shard_id: Optional[int] = None,

841 + loaded_shard_id: int | None = None,

842 ):

843 if loaded_shard_id is None:

844 if isinstance(param, PerTensorScaleParameter):

@@ -914,11 +914,11 @@ class QKVParallelLinear(ColumnParallelLinear):

914 hidden_size: int,

915 head_size: int,

916 total_num_heads: int,

917 - total_num_kv_heads: Optional[int] = None,

917 + total_num_kv_heads: int | None = None,

918 bias: bool = True,

919 skip_bias_add: bool = False,

920 - params_dtype: Optional[torch.dtype] = None,

921 - quant_config: Optional[QuantizationConfig] = None,

920 + params_dtype: torch.dtype | None = None,

921 + quant_config: QuantizationConfig | None = None,

922 prefix: str = "",

923 *,

924 return_bias: bool = True,

@@ -1027,7 +1027,7 @@ class QKVParallelLinear(ColumnParallelLinear):

1027 self,

1028 param: BasevLLMParameter,

1029 loaded_weight: torch.Tensor,

1030 - loaded_shard_id: Optional[str] = None,

1030 + loaded_shard_id: str | None = None,

1031 ):

1032 if loaded_shard_id is None: # special case for certain models

1033 if isinstance(param, PerTensorScaleParameter):

@@ -1071,7 +1071,7 @@ class QKVParallelLinear(ColumnParallelLinear):

1071 self,

1072 param: Parameter,

1073 loaded_weight: torch.Tensor,

1074 - loaded_shard_id: Optional[str] = None,

1074 + loaded_shard_id: str | None = None,

1075 ):

1076 # Special case for GGUF

1077 # initialize GGUF param after we know the quantize type

@@ -1296,9 +1296,9 @@ class RowParallelLinear(LinearBase):

1296 bias: bool = True,

1297 input_is_parallel: bool = True,

1298 skip_bias_add: bool = False,

1299 - params_dtype: Optional[torch.dtype] = None,

1299 + params_dtype: torch.dtype | None = None,

1300 reduce_results: bool = True,

1301 - quant_config: Optional[QuantizationConfig] = None,

1301 + quant_config: QuantizationConfig | None = None,

1302 prefix: str = "",

1303 *,

1304 return_bias: bool = True,

@@ -1405,7 +1405,7 @@ class RowParallelLinear(LinearBase):

1405 def forward(

1406 self,

1407 input_,

1408 - ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:

1408 + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:

1409 if self.input_is_parallel:

1410 input_parallel = input_

1411 else:

vllm/model_executor/layers/logits_processor.py

+6 -8

@@ -2,8 +2,6 @@

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

3 """A layer that compute logits from hidden_stats."""

4

5 -from typing import Optional

6 -

5 import torch

6

7 from vllm.distributed import (

@@ -28,10 +26,10 @@ class LogitsProcessor(CustomOp):

26 def __init__(

27 self,

28 vocab_size: int,

31 - org_vocab_size: Optional[int] = None,

29 + org_vocab_size: int | None = None,

30 scale: float = 1.0,

31 logits_as_input: bool = False,

34 - soft_cap: Optional[float] = None,

32 + soft_cap: float | None = None,

33 ) -> None:

34 """

35 Args:

@@ -53,8 +51,8 @@ class LogitsProcessor(CustomOp):

51 self,

52 lm_head: VocabParallelEmbedding,

53 hidden_states: torch.Tensor,

56 - embedding_bias: Optional[torch.Tensor] = None,

57 - ) -> Optional[torch.Tensor]:

54 + embedding_bias: torch.Tensor | None = None,

55 + ) -> torch.Tensor | None:

56 if self.logits_as_input:

57 logits = hidden_states

58 else:

@@ -88,8 +86,8 @@ class LogitsProcessor(CustomOp):

86 self,

87 hidden_states: torch.Tensor,

88 lm_head: VocabParallelEmbedding,

91 - embedding_bias: Optional[torch.Tensor],

92 - ) -> Optional[torch.Tensor]:

89 + embedding_bias: torch.Tensor | None,

90 + ) -> torch.Tensor | None:

91 # Get the logits for the next tokens.

92 logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)

93

vllm/model_executor/layers/mamba/linear_attn.py

+7 -7

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

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

3

4 import math

5 -from typing import TYPE_CHECKING, Optional, Union

5 +from typing import TYPE_CHECKING

6

7 if TYPE_CHECKING:

8 from vllm.attention.backends.abstract import AttentionBackend

@@ -87,8 +87,8 @@ class MiniMaxText01RMSNormTP(CustomOp):

87 def forward(

88 self,

89 x: torch.Tensor,

90 - residual: Optional[torch.Tensor] = None,

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

90 + residual: torch.Tensor | None = None,

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

92 assert residual is None, "RMSNorm does not support residual connection."

93 return self._forward(x)

94

@@ -102,7 +102,7 @@ class MiniMaxText01LinearKernel:

102 kv_caches: torch.Tensor,

103 slope_rate: torch.Tensor,

104 block_size: int,

105 - layer_idx: Optional[int] = None,

105 + layer_idx: int | None = None,

106 **kwargs,

107 ) -> torch.Tensor:

108 slope_rate = slope_rate.to(torch.float32)

@@ -154,9 +154,9 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):

154 max_position: int,

155 block_size: int,

156 num_hidden_layer: int,

157 - model_config: Optional[ModelConfig] = None,

158 - cache_config: Optional[CacheConfig] = None,

159 - quant_config: Optional[QuantizationConfig] = None,

157 + model_config: ModelConfig | None = None,

158 + cache_config: CacheConfig | None = None,

159 + quant_config: QuantizationConfig | None = None,

160 layer_idx: int = 0,

161 linear_layer_idx: int = 0,

162 prefix: str = "linear_attn",

vllm/model_executor/layers/mamba/mamba_mixer.py

+7 -7

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import TYPE_CHECKING, NamedTuple, Optional

4 +from typing import TYPE_CHECKING, NamedTuple

5

6 if TYPE_CHECKING:

7 from vllm.attention.backends.abstract import AttentionBackend

@@ -68,8 +68,8 @@ class MambaMixer(MambaBase, CustomOp):

68 rms_norm_eps: float = 1e-5,

69 activation="silu",

70 is_lora_enabled: bool = False,

71 - model_config: Optional[ModelConfig] = None,

72 - cache_config: Optional[CacheConfig] = None,

71 + model_config: ModelConfig | None = None,

72 + cache_config: CacheConfig | None = None,

73 prefix: str = "",

74 ):

75 super().__init__()

@@ -410,7 +410,7 @@ class MambaMixer(MambaBase, CustomOp):

410

411 return Mamba1AttentionBackend

412

413 - def _time_proj_bias(self) -> Optional[torch.Tensor]:

413 + def _time_proj_bias(self) -> torch.Tensor | None:

414 if hasattr(self.dt_proj, "bias") and self.dt_proj.bias is not None:

415 return self.dt_proj.bias.float()

416 return None

@@ -423,8 +423,8 @@ class PrefillDecodeSplit(NamedTuple):

423 gate_d: torch.Tensor

424 state_indices_tensor_p: torch.Tensor

425 state_indices_tensor_d: torch.Tensor

426 - query_start_loc_p: Optional[torch.Tensor]

427 - has_initial_states_p: Optional[torch.Tensor]

426 + query_start_loc_p: torch.Tensor | None

427 + has_initial_states_p: torch.Tensor | None

428

429

430 def split_batch_to_prefill_and_decode(

@@ -432,7 +432,7 @@ def split_batch_to_prefill_and_decode(

432 gate: torch.Tensor,

433 state_indices_tensor: torch.Tensor,

434 query_start_loc: torch.Tensor,

435 - has_initial_states: Optional[torch.Tensor],

435 + has_initial_states: torch.Tensor | None,

436 num_prefill_tokens: int,

437 num_decode_tokens: int,

438 num_prefills: int,

vllm/model_executor/layers/mamba/mamba_mixer2.py

+10 -10

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import TYPE_CHECKING, Optional, Union

4 +from typing import TYPE_CHECKING

5

6 if TYPE_CHECKING:

7 from vllm.attention.backends.abstract import AttentionBackend

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

138 self,

139 x: torch.Tensor,

140 gate: torch.Tensor,

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

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

142 input_dtype = x.dtype

143 if not self.use_rms_norm:

144 # Keep gate in float32 for numerical stability during silu

@@ -244,9 +244,9 @@ class MambaMixer2(MambaBase, CustomOp):

244 rms_norm_eps: float = 1e-5,

245 activation: str = "silu",

246 use_rms_norm: bool = True,

247 - model_config: Optional[ModelConfig] = None,

248 - cache_config: Optional[CacheConfig] = None,

249 - quant_config: Optional[QuantizationConfig] = None,

247 + model_config: ModelConfig | None = None,

248 + cache_config: CacheConfig | None = None,

249 + quant_config: QuantizationConfig | None = None,

250 prefix: str = "",

251 ):

252 super().__init__()

@@ -474,7 +474,7 @@ class MambaMixer2(MambaBase, CustomOp):

474 self,

475 hidden_states: torch.Tensor,

476 output: torch.Tensor,

477 - mup_vector: Optional[torch.Tensor] = None,

477 + mup_vector: torch.Tensor | None = None,

478 ):

479 pass

480

@@ -482,7 +482,7 @@ class MambaMixer2(MambaBase, CustomOp):

482 self,

483 hidden_states: torch.Tensor,

484 output: torch.Tensor,

485 - mup_vector: Optional[torch.Tensor] = None,

485 + mup_vector: torch.Tensor | None = None,

486 ):

487 torch.ops.vllm.mamba_mixer2(

488 hidden_states,

@@ -495,7 +495,7 @@ class MambaMixer2(MambaBase, CustomOp):

495 self,

496 hidden_states: torch.Tensor,

497 output: torch.Tensor,

498 - mup_vector: Optional[torch.Tensor] = None,

498 + mup_vector: torch.Tensor | None = None,

499 ):

500 forward_context = get_forward_context()

501 # attn_metadata contains metadata necessary for the mamba2 triton

@@ -904,7 +904,7 @@ def mamba_mixer2(

904 hidden_states: torch.Tensor,

905 output: torch.Tensor,

906 layer_name: str,

907 - mup_vector: Optional[torch.Tensor] = None,

907 + mup_vector: torch.Tensor | None = None,

908 ) -> None:

909 forward_context: ForwardContext = get_forward_context()

910 self = forward_context.no_compile_layers[layer_name]

@@ -915,7 +915,7 @@ def mamba_mixer2_fake(

915 hidden_states: torch.Tensor,

916 output: torch.Tensor,

917 layer_name: str,

918 - mup_vector: Optional[torch.Tensor] = None,

918 + mup_vector: torch.Tensor | None = None,

919 ) -> None:

920 return

921

vllm/model_executor/layers/mamba/mamba_utils.py

+6 -7

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Union

3

4 import torch

5

@@ -14,7 +13,7 @@ class MambaStateDtypeCalculator:

13 @classmethod

14 def linear_attention_state_dtype(

15 cls,

17 - model_dtype: Union[ModelDType, torch.dtype],

16 + model_dtype: ModelDType | torch.dtype,

17 mamba_cache_dtype: MambaDType,

18 ) -> tuple[torch.dtype, ...]:

19 # TODO (tdoublep) requires testing

@@ -26,7 +25,7 @@ class MambaStateDtypeCalculator:

25 @classmethod

26 def mamba1_state_dtype(

27 cls,

29 - model_dtype: Union[ModelDType, torch.dtype],

28 + model_dtype: ModelDType | torch.dtype,

29 mamba_cache_dtype: MambaDType,

30 mamba_ssm_cache_dtype: MambaDType,

31 ) -> tuple[torch.dtype, ...]:

@@ -37,7 +36,7 @@ class MambaStateDtypeCalculator:

36 @classmethod

37 def mamba2_state_dtype(

38 cls,

40 - model_dtype: Union[ModelDType, torch.dtype],

39 + model_dtype: ModelDType | torch.dtype,

40 mamba_cache_dtype: MambaDType,

41 mamba_ssm_cache_dtype: MambaDType,

42 ) -> tuple[torch.dtype, ...]:

@@ -48,7 +47,7 @@ class MambaStateDtypeCalculator:

47 @classmethod

48 def _mamba_state_dtype(

49 cls,

51 - model_dtype: Union[ModelDType, torch.dtype],

50 + model_dtype: ModelDType | torch.dtype,

51 mamba_cache_dtype: MambaDType,

52 mamba_ssm_cache_dtype: MambaDType,

53 ) -> tuple[torch.dtype, ...]:

@@ -63,7 +62,7 @@ class MambaStateDtypeCalculator:

62 @classmethod

63 def short_conv_state_dtype(

64 cls,

66 - model_dtype: Union[ModelDType, torch.dtype],

65 + model_dtype: ModelDType | torch.dtype,

66 mamba_cache_dtype: MambaDType,

67 ) -> tuple[torch.dtype, ...]:

68 conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)

@@ -72,7 +71,7 @@ class MambaStateDtypeCalculator:

71 @classmethod

72 def gated_delta_net_state_dtype(

73 cls,

75 - model_dtype: Union[ModelDType, torch.dtype],

74 + model_dtype: ModelDType | torch.dtype,

75 mamba_cache_dtype: MambaDType,

76 ) -> tuple[torch.dtype, torch.dtype]:

77 state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

+15 -16

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

4 # Copyright (c) 2024, Tri Dao.

5 # Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py

6

7 -from typing import Optional, Union

7

8 import numpy as np

9 import torch

@@ -469,17 +468,17 @@ def _causal_conv1d_fwd_kernel( # continuous batching

468 def causal_conv1d_fn(

469 x: torch.Tensor,

470 weight: torch.Tensor,

472 - bias: Union[torch.Tensor, None],

471 + bias: torch.Tensor | None,

472 conv_states: torch.Tensor,

473 query_start_loc: torch.Tensor,

475 - cache_indices: Optional[torch.Tensor] = None,

476 - has_initial_state: Optional[torch.Tensor] = None,

477 - activation: Optional[str] = "silu",

474 + cache_indices: torch.Tensor | None = None,

475 + has_initial_state: torch.Tensor | None = None,

476 + activation: str | None = "silu",

477 pad_slot_id: int = PAD_SLOT_ID,

479 - block_idx_first_scheduled_token: Optional[torch.Tensor] = None,

480 - block_idx_last_scheduled_token: Optional[torch.Tensor] = None,

481 - initial_state_idx: Optional[torch.Tensor] = None,

482 - num_computed_tokens: Optional[torch.Tensor] = None,

478 + block_idx_first_scheduled_token: torch.Tensor | None = None,

479 + block_idx_last_scheduled_token: torch.Tensor | None = None,

480 + initial_state_idx: torch.Tensor | None = None,

481 + num_computed_tokens: torch.Tensor | None = None,

482 block_size_to_align=0,

483 metadata=None,

484 validate_data=False,

@@ -1071,15 +1070,15 @@ def causal_conv1d_update(

1070 x: torch.Tensor,

1071 conv_state: torch.Tensor,

1072 weight: torch.Tensor,

1074 - bias: Optional[torch.Tensor] = None,

1075 - activation: Union[bool, str, None] = None,

1076 - conv_state_indices: Optional[torch.Tensor] = None,

1077 - num_accepted_tokens: Optional[torch.Tensor] = None,

1078 - query_start_loc: Optional[torch.Tensor] = None,

1073 + bias: torch.Tensor | None = None,

1074 + activation: bool | str | None = None,

1075 + conv_state_indices: torch.Tensor | None = None,

1076 + num_accepted_tokens: torch.Tensor | None = None,

1077 + query_start_loc: torch.Tensor | None = None,

1078 max_query_len: int = -1,

1079 pad_slot_id: int = PAD_SLOT_ID,

1081 - block_idx_last_scheduled_token: Optional[torch.Tensor] = None,

1082 - initial_state_idx: Optional[torch.Tensor] = None,

1080 + block_idx_last_scheduled_token: torch.Tensor | None = None,

1081 + initial_state_idx: torch.Tensor | None = None,

1082 validate_data=False,

1083 ):

1084 """

vllm/model_executor/layers/mamba/short_conv.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import TYPE_CHECKING, Optional

4 +from typing import TYPE_CHECKING

5

6 if TYPE_CHECKING:

7 from vllm.attention.backends.abstract import AttentionBackend

@@ -38,8 +38,8 @@ class ShortConv(MambaBase, CustomOp):

38 config,

39 dim: int,

40 layer_idx: int,

41 - model_config: Optional[ModelConfig] = None,

42 - cache_config: Optional[CacheConfig] = None,

41 + model_config: ModelConfig | None = None,

42 + cache_config: CacheConfig | None = None,

43 prefix: str = "",

44 ):

45 super().__init__()

vllm/model_executor/layers/mla.py

+10 -11

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass

4 -from typing import Optional

4

5 import torch

6

@@ -19,14 +18,14 @@ class MLAModules:

18 kv_b_proj: torch.nn.Module

19 rotary_emb: torch.nn.Module

20 o_proj: torch.nn.Module

22 - fused_qkv_a_proj: Optional[torch.nn.Module]

23 - kv_a_proj_with_mqa: Optional[torch.nn.Module]

24 - q_a_layernorm: Optional[torch.nn.Module]

25 - q_b_proj: Optional[torch.nn.Module]

26 - q_proj: Optional[torch.nn.Module]

27 - indexer: Optional[torch.nn.Module]

21 + fused_qkv_a_proj: torch.nn.Module | None

22 + kv_a_proj_with_mqa: torch.nn.Module | None

23 + q_a_layernorm: torch.nn.Module | None

24 + q_b_proj: torch.nn.Module | None

25 + q_proj: torch.nn.Module | None

26 + indexer: torch.nn.Module | None

27 is_sparse: bool

29 - topk_indices_buffer: Optional[torch.Tensor]

28 + topk_indices_buffer: torch.Tensor | None

29

30

31 @CustomOp.register("multi_head_latent_attention")

@@ -55,11 +54,11 @@ class MultiHeadLatentAttentionWrapper(CustomOp):

54 qk_nope_head_dim: int,

55 qk_rope_head_dim: int,

56 v_head_dim: int,

58 - q_lora_rank: Optional[int],

57 + q_lora_rank: int | None,

58 kv_lora_rank: int,

59 mla_modules: MLAModules,

61 - cache_config: Optional[CacheConfig] = None,

62 - quant_config: Optional[QuantizationConfig] = None,

60 + cache_config: CacheConfig | None = None,

61 + quant_config: QuantizationConfig | None = None,

62 prefix: str = "",

63 ) -> None:

64 super().__init__()

vllm/model_executor/layers/pooler.py

+27 -27

@@ -1,11 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 from abc import ABC, abstractmethod

4 -from collections.abc import Mapping, Set

4 +from collections.abc import Callable, Mapping, Set

5 from dataclasses import dataclass

6 from enum import IntEnum

7 from itertools import groupby

8 -from typing import Callable, Optional, TypeVar, Union

8 +from typing import TypeVar

9

10 import torch

11 import torch.nn as nn

@@ -24,8 +24,8 @@ from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata

24 logger = init_logger(__name__)

25

26 PoolingFn = Callable[

27 - [Union[torch.Tensor, list[torch.Tensor]], PoolingMetadata],

28 - Union[torch.Tensor, list[torch.Tensor]],

27 + [torch.Tensor | list[torch.Tensor], PoolingMetadata],

28 + torch.Tensor | list[torch.Tensor],

29 ]

30 ClassifierFn = Callable[[torch.Tensor], torch.Tensor]

31

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

90 @staticmethod

91 def for_classify(

92 pooler_config: PoolerConfig,

93 - classifier: Optional[ClassifierFn],

93 + classifier: ClassifierFn | None,

94 ):

95 resolved_config = ResolvedPoolingConfig.from_config(

96 task="classify",

@@ -118,14 +118,14 @@ class Pooler(nn.Module, ABC):

118 @abstractmethod

119 def forward(

120 self,

121 - hidden_states: Union[list[torch.Tensor], torch.Tensor],

121 + hidden_states: list[torch.Tensor] | torch.Tensor,

122 pooling_metadata: PoolingMetadata,

123 ) -> PoolerOutput:

124 raise NotImplementedError

125

126

127 def get_prompt_lens(

128 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

128 + hidden_states: torch.Tensor | list[torch.Tensor],

129 pooling_metadata: PoolingMetadata,

130 ) -> torch.Tensor:

131 return pooling_metadata.prompt_lens

@@ -174,7 +174,7 @@ def get_classification_activation_function(config: PretrainedConfig):

174

175

176 def get_cross_encoder_activation_function(config: PretrainedConfig):

177 - function_name: Optional[str] = None

177 + function_name: str | None = None

178 if (

179 hasattr(config, "sentence_transformers")

180 and "activation_fn" in config.sentence_transformers

@@ -223,14 +223,14 @@ class PoolingMethod(nn.Module, ABC):

223 self,

224 hidden_states: torch.Tensor,

225 pooling_cursor: PoolingCursor,

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

226 + ) -> list[torch.Tensor] | torch.Tensor:

227 raise NotImplementedError

228

229 def forward(

230 self,

231 hidden_states: torch.Tensor,

232 pooling_metadata: PoolingMetadata,

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

233 + ) -> list[torch.Tensor] | torch.Tensor:

234 pooling_cursor = pooling_metadata.pooling_cursor

235 return self.forward_all(hidden_states, pooling_cursor)

236

@@ -243,7 +243,7 @@ class CLSPool(PoolingMethod):

243 self,

244 hidden_states: torch.Tensor,

245 pooling_cursor: PoolingCursor,

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

246 + ) -> list[torch.Tensor] | torch.Tensor:

247 assert not pooling_cursor.is_partial_prefill(), (

248 "partial prefill not supported with CLS pooling"

249 )

@@ -259,7 +259,7 @@ class LastPool(PoolingMethod):

259 self,

260 hidden_states: torch.Tensor,

261 pooling_cursor: PoolingCursor,

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

262 + ) -> list[torch.Tensor] | torch.Tensor:

263 return hidden_states[pooling_cursor.last_token_indices_gpu]

264

265

@@ -271,7 +271,7 @@ class AllPool(PoolingMethod):

271 self,

272 hidden_states: torch.Tensor,

273 pooling_cursor: PoolingCursor,

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

274 + ) -> list[torch.Tensor] | torch.Tensor:

275 assert not pooling_cursor.is_partial_prefill(), (

276 "partial prefill not supported with ALL pooling"

277 )

@@ -290,7 +290,7 @@ class MeanPool(PoolingMethod):

290 self,

291 hidden_states: torch.Tensor,

292 pooling_cursor: PoolingCursor,

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

293 + ) -> list[torch.Tensor] | torch.Tensor:

294 assert not pooling_cursor.is_partial_prefill(), (

295 "partial prefill not supported with MEAN pooling"

296 )

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

405

406 def forward(

407 self,

408 - pooled_data: Union[list[torch.Tensor], torch.Tensor],

408 + pooled_data: list[torch.Tensor] | torch.Tensor,

409 pooling_metadata: PoolingMetadata,

410 ):

411 return self.activation(pooled_data)

@@ -418,14 +418,14 @@ class EmbeddingPoolerHead(PoolerHead):

418 # Load ST projector if available

419

420 vllm_config = get_current_vllm_config()

421 - self.projector: Optional[nn.Module] = (

421 + self.projector: nn.Module | None = (

422 _load_st_projector(vllm_config.model_config) if vllm_config else None

423 )

424 self.head_dtype = vllm_config.model_config.head_dtype

425

426 def forward(

427 self,

428 - pooled_data: Union[list[torch.Tensor], torch.Tensor],

428 + pooled_data: list[torch.Tensor] | torch.Tensor,

429 pooling_metadata: PoolingMetadata,

430 ):

431 if isinstance(pooled_data, list):

@@ -480,7 +480,7 @@ class RewardPoolerHead(PoolerHead):

480

481 def forward(

482 self,

483 - pooled_data: Union[list[torch.Tensor], torch.Tensor],

483 + pooled_data: list[torch.Tensor] | torch.Tensor,

484 pooling_metadata: PoolingMetadata,

485 ):

486 if isinstance(pooled_data, list):

@@ -541,7 +541,7 @@ class SimplePooler(Pooler):

541

542 def forward(

543 self,

544 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

544 + hidden_states: torch.Tensor | list[torch.Tensor],

545 pooling_metadata: PoolingMetadata,

546 ) -> PoolerOutput:

547 pooled_data = self.pooling(hidden_states, pooling_metadata)

@@ -560,9 +560,9 @@ class StepPooler(Pooler):

560

561 def extract_states(

562 self,

563 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

563 + hidden_states: torch.Tensor | list[torch.Tensor],

564 pooling_metadata: PoolingMetadata,

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

565 + ) -> list[torch.Tensor] | torch.Tensor:

566 pooled_data_lst = self.pooling(hidden_states, pooling_metadata)

567 prompt_token_ids = get_prompt_token_ids(pooling_metadata)

568

@@ -593,7 +593,7 @@ class StepPooler(Pooler):

593

594 def forward(

595 self,

596 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

596 + hidden_states: torch.Tensor | list[torch.Tensor],

597 pooling_metadata: PoolingMetadata,

598 ) -> PoolerOutput:

599 pooled_data = self.extract_states(hidden_states, pooling_metadata)

@@ -621,8 +621,8 @@ class ClassifierPooler(Pooler):

621 def __init__(

622 self,

623 pooling: PoolingFn,

624 - classifier: Optional[ClassifierFn],

625 - act_fn: Optional[PoolerActivation] = None,

624 + classifier: ClassifierFn | None,

625 + act_fn: PoolerActivation | None = None,

626 ) -> None:

627 super().__init__()

628

@@ -631,7 +631,7 @@ class ClassifierPooler(Pooler):

631 self.pooling = pooling

632 self.classifier = classifier

633 self.act_fn = act_fn or PoolerClassify()

634 - self.logit_bias: Optional[float] = (

634 + self.logit_bias: float | None = (

635 vllm_config.model_config.pooler_config.logit_bias

636 )

637 self.head_dtype = vllm_config.model_config.head_dtype

@@ -641,7 +641,7 @@ class ClassifierPooler(Pooler):

641

642 def forward(

643 self,

644 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

644 + hidden_states: torch.Tensor | list[torch.Tensor],

645 pooling_metadata: PoolingMetadata,

646 ) -> PoolerOutput:

647 pooled_data = self.pooling(hidden_states, pooling_metadata)

@@ -695,7 +695,7 @@ class DispatchPooler(Pooler):

695

696 def forward(

697 self,

698 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

698 + hidden_states: torch.Tensor | list[torch.Tensor],

699 pooling_metadata: PoolingMetadata,

700 ) -> PoolerOutput:

701 poolers_by_task = self.poolers_by_task

vllm/model_executor/layers/quantization/auto_round.py

+3 -3

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

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

3

4 from fractions import Fraction

5 -from typing import TYPE_CHECKING, Any, Optional, Union

5 +from typing import TYPE_CHECKING, Any

6

7 import torch

8

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

46 group_size: int,

47 sym: bool = True,

48 packing_format: str = "auto_round:auto_gptq",

49 - block_name_to_quantize: Optional[Union[str, list[str]]] = None,

50 - extra_config: Optional[dict[str, Any]] = None,

49 + block_name_to_quantize: str | list[str] | None = None,

50 + extra_config: dict[str, Any] | None = None,

51 data_type: str = "int",

52 backend: str = "auto",

53 ) -> None:

vllm/model_executor/layers/quantization/awq.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Optional, Union

4 +from typing import Any, Union

5

6 import torch

7

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

34 weight_bits: int,

35 group_size: int,

36 zero_point: bool,

37 - modules_to_not_convert: Optional[list[str]] = None,

37 + modules_to_not_convert: list[str] | None = None,

38 ) -> None:

39 super().__init__()

40 self.weight_bits = weight_bits

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

88

89 def get_quant_method(

90 self, layer: torch.nn.Module, prefix: str

91 - ) -> Optional[Union["LinearMethodBase", "QuantizeMethodBase"]]:

91 + ) -> Union["LinearMethodBase", "QuantizeMethodBase"] | None:

92 if isinstance(layer, LinearBase):

93 if is_layer_skipped_awq(prefix, self.modules_to_not_convert):

94 return UnquantizedLinearMethod()

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

227 self,

228 layer: torch.nn.Module,

229 x: torch.Tensor,

230 - bias: Optional[torch.Tensor] = None,

230 + bias: torch.Tensor | None = None,

231 ) -> torch.Tensor:

232 qweight = layer.qweight

233 scales = layer.scales

vllm/model_executor/layers/quantization/awq_marlin.py

+15 -14

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Callable, Optional, Union

4 +from collections.abc import Callable

5 +from typing import Any, Optional

6

7 import torch

8 from torch.nn import Parameter

@@ -70,7 +71,7 @@ class AWQMarlinConfig(QuantizationConfig):

71 group_size: int,

72 zero_point: bool,

73 lm_head_quantized: bool,

73 - modules_to_not_convert: Optional[list[str]],

74 + modules_to_not_convert: list[str] | None,

75 full_config: dict[str, Any],

76 ) -> None:

77 super().__init__()

@@ -140,7 +141,7 @@ class AWQMarlinConfig(QuantizationConfig):

141 @classmethod

142 def override_quantization_method(

143 cls, hf_quant_cfg, user_quant

143 - ) -> Optional[QuantizationMethods]:

144 + ) -> QuantizationMethods | None:

145 can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg)

146 is_valid_user_quant = (

147 user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin"

@@ -360,7 +361,7 @@ class AWQMarlinLinearMethod(LinearMethodBase):

361 self,

362 layer: torch.nn.Module,

363 x: torch.Tensor,

363 - bias: Optional[torch.Tensor] = None,

364 + bias: torch.Tensor | None = None,

365 ) -> torch.Tensor:

366 return apply_awq_marlin_linear(

367 input=x,

@@ -555,7 +556,7 @@ class AWQMoEMethod(FusedMoEMethodBase):

556

557 def get_fused_moe_quant_config(

558 self, layer: torch.nn.Module

558 - ) -> Optional[FusedMoEQuantConfig]:

559 + ) -> FusedMoEQuantConfig | None:

560 return None

561

562 def apply(

@@ -566,21 +567,21 @@ class AWQMoEMethod(FusedMoEMethodBase):

567 top_k: int,

568 renormalize: bool,

569 use_grouped_topk: bool = False,

569 - topk_group: Optional[int] = None,

570 - num_expert_group: Optional[int] = None,

570 + topk_group: int | None = None,

571 + num_expert_group: int | None = None,

572 global_num_experts: int = -1,

572 - expert_map: Optional[torch.Tensor] = None,

573 - custom_routing_function: Optional[Callable] = None,

573 + expert_map: torch.Tensor | None = None,

574 + custom_routing_function: Callable | None = None,

575 scoring_func: str = "softmax",

576 routed_scaling_factor: float = 1.0,

576 - e_score_correction_bias: Optional[torch.Tensor] = None,

577 + e_score_correction_bias: torch.Tensor | None = None,

578 apply_router_weight_on_input: bool = False,

579 activation: str = "silu",

580 enable_eplb: bool = False,

580 - expert_load_view: Optional[torch.Tensor] = None,

581 - logical_to_physical_map: Optional[torch.Tensor] = None,

582 - logical_replica_count: Optional[torch.Tensor] = None,

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

581 + expert_load_view: torch.Tensor | None = None,

582 + logical_to_physical_map: torch.Tensor | None = None,

583 + logical_replica_count: torch.Tensor | None = None,

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

585 assert self.fused_experts is None

586

587 if enable_eplb:

vllm/model_executor/layers/quantization/base_config.py

+4 -4

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

3

4 import inspect

5 from abc import ABC, abstractmethod

6 -from typing import TYPE_CHECKING, Any, Optional

6 +from typing import TYPE_CHECKING, Any

7

8 import torch

9 from torch import nn

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

105 @classmethod

106 def override_quantization_method(

107 cls, hf_quant_cfg, user_quant

108 - ) -> Optional[QuantizationMethods]:

108 + ) -> QuantizationMethods | None:

109 """

110 Detects if this quantization method can support a given checkpoint

111 format by overriding the user specified quantization method --

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

135 @abstractmethod

136 def get_quant_method(

137 self, layer: torch.nn.Module, prefix: str

138 - ) -> Optional[QuantizeMethodBase]:

138 + ) -> QuantizeMethodBase | None:

139 """Get the quantize method to use for the quantized layer.

140

141 Args:

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

147 """

148 raise NotImplementedError

149

150 - def get_cache_scale(self, name: str) -> Optional[str]:

150 + def get_cache_scale(self, name: str) -> str | None:

151 return None

152

153 def apply_vllm_mapper( # noqa: B027

vllm/model_executor/layers/quantization/bitblas.py

+6 -6

@@ -45,10 +45,10 @@ class BitBLASConfig(QuantizationConfig):

45 def __init__(

46 self,

47 weight_bits: int,

48 - group_size: Optional[int],

49 - desc_act: Optional[bool],

50 - is_sym: Optional[bool],

51 - quant_method: Optional[str],

48 + group_size: int | None,

49 + desc_act: bool | None,

50 + is_sym: bool | None,

51 + quant_method: str | None,

52 lm_head_quantized: bool,

53 ) -> None:

54 try:

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

160 @classmethod

161 def override_quantization_method(

162 cls, hf_quant_cfg, user_quant

163 - ) -> Optional[QuantizationMethods]:

163 + ) -> QuantizationMethods | None:

164 # compat: autogptq >=0.8.0 use checkpoint_format: str

165 # compat: autogptq <=0.7.1 is_bitblas_format: bool

166 is_bitblas_format = hf_quant_cfg.get(

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

469 self,

470 layer: torch.nn.Module,

471 x: torch.Tensor,

472 - bias: Optional[torch.Tensor] = None,

472 + bias: torch.Tensor | None = None,

473 ) -> torch.Tensor:

474 qweight = layer.qweight

475 scales = layer.scales

vllm/model_executor/layers/quantization/bitsandbytes.py

+17 -16

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Callable, Optional, Union

4 +from collections.abc import Callable

5 +from typing import Any, Union

6

7 import torch

8 from packaging import version

@@ -41,7 +42,7 @@ class BitsAndBytesConfig(QuantizationConfig):

42 bnb_4bit_use_double_quant: bool = False,

43 llm_int8_enable_fp32_cpu_offload: bool = False,

44 llm_int8_has_fp16_weight: bool = False,

44 - llm_int8_skip_modules: Optional[list[str]] = None,

45 + llm_int8_skip_modules: list[str] | None = None,

46 llm_int8_threshold: float = 6.0,

47 ) -> None:

48 super().__init__()

@@ -138,7 +139,7 @@ class BitsAndBytesConfig(QuantizationConfig):

139

140 def get_quant_method(

141 self, layer: torch.nn.Module, prefix: str

141 - ) -> Optional[Union["LinearMethodBase", "BitsAndBytesMoEMethod"]]:

142 + ) -> Union["LinearMethodBase", "BitsAndBytesMoEMethod"] | None:

143 if isinstance(layer, LinearBase):

144 if is_layer_skipped_bnb(prefix, self.llm_int8_skip_modules):

145 return UnquantizedLinearMethod()

@@ -268,7 +269,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):

269 self,

270 layer: torch.nn.Module,

271 x: torch.Tensor,

271 - bias: Optional[torch.Tensor] = None,

272 + bias: torch.Tensor | None = None,

273 ) -> torch.Tensor:

274 if self.quant_config.load_in_8bit:

275 return self._apply_8bit_weight(layer, x, bias)

@@ -279,7 +280,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):

280 self,

281 layer: torch.nn.Module,

282 x: torch.Tensor,

282 - bias: Optional[torch.Tensor] = None,

283 + bias: torch.Tensor | None = None,

284 ) -> torch.Tensor:

285 # only load the bitsandbytes module when needed

286 from bitsandbytes import MatmulLtState, matmul

@@ -359,7 +360,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):

360 self,

361 layer: torch.nn.Module,

362 x: torch.Tensor,

362 - bias: Optional[torch.Tensor] = None,

363 + bias: torch.Tensor | None = None,

364 ) -> torch.Tensor:

365 original_type = x.dtype

366 original_shape = x.shape

@@ -489,7 +490,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

490

491 def get_fused_moe_quant_config(

492 self, layer: torch.nn.Module

492 - ) -> Optional[FusedMoEQuantConfig]:

493 + ) -> FusedMoEQuantConfig | None:

494 return None

495

496 def apply(

@@ -500,21 +501,21 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

501 top_k: int,

502 renormalize: bool,

503 use_grouped_topk: bool = False,

503 - topk_group: Optional[int] = None,

504 - num_expert_group: Optional[int] = None,

504 + topk_group: int | None = None,

505 + num_expert_group: int | None = None,

506 global_num_experts: int = -1,

506 - expert_map: Optional[torch.Tensor] = None,

507 - custom_routing_function: Optional[Callable] = None,

507 + expert_map: torch.Tensor | None = None,

508 + custom_routing_function: Callable | None = None,

509 scoring_func: str = "softmax",

510 routed_scaling_factor: float = 1.0,

510 - e_score_correction_bias: Optional[torch.Tensor] = None,

511 + e_score_correction_bias: torch.Tensor | None = None,

512 apply_router_weight_on_input: bool = False,

513 activation: str = "silu",

514 enable_eplb: bool = False,

514 - expert_load_view: Optional[torch.Tensor] = None,

515 - logical_to_physical_map: Optional[torch.Tensor] = None,

516 - logical_replica_count: Optional[torch.Tensor] = None,

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

515 + expert_load_view: torch.Tensor | None = None,

516 + logical_to_physical_map: torch.Tensor | None = None,

517 + logical_replica_count: torch.Tensor | None = None,

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

519 from vllm.model_executor.layers.fused_moe import fused_experts

520

521 assert self.fused_experts is None

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

+13 -13

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

71 __all__ = ["CompressedTensorsLinearMethod"]

72

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

74 -QUANTIZATION_SCHEME_MAP_TYPE = dict[str, Optional[dict[str, QuantizationArgs]]]

74 +QUANTIZATION_SCHEME_MAP_TYPE = dict[str, dict[str, QuantizationArgs] | None]

75

76

77 class CompressedTensorsConfig(QuantizationConfig):

@@ -82,9 +82,9 @@ class CompressedTensorsConfig(QuantizationConfig):

82 quant_format: str,

83 sparsity_scheme_map: dict[str, SparsityCompressionConfig],

84 sparsity_ignore_list: list[str],

85 - kv_cache_scheme: Optional[dict[str, Any]] = None,

86 - config: Optional[dict[str, Any]] = None,

87 - transform_config: Optional[dict[str, Any]] = None,

85 + kv_cache_scheme: dict[str, Any] | None = None,

86 + config: dict[str, Any] | None = None,

87 + transform_config: dict[str, Any] | None = None,

88 ):

89 super().__init__()

90 self.ignore = ignore

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

524 self,

525 weight_quant: QuantizationArgs,

526 input_quant: QuantizationArgs,

527 - format: Optional[str] = None,

527 + format: str | None = None,

528 ) -> "CompressedTensorsScheme":

529 # use the per-layer format if defined, otherwise, use global format

530 format = format if format is not None else self.quant_format

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

631 raise NotImplementedError("No compressed-tensors compatible scheme was found.")

632

633 def get_scheme(

634 - self, layer: torch.nn.Module, layer_name: Optional[str] = None

634 + self, layer: torch.nn.Module, layer_name: str | None = None

635 ) -> Optional["CompressedTensorsScheme"]:

636 """

637 compressed-tensors supports non uniform in the following way:

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

674 sparsity_targets = self.sparsity_scheme_map.keys() - set(

675 self.sparsity_ignore_list

676 )

677 - sparsity_scheme: Optional[SparsityCompressionConfig] = None

677 + sparsity_scheme: SparsityCompressionConfig | None = None

678 with suppress(ValueError):

679 matched_target = find_matched_target(

680 layer_name=layer_name,

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

723 logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name)

724 return scheme

725

726 - def get_cache_scale(self, name: str) -> Optional[str]:

726 + def get_cache_scale(self, name: str) -> str | None:

727 """

728 Check whether the param name matches the format for k/v cache scales

729 in compressed-tensors. If this is the case, return its equivalent

@@ -751,9 +751,9 @@ class CompressedTensorsConfig(QuantizationConfig):

751

752 @staticmethod

753 def supports_cutlass_24(

754 - weight_quant: Optional[QuantizationArgs],

755 - input_quant: Optional[QuantizationArgs],

756 - sparsity_scheme: Optional[SparsityCompressionConfig] = None,

754 + weight_quant: QuantizationArgs | None,

755 + input_quant: QuantizationArgs | None,

756 + sparsity_scheme: SparsityCompressionConfig | None = None,

757 ) -> bool:

758 """

759 Check if the layer is supported by the Cutlass 2:4 Kernel

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

853 self,

854 layer: torch.nn.Module,

855 x: torch.Tensor,

856 - bias: Optional[torch.Tensor] = None,

856 + bias: torch.Tensor | None = None,

857 ):

858 """

859 Use the output of create_weights and the CompressedTensorsScheme

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

878 super().__init__(quant_config)

879

880 @staticmethod

881 - def validate_kv_cache_scheme(kv_cache_scheme: Optional[dict[str, Any]]):

881 + def validate_kv_cache_scheme(kv_cache_scheme: dict[str, Any] | None):

882 """

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

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

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

+63 -63

@@ -2,8 +2,8 @@

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

3

4 import enum

5 +from collections.abc import Callable

6 from enum import Enum

6 -from typing import Callable, Optional, Union

7

8 import torch

9 from compressed_tensors import CompressionFormat

@@ -372,7 +372,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

372 (layer.w2_input_global_scale), requires_grad=False

373 )

374

375 - def maybe_make_prepare_finalize(self) -> Optional[mk.FusedMoEPrepareAndFinalize]:

375 + def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

376 if self.use_marlin:

377 return None

378 elif not self.allow_flashinfer:

@@ -399,7 +399,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

399

400 def get_fused_moe_quant_config(

401 self, layer: torch.nn.Module

402 - ) -> Optional[FusedMoEQuantConfig]:

402 + ) -> FusedMoEQuantConfig | None:

403 if self.use_marlin:

404 return None

405

@@ -420,21 +420,21 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

420 top_k: int,

421 renormalize: bool,

422 use_grouped_topk: bool = False,

423 - topk_group: Optional[int] = None,

424 - num_expert_group: Optional[int] = None,

423 + topk_group: int | None = None,

424 + num_expert_group: int | None = None,

425 global_num_experts: int = -1,

426 - expert_map: Optional[torch.Tensor] = None,

427 - custom_routing_function: Optional[Callable] = None,

426 + expert_map: torch.Tensor | None = None,

427 + custom_routing_function: Callable | None = None,

428 scoring_func: str = "softmax",

429 routed_scaling_factor: float = 1.0,

430 - e_score_correction_bias: Optional[torch.Tensor] = None,

430 + e_score_correction_bias: torch.Tensor | None = None,

431 apply_router_weight_on_input: bool = False,

432 activation: str = "silu",

433 enable_eplb: bool = False,

434 - expert_load_view: Optional[torch.Tensor] = None,

435 - logical_to_physical_map: Optional[torch.Tensor] = None,

436 - logical_replica_count: Optional[torch.Tensor] = None,

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

434 + expert_load_view: torch.Tensor | None = None,

435 + logical_to_physical_map: torch.Tensor | None = None,

436 + logical_replica_count: torch.Tensor | None = None,

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

438 if enable_eplb:

439 raise NotImplementedError(

440 "EPLB not supported for `CompressedTensorsW4A4MoeMethod` yet."

@@ -913,7 +913,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

913 layer.w2_weight_scale

914 )

915

916 - def maybe_make_prepare_finalize(self) -> Optional[mk.FusedMoEPrepareAndFinalize]:

916 + def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

917 if self.use_marlin or self.rocm_aiter_moe_enabled:

918 return None

919 else:

@@ -997,7 +997,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

997

998 def get_fused_moe_quant_config(

999 self, layer: torch.nn.Module

1000 - ) -> Optional[FusedMoEQuantConfig]:

1000 + ) -> FusedMoEQuantConfig | None:

1001 if self.use_marlin:

1002 return None

1003

@@ -1022,21 +1022,21 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1022 top_k: int,

1023 renormalize: bool,

1024 use_grouped_topk: bool = False,

1025 - topk_group: Optional[int] = None,

1026 - num_expert_group: Optional[int] = None,

1025 + topk_group: int | None = None,

1026 + num_expert_group: int | None = None,

1027 global_num_experts: int = -1,

1028 - expert_map: Optional[torch.Tensor] = None,

1029 - custom_routing_function: Optional[Callable] = None,

1028 + expert_map: torch.Tensor | None = None,

1029 + custom_routing_function: Callable | None = None,

1030 scoring_func: str = "softmax",

1031 routed_scaling_factor: float = 1.0,

1032 - e_score_correction_bias: Optional[torch.Tensor] = None,

1032 + e_score_correction_bias: torch.Tensor | None = None,

1033 apply_router_weight_on_input: bool = False,

1034 activation: str = "silu",

1035 enable_eplb: bool = False,

1036 - expert_load_view: Optional[torch.Tensor] = None,

1037 - logical_to_physical_map: Optional[torch.Tensor] = None,

1038 - logical_replica_count: Optional[torch.Tensor] = None,

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

1036 + expert_load_view: torch.Tensor | None = None,

1037 + logical_to_physical_map: torch.Tensor | None = None,

1038 + logical_replica_count: torch.Tensor | None = None,

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

1040 if enable_eplb:

1041 raise NotImplementedError(

1042 "EPLB not supported for `CompressedTensorsW8A8Fp8MoEMethod` yet."

@@ -1280,7 +1280,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1280

1281 def get_fused_moe_quant_config(

1282 self, layer: torch.nn.Module

1283 - ) -> Optional[FusedMoEQuantConfig]:

1283 + ) -> FusedMoEQuantConfig | None:

1284 return int8_w8a8_moe_quant_config(

1285 w1_scale=layer.w13_weight_scale,

1286 w2_scale=layer.w2_weight_scale,

@@ -1297,21 +1297,21 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1297 top_k: int,

1298 renormalize: bool,

1299 use_grouped_topk: bool = False,

1300 - topk_group: Optional[int] = None,

1301 - num_expert_group: Optional[int] = None,

1300 + topk_group: int | None = None,

1301 + num_expert_group: int | None = None,

1302 global_num_experts: int = -1,

1303 - expert_map: Optional[torch.Tensor] = None,

1304 - custom_routing_function: Optional[Callable] = None,

1303 + expert_map: torch.Tensor | None = None,

1304 + custom_routing_function: Callable | None = None,

1305 scoring_func: str = "softmax",

1306 routed_scaling_factor: float = 1.0,

1307 - e_score_correction_bias: Optional[torch.Tensor] = None,

1307 + e_score_correction_bias: torch.Tensor | None = None,

1308 apply_router_weight_on_input: bool = False,

1309 activation: str = "silu",

1310 enable_eplb: bool = False,

1311 - expert_load_view: Optional[torch.Tensor] = None,

1312 - logical_to_physical_map: Optional[torch.Tensor] = None,

1313 - logical_replica_count: Optional[torch.Tensor] = None,

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

1311 + expert_load_view: torch.Tensor | None = None,

1312 + logical_to_physical_map: torch.Tensor | None = None,

1313 + logical_replica_count: torch.Tensor | None = None,

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

1315 assert self.fused_experts is None

1316

1317 if enable_eplb:

@@ -1604,7 +1604,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1604

1605 def get_fused_moe_quant_config(

1606 self, layer: torch.nn.Module

1607 - ) -> Optional[FusedMoEQuantConfig]:

1607 + ) -> FusedMoEQuantConfig | None:

1608 return None

1609

1610 def apply(

@@ -1615,21 +1615,21 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1615 top_k: int,

1616 renormalize: bool,

1617 use_grouped_topk: bool = False,

1618 - topk_group: Optional[int] = None,

1619 - num_expert_group: Optional[int] = None,

1618 + topk_group: int | None = None,

1619 + num_expert_group: int | None = None,

1620 global_num_experts: int = -1,

1621 - expert_map: Optional[torch.Tensor] = None,

1622 - custom_routing_function: Optional[Callable] = None,

1621 + expert_map: torch.Tensor | None = None,

1622 + custom_routing_function: Callable | None = None,

1623 scoring_func: str = "softmax",

1624 routed_scaling_factor: float = 1.0,

1625 - e_score_correction_bias: Optional[torch.Tensor] = None,

1625 + e_score_correction_bias: torch.Tensor | None = None,

1626 apply_router_weight_on_input: bool = False,

1627 activation: str = "silu",

1628 enable_eplb: bool = False,

1629 - expert_load_view: Optional[torch.Tensor] = None,

1630 - logical_to_physical_map: Optional[torch.Tensor] = None,

1631 - logical_replica_count: Optional[torch.Tensor] = None,

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

1629 + expert_load_view: torch.Tensor | None = None,

1630 + logical_to_physical_map: torch.Tensor | None = None,

1631 + logical_replica_count: torch.Tensor | None = None,

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

1633 assert self.fused_experts is None

1634

1635 if enable_eplb:

@@ -1856,7 +1856,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1856

1857 def get_fused_moe_quant_config(

1858 self, layer: torch.nn.Module

1859 - ) -> Optional[FusedMoEQuantConfig]:

1859 + ) -> FusedMoEQuantConfig | None:

1860 assert self.num_bits == 4 or self.num_bits == 8

1861 config_builder = (

1862 int4_w4a16_moe_quant_config

@@ -1880,21 +1880,21 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1880 top_k: int,

1881 renormalize: bool,

1882 use_grouped_topk: bool = False,

1883 - topk_group: Optional[int] = None,

1884 - num_expert_group: Optional[int] = None,

1883 + topk_group: int | None = None,

1884 + num_expert_group: int | None = None,

1885 global_num_experts: int = -1,

1886 - expert_map: Optional[torch.Tensor] = None,

1887 - custom_routing_function: Optional[Callable] = None,

1886 + expert_map: torch.Tensor | None = None,

1887 + custom_routing_function: Callable | None = None,

1888 scoring_func: str = "softmax",

1889 routed_scaling_factor: float = 1.0,

1890 - e_score_correction_bias: Optional[torch.Tensor] = None,

1890 + e_score_correction_bias: torch.Tensor | None = None,

1891 apply_router_weight_on_input: bool = False,

1892 activation: str = "silu",

1893 enable_eplb: bool = False,

1894 - expert_load_view: Optional[torch.Tensor] = None,

1895 - logical_to_physical_map: Optional[torch.Tensor] = None,

1896 - logical_replica_count: Optional[torch.Tensor] = None,

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

1894 + expert_load_view: torch.Tensor | None = None,

1895 + logical_to_physical_map: torch.Tensor | None = None,

1896 + logical_replica_count: torch.Tensor | None = None,

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

1898 assert self.fused_experts is None

1899

1900 if enable_eplb:

@@ -2092,7 +2092,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):

2092 def _pack_matrix(

2093 int4_as_int8_2d: torch.Tensor,

2094 scales_2d: torch.Tensor,

2095 - bias_1d: Optional[torch.Tensor],

2095 + bias_1d: torch.Tensor | None,

2096 in_features: int,

2097 out_features: int,

2098 ) -> torch.Tensor:

@@ -2192,7 +2192,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):

2192

2193 def get_fused_moe_quant_config(

2194 self, layer: torch.nn.Module

2195 - ) -> Optional[FusedMoEQuantConfig]:

2195 + ) -> FusedMoEQuantConfig | None:

2196 # CPU dynamic 4-bit MoE path does not use modular kernels or

2197 # fused_experts; quant config is not needed.

2198 return None

@@ -2205,20 +2205,20 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):

2205 top_k: int,

2206 renormalize: bool,

2207 use_grouped_topk: bool = False,

2208 - topk_group: Optional[int] = None,

2209 - num_expert_group: Optional[int] = None,

2208 + topk_group: int | None = None,

2209 + num_expert_group: int | None = None,

2210 global_num_experts: int = -1,

2211 - expert_map: Optional[torch.Tensor] = None,

2212 - custom_routing_function: Optional[Callable] = None,

2211 + expert_map: torch.Tensor | None = None,

2212 + custom_routing_function: Callable | None = None,

2213 scoring_func: str = "softmax",

2214 routed_scaling_factor: float = 1.0,

2215 - e_score_correction_bias: Optional[torch.Tensor] = None,

2215 + e_score_correction_bias: torch.Tensor | None = None,

2216 apply_router_weight_on_input: bool = False,

2217 activation: str = "silu",

2218 enable_eplb: bool = False,

2219 - expert_load_view: Optional[torch.Tensor] = None,

2220 - logical_to_physical_map: Optional[torch.Tensor] = None,

2221 - logical_replica_count: Optional[torch.Tensor] = None,

2219 + expert_load_view: torch.Tensor | None = None,

2220 + logical_to_physical_map: torch.Tensor | None = None,

2221 + logical_replica_count: torch.Tensor | None = None,

2222 ) -> torch.Tensor:

2223 assert not enable_eplb, "EPLB not supported for W4A8-int MoE yet."

2224 assert activation in ("silu", "swigluoai", "swiglu"), (

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

+6 -5

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Callable, Optional

4 +from collections.abc import Callable

5 +from typing import Any

6

7 import torch

8 from compressed_tensors import CompressionFormat, ModelCompressor

@@ -42,9 +43,9 @@ class CompressedTensors24(CompressedTensorsScheme):

43 def __init__(

44 self,

45 quantized: bool = False,

45 - weight_quant: Optional[QuantizationArgs] = None,

46 - input_quant: Optional[QuantizationArgs] = None,

47 - model_compression_config: Optional[dict[str, Any]] = None,

46 + weight_quant: QuantizationArgs | None = None,

47 + input_quant: QuantizationArgs | None = None,

48 + model_compression_config: dict[str, Any] | None = None,

49 ):

50 self.quantized = quantized

51 self.weight_quant = weight_quant

@@ -247,7 +248,7 @@ class CompressedTensors24(CompressedTensorsScheme):

248 self,

249 layer: torch.nn.Module,

250 x: torch.Tensor,

250 - bias: Optional[torch.Tensor] = None,

251 + bias: torch.Tensor | None = None,

252 ) -> torch.Tensor:

253 """

254 Returns the output tensor for the layer with 2:4

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

+1 -2

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

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

3

4 from abc import ABC, abstractmethod

5 -from typing import Optional

5

6 import torch

7

@@ -33,7 +32,7 @@ class CompressedTensorsScheme(ABC):

32

33 @abstractmethod

34 def apply_weights(

36 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

35 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

36 ):

37 """

38 Run the forward pass for the particular scheme. This is where

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

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7 from torch.nn import Parameter

@@ -30,7 +30,7 @@ W4A16SPARSE24_SUPPORTED_BITS = list(W4A16SPARSE24_SUPPORTED_TYPES_MAP.keys())

30

31

32 class CompressedTensorsW4A16Sparse24(CompressedTensorsScheme):

33 - def __init__(self, strategy: str, num_bits: int, group_size: Optional[int] = None):

33 + def __init__(self, strategy: str, num_bits: int, group_size: int | None = None):

34 self.strategy = strategy

35 self.group_size = group_size

36 self.tile_size = 16

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

143 layer.workspace = workspace

144

145 def apply_weights(

146 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

146 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

147 ) -> torch.Tensor:

148 qweight = layer.weight_packed

149 meta = layer.meta

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

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Callable, Optional

3 +from collections.abc import Callable

4

5 import torch

6 from torch.nn.parameter import Parameter

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

110 self,

111 layer: torch.nn.Module,

112 x: torch.Tensor,

113 - bias: Optional[torch.Tensor] = None,

113 + bias: torch.Tensor | None = None,

114 ) -> torch.Tensor:

115 return apply_fp4_marlin_linear(

116 input=x,

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

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Callable, Optional

3 +from collections.abc import Callable

4

5 import torch

6 from torch.nn.parameter import Parameter

@@ -156,7 +156,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

156 self,

157 layer: torch.nn.Module,

158 x: torch.Tensor,

159 - bias: Optional[torch.Tensor] = None,

159 + bias: torch.Tensor | None = None,

160 ) -> torch.Tensor:

161 if envs.VLLM_USE_NVFP4_CT_EMULATIONS:

162 out = run_nvfp4_emulations(

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

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7 from compressed_tensors.quantization import ActivationOrdering

@@ -41,9 +41,9 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):

41 self,

42 strategy: str,

43 num_bits: int,

44 - group_size: Optional[int] = None,

45 - symmetric: Optional[bool] = True,

46 - actorder: Optional[ActivationOrdering] = None,

44 + group_size: int | None = None,

45 + symmetric: bool | None = True,

46 + actorder: ActivationOrdering | None = None,

47 ):

48 self.pack_factor = 32 // num_bits

49 self.strategy = strategy

@@ -178,6 +178,6 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):

178 self.kernel.process_weights_after_loading(layer)

179

180 def apply_weights(

181 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

181 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

182 ) -> torch.Tensor:

183 return self.kernel.apply_weights(layer, x, bias)

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

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7

@@ -36,7 +36,7 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme):

36 self,

37 strategy: str,

38 num_bits: int,

39 - group_size: Optional[int] = None,

39 + group_size: int | None = None,

40 is_static_input_scheme: bool = False,

41 input_symmetric: bool = True,

42 ):

@@ -148,6 +148,6 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme):

148 self.kernel.process_weights_after_loading(layer)

149

150 def apply_weights(

151 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

151 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

152 ) -> torch.Tensor:

153 return self.kernel.apply_weights(layer, x, bias)

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

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7 from compressed_tensors.quantization import QuantizationStrategy

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

125 self,

126 layer: torch.nn.Module,

127 x: torch.Tensor,

128 - bias: Optional[torch.Tensor] = None,

128 + bias: torch.Tensor | None = None,

129 ) -> torch.Tensor:

130 return apply_fp8_marlin_linear(

131 input=x,

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

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7 from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy

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

179 self,

180 layer: torch.nn.Module,

181 x: torch.Tensor,

182 - bias: Optional[torch.Tensor] = None,

182 + bias: torch.Tensor | None = None,

183 ) -> torch.Tensor:

184 if self.weight_block_size is not None:

185 return self.w8a8_block_fp8_linear.apply(

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

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7 from compressed_tensors.quantization import QuantizationStrategy

@@ -120,6 +120,6 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):

120 self.kernel.process_weights_after_loading(layer)

121

122 def apply_weights(

123 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

123 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

124 ) -> torch.Tensor:

125 return self.kernel.apply_weights(layer, x, bias)

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

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7 from compressed_tensors.quantization import ActivationOrdering

@@ -42,9 +42,9 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):

42 self,

43 strategy: str,

44 num_bits: int,

45 - group_size: Optional[int] = None,

46 - symmetric: Optional[bool] = True,

47 - actorder: Optional[ActivationOrdering] = None,

45 + group_size: int | None = None,

46 + symmetric: bool | None = True,

47 + actorder: ActivationOrdering | None = None,

48 ):

49 self.pack_factor = 32 // num_bits

50 self.strategy = strategy

@@ -214,6 +214,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):

214 self.kernel.process_weights_after_loading(layer)

215

216 def apply_weights(

217 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

217 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

218 ) -> torch.Tensor:

219 return self.kernel.apply_weights(layer, x, bias)

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

+7 -8

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from collections.abc import Generator

3 +from collections.abc import Callable, Generator

4 from itertools import accumulate

5 -from typing import Callable, Optional

5

6 import torch

7 from compressed_tensors.transform import (

@@ -38,7 +37,7 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase):

37 def from_schemes(

38 cls,

39 quant_method: LinearMethodBase,

41 - quant_scheme: Optional[CompressedTensorsScheme],

40 + quant_scheme: CompressedTensorsScheme | None,

41 input_tfms: dict[int, TransformTuple],

42 output_tfms: dict[int, TransformTuple],

43 ) -> "CompressedTensorsLinearTransformMethod":

@@ -66,8 +65,8 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase):

65 self.input_tfms = input_tfms

66 self.output_tfms = output_tfms

67

69 - self.input_transform: Optional[HadamardTransform] = None

70 - self.output_transform: Optional[HadamardTransform] = None

68 + self.input_transform: HadamardTransform | None = None

69 + self.output_transform: HadamardTransform | None = None

70

71 def create_weights(

72 self,

@@ -151,7 +150,7 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase):

150 self,

151 layer: torch.nn.Module,

152 x: torch.Tensor,

154 - bias: Optional[torch.Tensor] = None,

153 + bias: torch.Tensor | None = None,

154 ) -> torch.Tensor:

155 if self.input_transform is not None:

156 x = self.input_transform(x)

@@ -194,7 +193,7 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase):

193 def get_linear_transform_schemes(

194 layer: torch.nn.Module,

195 layer_name: str,

197 - transform_config: Optional[TransformConfig],

196 + transform_config: TransformConfig | None,

197 packed_modules_mapping: dict[str, list[str]],

198 ) -> tuple[

199 dict[int, TransformTuple], dict[int, TransformTuple]

@@ -226,7 +225,7 @@ def get_linear_transform_schemes(

225

226

227 def get_schemes_args(

229 - transform_config: Optional[TransformConfig],

228 + transform_config: TransformConfig | None,

229 ) -> Generator[tuple[str, TransformScheme, TransformArgs]]:

230 if transform_config is None:

231 return

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

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import math

4 -from collections.abc import Hashable

5 -from typing import Callable

4 +from collections.abc import Callable, Hashable

5

6 import torch

7 from compressed_tensors.transform import (

vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py

+2 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5

@@ -17,7 +16,7 @@ __all__ = ["is_qutlass_fp4_scheme", "QutlassNvFP4LinearMethod"]

16

17

18 def is_qutlass_fp4_scheme(

20 - quant_scheme: Optional[CompressedTensorsScheme],

19 + quant_scheme: CompressedTensorsScheme | None,

20 input_tfms: dict[int, TransformTuple],

21 ) -> bool:

22 return (

@@ -60,6 +59,6 @@ class QutlassNvFP4LinearMethod(CompressedTensorsLinearTransformMethod):

59 self,

60 layer: torch.nn.Module,

61 x: torch.Tensor,

63 - bias: Optional[torch.Tensor] = None,

62 + bias: torch.Tensor | None = None,

63 ) -> torch.Tensor:

64 raise NotImplementedError()

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

+1 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Optional

4

5 import torch

6

@@ -145,7 +144,7 @@ def triton_scaled_mm(

144 scale_a: torch.Tensor,

145 scale_b: torch.Tensor,

146 out_dtype: type[torch.dtype],

148 - bias: Optional[torch.Tensor] = None,

147 + bias: torch.Tensor | None = None,

148 block_size_m: int = 32,

149 block_size_n: int = 32,

150 block_size_k: int = 32,

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

+5 -6

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

3

4 from collections.abc import Iterable, Mapping

5 from types import MappingProxyType

6 -from typing import Optional

6

7 import regex as re

8 from compressed_tensors import CompressionFormat

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

20

21

22 def should_ignore_layer(

24 - layer_name: Optional[str],

23 + layer_name: str | None,

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

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

26 ) -> bool:

@@ -84,7 +83,7 @@ def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:

83

84

85 def find_matched_target(

87 - layer_name: Optional[str],

86 + layer_name: str | None,

87 module: Module,

88 targets: Iterable[str],

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

@@ -134,7 +133,7 @@ def find_matched_target(

133

134 def _find_first_match(

135 value: str, targets: Iterable[str], check_contains: bool = False

137 -) -> Optional[str]:

136 +) -> str | None:

137 """

138 Returns first element of target that matches value either

139 exactly or as a regex after 're:'. If check_contains is set to True,

@@ -176,7 +175,7 @@ def _match_fused_layer(

175 layer_name: str,

176 target_layers: Iterable[str],

177 fused_mapping: Mapping[str, list[str]],

179 -) -> Optional[str]:

178 +) -> str | None:

179 """

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

181 target_layers. Returns first value in fused_mapping which matches targets

@@ -205,7 +204,7 @@ def _match_fused_layer(

204 ]

205

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

208 - unfused_matches: list[Optional[str]] = []

207 + unfused_matches: list[str | None] = []

208 for unfused in unfused_paths:

209 for target in target_layers:

210 if _is_equal_or_regex_match(unfused, target):

vllm/model_executor/layers/quantization/deepspeedfp.py

+1 -1

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

140 self,

141 layer: torch.nn.Module,

142 x: torch.Tensor,

143 - bias: Optional[torch.Tensor] = None,

143 + bias: torch.Tensor | None = None,

144 ) -> torch.Tensor:

145 weight = layer.weight

146 y = weight.ds_dequantize()

vllm/model_executor/layers/quantization/experts_int8.py

+12 -11

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Callable, Optional, Union

4 +from collections.abc import Callable

5 +from typing import Any, Optional

6

7 import torch

8

@@ -129,7 +130,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

130

131 def get_fused_moe_quant_config(

132 self, layer: torch.nn.Module

132 - ) -> Optional[FusedMoEQuantConfig]:

133 + ) -> FusedMoEQuantConfig | None:

134 return int8_w8a16_moe_quant_config(

135 w1_scale=layer.w13_scale, w2_scale=layer.w2_scale, w1_zp=None, w2_zp=None

136 )

@@ -142,21 +143,21 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

143 top_k: int,

144 renormalize: bool,

145 use_grouped_topk: bool = False,

145 - topk_group: Optional[int] = None,

146 - num_expert_group: Optional[int] = None,

146 + topk_group: int | None = None,

147 + num_expert_group: int | None = None,

148 global_num_experts: int = -1,

148 - expert_map: Optional[torch.Tensor] = None,

149 - custom_routing_function: Optional[Callable] = None,

149 + expert_map: torch.Tensor | None = None,

150 + custom_routing_function: Callable | None = None,

151 scoring_func: str = "softmax",

152 routed_scaling_factor: float = 1.0,

152 - e_score_correction_bias: Optional[torch.Tensor] = None,

153 + e_score_correction_bias: torch.Tensor | None = None,

154 apply_router_weight_on_input: bool = False,

155 activation: str = "silu",

156 enable_eplb: bool = False,

156 - expert_load_view: Optional[torch.Tensor] = None,

157 - logical_to_physical_map: Optional[torch.Tensor] = None,

158 - logical_replica_count: Optional[torch.Tensor] = None,

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

157 + expert_load_view: torch.Tensor | None = None,

158 + logical_to_physical_map: torch.Tensor | None = None,

159 + logical_replica_count: torch.Tensor | None = None,

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

161 assert self.fused_experts is None

162

163 if enable_eplb:

vllm/model_executor/layers/quantization/fbgemm_fp8.py

+1 -1

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

171 self,

172 layer: torch.nn.Module,

173 x: torch.Tensor,

174 - bias: Optional[torch.Tensor] = None,

174 + bias: torch.Tensor | None = None,

175 ) -> torch.Tensor:

176 if self.quant_config.use_marlin:

177 return apply_fp8_marlin_linear(

vllm/model_executor/layers/quantization/fp8.py

+19 -18

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 +from collections.abc import Callable

5 from enum import Enum

5 -from typing import TYPE_CHECKING, Any, Callable, Optional, Union

6 +from typing import TYPE_CHECKING, Any, Optional

7

8 import torch

9 from torch.nn import Module

@@ -173,8 +174,8 @@ class Fp8Config(QuantizationConfig):

174 self,

175 is_checkpoint_fp8_serialized: bool = False,

176 activation_scheme: str = "dynamic",

176 - ignored_layers: Optional[list[str]] = None,

177 - weight_block_size: Optional[list[int]] = None,

177 + ignored_layers: list[str] | None = None,

178 + weight_block_size: list[int] | None = None,

179 ) -> None:

180 super().__init__()

181

@@ -298,7 +299,7 @@ class Fp8Config(QuantizationConfig):

299 return Fp8KVCacheMethod(self)

300 return None

301

301 - def get_cache_scale(self, name: str) -> Optional[str]:

302 + def get_cache_scale(self, name: str) -> str | None:

303 """

304 Check whether the param name matches the format for k/v cache scales

305 in compressed-tensors. If this is the case, return its equivalent

@@ -530,7 +531,7 @@ class Fp8LinearMethod(LinearMethodBase):

531 self,

532 layer: torch.nn.Module,

533 x: torch.Tensor,

533 - bias: Optional[torch.Tensor] = None,

534 + bias: torch.Tensor | None = None,

535 ) -> torch.Tensor:

536 if self.use_marlin:

537 return apply_fp8_marlin_linear(

@@ -584,12 +585,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):

585 self.weight_block_size = self.quant_config.weight_block_size

586 self.block_quant: bool = self.weight_block_size is not None

587

587 - self.fused_experts: Optional[mk.FusedMoEModularKernel] = None # type: ignore

588 + self.fused_experts: mk.FusedMoEModularKernel | None = None # type: ignore

589

590 self.fp8_backend = get_fp8_moe_backend(self.block_quant)

591

592 self.use_marlin = self.fp8_backend == Fp8MoeBackend.MARLIN

592 - self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None

593 + self.flashinfer_moe_backend: FlashinferMoeBackend | None = None

594 if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

595 self.flashinfer_moe_backend = FlashinferMoeBackend.TENSORRT_LLM

596 elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

@@ -970,7 +971,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

971 layer.w2_weight_scale_inv

972 )

973

973 - def maybe_make_prepare_finalize(self) -> Optional[mk.FusedMoEPrepareAndFinalize]:

974 + def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

975 if (

976 self.rocm_aiter_moe_enabled

977 or self.use_marlin

@@ -1043,7 +1044,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1044

1045 def get_fused_moe_quant_config(

1046 self, layer: torch.nn.Module

1046 - ) -> Optional[FusedMoEQuantConfig]:

1047 + ) -> FusedMoEQuantConfig | None:

1048 if self.use_marlin:

1049 return None

1050

@@ -1069,21 +1070,21 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1070 top_k: int,

1071 renormalize: bool,

1072 use_grouped_topk: bool = False,

1072 - topk_group: Optional[int] = None,

1073 - num_expert_group: Optional[int] = None,

1073 + topk_group: int | None = None,

1074 + num_expert_group: int | None = None,

1075 global_num_experts: int = -1,

1075 - expert_map: Optional[torch.Tensor] = None,

1076 - custom_routing_function: Optional[Callable] = None,

1076 + expert_map: torch.Tensor | None = None,

1077 + custom_routing_function: Callable | None = None,

1078 scoring_func: str = "softmax",

1079 routed_scaling_factor: float = 1.0,

1079 - e_score_correction_bias: Optional[torch.Tensor] = None,

1080 + e_score_correction_bias: torch.Tensor | None = None,

1081 apply_router_weight_on_input: bool = False,

1082 activation: str = "silu",

1083 enable_eplb: bool = False,

1083 - expert_load_view: Optional[torch.Tensor] = None,

1084 - logical_to_physical_map: Optional[torch.Tensor] = None,

1085 - logical_replica_count: Optional[torch.Tensor] = None,

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

1084 + expert_load_view: torch.Tensor | None = None,

1085 + logical_to_physical_map: torch.Tensor | None = None,

1086 + logical_replica_count: torch.Tensor | None = None,

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

1088 if enable_eplb:

1089 assert expert_load_view is not None

1090 assert logical_to_physical_map is not None

vllm/model_executor/layers/quantization/fp_quant.py

+5 -5

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

3

4 # Supports FP-Quant compression, see https://arxiv.org/abs/2509.23202

5

6 -from typing import Any, Optional

6 +from typing import Any

7

8 import torch

9 from torch.nn.parameter import Parameter

@@ -36,7 +36,7 @@ class FPQuantConfig(QuantizationConfig):

36 forward_dtype: str = "mxfp4",

37 forward_method: str = "abs_max",

38 pseudoquantization: bool = False,

39 - modules_to_not_convert: Optional[list[str]] = None,

39 + modules_to_not_convert: list[str] | None = None,

40 ) -> None:

41 super().__init__()

42 self.hadamard_group_size = hadamard_group_size

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

90

91 def get_quant_method(

92 self, layer: torch.nn.Module, prefix: str

93 - ) -> Optional[LinearMethodBase]:

93 + ) -> LinearMethodBase | None:

94 if self.modules_to_not_convert is not None and any(

95 prefix.endswith(module) for module in self.modules_to_not_convert

96 ):

@@ -233,7 +233,7 @@ class FPQuantLinearMethod(LinearMethodBase):

233 self,

234 layer: torch.nn.Module,

235 x: torch.Tensor,

236 - bias: Optional[torch.Tensor] = None,

236 + bias: torch.Tensor | None = None,

237 ) -> torch.Tensor:

238 return quantized_forward(

239 x,

@@ -381,7 +381,7 @@ def quantized_forward(

381 weight_scales: torch.Tensor,

382 weight_global_scale: torch.Tensor,

383 act_global_scale: torch.Tensor,

384 - bias: Optional[torch.Tensor],

384 + bias: torch.Tensor | None,

385 forward_hadamard_matrix: torch.Tensor,

386 forward_method: str,

387 forward_dtype: str,

vllm/model_executor/layers/quantization/gguf.py

+16 -15

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Callable, Optional, Union

4 +from collections.abc import Callable

5 +from typing import Any, Optional

6

7 import gguf

8 import torch

@@ -35,7 +36,7 @@ logger = init_logger(__name__)

36 class GGUFConfig(QuantizationConfig):

37 """Config class for GGUF."""

38

38 - def __init__(self, unquantized_modules: Optional[list[str]] = None) -> None:

39 + def __init__(self, unquantized_modules: list[str] | None = None) -> None:

40 super().__init__()

41 self.unquantized_modules = unquantized_modules or []

42

@@ -307,7 +308,7 @@ def _apply_gguf_embedding(

308 qweight: torch.Tensor,

309 qweight_type: int,

310 hidden_size: int,

310 - dtype: Optional[torch.dtype] = None,

311 + dtype: torch.dtype | None = None,

312 ) -> torch.Tensor:

313 if qweight_type in UNQUANTIZED_TYPES:

314 return torch.embedding(qweight, x)

@@ -330,7 +331,7 @@ def _apply_gguf_embedding_fake(

331 qweight: torch.Tensor,

332 qweight_type: int,

333 hidden_size: int,

333 - dtype: Optional[torch.dtype] = None,

334 + dtype: torch.dtype | None = None,

335 ) -> torch.Tensor:

336 return torch.empty(x.shape[0], hidden_size, dtype=dtype, device=x.device)

337

@@ -452,7 +453,7 @@ class GGUFLinearMethod(LinearMethodBase):

453 self,

454 layer: torch.nn.Module,

455 x: torch.Tensor,

455 - bias: Optional[torch.Tensor] = None,

456 + bias: torch.Tensor | None = None,

457 ) -> torch.Tensor:

458 shard_id = layer.qweight.shard_id

459

@@ -558,7 +559,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):

559

560 def get_fused_moe_quant_config(

561 self, layer: torch.nn.Module

561 - ) -> Optional[FusedMoEQuantConfig]:

562 + ) -> FusedMoEQuantConfig | None:

563 return None

564

565 def apply(

@@ -569,21 +570,21 @@ class GGUFMoEMethod(FusedMoEMethodBase):

570 top_k: int,

571 renormalize: bool,

572 use_grouped_topk: bool = False,

572 - topk_group: Optional[int] = None,

573 - num_expert_group: Optional[int] = None,

573 + topk_group: int | None = None,

574 + num_expert_group: int | None = None,

575 global_num_experts: int = -1,

575 - expert_map: Optional[torch.Tensor] = None,

576 - custom_routing_function: Optional[Callable] = None,

576 + expert_map: torch.Tensor | None = None,

577 + custom_routing_function: Callable | None = None,

578 scoring_func: str = "softmax",

579 routed_scaling_factor: float = 1.0,

579 - e_score_correction_bias: Optional[torch.Tensor] = None,

580 + e_score_correction_bias: torch.Tensor | None = None,

581 apply_router_weight_on_input: bool = False,

582 activation: str = "silu",

583 enable_eplb: bool = False,

583 - expert_load_view: Optional[torch.Tensor] = None,

584 - logical_to_physical_map: Optional[torch.Tensor] = None,

585 - logical_replica_count: Optional[torch.Tensor] = None,

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

584 + expert_load_view: torch.Tensor | None = None,

585 + logical_to_physical_map: torch.Tensor | None = None,

586 + logical_replica_count: torch.Tensor | None = None,

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

588 assert self.fused_experts is None

589

590 if enable_eplb:

vllm/model_executor/layers/quantization/gptq.py

+6 -6

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

4 import enum

5 from enum import Enum

6 from fractions import Fraction

7 -from typing import TYPE_CHECKING, Any, Optional, Union

7 +from typing import TYPE_CHECKING, Any, Union

8

9 import torch

10 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE

@@ -48,9 +48,9 @@ class GPTQConfig(QuantizationConfig):

48 group_size: int,

49 desc_act: bool,

50 lm_head_quantized: bool,

51 - dynamic: dict[str, dict[str, Union[int, bool]]],

51 + dynamic: dict[str, dict[str, int | bool]],

52 autoround_version: str = "",

53 - modules_in_block_to_quantize: Optional[list[str]] = None,

53 + modules_in_block_to_quantize: list[str] | None = None,

54 ) -> None:

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

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

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

148

149 def get_quant_method(

150 self, layer: torch.nn.Module, prefix: str

151 - ) -> Optional[Union["GPTQLinearMethod", "QuantizeMethodBase"]]:

151 + ) -> Union["GPTQLinearMethod", "QuantizeMethodBase"] | None:

152 if isinstance(layer, FusedMoE):

153 # GPTQ MoE support: fall back to MoeWNA16 for broad compatibility

154 from .moe_wna16 import MoeWNA16Config

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

170 self.modules_in_block_to_quantize

171 )

172

173 - def maybe_update_config(self, model_name: str, revision: Optional[str] = None):

173 + def maybe_update_config(self, model_name: str, revision: str | None = None):

174 if self.modules_in_block_to_quantize:

175 if is_list_of(self.modules_in_block_to_quantize, list):

176 # original modules_in_block_to_quantize: list[list[str]]

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

345 self,

346 layer: torch.nn.Module,

347 x: torch.Tensor,

348 - bias: Optional[torch.Tensor] = None,

348 + bias: torch.Tensor | None = None,

349 ) -> torch.Tensor:

350 out_shape = x.shape[:-1] + (layer.qweight.shape[-1],)

351 reshaped_x = x.reshape(-1, x.shape[-1])

vllm/model_executor/layers/quantization/gptq_bitblas.py

+3 -3

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

71 group_size: int,

72 desc_act: bool,

73 is_sym: bool,

74 - quant_method: Optional[str],

74 + quant_method: str | None,

75 lm_head_quantized: bool,

76 ) -> None:

77 try:

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

180 @classmethod

181 def override_quantization_method(

182 cls, hf_quant_cfg, user_quant

183 - ) -> Optional[QuantizationMethods]:

183 + ) -> QuantizationMethods | None:

184 can_convert = cls.is_gptq_bitblas_compatible(hf_quant_cfg)

185

186 is_valid_user_quant = (

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

474 self,

475 layer: torch.nn.Module,

476 x: torch.Tensor,

477 - bias: Optional[torch.Tensor] = None,

477 + bias: torch.Tensor | None = None,

478 ) -> torch.Tensor:

479 out = self.kernel.apply_gptq_bitblas_linear(layer, x)

480 if bias is not None:

vllm/model_executor/layers/quantization/gptq_marlin.py

+17 -16

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 +from collections.abc import Callable

5 from copy import deepcopy

5 -from typing import Any, Callable, Optional, Union

6 +from typing import Any, Optional

7

8 import torch

9 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE

@@ -103,9 +104,9 @@ class GPTQMarlinConfig(QuantizationConfig):

104 desc_act: bool,

105 is_sym: bool,

106 lm_head_quantized: bool,

106 - dynamic: dict[str, dict[str, Union[int, bool]]],

107 + dynamic: dict[str, dict[str, int | bool]],

108 full_config: dict[str, Any],

108 - modules_in_block_to_quantize: Optional[list[str]] = None,

109 + modules_in_block_to_quantize: list[str] | None = None,

110 ) -> None:

111 super().__init__()

112 if desc_act and group_size == -1:

@@ -211,7 +212,7 @@ class GPTQMarlinConfig(QuantizationConfig):

212 @classmethod

213 def override_quantization_method(

214 cls, hf_quant_cfg, user_quant

214 - ) -> Optional[QuantizationMethods]:

215 + ) -> QuantizationMethods | None:

216 can_convert = cls.is_gptq_marlin_compatible(hf_quant_cfg)

217

218 is_valid_user_quant = (

@@ -283,7 +284,7 @@ class GPTQMarlinConfig(QuantizationConfig):

284 self.modules_in_block_to_quantize

285 )

286

286 - def maybe_update_config(self, model_name: str, revision: Optional[str] = None):

287 + def maybe_update_config(self, model_name: str, revision: str | None = None):

288 if self.modules_in_block_to_quantize:

289 if is_list_of(self.modules_in_block_to_quantize, list):

290 # original modules_in_block_to_quantize: list[list[str]]

@@ -459,7 +460,7 @@ class GPTQMarlinLinearMethod(LinearMethodBase):

460 self,

461 layer: torch.nn.Module,

462 x: torch.Tensor,

462 - bias: Optional[torch.Tensor] = None,

463 + bias: torch.Tensor | None = None,

464 ) -> torch.Tensor:

465 return self.kernel.apply_weights(layer, x, bias)

466

@@ -714,7 +715,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

715

716 def get_fused_moe_quant_config(

717 self, layer: torch.nn.Module

717 - ) -> Optional[FusedMoEQuantConfig]:

718 + ) -> FusedMoEQuantConfig | None:

719 return None

720

721 def apply(

@@ -725,21 +726,21 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

726 top_k: int,

727 renormalize: bool,

728 use_grouped_topk: bool = False,

728 - topk_group: Optional[int] = None,

729 - num_expert_group: Optional[int] = None,

729 + topk_group: int | None = None,

730 + num_expert_group: int | None = None,

731 global_num_experts: int = -1,

731 - expert_map: Optional[torch.Tensor] = None,

732 - custom_routing_function: Optional[Callable] = None,

732 + expert_map: torch.Tensor | None = None,

733 + custom_routing_function: Callable | None = None,

734 scoring_func: str = "softmax",

735 routed_scaling_factor: float = 1.0,

735 - e_score_correction_bias: Optional[torch.Tensor] = None,

736 + e_score_correction_bias: torch.Tensor | None = None,

737 apply_router_weight_on_input: bool = False,

738 activation: str = "silu",

739 enable_eplb: bool = False,

739 - expert_load_view: Optional[torch.Tensor] = None,

740 - logical_to_physical_map: Optional[torch.Tensor] = None,

741 - logical_replica_count: Optional[torch.Tensor] = None,

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

740 + expert_load_view: torch.Tensor | None = None,

741 + logical_to_physical_map: torch.Tensor | None = None,

742 + logical_replica_count: torch.Tensor | None = None,

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

744 assert self.fused_experts is None

745

746 if enable_eplb:

vllm/model_executor/layers/quantization/gptq_marlin_24.py

+2 -2

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

114 @classmethod

115 def override_quantization_method(

116 cls, hf_quant_cfg, user_quant

117 - ) -> Optional[QuantizationMethods]:

117 + ) -> QuantizationMethods | None:

118 is_marlin_24_format = hf_quant_cfg.get("checkpoint_format") == "marlin_24"

119

120 is_valid_user_quant = (

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

287 self,

288 layer: torch.nn.Module,

289 x: torch.Tensor,

290 - bias: Optional[torch.Tensor] = None,

290 + bias: torch.Tensor | None = None,

291 ) -> torch.Tensor:

292 qweight = layer.B_24

293 meta = layer.B_meta

vllm/model_executor/layers/quantization/hqq_marlin.py

+2 -2

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

45 self,

46 weight_bits: int,

47 group_size: int,

48 - skip_modules: Optional[list[str]] = None,

48 + skip_modules: list[str] | None = None,

49 ) -> None:

50 super().__init__()

51 assert group_size == 64, "The only supported HQQ group size is currently 64."

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

327 self,

328 layer: torch.nn.Module,

329 x: torch.Tensor,

330 - bias: Optional[torch.Tensor] = None,

330 + bias: torch.Tensor | None = None,

331 ) -> torch.Tensor:

332 workspace = MarlinWorkspace(

333 self.output_size_per_partition,

vllm/model_executor/layers/quantization/input_quant_fp8.py

+6 -7

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional

3

4 import torch

5 import torch.nn.functional as F

@@ -30,9 +29,9 @@ class QuantFP8(CustomOp):

29 self,

30 static: bool,

31 group_shape: GroupShape,

33 - num_token_padding: Optional[int] = None,

32 + num_token_padding: int | None = None,

33 column_major_scales: bool = False,

35 - use_ue8m0: Optional[bool] = None, # for Torch compile

34 + use_ue8m0: bool | None = None, # for Torch compile

35 ):

36 """

37 :param static: static or dynamic quantization

@@ -64,8 +63,8 @@ class QuantFP8(CustomOp):

63 def forward_cuda(

64 self,

65 x: torch.Tensor,

67 - scale: Optional[torch.Tensor] = None,

68 - scale_ub: Optional[torch.Tensor] = None,

66 + scale: torch.Tensor | None = None,

67 + scale_ub: torch.Tensor | None = None,

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

69 if self.is_group_quant:

70 assert scale is None, "Group quantization is always dynamic"

@@ -96,8 +95,8 @@ class QuantFP8(CustomOp):

95 def forward_native(

96 self,

97 x: torch.Tensor,

99 - scale: Optional[torch.Tensor] = None,

100 - scale_ub: Optional[torch.Tensor] = None,

98 + scale: torch.Tensor | None = None,

99 + scale_ub: torch.Tensor | None = None,

100 ):

101 if self.is_group_quant:

102 assert scale is None, "Group quantization is always dynamic"

vllm/model_executor/layers/quantization/ipex_quant.py

+18 -17

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

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 -from typing import Any, Callable, Optional

4 +from collections.abc import Callable

5 +from typing import Any, Optional

6

7 import torch

8 from packaging import version

@@ -50,9 +51,9 @@ class IPEXConfig(QuantizationConfig):

51 method: str,

52 weight_bits: int,

53 group_size: int,

53 - modules_to_not_convert: Optional[list[str]] = None,

54 - desc_act: Optional[bool] = None,

55 - lm_head_quantized: Optional[bool] = None,

54 + modules_to_not_convert: list[str] | None = None,

55 + desc_act: bool | None = None,

56 + lm_head_quantized: bool | None = None,

57 ) -> None:

58 super().__init__()

59 self.method = method

@@ -122,7 +123,7 @@ class IPEXConfig(QuantizationConfig):

123 @classmethod

124 def override_quantization_method(

125 cls, hf_quant_cfg, user_quant

125 - ) -> Optional[QuantizationMethods]:

126 + ) -> QuantizationMethods | None:

127 if not current_platform.is_cpu() and not current_platform.is_xpu():

128 return None

129

@@ -206,7 +207,7 @@ class IPEXGPTQLinearMethod(GPTQLinearMethod):

207 self,

208 layer: torch.nn.Module,

209 x: torch.Tensor,

209 - bias: Optional[torch.Tensor] = None,

210 + bias: torch.Tensor | None = None,

211 ) -> torch.Tensor:

212 reshaped_x = x.reshape(-1, x.shape[-1])

213 out = layer.ipex_qlinear(reshaped_x)

@@ -275,7 +276,7 @@ class IPEXAWQLinearMethod(AWQLinearMethod):

276 self,

277 layer: torch.nn.Module,

278 x: torch.Tensor,

278 - bias: Optional[torch.Tensor] = None,

279 + bias: torch.Tensor | None = None,

280 ) -> torch.Tensor:

281 reshaped_x = x.reshape(-1, x.shape[-1])

282 out = layer.ipex_qlinear(reshaped_x)

@@ -299,7 +300,7 @@ class XPUFp8LinearMethod(Fp8LinearMethod):

300 self,

301 layer: torch.nn.Module,

302 x: torch.Tensor,

302 - bias: Optional[torch.Tensor] = None,

303 + bias: torch.Tensor | None = None,

304 ) -> torch.Tensor:

305 weight = layer.weight.data

306 weight_scale = layer.weight_scale.data

@@ -410,7 +411,7 @@ class XPUFp8MoEMethod(FusedMoEMethodBase):

411

412 def get_fused_moe_quant_config(

413 self, layer: torch.nn.Module

413 - ) -> Optional[FusedMoEQuantConfig]:

414 + ) -> FusedMoEQuantConfig | None:

415 return None

416

417 def apply(

@@ -421,20 +422,20 @@ class XPUFp8MoEMethod(FusedMoEMethodBase):

422 top_k: int,

423 renormalize: bool,

424 use_grouped_topk: bool = False,

424 - topk_group: Optional[int] = None,

425 - num_expert_group: Optional[int] = None,

425 + topk_group: int | None = None,

426 + num_expert_group: int | None = None,

427 global_num_experts: int = -1,

427 - expert_map: Optional[torch.Tensor] = None,

428 - custom_routing_function: Optional[Callable] = None,

428 + expert_map: torch.Tensor | None = None,

429 + custom_routing_function: Callable | None = None,

430 scoring_func: str = "softmax",

431 routed_scaling_factor: float = 1.0,

431 - e_score_correction_bias: Optional[torch.Tensor] = None,

432 + e_score_correction_bias: torch.Tensor | None = None,

433 apply_router_weight_on_input: bool = False,

434 activation: str = "silu",

435 enable_eplb: bool = False,

435 - expert_load_view: Optional[torch.Tensor] = None,

436 - logical_to_physical_map: Optional[torch.Tensor] = None,

437 - logical_replica_count: Optional[torch.Tensor] = None,

436 + expert_load_view: torch.Tensor | None = None,

437 + logical_to_physical_map: torch.Tensor | None = None,

438 + logical_replica_count: torch.Tensor | None = None,

439 ) -> torch.Tensor:

440 return layer.ipex_fusion(

441 x,

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

+9 -9

@@ -2,8 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from abc import ABC, abstractmethod

5 +from collections.abc import Callable

6 from dataclasses import dataclass

6 -from typing import Callable, Optional

7

8 import torch

9

@@ -20,7 +20,7 @@ class MPLinearLayerConfig:

20 group_size: int

21 zero_points: bool

22 has_g_idx: bool

23 - out_type: Optional[torch.dtype] = None

23 + out_type: torch.dtype | None = None

24

25

26 class MPLinearKernel(ABC):

@@ -31,7 +31,7 @@ class MPLinearKernel(ABC):

31

32 @classmethod

33 @abstractmethod

34 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

34 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

35 raise NotImplementedError

36

37 def __init__(

@@ -39,8 +39,8 @@ class MPLinearKernel(ABC):

39 c: MPLinearLayerConfig,

40 w_q_param_name: str,

41 w_s_param_name: str,

42 - w_zp_param_name: Optional[str] = None,

43 - w_gidx_param_name: Optional[str] = None,

42 + w_zp_param_name: str | None = None,

43 + w_gidx_param_name: str | None = None,

44 ) -> None:

45 assert self.can_implement(c)

46 self.config = c

@@ -62,12 +62,12 @@ class MPLinearKernel(ABC):

62 self,

63 layer: torch.nn.Module,

64 x: torch.Tensor,

65 - bias: Optional[torch.Tensor] = None,

65 + bias: torch.Tensor | None = None,

66 ) -> torch.Tensor:

67 raise NotImplementedError

68

69 def _transform_param(

70 - self, layer: torch.nn.Module, name: Optional[str], fn: Callable

70 + self, layer: torch.nn.Module, name: str | None, fn: Callable

71 ) -> None:

72 if name is not None and getattr(layer, name, None) is not None:

73 old_param = getattr(layer, name)

@@ -83,8 +83,8 @@ class MPLinearKernel(ABC):

83 ) -> tuple[

84 torch.Tensor, # w_q

85 torch.Tensor, # w_s

86 - Optional[torch.Tensor], # w_zp,

87 - Optional[torch.Tensor], # w_gidx

86 + torch.Tensor | None, # w_zp,

87 + torch.Tensor | None, # w_gidx

88 ]:

89 return (

90 getattr(layer, self.w_q_name),

vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py

+1 -3

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

5 -

4 import vllm.envs as envs

5 from vllm.model_executor.layers.quantization.kernels.mixed_precision.allspark import ( # noqa: E501

6 AllSparkLinearKernel,

@@ -48,7 +46,7 @@ _POSSIBLE_KERNELS: list[type[MPLinearKernel]] = [

46

47

48 def choose_mp_linear_kernel(

51 - config: MPLinearLayerConfig, compute_capability: Optional[int] = None

49 + config: MPLinearLayerConfig, compute_capability: int | None = None

50 ) -> type[MPLinearKernel]:

51 """

52 Choose an MPLinearKernel that can implement the given config for the given

vllm/model_executor/layers/quantization/kernels/mixed_precision/allspark.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -22,7 +21,7 @@ class AllSparkLinearKernel(MPLinearKernel):

21 return 80

22

23 @classmethod

25 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

24 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

25 if c.has_g_idx:

26 return False, "Act reordering currently not supported by AllSpark"

27

@@ -87,7 +86,7 @@ class AllSparkLinearKernel(MPLinearKernel):

86 self,

87 layer: torch.nn.Module,

88 x: torch.Tensor,

90 - bias: Optional[torch.Tensor] = None,

89 + bias: torch.Tensor | None = None,

90 ) -> torch.Tensor:

91 c = self.config

92 gemm_args = self.gemm_args

vllm/model_executor/layers/quantization/kernels/mixed_precision/bitblas.py

+6 -7

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6 from packaging import version

@@ -44,9 +43,9 @@ class BitBLASLinearKernel(MPLinearKernel):

43 c: MPLinearLayerConfig,

44 w_q_param_name: str,

45 w_s_param_name: str,

47 - w_zp_param_name: Optional[str] = None,

48 - w_gidx_param_name: Optional[str] = None,

49 - bitblas_quant_config: Optional[QuantizationConfig] = None,

46 + w_zp_param_name: str | None = None,

47 + w_gidx_param_name: str | None = None,

48 + bitblas_quant_config: QuantizationConfig | None = None,

49 ):

50 self.quant_config = bitblas_quant_config

51 super().__init__(

@@ -57,7 +56,7 @@ class BitBLASLinearKernel(MPLinearKernel):

56 self,

57 b_q_weight: torch.Tensor,

58 scales: torch.Tensor,

60 - qzeros: Optional[torch.Tensor] = None,

59 + qzeros: torch.Tensor | None = None,

60 ):

61 from bitblas.quantization.utils import general_compress

62

@@ -82,7 +81,7 @@ class BitBLASLinearKernel(MPLinearKernel):

81 # qzeros should be de-quantized to int zeros.

82 weight_bits = quant_config.weight_bits # type: ignore[union-attr]

83 intzeros = unpack_gptq_qzeros(qzeros, weight_bits).T.contiguous()

85 - zeros: Optional[torch.Tensor] = None

84 + zeros: torch.Tensor | None = None

85 zeros_mode = self.bitblas_matmul.config.zeros_mode # type: ignore[attr-defined]

86 if zeros_mode == "original":

87 zeros = intzeros.to(torch.float16).contiguous()

@@ -113,7 +112,7 @@ class BitBLASLinearKernel(MPLinearKernel):

112 return 70

113

114 @classmethod

116 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

115 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

116 is_bitblas_installed = True

117

118 try:

vllm/model_executor/layers/quantization/kernels/mixed_precision/conch.py

+3 -3

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from importlib.util import find_spec

5 -from typing import Final, Optional

5 +from typing import Final

6

7 import torch

8

@@ -26,7 +26,7 @@ class ConchLinearKernel(MPLinearKernel):

26 return 80

27

28 @classmethod

29 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

29 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

30 if c.weight_type not in _CONCH_SUPPORTED_WEIGHT_TYPES:

31 error_msg = (

32 f"Weight type ({c.weight_type}) not supported by "

@@ -76,7 +76,7 @@ class ConchLinearKernel(MPLinearKernel):

76 self,

77 layer: torch.nn.Module,

78 x: torch.Tensor,

79 - bias: Optional[torch.Tensor] = None,

79 + bias: torch.Tensor | None = None,

80 ) -> torch.Tensor:

81 from conch.ops.quantization.gemm import mixed_precision_gemm

82

vllm/model_executor/layers/quantization/kernels/mixed_precision/cutlass.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -26,7 +25,7 @@ class CutlassW4A8LinearKernel(MPLinearKernel):

25 return 90

26

27 @classmethod

29 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

28 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

29 if not current_platform.is_cuda():

30 return False, "CUTLASS only supported on CUDA"

31

@@ -95,7 +94,7 @@ class CutlassW4A8LinearKernel(MPLinearKernel):

94 self,

95 layer: torch.nn.Module,

96 x: torch.Tensor,

98 - bias: Optional[torch.Tensor] = None,

97 + bias: torch.Tensor | None = None,

98 ) -> torch.Tensor:

99 c = self.config

100 w_q, w_s, _, _ = self._get_weight_params(layer)

vllm/model_executor/layers/quantization/kernels/mixed_precision/dynamic_4bit.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -20,7 +19,7 @@ class Dynamic4bitLinearKernel(MPLinearKernel):

19 return 1

20

21 @classmethod

23 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

22 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

23 if not current_platform.is_cpu():

24 return False, "Only CPU is supported"

25 if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:

@@ -95,7 +94,7 @@ class Dynamic4bitLinearKernel(MPLinearKernel):

94 self,

95 layer: torch.nn.Module,

96 x: torch.Tensor,

98 - bias: Optional[torch.Tensor] = None,

97 + bias: torch.Tensor | None = None,

98 ) -> torch.Tensor:

99 c = self.config

100 x_2d = x.reshape(-1, x.shape[-1])

vllm/model_executor/layers/quantization/kernels/mixed_precision/exllama.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -25,7 +24,7 @@ class ExllamaLinearKernel(MPLinearKernel):

24 return 60

25

26 @classmethod

28 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

27 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

28 if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:

29 return (

30 False,

@@ -137,7 +136,7 @@ class ExllamaLinearKernel(MPLinearKernel):

136 self,

137 layer: torch.nn.Module,

138 x: torch.Tensor,

140 - bias: Optional[torch.Tensor] = None,

139 + bias: torch.Tensor | None = None,

140 ) -> torch.Tensor:

141 c = self.config

142

vllm/model_executor/layers/quantization/kernels/mixed_precision/machete.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from functools import partial

5 -from typing import Optional

5

6 import torch

7

@@ -28,7 +27,7 @@ class MacheteLinearKernel(MPLinearKernel):

27 return 90

28

29 @classmethod

31 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

30 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

31 # Machete uses CUTLASS, so it can only be compatible with Nvidia

32 if not current_platform.is_cuda():

33 return False, "Machete only supported on CUDA"

@@ -129,7 +128,7 @@ class MacheteLinearKernel(MPLinearKernel):

128 self,

129 layer: torch.nn.Module,

130 x: torch.Tensor,

132 - bias: Optional[torch.Tensor] = None,

131 + bias: torch.Tensor | None = None,

132 ) -> torch.Tensor:

133 c = self.config

134 w_q, w_s, w_zp, _ = self._get_weight_params(layer)

vllm/model_executor/layers/quantization/kernels/mixed_precision/marlin.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -32,7 +31,7 @@ class MarlinLinearKernel(MPLinearKernel):

31 return 80

32

33 @classmethod

35 - def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:

34 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:

35 # Marlin uses inline PTX, so it can only be compatible with Nvidia

36 if not current_platform.is_cuda():

37 return False, "Marlin only supported on CUDA"

@@ -144,7 +143,7 @@ class MarlinLinearKernel(MPLinearKernel):

143 self,

144 layer: torch.nn.Module,

145 x: torch.Tensor,

147 - bias: Optional[torch.Tensor] = None,

146 + bias: torch.Tensor | None = None,

147 ) -> torch.Tensor:

148 c = self.config

149 w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)

vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py

+5 -6

@@ -3,7 +3,6 @@

3

4 from abc import ABC, abstractmethod

5 from dataclasses import dataclass

6 -from typing import Optional

6

7 import torch

8

@@ -23,7 +22,7 @@ class ScaledMMLinearKernel(ABC):

22

23 @classmethod

24 @abstractmethod

26 - def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

25 + def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

26 raise NotImplementedError

27

28 def __init__(

@@ -52,7 +51,7 @@ class ScaledMMLinearKernel(ABC):

51 self,

52 layer: torch.nn.Module,

53 x: torch.Tensor,

55 - bias: Optional[torch.Tensor] = None,

54 + bias: torch.Tensor | None = None,

55 ) -> torch.Tensor:

56 raise NotImplementedError

57

@@ -61,9 +60,9 @@ class ScaledMMLinearKernel(ABC):

60 ) -> tuple[

61 torch.Tensor, # weight

62 torch.Tensor, # weight_scale

64 - Optional[torch.Tensor], # input_scale,

65 - Optional[torch.Tensor], # input_zp

66 - Optional[torch.Tensor], # azp_adj

63 + torch.Tensor | None, # input_scale,

64 + torch.Tensor | None, # input_zp

65 + torch.Tensor | None, # azp_adj

66 ]:

67 return (

68 getattr(layer, self.w_q_name),

vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import os

5 -from typing import Optional

5

6 from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (

7 AiterScaledMMLinearKernel,

@@ -35,7 +34,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[ScaledMMLinearKernel]]] = {

34

35

36 def choose_scaled_mm_linear_kernel(

38 - config: ScaledMMLinearLayerConfig, compute_capability: Optional[int] = None

37 + config: ScaledMMLinearLayerConfig, compute_capability: int | None = None

38 ) -> type[ScaledMMLinearKernel]:

39 """

40 Choose an ScaledMMLinearKernel that can implement the given config for the

vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py

+4 -5

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -19,7 +18,7 @@ def rocm_aiter_gemm_w8a8_impl(

18 B: torch.Tensor,

19 As: torch.Tensor,

20 Bs: torch.Tensor,

22 - bias: Optional[torch.Tensor] = None,

21 + bias: torch.Tensor | None = None,

22 output_dtype: torch.dtype = torch.float16,

23 ) -> torch.Tensor:

24 from aiter import gemm_a8w8_CK

@@ -36,7 +35,7 @@ def rocm_aiter_gemm_w8a8_fake(

35 B: torch.Tensor,

36 As: torch.Tensor,

37 Bs: torch.Tensor,

39 - bias: Optional[torch.Tensor] = None,

38 + bias: torch.Tensor | None = None,

39 output_dtype: torch.dtype = torch.float16,

40 ) -> torch.Tensor:

41 m = A.shape[0]

@@ -59,7 +58,7 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):

58 return 90

59

60 @classmethod

62 - def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

61 + def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

62 if not current_platform.is_rocm():

63 return (

64 False,

@@ -99,7 +98,7 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):

98 self,

99 layer: torch.nn.Module,

100 x: torch.Tensor,

102 - bias: Optional[torch.Tensor] = None,

101 + bias: torch.Tensor | None = None,

102 ) -> torch.Tensor:

103 """

104 `AiterScaledMMLinearKernel` implements a fused version of

vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py

+4 -5

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -24,7 +23,7 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):

23 return 75

24

25 @classmethod

27 - def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

26 + def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

27 if not current_platform.is_cpu():

28 return False, "CPUScaledMM requires running on CPU."

29

@@ -173,7 +172,7 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):

172 self,

173 layer: torch.nn.Module,

174 x: torch.Tensor,

176 - bias: Optional[torch.Tensor] = None,

175 + bias: torch.Tensor | None = None,

176 ) -> torch.Tensor:

177 return self.linear_method(

178 layer,

@@ -185,7 +184,7 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):

184 self,

185 layer: torch.nn.Module,

186 x: torch.Tensor,

188 - bias: Optional[torch.Tensor] = None,

187 + bias: torch.Tensor | None = None,

188 ) -> torch.Tensor:

189 w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)

190

@@ -207,7 +206,7 @@ class CPUScaledMMLinearKernel(ScaledMMLinearKernel):

206 self,

207 layer: torch.nn.Module,

208 x: torch.Tensor,

210 - bias: Optional[torch.Tensor] = None,

209 + bias: torch.Tensor | None = None,

210 ) -> torch.Tensor:

211 w_q, w_s, _, _, _ = self._get_weight_params(layer)

212 return torch.ops._C.int8_scaled_mm_with_quant(

vllm/model_executor/layers/quantization/kernels/scaled_mm/cutlass.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -21,7 +20,7 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):

20 return 75

21

22 @classmethod

24 - def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

23 + def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

24 if not current_platform.is_cuda():

25 return False, "CutlassScaledMM requires running on CUDA."

26

@@ -110,7 +109,7 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):

109 self,

110 layer: torch.nn.Module,

111 x: torch.Tensor,

113 - bias: Optional[torch.Tensor] = None,

112 + bias: torch.Tensor | None = None,

113 ) -> torch.Tensor:

114 w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)

115

vllm/model_executor/layers/quantization/kernels/scaled_mm/triton.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -17,7 +16,7 @@ class TritonScaledMMLinearKernel(CutlassScaledMMLinearKernel):

16 return 75

17

18 @classmethod

20 - def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

19 + def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

20 if current_platform.is_cpu():

21 return (

22 False,

@@ -38,6 +37,6 @@ class TritonScaledMMLinearKernel(CutlassScaledMMLinearKernel):

37 self,

38 layer: torch.nn.Module,

39 x: torch.Tensor,

41 - bias: Optional[torch.Tensor] = None,

40 + bias: torch.Tensor | None = None,

41 ) -> torch.Tensor:

42 return super().apply_weights(layer, x, bias)

vllm/model_executor/layers/quantization/kernels/scaled_mm/xla.py

+4 -5

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import warnings

5 -from typing import Optional

5

6 import torch

7 from functorch.experimental.control_flow import cond # noqa: F401

@@ -25,7 +24,7 @@ class XLAScaledMMLinearKernel(ScaledMMLinearKernel):

24 )

25

26 @classmethod

28 - def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:

27 + def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

28 if not current_platform.is_tpu():

29 return False, "ScaledMMXLA requires running on TPU."

30

@@ -77,17 +76,17 @@ class XLAScaledMMLinearKernel(ScaledMMLinearKernel):

76 message="Pred is a Python constant. When used with torch.cond, it specializes on one of the branches.", # noqa: E501

77 )

78

80 - def no_add_bias(self, x: torch.Tensor, bias: Optional[torch.Tensor]):

79 + def no_add_bias(self, x: torch.Tensor, bias: torch.Tensor | None):

80 return x

81

83 - def add_bias(self, x: torch.Tensor, bias: Optional[torch.Tensor]):

82 + def add_bias(self, x: torch.Tensor, bias: torch.Tensor | None):

83 return x + bias

84

85 def apply_weights(

86 self,

87 layer: torch.nn.Module,

88 x: torch.Tensor,

90 - bias: Optional[torch.Tensor] = None,

89 + bias: torch.Tensor | None = None,

90 ) -> torch.Tensor:

91 w_q, w_s, _, _, _ = self._get_weight_params(layer)

92

vllm/model_executor/layers/quantization/modelopt.py

+33 -32

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import TYPE_CHECKING, Any, Callable, Optional, Union

4 +from collections.abc import Callable

5 +from typing import TYPE_CHECKING, Any, Optional

6

7 import torch

8 from torch.nn import Module

@@ -92,8 +93,8 @@ class ModelOptFp8Config(QuantizationConfig):

93 def __init__(

94 self,

95 is_checkpoint_fp8_serialized: bool = False,

95 - kv_cache_quant_method: Optional[str] = None,

96 - exclude_modules: Optional[list[str]] = None,

96 + kv_cache_quant_method: str | None = None,

97 + exclude_modules: list[str] | None = None,

98 ) -> None:

99 super().__init__()

100 self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized

@@ -128,7 +129,7 @@ class ModelOptFp8Config(QuantizationConfig):

129 @classmethod

130 def override_quantization_method(

131 cls, hf_quant_cfg, user_quant

131 - ) -> Optional[QuantizationMethods]:

132 + ) -> QuantizationMethods | None:

133 """Detect if this ModelOpt config should be used based on

134 quantization config."""

135

@@ -319,7 +320,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):

320 self,

321 layer: torch.nn.Module,

322 x: torch.Tensor,

322 - bias: Optional[torch.Tensor] = None,

323 + bias: torch.Tensor | None = None,

324 ) -> torch.Tensor:

325 return self.fp8_linear.apply(

326 input=x,

@@ -351,7 +352,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

352 )

353

354 self.cutlass_fp8_supported = cutlass_fp8_supported()

354 - self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None

355 + self.flashinfer_moe_backend: FlashinferMoeBackend | None = None

356 if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():

357 self.flashinfer_moe_backend = get_flashinfer_moe_backend()

358 logger.info_once(

@@ -360,7 +361,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

361

362 def maybe_make_prepare_finalize(

363 self,

363 - ) -> Optional[mk.FusedMoEPrepareAndFinalize]:

364 + ) -> mk.FusedMoEPrepareAndFinalize | None:

365 # TRT LLM not supported with all2all yet.

366 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

367 return None

@@ -541,7 +542,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

542

543 def get_fused_moe_quant_config(

544 self, layer: torch.nn.Module

544 - ) -> Optional[FusedMoEQuantConfig]:

545 + ) -> FusedMoEQuantConfig | None:

546 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

547 return None

548

@@ -561,21 +562,21 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

562 top_k: int,

563 renormalize: bool,

564 use_grouped_topk: bool = False,

564 - topk_group: Optional[int] = None,

565 - num_expert_group: Optional[int] = None,

565 + topk_group: int | None = None,

566 + num_expert_group: int | None = None,

567 global_num_experts: int = -1,

567 - expert_map: Optional[torch.Tensor] = None,

568 - custom_routing_function: Optional[Callable] = None,

568 + expert_map: torch.Tensor | None = None,

569 + custom_routing_function: Callable | None = None,

570 scoring_func: str = "softmax",

571 routed_scaling_factor: float = 1.0,

571 - e_score_correction_bias: Optional[torch.Tensor] = None,

572 + e_score_correction_bias: torch.Tensor | None = None,

573 apply_router_weight_on_input: bool = False,

574 activation: str = "silu",

575 enable_eplb: bool = False,

575 - expert_load_view: Optional[torch.Tensor] = None,

576 - logical_to_physical_map: Optional[torch.Tensor] = None,

577 - logical_replica_count: Optional[torch.Tensor] = None,

578 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

576 + expert_load_view: torch.Tensor | None = None,

577 + logical_to_physical_map: torch.Tensor | None = None,

578 + logical_replica_count: torch.Tensor | None = None,

579 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

580 if enable_eplb:

581 raise NotImplementedError(

582 "EPLB not supported for `ModelOptFp8MoEMethod` yet."

@@ -674,7 +675,7 @@ class ModelOptNvFp4Config(QuantizationConfig):

675 def __init__(

676 self,

677 is_checkpoint_nvfp4_serialized: bool,

677 - kv_cache_quant_algo: Optional[str],

678 + kv_cache_quant_algo: str | None,

679 exclude_modules: list[str],

680 group_size: int = 16,

681 ) -> None:

@@ -713,7 +714,7 @@ class ModelOptNvFp4Config(QuantizationConfig):

714 @classmethod

715 def override_quantization_method(

716 cls, hf_quant_cfg, user_quant

716 - ) -> Optional[QuantizationMethods]:

717 + ) -> QuantizationMethods | None:

718 """Detect if this ModelOpt FP4 config should be used based on

719 quantization config."""

720 if hf_quant_cfg is None:

@@ -906,7 +907,7 @@ class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):

907 Supports loading kv-cache scaling factors from FP8 checkpoints.

908 """

909

909 - def __init__(self, quant_config: Union[ModelOptFp8Config, ModelOptNvFp4Config]):

910 + def __init__(self, quant_config: ModelOptFp8Config | ModelOptNvFp4Config):

911 super().__init__(quant_config)

912

913

@@ -1071,7 +1072,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

1072 self,

1073 layer: torch.nn.Module,

1074 x: torch.Tensor,

1074 - bias: Optional[torch.Tensor] = None,

1075 + bias: torch.Tensor | None = None,

1076 ) -> torch.Tensor:

1077 if self.backend == "marlin":

1078 return apply_fp4_marlin_linear(

@@ -1162,7 +1163,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1163 " for ModelOptNvFp4FusedMoE."

1164 )

1165

1165 - def maybe_make_prepare_finalize(self) -> Optional[mk.FusedMoEPrepareAndFinalize]:

1166 + def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

1167 if self.use_marlin or (

1168 self.allow_flashinfer

1169 and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

@@ -1565,7 +1566,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1566

1567 def get_fused_moe_quant_config(

1568 self, layer: torch.nn.Module

1568 - ) -> Optional[FusedMoEQuantConfig]:

1569 + ) -> FusedMoEQuantConfig | None:

1570 if (

1571 self.use_marlin

1572 or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

@@ -1589,21 +1590,21 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1590 top_k: int,

1591 renormalize: bool,

1592 use_grouped_topk: bool = False,

1592 - topk_group: Optional[int] = None,

1593 - num_expert_group: Optional[int] = None,

1593 + topk_group: int | None = None,

1594 + num_expert_group: int | None = None,

1595 global_num_experts: int = -1,

1595 - expert_map: Optional[torch.Tensor] = None,

1596 - custom_routing_function: Optional[Callable] = None,

1596 + expert_map: torch.Tensor | None = None,

1597 + custom_routing_function: Callable | None = None,

1598 scoring_func: str = "softmax",

1599 routed_scaling_factor: float = 1.0,

1599 - e_score_correction_bias: Optional[torch.Tensor] = None,

1600 + e_score_correction_bias: torch.Tensor | None = None,

1601 apply_router_weight_on_input: bool = False,

1602 activation: str = "silu",

1603 enable_eplb: bool = False,

1603 - expert_load_view: Optional[torch.Tensor] = None,

1604 - logical_to_physical_map: Optional[torch.Tensor] = None,

1605 - logical_replica_count: Optional[torch.Tensor] = None,

1606 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

1604 + expert_load_view: torch.Tensor | None = None,

1605 + logical_to_physical_map: torch.Tensor | None = None,

1606 + logical_replica_count: torch.Tensor | None = None,

1607 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1608 if enable_eplb:

1609 raise NotImplementedError(

1610 "EPLB not supported for `ModelOptNvFp4FusedMoE` yet."

vllm/model_executor/layers/quantization/moe_wna16.py

+14 -13

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Any, Callable, Optional, Union

4 +from collections.abc import Callable

5 +from typing import Any, Optional

6

7 import torch

8

@@ -40,7 +41,7 @@ class MoeWNA16Config(QuantizationConfig):

41 group_size: int,

42 has_zp: bool,

43 lm_head_quantized: bool,

43 - modules_to_not_convert: Optional[list[str]],

44 + modules_to_not_convert: list[str] | None,

45 full_config: dict[str, Any],

46 ) -> None:

47 super().__init__()

@@ -127,7 +128,7 @@ class MoeWNA16Config(QuantizationConfig):

128 @classmethod

129 def override_quantization_method(

130 cls, hf_quant_cfg, user_quant

130 - ) -> Optional[QuantizationMethods]:

131 + ) -> QuantizationMethods | None:

132 can_convert = cls.is_moe_wna16_compatible(hf_quant_cfg)

133 if can_convert and user_quant == "moe_wna16":

134 return cls.get_name()

@@ -339,7 +340,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

340

341 def get_fused_moe_quant_config(

342 self, layer: torch.nn.Module

342 - ) -> Optional[FusedMoEQuantConfig]:

343 + ) -> FusedMoEQuantConfig | None:

344 weight_bits = self.quant_config.weight_bits

345 has_zp = self.quant_config.has_zp

346 assert weight_bits == 4 or weight_bits == 8

@@ -365,21 +366,21 @@ class MoeWNA16Method(FusedMoEMethodBase):

366 top_k: int,

367 renormalize: bool,

368 use_grouped_topk: bool = False,

368 - topk_group: Optional[int] = None,

369 - num_expert_group: Optional[int] = None,

369 + topk_group: int | None = None,

370 + num_expert_group: int | None = None,

371 global_num_experts: int = -1,

371 - expert_map: Optional[torch.Tensor] = None,

372 - custom_routing_function: Optional[Callable] = None,

372 + expert_map: torch.Tensor | None = None,

373 + custom_routing_function: Callable | None = None,

374 scoring_func: str = "softmax",

375 routed_scaling_factor: float = 1.0,

375 - e_score_correction_bias: Optional[torch.Tensor] = None,

376 + e_score_correction_bias: torch.Tensor | None = None,

377 apply_router_weight_on_input: bool = False,

378 activation: str = "silu",

379 enable_eplb: bool = False,

379 - expert_load_view: Optional[torch.Tensor] = None,

380 - logical_to_physical_map: Optional[torch.Tensor] = None,

381 - logical_replica_count: Optional[torch.Tensor] = None,

382 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

380 + expert_load_view: torch.Tensor | None = None,

381 + logical_to_physical_map: torch.Tensor | None = None,

382 + logical_replica_count: torch.Tensor | None = None,

383 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

384 assert self.fused_experts is None

385 if enable_eplb:

386 raise NotImplementedError("EPLB not supported for `MoeWNA16Method` yet.")

vllm/model_executor/layers/quantization/mxfp4.py

+21 -20

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Callable

4 from enum import Enum

4 -from typing import Callable, Optional, Union

5 +from typing import Optional

6

7 import torch

8 from torch.nn.parameter import Parameter

@@ -137,7 +138,7 @@ def get_mxfp4_backend():

138

139

140 class Mxfp4Config(QuantizationConfig):

140 - def __init__(self, ignored_layers: Optional[list[str]] = None):

141 + def __init__(self, ignored_layers: list[str] | None = None):

142 super().__init__()

143 self.ignored_layers = ignored_layers

144

@@ -756,7 +757,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

757

758 def get_fused_moe_quant_config(

759 self, layer: torch.nn.Module

759 - ) -> Optional[FusedMoEQuantConfig]:

760 + ) -> FusedMoEQuantConfig | None:

761 if self.mxfp4_backend == Mxfp4Backend.MARLIN:

762 return mxfp4_w4a16_moe_quant_config(

763 w1_bias=layer.w13_bias,

@@ -824,19 +825,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

825 top_k: int,

826 renormalize: bool,

827 use_grouped_topk: bool = False,

827 - topk_group: Optional[int] = None,

828 - num_expert_group: Optional[int] = None,

828 + topk_group: int | None = None,

829 + num_expert_group: int | None = None,

830 global_num_experts: int = -1,

830 - expert_map: Optional[torch.Tensor] = None,

831 - custom_routing_function: Optional[Callable] = None,

831 + expert_map: torch.Tensor | None = None,

832 + custom_routing_function: Callable | None = None,

833 scoring_func: str = "softmax",

833 - e_score_correction_bias: Optional[torch.Tensor] = None,

834 + e_score_correction_bias: torch.Tensor | None = None,

835 apply_router_weight_on_input: bool = False,

836 activation: str = "silu",

837 enable_eplb: bool = False,

837 - expert_load_view: Optional[torch.Tensor] = None,

838 - logical_to_physical_map: Optional[torch.Tensor] = None,

839 - logical_replica_count: Optional[torch.Tensor] = None,

838 + expert_load_view: torch.Tensor | None = None,

839 + logical_to_physical_map: torch.Tensor | None = None,

840 + logical_replica_count: torch.Tensor | None = None,

841 ) -> torch.Tensor:

842 assert isinstance(self.fused_experts, mk.FusedMoEModularKernel)

843

@@ -890,21 +891,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

891 top_k: int,

892 renormalize: bool,

893 use_grouped_topk: bool = False,

893 - topk_group: Optional[int] = None,

894 - num_expert_group: Optional[int] = None,

894 + topk_group: int | None = None,

895 + num_expert_group: int | None = None,

896 global_num_experts: int = -1,

896 - expert_map: Optional[torch.Tensor] = None,

897 - custom_routing_function: Optional[Callable] = None,

897 + expert_map: torch.Tensor | None = None,

898 + custom_routing_function: Callable | None = None,

899 scoring_func: str = "softmax",

900 routed_scaling_factor: float = 1.0,

900 - e_score_correction_bias: Optional[torch.Tensor] = None,

901 + e_score_correction_bias: torch.Tensor | None = None,

902 apply_router_weight_on_input: bool = False,

903 activation: str = "silu",

904 enable_eplb: bool = False,

904 - expert_load_view: Optional[torch.Tensor] = None,

905 - logical_to_physical_map: Optional[torch.Tensor] = None,

906 - logical_replica_count: Optional[torch.Tensor] = None,

907 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

905 + expert_load_view: torch.Tensor | None = None,

906 + logical_to_physical_map: torch.Tensor | None = None,

907 + logical_replica_count: torch.Tensor | None = None,

908 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

909 if enable_eplb:

910 raise NotImplementedError("EPLB is not supported for mxfp4")

911

vllm/model_executor/layers/quantization/petit.py

+5 -5

@@ -41,9 +41,9 @@ class PetitNvFp4Config(QuantizationConfig):

41 def __init__(

42 self,

43 is_checkpoint_nvfp4_serialized: bool = False,

44 - kv_cache_quant_algo: Optional[str] = None,

45 - group_size: Optional[int] = None,

46 - exclude_modules: Optional[list[str]] = None,

44 + kv_cache_quant_algo: str | None = None,

45 + group_size: int | None = None,

46 + exclude_modules: list[str] | None = None,

47 ) -> None:

48 self._check_hardware_support()

49 self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized

@@ -133,7 +133,7 @@ class PetitNvFp4Config(QuantizationConfig):

133 @classmethod

134 def override_quantization_method(

135 cls, hf_quant_cfg, user_quant

136 - ) -> Optional[QuantizationMethods]:

136 + ) -> QuantizationMethods | None:

137 if not current_platform.is_rocm():

138 return None

139

@@ -307,7 +307,7 @@ class PetitNvFp4LinearMethod(LinearMethodBase):

307 self,

308 layer: torch.nn.Module,

309 x: torch.Tensor,

310 - bias: Optional[torch.Tensor] = None,

310 + bias: torch.Tensor | None = None,

311 ) -> torch.Tensor:

312 return apply_petit_nvfp4_linear(

313 input=x,

vllm/model_executor/layers/quantization/ptpc_fp8.py

+2 -2

@@ -34,7 +34,7 @@ class PTPCFp8Config(Fp8Config):

34 def __init__(

35 self,

36 activation_scheme: str = "dynamic",

37 - ignored_layers: Optional[list[str]] = None,

37 + ignored_layers: list[str] | None = None,

38 ) -> None:

39 if not current_platform.is_rocm():

40 raise ValueError("ptpc_fp8 quantization is supported only on ROCm.")

@@ -125,7 +125,7 @@ class PTPCFp8LinearMethod(Fp8LinearMethod):

125 self,

126 layer: torch.nn.Module,

127 x: torch.Tensor,

128 - bias: Optional[torch.Tensor] = None,

128 + bias: torch.Tensor | None = None,

129 ) -> torch.Tensor:

130 return self.fp8_linear.apply(

131 input=x,

vllm/model_executor/layers/quantization/quark/quark.py

+11 -11

@@ -43,8 +43,8 @@ class QuarkConfig(QuantizationConfig):

43 def __init__(

44 self,

45 quant_config: dict[str, Any],

46 - kv_cache_group: Optional[list[str]] = None,

47 - kv_cache_config: Optional[dict[str, Any]] = None,

46 + kv_cache_group: list[str] | None = None,

47 + kv_cache_config: dict[str, Any] | None = None,

48 pack_method: str = "reorder",

49 ):

50 super().__init__()

@@ -178,8 +178,8 @@ class QuarkConfig(QuantizationConfig):

178

179 def _is_fp8_w8a8(

180 self,

181 - weight_quant: Optional[dict[str, Any]],

182 - input_quant: Optional[dict[str, Any]],

181 + weight_quant: dict[str, Any] | None,

182 + input_quant: dict[str, Any] | None,

183 ) -> bool:

184 # Confirm weights and input quantized.

185 if weight_quant is None or input_quant is None:

@@ -209,8 +209,8 @@ class QuarkConfig(QuantizationConfig):

209

210 def _is_static_tensor_w8a8(

211 self,

212 - weight_quant: Optional[dict[str, Any]],

213 - input_quant: Optional[dict[str, Any]],

212 + weight_quant: dict[str, Any] | None,

213 + input_quant: dict[str, Any] | None,

214 ) -> bool:

215 # Confirm weights and input quantized.

216 if weight_quant is None or input_quant is None:

@@ -237,8 +237,8 @@ class QuarkConfig(QuantizationConfig):

237

238 def _is_ocp_mx(

239 self,

240 - weight_quant: Optional[dict[str, Any]],

241 - input_quant: Optional[dict[str, Any]],

240 + weight_quant: dict[str, Any] | None,

241 + input_quant: dict[str, Any] | None,

242 ) -> bool:

243 # Confirm weights and input quantized.

244 if weight_quant is None or input_quant is None:

@@ -370,7 +370,7 @@ class QuarkConfig(QuantizationConfig):

370

371 return scheme

372

373 - def get_cache_scale(self, name: str) -> Optional[str]:

373 + def get_cache_scale(self, name: str) -> str | None:

374 """

375 Check whether the param name matches the format for k/v cache scales

376 in quark. If this is the case, return its equivalent param name

@@ -429,7 +429,7 @@ class QuarkLinearMethod(LinearMethodBase):

429 self,

430 layer: torch.nn.Module,

431 x: torch.Tensor,

432 - bias: Optional[torch.Tensor] = None,

432 + bias: torch.Tensor | None = None,

433 ):

434 """

435 Use the output of create_weights and the CompressedTensorsScheme

@@ -454,7 +454,7 @@ class QuarkKVCacheMethod(BaseKVCacheMethod):

454 super().__init__(quant_config)

455

456 @staticmethod

457 - def validate_kv_cache_config(kv_cache_config: Optional[dict[str, Any]]):

457 + def validate_kv_cache_config(kv_cache_config: dict[str, Any] | None):

458 """

459 Validator for the kv cache configuration. Useful for controlling the

460 kv cache quantization schemes, that are being supported in vLLM

vllm/model_executor/layers/quantization/quark/quark_moe.py

+22 -21

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Any, Callable, Optional, Union

4 +from collections.abc import Callable

5 +from typing import Any

6

7 import torch

8

@@ -333,7 +334,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

334

335 def get_fused_moe_quant_config(

336 self, layer: torch.nn.Module

336 - ) -> Optional[FusedMoEQuantConfig]:

337 + ) -> FusedMoEQuantConfig | None:

338 return fp8_w8a8_moe_quant_config(

339 w1_scale=layer.w13_weight_scale,

340 w2_scale=layer.w2_weight_scale,

@@ -350,21 +351,21 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

351 top_k: int,

352 renormalize: bool,

353 use_grouped_topk: bool = False,

353 - topk_group: Optional[int] = None,

354 - num_expert_group: Optional[int] = None,

354 + topk_group: int | None = None,

355 + num_expert_group: int | None = None,

356 global_num_experts: int = -1,

356 - expert_map: Optional[torch.Tensor] = None,

357 - custom_routing_function: Optional[Callable] = None,

357 + expert_map: torch.Tensor | None = None,

358 + custom_routing_function: Callable | None = None,

359 scoring_func: str = "softmax",

360 routed_scaling_factor: float = 1.0,

360 - e_score_correction_bias: Optional[torch.Tensor] = None,

361 + e_score_correction_bias: torch.Tensor | None = None,

362 apply_router_weight_on_input: bool = False,

363 activation: str = "silu",

364 enable_eplb: bool = False,

364 - expert_load_view: Optional[torch.Tensor] = None,

365 - logical_to_physical_map: Optional[torch.Tensor] = None,

366 - logical_replica_count: Optional[torch.Tensor] = None,

367 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

365 + expert_load_view: torch.Tensor | None = None,

366 + logical_to_physical_map: torch.Tensor | None = None,

367 + logical_replica_count: torch.Tensor | None = None,

368 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

369 assert self.fused_experts is None

370

371 if enable_eplb:

@@ -568,7 +569,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

569

570 def get_fused_moe_quant_config(

571 self, layer: torch.nn.Module

571 - ) -> Optional[FusedMoEQuantConfig]:

572 + ) -> FusedMoEQuantConfig | None:

573 return ocp_mx_moe_quant_config(

574 quant_dtype=self.input_dtype,

575 weight_dtype=self.weight_dtype,

@@ -587,21 +588,21 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

588 top_k: int,

589 renormalize: bool,

590 use_grouped_topk: bool = False,

590 - topk_group: Optional[int] = None,

591 - num_expert_group: Optional[int] = None,

591 + topk_group: int | None = None,

592 + num_expert_group: int | None = None,

593 global_num_experts: int = -1,

593 - expert_map: Optional[torch.Tensor] = None,

594 - custom_routing_function: Optional[Callable] = None,

594 + expert_map: torch.Tensor | None = None,

595 + custom_routing_function: Callable | None = None,

596 scoring_func: str = "softmax",

597 routed_scaling_factor: float = 1.0,

597 - e_score_correction_bias: Optional[torch.Tensor] = None,

598 + e_score_correction_bias: torch.Tensor | None = None,

599 apply_router_weight_on_input: bool = False,

600 activation: str = "silu",

601 enable_eplb: bool = False,

601 - expert_load_view: Optional[torch.Tensor] = None,

602 - logical_to_physical_map: Optional[torch.Tensor] = None,

603 - logical_replica_count: Optional[torch.Tensor] = None,

604 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

602 + expert_load_view: torch.Tensor | None = None,

603 + logical_to_physical_map: torch.Tensor | None = None,

604 + logical_replica_count: torch.Tensor | None = None,

605 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

606 assert self.fused_experts is None

607

608 if enable_eplb:

vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py

+7 -6

@@ -1,9 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from collections.abc import Callable

5 from fractions import Fraction

6 from functools import cache, partial

6 -from typing import Any, Callable, Optional, Union

7 +from typing import Any

8

9 import torch

10 import torch.nn.functional as F

@@ -54,8 +55,8 @@ try:

55 weight: torch.Tensor,

56 weight_scale: torch.Tensor,

57 rocm_use_aiter_fp4_asm_gemm: bool = False,

57 - out_dtype: Optional[torch.dtype] = torch.bfloat16,

58 - x_scales: Optional[torch.Tensor] = None,

58 + out_dtype: torch.dtype | None = torch.bfloat16,

59 + x_scales: torch.Tensor | None = None,

60 ) -> torch.Tensor:

61 M = x.shape[0]

62 if rocm_use_aiter_fp4_asm_gemm:

@@ -95,7 +96,7 @@ try:

96 weight_scale: torch.Tensor,

97 x_scales: torch.Tensor = None,

98 rocm_use_aiter_fp4_asm_gemm: bool = False,

98 - out_dtype: Optional[torch.dtype] = torch.bfloat16,

99 + out_dtype: torch.dtype | None = torch.bfloat16,

100 ) -> torch.Tensor:

101 return torch.empty(

102 (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device

@@ -129,7 +130,7 @@ class QuarkOCP_MX(QuarkScheme):

130 )

131

132 if self.weight_dtype == "mxfp4":

132 - self.packed_factor: Union[int, Fraction] = 2

133 + self.packed_factor: int | Fraction = 2

134 self.dequant_func = dequant_mxfp4

135 else:

136 self.packed_factor = Fraction(numerator=8, denominator=6)

@@ -282,7 +283,7 @@ class QuarkOCP_MX(QuarkScheme):

283 self,

284 layer: torch.nn.Module,

285 x: torch.Tensor,

285 - bias: Optional[torch.Tensor] = None,

286 + bias: torch.Tensor | None = None,

287 ) -> torch.Tensor:

288 if self.emulate:

289 dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype)

vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from abc import ABC, abstractmethod

5 -from typing import Optional

5

6 import torch

7

@@ -33,7 +32,7 @@ class QuarkScheme(ABC):

32

33 @abstractmethod

34 def apply_weights(

36 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

35 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

36 ):

37 """

38 Run the forward pass for the particular scheme. This is where

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py

+5 -4

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Any, Callable, Optional, cast

4 +from collections.abc import Callable

5 +from typing import Any, cast

6

7 import torch

8 from torch.nn import Parameter

@@ -25,11 +26,11 @@ __all__ = ["QuarkW8A8Fp8"]

26

27 class QuarkW8A8Fp8(QuarkScheme):

28 def __init__(

28 - self, weight_config: dict[str, Any], input_config: Optional[dict[str, Any]]

29 + self, weight_config: dict[str, Any], input_config: dict[str, Any] | None

30 ):

31 self.weight_qscheme = cast(str, weight_config.get("qscheme"))

32 self.is_static_input_scheme: bool = False

32 - self.input_qscheme: Optional[str] = None

33 + self.input_qscheme: str | None = None

34 if input_config is not None:

35 self.is_static_input_scheme = not cast(bool, input_config.get("is_dynamic"))

36 self.input_qscheme = cast(str, input_config.get("qscheme"))

@@ -166,7 +167,7 @@ class QuarkW8A8Fp8(QuarkScheme):

167 self,

168 layer: torch.nn.Module,

169 x: torch.Tensor,

169 - bias: Optional[torch.Tensor] = None,

170 + bias: torch.Tensor | None = None,

171 ) -> torch.Tensor:

172 return self.fp8_linear.apply(

173 input=x,

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py

+4 -4

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Callable, Optional

4 +from collections.abc import Callable

5

6 import torch

7

@@ -27,8 +27,8 @@ class QuarkW8A8Int8(QuarkScheme):

27 def __init__(

28 self,

29 qscheme: str,

30 - is_static_input_scheme: Optional[bool],

31 - input_symmetric: Optional[bool],

30 + is_static_input_scheme: bool | None,

31 + input_symmetric: bool | None,

32 ):

33 self.qscheme = qscheme

34 self.is_static_input_scheme = is_static_input_scheme

@@ -134,6 +134,6 @@ class QuarkW8A8Int8(QuarkScheme):

134 self.kernel.process_weights_after_loading(layer)

135

136 def apply_weights(

137 - self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]

137 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None

138 ) -> torch.Tensor:

139 return self.kernel.apply_weights(layer, x, bias)

vllm/model_executor/layers/quantization/quark/utils.py

+2 -2

@@ -3,7 +3,7 @@

3

4 from collections.abc import Iterable, Mapping

5 from types import MappingProxyType

6 -from typing import Any, Optional

6 +from typing import Any

7

8 import regex as re

9

@@ -22,7 +22,7 @@ def deep_compare(dict1: Any, dict2: Any) -> bool:

22

23

24 def should_ignore_layer(

25 - layer_name: Optional[str],

25 + layer_name: str | None,

26 ignore: Iterable[str],

27 fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),

28 ) -> bool:

vllm/model_executor/layers/quantization/rtn.py

+13 -12

@@ -3,7 +3,8 @@

3 # Copyright © 2025, Oracle and/or its affiliates.

4

5 import os

6 -from typing import Any, Callable, Optional, Union

6 +from collections.abc import Callable

7 +from typing import Any, Optional

8

9 import torch

10 import torch.nn.functional as F

@@ -226,7 +227,7 @@ class RTNLinearMethod(LinearMethodBase):

227 self,

228 layer: torch.nn.Module,

229 x: torch.Tensor,

229 - bias: Optional[torch.Tensor] = None,

230 + bias: torch.Tensor | None = None,

231 ) -> torch.Tensor:

232 qweight = layer.weight

233 scale = layer.scale

@@ -320,7 +321,7 @@ class RTNMoEMethod(FusedMoEMethodBase):

321

322 def get_fused_moe_quant_config(

323 self, layer: torch.nn.Module

323 - ) -> Optional[FusedMoEQuantConfig]:

324 + ) -> FusedMoEQuantConfig | None:

325 weight_bits = self.quant_config.weight_bits

326 group_size = self.quant_config.group_size

327 assert weight_bits == 4 or weight_bits == 8

@@ -345,21 +346,21 @@ class RTNMoEMethod(FusedMoEMethodBase):

346 top_k: int,

347 renormalize: bool,

348 use_grouped_topk: bool = False,

348 - topk_group: Optional[int] = None,

349 - num_expert_group: Optional[int] = None,

349 + topk_group: int | None = None,

350 + num_expert_group: int | None = None,

351 global_num_experts: int = -1,

351 - expert_map: Optional[torch.Tensor] = None,

352 - custom_routing_function: Optional[Callable] = None,

352 + expert_map: torch.Tensor | None = None,

353 + custom_routing_function: Callable | None = None,

354 scoring_func: str = "softmax",

355 routed_scaling_factor: float = 1.0,

355 - e_score_correction_bias: Optional[torch.Tensor] = None,

356 + e_score_correction_bias: torch.Tensor | None = None,

357 apply_router_weight_on_input: bool = False,

358 activation: str = "silu",

359 enable_eplb: bool = False,

359 - expert_load_view: Optional[torch.Tensor] = None,

360 - logical_to_physical_map: Optional[torch.Tensor] = None,

361 - logical_replica_count: Optional[torch.Tensor] = None,

362 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

360 + expert_load_view: torch.Tensor | None = None,

361 + logical_to_physical_map: torch.Tensor | None = None,

362 + logical_replica_count: torch.Tensor | None = None,

363 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

364 assert self.fused_experts is None

365

366 if enable_eplb:

vllm/model_executor/layers/quantization/schema.py

+1 -3

@@ -13,8 +13,6 @@ possible on ROCm), the model can be optionally augmented with KV cache

13 scaling factors.

14 """

15

16 -from typing import Optional

17 -

16 from pydantic import BaseModel, ConfigDict, ValidationInfo, model_validator

17

18

@@ -75,7 +73,7 @@ class QuantParamSchema(BaseModel):

73 # TODO: Generalize and extend with more fields

74 # (e.g. weights/activations params) once functionality is enabled

75 model_config = ConfigDict(protected_namespaces=())

78 - model_type: Optional[str]

76 + model_type: str | None

77 kv_cache: KVCacheQuantSchema

78

79 @model_validator(mode="after")

vllm/model_executor/layers/quantization/torchao.py

+2 -2

@@ -63,7 +63,7 @@ class TorchAOConfig(QuantizationConfig):

63 def __init__(

64 self,

65 torchao_config,

66 - skip_modules: Optional[list[str]] = None,

66 + skip_modules: list[str] | None = None,

67 is_checkpoint_torchao_serialized: bool = False,

68 ) -> None:

69 """

@@ -301,7 +301,7 @@ class TorchAOLinearMethod(LinearMethodBase):

301 self,

302 layer: torch.nn.Module,

303 x: torch.Tensor,

304 - bias: Optional[torch.Tensor] = None,

304 + bias: torch.Tensor | None = None,

305 ) -> torch.Tensor:

306 return F.linear(x, layer.weight, bias)

307

vllm/model_executor/layers/quantization/tpu_int8.py

+1 -1

@@ -119,7 +119,7 @@ class TPUInt8LinearMethod(LinearMethodBase):

119 self,

120 layer: torch.nn.Module,

121 x: torch.Tensor,

122 - bias: Optional[torch.Tensor] = None,

122 + bias: torch.Tensor | None = None,

123 ) -> torch.Tensor:

124 try:

125 import torch_xla.experimental.custom_kernel # noqa: F401

vllm/model_executor/layers/quantization/utils/bitblas_utils.py

+6 -7

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Optional

3

4 import torch

5 from packaging import version

@@ -29,7 +28,7 @@ BITBLAS_SUPPORTED_SYM = [False, True]

28 # Determines the supported quantization types for BitBLAS based on the

29 # device's capability and whether zero-point (zp) is used.

30 def query_bitblas_supported_quant_types(

32 - has_zp: bool, device_capability: Optional[int] = None

31 + has_zp: bool, device_capability: int | None = None

32 ):

33 if device_capability is None:

34 capability_tuple = current_platform.get_device_capability()

@@ -52,10 +51,10 @@ def query_bitblas_supported_quant_types(

51

52 def _check_bitblas_supported(

53 quant_type: ScalarType,

55 - group_size: Optional[int],

54 + group_size: int | None,

55 has_zp: bool,

57 - device_capability: Optional[int] = None,

58 -) -> tuple[bool, Optional[str]]:

56 + device_capability: int | None = None,

57 +) -> tuple[bool, str | None]:

58 if device_capability is None:

59 capability_tuple = current_platform.get_device_capability()

60 device_capability = (

@@ -99,7 +98,7 @@ def check_bitblas_supported(

98 quant_type: ScalarType,

99 group_size: int,

100 has_zp: bool = False,

102 - device_capability: Optional[int] = None,

101 + device_capability: int | None = None,

102 ) -> bool:

103 cond, _ = _check_bitblas_supported(

104 quant_type, group_size, has_zp, device_capability

@@ -156,7 +155,7 @@ def check_bitblas_supports_shape(

155 input_size_per_partition: int,

156 input_size: int,

157 group_size: int,

159 -) -> tuple[bool, Optional[str]]:

158 +) -> tuple[bool, str | None]:

159 try:

160 verify_bitblas_supports_shape(

161 output_size_per_partition, input_size_per_partition, input_size, group_size

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

-2

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Utility helpers for NVFP4 + FlashInfer fused-MoE path"""

4

5 -from __future__ import annotations

6 -

5 import torch

6

7 import vllm.envs as envs

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

+7 -8

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from enum import Enum

4 -from typing import Optional

4

5 import torch

6

@@ -101,10 +100,10 @@ def apply_flashinfer_per_tensor_scale_fp8(

100 layer: torch.nn.Module,

101 hidden_states: torch.Tensor,

102 router_logits: torch.Tensor,

104 - routing_bias: Optional[torch.Tensor],

103 + routing_bias: torch.Tensor | None,

104 top_k: int,

106 - num_expert_group: Optional[int],

107 - topk_group: Optional[int],

105 + num_expert_group: int | None,

106 + topk_group: int | None,

107 global_num_experts: int,

108 apply_router_weight_on_input: bool,

109 ) -> torch.Tensor:

@@ -186,7 +185,7 @@ def register_moe_scaling_factors(layer: torch.nn.Module) -> None:

185

186

187 def build_flashinfer_fp8_cutlass_moe_prepare_finalize(

189 - moe: Optional[FusedMoEConfig],

188 + moe: FusedMoEConfig | None,

189 ) -> mk.FusedMoEPrepareAndFinalize:

190 """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""

191 use_dp = moe.moe_parallel_config.dp_size > 1 if moe is not None else False

@@ -194,9 +193,9 @@ def build_flashinfer_fp8_cutlass_moe_prepare_finalize(

193

194

195 def select_cutlass_fp8_gemm_impl(

197 - moe: Optional[FusedMoEConfig],

196 + moe: FusedMoEConfig | None,

197 quant_config: FusedMoEQuantConfig,

199 - out_dtype: Optional[torch.dtype] = None,

198 + out_dtype: torch.dtype | None = None,

199 ) -> mk.FusedMoEPermuteExpertsUnpermute:

200 """Return a GEMM *experts* implementation for fused-MoE layers"""

201

@@ -225,7 +224,7 @@ def flashinfer_cutlass_moe_fp8(

224 inplace: bool = False,

225 activation: str = "silu",

226 global_num_experts: int = -1,

228 - expert_map: Optional[torch.Tensor] = None,

227 + expert_map: torch.Tensor | None = None,

228 apply_router_weight_on_input: bool = False,

229 ) -> torch.Tensor:

230 quant_config = layer.quant_method.get_fused_moe_quant_config(layer)

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+20 -20

@@ -5,8 +5,8 @@

5 import functools

6 import json

7 import os

8 -from collections.abc import Sequence

9 -from typing import Any, Callable, Optional, Union

8 +from collections.abc import Callable, Sequence

9 +from typing import Any

10

11 import torch

12

@@ -39,7 +39,7 @@ from vllm.utils.deep_gemm import (

39 logger = init_logger(__name__)

40

41

42 -def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:

42 +def is_fp8(x: torch.dtype | torch.Tensor) -> bool:

43 if isinstance(x, torch.Tensor):

44 x = x.dtype

45 return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz

@@ -54,7 +54,7 @@ def cutlass_scaled_mm(

54 Bs: torch.Tensor,

55 block_size: list[int],

56 output_dtype: torch.dtype = torch.float16,

57 - is_hopper: Optional[bool] = None,

57 + is_hopper: bool | None = None,

58 ) -> torch.Tensor:

59 if is_hopper is None:

60 is_hopper = current_platform.is_device_capability(90)

@@ -279,8 +279,8 @@ class W8A8BlockFp8LinearOp:

279 input: torch.Tensor,

280 weight: torch.Tensor,

281 weight_scale: torch.Tensor,

282 - input_scale: Optional[torch.Tensor] = None,

283 - bias: Optional[torch.Tensor] = None,

282 + input_scale: torch.Tensor | None = None,

283 + bias: torch.Tensor | None = None,

284 ) -> torch.Tensor:

285 assert input_scale is None

286 # View input as 2D matrix for fp8 methods

@@ -394,7 +394,7 @@ class W8A8BlockFp8LinearOp:

394 ],

395 torch.Tensor,

396 ],

397 - Optional[QuantFP8],

397 + QuantFP8 | None,

398 ]:

399 if use_cutlass:

400 return self._run_cutlass, (

@@ -418,7 +418,7 @@ class W8A8BlockFp8LinearOp:

418

419

420 def input_to_float8(

421 - x: torch.Tensor, dtype: Optional[torch.dtype] = None

421 + x: torch.Tensor, dtype: torch.dtype | None = None

422 ) -> tuple[torch.Tensor, torch.Tensor]:

423 """This function quantizes input values to float8 values "

424 "with tensor-wise quantization."""

@@ -568,10 +568,10 @@ def per_token_group_quant_fp8(

568 x: torch.Tensor,

569 group_size: int,

570 eps: float = 1e-10,

571 - dtype: Optional[torch.dtype] = None,

571 + dtype: torch.dtype | None = None,

572 column_major_scales: bool = False,

573 - out_q: Optional[torch.Tensor] = None,

574 - use_ue8m0: Optional[bool] = None,

573 + out_q: torch.Tensor | None = None,

574 + use_ue8m0: bool | None = None,

575 ) -> tuple[torch.Tensor, torch.Tensor]:

576 """Function to perform per-token-group quantization on an input tensor `x`.

577 It converts the tensor values into signed float8 values and returns the

@@ -754,7 +754,7 @@ def _w8a8_triton_block_scaled_mm(

754 @functools.lru_cache

755 def get_w8a8_block_fp8_configs(

756 N: int, K: int, block_n: int, block_k: int

757 -) -> Optional[dict[int, Any]]:

757 +) -> dict[int, Any] | None:

758 """

759 Return optimized configurations for the w8a8 block fp8 kernel.

760 The return value will be a dictionary that maps an irregular grid of

@@ -1012,7 +1012,7 @@ def validate_fp8_block_shape(

1012 def create_fp8_weight_parameter(

1013 output_size_per_partition: int,

1014 input_size_per_partition: int,

1015 - weight_loader: Optional[Callable],

1015 + weight_loader: Callable | None,

1016 ) -> torch.nn.Parameter:

1017 """Create FP8 weight parameter."""

1018 from vllm.model_executor.parameter import ModelWeightParameter

@@ -1033,8 +1033,8 @@ def create_fp8_scale_parameter(

1033 parameter_type: torch.nn.Parameter,

1034 output_partition_sizes: list[int],

1035 input_size_per_partition: int,

1036 - block_size: Optional[list[int]],

1037 - weight_loader: Optional[Callable],

1036 + block_size: list[int] | None,

1037 + weight_loader: Callable | None,

1038 ) -> torch.nn.Parameter:

1039 """Create scale parameter based on quantization strategy."""

1040 if parameter_type == ChannelQuantScaleParameter:

@@ -1070,7 +1070,7 @@ def create_fp8_scale_parameter(

1070

1071

1072 def create_fp8_input_scale(

1073 - output_partition_sizes: list[int], weight_loader: Optional[Callable]

1073 + output_partition_sizes: list[int], weight_loader: Callable | None

1074 ) -> torch.nn.Parameter:

1075 """Create input scale parameter for static activation quantization."""

1076 from vllm.model_executor.parameter import PerTensorScaleParameter

@@ -1087,8 +1087,8 @@ def process_fp8_weight_tensor_strategy(

1087 weight: torch.Tensor,

1088 weight_scale: torch.Tensor,

1089 logical_widths: list[int],

1090 - input_scale: Optional[torch.Tensor] = None,

1091 -) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:

1090 + input_scale: torch.Tensor | None = None,

1091 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

1092 """Process weights for tensor-wise quantization strategy."""

1093 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

1094 normalize_e4m3fn_to_e4m3fnuz,

@@ -1114,8 +1114,8 @@ def process_fp8_weight_tensor_strategy(

1114 def process_fp8_weight_channel_strategy(

1115 weight: torch.Tensor,

1116 weight_scale: torch.Tensor,

1117 - input_scale: Optional[torch.Tensor] = None,

1118 -) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:

1117 + input_scale: torch.Tensor | None = None,

1118 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

1119 """Process weights for channel-wise quantization strategy."""

1120 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

1121 normalize_e4m3fn_to_e4m3fnuz,

vllm/model_executor/layers/quantization/utils/gptq_utils.py

+7 -7

@@ -4,7 +4,7 @@ from collections.abc import Mapping

4 from copy import deepcopy

5 from fractions import Fraction

6 from types import MappingProxyType

7 -from typing import TYPE_CHECKING, Optional, Union

7 +from typing import TYPE_CHECKING

8

9 import regex as re

10 import torch

@@ -25,7 +25,7 @@ else:

25

26 # Match dynamic rules with module name (prefix) and override quantize

27 # config if module (prefix) matches a rule

28 -def override_config(config: Union[GPTQConfig, GPTQMarlinConfig], prefix: str):

28 +def override_config(config: GPTQConfig | GPTQMarlinConfig, prefix: str):

29 weight_bits = get_dynamic_override(config, prefix, "bits", config.weight_bits)

30 if isinstance(weight_bits, int):

31 config.weight_bits = weight_bits

@@ -60,11 +60,11 @@ def override_config(config: Union[GPTQConfig, GPTQMarlinConfig], prefix: str):

60

61

62 def get_dynamic_override(

63 - config: Union[GPTQConfig, GPTQMarlinConfig],

63 + config: GPTQConfig | GPTQMarlinConfig,

64 layer_name: str,

65 - key: Optional[str] = None,

66 - default_value: Union[int, bool, None] = None,

67 -) -> Union[dict, int, bool, None]:

65 + key: str | None = None,

66 + default_value: int | bool | None = None,

67 +) -> dict | int | bool | None:

68 for pattern, pattern_dict in config.dynamic.items():

69 # Negative match: matched modules are excluded from quantized init

70 if pattern.startswith("-:"):

@@ -126,7 +126,7 @@ def is_layer_gptq_quantized(

126

127

128 def get_linear_quant_method(

129 - config: Union[GPTQConfig, GPTQMarlinConfig],

129 + config: GPTQConfig | GPTQMarlinConfig,

130 layer: torch.nn.Module,

131 prefix: str,

132 linear_method_cls: type,

vllm/model_executor/layers/quantization/utils/int8_utils.py

+4 -4

@@ -6,7 +6,7 @@ import functools

6 import json

7 import logging

8 import os

9 -from typing import Any, Optional

9 +from typing import Any

10

11 import torch

12

@@ -21,8 +21,8 @@ def apply_w8a8_block_int8_linear(

21 weight: torch.Tensor,

22 block_size: list[int],

23 weight_scale: torch.Tensor,

24 - input_scale: Optional[torch.Tensor] = None,

25 - bias: Optional[torch.Tensor] = None,

24 + input_scale: torch.Tensor | None = None,

25 + bias: torch.Tensor | None = None,

26 ) -> torch.Tensor:

27 assert input_scale is None

28 # View input as 2D matrix for fp8 methods

@@ -359,7 +359,7 @@ def _w8a8_block_int8_matmul(

359 @functools.lru_cache

360 def get_w8a8_block_int8_configs(

361 N: int, K: int, block_n: int, block_k: int

362 -) -> Optional[dict[int, Any]]:

362 +) -> dict[int, Any] | None:

363 """

364 Return optimized configurations for the w8a8 block fp8 kernel.

365

vllm/model_executor/layers/quantization/utils/layer_utils.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Union

4

5 import torch

6

@@ -21,7 +20,7 @@ def update_tensor_inplace(dst: torch.Tensor, src: torch.Tensor):

20 # Newly generated tensors need to replace existing tensors that are

21 # already registered as parameters by vLLM (and won't be freed)

22 def replace_parameter(

24 - mod: torch.nn.Module, name: str, new: Union[torch.Tensor, torch.nn.Parameter]

23 + mod: torch.nn.Module, name: str, new: torch.Tensor | torch.nn.Parameter

24 ) -> None:

25 old = getattr(mod, name)

26 if (

vllm/model_executor/layers/quantization/utils/machete_utils.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -41,7 +40,7 @@ def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]:

40

41 def check_machete_supports_shape(

42 in_features: int, out_featrues: int

44 -) -> tuple[bool, Optional[str]]:

43 +) -> tuple[bool, str | None]:

44 if in_features % MACHETE_PREPACKED_BLOCK_SHAPE[0] != 0:

45 return (

46 False,

vllm/model_executor/layers/quantization/utils/marlin_utils.py

+9 -10

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import numpy

6 import torch

@@ -34,9 +33,9 @@ USE_FP32_REDUCE_DEFAULT = True

33 # without runtime zero-point. We support common cases, i.e. AWQ and GPTQ.

34 # TODO: we may want to move this into the C++ so its closer to the actual impl

35 def query_marlin_supported_quant_types(

37 - has_zp: Optional[bool] = None,

36 + has_zp: bool | None = None,

37 include_fp_type: bool = True,

39 - device_capability: Optional[int] = None,

38 + device_capability: int | None = None,

39 ):

40 if device_capability is None:

41 capability_tuple = current_platform.get_device_capability()

@@ -72,10 +71,10 @@ def query_marlin_supported_quant_types(

71

72 def _check_marlin_supported(

73 quant_type: ScalarType,

75 - group_size: Optional[int],

74 + group_size: int | None,

75 has_zp: bool,

77 - device_capability: Optional[int] = None,

78 -) -> tuple[bool, Optional[str]]:

76 + device_capability: int | None = None,

77 +) -> tuple[bool, str | None]:

78 if device_capability is None:

79 capability_tuple = current_platform.get_device_capability()

80 device_capability = (

@@ -109,7 +108,7 @@ def check_marlin_supported(

108 quant_type: ScalarType,

109 group_size: int,

110 has_zp: bool = False,

112 - device_capability: Optional[int] = None,

111 + device_capability: int | None = None,

112 ) -> bool:

113 cond, _ = _check_marlin_supported(quant_type, group_size, has_zp, device_capability)

114 return cond

@@ -164,7 +163,7 @@ def check_marlin_supports_shape(

163 input_size_per_partition: int,

164 input_size: int,

165 group_size: int,

167 -) -> tuple[bool, Optional[str]]:

166 +) -> tuple[bool, str | None]:

167 try:

168 verify_marlin_supports_shape(

169 output_size_per_partition, input_size_per_partition, input_size, group_size

@@ -445,7 +444,7 @@ def apply_gptq_marlin_linear(

444 output_size_per_partition: int,

445 input_size_per_partition: int,

446 is_k_full: bool,

448 - bias: Optional[torch.Tensor] = None,

447 + bias: torch.Tensor | None = None,

448 use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,

449 ) -> torch.Tensor:

450 reshaped_x = input.reshape(-1, input.shape[-1])

@@ -494,7 +493,7 @@ def apply_awq_marlin_linear(

493 quant_type: ScalarType,

494 output_size_per_partition: int,

495 input_size_per_partition: int,

497 - bias: Optional[torch.Tensor] = None,

496 + bias: torch.Tensor | None = None,

497 use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,

498 ) -> torch.Tensor:

499 reshaped_x = input.reshape(-1, input.shape[-1])

vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -95,11 +94,11 @@ def apply_fp4_marlin_linear(

94 input: torch.Tensor,

95 weight: torch.Tensor,

96 weight_scale: torch.Tensor,

98 - weight_scale_2: Optional[torch.Tensor],

97 + weight_scale_2: torch.Tensor | None,

98 workspace: torch.Tensor,

99 size_n: int,

100 size_k: int,

102 - bias: Optional[torch.Tensor] = None,

101 + bias: torch.Tensor | None = None,

102 use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,

103 ) -> torch.Tensor:

104 # For GPUs that lack FP4 hardware support, we can leverage the

vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -45,7 +44,7 @@ def apply_fp8_marlin_linear(

44 workspace: torch.Tensor,

45 size_n: int,

46 size_k: int,

48 - bias: Optional[torch.Tensor],

47 + bias: torch.Tensor | None,

48 use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,

49 ) -> torch.Tensor:

50 # For GPUs that lack FP8 hardware support, we can leverage the

vllm/model_executor/layers/quantization/utils/marlin_utils_test.py

+1 -3

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Utility functions used for tests and benchmarks"""

4

5 -from typing import Optional

6 -

5 import numpy as np

6 import torch

7

@@ -100,7 +98,7 @@ def marlin_quantize(

98 quant_type: ScalarType,

99 group_size: int,

100 act_order: bool,

103 - test_perm: Optional[torch.Tensor] = None,

101 + test_perm: torch.Tensor | None = None,

102 ):

103 size_k, size_n = w.shape

104 num_bits = quant_type.size_bits

vllm/model_executor/layers/quantization/utils/mxfp4_utils.py

+10 -9

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Any, Callable, Optional

3 +from collections.abc import Callable

4 +from typing import Any

5

6 import torch

7

@@ -71,17 +72,17 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):

72

73 def _can_support_mxfp4(

74 use_grouped_topk: bool = False,

74 - topk_group: Optional[int] = None,

75 - num_expert_group: Optional[int] = None,

76 - expert_map: Optional[torch.Tensor] = None,

77 - custom_routing_function: Optional[Callable] = None,

78 - e_score_correction_bias: Optional[torch.Tensor] = None,

75 + topk_group: int | None = None,

76 + num_expert_group: int | None = None,

77 + expert_map: torch.Tensor | None = None,

78 + custom_routing_function: Callable | None = None,

79 + e_score_correction_bias: torch.Tensor | None = None,

80 apply_router_weight_on_input: bool = False,

81 scoring_func: str = "softmax",

82 activation: str = "swigluoai",

82 - expert_load_view: Optional[torch.Tensor] = None,

83 - logical_to_physical_map: Optional[torch.Tensor] = None,

84 - logical_replica_count: Optional[torch.Tensor] = None,

83 + expert_load_view: torch.Tensor | None = None,

84 + logical_to_physical_map: torch.Tensor | None = None,

85 + logical_replica_count: torch.Tensor | None = None,

86 ):

87 return not (

88 use_grouped_topk

vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py

+1 -4

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from enum import Enum

4 -from typing import Union

4

5 from vllm.logger import init_logger

6

@@ -28,9 +27,7 @@ class OCP_MX_Scheme(str, Enum):

27 w_mxfp6_e2m3_a_mxfp6_e2m3 = "w_mxfp6_e2m3_a_mxfp6_e2m3"

28

29 @classmethod

31 - def from_quant_dtype(

32 - cls, input_dtype: Union[str, None], weight_dtype: Union[str, None]

33 - ):

30 + def from_quant_dtype(cls, input_dtype: str | None, weight_dtype: str | None):

31 if input_dtype not in OCP_MX_DTYPES or weight_dtype not in OCP_MX_DTYPES:

32 return None

33 elif input_dtype == "mxfp4" and weight_dtype == "mxfp4":

vllm/model_executor/layers/quantization/utils/petit_utils.py

+4 -4

@@ -43,8 +43,8 @@ _require_petit = _import_petit_kernel

43

44

45 def _check_petit_nvfp4_supported(

46 - quant_method: str, group_size: Optional[int]

47 -) -> tuple[bool, Optional[str]]:

46 + quant_method: str, group_size: int | None

47 +) -> tuple[bool, str | None]:

48 if quant_method != "NVFP4":

49 return (

50 False,

@@ -62,7 +62,7 @@ def _check_petit_nvfp4_supported(

62 return (True, None)

63

64

65 -def verify_petit_nvfp4_supported(quant_method: str, group_size: Optional[int]) -> None:

65 +def verify_petit_nvfp4_supported(quant_method: str, group_size: int | None) -> None:

66 supported, error_msg = _check_petit_nvfp4_supported(quant_method, group_size)

67 if not supported:

68 assert error_msg is not None

@@ -98,7 +98,7 @@ def apply_petit_nvfp4_linear(

98 weight_scale_2: torch.Tensor,

99 size_n: int,

100 size_k: int,

101 - bias: Optional[torch.Tensor] = None,

101 + bias: torch.Tensor | None = None,

102 ) -> torch.Tensor:

103 # Trigger (or get) the import here as well.

104 petit_kernel = _import_petit_kernel()

vllm/model_executor/layers/quantization/utils/quant_utils.py

+6 -6

@@ -5,7 +5,7 @@

5 from collections.abc import Mapping

6 from dataclasses import dataclass

7 from types import MappingProxyType

8 -from typing import ClassVar, NamedTuple, Optional

8 +from typing import ClassVar, NamedTuple

9

10 import numpy

11 import torch

@@ -91,7 +91,7 @@ class QuantKey:

91

92 dtype: torch.dtype

93 scale: ScaleDesc

94 - scale2: Optional[ScaleDesc] = None

94 + scale2: ScaleDesc | None = None

95 symmetric: bool = True

96

97 def __str__(self):

@@ -205,7 +205,7 @@ def scaled_quantize(

205 def scaled_dequantize(

206 x_q: torch.Tensor,

207 x_s: torch.Tensor,

208 - group_shape: Optional[GroupShape] = None,

208 + group_shape: GroupShape | None = None,

209 out_dtype: torch.dtype = torch.float32,

210 ) -> tuple[torch.Tensor, torch.Tensor]:

211 if group_shape is not None:

@@ -336,7 +336,7 @@ def permute_rows(

336 q_w: torch.Tensor,

337 w_ref: torch.Tensor,

338 group_size: int,

339 - test_perm: Optional[torch.Tensor] = None,

339 + test_perm: torch.Tensor | None = None,

340 ):

341 assert q_w.shape == w_ref.shape

342

@@ -365,7 +365,7 @@ def permute_rows(

365 def quantize_weights(

366 w: torch.Tensor,

367 quant_type: ScalarType,

368 - group_size: Optional[int],

368 + group_size: int | None,

369 zero_points: bool = False,

370 ref_zero_points_after_scales: bool = False,

371 ):

@@ -466,7 +466,7 @@ def gptq_quantize_weights(

466 quant_type: ScalarType,

467 group_size: int,

468 act_order: bool,

469 - test_perm: Optional[torch.Tensor] = None,

469 + test_perm: torch.Tensor | None = None,

470 ):

471 size_k, _ = w.shape

472

vllm/model_executor/layers/quantization/utils/w8a8_utils.py

+9 -9

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Callable, Optional, Union

4 +from collections.abc import Callable

5

6 import torch

7 from packaging import version

@@ -75,7 +75,7 @@ CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported()

75

76

77 def per_tensor_dequantize(

78 - tensor: torch.Tensor, inv_scale: Union[float, torch.Tensor]

78 + tensor: torch.Tensor, inv_scale: float | torch.Tensor

79 ) -> torch.Tensor:

80 fake_qweight = tensor.to(torch.float16)

81 dq_weight = fake_qweight * inv_scale

@@ -399,7 +399,7 @@ class Fp8LinearOp:

399 self,

400 act_quant_static: bool,

401 act_quant_group_shape: GroupShape = GroupShape.PER_TENSOR,

402 - pad_output: Optional[bool] = None,

402 + pad_output: bool | None = None,

403 ):

404 if current_platform.is_rocm():

405 self.preferred_backend = "rocm"

@@ -437,10 +437,10 @@ class Fp8LinearOp:

437 input: torch.Tensor,

438 weight: torch.Tensor,

439 weight_scale: torch.Tensor,

440 - out_dtype: Optional[torch.dtype] = None,

441 - input_scale: Optional[torch.Tensor] = None,

442 - input_scale_ub: Optional[torch.Tensor] = None,

443 - bias: Optional[torch.Tensor] = None,

440 + out_dtype: torch.dtype | None = None,

441 + input_scale: torch.Tensor | None = None,

442 + input_scale_ub: torch.Tensor | None = None,

443 + bias: torch.Tensor | None = None,

444 ) -> torch.Tensor:

445 # ops.scaled_fp8_quant supports both dynamic and static quant.

446 # If dynamic, layer.input_scale is None and x_scale computed from x.

@@ -486,8 +486,8 @@ class Fp8LinearOp:

486 def normalize_e4m3fn_to_e4m3fnuz(

487 weight: torch.Tensor,

488 weight_scale: torch.Tensor,

489 - input_scale: Optional[torch.Tensor] = None,

490 -) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:

489 + input_scale: torch.Tensor | None = None,

490 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

491 assert weight.dtype == torch.float8_e4m3fn

492 # The bits pattern 10000000(-128) represents zero in e4m3fn

493 # but NaN in e4m3fnuz. So here we set it to 0.

vllm/model_executor/layers/resampler.py

+9 -11

@@ -34,8 +34,8 @@ Example models: Qwen (Qwen-VL), MiniCPM-V 2.0

34 """

35

36 import math

37 +from collections.abc import Callable

38 from functools import partial

38 -from typing import Callable, Optional, Union

39

40 import numpy as np

41 import torch

@@ -48,9 +48,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

48 DEFAULT_LN = partial(nn.LayerNorm, eps=1e-6)

49

50

51 -def get_abs_pos(

52 - abs_pos: torch.Tensor, tgt_size: Union[torch.Tensor, int]

53 -) -> torch.Tensor:

51 +def get_abs_pos(abs_pos: torch.Tensor, tgt_size: torch.Tensor | int) -> torch.Tensor:

52 # abs_pos: L, C

53 # tgt_size: (H, W)

54 # return: M, C

@@ -124,7 +122,7 @@ def get_2d_sincos_pos_embed_from_grid(

122

123 def get_2d_sincos_pos_embed(

124 embed_dim: int,

127 - grid_size: Union[int, tuple[int, int]],

125 + grid_size: int | tuple[int, int],

126 cls_token: bool = False,

127 version: tuple[int, int] = (2, 0),

128 ) -> torch.Tensor:

@@ -168,10 +166,10 @@ class BaseResampler(nn.Module):

166 num_queries: int,

167 embed_dim: int,

168 num_heads: int,

171 - kv_dim: Optional[int] = None,

169 + kv_dim: int | None = None,

170 norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,

171 do_post_projection: bool = True,

174 - quant_config: Optional[QuantizationConfig] = None,

172 + quant_config: QuantizationConfig | None = None,

173 prefix: str = "",

174 ) -> None:

175 super().__init__()

@@ -222,11 +220,11 @@ class Resampler2(BaseResampler):

220 grid_size: int,

221 embed_dim: int,

222 num_heads: int,

225 - kv_dim: Optional[int] = None,

223 + kv_dim: int | None = None,

224 norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,

225 adaptive: bool = False,

226 do_post_projection: bool = True,

229 - quant_config: Optional[QuantizationConfig] = None,

227 + quant_config: QuantizationConfig | None = None,

228 prefix: str = "",

229 ) -> None:

230 super().__init__(

@@ -250,8 +248,8 @@ class Resampler2(BaseResampler):

248 def forward(

249 self,

250 x: torch.Tensor,

253 - tgt_sizes: Optional[torch.Tensor] = None,

254 - attn_mask: Optional[torch.Tensor] = None,

251 + tgt_sizes: torch.Tensor | None = None,

252 + attn_mask: torch.Tensor | None = None,

253 ) -> torch.Tensor:

254 if tgt_sizes is None:

255 tgt_sizes = int(math.sqrt(x.size(1)))

vllm/model_executor/layers/rotary_embedding/__init__.py

+4 -4

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Rotary Positional Embeddings."""

4

5 -from typing import Any, Optional

5 +from typing import Any

6

7 import torch

8

@@ -28,10 +28,10 @@ def get_rope(

28 max_position: int,

29 base: float,

30 is_neox_style: bool = True,

31 - rope_scaling: Optional[dict[str, Any]] = None,

32 - dtype: Optional[torch.dtype] = None,

31 + rope_scaling: dict[str, Any] | None = None,

32 + dtype: torch.dtype | None = None,

33 partial_rotary_factor: float = 1.0,

34 - dual_chunk_attention_config: Optional[dict[str, Any]] = None,

34 + dual_chunk_attention_config: dict[str, Any] | None = None,

35 ) -> RotaryEmbedding:

36 if dtype is None:

37 dtype = torch.get_default_dtype()

vllm/model_executor/layers/rotary_embedding/base.py

+8 -10

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Rotary Positional Embeddings Base Class."""

4

5 -from typing import Optional

6 -

5 import torch

6

7 from vllm.model_executor.custom_op import CustomOp

@@ -92,8 +90,8 @@ class RotaryEmbedding(CustomOp):

90 self,

91 positions: torch.Tensor,

92 query: torch.Tensor,

95 - key: Optional[torch.Tensor] = None,

96 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

93 + key: torch.Tensor | None = None,

94 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

95 """A PyTorch-native implementation of forward()."""

96 positions = positions.flatten()

97 num_tokens = positions.shape[0]

@@ -121,8 +119,8 @@ class RotaryEmbedding(CustomOp):

119 self,

120 positions: torch.Tensor,

121 query: torch.Tensor,

124 - key: Optional[torch.Tensor] = None,

125 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

122 + key: torch.Tensor | None = None,

123 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

124 if self.use_flashinfer:

125 torch.ops.vllm.flashinfer_rotary_embedding(

126 positions,

@@ -154,8 +152,8 @@ class RotaryEmbedding(CustomOp):

152 self,

153 positions: torch.Tensor,

154 query: torch.Tensor,

157 - key: Optional[torch.Tensor] = None,

158 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

155 + key: torch.Tensor | None = None,

156 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

157 if self.is_rocm_triton_rotary_embedding_enabled:

158 self._match_cos_sin_cache_dtype(query)

159 rocm_aiter_rotary_emb(

@@ -177,8 +175,8 @@ class RotaryEmbedding(CustomOp):

175 self,

176 positions: torch.Tensor,

177 query: torch.Tensor,

180 - key: Optional[torch.Tensor] = None,

181 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

178 + key: torch.Tensor | None = None,

179 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

180 from vllm._ipex_ops import ipex_ops as ops

181

182 self._match_cos_sin_cache_dtype(query)

vllm/model_executor/layers/rotary_embedding/common.py

+2 -2

@@ -2,9 +2,9 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import math

5 +from collections.abc import Callable

6 from functools import cache

7 from importlib.util import find_spec

7 -from typing import Callable, Optional

8

9 import torch

10

@@ -72,7 +72,7 @@ def apply_rotary_emb_dispatch(

72

73 @cache

74 def dispatch_rotary_emb_function(

75 - default: Optional[Callable[..., torch.Tensor]] = None,

75 + default: Callable[..., torch.Tensor] | None = None,

76 ) -> Callable[..., torch.Tensor]:

77 if current_platform.is_cuda():

78 return apply_rotary_emb

vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py

+6 -7

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import math

5 -from typing import Optional

5

6 import torch

7

@@ -110,9 +109,9 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):

109 self,

110 positions: torch.Tensor,

111 query: torch.Tensor,

113 - key: Optional[torch.Tensor] = None,

114 - offsets: Optional[torch.Tensor] = None,

115 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

112 + key: torch.Tensor | None = None,

113 + offsets: torch.Tensor | None = None,

114 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

115 """PyTorch-native implementation equivalent to forward()."""

116 assert key is not None

117 self._match_cos_sin_cache_dtype(query)

@@ -151,7 +150,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):

150 self,

151 positions: torch.Tensor,

152 query: torch.Tensor,

154 - key: Optional[torch.Tensor] = None,

155 - offsets: Optional[torch.Tensor] = None,

156 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

153 + key: torch.Tensor | None = None,

154 + offsets: torch.Tensor | None = None,

155 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

156 return self.forward_native(positions, query, key, offsets)

vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -121,7 +120,7 @@ class DualChunkRotaryEmbedding(CustomOp):

120 positions: torch.Tensor,

121 query: torch.Tensor,

122 key: torch.Tensor,

124 - offsets: Optional[torch.Tensor] = None,

123 + offsets: torch.Tensor | None = None,

124 ) -> tuple[torch.Tensor, torch.Tensor]:

125 query = query.view(*query.shape[:-1], -1, self.head_size)

126 key = key.view(*key.shape[:-1], -1, self.head_size)

@@ -185,7 +184,7 @@ class DualChunkRotaryEmbedding(CustomOp):

184 positions: torch.Tensor,

185 query: torch.Tensor,

186 key: torch.Tensor,

188 - offsets: Optional[torch.Tensor] = None,

187 + offsets: torch.Tensor | None = None,

188 ) -> tuple[torch.Tensor, torch.Tensor]:

189 return self.forward_native(positions, query, key, offsets)

190

vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py

+4 -5

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -16,8 +15,8 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):

15 self,

16 positions: torch.Tensor,

17 query: torch.Tensor,

19 - key: Optional[torch.Tensor] = None,

20 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

18 + key: torch.Tensor | None = None,

19 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

20 assert positions.ndim == 1 or positions.ndim == 2

21 assert key is not None

22

@@ -71,6 +70,6 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):

70 self,

71 positions: torch.Tensor,

72 query: torch.Tensor,

74 - key: Optional[torch.Tensor] = None,

75 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

73 + key: torch.Tensor | None = None,

74 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

75 return self.forward_native(positions, query, key)

vllm/model_executor/layers/rotary_embedding/linear_scaling_rope.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Union

4

5 # Adapted from

6 # https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py

@@ -65,7 +64,7 @@ class LinearScalingRotaryEmbedding(RotaryEmbedding):

64 max_position_embeddings: int,

65 base: float,

66 is_neox_style: bool,

68 - scaling_factors: Union[list[float], float],

67 + scaling_factors: list[float] | float,

68 dtype: torch.dtype,

69 ) -> None:

70 if isinstance(scaling_factors, float):

vllm/model_executor/layers/rotary_embedding/llama4_vision_rope.py

+4 -5

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import math

5 -from typing import Optional

5

6 import torch

7

@@ -56,8 +55,8 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):

55 def forward_native( # type: ignore[override]

56 self,

57 query: torch.Tensor,

59 - key: Optional[torch.Tensor] = None,

60 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

58 + key: torch.Tensor | None = None,

59 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

60 assert key is not None

61 # self.cos_sin_cache here is complex tensor so we cannot cast into

62 # query's dtype directly with self._match_cos_sin_cache_dtype

@@ -76,6 +75,6 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):

75 def forward_cuda( # type: ignore[override]

76 self,

77 query: torch.Tensor,

79 - key: Optional[torch.Tensor] = None,

80 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

78 + key: torch.Tensor | None = None,

79 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

80 return self.forward_native(query, key)

vllm/model_executor/layers/rotary_embedding/mrope.py

+19 -20

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional, Union

4

5 import numpy as np

6 import torch

@@ -212,11 +211,11 @@ class MRotaryEmbedding(RotaryEmbedding):

211 base: float,

212 is_neox_style: bool,

213 dtype: torch.dtype,

215 - mrope_section: Optional[list[int]] = None,

214 + mrope_section: list[int] | None = None,

215 mrope_interleaved: bool = False,

216 # YaRN parameters.

217 *,

219 - scaling_factor: Optional[float] = None,

218 + scaling_factor: float | None = None,

219 extrapolation_factor: float = 1,

220 attn_factor: float = 1,

221 beta_fast: int = 32,

@@ -265,9 +264,9 @@ class MRotaryEmbedding(RotaryEmbedding):

264 self,

265 positions: torch.Tensor,

266 query: torch.Tensor,

268 - key: Optional[torch.Tensor] = None,

269 - offsets: Optional[torch.Tensor] = None,

270 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

267 + key: torch.Tensor | None = None,

268 + offsets: torch.Tensor | None = None,

269 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

270 """PyTorch-native implementation equivalent to forward().

271

272 Args:

@@ -318,9 +317,9 @@ class MRotaryEmbedding(RotaryEmbedding):

317 self,

318 positions: torch.Tensor,

319 query: torch.Tensor,

321 - key: Optional[torch.Tensor] = None,

322 - offsets: Optional[torch.Tensor] = None,

323 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

320 + key: torch.Tensor | None = None,

321 + offsets: torch.Tensor | None = None,

322 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

323 assert positions.ndim == 1 or positions.ndim == 2

324 assert key is not None

325

@@ -363,18 +362,18 @@ class MRotaryEmbedding(RotaryEmbedding):

362 self,

363 positions: torch.Tensor,

364 query: torch.Tensor,

366 - key: Optional[torch.Tensor] = None,

367 - offsets: Optional[torch.Tensor] = None,

368 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

365 + key: torch.Tensor | None = None,

366 + offsets: torch.Tensor | None = None,

367 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

368 return self.forward_native(positions, query, key, offsets)

369

370 def forward_cpu(

371 self,

372 positions: torch.Tensor,

373 query: torch.Tensor,

375 - key: Optional[torch.Tensor] = None,

376 - offsets: Optional[torch.Tensor] = None,

377 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

374 + key: torch.Tensor | None = None,

375 + offsets: torch.Tensor | None = None,

376 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

377 return self.forward_native(positions, query, key, offsets)

378

379 @classmethod

@@ -382,12 +381,12 @@ class MRotaryEmbedding(RotaryEmbedding):

381 cls,

382 input_tokens: list[int],

383 hf_config: PretrainedConfig,

385 - image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

386 - video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

387 - second_per_grid_ts: Optional[list[float]],

384 + image_grid_thw: list[list[int]] | torch.Tensor | None,

385 + video_grid_thw: list[list[int]] | torch.Tensor | None,

386 + second_per_grid_ts: list[float] | None,

387 context_len: int = 0,

389 - seq_len: Optional[int] = None,

390 - audio_feature_lengths: Optional[torch.Tensor] = None,

388 + seq_len: int | None = None,

389 + audio_feature_lengths: torch.Tensor | None = None,

390 use_audio_in_video: bool = False,

391 ) -> tuple[list[list[int]], int]:

392 """Get mrope input positions and delta value."""

vllm/model_executor/layers/rotary_embedding/ntk_scaling_rope.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6

@@ -21,7 +20,7 @@ class NTKScalingRotaryEmbedding(RotaryEmbedding):

20 is_neox_style: bool,

21 scaling_factor: float,

22 dtype: torch.dtype,

24 - mixed_b: Optional[float] = None,

23 + mixed_b: float | None = None,

24 ) -> None:

25 self.scaling_factor = scaling_factor

26 self.mixed_b = mixed_b

vllm/model_executor/layers/rotary_embedding/phi3_long_rope_scaled_rope.py

+5 -6

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 -from typing import Optional

4

5 import torch

6 import torch.nn as nn

@@ -26,8 +25,8 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):

25 dtype: torch.dtype,

26 short_factor: list[float],

27 long_factor: list[float],

29 - short_mscale: Optional[float] = None,

30 - long_mscale: Optional[float] = None,

28 + short_mscale: float | None = None,

29 + long_mscale: float | None = None,

30 ):

31 super().__init__()

32

@@ -106,9 +105,9 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):

105 self,

106 positions: torch.Tensor,

107 query: torch.Tensor,

109 - key: Optional[torch.Tensor] = None,

110 - offsets: Optional[torch.Tensor] = None,

111 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

108 + key: torch.Tensor | None = None,

109 + offsets: torch.Tensor | None = None,

110 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

111 assert key is not None

112 query = query.view(*query.shape[:-1], -1, self.head_size)

113 key = key.view(*key.shape[:-1], -1, self.head_size)

vllm/model_executor/layers/utils.py

+6 -6

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Utility methods for model layers."""

4

5 -from typing import Callable, Optional

5 +from collections.abc import Callable

6

7 import torch

8

@@ -95,13 +95,13 @@ def default_unquantized_gemm(

95 layer: torch.nn.Module,

96 x: torch.Tensor,

97 weight: torch.Tensor,

98 - bias: Optional[torch.Tensor] = None,

98 + bias: torch.Tensor | None = None,

99 ):

100 return torch.nn.functional.linear(x, weight, bias)

101

102

103 def rocm_unquantized_gemm_impl(

104 - x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None

104 + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None

105 ) -> torch.Tensor:

106 from vllm.platforms.rocm import on_gfx9

107

@@ -131,7 +131,7 @@ def rocm_unquantized_gemm_impl(

131

132

133 def rocm_unquantized_gemm_impl_fake(

134 - x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None

134 + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None

135 ) -> torch.Tensor:

136 return x.new_empty((*x.shape[:-1], weight.shape[0]))

137

@@ -140,7 +140,7 @@ def rocm_unquantized_gemm(

140 layer: torch.nn.Module,

141 x: torch.Tensor,

142 weight: torch.Tensor,

143 - bias: Optional[torch.Tensor] = None,

143 + bias: torch.Tensor | None = None,

144 ) -> torch.Tensor:

145 return torch.ops.vllm.rocm_unquantized_gemm_impl(x, weight, bias)

146

@@ -197,7 +197,7 @@ def cpu_unquantized_gemm(

197 layer: torch.nn.Module,

198 x: torch.Tensor,

199 weight: torch.Tensor,

200 - bias: Optional[torch.Tensor] = None,

200 + bias: torch.Tensor | None = None,

201 ):

202 return layer.cpu_linear(x, weight, bias)

203

vllm/model_executor/layers/vocab_parallel_embedding.py

+8 -9

@@ -3,7 +3,6 @@

3

4 from collections.abc import Sequence

5 from dataclasses import dataclass

6 -from typing import Optional

6

7 import torch

8 import torch.nn.functional as F

@@ -65,7 +64,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):

64 self,

65 layer: torch.nn.Module,

66 x: torch.Tensor,

68 - bias: Optional[torch.Tensor] = None,

67 + bias: torch.Tensor | None = None,

68 ) -> torch.Tensor:

69 return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)

70

@@ -226,10 +225,10 @@ class VocabParallelEmbedding(CustomOp):

225 self,

226 num_embeddings: int,

227 embedding_dim: int,

229 - params_dtype: Optional[torch.dtype] = None,

230 - org_num_embeddings: Optional[int] = None,

228 + params_dtype: torch.dtype | None = None,

229 + org_num_embeddings: int | None = None,

230 padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,

232 - quant_config: Optional[QuantizationConfig] = None,

231 + quant_config: QuantizationConfig | None = None,

232 prefix: str = "",

233 ):

234 super().__init__()

@@ -347,7 +346,7 @@ class VocabParallelEmbedding(CustomOp):

346 added_vocab_end_index,

347 )

348

350 - def get_sharded_to_full_mapping(self) -> Optional[list[int]]:

349 + def get_sharded_to_full_mapping(self) -> list[int] | None:

350 """Get a mapping that can be used to reindex the gathered

351 logits for sampling.

352

@@ -515,10 +514,10 @@ class ParallelLMHead(VocabParallelEmbedding):

514 num_embeddings: int,

515 embedding_dim: int,

516 bias: bool = False,

518 - params_dtype: Optional[torch.dtype] = None,

519 - org_num_embeddings: Optional[int] = None,

517 + params_dtype: torch.dtype | None = None,

518 + org_num_embeddings: int | None = None,

519 padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,

521 - quant_config: Optional[QuantizationConfig] = None,

520 + quant_config: QuantizationConfig | None = None,

521 prefix: str = "",

522 ):

523 super().__init__(

vllm/model_executor/model_loader/__init__.py

+2 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Literal, Optional

4 +from typing import Literal

5

6 from torch import nn

7

@@ -122,7 +122,7 @@ def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:

122

123

124 def get_model(

125 - *, vllm_config: VllmConfig, model_config: Optional[ModelConfig] = None

125 + *, vllm_config: VllmConfig, model_config: ModelConfig | None = None

126 ) -> nn.Module:

127 loader = get_model_loader(vllm_config.load_config)

128 if model_config is None:

vllm/model_executor/model_loader/bitsandbytes_loader.py

+5 -5

@@ -6,8 +6,8 @@ import glob

6 import itertools

7 import math

8 import os

9 -from collections.abc import Generator

10 -from typing import Any, Callable, Optional

9 +from collections.abc import Callable, Generator

10 +from typing import Any

11

12 import numpy as np

13 import torch

@@ -88,7 +88,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

88 self,

89 model_name_or_path: str,

90 allowed_patterns: list[str],

91 - revision: Optional[str] = None,

91 + revision: str | None = None,

92 ) -> tuple[str, list[str], str]:

93 """Retrieve weight files. Download the files if necessary.

94

@@ -122,7 +122,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

122 raise RuntimeError(f"No model weights found in: `{model_name_or_path}`")

123

124 def _prepare_weights(

125 - self, model_name_or_path: str, revision: Optional[str]

125 + self, model_name_or_path: str, revision: str | None

126 ) -> tuple[list[str], bool]:

127 """Prepare weight files for the model."""

128

@@ -196,7 +196,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

196 def _get_quantized_weights_iterator(

197 self,

198 model_name_or_path: str,

199 - revision: Optional[str],

199 + revision: str | None,

200 ) -> tuple[Generator[tuple[str, torch.Tensor], None, None], dict[str, Any]]:

201 """Get an iterator to the model weights with bitsandbytes quantization,

202 as well as the quantization state dictionary."""

vllm/model_executor/model_loader/default_loader.py

+5 -5

@@ -5,7 +5,7 @@ import glob

5 import os

6 import time

7 from collections.abc import Generator, Iterable

8 -from typing import Optional, cast

8 +from typing import cast

9

10 import torch

11 from torch import nn

@@ -47,7 +47,7 @@ class DefaultModelLoader(BaseModelLoader):

47 model_or_path: str

48 """The model ID or path."""

49

50 - revision: Optional[str]

50 + revision: str | None

51 """The optional model revision."""

52

53 prefix: str = ""

@@ -56,7 +56,7 @@ class DefaultModelLoader(BaseModelLoader):

56 fall_back_to_pt: bool = True

57 """Whether .pt weights can be used."""

58

59 - allow_patterns_overrides: Optional[list[str]] = None

59 + allow_patterns_overrides: list[str] | None = None

60 """If defined, weights will load exclusively using these patterns."""

61

62 counter_before_loading_weights: float = 0.0

@@ -79,9 +79,9 @@ class DefaultModelLoader(BaseModelLoader):

79 def _prepare_weights(

80 self,

81 model_name_or_path: str,

82 - revision: Optional[str],

82 + revision: str | None,

83 fall_back_to_pt: bool,

84 - allow_patterns_overrides: Optional[list[str]],

84 + allow_patterns_overrides: list[str] | None,

85 ) -> tuple[str, list[str], bool]:

86 """Prepare weights for the model.

87

vllm/model_executor/model_loader/runai_streamer_loader.py

+1 -2

@@ -3,7 +3,6 @@

3 # ruff: noqa: SIM117

4 import os

5 from collections.abc import Generator

6 -from typing import Optional

6

7 import torch

8 from torch import nn

@@ -51,7 +50,7 @@ class RunaiModelStreamerLoader(BaseModelLoader):

50 os.environ["RUNAI_STREAMER_S3_ENDPOINT"] = aws_endpoint_url

51

52 def _prepare_weights(

54 - self, model_name_or_path: str, revision: Optional[str]

53 + self, model_name_or_path: str, revision: str | None

54 ) -> list[str]:

55 """Prepare weights for the model.

56

vllm/model_executor/model_loader/sharded_state_loader.py

+4 -4

@@ -5,7 +5,7 @@ import collections

5 import glob

6 import os

7 from collections.abc import Generator

8 -from typing import Any, Optional

8 +from typing import Any

9

10 import torch

11 from torch import nn

@@ -89,7 +89,7 @@ class ShardedStateLoader(BaseModelLoader):

89 result[k] = t

90 return result

91

92 - def _prepare_weights(self, model_name_or_path: str, revision: Optional[str]):

92 + def _prepare_weights(self, model_name_or_path: str, revision: str | None):

93 if is_s3(model_name_or_path) or os.path.isdir(model_name_or_path):

94 return model_name_or_path

95 else:

@@ -171,8 +171,8 @@ class ShardedStateLoader(BaseModelLoader):

171 def save_model(

172 model: torch.nn.Module,

173 path: str,

174 - pattern: Optional[str] = None,

175 - max_size: Optional[int] = None,

174 + pattern: str | None = None,

175 + max_size: int | None = None,

176 ) -> None:

177 from safetensors.torch import save_file

178

vllm/model_executor/model_loader/tensorizer.py

+23 -25

@@ -12,7 +12,7 @@ import threading

12 import time

13 from collections.abc import Generator, MutableMapping

14 from dataclasses import asdict, dataclass, field, fields

15 -from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union

15 +from typing import TYPE_CHECKING, Any, ClassVar, Optional

16

17 import regex as re

18 import torch

@@ -67,7 +67,7 @@ __all__ = [

67 logger = init_logger(__name__)

68

69

70 -def is_valid_deserialization_uri(uri: Optional[str]) -> bool:

70 +def is_valid_deserialization_uri(uri: str | None) -> bool:

71 if uri:

72 scheme = uri.lower().split("://")[0]

73 return scheme in {"s3", "http", "https"} or os.path.exists(uri)

@@ -156,25 +156,23 @@ class _NoInitOrTensorImpl:

156

157 @dataclass

158 class TensorizerConfig(MutableMapping):

159 - tensorizer_uri: Optional[str] = None

160 - tensorizer_dir: Optional[str] = None

161 - vllm_tensorized: Optional[bool] = None

162 - verify_hash: Optional[bool] = None

163 - num_readers: Optional[int] = None

164 - encryption_keyfile: Optional[str] = None

165 - s3_access_key_id: Optional[str] = None

166 - s3_secret_access_key: Optional[str] = None

167 - s3_endpoint: Optional[str] = None

168 - lora_dir: Optional[str] = None

169 - stream_kwargs: Optional[dict[str, Any]] = None

170 - serialization_kwargs: Optional[dict[str, Any]] = None

171 - deserialization_kwargs: Optional[dict[str, Any]] = None

172 - _extra_serialization_attrs: Optional[dict[str, Any]] = field(

173 - init=False, default=None

174 - )

175 - model_class: Optional[type[torch.nn.Module]] = field(init=False, default=None)

176 - hf_config: Optional[PretrainedConfig] = field(init=False, default=None)

177 - dtype: Optional[Union[str, torch.dtype]] = field(init=False, default=None)

159 + tensorizer_uri: str | None = None

160 + tensorizer_dir: str | None = None

161 + vllm_tensorized: bool | None = None

162 + verify_hash: bool | None = None

163 + num_readers: int | None = None

164 + encryption_keyfile: str | None = None

165 + s3_access_key_id: str | None = None

166 + s3_secret_access_key: str | None = None

167 + s3_endpoint: str | None = None

168 + lora_dir: str | None = None

169 + stream_kwargs: dict[str, Any] | None = None

170 + serialization_kwargs: dict[str, Any] | None = None

171 + deserialization_kwargs: dict[str, Any] | None = None

172 + _extra_serialization_attrs: dict[str, Any] | None = field(init=False, default=None)

173 + model_class: type[torch.nn.Module] | None = field(init=False, default=None)

174 + hf_config: PretrainedConfig | None = field(init=False, default=None)

175 + dtype: str | torch.dtype | None = field(init=False, default=None)

176 _is_sharded: bool = field(init=False, default=False)

177 _fields: ClassVar[tuple[str, ...]]

178 _keys: ClassVar[frozenset[str]]

@@ -362,9 +360,9 @@ TensorizerConfig._keys = frozenset(TensorizerConfig._fields)

360

361 @dataclass

362 class TensorizerArgs:

365 - tensorizer_uri: Optional[str] = None

366 - tensorizer_dir: Optional[str] = None

367 - encryption_keyfile: Optional[str] = None

363 + tensorizer_uri: str | None = None

364 + tensorizer_dir: str | None = None

365 + encryption_keyfile: str | None = None

366

367 def __init__(self, tensorizer_config: TensorizerConfig):

368 for k, v in tensorizer_config.items():

@@ -621,7 +619,7 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:

619

620

621 def serialize_extra_artifacts(

624 - tensorizer_args: TensorizerArgs, served_model_name: Union[str, list[str], None]

622 + tensorizer_args: TensorizerArgs, served_model_name: str | list[str] | None

623 ) -> None:

624 if not isinstance(served_model_name, str):

625 raise ValueError(

vllm/model_executor/model_loader/tensorizer_loader.py

+1 -2

@@ -3,7 +3,6 @@

3 # ruff: noqa: SIM117

4 import copy

5 from collections.abc import Generator

6 -from typing import Union

6

7 import torch

8 from torch import nn

@@ -140,7 +139,7 @@ class TensorizerLoader(BaseModelLoader):

139 @staticmethod

140 def save_model(

141 model: torch.nn.Module,

143 - tensorizer_config: Union[TensorizerConfig, dict],

142 + tensorizer_config: TensorizerConfig | dict,

143 model_config: ModelConfig,

144 ) -> None:

145 if isinstance(tensorizer_config, dict):

vllm/model_executor/model_loader/tpu.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import time

4 -from typing import Optional

4

5 import torch

6 import torch.nn as nn

@@ -30,7 +29,7 @@ class TPUModelLoader(DefaultModelLoader):

29 self,

30 vllm_config: VllmConfig,

31 model_config: ModelConfig,

33 - mesh: Optional[xs.Mesh] = None,

32 + mesh: xs.Mesh | None = None,

33 ) -> nn.Module:

34 # Initialize model and load weights on CPU. Then, during SPMD partition,

35 # weights are sharded and transferred to TPUs.

@@ -90,7 +89,7 @@ class TPUModelLoader(DefaultModelLoader):

89 )

90 return model

91

93 - def _check_model_is_loaded(self, mesh: Optional[xs.Mesh], model: nn.Module) -> None:

92 + def _check_model_is_loaded(self, mesh: xs.Mesh | None, model: nn.Module) -> None:

93 """

94 Ensure the model is properly loaded.

95 1. All model parameters and buffers are on XLA device.

vllm/model_executor/model_loader/utils.py

+3 -4

@@ -7,7 +7,6 @@ import inspect

7 import warnings

8 from contextlib import contextmanager

9 from dataclasses import dataclass, field

10 -from typing import Optional

10

11 import torch

12 from torch import nn

@@ -46,8 +45,8 @@ def initialize_model(

45 vllm_config: VllmConfig,

46 *,

47 prefix: str = "",

49 - model_class: Optional[type[nn.Module]] = None,

50 - model_config: Optional[ModelConfig] = None,

48 + model_class: type[nn.Module] | None = None,

49 + model_config: ModelConfig | None = None,

50 ) -> nn.Module:

51 """Initialize a model with the given configurations."""

52 if model_config is None:

@@ -268,7 +267,7 @@ class ParamMapping:

267 index,

268 )

269

271 - def get_sub_modules(self, module_name: str) -> Optional[tuple[str, list[str]]]:

270 + def get_sub_modules(self, module_name: str) -> tuple[str, list[str]] | None:

271 for key, value in self.packed_mapping.items():

272 if module_name.endswith(key):

273 return key, value

vllm/model_executor/model_loader/weight_utils.py

+18 -18

@@ -11,10 +11,10 @@ import os

11 import tempfile

12 import time

13 from collections import defaultdict

14 -from collections.abc import Generator

14 +from collections.abc import Callable, Generator

15 from contextlib import contextmanager

16 from pathlib import Path

17 -from typing import IO, Any, Callable, Optional, Union

17 +from typing import IO, Any

18

19 import filelock

20 import huggingface_hub.constants

@@ -85,7 +85,7 @@ class DisabledTqdm(tqdm):

85 super().__init__(*args, **kwargs, disable=True)

86

87

88 -def get_lock(model_name_or_path: Union[str, Path], cache_dir: Optional[str] = None):

88 +def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):

89 lock_dir = cache_dir or temp_dir

90 model_name_or_path = str(model_name_or_path)

91 os.makedirs(os.path.dirname(lock_dir), exist_ok=True)

@@ -100,7 +100,7 @@ def get_lock(model_name_or_path: Union[str, Path], cache_dir: Optional[str] = No

100

101 @contextmanager

102 def atomic_writer(

103 - filepath: Union[str, Path], mode: str = "w", encoding: Optional[str] = None

103 + filepath: str | Path, mode: str = "w", encoding: str | None = None

104 ) -> Generator[IO]:

105 """

106 Context manager that provides an atomic file writing routine.

@@ -143,11 +143,11 @@ def atomic_writer(

143

144 def maybe_download_from_modelscope(

145 model: str,

146 - revision: Optional[str] = None,

147 - download_dir: Optional[str] = None,

148 - ignore_patterns: Optional[Union[str, list[str]]] = None,

149 - allow_patterns: Optional[Union[list[str], str]] = None,

150 -) -> Optional[str]:

146 + revision: str | None = None,

147 + download_dir: str | None = None,

148 + ignore_patterns: str | list[str] | None = None,

149 + allow_patterns: list[str] | str | None = None,

150 +) -> str | None:

151 """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.

152

153 Returns the path to the downloaded model, or None if the model is not

@@ -370,10 +370,10 @@ def get_sparse_attention_config(

370

371 def download_weights_from_hf(

372 model_name_or_path: str,

373 - cache_dir: Optional[str],

373 + cache_dir: str | None,

374 allow_patterns: list[str],

375 - revision: Optional[str] = None,

376 - ignore_patterns: Optional[Union[str, list[str]]] = None,

375 + revision: str | None = None,

376 + ignore_patterns: str | list[str] | None = None,

377 ) -> str:

378 """Download model weights from Hugging Face Hub.

379

@@ -448,8 +448,8 @@ def download_weights_from_hf(

448 def download_safetensors_index_file_from_hf(

449 model_name_or_path: str,

450 index_file: str,

451 - cache_dir: Optional[str],

452 - revision: Optional[str] = None,

451 + cache_dir: str | None,

452 + revision: str | None = None,

453 ) -> None:

454 """Download hf safetensors index file from Hugging Face Hub.

455

@@ -540,7 +540,7 @@ def enable_tqdm(use_tqdm_on_load: bool):

540

541 def np_cache_weights_iterator(

542 model_name_or_path: str,

543 - cache_dir: Optional[str],

543 + cache_dir: str | None,

544 hf_folder: str,

545 hf_weights_files: list[str],

546 use_tqdm_on_load: bool,

@@ -746,7 +746,7 @@ def fastsafetensors_weights_iterator(

746 def pt_weights_iterator(

747 hf_weights_files: list[str],

748 use_tqdm_on_load: bool,

749 - pt_load_map_location: Union[str, dict[str, str]] = "cpu",

749 + pt_load_map_location: str | dict[str, str] = "cpu",

750 ) -> Generator[tuple[str, torch.Tensor], None, None]:

751 """Iterate over the weights in the model bin/pt files."""

752 for bin_file in tqdm(

@@ -765,7 +765,7 @@ def pt_weights_iterator(

765 def multi_thread_pt_weights_iterator(

766 hf_weights_files: list[str],

767 use_tqdm_on_load: bool,

768 - pt_load_map_location: Union[str, dict[str, str]] = "cpu",

768 + pt_load_map_location: str | dict[str, str] = "cpu",

769 max_workers: int = 4,

770 ) -> Generator[tuple[str, torch.Tensor], None, None]:

771 """Multi-Thread iterate over the weights in the model bin/pt files."""

@@ -985,7 +985,7 @@ def initialize_dummy_weights(

985 param.uniform_(low, high, generator=generator)

986

987

988 -def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:

988 +def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:

989 """Remap the name of FP8 k/v_scale parameters.

990

991 This function handles the remapping of FP8 k/v_scale parameter names.

vllm/model_executor/models/adapters.py

+4 -4

@@ -4,7 +4,7 @@

4 import ast

5 import inspect

6 from collections.abc import Iterable

7 -from typing import TYPE_CHECKING, Any, Optional, TypeVar, cast

7 +from typing import TYPE_CHECKING, Any, TypeVar, cast

8

9 import torch

10 import torch.nn as nn

@@ -32,7 +32,7 @@ _GENERATE_SUFFIXES = [

32 ]

33

34

35 -def _load_st_projector(model_config: "ModelConfig") -> Optional[nn.Module]:

35 +def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None:

36 """Load Sentence-Transformers Dense projection layers."""

37

38 try:

@@ -357,8 +357,8 @@ def as_seq_cls_model(cls: _T) -> _T:

357 self,

358 input_ids: torch.Tensor,

359 positions: torch.Tensor,

360 - intermediate_tensors: Optional[IntermediateTensors] = None,

361 - inputs_embeds: Optional[torch.Tensor] = None,

360 + intermediate_tensors: IntermediateTensors | None = None,

361 + inputs_embeds: torch.Tensor | None = None,

362 ) -> torch.Tensor:

363 return super().forward(

364 input_ids, positions, intermediate_tensors, inputs_embeds

vllm/model_executor/models/aimv2.py

+2 -3

@@ -4,7 +4,6 @@

4 # A modified implementation of the AIMv2 Transformer

5 # inserted here also the image tokenizer used by Ovis2

6 from collections.abc import Iterable

7 -from typing import Optional

7

8 import torch

9 import torch.nn as nn

@@ -165,7 +164,7 @@ class AIMv2Transformer(nn.Module):

164 config: AIMv2Config,

165 quant_config: QuantizationConfig,

166 *,

168 - require_post_norm: Optional[bool] = None,

167 + require_post_norm: bool | None = None,

168 prefix: str = "",

169 ):

170 super().__init__()

@@ -196,7 +195,7 @@ class AIMv2Model(torch.nn.Module):

195 config: AIMv2Config,

196 quant_config: QuantizationConfig,

197 *,

199 - require_post_norm: Optional[bool] = None,

198 + require_post_norm: bool | None = None,

199 prefix: str = "",

200 ):

201 super().__init__()

vllm/model_executor/models/apertus.py

+18 -20

@@ -27,7 +27,7 @@

27

28 from collections.abc import Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import torch

33 from torch import nn

@@ -77,7 +77,7 @@ class ApertusMLP(nn.Module):

77 hidden_size: int,

78 intermediate_size: int,

79 hidden_act: str,

80 - quant_config: Optional[QuantizationConfig] = None,

80 + quant_config: QuantizationConfig | None = None,

81 bias: bool = False,

82 prefix: str = "",

83 reduce_results: bool = True,

@@ -120,12 +120,12 @@ class ApertusAttention(nn.Module):

120 num_heads: int,

121 num_kv_heads: int,

122 rope_theta: float = 10000,

123 - rope_scaling: Optional[dict[str, Any]] = None,

123 + rope_scaling: dict[str, Any] | None = None,

124 max_position_embeddings: int = 8192,

125 - quant_config: Optional[QuantizationConfig] = None,

125 + quant_config: QuantizationConfig | None = None,

126 bias: bool = False,

127 bias_o_proj: bool = False,

128 - cache_config: Optional[CacheConfig] = None,

128 + cache_config: CacheConfig | None = None,

129 prefix: str = "",

130 attn_type: str = AttentionType.DECODER,

131 ) -> None:

@@ -225,8 +225,8 @@ class ApertusAttention(nn.Module):

225 def _init_rotary_emb(

226 self,

227 config: ApertusConfig,

228 - rope_scaling: Optional[dict[str, Any]],

229 - quant_config: Optional[QuantizationConfig],

228 + rope_scaling: dict[str, Any] | None,

229 + quant_config: QuantizationConfig | None,

230 ) -> None:

231 is_neox_style = True

232 is_gguf = quant_config and quant_config.get_name() == "gguf"

@@ -248,8 +248,8 @@ class ApertusDecoderLayer(nn.Module):

248 def __init__(

249 self,

250 config: ApertusConfig,

251 - cache_config: Optional[CacheConfig] = None,

252 - quant_config: Optional[QuantizationConfig] = None,

251 + cache_config: CacheConfig | None = None,

252 + quant_config: QuantizationConfig | None = None,

253 prefix: str = "",

254 ) -> None:

255 super().__init__()

@@ -316,7 +316,7 @@ class ApertusDecoderLayer(nn.Module):

316 self,

317 positions: torch.Tensor,

318 hidden_states: torch.Tensor,

319 - residual: Optional[torch.Tensor],

319 + residual: torch.Tensor | None,

320 ) -> tuple[torch.Tensor, torch.Tensor]:

321 # Self Attention

322 if residual is None:

@@ -394,13 +394,11 @@ class ApertusModel(nn.Module):

394

395 def forward(

396 self,

397 - input_ids: Optional[torch.Tensor],

397 + input_ids: torch.Tensor | None,

398 positions: torch.Tensor,

399 - intermediate_tensors: Optional[IntermediateTensors],

400 - inputs_embeds: Optional[torch.Tensor] = None,

401 - ) -> Union[

402 - torch.Tensor, IntermediateTensors, tuple[torch.Tensor, list[torch.Tensor]]

403 - ]:

399 + intermediate_tensors: IntermediateTensors | None,

400 + inputs_embeds: torch.Tensor | None = None,

401 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:

402 if get_pp_group().is_first_rank:

403 if inputs_embeds is not None:

404 hidden_states = inputs_embeds

@@ -586,9 +584,9 @@ class ApertusForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

584 self,

585 input_ids: torch.Tensor,

586 positions: torch.Tensor,

589 - intermediate_tensors: Optional[IntermediateTensors] = None,

590 - inputs_embeds: Optional[torch.Tensor] = None,

591 - ) -> Union[torch.Tensor, IntermediateTensors]:

587 + intermediate_tensors: IntermediateTensors | None = None,

588 + inputs_embeds: torch.Tensor | None = None,

589 + ) -> torch.Tensor | IntermediateTensors:

590 model_output = self.model(

591 input_ids, positions, intermediate_tensors, inputs_embeds

592 )

@@ -597,7 +595,7 @@ class ApertusForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

595 def compute_logits(

596 self,

597 hidden_states: torch.Tensor,

600 - ) -> Optional[torch.Tensor]:

598 + ) -> torch.Tensor | None:

599 logits = self.logits_processor(self.lm_head, hidden_states)

600 return logits

601

vllm/model_executor/models/arcee.py

+13 -15

@@ -10,7 +10,7 @@

10

11 from collections.abc import Iterable

12 from itertools import islice

13 -from typing import Any, Optional, Union

13 +from typing import Any

14

15 import torch

16 from torch import nn

@@ -52,7 +52,7 @@ class ArceeMLP(nn.Module):

52 hidden_size: int,

53 intermediate_size: int,

54 hidden_act: str,

55 - quant_config: Optional[Any] = None,

55 + quant_config: Any | None = None,

56 bias: bool = False,

57 prefix: str = "",

58 reduce_results: bool = True,

@@ -98,8 +98,8 @@ class ArceeDecoderLayer(nn.Module):

98 def __init__(

99 self,

100 config: LlamaConfig,

101 - cache_config: Optional[Any] = None,

102 - quant_config: Optional[Any] = None,

101 + cache_config: Any | None = None,

102 + quant_config: Any | None = None,

103 prefix: str = "",

104 ) -> None:

105 super().__init__()

@@ -165,7 +165,7 @@ class ArceeDecoderLayer(nn.Module):

165 self,

166 positions: torch.Tensor,

167 hidden_states: torch.Tensor,

168 - residual: Optional[torch.Tensor],

168 + residual: torch.Tensor | None,

169 ) -> tuple[torch.Tensor, torch.Tensor]:

170 # Self-Attention block

171 if residual is None:

@@ -247,13 +247,11 @@ class ArceeModel(nn.Module):

247

248 def forward(

249 self,

250 - input_ids: Optional[torch.Tensor],

250 + input_ids: torch.Tensor | None,

251 positions: torch.Tensor,

252 - intermediate_tensors: Optional[IntermediateTensors],

253 - inputs_embeds: Optional[torch.Tensor] = None,

254 - ) -> Union[

255 - torch.Tensor, IntermediateTensors, tuple[torch.Tensor, list[torch.Tensor]]

256 - ]:

252 + intermediate_tensors: IntermediateTensors | None,

253 + inputs_embeds: torch.Tensor | None = None,

254 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:

255 # Embedding lookup (on first pipeline rank)

256 if get_pp_group().is_first_rank:

257 hidden_states = (

@@ -415,9 +413,9 @@ class ArceeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

413 self,

414 input_ids: torch.Tensor,

415 positions: torch.Tensor,

418 - intermediate_tensors: Optional[IntermediateTensors] = None,

419 - inputs_embeds: Optional[torch.Tensor] = None,

420 - ) -> Union[torch.Tensor, IntermediateTensors]:

416 + intermediate_tensors: IntermediateTensors | None = None,

417 + inputs_embeds: torch.Tensor | None = None,

418 + ) -> torch.Tensor | IntermediateTensors:

419 model_output = self.model(

420 input_ids=input_ids,

421 positions=positions,

@@ -426,7 +424,7 @@ class ArceeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

424 )

425 return model_output

426

429 - def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:

427 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

428 # Compute final logits from hidden states (last pipeline rank only)

429 logits = self.logits_processor(self.lm_head, hidden_states)

430 return logits

vllm/model_executor/models/arctic.py

+15 -16

@@ -4,7 +4,6 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import Optional, Union

7

8 import torch

9 from torch import nn

@@ -63,7 +62,7 @@ class ArcticMLP(nn.Module):

62 config: ArcticConfig,

63 expert_id: int = -1,

64 is_residual_mlp: bool = False,

66 - quant_config: Optional[QuantizationConfig] = None,

65 + quant_config: QuantizationConfig | None = None,

66 reduce_results: bool = True,

67 prefix: str = "",

68 ):

@@ -107,9 +106,9 @@ class ArcticMoE(nn.Module):

106 def __init__(

107 self,

108 config: ArcticConfig,

110 - tp_size: Optional[int] = None,

111 - params_dtype: Optional[torch.dtype] = None,

112 - quant_config: Optional[QuantizationConfig] = None,

109 + tp_size: int | None = None,

110 + params_dtype: torch.dtype | None = None,

111 + quant_config: QuantizationConfig | None = None,

112 reduce_results: bool = True,

113 prefix: str = "",

114 ):

@@ -265,8 +264,8 @@ class ArcticAttention(nn.Module):

264 def __init__(

265 self,

266 config: ArcticConfig,

268 - cache_config: Optional[CacheConfig] = None,

269 - quant_config: Optional[QuantizationConfig] = None,

267 + cache_config: CacheConfig | None = None,

268 + quant_config: QuantizationConfig | None = None,

269 prefix: str = "",

270 ):

271 super().__init__()

@@ -342,8 +341,8 @@ class ArcticDecoderLayer(nn.Module):

341 def __init__(

342 self,

343 config: ArcticConfig,

345 - cache_config: Optional[CacheConfig] = None,

346 - quant_config: Optional[QuantizationConfig] = None,

344 + cache_config: CacheConfig | None = None,

345 + quant_config: QuantizationConfig | None = None,

346 prefix: str = "",

347 ) -> None:

348 super().__init__()

@@ -443,9 +442,9 @@ class ArcticModel(nn.Module):

442 self,

443 input_ids: torch.Tensor,

444 positions: torch.Tensor,

446 - intermediate_tensors: Optional[IntermediateTensors],

447 - inputs_embeds: Optional[torch.Tensor] = None,

448 - ) -> Union[torch.Tensor, IntermediateTensors]:

445 + intermediate_tensors: IntermediateTensors | None,

446 + inputs_embeds: torch.Tensor | None = None,

447 + ) -> torch.Tensor | IntermediateTensors:

448 if get_pp_group().is_first_rank:

449 if inputs_embeds is not None:

450 hidden_states = inputs_embeds

@@ -499,9 +498,9 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):

498 self,

499 input_ids: torch.Tensor,

500 positions: torch.Tensor,

502 - intermediate_tensors: Optional[IntermediateTensors] = None,

503 - inputs_embeds: Optional[torch.Tensor] = None,

504 - ) -> Union[torch.Tensor, IntermediateTensors]:

501 + intermediate_tensors: IntermediateTensors | None = None,

502 + inputs_embeds: torch.Tensor | None = None,

503 + ) -> torch.Tensor | IntermediateTensors:

504 hidden_states = self.model(

505 input_ids, positions, intermediate_tensors, inputs_embeds

506 )

@@ -510,7 +509,7 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):

509 def compute_logits(

510 self,

511 hidden_states: torch.Tensor,

513 - ) -> Optional[torch.Tensor]:

512 + ) -> torch.Tensor | None:

513 logits = self.logits_processor(self.lm_head, hidden_states)

514 return logits

515

vllm/model_executor/models/aria.py

+14 -14

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Mapping, Sequence

4 -from typing import Annotated, Literal, Optional, Union

4 +from typing import Annotated, Literal

5

6 import torch

7 import torch.nn as nn

@@ -71,7 +71,7 @@ class AriaImagePixelInputs(TensorSchema):

71 ]

72

73 pixel_mask: Annotated[

74 - Optional[torch.Tensor],

74 + torch.Tensor | None,

75 TensorShape("bn", "h", "w"),

76 ]

77

@@ -82,7 +82,7 @@ class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant):

82 def __init__(

83 self,

84 config: Idefics2VisionConfig,

85 - quant_config: Optional[QuantizationConfig] = None,

85 + quant_config: QuantizationConfig | None = None,

86 prefix: str = "",

87 ) -> None:

88 super().__init__(config, quant_config=quant_config, prefix=prefix)

@@ -180,7 +180,7 @@ class AriaProjector(nn.Module):

180 def forward(

181 self,

182 x: torch.Tensor,

183 - attn_mask: Optional[torch.Tensor] = None,

183 + attn_mask: torch.Tensor | None = None,

184 ) -> torch.Tensor:

185 batch_size, num_patches = x.shape[0], x.shape[1]

186

@@ -250,7 +250,7 @@ class AriaTextMoELayer(nn.Module):

250 def __init__(

251 self,

252 config: AriaTextConfig,

253 - quant_config: Optional[QuantizationConfig],

253 + quant_config: QuantizationConfig | None,

254 prefix: str = "",

255 ) -> None:

256 super().__init__()

@@ -415,7 +415,7 @@ class AriaProcessingInfo(BaseProcessingInfo):

415 def get_hf_processor(self, **kwargs: object):

416 return self.ctx.get_hf_processor(AriaProcessor, **kwargs)

417

418 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

418 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

419 return {"image": None}

420

421 def get_num_image_tokens(self) -> int:

@@ -436,7 +436,7 @@ class AriaDummyInputsBuilder(BaseDummyInputsBuilder[AriaProcessingInfo]):

436 self,

437 seq_len: int,

438 mm_counts: Mapping[str, int],

439 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

439 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

440 ) -> MultiModalDataDict:

441 vision_config = self.info.get_vision_config()

442

@@ -517,7 +517,7 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

517 )

518

519 @classmethod

520 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

520 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

521 if modality.startswith("image"):

522 return "<|fim_prefix|><|img|><|fim_suffix|>"

523

@@ -562,7 +562,7 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

562

563 def _parse_and_validate_image_input(

564 self, **kwargs: object

565 - ) -> Optional[AriaImagePixelInputs]:

565 + ) -> AriaImagePixelInputs | None:

566 pixel_values = kwargs.pop("pixel_values", None)

567 pixel_mask = kwargs.pop("pixel_mask", None)

568

@@ -577,8 +577,8 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

577

578 def _create_patch_attention_mask(

579 self,

580 - pixel_mask: Optional[torch.Tensor],

581 - ) -> Optional[torch.Tensor]:

580 + pixel_mask: torch.Tensor | None,

581 + ) -> torch.Tensor | None:

582 if pixel_mask is None:

583 return None

584

@@ -628,10 +628,10 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

628 self,

629 input_ids: torch.Tensor,

630 positions: torch.Tensor,

631 - intermediate_tensors: Optional[IntermediateTensors] = None,

632 - inputs_embeds: Optional[torch.Tensor] = None,

631 + intermediate_tensors: IntermediateTensors | None = None,

632 + inputs_embeds: torch.Tensor | None = None,

633 **kwargs: object,

634 - ) -> Union[torch.Tensor, IntermediateTensors]:

634 + ) -> torch.Tensor | IntermediateTensors:

635 if inputs_embeds is None:

636 multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

637 inputs_embeds = self.get_input_embeddings(

vllm/model_executor/models/aya_vision.py

+10 -10

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 # Adapted from https://github.com/huggingface/transformers/tree/main/src/transformers/models/aya_vision

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import Annotated, Literal, Optional, Union

5 +from typing import Annotated, Literal

6

7 import torch

8 from torch import nn

@@ -139,7 +139,7 @@ class AyaVisionProcessingInfo(BaseProcessingInfo):

139 def get_image_processor(self, **kwargs: object) -> GotOcr2ImageProcessor:

140 return self.get_hf_processor(**kwargs).image_processor

141

142 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

142 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

143 return {"image": None}

144

145 def get_image_size_with_most_features(self) -> ImageSize:

@@ -187,7 +187,7 @@ class AyaVisionDummyInputsBuilder(BaseDummyInputsBuilder[AyaVisionProcessingInfo

187 self,

188 seq_len: int,

189 mm_counts: Mapping[str, int],

190 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

190 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

191 ) -> MultiModalDataDict:

192 num_images = mm_counts.get("image", 0)

193 image_size = self.info.get_image_size_with_most_features()

@@ -331,7 +331,7 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

331 )

332

333 @classmethod

334 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

334 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

335 if modality.startswith("image"):

336 return "<image>"

337

@@ -375,7 +375,7 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

375 self,

376 vision_tower: SiglipVisionModel,

377 pixel_values: torch.Tensor,

378 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

378 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

379 return vision_tower(

380 pixel_values.to(dtype=vision_tower.dtype),

381 feature_select_strategy=self.config.vision_feature_select_strategy,

@@ -395,7 +395,7 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

395

396 def _parse_and_validate_image_input(

397 self, **kwargs: object

398 - ) -> Optional[AyaVisionImagePixelInputs]:

398 + ) -> AyaVisionImagePixelInputs | None:

399 pixel_values = kwargs.pop("pixel_values", None)

400 num_patches = kwargs.pop("num_patches", None)

401 image_embeds = kwargs.pop("image_embeds", None)

@@ -428,10 +428,10 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

428 self,

429 input_ids: torch.Tensor,

430 positions: torch.Tensor,

431 - intermediate_tensors: Optional[IntermediateTensors] = None,

432 - inputs_embeds: Optional[torch.Tensor] = None,

431 + intermediate_tensors: IntermediateTensors | None = None,

432 + inputs_embeds: torch.Tensor | None = None,

433 **kwargs: object,

434 - ) -> Union[torch.Tensor, IntermediateTensors]:

434 + ) -> torch.Tensor | IntermediateTensors:

435 if intermediate_tensors is not None:

436 inputs_embeds = None

437

@@ -446,5 +446,5 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

446 def compute_logits(

447 self,

448 hidden_states: torch.Tensor,

449 - ) -> Optional[torch.Tensor]:

449 + ) -> torch.Tensor | None:

450 return self.language_model.compute_logits(hidden_states)

vllm/model_executor/models/baichuan.py

+13 -14

@@ -24,7 +24,6 @@

24 import math

25 from collections.abc import Iterable

26 from itertools import islice

27 -from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -98,7 +97,7 @@ class BaiChuanMLP(nn.Module):

97 hidden_size: int,

98 intermediate_size: int,

99 hidden_act: str,

101 - quant_config: Optional[QuantizationConfig] = None,

100 + quant_config: QuantizationConfig | None = None,

101 ):

102 super().__init__()

103 self.gate_up_proj = MergedColumnParallelLinear(

@@ -130,8 +129,8 @@ class BaiChuanAttention(nn.Module):

129 position_embedding: str,

130 rope_theta: float = 10000,

131 max_position_embeddings: int = 8192,

133 - cache_config: Optional[CacheConfig] = None,

134 - quant_config: Optional[QuantizationConfig] = None,

132 + cache_config: CacheConfig | None = None,

133 + quant_config: QuantizationConfig | None = None,

134 prefix: str = "",

135 ):

136 super().__init__()

@@ -213,8 +212,8 @@ class BaiChuanDecoderLayer(nn.Module):

212 self,

213 config: PretrainedConfig,

214 position_embedding: str,

216 - cache_config: Optional[CacheConfig] = None,

217 - quant_config: Optional[QuantizationConfig] = None,

215 + cache_config: CacheConfig | None = None,

216 + quant_config: QuantizationConfig | None = None,

217 prefix: str = "",

218 ):

219 super().__init__()

@@ -246,7 +245,7 @@ class BaiChuanDecoderLayer(nn.Module):

245 self,

246 positions: torch.Tensor,

247 hidden_states: torch.Tensor,

249 - residual: Optional[torch.Tensor],

248 + residual: torch.Tensor | None,

249 ) -> tuple[torch.Tensor, torch.Tensor]:

250 # Self Attention

251 if residual is None:

@@ -305,9 +304,9 @@ class BaiChuanModel(nn.Module):

304 self,

305 input_ids: torch.Tensor,

306 positions: torch.Tensor,

308 - intermediate_tensors: Optional[IntermediateTensors],

309 - inputs_embeds: Optional[torch.Tensor] = None,

310 - ) -> Union[torch.Tensor, IntermediateTensors]:

307 + intermediate_tensors: IntermediateTensors | None,

308 + inputs_embeds: torch.Tensor | None = None,

309 + ) -> torch.Tensor | IntermediateTensors:

310 if get_pp_group().is_first_rank:

311 if inputs_embeds is not None:

312 hidden_states = inputs_embeds

@@ -422,9 +421,9 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant

421 self,

422 input_ids: torch.Tensor,

423 positions: torch.Tensor,

425 - intermediate_tensors: Optional[IntermediateTensors] = None,

426 - inputs_embeds: Optional[torch.Tensor] = None,

427 - ) -> Union[torch.Tensor, IntermediateTensors]:

424 + intermediate_tensors: IntermediateTensors | None = None,

425 + inputs_embeds: torch.Tensor | None = None,

426 + ) -> torch.Tensor | IntermediateTensors:

427 hidden_states = self.model(

428 input_ids, positions, intermediate_tensors, inputs_embeds

429 )

@@ -433,7 +432,7 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant

432 def compute_logits(

433 self,

434 hidden_states: torch.Tensor,

436 - ) -> Optional[torch.Tensor]:

435 + ) -> torch.Tensor | None:

436 logits = self.logits_processor(self.lm_head, hidden_states)

437 return logits

438

vllm/model_executor/models/bailing_moe.py

+16 -17

@@ -26,7 +26,6 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Optional, Union

29

30 import torch

31 import torch.nn.functional as F

@@ -75,8 +74,8 @@ class BailingAttention(nn.Module):

74 def __init__(

75 self,

76 config: PretrainedConfig,

78 - cache_config: Optional[CacheConfig] = None,

79 - quant_config: Optional[QuantizationConfig] = None,

77 + cache_config: CacheConfig | None = None,

78 + quant_config: QuantizationConfig | None = None,

79 reduce_results: bool = True,

80 prefix: str = "",

81 ):

@@ -184,8 +183,8 @@ class BailingMLP(nn.Module):

183 self,

184 intermediate_size: int,

185 config: PretrainedConfig,

187 - quant_config: Optional[QuantizationConfig] = None,

188 - reduce_results: Optional[bool] = True,

186 + quant_config: QuantizationConfig | None = None,

187 + reduce_results: bool | None = True,

188 prefix: str = "",

189 ) -> None:

190 super().__init__()

@@ -218,8 +217,8 @@ class BailingMoE(nn.Module):

217 self,

218 intermediate_size: int,

219 config: PretrainedConfig,

221 - quant_config: Optional[QuantizationConfig] = None,

222 - reduce_results: Optional[bool] = True,

220 + quant_config: QuantizationConfig | None = None,

221 + reduce_results: bool | None = True,

222 prefix: str = "",

223 ):

224 super().__init__()

@@ -340,8 +339,8 @@ class BailingMoeBlock(nn.Module):

339 def __init__(

340 self,

341 config: PretrainedConfig,

343 - cache_config: Optional[CacheConfig] = None,

344 - quant_config: Optional[QuantizationConfig] = None,

342 + cache_config: CacheConfig | None = None,

343 + quant_config: QuantizationConfig | None = None,

344 prefix: str = "",

345 ):

346 super().__init__()

@@ -370,7 +369,7 @@ class BailingMoeBlock(nn.Module):

369 self,

370 hidden_states: torch.Tensor,

371 position_ids: torch.Tensor,

373 - residual: Optional[torch.Tensor],

372 + residual: torch.Tensor | None,

373 ) -> torch.Tensor:

374 if residual is None:

375 residual = hidden_states

@@ -447,9 +446,9 @@ class BailingMoeModel(nn.Module):

446 self,

447 input_ids: torch.Tensor,

448 position_ids: torch.Tensor,

450 - intermediate_tensors: Optional[IntermediateTensors],

451 - inputs_embeds: Optional[torch.Tensor] = None,

452 - ) -> Union[torch.Tensor, IntermediateTensors]:

449 + intermediate_tensors: IntermediateTensors | None,

450 + inputs_embeds: torch.Tensor | None = None,

451 + ) -> torch.Tensor | IntermediateTensors:

452 if get_pp_group().is_first_rank:

453 if inputs_embeds is not None:

454 hidden_states = inputs_embeds

@@ -619,9 +618,9 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

618 self,

619 input_ids: torch.Tensor,

620 positions: torch.Tensor,

622 - intermediate_tensors: Optional[IntermediateTensors] = None,

623 - inputs_embeds: Optional[torch.Tensor] = None,

624 - ) -> Union[torch.Tensor, IntermediateTensors]:

621 + intermediate_tensors: IntermediateTensors | None = None,

622 + inputs_embeds: torch.Tensor | None = None,

623 + ) -> torch.Tensor | IntermediateTensors:

624 model_output = self.model(

625 input_ids, positions, intermediate_tensors, inputs_embeds

626 )

@@ -630,7 +629,7 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

629 def compute_logits(

630 self,

631 hidden_states: torch.Tensor,

633 - ) -> Optional[torch.Tensor]:

632 + ) -> torch.Tensor | None:

633 logits = self.logits_processor(self.lm_head, hidden_states)

634 return logits

635

vllm/model_executor/models/bamba.py

+14 -15

@@ -4,7 +4,6 @@

4

5 # Added by the IBM Team, 2024

6 from collections.abc import Iterable

7 -from typing import Optional

7

8 import torch

9 from torch import nn

@@ -52,7 +51,7 @@ class BambaMLP(nn.Module):

51 def __init__(

52 self,

53 config: BambaConfig,

55 - quant_config: Optional[QuantizationConfig] = None,

54 + quant_config: QuantizationConfig | None = None,

55 bias: bool = False,

56 ) -> None:

57 super().__init__()

@@ -87,9 +86,9 @@ class BambaMixerDecoderLayer(nn.Module):

86 self,

87 config: BambaConfig,

88 layer_idx: int,

90 - model_config: Optional[ModelConfig] = None,

91 - cache_config: Optional[CacheConfig] = None,

92 - quant_config: Optional[QuantizationConfig] = None,

89 + model_config: ModelConfig | None = None,

90 + cache_config: CacheConfig | None = None,

91 + quant_config: QuantizationConfig | None = None,

92 prefix: str = "",

93 ) -> None:

94 super().__init__()

@@ -119,7 +118,7 @@ class BambaMixerDecoderLayer(nn.Module):

118 def forward(

119 self,

120 hidden_states: torch.Tensor,

122 - residual: Optional[torch.Tensor],

121 + residual: torch.Tensor | None,

122 **kwargs,

123 ):

124 if residual is None:

@@ -141,9 +140,9 @@ class BambaAttentionDecoderLayer(nn.Module):

140 self,

141 config: BambaConfig,

142 layer_idx: int,

144 - model_config: Optional[ModelConfig] = None,

145 - cache_config: Optional[CacheConfig] = None,

146 - quant_config: Optional[QuantizationConfig] = None,

143 + model_config: ModelConfig | None = None,

144 + cache_config: CacheConfig | None = None,

145 + quant_config: QuantizationConfig | None = None,

146 prefix: str = "",

147 ) -> None:

148 super().__init__()

@@ -235,7 +234,7 @@ class BambaAttentionDecoderLayer(nn.Module):

234 self,

235 positions: torch.Tensor,

236 hidden_states: torch.Tensor,

238 - residual: Optional[torch.Tensor],

237 + residual: torch.Tensor | None,

238 **kwargs,

239 ):

240 if residual is None:

@@ -314,8 +313,8 @@ class BambaModel(nn.Module):

313 self,

314 input_ids: torch.Tensor,

315 positions: torch.Tensor,

317 - intermediate_tensors: Optional[IntermediateTensors] = None,

318 - inputs_embeds: Optional[torch.Tensor] = None,

316 + intermediate_tensors: IntermediateTensors | None = None,

317 + inputs_embeds: torch.Tensor | None = None,

318 ) -> torch.Tensor:

319 if get_pp_group().is_first_rank:

320 if inputs_embeds is not None:

@@ -497,8 +496,8 @@ class BambaForCausalLM(

496 self,

497 input_ids: torch.Tensor,

498 positions: torch.Tensor,

500 - intermediate_tensors: Optional[IntermediateTensors] = None,

501 - inputs_embeds: Optional[torch.Tensor] = None,

499 + intermediate_tensors: IntermediateTensors | None = None,

500 + inputs_embeds: torch.Tensor | None = None,

501 **kwargs,

502 ):

503 hidden_states = self.model(

@@ -510,7 +509,7 @@ class BambaForCausalLM(

509 def compute_logits(

510 self,

511 hidden_states: torch.Tensor,

513 - ) -> Optional[torch.Tensor]:

512 + ) -> torch.Tensor | None:

513 logits = self.logits_processor(self.lm_head, hidden_states)

514 return logits

515

vllm/model_executor/models/bert.py

+24 -25

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable, Set

5 -from typing import Optional, Union

5

6 import torch

7 from torch import nn

@@ -66,7 +65,7 @@ class BertEmbedding(nn.Module):

65 self,

66 input_ids: torch.Tensor,

67 position_ids: torch.Tensor,

69 - inputs_embeds: Optional[torch.Tensor] = None,

68 + inputs_embeds: torch.Tensor | None = None,

69 ) -> torch.Tensor:

70 token_type_ids = _decode_token_type_ids(input_ids)

71

@@ -103,9 +102,9 @@ class BertPooler(Pooler):

102

103 def forward(

104 self,

106 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

105 + hidden_states: torch.Tensor | list[torch.Tensor],

106 pooling_metadata: PoolingMetadata,

108 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

107 + ) -> torch.Tensor | list[torch.Tensor]:

108 pooled_output = self.pooling(hidden_states, pooling_metadata)

109

110 if isinstance(pooled_output, list):

@@ -147,8 +146,8 @@ class BertLayer(nn.Module):

146 def __init__(

147 self,

148 config: BertConfig,

150 - cache_config: Optional[CacheConfig] = None,

151 - quant_config: Optional[QuantizationConfig] = None,

149 + cache_config: CacheConfig | None = None,

150 + quant_config: QuantizationConfig | None = None,

151 prefix: str = "",

152 ):

153 super().__init__()

@@ -191,8 +190,8 @@ class BertAttention(nn.Module):

190 hidden_size: int,

191 num_attention_heads: int,

192 layer_norm_eps: float,

194 - cache_config: Optional[CacheConfig] = None,

195 - quant_config: Optional[QuantizationConfig] = None,

193 + cache_config: CacheConfig | None = None,

194 + quant_config: QuantizationConfig | None = None,

195 prefix: str = "",

196 ):

197 super().__init__()

@@ -225,8 +224,8 @@ class BertSelfAttention(nn.Module):

224 self,

225 hidden_size: int,

226 num_attention_heads: int,

228 - cache_config: Optional[CacheConfig] = None,

229 - quant_config: Optional[QuantizationConfig] = None,

227 + cache_config: CacheConfig | None = None,

228 + quant_config: QuantizationConfig | None = None,

229 prefix: str = "",

230 ):

231 super().__init__()

@@ -281,7 +280,7 @@ class BertSelfOutput(nn.Module):

280 self,

281 hidden_size: int,

282 layer_norm_eps: float,

284 - quant_config: Optional[QuantizationConfig] = None,

283 + quant_config: QuantizationConfig | None = None,

284 prefix: str = "",

285 ):

286 super().__init__()

@@ -308,7 +307,7 @@ class BertIntermediate(nn.Module):

307 hidden_size: int,

308 intermediate_size: int,

309 hidden_act: str,

311 - quant_config: Optional[QuantizationConfig] = None,

310 + quant_config: QuantizationConfig | None = None,

311 prefix: str = "",

312 ):

313 super().__init__()

@@ -333,7 +332,7 @@ class BertOutput(nn.Module):

332 hidden_size: int,

333 intermediate_size: int,

334 layer_norm_eps: float,

336 - quant_config: Optional[QuantizationConfig] = None,

335 + quant_config: QuantizationConfig | None = None,

336 prefix: str = "",

337 ):

338 super().__init__()

@@ -383,8 +382,8 @@ class BertModel(nn.Module, SupportsQuant):

382 self,

383 input_ids: torch.Tensor,

384 positions: torch.Tensor,

386 - intermediate_tensors: Optional[IntermediateTensors] = None,

387 - inputs_embeds: Optional[torch.Tensor] = None,

385 + intermediate_tensors: IntermediateTensors | None = None,

386 + inputs_embeds: torch.Tensor | None = None,

387 ) -> torch.Tensor:

388 hidden_states = self.embeddings(

389 input_ids=input_ids,

@@ -494,8 +493,8 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):

493 self,

494 input_ids: torch.Tensor,

495 positions: torch.Tensor,

497 - intermediate_tensors: Optional[IntermediateTensors] = None,

498 - inputs_embeds: Optional[torch.Tensor] = None,

496 + intermediate_tensors: IntermediateTensors | None = None,

497 + inputs_embeds: torch.Tensor | None = None,

498 ) -> torch.Tensor:

499 return self.model(

500 input_ids=input_ids,

@@ -636,11 +635,11 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu

635

636 def forward(

637 self,

639 - input_ids: Optional[torch.Tensor],

638 + input_ids: torch.Tensor | None,

639 positions: torch.Tensor,

641 - intermediate_tensors: Optional[IntermediateTensors] = None,

642 - inputs_embeds: Optional[torch.Tensor] = None,

643 - token_type_ids: Optional[torch.Tensor] = None,

640 + intermediate_tensors: IntermediateTensors | None = None,

641 + inputs_embeds: torch.Tensor | None = None,

642 + token_type_ids: torch.Tensor | None = None,

643 ) -> torch.Tensor:

644 if token_type_ids is not None:

645 assert self.bert.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)

@@ -692,11 +691,11 @@ class BertForTokenClassification(nn.Module):

691

692 def forward(

693 self,

695 - input_ids: Optional[torch.Tensor],

694 + input_ids: torch.Tensor | None,

695 positions: torch.Tensor,

697 - intermediate_tensors: Optional[IntermediateTensors] = None,

698 - inputs_embeds: Optional[torch.Tensor] = None,

699 - token_type_ids: Optional[torch.Tensor] = None,

696 + intermediate_tensors: IntermediateTensors | None = None,

697 + inputs_embeds: torch.Tensor | None = None,

698 + token_type_ids: torch.Tensor | None = None,

699 ) -> torch.Tensor:

700 if token_type_ids is not None:

701 assert self.bert.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)

vllm/model_executor/models/bert_with_rope.py

+18 -19

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable

4 -from typing import Optional

4

5 import torch

6 from torch import nn

@@ -67,7 +66,7 @@ class BertWithRopeEmbedding(nn.Module):

66 def forward(

67 self,

68 input_ids: torch.Tensor,

70 - token_type_ids: Optional[torch.Tensor] = None,

69 + token_type_ids: torch.Tensor | None = None,

70 ) -> torch.Tensor:

71 input_shape = input_ids.size()

72 inputs_embeds = self.word_embeddings(input_ids)

@@ -91,10 +90,10 @@ class BertWithRopeAttention(nn.Module):

90 self,

91 hidden_size: int,

92 num_attention_heads: int,

94 - cache_config: Optional[CacheConfig] = None,

95 - quant_config: Optional[QuantizationConfig] = None,

93 + cache_config: CacheConfig | None = None,

94 + quant_config: QuantizationConfig | None = None,

95 bias: bool = True,

97 - rotary_kwargs: Optional[dict] = None,

96 + rotary_kwargs: dict | None = None,

97 prefix: str = "",

98 ):

99 super().__init__()

@@ -166,7 +165,7 @@ class BertWithRopeGatedMLP(nn.Module):

165 intermediate_size: int,

166 hidden_act: str,

167 bias: bool = True,

169 - quant_config: Optional[QuantizationConfig] = None,

168 + quant_config: QuantizationConfig | None = None,

169 prefix: str = "",

170 ):

171 super().__init__()

@@ -200,7 +199,7 @@ class BertWithRopeMLP(nn.Module):

199 intermediate_size: int,

200 hidden_act: str,

201 bias: bool = True,

203 - quant_config: Optional[QuantizationConfig] = None,

202 + quant_config: QuantizationConfig | None = None,

203 prefix: str = "",

204 ):

205 super().__init__()

@@ -235,8 +234,8 @@ class NomicMoE(nn.Module):

234 hidden_size: int,

235 intermediate_size: int,

236 hidden_act: str,

238 - params_dtype: Optional[torch.dtype] = None,

239 - tp_size: Optional[int] = None,

237 + params_dtype: torch.dtype | None = None,

238 + tp_size: int | None = None,

239 ):

240 super().__init__()

241

@@ -344,11 +343,11 @@ class BertWithRopeBlock(nn.Module):

343 def __init__(

344 self,

345 config: PretrainedConfig,

347 - cache_config: Optional[CacheConfig] = None,

348 - quant_config: Optional[QuantizationConfig] = None,

346 + cache_config: CacheConfig | None = None,

347 + quant_config: QuantizationConfig | None = None,

348 moe: bool = False,

349 bias: bool = True,

351 - rotary_kwargs: Optional[dict] = None,

350 + rotary_kwargs: dict | None = None,

351 prefix: str = "",

352 ):

353 super().__init__()

@@ -406,7 +405,7 @@ class BertWithRopeEncoder(nn.Module):

405 self,

406 vllm_config: VllmConfig,

407 bias: bool = True,

409 - rotary_kwargs: Optional[dict] = None,

408 + rotary_kwargs: dict | None = None,

409 prefix: str = "",

410 ):

411 super().__init__()

@@ -471,9 +470,9 @@ class BertWithRope(nn.Module, SupportsQuant):

470 self,

471 input_ids: torch.Tensor,

472 positions: torch.Tensor,

474 - intermediate_tensors: Optional[IntermediateTensors] = None,

475 - inputs_embeds: Optional[torch.Tensor] = None,

476 - token_type_ids: Optional[torch.Tensor] = None,

473 + intermediate_tensors: IntermediateTensors | None = None,

474 + inputs_embeds: torch.Tensor | None = None,

475 + token_type_ids: torch.Tensor | None = None,

476 ) -> torch.Tensor:

477 if inputs_embeds is not None:

478 hidden_states = inputs_embeds

@@ -724,10 +723,10 @@ class GteNewForSequenceClassification(nn.Module, SupportsCrossEncoding):

723

724 def forward(

725 self,

727 - input_ids: Optional[torch.Tensor],

726 + input_ids: torch.Tensor | None,

727 positions: torch.Tensor,

729 - intermediate_tensors: Optional[IntermediateTensors] = None,

730 - inputs_embeds: Optional[torch.Tensor] = None,

728 + intermediate_tensors: IntermediateTensors | None = None,

729 + inputs_embeds: torch.Tensor | None = None,

730 ) -> torch.Tensor:

731 return self.new(

732 input_ids=input_ids,

vllm/model_executor/models/blip.py

+10 -11

@@ -4,7 +4,6 @@

4 within a vision language model."""

5

6 from collections.abc import Iterable

7 -from typing import Optional, Union

7

8 import torch

9 import torch.nn as nn

@@ -38,7 +37,7 @@ def get_blip_num_patches(*, image_size: int, patch_size: int) -> int:

37

38 # Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/blip/modeling_blip.py#L164 # noqa

39 class BlipVisionEmbeddings(nn.Module):

41 - def __init__(self, config: Union[BlipVisionConfig, Blip2VisionConfig]):

40 + def __init__(self, config: BlipVisionConfig | Blip2VisionConfig):

41 super().__init__()

42

43 self.config = config

@@ -86,8 +85,8 @@ class BlipAttention(nn.Module):

85

86 def __init__(

87 self,

89 - config: Union[BlipVisionConfig, Blip2VisionConfig],

90 - quant_config: Optional[QuantizationConfig] = None,

88 + config: BlipVisionConfig | Blip2VisionConfig,

89 + quant_config: QuantizationConfig | None = None,

90 prefix: str = "",

91 ) -> None:

92 super().__init__()

@@ -151,7 +150,7 @@ class BlipMLP(nn.Module):

150 def __init__(

151 self,

152 config: BlipVisionConfig,

154 - quant_config: Optional[QuantizationConfig] = None,

153 + quant_config: QuantizationConfig | None = None,

154 prefix: str = "",

155 ) -> None:

156 super().__init__()

@@ -186,7 +185,7 @@ class BlipEncoderLayer(nn.Module):

185 def __init__(

186 self,

187 config: BlipVisionConfig,

189 - quant_config: Optional[QuantizationConfig] = None,

188 + quant_config: QuantizationConfig | None = None,

189 prefix: str = "",

190 ) -> None:

191 super().__init__()

@@ -228,8 +227,8 @@ class BlipEncoder(nn.Module):

227 def __init__(

228 self,

229 config: BlipVisionConfig,

231 - quant_config: Optional[QuantizationConfig] = None,

232 - num_hidden_layers_override: Optional[int] = None,

230 + quant_config: QuantizationConfig | None = None,

231 + num_hidden_layers_override: int | None = None,

232 prefix: str = "",

233 ) -> None:

234 super().__init__()

@@ -268,10 +267,10 @@ class BlipVisionModel(nn.Module, SupportsQuant):

267 def __init__(

268 self,

269 config: BlipVisionConfig,

271 - quant_config: Optional[QuantizationConfig] = None,

270 + quant_config: QuantizationConfig | None = None,

271 *,

273 - num_hidden_layers_override: Optional[int] = None,

274 - require_post_norm: Optional[bool] = None,

272 + num_hidden_layers_override: int | None = None,

273 + require_post_norm: bool | None = None,

274 prefix: str = "",

275 ) -> None:

276 super().__init__()

vllm/model_executor/models/blip2.py

+21 -21

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import Annotated, Literal, Optional, Union

5 +from typing import Annotated, Literal, TypeAlias

6

7 import torch

8 import torch.nn as nn

@@ -70,7 +70,7 @@ class Blip2ImageEmbeddingInputs(TensorSchema):

70 data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")]

71

72

73 -Blip2ImageInputs = Union[Blip2ImagePixelInputs, Blip2ImageEmbeddingInputs]

73 +Blip2ImageInputs: TypeAlias = Blip2ImagePixelInputs | Blip2ImageEmbeddingInputs

74

75

76 class Blip2QFormerMultiHeadAttention(nn.Module):

@@ -78,8 +78,8 @@ class Blip2QFormerMultiHeadAttention(nn.Module):

78 self,

79 config: Blip2QFormerConfig,

80 *,

81 - quant_config: Optional[QuantizationConfig],

82 - cache_config: Optional[CacheConfig],

81 + quant_config: QuantizationConfig | None,

82 + cache_config: CacheConfig | None,

83 is_cross_attention: bool = False,

84 prefix: str = "",

85 ) -> None:

@@ -123,7 +123,7 @@ class Blip2QFormerMultiHeadAttention(nn.Module):

123 def forward(

124 self,

125 hidden_states: torch.Tensor,

126 - encoder_hidden_states: Optional[torch.FloatTensor] = None,

126 + encoder_hidden_states: torch.FloatTensor | None = None,

127 ):

128 is_cross_attention = encoder_hidden_states is not None

129

@@ -179,8 +179,8 @@ class Blip2QFormerAttention(nn.Module):

179 self,

180 config: Blip2QFormerConfig,

181 *,

182 - quant_config: Optional[QuantizationConfig],

183 - cache_config: Optional[CacheConfig],

182 + quant_config: QuantizationConfig | None,

183 + cache_config: CacheConfig | None,

184 is_cross_attention: bool = False,

185 prefix: str = "",

186 ) -> None:

@@ -199,7 +199,7 @@ class Blip2QFormerAttention(nn.Module):

199 def forward(

200 self,

201 hidden_states: torch.Tensor,

202 - encoder_hidden_states: Optional[torch.FloatTensor] = None,

202 + encoder_hidden_states: torch.FloatTensor | None = None,

203 ) -> tuple[torch.Tensor]:

204 self_output = self.attention(

205 hidden_states,

@@ -247,8 +247,8 @@ class Blip2QFormerLayer(nn.Module):

247 self,

248 config: Blip2QFormerConfig,

249 *,

250 - quant_config: Optional[QuantizationConfig],

251 - cache_config: Optional[CacheConfig],

250 + quant_config: QuantizationConfig | None,

251 + cache_config: CacheConfig | None,

252 layer_idx: int,

253 prefix: str = "",

254 ) -> None:

@@ -340,8 +340,8 @@ class Blip2QFormerEncoder(nn.Module):

340 self,

341 config: Blip2QFormerConfig,

342 *,

343 - quant_config: Optional[QuantizationConfig],

344 - cache_config: Optional[CacheConfig],

343 + quant_config: QuantizationConfig | None,

344 + cache_config: CacheConfig | None,

345 prefix: str = "",

346 ) -> None:

347 super().__init__()

@@ -385,8 +385,8 @@ class Blip2QFormerModel(nn.Module):

385 self,

386 config: Blip2QFormerConfig,

387 *,

388 - quant_config: Optional[QuantizationConfig],

389 - cache_config: Optional[CacheConfig],

388 + quant_config: QuantizationConfig | None,

389 + cache_config: CacheConfig | None,

390 prefix: str = "",

391 ) -> None:

392 super().__init__()

@@ -426,7 +426,7 @@ class Blip2ProcessingInfo(BaseProcessingInfo):

426 def get_hf_config(self):

427 return self.ctx.get_hf_config(Blip2Config)

428

429 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

429 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

430 return {"image": 1}

431

432 def get_num_image_tokens(self) -> int:

@@ -442,7 +442,7 @@ class Blip2DummyInputsBuilder(BaseDummyInputsBuilder[Blip2ProcessingInfo]):

442 self,

443 seq_len: int,

444 mm_counts: Mapping[str, int],

445 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

445 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

446 ) -> MultiModalDataDict:

447 hf_config = self.info.get_hf_config()

448 vision_config = hf_config.vision_config

@@ -526,7 +526,7 @@ class Blip2ForConditionalGeneration(

526 merge_by_field_config = True

527

528 @classmethod

529 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

529 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

530 if modality.startswith("image"):

531 return None

532

@@ -573,7 +573,7 @@ class Blip2ForConditionalGeneration(

573

574 def _parse_and_validate_image_input(

575 self, **kwargs: object

576 - ) -> Optional[Blip2ImageInputs]:

576 + ) -> Blip2ImageInputs | None:

577 pixel_values = kwargs.pop("pixel_values", None)

578 image_embeds = kwargs.pop("image_embeds", None)

579

@@ -641,8 +641,8 @@ class Blip2ForConditionalGeneration(

641 self,

642 input_ids: torch.Tensor,

643 positions: torch.Tensor,

644 - intermediate_tensors: Optional[IntermediateTensors] = None,

645 - inputs_embeds: Optional[torch.Tensor] = None,

644 + intermediate_tensors: IntermediateTensors | None = None,

645 + inputs_embeds: torch.Tensor | None = None,

646 **kwargs: object,

647 ) -> IntermediateTensors:

648 """Run forward pass for BLIP-2.

@@ -687,7 +687,7 @@ class Blip2ForConditionalGeneration(

687 def compute_logits(

688 self,

689 hidden_states: torch.Tensor,

690 - ) -> Optional[torch.Tensor]:

690 + ) -> torch.Tensor | None:

691 return self.language_model.compute_logits(hidden_states)

692

693 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/bloom.py

+12 -13

@@ -22,7 +22,6 @@

22 import math

23 from collections.abc import Iterable

24 from itertools import islice

25 -from typing import Optional, Union

25

26 import torch

27 from torch import nn

@@ -89,8 +88,8 @@ class BloomAttention(nn.Module):

88 def __init__(

89 self,

90 config: BloomConfig,

92 - cache_config: Optional[CacheConfig] = None,

93 - quant_config: Optional[QuantizationConfig] = None,

91 + cache_config: CacheConfig | None = None,

92 + quant_config: QuantizationConfig | None = None,

93 prefix: str = "",

94 ):

95 super().__init__()

@@ -152,7 +151,7 @@ class BloomMLP(nn.Module):

151 def __init__(

152 self,

153 config: BloomConfig,

155 - quant_config: Optional[QuantizationConfig] = None,

154 + quant_config: QuantizationConfig | None = None,

155 ):

156 super().__init__()

157 hidden_size = config.hidden_size

@@ -179,8 +178,8 @@ class BloomBlock(nn.Module):

178 def __init__(

179 self,

180 config: BloomConfig,

182 - cache_config: Optional[CacheConfig] = None,

183 - quant_config: Optional[QuantizationConfig] = None,

181 + cache_config: CacheConfig | None = None,

182 + quant_config: QuantizationConfig | None = None,

183 prefix: str = "",

184 ):

185 super().__init__()

@@ -274,9 +273,9 @@ class BloomModel(nn.Module):

273 self,

274 input_ids: torch.Tensor,

275 position_ids: torch.Tensor,

277 - intermediate_tensors: Optional[IntermediateTensors],

278 - inputs_embeds: Optional[torch.Tensor] = None,

279 - ) -> Union[torch.Tensor, IntermediateTensors]:

276 + intermediate_tensors: IntermediateTensors | None,

277 + inputs_embeds: torch.Tensor | None = None,

278 + ) -> torch.Tensor | IntermediateTensors:

279 if get_pp_group().is_first_rank:

280 if inputs_embeds is not None:

281 hidden_states = inputs_embeds

@@ -356,9 +355,9 @@ class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant):

355 self,

356 input_ids: torch.Tensor,

357 positions: torch.Tensor,

359 - intermediate_tensors: Optional[IntermediateTensors] = None,

360 - inputs_embeds: Optional[torch.Tensor] = None,

361 - ) -> Union[torch.Tensor, IntermediateTensors]:

358 + intermediate_tensors: IntermediateTensors | None = None,

359 + inputs_embeds: torch.Tensor | None = None,

360 + ) -> torch.Tensor | IntermediateTensors:

361 hidden_states = self.transformer(

362 input_ids, positions, intermediate_tensors, inputs_embeds

363 )

@@ -367,7 +366,7 @@ class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant):

366 def compute_logits(

367 self,

368 hidden_states: torch.Tensor,

370 - ) -> Optional[torch.Tensor]:

369 + ) -> torch.Tensor | None:

370 logits = self.logits_processor(self.lm_head, hidden_states)

371 return logits

372

vllm/model_executor/models/chameleon.py

+24 -24

@@ -4,7 +4,7 @@

4 from collections.abc import Iterable, Mapping, Sequence

5 from functools import cached_property

6 from itertools import islice

7 -from typing import Annotated, Any, Literal, Optional, Union

7 +from typing import Annotated, Any, Literal

8

9 import torch

10 import torch.nn as nn

@@ -94,7 +94,7 @@ class ChameleonProcessingInfo(BaseProcessingInfo):

94 def get_hf_processor(self, **kwargs: object):

95 return self.ctx.get_hf_processor(ChameleonProcessor, **kwargs)

96

97 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

97 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

98 return {"image": 1}

99

100 def get_num_image_tokens(self) -> int:

@@ -115,7 +115,7 @@ class ChameleonDummyInputsBuilder(BaseDummyInputsBuilder[ChameleonProcessingInfo

115 self,

116 seq_len: int,

117 mm_counts: Mapping[str, int],

118 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

118 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

119 ) -> MultiModalDataDict:

120 config = self.info.get_hf_config()

121

@@ -225,7 +225,7 @@ class ChameleonMLP(nn.Module):

225 hidden_size: int,

226 intermediate_size: int,

227 hidden_act: str,

228 - quant_config: Optional[QuantizationConfig] = None,

228 + quant_config: QuantizationConfig | None = None,

229 bias: bool = False,

230 ) -> None:

231 super().__init__()

@@ -262,11 +262,11 @@ class ChameleonAttention(nn.Module):

262 num_heads: int,

263 num_kv_heads: int,

264 rope_theta: float = 10000,

265 - rope_scaling: Optional[dict[str, Any]] = None,

265 + rope_scaling: dict[str, Any] | None = None,

266 max_position_embeddings: int = 4096,

267 - quant_config: Optional[QuantizationConfig] = None,

267 + quant_config: QuantizationConfig | None = None,

268 bias: bool = False,

269 - cache_config: Optional[CacheConfig] = None,

269 + cache_config: CacheConfig | None = None,

270 prefix: str = "",

271 ) -> None:

272 super().__init__()

@@ -357,8 +357,8 @@ class ChameleonDecoderLayer(nn.Module):

357 def __init__(

358 self,

359 config: ChameleonConfig,

360 - cache_config: Optional[CacheConfig] = None,

361 - quant_config: Optional[QuantizationConfig] = None,

360 + cache_config: CacheConfig | None = None,

361 + quant_config: QuantizationConfig | None = None,

362 prefix: str = "",

363 ) -> None:

364 super().__init__()

@@ -403,8 +403,8 @@ class ChameleonDecoderLayer(nn.Module):

403 self,

404 positions: torch.Tensor,

405 hidden_states: torch.Tensor,

406 - residual: Optional[torch.Tensor],

407 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

406 + residual: torch.Tensor | None,

407 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

408 if residual is None:

409 residual = hidden_states

410 hidden_states = self.input_layernorm(hidden_states)

@@ -426,8 +426,8 @@ class ChameleonSwinDecoderLayer(nn.Module):

426 def __init__(

427 self,

428 config: ChameleonConfig,

429 - cache_config: Optional[CacheConfig] = None,

430 - quant_config: Optional[QuantizationConfig] = None,

429 + cache_config: CacheConfig | None = None,

430 + quant_config: QuantizationConfig | None = None,

431 prefix: str = "",

432 ) -> None:

433 super().__init__()

@@ -472,7 +472,7 @@ class ChameleonSwinDecoderLayer(nn.Module):

472 self,

473 positions: torch.Tensor,

474 hidden_states: torch.Tensor,

475 - residual: Optional[torch.Tensor],

475 + residual: torch.Tensor | None,

476 ) -> tuple[torch.Tensor, torch.Tensor]:

477 residual = hidden_states

478 hidden_states = self.self_attn(

@@ -896,11 +896,11 @@ class ChameleonModel(nn.Module):

896

897 def forward(

898 self,

899 - input_ids: Optional[torch.Tensor],

899 + input_ids: torch.Tensor | None,

900 positions: torch.Tensor,

901 - intermediate_tensors: Optional[IntermediateTensors],

902 - inputs_embeds: Optional[torch.Tensor] = None,

903 - ) -> Union[torch.Tensor, IntermediateTensors]:

901 + intermediate_tensors: IntermediateTensors | None,

902 + inputs_embeds: torch.Tensor | None = None,

903 + ) -> torch.Tensor | IntermediateTensors:

904 if get_pp_group().is_first_rank:

905 if inputs_embeds is not None:

906 hidden_states = inputs_embeds

@@ -941,7 +941,7 @@ class ChameleonForConditionalGeneration(

941 }

942

943 @classmethod

944 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

944 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

945 if modality.startswith("image"):

946 return "<image>"

947

@@ -975,7 +975,7 @@ class ChameleonForConditionalGeneration(

975

976 def _parse_and_validate_image_input(

977 self, **kwargs: object

978 - ) -> Optional[ChameleonImagePixelInputs]:

978 + ) -> ChameleonImagePixelInputs | None:

979 pixel_values = kwargs.pop("pixel_values", None)

980

981 if pixel_values is None:

@@ -1008,10 +1008,10 @@ class ChameleonForConditionalGeneration(

1008 self,

1009 input_ids: torch.Tensor,

1010 positions: torch.Tensor,

1011 - intermediate_tensors: Optional[IntermediateTensors] = None,

1012 - inputs_embeds: Optional[torch.Tensor] = None,

1011 + intermediate_tensors: IntermediateTensors | None = None,

1012 + inputs_embeds: torch.Tensor | None = None,

1013 **kwargs,

1014 - ) -> Union[torch.Tensor, IntermediateTensors]:

1014 + ) -> torch.Tensor | IntermediateTensors:

1015 if intermediate_tensors is not None:

1016 inputs_embeds = None

1017

@@ -1023,7 +1023,7 @@ class ChameleonForConditionalGeneration(

1023 def compute_logits(

1024 self,

1025 hidden_states: torch.Tensor,

1026 - ) -> Optional[torch.Tensor]:

1026 + ) -> torch.Tensor | None:

1027 logits = self.logits_processor(self.lm_head, hidden_states)

1028

1029 # Disallow image tokens which does not include special

vllm/model_executor/models/chatglm.py

+15 -16

@@ -7,7 +7,6 @@

7 import json

8 from collections.abc import Iterable

9 from itertools import islice

10 -from typing import Optional, Union

10

11 import torch

12 from torch import nn

@@ -50,8 +49,8 @@ class GLMAttention(nn.Module):

49 def __init__(

50 self,

51 config: ChatGLMConfig,

53 - cache_config: Optional[CacheConfig] = None,

54 - quant_config: Optional[QuantizationConfig] = None,

52 + cache_config: CacheConfig | None = None,

53 + quant_config: QuantizationConfig | None = None,

54 prefix: str = "",

55 ):

56 super().__init__()

@@ -144,7 +143,7 @@ class GLMMLP(nn.Module):

143 def __init__(

144 self,

145 config: ChatGLMConfig,

147 - quant_config: Optional[QuantizationConfig] = None,

146 + quant_config: QuantizationConfig | None = None,

147 prefix: str = "",

148 ):

149 super().__init__()

@@ -190,8 +189,8 @@ class GLMBlock(nn.Module):

189 def __init__(

190 self,

191 config: ChatGLMConfig,

193 - cache_config: Optional[CacheConfig] = None,

194 - quant_config: Optional[QuantizationConfig] = None,

192 + cache_config: CacheConfig | None = None,

193 + quant_config: QuantizationConfig | None = None,

194 prefix: str = "",

195 ):

196 super().__init__()

@@ -263,8 +262,8 @@ class GLMTransformer(nn.Module):

262 def __init__(

263 self,

264 config: ChatGLMConfig,

266 - cache_config: Optional[CacheConfig] = None,

267 - quant_config: Optional[QuantizationConfig] = None,

265 + cache_config: CacheConfig | None = None,

266 + quant_config: QuantizationConfig | None = None,

267 prefix: str = "",

268 ):

269 super().__init__()

@@ -295,7 +294,7 @@ class GLMTransformer(nn.Module):

294 self,

295 hidden_states: torch.Tensor,

296 position_ids: torch.Tensor,

298 - ) -> Union[torch.Tensor, IntermediateTensors]:

297 + ) -> torch.Tensor | IntermediateTensors:

298 for layer in islice(self.layers, self.start_layer, self.end_layer):

299 hidden_states = layer(

300 hidden_states=hidden_states, position_ids=position_ids

@@ -361,10 +360,10 @@ class ChatGLMModel(nn.Module, SupportsQuant):

360 self,

361 input_ids: torch.Tensor,

362 positions: torch.Tensor,

364 - intermediate_tensors: Optional[IntermediateTensors] = None,

365 - inputs_embeds: Optional[torch.Tensor] = None,

363 + intermediate_tensors: IntermediateTensors | None = None,

364 + inputs_embeds: torch.Tensor | None = None,

365 **kwargs: object,

367 - ) -> Union[torch.Tensor, IntermediateTensors]:

366 + ) -> torch.Tensor | IntermediateTensors:

367 if get_pp_group().is_first_rank:

368 if inputs_embeds is not None:

369 hidden_states = inputs_embeds

@@ -459,7 +458,7 @@ class ChatGLMBaseModel(nn.Module):

458 def compute_logits(

459 self,

460 hidden_states: torch.Tensor,

462 - ) -> Optional[torch.Tensor]:

461 + ) -> torch.Tensor | None:

462 logits = self.logits_processor(self.lm_head, hidden_states)

463 return logits

464

@@ -491,9 +490,9 @@ class ChatGLMForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP, SupportsQua

490 self,

491 input_ids: torch.Tensor,

492 positions: torch.Tensor,

494 - intermediate_tensors: Optional[IntermediateTensors] = None,

495 - inputs_embeds: Optional[torch.Tensor] = None,

496 - ) -> Union[torch.Tensor, IntermediateTensors]:

493 + intermediate_tensors: IntermediateTensors | None = None,

494 + inputs_embeds: torch.Tensor | None = None,

495 + ) -> torch.Tensor | IntermediateTensors:

496 hidden_states = self.transformer(

497 input_ids, positions, intermediate_tensors, inputs_embeds

498 )

vllm/model_executor/models/clip.py

+44 -44

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Mapping, Sequence

4 from functools import cached_property

5 -from typing import Annotated, Literal, Optional, Union

5 +from typing import Annotated, Literal

6

7 import torch

8 import torch.nn as nn

@@ -125,7 +125,7 @@ class CLIPProcessingInfo(BaseProcessingInfo):

125 def get_hf_processor(self, **kwargs: object):

126 return self.ctx.get_hf_processor(CLIPProcessor, **kwargs)

127

128 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

128 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

129 return {"image": 1}

130

131 def get_num_image_tokens(

@@ -169,7 +169,7 @@ class CLIPDummyInputsBuilder(BaseDummyInputsBuilder[CLIPProcessingInfo]):

169 self,

170 seq_len: int,

171 mm_counts: Mapping[str, int],

172 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

172 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

173 ) -> MultiModalDataDict:

174 num_images = mm_counts.get("image", 0)

175

@@ -199,12 +199,12 @@ class CLIPMultiModalProcessor(BaseMultiModalProcessor[CLIPProcessingInfo]):

199

200 def apply(

201 self,

202 - prompt: Union[str, list[int]],

202 + prompt: str | list[int],

203 mm_data: MultiModalDataDict,

204 hf_processor_mm_kwargs: Mapping[str, object],

205 - tokenization_kwargs: Optional[Mapping[str, object]] = None,

205 + tokenization_kwargs: Mapping[str, object] | None = None,

206 *,

207 - mm_uuids: Optional[MultiModalUUIDDict] = None,

207 + mm_uuids: MultiModalUUIDDict | None = None,

208 ) -> MultiModalInputs:

209 if prompt and mm_data:

210 raise ValueError(

@@ -286,9 +286,9 @@ class CLIPTextEmbeddings(nn.Module):

286

287 def forward(

288 self,

289 - input_ids: Optional[torch.Tensor],

289 + input_ids: torch.Tensor | None,

290 position_ids: torch.Tensor,

291 - inputs_embeds: Optional[torch.Tensor] = None,

291 + inputs_embeds: torch.Tensor | None = None,

292 ) -> torch.Tensor:

293 if inputs_embeds is None:

294 if input_ids is None:

@@ -350,11 +350,11 @@ class CLIPVisionEmbeddings(nn.Module):

350 class CLIPAttention(nn.Module):

351 def __init__(

352 self,

353 - config: Union[CLIPTextConfig, CLIPVisionConfig],

354 - quant_config: Optional[QuantizationConfig] = None,

353 + config: CLIPTextConfig | CLIPVisionConfig,

354 + quant_config: QuantizationConfig | None = None,

355 *,

356 prefix: str = "",

357 - attn_cls: Union[type[Attention], type[MultiHeadAttention]],

357 + attn_cls: type[Attention] | type[MultiHeadAttention],

358 ) -> None:

359 super().__init__()

360

@@ -412,8 +412,8 @@ class CLIPAttention(nn.Module):

412 class CLIPMLP(nn.Module):

413 def __init__(

414 self,

415 - config: Union[CLIPTextConfig, CLIPVisionConfig],

416 - quant_config: Optional[QuantizationConfig] = None,

415 + config: CLIPTextConfig | CLIPVisionConfig,

416 + quant_config: QuantizationConfig | None = None,

417 prefix: str = "",

418 ) -> None:

419 super().__init__()

@@ -445,11 +445,11 @@ class CLIPMLP(nn.Module):

445 class CLIPEncoderLayer(nn.Module):

446 def __init__(

447 self,

448 - config: Union[CLIPTextConfig, CLIPVisionConfig],

449 - quant_config: Optional[QuantizationConfig] = None,

448 + config: CLIPTextConfig | CLIPVisionConfig,

449 + quant_config: QuantizationConfig | None = None,

450 *,

451 prefix: str = "",

452 - attn_cls: Union[type[Attention], type[MultiHeadAttention]],

452 + attn_cls: type[Attention] | type[MultiHeadAttention],

453 ) -> None:

454 super().__init__()

455 self.self_attn = CLIPAttention(

@@ -488,12 +488,12 @@ class CLIPEncoder(nn.Module):

488

489 def __init__(

490 self,

491 - config: Union[CLIPTextConfig, CLIPVisionConfig],

492 - quant_config: Optional[QuantizationConfig] = None,

493 - num_hidden_layers_override: Optional[int] = None,

491 + config: CLIPTextConfig | CLIPVisionConfig,

492 + quant_config: QuantizationConfig | None = None,

493 + num_hidden_layers_override: int | None = None,

494 *,

495 prefix: str = "",

496 - attn_cls: Union[type[Attention], type[MultiHeadAttention]],

496 + attn_cls: type[Attention] | type[MultiHeadAttention],

497 ) -> None:

498 super().__init__()

499

@@ -519,7 +519,7 @@ class CLIPEncoder(nn.Module):

519 self,

520 inputs_embeds: torch.Tensor,

521 return_all_hidden_states: bool,

522 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

522 + ) -> torch.Tensor | list[torch.Tensor]:

523 hidden_states_pool = [inputs_embeds]

524 hidden_states = inputs_embeds

525

@@ -538,7 +538,7 @@ class CLIPTextTransformer(nn.Module):

538 def __init__(

539 self,

540 config: CLIPTextConfig,

541 - quant_config: Optional[QuantizationConfig] = None,

541 + quant_config: QuantizationConfig | None = None,

542 *,

543 prefix: str = "",

544 ) -> None:

@@ -566,9 +566,9 @@ class CLIPTextTransformer(nn.Module):

566

567 def forward(

568 self,

569 - input_ids: Optional[torch.Tensor],

569 + input_ids: torch.Tensor | None,

570 position_ids: torch.Tensor,

571 - inputs_embeds: Optional[torch.Tensor] = None,

571 + inputs_embeds: torch.Tensor | None = None,

572 ) -> torch.Tensor:

573 hidden_states = self.embeddings(

574 input_ids=input_ids,

@@ -616,10 +616,10 @@ class CLIPVisionTransformer(nn.Module):

616 def __init__(

617 self,

618 config: CLIPVisionConfig,

619 - quant_config: Optional[QuantizationConfig] = None,

619 + quant_config: QuantizationConfig | None = None,

620 *,

621 - num_hidden_layers_override: Optional[int] = None,

622 - require_post_norm: Optional[bool] = None,

621 + num_hidden_layers_override: int | None = None,

622 + require_post_norm: bool | None = None,

623 prefix: str = "",

624 ) -> None:

625 super().__init__()

@@ -669,8 +669,8 @@ class CLIPVisionTransformer(nn.Module):

669 self,

670 pixel_values: torch.Tensor,

671 *,

672 - select_layers: Optional[list[int]] = None,

673 - feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

672 + select_layers: list[int] | None = None,

673 + feature_select_strategy: VisionFeatureSelectStrategy | None = None,

674 ) -> torch.Tensor:

675 hidden_states = self.embeddings(pixel_values)

676 hidden_states = self.pre_layrnorm(hidden_states)

@@ -736,10 +736,10 @@ class CLIPVisionModel(nn.Module):

736 def __init__(

737 self,

738 config: CLIPVisionConfig,

739 - quant_config: Optional[QuantizationConfig] = None,

739 + quant_config: QuantizationConfig | None = None,

740 *,

741 - num_hidden_layers_override: Optional[int] = None,

742 - require_post_norm: Optional[bool] = None,

741 + num_hidden_layers_override: int | None = None,

742 + require_post_norm: bool | None = None,

743 prefix: str = "",

744 ) -> None:

745 super().__init__()

@@ -755,8 +755,8 @@ class CLIPVisionModel(nn.Module):

755 def forward(

756 self,

757 pixel_values: torch.Tensor,

758 - select_layers: Optional[list[int]] = None,

759 - feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

758 + select_layers: list[int] | None = None,

759 + feature_select_strategy: VisionFeatureSelectStrategy | None = None,

760 ) -> torch.Tensor:

761 return self.vision_model(

762 pixel_values,

@@ -787,7 +787,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

787 merge_by_field_config = True

788

789 @classmethod

790 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

790 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

791 if modality.startswith("image"):

792 return None

793

@@ -847,9 +847,9 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

847

848 def get_text_features(

849 self,

850 - input_ids: Optional[torch.Tensor],

850 + input_ids: torch.Tensor | None,

851 position_ids: torch.Tensor,

852 - inputs_embeds: Optional[torch.Tensor] = None,

852 + inputs_embeds: torch.Tensor | None = None,

853 ) -> torch.Tensor:

854 pooled_output = self.text_model(

855 input_ids=input_ids,

@@ -864,7 +864,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

864 def get_image_features(

865 self,

866 pixel_values: torch.Tensor,

867 - feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

867 + feature_select_strategy: VisionFeatureSelectStrategy | None = None,

868 ) -> torch.Tensor:

869 if feature_select_strategy is None:

870 feature_select_strategy = _get_vision_feature_select_strategy(

@@ -883,7 +883,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

883

884 def _parse_and_validate_image_input(

885 self, **kwargs: object

886 - ) -> Optional[CLIPImagePixelInputs]:

886 + ) -> CLIPImagePixelInputs | None:

887 pixel_values = kwargs.pop("pixel_values", None)

888 if pixel_values is None:

889 return None

@@ -906,9 +906,9 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

906 def get_input_embeddings(

907 self,

908 input_ids: torch.Tensor,

909 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

909 + multimodal_embeddings: MultiModalEmbeddings | None = None,

910 *,

911 - is_multimodal: Optional[torch.Tensor] = None,

911 + is_multimodal: torch.Tensor | None = None,

912 handle_oov_mm_token: bool = False,

913 ) -> torch.Tensor:

914 self._is_text_input = (

@@ -936,10 +936,10 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

936

937 def forward(

938 self,

939 - input_ids: Optional[torch.Tensor],

939 + input_ids: torch.Tensor | None,

940 positions: torch.Tensor,

941 - intermediate_tensors: Optional[IntermediateTensors] = None,

942 - inputs_embeds: Optional[torch.Tensor] = None,

941 + intermediate_tensors: IntermediateTensors | None = None,

942 + inputs_embeds: torch.Tensor | None = None,

943 **kwargs: object,

944 ) -> torch.Tensor:

945 if intermediate_tensors is not None:

vllm/model_executor/models/cohere2_vision.py

+9 -9

@@ -4,7 +4,7 @@

4 """Command-A-Vision (Cohere2Vision) multimodal model implementation for vLLM."""

5

6 from collections.abc import Iterable, Mapping, Sequence

7 -from typing import Annotated, Literal, Optional, Union

7 +from typing import Annotated, Literal

8

9 import torch

10 from torch import nn

@@ -148,7 +148,7 @@ class Cohere2VisionProcessingInfo(BaseProcessingInfo):

148 def get_image_processor(self, **kwargs: object):

149 return self.get_hf_processor(**kwargs).image_processor

150

151 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

151 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

152 return {"image": None}

153

154 def get_image_size_with_most_features(self) -> ImageSize:

@@ -163,7 +163,7 @@ class Cohere2VisionProcessingInfo(BaseProcessingInfo):

163 *,

164 image_width: int,

165 image_height: int,

166 - processor: Optional[Cohere2VisionProcessor],

166 + processor: Cohere2VisionProcessor | None,

167 ) -> int:

168 """

169 Calculate the number of image patches for a given image.

@@ -217,7 +217,7 @@ class Cohere2VisionDummyInputsBuilder(

217 self,

218 seq_len: int,

219 mm_counts: Mapping[str, int],

220 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

220 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

221 ) -> MultiModalDataDict:

222 num_images = mm_counts.get("image", 0)

223 image_size = self.info.get_image_size_with_most_features()

@@ -404,7 +404,7 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo

404

405 def _parse_and_validate_image_input(

406 self, **kwargs: object

407 - ) -> Optional[Cohere2VisionImagePixelInputs]:

407 + ) -> Cohere2VisionImagePixelInputs | None:

408 pixel_values = kwargs.pop("pixel_values", None)

409 num_patches = kwargs.pop("num_patches", None)

410 image_embeds = kwargs.pop("image_embeds", None)

@@ -450,10 +450,10 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo

450 self,

451 input_ids: torch.Tensor,

452 positions: torch.Tensor,

453 - intermediate_tensors: Optional[IntermediateTensors] = None,

454 - inputs_embeds: Optional[torch.Tensor] = None,

453 + intermediate_tensors: IntermediateTensors | None = None,

454 + inputs_embeds: torch.Tensor | None = None,

455 **kwargs: object,

456 - ) -> Union[torch.Tensor, IntermediateTensors]:

456 + ) -> torch.Tensor | IntermediateTensors:

457 if intermediate_tensors is not None:

458 inputs_embeds = None

459

@@ -468,5 +468,5 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo

468 def compute_logits(

469 self,

470 hidden_states: torch.Tensor,

471 - ) -> Optional[torch.Tensor]:

471 + ) -> torch.Tensor | None:

472 return self.language_model.compute_logits(hidden_states)

vllm/model_executor/models/commandr.py

+16 -17

@@ -25,7 +25,6 @@

25

26 from collections.abc import Iterable

27 from itertools import islice

28 -from typing import Optional, Union

28

29 import torch

30 from torch import nn

@@ -94,8 +93,8 @@ class LayerNorm(nn.Module):

93 class CohereMLP(nn.Module):

94 def __init__(

95 self,

97 - config: Union[CohereConfig, Cohere2Config],

98 - quant_config: Optional[QuantizationConfig] = None,

96 + config: CohereConfig | Cohere2Config,

97 + quant_config: QuantizationConfig | None = None,

98 prefix: str = "",

99 ):

100 super().__init__()

@@ -128,9 +127,9 @@ class CohereMLP(nn.Module):

127 class CohereAttention(nn.Module):

128 def __init__(

129 self,

131 - config: Union[CohereConfig, Cohere2Config],

132 - cache_config: Optional[CacheConfig] = None,

133 - quant_config: Optional[QuantizationConfig] = None,

130 + config: CohereConfig | Cohere2Config,

131 + cache_config: CacheConfig | None = None,

132 + quant_config: QuantizationConfig | None = None,

133 prefix: str = "",

134 ):

135 super().__init__()

@@ -241,9 +240,9 @@ class CohereAttention(nn.Module):

240 class CohereDecoderLayer(nn.Module):

241 def __init__(

242 self,

244 - config: Union[CohereConfig, Cohere2Config],

245 - cache_config: Optional[CacheConfig] = None,

246 - quant_config: Optional[QuantizationConfig] = None,

243 + config: CohereConfig | Cohere2Config,

244 + cache_config: CacheConfig | None = None,

245 + quant_config: QuantizationConfig | None = None,

246 prefix: str = "",

247 ):

248 super().__init__()

@@ -265,7 +264,7 @@ class CohereDecoderLayer(nn.Module):

264 self,

265 positions: torch.Tensor,

266 hidden_states: torch.Tensor,

268 - residual: Optional[torch.Tensor],

267 + residual: torch.Tensor | None,

268 ) -> tuple[torch.Tensor, torch.Tensor]:

269 # Self Attention

270 residual = hidden_states

@@ -324,9 +323,9 @@ class CohereModel(nn.Module):

323 self,

324 input_ids: torch.Tensor,

325 positions: torch.Tensor,

327 - intermediate_tensors: Optional[IntermediateTensors],

328 - inputs_embeds: Optional[torch.Tensor] = None,

329 - ) -> Union[torch.Tensor, IntermediateTensors]:

326 + intermediate_tensors: IntermediateTensors | None,

327 + inputs_embeds: torch.Tensor | None = None,

328 + ) -> torch.Tensor | IntermediateTensors:

329 if get_pp_group().is_first_rank:

330 if inputs_embeds is not None:

331 hidden_states = inputs_embeds

@@ -452,9 +451,9 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

451 self,

452 input_ids: torch.Tensor,

453 positions: torch.Tensor,

455 - intermediate_tensors: Optional[IntermediateTensors] = None,

456 - inputs_embeds: Optional[torch.Tensor] = None,

457 - ) -> Union[torch.Tensor, IntermediateTensors]:

454 + intermediate_tensors: IntermediateTensors | None = None,

455 + inputs_embeds: torch.Tensor | None = None,

456 + ) -> torch.Tensor | IntermediateTensors:

457 hidden_states = self.model(

458 input_ids, positions, intermediate_tensors, inputs_embeds

459 )

@@ -463,7 +462,7 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

462 def compute_logits(

463 self,

464 hidden_states: torch.Tensor,

466 - ) -> Optional[torch.Tensor]:

465 + ) -> torch.Tensor | None:

466 is_not_lora = hasattr(self.model.embed_tokens, "weight")

467 if is_not_lora:

468 logits = self.logits_processor(self.model.embed_tokens, hidden_states)

vllm/model_executor/models/dbrx.py

+18 -19

@@ -3,7 +3,6 @@

3

4 from collections.abc import Iterable

5 from itertools import islice

6 -from typing import Optional, Union

6

7 import torch

8 import torch.nn as nn

@@ -54,7 +53,7 @@ class DbrxRouter(nn.Module):

53 def __init__(

54 self,

55 config: DbrxConfig,

57 - params_dtype: Optional[torch.dtype] = None,

56 + params_dtype: torch.dtype | None = None,

57 ):

58 super().__init__()

59 self.tp_size = get_tensor_model_parallel_world_size()

@@ -77,8 +76,8 @@ class DbrxExperts(FusedMoE):

76 def __init__(

77 self,

78 config: DbrxConfig,

80 - quant_config: Optional[QuantizationConfig] = None,

81 - params_dtype: Optional[torch.dtype] = None,

79 + quant_config: QuantizationConfig | None = None,

80 + params_dtype: torch.dtype | None = None,

81 prefix: str = "",

82 ):

83 super().__init__(

@@ -157,8 +156,8 @@ class DbrxMoE(nn.Module):

156 def __init__(

157 self,

158 config: DbrxConfig,

160 - quant_config: Optional[QuantizationConfig] = None,

161 - params_dtype: Optional[torch.dtype] = None,

159 + quant_config: QuantizationConfig | None = None,

160 + params_dtype: torch.dtype | None = None,

161 prefix: str = "",

162 ):

163 super().__init__()

@@ -189,8 +188,8 @@ class DbrxAttention(nn.Module):

188 def __init__(

189 self,

190 config: DbrxConfig,

192 - cache_config: Optional[CacheConfig] = None,

193 - quant_config: Optional[QuantizationConfig] = None,

191 + cache_config: CacheConfig | None = None,

192 + quant_config: QuantizationConfig | None = None,

193 prefix: str = "",

194 ):

195 super().__init__()

@@ -270,8 +269,8 @@ class DbrxFusedNormAttention(nn.Module):

269 def __init__(

270 self,

271 config: DbrxConfig,

273 - cache_config: Optional[CacheConfig] = None,

274 - quant_config: Optional[QuantizationConfig] = None,

272 + cache_config: CacheConfig | None = None,

273 + quant_config: QuantizationConfig | None = None,

274 prefix: str = "",

275 ):

276 super().__init__()

@@ -303,8 +302,8 @@ class DbrxBlock(nn.Module):

302 def __init__(

303 self,

304 config: DbrxConfig,

306 - cache_config: Optional[CacheConfig] = None,

307 - quant_config: Optional[QuantizationConfig] = None,

305 + cache_config: CacheConfig | None = None,

306 + quant_config: QuantizationConfig | None = None,

307 prefix: str = "",

308 ):

309 super().__init__()

@@ -361,9 +360,9 @@ class DbrxModel(nn.Module):

360 self,

361 input_ids: torch.Tensor,

362 position_ids: torch.Tensor,

364 - intermediate_tensors: Optional[IntermediateTensors],

365 - inputs_embeds: Optional[torch.Tensor] = None,

366 - ) -> Union[torch.Tensor, IntermediateTensors]:

363 + intermediate_tensors: IntermediateTensors | None,

364 + inputs_embeds: torch.Tensor | None = None,

365 + ) -> torch.Tensor | IntermediateTensors:

366 if get_pp_group().is_first_rank:

367 if inputs_embeds is not None:

368 hidden_states = inputs_embeds

@@ -466,9 +465,9 @@ class DbrxForCausalLM(nn.Module, SupportsPP):

465 self,

466 input_ids: torch.Tensor,

467 positions: torch.Tensor,

469 - intermediate_tensors: Optional[IntermediateTensors] = None,

470 - inputs_embeds: Optional[torch.Tensor] = None,

471 - ) -> Union[torch.Tensor, IntermediateTensors]:

468 + intermediate_tensors: IntermediateTensors | None = None,

469 + inputs_embeds: torch.Tensor | None = None,

470 + ) -> torch.Tensor | IntermediateTensors:

471 hidden_states = self.transformer(

472 input_ids, positions, intermediate_tensors, inputs_embeds

473 )

@@ -477,7 +476,7 @@ class DbrxForCausalLM(nn.Module, SupportsPP):

476 def compute_logits(

477 self,

478 hidden_states: torch.Tensor,

480 - ) -> Optional[torch.Tensor]:

479 + ) -> torch.Tensor | None:

480 logits = self.logits_processor(self.lm_head, hidden_states)

481 return logits

482

vllm/model_executor/models/deepseek.py

+16 -16

@@ -26,7 +26,7 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Any, Optional, Union

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -76,7 +76,7 @@ class DeepseekMLP(nn.Module):

76 hidden_size: int,

77 intermediate_size: int,

78 hidden_act: str,

79 - quant_config: Optional[QuantizationConfig] = None,

79 + quant_config: QuantizationConfig | None = None,

80 reduce_results: bool = True,

81 prefix: str = "",

82 ) -> None:

@@ -108,7 +108,7 @@ class DeepseekMoE(nn.Module):

108 def __init__(

109 self,

110 config: PretrainedConfig,

111 - quant_config: Optional[QuantizationConfig] = None,

111 + quant_config: QuantizationConfig | None = None,

112 prefix: str = "",

113 ):

114 super().__init__()

@@ -203,10 +203,10 @@ class DeepseekAttention(nn.Module):

203 num_heads: int,

204 num_kv_heads: int,

205 rope_theta: float = 10000,

206 - rope_scaling: Optional[dict[str, Any]] = None,

206 + rope_scaling: dict[str, Any] | None = None,

207 max_position_embeddings: int = 8192,

208 - cache_config: Optional[CacheConfig] = None,

209 - quant_config: Optional[QuantizationConfig] = None,

208 + cache_config: CacheConfig | None = None,

209 + quant_config: QuantizationConfig | None = None,

210 prefix: str = "",

211 ) -> None:

212 super().__init__()

@@ -282,8 +282,8 @@ class DeepseekDecoderLayer(nn.Module):

282 def __init__(

283 self,

284 config: PretrainedConfig,

285 - cache_config: Optional[CacheConfig] = None,

286 - quant_config: Optional[QuantizationConfig] = None,

285 + cache_config: CacheConfig | None = None,

286 + quant_config: QuantizationConfig | None = None,

287 prefix: str = "",

288 ) -> None:

289 super().__init__()

@@ -328,7 +328,7 @@ class DeepseekDecoderLayer(nn.Module):

328 self,

329 positions: torch.Tensor,

330 hidden_states: torch.Tensor,

331 - residual: Optional[torch.Tensor],

331 + residual: torch.Tensor | None,

332 ) -> torch.Tensor:

333 # Self Attention

334 if residual is None:

@@ -382,9 +382,9 @@ class DeepseekModel(nn.Module):

382 self,

383 input_ids: torch.Tensor,

384 positions: torch.Tensor,

385 - intermediate_tensors: Optional[IntermediateTensors],

386 - inputs_embeds: Optional[torch.Tensor] = None,

387 - ) -> Union[torch.Tensor, IntermediateTensors]:

385 + intermediate_tensors: IntermediateTensors | None,

386 + inputs_embeds: torch.Tensor | None = None,

387 + ) -> torch.Tensor | IntermediateTensors:

388 if get_pp_group().is_first_rank:

389 if inputs_embeds is not None:

390 hidden_states = inputs_embeds

@@ -489,9 +489,9 @@ class DeepseekForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

489 self,

490 input_ids: torch.Tensor,

491 positions: torch.Tensor,

492 - intermediate_tensors: Optional[IntermediateTensors] = None,

493 - inputs_embeds: Optional[torch.Tensor] = None,

494 - ) -> Union[torch.Tensor, IntermediateTensors]:

492 + intermediate_tensors: IntermediateTensors | None = None,

493 + inputs_embeds: torch.Tensor | None = None,

494 + ) -> torch.Tensor | IntermediateTensors:

495 hidden_states = self.model(

496 input_ids, positions, intermediate_tensors, inputs_embeds

497 )

@@ -500,7 +500,7 @@ class DeepseekForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

500 def compute_logits(

501 self,

502 hidden_states: torch.Tensor,

503 - ) -> Optional[torch.Tensor]:

503 + ) -> torch.Tensor | None:

504 logits = self.logits_processor(self.lm_head, hidden_states)

505 return logits

506

vllm/model_executor/models/deepseek_eagle.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable

5 -from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -224,7 +223,7 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):

223 input_ids: torch.Tensor,

224 positions: torch.Tensor,

225 hidden_states: torch.Tensor,

227 - inputs_embeds: Optional[torch.Tensor] = None,

226 + inputs_embeds: torch.Tensor | None = None,

227 ) -> tuple[torch.Tensor, torch.Tensor]:

228 if inputs_embeds is not None:

229 raise NotImplementedError(

@@ -235,7 +234,7 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):

234 def compute_logits(

235 self,

236 hidden_states: torch.Tensor,

238 - ) -> Optional[torch.Tensor]:

237 + ) -> torch.Tensor | None:

238 logits = self.logits_processor(self.lm_head, hidden_states)

239 return logits

240

vllm/model_executor/models/deepseek_mtp.py

+6 -7

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable

4 -from typing import Optional

4

5 import torch

6 import torch.nn as nn

@@ -30,7 +29,7 @@ class SharedHead(nn.Module):

29 self,

30 config: PretrainedConfig,

31 prefix: str,

33 - quant_config: Optional[QuantizationConfig] = None,

32 + quant_config: QuantizationConfig | None = None,

33 ) -> None:

34 super().__init__()

35 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

@@ -84,7 +83,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):

83 input_ids: torch.Tensor,

84 positions: torch.Tensor,

85 previous_hidden_states: torch.Tensor,

87 - inputs_embeds: Optional[torch.Tensor] = None,

86 + inputs_embeds: torch.Tensor | None = None,

87 spec_step_index: int = 0,

88 ) -> torch.Tensor:

89 assert inputs_embeds is not None

@@ -136,7 +135,7 @@ class DeepSeekMultiTokenPredictor(nn.Module):

135 input_ids: torch.Tensor,

136 positions: torch.Tensor,

137 previous_hidden_states: torch.Tensor,

139 - inputs_embeds: Optional[torch.Tensor] = None,

138 + inputs_embeds: torch.Tensor | None = None,

139 spec_step_idx: int = 0,

140 ) -> torch.Tensor:

141 if inputs_embeds is None:

@@ -180,8 +179,8 @@ class DeepSeekMTP(nn.Module, SupportsPP):

179 input_ids: torch.Tensor,

180 positions: torch.Tensor,

181 hidden_states: torch.Tensor,

183 - intermediate_tensors: Optional[IntermediateTensors] = None,

184 - inputs_embeds: Optional[torch.Tensor] = None,

182 + intermediate_tensors: IntermediateTensors | None = None,

183 + inputs_embeds: torch.Tensor | None = None,

184 spec_step_idx: int = 0,

185 ) -> torch.Tensor:

186 hidden_states = self.model(

@@ -193,7 +192,7 @@ class DeepSeekMTP(nn.Module, SupportsPP):

192 self,

193 hidden_states: torch.Tensor,

194 spec_step_idx: int = 0,

196 - ) -> Optional[torch.Tensor]:

195 + ) -> torch.Tensor | None:

196 return self.model.compute_logits(hidden_states, spec_step_idx)

197

198 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/deepseek_v2.py

+35 -35

@@ -27,7 +27,7 @@

27 import typing

28 from collections.abc import Callable, Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import torch

33 from torch import nn

@@ -106,7 +106,7 @@ class DeepseekV2MLP(nn.Module):

106 hidden_size: int,

107 intermediate_size: int,

108 hidden_act: str,

109 - quant_config: Optional[QuantizationConfig] = None,

109 + quant_config: QuantizationConfig | None = None,

110 reduce_results: bool = True,

111 is_sequence_parallel=False,

112 prefix: str = "",

@@ -150,9 +150,9 @@ class DeepseekV2MLP(nn.Module):

150 class DeepseekV2MoE(nn.Module):

151 def __init__(

152 self,

153 - config: Union[DeepseekV2Config, DeepseekV3Config],

153 + config: DeepseekV2Config | DeepseekV3Config,

154 parallel_config: ParallelConfig,

155 - quant_config: Optional[QuantizationConfig] = None,

155 + quant_config: QuantizationConfig | None = None,

156 prefix: str = "",

157 ):

158 super().__init__()

@@ -301,7 +301,7 @@ class DeepseekV2Attention(nn.Module):

301 def __init__(

302 self,

303 vllm_config: VllmConfig,

304 - config: Union[DeepseekV2Config, DeepseekV3Config],

304 + config: DeepseekV2Config | DeepseekV3Config,

305 hidden_size: int,

306 num_heads: int,

307 qk_nope_head_dim: int,

@@ -310,11 +310,11 @@ class DeepseekV2Attention(nn.Module):

310 q_lora_rank: int,

311 kv_lora_rank: int,

312 rope_theta: float = 10000,

313 - rope_scaling: Optional[dict[str, Any]] = None,

313 + rope_scaling: dict[str, Any] | None = None,

314 max_position_embeddings: int = 8192,

315 - cache_config: Optional[CacheConfig] = None,

316 - quant_config: Optional[QuantizationConfig] = None,

317 - topk_indices_buffer: Optional[torch.Tensor] = None,

315 + cache_config: CacheConfig | None = None,

316 + quant_config: QuantizationConfig | None = None,

317 + topk_indices_buffer: torch.Tensor | None = None,

318 prefix: str = "",

319 ) -> None:

320 super().__init__()

@@ -554,12 +554,12 @@ def sparse_attn_indexer(

554 k: torch.Tensor,

555 weights: torch.Tensor,

556 quant_block_size: int,

557 - scale_fmt: Optional[str],

557 + scale_fmt: str | None,

558 topk_tokens: int,

559 head_dim: int,

560 max_model_len: int,

561 total_seq_lens: int,

562 - topk_indices_buffer: Optional[torch.Tensor],

562 + topk_indices_buffer: torch.Tensor | None,

563 ) -> torch.Tensor:

564 # careful! this will be None in dummy run

565 attn_metadata = get_forward_context().attn_metadata

@@ -727,12 +727,12 @@ def sparse_attn_indexer_fake(

727 k: torch.Tensor,

728 weights: torch.Tensor,

729 quant_block_size: int,

730 - scale_fmt: Optional[str],

730 + scale_fmt: str | None,

731 topk_tokens: int,

732 head_dim: int,

733 max_model_len: int,

734 total_seq_lens: int,

735 - topk_indices_buffer: Optional[torch.Tensor],

735 + topk_indices_buffer: torch.Tensor | None,

736 ) -> torch.Tensor:

737 # profile run

738 # NOTE(Chen): create the max possible flattened_kv. So that

@@ -758,12 +758,12 @@ class Indexer(nn.Module):

758 def __init__(

759 self,

760 vllm_config: VllmConfig,

761 - config: Union[DeepseekV2Config, DeepseekV3Config],

761 + config: DeepseekV2Config | DeepseekV3Config,

762 hidden_size: int,

763 q_lora_rank: int,

764 - quant_config: Optional[QuantizationConfig],

765 - cache_config: Optional[CacheConfig],

766 - topk_indices_buffer: Optional[torch.Tensor],

764 + quant_config: QuantizationConfig | None,

765 + cache_config: CacheConfig | None,

766 + topk_indices_buffer: torch.Tensor | None,

767 prefix: str = "",

768 ):

769 super().__init__()

@@ -880,21 +880,21 @@ class DeepseekV2MLAAttention(nn.Module):

880 def __init__(

881 self,

882 vllm_config: VllmConfig,

883 - config: Union[DeepseekV2Config, DeepseekV3Config],

883 + config: DeepseekV2Config | DeepseekV3Config,

884 hidden_size: int,

885 num_heads: int,

886 qk_nope_head_dim: int,

887 qk_rope_head_dim: int,

888 v_head_dim: int,

889 - q_lora_rank: Optional[int],

889 + q_lora_rank: int | None,

890 kv_lora_rank: int,

891 rope_theta: float = 10000,

892 - rope_scaling: Optional[dict[str, Any]] = None,

892 + rope_scaling: dict[str, Any] | None = None,

893 max_position_embeddings: int = 8192,

894 - cache_config: Optional[CacheConfig] = None,

895 - quant_config: Optional[QuantizationConfig] = None,

894 + cache_config: CacheConfig | None = None,

895 + quant_config: QuantizationConfig | None = None,

896 prefix: str = "",

897 - topk_indices_buffer: Optional[torch.Tensor] = None,

897 + topk_indices_buffer: torch.Tensor | None = None,

898 ) -> None:

899 super().__init__()

900 self.hidden_size = hidden_size

@@ -1045,8 +1045,8 @@ class DeepseekV2DecoderLayer(nn.Module):

1045 self,

1046 vllm_config: VllmConfig,

1047 prefix: str,

1048 - config: Optional[DeepseekV2Config] = None,

1049 - topk_indices_buffer: Optional[torch.Tensor] = None,

1048 + config: DeepseekV2Config | None = None,

1049 + topk_indices_buffer: torch.Tensor | None = None,

1050 ) -> None:

1051 super().__init__()

1052

@@ -1117,7 +1117,7 @@ class DeepseekV2DecoderLayer(nn.Module):

1117 self,

1118 positions: torch.Tensor,

1119 hidden_states: torch.Tensor,

1120 - residual: Optional[torch.Tensor],

1120 + residual: torch.Tensor | None,

1121 ) -> torch.Tensor:

1122 # Self Attention

1123 if residual is None:

@@ -1212,9 +1212,9 @@ class DeepseekV2Model(nn.Module):

1212 self,

1213 input_ids: torch.Tensor,

1214 positions: torch.Tensor,

1215 - intermediate_tensors: Optional[IntermediateTensors],

1216 - inputs_embeds: Optional[torch.Tensor] = None,

1217 - ) -> Union[torch.Tensor, IntermediateTensors]:

1215 + intermediate_tensors: IntermediateTensors | None,

1216 + inputs_embeds: torch.Tensor | None = None,

1217 + ) -> torch.Tensor | IntermediateTensors:

1218 if get_pp_group().is_first_rank:

1219 if inputs_embeds is not None:

1220 hidden_states = inputs_embeds

@@ -1347,9 +1347,9 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR

1347 self,

1348 input_ids: torch.Tensor,

1349 positions: torch.Tensor,

1350 - intermediate_tensors: Optional[IntermediateTensors] = None,

1351 - inputs_embeds: Optional[torch.Tensor] = None,

1352 - ) -> Union[torch.Tensor, IntermediateTensors]:

1350 + intermediate_tensors: IntermediateTensors | None = None,

1351 + inputs_embeds: torch.Tensor | None = None,

1352 + ) -> torch.Tensor | IntermediateTensors:

1353 hidden_states = self.model(

1354 input_ids, positions, intermediate_tensors, inputs_embeds

1355 )

@@ -1358,7 +1358,7 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR

1358 def compute_logits(

1359 self,

1360 hidden_states: torch.Tensor,

1361 - ) -> Optional[torch.Tensor]:

1361 + ) -> torch.Tensor | None:

1362 logits = self.logits_processor(self.lm_head, hidden_states)

1363 return logits

1364

@@ -1497,8 +1497,8 @@ class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):

1497 # Compatibility with

1498 # https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py

1499 def get_spec_layer_idx_from_weight_name(

1500 - config: Union[DeepseekV2Config, DeepseekV3Config], weight_name: str

1501 -) -> Optional[int]:

1500 + config: DeepseekV2Config | DeepseekV3Config, weight_name: str

1501 +) -> int | None:

1502 if (

1503 hasattr(config, "num_nextn_predict_layers")

1504 and config.num_nextn_predict_layers > 0

vllm/model_executor/models/deepseek_vl2.py

+15 -17

@@ -6,7 +6,7 @@

6

7 import math

8 from collections.abc import Iterable, Mapping, Sequence

9 -from typing import Annotated, Literal, Optional, Union

9 +from typing import Annotated, Literal, TypeAlias

10

11 import torch

12 import torch.nn as nn

@@ -88,14 +88,12 @@ class DeepseekVL2VImageEmbeddingInputs(TensorSchema):

88 """

89

90 type: Literal["image_embeds"]

91 - data: Annotated[

92 - Union[torch.Tensor, list[torch.Tensor]], TensorShape("bn", "f", "h")

93 - ]

91 + data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("bn", "f", "h")]

92

93

96 -DeepseekVL2ImageInputs = Union[

97 - DeepseekVL2ImagePixelInputs, DeepseekVL2VImageEmbeddingInputs

98 -]

94 +DeepseekVL2ImageInputs: TypeAlias = (

95 + DeepseekVL2ImagePixelInputs | DeepseekVL2VImageEmbeddingInputs

96 +)

97

98

99 class MlpProjector(nn.Module):

@@ -161,7 +159,7 @@ class DeepseekVL2ProcessingInfo(BaseProcessingInfo):

159 def get_hf_processor(self, **kwargs: object):

160 return self.ctx.get_hf_processor(DeepseekVLV2Processor, **kwargs)

161

164 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

162 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

163 return {"image": None}

164

165 def get_num_image_tokens(

@@ -214,7 +212,7 @@ class DeepseekVL2DummyInputsBuilder(BaseDummyInputsBuilder[DeepseekVL2Processing

212 self,

213 seq_len: int,

214 mm_counts: Mapping[str, int],

217 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

215 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

216 ) -> MultiModalDataDict:

217 num_images = mm_counts.get("image", 0)

218

@@ -310,11 +308,11 @@ class DeepseekVL2MultiModalProcessor(

308

309 def _cached_apply_hf_processor(

310 self,

313 - prompt: Union[str, list[int]],

311 + prompt: str | list[int],

312 mm_data_items: MultiModalDataItems,

313 hf_processor_mm_kwargs: Mapping[str, object],

314 tokenization_kwargs: Mapping[str, object],

317 - mm_uuids: Optional[MultiModalUUIDDict] = None,

315 + mm_uuids: MultiModalUUIDDict | None = None,

316 ) -> tuple[list[int], MultiModalProcessingInfo, bool]:

317 # The processor logic is different for len(images) <= 2 vs > 2

318 # Since the processing cache assumes that the processor output is

@@ -353,7 +351,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

351 )

352

353 @classmethod

356 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

354 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

355 if modality.startswith("image"):

356 return "<image>"

357

@@ -454,7 +452,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

452 def _init_vision_module(

453 self,

454 vision_config: VisionEncoderConfig,

457 - quant_config: Optional[QuantizationConfig],

455 + quant_config: QuantizationConfig | None,

456 prefix: str = "",

457 ) -> nn.Module:

458 # TODO: refactor vision model through timm wrapper from transformers

@@ -480,7 +478,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

478

479 def _parse_and_validate_image_input(

480 self, **kwargs: object

483 - ) -> Optional[DeepseekVL2ImageInputs]:

481 + ) -> DeepseekVL2ImageInputs | None:

482 pixel_values = kwargs.pop("pixel_values", None)

483 images_spatial_crop = kwargs.pop("images_spatial_crop", None)

484 image_embeds = kwargs.pop("image_embeds", None)

@@ -637,8 +635,8 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

635 self,

636 input_ids: torch.Tensor,

637 positions: torch.Tensor,

640 - intermediate_tensors: Optional[IntermediateTensors] = None,

641 - inputs_embeds: Optional[torch.Tensor] = None,

638 + intermediate_tensors: IntermediateTensors | None = None,

639 + inputs_embeds: torch.Tensor | None = None,

640 **kwargs: object,

641 ):

642 if intermediate_tensors is not None:

@@ -653,7 +651,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

651 def compute_logits(

652 self,

653 hidden_states: torch.Tensor,

656 - ) -> Optional[torch.Tensor]:

654 + ) -> torch.Tensor | None:

655 return self.language_model.compute_logits(hidden_states)

656

657 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/dots1.py

+16 -16

@@ -27,7 +27,7 @@

27

28 from collections.abc import Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import torch

33 from torch import nn

@@ -80,7 +80,7 @@ class Dots1MLP(nn.Module):

80 hidden_size: int,

81 intermediate_size: int,

82 hidden_act: str,

83 - quant_config: Optional[QuantizationConfig] = None,

83 + quant_config: QuantizationConfig | None = None,

84 reduce_results: bool = True,

85 prefix: str = "",

86 ) -> None:

@@ -117,7 +117,7 @@ class Dots1MoE(nn.Module):

117 def __init__(

118 self,

119 config: Dots1Config,

120 - quant_config: Optional[QuantizationConfig] = None,

120 + quant_config: QuantizationConfig | None = None,

121 prefix: str = "",

122 ):

123 super().__init__()

@@ -203,10 +203,10 @@ class Dots1Attention(nn.Module):

203 num_kv_heads: int,

204 config: Dots1Config,

205 rope_theta: float = 10000,

206 - rope_scaling: Optional[dict[str, Any]] = None,

206 + rope_scaling: dict[str, Any] | None = None,

207 max_position_embeddings: int = 8192,

208 - cache_config: Optional[CacheConfig] = None,

209 - quant_config: Optional[QuantizationConfig] = None,

208 + cache_config: CacheConfig | None = None,

209 + quant_config: QuantizationConfig | None = None,

210 prefix: str = "",

211 ) -> None:

212 super().__init__()

@@ -289,8 +289,8 @@ class Dots1DecoderLayer(nn.Module):

289 config: Dots1Config,

290 prefix: str,

291 model_config: ModelConfig,

292 - cache_config: Optional[CacheConfig] = None,

293 - quant_config: Optional[QuantizationConfig] = None,

292 + cache_config: CacheConfig | None = None,

293 + quant_config: QuantizationConfig | None = None,

294 ) -> None:

295 super().__init__()

296 self.hidden_size = config.hidden_size

@@ -338,7 +338,7 @@ class Dots1DecoderLayer(nn.Module):

338 self,

339 positions: torch.Tensor,

340 hidden_states: torch.Tensor,

341 - residual: Optional[torch.Tensor],

341 + residual: torch.Tensor | None,

342 ) -> torch.Tensor:

343 if residual is None:

344 residual = hidden_states

@@ -403,9 +403,9 @@ class Dots1Model(nn.Module):

403 self,

404 input_ids: torch.Tensor,

405 positions: torch.Tensor,

406 - intermediate_tensors: Optional[IntermediateTensors],

407 - inputs_embeds: Optional[torch.Tensor] = None,

408 - ) -> Union[torch.Tensor, IntermediateTensors]:

406 + intermediate_tensors: IntermediateTensors | None,

407 + inputs_embeds: torch.Tensor | None = None,

408 + ) -> torch.Tensor | IntermediateTensors:

409 if get_pp_group().is_first_rank:

410 if inputs_embeds is not None:

411 hidden_states = inputs_embeds

@@ -546,9 +546,9 @@ class Dots1ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

546 self,

547 input_ids: torch.Tensor,

548 positions: torch.Tensor,

549 - intermediate_tensors: Optional[IntermediateTensors] = None,

550 - inputs_embeds: Optional[torch.Tensor] = None,

551 - ) -> Union[torch.Tensor, IntermediateTensors]:

549 + intermediate_tensors: IntermediateTensors | None = None,

550 + inputs_embeds: torch.Tensor | None = None,

551 + ) -> torch.Tensor | IntermediateTensors:

552 hidden_states = self.model(

553 input_ids,

554 positions,

@@ -560,7 +560,7 @@ class Dots1ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

560 def compute_logits(

561 self,

562 hidden_states: torch.Tensor,

563 - ) -> Optional[torch.Tensor]:

563 + ) -> torch.Tensor | None:

564 logits = self.logits_processor(self.lm_head, hidden_states)

565 return logits

566

vllm/model_executor/models/dots_ocr.py

+22 -22

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Mapping

4 -from typing import Annotated, Literal, Optional, Union

4 +from typing import Annotated, Literal, TypeAlias

5

6 import torch

7 import torch.nn as nn

@@ -92,7 +92,7 @@ class DotsOCRImageEmbeddingInputs(TensorSchema):

92 image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

93

94

95 -DotsOCRImageInputs = Union[DotsOCRImagePixelInputs, DotsOCRImageEmbeddingInputs]

95 +DotsOCRImageInputs: TypeAlias = DotsOCRImagePixelInputs | DotsOCRImageEmbeddingInputs

96

97

98 class DotsOCRDummyInputsBuilder(Qwen2VLDummyInputsBuilder):

@@ -104,7 +104,7 @@ class DotsOCRDummyInputsBuilder(Qwen2VLDummyInputsBuilder):

104 self,

105 seq_len: int,

106 mm_counts: Mapping[str, int],

107 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

107 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

108 ) -> MultiModalDataDict:

109 num_images = mm_counts.get("image", 0)

110

@@ -134,7 +134,7 @@ class DotsOCRProcessingInfo(Qwen2VLProcessingInfo):

134

135 return config

136

137 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

137 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

138 return {"image": None}

139

140 def get_mm_max_tokens_per_item(

@@ -253,7 +253,7 @@ class DotsVisionAttention(nn.Module):

253 num_heads: int = 16,

254 bias: bool = True,

255 *,

256 - quant_config: Optional[QuantizationConfig] = None,

256 + quant_config: QuantizationConfig | None = None,

257 prefix: str = "",

258 use_data_parallel: bool = False,

259 ) -> None:

@@ -316,10 +316,10 @@ class DotsVisionAttention(nn.Module):

316 self,

317 hidden_states: torch.Tensor,

318 cu_seqlens: torch.Tensor,

319 - rotary_pos_emb: Optional[torch.Tensor] = None,

319 + rotary_pos_emb: torch.Tensor | None = None,

320 *,

321 - max_seqlen: Optional[int] = None,

322 - seqlens: Optional[list[int]] = None,

321 + max_seqlen: int | None = None,

322 + seqlens: list[int] | None = None,

323 ) -> torch.Tensor:

324 # [S, C] -> [S, B=1, C]

325 x = hidden_states.unsqueeze(1)

@@ -394,7 +394,7 @@ class DotsSwiGLUFFN(nn.Module):

394 self,

395 config,

396 *,

397 - quant_config: Optional[QuantizationConfig] = None,

397 + quant_config: QuantizationConfig | None = None,

398 prefix: str = "",

399 use_data_parallel: bool = False,

400 ):

@@ -507,7 +507,7 @@ class DotsVisionBlock(nn.Module):

507 self,

508 config,

509 *,

510 - quant_config: Optional[QuantizationConfig] = None,

510 + quant_config: QuantizationConfig | None = None,

511 prefix: str = "",

512 use_data_parallel: bool = False,

513 ):

@@ -537,8 +537,8 @@ class DotsVisionBlock(nn.Module):

537 *,

538 cu_seqlens: torch.Tensor,

539 rotary_pos_emb: torch.Tensor,

540 - max_seqlen: Optional[int] = None,

541 - seqlens: Optional[list[int]] = None,

540 + max_seqlen: int | None = None,

541 + seqlens: list[int] | None = None,

542 ) -> torch.Tensor:

543 hidden_states = hidden_states + self.attn(

544 self.norm1(hidden_states),

@@ -555,10 +555,10 @@ class DotsVisionTransformer(nn.Module):

555 def __init__(

556 self,

557 config: DotsVisionConfig,

558 - quant_config: Optional[QuantizationConfig] = None,

558 + quant_config: QuantizationConfig | None = None,

559 *,

560 - num_hidden_layers_override: Optional[int] = None,

561 - require_post_norm: Optional[bool] = None,

560 + num_hidden_layers_override: int | None = None,

561 + require_post_norm: bool | None = None,

562 prefix: str = "",

563 use_data_parallel: bool = False,

564 ) -> None:

@@ -653,7 +653,7 @@ class DotsVisionTransformer(nn.Module):

653

654 def compute_attn_mask_seqlen(

655 self, cu_seqlens: torch.Tensor

656 - ) -> tuple[Optional[int], Optional[list[int]]]:

656 + ) -> tuple[int | None, list[int] | None]:

657 max_seqlen, seqlens = None, None

658 if (

659 self.attn_backend == _Backend.FLASH_ATTN

@@ -734,7 +734,7 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

734 supports_encoder_tp_data = True

735

736 @classmethod

737 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

737 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

738 if modality.startswith("image"):

739 return "<|img|><|imgpad|><|endofimg|>"

740

@@ -765,7 +765,7 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

765

766 def _parse_and_validate_image_input(

767 self, **kwargs: object

768 - ) -> Optional[DotsOCRImageInputs]:

768 + ) -> DotsOCRImageInputs | None:

769 pixel_values = kwargs.pop("pixel_values", None)

770 image_embeds = kwargs.pop("image_embeds", None)

771 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -834,10 +834,10 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

834 self,

835 input_ids: torch.Tensor,

836 positions: torch.Tensor,

837 - intermediate_tensors: Optional[IntermediateTensors] = None,

838 - inputs_embeds: Optional[torch.Tensor] = None,

837 + intermediate_tensors: IntermediateTensors | None = None,

838 + inputs_embeds: torch.Tensor | None = None,

839 **kwargs,

840 - ) -> Union[torch.Tensor, IntermediateTensors]:

840 + ) -> torch.Tensor | IntermediateTensors:

841 if intermediate_tensors is not None:

842 inputs_embeds = None

843 elif inputs_embeds is None:

@@ -861,7 +861,7 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

861 def compute_logits(

862 self,

863 hidden_states: torch.Tensor,

864 - ) -> Optional[torch.Tensor]:

864 + ) -> torch.Tensor | None:

865 return self.language_model.compute_logits(hidden_states)

866

867 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/ernie45_moe.py

+17 -17

@@ -25,7 +25,7 @@

25

26 from collections.abc import Iterable

27 from itertools import islice

28 -from typing import Any, Optional, Union

28 +from typing import Any

29

30 import torch

31 from torch import nn

@@ -83,7 +83,7 @@ class Ernie4_5_MoeMLP(nn.Module):

83 intermediate_size: int,

84 hidden_act: str,

85 use_bias: bool = False,

86 - quant_config: Optional[QuantizationConfig] = None,

86 + quant_config: QuantizationConfig | None = None,

87 reduce_results: bool = True,

88 prefix: str = "",

89 ) -> None:

@@ -120,7 +120,7 @@ class Ernie4_5_MoeMoE(nn.Module):

120 def __init__(

121 self,

122 config: PretrainedConfig,

123 - quant_config: Optional[QuantizationConfig] = None,

123 + quant_config: QuantizationConfig | None = None,

124 prefix: str = "",

125 enable_eplb: bool = False,

126 ):

@@ -229,14 +229,14 @@ class Ernie4_5_MoeAttention(nn.Module):

229 hidden_size: int,

230 num_heads: int,

231 num_kv_heads: int,

232 - head_dim: Optional[int] = None,

232 + head_dim: int | None = None,

233 rope_theta: float = 500000,

234 - rope_scaling: Optional[dict[str, Any]] = None,

234 + rope_scaling: dict[str, Any] | None = None,

235 max_position_embeddings: int = 131072,

236 rms_norm_eps: float = 1e-05,

237 qkv_bias: bool = False,

238 - cache_config: Optional[CacheConfig] = None,

239 - quant_config: Optional[QuantizationConfig] = None,

238 + cache_config: CacheConfig | None = None,

239 + quant_config: QuantizationConfig | None = None,

240 prefix: str = "",

241 ) -> None:

242 super().__init__()

@@ -323,8 +323,8 @@ class Ernie4_5_MoeDecoderLayer(nn.Module):

323 def __init__(

324 self,

325 config: PretrainedConfig,

326 - cache_config: Optional[CacheConfig] = None,

327 - quant_config: Optional[QuantizationConfig] = None,

326 + cache_config: CacheConfig | None = None,

327 + quant_config: QuantizationConfig | None = None,

328 prefix: str = "",

329 enable_eplb: bool = False,

330 ) -> None:

@@ -391,7 +391,7 @@ class Ernie4_5_MoeDecoderLayer(nn.Module):

391 self,

392 positions: torch.Tensor,

393 hidden_states: torch.Tensor,

394 - residual: Optional[torch.Tensor],

394 + residual: torch.Tensor | None,

395 ) -> torch.Tensor:

396 # Self Attention

397 if residual is None:

@@ -467,9 +467,9 @@ class Ernie4_5_MoeModel(nn.Module):

467 self,

468 input_ids: torch.Tensor,

469 positions: torch.Tensor,

470 - intermediate_tensors: Optional[IntermediateTensors] = None,

471 - inputs_embeds: Optional[torch.Tensor] = None,

472 - ) -> Union[torch.Tensor, IntermediateTensors]:

470 + intermediate_tensors: IntermediateTensors | None = None,

471 + inputs_embeds: torch.Tensor | None = None,

472 + ) -> torch.Tensor | IntermediateTensors:

473 if get_pp_group().is_first_rank:

474 if inputs_embeds is not None:

475 hidden_states = inputs_embeds

@@ -737,9 +737,9 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe

737 self,

738 input_ids: torch.Tensor,

739 positions: torch.Tensor,

740 - intermediate_tensors: Optional[IntermediateTensors] = None,

741 - inputs_embeds: Optional[torch.Tensor] = None,

742 - ) -> Union[torch.Tensor, IntermediateTensors]:

740 + intermediate_tensors: IntermediateTensors | None = None,

741 + inputs_embeds: torch.Tensor | None = None,

742 + ) -> torch.Tensor | IntermediateTensors:

743 hidden_states = self.model(

744 input_ids, positions, intermediate_tensors, inputs_embeds

745 )

@@ -748,7 +748,7 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe

748 def compute_logits(

749 self,

750 hidden_states: torch.Tensor,

751 - ) -> Optional[torch.Tensor]:

751 + ) -> torch.Tensor | None:

752 logits = self.logits_processor(self.lm_head, hidden_states)

753 return logits

754

vllm/model_executor/models/ernie45_vl.py

+34 -34

@@ -25,9 +25,9 @@

25

26 import itertools

27 import math

28 -from collections.abc import Iterable, Mapping, Sequence

28 +from collections.abc import Callable, Iterable, Mapping, Sequence

29 from functools import partial

30 -from typing import Annotated, Any, Callable, Literal, Optional, Union

30 +from typing import Annotated, Any, Literal

31

32 import numpy as np

33 import torch

@@ -162,7 +162,7 @@ class Ernie4_5_VisionAttention(nn.Module):

162 embed_dim: int,

163 num_heads: int,

164 projection_size: int,

165 - quant_config: Optional[QuantizationConfig] = None,

165 + quant_config: QuantizationConfig | None = None,

166 prefix: str = "",

167 ) -> None:

168 super().__init__()

@@ -254,8 +254,8 @@ class Ernie4_5_VisionAttention(nn.Module):

254 x: torch.Tensor,

255 cu_seqlens: torch.Tensor,

256 rotary_pos_emb: torch.Tensor,

257 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

258 - seqlens: Optional[list[int]] = None, # Only used for xFormers

257 + max_seqlen: int | None = None, # Only used for Flash Attention

258 + seqlens: list[int] | None = None, # Only used for xFormers

259 ) -> torch.Tensor:

260 # [s, b, c] --> [s, b, head * 3 * head_dim]

261 x, _ = self.qkv(x)

@@ -332,7 +332,7 @@ class Ernie4_5_VisionMLP(nn.Module):

332 in_features: int,

333 hidden_features: int,

334 act_layer: type[nn.Module] = QuickGELU,

335 - quant_config: Optional[QuantizationConfig] = None,

335 + quant_config: QuantizationConfig | None = None,

336 prefix: str = "",

337 ):

338 super().__init__()

@@ -364,8 +364,8 @@ class Ernie4_5_VisionBlock(nn.Module):

364 num_heads: int,

365 mlp_ratio: float,

366 act_layer: type[nn.Module] = QuickGELU,

367 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

368 - quant_config: Optional[QuantizationConfig] = None,

367 + norm_layer: Callable[[int], nn.Module] | None = None,

368 + quant_config: QuantizationConfig | None = None,

369 prefix: str = "",

370 ) -> None:

371 super().__init__()

@@ -397,8 +397,8 @@ class Ernie4_5_VisionBlock(nn.Module):

397 hidden_states: torch.Tensor,

398 cu_seqlens: torch.Tensor,

399 rotary_pos_emb: torch.Tensor,

400 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

401 - seqlens: Optional[list[int]] = None, # Only used for xFormers

400 + max_seqlen: int | None = None, # Only used for Flash Attention

401 + seqlens: list[int] | None = None, # Only used for xFormers

402 ) -> torch.Tensor:

403 hidden_states = hidden_states + self.attn(

404 self.norm1(hidden_states),

@@ -456,7 +456,7 @@ class Ernie4_5_VisionTransformer(nn.Module):

456 self,

457 vision_config,

458 norm_eps: float = 1e-6,

459 - quant_config: Optional[QuantizationConfig] = None,

459 + quant_config: QuantizationConfig | None = None,

460 prefix: str = "",

461 ) -> None:

462 super().__init__()

@@ -553,7 +553,7 @@ class Ernie4_5_VisionTransformer(nn.Module):

553

554 def compute_attn_mask_seqlen(

555 self, cu_seqlens: torch.Tensor

556 - ) -> tuple[Optional[int], Optional[list[int]]]:

556 + ) -> tuple[int | None, list[int] | None]:

557 max_seqlen, seqlens = None, None

558 if (

559 self.attn_backend == _Backend.FLASH_ATTN

@@ -659,15 +659,15 @@ Ernie4_5_VLVideoInputs = Ernie4_5_VLVideoPixelInputs

659 # === Vision Processor === #

660

661

662 -def round_by_factor(number: Union[int, float], factor: int) -> int:

662 +def round_by_factor(number: int | float, factor: int) -> int:

663 return round(number / factor) * factor

664

665

666 -def ceil_by_factor(number: Union[int, float], factor: int) -> int:

666 +def ceil_by_factor(number: int | float, factor: int) -> int:

667 return math.ceil(number / factor) * factor

668

669

670 -def floor_by_factor(number: Union[int, float], factor: int) -> int:

670 +def floor_by_factor(number: int | float, factor: int) -> int:

671 return math.floor(number / factor) * factor

672

673

@@ -901,7 +901,7 @@ class Ernie4_5_VLProcessingInfo(BaseProcessingInfo):

901 def get_image_processor(self, **kwargs: object):

902 return self.get_hf_processor(**kwargs).image_processor

903

904 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

904 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

905 return {"image": None, "video": None}

906

907 def get_mm_max_tokens_per_item(

@@ -920,7 +920,7 @@ class Ernie4_5_VLProcessingInfo(BaseProcessingInfo):

920 image_height: int,

921 num_frames: int = 1,

922 do_resize: bool = True,

923 - image_processor: Optional[Any],

923 + image_processor: Any | None,

924 ) -> tuple[ImageSize, int]:

925 if image_processor is None:

926 image_processor = self.get_image_processor()

@@ -957,7 +957,7 @@ class Ernie4_5_VLProcessingInfo(BaseProcessingInfo):

957 *,

958 image_width: int,

959 image_height: int,

960 - image_processor: Optional[Any],

960 + image_processor: Any | None,

961 ) -> int:

962 _, num_image_tokens = self._get_vision_info(

963 image_width=image_width,

@@ -972,7 +972,7 @@ class Ernie4_5_VLProcessingInfo(BaseProcessingInfo):

972 image_width: int,

973 image_height: int,

974 num_frames: int,

975 - image_processor: Optional[Any],

975 + image_processor: Any | None,

976 ) -> int:

977 _, num_video_tokens = self._get_vision_info(

978 image_width=image_width,

@@ -1237,7 +1237,7 @@ class Ernie4_5_VLDummyInputsBuilder(BaseDummyInputsBuilder[Ernie4_5_VLProcessing

1237 self,

1238 seq_len: int,

1239 mm_counts: Mapping[str, int],

1240 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

1240 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

1241 ) -> MultiModalDataDict:

1242 num_images = mm_counts.get("image", 0)

1243 num_videos = mm_counts.get("video", 0)

@@ -1310,7 +1310,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1310 )

1311

1312 @classmethod

1313 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1313 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1314 if modality.startswith("image"):

1315 return "<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"

1316 if modality.startswith("video"):

@@ -1356,7 +1356,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1356 def compute_logits(

1357 self,

1358 hidden_states: torch.Tensor,

1359 - ) -> Optional[torch.Tensor]:

1359 + ) -> torch.Tensor | None:

1360 """compute logits"""

1361 return self.language_model.compute_logits(hidden_states)

1362

@@ -1395,12 +1395,12 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1395 cls,

1396 input_tokens: list[int],

1397 hf_config: PretrainedConfig,

1398 - image_grid_thw: Union[list[list[int]], torch.Tensor],

1399 - video_grid_thw: Union[list[list[int]], torch.Tensor],

1398 + image_grid_thw: list[list[int]] | torch.Tensor,

1399 + video_grid_thw: list[list[int]] | torch.Tensor,

1400 context_len: int = 0,

1401 - seq_len: Optional[int] = None,

1402 - second_per_grid_ts: Optional[list[float]] = None,

1403 - audio_feature_lengths: Optional[torch.Tensor] = None,

1401 + seq_len: int | None = None,

1402 + second_per_grid_ts: list[float] | None = None,

1403 + audio_feature_lengths: torch.Tensor | None = None,

1404 use_audio_in_video: bool = False,

1405 ) -> tuple[torch.Tensor, int]:

1406 """Get mrope input positions and delta value for Ernie VL."""

@@ -1540,7 +1540,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1540

1541 def _parse_and_validate_image_input(

1542 self, **kwargs: object

1543 - ) -> Optional[Ernie4_5_VLImageInputs]:

1543 + ) -> Ernie4_5_VLImageInputs | None:

1544 pixel_values = kwargs.pop("pixel_values", None)

1545 image_grid_thw = kwargs.pop("image_grid_thw", None)

1546

@@ -1556,7 +1556,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1556

1557 def _parse_and_validate_video_input(

1558 self, **kwargs: object

1559 - ) -> Optional[Ernie4_5_VLVideoInputs]:

1559 + ) -> Ernie4_5_VLVideoInputs | None:

1560 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1561 video_grid_thw = kwargs.pop("video_grid_thw", None)

1562

@@ -1631,7 +1631,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1631

1632 def get_multimodal_embeddings(

1633 self, **kwargs: object

1634 - ) -> Optional[MultiModalEmbeddings]:

1634 + ) -> MultiModalEmbeddings | None:

1635 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

1636 if not modalities:

1637 return None

@@ -1657,9 +1657,9 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1657 def get_input_embeddings(

1658 self,

1659 input_ids: torch.Tensor,

1660 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1660 + multimodal_embeddings: MultiModalEmbeddings | None = None,

1661 *,

1662 - is_multimodal: Optional[torch.Tensor] = None,

1662 + is_multimodal: torch.Tensor | None = None,

1663 handle_oov_mm_token: bool = False,

1664 ) -> torch.Tensor:

1665 if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:

@@ -1680,8 +1680,8 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1680 self,

1681 input_ids: torch.Tensor,

1682 positions: torch.Tensor,

1683 - intermediate_tensors: Optional[IntermediateTensors] = None,

1684 - inputs_embeds: Optional[torch.Tensor] = None,

1683 + intermediate_tensors: IntermediateTensors | None = None,

1684 + inputs_embeds: torch.Tensor | None = None,

1685 **kwargs,

1686 ):

1687 forward_kwargs = {

vllm/model_executor/models/ernie45_vl_moe.py

+19 -19

@@ -25,7 +25,7 @@

25

26 from collections.abc import Iterable

27 from itertools import islice

28 -from typing import Any, Optional, Union

28 +from typing import Any

29

30 import torch

31 from torch import nn

@@ -74,7 +74,7 @@ logger = init_logger(__name__)

74

75

76 class Ernie4_5_VLMoeMLP(Ernie4_5_MoeMLP):

77 - def __init__(self, shared_experts: Optional[torch.nn.Module] = None, **kwargs):

77 + def __init__(self, shared_experts: torch.nn.Module | None = None, **kwargs):

78 super().__init__(**kwargs)

79 self.shared_experts = shared_experts

80

@@ -91,15 +91,15 @@ class Ernie4_5_VLMoeAttention(nn.Module):

91 hidden_size: int,

92 num_heads: int,

93 num_kv_heads: int,

94 - head_dim: Optional[int] = None,

94 + head_dim: int | None = None,

95 rope_theta: float = 500000,

96 - rope_scaling: Optional[dict[str, Any]] = None,

96 + rope_scaling: dict[str, Any] | None = None,

97 freq_allocation: int = 20,

98 max_position_embeddings: int = 131072,

99 rms_norm_eps: float = 1e-05,

100 qkv_bias: bool = False,

101 - cache_config: Optional[CacheConfig] = None,

102 - quant_config: Optional[QuantizationConfig] = None,

101 + cache_config: CacheConfig | None = None,

102 + quant_config: QuantizationConfig | None = None,

103 prefix: str = "",

104 ) -> None:

105 super().__init__()

@@ -192,7 +192,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):

192 def __init__(

193 self,

194 config: PretrainedConfig,

195 - quant_config: Optional[QuantizationConfig] = None,

195 + quant_config: QuantizationConfig | None = None,

196 prefix: str = "",

197 ):

198 super().__init__()

@@ -390,8 +390,8 @@ class Ernie4_5_VLMoeDecoderLayer(nn.Module):

390 def __init__(

391 self,

392 config: PretrainedConfig,

393 - cache_config: Optional[CacheConfig] = None,

394 - quant_config: Optional[QuantizationConfig] = None,

393 + cache_config: CacheConfig | None = None,

394 + quant_config: QuantizationConfig | None = None,

395 prefix: str = "",

396 ) -> None:

397 super().__init__()

@@ -463,8 +463,8 @@ class Ernie4_5_VLMoeDecoderLayer(nn.Module):

463 self,

464 positions: torch.Tensor,

465 hidden_states: torch.Tensor,

466 - residual: Optional[torch.Tensor],

467 - visual_token_mask: Optional[torch.Tensor],

466 + residual: torch.Tensor | None,

467 + visual_token_mask: torch.Tensor | None,

468 **kwargs: object,

469 ) -> torch.Tensor:

470 # Self Attention

@@ -551,11 +551,11 @@ class Ernie4_5_VLMoeModel(nn.Module):

551 self,

552 input_ids: torch.Tensor,

553 positions: torch.Tensor,

554 - intermediate_tensors: Optional[IntermediateTensors] = None,

555 - inputs_embeds: Optional[torch.Tensor] = None,

556 - visual_token_mask: Optional[torch.Tensor] = None,

554 + intermediate_tensors: IntermediateTensors | None = None,

555 + inputs_embeds: torch.Tensor | None = None,

556 + visual_token_mask: torch.Tensor | None = None,

557 **kwargs: object,

558 - ) -> Union[torch.Tensor, IntermediateTensors]:

558 + ) -> torch.Tensor | IntermediateTensors:

559 if get_pp_group().is_first_rank:

560 if inputs_embeds is not None:

561 hidden_states = inputs_embeds

@@ -632,10 +632,10 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):

632 self,

633 input_ids: torch.Tensor,

634 positions: torch.Tensor,

635 - intermediate_tensors: Optional[IntermediateTensors] = None,

636 - inputs_embeds: Optional[torch.Tensor] = None,

635 + intermediate_tensors: IntermediateTensors | None = None,

636 + inputs_embeds: torch.Tensor | None = None,

637 **kwargs: object,

638 - ) -> Union[torch.Tensor, IntermediateTensors]:

638 + ) -> torch.Tensor | IntermediateTensors:

639 hidden_states = self.model(

640 input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs

641 )

@@ -644,7 +644,7 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):

644 def compute_logits(

645 self,

646 hidden_states: torch.Tensor,

647 - ) -> Optional[torch.Tensor]:

647 + ) -> torch.Tensor | None:

648 logits = self.logits_processor(self.lm_head, hidden_states)

649 return logits

650

vllm/model_executor/models/ernie_mtp.py

+4 -5

@@ -24,7 +24,6 @@

24 """Inference-only Ernie-MTP model."""

25

26 from collections.abc import Iterable

27 -from typing import Optional

27

28 import torch

29 import torch.nn as nn

@@ -121,7 +120,7 @@ class ErnieMultiTokenPredictor(nn.Module):

120 input_ids: torch.Tensor,

121 positions: torch.Tensor,

122 previous_hidden_states: torch.Tensor,

124 - inputs_embeds: Optional[torch.Tensor] = None,

123 + inputs_embeds: torch.Tensor | None = None,

124 spec_step_idx: int = 0,

125 ) -> torch.Tensor:

126 if inputs_embeds is None:

@@ -169,8 +168,8 @@ class ErnieMTP(nn.Module, SupportsPP):

168 input_ids: torch.Tensor,

169 positions: torch.Tensor,

170 hidden_states: torch.Tensor,

172 - intermediate_tensors: Optional[IntermediateTensors] = None,

173 - inputs_embeds: Optional[torch.Tensor] = None,

171 + intermediate_tensors: IntermediateTensors | None = None,

172 + inputs_embeds: torch.Tensor | None = None,

173 spec_step_idx: int = 0,

174 ) -> torch.Tensor:

175 assert spec_step_idx == 0, "ernie_mtp only support predict one token"

@@ -183,7 +182,7 @@ class ErnieMTP(nn.Module, SupportsPP):

182 self,

183 hidden_states: torch.Tensor,

184 spec_step_idx: int = 0,

186 - ) -> Optional[torch.Tensor]:

185 + ) -> torch.Tensor | None:

186 return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)

187

188 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/exaone.py

+19 -19

@@ -27,7 +27,7 @@

27

28 from collections.abc import Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import torch

33 from torch import nn

@@ -75,7 +75,7 @@ class ExaoneGatedMLP(nn.Module):

75 hidden_size: int,

76 intermediate_size: int,

77 hidden_act: str,

78 - quant_config: Optional[QuantizationConfig] = None,

78 + quant_config: QuantizationConfig | None = None,

79 bias: bool = False,

80 prefix: str = "",

81 ) -> None:

@@ -115,11 +115,11 @@ class ExaoneAttention(nn.Module):

115 num_heads: int,

116 num_kv_heads: int,

117 rope_theta: float = 10000,

118 - rope_scaling: Optional[dict[str, Any]] = None,

118 + rope_scaling: dict[str, Any] | None = None,

119 max_position_embeddings: int = 8192,

120 - quant_config: Optional[QuantizationConfig] = None,

120 + quant_config: QuantizationConfig | None = None,

121 bias: bool = False,

122 - cache_config: Optional[CacheConfig] = None,

122 + cache_config: CacheConfig | None = None,

123 prefix: str = "",

124 ) -> None:

125 super().__init__()

@@ -209,11 +209,11 @@ class ExaoneBlockAttention(nn.Module):

209 num_heads: int,

210 num_kv_heads: int,

211 rope_theta: float = 10000,

212 - rope_scaling: Optional[dict[str, Any]] = None,

212 + rope_scaling: dict[str, Any] | None = None,

213 max_position_embeddings: int = 8192,

214 - quant_config: Optional[QuantizationConfig] = None,

214 + quant_config: QuantizationConfig | None = None,

215 bias: bool = False,

216 - cache_config: Optional[CacheConfig] = None,

216 + cache_config: CacheConfig | None = None,

217 prefix: str = "",

218 ) -> None:

219 super().__init__()

@@ -246,8 +246,8 @@ class ExaoneDecoderLayer(nn.Module):

246 def __init__(

247 self,

248 config: PretrainedConfig,

249 - cache_config: Optional[CacheConfig] = None,

250 - quant_config: Optional[QuantizationConfig] = None,

249 + cache_config: CacheConfig | None = None,

250 + quant_config: QuantizationConfig | None = None,

251 prefix: str = "",

252 ) -> None:

253 super().__init__()

@@ -296,7 +296,7 @@ class ExaoneDecoderLayer(nn.Module):

296 self,

297 positions: torch.Tensor,

298 hidden_states: torch.Tensor,

299 - residual: Optional[torch.Tensor],

299 + residual: torch.Tensor | None,

300 ) -> tuple[torch.Tensor, torch.Tensor]:

301 # Self Attention

302 if residual is None:

@@ -369,11 +369,11 @@ class ExaoneModel(nn.Module):

369

370 def forward(

371 self,

372 - input_ids: Optional[torch.Tensor],

372 + input_ids: torch.Tensor | None,

373 positions: torch.Tensor,

374 - intermediate_tensors: Optional[IntermediateTensors],

375 - inputs_embeds: Optional[torch.Tensor] = None,

376 - ) -> Union[torch.Tensor, IntermediateTensors]:

374 + intermediate_tensors: IntermediateTensors | None,

375 + inputs_embeds: torch.Tensor | None = None,

376 + ) -> torch.Tensor | IntermediateTensors:

377 if get_pp_group().is_first_rank:

378 if inputs_embeds is not None:

379 hidden_states = inputs_embeds

@@ -536,9 +536,9 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

536 self,

537 input_ids: torch.Tensor,

538 positions: torch.Tensor,

539 - intermediate_tensors: Optional[IntermediateTensors] = None,

540 - inputs_embeds: Optional[torch.Tensor] = None,

541 - ) -> Union[torch.Tensor, IntermediateTensors]:

539 + intermediate_tensors: IntermediateTensors | None = None,

540 + inputs_embeds: torch.Tensor | None = None,

541 + ) -> torch.Tensor | IntermediateTensors:

542 model_output = self.transformer(

543 input_ids, positions, intermediate_tensors, inputs_embeds

544 )

@@ -547,7 +547,7 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

547 def compute_logits(

548 self,

549 hidden_states: torch.Tensor,

550 - ) -> Optional[torch.Tensor]:

550 + ) -> torch.Tensor | None:

551 logits = self.logits_processor(self.lm_head, hidden_states)

552 return logits

553

vllm/model_executor/models/exaone4.py

+16 -16

@@ -23,7 +23,7 @@

23

24 from collections.abc import Iterable

25 from itertools import islice

26 -from typing import Any, Optional, Union

26 +from typing import Any

27

28 import torch

29 from torch import nn

@@ -72,7 +72,7 @@ class Exaone4GatedMLP(nn.Module):

72 hidden_size: int,

73 intermediate_size: int,

74 hidden_act: str,

75 - quant_config: Optional[QuantizationConfig] = None,

75 + quant_config: QuantizationConfig | None = None,

76 bias: bool = False,

77 prefix: str = "",

78 ) -> None:

@@ -112,11 +112,11 @@ class Exaone4Attention(nn.Module):

112 num_heads: int,

113 num_kv_heads: int,

114 rope_theta: float = 1000000,

115 - rope_scaling: Optional[dict[str, Any]] = None,

115 + rope_scaling: dict[str, Any] | None = None,

116 max_position_embeddings: int = 8192,

117 - quant_config: Optional[QuantizationConfig] = None,

117 + quant_config: QuantizationConfig | None = None,

118 bias: bool = False,

119 - cache_config: Optional[CacheConfig] = None,

119 + cache_config: CacheConfig | None = None,

120 prefix: str = "",

121 ) -> None:

122 super().__init__()

@@ -222,8 +222,8 @@ class Exaone4DecoderLayer(nn.Module):

222 def __init__(

223 self,

224 config: Exaone4Config,

225 - cache_config: Optional[CacheConfig] = None,

226 - quant_config: Optional[QuantizationConfig] = None,

225 + cache_config: CacheConfig | None = None,

226 + quant_config: QuantizationConfig | None = None,

227 prefix: str = "",

228 ) -> None:

229 super().__init__()

@@ -277,7 +277,7 @@ class Exaone4DecoderLayer(nn.Module):

277 self,

278 positions: torch.Tensor,

279 hidden_states: torch.Tensor,

280 - residual: Optional[torch.Tensor],

280 + residual: torch.Tensor | None,

281 ) -> tuple[torch.Tensor, torch.Tensor]:

282 residual = hidden_states

283

@@ -356,11 +356,11 @@ class Exaone4Model(nn.Module):

356

357 def forward(

358 self,

359 - input_ids: Optional[torch.Tensor],

359 + input_ids: torch.Tensor | None,

360 positions: torch.Tensor,

361 - intermediate_tensors: Optional[IntermediateTensors],

362 - inputs_embeds: Optional[torch.Tensor] = None,

363 - ) -> Union[torch.Tensor, IntermediateTensors]:

361 + intermediate_tensors: IntermediateTensors | None,

362 + inputs_embeds: torch.Tensor | None = None,

363 + ) -> torch.Tensor | IntermediateTensors:

364 if get_pp_group().is_first_rank:

365 if inputs_embeds is not None:

366 hidden_states = inputs_embeds

@@ -523,9 +523,9 @@ class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

523 self,

524 input_ids: torch.Tensor,

525 positions: torch.Tensor,

526 - intermediate_tensors: Optional[IntermediateTensors] = None,

527 - inputs_embeds: Optional[torch.Tensor] = None,

528 - ) -> Union[torch.Tensor, IntermediateTensors]:

526 + intermediate_tensors: IntermediateTensors | None = None,

527 + inputs_embeds: torch.Tensor | None = None,

528 + ) -> torch.Tensor | IntermediateTensors:

529 model_output = self.model(

530 input_ids, positions, intermediate_tensors, inputs_embeds

531 )

@@ -534,7 +534,7 @@ class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

534 def compute_logits(

535 self,

536 hidden_states: torch.Tensor,

537 - ) -> Optional[torch.Tensor]:

537 + ) -> torch.Tensor | None:

538 logits = self.logits_processor(self.lm_head, hidden_states)

539 return logits

540

vllm/model_executor/models/falcon.py

+13 -13

@@ -23,7 +23,7 @@

23 import math

24 from collections.abc import Iterable

25 from itertools import islice

26 -from typing import Optional, Union

26 +from typing import TypeAlias

27

28 import torch

29 from torch import nn

@@ -65,7 +65,7 @@ from .utils import (

65 maybe_prefix,

66 )

67

68 -FalconConfig = Union[HF_FalconConfig, RWConfig]

68 +FalconConfig: TypeAlias = HF_FalconConfig | RWConfig

69

70

71 def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:

@@ -95,8 +95,8 @@ class FalconAttention(nn.Module):

95 def __init__(

96 self,

97 config: FalconConfig,

98 - cache_config: Optional[CacheConfig] = None,

99 - quant_config: Optional[QuantizationConfig] = None,

98 + cache_config: CacheConfig | None = None,

99 + quant_config: QuantizationConfig | None = None,

100 prefix: str = "",

101 ):

102 super().__init__()

@@ -226,7 +226,7 @@ class FalconMLP(nn.Module):

226 def __init__(

227 self,

228 config: FalconConfig,

229 - quant_config: Optional[QuantizationConfig] = None,

229 + quant_config: QuantizationConfig | None = None,

230 ):

231 super().__init__()

232 hidden_size = config.hidden_size

@@ -265,8 +265,8 @@ class FalconDecoderLayer(nn.Module):

265 def __init__(

266 self,

267 config: FalconConfig,

268 - cache_config: Optional[CacheConfig] = None,

269 - quant_config: Optional[QuantizationConfig] = None,

268 + cache_config: CacheConfig | None = None,

269 + quant_config: QuantizationConfig | None = None,

270 prefix: str = "",

271 ):

272 super().__init__()

@@ -401,9 +401,9 @@ class FalconModel(nn.Module):

401 self,

402 input_ids: torch.Tensor,

403 positions: torch.Tensor,

404 - intermediate_tensors: Optional[IntermediateTensors],

405 - inputs_embeds: Optional[torch.Tensor] = None,

406 - ) -> Union[torch.Tensor, IntermediateTensors]:

404 + intermediate_tensors: IntermediateTensors | None,

405 + inputs_embeds: torch.Tensor | None = None,

406 + ) -> torch.Tensor | IntermediateTensors:

407 if get_pp_group().is_first_rank:

408 if inputs_embeds is not None:

409 hidden_states = inputs_embeds

@@ -517,8 +517,8 @@ class FalconForCausalLM(nn.Module, SupportsPP):

517 self,

518 input_ids: torch.LongTensor,

519 positions: torch.Tensor,

520 - intermediate_tensors: Optional[IntermediateTensors] = None,

521 - inputs_embeds: Optional[torch.Tensor] = None,

520 + intermediate_tensors: IntermediateTensors | None = None,

521 + inputs_embeds: torch.Tensor | None = None,

522 ) -> torch.Tensor:

523 hidden_states = self.transformer(

524 input_ids, positions, intermediate_tensors, inputs_embeds

@@ -528,7 +528,7 @@ class FalconForCausalLM(nn.Module, SupportsPP):

528 def compute_logits(

529 self,

530 hidden_states: torch.Tensor,

531 - ) -> Optional[torch.Tensor]:

531 + ) -> torch.Tensor | None:

532 logits = self.logits_processor(self.lm_head, hidden_states)

533 return logits

534

vllm/model_executor/models/falcon_h1.py

+16 -17

@@ -4,7 +4,6 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import Optional

7

8 import torch

9 from torch import nn

@@ -52,7 +51,7 @@ class FalconH1MLP(nn.Module):

51 def __init__(

52 self,

53 config: FalconH1Config,

55 - quant_config: Optional[QuantizationConfig] = None,

54 + quant_config: QuantizationConfig | None = None,

55 bias: bool = False,

56 ) -> None:

57 super().__init__()

@@ -91,9 +90,9 @@ class FalconH1SSMDecoderLayer(nn.Module):

90 def __init__(

91 self,

92 config: FalconH1Config,

94 - model_config: Optional[ModelConfig] = None,

95 - cache_config: Optional[CacheConfig] = None,

96 - quant_config: Optional[QuantizationConfig] = None,

93 + model_config: ModelConfig | None = None,

94 + cache_config: CacheConfig | None = None,

95 + quant_config: QuantizationConfig | None = None,

96 prefix: str = "",

97 ) -> None:

98 super().__init__()

@@ -188,7 +187,7 @@ class FalconH1SSMDecoderLayer(nn.Module):

187 def forward(

188 self,

189 hidden_states: torch.Tensor,

191 - residual: Optional[torch.Tensor],

190 + residual: torch.Tensor | None,

191 **kwargs,

192 ):

193 output = torch.empty_like(hidden_states)

@@ -204,8 +203,8 @@ class FalconH1AttentionDecoderLayer(nn.Module):

203 def __init__(

204 self,

205 config: FalconH1Config,

207 - cache_config: Optional[CacheConfig] = None,

208 - quant_config: Optional[QuantizationConfig] = None,

206 + cache_config: CacheConfig | None = None,

207 + quant_config: QuantizationConfig | None = None,

208 prefix: str = "",

209 ) -> None:

210 super().__init__()

@@ -301,7 +300,7 @@ class FalconH1AttentionDecoderLayer(nn.Module):

300 self,

301 positions: torch.Tensor,

302 hidden_states: torch.Tensor,

304 - residual: Optional[torch.Tensor],

303 + residual: torch.Tensor | None,

304 **kwargs,

305 ):

306 hidden_states = self.self_attention(

@@ -326,9 +325,9 @@ class FalconH1ParallelHybrid(nn.Module):

325 self,

326 config: FalconH1Config,

327 layer_idx: int,

329 - model_config: Optional[ModelConfig] = None,

330 - cache_config: Optional[CacheConfig] = None,

331 - quant_config: Optional[QuantizationConfig] = None,

328 + model_config: ModelConfig | None = None,

329 + cache_config: CacheConfig | None = None,

330 + quant_config: QuantizationConfig | None = None,

331 prefix: str = "",

332 ) -> None:

333 super().__init__()

@@ -467,8 +466,8 @@ class FalconH1Model(nn.Module):

466 self,

467 input_ids: torch.Tensor,

468 positions: torch.Tensor,

470 - intermediate_tensors: Optional[IntermediateTensors] = None,

471 - inputs_embeds: Optional[torch.Tensor] = None,

469 + intermediate_tensors: IntermediateTensors | None = None,

470 + inputs_embeds: torch.Tensor | None = None,

471 ) -> torch.Tensor:

472 if get_pp_group().is_first_rank:

473 if inputs_embeds is not None:

@@ -610,8 +609,8 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, Is

609 self,

610 input_ids: torch.Tensor,

611 positions: torch.Tensor,

613 - intermediate_tensors: Optional[IntermediateTensors] = None,

614 - inputs_embeds: Optional[torch.Tensor] = None,

612 + intermediate_tensors: IntermediateTensors | None = None,

613 + inputs_embeds: torch.Tensor | None = None,

614 **kwargs,

615 ):

616 hidden_states = self.model(

@@ -626,7 +625,7 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, Is

625 def compute_logits(

626 self,

627 hidden_states: torch.Tensor,

629 - ) -> Optional[torch.Tensor]:

628 + ) -> torch.Tensor | None:

629 logits = self.logits_processor(self.lm_head, hidden_states)

630

631 return logits

vllm/model_executor/models/flex_olmo.py

+2 -4

@@ -14,8 +14,6 @@

14 # limitations under the License.

15 """Inference-only FlexOlmo model compatible with HuggingFace weights."""

16

17 -from typing import Optional

18 -

17 import torch

18 from torch import nn

19

@@ -128,8 +126,8 @@ class FlexOlmoDecoderLayer(nn.Module):

126 self,

127 positions: torch.Tensor,

128 hidden_states: torch.Tensor,

131 - residual: Optional[torch.Tensor],

132 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

129 + residual: torch.Tensor | None,

130 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

131 # Attention block.

132 residual = hidden_states

133 hidden_states = self.self_attn(positions, hidden_states)

vllm/model_executor/models/fuyu.py

+8 -8

@@ -20,7 +20,7 @@

20

21 import math

22 from collections.abc import Iterable, Mapping, Sequence

23 -from typing import Annotated, Literal, Optional

23 +from typing import Annotated, Literal

24

25 import torch

26 import torch.nn as nn

@@ -87,7 +87,7 @@ class FuyuProcessingInfo(BaseProcessingInfo):

87 def get_image_processor(self, **kwargs: object) -> FuyuImageProcessor:

88 return self.get_hf_processor(**kwargs).image_processor

89

90 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

90 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

91 return {"image": 1}

92

93 def get_image_feature_grid_size(

@@ -142,7 +142,7 @@ class FuyuDummyInputsBuilder(BaseDummyInputsBuilder[FuyuProcessingInfo]):

142 self,

143 seq_len: int,

144 mm_counts: Mapping[str, int],

145 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

145 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

146 ) -> MultiModalDataDict:

147 target_width, target_height = self.info.get_image_size_with_most_features()

148 num_images = mm_counts.get("image", 0)

@@ -271,7 +271,7 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

271 )

272

273 @classmethod

274 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

274 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

275 if modality.startswith("image"):

276 return None

277

@@ -305,7 +305,7 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

305

306 def _parse_and_validate_image_input(

307 self, **kwargs: object

308 - ) -> Optional[FuyuImagePatchInputs]:

308 + ) -> FuyuImagePatchInputs | None:

309 image_patches = kwargs.pop("image_patches", None)

310 patches_per_image = kwargs.pop("patches_per_image", None)

311

@@ -344,8 +344,8 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

344 self,

345 input_ids: torch.Tensor,

346 positions: torch.Tensor,

347 - intermediate_tensors: Optional[IntermediateTensors] = None,

348 - inputs_embeds: Optional[torch.Tensor] = None,

347 + intermediate_tensors: IntermediateTensors | None = None,

348 + inputs_embeds: torch.Tensor | None = None,

349 **kwargs: object,

350 ):

351 if intermediate_tensors is not None:

@@ -362,7 +362,7 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

362 def compute_logits(

363 self,

364 hidden_states: torch.Tensor,

365 - ) -> Optional[torch.Tensor]:

365 + ) -> torch.Tensor | None:

366 logits = self.language_model.logits_processor(

367 self.language_model.lm_head, hidden_states

368 )

vllm/model_executor/models/gemma.py

+17 -18

@@ -20,7 +20,6 @@

20 from collections.abc import Iterable

21 from functools import cache

22 from itertools import islice

23 -from typing import Optional, Union

23

24 import torch

25 from torch import nn

@@ -59,8 +58,8 @@ logger = init_logger(__name__)

58

59 @cache

60 def _get_gemma_act_fn(

62 - hidden_act: Optional[str],

63 - hidden_activation: Optional[str],

61 + hidden_act: str | None,

62 + hidden_activation: str | None,

63 ) -> nn.Module:

64 if hidden_activation is None:

65 if hidden_act is not None:

@@ -92,9 +91,9 @@ class GemmaMLP(nn.Module):

91 self,

92 hidden_size: int,

93 intermediate_size: int,

95 - hidden_act: Optional[str] = None,

96 - hidden_activation: Optional[str] = None,

97 - quant_config: Optional[QuantizationConfig] = None,

94 + hidden_act: str | None = None,

95 + hidden_activation: str | None = None,

96 + quant_config: QuantizationConfig | None = None,

97 prefix: str = "",

98 ) -> None:

99 super().__init__()

@@ -130,8 +129,8 @@ class GemmaAttention(nn.Module):

129 head_dim: int,

130 max_position_embeddings: int = 8192,

131 rope_theta: float = 10000,

133 - cache_config: Optional[CacheConfig] = None,

134 - quant_config: Optional[QuantizationConfig] = None,

132 + cache_config: CacheConfig | None = None,

133 + quant_config: QuantizationConfig | None = None,

134 prefix: str = "",

135 ) -> None:

136 super().__init__()

@@ -207,8 +206,8 @@ class GemmaDecoderLayer(nn.Module):

206 def __init__(

207 self,

208 config: GemmaConfig,

210 - cache_config: Optional[CacheConfig] = None,

211 - quant_config: Optional[QuantizationConfig] = None,

209 + cache_config: CacheConfig | None = None,

210 + quant_config: QuantizationConfig | None = None,

211 prefix: str = "",

212 ) -> None:

213 super().__init__()

@@ -241,7 +240,7 @@ class GemmaDecoderLayer(nn.Module):

240 self,

241 positions: torch.Tensor,

242 hidden_states: torch.Tensor,

244 - residual: Optional[torch.Tensor],

243 + residual: torch.Tensor | None,

244 ) -> tuple[torch.Tensor, torch.Tensor]:

245 # Self Attention

246 if residual is None:

@@ -301,9 +300,9 @@ class GemmaModel(nn.Module):

300 self,

301 input_ids: torch.Tensor,

302 positions: torch.Tensor,

304 - intermediate_tensors: Optional[IntermediateTensors],

305 - inputs_embeds: Optional[torch.Tensor] = None,

306 - ) -> Union[torch.Tensor, IntermediateTensors]:

303 + intermediate_tensors: IntermediateTensors | None,

304 + inputs_embeds: torch.Tensor | None = None,

305 + ) -> torch.Tensor | IntermediateTensors:

306 if get_pp_group().is_first_rank:

307 if inputs_embeds is not None:

308 hidden_states = inputs_embeds

@@ -406,9 +405,9 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

405 self,

406 input_ids: torch.Tensor,

407 positions: torch.Tensor,

409 - intermediate_tensors: Optional[IntermediateTensors] = None,

410 - inputs_embeds: Optional[torch.Tensor] = None,

411 - ) -> Union[torch.Tensor, IntermediateTensors]:

408 + intermediate_tensors: IntermediateTensors | None = None,

409 + inputs_embeds: torch.Tensor | None = None,

410 + ) -> torch.Tensor | IntermediateTensors:

411 hidden_states = self.model(

412 input_ids, positions, intermediate_tensors, inputs_embeds

413 )

@@ -417,7 +416,7 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

416 def compute_logits(

417 self,

418 hidden_states: torch.Tensor,

420 - ) -> Optional[torch.Tensor]:

419 + ) -> torch.Tensor | None:

420 logits = self.logits_processor(self.model.embed_tokens, hidden_states)

421 return logits

422

vllm/model_executor/models/gemma2.py

+15 -16

@@ -18,7 +18,6 @@

18 # limitations under the License.

19 from collections.abc import Iterable

20 from itertools import islice

21 -from typing import Optional, Union

21

22 import torch

23 from torch import nn

@@ -66,7 +65,7 @@ class Gemma2MLP(nn.Module):

65 intermediate_size: int,

66 hidden_act: str,

67 hidden_activation: str,

69 - quant_config: Optional[QuantizationConfig] = None,

68 + quant_config: QuantizationConfig | None = None,

69 ) -> None:

70 super().__init__()

71 self.gate_up_proj = MergedColumnParallelLinear(

@@ -100,9 +99,9 @@ class Gemma2Attention(nn.Module):

99 head_dim: int,

100 max_position_embeddings: int,

101 rope_theta: float,

103 - cache_config: Optional[CacheConfig] = None,

104 - quant_config: Optional[QuantizationConfig] = None,

105 - attn_logits_soft_cap: Optional[float] = None,

102 + cache_config: CacheConfig | None = None,

103 + quant_config: QuantizationConfig | None = None,

104 + attn_logits_soft_cap: float | None = None,

105 prefix: str = "",

106 ) -> None:

107 super().__init__()

@@ -183,8 +182,8 @@ class Gemma2DecoderLayer(nn.Module):

182 def __init__(

183 self,

184 config: Gemma2Config,

186 - cache_config: Optional[CacheConfig] = None,

187 - quant_config: Optional[QuantizationConfig] = None,

185 + cache_config: CacheConfig | None = None,

186 + quant_config: QuantizationConfig | None = None,

187 prefix: str = "",

188 ) -> None:

189 super().__init__()

@@ -225,7 +224,7 @@ class Gemma2DecoderLayer(nn.Module):

224 self,

225 positions: torch.Tensor,

226 hidden_states: torch.Tensor,

228 - residual: Optional[torch.Tensor],

227 + residual: torch.Tensor | None,

228 ) -> tuple[torch.Tensor, torch.Tensor]:

229 if residual is None:

230 residual = hidden_states

@@ -284,11 +283,11 @@ class Gemma2Model(nn.Module):

283

284 def forward(

285 self,

287 - input_ids: Optional[torch.Tensor],

286 + input_ids: torch.Tensor | None,

287 positions: torch.Tensor,

289 - intermediate_tensors: Optional[IntermediateTensors],

290 - inputs_embeds: Optional[torch.Tensor] = None,

291 - ) -> Union[torch.Tensor, IntermediateTensors]:

288 + intermediate_tensors: IntermediateTensors | None,

289 + inputs_embeds: torch.Tensor | None = None,

290 + ) -> torch.Tensor | IntermediateTensors:

291 if get_pp_group().is_first_rank:

292 if inputs_embeds is not None:

293 hidden_states = inputs_embeds

@@ -406,9 +405,9 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

405 self,

406 input_ids: torch.Tensor,

407 positions: torch.Tensor,

409 - intermediate_tensors: Optional[IntermediateTensors] = None,

410 - inputs_embeds: Optional[torch.Tensor] = None,

411 - ) -> Union[torch.Tensor, IntermediateTensors]:

408 + intermediate_tensors: IntermediateTensors | None = None,

409 + inputs_embeds: torch.Tensor | None = None,

410 + ) -> torch.Tensor | IntermediateTensors:

411 hidden_states = self.model(

412 input_ids, positions, intermediate_tensors, inputs_embeds

413 )

@@ -417,7 +416,7 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

416 def compute_logits(

417 self,

418 hidden_states: torch.Tensor,

420 - ) -> Optional[torch.Tensor]:

419 + ) -> torch.Tensor | None:

420 logits = self.logits_processor(self.model.embed_tokens, hidden_states)

421 return logits

422

vllm/model_executor/models/gemma3.py

+15 -16

@@ -17,7 +17,6 @@

17 # limitations under the License.

18 from collections.abc import Iterable

19 from itertools import islice

20 -from typing import Optional, Union

20

21 import torch

22 import torch.nn.functional as F

@@ -66,7 +65,7 @@ class Gemma3MLP(nn.Module):

65 hidden_size: int,

66 intermediate_size: int,

67 hidden_activation: str,

69 - quant_config: Optional[QuantizationConfig] = None,

68 + quant_config: QuantizationConfig | None = None,

69 prefix: str = "",

70 ) -> None:

71 super().__init__()

@@ -108,9 +107,9 @@ class Gemma3Attention(nn.Module):

107 num_kv_heads: int,

108 head_dim: int,

109 max_position_embeddings: int,

111 - cache_config: Optional[CacheConfig] = None,

112 - quant_config: Optional[QuantizationConfig] = None,

113 - attn_logits_soft_cap: Optional[float] = None,

110 + cache_config: CacheConfig | None = None,

111 + quant_config: QuantizationConfig | None = None,

112 + attn_logits_soft_cap: float | None = None,

113 prefix: str = "",

114 ) -> None:

115 super().__init__()

@@ -295,8 +294,8 @@ class Gemma3DecoderLayer(nn.Module):

294 def __init__(

295 self,

296 config: Gemma3TextConfig,

298 - cache_config: Optional[CacheConfig] = None,

299 - quant_config: Optional[QuantizationConfig] = None,

297 + cache_config: CacheConfig | None = None,

298 + quant_config: QuantizationConfig | None = None,

299 prefix: str = "",

300 ) -> None:

301 super().__init__()

@@ -336,7 +335,7 @@ class Gemma3DecoderLayer(nn.Module):

335 self,

336 positions: torch.Tensor,

337 hidden_states: torch.Tensor,

339 - residual: Optional[torch.Tensor],

338 + residual: torch.Tensor | None,

339 **kwargs,

340 ) -> tuple[torch.Tensor, torch.Tensor]:

341 if residual is None:

@@ -401,12 +400,12 @@ class Gemma3Model(nn.Module):

400

401 def forward(

402 self,

404 - input_ids: Optional[torch.Tensor],

403 + input_ids: torch.Tensor | None,

404 positions: torch.Tensor,

406 - intermediate_tensors: Optional[IntermediateTensors],

407 - inputs_embeds: Optional[torch.Tensor] = None,

405 + intermediate_tensors: IntermediateTensors | None,

406 + inputs_embeds: torch.Tensor | None = None,

407 **kwargs,

409 - ) -> Union[torch.Tensor, IntermediateTensors]:

408 + ) -> torch.Tensor | IntermediateTensors:

409 if get_pp_group().is_first_rank:

410 if inputs_embeds is not None:

411 hidden_states = inputs_embeds

@@ -549,10 +548,10 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

548 self,

549 input_ids: torch.Tensor,

550 positions: torch.Tensor,

552 - intermediate_tensors: Optional[IntermediateTensors] = None,

553 - inputs_embeds: Optional[torch.Tensor] = None,

551 + intermediate_tensors: IntermediateTensors | None = None,

552 + inputs_embeds: torch.Tensor | None = None,

553 **kwargs,

555 - ) -> Union[torch.Tensor, IntermediateTensors]:

554 + ) -> torch.Tensor | IntermediateTensors:

555 hidden_states = self.model(

556 input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs

557 )

@@ -561,7 +560,7 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

560 def compute_logits(

561 self,

562 hidden_states: torch.Tensor,

564 - ) -> Optional[torch.Tensor]:

563 + ) -> torch.Tensor | None:

564 logits = self.logits_processor(self.model.embed_tokens, hidden_states)

565 return logits

566

vllm/model_executor/models/gemma3_mm.py

+11 -11

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import Annotated, Any, Literal, Optional

5 +from typing import Annotated, Any, Literal

6

7 import torch

8 from torch import nn

@@ -82,7 +82,7 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):

82 def get_hf_processor(self, **kwargs: object):

83 return self.ctx.get_hf_processor(Gemma3Processor, **kwargs)

84

85 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

85 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

86 return {"image": None}

87

88 def _resolve_image_kwargs(

@@ -112,7 +112,7 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):

112 *,

113 image_width: int,

114 image_height: int,

115 - processor: Optional[Gemma3Processor],

115 + processor: Gemma3Processor | None,

116 ) -> int:

117 if processor is None:

118 processor = self.get_hf_processor()

@@ -182,7 +182,7 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):

182 *,

183 image_width: int,

184 image_height: int,

185 - processor: Optional[Gemma3Processor],

185 + processor: Gemma3Processor | None,

186 ) -> PromptUpdateDetails[str]:

187 if processor is None:

188 processor = self.get_hf_processor()

@@ -217,7 +217,7 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):

217 *,

218 image_width: int,

219 image_height: int,

220 - processor: Optional[Gemma3Processor],

220 + processor: Gemma3Processor | None,

221 ) -> int:

222 if processor is None:

223 processor = self.get_hf_processor()

@@ -256,7 +256,7 @@ class Gemma3DummyInputsBuilder(BaseDummyInputsBuilder[Gemma3ProcessingInfo]):

256 self,

257 seq_len: int,

258 mm_counts: Mapping[str, int],

259 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

259 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

260 ) -> MultiModalDataDict:

261 num_images = mm_counts.get("image", 0)

262

@@ -510,7 +510,7 @@ class Gemma3ForConditionalGeneration(

510 )

511

512 @classmethod

513 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

513 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

514 if modality.startswith("image"):

515 return "<start_of_image>"

516

@@ -555,7 +555,7 @@ class Gemma3ForConditionalGeneration(

555

556 def _parse_and_validate_image_input(

557 self, **kwargs: object

558 - ) -> Optional[Gemma3ImageInputs]:

558 + ) -> Gemma3ImageInputs | None:

559 pixel_values = kwargs.pop("pixel_values", None)

560 num_patches = kwargs.pop("num_patches", None)

561 image_embeds = kwargs.pop("image_embeds", None)

@@ -609,8 +609,8 @@ class Gemma3ForConditionalGeneration(

609 self,

610 input_ids: torch.Tensor,

611 positions: torch.Tensor,

612 - intermediate_tensors: Optional[IntermediateTensors] = None,

613 - inputs_embeds: Optional[torch.Tensor] = None,

612 + intermediate_tensors: IntermediateTensors | None = None,

613 + inputs_embeds: torch.Tensor | None = None,

614 **kwargs: object,

615 ) -> IntermediateTensors:

616 if intermediate_tensors is not None:

@@ -692,7 +692,7 @@ class Gemma3ForConditionalGeneration(

692 def compute_logits(

693 self,

694 hidden_states: torch.Tensor,

695 - ) -> Optional[torch.Tensor]:

695 + ) -> torch.Tensor | None:

696 return self.language_model.compute_logits(hidden_states)

697

698 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/gemma3n.py

+23 -24

@@ -16,7 +16,6 @@

16 # See the License for the specific language governing permissions and

17 # limitations under the License.

18 from collections.abc import Iterable

19 -from typing import Optional, Union

19

20 import torch

21 from torch import nn

@@ -196,7 +195,7 @@ class Gemma3nLaurelBlock(nn.Module):

195 laurel_rank: int,

196 rms_norm_eps: float,

197 *,

199 - quant_config: Optional[QuantizationConfig] = None,

198 + quant_config: QuantizationConfig | None = None,

199 prefix: str,

200 ) -> None:

201 super().__init__()

@@ -236,7 +235,7 @@ class Gemma3nMLP(nn.Module):

235 intermediate_size: int,

236 hidden_activation: str,

237 activation_sparsity: float = 0.0,

239 - quant_config: Optional[QuantizationConfig] = None,

238 + quant_config: QuantizationConfig | None = None,

239 prefix: str = "",

240 ) -> None:

241 super().__init__()

@@ -285,8 +284,8 @@ class Gemma3nAttention(nn.Module):

284 num_kv_heads: int,

285 head_dim: int,

286 max_position_embeddings: int,

288 - cache_config: Optional[CacheConfig] = None,

289 - quant_config: Optional[QuantizationConfig] = None,

287 + cache_config: CacheConfig | None = None,

288 + quant_config: QuantizationConfig | None = None,

289 prefix: str = "",

290 ) -> None:

291 super().__init__()

@@ -412,8 +411,8 @@ class Gemma3nDecoderLayer(nn.Module):

411 def __init__(

412 self,

413 config: Gemma3nTextConfig,

415 - cache_config: Optional[CacheConfig] = None,

416 - quant_config: Optional[QuantizationConfig] = None,

414 + cache_config: CacheConfig | None = None,

415 + quant_config: QuantizationConfig | None = None,

416 prefix: str = "",

417 ) -> None:

418 super().__init__()

@@ -650,7 +649,7 @@ class Gemma3nSelfDecoder(nn.Module):

649 def get_per_layer_inputs(

650 self,

651 hidden_states_0: torch.Tensor,

653 - per_layer_inputs: Optional[torch.Tensor],

652 + per_layer_inputs: torch.Tensor | None,

653 ) -> torch.Tensor:

654 per_layer_projection = self.per_layer_model_projection(hidden_states_0)

655 per_layer_projection = per_layer_projection.reshape(

@@ -687,8 +686,8 @@ class Gemma3nSelfDecoder(nn.Module):

686 self,

687 input_ids: torch.Tensor,

688 positions: torch.Tensor,

690 - inputs_embeds: Optional[torch.Tensor] = None,

691 - per_layer_inputs: Optional[torch.Tensor] = None,

689 + inputs_embeds: torch.Tensor | None = None,

690 + per_layer_inputs: torch.Tensor | None = None,

691 **kwargs,

692 ) -> tuple[torch.Tensor, torch.Tensor]:

693 if inputs_embeds is not None:

@@ -870,8 +869,8 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

869 self,

870 input_ids: torch.Tensor,

871 positions: torch.Tensor,

873 - inputs_embeds: Optional[torch.Tensor] = None,

874 - per_layer_inputs: Optional[torch.Tensor] = None,

872 + inputs_embeds: torch.Tensor | None = None,

873 + per_layer_inputs: torch.Tensor | None = None,

874 **kwargs,

875 ) -> torch.Tensor:

876 logits_indices_padded, num_logits_indices = None, None

@@ -947,8 +946,8 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

946 self,

947 input_ids: torch.Tensor,

948 positions: torch.Tensor,

950 - inputs_embeds: Optional[torch.Tensor] = None,

951 - per_layer_inputs: Optional[torch.Tensor] = None,

949 + inputs_embeds: torch.Tensor | None = None,

950 + per_layer_inputs: torch.Tensor | None = None,

951 **kwargs,

952 ) -> torch.Tensor:

953 hidden_states, per_layer_inputs = self.self_decoder(

@@ -990,13 +989,13 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

989

990 def forward(

991 self,

993 - input_ids: Optional[torch.Tensor],

992 + input_ids: torch.Tensor | None,

993 positions: torch.Tensor,

995 - per_layer_inputs: Optional[torch.Tensor] = None,

996 - intermediate_tensors: Optional[IntermediateTensors] = None,

997 - inputs_embeds: Optional[torch.Tensor] = None,

994 + per_layer_inputs: torch.Tensor | None = None,

995 + intermediate_tensors: IntermediateTensors | None = None,

996 + inputs_embeds: torch.Tensor | None = None,

997 **kwargs,

999 - ) -> Union[torch.Tensor, IntermediateTensors]:

998 + ) -> torch.Tensor | IntermediateTensors:

999 if self.fast_prefill_enabled:

1000 hidden_states = self.fast_prefill_forward(

1001 input_ids,

@@ -1116,11 +1115,11 @@ class Gemma3nForCausalLM(nn.Module):

1115 input_ids: torch.Tensor,

1116 positions: torch.Tensor,

1117 *,

1119 - per_layer_inputs: Optional[torch.Tensor] = None,

1120 - intermediate_tensors: Optional[IntermediateTensors] = None,

1121 - inputs_embeds: Optional[torch.Tensor] = None,

1118 + per_layer_inputs: torch.Tensor | None = None,

1119 + intermediate_tensors: IntermediateTensors | None = None,

1120 + inputs_embeds: torch.Tensor | None = None,

1121 **kwargs,

1123 - ) -> Union[torch.Tensor, IntermediateTensors]:

1122 + ) -> torch.Tensor | IntermediateTensors:

1123 hidden_states = self.model(

1124 input_ids,

1125 positions,

@@ -1134,7 +1133,7 @@ class Gemma3nForCausalLM(nn.Module):

1133 def compute_logits(

1134 self,

1135 hidden_states: torch.Tensor,

1137 - ) -> Optional[torch.Tensor]:

1136 + ) -> torch.Tensor | None:

1137 logits = self.logits_processor(self.model.embed_tokens, hidden_states)

1138 return logits

1139

vllm/model_executor/models/glm4.py

+10 -11

@@ -24,7 +24,6 @@

24 """Inference-only GLM-4-0414 model compatible with HuggingFace weights."""

25

26 from collections.abc import Iterable

27 -from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -56,12 +55,12 @@ class Glm4Attention(nn.Module):

55 num_heads: int,

56 num_kv_heads: int,

57 max_position: int = 4096 * 32,

59 - head_dim: Optional[int] = None,

58 + head_dim: int | None = None,

59 qkv_bias: bool = False,

60 rope_theta: float = 10000,

62 - cache_config: Optional[CacheConfig] = None,

63 - quant_config: Optional[QuantizationConfig] = None,

64 - rope_scaling: Optional[tuple] = None,

61 + cache_config: CacheConfig | None = None,

62 + quant_config: QuantizationConfig | None = None,

63 + rope_scaling: tuple | None = None,

64 prefix: str = "",

65 attn_type: str = AttentionType.DECODER,

66 ) -> None:

@@ -142,7 +141,7 @@ class Glm4DecoderLayer(nn.Module):

141 self,

142 vllm_config: VllmConfig,

143 prefix: str = "",

145 - config: Optional[Glm4Config] = None,

144 + config: Glm4Config | None = None,

145 ) -> None:

146 super().__init__()

147

@@ -189,7 +188,7 @@ class Glm4DecoderLayer(nn.Module):

188 self,

189 positions: torch.Tensor,

190 hidden_states: torch.Tensor,

192 - residual: Optional[torch.Tensor],

191 + residual: torch.Tensor | None,

192 ) -> tuple[torch.Tensor, torch.Tensor]:

193 # Self Attention

194 if residual is None:

@@ -285,9 +284,9 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

284 self,

285 input_ids: torch.Tensor,

286 positions: torch.Tensor,

288 - intermediate_tensors: Optional[IntermediateTensors] = None,

289 - inputs_embeds: Optional[torch.Tensor] = None,

290 - ) -> Union[torch.Tensor, IntermediateTensors]:

287 + intermediate_tensors: IntermediateTensors | None = None,

288 + inputs_embeds: torch.Tensor | None = None,

289 + ) -> torch.Tensor | IntermediateTensors:

290 hidden_states = self.model(

291 input_ids, positions, intermediate_tensors, inputs_embeds

292 )

@@ -296,7 +295,7 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

295 def compute_logits(

296 self,

297 hidden_states: torch.Tensor,

299 - ) -> Optional[torch.Tensor]:

298 + ) -> torch.Tensor | None:

299 logits = self.logits_processor(self.lm_head, hidden_states)

300 return logits

301

vllm/model_executor/models/glm4_1v.py

+26 -26

@@ -27,9 +27,9 @@

27 """Inference-only GLM-4V model compatible with HuggingFace weights."""

28

29 import math

30 -from collections.abc import Iterable, Mapping, Sequence

30 +from collections.abc import Callable, Iterable, Mapping, Sequence

31 from functools import partial

32 -from typing import Annotated, Any, Callable, Literal, Optional, Union

32 +from typing import Annotated, Any, Literal, TypeAlias

33

34 import numpy as np

35 import torch

@@ -140,7 +140,7 @@ class Glm4vImageEmbeddingInputs(TensorSchema):

140 image_grid_thw: Annotated[torch.Tensor, TensorShape("n", 3)]

141

142

143 -Glm4vImageInputs = Union[Glm4vImagePixelInputs, Glm4vImageEmbeddingInputs]

143 +Glm4vImageInputs: TypeAlias = Glm4vImagePixelInputs | Glm4vImageEmbeddingInputs

144

145

146 class Glm4vVideoPixelInputs(TensorSchema):

@@ -176,7 +176,7 @@ class Glm4vVideoEmbeddingInputs(TensorSchema):

176 video_grid_thw: Annotated[torch.Tensor, TensorShape("f", 3)]

177

178

179 -Glm4vVideoInputs = Union[Glm4vVideoPixelInputs, Glm4vVideoEmbeddingInputs]

179 +Glm4vVideoInputs: TypeAlias = Glm4vVideoPixelInputs | Glm4vVideoEmbeddingInputs

180

181 # ==== Vision Encoder ==== #

182

@@ -187,7 +187,7 @@ class Glm4vVisionMLP(nn.Module):

187 in_features: int,

188 hidden_features: int,

189 bias: bool = False,

190 - quant_config: Optional[QuantizationConfig] = None,

190 + quant_config: QuantizationConfig | None = None,

191 prefix: str = "",

192 use_data_parallel: bool = False,

193 ):

@@ -244,7 +244,7 @@ class Glm4vVisionAttention(nn.Module):

244 embed_dim: int,

245 num_heads: int,

246 projection_size: int,

247 - quant_config: Optional[QuantizationConfig] = None,

247 + quant_config: QuantizationConfig | None = None,

248 prefix: str = "",

249 use_data_parallel: bool = False,

250 ) -> None:

@@ -334,8 +334,8 @@ class Glm4vVisionAttention(nn.Module):

334 x: torch.Tensor,

335 cu_seqlens: torch.Tensor,

336 rotary_pos_emb: torch.Tensor,

337 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

338 - seqlens: Optional[list[int]] = None, # Only used for xFormers

337 + max_seqlen: int | None = None, # Only used for Flash Attention

338 + seqlens: list[int] | None = None, # Only used for xFormers

339 ) -> torch.Tensor:

340 # [s, b, c] --> [s, b, head * 3 * head_dim]

341 x, _ = self.qkv(x)

@@ -413,8 +413,8 @@ class Glm4vVisionBlock(nn.Module):

413 dim: int,

414 num_heads: int,

415 mlp_hidden_dim: int,

416 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

417 - quant_config: Optional[QuantizationConfig] = None,

416 + norm_layer: Callable[[int], nn.Module] | None = None,

417 + quant_config: QuantizationConfig | None = None,

418 prefix: str = "",

419 use_data_parallel: bool = False,

420 ) -> None:

@@ -445,8 +445,8 @@ class Glm4vVisionBlock(nn.Module):

445 x: torch.Tensor,

446 cu_seqlens: torch.Tensor,

447 rotary_pos_emb: torch.Tensor,

448 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

449 - seqlens: Optional[list[int]] = None, # Only used for xFormers

448 + max_seqlen: int | None = None, # Only used for Flash Attention

449 + seqlens: list[int] | None = None, # Only used for xFormers

450 ) -> torch.Tensor:

451 x_attn = self.attn(

452 self.norm1(x),

@@ -495,7 +495,7 @@ class Glm4vPatchMerger(nn.Module):

495 self,

496 d_model: int,

497 context_dim: int,

498 - quant_config: Optional[QuantizationConfig] = None,

498 + quant_config: QuantizationConfig | None = None,

499 bias: bool = False,

500 prefix: str = "",

501 use_data_parallel: bool = False,

@@ -693,7 +693,7 @@ class Glm4vVisionTransformer(nn.Module):

693 self,

694 vision_config: Glm4vVisionConfig,

695 norm_eps: float = 1e-6,

696 - quant_config: Optional[QuantizationConfig] = None,

696 + quant_config: QuantizationConfig | None = None,

697 prefix: str = "",

698 use_data_parallel: bool = False,

699 ) -> None:

@@ -809,7 +809,7 @@ class Glm4vVisionTransformer(nn.Module):

809 def compute_attn_mask_seqlen(

810 self,

811 cu_seqlens: torch.Tensor,

812 - ) -> tuple[Optional[int], Optional[list[int]]]:

812 + ) -> tuple[int | None, list[int] | None]:

813 max_seqlen, seqlens = None, None

814 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

815 if (

@@ -904,7 +904,7 @@ class Glm4vProcessingInfo(BaseProcessingInfo):

904 def get_tokenizer(self):

905 return self.ctx.tokenizer

906

907 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

907 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

908 return {"image": None, "video": 1}

909

910 def get_image_processor(self, **kwargs: object) -> Glm4vImageProcessor:

@@ -1141,7 +1141,7 @@ class Glm4vDummyInputsBuilder(BaseDummyInputsBuilder[Glm4vProcessingInfo]):

1141 self,

1142 seq_len: int,

1143 mm_counts: Mapping[str, int],

1144 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

1144 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

1145 ) -> MultiModalDataDict:

1146 num_images = mm_counts.get("image", 0)

1147 num_videos = mm_counts.get("video", 0)

@@ -1177,7 +1177,7 @@ class Glm4vDummyInputsBuilder(BaseDummyInputsBuilder[Glm4vProcessingInfo]):

1177 height: int,

1178 num_frames: int,

1179 num_videos: int,

1180 - overrides: Optional[VideoDummyOptions] = None,

1180 + overrides: VideoDummyOptions | None = None,

1181 ) -> list[VideoItem]:

1182 if overrides:

1183 if overrides.num_frames:

@@ -1419,7 +1419,7 @@ class Glm4vForConditionalGeneration(

1419 supports_encoder_tp_data = True

1420

1421 @classmethod

1422 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1422 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1423 if modality.startswith("image"):

1424 return "<|begin_of_image|><|image|><|end_of_image|>"

1425 if modality.startswith("video"):

@@ -1465,7 +1465,7 @@ class Glm4vForConditionalGeneration(

1465

1466 def _parse_and_validate_image_input(

1467 self, **kwargs: object

1468 - ) -> Optional[Glm4vImageInputs]:

1468 + ) -> Glm4vImageInputs | None:

1469 pixel_values = kwargs.pop("pixel_values", None)

1470 image_embeds = kwargs.pop("image_embeds", None)

1471 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -1489,7 +1489,7 @@ class Glm4vForConditionalGeneration(

1489

1490 def _parse_and_validate_video_input(

1491 self, **kwargs: object

1492 - ) -> Optional[Glm4vVideoInputs]:

1492 + ) -> Glm4vVideoInputs | None:

1493 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1494 video_embeds = kwargs.pop("video_embeds", None)

1495 video_grid_thw = kwargs.pop("video_grid_thw", None)

@@ -1594,7 +1594,7 @@ class Glm4vForConditionalGeneration(

1594

1595 def get_multimodal_embeddings(

1596 self, **kwargs: object

1597 - ) -> Optional[MultiModalEmbeddings]:

1597 + ) -> MultiModalEmbeddings | None:

1598 mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

1599 if not mm_input_by_modality:

1600 return None

@@ -1619,10 +1619,10 @@ class Glm4vForConditionalGeneration(

1619 self,

1620 input_ids: torch.Tensor,

1621 positions: torch.Tensor,

1622 - intermediate_tensors: Optional[IntermediateTensors] = None,

1623 - inputs_embeds: Optional[torch.Tensor] = None,

1622 + intermediate_tensors: IntermediateTensors | None = None,

1623 + inputs_embeds: torch.Tensor | None = None,

1624 **kwargs: object,

1625 - ) -> Union[torch.Tensor, IntermediateTensors]:

1625 + ) -> torch.Tensor | IntermediateTensors:

1626 """Run forward pass for GLM-4V.

1627

1628 Args:

@@ -1652,7 +1652,7 @@ class Glm4vForConditionalGeneration(

1652 def compute_logits(

1653 self,

1654 hidden_states: torch.Tensor,

1655 - ) -> Optional[torch.Tensor]:

1655 + ) -> torch.Tensor | None:

1656 return self.language_model.compute_logits(hidden_states)

1657

1658 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/glm4_moe.py

+18 -18

@@ -26,7 +26,7 @@

26 import typing

27 from collections.abc import Callable, Iterable

28 from itertools import islice

29 -from typing import Any, Optional, Union

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -81,7 +81,7 @@ class Glm4MoeMLP(nn.Module):

81 hidden_size: int,

82 intermediate_size: int,

83 hidden_act: str,

84 - quant_config: Optional[QuantizationConfig] = None,

84 + quant_config: QuantizationConfig | None = None,

85 reduce_results: bool = True,

86 prefix: str = "",

87 ) -> None:

@@ -118,7 +118,7 @@ class Glm4MoE(nn.Module):

118 def __init__(

119 self,

120 config: Glm4MoeConfig,

121 - quant_config: Optional[QuantizationConfig] = None,

121 + quant_config: QuantizationConfig | None = None,

122 prefix: str = "",

123 enable_eplb: bool = False,

124 ):

@@ -234,14 +234,14 @@ class Glm4MoeAttention(nn.Module):

234 num_heads: int,

235 num_kv_heads: int,

236 rope_theta: float = 10000,

237 - rope_scaling: Optional[dict[str, Any]] = None,

237 + rope_scaling: dict[str, Any] | None = None,

238 max_position_embeddings: int = 131072,

239 - head_dim: Optional[int] = None,

239 + head_dim: int | None = None,

240 rms_norm_eps: float = 1e-05,

241 qkv_bias: bool = False,

242 use_qk_norm: bool = False,

243 - cache_config: Optional[CacheConfig] = None,

244 - quant_config: Optional[QuantizationConfig] = None,

243 + cache_config: CacheConfig | None = None,

244 + quant_config: QuantizationConfig | None = None,

245 prefix: str = "",

246 ) -> None:

247 super().__init__()

@@ -334,8 +334,8 @@ class Glm4MoeDecoderLayer(nn.Module):

334 def __init__(

335 self,

336 config: Glm4MoeConfig,

337 - cache_config: Optional[CacheConfig] = None,

338 - quant_config: Optional[QuantizationConfig] = None,

337 + cache_config: CacheConfig | None = None,

338 + quant_config: QuantizationConfig | None = None,

339 prefix: str = "",

340 enable_eplb: bool = False,

341 ) -> None:

@@ -395,7 +395,7 @@ class Glm4MoeDecoderLayer(nn.Module):

395 self,

396 positions: torch.Tensor,

397 hidden_states: torch.Tensor,

398 - residual: Optional[torch.Tensor],

398 + residual: torch.Tensor | None,

399 ) -> tuple[torch.Tensor, torch.Tensor]:

400 if residual is None:

401 residual = hidden_states

@@ -462,9 +462,9 @@ class Glm4MoeModel(nn.Module):

462 self,

463 input_ids: torch.Tensor,

464 positions: torch.Tensor,

465 - intermediate_tensors: Optional[IntermediateTensors] = None,

466 - inputs_embeds: Optional[torch.Tensor] = None,

467 - ) -> Union[torch.Tensor, IntermediateTensors]:

465 + intermediate_tensors: IntermediateTensors | None = None,

466 + inputs_embeds: torch.Tensor | None = None,

467 + ) -> torch.Tensor | IntermediateTensors:

468 if get_pp_group().is_first_rank:

469 if inputs_embeds is not None:

470 hidden_states = inputs_embeds

@@ -704,9 +704,9 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

704 self,

705 input_ids: torch.Tensor,

706 positions: torch.Tensor,

707 - intermediate_tensors: Optional[IntermediateTensors] = None,

708 - inputs_embeds: Optional[torch.Tensor] = None,

709 - ) -> Union[torch.Tensor, IntermediateTensors]:

707 + intermediate_tensors: IntermediateTensors | None = None,

708 + inputs_embeds: torch.Tensor | None = None,

709 + ) -> torch.Tensor | IntermediateTensors:

710 hidden_states = self.model(

711 input_ids, positions, intermediate_tensors, inputs_embeds

712 )

@@ -715,7 +715,7 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

715 def compute_logits(

716 self,

717 hidden_states: torch.Tensor,

718 - ) -> Optional[torch.Tensor]:

718 + ) -> torch.Tensor | None:

719 logits = self.logits_processor(self.lm_head, hidden_states)

720 return logits

721

@@ -729,7 +729,7 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

729

730 def get_spec_layer_idx_from_weight_name(

731 config: Glm4MoeConfig, weight_name: str

732 -) -> Optional[int]:

732 +) -> int | None:

733 if hasattr(config, "num_nextn_predict_layers") and (

734 config.num_nextn_predict_layers > 0

735 ):

vllm/model_executor/models/glm4_moe_mtp.py

+8 -9

@@ -24,7 +24,6 @@

24 """Inference-only GLM-4.5 MTP model compatible with HuggingFace weights."""

25

26 from collections.abc import Iterable

27 -from typing import Optional

27

28 import torch

29 import torch.nn as nn

@@ -52,7 +51,7 @@ class SharedHead(nn.Module):

51 self,

52 config: PretrainedConfig,

53 prefix: str,

55 - quant_config: Optional[QuantizationConfig] = None,

54 + quant_config: QuantizationConfig | None = None,

55 ) -> None:

56 super().__init__()

57 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

@@ -72,8 +71,8 @@ class Glm4MoeMultiTokenPredictorLayer(nn.Module):

71 self,

72 config: PretrainedConfig,

73 prefix: str,

75 - cache_config: Optional[CacheConfig] = None,

76 - quant_config: Optional[QuantizationConfig] = None,

74 + cache_config: CacheConfig | None = None,

75 + quant_config: QuantizationConfig | None = None,

76 ) -> None:

77 super().__init__()

78 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

@@ -94,7 +93,7 @@ class Glm4MoeMultiTokenPredictorLayer(nn.Module):

93 input_ids: torch.Tensor,

94 positions: torch.Tensor,

95 previous_hidden_states: torch.Tensor,

97 - inputs_embeds: Optional[torch.Tensor] = None,

96 + inputs_embeds: torch.Tensor | None = None,

97 spec_step_index: int = 0,

98 ) -> torch.Tensor:

99 assert inputs_embeds is not None

@@ -149,7 +148,7 @@ class Glm4MoeMultiTokenPredictor(nn.Module):

148 input_ids: torch.Tensor,

149 positions: torch.Tensor,

150 previous_hidden_states: torch.Tensor,

152 - inputs_embeds: Optional[torch.Tensor] = None,

151 + inputs_embeds: torch.Tensor | None = None,

152 spec_step_idx: int = 0,

153 ) -> torch.Tensor:

154 if inputs_embeds is None:

@@ -192,8 +191,8 @@ class Glm4MoeMTP(nn.Module, SupportsPP):

191 input_ids: torch.Tensor,

192 positions: torch.Tensor,

193 hidden_states: torch.Tensor,

195 - intermediate_tensors: Optional[IntermediateTensors] = None,

196 - inputs_embeds: Optional[torch.Tensor] = None,

194 + intermediate_tensors: IntermediateTensors | None = None,

195 + inputs_embeds: torch.Tensor | None = None,

196 spec_step_idx: int = 0,

197 ) -> torch.Tensor:

198 hidden_states = self.model(

@@ -205,7 +204,7 @@ class Glm4MoeMTP(nn.Module, SupportsPP):

204 self,

205 hidden_states: torch.Tensor,

206 spec_step_idx: int = 0,

208 - ) -> Optional[torch.Tensor]:

207 + ) -> torch.Tensor | None:

208 return self.model.compute_logits(hidden_states, spec_step_idx)

209

210 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/glm4v.py

+22 -22

@@ -8,7 +8,7 @@

8 import itertools

9 from argparse import Namespace

10 from collections.abc import Mapping, Sequence

11 -from typing import Annotated, Literal, Optional, Union

11 +from typing import Annotated, Literal

12

13 import torch

14 from torch import nn

@@ -109,7 +109,7 @@ class EVA2CLIPAttention(nn.Module):

109 def __init__(

110 self,

111 config,

112 - quant_config: Optional[QuantizationConfig] = None,

112 + quant_config: QuantizationConfig | None = None,

113 prefix: str = "",

114 ):

115 super().__init__()

@@ -152,7 +152,7 @@ class EVA2CLIPMLP(nn.Module):

152 def __init__(

153 self,

154 config,

155 - quant_config: Optional[QuantizationConfig] = None,

155 + quant_config: QuantizationConfig | None = None,

156 prefix: str = "",

157 ):

158 super().__init__()

@@ -182,7 +182,7 @@ class EVA2CLIPTransformerLayer(nn.Module):

182 def __init__(

183 self,

184 config,

185 - quant_config: Optional[QuantizationConfig] = None,

185 + quant_config: QuantizationConfig | None = None,

186 prefix: str = "",

187 ):

188 super().__init__()

@@ -211,7 +211,7 @@ class EVA2CLIPTransformer(nn.Module):

211 def __init__(

212 self,

213 config,

214 - quant_config: Optional[QuantizationConfig] = None,

214 + quant_config: QuantizationConfig | None = None,

215 prefix: str = "",

216 ):

217 super().__init__()

@@ -237,7 +237,7 @@ class EVA2CLIPGLU(nn.Module):

237 self,

238 config,

239 in_features,

240 - quant_config: Optional[QuantizationConfig] = None,

240 + quant_config: QuantizationConfig | None = None,

241 prefix: str = "",

242 ):

243 """

@@ -317,7 +317,7 @@ class EVA2CLIPModel(nn.Module):

317 def __init__(

318 self,

319 config,

320 - quant_config: Optional[QuantizationConfig] = None,

320 + quant_config: QuantizationConfig | None = None,

321 prefix: str = "",

322 ):

323 super().__init__()

@@ -416,9 +416,9 @@ class GLM4VProcessor:

416

417 def __call__(

418 self,

419 - text: Optional[Union[TextInput, list[TextInput]]] = None,

420 - images: Optional[Union[ImageInput, list[ImageInput]]] = None,

421 - return_tensors: Optional[Union[str, TensorType]] = None,

419 + text: TextInput | list[TextInput] | None = None,

420 + images: ImageInput | list[ImageInput] | None = None,

421 + return_tensors: str | TensorType | None = None,

422 ) -> BatchFeature:

423 if text is None:

424 text = []

@@ -458,7 +458,7 @@ class GLM4VProcessingInfo(BaseProcessingInfo):

458 **kwargs,

459 )

460

461 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

461 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

462 return {"image": 1}

463

464 def get_num_image_tokens(self) -> int:

@@ -487,7 +487,7 @@ class GLM4VDummyInputsBuilder(BaseDummyInputsBuilder[GLM4VProcessingInfo]):

487 self,

488 seq_len: int,

489 mm_counts: Mapping[str, int],

490 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

490 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

491 ) -> MultiModalDataDict:

492 hf_config = self.info.get_hf_config()

493 vision_config = hf_config.vision_config

@@ -578,7 +578,7 @@ class GLM4VForCausalLM(

578 )

579

580 @classmethod

581 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

581 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

582 if modality.startswith("image"):

583 return "<|begin_of_image|><|endoftext|><|end_of_image|>"

584

@@ -601,7 +601,7 @@ class GLM4VForCausalLM(

601

602 def _parse_and_validate_image_input(

603 self, **kwargs: object

604 - ) -> Optional[GLMVImagePixelInputs]:

604 + ) -> GLMVImagePixelInputs | None:

605 pixel_values = kwargs.pop("pixel_values", None)

606

607 if pixel_values is not None:

@@ -624,12 +624,12 @@ class GLM4VForCausalLM(

624 cls,

625 input_tokens: list[int],

626 hf_config: PretrainedConfig,

627 - image_grid_thw: Union[list[list[int]], torch.Tensor],

628 - video_grid_thw: Union[list[list[int]], torch.Tensor],

627 + image_grid_thw: list[list[int]] | torch.Tensor,

628 + video_grid_thw: list[list[int]] | torch.Tensor,

629 context_len: int = 0,

630 - seq_len: Optional[int] = None,

631 - second_per_grid_ts: Optional[list[float]] = None,

632 - audio_feature_lengths: Optional[torch.Tensor] = None,

630 + seq_len: int | None = None,

631 + second_per_grid_ts: list[float] | None = None,

632 + audio_feature_lengths: torch.Tensor | None = None,

633 use_audio_in_video: bool = False,

634 ) -> tuple[torch.Tensor, int]:

635 """Get mrope input positions and delta value for GLM4V."""

@@ -780,10 +780,10 @@ class GLM4VForCausalLM(

780 self,

781 input_ids: torch.Tensor,

782 positions: torch.Tensor,

783 - intermediate_tensors: Optional[IntermediateTensors] = None,

784 - inputs_embeds: Optional[torch.Tensor] = None,

783 + intermediate_tensors: IntermediateTensors | None = None,

784 + inputs_embeds: torch.Tensor | None = None,

785 **kwargs: object,

786 - ) -> Union[torch.Tensor, IntermediateTensors]:

786 + ) -> torch.Tensor | IntermediateTensors:

787 if intermediate_tensors is not None:

788 inputs_embeds = None

789

vllm/model_executor/models/gpt2.py

+14 -15

@@ -22,7 +22,6 @@

22

23 from collections.abc import Iterable

24 from itertools import islice

25 -from typing import Optional, Union

25

26 import torch

27 from torch import nn

@@ -65,8 +64,8 @@ class GPT2Attention(nn.Module):

64 def __init__(

65 self,

66 config: GPT2Config,

68 - cache_config: Optional[CacheConfig] = None,

69 - quant_config: Optional[QuantizationConfig] = None,

67 + cache_config: CacheConfig | None = None,

68 + quant_config: QuantizationConfig | None = None,

69 prefix: str = "",

70 ):

71 super().__init__()

@@ -118,7 +117,7 @@ class GPT2MLP(nn.Module):

117 self,

118 intermediate_size: int,

119 config: GPT2Config,

121 - quant_config: Optional[QuantizationConfig] = None,

120 + quant_config: QuantizationConfig | None = None,

121 prefix: str = "",

122 ):

123 super().__init__()

@@ -150,8 +149,8 @@ class GPT2Block(nn.Module):

149 def __init__(

150 self,

151 config: GPT2Config,

153 - cache_config: Optional[CacheConfig] = None,

154 - quant_config: Optional[QuantizationConfig] = None,

152 + cache_config: CacheConfig | None = None,

153 + quant_config: QuantizationConfig | None = None,

154 prefix: str = "",

155 ):

156 super().__init__()

@@ -221,9 +220,9 @@ class GPT2Model(nn.Module):

220 self,

221 input_ids: torch.Tensor,

222 position_ids: torch.Tensor,

224 - intermediate_tensors: Optional[IntermediateTensors],

225 - inputs_embeds: Optional[torch.Tensor],

226 - ) -> Union[torch.Tensor, IntermediateTensors]:

223 + intermediate_tensors: IntermediateTensors | None,

224 + inputs_embeds: torch.Tensor | None,

225 + ) -> torch.Tensor | IntermediateTensors:

226 if get_pp_group().is_first_rank:

227 if inputs_embeds is None:

228 inputs_embeds = self.get_input_embeddings(input_ids)

@@ -301,9 +300,9 @@ class GPT2LMHeadModel(nn.Module, SupportsPP):

300 self,

301 input_ids: torch.Tensor,

302 positions: torch.Tensor,

304 - intermediate_tensors: Optional[IntermediateTensors] = None,

305 - inputs_embeds: Optional[torch.Tensor] = None,

306 - ) -> Union[torch.Tensor, IntermediateTensors]:

303 + intermediate_tensors: IntermediateTensors | None = None,

304 + inputs_embeds: torch.Tensor | None = None,

305 + ) -> torch.Tensor | IntermediateTensors:

306 hidden_states = self.transformer(

307 input_ids, positions, intermediate_tensors, inputs_embeds

308 )

@@ -312,7 +311,7 @@ class GPT2LMHeadModel(nn.Module, SupportsPP):

311 def compute_logits(

312 self,

313 hidden_states: torch.Tensor,

315 - ) -> Optional[torch.Tensor]:

314 + ) -> torch.Tensor | None:

315 logits = self.logits_processor(self.lm_head, hidden_states)

316 return logits

317

@@ -367,8 +366,8 @@ class GPT2ForSequenceClassification(nn.Module):

366 self,

367 input_ids: torch.Tensor,

368 positions: torch.Tensor,

370 - intermediate_tensors: Optional[IntermediateTensors] = None,

371 - inputs_embeds: Optional[torch.Tensor] = None,

369 + intermediate_tensors: IntermediateTensors | None = None,

370 + inputs_embeds: torch.Tensor | None = None,

371 ) -> torch.Tensor:

372 hidden_states = self.transformer(

373 input_ids=input_ids,

vllm/model_executor/models/gpt_bigcode.py

+12 -13

@@ -23,7 +23,6 @@

23

24 from collections.abc import Iterable

25 from itertools import islice

26 -from typing import Optional, Union

26

27 import torch

28 from torch import nn

@@ -62,8 +61,8 @@ class GPTBigCodeAttention(nn.Module):

61 def __init__(

62 self,

63 config: GPTBigCodeConfig,

65 - cache_config: Optional[CacheConfig] = None,

66 - quant_config: Optional[QuantizationConfig] = None,

64 + cache_config: CacheConfig | None = None,

65 + quant_config: QuantizationConfig | None = None,

66 prefix: str = "",

67 ):

68 super().__init__()

@@ -133,7 +132,7 @@ class GPTBigMLP(nn.Module):

132 self,

133 intermediate_size: int,

134 config: GPTBigCodeConfig,

136 - quant_config: Optional[QuantizationConfig] = None,

135 + quant_config: QuantizationConfig | None = None,

136 prefix: str = "",

137 ):

138 super().__init__()

@@ -165,8 +164,8 @@ class GPTBigCodeBlock(nn.Module):

164 def __init__(

165 self,

166 config: GPTBigCodeConfig,

168 - cache_config: Optional[CacheConfig] = None,

169 - quant_config: Optional[QuantizationConfig] = None,

167 + cache_config: CacheConfig | None = None,

168 + quant_config: QuantizationConfig | None = None,

169 prefix: str = "",

170 ):

171 super().__init__()

@@ -243,9 +242,9 @@ class GPTBigCodeModel(nn.Module):

242 self,

243 input_ids: torch.Tensor,

244 position_ids: torch.Tensor,

246 - intermediate_tensors: Optional[IntermediateTensors],

247 - inputs_embeds: Optional[torch.Tensor] = None,

248 - ) -> Union[torch.Tensor, IntermediateTensors]:

245 + intermediate_tensors: IntermediateTensors | None,

246 + inputs_embeds: torch.Tensor | None = None,

247 + ) -> torch.Tensor | IntermediateTensors:

248 if get_pp_group().is_first_rank:

249 if inputs_embeds is None:

250 inputs_embeds = self.get_input_embeddings(input_ids)

@@ -326,9 +325,9 @@ class GPTBigCodeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

325 self,

326 input_ids: torch.Tensor,

327 positions: torch.Tensor,

329 - intermediate_tensors: Optional[IntermediateTensors] = None,

330 - inputs_embeds: Optional[torch.Tensor] = None,

331 - ) -> Union[torch.Tensor, IntermediateTensors]:

328 + intermediate_tensors: IntermediateTensors | None = None,

329 + inputs_embeds: torch.Tensor | None = None,

330 + ) -> torch.Tensor | IntermediateTensors:

331 hidden_states = self.transformer(

332 input_ids, positions, intermediate_tensors, inputs_embeds

333 )

@@ -337,7 +336,7 @@ class GPTBigCodeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

336 def compute_logits(

337 self,

338 hidden_states: torch.Tensor,

340 - ) -> Optional[torch.Tensor]:

339 + ) -> torch.Tensor | None:

340 logits = self.logits_processor(self.lm_head, hidden_states)

341 return logits

342

vllm/model_executor/models/gpt_j.py

+12 -13

@@ -21,7 +21,6 @@

21

22 from collections.abc import Iterable

23 from itertools import islice

24 -from typing import Optional, Union

24

25 import torch

26 from torch import nn

@@ -64,8 +63,8 @@ class GPTJAttention(nn.Module):

63 def __init__(

64 self,

65 config: GPTJConfig,

67 - cache_config: Optional[CacheConfig] = None,

68 - quant_config: Optional[QuantizationConfig] = None,

66 + cache_config: CacheConfig | None = None,

67 + quant_config: QuantizationConfig | None = None,

68 prefix: str = "",

69 ):

70 super().__init__()

@@ -130,7 +129,7 @@ class GPTJMLP(nn.Module):

129 self,

130 intermediate_size: int,

131 config: GPTJConfig,

133 - quant_config: Optional[QuantizationConfig] = None,

132 + quant_config: QuantizationConfig | None = None,

133 ):

134 super().__init__()

135 hidden_size = config.n_embd

@@ -157,8 +156,8 @@ class GPTJBlock(nn.Module):

156 def __init__(

157 self,

158 config: GPTJConfig,

160 - cache_config: Optional[CacheConfig] = None,

161 - quant_config: Optional[QuantizationConfig] = None,

159 + cache_config: CacheConfig | None = None,

160 + quant_config: QuantizationConfig | None = None,

161 prefix: str = "",

162 ):

163 super().__init__()

@@ -218,9 +217,9 @@ class GPTJModel(nn.Module):

217 self,

218 input_ids: torch.Tensor,

219 position_ids: torch.Tensor,

221 - intermediate_tensors: Optional[IntermediateTensors],

222 - inputs_embeds: Optional[torch.Tensor] = None,

223 - ) -> Union[torch.Tensor, IntermediateTensors]:

220 + intermediate_tensors: IntermediateTensors | None,

221 + inputs_embeds: torch.Tensor | None = None,

222 + ) -> torch.Tensor | IntermediateTensors:

223 if get_pp_group().is_first_rank:

224 if inputs_embeds is not None:

225 hidden_states = inputs_embeds

@@ -322,9 +321,9 @@ class GPTJForCausalLM(nn.Module, SupportsPP):

321 self,

322 input_ids: torch.Tensor,

323 positions: torch.Tensor,

325 - intermediate_tensors: Optional[IntermediateTensors] = None,

326 - inputs_embeds: Optional[torch.Tensor] = None,

327 - ) -> Union[torch.Tensor, IntermediateTensors]:

324 + intermediate_tensors: IntermediateTensors | None = None,

325 + inputs_embeds: torch.Tensor | None = None,

326 + ) -> torch.Tensor | IntermediateTensors:

327 hidden_states = self.transformer(

328 input_ids, positions, intermediate_tensors, inputs_embeds

329 )

@@ -333,7 +332,7 @@ class GPTJForCausalLM(nn.Module, SupportsPP):

332 def compute_logits(

333 self,

334 hidden_states: torch.Tensor,

336 - ) -> Optional[torch.Tensor]:

335 + ) -> torch.Tensor | None:

336 logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias)

337 return logits

338

vllm/model_executor/models/gpt_neox.py

+12 -13

@@ -21,7 +21,6 @@

21

22 from collections.abc import Iterable

23 from itertools import islice

24 -from typing import Optional, Union

24

25 import torch

26 from torch import nn

@@ -61,8 +60,8 @@ class GPTNeoXAttention(nn.Module):

60 def __init__(

61 self,

62 config: GPTNeoXConfig,

64 - cache_config: Optional[CacheConfig] = None,

65 - quant_config: Optional[QuantizationConfig] = None,

63 + cache_config: CacheConfig | None = None,

64 + quant_config: QuantizationConfig | None = None,

65 prefix: str = "",

66 ):

67 super().__init__()

@@ -125,7 +124,7 @@ class GPTNeoXMLP(nn.Module):

124 def __init__(

125 self,

126 config: GPTNeoXConfig,

128 - quant_config: Optional[QuantizationConfig] = None,

127 + quant_config: QuantizationConfig | None = None,

128 ):

129 super().__init__()

130 self.dense_h_to_4h = ColumnParallelLinear(

@@ -151,8 +150,8 @@ class GPTNeoXLayer(nn.Module):

150 def __init__(

151 self,

152 config: GPTNeoXConfig,

154 - cache_config: Optional[CacheConfig] = None,

155 - quant_config: Optional[QuantizationConfig] = None,

153 + cache_config: CacheConfig | None = None,

154 + quant_config: QuantizationConfig | None = None,

155 prefix: str = "",

156 ):

157 super().__init__()

@@ -232,9 +231,9 @@ class GPTNeoXModel(nn.Module):

231 self,

232 input_ids: torch.Tensor,

233 position_ids: torch.Tensor,

235 - intermediate_tensors: Optional[IntermediateTensors],

236 - inputs_embeds: Optional[torch.Tensor] = None,

237 - ) -> Union[torch.Tensor, IntermediateTensors]:

234 + intermediate_tensors: IntermediateTensors | None,

235 + inputs_embeds: torch.Tensor | None = None,

236 + ) -> torch.Tensor | IntermediateTensors:

237 if get_pp_group().is_first_rank:

238 if inputs_embeds is not None:

239 hidden_states = inputs_embeds

@@ -320,9 +319,9 @@ class GPTNeoXForCausalLM(nn.Module, SupportsPP):

319 self,

320 input_ids: torch.Tensor,

321 positions: torch.Tensor,

323 - intermediate_tensors: Optional[IntermediateTensors] = None,

324 - inputs_embeds: Optional[torch.Tensor] = None,

325 - ) -> Union[torch.Tensor, IntermediateTensors]:

322 + intermediate_tensors: IntermediateTensors | None = None,

323 + inputs_embeds: torch.Tensor | None = None,

324 + ) -> torch.Tensor | IntermediateTensors:

325 hidden_states = self.gpt_neox(

326 input_ids, positions, intermediate_tensors, inputs_embeds

327 )

@@ -331,7 +330,7 @@ class GPTNeoXForCausalLM(nn.Module, SupportsPP):

330 def compute_logits(

331 self,

332 hidden_states: torch.Tensor,

334 - ) -> Optional[torch.Tensor]:

333 + ) -> torch.Tensor | None:

334 logits = self.logits_processor(self.embed_out, hidden_states)

335 return logits

336

vllm/model_executor/models/gpt_oss.py

+7 -8

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable

4 -from typing import Optional

4

5 import torch

6 import torch.distributed as dist

@@ -49,8 +48,8 @@ class OAIAttention(nn.Module):

48 def __init__(

49 self,

50 config: GptOssConfig,

52 - quant_config: Optional[QuantizationConfig] = None,

53 - cache_config: Optional[CacheConfig] = None,

51 + quant_config: QuantizationConfig | None = None,

52 + cache_config: CacheConfig | None = None,

53 prefix: str = "",

54 ):

55 super().__init__()

@@ -208,7 +207,7 @@ class TransformerBlock(torch.nn.Module):

207 self,

208 hidden_states: torch.Tensor,

209 positions: torch.Tensor,

211 - residual: Optional[torch.Tensor],

210 + residual: torch.Tensor | None,

211 ) -> torch.Tensor:

212 # Self Attention

213 if residual is None:

@@ -260,8 +259,8 @@ class GptOssModel(nn.Module):

259 self,

260 input_ids: torch.Tensor,

261 positions: torch.Tensor,

263 - intermediate_tensors: Optional[IntermediateTensors] = None,

264 - inputs_embeds: Optional[torch.Tensor] = None,

262 + intermediate_tensors: IntermediateTensors | None = None,

263 + inputs_embeds: torch.Tensor | None = None,

264 ) -> torch.Tensor:

265 if get_pp_group().is_first_rank:

266 if inputs_embeds is not None:

@@ -687,8 +686,8 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3):

686 self,

687 input_ids: torch.Tensor,

688 positions: torch.Tensor,

690 - intermediate_tensors: Optional[IntermediateTensors] = None,

691 - inputs_embeds: Optional[torch.Tensor] = None,

689 + intermediate_tensors: IntermediateTensors | None = None,

690 + inputs_embeds: torch.Tensor | None = None,

691 ) -> torch.Tensor:

692 return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)

693

vllm/model_executor/models/granite.py

+15 -15

@@ -26,7 +26,7 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Any, Optional, Union

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -73,7 +73,7 @@ class GraniteMLP(nn.Module):

73 hidden_size: int,

74 intermediate_size: int,

75 hidden_act: str,

76 - quant_config: Optional[QuantizationConfig] = None,

76 + quant_config: QuantizationConfig | None = None,

77 bias: bool = False,

78 prefix: str = "",

79 ) -> None:

@@ -113,11 +113,11 @@ class GraniteAttention(nn.Module):

113 num_heads: int,

114 num_kv_heads: int,

115 rope_theta: float = 10000,

116 - rope_scaling: Optional[dict[str, Any]] = None,

116 + rope_scaling: dict[str, Any] | None = None,

117 max_position_embeddings: int = 8192,

118 - quant_config: Optional[QuantizationConfig] = None,

118 + quant_config: QuantizationConfig | None = None,

119 bias: bool = False,

120 - cache_config: Optional[CacheConfig] = None,

120 + cache_config: CacheConfig | None = None,

121 prefix: str = "",

122 ) -> None:

123 super().__init__()

@@ -197,8 +197,8 @@ class GraniteDecoderLayer(nn.Module):

197 def __init__(

198 self,

199 config: GraniteConfig,

200 - cache_config: Optional[CacheConfig] = None,

201 - quant_config: Optional[QuantizationConfig] = None,

200 + cache_config: CacheConfig | None = None,

201 + quant_config: QuantizationConfig | None = None,

202 prefix: str = "",

203 ) -> None:

204 super().__init__()

@@ -323,11 +323,11 @@ class GraniteModel(nn.Module):

323

324 def forward(

325 self,

326 - input_ids: Optional[torch.Tensor],

326 + input_ids: torch.Tensor | None,

327 positions: torch.Tensor,

328 - intermediate_tensors: Optional[IntermediateTensors],

329 - inputs_embeds: Optional[torch.Tensor] = None,

330 - ) -> Union[torch.Tensor, IntermediateTensors]:

328 + intermediate_tensors: IntermediateTensors | None,

329 + inputs_embeds: torch.Tensor | None = None,

330 + ) -> torch.Tensor | IntermediateTensors:

331 if get_pp_group().is_first_rank:

332 if inputs_embeds is not None:

333 hidden_states = inputs_embeds

@@ -480,15 +480,15 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

480 self,

481 input_ids: torch.Tensor,

482 positions: torch.Tensor,

483 - intermediate_tensors: Optional[IntermediateTensors] = None,

484 - inputs_embeds: Optional[torch.Tensor] = None,

485 - ) -> Union[torch.Tensor, IntermediateTensors]:

483 + intermediate_tensors: IntermediateTensors | None = None,

484 + inputs_embeds: torch.Tensor | None = None,

485 + ) -> torch.Tensor | IntermediateTensors:

486 model_output = self.model(

487 input_ids, positions, intermediate_tensors, inputs_embeds

488 )

489 return model_output

490

491 - def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:

491 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

492 logits = self.logits_processor(self.lm_head, hidden_states)

493 return logits

494

vllm/model_executor/models/granite_speech.py

+14 -14

@@ -26,7 +26,7 @@

26

27 import math

28 from collections.abc import Iterable, Mapping

29 -from typing import Annotated, Optional, Union

29 +from typing import Annotated

30

31 import torch

32 import torch.nn.functional as F

@@ -92,7 +92,7 @@ class GraniteSpeechAudioInputs(TensorSchema):

92

93

94 class GraniteSpeechMultiModalProcessingInfo(BaseProcessingInfo):

95 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

95 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

96 return {"audio": 1}

97

98 # There is no limit to the maximum number of audio tokens that can be

@@ -196,7 +196,7 @@ class GraniteSpeechDummyInputsBuilder(

196 self,

197 seq_len: int,

198 mm_counts: Mapping[str, int],

199 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

199 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

200 ) -> MultiModalDataDict:

201 num_audios = mm_counts.get("audio", 0)

202 audio_overrides = mm_options.get("audio") if mm_options else None

@@ -222,7 +222,7 @@ class GraniteSpeechEncoderProjector(nn.Module):

222 self,

223 config: PretrainedConfig,

224 cache_config: CacheConfig,

225 - quant_config: Optional[QuantizationConfig] = None,

225 + quant_config: QuantizationConfig | None = None,

226 prefix: str = "",

227 ):

228 super().__init__()

@@ -279,7 +279,7 @@ class GraniteSpeechConformerFeedForward(nn.Module):

279 def __init__(

280 self,

281 config: PretrainedConfig,

282 - quant_config: Optional[QuantizationConfig] = None,

282 + quant_config: QuantizationConfig | None = None,

283 prefix: str = "",

284 ):

285 super().__init__()

@@ -479,7 +479,7 @@ class GraniteSpeechCTCEncoder(nn.Module):

479 self,

480 config: PretrainedConfig,

481 prefix: str,

482 - quant_config: Optional[QuantizationConfig] = None,

482 + quant_config: QuantizationConfig | None = None,

483 ):

484 super().__init__()

485 self.config = config

@@ -561,7 +561,7 @@ class GraniteSpeechForConditionalGeneration(

561 }

562

563 @classmethod

564 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

564 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

565 if modality.startswith("audio"):

566 return "<|audio|>"

567

@@ -606,7 +606,7 @@ class GraniteSpeechForConditionalGeneration(

606 def _parse_and_validate_audio_input(

607 self,

608 **kwargs: object,

609 - ) -> Optional[GraniteSpeechAudioInputs]:

609 + ) -> GraniteSpeechAudioInputs | None:

610 input_features = kwargs.pop("input_features", None)

611 input_features_mask = kwargs.pop("input_features_mask", None)

612 audio_embed_sizes = kwargs.pop("audio_embed_sizes", None)

@@ -763,9 +763,9 @@ class GraniteSpeechForConditionalGeneration(

763 def get_input_embeddings(

764 self,

765 input_ids: torch.Tensor,

766 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

766 + multimodal_embeddings: MultiModalEmbeddings | None = None,

767 *,

768 - is_multimodal: Optional[torch.Tensor] = None,

768 + is_multimodal: torch.Tensor | None = None,

769 # Multi-modal token ID may exceed vocab size

770 handle_oov_mm_token: bool = True,

771 ) -> torch.Tensor:

@@ -784,10 +784,10 @@ class GraniteSpeechForConditionalGeneration(

784 self,

785 input_ids: torch.Tensor,

786 positions: torch.Tensor,

787 - intermediate_tensors: Optional[IntermediateTensors] = None,

788 - inputs_embeds: Optional[torch.Tensor] = None,

787 + intermediate_tensors: IntermediateTensors | None = None,

788 + inputs_embeds: torch.Tensor | None = None,

789 **kwargs: object,

790 - ) -> Union[torch.Tensor, IntermediateTensors]:

790 + ) -> torch.Tensor | IntermediateTensors:

791 if intermediate_tensors is not None:

792 inputs_embeds = None

793

@@ -799,7 +799,7 @@ class GraniteSpeechForConditionalGeneration(

799 def compute_logits(

800 self,

801 hidden_states: torch.Tensor,

802 - ) -> Optional[torch.Tensor]:

802 + ) -> torch.Tensor | None:

803 return self.language_model.compute_logits(hidden_states)

804

805 def load_weights(

vllm/model_executor/models/granitemoe.py

+13 -13

@@ -26,7 +26,7 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Any, Optional

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -79,9 +79,9 @@ class GraniteMoeMoE(nn.Module):

79 top_k: int,

80 hidden_size: int,

81 intermediate_size: int,

82 - params_dtype: Optional[torch.dtype] = None,

83 - quant_config: Optional[QuantizationConfig] = None,

84 - tp_size: Optional[int] = None,

82 + params_dtype: torch.dtype | None = None,

83 + quant_config: QuantizationConfig | None = None,

84 + tp_size: int | None = None,

85 is_sequence_parallel=False,

86 prefix: str = "",

87 ):

@@ -143,10 +143,10 @@ class GraniteMoeAttention(nn.Module):

143 num_kv_heads: int,

144 max_position: int = 4096 * 32,

145 rope_theta: float = 10000,

146 - rope_scaling: Optional[dict[str, Any]] = None,

147 - cache_config: Optional[CacheConfig] = None,

148 - quant_config: Optional[QuantizationConfig] = None,

149 - attention_multiplier: Optional[float] = None,

146 + rope_scaling: dict[str, Any] | None = None,

147 + cache_config: CacheConfig | None = None,

148 + quant_config: QuantizationConfig | None = None,

149 + attention_multiplier: float | None = None,

150 prefix: str = "",

151 ) -> None:

152 super().__init__()

@@ -330,8 +330,8 @@ class GraniteMoeModel(nn.Module):

330 self,

331 input_ids: torch.Tensor,

332 positions: torch.Tensor,

333 - intermediate_tensors: Optional[IntermediateTensors],

334 - inputs_embeds: Optional[torch.Tensor] = None,

333 + intermediate_tensors: IntermediateTensors | None,

334 + inputs_embeds: torch.Tensor | None = None,

335 ) -> torch.Tensor:

336 if get_pp_group().is_first_rank:

337 if inputs_embeds is not None:

@@ -557,15 +557,15 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

557 self,

558 input_ids: torch.Tensor,

559 positions: torch.Tensor,

560 - intermediate_tensors: Optional[IntermediateTensors] = None,

561 - inputs_embeds: Optional[torch.Tensor] = None,

560 + intermediate_tensors: IntermediateTensors | None = None,

561 + inputs_embeds: torch.Tensor | None = None,

562 ) -> torch.Tensor:

563 hidden_states = self.model(

564 input_ids, positions, intermediate_tensors, inputs_embeds

565 )

566 return hidden_states

567

568 - def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:

568 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

569 logits = self.logits_processor(self.lm_head, hidden_states)

570 return logits

571

vllm/model_executor/models/granitemoehybrid.py

+16 -17

@@ -4,7 +4,6 @@

4

5 # Added by the IBM Team, 2025

6 from collections.abc import Iterable

7 -from typing import Optional

7

8 import torch

9 from torch import nn

@@ -50,9 +49,9 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):

49 self,

50 config: GraniteMoeHybridConfig,

51 layer_idx: int,

53 - model_config: Optional[ModelConfig] = None,

54 - cache_config: Optional[CacheConfig] = None,

55 - quant_config: Optional[QuantizationConfig] = None,

52 + model_config: ModelConfig | None = None,

53 + cache_config: CacheConfig | None = None,

54 + quant_config: QuantizationConfig | None = None,

55 prefix: str = "",

56 ) -> None:

57 super().__init__()

@@ -105,7 +104,7 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):

104 def forward(

105 self,

106 hidden_states: torch.Tensor,

108 - residual: Optional[torch.Tensor],

107 + residual: torch.Tensor | None,

108 **kwargs,

109 ):

110 residual = hidden_states

@@ -139,9 +138,9 @@ class GraniteMoeHybridAttentionDecoderLayer(nn.Module):

138 self,

139 config: GraniteMoeHybridConfig,

140 layer_idx: int,

142 - model_config: Optional[ModelConfig] = None,

143 - cache_config: Optional[CacheConfig] = None,

144 - quant_config: Optional[QuantizationConfig] = None,

141 + model_config: ModelConfig | None = None,

142 + cache_config: CacheConfig | None = None,

143 + quant_config: QuantizationConfig | None = None,

144 prefix: str = "",

145 ) -> None:

146 super().__init__()

@@ -183,7 +182,7 @@ class GraniteMoeHybridAttentionDecoderLayer(nn.Module):

182 self,

183 positions: torch.Tensor,

184 hidden_states: torch.Tensor,

186 - residual: Optional[torch.Tensor],

185 + residual: torch.Tensor | None,

186 ) -> torch.Tensor:

187 residual = hidden_states

188 hidden_states = self.input_layernorm(hidden_states)

@@ -218,9 +217,9 @@ class GraniteMoeHybridAttention(nn.Module):

217 def __init__(

218 self,

219 config: GraniteMoeHybridConfig,

221 - model_config: Optional[ModelConfig] = None,

222 - cache_config: Optional[CacheConfig] = None,

223 - quant_config: Optional[QuantizationConfig] = None,

220 + model_config: ModelConfig | None = None,

221 + cache_config: CacheConfig | None = None,

222 + quant_config: QuantizationConfig | None = None,

223 prefix: str = "",

224 ) -> None:

225 super().__init__()

@@ -374,8 +373,8 @@ class GraniteMoeHybridModel(nn.Module):

373 self,

374 input_ids: torch.Tensor,

375 positions: torch.Tensor,

377 - intermediate_tensors: Optional[IntermediateTensors] = None,

378 - inputs_embeds: Optional[torch.Tensor] = None,

376 + intermediate_tensors: IntermediateTensors | None = None,

377 + inputs_embeds: torch.Tensor | None = None,

378 ) -> torch.Tensor:

379 if get_pp_group().is_first_rank:

380 if inputs_embeds is not None:

@@ -614,8 +613,8 @@ class GraniteMoeHybridForCausalLM(

613 self,

614 input_ids: torch.Tensor,

615 positions: torch.Tensor,

617 - intermediate_tensors: Optional[IntermediateTensors] = None,

618 - inputs_embeds: Optional[torch.Tensor] = None,

616 + intermediate_tensors: IntermediateTensors | None = None,

617 + inputs_embeds: torch.Tensor | None = None,

618 **kwargs,

619 ):

620 hidden_states = self.model(

@@ -627,7 +626,7 @@ class GraniteMoeHybridForCausalLM(

626 def compute_logits(

627 self,

628 hidden_states: torch.Tensor,

630 - ) -> Optional[torch.Tensor]:

629 + ) -> torch.Tensor | None:

630 logits = self.logits_processor(self.lm_head, hidden_states)

631 return logits

632

vllm/model_executor/models/granitemoeshared.py

+8 -9

@@ -8,7 +8,6 @@ experts.

8

9 from collections.abc import Iterable

10 from itertools import islice

11 -from typing import Optional

11

12 import torch

13 from torch import nn

@@ -41,7 +40,7 @@ class GraniteMoeSharedMLP(nn.Module):

40 def __init__(

41 self,

42 config: GraniteMoeSharedConfig,

44 - quant_config: Optional[QuantizationConfig] = None,

43 + quant_config: QuantizationConfig | None = None,

44 prefix: str = "",

45 ):

46 super().__init__()

@@ -80,8 +79,8 @@ class GraniteMoeSharedDecoderLayer(nn.Module):

79 def __init__(

80 self,

81 config: GraniteMoeSharedConfig,

83 - cache_config: Optional[CacheConfig] = None,

84 - quant_config: Optional[QuantizationConfig] = None,

82 + cache_config: CacheConfig | None = None,

83 + quant_config: QuantizationConfig | None = None,

84 prefix: str = "",

85 ) -> None:

86 super().__init__()

@@ -198,8 +197,8 @@ class GraniteMoeSharedModel(nn.Module):

197 self,

198 input_ids: torch.Tensor,

199 positions: torch.Tensor,

201 - intermediate_tensors: Optional[IntermediateTensors],

202 - inputs_embeds: Optional[torch.Tensor] = None,

200 + intermediate_tensors: IntermediateTensors | None,

201 + inputs_embeds: torch.Tensor | None = None,

202 ) -> torch.Tensor:

203 if get_pp_group().is_first_rank:

204 if inputs_embeds is not None:

@@ -321,15 +320,15 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

320 self,

321 input_ids: torch.Tensor,

322 positions: torch.Tensor,

324 - intermediate_tensors: Optional[IntermediateTensors] = None,

325 - inputs_embeds: Optional[torch.Tensor] = None,

323 + intermediate_tensors: IntermediateTensors | None = None,

324 + inputs_embeds: torch.Tensor | None = None,

325 ) -> torch.Tensor:

326 hidden_states = self.model(

327 input_ids, positions, intermediate_tensors, inputs_embeds

328 )

329 return hidden_states

330

332 - def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:

331 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

332 logits = self.logits_processor(self.lm_head, hidden_states)

333 return logits

334

vllm/model_executor/models/gritlm.py

+6 -7

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Set

4 -from typing import Optional, Union

4

5 import numpy as np

6 import torch

@@ -62,7 +61,7 @@ class GritLMMeanPool(nn.Module):

61 arr: np.ndarray,

62 target: np.ndarray,

63 start_idx: int = 0,

65 - end_idx: Optional[int] = None,

64 + end_idx: int | None = None,

65 ) -> int:

66 """

67 Find the first occurrence of `target` in `arr` starting from

@@ -152,8 +151,8 @@ class GritLMMeanPool(nn.Module):

151 def forward_one(

152 self,

153 hidden_states: torch.Tensor,

155 - prompt_len: Optional[torch.Tensor] = None,

156 - instr_len: Optional[torch.Tensor] = None,

154 + prompt_len: torch.Tensor | None = None,

155 + instr_len: torch.Tensor | None = None,

156 ) -> torch.Tensor:

157 assert prompt_len is None or prompt_len == hidden_states.shape[0], (

158 "partial prefill not supported with MEAN pooling"

@@ -166,7 +165,7 @@ class GritLMMeanPool(nn.Module):

165 hidden_states: torch.Tensor,

166 prompt_lens: torch.Tensor,

167 instr_lens: torch.Tensor,

169 - ) -> Union[list[torch.Tensor], torch.Tensor]:

168 + ) -> list[torch.Tensor] | torch.Tensor:

169 offset = 0

170 pooled_data = list[torch.Tensor]()

171

@@ -182,9 +181,9 @@ class GritLMMeanPool(nn.Module):

181

182 def forward(

183 self,

185 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

184 + hidden_states: torch.Tensor | list[torch.Tensor],

185 pooling_metadata: PoolingMetadata,

187 - ) -> Union[list[torch.Tensor], torch.Tensor]:

186 + ) -> list[torch.Tensor] | torch.Tensor:

187 prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)

188 instr_lens = torch.tensor(

189 [

vllm/model_executor/models/grok1.py

+15 -16

@@ -25,7 +25,6 @@

25

26 from collections.abc import Iterable

27 from itertools import islice

28 -from typing import Optional, Union

28

29 import torch

30 import torch.nn.functional as F

@@ -86,9 +85,9 @@ class Grok1MoE(nn.Module):

85 top_k: int,

86 hidden_size: int,

87 intermediate_size: int,

89 - params_dtype: Optional[torch.dtype] = None,

90 - quant_config: Optional[QuantizationConfig] = None,

91 - tp_size: Optional[int] = None,

88 + params_dtype: torch.dtype | None = None,

89 + quant_config: QuantizationConfig | None = None,

90 + tp_size: int | None = None,

91 prefix: str = "",

92 ):

93 super().__init__()

@@ -137,8 +136,8 @@ class Grok1Attention(nn.Module):

136 num_kv_heads: int,

137 max_position: int = 4096 * 32,

138 rope_theta: float = 10000,

140 - cache_config: Optional[CacheConfig] = None,

141 - quant_config: Optional[QuantizationConfig] = None,

139 + cache_config: CacheConfig | None = None,

140 + quant_config: QuantizationConfig | None = None,

141 prefix: str = "",

142 config=None, # Added config parameter

143 ) -> None:

@@ -223,8 +222,8 @@ class Grok1DecoderLayer(nn.Module):

222 def __init__(

223 self,

224 config,

226 - cache_config: Optional[CacheConfig] = None,

227 - quant_config: Optional[QuantizationConfig] = None,

225 + cache_config: CacheConfig | None = None,

226 + quant_config: QuantizationConfig | None = None,

227 prefix: str = "",

228 ) -> None:

229 super().__init__()

@@ -273,7 +272,7 @@ class Grok1DecoderLayer(nn.Module):

272 self,

273 positions: torch.Tensor,

274 hidden_states: torch.Tensor,

276 - residual: Optional[torch.Tensor],

275 + residual: torch.Tensor | None,

276 ) -> tuple[torch.Tensor, torch.Tensor]:

277 # Self Attention

278 if residual is None:

@@ -351,9 +350,9 @@ class Grok1Model(nn.Module):

350 self,

351 input_ids: torch.Tensor,

352 positions: torch.Tensor,

354 - intermediate_tensors: Optional[IntermediateTensors],

355 - inputs_embeds: Optional[torch.Tensor] = None,

356 - ) -> Union[torch.Tensor, IntermediateTensors]:

353 + intermediate_tensors: IntermediateTensors | None,

354 + inputs_embeds: torch.Tensor | None = None,

355 + ) -> torch.Tensor | IntermediateTensors:

356 if get_pp_group().is_first_rank:

357 if inputs_embeds is not None:

358 hidden_states = inputs_embeds

@@ -544,9 +543,9 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

543 self,

544 input_ids: torch.Tensor,

545 positions: torch.Tensor,

547 - intermediate_tensors: Optional[IntermediateTensors] = None,

548 - inputs_embeds: Optional[torch.Tensor] = None,

549 - ) -> Union[torch.Tensor, IntermediateTensors]:

546 + intermediate_tensors: IntermediateTensors | None = None,

547 + inputs_embeds: torch.Tensor | None = None,

548 + ) -> torch.Tensor | IntermediateTensors:

549 hidden_states = self.model(

550 input_ids, positions, intermediate_tensors, inputs_embeds

551 )

@@ -555,7 +554,7 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

554 def compute_logits(

555 self,

556 hidden_states: torch.Tensor,

558 - ) -> Optional[torch.Tensor]:

557 + ) -> torch.Tensor | None:

558 logits = self.logits_processor(self.lm_head, hidden_states)

559 return logits

560

vllm/model_executor/models/h2ovl.py

+26 -27

@@ -9,7 +9,6 @@

9 # Licensed under Apache 2.0 License [see LICENSE for details]

10 # --------------------------------------------------------

11 from collections.abc import Mapping, Sequence

12 -from typing import Optional, Union

12

13 import torch

14 from PIL import Image

@@ -67,7 +66,7 @@ def get_h2ovl_target_ratios(

66 min_num: int,

67 max_num: int,

68 *,

70 - prior_aspect_ratio: Optional[tuple[int, int]],

69 + prior_aspect_ratio: tuple[int, int] | None,

70 ) -> list[tuple[int, int]]:

71 target_ratios = get_internvl_target_ratios(min_num, max_num)

72

@@ -170,7 +169,7 @@ def _preprocess_image(

169 min_num: int,

170 max_num: int,

171 use_thumbnail: bool,

173 - prior_aspect_ratio: Optional[tuple[int, int]],

172 + prior_aspect_ratio: tuple[int, int] | None,

173 ) -> tuple[torch.Tensor, tuple[int, int]]:

174 target_ratios = get_h2ovl_target_ratios(

175 min_num,

@@ -244,10 +243,10 @@ class H2OVLProcessor(BaseInternVLProcessor):

243 config: PretrainedConfig,

244 tokenizer: AnyTokenizer,

245 *,

247 - min_dynamic_patch: Optional[int] = None,

248 - max_dynamic_patch: Optional[int] = None,

249 - dynamic_image_size: Optional[bool] = None,

250 - use_msac: Optional[bool] = None,

246 + min_dynamic_patch: int | None = None,

247 + max_dynamic_patch: int | None = None,

248 + dynamic_image_size: bool | None = None,

249 + use_msac: bool | None = None,

250 ) -> None:

251 super().__init__(

252 config,

@@ -270,7 +269,7 @@ class H2OVLProcessor(BaseInternVLProcessor):

269 def get_image_repl(

270 self,

271 feature_size: int,

273 - num_patches: Optional[int],

272 + num_patches: int | None,

273 ) -> PromptUpdateDetails[str]:

274 repl_features = IMG_CONTEXT * feature_size

275 repl_full = IMG_START + repl_features + IMG_END

@@ -280,10 +279,10 @@ class H2OVLProcessor(BaseInternVLProcessor):

279 def resolve_min_max_num(

280 self,

281 *,

283 - min_dynamic_patch: Optional[int] = None,

284 - max_dynamic_patch: Optional[int] = None,

285 - dynamic_image_size: Optional[bool] = None,

286 - use_thumbnail: Optional[bool] = None,

282 + min_dynamic_patch: int | None = None,

283 + max_dynamic_patch: int | None = None,

284 + dynamic_image_size: bool | None = None,

285 + use_thumbnail: bool | None = None,

286 ) -> tuple[int, int]:

287 min_dynamic_patch = (

288 self.min_dynamic_patch if min_dynamic_patch is None else min_dynamic_patch

@@ -308,12 +307,12 @@ class H2OVLProcessor(BaseInternVLProcessor):

307 def resolve_target_ratios(

308 self,

309 *,

311 - min_dynamic_patch: Optional[int] = None,

312 - max_dynamic_patch: Optional[int] = None,

313 - dynamic_image_size: Optional[bool] = None,

314 - use_thumbnail: Optional[bool] = None,

315 - prior_aspect_ratio: Optional[tuple[int, int]] = None,

316 - override_min_num: Optional[int] = None,

310 + min_dynamic_patch: int | None = None,

311 + max_dynamic_patch: int | None = None,

312 + dynamic_image_size: bool | None = None,

313 + use_thumbnail: bool | None = None,

314 + prior_aspect_ratio: tuple[int, int] | None = None,

315 + override_min_num: int | None = None,

316 ) -> list[tuple[int, int]]:

317 min_num, max_num = self.resolve_min_max_num(

318 min_dynamic_patch=min_dynamic_patch,

@@ -335,7 +334,7 @@ class H2OVLProcessor(BaseInternVLProcessor):

334 *,

335 image_width: int,

336 image_height: int,

338 - use_msac: Optional[bool] = None,

337 + use_msac: bool | None = None,

338 ) -> int:

339 use_msac = self.use_msac if use_msac is None else use_msac

340

@@ -385,9 +384,9 @@ class H2OVLProcessor(BaseInternVLProcessor):

384 def _images_to_pixel_values_lst(

385 self,

386 images: list[Image.Image],

388 - min_dynamic_patch: Optional[int] = None,

389 - max_dynamic_patch: Optional[int] = None,

390 - dynamic_image_size: Optional[bool] = None,

387 + min_dynamic_patch: int | None = None,

388 + max_dynamic_patch: int | None = None,

389 + dynamic_image_size: bool | None = None,

390 ) -> list[torch.Tensor]:

391 use_msac = self.use_msac if len(images) == 1 else False

392

@@ -425,8 +424,8 @@ class H2OVLProcessingInfo(BaseInternVLProcessingInfo):

424 *,

425 image_width: int,

426 image_height: int,

428 - processor: Optional[H2OVLProcessor],

429 - use_msac: Optional[bool] = None,

427 + processor: H2OVLProcessor | None,

428 + use_msac: bool | None = None,

429 ) -> int:

430 if processor is None:

431 processor = self.get_hf_processor()

@@ -493,11 +492,11 @@ class H2OVLMultiModalProcessor(BaseInternVLMultiModalProcessor[H2OVLProcessingIn

492

493 def _cached_apply_hf_processor(

494 self,

496 - prompt: Union[str, list[int]],

495 + prompt: str | list[int],

496 mm_data_items: MultiModalDataItems,

497 hf_processor_mm_kwargs: Mapping[str, object],

498 tokenization_kwargs: Mapping[str, object],

500 - mm_uuids: Optional[MultiModalUUIDDict] = None,

499 + mm_uuids: MultiModalUUIDDict | None = None,

500 ) -> tuple[list[int], MultiModalProcessingInfo, bool]:

501 # The processor logic is different for len(images) <= 1 vs > 1

502 # Since the processing cache assumes that the processor output is

@@ -530,7 +529,7 @@ class H2OVLChatModel(InternVLChatModel):

529 def _init_vision_model(

530 self,

531 config: PretrainedConfig,

533 - quant_config: Optional[QuantizationConfig],

532 + quant_config: QuantizationConfig | None,

533 *,

534 is_mono: bool,

535 prefix: str,

vllm/model_executor/models/hunyuan_v1.py

+23 -23

@@ -27,7 +27,7 @@

27 import typing

28 from collections.abc import Callable, Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import regex as re

33 import torch

@@ -102,7 +102,7 @@ class HunYuanMLP(nn.Module):

102 hidden_size: int,

103 intermediate_size: int,

104 hidden_act: str,

105 - quant_config: Optional[QuantizationConfig] = None,

105 + quant_config: QuantizationConfig | None = None,

106 bias: bool = False,

107 prefix: str = "",

108 reduce_results: bool = True,

@@ -144,11 +144,11 @@ class HunYuanAttention(nn.Module):

144 num_heads: int,

145 num_kv_heads: int,

146 rope_theta: float = 10000,

147 - rope_scaling: Optional[dict[str, Any]] = None,

147 + rope_scaling: dict[str, Any] | None = None,

148 max_position_embeddings: int = 8192,

149 - quant_config: Optional[QuantizationConfig] = None,

149 + quant_config: QuantizationConfig | None = None,

150 bias: bool = False,

151 - cache_config: Optional[CacheConfig] = None,

151 + cache_config: CacheConfig | None = None,

152 prefix: str = "",

153 layer_id: int = -1,

154 ) -> None:

@@ -227,7 +227,7 @@ class HunYuanAttention(nn.Module):

227 self,

228 positions: torch.Tensor,

229 hidden_states: torch.Tensor,

230 - kv_states: Optional[tuple[torch.Tensor]] = None,

230 + kv_states: tuple[torch.Tensor] | None = None,

231 ) -> torch.Tensor:

232 qkv, _ = self.qkv_proj(hidden_states)

233 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

@@ -256,11 +256,11 @@ class HunYuanCrossAttention(nn.Module):

256 num_heads: int,

257 num_kv_heads: int,

258 rope_theta: float = 10000,

259 - rope_scaling: Optional[dict[str, Any]] = None,

259 + rope_scaling: dict[str, Any] | None = None,

260 max_position_embeddings: int = 8192,

261 - quant_config: Optional[QuantizationConfig] = None,

261 + quant_config: QuantizationConfig | None = None,

262 bias: bool = False,

263 - cache_config: Optional[CacheConfig] = None,

263 + cache_config: CacheConfig | None = None,

264 prefix: str = "",

265 layer_id: int = -1,

266 ) -> None:

@@ -338,7 +338,7 @@ class HunYuanCrossAttention(nn.Module):

338 self,

339 positions: torch.Tensor,

340 hidden_states: torch.Tensor,

341 - kv_states: Optional[tuple[torch.Tensor]] = None,

341 + kv_states: tuple[torch.Tensor] | None = None,

342 ) -> torch.Tensor:

343 assert kv_states is not None

344 ori_k, v = kv_states # use last layer kv,

@@ -365,7 +365,7 @@ class HunYuanSparseMoeBlock(nn.Module):

365 def __init__(

366 self,

367 config: PretrainedConfig,

368 - quant_config: Optional[QuantizationConfig] = None,

368 + quant_config: QuantizationConfig | None = None,

369 layer_id: int = -1,

370 prefix: str = "",

371 enable_eplb: bool = False,

@@ -480,8 +480,8 @@ class HunYuanDecoderLayer(nn.Module):

480 def __init__(

481 self,

482 config: PretrainedConfig,

483 - cache_config: Optional[CacheConfig] = None,

484 - quant_config: Optional[QuantizationConfig] = None,

483 + cache_config: CacheConfig | None = None,

484 + quant_config: QuantizationConfig | None = None,

485 prefix: str = "",

486 layer_id: int = -1,

487 enable_eplb: bool = False,

@@ -577,8 +577,8 @@ class HunYuanDecoderLayer(nn.Module):

577 self,

578 positions: torch.Tensor,

579 hidden_states: torch.Tensor,

580 - residual: Optional[torch.Tensor],

581 - kv_states: Optional[tuple[torch.Tensor]] = None,

580 + residual: torch.Tensor | None,

581 + kv_states: tuple[torch.Tensor] | None = None,

582 ) -> tuple[torch.Tensor, torch.Tensor]:

583 # Self Attention

584 if residual is None:

@@ -654,11 +654,11 @@ class HunYuanModel(nn.Module):

654

655 def forward(

656 self,

657 - input_ids: Optional[torch.Tensor],

657 + input_ids: torch.Tensor | None,

658 positions: torch.Tensor,

659 - intermediate_tensors: Optional[IntermediateTensors],

660 - inputs_embeds: Optional[torch.Tensor] = None,

661 - ) -> Union[torch.Tensor, IntermediateTensors]:

659 + intermediate_tensors: IntermediateTensors | None,

660 + inputs_embeds: torch.Tensor | None = None,

661 + ) -> torch.Tensor | IntermediateTensors:

662 if get_pp_group().is_first_rank:

663 if inputs_embeds is not None:

664 hidden_states = inputs_embeds

@@ -960,9 +960,9 @@ class HunyuanV1ModelBase(nn.Module, SupportsLoRA, SupportsPP):

960 self,

961 input_ids: torch.Tensor,

962 positions: torch.Tensor,

963 - intermediate_tensors: Optional[IntermediateTensors] = None,

964 - inputs_embeds: Optional[torch.Tensor] = None,

965 - ) -> Union[torch.Tensor, IntermediateTensors]:

963 + intermediate_tensors: IntermediateTensors | None = None,

964 + inputs_embeds: torch.Tensor | None = None,

965 + ) -> torch.Tensor | IntermediateTensors:

966 model_output = self.model(

967 input_ids, positions, intermediate_tensors, inputs_embeds

968 )

@@ -971,7 +971,7 @@ class HunyuanV1ModelBase(nn.Module, SupportsLoRA, SupportsPP):

971 def compute_logits(

972 self,

973 hidden_states: torch.Tensor,

974 - ) -> Optional[torch.Tensor]:

974 + ) -> torch.Tensor | None:

975 logits = self.logits_processor(self.lm_head, hidden_states)

976 return logits

977

vllm/model_executor/models/hyperclovax_vision.py

+25 -25

@@ -6,7 +6,7 @@ from collections import defaultdict

6 from collections.abc import Iterable, Mapping, Sequence

7 from functools import partial

8 from itertools import accumulate

9 -from typing import Annotated, Any, Literal, Optional, Union

9 +from typing import Annotated, Any, Literal

10

11 import numpy as np

12 import torch

@@ -115,13 +115,13 @@ class HCXVisionProcessingInfo(BaseProcessingInfo):

115 def get_vision_encoder_info(self):

116 return get_vision_encoder_info(self.get_hf_config())

117

118 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

118 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

119 return {"image": None, "video": None}

120

121 def get_num_image_tokens(

122 self,

123 *,

124 - vision_query_length: Union[int, list[int]],

124 + vision_query_length: int | list[int],

125 ) -> int:

126 if isinstance(vision_query_length, int):

127 return vision_query_length

@@ -131,7 +131,7 @@ class HCXVisionProcessingInfo(BaseProcessingInfo):

131 def get_num_video_tokens(

132 self,

133 *,

134 - vision_query_length: Union[int, list[int]],

134 + vision_query_length: int | list[int],

135 ) -> int:

136 if isinstance(vision_query_length, int):

137 return vision_query_length

@@ -166,7 +166,7 @@ class HCXVisionDummyInputsBuilder(BaseDummyInputsBuilder[HCXVisionProcessingInfo

166 self,

167 seq_len: int,

168 mm_counts: Mapping[str, int],

169 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

169 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

170 ) -> MultiModalDataDict:

171 num_images = mm_counts.get("image", 0)

172 num_videos = mm_counts.get("video", 0)

@@ -346,7 +346,7 @@ def _build_hcxvision_hf_processor(

346 info: HCXVisionProcessingInfo,

347 dummy_inputs: BaseDummyInputsBuilder[HCXVisionProcessingInfo],

348 *,

349 - cache: Optional[BaseMultiModalProcessorCache] = None,

349 + cache: BaseMultiModalProcessorCache | None = None,

350 ) -> BaseMultiModalProcessor:

351 if isinstance(info, HCXVisionProcessingInfo):

352 return HCXVisionMultiModalProcessor(

@@ -360,12 +360,12 @@ def _build_hcxvision_hf_processor(

360

361 def init_vision_tower_for_hcxvision(

362 vision_config,

363 - quant_config: Optional[QuantizationConfig],

363 + quant_config: QuantizationConfig | None,

364 *,

365 - use_nth_layer: Optional[int] = None,

366 - require_post_norm: Optional[bool] = None,

365 + use_nth_layer: int | None = None,

366 + require_post_norm: bool | None = None,

367 prefix: str = "",

368 -) -> Union[CLIPVisionModel, SiglipVisionModel]:

368 +) -> CLIPVisionModel | SiglipVisionModel:

369 num_hidden_layers = vision_config.num_hidden_layers

370 if not isinstance(use_nth_layer, int):

371 pass

@@ -473,8 +473,8 @@ class HCXVisionCAbstractor(nn.Module):

473 def forward(

474 self,

475 x: torch.Tensor,

476 - num_queries_vis_abstractors: Optional[list[list[int]]] = None,

477 - num_grids: Optional[list[int]] = None,

476 + num_queries_vis_abstractors: list[list[int]] | None = None,

477 + num_grids: list[int] | None = None,

478 ) -> torch.Tensor:

479 if self.prenorm is not None:

480 x = self.prenorm(x)

@@ -493,8 +493,8 @@ class HCXVisionCAbstractor(nn.Module):

493 def _forward(

494 self,

495 x: torch.Tensor,

496 - num_queries_vis_abstractors: Optional[list[list[int]]] = None,

497 - num_grids: Optional[list[int]] = None,

496 + num_queries_vis_abstractors: list[list[int]] | None = None,

497 + num_grids: list[int] | None = None,

498 ) -> torch.Tensor:

499 # x: [B, L, dim]

500 B, L, dim = x.shape

@@ -515,8 +515,8 @@ class HCXVisionCAbstractor(nn.Module):

515 def _forward_adaptive_num_query(

516 self,

517 x: torch.Tensor,

518 - num_queries_vis_abstractors: Optional[list[list[int]]] = None,

519 - num_grids: Optional[list[int]] = None,

518 + num_queries_vis_abstractors: list[list[int]] | None = None,

519 + num_grids: list[int] | None = None,

520 ) -> list[torch.Tensor]:

521 # self.net is consisted by 3 layers (s1, sampler, s2)

522 assert len(self.net) == 3

@@ -604,7 +604,7 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

604 *,

605 vllm_config: VllmConfig,

606 prefix: str = "",

607 - **kwargs: Optional[Any],

607 + **kwargs: Any | None,

608 ) -> None:

609 super().__init__()

610

@@ -662,7 +662,7 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

662 # self.reduction = self._init_reduction_type(use_sum_loss)

663

664 @classmethod

665 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

665 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

666 if modality.startswith("image"):

667 return IMAGE_TOKEN

668 if modality.startswith("video"):

@@ -673,7 +673,7 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

673 def _parse_and_validate_image_input(

674 self,

675 **kwargs: object,

676 - ) -> Optional[HCXVisionImageInputs]:

676 + ) -> HCXVisionImageInputs | None:

677 pixel_values_images = kwargs.pop("pixel_values_images", None)

678

679 if pixel_values_images is None:

@@ -689,7 +689,7 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

689 def _parse_and_validate_video_input(

690 self,

691 **kwargs: object,

692 - ) -> Optional[HCXVisionVideoInputs]:

692 + ) -> HCXVisionVideoInputs | None:

693 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

694

695 if pixel_values_videos is None:

@@ -762,10 +762,10 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

762 self,

763 input_ids: torch.Tensor,

764 positions: torch.Tensor,

765 - intermediate_tensors: Optional[IntermediateTensors] = None,

766 - inputs_embeds: Optional[torch.Tensor] = None,

765 + intermediate_tensors: IntermediateTensors | None = None,

766 + inputs_embeds: torch.Tensor | None = None,

767 **kwargs: object,

768 - ) -> Union[torch.Tensor, IntermediateTensors]:

768 + ) -> torch.Tensor | IntermediateTensors:

769 if intermediate_tensors is not None:

770 inputs_embeds = None

771

@@ -946,7 +946,7 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

946 def compute_logits(

947 self,

948 hidden_states: torch.Tensor,

949 - ) -> Optional[torch.Tensor]:

949 + ) -> torch.Tensor | None:

950 return self.language_model.compute_logits(hidden_states)

951

952 def load_weights(

@@ -1062,7 +1062,7 @@ def select_best_resolution(original_size: tuple, possible_resolutions: list) ->

1062

1063 def get_anyres_image_grid_shape(

1064 image_size: tuple[int, int],

1065 - grid_pinpoints: Union[str, list[tuple[int, int]]],

1065 + grid_pinpoints: str | list[tuple[int, int]],

1066 patch_size: int,

1067 ) -> tuple[int, int]:

1068 possible_resolutions = (

vllm/model_executor/models/idefics2_vision_model.py

+10 -11

@@ -19,7 +19,6 @@

19 """PyTorch Idefics2 model."""

20

21 from collections.abc import Iterable

22 -from typing import Optional

22

23 import torch

24 from torch import nn

@@ -77,7 +76,7 @@ class Idefics2VisionEmbeddings(nn.Module):

76 self,

77 pixel_values: torch.FloatTensor,

78 patch_attention_mask: torch.BoolTensor,

80 - tgt_sizes: Optional[torch.IntTensor] = None,

79 + tgt_sizes: torch.IntTensor | None = None,

80 ) -> torch.Tensor:

81 batch_size, _, max_im_h, max_im_w = pixel_values.shape

82 target_dtype = self.patch_embedding.weight.dtype

@@ -124,7 +123,7 @@ class Idefics2VisionAttention(nn.Module):

123 def __init__(

124 self,

125 config: Idefics2VisionConfig,

127 - quant_config: Optional[QuantizationConfig] = None,

126 + quant_config: QuantizationConfig | None = None,

127 prefix: str = "",

128 use_data_parallel: bool = False,

129 ) -> None:

@@ -185,7 +184,7 @@ class Idefics2VisionMLP(nn.Module):

184 def __init__(

185 self,

186 config: Idefics2VisionConfig,

188 - quant_config: Optional[QuantizationConfig] = None,

187 + quant_config: QuantizationConfig | None = None,

188 prefix: str = "",

189 use_data_parallel: bool = False,

190 ) -> None:

@@ -220,7 +219,7 @@ class Idefics2EncoderLayer(nn.Module):

219 def __init__(

220 self,

221 config: Idefics2Config,

223 - quant_config: Optional[QuantizationConfig] = None,

222 + quant_config: QuantizationConfig | None = None,

223 prefix: str = "",

224 use_data_parallel: bool = False,

225 ) -> None:

@@ -275,9 +274,9 @@ class Idefics2Encoder(nn.Module):

274 def __init__(

275 self,

276 config: Idefics2Config,

278 - quant_config: Optional[QuantizationConfig] = None,

277 + quant_config: QuantizationConfig | None = None,

278 *,

280 - num_hidden_layers_override: Optional[int] = None,

279 + num_hidden_layers_override: int | None = None,

280 prefix: str = "",

281 use_data_parallel: bool = False,

282 ) -> None:

@@ -326,9 +325,9 @@ class Idefics2VisionTransformer(nn.Module):

325 def __init__(

326 self,

327 config: Idefics2VisionConfig,

329 - quant_config: Optional[QuantizationConfig] = None,

328 + quant_config: QuantizationConfig | None = None,

329 *,

331 - num_hidden_layers_override: Optional[int] = None,

330 + num_hidden_layers_override: int | None = None,

331 require_post_norm: bool = True,

332 prefix: str = "",

333 use_data_parallel: bool = False,

@@ -370,8 +369,8 @@ class Idefics2VisionTransformer(nn.Module):

369 def forward(

370 self,

371 pixel_values,

373 - patch_attention_mask: Optional[torch.BoolTensor] = None,

374 - tgt_sizes: Optional[torch.IntTensor] = None,

372 + patch_attention_mask: torch.BoolTensor | None = None,

373 + tgt_sizes: torch.IntTensor | None = None,

374 ) -> torch.Tensor:

375 hidden_states = self.embeddings(

376 pixel_values=pixel_values,

vllm/model_executor/models/idefics3.py

+22 -24

@@ -18,7 +18,7 @@

18

19 import math

20 from collections.abc import Iterable, Mapping, Sequence

21 -from typing import Annotated, Literal, Optional, Union

21 +from typing import Annotated, Literal, TypeAlias

22

23 import torch

24 from torch import nn

@@ -91,14 +91,14 @@ class Idefics3ImageEmbeddingInputs(TensorSchema):

91 data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")]

92

93

94 -ImageInputs = Union[Idefics3ImagePixelInputs, Idefics3ImageEmbeddingInputs]

94 +ImageInputs: TypeAlias = Idefics3ImagePixelInputs | Idefics3ImageEmbeddingInputs

95

96

97 class Idefics3ProcessingInfo(BaseProcessingInfo):

98 def get_hf_processor(self, **kwargs: object) -> Idefics3Processor:

99 return self.ctx.get_hf_processor(Idefics3Processor, **kwargs)

100

101 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

101 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

102 return {"image": None}

103

104 def _resize_output_size(

@@ -106,9 +106,9 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):

106 *,

107 height: int,

108 width: int,

109 - max_len: Optional[int] = None,

109 + max_len: int | None = None,

110 min_len: int = 1,

111 - max_size: Optional[int] = None,

111 + max_size: int | None = None,

112 ) -> tuple[int, int]:

113 # Set default value for max_len if not provided

114 max_len = max(height, width) if max_len is None else max_len

@@ -165,7 +165,7 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):

165 *,

166 image_width: int,

167 image_height: int,

168 - processor: Optional[Idefics3Processor],

168 + processor: Idefics3Processor | None,

169 ) -> tuple[int, int]:

170 if processor is None:

171 processor = self.get_hf_processor()

@@ -197,7 +197,7 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):

197 *,

198 image_width: int,

199 image_height: int,

200 - processor: Optional[Idefics3Processor],

200 + processor: Idefics3Processor | None,

201 ) -> int:

202 grid_w, grid_h = self._get_image_feature_grid_size(

203 image_width=image_width,

@@ -208,7 +208,7 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):

208 return grid_w * grid_h + 1

209

210 def _get_image_token(

211 - self, processor: Optional[Idefics3Processor]

211 + self, processor: Idefics3Processor | None

212 ) -> tuple[str, str, str]:

213 if processor is None:

214 processor = self.get_hf_processor()

@@ -223,7 +223,7 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):

223 *,

224 image_width: int,

225 image_height: int,

226 - processor: Optional[Idefics3Processor],

226 + processor: Idefics3Processor | None,

227 ) -> str:

228 if processor is None:

229 processor = self.get_hf_processor()

@@ -269,7 +269,7 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):

269 *,

270 image_width: int,

271 image_height: int,

272 - processor: Optional[Idefics3Processor],

272 + processor: Idefics3Processor | None,

273 ) -> int:

274 if processor is None:

275 processor = self.get_hf_processor()

@@ -305,7 +305,7 @@ class Idefics3DummyInputsBuilder(BaseDummyInputsBuilder[Idefics3ProcessingInfo])

305 self,

306 seq_len: int,

307 mm_counts: Mapping[str, int],

308 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

308 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

309 ) -> MultiModalDataDict:

310 num_images = mm_counts.get("image", 0)

311 hf_processor = self.info.get_hf_processor()

@@ -425,7 +425,7 @@ class Idefics3SimpleMLP(nn.Module):

425 def __init__(

426 self,

427 config: Idefics3Config,

428 - quant_config: Optional[QuantizationConfig] = None,

428 + quant_config: QuantizationConfig | None = None,

429 prefix: str = "",

430 ):

431 super().__init__()

@@ -448,7 +448,7 @@ class Idefics3Connector(nn.Module):

448 def __init__(

449 self,

450 config: Idefics3Config,

451 - quant_config: Optional[QuantizationConfig] = None,

451 + quant_config: QuantizationConfig | None = None,

452 prefix: str = "",

453 ):

454 super().__init__()

@@ -557,9 +557,9 @@ class Idefics3Model(nn.Module):

557 self,

558 input_ids: torch.Tensor,

559 positions: torch.Tensor,

560 - intermediate_tensors: Optional[IntermediateTensors] = None,

561 - inputs_embeds: Optional[torch.Tensor] = None,

562 - ) -> Union[torch.Tensor, IntermediateTensors]:

560 + intermediate_tensors: IntermediateTensors | None = None,

561 + inputs_embeds: torch.Tensor | None = None,

562 + ) -> torch.Tensor | IntermediateTensors:

563 hidden_states = self.text_model(

564 input_ids,

565 positions,

@@ -590,7 +590,7 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo

590 }

591

592 @classmethod

593 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

593 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

594 if modality.startswith("image"):

595 return "<image>"

596

@@ -621,9 +621,7 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo

621 self.lm_head.weight = self.model.text_model.embed_tokens.weight

622 self.logits_processor = LogitsProcessor(config.text_config.vocab_size)

623

624 - def _parse_and_validate_image_input(

625 - self, **kwargs: object

626 - ) -> Optional[ImageInputs]:

624 + def _parse_and_validate_image_input(self, **kwargs: object) -> ImageInputs | None:

625 pixel_values = kwargs.pop("pixel_values", None)

626 image_embeds = kwargs.pop("image_embeds", None)

627

@@ -663,7 +661,7 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo

661 def _process_image_input(

662 self,

663 image_input: ImageInputs,

666 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

664 + ) -> torch.Tensor | list[torch.Tensor]:

665 if image_input["type"] == "image_embeds":

666 return image_input["data"]

667

@@ -687,10 +685,10 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo

685 self,

686 input_ids: torch.Tensor,

687 positions: torch.Tensor,

690 - intermediate_tensors: Optional[IntermediateTensors] = None,

691 - inputs_embeds: Optional[torch.Tensor] = None,

688 + intermediate_tensors: IntermediateTensors | None = None,

689 + inputs_embeds: torch.Tensor | None = None,

690 **kwargs: object,

693 - ) -> Union[torch.Tensor, IntermediateTensors]:

691 + ) -> torch.Tensor | IntermediateTensors:

692 if intermediate_tensors is not None:

693 inputs_embeds = None

694

vllm/model_executor/models/interfaces.py

+67 -65

@@ -1,15 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from collections.abc import Iterable, Mapping, MutableSequence

4 +from collections.abc import Callable, Iterable, Mapping, MutableSequence

5 from typing import (

6 TYPE_CHECKING,

7 - Callable,

7 ClassVar,

8 Literal,

10 - Optional,

9 Protocol,

12 - Union,

10 + TypeAlias,

11 overload,

12 runtime_checkable,

13 )

@@ -34,10 +32,14 @@ if TYPE_CHECKING:

32 from vllm.config import VllmConfig

33 from vllm.model_executor.models.utils import WeightsMapper

34 from vllm.sequence import IntermediateTensors

35 +else:

36 + VllmConfig = object

37 + WeightsMapper = object

38 + IntermediateTensors = object

39

40 logger = init_logger(__name__)

41

40 -MultiModalEmbeddings = Union[list[Tensor], Tensor, tuple[Tensor, ...]]

42 +MultiModalEmbeddings: TypeAlias = list[Tensor] | Tensor | tuple[Tensor, ...]

43 """

44 The output embeddings must be one of the following formats:

45

@@ -79,7 +81,7 @@ class SupportsMultiModal(Protocol):

81 """

82

83 @classmethod

82 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

84 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

85 """

86 Get the placeholder text for the `i`th `modality` item in the prompt.

87 """

@@ -127,7 +129,7 @@ class SupportsMultiModal(Protocol):

129 input_ids: Tensor,

130 get_input_embeddings: Callable[[Tensor], Tensor],

131 *,

130 - is_multimodal: Optional[Tensor],

132 + is_multimodal: Tensor | None,

133 handle_oov_mm_token: bool,

134 ) -> Tensor:

135 if handle_oov_mm_token and is_multimodal is not None:

@@ -145,9 +147,9 @@ class SupportsMultiModal(Protocol):

147 def get_input_embeddings(

148 self,

149 input_ids: Tensor,

148 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

150 + multimodal_embeddings: MultiModalEmbeddings | None = None,

151 *,

150 - is_multimodal: Optional[Tensor] = None,

152 + is_multimodal: Tensor | None = None,

153 handle_oov_mm_token: bool = False,

154 ) -> Tensor:

155 """

@@ -236,16 +238,16 @@ def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]: ...

238

239

240 def supports_multimodal(

239 - model: Union[type[object], object],

240 -) -> Union[TypeIs[type[SupportsMultiModal]], TypeIs[SupportsMultiModal]]:

241 + model: type[object] | object,

242 +) -> TypeIs[type[SupportsMultiModal]] | TypeIs[SupportsMultiModal]:

243 return getattr(model, "supports_multimodal", False)

244

245

244 -def supports_multimodal_raw_input_only(model: Union[type[object], object]) -> bool:

246 +def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:

247 return getattr(model, "supports_multimodal_raw_input_only", False)

248

249

248 -def supports_multimodal_encoder_tp_data(model: Union[type[object], object]) -> bool:

250 +def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:

251 return getattr(model, "supports_encoder_tp_data", False)

252

253

@@ -260,8 +262,8 @@ def supports_multimodal_pruning(model: object) -> TypeIs[SupportsMultiModalPruni

262

263

264 def supports_multimodal_pruning(

263 - model: Union[type[object], object],

264 -) -> Union[TypeIs[type[SupportsMultiModalPruning]], TypeIs[SupportsMultiModalPruning]]:

265 + model: type[object] | object,

266 +) -> TypeIs[type[SupportsMultiModalPruning]] | TypeIs[SupportsMultiModalPruning]:

267 return getattr(model, "supports_multimodal_pruning", False)

268

269

@@ -279,7 +281,7 @@ class SupportsScoreTemplate(Protocol):

281 """

282

283 @classmethod

282 - def get_score_template(cls, query: str, document: str) -> Optional[str]:

284 + def get_score_template(cls, query: str, document: str) -> str | None:

285 """

286 Generate a full prompt by populating the score template with query and document content.

287 """ # noqa: E501

@@ -304,8 +306,8 @@ def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]: ...

306

307

308 def supports_score_template(

307 - model: Union[type[object], object],

308 -) -> Union[TypeIs[type[SupportsScoreTemplate]], TypeIs[SupportsScoreTemplate]]:

309 + model: type[object] | object,

310 +) -> TypeIs[type[SupportsScoreTemplate]] | TypeIs[SupportsScoreTemplate]:

311 return getattr(model, "supports_score_template", False)

312

313

@@ -348,8 +350,8 @@ def supports_lora(model: object) -> TypeIs[SupportsLoRA]: ...

350

351

352 def supports_lora(

351 - model: Union[type[object], object],

352 -) -> Union[TypeIs[type[SupportsLoRA]], TypeIs[SupportsLoRA]]:

353 + model: type[object] | object,

354 +) -> TypeIs[type[SupportsLoRA]] | TypeIs[SupportsLoRA]:

355 result = _supports_lora(model)

356

357 if not result:

@@ -379,7 +381,7 @@ def supports_lora(

381 return result

382

383

382 -def _supports_lora(model: Union[type[object], object]) -> bool:

384 +def _supports_lora(model: type[object] | object) -> bool:

385 if isinstance(model, type):

386 return isinstance(model, _SupportsLoRAType)

387

@@ -404,15 +406,15 @@ class SupportsPP(Protocol):

406 batch_size: int,

407 dtype: torch.dtype,

408 device: torch.device,

407 - ) -> "IntermediateTensors":

409 + ) -> IntermediateTensors:

410 """Called when PP rank > 0 for profiling purposes."""

411 ...

412

413 def forward(

414 self,

415 *,

414 - intermediate_tensors: Optional["IntermediateTensors"],

415 - ) -> Union[Tensor, "IntermediateTensors"]:

416 + intermediate_tensors: IntermediateTensors | None,

417 + ) -> IntermediateTensors | None:

418 """

419 Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when

420 PP rank > 0.

@@ -434,13 +436,13 @@ class _SupportsPPType(Protocol):

436 batch_size: int,

437 dtype: torch.dtype,

438 device: torch.device,

437 - ) -> "IntermediateTensors": ...

439 + ) -> IntermediateTensors: ...

440

441 def forward(

442 self,

443 *,

442 - intermediate_tensors: Optional["IntermediateTensors"],

443 - ) -> Union[Tensor, "IntermediateTensors"]: ...

444 + intermediate_tensors: IntermediateTensors | None,

445 + ) -> Tensor | IntermediateTensors: ...

446

447

448 @overload

@@ -452,8 +454,8 @@ def supports_pp(model: object) -> TypeIs[SupportsPP]: ...

454

455

456 def supports_pp(

455 - model: Union[type[object], object],

456 -) -> Union[bool, TypeIs[type[SupportsPP]], TypeIs[SupportsPP]]:

457 + model: type[object] | object,

458 +) -> bool | TypeIs[type[SupportsPP]] | TypeIs[SupportsPP]:

459 supports_attributes = _supports_pp_attributes(model)

460 supports_inspect = _supports_pp_inspect(model)

461

@@ -487,14 +489,14 @@ def supports_pp(

489 return supports_attributes and supports_inspect

490

491

490 -def _supports_pp_attributes(model: Union[type[object], object]) -> bool:

492 +def _supports_pp_attributes(model: type[object] | object) -> bool:

493 if isinstance(model, type):

494 return isinstance(model, _SupportsPPType)

495

496 return isinstance(model, SupportsPP)

497

498

497 -def _supports_pp_inspect(model: Union[type[object], object]) -> bool:

499 +def _supports_pp_inspect(model: type[object] | object) -> bool:

500 model_forward = getattr(model, "forward", None)

501 if not callable(model_forward):

502 return False

@@ -523,8 +525,8 @@ def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...

525

526

527 def has_inner_state(

526 - model: Union[type[object], object],

527 -) -> Union[TypeIs[type[HasInnerState]], TypeIs[HasInnerState]]:

528 + model: type[object] | object,

529 +) -> TypeIs[type[HasInnerState]] | TypeIs[HasInnerState]:

530 return getattr(model, "has_inner_state", False)

531

532

@@ -550,8 +552,8 @@ def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...

552

553

554 def is_attention_free(

553 - model: Union[type[object], object],

554 -) -> Union[TypeIs[type[IsAttentionFree]], TypeIs[IsAttentionFree]]:

555 + model: type[object] | object,

556 +) -> TypeIs[type[IsAttentionFree]] | TypeIs[IsAttentionFree]:

557 return getattr(model, "is_attention_free", False)

558

559

@@ -570,7 +572,7 @@ class IsHybrid(Protocol):

572 @classmethod

573 def get_mamba_state_shape_from_config(

574 cls,

573 - vllm_config: "VllmConfig",

575 + vllm_config: VllmConfig,

576 use_v1: bool = True,

577 ) -> tuple[tuple[int, int], tuple[int, int, int]]:

578 """Calculate shapes for Mamba's convolutional and state caches.

@@ -596,8 +598,8 @@ def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...

598

599

600 def is_hybrid(

599 - model: Union[type[object], object],

600 -) -> Union[TypeIs[type[IsHybrid]], TypeIs[IsHybrid]]:

601 + model: type[object] | object,

602 +) -> TypeIs[type[IsHybrid]] | TypeIs[IsHybrid]:

603 return getattr(model, "is_hybrid", False)

604

605

@@ -688,8 +690,8 @@ def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...

690

691

692 def has_noops(

691 - model: Union[type[object], object],

692 -) -> Union[TypeIs[type[HasNoOps]], TypeIs[HasNoOps]]:

693 + model: type[object] | object,

694 +) -> TypeIs[type[HasNoOps]] | TypeIs[HasNoOps]:

695 return getattr(model, "has_noops", False)

696

697

@@ -711,23 +713,23 @@ def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]: ...

713

714

715 def _supports_cross_encoding(

714 - model: Union[type[object], object],

715 -) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:

716 + model: type[object] | object,

717 +) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:

718 return getattr(model, "supports_cross_encoding", False)

719

720

721 def supports_cross_encoding(

720 - model: Union[type[object], object],

721 -) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:

722 + model: type[object] | object,

723 +) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:

724 return is_pooling_model(model) and _supports_cross_encoding(model)

725

726

727 class SupportsQuant:

728 """The interface required for all models that support quantization."""

729

728 - hf_to_vllm_mapper: ClassVar[Optional["WeightsMapper"]] = None

729 - packed_modules_mapping: ClassVar[Optional[dict[str, list[str]]]] = None

730 - quant_config: Optional[QuantizationConfig] = None

730 + hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None

731 + packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None

732 + quant_config: QuantizationConfig | None = None

733

734 def __new__(cls, *args, **kwargs) -> Self:

735 instance = super().__new__(cls)

@@ -749,7 +751,7 @@ class SupportsQuant:

751 return instance

752

753 @staticmethod

752 - def _find_quant_config(*args, **kwargs) -> Optional[QuantizationConfig]:

754 + def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:

755 """Find quant config passed through model constructor args"""

756 from vllm.config import VllmConfig # avoid circular import

757

@@ -797,10 +799,10 @@ class SupportsTranscription(Protocol):

799 audio: np.ndarray,

800 stt_config: SpeechToTextConfig,

801 model_config: ModelConfig,

800 - language: Optional[str],

802 + language: str | None,

803 task_type: Literal["transcribe", "translate"],

804 request_prompt: str,

803 - to_language: Optional[str],

805 + to_language: str | None,

806 ) -> PromptType:

807 """Get the prompt for the ASR model.

808 The model has control over the construction, as long as it

@@ -813,7 +815,7 @@ class SupportsTranscription(Protocol):

815 return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}

816

817 @classmethod

816 - def validate_language(cls, language: Optional[str]) -> Optional[str]:

818 + def validate_language(cls, language: str | None) -> str | None:

819 """

820 Ensure the language specified in the transcription request

821 is a valid ISO 639-1 language code. If the request language is

@@ -850,7 +852,7 @@ class SupportsTranscription(Protocol):

852 audio_duration_s: float,

853 stt_config: SpeechToTextConfig,

854 model_config: ModelConfig,

853 - ) -> Optional[int]:

855 + ) -> int | None:

856 """

857 Map from audio duration to number of audio tokens produced by the ASR

858 model, without running a forward pass.

@@ -870,8 +872,8 @@ def supports_transcription(model: object) -> TypeIs[SupportsTranscription]: ...

872

873

874 def supports_transcription(

873 - model: Union[type[object], object],

874 -) -> Union[TypeIs[type[SupportsTranscription]], TypeIs[SupportsTranscription]]:

875 + model: type[object] | object,

876 +) -> TypeIs[type[SupportsTranscription]] | TypeIs[SupportsTranscription]:

877 return getattr(model, "supports_transcription", False)

878

879

@@ -891,8 +893,8 @@ def supports_v0_only(model: object) -> TypeIs[SupportsV0Only]: ...

893

894

895 def supports_v0_only(

894 - model: Union[type[object], object],

895 -) -> Union[TypeIs[type[SupportsV0Only]], TypeIs[SupportsV0Only]]:

896 + model: type[object] | object,

897 +) -> TypeIs[type[SupportsV0Only]] | TypeIs[SupportsV0Only]:

898 return getattr(model, "supports_v0_only", False)

899

900

@@ -942,8 +944,8 @@ def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]: ...

944

945

946 def supports_eagle3(

945 - model: Union[type[object], object],

946 -) -> Union[TypeIs[type[SupportsEagle3]], TypeIs[SupportsEagle3]]:

947 + model: type[object] | object,

948 +) -> TypeIs[type[SupportsEagle3]] | TypeIs[SupportsEagle3]:

949 return isinstance(model, SupportsEagle3)

950

951

@@ -964,12 +966,12 @@ class SupportsMRoPE(Protocol):

966 self,

967 input_tokens: list[int],

968 hf_config: PretrainedConfig,

967 - image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

968 - video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

969 - second_per_grid_ts: Optional[list[float]] = None,

969 + image_grid_thw: list[list[int]] | torch.Tensor | None,

970 + video_grid_thw: list[list[int]] | torch.Tensor | None,

971 + second_per_grid_ts: list[float] | None = None,

972 context_len: int = 0,

971 - seq_len: Optional[int] = None,

972 - audio_feature_lengths: Optional[torch.Tensor] = None,

973 + seq_len: int | None = None,

974 + audio_feature_lengths: torch.Tensor | None = None,

975 use_audio_in_video: bool = False,

976 ) -> tuple[torch.Tensor, int]:

977 """

@@ -1007,6 +1009,6 @@ def supports_mrope(model: object) -> TypeIs[SupportsMRoPE]: ...

1009

1010

1011 def supports_mrope(

1010 - model: Union[type[object], object],

1011 -) -> Union[TypeIs[type[SupportsMRoPE]], TypeIs[SupportsMRoPE]]:

1012 + model: type[object] | object,

1013 +) -> TypeIs[type[SupportsMRoPE]] | TypeIs[SupportsMRoPE]:

1014 return isinstance(model, SupportsMRoPE)

vllm/model_executor/models/interfaces_base.py

+11 -15

@@ -5,9 +5,7 @@ from typing import (

5 Any,

6 ClassVar,

7 Literal,

8 - Optional,

8 Protocol,

10 - Union,

9 overload,

10 runtime_checkable,

11 )

@@ -63,12 +61,12 @@ class VllmModel(Protocol[T_co]):

61 ) -> T_co: ...

62

63

66 -def _check_vllm_model_init(model: Union[type[object], object]) -> bool:

64 +def _check_vllm_model_init(model: type[object] | object) -> bool:

65 model_init = model.__init__

66 return supports_kw(model_init, "vllm_config")

67

68

71 -def _check_vllm_model_get_input_embeddings(model: Union[type[object], object]) -> bool:

69 +def _check_vllm_model_get_input_embeddings(model: type[object] | object) -> bool:

70 model_get_input_embeddings = getattr(model, "get_input_embeddings", None)

71 if not callable(model_get_input_embeddings):

72 logger.warning(

@@ -80,7 +78,7 @@ def _check_vllm_model_get_input_embeddings(model: Union[type[object], object]) -

78 return True

79

80

83 -def _check_vllm_model_forward(model: Union[type[object], object]) -> bool:

81 +def _check_vllm_model_forward(model: type[object] | object) -> bool:

82 model_forward = getattr(model, "forward", None)

83 if not callable(model_forward):

84 return False

@@ -108,8 +106,8 @@ def is_vllm_model(model: object) -> TypeIs[VllmModel]: ...

106

107

108 def is_vllm_model(

111 - model: Union[type[object], object],

112 -) -> Union[TypeIs[type[VllmModel]], TypeIs[VllmModel]]:

109 + model: type[object] | object,

110 +) -> TypeIs[type[VllmModel]] | TypeIs[VllmModel]:

111 return (

112 _check_vllm_model_init(model)

113 and _check_vllm_model_get_input_embeddings(model)

@@ -124,7 +122,7 @@ class VllmModelForTextGeneration(VllmModel[T], Protocol[T]):

122 def compute_logits(

123 self,

124 hidden_states: T,

127 - ) -> Optional[T]:

125 + ) -> T | None:

126 """Return `None` if TP rank > 0."""

127 ...

128

@@ -140,10 +138,8 @@ def is_text_generation_model(model: object) -> TypeIs[VllmModelForTextGeneration

138

139

140 def is_text_generation_model(

143 - model: Union[type[object], object],

144 -) -> Union[

145 - TypeIs[type[VllmModelForTextGeneration]], TypeIs[VllmModelForTextGeneration]

146 -]:

141 + model: type[object] | object,

142 +) -> TypeIs[type[VllmModelForTextGeneration]] | TypeIs[VllmModelForTextGeneration]:

143 if not is_vllm_model(model):

144 return False

145

@@ -190,8 +186,8 @@ def is_pooling_model(model: object) -> TypeIs[VllmModelForPooling]: ...

186

187

188 def is_pooling_model(

193 - model: Union[type[object], object],

194 -) -> Union[TypeIs[type[VllmModelForPooling]], TypeIs[VllmModelForPooling]]:

189 + model: type[object] | object,

190 +) -> TypeIs[type[VllmModelForPooling]] | TypeIs[VllmModelForPooling]:

191 if not is_vllm_model(model):

192 return False

193

@@ -211,5 +207,5 @@ def default_pooling_type(pooling_type: str):

207 return func

208

209

214 -def get_default_pooling_type(model: Union[type[object], object]) -> str:

210 +def get_default_pooling_type(model: type[object] | object) -> str:

211 return getattr(model, "default_pooling_type", "LAST")

vllm/model_executor/models/intern_vit.py

+12 -13

@@ -9,7 +9,6 @@

9 # --------------------------------------------------------

10 from collections.abc import Iterable

11 from functools import partial

12 -from typing import Optional

12

13 import torch

14 import torch.nn as nn

@@ -121,8 +120,8 @@ class InternVisionPatchModel(nn.Module):

120

121 def forward(

122 self,

124 - pixel_values: Optional[torch.Tensor] = None,

125 - pixel_embeds: Optional[torch.Tensor] = None,

123 + pixel_values: torch.Tensor | None = None,

124 + pixel_embeds: torch.Tensor | None = None,

125 ) -> torch.FloatTensor:

126 if pixel_values is None and pixel_embeds is None:

127 raise ValueError("You have to specify pixel_values or pixel_embeds")

@@ -144,7 +143,7 @@ class InternParallelAttention(nn.Module):

143 def __init__(

144 self,

145 config: PretrainedConfig,

147 - quant_config: Optional[QuantizationConfig] = None,

146 + quant_config: QuantizationConfig | None = None,

147 *,

148 num_dummy_heads: int = 0,

149 prefix: str = "",

@@ -240,7 +239,7 @@ class InternMLP(nn.Module):

239 def __init__(

240 self,

241 config: PretrainedConfig,

243 - quant_config: Optional[QuantizationConfig] = None,

242 + quant_config: QuantizationConfig | None = None,

243 prefix: str = "",

244 use_data_parallel: bool = False,

245 ) -> None:

@@ -277,7 +276,7 @@ class InternVisionEncoderLayer(nn.Module):

276 def __init__(

277 self,

278 config: PretrainedConfig,

280 - quant_config: Optional[QuantizationConfig] = None,

279 + quant_config: QuantizationConfig | None = None,

280 *,

281 num_dummy_heads: int = 0,

282 prefix: str = "",

@@ -312,7 +311,7 @@ class InternVisionEncoderLayer(nn.Module):

311 def _init_attn(

312 self,

313 config: PretrainedConfig,

315 - quant_config: Optional[QuantizationConfig],

314 + quant_config: QuantizationConfig | None,

315 *,

316 num_dummy_heads: int,

317 prefix: str = "",

@@ -350,9 +349,9 @@ class InternVisionEncoder(nn.Module):

349 def __init__(

350 self,

351 config: PretrainedConfig,

353 - quant_config: Optional[QuantizationConfig] = None,

352 + quant_config: QuantizationConfig | None = None,

353 *,

355 - num_hidden_layers_override: Optional[int] = None,

354 + num_hidden_layers_override: int | None = None,

355 num_dummy_heads: int = 0,

356 prefix: str = "",

357 use_data_parallel: bool = False,

@@ -395,9 +394,9 @@ class InternVisionModel(nn.Module):

394 def __init__(

395 self,

396 config: PretrainedConfig,

398 - quant_config: Optional[QuantizationConfig] = None,

397 + quant_config: QuantizationConfig | None = None,

398 *,

400 - num_hidden_layers_override: Optional[int] = None,

399 + num_hidden_layers_override: int | None = None,

400 num_dummy_heads: int = 0,

401 prefix: str = "",

402 use_data_parallel: bool = False,

@@ -422,8 +421,8 @@ class InternVisionModel(nn.Module):

421

422 def forward(

423 self,

425 - pixel_values: Optional[torch.Tensor] = None,

426 - pixel_embeds: Optional[torch.Tensor] = None,

424 + pixel_values: torch.Tensor | None = None,

425 + pixel_embeds: torch.Tensor | None = None,

426 ) -> torch.FloatTensor:

427 if pixel_values is None and pixel_embeds is None:

428 raise ValueError("You have to specify pixel_values or pixel_embeds")

vllm/model_executor/models/internlm2.py

+17 -17

@@ -4,7 +4,7 @@

4 from collections.abc import Iterable

5 from functools import partial

6 from itertools import islice

7 -from typing import Any, Optional, Union

7 +from typing import Any

8

9 import torch

10 from torch import nn

@@ -54,7 +54,7 @@ class InternLM2MLP(nn.Module):

54 hidden_size: int,

55 intermediate_size: int,

56 hidden_act: str,

57 - quant_config: Optional[QuantizationConfig] = None,

57 + quant_config: QuantizationConfig | None = None,

58 prefix: str = "",

59 ) -> None:

60 super().__init__()

@@ -92,10 +92,10 @@ class InternLM2Attention(nn.Module):

92 num_heads: int,

93 num_kv_heads: int,

94 rope_theta: float = 10000,

95 - rope_scaling: Optional[dict[str, Any]] = None,

95 + rope_scaling: dict[str, Any] | None = None,

96 max_position_embeddings: int = 8192,

97 - cache_config: Optional[CacheConfig] = None,

98 - quant_config: Optional[QuantizationConfig] = None,

97 + cache_config: CacheConfig | None = None,

98 + quant_config: QuantizationConfig | None = None,

99 prefix: str = "",

100 ) -> None:

101 super().__init__()

@@ -198,8 +198,8 @@ class InternLMDecoderLayer(nn.Module):

198 def __init__(

199 self,

200 config: PretrainedConfig,

201 - cache_config: Optional[CacheConfig] = None,

202 - quant_config: Optional[QuantizationConfig] = None,

201 + cache_config: CacheConfig | None = None,

202 + quant_config: QuantizationConfig | None = None,

203 prefix: str = "",

204 ) -> None:

205 super().__init__()

@@ -232,7 +232,7 @@ class InternLMDecoderLayer(nn.Module):

232 self,

233 positions: torch.Tensor,

234 hidden_states: torch.Tensor,

235 - residual: Optional[torch.Tensor],

235 + residual: torch.Tensor | None,

236 ) -> tuple[torch.Tensor, torch.Tensor]:

237 # Self Attention

238 if residual is None:

@@ -291,9 +291,9 @@ class InternLM2Model(nn.Module):

291 self,

292 input_ids: torch.Tensor,

293 positions: torch.Tensor,

294 - intermediate_tensors: Optional[IntermediateTensors] = None,

295 - inputs_embeds: Optional[torch.Tensor] = None,

296 - ) -> Union[torch.Tensor, IntermediateTensors]:

294 + intermediate_tensors: IntermediateTensors | None = None,

295 + inputs_embeds: torch.Tensor | None = None,

296 + ) -> torch.Tensor | IntermediateTensors:

297 if get_pp_group().is_first_rank:

298 if inputs_embeds is not None:

299 hidden_states = inputs_embeds

@@ -359,8 +359,8 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

359 self,

360 input_ids: torch.Tensor,

361 positions: torch.Tensor,

362 - intermediate_tensors: Optional[IntermediateTensors],

363 - inputs_embeds: Optional[torch.Tensor] = None,

362 + intermediate_tensors: IntermediateTensors | None,

363 + inputs_embeds: torch.Tensor | None = None,

364 ) -> torch.Tensor:

365 hidden_states = self.model(

366 input_ids, positions, intermediate_tensors, inputs_embeds

@@ -370,7 +370,7 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

370 def compute_logits(

371 self,

372 hidden_states: torch.Tensor,

373 - ) -> Optional[torch.Tensor]:

373 + ) -> torch.Tensor | None:

374 logits = self.logits_processor(self.output, hidden_states)

375 return logits

376

@@ -451,9 +451,9 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):

451 self,

452 input_ids: torch.Tensor,

453 positions: torch.Tensor,

454 - intermediate_tensors: Optional[IntermediateTensors] = None,

455 - inputs_embeds: Optional[torch.Tensor] = None,

456 - ) -> Union[torch.Tensor, IntermediateTensors]:

454 + intermediate_tensors: IntermediateTensors | None = None,

455 + inputs_embeds: torch.Tensor | None = None,

456 + ) -> torch.Tensor | IntermediateTensors:

457 hidden_states = self.model(

458 input_ids, positions, intermediate_tensors, inputs_embeds

459 )

vllm/model_executor/models/internlm2_ve.py

+8 -9

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from itertools import islice

5 -from typing import Optional, Union

5

6 import torch

7 from torch import nn

@@ -25,8 +24,8 @@ class InternLM2VEDecoderLayer(nn.Module):

24 def __init__(

25 self,

26 config: PretrainedConfig,

28 - cache_config: Optional[CacheConfig] = None,

29 - quant_config: Optional[QuantizationConfig] = None,

27 + cache_config: CacheConfig | None = None,

28 + quant_config: QuantizationConfig | None = None,

29 prefix: str = "",

30 ) -> None:

31 super().__init__()

@@ -66,8 +65,8 @@ class InternLM2VEDecoderLayer(nn.Module):

65 self,

66 positions: torch.Tensor,

67 hidden_states: torch.Tensor,

69 - residual: Optional[torch.Tensor],

70 - visual_token_mask: Optional[torch.Tensor] = None,

68 + residual: torch.Tensor | None,

69 + visual_token_mask: torch.Tensor | None = None,

70 ) -> tuple[torch.Tensor, torch.Tensor]:

71 # Self Attention

72 if residual is None:

@@ -107,10 +106,10 @@ class InternLM2VEModel(InternLM2Model):

106 self,

107 input_ids: torch.Tensor,

108 positions: torch.Tensor,

110 - intermediate_tensors: Optional[IntermediateTensors] = None,

111 - inputs_embeds: Optional[torch.Tensor] = None,

112 - visual_token_mask: Optional[torch.Tensor] = None,

113 - ) -> Union[torch.Tensor, IntermediateTensors]:

109 + intermediate_tensors: IntermediateTensors | None = None,

110 + inputs_embeds: torch.Tensor | None = None,

111 + visual_token_mask: torch.Tensor | None = None,

112 + ) -> torch.Tensor | IntermediateTensors:

113 if get_pp_group().is_first_rank:

114 if inputs_embeds is not None:

115 hidden_states = inputs_embeds

vllm/model_executor/models/interns1.py

+19 -23

@@ -7,7 +7,7 @@

7 # Licensed under The MIT License [see LICENSE for details]

8 # --------------------------------------------------------

9 from collections.abc import Iterable, Mapping, Sequence

10 -from typing import Annotated, Literal, Optional, Union

10 +from typing import Annotated, Literal, TypeAlias

11

12 import regex as re

13 import torch

@@ -111,12 +111,10 @@ class InternS1ImageEmbeddingInputs(TensorSchema):

111 """

112

113 type: Literal["image_embeds"] = "image_embeds"

114 - data: Annotated[

115 - Union[torch.Tensor, list[torch.Tensor]], TensorShape("ni", "tifs", "hs")

116 - ]

114 + data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("ni", "tifs", "hs")]

115

116

119 -InternS1ImageInputs = Union[InternS1ImagePixelInputs, InternS1ImageEmbeddingInputs]

117 +InternS1ImageInputs: TypeAlias = InternS1ImagePixelInputs | InternS1ImageEmbeddingInputs

118

119

120 class InternS1VideoPixelInputs(TensorSchema):

@@ -143,12 +141,10 @@ class InternS1VideoEmbeddingInputs(TensorSchema):

141 """

142

143 type: Literal["video_embeds"] = "video_embeds"

146 - data: Annotated[

147 - Union[torch.Tensor, list[torch.Tensor]], TensorShape("nv", "tvfs", "hs")

148 - ]

144 + data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("nv", "tvfs", "hs")]

145

146

151 -InternS1VideoInputs = Union[InternS1VideoPixelInputs, InternS1VideoEmbeddingInputs]

147 +InternS1VideoInputs: TypeAlias = InternS1VideoPixelInputs | InternS1VideoEmbeddingInputs

148

149

150 def resolve_interns1_min_max_num(

@@ -190,7 +186,7 @@ class InternS1ProcessingInfo(BaseProcessingInfo):

186 )

187 return hf_processor

188

193 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

189 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

190 return {"image": None, "video": None}

191

192 def get_num_image_tokens(

@@ -198,7 +194,7 @@ class InternS1ProcessingInfo(BaseProcessingInfo):

194 *,

195 image_width: int,

196 image_height: int,

201 - processor: Optional["GotOcr2ImageProcessorFast"] = None,

197 + processor: GotOcr2ImageProcessorFast | None = None,

198 ) -> int:

199 if processor is None:

200 processor = self.get_hf_processor().image_processor

@@ -213,7 +209,7 @@ class InternS1ProcessingInfo(BaseProcessingInfo):

209 num_image_tokens = self.get_hf_processor().image_seq_length * num_image_patches

210 return num_image_tokens

211

216 - def resolve_target_ratios(self, use_thumbnail: Optional[bool] = None):

212 + def resolve_target_ratios(self, use_thumbnail: bool | None = None):

213 image_processor = self.get_hf_processor().image_processor

214 min_dynamic_patch = image_processor.min_patches

215 max_dynamic_patch = image_processor.max_patches

@@ -298,7 +294,7 @@ class InternS1DummyInputsBuilder(BaseDummyInputsBuilder[InternS1ProcessingInfo])

294 self,

295 seq_len: int,

296 mm_counts: Mapping[str, int],

301 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

297 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

298 ) -> MultiModalDataDict:

299 target_width, target_height = self.info.get_image_size_with_most_features()

300 target_num_frames = self.info.get_num_frames_with_most_features(

@@ -523,7 +519,7 @@ class InternS1ForConditionalGeneration(

519 )

520

521 @classmethod

526 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

522 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

523 # transformers InternVLProcessor uses <IMG_CONTEXT> as the separator

524 # refer to https://github.com/huggingface/transformers/blob/f90de364c2484c7c325bbe05befdcf487bd75b63/src/transformers/models/internvl/processing_internvl.py#L116

525 if modality.startswith("image"):

@@ -576,7 +572,7 @@ class InternS1ForConditionalGeneration(

572 def _init_vision_model(

573 self,

574 config: PretrainedConfig,

579 - quant_config: Optional[QuantizationConfig],

575 + quant_config: QuantizationConfig | None,

576 *,

577 prefix: str,

578 ):

@@ -620,7 +616,7 @@ class InternS1ForConditionalGeneration(

616

617 def _parse_and_validate_image_input(

618 self, **kwargs: object

623 - ) -> Optional[InternS1ImageInputs]:

619 + ) -> InternS1ImageInputs | None:

620 pixel_values = kwargs.pop("pixel_values", None)

621 image_num_patches = kwargs.pop("image_num_patches", None)

622 image_embeds = kwargs.pop("image_embeds", None)

@@ -654,7 +650,7 @@ class InternS1ForConditionalGeneration(

650

651 def _parse_and_validate_video_input(

652 self, **kwargs: object

657 - ) -> Optional[InternS1VideoInputs]:

653 + ) -> InternS1VideoInputs | None:

654 pixel_values_flat_video = kwargs.pop("pixel_values_videos", None)

655 video_num_patches = kwargs.pop("video_num_patches", None)

656 video_embeds = kwargs.pop("video_embeds", None)

@@ -688,7 +684,7 @@ class InternS1ForConditionalGeneration(

684

685 def _process_vision_input(

686 self,

691 - image_input: Union[InternS1ImageInputs, InternS1VideoInputs],

687 + image_input: InternS1ImageInputs | InternS1VideoInputs,

688 ) -> tuple[torch.Tensor, ...]:

689 if (

690 image_input["type"] == "image_embeds"

@@ -763,9 +759,9 @@ class InternS1ForConditionalGeneration(

759 def get_input_embeddings(

760 self,

761 input_ids: torch.Tensor,

766 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

762 + multimodal_embeddings: MultiModalEmbeddings | None = None,

763 *,

768 - is_multimodal: Optional[torch.Tensor] = None,

764 + is_multimodal: torch.Tensor | None = None,

765 handle_oov_mm_token: bool = False,

766 ) -> torch.Tensor:

767 if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:

@@ -786,8 +782,8 @@ class InternS1ForConditionalGeneration(

782 self,

783 input_ids: torch.Tensor,

784 positions: torch.Tensor,

789 - intermediate_tensors: Optional[IntermediateTensors] = None,

790 - inputs_embeds: Optional[torch.Tensor] = None,

785 + intermediate_tensors: IntermediateTensors | None = None,

786 + inputs_embeds: torch.Tensor | None = None,

787 **kwargs: object,

788 ) -> IntermediateTensors:

789 if intermediate_tensors is not None:

@@ -807,7 +803,7 @@ class InternS1ForConditionalGeneration(

803 def compute_logits(

804 self,

805 hidden_states: torch.Tensor,

810 - ) -> Optional[torch.Tensor]:

806 + ) -> torch.Tensor | None:

807 return self.language_model.compute_logits(hidden_states)

808

809 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/interns1_vit.py

+10 -11

@@ -8,7 +8,6 @@

8 # Licensed under The MIT License [see LICENSE for details]

9 # --------------------------------------------------------

10 from collections.abc import Iterable

11 -from typing import Optional

11

12 import torch

13 import torch.nn as nn

@@ -139,7 +138,7 @@ class InternS1VisionEmbeddings(nn.Module):

138 def forward(

139 self,

140 pixel_values: torch.Tensor,

142 - bool_masked_pos: Optional[torch.BoolTensor] = None,

141 + bool_masked_pos: torch.BoolTensor | None = None,

142 ) -> torch.Tensor:

143 _, _, height, width = pixel_values.shape

144 embeddings, (patch_height, patch_width) = self.patch_embeddings(pixel_values)

@@ -240,7 +239,7 @@ class InternS1VisionMLP(nn.Module):

239 def __init__(

240 self,

241 config: PretrainedConfig,

243 - quant_config: Optional[QuantizationConfig] = None,

242 + quant_config: QuantizationConfig | None = None,

243 prefix: str = "",

244 ) -> None:

245 super().__init__()

@@ -274,7 +273,7 @@ class InternS1VisionLayer(nn.Module):

273 def __init__(

274 self,

275 config: PretrainedConfig,

277 - quant_config: Optional[QuantizationConfig] = None,

276 + quant_config: QuantizationConfig | None = None,

277 *,

278 num_dummy_heads: int = 0,

279 prefix: str = "",

@@ -309,7 +308,7 @@ class InternS1VisionLayer(nn.Module):

308 def _init_attn(

309 self,

310 config: PretrainedConfig,

312 - quant_config: Optional[QuantizationConfig],

311 + quant_config: QuantizationConfig | None,

312 *,

313 num_dummy_heads: int,

314 prefix: str = "",

@@ -337,9 +336,9 @@ class InternS1VisionEncoder(nn.Module):

336 def __init__(

337 self,

338 config: PretrainedConfig,

340 - quant_config: Optional[QuantizationConfig] = None,

339 + quant_config: QuantizationConfig | None = None,

340 *,

342 - num_hidden_layers_override: Optional[int] = None,

341 + num_hidden_layers_override: int | None = None,

342 num_dummy_heads: int = 0,

343 prefix: str = "",

344 ):

@@ -376,9 +375,9 @@ class InternS1VisionModel(nn.Module):

375 def __init__(

376 self,

377 config: PretrainedConfig,

379 - quant_config: Optional[QuantizationConfig] = None,

378 + quant_config: QuantizationConfig | None = None,

379 *,

381 - num_hidden_layers_override: Optional[int] = None,

380 + num_hidden_layers_override: int | None = None,

381 num_dummy_heads: int = 0,

382 prefix: str = "",

383 ) -> None:

@@ -404,8 +403,8 @@ class InternS1VisionModel(nn.Module):

403

404 def forward(

405 self,

407 - pixel_values: Optional[torch.Tensor] = None,

408 - pixel_embeds: Optional[torch.Tensor] = None,

406 + pixel_values: torch.Tensor | None = None,

407 + pixel_embeds: torch.Tensor | None = None,

408 ) -> torch.FloatTensor:

409 if pixel_values is None and pixel_embeds is None:

410 raise ValueError("You have to specify pixel_values or pixel_embeds")

vllm/model_executor/models/internvl.py

+61 -61

@@ -10,7 +10,7 @@

10 import os

11 from abc import ABC, abstractmethod

12 from collections.abc import Iterable, Mapping, Sequence

13 -from typing import Annotated, Any, Literal, Optional, TypeVar, Union

13 +from typing import Annotated, Any, Literal, TypeAlias, TypeVar

14

15 import numpy.typing as npt

16 import torch

@@ -94,10 +94,10 @@ class InternVLImageEmbeddingInputs(TensorSchema):

94 """

95

96 type: Literal["image_embeds"]

97 - data: Annotated[Union[torch.Tensor, list[torch.Tensor]], TensorShape("n", "f", "h")]

97 + data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("n", "f", "h")]

98

99

100 -InternVLImageInputs = Union[InternVLImagePixelInputs, InternVLImageEmbeddingInputs]

100 +InternVLImageInputs: TypeAlias = InternVLImagePixelInputs | InternVLImageEmbeddingInputs

101

102

103 class InternVLVideoPixelInputs(TensorSchema):

@@ -124,10 +124,10 @@ class InternVLVideoEmbeddingInputs(TensorSchema):

124 """

125

126 type: Literal["video_embeds"]

127 - data: Annotated[Union[torch.Tensor, list[torch.Tensor]], TensorShape("n", "f", "h")]

127 + data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("n", "f", "h")]

128

129

130 -InternVLVideoInputs = Union[InternVLVideoPixelInputs, InternVLVideoEmbeddingInputs]

130 +InternVLVideoInputs: TypeAlias = InternVLVideoPixelInputs | InternVLVideoEmbeddingInputs

131

132

133 # adapted from https://huggingface.co/OpenGVLab/InternVL2-1B

@@ -349,9 +349,9 @@ class BaseInternVLProcessor(ABC):

349 config: PretrainedConfig,

350 tokenizer: AnyTokenizer,

351 *,

352 - min_dynamic_patch: Optional[int] = None,

353 - max_dynamic_patch: Optional[int] = None,

354 - dynamic_image_size: Optional[bool] = None,

352 + min_dynamic_patch: int | None = None,

353 + max_dynamic_patch: int | None = None,

354 + dynamic_image_size: bool | None = None,

355 ) -> None:

356 super().__init__()

357

@@ -391,17 +391,17 @@ class BaseInternVLProcessor(ABC):

391 def get_image_repl(

392 self,

393 feature_size: int,

394 - num_patches: Optional[int],

394 + num_patches: int | None,

395 ) -> PromptUpdateDetails[str]:

396 raise NotImplementedError

397

398 def resolve_min_max_num(

399 self,

400 *,

401 - min_dynamic_patch: Optional[int] = None,

402 - max_dynamic_patch: Optional[int] = None,

403 - dynamic_image_size: Optional[bool] = None,

404 - use_thumbnail: Optional[bool] = None,

401 + min_dynamic_patch: int | None = None,

402 + max_dynamic_patch: int | None = None,

403 + dynamic_image_size: bool | None = None,

404 + use_thumbnail: bool | None = None,

405 ) -> tuple[int, int]:

406 min_dynamic_patch = (

407 self.min_dynamic_patch if min_dynamic_patch is None else min_dynamic_patch

@@ -426,10 +426,10 @@ class BaseInternVLProcessor(ABC):

426 def resolve_target_ratios(

427 self,

428 *,

429 - min_dynamic_patch: Optional[int] = None,

430 - max_dynamic_patch: Optional[int] = None,

431 - dynamic_image_size: Optional[bool] = None,

432 - use_thumbnail: Optional[bool] = None,

429 + min_dynamic_patch: int | None = None,

430 + max_dynamic_patch: int | None = None,

431 + dynamic_image_size: bool | None = None,

432 + use_thumbnail: bool | None = None,

433 ) -> list[tuple[int, int]]:

434 min_num, max_num = self.resolve_min_max_num(

435 min_dynamic_patch=min_dynamic_patch,

@@ -463,9 +463,9 @@ class BaseInternVLProcessor(ABC):

463 def _images_to_pixel_values_lst(

464 self,

465 images: list[Image.Image],

466 - min_dynamic_patch: Optional[int] = None,

467 - max_dynamic_patch: Optional[int] = None,

468 - dynamic_image_size: Optional[bool] = None,

466 + min_dynamic_patch: int | None = None,

467 + max_dynamic_patch: int | None = None,

468 + dynamic_image_size: bool | None = None,

469 ) -> list[torch.Tensor]:

470 min_num, max_num = self.resolve_min_max_num(

471 min_dynamic_patch=min_dynamic_patch,

@@ -489,9 +489,9 @@ class BaseInternVLProcessor(ABC):

489 self,

490 text: list[str],

491 images: list[Image.Image],

492 - min_dynamic_patch: Optional[int] = None,

493 - max_dynamic_patch: Optional[int] = None,

494 - dynamic_image_size: Optional[bool] = None,

492 + min_dynamic_patch: int | None = None,

493 + max_dynamic_patch: int | None = None,

494 + dynamic_image_size: bool | None = None,

495 ) -> tuple[list[str], dict[str, torch.Tensor]]:

496 if len(images) == 0:

497 image_inputs = {}

@@ -517,7 +517,7 @@ class BaseInternVLProcessor(ABC):

517 text = [t.replace("<image>", image_repl.full, 1) for t in text]

518 return text, image_inputs

519

520 - def _make_batch_input(self, input_item: Optional[Union[Any, list[Any]]] = None):

520 + def _make_batch_input(self, input_item: Any | list[Any] | None = None):

521 if input_item is None:

522 input_item = []

523 if not isinstance(input_item, list):

@@ -526,12 +526,12 @@ class BaseInternVLProcessor(ABC):

526

527 def __call__(

528 self,

529 - text: Optional[Union[str, list[str]]] = None,

530 - images: Optional[Union[Image.Image, list[Image.Image]]] = None,

531 - min_dynamic_patch: Optional[int] = None,

532 - max_dynamic_patch: Optional[int] = None,

533 - dynamic_image_size: Optional[bool] = None,

534 - return_tensors: Optional[Union[str, TensorType]] = None,

529 + text: str | list[str] | None = None,

530 + images: Image.Image | list[Image.Image] | None = None,

531 + min_dynamic_patch: int | None = None,

532 + max_dynamic_patch: int | None = None,

533 + dynamic_image_size: bool | None = None,

534 + return_tensors: str | TensorType | None = None,

535 ) -> BatchFeature:

536 text, images = [self._make_batch_input(x) for x in (text, images)]

537

@@ -563,10 +563,10 @@ class InternVLProcessor(BaseInternVLProcessor):

563 config: PretrainedConfig,

564 tokenizer: AnyTokenizer,

565 *,

566 - min_dynamic_patch: Optional[int] = None,

567 - max_dynamic_patch: Optional[int] = None,

568 - dynamic_image_size: Optional[bool] = None,

569 - video_token: Optional[str] = None,

566 + min_dynamic_patch: int | None = None,

567 + max_dynamic_patch: int | None = None,

568 + dynamic_image_size: bool | None = None,

569 + video_token: str | None = None,

570 ) -> None:

571 super().__init__(

572 config=config,

@@ -583,7 +583,7 @@ class InternVLProcessor(BaseInternVLProcessor):

583 return self.tokenizer.get_vocab()[IMG_CONTEXT]

584

585 @property

586 - def video_token_id(self) -> Optional[int]:

586 + def video_token_id(self) -> int | None:

587 if self.video_token is None:

588 return None

589 return self.tokenizer.get_vocab().get(self.video_token, None)

@@ -595,7 +595,7 @@ class InternVLProcessor(BaseInternVLProcessor):

595 def _videos_to_pixel_values_lst(

596 self,

597 videos: list[npt.NDArray],

598 - dynamic_image_size: Optional[bool] = None,

598 + dynamic_image_size: bool | None = None,

599 ) -> list[torch.Tensor]:

600 min_num, max_num = self.resolve_min_max_num(

601 min_dynamic_patch=1,

@@ -619,7 +619,7 @@ class InternVLProcessor(BaseInternVLProcessor):

619 self,

620 text: list[str],

621 videos: list[npt.NDArray],

622 - dynamic_image_size: Optional[bool] = None,

622 + dynamic_image_size: bool | None = None,

623 ):

624 if len(videos) == 0 or not self.supports_video:

625 video_inputs = {}

@@ -646,13 +646,13 @@ class InternVLProcessor(BaseInternVLProcessor):

646

647 def __call__(

648 self,

649 - text: Optional[Union[str, list[str]]] = None,

650 - images: Optional[Union[Image.Image, list[Image.Image]]] = None,

651 - videos: Optional[Union[npt.NDArray, list[npt.NDArray]]] = None,

652 - min_dynamic_patch: Optional[int] = None,

653 - max_dynamic_patch: Optional[int] = None,

654 - dynamic_image_size: Optional[bool] = None,

655 - return_tensors: Optional[Union[str, TensorType]] = None,

649 + text: str | list[str] | None = None,

650 + images: Image.Image | list[Image.Image] | None = None,

651 + videos: npt.NDArray | list[npt.NDArray] | None = None,

652 + min_dynamic_patch: int | None = None,

653 + max_dynamic_patch: int | None = None,

654 + dynamic_image_size: bool | None = None,

655 + return_tensors: str | TensorType | None = None,

656 ) -> BatchFeature:

657 text, images, videos = [

658 self._make_batch_input(x) for x in (text, images, videos)

@@ -681,7 +681,7 @@ class InternVLProcessor(BaseInternVLProcessor):

681 def get_image_repl(

682 self,

683 feature_size: int,

684 - num_patches: Optional[int],

684 + num_patches: int | None,

685 ) -> PromptUpdateDetails[str]:

686 repl_features = IMG_CONTEXT * feature_size

687 repl_full = IMG_START + repl_features + IMG_END

@@ -691,7 +691,7 @@ class InternVLProcessor(BaseInternVLProcessor):

691 def get_video_repl(

692 self,

693 feature_size: int,

694 - num_patches: Optional[int] = None,

694 + num_patches: int | None = None,

695 video_context_token: str = IMG_CONTEXT,

696 ) -> PromptUpdateDetails[str]:

697 repl_features = video_context_token * self.num_image_token

@@ -711,7 +711,7 @@ class BaseInternVLProcessingInfo(BaseProcessingInfo):

711 def get_hf_processor(self, **kwargs: object) -> BaseInternVLProcessor:

712 raise NotImplementedError

713

714 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

714 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

715 return {"image": None}

716

717 def get_num_image_tokens(

@@ -719,7 +719,7 @@ class BaseInternVLProcessingInfo(BaseProcessingInfo):

719 *,

720 image_width: int,

721 image_height: int,

722 - processor: Optional[BaseInternVLProcessor],

722 + processor: BaseInternVLProcessor | None,

723 ) -> int:

724 if processor is None:

725 processor = self.get_hf_processor()

@@ -779,7 +779,7 @@ class BaseInternVLDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

779 self,

780 seq_len: int,

781 mm_counts: Mapping[str, int],

782 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

782 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

783 ) -> MultiModalDataDict:

784 target_width, target_height = self.info.get_image_size_with_most_features()

785 num_images = mm_counts.get("image", 0)

@@ -901,7 +901,7 @@ class InternVLProcessingInfo(BaseInternVLProcessingInfo):

901 video_limit = {"video": None} if self.supports_video else {}

902 return {**super().get_supported_mm_limits(), **video_limit}

903

904 - def get_video_token(self) -> Optional[str]:

904 + def get_video_token(self) -> str | None:

905 text_model_type = self.get_hf_config().get_text_config().model_type

906 video_token_map = {

907 "qwen2": "<|video_pad|>",

@@ -951,7 +951,7 @@ class InternVLDummyInputsBuilder(

951 self,

952 seq_len: int,

953 mm_counts: Mapping[str, int],

954 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

954 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

955 ) -> MultiModalDataDict:

956 dummy_image = super().get_dummy_mm_data(

957 seq_len=seq_len, mm_counts=mm_counts, mm_options=mm_options

@@ -1079,7 +1079,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1079 supports_encoder_tp_data = True

1080

1081 @classmethod

1082 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1082 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1083 if modality.startswith("image"):

1084 return "<image>"

1085 if modality.startswith("video"):

@@ -1149,7 +1149,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1149 def _init_vision_model(

1150 self,

1151 config: PretrainedConfig,

1152 - quant_config: Optional[QuantizationConfig],

1152 + quant_config: QuantizationConfig | None,

1153 *,

1154 is_mono: bool,

1155 prefix: str,

@@ -1217,7 +1217,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1217

1218 def _parse_and_validate_image_input(

1219 self, **kwargs: object

1220 - ) -> Optional[InternVLImageInputs]:

1220 + ) -> InternVLImageInputs | None:

1221 pixel_values_flat = kwargs.pop("pixel_values_flat", None)

1222 image_num_patches = kwargs.pop("image_num_patches", None)

1223 image_embeds = kwargs.pop("image_embeds", None)

@@ -1250,7 +1250,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1250

1251 def _parse_and_validate_video_input(

1252 self, **kwargs: object

1253 - ) -> Optional[InternVLVideoPixelInputs]:

1253 + ) -> InternVLVideoPixelInputs | None:

1254 pixel_values_flat_video = kwargs.pop("pixel_values_flat_video", None)

1255 video_num_patches = kwargs.pop("video_num_patches", None)

1256 video_embeds = kwargs.pop("image_embeds", None)

@@ -1283,7 +1283,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1283

1284 def _process_vision_input(

1285 self,

1286 - image_input: Union[InternVLImageInputs, InternVLVideoInputs],

1286 + image_input: InternVLImageInputs | InternVLVideoInputs,

1287 ) -> tuple[torch.Tensor, ...]:

1288 if (

1289 image_input["type"] == "image_embeds"

@@ -1364,9 +1364,9 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1364 def get_input_embeddings(

1365 self,

1366 input_ids: torch.Tensor,

1367 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1367 + multimodal_embeddings: MultiModalEmbeddings | None = None,

1368 *,

1369 - is_multimodal: Optional[torch.Tensor] = None,

1369 + is_multimodal: torch.Tensor | None = None,

1370 handle_oov_mm_token: bool = False,

1371 ) -> torch.Tensor:

1372 if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:

@@ -1387,8 +1387,8 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1387 self,

1388 input_ids: torch.Tensor,

1389 positions: torch.Tensor,

1390 - intermediate_tensors: Optional[IntermediateTensors] = None,

1391 - inputs_embeds: Optional[torch.Tensor] = None,

1390 + intermediate_tensors: IntermediateTensors | None = None,

1391 + inputs_embeds: torch.Tensor | None = None,

1392 **kwargs: object,

1393 ) -> IntermediateTensors:

1394 if intermediate_tensors is not None:

@@ -1413,7 +1413,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1413 def compute_logits(

1414 self,

1415 hidden_states: torch.Tensor,

1416 - ) -> Optional[torch.Tensor]:

1416 + ) -> torch.Tensor | None:

1417 return self.language_model.compute_logits(hidden_states)

1418

1419 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/jais.py

+12 -13

@@ -24,7 +24,6 @@

24 import math

25 from collections.abc import Iterable

26 from itertools import islice

27 -from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -86,8 +85,8 @@ class JAISAttention(nn.Module):

85 def __init__(

86 self,

87 config: JAISConfig,

89 - cache_config: Optional[CacheConfig] = None,

90 - quant_config: Optional[QuantizationConfig] = None,

88 + cache_config: CacheConfig | None = None,

89 + quant_config: QuantizationConfig | None = None,

90 prefix: str = "",

91 ):

92 super().__init__()

@@ -147,7 +146,7 @@ class JAISMLP(nn.Module):

146 self,

147 intermediate_size: int,

148 config: JAISConfig,

150 - quant_config: Optional[QuantizationConfig] = None,

149 + quant_config: QuantizationConfig | None = None,

150 ):

151 super().__init__()

152 hidden_size = config.hidden_size

@@ -194,8 +193,8 @@ class JAISBlock(nn.Module):

193 def __init__(

194 self,

195 config: JAISConfig,

197 - cache_config: Optional[CacheConfig] = None,

198 - quant_config: Optional[QuantizationConfig] = None,

196 + cache_config: CacheConfig | None = None,

197 + quant_config: QuantizationConfig | None = None,

198 prefix: str = "",

199 ):

200 super().__init__()

@@ -277,9 +276,9 @@ class JAISModel(nn.Module):

276 self,

277 input_ids: torch.Tensor,

278 position_ids: torch.Tensor,

280 - intermediate_tensors: Optional[IntermediateTensors] = None,

281 - inputs_embeds: Optional[torch.Tensor] = None,

282 - ) -> Union[IntermediateTensors, torch.Tensor]:

279 + intermediate_tensors: IntermediateTensors | None = None,

280 + inputs_embeds: torch.Tensor | None = None,

281 + ) -> IntermediateTensors | torch.Tensor:

282 if get_pp_group().is_first_rank:

283 if inputs_embeds is None:

284 inputs_embeds = self.get_input_embeddings(input_ids)

@@ -341,9 +340,9 @@ class JAISLMHeadModel(nn.Module, SupportsPP):

340 self,

341 input_ids: torch.Tensor,

342 positions: torch.Tensor,

344 - intermediate_tensors: Optional[IntermediateTensors] = None,

345 - inputs_embeds: Optional[torch.Tensor] = None,

346 - ) -> Union[IntermediateTensors, torch.Tensor]:

343 + intermediate_tensors: IntermediateTensors | None = None,

344 + inputs_embeds: torch.Tensor | None = None,

345 + ) -> IntermediateTensors | torch.Tensor:

346 hidden_states = self.transformer(

347 input_ids, positions, intermediate_tensors, inputs_embeds

348 )

@@ -352,7 +351,7 @@ class JAISLMHeadModel(nn.Module, SupportsPP):

351 def compute_logits(

352 self,

353 hidden_states: torch.Tensor,

355 - ) -> Optional[torch.Tensor]:

354 + ) -> torch.Tensor | None:

355 logits = self.logits_processor(self.lm_head, hidden_states)

356 return logits

357

vllm/model_executor/models/jamba.py

+19 -20

@@ -4,7 +4,6 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import Optional

7

8 import torch

9 from torch import nn

@@ -54,11 +53,11 @@ class JambaMoE(nn.Module):

53 def __init__(

54 self,

55 config: JambaConfig,

57 - num_experts: Optional[int] = None,

58 - top_k: Optional[int] = None,

59 - params_dtype: Optional[torch.dtype] = None,

60 - tp_size: Optional[int] = None,

61 - quant_config: Optional[QuantizationConfig] = None,

56 + num_experts: int | None = None,

57 + top_k: int | None = None,

58 + params_dtype: torch.dtype | None = None,

59 + tp_size: int | None = None,

60 + quant_config: QuantizationConfig | None = None,

61 prefix: str = "",

62 ):

63 super().__init__()

@@ -111,10 +110,10 @@ class JambaMambaDecoderLayer(nn.Module):

110 self,

111 config: JambaConfig,

112 layer_idx: int,

114 - model_config: Optional[ModelConfig] = None,

115 - cache_config: Optional[CacheConfig] = None,

116 - quant_config: Optional[QuantizationConfig] = None,

117 - is_lora_enabled: Optional[bool] = False,

113 + model_config: ModelConfig | None = None,

114 + cache_config: CacheConfig | None = None,

115 + quant_config: QuantizationConfig | None = None,

116 + is_lora_enabled: bool | None = False,

117 prefix: str = "",

118 **kwargs,

119 ) -> None:

@@ -159,7 +158,7 @@ class JambaMambaDecoderLayer(nn.Module):

158 def forward(

159 self,

160 hidden_states: torch.Tensor,

162 - residual: Optional[torch.Tensor],

161 + residual: torch.Tensor | None,

162 **kwargs,

163 ):

164 if residual is None:

@@ -181,9 +180,9 @@ class JambaAttentionDecoderLayer(nn.Module):

180 self,

181 config: JambaConfig,

182 layer_idx: int,

184 - model_config: Optional[ModelConfig] = None,

185 - cache_config: Optional[CacheConfig] = None,

186 - quant_config: Optional[QuantizationConfig] = None,

183 + model_config: ModelConfig | None = None,

184 + cache_config: CacheConfig | None = None,

185 + quant_config: QuantizationConfig | None = None,

186 prefix: str = "",

187 **kwargs,

188 ) -> None:

@@ -266,7 +265,7 @@ class JambaAttentionDecoderLayer(nn.Module):

265 self,

266 positions: torch.Tensor,

267 hidden_states: torch.Tensor,

269 - residual: Optional[torch.Tensor],

268 + residual: torch.Tensor | None,

269 **kwargs,

270 ):

271 if residual is None:

@@ -348,8 +347,8 @@ class JambaModel(nn.Module):

347 self,

348 input_ids: torch.Tensor,

349 positions: torch.Tensor,

351 - intermediate_tensors: Optional[IntermediateTensors] = None,

352 - inputs_embeds: Optional[torch.Tensor] = None,

350 + intermediate_tensors: IntermediateTensors | None = None,

351 + inputs_embeds: torch.Tensor | None = None,

352 ) -> torch.Tensor:

353 if get_pp_group().is_first_rank:

354 if inputs_embeds is not None:

@@ -523,8 +522,8 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHyb

522 self,

523 input_ids: torch.Tensor,

524 positions: torch.Tensor,

526 - intermediate_tensors: Optional[IntermediateTensors] = None,

527 - inputs_embeds: Optional[torch.Tensor] = None,

525 + intermediate_tensors: IntermediateTensors | None = None,

526 + inputs_embeds: torch.Tensor | None = None,

527 **kwargs,

528 ):

529 hidden_states = self.model(

@@ -568,7 +567,7 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHyb

567 def compute_logits(

568 self,

569 hidden_states: torch.Tensor,

571 - ) -> Optional[torch.Tensor]:

570 + ) -> torch.Tensor | None:

571 logits = self.logits_processor(self.lm_head, hidden_states)

572 return logits

573

vllm/model_executor/models/jina_vl.py

+4 -5

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Mapping

4 -from typing import Optional

4

5 import torch

6 import torch.nn as nn

@@ -105,14 +104,14 @@ class JinaVLForSequenceClassification(

104 )

105

106 @classmethod

108 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

107 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

108 if modality.startswith("image"):

109 return "<|vision_start|><|image_pad|><|vision_end|>"

110

111 raise ValueError("Only image modality is supported")

112

113 @classmethod

115 - def get_score_template(cls, query: str, document: str) -> Optional[str]:

114 + def get_score_template(cls, query: str, document: str) -> str | None:

115 return f"**Document**:\n{document}\n**Query**:\n{query}"

116

117 @classmethod

@@ -124,8 +123,8 @@ class JinaVLForSequenceClassification(

123 self,

124 input_ids: torch.Tensor,

125 positions: torch.Tensor,

127 - intermediate_tensors: Optional[IntermediateTensors] = None,

128 - inputs_embeds: Optional[torch.Tensor] = None,

126 + intermediate_tensors: IntermediateTensors | None = None,

127 + inputs_embeds: torch.Tensor | None = None,

128 **kwargs: object,

129 ) -> torch.Tensor:

130 hidden_states = super().forward(

vllm/model_executor/models/keye.py

+76 -100

@@ -4,7 +4,7 @@ import math

4 from abc import abstractmethod

5 from collections.abc import Iterable, Mapping, Sequence

6 from functools import partial

7 -from typing import Annotated, Any, Literal, Optional, TypeVar, Union

7 +from typing import Annotated, Any, Literal, TypeAlias, TypeVar

8

9 import numpy as np

10 import torch

@@ -153,7 +153,7 @@ class KeyeImageEmbeddingInputs(TensorSchema):

153 image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

154

155

156 -KeyeImageInputs = Union[KeyeImagePixelInputs, KeyeImageEmbeddingInputs]

156 +KeyeImageInputs: TypeAlias = KeyeImagePixelInputs | KeyeImageEmbeddingInputs

157

158

159 class KeyeVideoPixelInputs(TensorSchema):

@@ -188,7 +188,7 @@ class KeyeVideoEmbeddingInputs(TensorSchema):

188 video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]

189

190

191 -KeyeVideoInputs = Union[KeyeVideoPixelInputs, KeyeVideoEmbeddingInputs]

191 +KeyeVideoInputs: TypeAlias = KeyeVideoPixelInputs | KeyeVideoEmbeddingInputs

192

193

194 class KeyeVisionEmbeddings(nn.Module):

@@ -278,15 +278,9 @@ class KeyeVisionEmbeddings(nn.Module):

278 def forward(

279 self,

280 pixel_values: torch.FloatTensor,

281 - position_ids: Optional[torch.Tensor] = None,

282 - image_grid_thw: Optional[

283 - list[

284 - Union[

285 - tuple[int, int, int],

286 - list[tuple[int, int, int]],

287 - ]

288 - ]

289 - ] = None,

281 + position_ids: torch.Tensor | None = None,

282 + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]

283 + | None = None,

284 interpolate_pos_encoding=False,

285 ) -> torch.Tensor:

286 if pixel_values.dim() == 4:

@@ -357,7 +351,7 @@ class KeyeSiglipAttention(nn.Module):

351 def __init__(

352 self,

353 config: PretrainedConfig,

360 - quant_config: Optional[QuantizationConfig] = None,

354 + quant_config: QuantizationConfig | None = None,

355 prefix: str = "",

356 ):

357 super().__init__()

@@ -416,10 +410,10 @@ class KeyeSiglipAttention(nn.Module):

410 def forward(

411 self,

412 hidden_states: torch.Tensor,

419 - attention_mask: Optional[torch.Tensor] = None,

420 - output_attentions: Optional[bool] = False,

421 - cu_seqlens: Optional[list[torch.Tensor]] = None,

422 - rope_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,

413 + attention_mask: torch.Tensor | None = None,

414 + output_attentions: bool | None = False,

415 + cu_seqlens: list[torch.Tensor] | None = None,

416 + rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,

417 ) -> torch.Tensor:

418 qkv, _ = self.qkv_proj(hidden_states)

419 q, k, v = qkv.split(

@@ -524,8 +518,8 @@ class SigLIPRotaryEmbedding(nn.Module):

518 class KeyeSiglipEncoderLayer(nn.Module):

519 def __init__(

520 self,

527 - config: Union[PretrainedConfig],

528 - quant_config: Optional[QuantizationConfig] = None,

521 + config: PretrainedConfig,

522 + quant_config: QuantizationConfig | None = None,

523 prefix: str = "",

524 ):

525 super().__init__()

@@ -547,9 +541,9 @@ class KeyeSiglipEncoderLayer(nn.Module):

541 self,

542 hidden_states: torch.Tensor,

543 attention_mask: torch.Tensor,

550 - output_attentions: Optional[bool] = False,

551 - cu_seqlens: Optional[list[torch.Tensor]] = None,

552 - rope_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,

544 + output_attentions: bool | None = False,

545 + cu_seqlens: list[torch.Tensor] | None = None,

546 + rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,

547 ) -> tuple[torch.FloatTensor]:

548 residual = hidden_states

549

@@ -577,7 +571,7 @@ class KeyeSiglipEncoder(nn.Module):

571 def __init__(

572 self,

573 config: PretrainedConfig,

580 - quant_config: Optional[QuantizationConfig] = None,

574 + quant_config: QuantizationConfig | None = None,

575 prefix: str = "",

576 ):

577 super().__init__()

@@ -610,22 +604,16 @@ class KeyeSiglipEncoder(nn.Module):

604 def forward(

605 self,

606 inputs_embeds,

613 - attention_mask: Optional[torch.Tensor] = None,

614 - output_attentions: Optional[bool] = None,

615 - output_hidden_states: Optional[bool] = None,

616 - cu_seqlens: Optional[list[torch.Tensor]] = None,

617 - image_grid_thw: Optional[

618 - list[

619 - Union[

620 - tuple[int, int, int],

621 - list[tuple[int, int, int]],

622 - ]

623 - ]

624 - ] = None,

625 - height_position_ids: Optional[torch.Tensor] = None,

626 - width_position_ids: Optional[torch.Tensor] = None,

627 - use_rope: Optional[bool] = False,

628 - window_size: Optional[bool] = -1,

607 + attention_mask: torch.Tensor | None = None,

608 + output_attentions: bool | None = None,

609 + output_hidden_states: bool | None = None,

610 + cu_seqlens: list[torch.Tensor] | None = None,

611 + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]

612 + | None = None,

613 + height_position_ids: torch.Tensor | None = None,

614 + width_position_ids: torch.Tensor | None = None,

615 + use_rope: bool | None = False,

616 + window_size: bool | None = -1,

617 vision_or_text: str = "vision",

618 ) -> BaseModelOutput:

619 device = inputs_embeds.device

@@ -676,7 +664,7 @@ class KeyeSiglipVisionTransformer(nn.Module):

664 def __init__(

665 self,

666 config: PretrainedConfig,

679 - quant_config: Optional[QuantizationConfig] = None,

667 + quant_config: QuantizationConfig | None = None,

668 prefix: str = "",

669 ):

670 super().__init__()

@@ -694,29 +682,23 @@ class KeyeSiglipVisionTransformer(nn.Module):

682 def forward(

683 self,

684 pixel_values,

697 - output_attentions: Optional[bool] = None,

698 - output_hidden_states: Optional[bool] = None,

699 - interpolate_pos_encoding: Optional[bool] = False,

700 - attention_mask: Optional[torch.Tensor] = None,

701 - sample_indices: Optional[torch.Tensor] = None,

702 - image_indices: Optional[torch.Tensor] = None,

703 - position_ids: Optional[torch.Tensor] = None,

704 - height_position_ids: Optional[torch.Tensor] = None,

705 - width_position_ids: Optional[torch.Tensor] = None,

706 - cu_seqlens: Optional[list[torch.Tensor]] = None,

707 - padding_mask: Optional[torch.Tensor] = None,

708 - vision_return_embed_list: Optional[bool] = False,

709 - image_grid_thw: Optional[

710 - list[

711 - Union[

712 - tuple[int, int, int],

713 - list[tuple[int, int, int]],

714 - ]

715 - ]

716 - ] = None,

717 - return_pooler_output: Optional[bool] = True,

718 - use_rope: Optional[bool] = False,

719 - window_size: Optional[bool] = -1,

685 + output_attentions: bool | None = None,

686 + output_hidden_states: bool | None = None,

687 + interpolate_pos_encoding: bool | None = False,

688 + attention_mask: torch.Tensor | None = None,

689 + sample_indices: torch.Tensor | None = None,

690 + image_indices: torch.Tensor | None = None,

691 + position_ids: torch.Tensor | None = None,

692 + height_position_ids: torch.Tensor | None = None,

693 + width_position_ids: torch.Tensor | None = None,

694 + cu_seqlens: list[torch.Tensor] | None = None,

695 + padding_mask: torch.Tensor | None = None,

696 + vision_return_embed_list: bool | None = False,

697 + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]

698 + | None = None,

699 + return_pooler_output: bool | None = True,

700 + use_rope: bool | None = False,

701 + window_size: bool | None = -1,

702 ) -> BaseModelOutputWithPooling:

703 hidden_states = self.embeddings(

704 pixel_values,

@@ -763,7 +745,7 @@ class KeyeSiglipVisionModel(nn.Module):

745 def __init__(

746 self,

747 config: PretrainedConfig,

766 - quant_config: Optional[QuantizationConfig] = None,

748 + quant_config: QuantizationConfig | None = None,

749 prefix: str = "",

750 ):

751 super().__init__()

@@ -789,24 +771,18 @@ class KeyeSiglipVisionModel(nn.Module):

771 def forward(

772 self,

773 pixel_values,

792 - sample_indices: Optional[torch.Tensor] = None,

793 - output_attentions: Optional[bool] = None,

794 - output_hidden_states: Optional[bool] = None,

774 + sample_indices: torch.Tensor | None = None,

775 + output_attentions: bool | None = None,

776 + output_hidden_states: bool | None = None,

777 interpolate_pos_encoding: bool = False,

796 - position_ids: Optional[torch.Tensor] = None,

797 - vision_return_embed_list: Optional[bool] = False,

798 - image_grid_thw: Optional[

799 - list[

800 - Union[

801 - tuple[int, int, int],

802 - list[tuple[int, int, int]],

803 - ]

804 - ]

805 - ] = None,

806 - cu_seqlens: Optional[list[torch.Tensor]] = None,

807 - return_pooler_output: Optional[bool] = True,

808 - use_rope: Optional[bool] = False,

809 - window_size: Optional[bool] = -1,

778 + position_ids: torch.Tensor | None = None,

779 + vision_return_embed_list: bool | None = False,

780 + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]

781 + | None = None,

782 + cu_seqlens: list[torch.Tensor] | None = None,

783 + return_pooler_output: bool | None = True,

784 + use_rope: bool | None = False,

785 + window_size: bool | None = -1,

786 ) -> BaseModelOutputWithPooling:

787 return self.vision_model(

788 pixel_values=pixel_values,

@@ -893,7 +869,7 @@ class Projector(nn.Module):

869 self,

870 text_config: PretrainedConfig,

871 vision_config: PretrainedConfig,

896 - quant_config: Optional[QuantizationConfig] = None,

872 + quant_config: QuantizationConfig | None = None,

873 prefix: str = "",

874 ):

875 super().__init__()

@@ -927,9 +903,9 @@ class Projector(nn.Module):

903

904 def forward(

905 self,

930 - image_features: Union[torch.Tensor, list[torch.Tensor]],

906 + image_features: torch.Tensor | list[torch.Tensor],

907 image_grid_thw: list[tuple[int, int, int]],

932 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

908 + ) -> torch.Tensor | list[torch.Tensor]:

909 m1, m2 = self.merge_kernel_size

910 if isinstance(image_features, (list, tuple)):

911 processed_features = list()

@@ -988,7 +964,7 @@ def _keye_field_config(

964 class KeyeMultiModalDataParser(MultiModalDataParser):

965 def _parse_image_data(

966 self,

991 - data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],

967 + data: dict[str, torch.Tensor] | ModalityData[ImageItem],

968 ) -> ModalityDataItems[Any, Any]:

969 if isinstance(data, dict):

970 return DictEmbeddingItems(

@@ -1005,7 +981,7 @@ class KeyeMultiModalDataParser(MultiModalDataParser):

981

982 def _parse_video_data(

983 self,

1008 - data: Union[dict[str, torch.Tensor], ModalityData[VideoItem]],

984 + data: dict[str, torch.Tensor] | ModalityData[VideoItem],

985 ) -> ModalityDataItems[Any, Any]:

986 if isinstance(data, dict):

987 return DictEmbeddingItems(

@@ -1033,7 +1009,7 @@ class KeyeProcessingInfo(BaseProcessingInfo):

1009

1010 def get_supported_mm_limits(

1011 self,

1036 - ) -> Mapping[str, Optional[int]]:

1012 + ) -> Mapping[str, int | None]:

1013 return {"image": None, "video": None}

1014

1015 def get_mm_max_tokens_per_item(

@@ -1200,7 +1176,7 @@ class KeyeBaseDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

1176 self,

1177 seq_len: int,

1178 mm_counts: Mapping[str, int],

1203 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

1179 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

1180 ) -> MultiModalDataDict:

1181 num_images = mm_counts.get("image", 0)

1182 num_videos = mm_counts.get("video", 0)

@@ -1303,7 +1279,7 @@ class BaseKeyeModule(nn.Module):

1279 )

1280

1281 @classmethod

1306 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1282 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1283 if modality.startswith("image"):

1284 return "<|vision_start|><|image_pad|><|vision_end|>"

1285 if modality.startswith("video"):

@@ -1348,7 +1324,7 @@ class BaseKeyeModule(nn.Module):

1324 self,

1325 text_config: PretrainedConfig,

1326 vision_config: PretrainedConfig,

1351 - quant_config: Optional[QuantizationConfig] = None,

1327 + quant_config: QuantizationConfig | None = None,

1328 prefix: str = "",

1329 ) -> nn.Module:

1330 raise ValueError("Need projector")

@@ -1403,8 +1379,8 @@ class BaseKeyeModule(nn.Module):

1379 self,

1380 video_type: Literal["video_embeds", "pixel_values_videos"],

1381 video_grid_thw: list[torch.Tensor],

1406 - pixel_values_videos: Optional[torch.Tensor] = None,

1407 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

1382 + pixel_values_videos: torch.Tensor | None = None,

1383 + ) -> torch.Tensor | list[torch.Tensor]:

1384 siglip_position_ids = list()

1385 video_grid_hws = list()

1386 sample_indices = list()

@@ -1473,7 +1449,7 @@ class BaseKeyeModule(nn.Module):

1449

1450 def get_multimodal_embeddings(

1451 self, **kwargs: object

1476 - ) -> Optional[MultiModalEmbeddings]:

1452 + ) -> MultiModalEmbeddings | None:

1453 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

1454 if not modalities:

1455 return None

@@ -1495,10 +1471,10 @@ class BaseKeyeModule(nn.Module):

1471 self,

1472 input_ids: torch.Tensor,

1473 positions: torch.Tensor,

1498 - intermediate_tensors: Optional[IntermediateTensors] = None,

1499 - inputs_embeds: Optional[torch.Tensor] = None,

1474 + intermediate_tensors: IntermediateTensors | None = None,

1475 + inputs_embeds: torch.Tensor | None = None,

1476 **kwargs: object,

1501 - ) -> Union[torch.Tensor, IntermediateTensors]:

1477 + ) -> torch.Tensor | IntermediateTensors:

1478 """Run forward pass for Keye-VL.

1479

1480 Args:

@@ -1527,7 +1503,7 @@ class BaseKeyeModule(nn.Module):

1503 def compute_logits(

1504 self,

1505 hidden_states: torch.Tensor,

1530 - ) -> Optional[torch.Tensor]:

1506 + ) -> torch.Tensor | None:

1507 return self.language_model.compute_logits(hidden_states)

1508

1509 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -1555,14 +1531,14 @@ class KeyeForConditionalGeneration(

1531 self,

1532 text_config: PretrainedConfig,

1533 vision_config: PretrainedConfig,

1558 - quant_config: Optional[QuantizationConfig] = None,

1534 + quant_config: QuantizationConfig | None = None,

1535 prefix: str = "",

1536 ) -> nn.Module:

1537 return Projector(text_config, vision_config, quant_config, prefix)

1538

1539 def _parse_and_validate_image_input(

1540 self, **kwargs: object

1565 - ) -> Optional[KeyeImageInputs]:

1541 + ) -> KeyeImageInputs | None:

1542 pixel_values = kwargs.pop("pixel_values", None)

1543 image_embeds = kwargs.pop("image_embeds", None)

1544 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -1586,7 +1562,7 @@ class KeyeForConditionalGeneration(

1562

1563 def _parse_and_validate_video_input(

1564 self, **kwargs: object

1589 - ) -> Optional[KeyeVideoInputs]:

1565 + ) -> KeyeVideoInputs | None:

1566 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1567 video_embeds = kwargs.pop("video_embeds", None)

1568 video_grid_thw = kwargs.pop("video_grid_thw", None)

vllm/model_executor/models/keye_vl1_5.py

+23 -19

@@ -3,7 +3,7 @@

3 import itertools

4 from collections.abc import Mapping, Sequence

5 from functools import partial

6 -from typing import Annotated, Any, Literal, Optional, Union

6 +from typing import Annotated, Any, Literal, TypeAlias

7

8 import numpy as np

9 import torch

@@ -73,7 +73,7 @@ def split_thw(grid_thw: torch.Tensor) -> torch.Tensor:

73

74

75 def get_num_patches(

76 - grid_thw: torch.Tensor, num_frames: Union[list[int], torch.Tensor]

76 + grid_thw: torch.Tensor, num_frames: list[int] | torch.Tensor

77 ) -> list[int]:

78 """

79 Return num_patches per video.

@@ -153,7 +153,9 @@ class KeyeVL1_5ImageEmbeddingInputs(TensorSchema):

153 image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

154

155

156 -KeyeVL1_5ImageInputs = Union[KeyeVL1_5ImagePixelInputs, KeyeVL1_5ImageEmbeddingInputs]

156 +KeyeVL1_5ImageInputs: TypeAlias = (

157 + KeyeVL1_5ImagePixelInputs | KeyeVL1_5ImageEmbeddingInputs

158 +)

159

160

161 class KeyeVL1_5VideoPixelInputs(TensorSchema):

@@ -191,7 +193,9 @@ class KeyeVL1_5VideoEmbeddingInputs(TensorSchema):

193 num_frames: torch.Tensor

194

195

194 -KeyeVL1_5VideoInputs = Union[KeyeVL1_5VideoPixelInputs, KeyeVL1_5VideoEmbeddingInputs]

196 +KeyeVL1_5VideoInputs: TypeAlias = (

197 + KeyeVL1_5VideoPixelInputs | KeyeVL1_5VideoEmbeddingInputs

198 +)

199

200

201 class KeyeVL1_5Projector(nn.Module):

@@ -199,7 +203,7 @@ class KeyeVL1_5Projector(nn.Module):

203 self,

204 text_config: PretrainedConfig,

205 vision_config: PretrainedConfig,

202 - quant_config: Optional[QuantizationConfig] = None,

206 + quant_config: QuantizationConfig | None = None,

207 prefix: str = "",

208 ):

209 super().__init__()

@@ -233,9 +237,9 @@ class KeyeVL1_5Projector(nn.Module):

237

238 def forward(

239 self,

236 - image_features: Union[torch.Tensor, tuple[torch.Tensor], list[torch.Tensor]],

240 + image_features: torch.Tensor | tuple[torch.Tensor] | list[torch.Tensor],

241 image_grid_thw: list[tuple[int, int, int]],

238 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

242 + ) -> torch.Tensor | list[torch.Tensor]:

243 m1, m2 = self.merge_kernel_size

244 if isinstance(image_features, (list, tuple)):

245 processed_features = list()

@@ -275,7 +279,7 @@ class KeyeVL1_5ProcessingInfo(KeyeProcessingInfo):

279

280 def get_supported_mm_limits(

281 self,

278 - ) -> Mapping[str, Optional[int]]:

282 + ) -> Mapping[str, int | None]:

283 return {"image": None, "video": 1}

284

285

@@ -327,7 +331,7 @@ def _keye_field_config(

331 class KeyeVL1_5MultiModalDataParser(MultiModalDataParser):

332 def _parse_image_data(

333 self,

330 - data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],

334 + data: dict[str, torch.Tensor] | ModalityData[ImageItem],

335 ) -> ModalityDataItems[Any, Any]:

336 if isinstance(data, dict):

337 return DictEmbeddingItems(

@@ -344,7 +348,7 @@ class KeyeVL1_5MultiModalDataParser(MultiModalDataParser):

348

349 def _parse_video_data(

350 self,

347 - data: Union[dict[str, torch.Tensor], ModalityData[VideoItem]],

351 + data: dict[str, torch.Tensor] | ModalityData[VideoItem],

352 ) -> ModalityDataItems[Any, Any]:

353 if isinstance(data, dict):

354 return DictEmbeddingItems(

@@ -499,7 +503,7 @@ class KeyeVL1_5ForConditionalGeneration(

503 self,

504 text_config: PretrainedConfig,

505 vision_config: PretrainedConfig,

502 - quant_config: Optional[QuantizationConfig] = None,

506 + quant_config: QuantizationConfig | None = None,

507 prefix: str = "",

508 ) -> nn.Module:

509 return KeyeVL1_5Projector(text_config, vision_config, quant_config, prefix)

@@ -511,7 +515,7 @@ class KeyeVL1_5ForConditionalGeneration(

515

516 def _parse_and_validate_image_input(

517 self, **kwargs: object

514 - ) -> Optional[KeyeVL1_5ImageInputs]:

518 + ) -> KeyeVL1_5ImageInputs | None:

519 pixel_values = kwargs.pop("pixel_values", None)

520 image_embeds = kwargs.pop("image_embeds", None)

521 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -535,7 +539,7 @@ class KeyeVL1_5ForConditionalGeneration(

539

540 def _parse_and_validate_video_input(

541 self, **kwargs: object

538 - ) -> Optional[KeyeVL1_5VideoInputs]:

542 + ) -> KeyeVL1_5VideoInputs | None:

543 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

544 video_embeds = kwargs.pop("video_embeds", None)

545 video_grid_thw = kwargs.pop("video_grid_thw", None)

@@ -595,19 +599,19 @@ class KeyeVL1_5ForConditionalGeneration(

599 cls,

600 input_tokens: list[int],

601 hf_config: PretrainedConfig,

598 - image_grid_thw: Union[list[list[int]], torch.Tensor],

599 - video_grid_thw: Union[list[list[int]], torch.Tensor],

602 + image_grid_thw: list[list[int]] | torch.Tensor,

603 + video_grid_thw: list[list[int]] | torch.Tensor,

604 context_len: int = 0,

601 - seq_len: Optional[int] = None,

602 - second_per_grid_ts: Optional[list[float]] = None,

603 - audio_feature_lengths: Optional[torch.Tensor] = None,

605 + seq_len: int | None = None,

606 + second_per_grid_ts: list[float] | None = None,

607 + audio_feature_lengths: torch.Tensor | None = None,

608 use_audio_in_video: bool = False,

609 ) -> tuple[torch.Tensor, int]:

610 if isinstance(video_grid_thw, list) and len(video_grid_thw) > 0:

611 video_grid_thw = video_grid_thw[0]

612 """Get mrope input positions and delta value (Keye series)."""

613

610 - def split_thw(grid_thw: Union[torch.Tensor, list[int]]) -> list[list[int]]:

614 + def split_thw(grid_thw: torch.Tensor | list[int]) -> list[list[int]]:

615 """

616 Split grid_thw along the t dimension.

617

vllm/model_executor/models/kimi_vl.py

+10 -10

@@ -46,7 +46,7 @@ import copy

46 import math

47 from collections.abc import Iterable, Mapping, Sequence

48 from dataclasses import dataclass

49 -from typing import Annotated, Any, Literal, Optional, Union

49 +from typing import Annotated, Any, Literal

50

51 import torch

52 from torch import nn

@@ -153,7 +153,7 @@ class KimiVLImagePixelInputs(TensorSchema):

153 type: Literal["pixel_values"] = "pixel_values"

154

155 pixel_values: Annotated[

156 - Union[torch.Tensor, list[torch.Tensor]],

156 + torch.Tensor | list[torch.Tensor],

157 TensorShape("np", 3, "ps", "ps"),

158 ]

159

@@ -169,7 +169,7 @@ class KimiVLProcessingInfo(BaseProcessingInfo):

169 def get_hf_config(self):

170 return self.ctx.get_hf_config(KimiVLConfig)

171

172 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

172 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

173 return {"image": None}

174

175 def get_num_image_tokens(

@@ -227,7 +227,7 @@ class KimiVLDummyInputsBuilder(BaseDummyInputsBuilder[KimiVLProcessingInfo]):

227 self,

228 seq_len: int,

229 mm_counts: Mapping[str, int],

230 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

230 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

231 ) -> MultiModalDataDict:

232 num_images = mm_counts.get("image", 0)

233

@@ -305,7 +305,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

305 supports_encoder_tp_data = True

306

307 @classmethod

308 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

308 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

309 if modality.startswith("image"):

310 return "<|media_start|>image<|media_content|><|media_pad|><|media_end|>"

311

@@ -370,7 +370,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

370

371 def _parse_and_validate_image_input(

372 self, **kwargs: object

373 - ) -> Optional[KimiVLImageInputs]:

373 + ) -> KimiVLImageInputs | None:

374 # image input type must be pixel values now

375 pixel_values = kwargs.pop("pixel_values", None)

376 image_grid_hws = kwargs.pop("image_grid_hws", None)

@@ -411,7 +411,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

411 def get_language_model(self) -> torch.nn.Module:

412 return self.language_model

413

414 - def get_multimodal_embeddings(self, **kwargs: object) -> Optional[NestedTensors]:

414 + def get_multimodal_embeddings(self, **kwargs: object) -> NestedTensors | None:

415 # Validate the multimodal input keyword arguments

416 image_input = self._parse_and_validate_image_input(**kwargs)

417 if image_input is None:

@@ -425,8 +425,8 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

425 self,

426 input_ids: torch.Tensor,

427 positions: torch.Tensor,

428 - intermediate_tensors: Optional[IntermediateTensors] = None,

429 - inputs_embeds: Optional[torch.Tensor] = None,

428 + intermediate_tensors: IntermediateTensors | None = None,

429 + inputs_embeds: torch.Tensor | None = None,

430 **kwargs: object,

431 ) -> IntermediateTensors:

432 if intermediate_tensors is not None:

@@ -570,7 +570,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

570

571 def get_spec_layer_idx_from_weight_name(

572 config: DeepseekV2Config, weight_name: str

573 -) -> Optional[int]:

573 +) -> int | None:

574 if hasattr(config, "num_nextn_predict_layers") and (

575 config.num_nextn_predict_layers > 0

576 ):

vllm/model_executor/models/lfm2.py

+18 -18

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable

4 from itertools import islice

5 -from typing import Any, Optional

5 +from typing import Any

6

7 import torch

8 import torch.nn as nn

@@ -54,8 +54,8 @@ class Lfm2MLP(nn.Module):

54 ff_dim: int,

55 multiple_of: int,

56 auto_adjust_ff_dim: bool,

57 - ffn_dim_multiplier: Optional[float],

58 - quant_config: Optional[QuantizationConfig] = None,

57 + ffn_dim_multiplier: float | None,

58 + quant_config: QuantizationConfig | None = None,

59 prefix: str = "",

60 ):

61 super().__init__()

@@ -98,10 +98,10 @@ class Lfm2Attention(nn.Module):

98 num_heads: int,

99 num_kv_heads: int,

100 rope_theta: float = 10000,

101 - rope_scaling: Optional[dict[str, Any]] = None,

101 + rope_scaling: dict[str, Any] | None = None,

102 max_position_embeddings: int = 8192,

103 - cache_config: Optional[CacheConfig] = None,

104 - quant_config: Optional[QuantizationConfig] = None,

103 + cache_config: CacheConfig | None = None,

104 + quant_config: QuantizationConfig | None = None,

105 prefix: str = "",

106 ) -> None:

107 super().__init__()

@@ -190,9 +190,9 @@ class Lfm2AttentionDecoderLayer(nn.Module):

190 self,

191 config: Lfm2Config,

192 layer_idx: int,

193 - model_config: Optional[ModelConfig] = None,

194 - cache_config: Optional[CacheConfig] = None,

195 - quant_config: Optional[QuantizationConfig] = None,

193 + model_config: ModelConfig | None = None,

194 + cache_config: CacheConfig | None = None,

195 + quant_config: QuantizationConfig | None = None,

196 prefix: str = "",

197 ) -> None:

198 super().__init__()

@@ -240,7 +240,7 @@ class Lfm2AttentionDecoderLayer(nn.Module):

240 self,

241 positions: torch.Tensor,

242 hidden_states: torch.Tensor,

243 - residual: Optional[torch.Tensor],

243 + residual: torch.Tensor | None,

244 **kwargs,

245 ) -> tuple[torch.Tensor, torch.Tensor]:

246 if residual is None:

@@ -258,9 +258,9 @@ class Lfm2ShortConvDecoderLayer(nn.Module):

258 self,

259 config: Lfm2Config,

260 layer_idx: int,

261 - model_config: Optional[ModelConfig] = None,

262 - cache_config: Optional[CacheConfig] = None,

263 - quant_config: Optional[QuantizationConfig] = None,

261 + model_config: ModelConfig | None = None,

262 + cache_config: CacheConfig | None = None,

263 + quant_config: QuantizationConfig | None = None,

264 prefix: str = "",

265 ) -> None:

266 super().__init__()

@@ -289,7 +289,7 @@ class Lfm2ShortConvDecoderLayer(nn.Module):

289 def forward(

290 self,

291 hidden_states: torch.Tensor,

292 - residual: Optional[torch.Tensor],

292 + residual: torch.Tensor | None,

293 **kwargs,

294 ):

295 if residual is None:

@@ -365,8 +365,8 @@ class Lfm2Model(nn.Module):

365 self,

366 input_ids: torch.Tensor,

367 positions: torch.Tensor,

368 - intermediate_tensors: Optional[IntermediateTensors] = None,

369 - inputs_embeds: Optional[torch.Tensor] = None,

368 + intermediate_tensors: IntermediateTensors | None = None,

369 + inputs_embeds: torch.Tensor | None = None,

370 ) -> torch.Tensor:

371 if get_pp_group().is_first_rank:

372 if inputs_embeds is not None:

@@ -532,8 +532,8 @@ class Lfm2ForCausalLM(

532 self,

533 input_ids: torch.Tensor,

534 positions: torch.Tensor,

535 - intermediate_tensors: Optional[IntermediateTensors] = None,

536 - inputs_embeds: Optional[torch.Tensor] = None,

535 + intermediate_tensors: IntermediateTensors | None = None,

536 + inputs_embeds: torch.Tensor | None = None,

537 **kwargs,

538 ) -> torch.Tensor:

539 hidden_states = self.model(

vllm/model_executor/models/lfm2_moe.py

+18 -18

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable

4 from itertools import islice

5 -from typing import Any, Optional

5 +from typing import Any

6

7 import torch

8 import torch.nn as nn

@@ -65,7 +65,7 @@ class Lfm2MoeMlp(nn.Module):

65 self,

66 dim: int,

67 ff_dim: int,

68 - quant_config: Optional[QuantizationConfig] = None,

68 + quant_config: QuantizationConfig | None = None,

69 prefix: str = "",

70 ):

71 super().__init__()

@@ -96,7 +96,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module):

96 def __init__(

97 self,

98 config: Lfm2MoeConfig,

99 - quant_config: Optional[QuantizationConfig] = None,

99 + quant_config: QuantizationConfig | None = None,

100 prefix: str = "",

101 enable_eplb: bool = False,

102 ):

@@ -191,10 +191,10 @@ class Lfm2MoeAttention(nn.Module):

191 num_heads: int,

192 num_kv_heads: int,

193 rope_theta: float = 10000,

194 - rope_scaling: Optional[dict[str, Any]] = None,

194 + rope_scaling: dict[str, Any] | None = None,

195 max_position_embeddings: int = 8192,

196 - cache_config: Optional[CacheConfig] = None,

197 - quant_config: Optional[QuantizationConfig] = None,

196 + cache_config: CacheConfig | None = None,

197 + quant_config: QuantizationConfig | None = None,

198 prefix: str = "",

199 ) -> None:

200 super().__init__()

@@ -283,9 +283,9 @@ class Lfm2MoeAttentionDecoderLayer(nn.Module):

283 self,

284 config: Lfm2MoeConfig,

285 layer_idx: int,

286 - model_config: Optional[ModelConfig] = None,

287 - cache_config: Optional[CacheConfig] = None,

288 - quant_config: Optional[QuantizationConfig] = None,

286 + model_config: ModelConfig | None = None,

287 + cache_config: CacheConfig | None = None,

288 + quant_config: QuantizationConfig | None = None,

289 prefix: str = "",

290 enable_eplb: bool = False,

291 ) -> None:

@@ -340,7 +340,7 @@ class Lfm2MoeAttentionDecoderLayer(nn.Module):

340 self,

341 positions: torch.Tensor,

342 hidden_states: torch.Tensor,

343 - residual: Optional[torch.Tensor],

343 + residual: torch.Tensor | None,

344 **kwargs,

345 ) -> tuple[torch.Tensor, torch.Tensor]:

346 if residual is None:

@@ -358,9 +358,9 @@ class Lfm2MoeShortConvDecoderLayer(nn.Module):

358 self,

359 config: Lfm2MoeConfig,

360 layer_idx: int,

361 - model_config: Optional[ModelConfig] = None,

362 - cache_config: Optional[CacheConfig] = None,

363 - quant_config: Optional[QuantizationConfig] = None,

361 + model_config: ModelConfig | None = None,

362 + cache_config: CacheConfig | None = None,

363 + quant_config: QuantizationConfig | None = None,

364 prefix: str = "",

365 enable_eplb: bool = False,

366 ) -> None:

@@ -396,7 +396,7 @@ class Lfm2MoeShortConvDecoderLayer(nn.Module):

396 def forward(

397 self,

398 hidden_states: torch.Tensor,

399 - residual: Optional[torch.Tensor],

399 + residual: torch.Tensor | None,

400 **kwargs,

401 ):

402 if residual is None:

@@ -479,8 +479,8 @@ class Lfm2MoeModel(nn.Module):

479 self,

480 input_ids: torch.Tensor,

481 positions: torch.Tensor,

482 - intermediate_tensors: Optional[IntermediateTensors] = None,

483 - inputs_embeds: Optional[torch.Tensor] = None,

482 + intermediate_tensors: IntermediateTensors | None = None,

483 + inputs_embeds: torch.Tensor | None = None,

484 ) -> torch.Tensor:

485 if get_pp_group().is_first_rank:

486 if inputs_embeds is not None:

@@ -774,8 +774,8 @@ class Lfm2MoeForCausalLM(

774 self,

775 input_ids: torch.Tensor,

776 positions: torch.Tensor,

777 - intermediate_tensors: Optional[IntermediateTensors] = None,

778 - inputs_embeds: Optional[torch.Tensor] = None,

777 + intermediate_tensors: IntermediateTensors | None = None,

778 + inputs_embeds: torch.Tensor | None = None,

779 **kwargs,

780 ) -> torch.Tensor:

781 hidden_states = self.model(

vllm/model_executor/models/llama.py

+18 -20

@@ -26,7 +26,7 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Any, Optional, Union

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -76,7 +76,7 @@ class LlamaMLP(nn.Module):

76 hidden_size: int,

77 intermediate_size: int,

78 hidden_act: str,

79 - quant_config: Optional[QuantizationConfig] = None,

79 + quant_config: QuantizationConfig | None = None,

80 bias: bool = False,

81 prefix: str = "",

82 reduce_results: bool = True,

@@ -121,12 +121,12 @@ class LlamaAttention(nn.Module):

121 num_heads: int,

122 num_kv_heads: int,

123 rope_theta: float = 10000,

124 - rope_scaling: Optional[dict[str, Any]] = None,

124 + rope_scaling: dict[str, Any] | None = None,

125 max_position_embeddings: int = 8192,

126 - quant_config: Optional[QuantizationConfig] = None,

126 + quant_config: QuantizationConfig | None = None,

127 bias: bool = False,

128 bias_o_proj: bool = False,

129 - cache_config: Optional[CacheConfig] = None,

129 + cache_config: CacheConfig | None = None,

130 prefix: str = "",

131 attn_type: str = AttentionType.DECODER,

132 ) -> None:

@@ -236,8 +236,8 @@ class LlamaAttention(nn.Module):

236 def _init_rotary_emb(

237 self,

238 config: LlamaConfig,

239 - rope_scaling: Optional[dict[str, Any]],

240 - quant_config: Optional[QuantizationConfig],

239 + rope_scaling: dict[str, Any] | None,

240 + quant_config: QuantizationConfig | None,

241 ) -> None:

242 is_neox_style = True

243 is_gguf = quant_config and quant_config.get_name() == "gguf"

@@ -260,7 +260,7 @@ class LlamaDecoderLayer(nn.Module):

260 self,

261 vllm_config: VllmConfig,

262 prefix: str = "",

263 - config: Optional[LlamaConfig] = None,

263 + config: LlamaConfig | None = None,

264 ) -> None:

265 super().__init__()

266

@@ -331,7 +331,7 @@ class LlamaDecoderLayer(nn.Module):

331 self,

332 positions: torch.Tensor,

333 hidden_states: torch.Tensor,

334 - residual: Optional[torch.Tensor],

334 + residual: torch.Tensor | None,

335 ) -> tuple[torch.Tensor, torch.Tensor]:

336 # Self Attention

337 if residual is None:

@@ -346,7 +346,7 @@ class LlamaDecoderLayer(nn.Module):

346 hidden_states = self.mlp(hidden_states)

347 return hidden_states, residual

348

349 - def get_quant_config(self, vllm_config: VllmConfig) -> Optional[QuantizationConfig]:

349 + def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:

350 """Get quantization config for this layer. Override in subclasses."""

351 return vllm_config.quant_config

352

@@ -407,13 +407,11 @@ class LlamaModel(nn.Module):

407

408 def forward(

409 self,

410 - input_ids: Optional[torch.Tensor],

410 + input_ids: torch.Tensor | None,

411 positions: torch.Tensor,

412 - intermediate_tensors: Optional[IntermediateTensors],

413 - inputs_embeds: Optional[torch.Tensor] = None,

414 - ) -> Union[

415 - torch.Tensor, IntermediateTensors, tuple[torch.Tensor, list[torch.Tensor]]

416 - ]:

412 + intermediate_tensors: IntermediateTensors | None,

413 + inputs_embeds: torch.Tensor | None = None,

414 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:

415 if get_pp_group().is_first_rank:

416 if inputs_embeds is not None:

417 hidden_states = inputs_embeds

@@ -627,9 +625,9 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

625 self,

626 input_ids: torch.Tensor,

627 positions: torch.Tensor,

630 - intermediate_tensors: Optional[IntermediateTensors] = None,

631 - inputs_embeds: Optional[torch.Tensor] = None,

632 - ) -> Union[torch.Tensor, IntermediateTensors]:

628 + intermediate_tensors: IntermediateTensors | None = None,

629 + inputs_embeds: torch.Tensor | None = None,

630 + ) -> torch.Tensor | IntermediateTensors:

631 model_output = self.model(

632 input_ids, positions, intermediate_tensors, inputs_embeds

633 )

@@ -638,7 +636,7 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

636 def compute_logits(

637 self,

638 hidden_states: torch.Tensor,

641 - ) -> Optional[torch.Tensor]:

639 + ) -> torch.Tensor | None:

640 logits = self.logits_processor(self.lm_head, hidden_states)

641 return logits

642

vllm/model_executor/models/llama4.py

+6 -6

@@ -19,7 +19,7 @@

19 """Inference-only LLaMA model compatible with HuggingFace weights."""

20

21 from collections.abc import Iterable

22 -from typing import Any, Optional

22 +from typing import Any

23

24 import torch

25 from torch import nn

@@ -148,12 +148,12 @@ class Llama4Attention(nn.Module):

148 num_heads: int,

149 num_kv_heads: int,

150 rope_theta: float = 10000,

151 - rope_scaling: Optional[dict[str, Any]] = None,

151 + rope_scaling: dict[str, Any] | None = None,

152 max_position_embeddings: int = 8192,

153 - quant_config: Optional[QuantizationConfig] = None,

153 + quant_config: QuantizationConfig | None = None,

154 bias: bool = False,

155 bias_o_proj: bool = False,

156 - cache_config: Optional[CacheConfig] = None,

156 + cache_config: CacheConfig | None = None,

157 prefix: str = "",

158 ) -> None:

159 super().__init__()

@@ -296,7 +296,7 @@ class Llama4DecoderLayer(nn.Module):

296 self,

297 vllm_config: VllmConfig,

298 prefix: str = "",

299 - config: Optional[Llama4TextConfig] = None,

299 + config: Llama4TextConfig | None = None,

300 ) -> None:

301 super().__init__()

302

@@ -352,7 +352,7 @@ class Llama4DecoderLayer(nn.Module):

352 self,

353 positions: torch.Tensor,

354 hidden_states: torch.Tensor,

355 - residual: Optional[torch.Tensor],

355 + residual: torch.Tensor | None,

356 ) -> tuple[torch.Tensor, torch.Tensor]:

357 # Self Attention

358 if residual is None:

vllm/model_executor/models/llama4_eagle.py

+5 -6

@@ -17,7 +17,6 @@

17 # limitations under the License.

18

19 from collections.abc import Iterable

20 -from typing import Optional

20

21 import torch

22 import torch.nn as nn

@@ -49,7 +48,7 @@ class LlamaModel(nn.Module):

48 vllm_config: VllmConfig,

49 prefix: str = "",

50 start_layer_id: int = 0,

52 - quant_config: Optional[QuantizationConfig] = None,

51 + quant_config: QuantizationConfig | None = None,

52 ) -> None:

53 super().__init__()

54 self.config = vllm_config.speculative_config.draft_model_config.hf_config

@@ -81,10 +80,10 @@ class LlamaModel(nn.Module):

80

81 def forward(

82 self,

84 - input_ids: Optional[torch.Tensor],

83 + input_ids: torch.Tensor | None,

84 positions: torch.Tensor,

85 hidden_states: torch.Tensor,

87 - inputs_embeds: Optional[torch.Tensor] = None,

86 + inputs_embeds: torch.Tensor | None = None,

87 ) -> tuple[torch.Tensor, torch.Tensor]:

88 if inputs_embeds is None:

89 inputs_embeds = self.get_input_embeddings(input_ids)

@@ -136,7 +135,7 @@ class LlamaModel(nn.Module):

135 return loaded_params

136

137 def validate_and_update_config(

139 - self, start_layer_id: int, quant_config: Optional[QuantizationConfig] = None

138 + self, start_layer_id: int, quant_config: QuantizationConfig | None = None

139 ) -> None:

140 # yoco and moe is not supported by draft model yet

141 assert self.config.yoco_global_kv_layer is None

@@ -193,7 +192,7 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):

192 input_ids: torch.Tensor,

193 positions: torch.Tensor,

194 hidden_states: torch.Tensor,

196 - inputs_embeds: Optional[torch.Tensor] = None,

195 + inputs_embeds: torch.Tensor | None = None,

196 ) -> tuple[torch.Tensor, torch.Tensor]:

197 return self.model(input_ids, positions, hidden_states, inputs_embeds)

198

vllm/model_executor/models/llama_eagle.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable

5 -from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -28,7 +27,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):

27 vllm_config: VllmConfig,

28 disable_input_layernorm: bool,

29 prefix: str = "",

31 - config: Optional[LlamaConfig] = None,

30 + config: LlamaConfig | None = None,

31 ) -> None:

32 super().__init__(vllm_config, prefix=prefix, config=config)

33

@@ -155,7 +154,7 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):

154 input_ids: torch.Tensor,

155 positions: torch.Tensor,

156 hidden_states: torch.Tensor,

158 - inputs_embeds: Optional[torch.Tensor] = None,

157 + inputs_embeds: torch.Tensor | None = None,

158 ) -> tuple[torch.Tensor, torch.Tensor]:

159 if inputs_embeds is not None:

160 raise NotImplementedError(

vllm/model_executor/models/llama_eagle3.py

+8 -9

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable

5 -from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -35,7 +34,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):

34 self,

35 vllm_config: VllmConfig,

36 prefix: str = "",

38 - config: Optional[LlamaConfig] = None,

37 + config: LlamaConfig | None = None,

38 layer_idx: int = 0,

39 ) -> None:

40 super().__init__(vllm_config, prefix=prefix, config=config)

@@ -66,7 +65,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):

65 else:

66 self._residual_norm = self._norm_after_residual

67

69 - def get_quant_config(self, vllm_config: VllmConfig) -> Optional[QuantizationConfig]:

68 + def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:

69 """Use drafter's quantization config instead of verifier's."""

70 draft_model_config = vllm_config.speculative_config.draft_model_config

71 draft_load_config = vllm_config.load_config

@@ -96,7 +95,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):

95 positions: torch.Tensor,

96 embeds: torch.Tensor,

97 hidden_states: torch.Tensor,

99 - residual: Optional[torch.Tensor],

98 + residual: torch.Tensor | None,

99 ) -> tuple[torch.Tensor, torch.Tensor]:

100 if self.layer_idx == 0:

101 # First layer: concatenate embeds with hidden_states

@@ -182,7 +181,7 @@ class LlamaModel(nn.Module):

181 input_ids: torch.Tensor,

182 positions: torch.Tensor,

183 hidden_states: torch.Tensor,

185 - input_embeds: Optional[torch.Tensor] = None,

184 + input_embeds: torch.Tensor | None = None,

185 ) -> tuple[torch.Tensor, torch.Tensor]:

186 if input_embeds is None:

187 input_embeds = self.get_input_embeddings(input_ids)

@@ -268,8 +267,8 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

267 def get_input_embeddings(

268 self,

269 input_ids: torch.Tensor,

271 - multimodal_embeddings: Optional[NestedTensors] = None,

272 - is_multimodal: Optional[torch.Tensor] = None,

270 + multimodal_embeddings: NestedTensors | None = None,

271 + is_multimodal: torch.Tensor | None = None,

272 ) -> torch.Tensor:

273 return self.model.get_input_embeddings(input_ids)

274

@@ -278,14 +277,14 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

277 input_ids: torch.Tensor,

278 positions: torch.Tensor,

279 hidden_states: torch.Tensor,

281 - inputs_embeds: Optional[torch.Tensor] = None,

280 + inputs_embeds: torch.Tensor | None = None,

281 ) -> tuple[torch.Tensor, torch.Tensor]:

282 return self.model(input_ids, positions, hidden_states, inputs_embeds)

283

284 def compute_logits(

285 self,

286 hidden_states: torch.Tensor,

288 - ) -> Optional[torch.Tensor]:

287 + ) -> torch.Tensor | None:

288 logits = self.logits_processor(self.lm_head, hidden_states)

289 if self.draft_id_to_target_id is None:

290 assert logits.shape[1] == self.config.vocab_size, (

vllm/model_executor/models/llava.py

+28 -28

@@ -3,7 +3,7 @@

3

4 from abc import abstractmethod

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Annotated, Final, Literal, Optional, Protocol, TypeVar, Union

6 +from typing import Annotated, Final, Literal, Protocol, TypeAlias, TypeVar

7

8 import torch

9 import torch.nn as nn

@@ -93,7 +93,7 @@ class PixtralHFImagePixelInputs(TensorSchema):

93

94 type: Literal["pixel_values_pixtral"] = "pixel_values_pixtral"

95 pixel_values: Annotated[

96 - Union[torch.Tensor, list[torch.Tensor]],

96 + torch.Tensor | list[torch.Tensor],

97 TensorShape("bn", "c", "h", "w", dynamic_dims={"h", "w"}),

98 ]

99

@@ -110,9 +110,9 @@ class LlavaImageEmbeddingInputs(TensorSchema):

110 data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]

111

112

113 -LlavaImageInputs = Union[

114 - LlavaImagePixelInputs, PixtralHFImagePixelInputs, LlavaImageEmbeddingInputs

115 -]

113 +LlavaImageInputs: TypeAlias = (

114 + LlavaImagePixelInputs | PixtralHFImagePixelInputs | LlavaImageEmbeddingInputs

115 +)

116

117

118 class LlavaMultiModalProjector(nn.Module):

@@ -122,7 +122,7 @@ class LlavaMultiModalProjector(nn.Module):

122 text_hidden_size: int,

123 projector_hidden_act: str,

124 multimodal_projector_bias: bool,

125 - quant_config: Optional[QuantizationConfig] = None,

125 + quant_config: QuantizationConfig | None = None,

126 prefix: str = "",

127 ):

128 super().__init__()

@@ -154,7 +154,7 @@ class LlavaLikeConfig(Protocol):

154 vision_config: Final[PretrainedConfig]

155 image_token_index: Final[int]

156 vision_feature_select_strategy: Final[str]

157 - vision_feature_layer: Final[Union[int, list[int]]]

157 + vision_feature_layer: Final[int | list[int]]

158

159

160 class LlavaLikeProcessor(Protocol):

@@ -172,7 +172,7 @@ class BaseLlavaProcessingInfo(BaseProcessingInfo):

172 def get_hf_processor(self, **kwargs: object) -> LlavaLikeProcessor:

173 raise NotImplementedError

174

175 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

175 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

176 return {"image": None}

177

178 def get_num_image_tokens(

@@ -222,7 +222,7 @@ class LlavaDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

222 self,

223 seq_len: int,

224 mm_counts: Mapping[str, int],

225 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

225 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

226 ) -> MultiModalDataDict:

227 num_images = mm_counts.get("image", 0)

228

@@ -406,7 +406,7 @@ def _build_llava_or_pixtral_hf_processor(

406 info: _I,

407 dummy_inputs: BaseDummyInputsBuilder[_I],

408 *,

409 - cache: Optional[BaseMultiModalProcessorCache] = None,

409 + cache: BaseMultiModalProcessorCache | None = None,

410 ) -> BaseMultiModalProcessor:

411 if isinstance(info, PixtralHFProcessingInfo):

412 return PixtralHFMultiModalProcessor(

@@ -461,11 +461,11 @@ def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:

461

462 def init_vision_tower_for_llava(

463 hf_config: LlavaLikeConfig,

464 - quant_config: Optional[QuantizationConfig],

464 + quant_config: QuantizationConfig | None,

465 *,

466 - require_post_norm: Optional[bool] = None,

466 + require_post_norm: bool | None = None,

467 prefix: str = "",

468 -) -> Union[CLIPVisionModel, SiglipVisionModel, PixtralHFVisionModel]:

468 +) -> CLIPVisionModel | SiglipVisionModel | PixtralHFVisionModel:

469 vision_config = hf_config.vision_config

470

471 # Initialize the vision tower only up to the deepest required feature layer

@@ -524,7 +524,7 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

524 )

525

526 @classmethod

527 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

527 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

528 if modality.startswith("image"):

529 return "<image>"

530

@@ -585,7 +585,7 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

585

586 def _parse_and_validate_image_input(

587 self, **kwargs: object

588 - ) -> Optional[LlavaImageInputs]:

588 + ) -> LlavaImageInputs | None:

589 pixel_values = kwargs.pop("pixel_values", None)

590 image_embeds = kwargs.pop("image_embeds", None)

591

@@ -619,9 +619,9 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

619

620 def _image_pixels_to_features(

621 self,

622 - vision_tower: Union[CLIPVisionModel, SiglipVisionModel, PixtralHFVisionModel],

623 - pixel_values: Union[torch.Tensor, list[torch.Tensor]],

624 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

622 + vision_tower: CLIPVisionModel | SiglipVisionModel | PixtralHFVisionModel,

623 + pixel_values: torch.Tensor | list[torch.Tensor],

624 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

625 # NOTE: we skip the step to select the vision feature layer since

626 # this is already done inside the vision tower

627 return vision_tower(

@@ -631,8 +631,8 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

631

632 def _process_image_pixels(

633 self,

634 - inputs: Union[LlavaImagePixelInputs, PixtralHFImagePixelInputs],

635 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

634 + inputs: LlavaImagePixelInputs | PixtralHFImagePixelInputs,

635 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

636 assert self.vision_tower is not None

637

638 pixel_values = inputs["pixel_values"]

@@ -642,7 +642,7 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

642 def _process_image_input(

643 self,

644 image_input: LlavaImageInputs,

645 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

645 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

646 if image_input["type"] == "image_embeds":

647 return image_input["data"]

648

@@ -672,10 +672,10 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

672 self,

673 input_ids: torch.Tensor,

674 positions: torch.Tensor,

675 - intermediate_tensors: Optional[IntermediateTensors] = None,

676 - inputs_embeds: Optional[torch.Tensor] = None,

675 + intermediate_tensors: IntermediateTensors | None = None,

676 + inputs_embeds: torch.Tensor | None = None,

677 **kwargs: object,

678 - ) -> Union[torch.Tensor, IntermediateTensors]:

678 + ) -> torch.Tensor | IntermediateTensors:

679 """Run forward pass for LLaVA-1.5.

680

681 One key thing to understand is the `input_ids` already accounts for the

@@ -725,7 +725,7 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

725 def compute_logits(

726 self,

727 hidden_states: torch.Tensor,

728 - ) -> Optional[torch.Tensor]:

728 + ) -> torch.Tensor | None:

729 return self.language_model.compute_logits(hidden_states)

730

731 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -754,11 +754,11 @@ class MantisProcessingInfo(LlavaProcessingInfo):

754 class MantisMultiModalProcessor(LlavaMultiModalProcessor):

755 def apply(

756 self,

757 - prompt: Union[str, list[int]],

757 + prompt: str | list[int],

758 mm_data: MultiModalDataDict,

759 hf_processor_mm_kwargs: Mapping[str, object],

760 - tokenization_kwargs: Optional[Mapping[str, object]] = None,

761 - mm_uuids: Optional[MultiModalUUIDDict] = None,

760 + tokenization_kwargs: Mapping[str, object] | None = None,

761 + mm_uuids: MultiModalUUIDDict | None = None,

762 ) -> MultiModalInputs:

763 hf_config = self.info.get_hf_config()

764 image_token_id = hf_config.image_token_index

vllm/model_executor/models/llava_next.py

+17 -15

@@ -3,7 +3,7 @@

3

4 from abc import abstractmethod

5 from collections.abc import Iterable, Mapping

6 -from typing import Annotated, Final, Literal, Optional, Protocol, TypeVar, Union

6 +from typing import Annotated, Final, Literal, Protocol, TypeAlias, TypeVar

7

8 import torch

9 import torch.nn as nn

@@ -55,11 +55,11 @@ class LlavaNextImagePixelInputs(TensorSchema):

55

56 type: Literal["pixel_values"] = "pixel_values"

57 pixel_values: Annotated[

58 - Union[torch.Tensor, list[torch.Tensor]],

58 + torch.Tensor | list[torch.Tensor],

59 TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np"}),

60 ]

61

62 - image_sizes: Annotated[Optional[torch.Tensor], TensorShape("bn", 2)]

62 + image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)]

63 # This should be in `(height, width)` format.

64

65

@@ -75,7 +75,9 @@ class LlavaNextImageEmbeddingInputs(TensorSchema):

75 data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]

76

77

78 -LlavaNextImageInputs = Union[LlavaNextImagePixelInputs, LlavaNextImageEmbeddingInputs]

78 +LlavaNextImageInputs: TypeAlias = (

79 + LlavaNextImagePixelInputs | LlavaNextImageEmbeddingInputs

80 +)

81

82

83 class LlavaNextLikeConfig(LlavaLikeConfig, Protocol):

@@ -235,7 +237,7 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

237 )

238

239 @classmethod

238 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

240 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

241 if modality.startswith("image"):

242 return "<image>"

243

@@ -294,7 +296,7 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

296

297 def _parse_and_validate_image_input(

298 self, **kwargs: object

297 - ) -> Optional[LlavaNextImageInputs]:

299 + ) -> LlavaNextImageInputs | None:

300 pixel_values = kwargs.pop("pixel_values", None)

301 image_sizes = kwargs.pop("image_sizes", None)

302 image_embeds = kwargs.pop("image_embeds", None)

@@ -324,7 +326,7 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

326

327 def _image_pixels_to_features(

328 self,

327 - vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

329 + vision_tower: CLIPVisionModel | SiglipVisionModel,

330 pixel_values: torch.Tensor,

331 ) -> torch.Tensor:

332 # NOTE: we skip the step to select the vision feature layer since

@@ -424,7 +426,7 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

426 def _process_image_pixels(

427 self,

428 inputs: LlavaNextImagePixelInputs,

427 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

429 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

430 assert self.vision_tower is not None

431

432 pixel_values = inputs["pixel_values"]

@@ -456,7 +458,7 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

458 def _process_image_input(

459 self,

460 image_input: LlavaNextImageInputs,

459 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

461 + ) -> torch.Tensor | list[torch.Tensor]:

462 if image_input["type"] == "image_embeds":

463 return [image_input["data"]]

464

@@ -491,9 +493,9 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

493 def get_input_embeddings(

494 self,

495 input_ids: torch.Tensor,

494 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

496 + multimodal_embeddings: MultiModalEmbeddings | None = None,

497 *,

496 - is_multimodal: Optional[torch.Tensor] = None,

498 + is_multimodal: torch.Tensor | None = None,

499 # Multi-modal token ID may exceed vocab size

500 handle_oov_mm_token: bool = True,

501 ) -> torch.Tensor:

@@ -512,10 +514,10 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

514 self,

515 input_ids: torch.Tensor,

516 positions: torch.Tensor,

515 - intermediate_tensors: Optional[IntermediateTensors] = None,

516 - inputs_embeds: Optional[torch.Tensor] = None,

517 + intermediate_tensors: IntermediateTensors | None = None,

518 + inputs_embeds: torch.Tensor | None = None,

519 **kwargs: object,

518 - ) -> Union[torch.Tensor, IntermediateTensors]:

520 + ) -> torch.Tensor | IntermediateTensors:

521 """Run forward pass for LlaVA-NeXT.

522

523 One key thing to understand is the `input_ids` already accounts for the

@@ -573,7 +575,7 @@ model_executor.models.llava_next.LlavaNextProcessingInfo.get_num_image_tokens].

575 def compute_logits(

576 self,

577 hidden_states: torch.Tensor,

576 - ) -> Optional[torch.Tensor]:

578 + ) -> torch.Tensor | None:

579 return self.language_model.compute_logits(hidden_states)

580

581 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/llava_next_video.py

+11 -11

@@ -3,7 +3,7 @@

3

4 import math

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Annotated, Literal, Optional, Union

6 +from typing import Annotated, Literal

7

8 import torch

9 import torch.nn as nn

@@ -66,7 +66,7 @@ class LlavaNextVideoPixelInputs(TensorSchema):

66 type: Literal["pixel_values_videos"] = "pixel_values_videos"

67

68 pixel_values_videos: Annotated[

69 - Union[torch.Tensor, list[torch.Tensor]],

69 + torch.Tensor | list[torch.Tensor],

70 TensorShape("bn", "f", 3, "h", "w", dynamic_dims={"f"}),

71 ]

72

@@ -81,7 +81,7 @@ class LlavaNextVideoProcessingInfo(BaseProcessingInfo):

81 def get_hf_processor(self, **kwargs: object):

82 return self.ctx.get_hf_processor(LlavaNextVideoProcessor, **kwargs)

83

84 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

84 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

85 return {"video": 1}

86

87 def get_image_size_with_most_features(self) -> ImageSize:

@@ -165,7 +165,7 @@ class LlavaNextVideoDummyInputsBuilder(

165 self,

166 seq_len: int,

167 mm_counts: Mapping[str, int],

168 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

168 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

169 ) -> MultiModalDataDict:

170 num_videos = mm_counts.get("video", 0)

171

@@ -313,7 +313,7 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

313 )

314

315 @classmethod

316 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

316 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

317 if modality.startswith("image"):

318 return "<image>"

319 if modality.startswith("video"):

@@ -356,7 +356,7 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

356

357 def _parse_and_validate_video_input(

358 self, **kwargs: object

359 - ) -> Optional[LlavaNextVideoPixelInputs]:

359 + ) -> LlavaNextVideoPixelInputs | None:

360 """

361 A legal video input should have the following dimensions:

362 {

@@ -381,7 +381,7 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

381

382 def _video_pixels_to_features(

383 self,

384 - vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

384 + vision_tower: CLIPVisionModel | SiglipVisionModel,

385 pixel_values: torch.Tensor,

386 ) -> torch.Tensor:

387 # NOTE: we skip the step to select the vision feature layer since

@@ -433,10 +433,10 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

433 self,

434 input_ids: torch.Tensor,

435 positions: torch.Tensor,

436 - intermediate_tensors: Optional[IntermediateTensors] = None,

437 - inputs_embeds: Optional[torch.Tensor] = None,

436 + intermediate_tensors: IntermediateTensors | None = None,

437 + inputs_embeds: torch.Tensor | None = None,

438 **kwargs: object,

439 - ) -> Union[torch.Tensor, IntermediateTensors]:

439 + ) -> torch.Tensor | IntermediateTensors:

440 """Run forward pass for LlaVA-NeXT-Video.

441 Args:

442 input_ids: Flattened (concatenated) input_ids corresponding to a

@@ -455,7 +455,7 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

455 def compute_logits(

456 self,

457 hidden_states: torch.Tensor,

458 - ) -> Optional[torch.Tensor]:

458 + ) -> torch.Tensor | None:

459 return self.language_model.compute_logits(hidden_states)

460

461 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/llava_onevision.py

+23 -23

@@ -3,7 +3,7 @@

3

4 import math

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Annotated, Final, Literal, Optional, Protocol, Union

6 +from typing import Annotated, Final, Literal, Protocol, TypeAlias

7

8 import torch

9 import torch.nn as nn

@@ -69,7 +69,7 @@ class LlavaOnevisionVideoPixelInputs(TensorSchema):

69 type: Literal["pixel_values_videos"] = "pixel_values_videos"

70

71 pixel_values_videos: Annotated[

72 - Union[torch.Tensor, list[torch.Tensor]],

72 + torch.Tensor | list[torch.Tensor],

73 TensorShape("bn", "f", 3, "h", "w", dynamic_dims={"f"}),

74 ]

75

@@ -90,11 +90,11 @@ class LlavaOnevisionImagePixelInputs(TensorSchema):

90 type: Literal["pixel_values"] = "pixel_values"

91

92 pixel_values: Annotated[

93 - Union[torch.Tensor, list[torch.Tensor]],

93 + torch.Tensor | list[torch.Tensor],

94 TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np"}),

95 ]

96

97 - image_sizes: Annotated[Optional[torch.Tensor], TensorShape("bn", 2)]

97 + image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)]

98

99

100 class LlavaOnevisionImageEmbeddingInputs(TensorSchema):

@@ -113,13 +113,13 @@ class LlavaOnevisionImageEmbeddingInputs(TensorSchema):

113 ]

114

115

116 -LlavaOnevisionImageInputs = Union[

117 - LlavaOnevisionImagePixelInputs, LlavaOnevisionImageEmbeddingInputs

118 -]

116 +LlavaOnevisionImageInputs: TypeAlias = (

117 + LlavaOnevisionImagePixelInputs | LlavaOnevisionImageEmbeddingInputs

118 +)

119

120 -LlavaOnevisionMultiInputs = Union[

121 - LlavaOnevisionImageInputs, LlavaOnevisionVideoPixelInputs

122 -]

120 +LlavaOnevisionMultiInputs: TypeAlias = (

121 + LlavaOnevisionImageInputs | LlavaOnevisionVideoPixelInputs

122 +)

123

124

125 class LlavaOnevisionLikeConfig(LlavaNextLikeConfig, Protocol):

@@ -133,7 +133,7 @@ class LlavaOnevisionProcessingInfo(LlavaNextProcessingInfo):

133 def get_hf_processor(self, **kwargs: object):

134 return self.ctx.get_hf_processor(LlavaOnevisionProcessor, **kwargs)

135

136 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

136 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

137 return {"image": None, "video": None}

138

139 # Based on: https://github.com/huggingface/text-generation-inference/blob/v3.0.1/server/text_generation_server/models/vlm_causal_lm.py#L86

@@ -276,7 +276,7 @@ class LlavaOnevisionDummyInputsBuilder(

276 self,

277 seq_len: int,

278 mm_counts: Mapping[str, int],

279 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

279 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

280 ) -> MultiModalDataDict:

281 num_images = mm_counts.get("image", 0)

282 num_videos = mm_counts.get("video", 0)

@@ -493,7 +493,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

493 )

494

495 @classmethod

496 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

496 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

497 if modality.startswith("image"):

498 return "<image>"

499 if modality.startswith("video"):

@@ -531,7 +531,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

531

532 def _parse_and_validate_image_input(

533 self, **kwargs: object

534 - ) -> Optional[LlavaOnevisionImageInputs]:

534 + ) -> LlavaOnevisionImageInputs | None:

535 pixel_values = kwargs.pop("pixel_values", None)

536 image_sizes = kwargs.pop("image_sizes", None)

537 image_embeds = kwargs.pop("image_embeds", None)

@@ -560,7 +560,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

560

561 def _parse_and_validate_video_input(

562 self, **kwargs: object

563 - ) -> Optional[LlavaOnevisionVideoPixelInputs]:

563 + ) -> LlavaOnevisionVideoPixelInputs | None:

564 """

565 A legal video input should have the following dimensions:

566 {

@@ -606,7 +606,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

606

607 def _image_pixels_to_features(

608 self,

609 - vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

609 + vision_tower: CLIPVisionModel | SiglipVisionModel,

610 pixel_values: torch.Tensor,

611 ) -> torch.Tensor:

612 # NOTE: we skip the step to select the vision feature layer since

@@ -726,7 +726,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

726 def _process_image_pixels(

727 self,

728 inputs: LlavaOnevisionImagePixelInputs,

729 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

729 + ) -> torch.Tensor | list[torch.Tensor]:

730 assert self.vision_tower is not None

731

732 pixel_values = inputs["pixel_values"]

@@ -761,7 +761,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

761 def _process_image_input(

762 self,

763 image_input: LlavaOnevisionImageInputs,

764 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

764 + ) -> torch.Tensor | list[torch.Tensor]:

765 if image_input["type"] == "image_embeds":

766 return [image_input["data"]]

767

@@ -788,7 +788,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

788

789 def _video_pixels_to_features(

790 self,

791 - vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

791 + vision_tower: CLIPVisionModel | SiglipVisionModel,

792 pixel_values: torch.Tensor,

793 ) -> torch.Tensor:

794 # NOTE: we skip the step to select the vision feature layer since

@@ -893,10 +893,10 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

893 self,

894 input_ids: torch.Tensor,

895 positions: torch.Tensor,

896 - intermediate_tensors: Optional[IntermediateTensors] = None,

897 - inputs_embeds: Optional[torch.Tensor] = None,

896 + intermediate_tensors: IntermediateTensors | None = None,

897 + inputs_embeds: torch.Tensor | None = None,

898 **kwargs: object,

899 - ) -> Union[torch.Tensor, IntermediateTensors]:

899 + ) -> torch.Tensor | IntermediateTensors:

900 """Run forward pass for LlaVA-Onevision.

901 Args:

902 input_ids: Flattened (concatenated) input_ids corresponding to a

@@ -915,7 +915,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp

915 def compute_logits(

916 self,

917 hidden_states: torch.Tensor,

918 - ) -> Optional[torch.Tensor]:

918 + ) -> torch.Tensor | None:

919 return self.language_model.compute_logits(hidden_states)

920

921 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/longcat_flash.py

+13 -14

@@ -36,7 +36,6 @@

36 import typing

37 from collections.abc import Callable, Iterable

38 from itertools import islice

39 -from typing import Optional, Union

39

40 import torch

41 from torch import nn

@@ -194,7 +193,7 @@ class FlashMLP(nn.Module):

193 hidden_size: int,

194 intermediate_size: int,

195 hidden_act: str,

197 - quant_config: Optional[QuantizationConfig] = None,

196 + quant_config: QuantizationConfig | None = None,

197 reduce_results: bool = True,

198 prefix: str = "",

199 ) -> None:

@@ -270,8 +269,8 @@ class LongcatMoe(nn.Module):

269 top_k: int,

270 hidden_size: int,

271 intermediate_size: int,

273 - params_dtype: Optional[torch.dtype] = None,

274 - quant_config: Optional[QuantizationConfig] = None,

272 + params_dtype: torch.dtype | None = None,

273 + quant_config: QuantizationConfig | None = None,

274 prefix: str = "",

275 enable_eplb: bool = False,

276 ):

@@ -329,8 +328,8 @@ class FlashDecoderLayer(nn.Module):

328 self,

329 vllm_config: VllmConfig,

330 config: FlashConfig,

332 - cache_config: Optional[CacheConfig] = None,

333 - quant_config: Optional[QuantizationConfig] = None,

331 + cache_config: CacheConfig | None = None,

332 + quant_config: QuantizationConfig | None = None,

333 prefix: str = "",

334 enable_eplb: bool = False,

335 ) -> None:

@@ -415,7 +414,7 @@ class FlashDecoderLayer(nn.Module):

414 self,

415 positions: torch.Tensor,

416 hidden_states: torch.Tensor,

418 - residual: Optional[torch.Tensor],

417 + residual: torch.Tensor | None,

418 ) -> tuple[torch.Tensor, torch.Tensor]:

419 if residual is None:

420 residual = hidden_states

@@ -506,9 +505,9 @@ class FlashModel(nn.Module):

505 self,

506 input_ids: torch.Tensor,

507 positions: torch.Tensor,

509 - intermediate_tensors: Optional[IntermediateTensors] = None,

510 - inputs_embeds: Optional[torch.Tensor] = None,

511 - ) -> Union[torch.Tensor, IntermediateTensors]:

508 + intermediate_tensors: IntermediateTensors | None = None,

509 + inputs_embeds: torch.Tensor | None = None,

510 + ) -> torch.Tensor | IntermediateTensors:

511 if get_pp_group().is_first_rank:

512 if inputs_embeds is not None:

513 hidden_states = inputs_embeds

@@ -592,9 +591,9 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

591 self,

592 input_ids: torch.Tensor,

593 positions: torch.Tensor,

595 - intermediate_tensors: Optional[IntermediateTensors] = None,

596 - inputs_embeds: Optional[torch.Tensor] = None,

597 - ) -> Union[torch.Tensor, IntermediateTensors]:

594 + intermediate_tensors: IntermediateTensors | None = None,

595 + inputs_embeds: torch.Tensor | None = None,

596 + ) -> torch.Tensor | IntermediateTensors:

597 hidden_states = self.model(

598 input_ids, positions, intermediate_tensors, inputs_embeds

599 )

@@ -603,7 +602,7 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

602 def compute_logits(

603 self,

604 hidden_states: torch.Tensor,

606 - ) -> Optional[torch.Tensor]:

605 + ) -> torch.Tensor | None:

606 logits = self.logits_processor(self.lm_head, hidden_states)

607 return logits

608

vllm/model_executor/models/longcat_flash_mtp.py

+8 -9

@@ -4,7 +4,6 @@

4 # Adapted from

5 # https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/deepseek_mtp.py

6 from collections.abc import Iterable

7 -from typing import Optional

7

8 import torch

9 import torch.nn as nn

@@ -35,7 +34,7 @@ class LongCatMultiTokenPredictorLayer(nn.Module):

34 config: PretrainedConfig,

35 prefix: str,

36 vllm_config: VllmConfig,

38 - quant_config: Optional[QuantizationConfig] = None,

37 + quant_config: QuantizationConfig | None = None,

38 ) -> None:

39 super().__init__()

40 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

@@ -55,7 +54,7 @@ class LongCatMultiTokenPredictorLayer(nn.Module):

54 input_ids: torch.Tensor,

55 positions: torch.Tensor,

56 previous_hidden_states: torch.Tensor,

58 - inputs_embeds: Optional[torch.Tensor] = None,

57 + inputs_embeds: torch.Tensor | None = None,

58 spec_step_index: int = 0,

59 ) -> torch.Tensor:

60 assert inputs_embeds is not None

@@ -78,7 +77,7 @@ class LongCatMultiTokenPredictor(nn.Module):

77 self,

78 *,

79 vllm_config: VllmConfig,

81 - quant_config: Optional[QuantizationConfig] = None,

80 + quant_config: QuantizationConfig | None = None,

81 prefix: str = "",

82 ):

83 super().__init__()

@@ -110,7 +109,7 @@ class LongCatMultiTokenPredictor(nn.Module):

109 input_ids: torch.Tensor,

110 positions: torch.Tensor,

111 previous_hidden_states: torch.Tensor,

113 - inputs_embeds: Optional[torch.Tensor] = None,

112 + inputs_embeds: torch.Tensor | None = None,

113 spec_step_idx: int = 0,

114 ) -> torch.Tensor:

115 if inputs_embeds is None:

@@ -155,8 +154,8 @@ class LongCatFlashMTP(nn.Module, SupportsPP):

154 input_ids: torch.Tensor,

155 positions: torch.Tensor,

156 hidden_states: torch.Tensor,

158 - intermediate_tensors: Optional[IntermediateTensors] = None,

159 - inputs_embeds: Optional[torch.Tensor] = None,

157 + intermediate_tensors: IntermediateTensors | None = None,

158 + inputs_embeds: torch.Tensor | None = None,

159 spec_step_idx: int = 0,

160 ) -> torch.Tensor:

161 hidden_states = self.model(

@@ -168,7 +167,7 @@ class LongCatFlashMTP(nn.Module, SupportsPP):

167 self,

168 hidden_states: torch.Tensor,

169 spec_step_idx: int = 0,

171 - ) -> Optional[torch.Tensor]:

170 + ) -> torch.Tensor | None:

171 logits = self.logits_processor(self.lm_head, hidden_states)

172 return logits

173

@@ -344,7 +343,7 @@ class LongCatFlashMTP(nn.Module, SupportsPP):

343

344 def get_spec_layer_idx_from_weight_name(

345 self, config: PretrainedConfig, weight_name: str

347 - ) -> Optional[int]:

346 + ) -> int | None:

347 if "model.mtp" in weight_name:

348 return config.num_hidden_layers * 2

349 return None

vllm/model_executor/models/mamba.py

+9 -10

@@ -4,7 +4,6 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import Optional

7

8 import torch

9 from torch import nn

@@ -49,10 +48,10 @@ class MambaDecoderLayer(nn.Module):

48 def __init__(

49 self,

50 config: MambaConfig,

52 - model_config: Optional[ModelConfig] = None,

53 - cache_config: Optional[CacheConfig] = None,

54 - quant_config: Optional[QuantizationConfig] = None,

55 - is_lora_enabled: Optional[bool] = False,

51 + model_config: ModelConfig | None = None,

52 + cache_config: CacheConfig | None = None,

53 + quant_config: QuantizationConfig | None = None,

54 + is_lora_enabled: bool | None = False,

55 prefix: str = "",

56 ) -> None:

57 super().__init__()

@@ -83,7 +82,7 @@ class MambaDecoderLayer(nn.Module):

82 def forward(

83 self,

84 hidden_states: torch.Tensor,

86 - residual: Optional[torch.Tensor],

85 + residual: torch.Tensor | None,

86 **kwargs,

87 ):

88 if residual is None:

@@ -149,8 +148,8 @@ class MambaModel(nn.Module):

148 self,

149 input_ids: torch.Tensor,

150 positions: torch.Tensor,

152 - intermediate_tensors: Optional[IntermediateTensors] = None,

153 - inputs_embeds: Optional[torch.Tensor] = None,

151 + intermediate_tensors: IntermediateTensors | None = None,

152 + inputs_embeds: torch.Tensor | None = None,

153 ) -> torch.Tensor:

154 if get_pp_group().is_first_rank:

155 if inputs_embeds is not None:

@@ -244,8 +243,8 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):

243 self,

244 input_ids: torch.Tensor,

245 positions: torch.Tensor,

247 - intermediate_tensors: Optional[IntermediateTensors] = None,

248 - inputs_embeds: Optional[torch.Tensor] = None,

246 + intermediate_tensors: IntermediateTensors | None = None,

247 + inputs_embeds: torch.Tensor | None = None,

248 **kwargs,

249 ):

250 hidden_states = self.backbone(

vllm/model_executor/models/mamba2.py

+8 -9

@@ -3,7 +3,6 @@

3 """PyTorch MAMBA2 model."""

4

5 from collections.abc import Iterable

6 -from typing import Optional

6

7 import torch

8 from torch import nn

@@ -44,9 +43,9 @@ class Mamba2DecoderLayer(nn.Module):

43 def __init__(

44 self,

45 config: MambaConfig,

47 - model_config: Optional[ModelConfig] = None,

48 - cache_config: Optional[CacheConfig] = None,

49 - quant_config: Optional[QuantizationConfig] = None,

46 + model_config: ModelConfig | None = None,

47 + cache_config: CacheConfig | None = None,

48 + quant_config: QuantizationConfig | None = None,

49 prefix: str = "",

50 ) -> None:

51 super().__init__()

@@ -76,7 +75,7 @@ class Mamba2DecoderLayer(nn.Module):

75 def forward(

76 self,

77 hidden_states: torch.Tensor,

79 - residual: Optional[torch.Tensor],

78 + residual: torch.Tensor | None,

79 **kwargs,

80 ):

81 if residual is None:

@@ -142,8 +141,8 @@ class Mamba2Model(nn.Module):

141 self,

142 input_ids: torch.Tensor,

143 positions: torch.Tensor,

145 - intermediate_tensors: Optional[IntermediateTensors] = None,

146 - inputs_embeds: Optional[torch.Tensor] = None,

144 + intermediate_tensors: IntermediateTensors | None = None,

145 + inputs_embeds: torch.Tensor | None = None,

146 ) -> torch.Tensor:

147 if get_pp_group().is_first_rank:

148 if inputs_embeds is not None:

@@ -277,8 +276,8 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):

276 self,

277 input_ids: torch.Tensor,

278 positions: torch.Tensor,

280 - intermediate_tensors: Optional[IntermediateTensors] = None,

281 - inputs_embeds: Optional[torch.Tensor] = None,

279 + intermediate_tensors: IntermediateTensors | None = None,

280 + inputs_embeds: torch.Tensor | None = None,

281 **kwargs,

282 ):

283 hidden_states = self.backbone(

vllm/model_executor/models/midashenglm.py

+26 -26

@@ -25,8 +25,8 @@

25

26 import collections

27 import collections.abc

28 -from collections.abc import Iterable, Mapping, Sequence

29 -from typing import Any, Callable, Optional, TypedDict, Union, cast

28 +from collections.abc import Callable, Iterable, Mapping, Sequence

29 +from typing import Any, TypeAlias, TypedDict, cast

30

31 import numpy as np

32 import torch

@@ -66,7 +66,7 @@ from vllm.transformers_utils.configs.midashenglm import DashengConfig

66 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

67 from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix

68

69 -_Tuple2 = Union[int, tuple[int, int], Sequence[int]]

69 +_Tuple2: TypeAlias = int | tuple[int, int] | Sequence[int]

70

71

72 def _resolve_tuple2(x: _Tuple2) -> tuple[int, int]:

@@ -105,7 +105,7 @@ class AudioPatchEmbed(nn.Module):

105 patch_stride: _Tuple2 = 16,

106 in_chans: int = 1,

107 embed_dim: int = 768,

108 - norm_layer: Optional[Callable] = None,

108 + norm_layer: Callable | None = None,

109 flatten: bool = False,

110 ):

111 super().__init__()

@@ -151,9 +151,9 @@ class DashengMlp(nn.Module):

151 def __init__(

152 self,

153 in_features: int,

154 - hidden_features: Optional[int] = None,

155 - out_features: Optional[int] = None,

156 - quant_config: Optional[QuantizationConfig] = None,

154 + hidden_features: int | None = None,

155 + out_features: int | None = None,

156 + quant_config: QuantizationConfig | None = None,

157 prefix: str = "",

158 ):

159 super().__init__()

@@ -186,7 +186,7 @@ class DashengAttention(nn.Module):

186 dim: int,

187 num_heads: int = 8,

188 qkv_bias: bool = False,

189 - quant_config: Optional[QuantizationConfig] = None,

189 + quant_config: QuantizationConfig | None = None,

190 prefix: str = "",

191 ):

192 super().__init__()

@@ -226,7 +226,7 @@ class DashengAttention(nn.Module):

226 prefix=f"{prefix}.proj",

227 )

228

229 - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None):

229 + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):

230 B, N, C = x.shape

231

232 qkv, _ = self.qkv(x)

@@ -253,8 +253,8 @@ class DashengBlock(nn.Module):

253 num_heads: int,

254 mlp_ratio: float = 4.0,

255 qkv_bias: bool = False,

256 - init_values: Optional[float] = None,

257 - quant_config: Optional[QuantizationConfig] = None,

256 + init_values: float | None = None,

257 + quant_config: QuantizationConfig | None = None,

258 prefix: str = "",

259 ):

260 super().__init__()

@@ -285,7 +285,7 @@ class DashengBlock(nn.Module):

285 def forward(

286 self,

287 x: torch.Tensor,

288 - mask: Optional[torch.Tensor] = None,

288 + mask: torch.Tensor | None = None,

289 ) -> torch.Tensor:

290 x = x + self.ls1(self.attn(self.norm1(x), mask))

291 x = x + self.ls2(self.mlp(self.norm2(x)))

@@ -349,7 +349,7 @@ class DashengAudioTransformer(nn.Module):

349 def __init__(

350 self,

351 config: DashengConfig,

352 - quant_config: Optional[QuantizationConfig] = None,

352 + quant_config: QuantizationConfig | None = None,

353 prefix: str = "",

354 ):

355 super().__init__()

@@ -393,7 +393,7 @@ class DashengAudioTransformer(nn.Module):

393 def forward_features(

394 self,

395 x: torch.Tensor,

396 - mask: Optional[torch.Tensor] = None,

396 + mask: torch.Tensor | None = None,

397 ) -> torch.Tensor:

398 t = x.shape[-1]

399 x = x + self.time_pos_embed[:, :, :, :t]

@@ -418,8 +418,8 @@ class DashengAudioTransformer(nn.Module):

418 def forward(

419 self,

420 x: torch.Tensor,

421 - x_length: Optional[torch.Tensor] = None,

422 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

421 + x_length: torch.Tensor | None = None,

422 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

423 x = self.front_end(x)

424 x = x.to(self.time_pos_embed.dtype)

425 target_length_in_patches = self.target_length // 4

@@ -462,8 +462,8 @@ class AudioProjectorSubsample(nn.Module):

462 in_dim: int,

463 out_dim: int,

464 downsample_rate=5,

465 - dtype: Optional[torch.dtype] = None,

466 - quant_config: Optional[QuantizationConfig] = None,

465 + dtype: torch.dtype | None = None,

466 + quant_config: QuantizationConfig | None = None,

467 prefix: str = "",

468 ):

469 super().__init__()

@@ -524,7 +524,7 @@ class MiDashengLMProcessingInfo(BaseProcessingInfo):

524 feature_extractor = hf_processor.feature_extractor

525 return feature_extractor

526

527 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

527 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

528 return {"audio": None}

529

530 def get_min_audio_len(self):

@@ -550,7 +550,7 @@ class MiDashengLMDummyInputsBuilder(BaseDummyInputsBuilder[MiDashengLMProcessing

550 self,

551 seq_len: int,

552 mm_counts: Mapping[str, int],

553 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

553 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

554 ) -> MultiModalDataDict:

555 num_audios = mm_counts.get("audio", 0)

556

@@ -689,7 +689,7 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

689 }

690

691 @classmethod

692 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

692 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

693 if modality.startswith("audio"):

694 return "<|audio_bos|><|AUDIO|><|audio_eos|>"

695

@@ -750,7 +750,7 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

750

751 def _parse_and_validate_audio_input(

752 self, **kwargs: object

753 - ) -> Optional[MiDashengLMAudioInputs]:

753 + ) -> MiDashengLMAudioInputs | None:

754 input_values = kwargs.pop("input_values", None)

755 audio_length = kwargs.pop("audio_length", None)

756

@@ -820,10 +820,10 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

820 self,

821 input_ids: torch.Tensor,

822 positions: torch.Tensor,

823 - intermediate_tensors: Optional[IntermediateTensors] = None,

824 - inputs_embeds: Optional[torch.Tensor] = None,

823 + intermediate_tensors: IntermediateTensors | None = None,

824 + inputs_embeds: torch.Tensor | None = None,

825 **kwargs: object,

826 - ) -> Union[torch.Tensor, IntermediateTensors]:

826 + ) -> torch.Tensor | IntermediateTensors:

827 if intermediate_tensors is not None:

828 inputs_embeds = None

829 elif inputs_embeds is None:

@@ -845,7 +845,7 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

845 def compute_logits(

846 self,

847 hidden_states: torch.Tensor,

848 - ) -> Optional[torch.Tensor]:

848 + ) -> torch.Tensor | None:

849 return self.decoder.compute_logits(hidden_states)

850

851 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/mimo.py

+4 -5

@@ -28,7 +28,6 @@

28

29 from collections.abc import Iterable

30 from itertools import islice

31 -from typing import Optional, Union

31

32 import torch

33 import torch.nn as nn

@@ -64,9 +63,9 @@ class MiMoModel(Qwen2Model):

63 self,

64 input_ids: torch.Tensor,

65 positions: torch.Tensor,

67 - intermediate_tensors: Optional[IntermediateTensors] = None,

68 - inputs_embeds: Optional[torch.Tensor] = None,

69 - ) -> Union[torch.Tensor, IntermediateTensors]:

66 + intermediate_tensors: IntermediateTensors | None = None,

67 + inputs_embeds: torch.Tensor | None = None,

68 + ) -> torch.Tensor | IntermediateTensors:

69 if get_pp_group().is_first_rank:

70 if inputs_embeds is not None:

71 hidden_states = inputs_embeds

@@ -185,7 +184,7 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):

184 def compute_logits(

185 self,

186 hidden_states: torch.Tensor,

188 - ) -> Optional[torch.Tensor]:

187 + ) -> torch.Tensor | None:

188 hidden_states = self.model.norm(hidden_states)

189 logits = self.logits_processor(self.lm_head, hidden_states)

190 return logits

vllm/model_executor/models/mimo_mtp.py

+6 -7

@@ -21,7 +21,6 @@

21 """Inference-only MiMo-MTP model."""

22

23 from collections.abc import Iterable

24 -from typing import Optional

24

25 import torch

26 import torch.nn as nn

@@ -48,8 +47,8 @@ class MiMoMultiTokenPredictorLayer(nn.Module):

47 config: PretrainedConfig,

48 prefix: str,

49 model_config: ModelConfig,

51 - cache_config: Optional[CacheConfig] = None,

52 - quant_config: Optional[QuantizationConfig] = None,

50 + cache_config: CacheConfig | None = None,

51 + quant_config: QuantizationConfig | None = None,

52 ) -> None:

53 super().__init__()

54

@@ -129,7 +128,7 @@ class MiMoMultiTokenPredictor(nn.Module):

128 input_ids: torch.Tensor,

129 positions: torch.Tensor,

130 previous_hidden_states: torch.Tensor,

132 - inputs_embeds: Optional[torch.Tensor] = None,

131 + inputs_embeds: torch.Tensor | None = None,

132 spec_step_idx: int = 0,

133 ) -> torch.Tensor:

134 if inputs_embeds is None:

@@ -173,8 +172,8 @@ class MiMoMTP(nn.Module):

172 input_ids: torch.Tensor,

173 positions: torch.Tensor,

174 hidden_states: torch.Tensor,

176 - intermediate_tensors: Optional[IntermediateTensors] = None,

177 - inputs_embeds: Optional[torch.Tensor] = None,

175 + intermediate_tensors: IntermediateTensors | None = None,

176 + inputs_embeds: torch.Tensor | None = None,

177 spec_step_idx: int = 0,

178 ) -> torch.Tensor:

179 assert spec_step_idx == 0, "mimo_mtp only support predict one token now"

@@ -187,7 +186,7 @@ class MiMoMTP(nn.Module):

186 self,

187 hidden_states: torch.Tensor,

188 spec_step_idx: int = 0,

190 - ) -> Optional[torch.Tensor]:

189 + ) -> torch.Tensor | None:

190 return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)

191

192 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/minicpm.py

+19 -23

@@ -27,7 +27,7 @@

27 import math

28 from collections.abc import Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import torch

33 from torch import nn

@@ -89,8 +89,8 @@ class MiniCPMMoE(nn.Module):

89 top_k: int,

90 hidden_size: int,

91 intermediate_size: int,

92 - params_dtype: Optional[torch.dtype] = None,

93 - tp_size: Optional[int] = None,

92 + params_dtype: torch.dtype | None = None,

93 + tp_size: int | None = None,

94 ):

95 super().__init__()

96 self.tp_size = tp_size or get_tensor_model_parallel_world_size()

@@ -190,7 +190,7 @@ class MiniCPMMLP(nn.Module):

190 intermediate_size: int,

191 hidden_act: str,

192 hidden_act_param: float,

193 - quant_config: Optional[QuantizationConfig] = None,

193 + quant_config: QuantizationConfig | None = None,

194 ) -> None:

195 super().__init__()

196 self.gate_up_proj = MergedColumnParallelLinear(

@@ -223,10 +223,10 @@ class MiniCPMAttention(nn.Module):

223 num_heads: int,

224 num_kv_heads: int,

225 rope_theta: float = 10000,

226 - rope_scaling: Optional[dict[str, Any]] = None,

226 + rope_scaling: dict[str, Any] | None = None,

227 max_position_embeddings: int = 8192,

228 - cache_config: Optional[CacheConfig] = None,

229 - quant_config: Optional[QuantizationConfig] = None,

228 + cache_config: CacheConfig | None = None,

229 + quant_config: QuantizationConfig | None = None,

230 prefix: str = "",

231 ) -> None:

232 super().__init__()

@@ -305,8 +305,8 @@ class MiniCPMDecoderLayer(nn.Module):

305 def __init__(

306 self,

307 config: PretrainedConfig,

308 - cache_config: Optional[CacheConfig] = None,

309 - quant_config: Optional[QuantizationConfig] = None,

308 + cache_config: CacheConfig | None = None,

309 + quant_config: QuantizationConfig | None = None,

310 prefix: str = "",

311 ) -> None:

312 super().__init__()

@@ -362,7 +362,7 @@ class MiniCPMDecoderLayer(nn.Module):

362 self,

363 positions: torch.Tensor,

364 hidden_states: torch.Tensor,

365 - residual: Optional[torch.Tensor],

365 + residual: torch.Tensor | None,

366 ) -> tuple[torch.Tensor, torch.Tensor]:

367 # Self Attention

368 residual = hidden_states

@@ -425,8 +425,8 @@ class MiniCPMModel(nn.Module):

425 self,

426 prefix: str,

427 config: PretrainedConfig,

428 - cache_config: Optional[CacheConfig],

429 - quant_config: Optional[QuantizationConfig],

428 + cache_config: CacheConfig | None,

429 + quant_config: QuantizationConfig | None,

430 ):

431 self.start_layer, self.end_layer, self.layers = make_layers(

432 config.num_hidden_layers,

@@ -444,11 +444,9 @@ class MiniCPMModel(nn.Module):

444 self,

445 input_ids: torch.Tensor,

446 positions: torch.Tensor,

447 - intermediate_tensors: Optional[IntermediateTensors] = None,

448 - inputs_embeds: Optional[torch.Tensor] = None,

449 - ) -> Union[

450 - torch.Tensor, IntermediateTensors, tuple[torch.Tensor, list[torch.Tensor]]

451 - ]:

447 + intermediate_tensors: IntermediateTensors | None = None,

448 + inputs_embeds: torch.Tensor | None = None,

449 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:

450 if get_pp_group().is_first_rank:

451 if inputs_embeds is not None:

452 hidden_states = inputs_embeds

@@ -633,11 +631,9 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

631 self,

632 input_ids: torch.Tensor,

633 positions: torch.Tensor,

636 - intermediate_tensors: Optional[IntermediateTensors] = None,

637 - inputs_embeds: Optional[torch.Tensor] = None,

638 - ) -> Union[

639 - torch.Tensor, IntermediateTensors, tuple[torch.Tensor, list[torch.Tensor]]

640 - ]:

634 + intermediate_tensors: IntermediateTensors | None = None,

635 + inputs_embeds: torch.Tensor | None = None,

636 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:

637 model_output = self.model(

638 input_ids, positions, intermediate_tensors, inputs_embeds

639 )

@@ -658,7 +654,7 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

654 def compute_logits(

655 self,

656 hidden_states: torch.Tensor,

661 - ) -> Optional[torch.Tensor]:

657 + ) -> torch.Tensor | None:

658 logits = self.logits_processor(self.lm_head, hidden_states)

659 return logits

660

vllm/model_executor/models/minicpm3.py

+6 -6

@@ -25,7 +25,7 @@

25 # limitations under the License.

26 """Inference-only MiniCPM3 model compatible with HuggingFace weights."""

27

28 -from typing import Any, Optional

28 +from typing import Any

29

30 import torch

31 from torch import nn

@@ -63,10 +63,10 @@ class MiniCPM3Attention(nn.Module):

63 q_lora_rank: int,

64 kv_lora_rank: int,

65 rope_theta: float = 10000,

66 - rope_scaling: Optional[dict[str, Any]] = None,

66 + rope_scaling: dict[str, Any] | None = None,

67 max_position_embeddings: int = 8192,

68 - cache_config: Optional[CacheConfig] = None,

69 - quant_config: Optional[QuantizationConfig] = None,

68 + cache_config: CacheConfig | None = None,

69 + quant_config: QuantizationConfig | None = None,

70 prefix: str = "",

71 ) -> None:

72 super().__init__()

@@ -214,8 +214,8 @@ class MiniCPM3Model(MiniCPMModel):

214 self,

215 prefix: str,

216 config: PretrainedConfig,

217 - cache_config: Optional[CacheConfig],

218 - quant_config: Optional[QuantizationConfig],

217 + cache_config: CacheConfig | None,

218 + quant_config: QuantizationConfig | None,

219 ):

220 self.start_layer, self.end_layer, self.layers = make_layers(

221 config.num_hidden_layers,

vllm/model_executor/models/minicpm_eagle.py

+7 -8

@@ -26,7 +26,6 @@

26

27 import math

28 from collections.abc import Iterable

29 -from typing import Optional, Union

29

30 import torch

31 from torch import nn

@@ -61,8 +60,8 @@ class EagleMiniCPMDecoderLayer(nn.Module):

60 def __init__(

61 self,

62 config: PretrainedConfig,

64 - cache_config: Optional[CacheConfig] = None,

65 - quant_config: Optional[QuantizationConfig] = None,

63 + cache_config: CacheConfig | None = None,

64 + quant_config: QuantizationConfig | None = None,

65 prefix: str = "",

66 ) -> None:

67 super().__init__()

@@ -118,7 +117,7 @@ class EagleMiniCPMDecoderLayer(nn.Module):

117 self,

118 positions: torch.Tensor,

119 hidden_states: torch.Tensor,

121 - residual: Optional[torch.Tensor],

120 + residual: torch.Tensor | None,

121 ) -> tuple[torch.Tensor, torch.Tensor]:

122 # Self Attention

123 residual = hidden_states

@@ -185,8 +184,8 @@ class EagleMiniCPMModel(nn.Module):

184 self,

185 prefix: str,

186 config: PretrainedConfig,

188 - cache_config: Optional[CacheConfig],

189 - quant_config: Optional[QuantizationConfig],

187 + cache_config: CacheConfig | None,

188 + quant_config: QuantizationConfig | None,

189 start_layer: int,

190 ):

191 self.eagle_layers = nn.ModuleList(

@@ -210,7 +209,7 @@ class EagleMiniCPMModel(nn.Module):

209 input_ids: torch.Tensor,

210 positions: torch.Tensor,

211 hidden_states: torch.Tensor,

213 - ) -> Union[torch.Tensor, IntermediateTensors]:

212 + ) -> torch.Tensor | IntermediateTensors:

213 input_embeds = self.get_input_embeddings(input_ids)

214 input_embeds = self.input_norm1(input_embeds)

215 hidden_states = self.input_norm2(hidden_states)

@@ -389,7 +388,7 @@ class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

388 def compute_logits(

389 self,

390 hidden_states: torch.Tensor,

392 - ) -> Optional[torch.Tensor]:

391 + ) -> torch.Tensor | None:

392 logits = self.logits_processor(self.lm_head, hidden_states)

393 return logits

394

vllm/model_executor/models/minicpmo.py

+16 -14

@@ -24,8 +24,8 @@

24 # limitations under the License.

25 """Inference-only MiniCPM-O model compatible with HuggingFace weights."""

26

27 -from collections.abc import Iterable, Mapping, Sequence

28 -from typing import Annotated, Any, Callable, Literal, Optional, Union

27 +from collections.abc import Callable, Iterable, Mapping, Sequence

28 +from typing import Annotated, Any, Literal, TypeAlias

29

30 import torch

31 from torch import nn

@@ -89,7 +89,7 @@ class MiniCPMOAudioFeatureInputs(TensorSchema):

89 type: Literal["audio_features"] = "audio_features"

90

91 audio_features: Annotated[

92 - Union[torch.Tensor, list[torch.Tensor]],

92 + torch.Tensor | list[torch.Tensor],

93 TensorShape("bns", "c", "l", dynamic_dims={"l"}),

94 ]

95 """

@@ -99,7 +99,7 @@ class MiniCPMOAudioFeatureInputs(TensorSchema):

99 """

100

101 audio_feature_lens: Annotated[

102 - Union[torch.Tensor, list[torch.Tensor]],

102 + torch.Tensor | list[torch.Tensor],

103 TensorShape("bn", "s"),

104 ]

105 """

@@ -121,12 +121,14 @@ class MiniCPMOAudioEmbeddingInputs(TensorSchema):

121 type: Literal["audio_embeds"] = "audio_embeds"

122

123 audio_embeds: Annotated[

124 - Union[torch.Tensor, list[torch.Tensor]],

124 + torch.Tensor | list[torch.Tensor],

125 TensorShape("bn", "s", "h", dynamic_dims={"s"}),

126 ]

127

128

129 -MiniCPMOAudioInputs = Union[MiniCPMOAudioFeatureInputs, MiniCPMOAudioEmbeddingInputs]

129 +MiniCPMOAudioInputs: TypeAlias = (

130 + MiniCPMOAudioFeatureInputs | MiniCPMOAudioEmbeddingInputs

131 +)

132

133

134 def _minicpmo_field_config(hf_inputs: Mapping[str, torch.Tensor]):

@@ -162,8 +164,8 @@ class MiniCPMOAudioEmbeddingItems(DictEmbeddingItems):

164 class MiniCPMOMultiModalDataParser(MiniCPMVMultiModalDataParser):

165 def _parse_audio_data(

166 self,

165 - data: Union[dict[str, torch.Tensor], ModalityData[AudioItem]],

166 - ) -> Optional[ModalityDataItems[Any, Any]]:

167 + data: dict[str, torch.Tensor] | ModalityData[AudioItem],

168 + ) -> ModalityDataItems[Any, Any] | None:

169 if isinstance(data, dict):

170 return MiniCPMOAudioEmbeddingItems(

171 data,

@@ -176,7 +178,7 @@ class MiniCPMOMultiModalDataParser(MiniCPMVMultiModalDataParser):

178 class MiniCPMOProcessingInfo(MiniCPMVProcessingInfo):

179 audio_pattern = "(<audio>./</audio>)"

180

179 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

181 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

182 return {**super().get_supported_mm_limits(), "audio": None}

183

184 def get_audio_placeholder(

@@ -253,7 +255,7 @@ class MiniCPMODummyInputsBuilder(MiniCPMVDummyInputsBuilder[MiniCPMOProcessingIn

255 self,

256 seq_len: int,

257 mm_counts: Mapping[str, int],

256 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

258 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

259 ) -> MultiModalDataDict:

260 num_audios = mm_counts.get("audio", 0)

261 audio_len = (

@@ -479,7 +481,7 @@ class MiniCPMWhisperEncoder(WhisperEncoder):

481 def forward(

482 self,

483 input_features: torch.Tensor,

482 - attention_mask: Optional[torch.Tensor] = None,

484 + attention_mask: torch.Tensor | None = None,

485 ) -> BaseModelOutputWithPast:

486 # Ignore copy

487 input_features = input_features.to(

@@ -549,7 +551,7 @@ class MiniCPMO(MiniCPMV2_6):

551 }

552

553 @classmethod

552 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

554 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

555 if modality.startswith("image"):

556 return "(<image>./</image>)"

557 if modality.startswith("video"):

@@ -722,7 +724,7 @@ class MiniCPMO(MiniCPMV2_6):

724

725 def _parse_and_validate_audio_input(

726 self, **kwargs: object

725 - ) -> Optional[MiniCPMOAudioInputs]:

727 + ) -> MiniCPMOAudioInputs | None:

728 audio_features = kwargs.pop("audio_features", None)

729 audio_embeds = kwargs.pop("audio_embeds", None)

730

@@ -785,7 +787,7 @@ class MiniCPMO(MiniCPMV2_6):

787 def _process_audio_input(

788 self,

789 audio_input: MiniCPMOAudioInputs,

788 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

790 + ) -> torch.Tensor | list[torch.Tensor]:

791 if audio_input["type"] == "audio_embeds":

792 return audio_input["audio_embeds"]

793

vllm/model_executor/models/minicpmv.py

+37 -37

@@ -26,10 +26,10 @@

26

27 import math

28 from collections import defaultdict

29 -from collections.abc import Iterable, Mapping, Sequence

29 +from collections.abc import Callable, Iterable, Mapping, Sequence

30 from functools import partial

31 from itertools import chain

32 -from typing import Annotated, Any, Callable, Literal, Optional, Union

32 +from typing import Annotated, Any, Literal, TypeAlias

33

34 import numpy as np

35 import torch

@@ -140,12 +140,12 @@ class MiniCPMVImageEmbeddingInputs(TensorSchema):

140

141 type: Literal["image_embeds"]

142 image_embeds: Annotated[

143 - Union[torch.Tensor, list[torch.Tensor]],

143 + torch.Tensor | list[torch.Tensor],

144 TensorShape("bn", "ns", "hs"),

145 ]

146

147

148 -MiniCPMVImageInputs = Union[MiniCPMVImagePixelInputs, MiniCPMVImageEmbeddingInputs]

148 +MiniCPMVImageInputs: TypeAlias = MiniCPMVImagePixelInputs | MiniCPMVImageEmbeddingInputs

149

150 DEFAULT_LN = partial(nn.LayerNorm, eps=1e-6)

151

@@ -156,10 +156,10 @@ class Resampler2_5(BaseResampler):

156 num_queries: int,

157 embed_dim: int,

158 num_heads: int,

159 - kv_dim: Optional[int] = None,

159 + kv_dim: int | None = None,

160 norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,

161 max_size: tuple[int, int] = (70, 70),

162 - quant_config: Optional[QuantizationConfig] = None,

162 + quant_config: QuantizationConfig | None = None,

163 prefix: str = "",

164 ) -> None:

165 super().__init__(

@@ -251,11 +251,11 @@ class Resampler4_5(Resampler2_5):

251 num_queries: int,

252 embed_dim: int,

253 num_heads: int,

254 - kv_dim: Optional[int] = None,

254 + kv_dim: int | None = None,

255 norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,

256 max_size: tuple[int, int] = (70, 70),

257 max_temporal_size: int = 36000,

258 - quant_config: Optional[QuantizationConfig] = None,

258 + quant_config: QuantizationConfig | None = None,

259 prefix: str = "",

260 ) -> None:

261 super().__init__(

@@ -318,7 +318,7 @@ class Resampler4_5(Resampler2_5):

318 self.max_temporal_size = max_temporal_size

319 self._set_temporal_pos_cache(self.max_temporal_size, device)

320

321 - def _init_weights(self, m: Union[nn.Linear, nn.LayerNorm]):

321 + def _init_weights(self, m: nn.Linear | nn.LayerNorm):

322 if isinstance(m, nn.Linear):

323 trunc_normal_(m.weight, std=0.02)

324 if isinstance(m, nn.Linear) and m.bias is not None:

@@ -521,8 +521,8 @@ class MiniCPMVVideoEmbeddingItems(DictEmbeddingItems):

521 class MiniCPMVMultiModalDataParser(MultiModalDataParser):

522 def _parse_image_data(

523 self,

524 - data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],

525 - ) -> Optional[ModalityDataItems[Any, Any]]:

524 + data: dict[str, torch.Tensor] | ModalityData[ImageItem],

525 + ) -> ModalityDataItems[Any, Any] | None:

526 if isinstance(data, dict):

527 return MiniCPMVImageEmbeddingItems(

528 data,

@@ -533,8 +533,8 @@ class MiniCPMVMultiModalDataParser(MultiModalDataParser):

533

534 def _parse_video_data(

535 self,

536 - data: Union[dict[str, torch.Tensor], ModalityData[VideoItem]],

537 - ) -> Optional[ModalityDataItems[Any, Any]]:

536 + data: dict[str, torch.Tensor] | ModalityData[VideoItem],

537 + ) -> ModalityDataItems[Any, Any] | None:

538 if isinstance(data, dict):

539 return MiniCPMVVideoEmbeddingItems(

540 data,

@@ -570,7 +570,7 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo):

570 def get_model_version(self):

571 return get_version_by_config(self.get_hf_config())

572

573 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

573 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

574 mm_limits = {"image": None}

575 if self.get_model_version() in {(2, 6), (4, 0), (4, 5)}:

576 mm_limits["video"] = None

@@ -582,7 +582,7 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo):

582 image_size: ImageSize,

583 # For MiniCPM V/O 2.6

584 image_idx: int = 0,

585 - max_slice_nums: Optional[int] = None,

585 + max_slice_nums: int | None = None,

586 use_image_id: bool = True,

587 ) -> str:

588 image_processor = self.get_image_processor()

@@ -602,8 +602,8 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo):

602 self,

603 image_size: ImageSize,

604 # For MiniCPM V/O 2.6

605 - max_slice_nums: Optional[int] = None,

606 - ) -> Optional[tuple[int, int]]:

605 + max_slice_nums: int | None = None,

606 + ) -> tuple[int, int] | None:

607 image_processor = self.get_image_processor()

608 version = self.get_model_version()

609

@@ -621,7 +621,7 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo):

621 def get_num_image_tokens(

622 self,

623 image_size: ImageSize,

624 - max_slice_nums: Optional[int] = None,

624 + max_slice_nums: int | None = None,

625 ) -> int:

626 image_processor = self.get_image_processor()

627

@@ -712,7 +712,7 @@ class MiniCPMVDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

712 self,

713 seq_len: int,

714 mm_counts: Mapping[str, int],

715 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

715 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

716 ) -> MultiModalDataDict:

717 num_images = mm_counts.get("image", 0)

718 num_videos = mm_counts.get("video", 0)

@@ -1024,7 +1024,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1024 supports_encoder_tp_data = True

1025

1026 @classmethod

1027 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1027 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1028 if modality.startswith("image"):

1029 return "(<image>./</image>)"

1030 if modality.startswith("video"):

@@ -1073,7 +1073,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1073 self,

1074 modality: str,

1075 **kwargs: object,

1076 - ) -> Optional[MiniCPMVImageInputs]:

1076 + ) -> MiniCPMVImageInputs | None:

1077 pixel_values = kwargs.pop("pixel_values", None)

1078 image_embeds = kwargs.pop("image_embeds", None)

1079

@@ -1158,7 +1158,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1158 def _process_vision_input(

1159 self,

1160 image_input: MiniCPMVImageInputs,

1161 - ) -> Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor, ...]]:

1161 + ) -> torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...]:

1162 if image_input["type"] == "image_embeds":

1163 return image_input["image_embeds"]

1164

@@ -1200,8 +1200,8 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1200 self,

1201 input_ids: torch.Tensor,

1202 positions: torch.Tensor,

1203 - intermediate_tensors: Optional[IntermediateTensors] = None,

1204 - inputs_embeds: Optional[torch.Tensor] = None,

1203 + intermediate_tensors: IntermediateTensors | None = None,

1204 + inputs_embeds: torch.Tensor | None = None,

1205 **kwargs: Any,

1206 ) -> torch.Tensor:

1207 if intermediate_tensors is not None:

@@ -1218,7 +1218,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1218 def compute_logits(

1219 self,

1220 hidden_states: torch.Tensor,

1221 - ) -> Optional[torch.Tensor]:

1221 + ) -> torch.Tensor | None:

1222 return self.llm.compute_logits(hidden_states)

1223

1224 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -1243,7 +1243,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1243 def init_vision_module(

1244 self,

1245 config: PretrainedConfig,

1246 - quant_config: Optional[QuantizationConfig],

1246 + quant_config: QuantizationConfig | None,

1247 prefix: str = "",

1248 ) -> nn.Module:

1249 raise NotImplementedError

@@ -1252,7 +1252,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1252 self,

1253 embed_dim: int,

1254 vision_dim: int,

1255 - quant_config: Optional[QuantizationConfig] = None,

1255 + quant_config: QuantizationConfig | None = None,

1256 prefix: str = "",

1257 ) -> nn.Module:

1258 raise NotImplementedError

@@ -1278,7 +1278,7 @@ class MiniCPMV2_0(MiniCPMVBaseModel):

1278 def init_vision_module(

1279 self,

1280 config: PretrainedConfig,

1281 - quant_config: Optional[QuantizationConfig],

1281 + quant_config: QuantizationConfig | None,

1282 prefix: str = "",

1283 ) -> nn.Module:

1284 # TODO: refactor vision model through timm wrapper from transformers

@@ -1313,7 +1313,7 @@ class MiniCPMV2_0(MiniCPMVBaseModel):

1313 self,

1314 embed_dim: int,

1315 vision_dim: int,

1316 - quant_config: Optional[QuantizationConfig] = None,

1316 + quant_config: QuantizationConfig | None = None,

1317 prefix: str = "",

1318 ) -> nn.Module:

1319 with set_default_torch_dtype(torch.float16):

@@ -1381,7 +1381,7 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):

1381 def init_vision_module(

1382 self,

1383 config: PretrainedConfig,

1384 - quant_config: Optional[QuantizationConfig],

1384 + quant_config: QuantizationConfig | None,

1385 prefix: str = "",

1386 ) -> nn.Module:

1387 model = Idefics2VisionTransformer(

@@ -1398,7 +1398,7 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):

1398 self,

1399 embed_dim: int,

1400 vision_dim: int,

1401 - quant_config: Optional[QuantizationConfig] = None,

1401 + quant_config: QuantizationConfig | None = None,

1402 prefix: str = "",

1403 ) -> nn.Module:

1404 with set_default_torch_dtype(torch.float16):

@@ -1474,7 +1474,7 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):

1474 def init_vision_module(

1475 self,

1476 config: PretrainedConfig,

1477 - quant_config: Optional[QuantizationConfig] = None,

1477 + quant_config: QuantizationConfig | None = None,

1478 prefix: str = "",

1479 ) -> nn.Module:

1480 model = Idefics2VisionTransformer(

@@ -1491,7 +1491,7 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):

1491 self,

1492 embed_dim: int,

1493 vision_dim: int,

1494 - quant_config: Optional[QuantizationConfig] = None,

1494 + quant_config: QuantizationConfig | None = None,

1495 prefix: str = "",

1496 ) -> nn.Module:

1497 with set_default_torch_dtype(torch.float16):

@@ -1577,7 +1577,7 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):

1577 def init_vision_module(

1578 self,

1579 config: PretrainedConfig,

1580 - quant_config: Optional[QuantizationConfig] = None,

1580 + quant_config: QuantizationConfig | None = None,

1581 prefix: str = "",

1582 ) -> nn.Module:

1583 quant_config = self._maybe_ignore_quant_config(quant_config)

@@ -1595,7 +1595,7 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):

1595 self,

1596 embed_dim: int,

1597 vision_dim: int,

1598 - quant_config: Optional[QuantizationConfig] = None,

1598 + quant_config: QuantizationConfig | None = None,

1599 prefix: str = "",

1600 ) -> nn.Module:

1601 quant_config = self._maybe_ignore_quant_config(quant_config)

@@ -1682,7 +1682,7 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):

1682 def init_vision_module(

1683 self,

1684 config: PretrainedConfig,

1685 - quant_config: Optional[QuantizationConfig] = None,

1685 + quant_config: QuantizationConfig | None = None,

1686 prefix: str = "",

1687 ) -> nn.Module:

1688 quant_config = self._maybe_ignore_quant_config(quant_config)

@@ -1700,7 +1700,7 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):

1700 self,

1701 embed_dim: int,

1702 vision_dim: int,

1703 - quant_config: Optional[QuantizationConfig] = None,

1703 + quant_config: QuantizationConfig | None = None,

1704 prefix: str = "",

1705 ) -> nn.Module:

1706 quant_config = self._maybe_ignore_quant_config(quant_config)

vllm/model_executor/models/minimax_text_01.py

+18 -18

@@ -4,7 +4,7 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import TYPE_CHECKING, Optional, Union

7 +from typing import TYPE_CHECKING

8

9 if TYPE_CHECKING:

10 pass

@@ -83,7 +83,7 @@ class MiniMaxText01MLP(nn.Module):

83 self,

84 hidden_size: int,

85 intermediate_size: int,

86 - quant_config: Optional[QuantizationConfig] = None,

86 + quant_config: QuantizationConfig | None = None,

87 layer_idx: int = None,

88 prefix: str = "mlp",

89 ) -> None:

@@ -121,9 +121,9 @@ class MiniMaxText01MoE(nn.Module):

121 top_k: int,

122 hidden_size: int,

123 intermediate_size: int,

124 - params_dtype: Optional[torch.dtype] = None,

124 + params_dtype: torch.dtype | None = None,

125 layer_idx: int = None,

126 - quant_config: Optional[QuantizationConfig] = None,

126 + quant_config: QuantizationConfig | None = None,

127 prefix: str = "moe",

128 ) -> None:

129 super().__init__()

@@ -191,10 +191,10 @@ class MiniMaxText01Attention(nn.Module):

191 rotary_dim: int,

192 max_position: int = 4096 * 32,

193 rope_theta: float = 10000,

194 - sliding_window: Optional[int] = None,

195 - quant_config: Optional[QuantizationConfig] = None,

194 + sliding_window: int | None = None,

195 + quant_config: QuantizationConfig | None = None,

196 layer_idx: int = None,

197 - cache_config: Optional[CacheConfig] = None,

197 + cache_config: CacheConfig | None = None,

198 prefix: str = "mha",

199 ) -> None:

200 super().__init__()

@@ -273,12 +273,12 @@ class MiniMaxText01DecoderLayer(nn.Module):

273 def __init__(

274 self,

275 config: MiniMaxConfig,

276 - model_config: Optional[ModelConfig] = None,

277 - cache_config: Optional[CacheConfig] = None,

278 - quant_config: Optional[QuantizationConfig] = None,

276 + model_config: ModelConfig | None = None,

277 + cache_config: CacheConfig | None = None,

278 + quant_config: QuantizationConfig | None = None,

279 expert_num: int = 1,

280 layer_id: int = None,

281 - linear_layer_id: Optional[int] = None,

281 + linear_layer_id: int | None = None,

282 prefix: str = "decoder",

283 ) -> None:

284 self._ilayer = layer_id

@@ -428,7 +428,7 @@ class MiniMaxText01DecoderLayer(nn.Module):

428 hidden_states: torch.Tensor,

429 positions: torch.Tensor,

430 attn_metadata: AttentionMetadata,

431 - residual: Optional[torch.Tensor],

431 + residual: torch.Tensor | None,

432 is_warmup: bool = False,

433 **kwargs,

434 ) -> tuple[torch.Tensor, torch.Tensor]:

@@ -627,12 +627,12 @@ class MiniMaxText01Model(nn.Module):

627

628 def forward(

629 self,

630 - input_ids: Optional[torch.Tensor],

630 + input_ids: torch.Tensor | None,

631 positions: torch.Tensor,

632 - intermediate_tensors: Optional[IntermediateTensors] = None,

633 - inputs_embeds: Optional[torch.Tensor] = None,

632 + intermediate_tensors: IntermediateTensors | None = None,

633 + inputs_embeds: torch.Tensor | None = None,

634 **kwargs,

635 - ) -> Union[torch.Tensor, IntermediateTensors]:

635 + ) -> torch.Tensor | IntermediateTensors:

636 forward_context = get_forward_context()

637 attn_metadata = forward_context.attn_metadata

638

@@ -722,8 +722,8 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):

722 self,

723 input_ids: torch.Tensor,

724 positions: torch.Tensor,

725 - intermediate_tensors: Optional[IntermediateTensors] = None,

726 - inputs_embeds: Optional[torch.Tensor] = None,

725 + intermediate_tensors: IntermediateTensors | None = None,

726 + inputs_embeds: torch.Tensor | None = None,

727 **kwargs,

728 ) -> torch.Tensor:

729 hidden_states = self.model(

vllm/model_executor/models/minimax_vl_01.py

+19 -19

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Mapping

4 -from typing import Annotated, Literal, Optional, Union

4 +from typing import Annotated, Literal, TypeAlias

5

6 import torch

7 import torch.nn as nn

@@ -52,11 +52,11 @@ class MiniMaxVL01ImagePixelInputs(TensorSchema):

52

53 type: Literal["pixel_values"] = "pixel_values"

54 pixel_values: Annotated[

55 - Union[torch.Tensor, list[torch.Tensor]],

55 + torch.Tensor | list[torch.Tensor],

56 TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np", "h", "w"}),

57 ]

58

59 - image_sizes: Annotated[Optional[torch.Tensor], TensorShape("bn", 2)]

59 + image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)]

60 # This should be in `(height, width)` format.

61

62

@@ -72,9 +72,9 @@ class MiniMaxVL01ImageEmbeddingInputs(TensorSchema):

72 data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]

73

74

75 -MiniMaxVL01ImageInputs = Union[

76 - MiniMaxVL01ImagePixelInputs, MiniMaxVL01ImageEmbeddingInputs

77 -]

75 +MiniMaxVL01ImageInputs: TypeAlias = (

76 + MiniMaxVL01ImagePixelInputs | MiniMaxVL01ImageEmbeddingInputs

77 +)

78

79

80 class MiniMaxVL01MultiModalProjector(nn.Module):

@@ -84,7 +84,7 @@ class MiniMaxVL01MultiModalProjector(nn.Module):

84 text_hidden_size: int,

85 projector_hidden_act: str,

86 multimodal_projector_bias: bool,

87 - quant_config: Optional[QuantizationConfig] = None,

87 + quant_config: QuantizationConfig | None = None,

88 prefix: str = "",

89 ):

90 super().__init__()

@@ -127,7 +127,7 @@ class MiniMaxVL01ProcessingInfo(LlavaNextProcessingInfo):

127

128 return hf_processor

129

130 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

130 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

131 return {"image": None}

132

133

@@ -187,7 +187,7 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, Support

187 }

188

189 @classmethod

190 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

190 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

191 if modality.startswith("image"):

192 return "<image>"

193

@@ -239,9 +239,9 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, Support

239

240 def _image_pixels_to_features(

241 self,

242 - vision_tower: Union[CLIPVisionModel, SiglipVisionModel, PixtralHFVisionModel],

243 - pixel_values: Union[torch.Tensor, list[torch.Tensor]],

244 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

242 + vision_tower: CLIPVisionModel | SiglipVisionModel | PixtralHFVisionModel,

243 + pixel_values: torch.Tensor | list[torch.Tensor],

244 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

245 # NOTE: we skip the step to select the vision feature layer since

246 # this is already done inside the vision tower

247 feature_select_strategy = self.config.vision_feature_select_strategy

@@ -302,7 +302,7 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, Support

302 def _process_image_pixels(

303 self,

304 inputs: MiniMaxVL01ImagePixelInputs,

305 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

305 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

306 assert self.vision_tower is not None

307

308 pixel_values = inputs["pixel_values"]

@@ -311,7 +311,7 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, Support

311 def _process_image_input(

312 self,

313 image_input: MiniMaxVL01ImageInputs,

314 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

314 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

315 if image_input["type"] == "image_embeds":

316 return image_input["data"]

317

@@ -330,7 +330,7 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, Support

330

331 def _parse_and_validate_image_input(

332 self, **kwargs: object

333 - ) -> Optional[MiniMaxVL01ImageInputs]:

333 + ) -> MiniMaxVL01ImageInputs | None:

334 pixel_values = kwargs.pop("pixel_values", None)

335 image_sizes = kwargs.pop("image_sizes", None)

336 image_embeds = kwargs.pop("image_embeds", None)

@@ -364,10 +364,10 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, Support

364 self,

365 input_ids: torch.Tensor,

366 positions: torch.Tensor,

367 - intermediate_tensors: Optional[IntermediateTensors] = None,

368 - inputs_embeds: Optional[torch.Tensor] = None,

367 + intermediate_tensors: IntermediateTensors | None = None,

368 + inputs_embeds: torch.Tensor | None = None,

369 **kwargs: object,

370 - ) -> Union[torch.Tensor, IntermediateTensors]:

370 + ) -> torch.Tensor | IntermediateTensors:

371 if intermediate_tensors is not None:

372 inputs_embeds = None

373 elif inputs_embeds is None:

@@ -388,7 +388,7 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, Support

388 def compute_logits(

389 self,

390 hidden_states: torch.Tensor,

391 - ) -> Optional[torch.Tensor]:

391 + ) -> torch.Tensor | None:

392 return self.language_model.compute_logits(hidden_states)

393

394 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/mistral3.py

+16 -16

@@ -3,7 +3,7 @@

3

4 from abc import abstractmethod

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Annotated, Final, Literal, Optional, Protocol, TypeVar, Union

6 +from typing import Annotated, Final, Literal, Protocol, TypeVar

7

8 import torch

9 import torch.nn as nn

@@ -72,7 +72,7 @@ class Mistral3ImagePixelInputs(TensorSchema):

72 # Note that `height` or `width` may be different per batch and image,

73 # in which case the data is passed as a list instead of a batched tensor.

74 pixel_values: Annotated[

75 - Union[torch.Tensor, list[torch.Tensor]],

75 + torch.Tensor | list[torch.Tensor],

76 TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}),

77 ]

78

@@ -136,7 +136,7 @@ class Mistral3MultiModalProjector(nn.Module):

136 patch_size: int,

137 projector_hidden_act: str,

138 multimodal_projector_bias: bool,

139 - quant_config: Optional[QuantizationConfig] = None,

139 + quant_config: QuantizationConfig | None = None,

140 prefix: str = "",

141 ):

142 super().__init__()

@@ -179,7 +179,7 @@ class LlavaLikeConfig(Protocol):

179 vision_config: Final[PretrainedConfig]

180 image_token_index: Final[int]

181 vision_feature_select_strategy: Final[str]

182 - vision_feature_layer: Final[Union[int, list[int]]]

182 + vision_feature_layer: Final[int | list[int]]

183

184

185 class LlavaLikeProcessor(Protocol):

@@ -197,7 +197,7 @@ class BaseLlavaProcessingInfo(BaseProcessingInfo):

197 def get_hf_processor(self, **kwargs: object) -> LlavaLikeProcessor:

198 raise NotImplementedError

199

200 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

200 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

201 return {"image": None}

202

203 def get_num_image_tokens(

@@ -234,7 +234,7 @@ class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder[_I]):

234 self,

235 seq_len: int,

236 mm_counts: Mapping[str, int],

237 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

237 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

238 ) -> MultiModalDataDict:

239 num_images = mm_counts.get("image", 0)

240

@@ -348,7 +348,7 @@ def _build_mistral3_processor(

348 info: _I,

349 dummy_inputs: BaseDummyInputsBuilder[_I],

350 *,

351 - cache: Optional[BaseMultiModalProcessorCache] = None,

351 + cache: BaseMultiModalProcessorCache | None = None,

352 ) -> BaseMultiModalProcessor:

353 assert isinstance(info, Mistral3ProcessingInfo)

354 return Mistral3MultiModalProcessor(

@@ -394,9 +394,9 @@ def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:

394

395 def init_vision_tower_for_llava(

396 hf_config: LlavaLikeConfig,

397 - quant_config: Optional[QuantizationConfig],

397 + quant_config: QuantizationConfig | None,

398 *,

399 - require_post_norm: Optional[bool] = None,

399 + require_post_norm: bool | None = None,

400 prefix: str = "",

401 ) -> PixtralHFVisionModel:

402 vision_config = hf_config.vision_config

@@ -441,7 +441,7 @@ class Mistral3ForConditionalGeneration(

441 )

442

443 @classmethod

444 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

444 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

445 if modality.startswith("image"):

446 return None

447

@@ -504,7 +504,7 @@ class Mistral3ForConditionalGeneration(

504

505 def _parse_and_validate_image_input(

506 self, **kwargs: object

507 - ) -> Optional[Mistral3ImagePixelInputs]:

507 + ) -> Mistral3ImagePixelInputs | None:

508 pixel_values = kwargs.pop("pixel_values", None)

509 image_embeds = kwargs.pop("image_embeds", None)

510

@@ -519,7 +519,7 @@ class Mistral3ForConditionalGeneration(

519 def _process_image_input(

520 self,

521 image_input: Mistral3ImagePixelInputs,

522 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

522 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

523 if image_input["type"] == "image_embeds":

524 return image_input["data"]

525

@@ -562,10 +562,10 @@ class Mistral3ForConditionalGeneration(

562 self,

563 input_ids: torch.Tensor,

564 positions: torch.Tensor,

565 - intermediate_tensors: Optional[IntermediateTensors] = None,

566 - inputs_embeds: Optional[torch.Tensor] = None,

565 + intermediate_tensors: IntermediateTensors | None = None,

566 + inputs_embeds: torch.Tensor | None = None,

567 **kwargs: object,

568 - ) -> Union[torch.Tensor, IntermediateTensors]:

568 + ) -> torch.Tensor | IntermediateTensors:

569 """Run forward pass for Mistral3.

570

571 One key thing to understand is the `input_ids` already accounts for the

@@ -615,7 +615,7 @@ class Mistral3ForConditionalGeneration(

615 def compute_logits(

616 self,

617 hidden_states: torch.Tensor,

618 - ) -> Optional[torch.Tensor]:

618 + ) -> torch.Tensor | None:

619 return self.language_model.compute_logits(hidden_states)

620

621 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/mixtral.py

+16 -17

@@ -27,7 +27,6 @@

27 import typing

28 from collections.abc import Callable, Iterable

29 from itertools import islice

30 -from typing import Optional, Union

30

31 import torch

32 from torch import nn

@@ -88,10 +87,10 @@ class MixtralMoE(nn.Module):

87 top_k: int,

88 hidden_size: int,

89 intermediate_size: int,

91 - params_dtype: Optional[torch.dtype] = None,

92 - quant_config: Optional[QuantizationConfig] = None,

93 - tp_size: Optional[int] = None,

94 - dp_size: Optional[int] = None,

90 + params_dtype: torch.dtype | None = None,

91 + quant_config: QuantizationConfig | None = None,

92 + tp_size: int | None = None,

93 + dp_size: int | None = None,

94 prefix: str = "",

95 enable_eplb: bool = False,

96 ):

@@ -163,8 +162,8 @@ class MixtralAttention(nn.Module):

162 num_kv_heads: int,

163 max_position: int = 4096 * 32,

164 rope_theta: float = 10000,

166 - cache_config: Optional[CacheConfig] = None,

167 - quant_config: Optional[QuantizationConfig] = None,

165 + cache_config: CacheConfig | None = None,

166 + quant_config: QuantizationConfig | None = None,

167 prefix: str = "",

168 ) -> None:

169 super().__init__()

@@ -242,8 +241,8 @@ class MixtralDecoderLayer(nn.Module):

241 def __init__(

242 self,

243 config: MixtralConfig,

245 - cache_config: Optional[CacheConfig] = None,

246 - quant_config: Optional[QuantizationConfig] = None,

244 + cache_config: CacheConfig | None = None,

245 + quant_config: QuantizationConfig | None = None,

246 prefix: str = "",

247 enable_eplb: bool = False,

248 ) -> None:

@@ -280,7 +279,7 @@ class MixtralDecoderLayer(nn.Module):

279 self,

280 positions: torch.Tensor,

281 hidden_states: torch.Tensor,

283 - residual: Optional[torch.Tensor],

282 + residual: torch.Tensor | None,

283 ) -> torch.Tensor:

284 # Self Attention

285 if residual is None:

@@ -353,9 +352,9 @@ class MixtralModel(nn.Module):

352 self,

353 input_ids: torch.Tensor,

354 positions: torch.Tensor,

356 - intermediate_tensors: Optional[IntermediateTensors],

357 - inputs_embeds: Optional[torch.Tensor] = None,

358 - ) -> Union[torch.Tensor, IntermediateTensors]:

355 + intermediate_tensors: IntermediateTensors | None,

356 + inputs_embeds: torch.Tensor | None = None,

357 + ) -> torch.Tensor | IntermediateTensors:

358 if get_pp_group().is_first_rank:

359 if inputs_embeds is not None:

360 hidden_states = inputs_embeds

@@ -615,9 +614,9 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):

614 self,

615 input_ids: torch.Tensor,

616 positions: torch.Tensor,

618 - intermediate_tensors: Optional[IntermediateTensors] = None,

619 - inputs_embeds: Optional[torch.Tensor] = None,

620 - ) -> Union[torch.Tensor, IntermediateTensors]:

617 + intermediate_tensors: IntermediateTensors | None = None,

618 + inputs_embeds: torch.Tensor | None = None,

619 + ) -> torch.Tensor | IntermediateTensors:

620 hidden_states = self.model(

621 input_ids, positions, intermediate_tensors, inputs_embeds

622 )

@@ -626,7 +625,7 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):

625 def compute_logits(

626 self,

627 hidden_states: torch.Tensor,

629 - ) -> Optional[torch.Tensor]:

628 + ) -> torch.Tensor | None:

629 logits = self.logits_processor(self.lm_head, hidden_states)

630 return logits

631

vllm/model_executor/models/mllama4.py

+17 -17

@@ -19,7 +19,7 @@

19 import math

20 from collections.abc import Iterable, Mapping

21 from itertools import tee

22 -from typing import Annotated, Literal, Optional, Union

22 +from typing import Annotated, Literal

23

24 import torch

25 from torch import nn

@@ -115,7 +115,7 @@ class Llama4VisionMLP(nn.Module):

115 output_size: int,

116 bias: bool,

117 output_activation: bool,

118 - quant_config: Optional[QuantizationConfig] = None,

118 + quant_config: QuantizationConfig | None = None,

119 prefix: str = "",

120 use_data_parallel: bool = False,

121 ):

@@ -152,7 +152,7 @@ class Llama4MultiModalProjector(nn.Module):

152 def __init__(

153 self,

154 config,

155 - quant_config: Optional[QuantizationConfig] = None,

155 + quant_config: QuantizationConfig | None = None,

156 prefix: str = "",

157 ):

158 super().__init__()

@@ -199,7 +199,7 @@ class Llama4VisionPixelShuffleMLP(nn.Module):

199 def __init__(

200 self,

201 config,

202 - quant_config: Optional[QuantizationConfig] = None,

202 + quant_config: QuantizationConfig | None = None,

203 prefix: str = "",

204 use_data_parallel: bool = False,

205 ):

@@ -229,7 +229,7 @@ class Llama4VisionAttention(nn.Module):

229 def __init__(

230 self,

231 config: Llama4VisionConfig,

232 - quant_config: Optional[QuantizationConfig],

232 + quant_config: QuantizationConfig | None,

233 prefix: str = "",

234 use_data_parallel: bool = False,

235 ):

@@ -323,7 +323,7 @@ class Llama4VisionEncoderLayer(nn.Module):

323 def __init__(

324 self,

325 config: Llama4VisionConfig,

326 - quant_config: Optional[QuantizationConfig],

326 + quant_config: QuantizationConfig | None,

327 prefix: str = "",

328 use_data_parallel: bool = False,

329 ):

@@ -376,7 +376,7 @@ class Llama4VisionEncoder(nn.Module):

376 def __init__(

377 self,

378 config: Llama4VisionConfig,

379 - quant_config: Optional[QuantizationConfig],

379 + quant_config: QuantizationConfig | None,

380 prefix: str = "",

381 use_data_parallel: bool = False,

382 ):

@@ -419,7 +419,7 @@ class Llama4UnfoldConvolution(nn.Module):

419 def __init__(

420 self,

421 config: Llama4VisionConfig,

422 - quant_config: Optional[QuantizationConfig] = None,

422 + quant_config: QuantizationConfig | None = None,

423 prefix: str = "",

424 use_data_parallel: bool = False,

425 ):

@@ -449,7 +449,7 @@ class Llama4VisionModel(nn.Module):

449 def __init__(

450 self,

451 config: Llama4VisionConfig,

452 - quant_config: Optional[QuantizationConfig] = None,

452 + quant_config: QuantizationConfig | None = None,

453 prefix: str = "",

454 use_data_parallel: bool = False,

455 ):

@@ -547,7 +547,7 @@ class Mllama4ProcessingInfo(BaseProcessingInfo):

547 Llama4Processor, use_fast=kwargs.pop("use_fast", True), **kwargs

548 )

549

550 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

550 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

551 # Although vLLM can support more images from an infra capability

552 # perspective, we do not recommend using >10 images in practice.

553 return {"image": None}

@@ -699,7 +699,7 @@ class Mllama4DummyInputsBuilder(BaseDummyInputsBuilder[Mllama4ProcessingInfo]):

699 self,

700 seq_len: int,

701 mm_counts: Mapping[str, int],

702 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

702 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

703 ) -> MultiModalDataDict:

704 num_images = mm_counts.get("image", 0)

705

@@ -733,7 +733,7 @@ class Llama4ForConditionalGeneration(

733 supports_encoder_tp_data = True

734

735 @classmethod

736 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

736 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

737 if modality.startswith("image"):

738 return "<|image|>"

739

@@ -792,7 +792,7 @@ class Llama4ForConditionalGeneration(

792

793 def _parse_and_validate_image_input(

794 self, **kwargs: object

795 - ) -> Optional[Llama4ImagePatchInputs]:

795 + ) -> Llama4ImagePatchInputs | None:

796 # num_images, 1, num_chunks, channel, image_size, image_size

797 pixel_values = kwargs.pop("pixel_values", None)

798 if pixel_values is None:

@@ -849,10 +849,10 @@ class Llama4ForConditionalGeneration(

849 self,

850 input_ids: torch.Tensor,

851 positions: torch.Tensor,

852 - intermediate_tensors: Optional[IntermediateTensors] = None,

853 - inputs_embeds: Optional[torch.Tensor] = None,

852 + intermediate_tensors: IntermediateTensors | None = None,

853 + inputs_embeds: torch.Tensor | None = None,

854 **kwargs: object,

855 - ) -> Union[torch.Tensor, IntermediateTensors]:

855 + ) -> torch.Tensor | IntermediateTensors:

856 if intermediate_tensors is not None:

857 inputs_embeds = None

858

@@ -863,7 +863,7 @@ class Llama4ForConditionalGeneration(

863 def compute_logits(

864 self,

865 hidden_states: torch.Tensor,

866 - ) -> Optional[torch.Tensor]:

866 + ) -> torch.Tensor | None:

867 return self.language_model.compute_logits(hidden_states)

868

869 def separate_weights(

vllm/model_executor/models/modernbert.py

+13 -14

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Set

4 -from typing import Optional, Union

4

5 import torch

6 from torch import nn

@@ -50,7 +49,7 @@ class ModernBertEmbeddings(nn.Module):

49 def forward(

50 self,

51 input_ids: torch.Tensor,

53 - inputs_embeds: Optional[torch.Tensor] = None,

52 + inputs_embeds: torch.Tensor | None = None,

53 ) -> torch.Tensor:

54 if inputs_embeds is not None:

55 return self.norm(inputs_embeds)

@@ -74,7 +73,7 @@ class ModernBertRotaryEmbedding(RotaryEmbedding):

73

74

75 class ModernBertAttention(nn.Module):

77 - def __init__(self, config: ModernBertConfig, layer_id: Optional[int] = None):

76 + def __init__(self, config: ModernBertConfig, layer_id: int | None = None):

77 super().__init__()

78 self.config = config

79 self.hidden_size = config.hidden_size

@@ -151,7 +150,7 @@ class ModernBertMLP(nn.Module):

150

151 class ModernBertLayer(nn.Module):

152 def __init__(

154 - self, config: ModernBertConfig, prefix: str = "", layer_id: Optional[int] = None

153 + self, config: ModernBertConfig, prefix: str = "", layer_id: int | None = None

154 ):

155 super().__init__()

156 self.config = config

@@ -243,8 +242,8 @@ class ModernBertModel(nn.Module):

242 self,

243 input_ids: torch.Tensor,

244 positions: torch.Tensor,

246 - intermediate_tensors: Optional[IntermediateTensors] = None,

247 - inputs_embeds: Optional[torch.Tensor] = None,

245 + intermediate_tensors: IntermediateTensors | None = None,

246 + inputs_embeds: torch.Tensor | None = None,

247 ) -> torch.Tensor:

248 if inputs_embeds is not None:

249 hidden_states = inputs_embeds

@@ -287,9 +286,9 @@ class ModernBertPooler(Pooler):

286

287 def forward(

288 self,

290 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],

289 + hidden_states: torch.Tensor | list[torch.Tensor],

290 pooling_metadata: PoolingMetadata,

292 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

291 + ) -> torch.Tensor | list[torch.Tensor]:

292 pooled_output = self.pooling(hidden_states, pooling_metadata)

293

294 if isinstance(pooled_output, list):

@@ -370,10 +369,10 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):

369

370 def forward(

371 self,

373 - input_ids: Optional[torch.LongTensor],

372 + input_ids: torch.LongTensor | None,

373 positions: torch.Tensor,

375 - intermediate_tensors: Optional[IntermediateTensors] = None,

376 - inputs_embeds: Optional[torch.Tensor] = None,

374 + intermediate_tensors: IntermediateTensors | None = None,

375 + inputs_embeds: torch.Tensor | None = None,

376 ) -> torch.Tensor:

377 return self.model(

378 input_ids=input_ids,

@@ -436,10 +435,10 @@ class ModernBertForTokenClassification(nn.Module):

435

436 def forward(

437 self,

439 - input_ids: Optional[torch.Tensor],

438 + input_ids: torch.Tensor | None,

439 positions: torch.Tensor,

441 - intermediate_tensors: Optional[IntermediateTensors] = None,

442 - inputs_embeds: Optional[torch.Tensor] = None,

440 + intermediate_tensors: IntermediateTensors | None = None,

441 + inputs_embeds: torch.Tensor | None = None,

442 ) -> torch.Tensor:

443 hidden_states = self.model(

444 input_ids=input_ids,

vllm/model_executor/models/module_mapping.py

+4 -5

@@ -5,7 +5,6 @@

5 # https://github.com/modelscope/ms-swift/blob/v2.4.2/swift/utils/module_mapping.py

6

7 from dataclasses import dataclass, field

8 -from typing import Union

8

9

10 @dataclass

@@ -55,10 +54,10 @@ class MultiModelKeys(ModelKeys):

54

55 @staticmethod

56 def from_string_field(

58 - language_model: Union[str, list[str]] = None,

59 - connector: Union[str, list[str]] = None,

60 - tower_model: Union[str, list[str]] = None,

61 - generator: Union[str, list[str]] = None,

57 + language_model: str | list[str] = None,

58 + connector: str | list[str] = None,

59 + tower_model: str | list[str] = None,

60 + generator: str | list[str] = None,

61 **kwargs,

62 ) -> "MultiModelKeys":

63 def to_list(value):

vllm/model_executor/models/molmo.py

+40 -40

@@ -6,7 +6,7 @@ from collections.abc import Iterable, Mapping, Sequence

6 from dataclasses import dataclass

7 from functools import cached_property, partial

8 from itertools import islice

9 -from typing import Annotated, Optional, Union

9 +from typing import Annotated

10

11 import numpy as np

12 import torch

@@ -104,18 +104,18 @@ class MolmoImageInputs(TensorSchema):

104 """

105

106 images: Annotated[

107 - Union[torch.Tensor, list[torch.Tensor]],

107 + torch.Tensor | list[torch.Tensor],

108 TensorShape("bn", "nc", "np", "pd", dynamic_dims={"nc"}),

109 ]

110 # Number of crops may vary per batch and image, so pass it as a list.

111

112 image_masks: Annotated[

113 - Optional[Union[torch.Tensor, list[torch.Tensor]]],

113 + torch.Tensor | list[torch.Tensor] | None,

114 TensorShape("bn", "nc", "np", dynamic_dims={"nc"}),

115 ]

116

117 image_input_idx: Annotated[

118 - Union[torch.Tensor, list[torch.Tensor]],

118 + torch.Tensor | list[torch.Tensor],

119 TensorShape("bn", "nc", "tp", dynamic_dims={"nc"}),

120 ]

121 # An index tensor that maps image features to their corresponding patch tokens.

@@ -151,7 +151,7 @@ class ViTMLP(nn.Module):

151 def __init__(

152 self,

153 config: VisionBackboneConfig,

154 - quant_config: Optional[QuantizationConfig] = None,

154 + quant_config: QuantizationConfig | None = None,

155 ):

156 super().__init__()

157 self.w1 = ColumnParallelLinear(

@@ -185,7 +185,7 @@ class MultiHeadDotProductAttention(nn.Module):

185 config: VisionBackboneConfig,

186 use_bias: bool = True,

187 nlayers: int = 1,

188 - quant_config: Optional[QuantizationConfig] = None,

188 + quant_config: QuantizationConfig | None = None,

189 ):

190 super().__init__()

191

@@ -238,7 +238,7 @@ class MultiHeadDotProductAttention(nn.Module):

238 )

239

240 def forward(

241 - self, inputs_q: torch.Tensor, inputs_kv: Optional[torch.Tensor] = None

241 + self, inputs_q: torch.Tensor, inputs_kv: torch.Tensor | None = None

242 ) -> torch.Tensor:

243 if inputs_kv is not None:

244 inputs_k = inputs_kv

@@ -263,7 +263,7 @@ class ResidualAttentionBlock(nn.Module):

263 def __init__(

264 self,

265 config: VisionBackboneConfig,

266 - quant_config: Optional[QuantizationConfig] = None,

266 + quant_config: QuantizationConfig | None = None,

267 ):

268 super().__init__()

269 self.attention = MultiHeadDotProductAttention(config, quant_config=quant_config)

@@ -289,7 +289,7 @@ class BlockCollection(nn.Module):

289 def __init__(

290 self,

291 config: VisionBackboneConfig,

292 - quant_config: Optional[QuantizationConfig] = None,

292 + quant_config: QuantizationConfig | None = None,

293 ):

294 super().__init__()

295 self.resblocks = nn.ModuleList(

@@ -317,7 +317,7 @@ class VisionTransformer(nn.Module):

317 def __init__(

318 self,

319 config: VisionBackboneConfig,

320 - quant_config: Optional[QuantizationConfig] = None,

320 + quant_config: QuantizationConfig | None = None,

321 ):

322 super().__init__()

323 scale = config.image_emb_dim**-0.5

@@ -367,7 +367,7 @@ class VisionTransformer(nn.Module):

367 return x

368

369 def forward(

370 - self, x: torch.Tensor, patch_num: Optional[int] = None

370 + self, x: torch.Tensor, patch_num: int | None = None

371 ) -> list[torch.Tensor]:

372 """

373 : param x: (batch_size, num_patch, n_pixels)

@@ -396,8 +396,8 @@ class MolmoAttention(nn.Module):

396 def __init__(

397 self,

398 config: PretrainedConfig,

399 - cache_config: Optional[CacheConfig] = None,

400 - quant_config: Optional[QuantizationConfig] = None,

399 + cache_config: CacheConfig | None = None,

400 + quant_config: QuantizationConfig | None = None,

401 prefix: str = "",

402 ) -> None:

403 super().__init__()

@@ -432,9 +432,9 @@ class MolmoAttention(nn.Module):

432 quant_config=quant_config,

433 )

434

435 - self.tp_rank: Optional[int] = None

436 - self.k_norm: Optional[nn.Module] = None

437 - self.q_norm: Optional[nn.Module] = None

435 + self.tp_rank: int | None = None

436 + self.k_norm: nn.Module | None = None

437 + self.q_norm: nn.Module | None = None

438 if config.attention_layer_norm:

439 self.tp_rank = get_tensor_model_parallel_rank()

440 self.k_norm = RMSNorm(

@@ -503,8 +503,8 @@ class LanguageModelMLP(nn.Module):

503 def __init__(

504 self,

505 config: PretrainedConfig,

506 - input_dim: Optional[int] = None,

507 - quant_config: Optional[QuantizationConfig] = None,

506 + input_dim: int | None = None,

507 + quant_config: QuantizationConfig | None = None,

508 ) -> None:

509 super().__init__()

510 self.hidden_size = config.hidden_size

@@ -542,8 +542,8 @@ class ImageProjectorMLP(nn.Module):

542 def __init__(

543 self,

544 config: PretrainedConfig,

545 - input_dim: Optional[int] = None,

546 - quant_config: Optional[QuantizationConfig] = None,

545 + input_dim: int | None = None,

546 + quant_config: QuantizationConfig | None = None,

547 ) -> None:

548 super().__init__()

549 self.hidden_size = config.hidden_size

@@ -580,8 +580,8 @@ class MolmoDecoderLayer(nn.Module):

580 def __init__(

581 self,

582 config: PretrainedConfig,

583 - cache_config: Optional[CacheConfig] = None,

584 - quant_config: Optional[QuantizationConfig] = None,

583 + cache_config: CacheConfig | None = None,

584 + quant_config: QuantizationConfig | None = None,

585 prefix: str = "",

586 ) -> None:

587 super().__init__()

@@ -604,8 +604,8 @@ class MolmoDecoderLayer(nn.Module):

604 self,

605 positions: torch.Tensor,

606 hidden_states: torch.Tensor,

607 - residual: Optional[torch.Tensor],

608 - ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:

607 + residual: torch.Tensor | None,

608 + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:

609 # Self Attention

610 if residual is None:

611 residual = hidden_states

@@ -627,8 +627,8 @@ class MolmoDecoderNormAfterLayer(MolmoDecoderLayer):

627 self,

628 positions: torch.Tensor,

629 hidden_states: torch.Tensor,

630 - residual: Optional[torch.Tensor],

631 - ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:

630 + residual: torch.Tensor | None,

631 + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:

632 # Self Attention

633 residual = hidden_states

634 hidden_states = self.self_attn(

@@ -654,7 +654,7 @@ class MolmoVisionBackbone(nn.Module, SupportsQuant):

654 self,

655 config: PretrainedConfig,

656 vision_config: VisionBackboneConfig,

657 - quant_config: Optional[QuantizationConfig] = None,

657 + quant_config: QuantizationConfig | None = None,

658 ) -> None:

659 super().__init__()

660 self.vit_layers = VIT_LAYERS

@@ -849,8 +849,8 @@ class MolmoModel(nn.Module, SupportsQuant):

849 self,

850 input_ids: torch.Tensor,

851 positions: torch.Tensor,

852 - intermediate_tensors: Optional[IntermediateTensors] = None,

853 - inputs_embeds: Optional[torch.Tensor] = None,

852 + intermediate_tensors: IntermediateTensors | None = None,

853 + inputs_embeds: torch.Tensor | None = None,

854 ) -> torch.Tensor:

855 if get_pp_group().is_first_rank:

856 if inputs_embeds is not None:

@@ -1064,7 +1064,7 @@ class MolmoProcessorWrapper:

1064 return image_token_length_h

1065

1066 @property

1067 - def message_format(self) -> Optional[str]:

1067 + def message_format(self) -> str | None:

1068 return "role"

1069

1070 @property

@@ -1145,9 +1145,9 @@ class MolmoProcessorWrapper:

1145

1146 def __call__(

1147 self,

1148 - text: Optional[Union[TextInput, list[TextInput]]] = None,

1149 - images: Optional[Union[ImageInput, list[ImageInput]]] = None,

1150 - return_tensors: Optional[Union[str, TensorType]] = None,

1148 + text: TextInput | list[TextInput] | None = None,

1149 + images: ImageInput | list[ImageInput] | None = None,

1150 + return_tensors: str | TensorType | None = None,

1151 **kwargs,

1152 ) -> BatchFeature:

1153 outputs = self.processor.process( # type: ignore

@@ -1189,7 +1189,7 @@ class MolmoProcessingInfo(BaseProcessingInfo):

1189 processor = self.ctx.get_hf_processor(**kwargs)

1190 return MolmoProcessorWrapper(processor)

1191

1192 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

1192 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

1193 return {"image": None}

1194

1195 def get_num_image_tokens(

@@ -1197,7 +1197,7 @@ class MolmoProcessingInfo(BaseProcessingInfo):

1197 *,

1198 image_width: int,

1199 image_height: int,

1200 - processor: Optional[MolmoProcessorWrapper],

1200 + processor: MolmoProcessorWrapper | None,

1201 ) -> int:

1202 if processor is None:

1203 processor = self.get_hf_processor()

@@ -1250,7 +1250,7 @@ class MolmoDummyInputsBuilder(BaseDummyInputsBuilder[MolmoProcessingInfo]):

1250 self,

1251 seq_len: int,

1252 mm_counts: Mapping[str, int],

1253 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

1253 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

1254 ) -> MultiModalDataDict:

1255 target_width, target_height = self.info.get_image_size_with_most_features()

1256 num_images = mm_counts.get("image", 0)

@@ -1398,7 +1398,7 @@ class MolmoForCausalLM(

1398 }

1399

1400 @classmethod

1401 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1401 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1402 if modality.startswith("image"):

1403 return None

1404

@@ -1442,7 +1442,7 @@ class MolmoForCausalLM(

1442 def _parse_and_validate_image_input(

1443 self,

1444 **kwargs: object,

1445 - ) -> Optional[MolmoImageInputs]:

1445 + ) -> MolmoImageInputs | None:

1446 images = kwargs.pop("images", None)

1447 image_masks = kwargs.pop("image_masks", None)

1448 image_input_idx = kwargs.pop("image_input_idx", None)

@@ -1522,8 +1522,8 @@ class MolmoForCausalLM(

1522 self,

1523 input_ids: torch.LongTensor,

1524 positions: torch.LongTensor,

1525 - intermediate_tensors: Optional[IntermediateTensors] = None,

1526 - inputs_embeds: Optional[torch.Tensor] = None,

1525 + intermediate_tensors: IntermediateTensors | None = None,

1526 + inputs_embeds: torch.Tensor | None = None,

1527 **kwargs: object,

1528 ) -> torch.Tensor:

1529 if intermediate_tensors is not None:

vllm/model_executor/models/moonvit.py

+7 -8

@@ -45,7 +45,6 @@

45 from collections.abc import Sequence

46 from copy import deepcopy

47 from functools import cached_property

48 -from typing import Optional, Union

48

49 import torch

50 import torch.nn as nn

@@ -68,8 +67,8 @@ def multihead_attention(

67 q: torch.Tensor,

68 k: torch.Tensor,

69 v: torch.Tensor,

71 - q_cu_seqlens: Optional[torch.Tensor] = None,

72 - k_cu_seqlens: Optional[torch.Tensor] = None,

70 + q_cu_seqlens: torch.Tensor | None = None,

71 + k_cu_seqlens: torch.Tensor | None = None,

72 ) -> torch.Tensor:

73 """Multi-head attention using flash attention 2.

74

@@ -121,8 +120,8 @@ def sdpa_attention(

120 q: torch.Tensor,

121 k: torch.Tensor,

122 v: torch.Tensor,

124 - q_cu_seqlens: Optional[torch.Tensor] = None,

125 - k_cu_seqlens: Optional[torch.Tensor] = None,

123 + q_cu_seqlens: torch.Tensor | None = None,

124 + k_cu_seqlens: torch.Tensor | None = None,

125 ) -> torch.Tensor:

126 """SDPA attention.

127

@@ -230,7 +229,7 @@ class MoonVisionPatchEmbed(nn.Module):

229 self,

230 out_dim: int,

231 in_dim: int = 3,

233 - patch_size: Union[int, tuple[int, int]] = (14, 14),

232 + patch_size: int | tuple[int, int] = (14, 14),

233 pos_emb_height: int = 14,

234 pos_emb_width: int = 14,

235 ):

@@ -460,7 +459,7 @@ class MoonVitEncoderLayer(nn.Module):

459 self,

460 x: torch.Tensor,

461 cu_seqlens: torch.Tensor,

463 - rope_freqs_cis: Optional[torch.Tensor] = None,

462 + rope_freqs_cis: torch.Tensor | None = None,

463 ):

464 """

465 Args:

@@ -491,7 +490,7 @@ class MoonVitEncoderLayer(nn.Module):

490 self,

491 hidden_states: torch.Tensor,

492 cu_seqlens: torch.Tensor,

494 - rope_freqs_cis: Union[torch.Tensor, None] = None,

493 + rope_freqs_cis: torch.Tensor | None = None,

494 ) -> torch.Tensor:

495 """

496 Args:

vllm/model_executor/models/mpt.py

+12 -13

@@ -5,7 +5,6 @@

5 import math

6 from collections.abc import Iterable

7 from itertools import islice

8 -from typing import Optional, Union

8

9 import torch

10 import torch.nn as nn

@@ -58,8 +57,8 @@ class MPTAttention(nn.Module):

57 def __init__(

58 self,

59 config: MptConfig,

61 - cache_config: Optional[CacheConfig] = None,

62 - quant_config: Optional[QuantizationConfig] = None,

60 + cache_config: CacheConfig | None = None,

61 + quant_config: QuantizationConfig | None = None,

62 prefix: str = "",

63 ):

64 super().__init__()

@@ -152,7 +151,7 @@ class MPTMLP(nn.Module):

151 def __init__(

152 self,

153 config: MptConfig,

155 - quant_config: Optional[QuantizationConfig] = None,

154 + quant_config: QuantizationConfig | None = None,

155 ):

156 super().__init__()

157 hidden_size = config.d_model

@@ -183,8 +182,8 @@ class MPTBlock(nn.Module):

182 def __init__(

183 self,

184 config: MptConfig,

186 - cache_config: Optional[CacheConfig] = None,

187 - quant_config: Optional[QuantizationConfig] = None,

185 + cache_config: CacheConfig | None = None,

186 + quant_config: QuantizationConfig | None = None,

187 prefix: str = "",

188 ):

189 super().__init__()

@@ -251,9 +250,9 @@ class MPTModel(nn.Module):

250 self,

251 input_ids: torch.Tensor,

252 position_ids: torch.Tensor,

254 - intermediate_tensors: Optional[IntermediateTensors],

255 - inputs_embeds: Optional[torch.Tensor] = None,

256 - ) -> Union[torch.Tensor, IntermediateTensors]:

253 + intermediate_tensors: IntermediateTensors | None,

254 + inputs_embeds: torch.Tensor | None = None,

255 + ) -> torch.Tensor | IntermediateTensors:

256 if get_pp_group().is_first_rank:

257 if inputs_embeds is not None:

258 hidden_states = inputs_embeds

@@ -311,9 +310,9 @@ class MPTForCausalLM(nn.Module, SupportsPP):

310 self,

311 input_ids: torch.Tensor,

312 positions: torch.Tensor,

314 - intermediate_tensors: Optional[IntermediateTensors] = None,

315 - inputs_embeds: Optional[torch.Tensor] = None,

316 - ) -> Union[torch.Tensor, IntermediateTensors]:

313 + intermediate_tensors: IntermediateTensors | None = None,

314 + inputs_embeds: torch.Tensor | None = None,

315 + ) -> torch.Tensor | IntermediateTensors:

316 hidden_states = self.transformer(

317 input_ids, positions, intermediate_tensors, inputs_embeds

318 )

@@ -322,7 +321,7 @@ class MPTForCausalLM(nn.Module, SupportsPP):

321 def compute_logits(

322 self,

323 hidden_states: torch.Tensor,

325 - ) -> Optional[torch.Tensor]:

324 + ) -> torch.Tensor | None:

325 logits = self.logits_processor(self.lm_head, hidden_states)

326 return logits

327

vllm/model_executor/models/nano_nemotron_vl.py

+45 -45

@@ -11,7 +11,7 @@ import copy

11 import warnings

12 from abc import ABC, abstractmethod

13 from collections.abc import Iterable, Mapping, Sequence

14 -from typing import Annotated, Any, Literal, Optional, TypedDict, TypeVar, Union

14 +from typing import Annotated, Any, Literal, TypeAlias, TypedDict, TypeVar

15

16 import numpy.typing as npt

17 import torch

@@ -110,7 +110,7 @@ class NanoNemotronVLImagePixelInputs(TypedDict):

110

111 class NanoNemotronVLImageEmbeddinInputs(TypedDict):

112 type: Literal["image_embeds"]

113 - data: Union[torch.Tensor, list[torch.Tensor]]

113 + data: torch.Tensor | list[torch.Tensor]

114 """

115 A tensor of shape `(num_images, total_image_feature_size, hidden_size)`

116 or a list of tensors of shape `(total_image_feature_size, hidden_size)`

@@ -119,9 +119,9 @@ class NanoNemotronVLImageEmbeddinInputs(TypedDict):

119 """

120

121

122 -NanoNemotronVLImageInputs = Union[

123 - NanoNemotronVLImagePixelInputs, NanoNemotronVLImageEmbeddinInputs

124 -]

122 +NanoNemotronVLImageInputs: TypeAlias = (

123 + NanoNemotronVLImagePixelInputs | NanoNemotronVLImageEmbeddinInputs

124 +)

125

126

127 class NanoNemotronVLVideoPixelInputs(TensorSchema):

@@ -148,12 +148,12 @@ class NanoNemotronVLVideoEmbeddingInputs(TensorSchema):

148 """

149

150 type: Literal["video_embeds"]

151 - data: Annotated[Union[torch.Tensor, list[torch.Tensor]], TensorShape("n", "f", "h")]

151 + data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("n", "f", "h")]

152

153

154 -NanoNemotronVLVideoInputs = Union[

155 - NanoNemotronVLVideoPixelInputs, NanoNemotronVLVideoEmbeddingInputs

156 -]

154 +NanoNemotronVLVideoInputs: TypeAlias = (

155 + NanoNemotronVLVideoPixelInputs | NanoNemotronVLVideoEmbeddingInputs

156 +)

157

158

159 def dynamic_preprocess(

@@ -262,7 +262,7 @@ class BaseNanoNemotronVLProcessor(ABC):

262 config: PretrainedConfig,

263 tokenizer: AnyTokenizer,

264 *args,

265 - max_num_tiles: Optional[int] = None,

265 + max_num_tiles: int | None = None,

266 **kwargs,

267 ) -> None:

268 super().__init__()

@@ -291,7 +291,7 @@ class BaseNanoNemotronVLProcessor(ABC):

291 def get_image_repl(

292 self,

293 feature_size: int,

294 - num_patches: Optional[int],

294 + num_patches: int | None,

295 ) -> PromptUpdateDetails[str]:

296 raise NotImplementedError

297

@@ -354,7 +354,7 @@ class BaseNanoNemotronVLProcessor(ABC):

354 text = [t.replace("<image>", image_repl.full, 1) for t in text]

355 return text, image_inputs

356

357 - def _make_batch_input(self, input_item: Optional[Union[Any, list[Any]]] = None):

357 + def _make_batch_input(self, input_item: Any | list[Any] | None = None):

358 if input_item is None:

359 input_item = []

360 if not isinstance(input_item, list):

@@ -363,10 +363,10 @@ class BaseNanoNemotronVLProcessor(ABC):

363

364 def __call__(

365 self,

366 - text: Optional[Union[str, list[str]]] = None,

367 - images: Optional[Union[Image.Image, list[Image.Image]]] = None,

368 - return_tensors: Optional[Union[str, TensorType]] = None,

369 - max_num_tiles: Optional[int] = None,

366 + text: str | list[str] | None = None,

367 + images: Image.Image | list[Image.Image] | None = None,

368 + return_tensors: str | TensorType | None = None,

369 + max_num_tiles: int | None = None,

370 ) -> BatchFeature:

371 # Use default if not provided

372 if max_num_tiles is None:

@@ -399,12 +399,12 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

399 config: PretrainedConfig,

400 tokenizer: AnyTokenizer,

401 *,

402 - max_num_tiles: Optional[int] = None,

403 - min_dynamic_patch: Optional[int] = None,

404 - max_dynamic_patch: Optional[int] = None,

405 - dynamic_image_size: Optional[bool] = None,

406 - video_token: Optional[str] = None,

407 - video_pruning_rate: Optional[float] = None,

402 + max_num_tiles: int | None = None,

403 + min_dynamic_patch: int | None = None,

404 + max_dynamic_patch: int | None = None,

405 + dynamic_image_size: bool | None = None,

406 + video_token: str | None = None,

407 + video_pruning_rate: float | None = None,

408 ) -> None:

409 super().__init__(

410 config=config,

@@ -423,7 +423,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

423 return self.video_token_id is not None

424

425 @property

426 - def video_token_id(self) -> Optional[int]:

426 + def video_token_id(self) -> int | None:

427 if self.video_token is None:

428 return None

429 return self.tokenizer.get_vocab().get(self.video_token, None)

@@ -436,7 +436,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

436 self,

437 videos: list[npt.NDArray],

438 max_num_tiles: int,

439 - dynamic_image_size: Optional[bool] = None,

439 + dynamic_image_size: bool | None = None,

440 ) -> list[torch.Tensor]:

441 return [

442 video_to_pixel_values(

@@ -453,7 +453,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

453 text: list[str],

454 videos: list[npt.NDArray],

455 max_num_tiles: int,

456 - dynamic_image_size: Optional[bool] = None,

456 + dynamic_image_size: bool | None = None,

457 ):

458 if len(videos) == 0 or not self.supports_video:

459 video_inputs = {}

@@ -508,12 +508,12 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

508

509 def __call__(

510 self,

511 - text: Optional[Union[str, list[str]]] = None,

512 - images: Optional[Union[Image.Image, list[Image.Image]]] = None,

513 - videos: Optional[Union[npt.NDArray, list[npt.NDArray]]] = None,

514 - return_tensors: Optional[Union[str, TensorType]] = None,

515 - max_num_tiles: Optional[int] = None,

516 - dynamic_image_size: Optional[bool] = None,

511 + text: str | list[str] | None = None,

512 + images: Image.Image | list[Image.Image] | None = None,

513 + videos: npt.NDArray | list[npt.NDArray] | None = None,

514 + return_tensors: str | TensorType | None = None,

515 + max_num_tiles: int | None = None,

516 + dynamic_image_size: bool | None = None,

517 ) -> BatchFeature:

518 # Use default if not provided

519 if max_num_tiles is None:

@@ -545,7 +545,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

545 def get_image_repl(

546 self,

547 feature_size: int,

548 - num_patches: Optional[int],

548 + num_patches: int | None,

549 ) -> PromptUpdateDetails[str]:

550 repl_features = IMG_CONTEXT * feature_size

551 repl_full = IMG_START + repl_features + IMG_END

@@ -598,7 +598,7 @@ class BaseNanoNemotronVLProcessingInfo(BaseProcessingInfo):

598 ) -> BaseNanoNemotronVLProcessor:

599 raise NotImplementedError

600

601 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

601 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

602 return {"image": None}

603

604 def get_num_image_tokens(

@@ -607,7 +607,7 @@ class BaseNanoNemotronVLProcessingInfo(BaseProcessingInfo):

607 image_width: int,

608 image_height: int,

609 max_num_tiles: int,

610 - processor: Optional[BaseNanoNemotronVLProcessor],

610 + processor: BaseNanoNemotronVLProcessor | None,

611 ) -> int:

612 if processor is None:

613 processor = self.get_hf_processor()

@@ -673,10 +673,10 @@ class NanoNemotronVLProcessingInfo(BaseNanoNemotronVLProcessingInfo):

673 video_limit = {"video": None} if self.supports_video else {}

674 return {**super().get_supported_mm_limits(), **video_limit}

675

676 - def get_video_token(self) -> Optional[str]:

676 + def get_video_token(self) -> str | None:

677 return IMG_CONTEXT

678

679 - def get_video_pruning_rate(self) -> Optional[float]:

679 + def get_video_pruning_rate(self) -> float | None:

680 return self.ctx.get_mm_config().video_pruning_rate

681

682 def get_num_frames_with_most_features(

@@ -929,7 +929,7 @@ class NanoNemotronVLDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

929 self,

930 seq_len: int,

931 mm_counts: Mapping[str, int],

932 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

932 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

933 ) -> MultiModalDataDict:

934 # Use default max_num_tiles for dummy data generation

935 max_num_tiles = 12

@@ -964,7 +964,7 @@ class NanoNemotronVLDummyInputsBuilder(

964 self,

965 seq_len: int,

966 mm_counts: Mapping[str, int],

967 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

967 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

968 ) -> MultiModalDataDict:

969 dummy_image = super().get_dummy_mm_data(

970 seq_len=seq_len, mm_counts=mm_counts, mm_options=mm_options

@@ -1000,7 +1000,7 @@ class NemotronH_Nano_VL_V2(

1000 nn.Module, HasInnerState, IsHybrid, SupportsMultiModal, SupportsMultiModalPruning

1001 ):

1002 @classmethod

1003 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1003 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1004 if modality.startswith("image"):

1005 return "<image>"

1006 if modality.startswith("video"):

@@ -1097,7 +1097,7 @@ class NemotronH_Nano_VL_V2(

1097

1098 def _parse_and_validate_image_input(

1099 self, **kwargs: object

1100 - ) -> Optional[NanoNemotronVLImageInputs]:

1100 + ) -> NanoNemotronVLImageInputs | None:

1101 pixel_values_flat = kwargs.pop("pixel_values_flat", None)

1102 image_num_patches = kwargs.pop("image_num_patches", None)

1103 image_embeds = kwargs.pop("image_embeds", None)

@@ -1274,7 +1274,7 @@ class NemotronH_Nano_VL_V2(

1274

1275 def _parse_and_validate_video_input(

1276 self, **kwargs: object

1277 - ) -> Optional[NanoNemotronVLVideoPixelInputs]:

1277 + ) -> NanoNemotronVLVideoPixelInputs | None:

1278 pixel_values_flat_video = kwargs.pop("pixel_values_flat_video", None)

1279 video_num_patches = kwargs.pop("video_num_patches", None)

1280 video_embeds = kwargs.pop("video_embeds", None)

@@ -1365,10 +1365,10 @@ class NemotronH_Nano_VL_V2(

1365 self,

1366 input_ids: torch.Tensor,

1367 positions: torch.Tensor,

1368 - intermediate_tensors: Optional[IntermediateTensors] = None,

1369 - inputs_embeds: Optional[torch.Tensor] = None,

1368 + intermediate_tensors: IntermediateTensors | None = None,

1369 + inputs_embeds: torch.Tensor | None = None,

1370 **kwargs: object,

1371 - ) -> Union[torch.Tensor, IntermediateTensors]:

1371 + ) -> torch.Tensor | IntermediateTensors:

1372 if intermediate_tensors is not None:

1373 input_ids = None

1374 inputs_embeds = None

@@ -1396,7 +1396,7 @@ class NemotronH_Nano_VL_V2(

1396 def compute_logits(

1397 self,

1398 hidden_states: torch.Tensor,

1399 - ) -> Optional[torch.Tensor]:

1399 + ) -> torch.Tensor | None:

1400 return self.language_model.compute_logits(hidden_states)

1401

1402 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

vllm/model_executor/models/nemotron.py

+18 -18

@@ -26,7 +26,7 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Any, Optional, Union

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -85,7 +85,7 @@ def _cast_if_autocast_enabled(*args):

85 class NemotronLayerNorm1P(nn.LayerNorm):

86 def __init__(

87 self,

88 - normalized_shape: Union[int, list[int], torch.Size],

88 + normalized_shape: int | list[int] | torch.Size,

89 eps: float = 1e-5,

90 elementwise_affine: bool = True,

91 bias: bool = True,

@@ -97,7 +97,7 @@ class NemotronLayerNorm1P(nn.LayerNorm):

97 def forward(

98 self,

99 x: torch.Tensor,

100 - residual: Optional[torch.Tensor] = None,

100 + residual: torch.Tensor | None = None,

101 ) -> torch.Tensor:

102 if residual is not None:

103 x = x + residual

@@ -116,7 +116,7 @@ class NemotronMLP(nn.Module):

116 hidden_size: int,

117 intermediate_size: int,

118 hidden_act: str,

119 - quant_config: Optional[QuantizationConfig] = None,

119 + quant_config: QuantizationConfig | None = None,

120 bias: bool = False,

121 prefix: str = "",

122 ) -> None:

@@ -152,11 +152,11 @@ class NemotronAttention(nn.Module):

152 num_heads: int,

153 num_kv_heads: int,

154 rope_theta: float = 10000,

155 - rope_scaling: Optional[dict[str, Any]] = None,

155 + rope_scaling: dict[str, Any] | None = None,

156 max_position_embeddings: int = 8192,

157 - quant_config: Optional[QuantizationConfig] = None,

157 + quant_config: QuantizationConfig | None = None,

158 bias: bool = False,

159 - cache_config: Optional[CacheConfig] = None,

159 + cache_config: CacheConfig | None = None,

160 prefix: str = "",

161 ) -> None:

162 super().__init__()

@@ -238,8 +238,8 @@ class NemotronDecoderLayer(nn.Module):

238 def __init__(

239 self,

240 config: NemotronConfig,

241 - cache_config: Optional[CacheConfig] = None,

242 - quant_config: Optional[QuantizationConfig] = None,

241 + cache_config: CacheConfig | None = None,

242 + quant_config: QuantizationConfig | None = None,

243 prefix: str = "",

244 ) -> None:

245 super().__init__()

@@ -292,7 +292,7 @@ class NemotronDecoderLayer(nn.Module):

292 self,

293 positions: torch.Tensor,

294 hidden_states: torch.Tensor,

295 - residual: Optional[torch.Tensor],

295 + residual: torch.Tensor | None,

296 ) -> tuple[torch.Tensor, torch.Tensor]:

297 # Self Attention

298 if residual is None:

@@ -363,11 +363,11 @@ class NemotronModel(nn.Module):

363

364 def forward(

365 self,

366 - input_ids: Optional[torch.Tensor],

366 + input_ids: torch.Tensor | None,

367 positions: torch.Tensor,

368 - intermediate_tensors: Optional[IntermediateTensors],

369 - inputs_embeds: Optional[torch.Tensor] = None,

370 - ) -> Union[torch.Tensor, IntermediateTensors]:

368 + intermediate_tensors: IntermediateTensors | None,

369 + inputs_embeds: torch.Tensor | None = None,

370 + ) -> torch.Tensor | IntermediateTensors:

371 if get_pp_group().is_first_rank:

372 if inputs_embeds is not None:

373 hidden_states = inputs_embeds

@@ -514,9 +514,9 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

514 self,

515 input_ids: torch.Tensor,

516 positions: torch.Tensor,

517 - intermediate_tensors: Optional[IntermediateTensors] = None,

518 - inputs_embeds: Optional[torch.Tensor] = None,

519 - ) -> Union[torch.Tensor, IntermediateTensors]:

517 + intermediate_tensors: IntermediateTensors | None = None,

518 + inputs_embeds: torch.Tensor | None = None,

519 + ) -> torch.Tensor | IntermediateTensors:

520 model_output = self.model(

521 input_ids, positions, intermediate_tensors, inputs_embeds

522 )

@@ -525,7 +525,7 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

525 def compute_logits(

526 self,

527 hidden_states: torch.Tensor,

528 - ) -> Optional[torch.Tensor]:

528 + ) -> torch.Tensor | None:

529 logits = self.logits_processor(self.lm_head, hidden_states)

530 return logits

531

vllm/model_executor/models/nemotron_h.py

+21 -22

@@ -19,7 +19,6 @@

19 """Inference-only NemotronH model."""

20

21 from collections.abc import Iterable

22 -from typing import Optional

22

23 import torch

24 from torch import nn

@@ -75,7 +74,7 @@ class NemotronHMLP(nn.Module):

74 self,

75 config: NemotronHConfig,

76 layer_idx: int,

78 - quant_config: Optional[QuantizationConfig] = None,

77 + quant_config: QuantizationConfig | None = None,

78 bias: bool = False,

79 prefix: str = "",

80 ) -> None:

@@ -119,9 +118,9 @@ class NemotronHMLPDecoderLayer(nn.Module):

118 self,

119 config: NemotronHConfig,

120 layer_idx: int,

122 - model_config: Optional[ModelConfig] = None,

123 - cache_config: Optional[CacheConfig] = None,

124 - quant_config: Optional[QuantizationConfig] = None,

121 + model_config: ModelConfig | None = None,

122 + cache_config: CacheConfig | None = None,

123 + quant_config: QuantizationConfig | None = None,

124 prefix: str = "",

125 ) -> None:

126 super().__init__()

@@ -140,7 +139,7 @@ class NemotronHMLPDecoderLayer(nn.Module):

139 def forward(

140 self,

141 hidden_states: torch.Tensor,

143 - residual: Optional[torch.Tensor],

142 + residual: torch.Tensor | None,

143 **kwargs,

144 ):

145 if residual is None:

@@ -158,9 +157,9 @@ class NemotronHMambaDecoderLayer(nn.Module):

157 self,

158 config: NemotronHConfig,

159 layer_idx: int,

161 - model_config: Optional[ModelConfig] = None,

162 - cache_config: Optional[CacheConfig] = None,

163 - quant_config: Optional[QuantizationConfig] = None,

160 + model_config: ModelConfig | None = None,

161 + cache_config: CacheConfig | None = None,

162 + quant_config: QuantizationConfig | None = None,

163 prefix: str = "",

164 ) -> None:

165 super().__init__()

@@ -188,7 +187,7 @@ class NemotronHMambaDecoderLayer(nn.Module):

187 def forward(

188 self,

189 hidden_states: torch.Tensor,

191 - residual: Optional[torch.Tensor],

190 + residual: torch.Tensor | None,

191 **kwargs,

192 ):

193 if residual is None:

@@ -207,9 +206,9 @@ class NemotronHAttention(nn.Module):

206 self,

207 config: NemotronHConfig,

208 layer_idx: int,

210 - model_config: Optional[ModelConfig] = None,

211 - cache_config: Optional[CacheConfig] = None,

212 - quant_config: Optional[QuantizationConfig] = None,

209 + model_config: ModelConfig | None = None,

210 + cache_config: CacheConfig | None = None,

211 + quant_config: QuantizationConfig | None = None,

212 prefix: str = "",

213 ) -> None:

214 super().__init__()

@@ -279,9 +278,9 @@ class NemotronHAttentionDecoderLayer(nn.Module):

278 self,

279 config: NemotronHConfig,

280 layer_idx: int,

282 - model_config: Optional[ModelConfig] = None,

283 - cache_config: Optional[CacheConfig] = None,

284 - quant_config: Optional[QuantizationConfig] = None,

281 + model_config: ModelConfig | None = None,

282 + cache_config: CacheConfig | None = None,

283 + quant_config: QuantizationConfig | None = None,

284 prefix: str = "",

285 ) -> None:

286 super().__init__()

@@ -301,7 +300,7 @@ class NemotronHAttentionDecoderLayer(nn.Module):

300 self,

301 positions: torch.Tensor,

302 hidden_states: torch.Tensor,

304 - residual: Optional[torch.Tensor],

303 + residual: torch.Tensor | None,

304 **kwargs,

305 ):

306 if residual is None:

@@ -377,8 +376,8 @@ class NemotronHModel(nn.Module):

376 self,

377 input_ids: torch.Tensor,

378 positions: torch.Tensor,

380 - intermediate_tensors: Optional[IntermediateTensors] = None,

381 - inputs_embeds: Optional[torch.Tensor] = None,

379 + intermediate_tensors: IntermediateTensors | None = None,

380 + inputs_embeds: torch.Tensor | None = None,

381 ) -> torch.Tensor:

382 if get_pp_group().is_first_rank:

383 if inputs_embeds is not None:

@@ -553,8 +552,8 @@ class NemotronHForCausalLM(

552 self,

553 input_ids: torch.Tensor,

554 positions: torch.Tensor,

556 - intermediate_tensors: Optional[IntermediateTensors] = None,

557 - inputs_embeds: Optional[torch.Tensor] = None,

555 + intermediate_tensors: IntermediateTensors | None = None,

556 + inputs_embeds: torch.Tensor | None = None,

557 **kwargs,

558 ):

559 hidden_states = self.model(

@@ -566,7 +565,7 @@ class NemotronHForCausalLM(

565 def compute_logits(

566 self,

567 hidden_states: torch.Tensor,

569 - ) -> Optional[torch.Tensor]:

568 + ) -> torch.Tensor | None:

569 logits = self.logits_processor(self.lm_head, hidden_states)

570 return logits

571

vllm/model_executor/models/nemotron_nas.py

+17 -17

@@ -26,7 +26,7 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Any, Optional, Union

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -84,12 +84,12 @@ class DeciLMAttention(LlamaAttention):

84 num_heads: int,

85 num_kv_heads: int,

86 rope_theta: float = 10000,

87 - rope_scaling: Optional[dict[str, Any]] = None,

87 + rope_scaling: dict[str, Any] | None = None,

88 max_position_embeddings: int = 8192,

89 - quant_config: Optional[QuantizationConfig] = None,

89 + quant_config: QuantizationConfig | None = None,

90 bias: bool = False,

91 bias_o_proj: bool = False,

92 - cache_config: Optional[CacheConfig] = None,

92 + cache_config: CacheConfig | None = None,

93 prefix: str = "",

94 attn_type: str = AttentionType.DECODER,

95 ) -> None:

@@ -112,8 +112,8 @@ class DeciLMAttention(LlamaAttention):

112 def _init_rotary_emb(

113 self,

114 config,

115 - rope_scaling: Optional[dict[str, Any]],

116 - quant_config: Optional[QuantizationConfig],

115 + rope_scaling: dict[str, Any] | None,

116 + quant_config: QuantizationConfig | None,

117 ) -> None:

118 # Enables YARN for Mistral and LLaMA4 derivatives.

119 is_neox_style = True

@@ -139,8 +139,8 @@ class DeciLMDecoderLayer(nn.Module):

139 self,

140 config: LlamaConfig,

141 layer_idx: int,

142 - cache_config: Optional[CacheConfig] = None,

143 - quant_config: Optional[QuantizationConfig] = None,

142 + cache_config: CacheConfig | None = None,

143 + quant_config: QuantizationConfig | None = None,

144 prefix: str = "",

145 ) -> None:

146 super().__init__()

@@ -210,7 +210,7 @@ class DeciLMDecoderLayer(nn.Module):

210 self,

211 positions: torch.Tensor,

212 hidden_states: torch.Tensor,

213 - residual: Optional[torch.Tensor],

213 + residual: torch.Tensor | None,

214 ) -> tuple[torch.Tensor, torch.Tensor]:

215 # Self Attention

216

@@ -303,11 +303,11 @@ class DeciModel(nn.Module):

303

304 def forward(

305 self,

306 - input_ids: Optional[torch.Tensor],

306 + input_ids: torch.Tensor | None,

307 positions: torch.Tensor,

308 - intermediate_tensors: Optional[IntermediateTensors],

309 - inputs_embeds: Optional[torch.Tensor] = None,

310 - ) -> Union[torch.Tensor, IntermediateTensors]:

308 + intermediate_tensors: IntermediateTensors | None,

309 + inputs_embeds: torch.Tensor | None = None,

310 + ) -> torch.Tensor | IntermediateTensors:

311 if get_pp_group().is_first_rank:

312 if inputs_embeds is not None:

313 hidden_states = inputs_embeds

@@ -487,9 +487,9 @@ class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):

487 self,

488 input_ids: torch.Tensor,

489 positions: torch.Tensor,

490 - intermediate_tensors: Optional[IntermediateTensors] = None,

491 - inputs_embeds: Optional[torch.Tensor] = None,

492 - ) -> Union[torch.Tensor, IntermediateTensors]:

490 + intermediate_tensors: IntermediateTensors | None = None,

491 + inputs_embeds: torch.Tensor | None = None,

492 + ) -> torch.Tensor | IntermediateTensors:

493 model_output = self.model(

494 input_ids, positions, intermediate_tensors, inputs_embeds

495 )

@@ -498,7 +498,7 @@ class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):

498 def compute_logits(

499 self,

500 hidden_states: torch.Tensor,

501 - ) -> Optional[torch.Tensor]:

501 + ) -> torch.Tensor | None:

502 logits = self.logits_processor(self.lm_head, hidden_states)

503 return logits

504

vllm/model_executor/models/nemotron_vl.py

+18 -19

@@ -9,7 +9,6 @@

9 # --------------------------------------------------------

10 from abc import ABC

11 from collections.abc import Iterable

12 -from typing import Optional

12

13 import torch

14 import torch.nn as nn

@@ -207,9 +206,9 @@ class NemotronVLProcessor(InternVLProcessor):

206 tokenizer: AnyTokenizer,

207 image_processor: BaseImageProcessorFast,

208 *,

210 - min_dynamic_patch: Optional[int] = None,

211 - max_dynamic_patch: Optional[int] = None,

212 - dynamic_image_size: Optional[bool] = None,

209 + min_dynamic_patch: int | None = None,

210 + max_dynamic_patch: int | None = None,

211 + dynamic_image_size: bool | None = None,

212 ) -> None:

213 ABC.__init__(self)

214 self.config = config

@@ -266,9 +265,9 @@ class NemotronVLProcessor(InternVLProcessor):

265 def _images_to_pixel_values_lst(

266 self,

267 images: list[Image.Image],

269 - min_dynamic_patch: Optional[int] = None,

270 - max_dynamic_patch: Optional[int] = None,

271 - dynamic_image_size: Optional[bool] = None,

268 + min_dynamic_patch: int | None = None,

269 + max_dynamic_patch: int | None = None,

270 + dynamic_image_size: bool | None = None,

271 ) -> list[torch.Tensor]:

272 min_num, max_num = self.resolve_min_max_num(

273 min_dynamic_patch=min_dynamic_patch,

@@ -292,9 +291,9 @@ class NemotronVLProcessor(InternVLProcessor):

291 self,

292 text: list[str],

293 images: list[Image.Image],

295 - min_dynamic_patch: Optional[int] = None,

296 - max_dynamic_patch: Optional[int] = None,

297 - dynamic_image_size: Optional[bool] = None,

294 + min_dynamic_patch: int | None = None,

295 + max_dynamic_patch: int | None = None,

296 + dynamic_image_size: bool | None = None,

297 ) -> tuple[list[str], dict[str, torch.Tensor]]:

298 if len(images) == 0:

299 image_inputs = {}

@@ -326,7 +325,7 @@ class NemotronVLProcessor(InternVLProcessor):

325 def get_image_repl(

326 self,

327 feature_size: int,

329 - num_patches: Optional[int],

328 + num_patches: int | None,

329 ) -> PromptUpdateDetails[str]:

330 repl_features = IMG_CONTEXT * feature_size

331 repl_full = IMG_START + repl_features + IMG_END

@@ -362,7 +361,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

361 merge_by_field_config = True

362

363 @classmethod

365 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

364 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

365 if modality.startswith("image"):

366 return "<image>"

367

@@ -426,7 +425,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

425 def _init_vision_model(

426 self,

427 config: PretrainedConfig,

429 - quant_config: Optional[QuantizationConfig],

428 + quant_config: QuantizationConfig | None,

429 *,

430 prefix: str,

431 ):

@@ -482,7 +481,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

481

482 def _parse_and_validate_image_input(

483 self, **kwargs: object

485 - ) -> Optional[InternVLImageInputs]:

484 + ) -> InternVLImageInputs | None:

485 pixel_values_flat = kwargs.pop("pixel_values_flat", None)

486 image_num_patches = kwargs.pop("image_num_patches", None)

487 image_embeds = kwargs.pop("image_embeds", None)

@@ -581,9 +580,9 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

580 def get_input_embeddings(

581 self,

582 input_ids: torch.Tensor,

584 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

583 + multimodal_embeddings: MultiModalEmbeddings | None = None,

584 *,

586 - is_multimodal: Optional[torch.Tensor] = None,

585 + is_multimodal: torch.Tensor | None = None,

586 handle_oov_mm_token: bool = False,

587 ) -> torch.Tensor:

588 if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:

@@ -604,8 +603,8 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

603 self,

604 input_ids: torch.Tensor,

605 positions: torch.Tensor,

607 - intermediate_tensors: Optional[IntermediateTensors] = None,

608 - inputs_embeds: Optional[torch.Tensor] = None,

606 + intermediate_tensors: IntermediateTensors | None = None,

607 + inputs_embeds: torch.Tensor | None = None,

608 **kwargs: object,

609 ) -> IntermediateTensors:

610 if intermediate_tensors is not None:

@@ -630,7 +629,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

629 def compute_logits(

630 self,

631 hidden_states: torch.Tensor,

633 - ) -> Optional[torch.Tensor]:

632 + ) -> torch.Tensor | None:

633 return self.language_model.compute_logits(hidden_states)

634

635 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/nvlm_d.py

+3 -4

@@ -8,7 +8,6 @@

8 # Licensed under Apache 2.0 License [see LICENSE for details]

9 # --------------------------------------------------------

10 from collections.abc import Mapping, Sequence

11 -from typing import Optional

11

12 import torch

13 import torch.nn as nn

@@ -49,7 +48,7 @@ class NVLMProcessor(BaseInternVLProcessor):

48 def get_image_repl(

49 self,

50 feature_size: int,

52 - num_patches: Optional[int],

51 + num_patches: int | None,

52 ) -> PromptUpdateDetails[str]:

53 if num_patches is None:

54 raise NotImplementedError("Embedding inputs are not supported")

@@ -93,7 +92,7 @@ class NVLMDummyInputsBuilder(BaseInternVLDummyInputsBuilder[NVLMProcessingInfo])

92 self,

93 seq_len: int,

94 mm_counts: Mapping[str, int],

96 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

95 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

96 ) -> MultiModalDataDict:

97 target_width, target_height = self.info.get_image_size_with_most_features()

98 num_images = mm_counts.get("image", 0)

@@ -189,7 +188,7 @@ class NVLM_D_Model(InternVLChatModel):

188 def _init_vision_model(

189 self,

190 config: PretrainedConfig,

192 - quant_config: Optional[QuantizationConfig],

191 + quant_config: QuantizationConfig | None,

192 *,

193 is_mono: bool,

194 prefix: str,

vllm/model_executor/models/olmo.py

+13 -14

@@ -26,7 +26,6 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Optional, Union

29

30 import torch

31 from torch import nn

@@ -72,8 +71,8 @@ class OlmoAttention(nn.Module):

71 def __init__(

72 self,

73 config: OlmoConfig,

75 - cache_config: Optional[CacheConfig] = None,

76 - quant_config: Optional[QuantizationConfig] = None,

74 + cache_config: CacheConfig | None = None,

75 + quant_config: QuantizationConfig | None = None,

76 prefix: str = "",

77 ):

78 super().__init__()

@@ -152,7 +151,7 @@ class OlmoMLP(nn.Module):

151 def __init__(

152 self,

153 config: OlmoConfig,

155 - quant_config: Optional[QuantizationConfig] = None,

154 + quant_config: QuantizationConfig | None = None,

155 prefix: str = "",

156 ):

157 super().__init__()

@@ -201,8 +200,8 @@ class OlmoDecoderLayer(nn.Module):

200 def __init__(

201 self,

202 config: OlmoConfig,

204 - cache_config: Optional[CacheConfig] = None,

205 - quant_config: Optional[QuantizationConfig] = None,

203 + cache_config: CacheConfig | None = None,

204 + quant_config: QuantizationConfig | None = None,

205 prefix: str = "",

206 ):

207 super().__init__()

@@ -226,7 +225,7 @@ class OlmoDecoderLayer(nn.Module):

225 self,

226 positions: torch.Tensor,

227 hidden_states: torch.Tensor,

229 - ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:

228 + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:

229 # Attention block.

230 residual = hidden_states

231 hidden_states = self.input_layernorm(hidden_states)

@@ -276,9 +275,9 @@ class OlmoModel(nn.Module):

275 self,

276 input_ids: torch.Tensor,

277 positions: torch.Tensor,

279 - intermediate_tensors: Optional[IntermediateTensors],

280 - inputs_embeds: Optional[torch.Tensor] = None,

281 - ) -> Union[torch.Tensor, IntermediateTensors]:

278 + intermediate_tensors: IntermediateTensors | None,

279 + inputs_embeds: torch.Tensor | None = None,

280 + ) -> torch.Tensor | IntermediateTensors:

281 """

282 :param input_ids: A tensor of shape `(batch_size, seq_len)`.

283 """

@@ -389,9 +388,9 @@ class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

388 self,

389 input_ids: torch.Tensor,

390 positions: torch.Tensor,

392 - intermediate_tensors: Optional[IntermediateTensors] = None,

393 - inputs_embeds: Optional[torch.Tensor] = None,

394 - ) -> Union[torch.Tensor, IntermediateTensors]:

391 + intermediate_tensors: IntermediateTensors | None = None,

392 + inputs_embeds: torch.Tensor | None = None,

393 + ) -> torch.Tensor | IntermediateTensors:

394 hidden_states = self.model(

395 input_ids=input_ids,

396 positions=positions,

@@ -403,7 +402,7 @@ class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

402 def compute_logits(

403 self,

404 hidden_states: torch.Tensor,

406 - ) -> Optional[torch.Tensor]:

405 + ) -> torch.Tensor | None:

406 logits = self.logits_processor(self.lm_head, hidden_states)

407 return logits

408

vllm/model_executor/models/olmo2.py

+7 -8

@@ -27,7 +27,6 @@

27 from collections.abc import Iterable

28 from functools import partial

29 from itertools import islice

30 -from typing import Optional, Union

30

31 import torch

32 from torch import nn

@@ -312,9 +311,9 @@ class Olmo2Model(nn.Module):

311 self,

312 input_ids: torch.Tensor,

313 positions: torch.Tensor,

315 - intermediate_tensors: Optional[IntermediateTensors],

316 - inputs_embeds: Optional[torch.Tensor] = None,

317 - ) -> Union[torch.Tensor, IntermediateTensors]:

314 + intermediate_tensors: IntermediateTensors | None,

315 + inputs_embeds: torch.Tensor | None = None,

316 + ) -> torch.Tensor | IntermediateTensors:

317 """

318 :param input_ids: A tensor of shape `(batch_size, seq_len)`.

319 """

@@ -429,9 +428,9 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

428 self,

429 input_ids: torch.Tensor,

430 positions: torch.Tensor,

432 - intermediate_tensors: Optional[IntermediateTensors] = None,

433 - inputs_embeds: Optional[torch.Tensor] = None,

434 - ) -> Union[torch.Tensor, IntermediateTensors]:

431 + intermediate_tensors: IntermediateTensors | None = None,

432 + inputs_embeds: torch.Tensor | None = None,

433 + ) -> torch.Tensor | IntermediateTensors:

434 hidden_states = self.model(

435 input_ids=input_ids,

436 positions=positions,

@@ -443,7 +442,7 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

442 def compute_logits(

443 self,

444 hidden_states: torch.Tensor,

446 - ) -> Optional[torch.Tensor]:

445 + ) -> torch.Tensor | None:

446 logits = self.logits_processor(self.lm_head, hidden_states)

447 return logits

448

vllm/model_executor/models/olmoe.py

+10 -11

@@ -17,7 +17,6 @@

17 from collections.abc import Iterable

18 from functools import partial

19 from itertools import islice

20 -from typing import Optional, Union

20

21 import torch

22 from torch import nn

@@ -77,9 +76,9 @@ class OlmoeMoE(nn.Module):

76 top_k: int,

77 hidden_size: int,

78 intermediate_size: int,

80 - params_dtype: Optional[torch.dtype] = None,

81 - quant_config: Optional[QuantizationConfig] = None,

82 - tp_size: Optional[int] = None,

79 + params_dtype: torch.dtype | None = None,

80 + quant_config: QuantizationConfig | None = None,

81 + tp_size: int | None = None,

82 prefix: str = "",

83 ):

84 super().__init__()

@@ -245,7 +244,7 @@ class OlmoeDecoderLayer(nn.Module):

244 self,

245 positions: torch.Tensor,

246 hidden_states: torch.Tensor,

248 - residual: Optional[torch.Tensor],

247 + residual: torch.Tensor | None,

248 ) -> torch.Tensor:

249 # Self Attention

250 if residual is None:

@@ -302,9 +301,9 @@ class OlmoeModel(nn.Module):

301 self,

302 input_ids: torch.Tensor,

303 positions: torch.Tensor,

305 - intermediate_tensors: Optional[IntermediateTensors],

306 - inputs_embeds: Optional[torch.Tensor] = None,

307 - ) -> Union[torch.Tensor, IntermediateTensors]:

304 + intermediate_tensors: IntermediateTensors | None,

305 + inputs_embeds: torch.Tensor | None = None,

306 + ) -> torch.Tensor | IntermediateTensors:

307 if get_pp_group().is_first_rank:

308 if inputs_embeds is not None:

309 hidden_states = inputs_embeds

@@ -483,9 +482,9 @@ class OlmoeForCausalLM(nn.Module, SupportsPP):

482 self,

483 input_ids: torch.Tensor,

484 positions: torch.Tensor,

486 - intermediate_tensors: Optional[IntermediateTensors] = None,

487 - inputs_embeds: Optional[torch.Tensor] = None,

488 - ) -> Union[torch.Tensor, IntermediateTensors]:

485 + intermediate_tensors: IntermediateTensors | None = None,

486 + inputs_embeds: torch.Tensor | None = None,

487 + ) -> torch.Tensor | IntermediateTensors:

488 hidden_states = self.model(

489 input_ids, positions, intermediate_tensors, inputs_embeds

490 )

vllm/model_executor/models/opt.py

+16 -17

@@ -22,7 +22,6 @@

22

23 from collections.abc import Iterable

24 from itertools import islice

25 -from typing import Optional, Union

25

26 import torch

27 from torch import nn

@@ -77,8 +76,8 @@ class OPTAttention(nn.Module):

76 embed_dim: int,

77 num_heads: int,

78 bias: bool = True,

80 - cache_config: Optional[CacheConfig] = None,

81 - quant_config: Optional[QuantizationConfig] = None,

79 + cache_config: CacheConfig | None = None,

80 + quant_config: QuantizationConfig | None = None,

81 prefix: str = "",

82 ) -> None:

83 super().__init__()

@@ -129,8 +128,8 @@ class OPTDecoderLayer(nn.Module):

128 def __init__(

129 self,

130 config: OPTConfig,

132 - cache_config: Optional[CacheConfig] = None,

133 - quant_config: Optional[QuantizationConfig] = None,

131 + cache_config: CacheConfig | None = None,

132 + quant_config: QuantizationConfig | None = None,

133 prefix: str = "",

134 ):

135 super().__init__()

@@ -202,8 +201,8 @@ class OPTDecoder(nn.Module):

201 def __init__(

202 self,

203 config: OPTConfig,

205 - cache_config: Optional[CacheConfig] = None,

206 - quant_config: Optional[QuantizationConfig] = None,

204 + cache_config: CacheConfig | None = None,

205 + quant_config: QuantizationConfig | None = None,

206 prefix: str = "",

207 ):

208 super().__init__()

@@ -270,9 +269,9 @@ class OPTDecoder(nn.Module):

269 self,

270 input_ids: torch.Tensor,

271 positions: torch.Tensor,

273 - intermediate_tensors: Optional[IntermediateTensors],

274 - inputs_embeds: Optional[torch.Tensor] = None,

275 - ) -> Union[torch.Tensor, IntermediateTensors]:

272 + intermediate_tensors: IntermediateTensors | None,

273 + inputs_embeds: torch.Tensor | None = None,

274 + ) -> torch.Tensor | IntermediateTensors:

275 if get_pp_group().is_first_rank:

276 if inputs_embeds is None:

277 inputs_embeds = self.get_input_embeddings(input_ids)

@@ -319,9 +318,9 @@ class OPTModel(nn.Module):

318 self,

319 input_ids: torch.Tensor,

320 positions: torch.Tensor,

322 - intermediate_tensors: Optional[IntermediateTensors],

323 - inputs_embeds: Optional[torch.Tensor] = None,

324 - ) -> Union[torch.Tensor, IntermediateTensors]:

321 + intermediate_tensors: IntermediateTensors | None,

322 + inputs_embeds: torch.Tensor | None = None,

323 + ) -> torch.Tensor | IntermediateTensors:

324 return self.decoder(

325 input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds

326 )

@@ -402,9 +401,9 @@ class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

401 self,

402 input_ids: torch.Tensor,

403 positions: torch.Tensor,

405 - intermediate_tensors: Optional[IntermediateTensors] = None,

406 - inputs_embeds: Optional[torch.Tensor] = None,

407 - ) -> Union[torch.Tensor, IntermediateTensors]:

404 + intermediate_tensors: IntermediateTensors | None = None,

405 + inputs_embeds: torch.Tensor | None = None,

406 + ) -> torch.Tensor | IntermediateTensors:

407 hidden_states = self.model(

408 input_ids, positions, intermediate_tensors, inputs_embeds

409 )

@@ -413,7 +412,7 @@ class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

412 def compute_logits(

413 self,

414 hidden_states: torch.Tensor,

416 - ) -> Optional[torch.Tensor]:

415 + ) -> torch.Tensor | None:

416 logits = self.logits_processor(self.lm_head, hidden_states)

417 return logits

418

vllm/model_executor/models/orion.py

+14 -14

@@ -9,7 +9,7 @@

9

10 from collections.abc import Iterable

11 from itertools import islice

12 -from typing import Any, Optional, Union

12 +from typing import Any

13

14 import torch

15 from torch import nn

@@ -51,7 +51,7 @@ class OrionMLP(nn.Module):

51 hidden_size: int,

52 intermediate_size: int,

53 hidden_act: str,

54 - quant_config: Optional[QuantizationConfig] = None,

54 + quant_config: QuantizationConfig | None = None,

55 ) -> None:

56 super().__init__()

57 self.gate_up_proj = MergedColumnParallelLinear(

@@ -80,10 +80,10 @@ class OrionAttention(nn.Module):

80 num_heads: int,

81 num_kv_heads: int,

82 rope_theta: float = 10000,

83 - rope_scaling: Optional[dict[str, Any]] = None,

83 + rope_scaling: dict[str, Any] | None = None,

84 max_position_embeddings: int = 8192,

85 - cache_config: Optional[CacheConfig] = None,

86 - quant_config: Optional[QuantizationConfig] = None,

85 + cache_config: CacheConfig | None = None,

86 + quant_config: QuantizationConfig | None = None,

87 prefix: str = "",

88 ) -> None:

89 super().__init__()

@@ -158,8 +158,8 @@ class OrionDecoderLayer(nn.Module):

158 def __init__(

159 self,

160 config: PretrainedConfig,

161 - cache_config: Optional[CacheConfig] = None,

162 - quant_config: Optional[QuantizationConfig] = None,

161 + cache_config: CacheConfig | None = None,

162 + quant_config: QuantizationConfig | None = None,

163 prefix: str = "",

164 ) -> None:

165 super().__init__()

@@ -250,9 +250,9 @@ class OrionModel(nn.Module):

250 self,

251 input_ids: torch.Tensor,

252 positions: torch.Tensor,

253 - intermediate_tensors: Optional[IntermediateTensors],

254 - inputs_embeds: Optional[torch.Tensor] = None,

255 - ) -> Union[torch.Tensor, IntermediateTensors]:

253 + intermediate_tensors: IntermediateTensors | None,

254 + inputs_embeds: torch.Tensor | None = None,

255 + ) -> torch.Tensor | IntermediateTensors:

256 if get_pp_group().is_first_rank:

257 if inputs_embeds is not None:

258 hidden_states = inputs_embeds

@@ -340,9 +340,9 @@ class OrionForCausalLM(nn.Module, SupportsPP):

340 self,

341 input_ids: torch.Tensor,

342 positions: torch.Tensor,

343 - intermediate_tensors: Optional[IntermediateTensors] = None,

344 - inputs_embeds: Optional[torch.Tensor] = None,

345 - ) -> Union[torch.Tensor, IntermediateTensors]:

343 + intermediate_tensors: IntermediateTensors | None = None,

344 + inputs_embeds: torch.Tensor | None = None,

345 + ) -> torch.Tensor | IntermediateTensors:

346 hidden_states = self.model(

347 input_ids, positions, intermediate_tensors, inputs_embeds

348 )

@@ -351,7 +351,7 @@ class OrionForCausalLM(nn.Module, SupportsPP):

351 def compute_logits(

352 self,

353 hidden_states: torch.Tensor,

354 - ) -> Optional[torch.Tensor]:

354 + ) -> torch.Tensor | None:

355 logits = self.logits_processor(self.lm_head, hidden_states)

356 return logits

357

vllm/model_executor/models/ovis.py

+11 -11

@@ -20,7 +20,7 @@

20

21 import math

22 from collections.abc import Iterable, Mapping

23 -from typing import Annotated, Literal, Optional, Union

23 +from typing import Annotated, Literal

24

25 import torch

26 import torch.nn as nn

@@ -87,7 +87,7 @@ class VisualTokenizer(torch.nn.Module):

87 def __init__(

88 self,

89 config: PretrainedConfig,

90 - quant_config: Optional[QuantizationConfig] = None,

90 + quant_config: QuantizationConfig | None = None,

91 prefix: str = "",

92 ):

93 super().__init__()

@@ -114,7 +114,7 @@ class VisualTokenizer(torch.nn.Module):

114 def _init_backbone(

115 self,

116 config: PretrainedConfig,

117 - quant_config: Optional[QuantizationConfig] = None,

117 + quant_config: QuantizationConfig | None = None,

118 prefix: str = "",

119 ) -> nn.Module:

120 model_type = config.backbone_config.model_type

@@ -282,7 +282,7 @@ class OvisProcessingInfo(BaseProcessingInfo):

282 text_model_type = hf_text_config.model_type

283 return IMAGE_PAD_TOKEN_MAP.get(text_model_type)

284

285 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

285 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

286 return {"image": None}

287

288 def get_image_size_with_most_features(self) -> ImageSize:

@@ -302,7 +302,7 @@ class OvisDummyInputsBuilder(BaseDummyInputsBuilder[OvisProcessingInfo]):

302 self,

303 seq_len: int,

304 mm_counts: Mapping[str, int],

305 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

305 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

306 ) -> MultiModalDataDict:

307 num_images = mm_counts.get("image", 0)

308

@@ -417,7 +417,7 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):

417 merge_by_field_config = True

418

419 @classmethod

420 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

420 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

421 if modality.startswith("image"):

422 return "<image>"

423

@@ -453,7 +453,7 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):

453

454 def _parse_and_validate_image_input(

455 self, **kwargs: object

456 - ) -> Optional[OvisImagePatchInputs]:

456 + ) -> OvisImagePatchInputs | None:

457 pixel_values = kwargs.pop("pixel_values", None)

458 indicator_tokens = kwargs.pop("indicator_tokens", None)

459

@@ -527,10 +527,10 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):

527 self,

528 input_ids: torch.Tensor,

529 positions: torch.Tensor,

530 - intermediate_tensors: Optional[IntermediateTensors] = None,

531 - inputs_embeds: Optional[torch.Tensor] = None,

530 + intermediate_tensors: IntermediateTensors | None = None,

531 + inputs_embeds: torch.Tensor | None = None,

532 **kwargs: object,

533 - ) -> Union[torch.Tensor, IntermediateTensors]:

533 + ) -> torch.Tensor | IntermediateTensors:

534 if intermediate_tensors is not None:

535 inputs_embeds = None

536

@@ -547,7 +547,7 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):

547 def compute_logits(

548 self,

549 hidden_states: torch.Tensor,

550 - ) -> Optional[torch.Tensor]:

550 + ) -> torch.Tensor | None:

551 logits = self.llm.compute_logits(hidden_states)

552 return logits

553

vllm/model_executor/models/ovis2_5.py

+13 -13

@@ -4,7 +4,7 @@

4

5 from collections.abc import Iterable, Mapping

6 from functools import partial

7 -from typing import Annotated, Literal, Optional, Union

7 +from typing import Annotated, Literal

8

9 import torch

10 import torch.nn as nn

@@ -102,7 +102,7 @@ class VisualTokenizer(torch.nn.Module):

102 self,

103 config: PretrainedConfig,

104 visual_vocab_size: int,

105 - quant_config: Optional[QuantizationConfig] = None,

105 + quant_config: QuantizationConfig | None = None,

106 prefix: str = "",

107 use_data_parallel: bool = False,

108 ):

@@ -129,7 +129,7 @@ class VisualTokenizer(torch.nn.Module):

129 def _init_backbone(

130 self,

131 config: PretrainedConfig,

132 - quant_config: Optional[QuantizationConfig] = None,

132 + quant_config: QuantizationConfig | None = None,

133 prefix: str = "",

134 use_data_parallel: bool = False,

135 ):

@@ -205,7 +205,7 @@ class Ovis2_5ProcessingInfo(BaseProcessingInfo):

205 def get_image_processor(self) -> BaseImageProcessor:

206 return self.get_hf_processor().image_processor # type: ignore

207

208 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

208 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

209 return {"image": None, "video": 1}

210

211 def get_image_size_with_most_features(self) -> ImageSize:

@@ -274,7 +274,7 @@ class Ovis2_5ProcessingInfo(BaseProcessingInfo):

274 image_width: int,

275 image_height: int,

276 num_frames: int,

277 - image_processor: Optional[BaseImageProcessor],

277 + image_processor: BaseImageProcessor | None,

278 ) -> int:

279 num_video_tokens = self.get_num_image_tokens(

280 image_width=image_width, image_height=image_height, num_frames=num_frames

@@ -305,7 +305,7 @@ class Ovis2_5DummyInputsBuilder(BaseDummyInputsBuilder[Ovis2_5ProcessingInfo]):

305 self,

306 seq_len: int,

307 mm_counts: Mapping[str, int],

308 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

308 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

309 ) -> MultiModalDataDict:

310 num_images = mm_counts.get("image", 0)

311 num_videos = mm_counts.get("video", 0)

@@ -482,7 +482,7 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

482

483 def _parse_and_validate_image_input(

484 self, **kwargs: object

485 - ) -> Optional[Ovis2_5ImagePatchInputs]:

485 + ) -> Ovis2_5ImagePatchInputs | None:

486 pixel_values = kwargs.pop("pixel_values", None)

487 indicator_tokens = kwargs.pop("indicator_tokens", None)

488 grids = kwargs.pop("grids", None)

@@ -516,7 +516,7 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

516

517 def _parse_and_validate_video_input(

518 self, **kwargs: object

519 - ) -> Optional[Ovis2_5VideoPatchInputs]:

519 + ) -> Ovis2_5VideoPatchInputs | None:

520 pixel_values = kwargs.pop("video_pixel_values", None)

521 indicator_tokens = kwargs.pop("video_indicator_tokens", None)

522 grids = kwargs.pop("video_grids", None)

@@ -549,7 +549,7 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

549 raise AssertionError("This line should be unreachable.")

550

551 def _process_visual_input(

552 - self, visual_input: Union[Ovis2_5ImagePatchInputs, Ovis2_5VideoPatchInputs]

552 + self, visual_input: Ovis2_5ImagePatchInputs | Ovis2_5VideoPatchInputs

553 ) -> MultiModalEmbeddings:

554 image_patches_flat = visual_input["flat_data"]

555 patches_per_image = visual_input["patches_per_item"]

@@ -629,10 +629,10 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

629 self,

630 input_ids: torch.Tensor,

631 positions: torch.Tensor,

632 - intermediate_tensors: Optional[IntermediateTensors] = None,

633 - inputs_embeds: Optional[torch.Tensor] = None,

632 + intermediate_tensors: IntermediateTensors | None = None,

633 + inputs_embeds: torch.Tensor | None = None,

634 **kwargs: object,

635 - ) -> Union[torch.Tensor, IntermediateTensors]:

635 + ) -> torch.Tensor | IntermediateTensors:

636 if intermediate_tensors is not None:

637 inputs_embeds = None

638

@@ -649,7 +649,7 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

649 def compute_logits(

650 self,

651 hidden_states: torch.Tensor,

652 - ) -> Optional[torch.Tensor]:

652 + ) -> torch.Tensor | None:

653 logits = self.llm.compute_logits(hidden_states)

654 return logits

655

vllm/model_executor/models/paligemma.py

+14 -12

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Mapping, Sequence

4 -from typing import Annotated, Literal, Optional, Union

4 +from typing import Annotated, Literal, TypeAlias

5

6 import torch

7 from torch import nn

@@ -74,7 +74,9 @@ class PaliGemmaImageEmbeddingInputs(TensorSchema):

74 data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]

75

76

77 -PaliGemmaImageInputs = Union[PaliGemmaImagePixelInputs, PaliGemmaImageEmbeddingInputs]

77 +PaliGemmaImageInputs: TypeAlias = (

78 + PaliGemmaImagePixelInputs | PaliGemmaImageEmbeddingInputs

79 +)

80

81

82 class PaliGemmaMultiModalProjector(nn.Module):

@@ -95,7 +97,7 @@ class PaliGemmaProcessingInfo(BaseProcessingInfo):

97 def get_vision_encoder_info(self):

98 return get_vision_encoder_info(self.get_hf_config())

99

98 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

100 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

101 return {"image": 1}

102

103 def get_num_image_tokens(

@@ -120,7 +122,7 @@ class PaliGemmaDummyInputsBuilder(BaseDummyInputsBuilder[PaliGemmaProcessingInfo

122 self,

123 seq_len: int,

124 mm_counts: Mapping[str, int],

123 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

125 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

126 ) -> MultiModalDataDict:

127 hf_config = self.info.get_hf_config()

128 vision_config = hf_config.vision_config

@@ -217,11 +219,11 @@ class PaliGemmaMultiModalProcessor(BaseMultiModalProcessor[PaliGemmaProcessingIn

219

220 def apply(

221 self,

220 - prompt: Union[str, list[int]],

222 + prompt: str | list[int],

223 mm_data: MultiModalDataDict,

224 hf_processor_mm_kwargs: Mapping[str, object],

223 - tokenization_kwargs: Optional[Mapping[str, object]] = None,

224 - mm_uuids: Optional[MultiModalUUIDDict] = None,

225 + tokenization_kwargs: Mapping[str, object] | None = None,

226 + mm_uuids: MultiModalUUIDDict | None = None,

227 ) -> MultiModalInputs:

228 mm_inputs = super().apply(

229 prompt,

@@ -273,7 +275,7 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

275 )

276

277 @classmethod

276 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

278 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

279 if modality.startswith("image"):

280 return None

281

@@ -317,7 +319,7 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

319

320 def _parse_and_validate_image_input(

321 self, **kwargs: object

320 - ) -> Optional[PaliGemmaImageInputs]:

322 + ) -> PaliGemmaImageInputs | None:

323 pixel_values = kwargs.pop("pixel_values", None)

324 image_embeds = kwargs.pop("image_embeds", None)

325

@@ -386,8 +388,8 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

388 self,

389 input_ids: torch.Tensor,

390 positions: torch.Tensor,

389 - intermediate_tensors: Optional[IntermediateTensors] = None,

390 - inputs_embeds: Optional[torch.Tensor] = None,

391 + intermediate_tensors: IntermediateTensors | None = None,

392 + inputs_embeds: torch.Tensor | None = None,

393 **kwargs: object,

394 ) -> IntermediateTensors:

395 if intermediate_tensors is not None:

@@ -402,7 +404,7 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP

404 def compute_logits(

405 self,

406 hidden_states: torch.Tensor,

405 - ) -> Optional[torch.Tensor]:

407 + ) -> torch.Tensor | None:

408 return self.language_model.compute_logits(hidden_states)

409

410 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/persimmon.py

+11 -12

@@ -25,7 +25,6 @@

25

26 from collections.abc import Iterable

27 from itertools import islice

28 -from typing import Optional, Union

28

29 import torch

30 from torch import nn

@@ -63,7 +62,7 @@ from .utils import (

62

63 class PersimmonMLP(nn.Module):

64 def __init__(

66 - self, config: PersimmonConfig, quant_config: Optional[QuantizationConfig] = None

65 + self, config: PersimmonConfig, quant_config: QuantizationConfig | None = None

66 ):

67 super().__init__()

68 self.dense_h_to_4h = ColumnParallelLinear(

@@ -85,8 +84,8 @@ class PersimmonAttention(nn.Module):

84 def __init__(

85 self,

86 config: PersimmonConfig,

88 - cache_config: Optional[CacheConfig] = None,

89 - quant_config: Optional[QuantizationConfig] = None,

87 + cache_config: CacheConfig | None = None,

88 + quant_config: QuantizationConfig | None = None,

89 prefix: str = "",

90 ):

91 super().__init__()

@@ -181,8 +180,8 @@ class PersimmonDecoderLayer(nn.Module):

180 def __init__(

181 self,

182 config: PersimmonConfig,

184 - cache_config: Optional[CacheConfig] = None,

185 - quant_config: Optional[QuantizationConfig] = None,

183 + cache_config: CacheConfig | None = None,

184 + quant_config: QuantizationConfig | None = None,

185 prefix: str = "",

186 ):

187 super().__init__()

@@ -263,9 +262,9 @@ class PersimmonModel(nn.Module):

262 self,

263 input_ids: torch.Tensor,

264 positions: torch.Tensor,

266 - intermediate_tensors: Optional[IntermediateTensors],

267 - inputs_embeds: Optional[torch.Tensor] = None,

268 - ) -> Union[torch.Tensor, IntermediateTensors]:

265 + intermediate_tensors: IntermediateTensors | None,

266 + inputs_embeds: torch.Tensor | None = None,

267 + ) -> torch.Tensor | IntermediateTensors:

268 if get_pp_group().is_first_rank:

269 if inputs_embeds is not None:

270 hidden_states = inputs_embeds

@@ -340,8 +339,8 @@ class PersimmonForCausalLM(nn.Module, SupportsPP):

339 self,

340 input_ids: torch.Tensor,

341 positions: torch.Tensor,

343 - intermediate_tensors: Optional[IntermediateTensors] = None,

344 - inputs_embeds: Optional[torch.Tensor] = None,

342 + intermediate_tensors: IntermediateTensors | None = None,

343 + inputs_embeds: torch.Tensor | None = None,

344 ):

345 hidden_states = self.model(

346 input_ids=input_ids,

@@ -354,7 +353,7 @@ class PersimmonForCausalLM(nn.Module, SupportsPP):

353 def compute_logits(

354 self,

355 hidden_states: torch.Tensor,

357 - ) -> Optional[torch.Tensor]:

356 + ) -> torch.Tensor | None:

357 logits = self.logits_processor(self.lm_head, hidden_states)

358 return logits

359

vllm/model_executor/models/phi.py

+12 -13

@@ -40,7 +40,6 @@

40

41 from collections.abc import Iterable

42 from itertools import islice

43 -from typing import Optional, Union

43

44 import torch

45 from torch import nn

@@ -80,8 +79,8 @@ class PhiAttention(nn.Module):

79 def __init__(

80 self,

81 config: PhiConfig,

83 - cache_config: Optional[CacheConfig] = None,

84 - quant_config: Optional[QuantizationConfig] = None,

82 + cache_config: CacheConfig | None = None,

83 + quant_config: QuantizationConfig | None = None,

84 prefix: str = "",

85 ):

86 super().__init__()

@@ -149,7 +148,7 @@ class PhiAttention(nn.Module):

148

149 class PhiMLP(nn.Module):

150 def __init__(

152 - self, config: PhiConfig, quant_config: Optional[QuantizationConfig] = None

151 + self, config: PhiConfig, quant_config: QuantizationConfig | None = None

152 ):

153 super().__init__()

154

@@ -179,8 +178,8 @@ class PhiLayer(nn.Module):

178 def __init__(

179 self,

180 config: PhiConfig,

182 - cache_config: Optional[CacheConfig] = None,

183 - quant_config: Optional[QuantizationConfig] = None,

181 + cache_config: CacheConfig | None = None,

182 + quant_config: QuantizationConfig | None = None,

183 prefix: str = "",

184 ):

185 super().__init__()

@@ -241,9 +240,9 @@ class PhiModel(nn.Module):

240 self,

241 input_ids: torch.Tensor,

242 positions: torch.Tensor,

244 - intermediate_tensors: Optional[IntermediateTensors],

245 - inputs_embeds: Optional[torch.Tensor] = None,

246 - ) -> Union[torch.Tensor, IntermediateTensors]:

243 + intermediate_tensors: IntermediateTensors | None,

244 + inputs_embeds: torch.Tensor | None = None,

245 + ) -> torch.Tensor | IntermediateTensors:

246 if get_pp_group().is_first_rank:

247 if inputs_embeds is not None:

248 hidden_states = inputs_embeds

@@ -348,9 +347,9 @@ class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

347 self,

348 input_ids: torch.Tensor,

349 positions: torch.Tensor,

351 - intermediate_tensors: Optional[IntermediateTensors] = None,

352 - inputs_embeds: Optional[torch.Tensor] = None,

353 - ) -> Union[torch.Tensor, IntermediateTensors]:

350 + intermediate_tensors: IntermediateTensors | None = None,

351 + inputs_embeds: torch.Tensor | None = None,

352 + ) -> torch.Tensor | IntermediateTensors:

353 hidden_states = self.model(

354 input_ids, positions, intermediate_tensors, inputs_embeds

355 )

@@ -360,7 +359,7 @@ class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

359 def compute_logits(

360 self,

361 hidden_states: torch.Tensor,

363 - ) -> Optional[torch.Tensor]:

362 + ) -> torch.Tensor | None:

363 logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias)

364 return logits

365

vllm/model_executor/models/phi3v.py

+17 -17

@@ -16,7 +16,7 @@

16 # See the License for the specific language governing permissions and

17 # limitations under the License.

18 from collections.abc import Iterable, Mapping, Sequence

19 -from typing import Annotated, Any, Literal, Optional, Union

19 +from typing import Annotated, Any, Literal, TypeAlias

20

21 import regex as re

22 import torch

@@ -96,7 +96,7 @@ CLIP_VIT_LARGE_PATCH14_336_CONFIG = CLIPVisionConfig(

96

97 def _init_img_processor(

98 hf_config: PretrainedConfig,

99 - quant_config: Optional[QuantizationConfig],

99 + quant_config: QuantizationConfig | None,

100 prefix: str = "",

101 ) -> CLIPVisionModel:

102 clip_config = CLIP_VIT_LARGE_PATCH14_336_CONFIG

@@ -132,14 +132,14 @@ class Phi3VImagePixelInputs(TensorSchema):

132

133 # Supports either a stacked tensor or a list of (p, 3, h, w) tensors

134 pixel_values: Annotated[

135 - Union[torch.Tensor, list[torch.Tensor]],

135 + torch.Tensor | list[torch.Tensor],

136 TensorShape(

137 "bn", "p", 3, "h", "w", dynamic_dims={"p"}

138 ), # 'p' may vary across items

139 ]

140

141 # Stacked tensor with height and width for each image

142 - image_sizes: Annotated[Optional[torch.Tensor], TensorShape("bn", 2)]

142 + image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)]

143

144

145 class Phi3VImageEmbeddingInputs(TensorSchema):

@@ -153,12 +153,12 @@ class Phi3VImageEmbeddingInputs(TensorSchema):

153

154 type: Literal["image_embeds"] = "image_embeds"

155 data: Annotated[

156 - Union[torch.Tensor, list[torch.Tensor]],

156 + torch.Tensor | list[torch.Tensor],

157 TensorShape("bn", "f", "h"),

158 ]

159

160

161 -Phi3VImageInputs = Union[Phi3VImagePixelInputs, Phi3VImageEmbeddingInputs]

161 +Phi3VImageInputs: TypeAlias = Phi3VImagePixelInputs | Phi3VImageEmbeddingInputs

162

163

164 class Phi3ImageEmbeddingBase(nn.Module):

@@ -192,7 +192,7 @@ class Phi3HDImageEmbedding(Phi3ImageEmbeddingBase):

192 def __init__(

193 self,

194 config: PretrainedConfig,

195 - quant_config: Optional[QuantizationConfig],

195 + quant_config: QuantizationConfig | None,

196 prefix: str = "",

197 ) -> None:

198 super().__init__()

@@ -350,7 +350,7 @@ class Phi3HDImageEmbedding(Phi3ImageEmbeddingBase):

350

351

352 class Phi3VProcessingInfo(BaseProcessingInfo):

353 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

353 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

354 return {"image": None}

355

356 def get_num_image_tokens(

@@ -358,7 +358,7 @@ class Phi3VProcessingInfo(BaseProcessingInfo):

358 *,

359 image_width: int,

360 image_height: int,

361 - processor: Optional[ProcessorMixin] = None,

361 + processor: ProcessorMixin | None = None,

362 ) -> int:

363 if processor is None:

364 processor = self.get_hf_processor()

@@ -386,7 +386,7 @@ class Phi3VDummyInputsBuilder(BaseDummyInputsBuilder[Phi3VProcessingInfo]):

386 self,

387 seq_len: int,

388 mm_counts: Mapping[str, int],

389 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

389 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

390 ) -> MultiModalDataDict:

391 num_images = mm_counts.get("image", 0)

392

@@ -574,7 +574,7 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant)

574 )

575

576 @classmethod

577 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

577 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

578 if modality.startswith("image"):

579 return f"<|image_{i}|>"

580

@@ -620,7 +620,7 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant)

620

621 def _parse_and_validate_image_input(

622 self, **kwargs: object

623 - ) -> Optional[Phi3VImageInputs]:

623 + ) -> Phi3VImageInputs | None:

624 pixel_values = kwargs.pop("pixel_values", None)

625 image_sizes = kwargs.pop("image_sizes", None)

626 image_embeds = kwargs.pop("image_embeds", None)

@@ -684,9 +684,9 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant)

684 def get_input_embeddings(

685 self,

686 input_ids: torch.Tensor,

687 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

687 + multimodal_embeddings: MultiModalEmbeddings | None = None,

688 *,

689 - is_multimodal: Optional[torch.Tensor] = None,

689 + is_multimodal: torch.Tensor | None = None,

690 handle_oov_mm_token: bool = False,

691 ) -> torch.Tensor:

692 inputs_embeds = self._get_text_embeddings(

@@ -716,8 +716,8 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant)

716 self,

717 input_ids: torch.Tensor,

718 positions: torch.Tensor,

719 - intermediate_tensors: Optional[IntermediateTensors] = None,

720 - inputs_embeds: Optional[torch.Tensor] = None,

719 + intermediate_tensors: IntermediateTensors | None = None,

720 + inputs_embeds: torch.Tensor | None = None,

721 **kwargs: object,

722 ):

723 if intermediate_tensors is not None:

@@ -732,7 +732,7 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant)

732 def compute_logits(

733 self,

734 hidden_states: torch.Tensor,

735 - ) -> Optional[torch.Tensor]:

735 + ) -> torch.Tensor | None:

736 return self.language_model.compute_logits(hidden_states)

737

738 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/phi4_multimodal.py

+24 -24

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import Annotated, Any, Literal, Optional, Union

5 +from typing import Annotated, Any, Literal, TypeAlias

6

7 import numpy as np

8 import torch

@@ -147,7 +147,7 @@ class Phi4MMImageEmbedding(nn.Module):

147 def get_img_features(

148 self,

149 img_embeds: torch.FloatTensor,

150 - attention_mask: Optional[torch.Tensor] = None,

150 + attention_mask: torch.Tensor | None = None,

151 ) -> torch.FloatTensor:

152 img_feature = self.img_processor(

153 img_embeds, patch_attention_mask=attention_mask

@@ -172,8 +172,8 @@ class Phi4MMImageEmbedding(nn.Module):

172 def forward(

173 self,

174 image_pixel_values: torch.FloatTensor,

175 - image_sizes: Optional[torch.Tensor] = None,

176 - image_attention_mask: Optional[torch.Tensor] = None,

175 + image_sizes: torch.Tensor | None = None,

176 + image_attention_mask: torch.Tensor | None = None,

177 ) -> torch.FloatTensor:

178 image_pixel_values = image_pixel_values.to(

179 self.img_processor.embeddings.patch_embedding.weight.dtype

@@ -278,7 +278,7 @@ class Phi4MultimodalAudioMLP(nn.Module):

278 def __init__(

279 self,

280 config: Phi4MultimodalAudioConfig,

281 - quant_config: Optional[QuantizationConfig] = None,

281 + quant_config: QuantizationConfig | None = None,

282 prefix: str = "",

283 ):

284 super().__init__()

@@ -311,7 +311,7 @@ class Phi4MultimodalAudioAttention(nn.Module):

311 def __init__(

312 self,

313 config: Phi4MultimodalAudioConfig,

314 - quant_config: Optional[QuantizationConfig] = None,

314 + quant_config: QuantizationConfig | None = None,

315 prefix: str = "",

316 ):

317 super().__init__()

@@ -522,7 +522,7 @@ class Phi4MultimodalAudioModel(nn.Module):

522 pad_mask = pad_mask & enc_streaming_mask

523 return pad_mask

524

525 - def forward(self, hidden_states: torch.Tensor, mask: Optional[torch.Tensor] = None):

525 + def forward(self, hidden_states: torch.Tensor, mask: torch.Tensor | None = None):

526 hidden_states = self.encoder_embedding(hidden_states)

527 hidden_states, hs_mask, mask = self.forward_embeddings(hidden_states, mask)

528

@@ -673,7 +673,7 @@ class Phi4MMImagePixelInputs(TensorSchema):

673 type: Literal["pixel_values"]

674

675 data: Annotated[

676 - Union[torch.Tensor, list[torch.Tensor]],

676 + torch.Tensor | list[torch.Tensor],

677 TensorShape(

678 "bn", "p", 3, "h", "w", dynamic_dims={"p"}

679 ), # may be different per batch and image

@@ -706,7 +706,7 @@ class Phi4MMImageEmbeddingInputs(TensorSchema):

706 type: Literal["image_embeds"]

707

708 data: Annotated[

709 - Union[torch.Tensor, list[torch.Tensor]],

709 + torch.Tensor | list[torch.Tensor],

710 TensorShape("bn", "f", "h"),

711 ]

712

@@ -722,7 +722,7 @@ class Phi4MMAudioFeatureInputs(TensorSchema):

722 type: Literal["audio_features"]

723

724 data: Annotated[

725 - Union[torch.Tensor, list[torch.Tensor]],

725 + torch.Tensor | list[torch.Tensor],

726 TensorShape("bn", "t", 80, dynamic_dims={"t"}),

727 ]

728

@@ -744,8 +744,8 @@ class Phi4MMAudioEmbeddingInputs(TensorSchema):

744 ]

745

746

747 -Phi4MMImageInput = Union[Phi4MMImagePixelInputs, Phi4MMImageEmbeddingInputs]

748 -Phi4MMAudioInputs = Union[Phi4MMAudioFeatureInputs, Phi4MMAudioEmbeddingInputs]

747 +Phi4MMImageInput: TypeAlias = Phi4MMImagePixelInputs | Phi4MMImageEmbeddingInputs

748 +Phi4MMAudioInputs: TypeAlias = Phi4MMAudioFeatureInputs | Phi4MMAudioEmbeddingInputs

749

750

751 def cat_with_pad(tensors, dim, padding_value=0):

@@ -786,7 +786,7 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

786

787 def get_image_processor(

788 self,

789 - processor: Optional[Phi4MMProcessor] = None,

789 + processor: Phi4MMProcessor | None = None,

790 ) -> Phi4MultimodalImageProcessorFast:

791 if processor is None:

792 processor = self.get_hf_processor()

@@ -794,11 +794,11 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

794

795 def get_dynamic_hd(

796 self,

797 - processor: Optional[Phi4MMProcessor] = None,

797 + processor: Phi4MMProcessor | None = None,

798 ) -> int:

799 return self.get_image_processor(processor).dynamic_hd

800

801 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

801 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

802 return {"audio": None, "image": None}

803

804 def _find_target_aspect_ratio(

@@ -936,7 +936,7 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

936 *,

937 image_width: int,

938 image_height: int,

939 - processor: Optional[Phi4MMProcessor] = None,

939 + processor: Phi4MMProcessor | None = None,

940 ) -> int:

941 hf_config = self.get_hf_config()

942 vision_config = hf_config.vision_config

@@ -959,7 +959,7 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

959

960 def get_image_size_with_most_features(

961 self,

962 - processor: Optional[Phi4MMProcessor] = None,

962 + processor: Phi4MMProcessor | None = None,

963 ) -> ImageSize:

964 vit_image_size = self.get_hf_config().vision_config.image_size

965

@@ -1038,7 +1038,7 @@ class Phi4MMDummyInputsBuilder(BaseDummyInputsBuilder[Phi4MMProcessingInfo]):

1038 self,

1039 seq_len: int,

1040 mm_counts: Mapping[str, int],

1041 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

1041 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

1042 ) -> MultiModalDataDict:

1043 num_audios = mm_counts.get("audio", 0)

1044 num_images = mm_counts.get("image", 0)

@@ -1216,7 +1216,7 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1216 )

1217

1218 @classmethod

1219 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1219 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1220 if modality.startswith("image"):

1221 return "<|image|>"

1222 if modality.startswith("audio"):

@@ -1253,7 +1253,7 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1253

1254 def _parse_and_validate_audio_input(

1255 self, **kwargs: object

1256 - ) -> Optional[Phi4MMAudioInputs]:

1256 + ) -> Phi4MMAudioInputs | None:

1257 """

1258 Parse and validate the audio input to the model. This handles both

1259 audio features and audio embeddings, but only the former is used for

@@ -1314,7 +1314,7 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1314

1315 def _parse_and_validate_image_input(

1316 self, **kwargs: object

1317 - ) -> Optional[Phi4MMImagePixelInputs]:

1317 + ) -> Phi4MMImagePixelInputs | None:

1318 image_pixel_values: NestedTensors = kwargs.get("image_pixel_values")

1319 if image_pixel_values is None:

1320 return None

@@ -1445,8 +1445,8 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1445 self,

1446 input_ids: torch.Tensor,

1447 positions: torch.Tensor,

1448 - intermediate_tensors: Optional[IntermediateTensors] = None,

1449 - inputs_embeds: Optional[torch.Tensor] = None,

1448 + intermediate_tensors: IntermediateTensors | None = None,

1449 + inputs_embeds: torch.Tensor | None = None,

1450 **kwargs: object,

1451 ) -> torch.Tensor:

1452 if intermediate_tensors is not None:

@@ -1464,7 +1464,7 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1464 def compute_logits(

1465 self,

1466 hidden_states: torch.Tensor,

1467 - ) -> Optional[torch.Tensor]:

1467 + ) -> torch.Tensor | None:

1468 return self.language_model.compute_logits(hidden_states)

1469

1470 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/phi4mm.py

+16 -16

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 -from typing import Annotated, Any, Literal, Optional, Union

5 +from typing import Annotated, Any, Literal, TypeAlias

6

7 import numpy as np

8 import torch

@@ -122,7 +122,7 @@ class Phi4MMImageEncoder(nn.Module):

122 def __init__(

123 self,

124 config: PretrainedConfig,

125 - quant_config: Optional[QuantizationConfig],

125 + quant_config: QuantizationConfig | None,

126 prefix: str = "",

127 model_dir: str = "",

128 ) -> None:

@@ -468,7 +468,7 @@ class Phi4MMImagePixelInputs(TensorSchema):

468 type: Literal["pixel_values"]

469

470 data: Annotated[

471 - Union[torch.Tensor, list[torch.Tensor]],

471 + torch.Tensor | list[torch.Tensor],

472 TensorShape(

473 "bn", "p", 3, "h", "w", dynamic_dims={"p"}

474 ), # may be different per batch and image

@@ -500,7 +500,7 @@ class Phi4MMAudioFeatureInputs(TensorSchema):

500 type: Literal["audio_features"]

501

502 data: Annotated[

503 - Union[torch.Tensor, list[torch.Tensor]],

503 + torch.Tensor | list[torch.Tensor],

504 TensorShape("bn", "t", 80, dynamic_dims={"t"}),

505 ]

506

@@ -521,7 +521,7 @@ class Phi4MMAudioEmbeddingInputs(TensorSchema):

521 ]

522

523

524 -Phi4MMAudioInputs = Union[Phi4MMAudioFeatureInputs, Phi4MMAudioEmbeddingInputs]

524 +Phi4MMAudioInputs: TypeAlias = Phi4MMAudioFeatureInputs | Phi4MMAudioEmbeddingInputs

525

526

527 def cat_with_pad(tensors, dim, padding_value=0):

@@ -561,7 +561,7 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

561

562 def get_dynamic_hd(

563 self,

564 - processor: Optional[ProcessorMixin] = None,

564 + processor: ProcessorMixin | None = None,

565 ) -> int:

566 if processor is None:

567 processor = self.get_hf_processor()

@@ -571,7 +571,7 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

571 def get_feature_extractor(self, **kwargs: object) -> SequenceFeatureExtractor:

572 return self.get_hf_processor(**kwargs).audio_processor

573

574 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

574 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

575 return {"audio": None, "image": None}

576

577 def _find_target_aspect_ratio(

@@ -709,7 +709,7 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

709 *,

710 image_width: int,

711 image_height: int,

712 - processor: Optional[ProcessorMixin] = None,

712 + processor: ProcessorMixin | None = None,

713 ) -> int:

714 hf_config = self.get_hf_config()

715 vision_encoder_name = hf_config.img_processor

@@ -735,7 +735,7 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

735

736 def get_image_size_with_most_features(

737 self,

738 - processor: Optional[ProcessorMixin] = None,

738 + processor: ProcessorMixin | None = None,

739 ) -> ImageSize:

740 hf_config = self.get_hf_config()

741 vision_encoder_name = hf_config.img_processor

@@ -819,7 +819,7 @@ class Phi4MMDummyInputsBuilder(BaseDummyInputsBuilder[Phi4MMProcessingInfo]):

819 self,

820 seq_len: int,

821 mm_counts: Mapping[str, int],

822 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

822 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

823 ) -> MultiModalDataDict:

824 num_audios = mm_counts.get("audio", 0)

825 num_images = mm_counts.get("image", 0)

@@ -1008,7 +1008,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1008 )

1009

1010 @classmethod

1011 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1011 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1012 if modality.startswith("image"):

1013 return f"<|image_{i}|>"

1014 if modality.startswith("audio"):

@@ -1074,7 +1074,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1074

1075 def _parse_and_validate_audio_input(

1076 self, **kwargs: object

1077 - ) -> Optional[Phi4MMAudioInputs]:

1077 + ) -> Phi4MMAudioInputs | None:

1078 """

1079 Parse and validate the audio input to the model. This handles both

1080 audio features and audio embeddings, but only the former is used for

@@ -1135,7 +1135,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1135

1136 def _parse_and_validate_image_input(

1137 self, **kwargs: object

1138 - ) -> Optional[Phi4MMImagePixelInputs]:

1138 + ) -> Phi4MMImagePixelInputs | None:

1139 input_image_embeds: NestedTensors = kwargs.get("input_image_embeds")

1140 if input_image_embeds is None:

1141 return None

@@ -1263,8 +1263,8 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1263 self,

1264 input_ids: torch.Tensor,

1265 positions: torch.Tensor,

1266 - intermediate_tensors: Optional[IntermediateTensors] = None,

1267 - inputs_embeds: Optional[torch.Tensor] = None,

1266 + intermediate_tensors: IntermediateTensors | None = None,

1267 + inputs_embeds: torch.Tensor | None = None,

1268 **kwargs: object,

1269 ) -> torch.Tensor:

1270 if intermediate_tensors is not None:

@@ -1282,7 +1282,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1282 def compute_logits(

1283 self,

1284 hidden_states: torch.Tensor,

1285 - ) -> Optional[torch.Tensor]:

1285 + ) -> torch.Tensor | None:

1286 logits = self.logits_processor(self.lm_head, hidden_states)

1287 return logits

1288

vllm/model_executor/models/phi4mm_audio.py

+38 -38

@@ -7,7 +7,7 @@

7 #!/usr/bin/env python3

8 import abc

9 import math

10 -from typing import Any, Literal, Optional, Union

10 +from typing import Any, Literal

11

12 import numpy as np

13 import torch

@@ -221,7 +221,7 @@ class ConformerEncoderLayer(nn.Module):

221 pos_k: torch.Tensor,

222 pos_v: torch.Tensor,

223 mask: torch.Tensor,

224 - relative_attention_bias: Optional[Tensor] = None,

224 + relative_attention_bias: Tensor | None = None,

225 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

226 """ConformerEncoder forward.

227

@@ -329,8 +329,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

329 def __init__(

330 self,

331 input_size: int,

332 - chunk_size: Union[int, list[int]],

333 - left_chunk: Union[int, list[int]],

332 + chunk_size: int | list[int],

333 + left_chunk: int | list[int],

334 attention_dim: int = 256,

335 attention_heads: int = 4,

336 input_layer: str = "nemo_conv",

@@ -339,12 +339,12 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

339 time_reduction: int = 4,

340 dropout_rate: float = 0.0,

341 padding_idx: int = -1,

342 - relative_attention_bias_args: Optional[dict[str, Any]] = None,

342 + relative_attention_bias_args: dict[str, Any] | None = None,

343 positional_dropout_rate: float = 0.0,

344 - nemo_conv_settings: Optional[dict[str, Any]] = None,

344 + nemo_conv_settings: dict[str, Any] | None = None,

345 conv2d_extra_padding: Literal["feat", "feat_time", "none", True] = "none",

346 attention_group_size: int = 1,

347 - encoder_embedding_config: Optional[dict[str, Any]] = None,

347 + encoder_embedding_config: dict[str, Any] | None = None,

348 ) -> None:

349 super().__init__()

350 self.input_size = input_size

@@ -411,8 +411,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

411 )

412

413 def compute_lens_change(

414 - self, feature_lens: Union[int, torch.Tensor]

415 - ) -> Union[int, torch.Tensor]:

414 + self, feature_lens: int | torch.Tensor

415 + ) -> int | torch.Tensor:

416 """feature_lens: int

417 return updated feature lens.

418

@@ -452,8 +452,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

452

453 def _chunk_size_selection(

454 self,

455 - chunk_size: Optional[Union[int, list[int]]] = None,

456 - left_chunk: Optional[Union[int, list[int]]] = None,

455 + chunk_size: int | list[int] | None = None,

456 + left_chunk: int | list[int] | None = None,

457 ) -> tuple[int, int]:

458 """If chunk size is a list, we will randomly select a chunk size."""

459

@@ -503,7 +503,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

503

504 def _position_embedding(

505 self, input_tensor: torch.Tensor

506 - ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:

506 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:

507 pos_k = None

508 pos_v = None

509 if self.relative_attention_bias_layer is None:

@@ -516,8 +516,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

516 self,

517 seq_len: int,

518 batch_size: int,

519 - chunk_size: Union[int, list[int]],

520 - left_chunk: Union[int, list[int]],

519 + chunk_size: int | list[int],

520 + left_chunk: int | list[int],

521 ) -> torch.Tensor:

522 chunk_size_train_eff, left_chunk_train_eff = self._chunk_size_selection(

523 chunk_size, left_chunk

@@ -540,25 +540,25 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

540 self,

541 xs_pad: torch.Tensor,

542 masks: torch.Tensor,

543 - chunk_size_nc: Optional[Union[int, list[int]]] = None,

544 - left_chunk_nc: Optional[Union[int, list[int]]] = None,

545 - ) -> Union[

543 + chunk_size_nc: int | list[int] | None = None,

544 + left_chunk_nc: int | list[int] | None = None,

545 + ) -> (

546 tuple[

547 torch.Tensor,

548 - Optional[torch.Tensor],

549 - Optional[torch.Tensor],

548 + torch.Tensor | None,

549 + torch.Tensor | None,

550 torch.Tensor,

551 torch.Tensor,

552 - ],

553 - tuple[

552 + ]

553 + | tuple[

554 torch.Tensor,

555 - Optional[torch.Tensor],

556 - Optional[torch.Tensor],

555 + torch.Tensor | None,

556 + torch.Tensor | None,

557 torch.Tensor,

558 torch.Tensor,

559 torch.Tensor,

560 - ],

561 - ]:

560 + ]

561 + ):

562 """Forwarding the inputs through the top embedding layers

563

564 Args:

@@ -803,9 +803,9 @@ class ConformerEncoder(TransformerEncoderBase):

803 def __init__( # pylint: disable-all

804 self,

805 input_size: int,

806 - chunk_size: Union[int, list[int]],

807 - left_chunk: Union[int, list[int]],

808 - num_lang: Optional[int] = None,

806 + chunk_size: int | list[int],

807 + left_chunk: int | list[int],

808 + num_lang: int | None = None,

809 attention_dim: int = 256,

810 attention_heads: int = 4,

811 linear_units: int = 2048,

@@ -832,14 +832,14 @@ class ConformerEncoder(TransformerEncoderBase):

832 extra_layer_output_idx: int = -1,

833 extra_multi_layer_output_idxs: list[int] = [], # noqa

834 activation_checkpointing: str = "",

835 - relative_attention_bias_args: Optional[dict[str, Any]] = None,

835 + relative_attention_bias_args: dict[str, Any] | None = None,

836 time_reduction: int = 4,

837 use_pt_scaled_dot_product_attention: bool = False,

838 - nemo_conv_settings: Optional[dict[str, Any]] = None,

838 + nemo_conv_settings: dict[str, Any] | None = None,

839 conv2d_extra_padding: Literal["feat", "feat_time", "none", True] = "none",

840 replication_pad_for_subsample_embedding: bool = False,

841 attention_group_size: int = 1,

842 - encoder_embedding_config: Optional[dict[str, Any]] = None,

842 + encoder_embedding_config: dict[str, Any] | None = None,

843 ) -> None:

844 super().__init__(

845 input_size,

@@ -908,12 +908,12 @@ class ConformerEncoder(TransformerEncoderBase):

908

909 def init_relative_attention_bias(

910 self, input_tensor: torch.Tensor

911 - ) -> Optional[torch.Tensor]:

911 + ) -> torch.Tensor | None:

912 if self.relative_attention_bias_layer:

913 return self.relative_attention_bias_layer(input_tensor)

914

915 def calculate_hs_mask(

916 - self, xs_pad: torch.Tensor, device: torch.device, mask: Optional[torch.Tensor]

916 + self, xs_pad: torch.Tensor, device: torch.device, mask: torch.Tensor | None

917 ) -> torch.Tensor:

918 max_audio_length = xs_pad.shape[1]

919 batch_size = xs_pad.shape[0]

@@ -1066,9 +1066,9 @@ class WindowQformer(nn.Module):

1066 def forward(

1067 self,

1068 audio_embed: torch.Tensor,

1069 - mask: Optional[torch.Tensor],

1070 - embed_len: Optional[int] = None,

1071 - ) -> tuple[torch.Tensor, Optional[int]]:

1069 + mask: torch.Tensor | None,

1070 + embed_len: int | None = None,

1071 + ) -> tuple[torch.Tensor, int | None]:

1072 """forward decoder"""

1073 # audio_embed: N x T x D => N x D x T

1074

@@ -1224,7 +1224,7 @@ class AudioEmbedding(nn.Module):

1224 def get_audio_features(

1225 self,

1226 input_embeds: torch.Tensor,

1227 - audio_attention_mask: Optional[torch.Tensor] = None,

1227 + audio_attention_mask: torch.Tensor | None = None,

1228 audio_projection_mode: str = "speech",

1229 ) -> torch.Tensor:

1230 """

@@ -1278,7 +1278,7 @@ class AudioEmbedding(nn.Module):

1278 def forward(

1279 self,

1280 audio_features: torch.Tensor,

1281 - audio_attention_mask: Optional[torch.Tensor] = None,

1281 + audio_attention_mask: torch.Tensor | None = None,

1282 audio_projection_mode: str = "speech",

1283 ) -> torch.Tensor:

1284 """

vllm/model_executor/models/phi4mm_utils.py

+17 -20

@@ -6,7 +6,6 @@

6 # but implemented by the Phi-Speech team

7 #!/usr/bin/env python3

8 import math

9 -from typing import Optional, Union

9

10 import torch

11 import torch.nn.functional as F

@@ -917,7 +916,7 @@ class CausalConv1D(nn.Conv1d):

916 out_channels: int,

917 kernel_size: int,

918 stride: int = 1,

920 - padding: Union[str, int] = 0,

919 + padding: str | int = 0,

920 dilation: int = 1,

921 groups: int = 1,

922 bias: bool = True,

@@ -962,8 +961,8 @@ class CausalConv1D(nn.Conv1d):

961 )

962

963 def update_cache(

965 - self, x: Tensor, cache: Optional[Tensor] = None

966 - ) -> tuple[Tensor, Optional[Tensor]]:

964 + self, x: Tensor, cache: Tensor | None = None

965 + ) -> tuple[Tensor, Tensor | None]:

966 if cache is None:

967 new_x = F.pad(x, pad=(self._left_padding, self._right_padding))

968 next_cache = cache

@@ -978,8 +977,8 @@ class CausalConv1D(nn.Conv1d):

977 return new_x, next_cache

978

979 def forward(

981 - self, x: Tensor, cache: Optional[Tensor] = None

982 - ) -> Union[Tensor, tuple[Tensor, Optional[Tensor]]]:

980 + self, x: Tensor, cache: Tensor | None = None

981 + ) -> Tensor | tuple[Tensor, Tensor | None]:

982 x, cache = self.update_cache(x, cache=cache)

983 x = super().forward(x)

984 if cache is None:

@@ -1002,7 +1001,7 @@ class CausalConv2D(nn.Conv2d):

1001 out_channels: int,

1002 kernel_size: int,

1003 stride: int = 1,

1005 - padding: Union[str, int] = 0,

1004 + padding: str | int = 0,

1005 dilation: int = 1,

1006 groups: int = 1,

1007 bias: bool = True,

@@ -1371,9 +1370,7 @@ class NemoConvSubsampling(torch.nn.Module):

1370 def get_streaming_cache_size(self) -> list[int]:

1371 return [0, self.subsampling_factor + 1]

1372

1374 - def forward(

1375 - self, x: Tensor, mask: Optional[Tensor]

1376 - ) -> tuple[Tensor, Optional[Tensor]]:

1373 + def forward(self, x: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor | None]:

1374 """

1375 Forward method for NeMo subsampling.

1376

@@ -1615,10 +1612,10 @@ class AttModule(nn.Module):

1612 def forward(

1613 self,

1614 x: Tensor,

1618 - memory: Optional[Tensor] = None,

1619 - pos_emb: Optional[Tensor] = None,

1620 - att_mask: Optional[Tensor] = None,

1621 - ) -> tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:

1615 + memory: Tensor | None = None,

1616 + pos_emb: Tensor | None = None,

1617 + att_mask: Tensor | None = None,

1618 + ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:

1619 """AttModule forward

1620

1621 Args:

@@ -1640,7 +1637,7 @@ class AttBlock(BlockBase, AttModule):

1637

1638 def masked_softmax(

1639 scores: Tensor,

1643 - mask: Optional[Tensor],

1640 + mask: Tensor | None,

1641 ) -> Tensor:

1642 if mask is not None:

1643 mask = mask.unsqueeze(1).eq(0) # (batch, 1, time1, time2)

@@ -1720,7 +1717,7 @@ class MultiHeadedAttention(nn.Module):

1717 self.linear_v = nn.Linear(n_value, attention_inner_dim // group_size)

1718 self.linear_out = nn.Linear(attention_inner_dim // group_size, n_value)

1719

1723 - self.attn = torch.jit.Attribute(None, Optional[Tensor])

1720 + self.attn = torch.jit.Attribute(None, Tensor | None)

1721 self.dropout = nn.Dropout(p=dropout_rate)

1722 self.dropout_rate = dropout_rate

1723 self.use_pt_scaled_dot_product_attention = use_pt_scaled_dot_product_attention

@@ -1741,10 +1738,10 @@ class MultiHeadedAttention(nn.Module):

1738 query: Tensor,

1739 key: Tensor,

1740 value: Tensor,

1744 - pos_k: Optional[Tensor],

1745 - pos_v: Optional[Tensor],

1746 - mask: Optional[Tensor],

1747 - relative_attention_bias: Optional[Tensor] = None,

1741 + pos_k: Tensor | None,

1742 + pos_v: Tensor | None,

1743 + mask: Tensor | None,

1744 + relative_attention_bias: Tensor | None = None,

1745 ) -> Tensor:

1746 """Compute 'Scaled Dot Product Attention'.

1747

vllm/model_executor/models/phimoe.py

+16 -17

@@ -26,7 +26,6 @@

26

27 from collections.abc import Iterable

28 from itertools import islice

29 -from typing import Optional, Union

29

30 import torch

31 from torch import nn

@@ -257,9 +256,9 @@ class PhiMoE(nn.Module):

256 top_k: int,

257 hidden_size: int,

258 intermediate_size: int,

260 - params_dtype: Optional[torch.dtype] = None,

261 - quant_config: Optional[QuantizationConfig] = None,

262 - tp_size: Optional[int] = None,

259 + params_dtype: torch.dtype | None = None,

260 + quant_config: QuantizationConfig | None = None,

261 + tp_size: int | None = None,

262 prefix: str = "",

263 ):

264 super().__init__()

@@ -304,12 +303,12 @@ class PhiMoEAttention(nn.Module):

303 hidden_size: int,

304 num_heads: int,

305 num_kv_heads: int,

307 - head_dim: Optional[int] = None,

306 + head_dim: int | None = None,

307 max_position: int = 4096 * 32,

308 rope_theta: float = 10000,

310 - cache_config: Optional[CacheConfig] = None,

311 - quant_config: Optional[QuantizationConfig] = None,

312 - rope_scaling: Optional[dict] = None,

309 + cache_config: CacheConfig | None = None,

310 + quant_config: QuantizationConfig | None = None,

311 + rope_scaling: dict | None = None,

312 prefix: str = "",

313 ) -> None:

314 super().__init__()

@@ -386,8 +385,8 @@ class PhiMoEDecoderLayer(nn.Module):

385 def __init__(

386 self,

387 config: PhiMoEConfig,

389 - cache_config: Optional[CacheConfig] = None,

390 - quant_config: Optional[QuantizationConfig] = None,

388 + cache_config: CacheConfig | None = None,

389 + quant_config: QuantizationConfig | None = None,

390 prefix: str = "",

391 ) -> None:

392 super().__init__()

@@ -427,7 +426,7 @@ class PhiMoEDecoderLayer(nn.Module):

426 self,

427 positions: torch.Tensor,

428 hidden_states: torch.Tensor,

430 - residual: Optional[torch.Tensor],

429 + residual: torch.Tensor | None,

430 ) -> torch.Tensor:

431 residual = hidden_states

432

@@ -496,9 +495,9 @@ class PhiMoEModel(nn.Module):

495 self,

496 input_ids: torch.Tensor,

497 positions: torch.Tensor,

499 - intermediate_tensors: Optional[IntermediateTensors],

500 - inputs_embeds: Optional[torch.Tensor] = None,

501 - ) -> Union[torch.Tensor, IntermediateTensors]:

498 + intermediate_tensors: IntermediateTensors | None,

499 + inputs_embeds: torch.Tensor | None = None,

500 + ) -> torch.Tensor | IntermediateTensors:

501 if get_pp_group().is_first_rank:

502 if inputs_embeds is not None:

503 hidden_states = inputs_embeds

@@ -674,9 +673,9 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

673 self,

674 input_ids: torch.Tensor,

675 positions: torch.Tensor,

677 - intermediate_tensors: Optional[IntermediateTensors] = None,

678 - inputs_embeds: Optional[torch.Tensor] = None,

679 - ) -> Union[torch.Tensor, IntermediateTensors]:

676 + intermediate_tensors: IntermediateTensors | None = None,

677 + inputs_embeds: torch.Tensor | None = None,

678 + ) -> torch.Tensor | IntermediateTensors:

679 hidden_states = self.model(

680 input_ids, positions, intermediate_tensors, inputs_embeds

681 )

vllm/model_executor/models/pixtral.py

+31 -31

@@ -5,7 +5,7 @@ import math

5 from collections.abc import Iterable, Mapping, Sequence

6 from dataclasses import dataclass, fields

7 from functools import cached_property

8 -from typing import Annotated, Literal, Optional, Union

8 +from typing import Annotated, Literal

9

10 import torch

11 import torch.nn as nn

@@ -100,7 +100,7 @@ class PixtralImagePixelInputs(TensorSchema):

100 type: Literal["pixel_values"] = "pixel_values"

101

102 images: Annotated[

103 - Union[torch.Tensor, list[torch.Tensor]],

103 + torch.Tensor | list[torch.Tensor],

104 TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}),

105 ]

106

@@ -144,9 +144,9 @@ class PixtralProcessorAdapter:

144

145 def __call__(

146 self,

147 - text: Optional[Union[TextInput, list[TextInput]]] = None,

148 - images: Optional[Union[ImageInput, list[ImageInput]]] = None,

149 - return_tensors: Optional[Union[str, TensorType]] = None,

147 + text: TextInput | list[TextInput] | None = None,

148 + images: ImageInput | list[ImageInput] | None = None,

149 + return_tensors: str | TensorType | None = None,

150 **kwargs,

151 ) -> Mapping[str, NestedTensors]:

152 if text is None:

@@ -203,12 +203,12 @@ class PixtralProcessingInfo(BaseProcessingInfo):

203 def get_hf_processor(self) -> PixtralProcessorAdapter:

204 return PixtralProcessorAdapter(self.get_tokenizer())

205

206 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

206 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

207 return {"image": None}

208

209 def get_vision_config(

210 self,

211 - processor: Optional[PixtralProcessorAdapter] = None,

211 + processor: PixtralProcessorAdapter | None = None,

212 ):

213 if processor is None:

214 processor = self.get_hf_processor()

@@ -223,7 +223,7 @@ class PixtralProcessingInfo(BaseProcessingInfo):

223 *,

224 image_width: int,

225 image_height: int,

226 - processor: Optional[PixtralProcessorAdapter] = None,

226 + processor: PixtralProcessorAdapter | None = None,

227 ) -> int:

228 if processor is None:

229 processor = self.get_hf_processor()

@@ -249,7 +249,7 @@ class PixtralDummyInputsBuilder(BaseDummyInputsBuilder[PixtralProcessingInfo]):

249 self,

250 seq_len: int,

251 mm_counts: Mapping[str, int],

252 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

252 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

253 ) -> MultiModalDataDict:

254 num_images = mm_counts.get("image", 0)

255

@@ -270,7 +270,7 @@ class PixtralDummyInputsBuilder(BaseDummyInputsBuilder[PixtralProcessingInfo]):

270 self,

271 seq_len: int,

272 mm_counts: Mapping[str, int],

273 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

273 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

274 ) -> ProcessorInputs:

275 tokenizer = self.info.get_tokenizer()

276

@@ -342,11 +342,11 @@ class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo])

342

343 def _cached_apply_hf_processor(

344 self,

345 - prompt: Union[str, list[int]],

345 + prompt: str | list[int],

346 mm_data_items: MultiModalDataItems,

347 hf_processor_mm_kwargs: Mapping[str, object],

348 tokenization_kwargs: Mapping[str, object],

349 - mm_uuids: Optional[MultiModalUUIDDict] = None,

349 + mm_uuids: MultiModalUUIDDict | None = None,

350 ) -> tuple[list[int], MultiModalProcessingInfo, bool]:

351 prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(

352 prompt=prompt,

@@ -369,7 +369,7 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

369 merge_by_field_config = True

370

371 @classmethod

372 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

372 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

373 if modality.startswith("image"):

374 return None

375

@@ -420,7 +420,7 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

420

421 def _parse_and_validate_image_input(

422 self, **kwargs: object

423 - ) -> Optional[PixtralImagePixelInputs]:

423 + ) -> PixtralImagePixelInputs | None:

424 images = kwargs.pop("images", None)

425 if images is None:

426 return None

@@ -472,10 +472,10 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

472 self,

473 input_ids: torch.Tensor,

474 positions: torch.Tensor,

475 - intermediate_tensors: Optional[IntermediateTensors] = None,

476 - inputs_embeds: Optional[torch.Tensor] = None,

475 + intermediate_tensors: IntermediateTensors | None = None,

476 + inputs_embeds: torch.Tensor | None = None,

477 **kwargs: object,

478 - ) -> Union[torch.Tensor, IntermediateTensors]:

478 + ) -> torch.Tensor | IntermediateTensors:

479 """Run forward pass for pixtral."""

480 if intermediate_tensors is not None:

481 inputs_embeds = None

@@ -489,7 +489,7 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

489 def compute_logits(

490 self,

491 hidden_states: torch.Tensor,

492 - ) -> Optional[torch.Tensor]:

492 + ) -> torch.Tensor | None:

493 return self.language_model.compute_logits(hidden_states)

494

495 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

@@ -717,7 +717,7 @@ class Transformer(nn.Module):

717 self,

718 x: torch.Tensor,

719 mask: torch.Tensor,

720 - freqs_cis: Optional[torch.Tensor],

720 + freqs_cis: torch.Tensor | None,

721 ) -> torch.Tensor:

722 for layer in self.layers:

723 x = layer(x, mask=mask, freqs_cis=freqs_cis)

@@ -759,7 +759,7 @@ class VisionTransformer(nn.Module):

759

760 head_dim = self.args.hidden_size // self.args.num_attention_heads

761 assert head_dim % 2 == 0, "ROPE requires even head_dim"

762 - self._freqs_cis: Optional[torch.Tensor] = None

762 + self._freqs_cis: torch.Tensor | None = None

763

764 @property

765 def max_patches_per_side(self) -> int:

@@ -1015,7 +1015,7 @@ class PixtralHFMLP(nn.Module):

1015 def __init__(

1016 self,

1017 config: PixtralVisionConfig,

1018 - quant_config: Optional[QuantizationConfig] = None,

1018 + quant_config: QuantizationConfig | None = None,

1019 *,

1020 prefix: str = "",

1021 ) -> None:

@@ -1049,7 +1049,7 @@ class PixtralHFAttention(nn.Module):

1049 def __init__(

1050 self,

1051 config: PixtralVisionConfig,

1052 - quant_config: Optional[QuantizationConfig] = None,

1052 + quant_config: QuantizationConfig | None = None,

1053 *,

1054 prefix: str = "",

1055 ) -> None:

@@ -1084,7 +1084,7 @@ class PixtralHFAttention(nn.Module):

1084 hidden_states: torch.Tensor,

1085 attention_mask: torch.Tensor,

1086 position_embeddings: torch.Tensor,

1087 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

1087 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

1088 batch, patches, _ = hidden_states.size()

1089

1090 qkv_states, _ = self.qkv_proj(hidden_states)

@@ -1119,7 +1119,7 @@ class PixtralHFTransformerBlock(nn.Module):

1119 def __init__(

1120 self,

1121 config: PixtralVisionConfig,

1122 - quant_config: Optional[QuantizationConfig] = None,

1122 + quant_config: QuantizationConfig | None = None,

1123 *,

1124 prefix: str = "",

1125 ) -> None:

@@ -1155,9 +1155,9 @@ class PixtralHFTransformer(nn.Module):

1155 def __init__(

1156 self,

1157 config: PixtralVisionConfig,

1158 - quant_config: Optional[QuantizationConfig] = None,

1158 + quant_config: QuantizationConfig | None = None,

1159 *,

1160 - num_hidden_layers_override: Optional[int] = None,

1160 + num_hidden_layers_override: int | None = None,

1161 prefix: str = "",

1162 ) -> None:

1163 super().__init__()

@@ -1202,10 +1202,10 @@ class PixtralHFVisionModel(nn.Module):

1202 def __init__(

1203 self,

1204 config: PixtralVisionConfig,

1205 - quant_config: Optional[QuantizationConfig] = None,

1205 + quant_config: QuantizationConfig | None = None,

1206 *,

1207 - num_hidden_layers_override: Optional[int] = None,

1208 - require_post_norm: Optional[bool] = None,

1207 + num_hidden_layers_override: int | None = None,

1208 + require_post_norm: bool | None = None,

1209 prefix: str = "",

1210 ) -> None:

1211 super().__init__()

@@ -1247,8 +1247,8 @@ class PixtralHFVisionModel(nn.Module):

1247 self,

1248 pixel_values: list[torch.Tensor],

1249 *,

1250 - select_layers: Optional[list[int]] = None,

1251 - feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

1250 + select_layers: list[int] | None = None,

1251 + feature_select_strategy: VisionFeatureSelectStrategy | None = None,

1252 ) -> tuple[torch.Tensor, ...]:

1253 """

1254 Args:

vllm/model_executor/models/plamo2.py

+9 -9

@@ -4,7 +4,7 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import TYPE_CHECKING, Optional

7 +from typing import TYPE_CHECKING

8

9 if TYPE_CHECKING:

10 from vllm.attention.backends.abstract import AttentionBackend

@@ -504,7 +504,7 @@ class DenseMLP(nn.Module):

504 def __init__(

505 self,

506 config: Plamo2Config,

507 - quant_config: Optional[QuantizationConfig] = None,

507 + quant_config: QuantizationConfig | None = None,

508 prefix: str = "",

509 ) -> None:

510 super().__init__()

@@ -672,7 +672,7 @@ class Plamo2DecoderLayer(nn.Module):

672 self,

673 positions: torch.Tensor,

674 hidden_states: torch.Tensor,

675 - residual: Optional[torch.Tensor],

675 + residual: torch.Tensor | None,

676 **kwargs,

677 ):

678 if residual is None:

@@ -728,7 +728,7 @@ class Plamo2Decoder(torch.nn.Module):

728 self,

729 positions: torch.Tensor,

730 hidden_states: torch.Tensor,

731 - residual: Optional[torch.Tensor],

731 + residual: torch.Tensor | None,

732 ) -> torch.Tensor:

733 for layer in islice(self.layers, self.start_layer, self.end_layer):

734 hidden_states, residual = layer(

@@ -770,8 +770,8 @@ class Plamo2Model(torch.nn.Module):

770 self,

771 input_ids: torch.Tensor,

772 positions: torch.Tensor,

773 - intermediate_tensors: Optional[IntermediateTensors] = None,

774 - inputs_embeds: Optional[torch.Tensor] = None,

773 + intermediate_tensors: IntermediateTensors | None = None,

774 + inputs_embeds: torch.Tensor | None = None,

775 ) -> torch.Tensor:

776 if get_pp_group().is_first_rank:

777 if inputs_embeds is not None:

@@ -851,8 +851,8 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):

851 self,

852 input_ids: torch.Tensor,

853 positions: torch.Tensor,

854 - intermediate_tensors: Optional[IntermediateTensors] = None,

855 - inputs_embeds: Optional[torch.Tensor] = None,

854 + intermediate_tensors: IntermediateTensors | None = None,

855 + inputs_embeds: torch.Tensor | None = None,

856 **kwargs,

857 ):

858 hidden_states = self.model(

@@ -901,7 +901,7 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):

901 def compute_logits(

902 self,

903 hidden_states: torch.Tensor,

904 - ) -> Optional[torch.Tensor]:

904 + ) -> torch.Tensor | None:

905 logits = self.logits_processor(self.lm_head, hidden_states)

906 return logits

907

vllm/model_executor/models/qwen.py

+15 -15

@@ -10,7 +10,7 @@

10 import json

11 from collections.abc import Iterable

12 from itertools import islice

13 -from typing import Any, Optional, Union

13 +from typing import Any

14

15 import torch

16 from torch import nn

@@ -55,7 +55,7 @@ class QWenMLP(nn.Module):

55 hidden_size: int,

56 intermediate_size: int,

57 hidden_act: str = "silu",

58 - quant_config: Optional[QuantizationConfig] = None,

58 + quant_config: QuantizationConfig | None = None,

59 ):

60 super().__init__()

61 self.gate_up_proj = MergedColumnParallelLinear(

@@ -84,9 +84,9 @@ class QWenAttention(nn.Module):

84 num_heads: int,

85 max_position_embeddings: int,

86 rope_theta: float = 10000,

87 - rope_scaling: Optional[dict[str, Any]] = None,

88 - cache_config: Optional[CacheConfig] = None,

89 - quant_config: Optional[QuantizationConfig] = None,

87 + rope_scaling: dict[str, Any] | None = None,

88 + cache_config: CacheConfig | None = None,

89 + quant_config: QuantizationConfig | None = None,

90 prefix: str = "",

91 ):

92 super().__init__()

@@ -144,8 +144,8 @@ class QWenBlock(nn.Module):

144 def __init__(

145 self,

146 config: PretrainedConfig,

147 - cache_config: Optional[CacheConfig] = None,

148 - quant_config: Optional[QuantizationConfig] = None,

147 + cache_config: CacheConfig | None = None,

148 + quant_config: QuantizationConfig | None = None,

149 prefix: str = "",

150 ):

151 super().__init__()

@@ -174,7 +174,7 @@ class QWenBlock(nn.Module):

174 self,

175 positions: torch.Tensor,

176 hidden_states: torch.Tensor,

177 - residual: Optional[torch.Tensor],

177 + residual: torch.Tensor | None,

178 ) -> tuple[torch.Tensor, torch.Tensor]:

179 # Self Attention

180 if residual is None:

@@ -226,9 +226,9 @@ class QWenModel(nn.Module):

226 self,

227 input_ids: torch.Tensor,

228 positions: torch.Tensor,

229 - intermediate_tensors: Optional[IntermediateTensors],

230 - inputs_embeds: Optional[torch.Tensor] = None,

231 - ) -> Union[torch.Tensor, IntermediateTensors]:

229 + intermediate_tensors: IntermediateTensors | None,

230 + inputs_embeds: torch.Tensor | None = None,

231 + ) -> torch.Tensor | IntermediateTensors:

232 if get_pp_group().is_first_rank:

233 if inputs_embeds is not None:

234 hidden_states = inputs_embeds

@@ -288,7 +288,7 @@ class QWenBaseModel(nn.Module):

288 def compute_logits(

289 self,

290 hidden_states: torch.Tensor,

291 - ) -> Optional[torch.Tensor]:

291 + ) -> torch.Tensor | None:

292 logits = self.logits_processor(self.lm_head, hidden_states)

293 return logits

294

@@ -357,9 +357,9 @@ class QWenLMHeadModel(QWenBaseModel, SupportsPP, SupportsLoRA):

357 self,

358 input_ids: torch.Tensor,

359 positions: torch.Tensor,

360 - intermediate_tensors: Optional[IntermediateTensors] = None,

361 - inputs_embeds: Optional[torch.Tensor] = None,

362 - ) -> Union[torch.Tensor, IntermediateTensors]:

360 + intermediate_tensors: IntermediateTensors | None = None,

361 + inputs_embeds: torch.Tensor | None = None,

362 + ) -> torch.Tensor | IntermediateTensors:

363 hidden_states = self.transformer(

364 input_ids, positions, intermediate_tensors, inputs_embeds

365 )

vllm/model_executor/models/qwen2.py

+16 -16

@@ -27,7 +27,7 @@

27

28 from collections.abc import Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import torch

33 from torch import nn

@@ -77,7 +77,7 @@ class Qwen2MLP(nn.Module):

77 hidden_size: int,

78 intermediate_size: int,

79 hidden_act: str,

80 - quant_config: Optional[QuantizationConfig] = None,

80 + quant_config: QuantizationConfig | None = None,

81 prefix: str = "",

82 ) -> None:

83 super().__init__()

@@ -116,12 +116,12 @@ class Qwen2Attention(nn.Module):

116 num_kv_heads: int,

117 max_position: int = 4096 * 32,

118 rope_theta: float = 10000,

119 - cache_config: Optional[CacheConfig] = None,

120 - quant_config: Optional[QuantizationConfig] = None,

121 - rope_scaling: Optional[tuple] = None,

119 + cache_config: CacheConfig | None = None,

120 + quant_config: QuantizationConfig | None = None,

121 + rope_scaling: tuple | None = None,

122 prefix: str = "",

123 attn_type: str = AttentionType.DECODER,

124 - dual_chunk_attention_config: Optional[dict[str, Any]] = None,

124 + dual_chunk_attention_config: dict[str, Any] | None = None,

125 ) -> None:

126 super().__init__()

127 self.hidden_size = hidden_size

@@ -210,8 +210,8 @@ class Qwen2DecoderLayer(nn.Module):

210 def __init__(

211 self,

212 config: Qwen2Config,

213 - cache_config: Optional[CacheConfig] = None,

214 - quant_config: Optional[QuantizationConfig] = None,

213 + cache_config: CacheConfig | None = None,

214 + quant_config: QuantizationConfig | None = None,

215 prefix: str = "",

216 ) -> None:

217 super().__init__()

@@ -261,7 +261,7 @@ class Qwen2DecoderLayer(nn.Module):

261 self,

262 positions: torch.Tensor,

263 hidden_states: torch.Tensor,

264 - residual: Optional[torch.Tensor],

264 + residual: torch.Tensor | None,

265 ) -> tuple[torch.Tensor, torch.Tensor]:

266 # Self Attention

267 if residual is None:

@@ -362,9 +362,9 @@ class Qwen2Model(nn.Module):

362 self,

363 input_ids: torch.Tensor,

364 positions: torch.Tensor,

365 - intermediate_tensors: Optional[IntermediateTensors] = None,

366 - inputs_embeds: Optional[torch.Tensor] = None,

367 - ) -> Union[torch.Tensor, IntermediateTensors]:

365 + intermediate_tensors: IntermediateTensors | None = None,

366 + inputs_embeds: torch.Tensor | None = None,

367 + ) -> torch.Tensor | IntermediateTensors:

368 if get_pp_group().is_first_rank:

369 if inputs_embeds is not None:

370 hidden_states = inputs_embeds

@@ -520,9 +520,9 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

520 self,

521 input_ids: torch.Tensor,

522 positions: torch.Tensor,

523 - intermediate_tensors: Optional[IntermediateTensors] = None,

524 - inputs_embeds: Optional[torch.Tensor] = None,

525 - ) -> Union[torch.Tensor, IntermediateTensors]:

523 + intermediate_tensors: IntermediateTensors | None = None,

524 + inputs_embeds: torch.Tensor | None = None,

525 + ) -> torch.Tensor | IntermediateTensors:

526 hidden_states = self.model(

527 input_ids, positions, intermediate_tensors, inputs_embeds

528 )

@@ -531,7 +531,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

531 def compute_logits(

532 self,

533 hidden_states: torch.Tensor,

534 - ) -> Optional[torch.Tensor]:

534 + ) -> torch.Tensor | None:

535 logits = self.logits_processor(self.lm_head, hidden_states)

536 return logits

537

vllm/model_executor/models/qwen2_5_omni_thinker.py

+24 -24

@@ -22,10 +22,10 @@

22 # limitations under the License.

23 """Inference-only Qwen2.5-Omni model (thinker part)."""

24

25 -from collections.abc import Iterable, Mapping, Sequence

25 +from collections.abc import Callable, Iterable, Mapping, Sequence

26 from copy import copy

27 from functools import partial

28 -from typing import Annotated, Any, Callable, Literal, Optional, Union

28 +from typing import Annotated, Any, Literal

29

30 import torch

31 import torch.nn as nn

@@ -125,7 +125,7 @@ class Qwen2_5OmniAudioFeatureInputs(TensorSchema):

125

126 type: Literal["audio_features"]

127 input_features: Annotated[

128 - Union[torch.Tensor, list[torch.Tensor]],

128 + torch.Tensor | list[torch.Tensor],

129 TensorShape("nmb", "tsl"),

130 ]

131

@@ -191,7 +191,7 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):

191

192 def _parse_audio_data(

193 self,

194 - data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],

194 + data: dict[str, torch.Tensor] | ModalityData[ImageItem],

195 ) -> ModalityDataItems[Any, Any]:

196 if isinstance(data, dict):

197 return DictEmbeddingItems(

@@ -225,7 +225,7 @@ class Qwen2_5OmniThinkerProcessingInfo(

225 assert isinstance(feature_extractor, WhisperFeatureExtractor)

226 return feature_extractor

227

228 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

228 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

229 return {"audio": None, "image": None, "video": None}

230

231

@@ -253,7 +253,7 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(

253 self,

254 seq_len: int,

255 mm_counts: Mapping[str, int],

256 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

256 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

257 ) -> MultiModalDataDict:

258 num_audios = mm_counts.get("audio", 0)

259 num_images = mm_counts.get("image", 0)

@@ -420,7 +420,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(

420 cls,

421 thinker_config: PretrainedConfig,

422 audio_len: int,

423 - video_grid_thw: Union[list[int], torch.Tensor],

423 + video_grid_thw: list[int] | torch.Tensor,

424 video_second_per_grid_t: float,

425 ) -> list[int]:

426 """Get video prompt updates when `use_audio_in_video` is True.

@@ -580,7 +580,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(

580

581 def _apply_hf_processor_main(

582 self,

583 - prompt: Union[str, list[int]],

583 + prompt: str | list[int],

584 mm_items: MultiModalDataItems,

585 hf_processor_mm_kwargs: Mapping[str, object],

586 tokenization_kwargs: Mapping[str, object],

@@ -665,7 +665,7 @@ class Qwen2_5OmniConditionalGenerationMixin:

665

666 def _parse_and_validate_audio_input(

667 self, **kwargs: object

668 - ) -> Optional[Qwen2_5OmniAudioFeatureInputs]:

668 + ) -> Qwen2_5OmniAudioFeatureInputs | None:

669 input_audio_features = kwargs.pop("input_audio_features", None)

670 audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)

671 feature_attention_mask = kwargs.pop("feature_attention_mask", None)

@@ -693,7 +693,7 @@ class Qwen2_5OmniConditionalGenerationMixin:

693 def _parse_and_validate_image_input(

694 self,

695 **kwargs: dict[str, Any],

696 - ) -> Optional[Qwen2_5_VLImageInputs]:

696 + ) -> Qwen2_5_VLImageInputs | None:

697 pixel_values = kwargs.pop("pixel_values", None)

698 image_embeds = kwargs.pop("image_embeds", None)

699 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -743,7 +743,7 @@ class Qwen2_5OmniConditionalGenerationMixin:

743 def _parse_and_validate_video_input(

744 self,

745 **kwargs: dict[str, Any],

746 - ) -> Optional[Qwen2_5_VLVideoInputs]:

746 + ) -> Qwen2_5_VLVideoInputs | None:

747 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

748 video_embeds = kwargs.pop("video_embeds", None)

749 video_grid_thw = kwargs.pop("video_grid_thw", None)

@@ -892,7 +892,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

892 }

893

894 @classmethod

895 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

895 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

896 if modality.startswith("image"):

897 return "<|vision_start|><|IMAGE|><|vision_end|>"

898 if modality.startswith("video"):

@@ -991,12 +991,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

991 cls,

992 input_tokens: list[int],

993 hf_config: PretrainedConfig,

994 - image_grid_thw: Union[list[list[int]], torch.Tensor],

995 - video_grid_thw: Union[list[list[int]], torch.Tensor],

996 - second_per_grid_ts: Optional[list[float]] = None,

994 + image_grid_thw: list[list[int]] | torch.Tensor,

995 + video_grid_thw: list[list[int]] | torch.Tensor,

996 + second_per_grid_ts: list[float] | None = None,

997 context_len: int = 0,

998 - seq_len: Optional[int] = None,

999 - audio_feature_lengths: Optional[torch.Tensor] = None,

998 + seq_len: int | None = None,

999 + audio_feature_lengths: torch.Tensor | None = None,

1000 use_audio_in_video: bool = False,

1001 ) -> tuple[torch.Tensor, int]:

1002 """Get mrope input positions and delta value (Qwen2.5-Omni version).

@@ -1225,9 +1225,9 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

1225 def get_input_embeddings(

1226 self,

1227 input_ids: torch.Tensor,

1228 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1228 + multimodal_embeddings: MultiModalEmbeddings | None = None,

1229 *,

1230 - is_multimodal: Optional[torch.Tensor] = None,

1230 + is_multimodal: torch.Tensor | None = None,

1231 handle_oov_mm_token: bool = False,

1232 ) -> torch.Tensor:

1233 # This is to satisfy the type checker for each overload

@@ -1241,7 +1241,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

1241 handle_oov_mm_token=handle_oov_mm_token,

1242 )

1243

1244 - def get_multimodal_embeddings_v0(self, **kwargs: object) -> Optional[NestedTensors]:

1244 + def get_multimodal_embeddings_v0(self, **kwargs: object) -> NestedTensors | None:

1245 audio_input = self._parse_and_validate_audio_input(**kwargs)

1246 image_input = self._parse_and_validate_image_input(**kwargs)

1247 video_input = self._parse_and_validate_video_input(**kwargs)

@@ -1266,10 +1266,10 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

1266 self,

1267 input_ids: torch.Tensor,

1268 positions: torch.Tensor,

1269 - intermediate_tensors: Optional[IntermediateTensors] = None,

1270 - inputs_embeds: Optional[torch.Tensor] = None,

1269 + intermediate_tensors: IntermediateTensors | None = None,

1270 + inputs_embeds: torch.Tensor | None = None,

1271 **kwargs: object,

1272 - ) -> Union[torch.Tensor, IntermediateTensors]:

1272 + ) -> torch.Tensor | IntermediateTensors:

1273 if intermediate_tensors is not None:

1274 inputs_embeds = None

1275

@@ -1281,7 +1281,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

1281 def compute_logits(

1282 self,

1283 hidden_states: torch.Tensor,

1284 - ) -> Optional[torch.Tensor]:

1284 + ) -> torch.Tensor | None:

1285 return self.language_model.compute_logits(hidden_states)

1286

1287 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/qwen2_5_vl.py

+32 -32

@@ -26,9 +26,9 @@

26 # limitations under the License.

27 """Inference-only Qwen2.5-VL model compatible with HuggingFace weights."""

28

29 -from collections.abc import Iterable, Mapping, Sequence

29 +from collections.abc import Callable, Iterable, Mapping, Sequence

30 from functools import lru_cache, partial

31 -from typing import Annotated, Any, Callable, Literal, Optional, Union

31 +from typing import Annotated, Any, Literal, TypeAlias

32

33 import torch

34 import torch.nn as nn

@@ -161,9 +161,9 @@ class Qwen2_5_VLImageEmbeddingInputs(TensorSchema):

161 ]

162

163

164 -Qwen2_5_VLImageInputs = Union[

165 - Qwen2_5_VLImagePixelInputs, Qwen2_5_VLImageEmbeddingInputs

166 -]

164 +Qwen2_5_VLImageInputs: TypeAlias = (

165 + Qwen2_5_VLImagePixelInputs | Qwen2_5_VLImageEmbeddingInputs

166 +)

167

168

169 class Qwen2_5_VLVideoPixelInputs(TensorSchema):

@@ -197,7 +197,7 @@ class Qwen2_5_VLVideoPixelInputs(TensorSchema):

197 ]

198

199 second_per_grid_ts: Annotated[

200 - Optional[torch.Tensor],

200 + torch.Tensor | None,

201 TensorShape("nv"),

202 ]

203

@@ -231,9 +231,9 @@ class Qwen2_5_VLVideoEmbeddingInputs(TensorSchema):

231 ]

232

233

234 -Qwen2_5_VLVideoInputs = Union[

235 - Qwen2_5_VLVideoPixelInputs, Qwen2_5_VLVideoEmbeddingInputs

236 -]

234 +Qwen2_5_VLVideoInputs: TypeAlias = (

235 + Qwen2_5_VLVideoPixelInputs | Qwen2_5_VLVideoEmbeddingInputs

236 +)

237

238 # === Vision Encoder === #

239

@@ -245,7 +245,7 @@ class Qwen2_5_VisionMLP(nn.Module):

245 hidden_features: int,

246 bias: bool = False,

247 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

248 - quant_config: Optional[QuantizationConfig] = None,

248 + quant_config: QuantizationConfig | None = None,

249 prefix: str = "",

250 use_data_parallel: bool = False,

251 ):

@@ -301,7 +301,7 @@ class Qwen2_5_VisionAttention(nn.Module):

301 embed_dim: int,

302 num_heads: int,

303 projection_size: int,

304 - quant_config: Optional[QuantizationConfig] = None,

304 + quant_config: QuantizationConfig | None = None,

305 prefix: str = "",

306 use_data_parallel: bool = False,

307 attn_backend: _Backend = _Backend.TORCH_SDPA,

@@ -386,8 +386,8 @@ class Qwen2_5_VisionAttention(nn.Module):

386 x: torch.Tensor,

387 cu_seqlens: torch.Tensor,

388 rotary_pos_emb: torch.Tensor,

389 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

390 - seqlens: Optional[list[int]] = None, # Only used for xFormers

389 + max_seqlen: int | None = None, # Only used for Flash Attention

390 + seqlens: list[int] | None = None, # Only used for xFormers

391 ) -> torch.Tensor:

392 # [s, b, c] --> [s, b, head * 3 * head_dim]

393 x, _ = self.qkv(x)

@@ -466,8 +466,8 @@ class Qwen2_5_VisionBlock(nn.Module):

466 num_heads: int,

467 mlp_hidden_dim: int,

468 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

469 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

470 - quant_config: Optional[QuantizationConfig] = None,

469 + norm_layer: Callable[[int], nn.Module] | None = None,

470 + quant_config: QuantizationConfig | None = None,

471 prefix: str = "",

472 use_data_parallel: bool = False,

473 attn_backend: _Backend = _Backend.TORCH_SDPA,

@@ -503,8 +503,8 @@ class Qwen2_5_VisionBlock(nn.Module):

503 x: torch.Tensor,

504 cu_seqlens: torch.Tensor,

505 rotary_pos_emb: torch.Tensor,

506 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

507 - seqlens: Optional[list[int]] = None, # Only used for xFormers

506 + max_seqlen: int | None = None, # Only used for Flash Attention

507 + seqlens: list[int] | None = None, # Only used for xFormers

508 ) -> torch.Tensor:

509 x_attn = self.attn(

510 self.norm1(x),

@@ -552,9 +552,9 @@ class Qwen2_5_VisionPatchMerger(nn.Module):

552 self,

553 d_model: int,

554 context_dim: int,

555 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

555 + norm_layer: Callable[[int], nn.Module] | None = None,

556 spatial_merge_size: int = 2,

557 - quant_config: Optional[QuantizationConfig] = None,

557 + quant_config: QuantizationConfig | None = None,

558 prefix: str = "",

559 use_data_parallel: bool = False,

560 ) -> None:

@@ -634,7 +634,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

634 self,

635 vision_config: Qwen2_5_VLVisionConfig,

636 norm_eps: float = 1e-6,

637 - quant_config: Optional[QuantizationConfig] = None,

637 + quant_config: QuantizationConfig | None = None,

638 prefix: str = "",

639 use_data_parallel: bool = False,

640 ) -> None:

@@ -815,7 +815,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

815 def compute_attn_mask_seqlen(

816 self,

817 cu_seqlens: torch.Tensor,

818 - ) -> tuple[Optional[int], Optional[list[int]]]:

818 + ) -> tuple[int | None, list[int] | None]:

819 max_seqlen, seqlens = None, None

820 if (

821 self.attn_backend == _Backend.FLASH_ATTN

@@ -1080,12 +1080,12 @@ class Qwen2_5_VLForConditionalGeneration(

1080 cls,

1081 input_tokens: list[int],

1082 hf_config: PretrainedConfig,

1083 - image_grid_thw: Union[list[list[int]], torch.Tensor],

1084 - video_grid_thw: Union[list[list[int]], torch.Tensor],

1083 + image_grid_thw: list[list[int]] | torch.Tensor,

1084 + video_grid_thw: list[list[int]] | torch.Tensor,

1085 second_per_grid_ts: list[float],

1086 context_len: int = 0,

1087 - seq_len: Optional[int] = None,

1088 - audio_feature_lengths: Optional[torch.Tensor] = None,

1087 + seq_len: int | None = None,

1088 + audio_feature_lengths: torch.Tensor | None = None,

1089 use_audio_in_video: bool = False,

1090 ) -> tuple[torch.Tensor, int]:

1091 """Get mrope input positions and delta value."""

@@ -1202,7 +1202,7 @@ class Qwen2_5_VLForConditionalGeneration(

1202 return llm_positions, mrope_position_delta

1203

1204 @classmethod

1205 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1205 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1206 if modality.startswith("image"):

1207 return "<|vision_start|><|image_pad|><|vision_end|>"

1208 if modality.startswith("video"):

@@ -1273,7 +1273,7 @@ class Qwen2_5_VLForConditionalGeneration(

1273

1274 def _parse_and_validate_image_input(

1275 self, **kwargs: object

1276 - ) -> Optional[Qwen2_5_VLImageInputs]:

1276 + ) -> Qwen2_5_VLImageInputs | None:

1277 pixel_values = kwargs.pop("pixel_values", None)

1278 image_embeds = kwargs.pop("image_embeds", None)

1279 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -1311,7 +1311,7 @@ class Qwen2_5_VLForConditionalGeneration(

1311

1312 def _parse_and_validate_video_input(

1313 self, **kwargs: object

1314 - ) -> Optional[Qwen2_5_VLVideoInputs]:

1314 + ) -> Qwen2_5_VLVideoInputs | None:

1315 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1316 video_embeds = kwargs.pop("video_embeds", None)

1317 video_grid_thw = kwargs.pop("video_grid_thw", None)

@@ -1605,10 +1605,10 @@ class Qwen2_5_VLForConditionalGeneration(

1605 self,

1606 input_ids: torch.Tensor,

1607 positions: torch.Tensor,

1608 - intermediate_tensors: Optional[IntermediateTensors] = None,

1609 - inputs_embeds: Optional[torch.Tensor] = None,

1608 + intermediate_tensors: IntermediateTensors | None = None,

1609 + inputs_embeds: torch.Tensor | None = None,

1610 **kwargs: object,

1611 - ) -> Union[torch.Tensor, IntermediateTensors]:

1611 + ) -> torch.Tensor | IntermediateTensors:

1612 """Run forward pass for Qwen2.5-VL.

1613

1614 Args:

@@ -1634,7 +1634,7 @@ class Qwen2_5_VLForConditionalGeneration(

1634 def compute_logits(

1635 self,

1636 hidden_states: torch.Tensor,

1637 - ) -> Optional[torch.Tensor]:

1637 + ) -> torch.Tensor | None:

1638 return self.language_model.compute_logits(hidden_states)

1639

1640 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/qwen2_audio.py

+14 -14

@@ -24,7 +24,7 @@

24 """Inference-only Qwen2-Audio model compatible with HuggingFace weights."""

25

26 from collections.abc import Iterable, Mapping, Sequence

27 -from typing import Annotated, Any, Literal, Optional, Union

27 +from typing import Annotated, Any, Literal, TypeAlias

28

29 import torch

30 import torch.nn as nn

@@ -78,7 +78,7 @@ class Qwen2AudioFeatureInputs(TensorSchema):

78

79 type: Literal["audio_features"]

80 input_features: Annotated[

81 - Union[torch.Tensor, list[torch.Tensor]],

81 + torch.Tensor | list[torch.Tensor],

82 TensorShape("na", "nmb", 3000),

83 ]

84

@@ -105,7 +105,7 @@ class Qwen2AudioEmbeddingInputs(TensorSchema):

105 ]

106

107

108 -Qwen2AudioInputs = Union[Qwen2AudioFeatureInputs, Qwen2AudioEmbeddingInputs]

108 +Qwen2AudioInputs: TypeAlias = Qwen2AudioFeatureInputs | Qwen2AudioEmbeddingInputs

109

110 # === Audio Encoder === #

111

@@ -140,7 +140,7 @@ class Qwen2AudioProcessingInfo(BaseProcessingInfo):

140 assert isinstance(feature_extractor, WhisperFeatureExtractor)

141 return feature_extractor

142

143 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

143 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

144 return {"audio": None}

145

146

@@ -157,7 +157,7 @@ class Qwen2AudioDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2AudioProcessingIn

157 self,

158 seq_len: int,

159 mm_counts: Mapping[str, int],

160 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

160 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

161 ) -> MultiModalDataDict:

162 feature_extractor = self.info.get_feature_extractor()

163

@@ -185,8 +185,8 @@ def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):

185 class Qwen2AudioMultiModalDataParser(MultiModalDataParser):

186 def _parse_audio_data(

187 self,

188 - data: Union[dict[str, torch.Tensor], ModalityData[AudioItem]],

189 - ) -> Optional[ModalityDataItems[Any, Any]]:

188 + data: dict[str, torch.Tensor] | ModalityData[AudioItem],

189 + ) -> ModalityDataItems[Any, Any] | None:

190 if isinstance(data, dict):

191 return DictEmbeddingItems(

192 data,

@@ -314,7 +314,7 @@ class Qwen2AudioMultiModalProcessor(BaseMultiModalProcessor[Qwen2AudioProcessing

314 )

315 class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

316 @classmethod

317 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

317 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

318 if modality.startswith("audio"):

319 return f"Audio {i}: <|audio_bos|><|AUDIO|><|audio_eos|>"

320

@@ -358,7 +358,7 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports

358

359 def _parse_and_validate_audio_input(

360 self, **kwargs: object

361 - ) -> Optional[Qwen2AudioInputs]:

361 + ) -> Qwen2AudioInputs | None:

362 input_features = kwargs.pop("input_features", None)

363 audio_embeds = kwargs.pop("audio_embeds", None)

364 feature_attention_mask = kwargs.pop("feature_attention_mask", None)

@@ -395,7 +395,7 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports

395

396 def _process_audio_input(

397 self, audio_input: Qwen2AudioInputs

398 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

398 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

399 if audio_input["type"] == "audio_embeds":

400 audio_embeds = audio_input["audio_embeds"]

401 return tuple(audio_embeds)

@@ -471,10 +471,10 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports

471 self,

472 input_ids: torch.Tensor,

473 positions: torch.Tensor,

474 - intermediate_tensors: Optional[IntermediateTensors] = None,

475 - inputs_embeds: Optional[torch.Tensor] = None,

474 + intermediate_tensors: IntermediateTensors | None = None,

475 + inputs_embeds: torch.Tensor | None = None,

476 **kwargs: object,

477 - ) -> Union[torch.Tensor, IntermediateTensors]:

477 + ) -> torch.Tensor | IntermediateTensors:

478 if intermediate_tensors is not None:

479 inputs_embeds = None

480

@@ -486,7 +486,7 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports

486 def compute_logits(

487 self,

488 hidden_states: torch.Tensor,

489 - ) -> Optional[torch.Tensor]:

489 + ) -> torch.Tensor | None:

490 return self.language_model.compute_logits(hidden_states)

491

492 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/qwen2_moe.py

+19 -24

@@ -27,7 +27,7 @@

27

28 from collections.abc import Iterable

29 from itertools import islice

30 -from typing import Any, Optional, Union

30 +from typing import Any

31

32 import torch

33 import torch.nn.functional as F

@@ -77,9 +77,9 @@ class Qwen2MoeMLP(nn.Module):

77 hidden_size: int,

78 intermediate_size: int,

79 hidden_act: str,

80 - quant_config: Optional[QuantizationConfig] = None,

80 + quant_config: QuantizationConfig | None = None,

81 reduce_results: bool = True,

82 - expert_gate: Optional[torch.nn.Linear] = None,

82 + expert_gate: torch.nn.Linear | None = None,

83 prefix: str = "",

84 ) -> None:

85 super().__init__()

@@ -120,7 +120,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):

120 def __init__(

121 self,

122 config: Qwen2MoeConfig,

123 - quant_config: Optional[QuantizationConfig] = None,

123 + quant_config: QuantizationConfig | None = None,

124 prefix: str = "",

125 ):

126 super().__init__()

@@ -195,12 +195,12 @@ class Qwen2MoeAttention(nn.Module):

195 num_heads: int,

196 num_kv_heads: int,

197 rope_theta: float = 10000,

198 - rope_scaling: Optional[dict[str, Any]] = None,

198 + rope_scaling: dict[str, Any] | None = None,

199 max_position_embeddings: int = 8192,

200 - cache_config: Optional[CacheConfig] = None,

201 - quant_config: Optional[QuantizationConfig] = None,

200 + cache_config: CacheConfig | None = None,

201 + quant_config: QuantizationConfig | None = None,

202 prefix: str = "",

203 - dual_chunk_attention_config: Optional[dict[str, Any]] = None,

203 + dual_chunk_attention_config: dict[str, Any] | None = None,

204 ) -> None:

205 super().__init__()

206 self.hidden_size = hidden_size

@@ -285,8 +285,8 @@ class Qwen2MoeDecoderLayer(nn.Module):

285 def __init__(

286 self,

287 config: Qwen2MoeConfig,

288 - cache_config: Optional[CacheConfig] = None,

289 - quant_config: Optional[QuantizationConfig] = None,

288 + cache_config: CacheConfig | None = None,

289 + quant_config: QuantizationConfig | None = None,

290 prefix: str = "",

291 ) -> None:

292 super().__init__()

@@ -339,7 +339,7 @@ class Qwen2MoeDecoderLayer(nn.Module):

339 self,

340 positions: torch.Tensor,

341 hidden_states: torch.Tensor,

342 - residual: Optional[torch.Tensor],

342 + residual: torch.Tensor | None,

343 ) -> torch.Tensor:

344 # Self Attention

345 if residual is None:

@@ -396,9 +396,9 @@ class Qwen2MoeModel(nn.Module):

396 self,

397 input_ids: torch.Tensor,

398 positions: torch.Tensor,

399 - intermediate_tensors: Optional[IntermediateTensors] = None,

400 - inputs_embeds: Optional[torch.Tensor] = None,

401 - ) -> Union[torch.Tensor, IntermediateTensors]:

399 + intermediate_tensors: IntermediateTensors | None = None,

400 + inputs_embeds: torch.Tensor | None = None,

401 + ) -> torch.Tensor | IntermediateTensors:

402 if get_pp_group().is_first_rank:

403 if inputs_embeds is not None:

404 hidden_states = inputs_embeds

@@ -548,12 +548,7 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

548 getattr(config, "mlp_only_layers", [])

549 or config.shared_expert_intermediate_size > 0

550 ):

551 - self.packed_modules_mapping["gate_up_proj"] = (

552 - [

553 - "gate_proj",

554 - "up_proj",

555 - ],

556 - )

551 + self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]

552

553 self.model = Qwen2MoeModel(

554 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

@@ -578,9 +573,9 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

573 self,

574 input_ids: torch.Tensor,

575 positions: torch.Tensor,

581 - intermediate_tensors: Optional[IntermediateTensors] = None,

582 - inputs_embeds: Optional[torch.Tensor] = None,

583 - ) -> Union[torch.Tensor, IntermediateTensors]:

576 + intermediate_tensors: IntermediateTensors | None = None,

577 + inputs_embeds: torch.Tensor | None = None,

578 + ) -> torch.Tensor | IntermediateTensors:

579 hidden_states = self.model(

580 input_ids, positions, intermediate_tensors, inputs_embeds

581 )

@@ -589,7 +584,7 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

584 def compute_logits(

585 self,

586 hidden_states: torch.Tensor,

592 - ) -> Optional[torch.Tensor]:

587 + ) -> torch.Tensor | None:

588 logits = self.logits_processor(self.lm_head, hidden_states)

589 return logits

590

vllm/model_executor/models/qwen2_rm.py

+3 -4

@@ -8,7 +8,6 @@

8 """Inference-only Qwen2-RM model compatible with HuggingFace weights."""

9

10 from collections.abc import Iterable

11 -from typing import Optional, Union

11

12 import torch

13 from torch import nn

@@ -83,9 +82,9 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):

82 self,

83 input_ids: torch.Tensor,

84 positions: torch.Tensor,

86 - intermediate_tensors: Optional[IntermediateTensors] = None,

87 - inputs_embeds: Optional[torch.Tensor] = None,

88 - ) -> Union[torch.Tensor, IntermediateTensors]:

85 + intermediate_tensors: IntermediateTensors | None = None,

86 + inputs_embeds: torch.Tensor | None = None,

87 + ) -> torch.Tensor | IntermediateTensors:

88 hidden_states = self.model(

89 input_ids, positions, intermediate_tensors, inputs_embeds

90 )

vllm/model_executor/models/qwen2_vl.py

+38 -38

@@ -25,9 +25,9 @@

25 # limitations under the License.

26 """Inference-only Qwen2-VL model compatible with HuggingFace weights."""

27

28 -from collections.abc import Iterable, Mapping, Sequence

28 +from collections.abc import Callable, Iterable, Mapping, Sequence

29 from functools import partial

30 -from typing import Annotated, Any, Callable, Literal, Optional, Union

30 +from typing import Annotated, Any, Literal, TypeAlias

31

32 import torch

33 import torch.nn as nn

@@ -167,7 +167,7 @@ class Qwen2VLImageEmbeddingInputs(TensorSchema):

167 ]

168

169

170 -Qwen2VLImageInputs = Union[Qwen2VLImagePixelInputs, Qwen2VLImageEmbeddingInputs]

170 +Qwen2VLImageInputs: TypeAlias = Qwen2VLImagePixelInputs | Qwen2VLImageEmbeddingInputs

171

172

173 class Qwen2VLVideoPixelInputs(TensorSchema):

@@ -228,7 +228,7 @@ class Qwen2VLVideoEmbeddingInputs(TensorSchema):

228 ]

229

230

231 -Qwen2VLVideoInputs = Union[Qwen2VLVideoPixelInputs, Qwen2VLVideoEmbeddingInputs]

231 +Qwen2VLVideoInputs: TypeAlias = Qwen2VLVideoPixelInputs | Qwen2VLVideoEmbeddingInputs

232

233 # === Vision Encoder === #

234

@@ -239,7 +239,7 @@ class Qwen2VisionMLP(nn.Module):

239 in_features: int,

240 hidden_features: int,

241 act_layer: type[nn.Module] = QuickGELU,

242 - quant_config: Optional[QuantizationConfig] = None,

242 + quant_config: QuantizationConfig | None = None,

243 prefix: str = "",

244 use_data_parallel: bool = False,

245 ):

@@ -317,7 +317,7 @@ class Qwen2VisionAttention(nn.Module):

317 embed_dim: int,

318 num_heads: int,

319 projection_size: int,

320 - quant_config: Optional[QuantizationConfig] = None,

320 + quant_config: QuantizationConfig | None = None,

321 prefix: str = "",

322 use_data_parallel: bool = False,

323 ) -> None:

@@ -413,8 +413,8 @@ class Qwen2VisionAttention(nn.Module):

413 x: torch.Tensor,

414 cu_seqlens: torch.Tensor,

415 rotary_pos_emb: torch.Tensor,

416 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

417 - seqlens: Optional[list[int]] = None, # Only used for xFormers

416 + max_seqlen: int | None = None, # Only used for Flash Attention

417 + seqlens: list[int] | None = None, # Only used for xFormers

418 ) -> torch.Tensor:

419 # [s, b, c] --> [s, b, 3 * head * head_dim]

420 x, _ = self.qkv(x)

@@ -493,8 +493,8 @@ class Qwen2VisionBlock(nn.Module):

493 num_heads: int,

494 mlp_ratio: float,

495 act_layer: type[nn.Module] = QuickGELU,

496 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

497 - quant_config: Optional[QuantizationConfig] = None,

496 + norm_layer: Callable[[int], nn.Module] | None = None,

497 + quant_config: QuantizationConfig | None = None,

498 prefix: str = "",

499 use_data_parallel: bool = False,

500 ) -> None:

@@ -527,8 +527,8 @@ class Qwen2VisionBlock(nn.Module):

527 x: torch.Tensor,

528 cu_seqlens: torch.Tensor,

529 rotary_pos_emb: torch.Tensor,

530 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

531 - seqlens: Optional[list[int]] = None, # Only used for xFormers

530 + max_seqlen: int | None = None, # Only used for Flash Attention

531 + seqlens: list[int] | None = None, # Only used for xFormers

532 ) -> torch.Tensor:

533 x = x + self.attn(

534 self.norm1(x),

@@ -576,9 +576,9 @@ class Qwen2VisionPatchMerger(nn.Module):

576 self,

577 d_model: int,

578 context_dim: int,

579 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

579 + norm_layer: Callable[[int], nn.Module] | None = None,

580 spatial_merge_size: int = 2,

581 - quant_config: Optional[QuantizationConfig] = None,

581 + quant_config: QuantizationConfig | None = None,

582 prefix: str = "",

583 use_data_parallel: bool = False,

584 ) -> None:

@@ -659,7 +659,7 @@ class Qwen2VisionTransformer(nn.Module):

659 self,

660 vision_config: Qwen2VLVisionConfig,

661 norm_eps: float = 1e-6,

662 - quant_config: Optional[QuantizationConfig] = None,

662 + quant_config: QuantizationConfig | None = None,

663 prefix: str = "",

664 use_data_parallel: bool = False,

665 ) -> None:

@@ -766,7 +766,7 @@ class Qwen2VisionTransformer(nn.Module):

766

767 def compute_attn_mask_seqlen(

768 self, cu_seqlens: torch.Tensor

769 - ) -> tuple[Optional[int], Optional[list[int]]]:

769 + ) -> tuple[int | None, list[int] | None]:

770 max_seqlen, seqlens = None, None

771 if (

772 self.attn_backend == _Backend.FLASH_ATTN

@@ -889,8 +889,8 @@ class Qwen2VLMultiModalDataParser(MultiModalDataParser):

889

890 def _parse_image_data(

891 self,

892 - data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],

893 - ) -> Optional[ModalityDataItems[Any, Any]]:

892 + data: dict[str, torch.Tensor] | ModalityData[ImageItem],

893 + ) -> ModalityDataItems[Any, Any] | None:

894 if isinstance(data, dict):

895 return DictEmbeddingItems(

896 data,

@@ -903,8 +903,8 @@ class Qwen2VLMultiModalDataParser(MultiModalDataParser):

903

904 def _parse_video_data(

905 self,

906 - data: Union[dict[str, torch.Tensor], ModalityData[VideoItem]],

907 - ) -> Optional[ModalityDataItems[Any, Any]]:

906 + data: dict[str, torch.Tensor] | ModalityData[VideoItem],

907 + ) -> ModalityDataItems[Any, Any] | None:

908 if isinstance(data, dict):

909 return DictEmbeddingItems(

910 data,

@@ -930,7 +930,7 @@ class Qwen2VLProcessingInfo(BaseProcessingInfo):

930 def get_image_processor(self, **kwargs: object) -> Qwen2VLImageProcessor:

931 return self.get_hf_processor(**kwargs).image_processor

932

933 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

933 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

934 return {"image": None, "video": None}

935

936 def get_mm_max_tokens_per_item(

@@ -949,7 +949,7 @@ class Qwen2VLProcessingInfo(BaseProcessingInfo):

949 image_height: int,

950 num_frames: int = 1,

951 do_resize: bool = True,

952 - image_processor: Optional[Qwen2VLImageProcessor],

952 + image_processor: Qwen2VLImageProcessor | None,

953 ) -> tuple[ImageSize, int]:

954 if image_processor is None:

955 image_processor = self.get_image_processor()

@@ -990,7 +990,7 @@ class Qwen2VLProcessingInfo(BaseProcessingInfo):

990 *,

991 image_width: int,

992 image_height: int,

993 - image_processor: Optional[Qwen2VLImageProcessor],

993 + image_processor: Qwen2VLImageProcessor | None,

994 ) -> int:

995 _, num_image_tokens = self._get_vision_info(

996 image_width=image_width,

@@ -1006,7 +1006,7 @@ class Qwen2VLProcessingInfo(BaseProcessingInfo):

1006 image_width: int,

1007 image_height: int,

1008 num_frames: int,

1009 - image_processor: Optional[Qwen2VLImageProcessor],

1009 + image_processor: Qwen2VLImageProcessor | None,

1010 ) -> int:

1011 _, num_video_tokens = self._get_vision_info(

1012 image_width=image_width,

@@ -1100,7 +1100,7 @@ class Qwen2VLDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2VLProcessingInfo]):

1100 self,

1101 seq_len: int,

1102 mm_counts: Mapping[str, int],

1103 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

1103 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

1104 ) -> MultiModalDataDict:

1105 num_images = mm_counts.get("image", 0)

1106 num_videos = mm_counts.get("video", 0)

@@ -1207,12 +1207,12 @@ class Qwen2VLForConditionalGeneration(

1207 self,

1208 input_tokens: list[int],

1209 hf_config: PretrainedConfig,

1210 - image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

1211 - video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

1212 - second_per_grid_ts: Optional[list[float]] = None,

1210 + image_grid_thw: list[list[int]] | torch.Tensor | None,

1211 + video_grid_thw: list[list[int]] | torch.Tensor | None,

1212 + second_per_grid_ts: list[float] | None = None,

1213 context_len: int = 0,

1214 - seq_len: Optional[int] = None,

1215 - audio_feature_lengths: Optional[torch.Tensor] = None,

1214 + seq_len: int | None = None,

1215 + audio_feature_lengths: torch.Tensor | None = None,

1216 use_audio_in_video: bool = False,

1217 ) -> tuple[torch.Tensor, int]:

1218 """Get M-RoPE input positions for Qwen2-VL model."""

@@ -1335,7 +1335,7 @@ class Qwen2VLForConditionalGeneration(

1335 return llm_positions, mrope_position_delta

1336

1337 @classmethod

1338 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1338 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1339 if modality.startswith("image"):

1340 return "<|vision_start|><|image_pad|><|vision_end|>"

1341 if modality.startswith("video"):

@@ -1396,7 +1396,7 @@ class Qwen2VLForConditionalGeneration(

1396

1397 def _parse_and_validate_image_input(

1398 self, **kwargs: object

1399 - ) -> Optional[Qwen2VLImageInputs]:

1399 + ) -> Qwen2VLImageInputs | None:

1400 pixel_values = kwargs.pop("pixel_values", None)

1401 image_embeds = kwargs.pop("image_embeds", None)

1402 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -1434,7 +1434,7 @@ class Qwen2VLForConditionalGeneration(

1434

1435 def _parse_and_validate_video_input(

1436 self, **kwargs: object

1437 - ) -> Optional[Qwen2VLVideoInputs]:

1437 + ) -> Qwen2VLVideoInputs | None:

1438 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1439 video_embeds = kwargs.pop("video_embeds", None)

1440 video_grid_thw = kwargs.pop("video_grid_thw", None)

@@ -1574,10 +1574,10 @@ class Qwen2VLForConditionalGeneration(

1574 self,

1575 input_ids: torch.Tensor,

1576 positions: torch.Tensor,

1577 - intermediate_tensors: Optional[IntermediateTensors] = None,

1578 - inputs_embeds: Optional[torch.Tensor] = None,

1577 + intermediate_tensors: IntermediateTensors | None = None,

1578 + inputs_embeds: torch.Tensor | None = None,

1579 **kwargs: object,

1580 - ) -> Union[torch.Tensor, IntermediateTensors]:

1580 + ) -> torch.Tensor | IntermediateTensors:

1581 """Run forward pass for Qwen2-VL.

1582

1583 Args:

@@ -1606,7 +1606,7 @@ class Qwen2VLForConditionalGeneration(

1606 def compute_logits(

1607 self,

1608 hidden_states: torch.Tensor,

1609 - ) -> Optional[torch.Tensor]:

1609 + ) -> torch.Tensor | None:

1610 return self.language_model.compute_logits(hidden_states)

1611

1612 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -1634,7 +1634,7 @@ class Tarsier2MultiModalProcessor(Qwen2VLMultiModalProcessor):

1634 class Tarsier2ImageProcessor(Qwen2VLImageProcessor):

1635 def __init__(

1636 self,

1637 - size: Optional[dict[str, int]] = None,

1637 + size: dict[str, int] | None = None,

1638 **kwargs,

1639 ) -> None:

1640 if size is not None and "min_pixels" in size and "max_pixels" in size:

vllm/model_executor/models/qwen3.py

+13 -13

@@ -24,7 +24,7 @@

24 """Inference-only Qwen3 model compatible with HuggingFace weights."""

25

26 from collections.abc import Iterable

27 -from typing import Any, Optional, Union

27 +from typing import Any

28

29 import torch

30 from torch import nn

@@ -58,16 +58,16 @@ class Qwen3Attention(nn.Module):

58 num_heads: int,

59 num_kv_heads: int,

60 max_position: int = 4096 * 32,

61 - head_dim: Optional[int] = None,

61 + head_dim: int | None = None,

62 rms_norm_eps: float = 1e-06,

63 qkv_bias: bool = False,

64 rope_theta: float = 10000,

65 - cache_config: Optional[CacheConfig] = None,

66 - quant_config: Optional[QuantizationConfig] = None,

67 - rope_scaling: Optional[tuple] = None,

65 + cache_config: CacheConfig | None = None,

66 + quant_config: QuantizationConfig | None = None,

67 + rope_scaling: tuple | None = None,

68 prefix: str = "",

69 attn_type: str = AttentionType.DECODER,

70 - dual_chunk_attention_config: Optional[dict[str, Any]] = None,

70 + dual_chunk_attention_config: dict[str, Any] | None = None,

71 ) -> None:

72 super().__init__()

73 self.hidden_size = hidden_size

@@ -160,8 +160,8 @@ class Qwen3DecoderLayer(nn.Module):

160 def __init__(

161 self,

162 config: Qwen3Config,

163 - cache_config: Optional[CacheConfig] = None,

164 - quant_config: Optional[QuantizationConfig] = None,

163 + cache_config: CacheConfig | None = None,

164 + quant_config: QuantizationConfig | None = None,

165 prefix: str = "",

166 ) -> None:

167 super().__init__()

@@ -214,7 +214,7 @@ class Qwen3DecoderLayer(nn.Module):

214 self,

215 positions: torch.Tensor,

216 hidden_states: torch.Tensor,

217 - residual: Optional[torch.Tensor],

217 + residual: torch.Tensor | None,

218 ) -> tuple[torch.Tensor, torch.Tensor]:

219 # Self Attention

220 if residual is None:

@@ -315,9 +315,9 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

315 self,

316 input_ids: torch.Tensor,

317 positions: torch.Tensor,

318 - intermediate_tensors: Optional[IntermediateTensors] = None,

319 - inputs_embeds: Optional[torch.Tensor] = None,

320 - ) -> Union[torch.Tensor, IntermediateTensors]:

318 + intermediate_tensors: IntermediateTensors | None = None,

319 + inputs_embeds: torch.Tensor | None = None,

320 + ) -> torch.Tensor | IntermediateTensors:

321 hidden_states = self.model(

322 input_ids, positions, intermediate_tensors, inputs_embeds

323 )

@@ -326,7 +326,7 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

326 def compute_logits(

327 self,

328 hidden_states: torch.Tensor,

329 - ) -> Optional[torch.Tensor]:

329 + ) -> torch.Tensor | None:

330 logits = self.logits_processor(self.lm_head, hidden_states)

331 return logits

332

vllm/model_executor/models/qwen3_moe.py

+16 -23

@@ -26,7 +26,7 @@

26 import typing

27 from collections.abc import Callable, Iterable

28 from itertools import islice

29 -from typing import Any, Optional, Union

29 +from typing import Any

30

31 import torch

32 from torch import nn

@@ -84,7 +84,7 @@ class Qwen3MoeMLP(nn.Module):

84 hidden_size: int,

85 intermediate_size: int,

86 hidden_act: str,

87 - quant_config: Optional[QuantizationConfig] = None,

87 + quant_config: QuantizationConfig | None = None,

88 reduce_results: bool = True,

89 prefix: str = "",

90 ) -> None:

@@ -215,15 +215,15 @@ class Qwen3MoeAttention(nn.Module):

215 num_heads: int,

216 num_kv_heads: int,

217 rope_theta: float = 10000,

218 - rope_scaling: Optional[dict[str, Any]] = None,

218 + rope_scaling: dict[str, Any] | None = None,

219 max_position_embeddings: int = 8192,

220 - head_dim: Optional[int] = None,

220 + head_dim: int | None = None,

221 rms_norm_eps: float = 1e-06,

222 qkv_bias: bool = False,

223 - cache_config: Optional[CacheConfig] = None,

224 - quant_config: Optional[QuantizationConfig] = None,

223 + cache_config: CacheConfig | None = None,

224 + quant_config: QuantizationConfig | None = None,

225 prefix: str = "",

226 - dual_chunk_attention_config: Optional[dict[str, Any]] = None,

226 + dual_chunk_attention_config: dict[str, Any] | None = None,

227 ) -> None:

228 super().__init__()

229 self.hidden_size = hidden_size

@@ -374,7 +374,7 @@ class Qwen3MoeDecoderLayer(nn.Module):

374 self,

375 positions: torch.Tensor,

376 hidden_states: torch.Tensor,

377 - residual: Optional[torch.Tensor],

377 + residual: torch.Tensor | None,

378 ) -> tuple[torch.Tensor, torch.Tensor]:

379 # Self Attention

380 if residual is None:

@@ -432,11 +432,9 @@ class Qwen3MoeModel(nn.Module):

432 self,

433 input_ids: torch.Tensor,

434 positions: torch.Tensor,

435 - intermediate_tensors: Optional[IntermediateTensors] = None,

436 - inputs_embeds: Optional[torch.Tensor] = None,

437 - ) -> Union[

438 - torch.Tensor, IntermediateTensors, tuple[torch.Tensor, list[torch.Tensor]]

439 - ]:

435 + intermediate_tensors: IntermediateTensors | None = None,

436 + inputs_embeds: torch.Tensor | None = None,

437 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:

438 if get_pp_group().is_first_rank:

439 if inputs_embeds is not None:

440 hidden_states = inputs_embeds

@@ -647,12 +645,7 @@ class Qwen3MoeForCausalLM(

645 self.quant_config = quant_config

646 # Only perform the following mapping when Qwen3MoeMLP exists

647 if getattr(config, "mlp_only_layers", []):

650 - self.packed_modules_mapping["gate_up_proj"] = (

651 - [

652 - "gate_proj",

653 - "up_proj",

654 - ],

655 - )

648 + self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]

649 self.model = Qwen3MoeModel(

650 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

651 )

@@ -742,9 +735,9 @@ class Qwen3MoeForCausalLM(

735 self,

736 input_ids: torch.Tensor,

737 positions: torch.Tensor,

745 - intermediate_tensors: Optional[IntermediateTensors] = None,

746 - inputs_embeds: Optional[torch.Tensor] = None,

747 - ) -> Union[torch.Tensor, IntermediateTensors]:

738 + intermediate_tensors: IntermediateTensors | None = None,

739 + inputs_embeds: torch.Tensor | None = None,

740 + ) -> torch.Tensor | IntermediateTensors:

741 hidden_states = self.model(

742 input_ids, positions, intermediate_tensors, inputs_embeds

743 )

@@ -753,7 +746,7 @@ class Qwen3MoeForCausalLM(

746 def compute_logits(

747 self,

748 hidden_states: torch.Tensor,

756 - ) -> Optional[torch.Tensor]:

749 + ) -> torch.Tensor | None:

750 logits = self.logits_processor(self.lm_head, hidden_states)

751 return logits

752

vllm/model_executor/models/qwen3_next.py

+13 -14

@@ -4,7 +4,6 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import Optional

7

8 import torch

9 from einops import rearrange

@@ -233,10 +232,10 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

232 def __init__(

233 self,

234 config: Qwen3NextConfig,

236 - model_config: Optional[ModelConfig] = None,

237 - cache_config: Optional[CacheConfig] = None,

238 - quant_config: Optional[QuantizationConfig] = None,

239 - speculative_config: Optional[SpeculativeConfig] = None,

235 + model_config: ModelConfig | None = None,

236 + cache_config: CacheConfig | None = None,

237 + quant_config: QuantizationConfig | None = None,

238 + speculative_config: SpeculativeConfig | None = None,

239 prefix: str = "",

240 ) -> None:

241 super().__init__()

@@ -661,9 +660,9 @@ class Qwen3NextAttention(nn.Module):

660 def __init__(

661 self,

662 config: Qwen3NextConfig,

664 - model_config: Optional[ModelConfig] = None,

665 - cache_config: Optional[CacheConfig] = None,

666 - quant_config: Optional[QuantizationConfig] = None,

663 + model_config: ModelConfig | None = None,

664 + cache_config: CacheConfig | None = None,

665 + quant_config: QuantizationConfig | None = None,

666 prefix: str = "",

667 ) -> None:

668 super().__init__()

@@ -863,7 +862,7 @@ class Qwen3NextDecoderLayer(nn.Module):

862 def forward(

863 self,

864 hidden_states: torch.Tensor,

866 - residual: Optional[torch.Tensor],

865 + residual: torch.Tensor | None,

866 positions: torch.Tensor = None,

867 **kwargs: object,

868 ):

@@ -971,8 +970,8 @@ class Qwen3NextModel(nn.Module):

970 self,

971 input_ids: torch.Tensor,

972 positions: torch.Tensor,

974 - intermediate_tensors: Optional[IntermediateTensors] = None,

975 - inputs_embeds: Optional[torch.Tensor] = None,

973 + intermediate_tensors: IntermediateTensors | None = None,

974 + inputs_embeds: torch.Tensor | None = None,

975 ) -> torch.Tensor:

976 if get_pp_group().is_first_rank:

977 if inputs_embeds is not None:

@@ -1207,8 +1206,8 @@ class Qwen3NextForCausalLM(

1206 self,

1207 input_ids: torch.Tensor,

1208 positions: torch.Tensor,

1210 - intermediate_tensors: Optional[IntermediateTensors] = None,

1211 - inputs_embeds: Optional[torch.Tensor] = None,

1209 + intermediate_tensors: IntermediateTensors | None = None,

1210 + inputs_embeds: torch.Tensor | None = None,

1211 **kwargs: object,

1212 ):

1213 hidden_states = self.model(

@@ -1251,7 +1250,7 @@ class Qwen3NextForCausalLM(

1250 def compute_logits(

1251 self,

1252 hidden_states: torch.Tensor,

1254 - ) -> Optional[torch.Tensor]:

1253 + ) -> torch.Tensor | None:

1254 return self.logits_processor(self.lm_head, hidden_states)

1255

1256 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/qwen3_next_mtp.py

+5 -6

@@ -3,7 +3,6 @@

3 """Inference-only Qwen3Next MTP model."""

4

5 from collections.abc import Iterable

6 -from typing import Optional

6

7 import torch

8 from torch import nn

@@ -108,8 +107,8 @@ class Qwen3NextMultiTokenPredictor(nn.Module):

107 input_ids: torch.Tensor,

108 positions: torch.Tensor,

109 hidden_states: torch.Tensor,

111 - intermediate_tensors: Optional[IntermediateTensors] = None,

112 - inputs_embeds: Optional[torch.Tensor] = None,

110 + intermediate_tensors: IntermediateTensors | None = None,

111 + inputs_embeds: torch.Tensor | None = None,

112 spec_step_idx: int = 0,

113 ) -> torch.Tensor:

114 if get_pp_group().is_first_rank:

@@ -275,8 +274,8 @@ class Qwen3NextMTP(nn.Module, SupportsPP):

274 input_ids: torch.Tensor,

275 positions: torch.Tensor,

276 hidden_states: torch.Tensor,

278 - intermediate_tensors: Optional[IntermediateTensors] = None,

279 - inputs_embeds: Optional[torch.Tensor] = None,

277 + intermediate_tensors: IntermediateTensors | None = None,

278 + inputs_embeds: torch.Tensor | None = None,

279 **kwargs: object,

280 ):

281 hidden_states = self.model(

@@ -288,7 +287,7 @@ class Qwen3NextMTP(nn.Module, SupportsPP):

287 self,

288 hidden_states: torch.Tensor,

289 spec_step_idx: int = 0,

291 - ) -> Optional[torch.Tensor]:

290 + ) -> torch.Tensor | None:

291 return self.logits_processor(self.lm_head, hidden_states)

292

293 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+30 -30

@@ -22,9 +22,9 @@

22 # limitations under the License.

23 """Inference-only Qwen3-Omni-Moe model (thinker part)."""

24

25 -from collections.abc import Iterable, Mapping, Sequence

25 +from collections.abc import Callable, Iterable, Mapping, Sequence

26 from functools import partial

27 -from typing import Any, Callable, Optional, Union

27 +from typing import Any

28

29 import numpy as np

30 import torch

@@ -156,7 +156,7 @@ class Qwen3_VisionMLP(nn.Module):

156 hidden_features: int,

157 bias: bool = False,

158 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

159 - quant_config: Optional[QuantizationConfig] = None,

159 + quant_config: QuantizationConfig | None = None,

160 prefix: str = "",

161 ):

162 super().__init__()

@@ -190,8 +190,8 @@ class Qwen3_VisionBlock(nn.Module):

190 num_heads: int,

191 mlp_hidden_dim: int,

192 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

193 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

194 - quant_config: Optional[QuantizationConfig] = None,

193 + norm_layer: Callable[[int], nn.Module] | None = None,

194 + quant_config: QuantizationConfig | None = None,

195 prefix: str = "",

196 ) -> None:

197 super().__init__()

@@ -220,8 +220,8 @@ class Qwen3_VisionBlock(nn.Module):

220 x: torch.Tensor,

221 cu_seqlens: torch.Tensor,

222 rotary_pos_emb: torch.Tensor,

223 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

224 - seqlens: Optional[list[int]] = None, # Only used for xFormers

223 + max_seqlen: int | None = None, # Only used for Flash Attention

224 + seqlens: list[int] | None = None, # Only used for xFormers

225 ) -> torch.Tensor:

226 x = x + self.attn(

227 self.norm1(x),

@@ -240,10 +240,10 @@ class Qwen3_VisionPatchMerger(nn.Module):

240 self,

241 d_model: int,

242 context_dim: int,

243 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

243 + norm_layer: Callable[[int], nn.Module] | None = None,

244 spatial_merge_size: int = 2,

245 use_postshuffle_norm: bool = False,

246 - quant_config: Optional[QuantizationConfig] = None,

246 + quant_config: QuantizationConfig | None = None,

247 prefix: str = "",

248 ) -> None:

249 super().__init__()

@@ -297,7 +297,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

297 self,

298 vision_config,

299 norm_eps: float = 1e-6,

300 - quant_config: Optional[QuantizationConfig] = None,

300 + quant_config: QuantizationConfig | None = None,

301 prefix: str = "",

302 ) -> None:

303 super().__init__()

@@ -482,7 +482,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

482 def compute_attn_mask_seqlen(

483 self,

484 cu_seqlens: torch.Tensor,

485 - ) -> tuple[Optional[int], Optional[list[int]]]:

485 + ) -> tuple[int | None, list[int] | None]:

486 max_seqlen, seqlens = None, None

487 if self.attn_backend == _Backend.FLASH_ATTN:

488 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

@@ -595,10 +595,10 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):

595 self,

596 input_ids: torch.Tensor,

597 positions: torch.Tensor,

598 - intermediate_tensors: Optional[IntermediateTensors] = None,

599 - inputs_embeds: Optional[torch.Tensor] = None,

600 - deepstack_input_embeds: Optional[IntermediateTensors] = None,

601 - ) -> Union[torch.Tensor, IntermediateTensors]:

598 + intermediate_tensors: IntermediateTensors | None = None,

599 + inputs_embeds: torch.Tensor | None = None,

600 + deepstack_input_embeds: IntermediateTensors | None = None,

601 + ) -> torch.Tensor | IntermediateTensors:

602 if get_pp_group().is_first_rank:

603 if inputs_embeds is not None:

604 hidden_states = inputs_embeds

@@ -683,7 +683,7 @@ class Qwen3OmniMoeThinkerProcessingInfo(

683 assert isinstance(feature_extractor, WhisperFeatureExtractor)

684 return feature_extractor

685

686 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

686 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

687 return {"audio": None, "image": None, "video": None}

688

689

@@ -818,7 +818,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(

818 self,

819 thinker_config: PretrainedConfig,

820 audio_len: int,

821 - video_grid_thw: Union[list[int], torch.Tensor],

821 + video_grid_thw: list[int] | torch.Tensor,

822 video_second_per_grid_t: float,

823 ) -> list[int]:

824 shift = 0

@@ -1100,7 +1100,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1100 )

1101

1102 @classmethod

1103 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1103 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1104 if modality.startswith("image"):

1105 return "<|vision_start|><|image_pad|><|vision_end|>"

1106 if modality.startswith("video"):

@@ -1246,7 +1246,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1246

1247 def get_multimodal_embeddings(

1248 self, **kwargs: object

1249 - ) -> Optional[MultiModalEmbeddings]:

1249 + ) -> MultiModalEmbeddings | None:

1250 mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

1251 if not mm_input_by_modality:

1252 return []

@@ -1273,9 +1273,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1273 def get_input_embeddings(

1274 self,

1275 input_ids: torch.Tensor,

1276 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1276 + multimodal_embeddings: MultiModalEmbeddings | None = None,

1277 *,

1278 - is_multimodal: Optional[torch.Tensor] = None,

1278 + is_multimodal: torch.Tensor | None = None,

1279 handle_oov_mm_token: bool = False,

1280 ) -> torch.Tensor:

1281 inputs_embeds = self._get_text_embeddings(

@@ -1356,10 +1356,10 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1356 self,

1357 input_ids: torch.Tensor,

1358 positions: torch.Tensor,

1359 - intermediate_tensors: Optional[IntermediateTensors] = None,

1360 - inputs_embeds: Optional[torch.Tensor] = None,

1359 + intermediate_tensors: IntermediateTensors | None = None,

1360 + inputs_embeds: torch.Tensor | None = None,

1361 **kwargs: object,

1362 - ) -> Union[torch.Tensor, IntermediateTensors]:

1362 + ) -> torch.Tensor | IntermediateTensors:

1363 if intermediate_tensors is not None:

1364 inputs_embeds = None

1365

@@ -1391,7 +1391,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1391 def compute_logits(

1392 self,

1393 hidden_states: torch.Tensor,

1394 - ) -> Optional[torch.Tensor]:

1394 + ) -> torch.Tensor | None:

1395 return self.language_model.compute_logits(hidden_states)

1396

1397 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -1408,12 +1408,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1408 self,

1409 input_tokens: list[int],

1410 hf_config: PretrainedConfig,

1411 - image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

1412 - video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],

1413 - second_per_grid_ts: Optional[list[float]] = None,

1411 + image_grid_thw: list[list[int]] | torch.Tensor | None,

1412 + video_grid_thw: list[list[int]] | torch.Tensor | None,

1413 + second_per_grid_ts: list[float] | None = None,

1414 context_len: int = 0,

1415 - seq_len: Optional[int] = None,

1416 - audio_feature_lengths: Optional[torch.Tensor] = None,

1415 + seq_len: int | None = None,

1416 + audio_feature_lengths: torch.Tensor | None = None,

1417 use_audio_in_video: bool = False,

1418 ) -> tuple[torch.Tensor, int]:

1419 config = hf_config.thinker_config

vllm/model_executor/models/qwen3_vl.py

+34 -36

@@ -24,10 +24,10 @@

24 # limitations under the License.

25 """Inference-only Qwen3VL model compatible with HuggingFace weights."""

26

27 -from collections.abc import Iterable, Mapping, Sequence

27 +from collections.abc import Callable, Iterable, Mapping, Sequence

28 from functools import partial

29 from itertools import islice

30 -from typing import Any, Callable, Optional, Union

30 +from typing import Any

31

32 import numpy as np

33 import torch

@@ -151,7 +151,7 @@ class Qwen3_VisionMLP(nn.Module):

151 hidden_features: int,

152 bias: bool = False,

153 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

154 - quant_config: Optional[QuantizationConfig] = None,

154 + quant_config: QuantizationConfig | None = None,

155 prefix: str = "",

156 use_data_parallel: bool = False,

157 ):

@@ -188,8 +188,8 @@ class Qwen3_VisionBlock(nn.Module):

188 num_heads: int,

189 mlp_hidden_dim: int,

190 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

191 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

192 - quant_config: Optional[QuantizationConfig] = None,

191 + norm_layer: Callable[[int], nn.Module] | None = None,

192 + quant_config: QuantizationConfig | None = None,

193 prefix: str = "",

194 use_data_parallel: bool = False,

195 attn_backend: _Backend = _Backend.TORCH_SDPA,

@@ -225,8 +225,8 @@ class Qwen3_VisionBlock(nn.Module):

225 x: torch.Tensor,

226 cu_seqlens: torch.Tensor,

227 rotary_pos_emb: torch.Tensor,

228 - max_seqlen: Optional[int] = None, # Only used for Flash Attention

229 - seqlens: Optional[list[int]] = None, # Only used for xFormers

228 + max_seqlen: int | None = None, # Only used for Flash Attention

229 + seqlens: list[int] | None = None, # Only used for xFormers

230 ) -> torch.Tensor:

231 x = x + self.attn(

232 self.norm1(x),

@@ -245,10 +245,10 @@ class Qwen3_VisionPatchMerger(nn.Module):

245 self,

246 d_model: int,

247 context_dim: int,

248 - norm_layer: Optional[Callable[[int], nn.Module]] = None,

248 + norm_layer: Callable[[int], nn.Module] | None = None,

249 spatial_merge_size: int = 2,

250 use_postshuffle_norm: bool = False,

251 - quant_config: Optional[QuantizationConfig] = None,

251 + quant_config: QuantizationConfig | None = None,

252 prefix: str = "",

253 use_data_parallel: bool = False,

254 ) -> None:

@@ -297,7 +297,7 @@ class Qwen3_VisionTransformer(nn.Module):

297 self,

298 vision_config: Qwen3VLVisionConfig,

299 norm_eps: float = 1e-6,

300 - quant_config: Optional[QuantizationConfig] = None,

300 + quant_config: QuantizationConfig | None = None,

301 prefix: str = "",

302 use_data_parallel: bool = False,

303 ) -> None:

@@ -511,7 +511,7 @@ class Qwen3_VisionTransformer(nn.Module):

511 def compute_attn_mask_seqlen(

512 self,

513 cu_seqlens: torch.Tensor,

514 - ) -> tuple[Optional[int], Optional[list[int]]]:

514 + ) -> tuple[int | None, list[int] | None]:

515 max_seqlen, seqlens = None, None

516 if (

517 self.attn_backend == _Backend.FLASH_ATTN

@@ -625,9 +625,7 @@ class Qwen3VLProcessingInfo(Qwen2VLProcessingInfo):

625 image_height: int,

626 num_frames: int = 2,

627 do_resize: bool = True,

628 - image_processor: Optional[

629 - Union[Qwen2VLImageProcessorFast, Qwen3VLVideoProcessor]

630 - ],

628 + image_processor: Qwen2VLImageProcessorFast | Qwen3VLVideoProcessor | None,

629 ) -> tuple[ImageSize, int]:

630 if image_processor is None and num_frames > 1:

631 image_processor = self.get_video_processor()

@@ -726,8 +724,8 @@ class Qwen3VLProcessingInfo(Qwen2VLProcessingInfo):

724 self,

725 metadata: dict[str, Any],

726 out_item: MultiModalKwargsItem,

729 - do_sample_frames: Optional[bool] = None,

730 - sampled_fps: Optional[float] = None,

727 + do_sample_frames: bool | None = None,

728 + sampled_fps: float | None = None,

729 ) -> list[int]:

730 video_processor = self.get_video_processor()

731 merge_size = video_processor.merge_size

@@ -778,7 +776,7 @@ class Qwen3VLDummyInputsBuilder(BaseDummyInputsBuilder[Qwen3VLProcessingInfo]):

776 self,

777 seq_len: int,

778 mm_counts: Mapping[str, int],

781 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

779 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

780 ) -> MultiModalDataDict:

781 num_images = mm_counts.get("image", 0)

782 num_videos = mm_counts.get("video", 0)

@@ -1096,11 +1094,11 @@ class Qwen3LLMModel(Qwen3Model):

1094 self,

1095 input_ids: torch.Tensor,

1096 positions: torch.Tensor,

1099 - intermediate_tensors: Optional[IntermediateTensors] = None,

1100 - inputs_embeds: Optional[torch.Tensor] = None,

1097 + intermediate_tensors: IntermediateTensors | None = None,

1098 + inputs_embeds: torch.Tensor | None = None,

1099 # args for deepstack

1102 - deepstack_input_embeds: Optional[IntermediateTensors] = None,

1103 - ) -> Union[torch.Tensor, IntermediateTensors]:

1100 + deepstack_input_embeds: IntermediateTensors | None = None,

1101 + ) -> torch.Tensor | IntermediateTensors:

1102 if get_pp_group().is_first_rank:

1103 if inputs_embeds is not None:

1104 hidden_states = inputs_embeds

@@ -1201,7 +1199,7 @@ class Qwen3VLForConditionalGeneration(

1199 )

1200

1201 @classmethod

1204 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1202 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1203 if modality.startswith("image"):

1204 return "<|vision_start|><|image_pad|><|vision_end|>"

1205 if modality.startswith("video"):

@@ -1314,7 +1312,7 @@ class Qwen3VLForConditionalGeneration(

1312

1313 def _parse_and_validate_image_input(

1314 self, **kwargs: object

1317 - ) -> Optional[Qwen2_5_VLImageInputs]:

1315 + ) -> Qwen2_5_VLImageInputs | None:

1316 pixel_values = kwargs.pop("pixel_values", None)

1317 image_embeds = kwargs.pop("image_embeds", None)

1318 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -1363,7 +1361,7 @@ class Qwen3VLForConditionalGeneration(

1361

1362 def _parse_and_validate_video_input(

1363 self, **kwargs: object

1366 - ) -> Optional[Qwen2_5_VLVideoInputs]:

1364 + ) -> Qwen2_5_VLVideoInputs | None:

1365 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1366 video_embeds = kwargs.pop("video_embeds", None)

1367 video_grid_thw = kwargs.pop("video_grid_thw", None)

@@ -1486,12 +1484,12 @@ class Qwen3VLForConditionalGeneration(

1484 cls,

1485 input_tokens: list[int],

1486 hf_config: PretrainedConfig,

1489 - image_grid_thw: Union[list[list[int]], torch.Tensor],

1490 - video_grid_thw: Union[list[list[int]], torch.Tensor],

1487 + image_grid_thw: list[list[int]] | torch.Tensor,

1488 + video_grid_thw: list[list[int]] | torch.Tensor,

1489 context_len: int = 0,

1492 - seq_len: Optional[int] = None,

1493 - second_per_grid_ts: Optional[list[float]] = None,

1494 - audio_feature_lengths: Optional[torch.Tensor] = None,

1490 + seq_len: int | None = None,

1491 + second_per_grid_ts: list[float] | None = None,

1492 + audio_feature_lengths: torch.Tensor | None = None,

1493 use_audio_in_video: bool = False,

1494 ) -> tuple[torch.Tensor, int]:

1495 """Get mrope input positions and delta value."""

@@ -1596,7 +1594,7 @@ class Qwen3VLForConditionalGeneration(

1594

1595 def get_multimodal_embeddings(

1596 self, **kwargs: object

1599 - ) -> Optional[MultiModalEmbeddings]:

1597 + ) -> MultiModalEmbeddings | None:

1598 mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

1599 if not mm_input_by_modality:

1600 return None

@@ -1661,9 +1659,9 @@ class Qwen3VLForConditionalGeneration(

1659 def get_input_embeddings(

1660 self,

1661 input_ids: torch.Tensor,

1664 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1662 + multimodal_embeddings: MultiModalEmbeddings | None = None,

1663 *,

1666 - is_multimodal: Optional[torch.Tensor] = None,

1664 + is_multimodal: torch.Tensor | None = None,

1665 handle_oov_mm_token: bool = False,

1666 ) -> torch.Tensor:

1667 inputs_embeds = self._get_text_embeddings(

@@ -1710,10 +1708,10 @@ class Qwen3VLForConditionalGeneration(

1708 self,

1709 input_ids: torch.Tensor,

1710 positions: torch.Tensor,

1713 - intermediate_tensors: Optional[IntermediateTensors] = None,

1714 - inputs_embeds: Optional[torch.Tensor] = None,

1711 + intermediate_tensors: IntermediateTensors | None = None,

1712 + inputs_embeds: torch.Tensor | None = None,

1713 **kwargs: object,

1716 - ) -> Union[torch.Tensor, IntermediateTensors]:

1714 + ) -> torch.Tensor | IntermediateTensors:

1715 """Run forward pass for Qwen3VL.

1716

1717 Args:

@@ -1769,7 +1767,7 @@ class Qwen3VLForConditionalGeneration(

1767 def compute_logits(

1768 self,

1769 hidden_states: torch.Tensor,

1772 - ) -> Optional[torch.Tensor]:

1770 + ) -> torch.Tensor | None:

1771 return self.language_model.compute_logits(hidden_states)

1772

1773 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/qwen3_vl_moe.py

+5 -6

@@ -25,9 +25,8 @@

25 """Inference-only Qwen3-VL-MoE model compatible with HuggingFace weights."""

26

27 import typing

28 -from collections.abc import Iterable

28 +from collections.abc import Callable, Iterable

29 from itertools import islice

30 -from typing import Callable, Optional, Union

30

31 import torch

32 from transformers.models.qwen3_vl_moe.configuration_qwen3_vl_moe import Qwen3VLMoeConfig

@@ -90,10 +89,10 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):

89 self,

90 input_ids: torch.Tensor,

91 positions: torch.Tensor,

93 - intermediate_tensors: Optional[IntermediateTensors] = None,

94 - inputs_embeds: Optional[torch.Tensor] = None,

95 - deepstack_input_embeds: Optional[IntermediateTensors] = None,

96 - ) -> Union[torch.Tensor, IntermediateTensors]:

92 + intermediate_tensors: IntermediateTensors | None = None,

93 + inputs_embeds: torch.Tensor | None = None,

94 + deepstack_input_embeds: IntermediateTensors | None = None,

95 + ) -> torch.Tensor | IntermediateTensors:

96 if get_pp_group().is_first_rank:

97 if inputs_embeds is not None:

98 hidden_states = inputs_embeds

vllm/model_executor/models/qwen_vl.py

+28 -28

@@ -9,9 +9,9 @@

9 import copy

10 import math

11 import unicodedata

12 -from collections.abc import Collection, Mapping, Sequence, Set

12 +from collections.abc import Callable, Collection, Mapping, Sequence, Set

13 from functools import lru_cache, partial

14 -from typing import Annotated, Callable, Literal, Optional, Union

14 +from typing import Annotated, Literal, TypeAlias

15

16 import regex as re

17 import torch

@@ -93,7 +93,7 @@ class QwenImageEmbeddingInputs(TensorSchema):

93 data: Annotated[torch.Tensor, TensorShape("bn", 256, "hs")]

94

95

96 -QwenImageInputs = Union[QwenImagePixelInputs, QwenImageEmbeddingInputs]

96 +QwenImageInputs: TypeAlias = QwenImagePixelInputs | QwenImageEmbeddingInputs

97

98

99 class VisualAttention(nn.Module):

@@ -107,8 +107,8 @@ class VisualAttention(nn.Module):

107 embed_dim: int,

108 num_heads: int,

109 bias: bool = True,

110 - kdim: Optional[int] = None,

111 - vdim: Optional[int] = None,

110 + kdim: int | None = None,

111 + vdim: int | None = None,

112 ):

113 super().__init__()

114 self.embed_dim = embed_dim

@@ -135,7 +135,7 @@ class VisualAttention(nn.Module):

135 def forward(

136 self,

137 x: torch.Tensor,

138 - attn_mask: Optional[torch.Tensor] = None,

138 + attn_mask: torch.Tensor | None = None,

139 ) -> torch.Tensor:

140 # query/key/value: [sq, b, h]

141 sq, b, _ = x.size()

@@ -213,7 +213,7 @@ class QwenVLMLP(nn.Module):

213 self,

214 hidden_size: int,

215 intermediate_size: int,

216 - quant_config: Optional[QuantizationConfig] = None,

216 + quant_config: QuantizationConfig | None = None,

217 ):

218 super().__init__()

219 self.c_fc = ColumnParallelLinear(

@@ -241,7 +241,7 @@ class VisualAttentionBlock(nn.Module):

241 n_head: int,

242 mlp_ratio: float = 4.0,

243 norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,

244 - quant_config: Optional[QuantizationConfig] = None,

244 + quant_config: QuantizationConfig | None = None,

245 ):

246 super().__init__()

247

@@ -258,7 +258,7 @@ class VisualAttentionBlock(nn.Module):

258 def attention(

259 self,

260 x: torch.Tensor,

261 - attn_mask: Optional[torch.Tensor] = None,

261 + attn_mask: torch.Tensor | None = None,

262 ) -> torch.Tensor:

263 attn_mask = attn_mask.to(x.dtype) if attn_mask is not None else None

264 return self.attn(x, attn_mask=attn_mask)

@@ -266,7 +266,7 @@ class VisualAttentionBlock(nn.Module):

266 def forward(

267 self,

268 x: torch.Tensor,

269 - attn_mask: Optional[torch.Tensor] = None,

269 + attn_mask: torch.Tensor | None = None,

270 ) -> torch.Tensor:

271 x = x + self.attention(self.ln_1(x), attn_mask=attn_mask)

272 x = x + self.mlp(self.ln_2(x))

@@ -281,7 +281,7 @@ class TransformerBlock(nn.Module):

281 heads: int,

282 mlp_ratio: float = 4.0,

283 norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,

284 - quant_config: Optional[QuantizationConfig] = None,

284 + quant_config: QuantizationConfig | None = None,

285 ):

286 super().__init__()

287 self.width = width

@@ -307,7 +307,7 @@ class TransformerBlock(nn.Module):

307 return self.resblocks[0].mlp.c_fc.weight.device

308

309 def forward(

310 - self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None

310 + self, x: torch.Tensor, attn_mask: torch.Tensor | None = None

311 ) -> torch.Tensor:

312 for r in self.resblocks:

313 x = r(x, attn_mask=attn_mask)

@@ -326,7 +326,7 @@ class VisionTransformer(nn.Module):

326 n_queries: int = 256,

327 output_dim: int = 512,

328 image_start_id: int = 151857,

329 - quant_config: Optional[QuantizationConfig] = None,

329 + quant_config: QuantizationConfig | None = None,

330 **kwargs,

331 ):

332 super().__init__()

@@ -434,10 +434,10 @@ def _get_tokenizer_without_image_pad(

434 def tokenize(

435 self,

436 text: str,

437 - allowed_special: Union[Set[str], str] = "all",

438 - disallowed_special: Union[Collection[str], str] = (),

437 + allowed_special: Set[str] | str = "all",

438 + disallowed_special: Collection[str] | str = (),

439 **kwargs,

440 - ) -> list[Union[bytes, str]]:

440 + ) -> list[bytes | str]:

441 text = unicodedata.normalize("NFC", text)

442

443 return [

@@ -451,9 +451,9 @@ def _get_tokenizer_without_image_pad(

451

452 def _decode(

453 self,

454 - token_ids: Union[int, list[int]],

454 + token_ids: int | list[int],

455 skip_special_tokens: bool = False,

456 - errors: Optional[str] = None,

456 + errors: str | None = None,

457 **kwargs,

458 ) -> str:

459 if isinstance(token_ids, int):

@@ -523,9 +523,9 @@ class QwenVLProcessor:

523

524 def __call__(

525 self,

526 - text: Optional[Union[TextInput, list[TextInput]]] = None,

527 - images: Optional[Union[ImageInput, list[ImageInput]]] = None,

528 - return_tensors: Optional[Union[str, TensorType]] = None,

526 + text: TextInput | list[TextInput] | None = None,

527 + images: ImageInput | list[ImageInput] | None = None,

528 + return_tensors: str | TensorType | None = None,

529 ) -> BatchFeature:

530 if text is None:

531 text = []

@@ -568,7 +568,7 @@ class QwenVLProcessingInfo(BaseProcessingInfo):

568 **kwargs,

569 )

570

571 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

571 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

572 return {"image": None}

573

574 def get_num_image_tokens(self) -> int:

@@ -597,7 +597,7 @@ class QwenVLDummyInputsBuilder(BaseDummyInputsBuilder[QwenVLProcessingInfo]):

597 self,

598 seq_len: int,

599 mm_counts: Mapping[str, int],

600 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

600 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

601 ) -> MultiModalDataDict:

602 hf_config = self.info.get_hf_config()

603 vision_config = hf_config.visual

@@ -722,7 +722,7 @@ class QwenVLForConditionalGeneration(

722 )

723

724 @classmethod

725 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

725 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

726 if modality.startswith("image"):

727 return f"Picture {i}: <img></img>"

728

@@ -745,7 +745,7 @@ class QwenVLForConditionalGeneration(

745

746 def _parse_and_validate_image_input(

747 self, **kwargs: object

748 - ) -> Optional[QwenImageInputs]:

748 + ) -> QwenImageInputs | None:

749 pixel_values = kwargs.pop("pixel_values", None)

750 image_embeds = kwargs.pop("image_embeds", None)

751

@@ -799,10 +799,10 @@ class QwenVLForConditionalGeneration(

799 self,

800 input_ids: torch.Tensor,

801 positions: torch.Tensor,

802 - intermediate_tensors: Optional[IntermediateTensors] = None,

803 - inputs_embeds: Optional[torch.Tensor] = None,

802 + intermediate_tensors: IntermediateTensors | None = None,

803 + inputs_embeds: torch.Tensor | None = None,

804 **kwargs: object,

805 - ) -> Union[torch.Tensor, IntermediateTensors]:

805 + ) -> torch.Tensor | IntermediateTensors:

806 if intermediate_tensors is not None:

807 inputs_embeds = None

808

vllm/model_executor/models/radio.py

+19 -19

@@ -11,7 +11,7 @@

11 import math

12 from collections.abc import Iterable

13 from itertools import repeat

14 -from typing import Optional, Union

14 +from typing import TypeAlias

15

16 import torch

17 import torch.nn as nn

@@ -23,8 +23,8 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

23 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

24 from vllm.model_executor.models.intern_vit import InternVisionEncoder

25

26 -input_dim_t = Union[int, tuple[int, int]]

27 -norm_t = Union[tuple[float, float, float], torch.Tensor]

26 +input_dim_t: TypeAlias = int | tuple[int, int]

27 +norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor

28

29

30 def _ntuple(n):

@@ -75,8 +75,8 @@ class ClsToken(nn.Module):

75 ndim: int,

76 num_tokens: int = 1,

77 enabled: bool = True,

78 - register_multiple: Optional[int] = None,

79 - num_registers: Optional[int] = None,

78 + register_multiple: int | None = None,

79 + num_registers: int | None = None,

80 ):

81 super().__init__()

82

@@ -128,12 +128,12 @@ class ViTPatchGenerator(nn.Module):

128 abs_pos: bool = True,

129 normalize_patches: bool = False,

130 cls_token: bool = False,

131 - max_input_dims: Optional[input_dim_t] = None,

131 + max_input_dims: input_dim_t | None = None,

132 pos_dropout: float = 0.0,

133 return_pos_enc: bool = False,

134 num_cls_tokens: int = 1,

135 - register_multiple: Optional[int] = None,

136 - num_registers: Optional[int] = None,

135 + register_multiple: int | None = None,

136 + num_registers: int | None = None,

137 patch_bias: bool = False,

138 device=None,

139 dtype=None,

@@ -275,8 +275,8 @@ class ViTPatchGenerator(nn.Module):

275 def apply_pos_enc(

276 self,

277 patches: torch.Tensor,

278 - patch_idxs: Optional[torch.Tensor] = None,

279 - input_size: Optional[tuple[int, int]] = None,

278 + patch_idxs: torch.Tensor | None = None,

279 + input_size: tuple[int, int] | None = None,

280 ) -> torch.Tensor:

281 if not self.abs_pos:

282 return patches

@@ -299,8 +299,8 @@ class ViTPatchGenerator(nn.Module):

299 def get_pos_enc(

300 self,

301 batch_size: int,

302 - patch_idxs: Optional[torch.Tensor] = None,

303 - input_size: Optional[tuple[int, int]] = None,

302 + patch_idxs: torch.Tensor | None = None,

303 + input_size: tuple[int, int] | None = None,

304 ) -> torch.Tensor:

305 if input_size is None:

306 input_dims = self.input_dims

@@ -440,9 +440,9 @@ class RadioInternVisionModel(nn.Module):

440 def __init__(

441 self,

442 config: PretrainedConfig = None,

443 - quant_config: Optional[QuantizationConfig] = None,

443 + quant_config: QuantizationConfig | None = None,

444 *,

445 - num_hidden_layers_override: Optional[int] = None,

445 + num_hidden_layers_override: int | None = None,

446 num_dummy_heads: int = 0,

447 prefix: str = "",

448 ) -> None:

@@ -472,7 +472,7 @@ class RadioInternVisionModel(nn.Module):

472 prefix=f"{prefix}.encoder",

473 )

474

475 - def _init_img_size(self, patch_size, img_size: Union[int, tuple[int, int]]):

475 + def _init_img_size(self, patch_size, img_size: int | tuple[int, int]):

476 if img_size is None:

477 return None, None, None

478 img_size = to_2tuple(img_size)

@@ -498,9 +498,9 @@ class RadioModel(nn.Module):

498 def __init__(

499 self,

500 config: PretrainedConfig,

501 - quant_config: Optional[QuantizationConfig] = None,

501 + quant_config: QuantizationConfig | None = None,

502 *,

503 - num_hidden_layers_override: Optional[int] = None,

503 + num_hidden_layers_override: int | None = None,

504 num_dummy_heads: int = 0,

505 prefix: str = "",

506 ) -> None:

@@ -522,8 +522,8 @@ class RadioModel(nn.Module):

522

523 def forward(

524 self,

525 - pixel_values: Optional[torch.Tensor] = None,

526 - pixel_embeds: Optional[torch.Tensor] = None,

525 + pixel_values: torch.Tensor | None = None,

526 + pixel_embeds: torch.Tensor | None = None,

527 ) -> torch.FloatTensor:

528 x = self.input_conditioner(pixel_values)

529 y = self.model(x)

vllm/model_executor/models/registry.py

+23 -23

@@ -14,11 +14,11 @@ import subprocess

14 import sys

15 import tempfile

16 from abc import ABC, abstractmethod

17 -from collections.abc import Set

17 +from collections.abc import Callable, Set

18 from dataclasses import asdict, dataclass, field

19 from functools import lru_cache

20 from pathlib import Path

21 -from typing import Callable, Optional, TypeVar, Union

21 +from typing import TypeVar

22

23 import torch.nn as nn

24 import transformers

@@ -654,7 +654,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel):

654 def _try_load_model_cls(

655 model_arch: str,

656 model: _BaseRegisteredModel,

657 -) -> Optional[type[nn.Module]]:

657 +) -> type[nn.Module] | None:

658 from vllm.platforms import current_platform

659

660 current_platform.verify_model_arch(model_arch)

@@ -669,7 +669,7 @@ def _try_load_model_cls(

669 def _try_inspect_model_cls(

670 model_arch: str,

671 model: _BaseRegisteredModel,

672 -) -> Optional[_ModelInfo]:

672 +) -> _ModelInfo | None:

673 try:

674 return model.inspect_model_cls()

675 except Exception:

@@ -688,7 +688,7 @@ class _ModelRegistry:

688 def register_model(

689 self,

690 model_arch: str,

691 - model_cls: Union[type[nn.Module], str],

691 + model_cls: type[nn.Module] | str,

692 ) -> None:

693 """

694 Register an external model to be used in vLLM.

@@ -756,13 +756,13 @@ class _ModelRegistry:

756 f"Supported architectures: {all_supported_archs}"

757 )

758

759 - def _try_load_model_cls(self, model_arch: str) -> Optional[type[nn.Module]]:

759 + def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None:

760 if model_arch not in self.models:

761 return None

762

763 return _try_load_model_cls(model_arch, self.models[model_arch])

764

765 - def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:

765 + def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None:

766 if model_arch not in self.models:

767 return None

768

@@ -772,7 +772,7 @@ class _ModelRegistry:

772 self,

773 architecture: str,

774 model_config: ModelConfig,

775 - ) -> Optional[str]:

775 + ) -> str | None:

776 if architecture in _TRANSFORMERS_BACKEND_MODELS:

777 return architecture

778

@@ -862,7 +862,7 @@ class _ModelRegistry:

862

863 def inspect_model_cls(

864 self,

865 - architectures: Union[str, list[str]],

865 + architectures: str | list[str],

866 model_config: ModelConfig,

867 ) -> tuple[_ModelInfo, str]:

868 if isinstance(architectures, str):

@@ -914,7 +914,7 @@ class _ModelRegistry:

914

915 def resolve_model_cls(

916 self,

917 - architectures: Union[str, list[str]],

917 + architectures: str | list[str],

918 model_config: ModelConfig,

919 ) -> tuple[type[nn.Module], str]:

920 if isinstance(architectures, str):

@@ -968,7 +968,7 @@ class _ModelRegistry:

968

969 def is_text_generation_model(

970 self,

971 - architectures: Union[str, list[str]],

971 + architectures: str | list[str],

972 model_config: ModelConfig,

973 ) -> bool:

974 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -976,7 +976,7 @@ class _ModelRegistry:

976

977 def is_pooling_model(

978 self,

979 - architectures: Union[str, list[str]],

979 + architectures: str | list[str],

980 model_config: ModelConfig,

981 ) -> bool:

982 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -984,7 +984,7 @@ class _ModelRegistry:

984

985 def is_cross_encoder_model(

986 self,

987 - architectures: Union[str, list[str]],

987 + architectures: str | list[str],

988 model_config: ModelConfig,

989 ) -> bool:

990 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -992,7 +992,7 @@ class _ModelRegistry:

992

993 def is_multimodal_model(

994 self,

995 - architectures: Union[str, list[str]],

995 + architectures: str | list[str],

996 model_config: ModelConfig,

997 ) -> bool:

998 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1000,7 +1000,7 @@ class _ModelRegistry:

1000

1001 def is_multimodal_raw_input_only_model(

1002 self,

1003 - architectures: Union[str, list[str]],

1003 + architectures: str | list[str],

1004 model_config: ModelConfig,

1005 ) -> bool:

1006 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1008,7 +1008,7 @@ class _ModelRegistry:

1008

1009 def is_pp_supported_model(

1010 self,

1011 - architectures: Union[str, list[str]],

1011 + architectures: str | list[str],

1012 model_config: ModelConfig,

1013 ) -> bool:

1014 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1016,7 +1016,7 @@ class _ModelRegistry:

1016

1017 def model_has_inner_state(

1018 self,

1019 - architectures: Union[str, list[str]],

1019 + architectures: str | list[str],

1020 model_config: ModelConfig,

1021 ) -> bool:

1022 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1024,7 +1024,7 @@ class _ModelRegistry:

1024

1025 def is_attention_free_model(

1026 self,

1027 - architectures: Union[str, list[str]],

1027 + architectures: str | list[str],

1028 model_config: ModelConfig,

1029 ) -> bool:

1030 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1032,7 +1032,7 @@ class _ModelRegistry:

1032

1033 def is_hybrid_model(

1034 self,

1035 - architectures: Union[str, list[str]],

1035 + architectures: str | list[str],

1036 model_config: ModelConfig,

1037 ) -> bool:

1038 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1040,7 +1040,7 @@ class _ModelRegistry:

1040

1041 def is_noops_model(

1042 self,

1043 - architectures: Union[str, list[str]],

1043 + architectures: str | list[str],

1044 model_config: ModelConfig,

1045 ) -> bool:

1046 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1048,7 +1048,7 @@ class _ModelRegistry:

1048

1049 def is_transcription_model(

1050 self,

1051 - architectures: Union[str, list[str]],

1051 + architectures: str | list[str],

1052 model_config: ModelConfig,

1053 ) -> bool:

1054 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1056,7 +1056,7 @@ class _ModelRegistry:

1056

1057 def is_transcription_only_model(

1058 self,

1059 - architectures: Union[str, list[str]],

1059 + architectures: str | list[str],

1060 model_config: ModelConfig,

1061 ) -> bool:

1062 model_cls, _ = self.inspect_model_cls(architectures, model_config)

@@ -1064,7 +1064,7 @@ class _ModelRegistry:

1064

1065 def is_v1_compatible(

1066 self,

1067 - architectures: Union[str, list[str]],

1067 + architectures: str | list[str],

1068 model_config: ModelConfig,

1069 ) -> bool:

1070 model_cls, _ = self.inspect_model_cls(architectures, model_config)

vllm/model_executor/models/roberta.py

+8 -9

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable

5 -from typing import Optional, Union

5

6 import torch

7 from torch import nn

@@ -68,7 +67,7 @@ class RobertaEmbedding(nn.Module):

67 self,

68 input_ids: torch.Tensor,

69 position_ids: torch.Tensor,

71 - inputs_embeds: Optional[torch.Tensor] = None,

70 + inputs_embeds: torch.Tensor | None = None,

71 ) -> torch.Tensor:

72 token_type_ids = _decode_token_type_ids(input_ids)

73

@@ -124,8 +123,8 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

123 self,

124 input_ids: torch.Tensor,

125 positions: torch.Tensor,

127 - intermediate_tensors: Optional[IntermediateTensors] = None,

128 - inputs_embeds: Optional[torch.Tensor] = None,

126 + intermediate_tensors: IntermediateTensors | None = None,

127 + inputs_embeds: torch.Tensor | None = None,

128 ) -> torch.Tensor:

129 # Fix Roberta positions here outside of the CUDA graph.

130 # Because we need the to extract the sequences from

@@ -143,7 +142,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

142

143 def _build_model(

144 self, vllm_config: VllmConfig, prefix: str = ""

146 - ) -> Union[BertModel, BertWithRope]:

145 + ) -> BertModel | BertWithRope:

146 if vllm_config.model_config.hf_config.position_embedding_type == "rotary":

147 return JinaRobertaModel(vllm_config=vllm_config, prefix=prefix)

148 else:

@@ -240,11 +239,11 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):

239

240 def forward(

241 self,

243 - input_ids: Optional[torch.Tensor],

242 + input_ids: torch.Tensor | None,

243 positions: torch.Tensor,

245 - intermediate_tensors: Optional[IntermediateTensors] = None,

246 - inputs_embeds: Optional[torch.Tensor] = None,

247 - token_type_ids: Optional[torch.Tensor] = None,

244 + intermediate_tensors: IntermediateTensors | None = None,

245 + inputs_embeds: torch.Tensor | None = None,

246 + token_type_ids: torch.Tensor | None = None,

247 ) -> torch.Tensor:

248 replace_roberta_positions(

249 input_ids=input_ids, position_ids=positions, padding_idx=self.padding_idx

vllm/model_executor/models/rvl.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Mapping

5 -from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -41,7 +40,7 @@ class RVLDummyInputsBuilder(LlavaDummyInputsBuilder[RVLProcessingInfo]):

40 self,

41 seq_len: int,

42 mm_counts: Mapping[str, int],

44 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

43 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

44 ) -> MultiModalDataDict:

45 num_images = mm_counts.get("image", 0)

46

vllm/model_executor/models/seed_oss.py

+14 -15

@@ -25,7 +25,6 @@

25

26 from collections.abc import Iterable

27 from itertools import islice

28 -from typing import Optional, Union

28

29 import torch

30 from torch import nn

@@ -75,7 +74,7 @@ class SeedOssMLP(nn.Module):

74 hidden_size: int,

75 intermediate_size: int,

76 hidden_act: str,

78 - quant_config: Optional[QuantizationConfig] = None,

77 + quant_config: QuantizationConfig | None = None,

78 prefix: str = "",

79 ) -> None:

80 super().__init__()

@@ -115,9 +114,9 @@ class SeedOssAttention(nn.Module):

114 head_dim: int,

115 max_position: int = 4096 * 32,

116 rope_theta: float = 10000,

118 - cache_config: Optional[CacheConfig] = None,

119 - quant_config: Optional[QuantizationConfig] = None,

120 - rope_scaling: Optional[tuple] = None,

117 + cache_config: CacheConfig | None = None,

118 + quant_config: QuantizationConfig | None = None,

119 + rope_scaling: tuple | None = None,

120 prefix: str = "",

121 attn_type: str = AttentionType.DECODER,

122 ) -> None:

@@ -195,8 +194,8 @@ class SeedOssDecoderLayer(nn.Module):

194 def __init__(

195 self,

196 config: SeedOssConfig,

198 - cache_config: Optional[CacheConfig] = None,

199 - quant_config: Optional[QuantizationConfig] = None,

197 + cache_config: CacheConfig | None = None,

198 + quant_config: QuantizationConfig | None = None,

199 prefix: str = "",

200 ) -> None:

201 super().__init__()

@@ -243,7 +242,7 @@ class SeedOssDecoderLayer(nn.Module):

242 self,

243 positions: torch.Tensor,

244 hidden_states: torch.Tensor,

246 - residual: Optional[torch.Tensor],

245 + residual: torch.Tensor | None,

246 ) -> tuple[torch.Tensor, torch.Tensor]:

247 # Self Attention

248 if residual is None:

@@ -342,9 +341,9 @@ class SeedOssModel(nn.Module):

341 self,

342 input_ids: torch.Tensor,

343 positions: torch.Tensor,

345 - intermediate_tensors: Optional[IntermediateTensors] = None,

346 - inputs_embeds: Optional[torch.Tensor] = None,

347 - ) -> Union[torch.Tensor, IntermediateTensors]:

344 + intermediate_tensors: IntermediateTensors | None = None,

345 + inputs_embeds: torch.Tensor | None = None,

346 + ) -> torch.Tensor | IntermediateTensors:

347 if get_pp_group().is_first_rank:

348 if inputs_embeds is not None:

349 hidden_states = inputs_embeds

@@ -477,9 +476,9 @@ class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

476 self,

477 input_ids: torch.Tensor,

478 positions: torch.Tensor,

480 - intermediate_tensors: Optional[IntermediateTensors] = None,

481 - inputs_embeds: Optional[torch.Tensor] = None,

482 - ) -> Union[torch.Tensor, IntermediateTensors]:

479 + intermediate_tensors: IntermediateTensors | None = None,

480 + inputs_embeds: torch.Tensor | None = None,

481 + ) -> torch.Tensor | IntermediateTensors:

482 hidden_states = self.model(

483 input_ids, positions, intermediate_tensors, inputs_embeds

484 )

@@ -488,7 +487,7 @@ class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

487 def compute_logits(

488 self,

489 hidden_states: torch.Tensor,

491 - ) -> Optional[torch.Tensor]:

490 + ) -> torch.Tensor | None:

491 logits = self.logits_processor(self.lm_head, hidden_states)

492 return logits

493

vllm/model_executor/models/siglip.py

+17 -18

@@ -5,7 +5,6 @@ within a vision language model."""

5

6 import math

7 from collections.abc import Iterable

8 -from typing import Optional, Union

8

9 import torch

10 from torch import nn

@@ -153,7 +152,7 @@ class SiglipAttention(nn.Module):

152 def __init__(

153 self,

154 config: SiglipVisionConfig,

156 - quant_config: Optional[QuantizationConfig] = None,

155 + quant_config: QuantizationConfig | None = None,

156 prefix: str = "",

157 ) -> None:

158 super().__init__()

@@ -211,7 +210,7 @@ class SiglipMLP(nn.Module):

210 def __init__(

211 self,

212 config: SiglipVisionConfig,

214 - quant_config: Optional[QuantizationConfig] = None,

213 + quant_config: QuantizationConfig | None = None,

214 prefix: str = "",

215 ) -> None:

216 super().__init__()

@@ -251,7 +250,7 @@ class SiglipEncoderLayer(nn.Module):

250 def __init__(

251 self,

252 config: SiglipVisionConfig,

254 - quant_config: Optional[QuantizationConfig] = None,

253 + quant_config: QuantizationConfig | None = None,

254 prefix: str = "",

255 ) -> None:

256 super().__init__()

@@ -293,8 +292,8 @@ class SiglipEncoder(nn.Module):

292 def __init__(

293 self,

294 config: SiglipVisionConfig,

296 - quant_config: Optional[QuantizationConfig] = None,

297 - num_hidden_layers_override: Optional[int] = None,

295 + quant_config: QuantizationConfig | None = None,

296 + num_hidden_layers_override: int | None = None,

297 prefix: str = "",

298 ) -> None:

299 super().__init__()

@@ -321,7 +320,7 @@ class SiglipEncoder(nn.Module):

320 self,

321 inputs_embeds: torch.Tensor,

322 return_all_hidden_states: bool,

324 - ) -> Union[torch.Tensor, list[torch.Tensor]]:

323 + ) -> torch.Tensor | list[torch.Tensor]:

324 hidden_states_pool = [inputs_embeds]

325 hidden_states = inputs_embeds

326

@@ -342,7 +341,7 @@ class SiglipMultiheadAttentionPoolingHead(nn.Module):

341 def __init__(

342 self,

343 config: SiglipVisionConfig,

345 - quant_config: Optional[QuantizationConfig] = None,

344 + quant_config: QuantizationConfig | None = None,

345 prefix: str = "",

346 ) -> None:

347 super().__init__()

@@ -375,10 +374,10 @@ class SiglipVisionTransformer(nn.Module):

374 def __init__(

375 self,

376 config: SiglipVisionConfig,

378 - quant_config: Optional[QuantizationConfig] = None,

377 + quant_config: QuantizationConfig | None = None,

378 *,

380 - num_hidden_layers_override: Optional[int] = None,

381 - require_post_norm: Optional[bool] = None,

379 + num_hidden_layers_override: int | None = None,

380 + require_post_norm: bool | None = None,

381 prefix: str = "",

382 ) -> None:

383 super().__init__()

@@ -426,8 +425,8 @@ class SiglipVisionTransformer(nn.Module):

425 pixel_values: torch.Tensor,

426 *,

427 interpolate_pos_encoding: bool = False,

429 - select_layers: Optional[list[int]] = None,

430 - feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

428 + select_layers: list[int] | None = None,

429 + feature_select_strategy: VisionFeatureSelectStrategy | None = None,

430 ) -> torch.Tensor:

431 hidden_states = self.embeddings(

432 pixel_values,

@@ -464,10 +463,10 @@ class SiglipVisionModel(nn.Module):

463 def __init__(

464 self,

465 config: SiglipVisionConfig,

467 - quant_config: Optional[QuantizationConfig] = None,

466 + quant_config: QuantizationConfig | None = None,

467 *,

469 - num_hidden_layers_override: Optional[int] = None,

470 - require_post_norm: Optional[bool] = None,

468 + num_hidden_layers_override: int | None = None,

469 + require_post_norm: bool | None = None,

470 prefix: str = "",

471 ) -> None:

472 super().__init__()

@@ -491,8 +490,8 @@ class SiglipVisionModel(nn.Module):

490 self,

491 pixel_values: torch.Tensor,

492 interpolate_pos_encoding: bool = False,

494 - select_layers: Optional[list[int]] = None,

495 - feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

493 + select_layers: list[int] | None = None,

494 + feature_select_strategy: VisionFeatureSelectStrategy | None = None,

495 ) -> torch.Tensor:

496 return self.vision_model(

497 pixel_values=pixel_values,

vllm/model_executor/models/siglip2navit.py

+9 -10

@@ -4,7 +4,6 @@

4 within a vision language model."""

5

6 from collections.abc import Iterable

7 -from typing import Optional

7

8 import torch

9 from einops import rearrange, repeat

@@ -82,7 +81,7 @@ class Siglip2VisionEmbeddings(nn.Module):

81 def forward(

82 self,

83 pixel_values: torch.FloatTensor,

85 - grid_thws: Optional[torch.LongTensor] = None,

84 + grid_thws: torch.LongTensor | None = None,

85 ) -> torch.Tensor:

86 """

87 Args:

@@ -206,7 +205,7 @@ class Siglip2Attention(nn.Module):

205 def __init__(

206 self,

207 config: Siglip2VisionConfig,

209 - quant_config: Optional[QuantizationConfig] = None,

208 + quant_config: QuantizationConfig | None = None,

209 prefix: str = "",

210 use_data_parallel: bool = False,

211 ):

@@ -275,8 +274,8 @@ class Siglip2Attention(nn.Module):

274 self,

275 hidden_states: torch.Tensor,

276 cu_seqlens: torch.Tensor,

278 - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,

279 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

277 + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,

278 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

279 """Input shape: Batch x Time x Channel"""

280

281 seq_length, embed_dim = hidden_states.shape

@@ -337,7 +336,7 @@ class Siglip2MLP(nn.Module):

336 def __init__(

337 self,

338 config: Siglip2VisionConfig,

340 - quant_config: Optional[QuantizationConfig] = None,

339 + quant_config: QuantizationConfig | None = None,

340 prefix: str = "",

341 use_data_parallel: bool = False,

342 ):

@@ -370,7 +369,7 @@ class Siglip2EncoderLayer(nn.Module):

369 def __init__(

370 self,

371 config: Siglip2VisionConfig,

373 - quant_config: Optional[QuantizationConfig] = None,

372 + quant_config: QuantizationConfig | None = None,

373 prefix: str = "",

374 use_data_parallel: bool = False,

375 ):

@@ -432,7 +431,7 @@ class Siglip2Encoder(nn.Module):

431 def __init__(

432 self,

433 config: Siglip2VisionConfig,

435 - quant_config: Optional[QuantizationConfig] = None,

434 + quant_config: QuantizationConfig | None = None,

435 prefix: str = "",

436 use_data_parallel: bool = False,

437 ):

@@ -616,7 +615,7 @@ class Siglip2VisionTransformer(nn.Module):

615 def __init__(

616 self,

617 config: Siglip2VisionConfig,

619 - quant_config: Optional[QuantizationConfig] = None,

618 + quant_config: QuantizationConfig | None = None,

619 prefix: str = "",

620 use_data_parallel: bool = False,

621 ):

@@ -655,7 +654,7 @@ class Siglip2NavitModel(torch.nn.Module):

654 def __init__(

655 self,

656 config: Siglip2VisionConfig,

658 - quant_config: Optional[QuantizationConfig] = None,

657 + quant_config: QuantizationConfig | None = None,

658 prefix: str = "",

659 use_data_parallel: bool = False,

660 ):

vllm/model_executor/models/skyworkr1v.py

+38 -38

@@ -8,7 +8,7 @@

8 # Licensed under The MIT License [see LICENSE for details]

9 # --------------------------------------------------------

10 from collections.abc import Iterable, Mapping, Sequence

11 -from typing import Annotated, Literal, Optional, Union

11 +from typing import Annotated, Literal, TypeAlias

12

13 import torch

14 import torch.nn as nn

@@ -96,14 +96,14 @@ class SkyworkR1VImageEmbeddingInputs(TensorSchema):

96 type: Literal["image_embeds"] = "image_embeds"

97

98 data: Annotated[

99 - Union[torch.Tensor, list[torch.Tensor]],

99 + torch.Tensor | list[torch.Tensor],

100 TensorShape("ni", "ifs", "hs"),

101 ]

102

103

104 -SkyworkR1VImageInputs = Union[

105 - SkyworkR1VImagePixelInputs, SkyworkR1VImageEmbeddingInputs

106 -]

104 +SkyworkR1VImageInputs: TypeAlias = (

105 + SkyworkR1VImagePixelInputs | SkyworkR1VImageEmbeddingInputs

106 +)

107

108

109 # adapted from https://huggingface.co/Skywork/Skywork-R1V-38B/

@@ -284,9 +284,9 @@ class SkyworkR1VProcessor:

284 config: PretrainedConfig,

285 tokenizer: AnyTokenizer,

286 *,

287 - min_dynamic_patch: Optional[int] = None,

288 - max_dynamic_patch: Optional[int] = None,

289 - dynamic_image_size: Optional[bool] = None,

287 + min_dynamic_patch: int | None = None,

288 + max_dynamic_patch: int | None = None,

289 + dynamic_image_size: bool | None = None,

290 ) -> None:

291 super().__init__()

292

@@ -324,7 +324,7 @@ class SkyworkR1VProcessor:

324 def get_image_repl(

325 self,

326 feature_size: int,

327 - num_patches: Optional[int],

327 + num_patches: int | None,

328 ) -> PromptUpdateDetails[str]:

329 repl_features = IMG_CONTEXT * feature_size

330 repl_full = IMG_START + repl_features + IMG_END

@@ -334,10 +334,10 @@ class SkyworkR1VProcessor:

334 def resolve_min_max_num(

335 self,

336 *,

337 - min_dynamic_patch: Optional[int] = None,

338 - max_dynamic_patch: Optional[int] = None,

339 - dynamic_image_size: Optional[bool] = None,

340 - use_thumbnail: Optional[bool] = None,

337 + min_dynamic_patch: int | None = None,

338 + max_dynamic_patch: int | None = None,

339 + dynamic_image_size: bool | None = None,

340 + use_thumbnail: bool | None = None,

341 ) -> tuple[int, int]:

342 min_dynamic_patch = (

343 self.min_dynamic_patch if min_dynamic_patch is None else min_dynamic_patch

@@ -362,10 +362,10 @@ class SkyworkR1VProcessor:

362 def resolve_target_ratios(

363 self,

364 *,

365 - min_dynamic_patch: Optional[int] = None,

366 - max_dynamic_patch: Optional[int] = None,

367 - dynamic_image_size: Optional[bool] = None,

368 - use_thumbnail: Optional[bool] = None,

365 + min_dynamic_patch: int | None = None,

366 + max_dynamic_patch: int | None = None,

367 + dynamic_image_size: bool | None = None,

368 + use_thumbnail: bool | None = None,

369 ) -> list[tuple[int, int]]:

370 min_num, max_num = self.resolve_min_max_num(

371 min_dynamic_patch=min_dynamic_patch,

@@ -399,9 +399,9 @@ class SkyworkR1VProcessor:

399 def _images_to_pixel_values_lst(

400 self,

401 images: list[Image.Image],

402 - min_dynamic_patch: Optional[int] = None,

403 - max_dynamic_patch: Optional[int] = None,

404 - dynamic_image_size: Optional[bool] = None,

402 + min_dynamic_patch: int | None = None,

403 + max_dynamic_patch: int | None = None,

404 + dynamic_image_size: bool | None = None,

405 ) -> list[torch.Tensor]:

406 min_num, max_num = self.resolve_min_max_num(

407 min_dynamic_patch=min_dynamic_patch,

@@ -423,12 +423,12 @@ class SkyworkR1VProcessor:

423

424 def __call__(

425 self,

426 - text: Optional[Union[str, list[str]]] = None,

427 - images: Optional[Union[Image.Image, list[Image.Image]]] = None,

428 - min_dynamic_patch: Optional[int] = None,

429 - max_dynamic_patch: Optional[int] = None,

430 - dynamic_image_size: Optional[bool] = None,

431 - return_tensors: Optional[Union[str, TensorType]] = None,

426 + text: str | list[str] | None = None,

427 + images: Image.Image | list[Image.Image] | None = None,

428 + min_dynamic_patch: int | None = None,

429 + max_dynamic_patch: int | None = None,

430 + dynamic_image_size: bool | None = None,

431 + return_tensors: str | TensorType | None = None,

432 ) -> BatchFeature:

433 if text is None:

434 text = []

@@ -479,7 +479,7 @@ class SkyworkR1VProcessingInfo(BaseProcessingInfo):

479 **kwargs,

480 )

481

482 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

482 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

483 return {"image": None}

484

485 def get_num_image_tokens(

@@ -487,7 +487,7 @@ class SkyworkR1VProcessingInfo(BaseProcessingInfo):

487 *,

488 image_width: int,

489 image_height: int,

490 - processor: Optional[SkyworkR1VProcessor],

490 + processor: SkyworkR1VProcessor | None,

491 ) -> int:

492 if processor is None:

493 processor = self.get_hf_processor()

@@ -532,7 +532,7 @@ class SkyworkR1VDummyInputsBuilder(BaseDummyInputsBuilder[SkyworkR1VProcessingIn

532 self,

533 seq_len: int,

534 mm_counts: Mapping[str, int],

535 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

535 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

536 ) -> MultiModalDataDict:

537 target_width, target_height = self.info.get_image_size_with_most_features()

538 num_images = mm_counts.get("image", 0)

@@ -650,7 +650,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

650 merge_by_field_config = True

651

652 @classmethod

653 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

653 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

654 if modality.startswith("image"):

655 return "<image>"

656

@@ -715,7 +715,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

715 def _init_vision_model(

716 self,

717 config: PretrainedConfig,

718 - quant_config: Optional[QuantizationConfig],

718 + quant_config: QuantizationConfig | None,

719 *,

720 is_mono: bool,

721 prefix: str,

@@ -784,7 +784,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

784

785 def _parse_and_validate_image_input(

786 self, **kwargs: object

787 - ) -> Optional[SkyworkR1VImageInputs]:

787 + ) -> SkyworkR1VImageInputs | None:

788 pixel_values_flat = kwargs.pop("pixel_values_flat", None)

789 image_num_patches = kwargs.pop("image_num_patches", None)

790 image_embeds = kwargs.pop("image_embeds", None)

@@ -818,7 +818,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

818 def _process_image_input(

819 self,

820 image_input: SkyworkR1VImageInputs,

821 - ) -> Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor, ...]]:

821 + ) -> torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...]:

822 if image_input["type"] == "image_embeds":

823 return image_input["data"]

824

@@ -864,9 +864,9 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

864 def get_input_embeddings(

865 self,

866 input_ids: torch.Tensor,

867 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

867 + multimodal_embeddings: MultiModalEmbeddings | None = None,

868 *,

869 - is_multimodal: Optional[torch.Tensor] = None,

869 + is_multimodal: torch.Tensor | None = None,

870 handle_oov_mm_token: bool = False,

871 ) -> torch.Tensor:

872 if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:

@@ -887,8 +887,8 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

887 self,

888 input_ids: torch.Tensor,

889 positions: torch.Tensor,

890 - intermediate_tensors: Optional[IntermediateTensors] = None,

891 - inputs_embeds: Optional[torch.Tensor] = None,

890 + intermediate_tensors: IntermediateTensors | None = None,

891 + inputs_embeds: torch.Tensor | None = None,

892 **kwargs: object,

893 ) -> IntermediateTensors:

894 if intermediate_tensors is not None:

@@ -913,7 +913,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

913 def compute_logits(

914 self,

915 hidden_states: torch.Tensor,

916 - ) -> Optional[torch.Tensor]:

916 + ) -> torch.Tensor | None:

917 return self.language_model.compute_logits(hidden_states)

918

919 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/smolvlm.py

+1 -4

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 from transformers import SmolVLMProcessor

6

@@ -17,9 +16,7 @@ class SmolVLMProcessingInfo(Idefics3ProcessingInfo):

16 def get_hf_processor(self, **kwargs: object) -> SmolVLMProcessor:

17 return self.ctx.get_hf_processor(SmolVLMProcessor, **kwargs)

18

20 - def _get_image_token(

21 - self, processor: Optional[SmolVLMProcessor]

22 - ) -> tuple[str, str]:

19 + def _get_image_token(self, processor: SmolVLMProcessor | None) -> tuple[str, str]:

20 if processor is None:

21 processor = self.get_hf_processor()

22 image_token = processor.image_token

vllm/model_executor/models/solar.py

+15 -15

@@ -25,7 +25,7 @@

25 """Inference-only Solar model compatible with HuggingFace weights."""

26

27 from collections.abc import Iterable

28 -from typing import Any, Optional, Union

28 +from typing import Any

29

30 import torch

31 from torch import nn

@@ -73,7 +73,7 @@ class SolarMLP(nn.Module):

73 hidden_size: int,

74 intermediate_size: int,

75 hidden_act: str,

76 - quant_config: Optional[QuantizationConfig] = None,

76 + quant_config: QuantizationConfig | None = None,

77 bias: bool = False,

78 prefix: str = "",

79 ) -> None:

@@ -113,11 +113,11 @@ class SolarAttention(nn.Module):

113 num_heads: int,

114 num_kv_heads: int,

115 rope_theta: float = 10000,

116 - rope_scaling: Optional[dict[str, Any]] = None,

116 + rope_scaling: dict[str, Any] | None = None,

117 max_position_embeddings: int = 8192,

118 - quant_config: Optional[QuantizationConfig] = None,

118 + quant_config: QuantizationConfig | None = None,

119 bias: bool = False,

120 - cache_config: Optional[CacheConfig] = None,

120 + cache_config: CacheConfig | None = None,

121 prefix: str = "",

122 ) -> None:

123 super().__init__()

@@ -197,8 +197,8 @@ class SolarDecoderLayer(nn.Module):

197 def __init__(

198 self,

199 config: PretrainedConfig,

200 - cache_config: Optional[CacheConfig] = None,

201 - quant_config: Optional[QuantizationConfig] = None,

200 + cache_config: CacheConfig | None = None,

201 + quant_config: QuantizationConfig | None = None,

202 prefix: str = "",

203 ) -> None:

204 super().__init__()

@@ -250,7 +250,7 @@ class SolarDecoderLayer(nn.Module):

250 self,

251 positions: torch.Tensor,

252 hidden_states: torch.Tensor,

253 - residual: Optional[torch.Tensor],

253 + residual: torch.Tensor | None,

254 ) -> tuple[torch.Tensor, torch.Tensor]:

255 # Self Attention

256 if residual is None:

@@ -322,11 +322,11 @@ class SolarModel(nn.Module):

322

323 def forward(

324 self,

325 - input_ids: Optional[torch.Tensor],

325 + input_ids: torch.Tensor | None,

326 positions: torch.Tensor,

327 - intermediate_tensors: Optional[IntermediateTensors],

328 - inputs_embeds: Optional[torch.Tensor] = None,

329 - ) -> Union[torch.Tensor, IntermediateTensors]:

327 + intermediate_tensors: IntermediateTensors | None,

328 + inputs_embeds: torch.Tensor | None = None,

329 + ) -> torch.Tensor | IntermediateTensors:

330 if get_pp_group().is_first_rank:

331 if inputs_embeds is not None:

332 hidden_states = inputs_embeds

@@ -501,9 +501,9 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

501 self,

502 input_ids: torch.Tensor,

503 positions: torch.Tensor,

504 - intermediate_tensors: Optional[IntermediateTensors] = None,

505 - inputs_embeds: Optional[torch.Tensor] = None,

506 - ) -> Union[torch.Tensor, IntermediateTensors]:

504 + intermediate_tensors: IntermediateTensors | None = None,

505 + inputs_embeds: torch.Tensor | None = None,

506 + ) -> torch.Tensor | IntermediateTensors:

507 model_output = self.model(

508 input_ids, positions, intermediate_tensors, inputs_embeds

509 )

vllm/model_executor/models/stablelm.py

+12 -13

@@ -24,7 +24,6 @@ model compatible with HuggingFace weights."""

24

25 from collections.abc import Iterable

26 from itertools import islice

27 -from typing import Optional, Union

27

28 import torch

29 from torch import nn

@@ -63,7 +62,7 @@ class StablelmMLP(nn.Module):

62 def __init__(

63 self,

64 config: StableLmConfig,

66 - quant_config: Optional[QuantizationConfig] = None,

65 + quant_config: QuantizationConfig | None = None,

66 prefix: str = "",

67 ) -> None:

68 super().__init__()

@@ -97,8 +96,8 @@ class StablelmAttention(nn.Module):

96 def __init__(

97 self,

98 config: StableLmConfig,

100 - cache_config: Optional[CacheConfig] = None,

101 - quant_config: Optional[QuantizationConfig] = None,

99 + cache_config: CacheConfig | None = None,

100 + quant_config: QuantizationConfig | None = None,

101 prefix: str = "",

102 ) -> None:

103 super().__init__()

@@ -184,8 +183,8 @@ class StablelmDecoderLayer(nn.Module):

183 def __init__(

184 self,

185 config: StableLmConfig,

187 - cache_config: Optional[CacheConfig] = None,

188 - quant_config: Optional[QuantizationConfig] = None,

186 + cache_config: CacheConfig | None = None,

187 + quant_config: QuantizationConfig | None = None,

188 prefix: str = "",

189 ) -> None:

190 super().__init__()

@@ -254,9 +253,9 @@ class StableLMEpochModel(nn.Module):

253 self,

254 input_ids: torch.Tensor,

255 positions: torch.Tensor,

257 - intermediate_tensors: Optional[IntermediateTensors],

258 - inputs_embeds: Optional[torch.Tensor] = None,

259 - ) -> Union[torch.Tensor, IntermediateTensors]:

256 + intermediate_tensors: IntermediateTensors | None,

257 + inputs_embeds: torch.Tensor | None = None,

258 + ) -> torch.Tensor | IntermediateTensors:

259 if get_pp_group().is_first_rank:

260 if inputs_embeds is not None:

261 hidden_states = inputs_embeds

@@ -340,9 +339,9 @@ class StablelmForCausalLM(nn.Module, SupportsPP):

339 self,

340 input_ids: torch.Tensor,

341 positions: torch.Tensor,

343 - intermediate_tensors: Optional[IntermediateTensors] = None,

344 - inputs_embeds: Optional[torch.Tensor] = None,

345 - ) -> Union[torch.Tensor, IntermediateTensors]:

342 + intermediate_tensors: IntermediateTensors | None = None,

343 + inputs_embeds: torch.Tensor | None = None,

344 + ) -> torch.Tensor | IntermediateTensors:

345 hidden_states = self.model(

346 input_ids, positions, intermediate_tensors, inputs_embeds

347 )

@@ -351,7 +350,7 @@ class StablelmForCausalLM(nn.Module, SupportsPP):

350 def compute_logits(

351 self,

352 hidden_states: torch.Tensor,

354 - ) -> Optional[torch.Tensor]:

353 + ) -> torch.Tensor | None:

354 logits = self.logits_processor(self.lm_head, hidden_states)

355 return logits

356

vllm/model_executor/models/starcoder2.py

+12 -13

@@ -23,7 +23,6 @@

23

24 from collections.abc import Iterable

25 from itertools import islice

26 -from typing import Optional, Union

26

27 import torch

28 from torch import nn

@@ -67,8 +66,8 @@ class Starcoder2Attention(nn.Module):

66 def __init__(

67 self,

68 config: Starcoder2Config,

70 - cache_config: Optional[CacheConfig] = None,

71 - quant_config: Optional[QuantizationConfig] = None,

69 + cache_config: CacheConfig | None = None,

70 + quant_config: QuantizationConfig | None = None,

71 prefix: str = "",

72 ):

73 super().__init__()

@@ -147,7 +146,7 @@ class Starcoder2MLP(nn.Module):

146 def __init__(

147 self,

148 config: Starcoder2Config,

150 - quant_config: Optional[QuantizationConfig] = None,

149 + quant_config: QuantizationConfig | None = None,

150 prefix: str = "",

151 ):

152 super().__init__()

@@ -178,8 +177,8 @@ class Starcoder2DecoderLayer(nn.Module):

177 def __init__(

178 self,

179 config: Starcoder2Config,

181 - cache_config: Optional[CacheConfig] = None,

182 - quant_config: Optional[QuantizationConfig] = None,

180 + cache_config: CacheConfig | None = None,

181 + quant_config: QuantizationConfig | None = None,

182 prefix: str = "",

183 ):

184 super().__init__()

@@ -258,9 +257,9 @@ class Starcoder2Model(nn.Module):

257 self,

258 input_ids: torch.Tensor,

259 positions: torch.Tensor,

261 - intermediate_tensors: Optional[IntermediateTensors],

262 - inputs_embeds: Optional[torch.Tensor] = None,

263 - ) -> Union[torch.Tensor, IntermediateTensors]:

260 + intermediate_tensors: IntermediateTensors | None,

261 + inputs_embeds: torch.Tensor | None = None,

262 + ) -> torch.Tensor | IntermediateTensors:

263 if get_pp_group().is_first_rank:

264 if inputs_embeds is not None:

265 hidden_states = inputs_embeds

@@ -347,9 +346,9 @@ class Starcoder2ForCausalLM(nn.Module, SupportsPP):

346 self,

347 input_ids: torch.Tensor,

348 positions: torch.Tensor,

350 - intermediate_tensors: Optional[IntermediateTensors] = None,

351 - inputs_embeds: Optional[torch.Tensor] = None,

352 - ) -> Union[torch.Tensor, IntermediateTensors]:

349 + intermediate_tensors: IntermediateTensors | None = None,

350 + inputs_embeds: torch.Tensor | None = None,

351 + ) -> torch.Tensor | IntermediateTensors:

352 hidden_states = self.model(

353 input_ids, positions, intermediate_tensors, inputs_embeds

354 )

@@ -358,7 +357,7 @@ class Starcoder2ForCausalLM(nn.Module, SupportsPP):

357 def compute_logits(

358 self,

359 hidden_states: torch.Tensor,

361 - ) -> Optional[torch.Tensor]:

360 + ) -> torch.Tensor | None:

361 logits = self.logits_processor(self.lm_head, hidden_states)

362 return logits

363

vllm/model_executor/models/step3_text.py

+14 -14

@@ -4,7 +4,7 @@

4

5 from collections.abc import Iterable

6 from itertools import islice

7 -from typing import Any, Optional

7 +from typing import Any

8

9 import torch

10 from torch import nn

@@ -54,7 +54,7 @@ class FusedMoEBlock(nn.Module):

54 def __init__(

55 self,

56 config: ModelConfig,

57 - quant_config: Optional[QuantizationConfig] = None,

57 + quant_config: QuantizationConfig | None = None,

58 prefix: str = "",

59 ):

60 super().__init__()

@@ -106,7 +106,7 @@ class Step3TextMLP(nn.Module):

106 hidden_size: int,

107 intermediate_size: int,

108 hidden_act: str,

109 - quant_config: Optional[QuantizationConfig] = None,

109 + quant_config: QuantizationConfig | None = None,

110 prefix: str = "",

111 ) -> None:

112 super().__init__()

@@ -146,12 +146,12 @@ class Step3TextAttention(nn.Module):

146 num_kv_heads: int,

147 norm_eps: float,

148 rope_theta: int,

149 - share_q_dim: Optional[int] = None,

150 - rope_scaling: Optional[dict[str, Any]] = None,

149 + share_q_dim: int | None = None,

150 + rope_scaling: dict[str, Any] | None = None,

151 max_position_embedding: int = 8192,

152 head_dim: int = 256,

153 - cache_config: Optional[CacheConfig] = None,

154 - quant_config: Optional[QuantizationConfig] = None,

153 + cache_config: CacheConfig | None = None,

154 + quant_config: QuantizationConfig | None = None,

155 prefix: str = "",

156 ):

157 super().__init__()

@@ -229,8 +229,8 @@ class Step3TextDecoderLayer(nn.Module):

229 def __init__(

230 self,

231 config: ModelConfig,

232 - cache_config: Optional[CacheConfig] = None,

233 - quant_config: Optional[QuantizationConfig] = None,

232 + cache_config: CacheConfig | None = None,

233 + quant_config: QuantizationConfig | None = None,

234 prefix: str = "",

235 ) -> None:

236 super().__init__()

@@ -291,7 +291,7 @@ class Step3TextDecoderLayer(nn.Module):

291 self,

292 positions: torch.Tensor,

293 hidden_states: torch.Tensor,

294 - residual: Optional[torch.Tensor],

294 + residual: torch.Tensor | None,

295 ) -> tuple[torch.Tensor, torch.Tensor]:

296 if residual is None:

297 residual = hidden_states

@@ -362,8 +362,8 @@ class Step3TextModel(nn.Module):

362 self,

363 input_ids: torch.Tensor,

364 positions: torch.Tensor,

365 - intermediate_tensors: Optional[IntermediateTensors] = None,

366 - inputs_embeds: Optional[torch.Tensor] = None,

365 + intermediate_tensors: IntermediateTensors | None = None,

366 + inputs_embeds: torch.Tensor | None = None,

367 ) -> torch.Tensor:

368 if get_pp_group().is_first_rank:

369 if inputs_embeds is not None:

@@ -436,8 +436,8 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):

436 self,

437 input_ids: torch.Tensor,

438 positions: torch.Tensor,

439 - intermediate_tensors: Optional[IntermediateTensors] = None,

440 - inputs_embeds: Optional[torch.Tensor] = None,

439 + intermediate_tensors: IntermediateTensors | None = None,

440 + inputs_embeds: torch.Tensor | None = None,

441 ):

442 hidden_states = self.model(

443 input_ids, positions, intermediate_tensors, inputs_embeds

vllm/model_executor/models/step3_vl.py

+22 -22

@@ -4,7 +4,7 @@ import math

4 from collections.abc import Iterable, Mapping, Sequence

5 from itertools import product

6 from math import ceil, sqrt

7 -from typing import Annotated, Any, Literal, Optional, Union

7 +from typing import Annotated, Any, Literal, TypeAlias

8

9 import numpy as np

10 import torch

@@ -71,7 +71,7 @@ class Step3VLImagePixelInputs(TensorSchema):

71 type: Literal["pixel_values"]

72 pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]

73 patch_pixel_values: Annotated[

74 - Optional[torch.Tensor], TensorShape("bnp", 3, "hp", "wp")

74 + torch.Tensor | None, TensorShape("bnp", 3, "hp", "wp")

75 ]

76 num_patches: Annotated[torch.Tensor, TensorShape("bn")]

77

@@ -88,7 +88,7 @@ class Step3VLImageEmbeddingInputs(TensorSchema):

88 data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")]

89

90

91 -Step3VLImageInputs = Union[Step3VLImagePixelInputs, Step3VLImageEmbeddingInputs]

91 +Step3VLImageInputs: TypeAlias = Step3VLImagePixelInputs | Step3VLImageEmbeddingInputs

92

93 ImageWithPatches = tuple[Image.Image, list[Image.Image], list[int] | None]

94

@@ -409,7 +409,7 @@ class Step3VLProcessor:

409 self,

410 num_images: int,

411 num_patches: int,

412 - patch_new_line_idx: Optional[list[bool]],

412 + patch_new_line_idx: list[bool] | None,

413 ) -> tuple[str, list[int]]:

414 if num_patches > 0:

415 patch_repl, patch_repl_ids = self._get_patch_repl(

@@ -438,9 +438,9 @@ class Step3VLProcessor:

438

439 def __call__(

440 self,

441 - text: Optional[Union[str, list[str]]] = None,

442 - images: Optional[Union[Image.Image, list[Image.Image]]] = None,

443 - return_tensors: Optional[Union[str, TensorType]] = None,

441 + text: str | list[str] | None = None,

442 + images: Image.Image | list[Image.Image] | None = None,

443 + return_tensors: str | TensorType | None = None,

444 ) -> BatchFeature:

445 if text is None:

446 text = []

@@ -513,7 +513,7 @@ class Step3VLProcessingInfo(BaseProcessingInfo):

513 self.get_tokenizer(),

514 )

515

516 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

516 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

517 return {"image": None}

518

519 def get_max_image_tokens(self) -> int:

@@ -556,7 +556,7 @@ class Step3VLDummyInputsBuilder(BaseDummyInputsBuilder[Step3VLProcessingInfo]):

556 self,

557 seq_len: int,

558 mm_counts: Mapping[str, int],

559 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

559 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

560 ) -> MultiModalDataDict:

561 target_width, target_height = self.info.get_image_size_with_most_features()

562 num_images = mm_counts.get("image", 0)

@@ -716,7 +716,7 @@ class Step3VisionAttention(nn.Module):

716 def __init__(

717 self,

718 config,

719 - quant_config: Optional[QuantizationConfig] = None,

719 + quant_config: QuantizationConfig | None = None,

720 prefix: str = "",

721 use_data_parallel: bool = False,

722 ):

@@ -778,7 +778,7 @@ class Step3VisionMLP(nn.Module):

778 def __init__(

779 self,

780 config,

781 - quant_config: Optional[QuantizationConfig] = None,

781 + quant_config: QuantizationConfig | None = None,

782 prefix: str = "",

783 use_data_parallel: bool = False,

784 ):

@@ -813,7 +813,7 @@ class Step3VisionEncoderLayer(nn.Module):

813 def __init__(

814 self,

815 config: Step3VisionEncoderConfig,

816 - quant_config: Optional[QuantizationConfig] = None,

816 + quant_config: QuantizationConfig | None = None,

817 prefix: str = "",

818 use_data_parallel: bool = False,

819 ):

@@ -848,7 +848,7 @@ class Step3VisionEncoder(nn.Module):

848 def __init__(

849 self,

850 config: Step3VisionEncoderConfig,

851 - quant_config: Optional[QuantizationConfig] = None,

851 + quant_config: QuantizationConfig | None = None,

852 prefix: str = "",

853 use_data_parallel: bool = False,

854 ):

@@ -881,7 +881,7 @@ class Step3VisionTransformer(nn.Module):

881 def __init__(

882 self,

883 config: Step3VisionEncoderConfig,

884 - quant_config: Optional[QuantizationConfig] = None,

884 + quant_config: QuantizationConfig | None = None,

885 prefix: str = "",

886 use_data_parallel: bool = False,

887 ):

@@ -927,7 +927,7 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

927 supports_encoder_tp_data = True

928

929 @classmethod

930 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

930 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

931 if modality.startswith("image"):

932 return "<im_patch>"

933

@@ -994,7 +994,7 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

994

995 def _parse_and_validate_image_input(

996 self, **kwargs: object

997 - ) -> Optional[Step3VLImageInputs]:

997 + ) -> Step3VLImageInputs | None:

998 pixel_values = kwargs.pop("pixel_values", None)

999 patch_pixel_values = kwargs.pop("patch_pixel_values", None)

1000 num_patches = kwargs.pop("num_patches", None)

@@ -1085,9 +1085,9 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

1085 def get_input_embeddings(

1086 self,

1087 input_ids: torch.Tensor,

1088 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1088 + multimodal_embeddings: MultiModalEmbeddings | None = None,

1089 *,

1090 - is_multimodal: Optional[torch.Tensor] = None,

1090 + is_multimodal: torch.Tensor | None = None,

1091 # Multi-modal token ID may exceed vocab size

1092 handle_oov_mm_token: bool = True,

1093 ) -> torch.Tensor:

@@ -1106,10 +1106,10 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

1106 self,

1107 input_ids: torch.Tensor,

1108 positions: torch.Tensor,

1109 - intermediate_tensors: Optional[IntermediateTensors] = None,

1110 - inputs_embeds: Optional[torch.Tensor] = None,

1109 + intermediate_tensors: IntermediateTensors | None = None,

1110 + inputs_embeds: torch.Tensor | None = None,

1111 **kwargs: object,

1112 - ) -> Union[torch.Tensor, IntermediateTensors]:

1112 + ) -> torch.Tensor | IntermediateTensors:

1113 if intermediate_tensors is not None:

1114 inputs_embeds = None

1115 elif inputs_embeds is None:

@@ -1130,7 +1130,7 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

1130 def compute_logits(

1131 self,

1132 hidden_states: torch.Tensor,

1133 - ) -> Optional[torch.Tensor]:

1133 + ) -> torch.Tensor | None:

1134 return self.language_model.compute_logits(hidden_states)

1135

1136 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

vllm/model_executor/models/swin.py

+25 -26

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Iterable

5 -from typing import Optional

5

6 import torch

7 import torch.nn as nn

@@ -28,7 +27,7 @@ class SwinSelfAttention(nn.Module):

27 dim: int,

28 num_heads: int,

29 window_size: int,

31 - quant_config: Optional[QuantizationConfig] = None,

30 + quant_config: QuantizationConfig | None = None,

31 prefix: str = "",

32 ) -> None:

33 super().__init__()

@@ -102,9 +101,9 @@ class SwinSelfAttention(nn.Module):

101 def forward(

102 self,

103 hidden_states: torch.Tensor,

105 - attention_mask: Optional[torch.FloatTensor] = None,

106 - head_mask: Optional[torch.FloatTensor] = None,

107 - output_attentions: Optional[bool] = False,

104 + attention_mask: torch.FloatTensor | None = None,

105 + head_mask: torch.FloatTensor | None = None,

106 + output_attentions: bool | None = False,

107 ) -> tuple[torch.Tensor, ...]:

108 batch_size, dim, num_channels = hidden_states.shape

109

@@ -155,7 +154,7 @@ class SwinSelfOutput(nn.Module):

154 self,

155 config: SwinConfig,

156 dim: int,

158 - quant_config: Optional[QuantizationConfig] = None,

157 + quant_config: QuantizationConfig | None = None,

158 prefix: str = "",

159 ) -> None:

160 super().__init__()

@@ -181,7 +180,7 @@ class SwinAttention(nn.Module):

180 dim: int,

181 num_heads: int,

182 window_size: int,

184 - quant_config: Optional[QuantizationConfig] = None,

183 + quant_config: QuantizationConfig | None = None,

184 prefix: str = "",

185 ) -> None:

186 super().__init__()

@@ -201,9 +200,9 @@ class SwinAttention(nn.Module):

200 def forward(

201 self,

202 hidden_states: torch.Tensor,

204 - attention_mask: Optional[torch.FloatTensor] = None,

205 - head_mask: Optional[torch.FloatTensor] = None,

206 - output_attentions: Optional[bool] = False,

203 + attention_mask: torch.FloatTensor | None = None,

204 + head_mask: torch.FloatTensor | None = None,

205 + output_attentions: bool | None = False,

206 ) -> tuple[torch.Tensor]:

207 self_outputs = self.self(

208 hidden_states, attention_mask, head_mask, output_attentions

@@ -218,7 +217,7 @@ class SwinIntermediate(nn.Module):

217 self,

218 config: SwinConfig,

219 dim: int,

221 - quant_config: Optional[QuantizationConfig] = None,

220 + quant_config: QuantizationConfig | None = None,

221 prefix: str = "",

222 ) -> None:

223 super().__init__()

@@ -241,7 +240,7 @@ class SwinOutput(nn.Module):

240 self,

241 config: SwinConfig,

242 dim: int,

244 - quant_config: Optional[QuantizationConfig] = None,

243 + quant_config: QuantizationConfig | None = None,

244 prefix: str = "",

245 ) -> None:

246 super().__init__()

@@ -266,7 +265,7 @@ class SwinLayer(HFSwinLayer):

265 num_heads: int,

266 drop_path_rate: float = 0.0,

267 shift_size: int = 0,

269 - quant_config: Optional[QuantizationConfig] = None,

268 + quant_config: QuantizationConfig | None = None,

269 prefix: str = "",

270 ) -> None:

271 super().__init__(

@@ -303,8 +302,8 @@ class SwinStage(nn.Module):

302 depth: int,

303 num_heads: int,

304 drop_path: list[float],

306 - downsample: Optional[SwinPatchMerging] = None,

307 - quant_config: Optional[QuantizationConfig] = None,

305 + downsample: SwinPatchMerging | None = None,

306 + quant_config: QuantizationConfig | None = None,

307 prefix: str = "",

308 ) -> None:

309 super().__init__()

@@ -340,9 +339,9 @@ class SwinStage(nn.Module):

339 self,

340 hidden_states: torch.Tensor,

341 input_dimensions: tuple[int, int],

343 - head_mask: Optional[torch.FloatTensor] = None,

344 - output_attentions: Optional[bool] = False,

345 - always_partition: Optional[bool] = False,

342 + head_mask: torch.FloatTensor | None = None,

343 + output_attentions: bool | None = False,

344 + always_partition: bool | None = False,

345 ) -> tuple[torch.Tensor]:

346 height, width = input_dimensions

347 for i, layer_module in enumerate(self.blocks):

@@ -384,7 +383,7 @@ class SwinEncoder(nn.Module):

383 self,

384 config: SwinConfig,

385 grid_size: int,

387 - quant_config: Optional[QuantizationConfig] = None,

386 + quant_config: QuantizationConfig | None = None,

387 prefix: str = "",

388 ) -> None:

389 super().__init__()

@@ -426,9 +425,9 @@ class SwinEncoder(nn.Module):

425 self,

426 hidden_states: torch.Tensor,

427 input_dimensions: tuple[int, int],

429 - head_mask: Optional[torch.FloatTensor] = None,

430 - output_attentions: Optional[bool] = False,

431 - always_partition: Optional[bool] = False,

428 + head_mask: torch.FloatTensor | None = None,

429 + output_attentions: bool | None = False,

430 + always_partition: bool | None = False,

431 ) -> tuple[torch.Tensor]:

432 for i, layer_module in enumerate(self.layers):

433 layer_head_mask = head_mask[i] if head_mask is not None else None

@@ -455,7 +454,7 @@ class SwinModel(nn.Module):

454 def __init__(

455 self,

456 config: SwinConfig,

458 - quant_config: Optional[QuantizationConfig] = None,

457 + quant_config: QuantizationConfig | None = None,

458 prefix: str = "",

459 ) -> None:

460 super().__init__()

@@ -473,9 +472,9 @@ class SwinModel(nn.Module):

472

473 def forward(

474 self,

476 - pixel_values: Optional[torch.FloatTensor] = None,

477 - head_mask: Optional[torch.FloatTensor] = None,

478 - output_attentions: Optional[bool] = None,

475 + pixel_values: torch.FloatTensor | None = None,

476 + head_mask: torch.FloatTensor | None = None,

477 + output_attentions: bool | None = None,

478 ) -> tuple[torch.Tensor]:

479 embedding_output, input_dimensions = self.embeddings(pixel_values)

480

vllm/model_executor/models/tarsier.py

+24 -23

@@ -3,7 +3,7 @@

3

4 import math

5 from collections.abc import Iterable, Mapping, Sequence

6 -from typing import Annotated, Final, Literal, Optional, Protocol, TypeVar, Union

6 +from typing import Annotated, Final, Literal, Protocol, TypeAlias, TypeVar

7

8 import torch

9 import torch.nn as nn

@@ -81,7 +81,7 @@ class TarsierImageEmbeddingInputs(TensorSchema):

81 data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]

82

83

84 -TarsierImageInputs = Union[TarsierImagePixelInputs, TarsierImageEmbeddingInputs]

84 +TarsierImageInputs: TypeAlias = TarsierImagePixelInputs | TarsierImageEmbeddingInputs

85

86

87 class TarsierHfConfig(Protocol): # Based on the Tarsier's LlavaConfig

@@ -89,7 +89,7 @@ class TarsierHfConfig(Protocol): # Based on the Tarsier's LlavaConfig

89 text_config: Final[PretrainedConfig] # Added from Tarsier's LlavaConfig

90 image_token_index: Final[int]

91 vision_feature_select_strategy: Final[str]

92 - vision_feature_layer: Final[Union[int, list[int]]]

92 + vision_feature_layer: Final[int | list[int]]

93 projector_hidden_act: Final[str]

94 image_newline_idx: Final[int]

95 image_new_idx: Final[int]

@@ -109,9 +109,10 @@ class TarsierProcessor(LlavaProcessor):

109 def __call__(

110 self,

111 images: ImageInput = None,

112 - text: Union[

113 - TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]

114 - ] = None,

112 + text: TextInput

113 + | PreTokenizedInput

114 + | list[TextInput]

115 + | list[PreTokenizedInput] = None,

116 audio=None,

117 videos=None,

118 **kwargs: Unpack[TarsierProcessorKwargs],

@@ -173,7 +174,7 @@ class TarsierMultiModalProjector(nn.Module):

174 text_hidden_size: int,

175 projector_hidden_act: str,

176 multimodal_projector_bias: bool,

176 - quant_config: Optional[QuantizationConfig] = None,

177 + quant_config: QuantizationConfig | None = None,

178 prefix: str = "",

179 ):

180 super().__init__()

@@ -215,7 +216,7 @@ class TarsierProcessingInfo(BaseProcessingInfo):

216

217 return self.ctx.get_hf_processor(TarsierProcessor, **kwargs)

218

218 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

219 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

220 return {"image": None}

221

222 def get_num_image_tokens(

@@ -331,7 +332,7 @@ def _build_tarsier_hf_processor(

332 info: _I_Tarsier,

333 dummy_inputs: BaseDummyInputsBuilder[_I_Tarsier],

334 *,

334 - cache: Optional[BaseMultiModalProcessorCache] = None,

335 + cache: BaseMultiModalProcessorCache | None = None,

336 ) -> BaseMultiModalProcessor:

337 if isinstance(info, TarsierProcessingInfo):

338 return TarsierMultiModalProcessor(

@@ -344,11 +345,11 @@ def _build_tarsier_hf_processor(

345

346 def init_vision_tower_for_tarsier(

347 hf_config: TarsierHfConfig, # Use the Tarsier specific config protocol

347 - quant_config: Optional[QuantizationConfig],

348 + quant_config: QuantizationConfig | None,

349 *,

349 - require_post_norm: Optional[bool] = None,

350 + require_post_norm: bool | None = None,

351 prefix: str = "",

351 -) -> Union[CLIPVisionModel, SiglipVisionModel]:

352 +) -> CLIPVisionModel | SiglipVisionModel:

353 vision_config = hf_config.vision_config

354

355 feature_layers = hf_config.vision_feature_layer

@@ -407,7 +408,7 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

408 }

409

410 @classmethod

410 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

411 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

412 if modality.startswith("image"):

413 return "<image>"

414

@@ -456,7 +457,7 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

457

458 def _parse_and_validate_image_input(

459 self, **kwargs: object

459 - ) -> Optional[TarsierImageInputs]:

460 + ) -> TarsierImageInputs | None:

461 pixel_values = kwargs.pop("pixel_values", None)

462 image_embeds = kwargs.pop("image_embeds", None)

463

@@ -479,9 +480,9 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

480

481 def _image_pixels_to_features(

482 self,

482 - vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

483 - pixel_values: Union[torch.Tensor, list[torch.Tensor]],

484 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

483 + vision_tower: CLIPVisionModel | SiglipVisionModel,

484 + pixel_values: torch.Tensor | list[torch.Tensor],

485 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

486 # From vLLM LLaVA, vision tower output handling

487 return vision_tower(

488 pixel_values,

@@ -540,7 +541,7 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

541 def _process_image_pixels(

542 self,

543 inputs: TarsierImagePixelInputs,

543 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

544 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

545 assert self.vision_tower is not None

546 pixel_values = inputs["pixel_values"]

547 image_features_selected = self._image_pixels_to_features(

@@ -559,7 +560,7 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

560 def _process_image_input(

561 self,

562 image_input: TarsierImageInputs,

562 - ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

563 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:

564 if image_input["type"] == "image_embeds":

565 projected_features = image_input["data"]

566 if isinstance(projected_features, torch.Tensor):

@@ -585,10 +586,10 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

586 self,

587 input_ids: torch.Tensor,

588 positions: torch.Tensor,

588 - intermediate_tensors: Optional[IntermediateTensors] = None,

589 - inputs_embeds: Optional[torch.Tensor] = None,

589 + intermediate_tensors: IntermediateTensors | None = None,

590 + inputs_embeds: torch.Tensor | None = None,

591 **kwargs: object,

591 - ) -> Union[torch.Tensor, IntermediateTensors]:

592 + ) -> torch.Tensor | IntermediateTensors:

593 if intermediate_tensors is not None:

594 inputs_embeds = None

595 elif inputs_embeds is None:

@@ -610,7 +611,7 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

611 def compute_logits(

612 self,

613 hidden_states: torch.Tensor,

613 - ) -> Optional[torch.Tensor]:

614 + ) -> torch.Tensor | None:

615 return self.language_model.compute_logits(hidden_states)

616

617 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/terratorch.py

+16 -16

@@ -18,8 +18,8 @@

18 """Wrapper around `Terratorch` models"""

19

20 from collections import OrderedDict

21 -from collections.abc import Iterable, Mapping, Sequence

22 -from typing import Any, Callable, Optional, Union

21 +from collections.abc import Callable, Iterable, Mapping, Sequence

22 +from typing import Any

23

24 import torch

25 import torch.nn as nn

@@ -96,7 +96,7 @@ def _terratorch_field_factory(

96

97

98 class TerratorchProcessingInfo(BaseProcessingInfo):

99 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

99 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

100 return {"image": None}

101

102

@@ -114,7 +114,7 @@ class TerratorchInputBuilder(BaseDummyInputsBuilder[TerratorchProcessingInfo]):

114 self,

115 seq_len: int,

116 mm_counts: Mapping[str, int],

117 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

117 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

118 ) -> MultiModalDataDict:

119 # Dummy data is generated based on the 'input' section

120 # defined in the HF configuration file

@@ -136,8 +136,8 @@ class TerratorchMultiModalDataParser(MultiModalDataParser):

136

137 def _parse_image_data(

138 self,

139 - data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],

140 - ) -> Optional[ModalityDataItems[Any, Any]]:

139 + data: dict[str, torch.Tensor] | ModalityData[ImageItem],

140 + ) -> ModalityDataItems[Any, Any] | None:

141 if isinstance(data, dict):

142 terratorch_fields = _terratorch_field_names(self._pretrained_cfg)

143

@@ -157,7 +157,7 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor):

157 info: TerratorchProcessingInfo,

158 dummy_inputs: "BaseDummyInputsBuilder[TerratorchProcessingInfo]",

159 *,

160 - cache: Optional[MultiModalProcessorOnlyCache] = None,

160 + cache: MultiModalProcessorOnlyCache | None = None,

161 ) -> None:

162 self.pretrained_cfg = info.get_hf_config().to_dict()["pretrained_cfg"]

163 super().__init__(info=info, dummy_inputs=dummy_inputs, cache=cache)

@@ -182,11 +182,11 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor):

182

183 def apply(

184 self,

185 - prompt: Union[str, list[int]],

185 + prompt: str | list[int],

186 mm_data: MultiModalDataDict,

187 hf_processor_mm_kwargs: Mapping[str, object],

188 - tokenization_kwargs: Optional[Mapping[str, object]] = None,

189 - mm_uuids: Optional[MultiModalUUIDDict] = None,

188 + tokenization_kwargs: Mapping[str, object] | None = None,

189 + mm_uuids: MultiModalUUIDDict | None = None,

190 ) -> MultiModalInputs:

191 if "image" in mm_data:

192 image_data = mm_data["image"]

@@ -232,7 +232,7 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):

232 is_pooling_model = True

233

234 @classmethod

235 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

235 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

236 if modality.startswith("image"):

237 return None

238

@@ -256,9 +256,9 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):

256 def get_input_embeddings(

257 self,

258 input_ids: torch.Tensor,

259 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

259 + multimodal_embeddings: MultiModalEmbeddings | None = None,

260 *,

261 - is_multimodal: Optional[torch.Tensor] = None,

261 + is_multimodal: torch.Tensor | None = None,

262 handle_oov_mm_token: bool = False,

263 ) -> torch.Tensor:

264 # We do not really use any input tokens and therefore no embeddings

@@ -269,10 +269,10 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):

269

270 def forward(

271 self,

272 - input_ids: Optional[torch.Tensor],

272 + input_ids: torch.Tensor | None,

273 positions: torch.Tensor,

274 - intermediate_tensors: Optional[IntermediateTensors] = None,

275 - inputs_embeds: Optional[torch.Tensor] = None,

274 + intermediate_tensors: IntermediateTensors | None = None,

275 + inputs_embeds: torch.Tensor | None = None,

276 **kwargs: object,

277 ):

278 model_output = self.inference_runner.forward(**kwargs)

vllm/model_executor/models/transformers.py

+23 -23

@@ -19,7 +19,7 @@

19 from collections.abc import Iterable, Mapping

20 from contextlib import contextmanager

21 from pathlib import Path

22 -from typing import Literal, Optional, Union

22 +from typing import Literal

23

24 import regex as re

25 import torch

@@ -103,9 +103,9 @@ def vllm_flash_attention_forward(

103 value: torch.Tensor,

104 attention_mask: torch.Tensor,

105 # Transformers kwargs

106 - scaling: Optional[float] = None,

106 + scaling: float | None = None,

107 # vLLM kwargs

108 - attention_instances: Optional[dict[Attention]] = None,

108 + attention_instances: dict[Attention] | None = None,

109 **kwargs,

110 ):

111 self_attn = attention_instances[module.layer_idx]

@@ -147,10 +147,10 @@ Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep", "replicate"]

147 def replace_linear_class(

148 linear: nn.Linear,

149 style: Style = "replicate",

150 - quant_config: Optional[QuantizationConfig] = None,

150 + quant_config: QuantizationConfig | None = None,

151 *,

152 prefix: str = "",

153 -) -> Union[ColumnParallelLinear, RowParallelLinear, ReplicatedLinear]:

153 +) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear:

154 """

155 Replace nn.Linear with one of vLLM's tensor parallel linear classes.

156

@@ -312,7 +312,7 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn

312 self,

313 seq_len: int,

314 mm_counts: Mapping[str, int],

315 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

315 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

316 ) -> MultiModalDataDict:

317 num_images = mm_counts.get("image", 0)

318

@@ -388,11 +388,11 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):

388

389 def apply(

390 self,

391 - prompt: Union[str, list[int]],

391 + prompt: str | list[int],

392 mm_data: MultiModalDataDict,

393 hf_processor_mm_kwargs: Mapping[str, object],

394 - tokenization_kwargs: Optional[Mapping[str, object]] = None,

395 - mm_uuids: Optional[MultiModalUUIDDict] = None,

394 + tokenization_kwargs: Mapping[str, object] | None = None,

395 + mm_uuids: MultiModalUUIDDict | None = None,

396 ) -> MultiModalInputs:

397 """

398 Process multi-modal inputs to be used in vLLM.

@@ -498,7 +498,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

498 self.device_config: DeviceConfig = vllm_config.device_config

499 self.model_config: ModelConfig = vllm_config.model_config

500 self.parallel_config: ParallelConfig = vllm_config.parallel_config

501 - self.quant_config: Optional[QuantizationConfig] = vllm_config.quant_config

501 + self.quant_config: QuantizationConfig | None = vllm_config.quant_config

502

503 self.pp_group = get_pp_group()

504 self.tp_group = get_tp_group()

@@ -714,7 +714,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

714 )

715 return attention_instances

716

717 - def init_parameters(self, module: nn.Module, dtype: Optional[torch.dtype] = None):

717 + def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None):

718 """

719 If a `parameter` is on the `meta` device, then its parent

720 `module` is the original module created by:

@@ -725,7 +725,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

725 ```

726 """

727

728 - def _init_parameters(module: nn.Module, dtype: Optional[torch.dtype]):

728 + def _init_parameters(module: nn.Module, dtype: torch.dtype | None):

729 for name, param in module.named_parameters(recurse=False):

730 if param.device == torch.device("meta"):

731 new_param = nn.Parameter(

@@ -743,12 +743,12 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

743

744 def forward(

745 self,

746 - input_ids: Optional[torch.Tensor],

746 + input_ids: torch.Tensor | None,

747 positions: torch.Tensor,

748 - intermediate_tensors: Optional[IntermediateTensors] = None,

749 - inputs_embeds: Optional[torch.Tensor] = None,

748 + intermediate_tensors: IntermediateTensors | None = None,

749 + inputs_embeds: torch.Tensor | None = None,

750 **kwargs,

751 - ) -> Union[torch.Tensor, IntermediateTensors]:

751 + ) -> torch.Tensor | IntermediateTensors:

752 if not self.pp_group.is_first_rank:

753 assert intermediate_tensors is not None

754 input_ids = None

@@ -841,7 +841,7 @@ class TransformersForCausalLM(TransformersBase):

841 def compute_logits(

842 self,

843 hidden_states: torch.Tensor,

844 - ) -> Optional[torch.Tensor]:

844 + ) -> torch.Tensor | None:

845 logits = self.logits_processor(self.lm_head, hidden_states)

846 return logits

847

@@ -895,12 +895,12 @@ class TransformersForMultimodalLM(TransformersForCausalLM, SupportsMultiModal):

895

896 def forward(

897 self,

898 - input_ids: Optional[torch.Tensor],

898 + input_ids: torch.Tensor | None,

899 positions: torch.Tensor,

900 - intermediate_tensors: Optional[IntermediateTensors] = None,

901 - inputs_embeds: Optional[torch.Tensor] = None,

900 + intermediate_tensors: IntermediateTensors | None = None,

901 + inputs_embeds: torch.Tensor | None = None,

902 **kwargs: object,

903 - ) -> Union[torch.Tensor, IntermediateTensors]:

903 + ) -> torch.Tensor | IntermediateTensors:

904 # Gemma3 and PaliGemma needs `token_type_ids` to work correctly

905 # Other models will not have `token_type_ids` in kwargs

906 kwargs = {k: v for k, v in kwargs.items() if k == "token_type_ids"}

@@ -924,8 +924,8 @@ class TransformersForMultimodalLM(TransformersForCausalLM, SupportsMultiModal):

924 return LanguageModelWrapper(self)

925

926 def get_multimodal_embeddings(self, **kwargs):

927 - pixel_values: Optional[torch.Tensor] = kwargs.pop("pixel_values", None)

928 - image_embeds: Optional[torch.Tensor] = kwargs.pop("image_embeds", None)

927 + pixel_values: torch.Tensor | None = kwargs.pop("pixel_values", None)

928 + image_embeds: torch.Tensor | None = kwargs.pop("image_embeds", None)

929 # Model might use `image_patches` instead of `pixel_values`

930 if pixel_values is None:

931 pixel_values = kwargs.pop("image_patches", None)

vllm/model_executor/models/transformers_pooling.py

+4 -6

@@ -16,8 +16,6 @@

16 # limitations under the License.

17 """Wrapper around `transformers` models for pooling tasks."""

18

19 -from typing import Optional, Union

20 -

19 import torch

20 from transformers import AutoModelForSequenceClassification

21

@@ -109,11 +107,11 @@ class TransformersPoolingBase(TransformersBase, VllmModelForPooling):

107

108 def forward(

109 self,

112 - input_ids: Optional[torch.Tensor],

110 + input_ids: torch.Tensor | None,

111 positions: torch.Tensor,

114 - intermediate_tensors: Optional[IntermediateTensors] = None,

115 - inputs_embeds: Optional[torch.Tensor] = None,

116 - ) -> Union[torch.Tensor, IntermediateTensors]:

112 + intermediate_tensors: IntermediateTensors | None = None,

113 + inputs_embeds: torch.Tensor | None = None,

114 + ) -> torch.Tensor | IntermediateTensors:

115 if self.is_roberta:

116 # RoBERTa-specific positions padding

117 positions += self.padding_idx + 1

vllm/model_executor/models/ultravox.py

+19 -17

@@ -5,7 +5,7 @@

5 """PyTorch Ultravox model."""

6

7 from collections.abc import Iterable, Mapping, Sequence

8 -from typing import Annotated, Any, Literal, Optional, Union

8 +from typing import Annotated, Any, Literal, TypeAlias

9

10 import torch

11 from torch import nn

@@ -68,7 +68,7 @@ class UltravoxAudioFeatureInputs(TensorSchema):

68

69 type: Literal["audio_features"]

70 data: Annotated[

71 - Union[torch.Tensor, list[torch.Tensor], list[list[torch.Tensor]]],

71 + torch.Tensor | list[torch.Tensor] | list[list[torch.Tensor]],

72 TensorShape("bn", "nmb", "t"),

73 ]

74 lens: Annotated[torch.Tensor, TensorShape("bn")]

@@ -92,11 +92,13 @@ class UltravoxAudioEmbeddingInputs(TensorSchema):

92

93 type: Literal["audio_embeds"]

94 data: Annotated[

95 - Union[torch.Tensor, list[torch.Tensor]], TensorShape("b", "na", "afs", "hs")

95 + torch.Tensor | list[torch.Tensor], TensorShape("b", "na", "afs", "hs")

96 ]

97

98

99 -UltravoxAudioInputs = Union[UltravoxAudioFeatureInputs, UltravoxAudioEmbeddingInputs]

99 +UltravoxAudioInputs: TypeAlias = (

100 + UltravoxAudioFeatureInputs | UltravoxAudioEmbeddingInputs

101 +)

102

103

104 class UltravoxProcessingInfo(BaseProcessingInfo):

@@ -119,7 +121,7 @@ class UltravoxProcessingInfo(BaseProcessingInfo):

121 assert isinstance(feature_extractor, WhisperFeatureExtractor)

122 return feature_extractor

123

122 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

124 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

125 return {"audio": None}

126

127

@@ -133,7 +135,7 @@ class UltravoxDummyInputsBuilder(BaseDummyInputsBuilder[UltravoxProcessingInfo])

135 self,

136 seq_len: int,

137 mm_counts: Mapping[str, int],

136 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

138 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

139 ) -> MultiModalDataDict:

140 feature_extractor = self.info.get_feature_extractor()

141

@@ -346,7 +348,7 @@ class ModifiedWhisperEncoder(WhisperEncoder):

348 )

349

350 def get_attention_mask_by_audio_len(

349 - self, audio_lens: Optional[torch.Tensor], hidden_states: torch.Tensor

351 + self, audio_lens: torch.Tensor | None, hidden_states: torch.Tensor

352 ):

353 """

354 Create attention mask based on audio lengths to mask out padding tokens

@@ -376,7 +378,7 @@ class ModifiedWhisperEncoder(WhisperEncoder):

378 def forward(

379 self,

380 input_features: torch.Tensor,

379 - audio_lens: Optional[torch.Tensor] = None,

381 + audio_lens: torch.Tensor | None = None,

382 ):

383 expected_seq_length = self.max_context_length

384 if input_features.shape[-1] > expected_seq_length:

@@ -431,7 +433,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

433 )

434

435 @classmethod

434 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

436 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

437 if modality.startswith("audio"):

438 return "<|audio|>"

439

@@ -514,7 +516,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

516

517 def _parse_and_validate_audio_input(

518 self, **kwargs: object

517 - ) -> Optional[UltravoxAudioInputs]:

519 + ) -> UltravoxAudioInputs | None:

520 audio_features = kwargs.pop("audio_features", None)

521 audio_embeds = kwargs.pop("audio_embeds", None)

522 audio_lens = kwargs.pop("audio_lens", None)

@@ -541,7 +543,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

543 def _process_audio_input(

544 self,

545 audio_input: UltravoxAudioInputs,

544 - ) -> Union[NestedTensors, tuple[torch.Tensor, ...]]:

546 + ) -> NestedTensors | tuple[torch.Tensor, ...]:

547 if audio_input["type"] == "audio_embeds":

548 return audio_input["data"]

549

@@ -587,9 +589,9 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

589 def get_input_embeddings(

590 self,

591 input_ids: torch.Tensor,

590 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

592 + multimodal_embeddings: MultiModalEmbeddings | None = None,

593 *,

592 - is_multimodal: Optional[torch.Tensor] = None,

594 + is_multimodal: torch.Tensor | None = None,

595 # Multi-modal token ID may exceed vocab size

596 handle_oov_mm_token: bool = True,

597 ) -> torch.Tensor:

@@ -608,10 +610,10 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

610 self,

611 input_ids: torch.Tensor,

612 positions: torch.Tensor,

611 - intermediate_tensors: Optional[torch.Tensor] = None,

612 - inputs_embeds: Optional[torch.Tensor] = None,

613 + intermediate_tensors: torch.Tensor | None = None,

614 + inputs_embeds: torch.Tensor | None = None,

615 **kwargs,

614 - ) -> Union[torch.Tensor, IntermediateTensors]:

616 + ) -> torch.Tensor | IntermediateTensors:

617 """Run forward pass for Ultravox

618

619 One key thing to understand is the `input_ids` already accounts for the

@@ -651,7 +653,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

653

654

655 def pad_and_concat_to_dim3(

654 - features: Union[torch.Tensor, list[torch.Tensor], list[list[torch.Tensor]]],

656 + features: torch.Tensor | list[torch.Tensor] | list[list[torch.Tensor]],

657 ) -> torch.Tensor:

658 """

659 Pad and concatenate a list of tensors.

vllm/model_executor/models/utils.py

+16 -16

@@ -4,7 +4,7 @@

4 import itertools

5 from collections.abc import Iterable, Mapping

6 from dataclasses import dataclass, field

7 -from typing import Any, Literal, Optional, Protocol, Union, overload

7 +from typing import Any, Literal, Protocol, overload

8

9 import torch

10 import torch.nn as nn

@@ -32,7 +32,7 @@ from vllm.utils import (

32

33 logger = init_logger(__name__)

34

35 -WeightsMapping = Mapping[str, Optional[str]]

35 +WeightsMapping = Mapping[str, str | None]

36 """If a key maps to a value of `None`, the corresponding weight is ignored."""

37

38

@@ -44,7 +44,7 @@ class WeightsMapper:

44 orig_to_new_prefix: WeightsMapping = field(default_factory=dict)

45 orig_to_new_suffix: WeightsMapping = field(default_factory=dict)

46

47 - def _map_name(self, key: str) -> Optional[str]:

47 + def _map_name(self, key: str) -> str | None:

48 for substr, new_key in self.orig_to_new_substr.items():

49 if substr in key:

50 if new_key is None:

@@ -120,10 +120,10 @@ class AutoWeightsLoader:

120 self,

121 module: nn.Module,

122 *,

123 - skip_prefixes: Optional[list[str]] = None,

124 - skip_substrs: Optional[list[str]] = None,

125 - ignore_unexpected_prefixes: Optional[list[str]] = None,

126 - ignore_unexpected_suffixes: Optional[list[str]] = None,

123 + skip_prefixes: list[str] | None = None,

124 + skip_substrs: list[str] | None = None,

125 + ignore_unexpected_prefixes: list[str] | None = None,

126 + ignore_unexpected_suffixes: list[str] | None = None,

127 ) -> None:

128 super().__init__()

129

@@ -306,7 +306,7 @@ class AutoWeightsLoader:

306 self,

307 weights: Iterable[tuple[str, torch.Tensor]],

308 *,

309 - mapper: Optional[WeightsMapper] = None,

309 + mapper: WeightsMapper | None = None,

310 ) -> set[str]:

311 if mapper is not None:

312 weights = mapper.apply(weights)

@@ -323,8 +323,8 @@ def init_vllm_registered_model(

323 vllm_config: VllmConfig,

324 *,

325 prefix: str = "",

326 - hf_config: Optional[PretrainedConfig] = None,

327 - architectures: Optional[list[str]] = None,

326 + hf_config: PretrainedConfig | None = None,

327 + architectures: list[str] | None = None,

328 ) -> nn.Module:

329 """

330 Helper function to initialize an inner model registered to vLLM,

@@ -352,7 +352,7 @@ def flatten_bn(x: list[torch.Tensor]) -> list[torch.Tensor]: ...

352

353 @overload

354 def flatten_bn(

355 - x: Union[list[torch.Tensor], torch.Tensor],

355 + x: list[torch.Tensor] | torch.Tensor,

356 *,

357 concat: Literal[True],

358 ) -> torch.Tensor: ...

@@ -360,17 +360,17 @@ def flatten_bn(

360

361 @overload

362 def flatten_bn(

363 - x: Union[list[torch.Tensor], torch.Tensor],

363 + x: list[torch.Tensor] | torch.Tensor,

364 *,

365 concat: bool = False,

366 -) -> Union[list[torch.Tensor], torch.Tensor]: ...

366 +) -> list[torch.Tensor] | torch.Tensor: ...

367

368

369 def flatten_bn(

370 - x: Union[list[torch.Tensor], torch.Tensor],

370 + x: list[torch.Tensor] | torch.Tensor,

371 *,

372 concat: bool = False,

373 -) -> Union[list[torch.Tensor], torch.Tensor]:

373 +) -> list[torch.Tensor] | torch.Tensor:

374 """

375 Flatten the ``B`` and ``N`` dimensions of batched multimodal inputs.

376

@@ -472,7 +472,7 @@ def merge_multimodal_embeddings(

472 input_ids: torch.Tensor,

473 inputs_embeds: torch.Tensor,

474 multimodal_embeddings: NestedTensors,

475 - placeholder_token_id: Union[int, list[int]],

475 + placeholder_token_id: int | list[int],

476 ) -> torch.Tensor:

477 """

478 Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the

vllm/model_executor/models/vision.py

+13 -13

@@ -4,7 +4,8 @@

4 import itertools

5 import math

6 from abc import ABC, abstractmethod

7 -from typing import Callable, Final, Generic, Literal, Optional, Protocol, TypeVar, Union

7 +from collections.abc import Callable

8 +from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar

9

10 import torch

11 from transformers import PretrainedConfig

@@ -84,7 +85,7 @@ def get_vit_attn_backend(head_size: int, dtype: torch.dtype) -> _Backend:

85 # Lazy import to avoid circular dependency

86 from vllm.attention.selector import get_env_variable_attn_backend

87

87 - selected_backend: Optional[_Backend] = get_env_variable_attn_backend()

88 + selected_backend: _Backend | None = get_env_variable_attn_backend()

89 if selected_backend is not None:

90 return selected_backend

91

@@ -93,14 +94,13 @@ def get_vit_attn_backend(head_size: int, dtype: torch.dtype) -> _Backend:

94

95 VisionFeatureSelectStrategyStr = Literal["class", "default", "full"]

96

96 -VisionFeatureSelectStrategy = Union[

97 - VisionFeatureSelectStrategyStr,

98 - Callable[[torch.Tensor], torch.Tensor],

99 -]

97 +VisionFeatureSelectStrategy: TypeAlias = (

98 + VisionFeatureSelectStrategyStr | Callable[[torch.Tensor], torch.Tensor]

99 +)

100

101

102 def _get_vision_feature_selector(

103 - strategy: Union[VisionFeatureSelectStrategy, str],

103 + strategy: VisionFeatureSelectStrategy | str,

104 ) -> Callable[[torch.Tensor], torch.Tensor]:

105 if callable(strategy):

106 return strategy

@@ -121,7 +121,7 @@ def _get_vision_feature_selector(

121

122 def get_num_selected_vision_tokens(

123 num_vision_tokens: int,

124 - strategy: Union[VisionFeatureSelectStrategy, str],

124 + strategy: VisionFeatureSelectStrategy | str,

125 ) -> int:

126 if callable(strategy):

127 dummy_features = torch.empty(1, num_vision_tokens, 64) # [B, L, D]

@@ -141,12 +141,12 @@ def get_num_selected_vision_tokens(

141

142

143 def resolve_visual_encoder_outputs(

144 - encoder_outputs: Union[torch.Tensor, list[torch.Tensor]],

145 - post_layer_norm: Optional[torch.nn.LayerNorm],

144 + encoder_outputs: torch.Tensor | list[torch.Tensor],

145 + post_layer_norm: torch.nn.LayerNorm | None,

146 *,

147 - select_layers: Optional[list[int]] = None,

148 - max_possible_layers: Optional[int] = None,

149 - feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

147 + select_layers: list[int] | None = None,

148 + max_possible_layers: int | None = None,

149 + feature_select_strategy: VisionFeatureSelectStrategy | None = None,

150 ) -> torch.Tensor:

151 """Given the outputs a visual encoder module that may correspond to the

152 output of the last layer, or a list of hidden states to be stacked,

vllm/model_executor/models/voxtral.py

+19 -19

@@ -5,7 +5,7 @@ import math

5 from collections.abc import Iterable, Mapping, Sequence

6 from functools import cached_property

7 from math import ceil

8 -from typing import Literal, Optional, Union, cast

8 +from typing import Literal, cast

9

10 import numpy as np

11 import regex as re

@@ -125,9 +125,9 @@ class VoxtralProcessorAdapter:

125

126 def __call__(

127 self,

128 - text: Optional[Union[TextInput, list[TextInput]]] = None,

129 - audios: Optional[Union[np.ndarray, list[np.ndarray]]] = None,

130 - return_tensors: Optional[Union[str, TensorType]] = None,

128 + text: TextInput | list[TextInput] | None = None,

129 + audios: np.ndarray | list[np.ndarray] | None = None,

130 + return_tensors: str | TensorType | None = None,

131 **kwargs,

132 ) -> Mapping[str, NestedTensors]:

133 if text is None:

@@ -188,7 +188,7 @@ class VoxtralProcessingInfo(BaseProcessingInfo):

188 def get_hf_processor(self) -> VoxtralProcessorAdapter:

189 return VoxtralProcessorAdapter(self.get_tokenizer())

190

191 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

191 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

192 return {"audio": 5} # Performance tends to degrade after 5

193

194 def get_mm_max_tokens_per_item(

@@ -216,7 +216,7 @@ class VoxtralDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralProcessingInfo]):

216 self,

217 seq_len: int,

218 mm_counts: Mapping[str, int],

219 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

219 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

220 ) -> MultiModalDataDict:

221 num_audios = mm_counts.get("audio", 0)

222

@@ -234,7 +234,7 @@ class VoxtralDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralProcessingInfo]):

234 self,

235 seq_len: int,

236 mm_counts: Mapping[str, int],

237 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

237 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

238 ) -> ProcessorInputs:

239 tokenizer = self.info.get_tokenizer()

240

@@ -303,11 +303,11 @@ class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo])

303

304 def _cached_apply_hf_processor(

305 self,

306 - prompt: Union[str, list[int]],

306 + prompt: str | list[int],

307 mm_data_items: MultiModalDataItems,

308 hf_processor_mm_kwargs: Mapping[str, object],

309 tokenization_kwargs: Mapping[str, object],

310 - mm_uuids: Optional[MultiModalUUIDDict] = None,

310 + mm_uuids: MultiModalUUIDDict | None = None,

311 ) -> tuple[list[int], MultiModalProcessingInfo, bool]:

312 prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(

313 prompt=prompt,

@@ -386,10 +386,10 @@ class VoxtralForConditionalGeneration(

386 self,

387 input_ids: torch.Tensor,

388 positions: torch.Tensor,

389 - intermediate_tensors: Optional[IntermediateTensors] = None,

390 - inputs_embeds: Optional[torch.Tensor] = None,

389 + intermediate_tensors: IntermediateTensors | None = None,

390 + inputs_embeds: torch.Tensor | None = None,

391 **kwargs: object,

392 - ) -> Union[torch.Tensor, IntermediateTensors]:

392 + ) -> torch.Tensor | IntermediateTensors:

393 if intermediate_tensors is not None:

394 inputs_embeds = None

395

@@ -401,7 +401,7 @@ class VoxtralForConditionalGeneration(

401

402 def get_multimodal_embeddings(

403 self, **kwargs

404 - ) -> Union[list[torch.Tensor], torch.Tensor, tuple[torch.Tensor, ...], None]:

404 + ) -> list[torch.Tensor] | torch.Tensor | tuple[torch.Tensor, ...] | None:

405 audio_inputs = self._parse_and_validate_audio_arrays(**kwargs)

406 if audio_inputs is None:

407 return None

@@ -433,7 +433,7 @@ class VoxtralForConditionalGeneration(

433

434 def _parse_and_validate_audio_arrays(

435 self, **kwargs: object

436 - ) -> Union[list[torch.Tensor], None]:

436 + ) -> list[torch.Tensor] | None:

437 audio_arrays = kwargs.pop("audio_arrays", None)

438 if audio_arrays is None:

439 return None

@@ -450,7 +450,7 @@ class VoxtralForConditionalGeneration(

450 def compute_logits(

451 self,

452 hidden_states: torch.Tensor,

453 - ) -> Optional[torch.Tensor]:

453 + ) -> torch.Tensor | None:

454 return self.language_model.compute_logits(hidden_states)

455

456 @classmethod

@@ -475,10 +475,10 @@ class VoxtralForConditionalGeneration(

475 audio: np.ndarray,

476 model_config: ModelConfig,

477 stt_config: SpeechToTextConfig,

478 - language: Optional[str],

478 + language: str | None,

479 task_type: Literal["transcribe", "translate"],

480 request_prompt: str,

481 - to_language: Optional[str],

481 + to_language: str | None,

482 ) -> PromptType:

483 tokenizer = cached_tokenizer_from_config(model_config)

484 audio = Audio(audio, int(stt_config.sample_rate), format="wav") # lossless

@@ -500,7 +500,7 @@ class VoxtralForConditionalGeneration(

500 audio_duration_s: float,

501 stt_config: SpeechToTextConfig,

502 model_config: ModelConfig,

503 - ) -> Optional[int]:

503 + ) -> int | None:

504 """

505 Map from audio duration to number of audio tokens produced by the ASR

506 model, without running a forward pass.

@@ -793,7 +793,7 @@ class VoxtralEncoderModel(nn.Module):

793 return torch.stack(chunked_features), chunks_per_example

794

795 def forward(

796 - self, input_features: Union[torch.Tensor, list[torch.Tensor]]

796 + self, input_features: torch.Tensor | list[torch.Tensor]

797 ) -> list[torch.Tensor]:

798 if not isinstance(input_features, list):

799 input_features = [input_features]

vllm/model_executor/models/whisper.py

+28 -28

@@ -4,7 +4,7 @@

4 import math

5 from collections.abc import Iterable, Mapping, Sequence

6 from contextlib import nullcontext

7 -from typing import Annotated, Literal, Optional, Union, cast

7 +from typing import Annotated, Literal, cast

8

9 import numpy as np

10 import torch

@@ -137,7 +137,7 @@ class WhisperAudioInputs(TensorSchema):

137 """

138

139 input_features: Annotated[

140 - Optional[list[torch.Tensor]],

140 + list[torch.Tensor] | None,

141 TensorShape("b", "nmb", "t"),

142 ]

143

@@ -185,8 +185,8 @@ class WhisperAttention(nn.Module):

185 num_heads: int,

186 bias: bool = True,

187 attn_type: AttentionType = AttentionType.DECODER,

188 - cache_config: Optional[CacheConfig] = None,

189 - quant_config: Optional[QuantizationConfig] = None,

188 + cache_config: CacheConfig | None = None,

189 + quant_config: QuantizationConfig | None = None,

190 prefix: str = "",

191 ):

192 super().__init__()

@@ -258,7 +258,7 @@ class WhisperAttention(nn.Module):

258 self,

259 embed_dim: int,

260 bias: bool = True,

261 - quant_config: Optional[QuantizationConfig] = None,

261 + quant_config: QuantizationConfig | None = None,

262 prefix: str = "",

263 ) -> None:

264 self.qkv_proj = QKVParallelLinear(

@@ -291,8 +291,8 @@ class WhisperCrossAttention(WhisperAttention):

291 embed_dim: int,

292 num_heads: int,

293 bias: bool = True,

294 - cache_config: Optional[CacheConfig] = None,

295 - quant_config: Optional[QuantizationConfig] = None,

294 + cache_config: CacheConfig | None = None,

295 + quant_config: QuantizationConfig | None = None,

296 prefix: str = "",

297 ):

298 super().__init__(

@@ -309,7 +309,7 @@ class WhisperCrossAttention(WhisperAttention):

309 self,

310 embed_dim: int,

311 bias: bool = True,

312 - quant_config: Optional[QuantizationConfig] = None,

312 + quant_config: QuantizationConfig | None = None,

313 prefix: str = "",

314 ) -> None:

315 self.q_proj = ColumnParallelLinear(

@@ -332,7 +332,7 @@ class WhisperCrossAttention(WhisperAttention):

332 def forward(

333 self,

334 hidden_states: torch.Tensor,

335 - encoder_hidden_states: Optional[torch.Tensor],

335 + encoder_hidden_states: torch.Tensor | None,

336 ):

337 q, _ = self.q_proj(hidden_states)

338

@@ -357,7 +357,7 @@ class WhisperMLP(nn.Module):

357 embed_dim: int,

358 ffn_dim: int,

359 act_fn: str,

360 - quant_config: Optional[QuantizationConfig] = None,

360 + quant_config: QuantizationConfig | None = None,

361 prefix: str = "",

362 ):

363 super().__init__()

@@ -463,7 +463,7 @@ class WhisperDecoderLayer(nn.Module):

463 def forward(

464 self,

465 hidden_states: torch.Tensor,

466 - encoder_hidden_states: Optional[torch.Tensor],

466 + encoder_hidden_states: torch.Tensor | None,

467 ):

468 residual = hidden_states

469 hidden_states = self.self_attn_layer_norm(hidden_states)

@@ -521,7 +521,7 @@ class WhisperEncoder(nn.Module):

521 sinusoids(*self.embed_positions.weight.shape)

522 )

523

524 - def forward(self, input_features: Union[torch.Tensor, list[torch.Tensor]]):

524 + def forward(self, input_features: torch.Tensor | list[torch.Tensor]):

525 hidden_states = []

526 for features in input_features:

527 embeds = nn.functional.gelu(self.conv1(features))

@@ -569,7 +569,7 @@ class WhisperDecoder(nn.Module):

569 self,

570 input_ids,

571 positions: torch.Tensor,

572 - encoder_hidden_states: Optional[torch.Tensor],

572 + encoder_hidden_states: torch.Tensor | None,

573 ):

574 inputs_embeds = self.get_input_embeddings(input_ids)

575 positions = self.embed_positions(positions)

@@ -600,8 +600,8 @@ class WhisperModel(nn.Module):

600

601 def forward(

602 self,

603 - input_features: Optional[Union[torch.Tensor, list[torch.Tensor]]],

604 - input_ids: Optional[torch.Tensor],

603 + input_features: torch.Tensor | list[torch.Tensor] | None,

604 + input_ids: torch.Tensor | None,

605 positions: torch.Tensor,

606 ) -> torch.Tensor:

607 encoder_outputs = self.get_encoder_outputs(input_features)

@@ -614,8 +614,8 @@ class WhisperModel(nn.Module):

614

615 def get_encoder_outputs(

616 self,

617 - input_features: Optional[Union[torch.Tensor, list[torch.Tensor]]],

618 - ) -> Optional[torch.Tensor]:

617 + input_features: torch.Tensor | list[torch.Tensor] | None,

618 + ) -> torch.Tensor | None:

619 if input_features is None:

620 return None

621 return self.encoder(input_features)

@@ -670,7 +670,7 @@ class WhisperProcessingInfo(BaseProcessingInfo):

670 processor_class.tokenizer_class = tokenizer_class

671 return self.ctx.get_hf_processor(processor_class, **kwargs)

672

673 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

673 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

674 return {"audio": 1}

675

676 def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor:

@@ -693,7 +693,7 @@ class WhisperDummyInputsBuilder(BaseDummyInputsBuilder[WhisperProcessingInfo]):

693 self,

694 seq_len: int,

695 mm_counts: Mapping[str, int],

696 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

696 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

697 ) -> MultiModalDataDict:

698 feature_extractor = self.info.get_feature_extractor()

699

@@ -721,9 +721,9 @@ class WhisperMultiModalProcessor(EncDecMultiModalProcessor[WhisperProcessingInfo

721

722 def create_encoder_prompt(

723 self,

724 - prompt: Union[str, list[int]],

724 + prompt: str | list[int],

725 mm_data: MultiModalDataDict,

726 - ) -> Union[str, list[int]]:

726 + ) -> str | list[int]:

727 # Strictly speaking, whisper encoder only accept audio features.

728 # We create a dummy encoder prompt here which will be padded to

729 # num_audio_tokens. So that we can create dummy data from this

@@ -804,7 +804,7 @@ class WhisperForConditionalGeneration(

804 supported_languages = ISO639_1_SUPPORTED_LANGS

805

806 @classmethod

807 - def validate_language(cls, language: Optional[str]) -> Optional[str]:

807 + def validate_language(cls, language: str | None) -> str | None:

808 if language is None:

809 # TODO language should be optional and can be guessed.

810 # For now we default to en. See

@@ -823,10 +823,10 @@ class WhisperForConditionalGeneration(

823 audio: np.ndarray,

824 model_config: ModelConfig, # not needed here

825 stt_config: SpeechToTextConfig,

826 - language: Optional[str],

826 + language: str | None,

827 task_type: Literal["transcribe", "translate"],

828 request_prompt: str,

829 - to_language: Optional[str],

829 + to_language: str | None,

830 ) -> PromptType:

831 if language is None:

832 raise ValueError(

@@ -849,7 +849,7 @@ class WhisperForConditionalGeneration(

849 return cast(PromptType, prompt)

850

851 @classmethod

852 - def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

852 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

853 if modality.startswith("audio"):

854 return None

855

@@ -872,7 +872,7 @@ class WhisperForConditionalGeneration(

872 audio_duration_s: float,

873 stt_config: SpeechToTextConfig,

874 model_config: ModelConfig,

875 - ) -> Optional[int]:

875 + ) -> int | None:

876 processor = cached_get_processor(model_config.model)

877 hop_length = processor.feature_extractor.hop_length

878 assert hop_length is not None

@@ -928,9 +928,9 @@ class WhisperForConditionalGeneration(

928 def get_input_embeddings(

929 self,

930 input_ids: torch.Tensor,

931 - multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

931 + multimodal_embeddings: MultiModalEmbeddings | None = None,

932 *,

933 - is_multimodal: Optional[torch.Tensor] = None,

933 + is_multimodal: torch.Tensor | None = None,

934 handle_oov_mm_token: bool = False,

935 ) -> torch.Tensor:

936 # This method just returns the decoder sequence embeddings since

vllm/model_executor/models/zamba2.py

+21 -21

@@ -10,7 +10,7 @@ model alternates between state space model layers and attention-based layers.

10

11 from collections.abc import Iterable

12 from itertools import cycle

13 -from typing import Any, Optional, Union

13 +from typing import Any

14

15 import torch

16 from torch import nn

@@ -60,8 +60,8 @@ class Zamba2LoRA(nn.Module):

60 self,

61 input_dim: int,

62 rank: int,

63 - output_dim: Union[int, list[int]],

64 - quant_config: Optional[QuantizationConfig] = None,

63 + output_dim: int | list[int],

64 + quant_config: QuantizationConfig | None = None,

65 prefix: str = "",

66 ):

67 """Initialize the attention layer.

@@ -106,8 +106,8 @@ class Zamba2Attention(nn.Module):

106 config: Zamba2Config,

107 bare_block_idx: int,

108 num_hybrid_layers: int,

109 - cache_config: Optional[CacheConfig] = None,

110 - quant_config: Optional[QuantizationConfig] = None,

109 + cache_config: CacheConfig | None = None,

110 + quant_config: QuantizationConfig | None = None,

111 prefix: str = "",

112 ) -> None:

113 """Initialize the attention layer.

@@ -288,7 +288,7 @@ class Zamba2MLP(nn.Module):

288 config: Zamba2Config,

289 bare_block_idx: int,

290 num_hybrid_layers: dict[int, int],

291 - quant_config: Optional[QuantizationConfig] = None,

291 + quant_config: QuantizationConfig | None = None,

292 prefix: str = "",

293 ) -> None:

294 """Initialize the MLP layer.

@@ -386,8 +386,8 @@ class Zamba2AttentionDecoderLayer(nn.Module):

386 config: Zamba2Config,

387 bare_block_idx: int,

388 num_hybrid_layers: int,

389 - cache_config: Optional[CacheConfig] = None,

390 - quant_config: Optional[QuantizationConfig] = None,

389 + cache_config: CacheConfig | None = None,

390 + quant_config: QuantizationConfig | None = None,

391 prefix: str = "",

392 ) -> None:

393 """Initialize the decoder layer.

@@ -484,9 +484,9 @@ class Zamba2MambaDecoderLayer(nn.Module):

484 def __init__(

485 self,

486 config: Zamba2Config,

487 - model_config: Optional[ModelConfig] = None,

488 - cache_config: Optional[CacheConfig] = None,

489 - quant_config: Optional[QuantizationConfig] = None,

487 + model_config: ModelConfig | None = None,

488 + cache_config: CacheConfig | None = None,

489 + quant_config: QuantizationConfig | None = None,

490 prefix: str = "",

491 ) -> None:

492 """Initialize the Mamba decoder layer.

@@ -523,9 +523,9 @@ class Zamba2MambaDecoderLayer(nn.Module):

523 def forward(

524 self,

525 hidden_states: torch.Tensor,

526 - transformer_hidden_states: Optional[torch.Tensor] = None,

527 - positions: Optional[torch.Tensor] = None,

528 - original_hidden_states: Optional[torch.Tensor] = None,

526 + transformer_hidden_states: torch.Tensor | None = None,

527 + positions: torch.Tensor | None = None,

528 + original_hidden_states: torch.Tensor | None = None,

529 ) -> torch.Tensor:

530 """Forward pass through the Mamba decoder layer.

531

@@ -581,9 +581,9 @@ class Zamba2HybridLayer(nn.Module):

581 shared_transformer: Zamba2AttentionDecoderLayer,

582 config: Zamba2Config,

583 block_idx: int,

584 - model_config: Optional[ModelConfig] = None,

585 - cache_config: Optional[CacheConfig] = None,

586 - quant_config: Optional[QuantizationConfig] = None,

584 + model_config: ModelConfig | None = None,

585 + cache_config: CacheConfig | None = None,

586 + quant_config: QuantizationConfig | None = None,

587 prefix: str = "",

588 ) -> None:

589 """Initialize the hybrid layer.

@@ -764,8 +764,8 @@ class Zamba2Model(nn.Module):

764 self,

765 input_ids: torch.Tensor,

766 positions: torch.Tensor,

767 - inputs_embeds: Optional[torch.Tensor] = None,

768 - ) -> Union[torch.Tensor, IntermediateTensors]:

767 + inputs_embeds: torch.Tensor | None = None,

768 + ) -> torch.Tensor | IntermediateTensors:

769 """Forward pass through the model.

770

771 Args:

@@ -947,7 +947,7 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):

947 self,

948 input_ids: torch.Tensor,

949 positions: torch.Tensor,

950 - inputs_embeds: Optional[torch.Tensor] = None,

950 + inputs_embeds: torch.Tensor | None = None,

951 **kwargs: Any,

952 ) -> torch.Tensor:

953 """Forward pass through the model.

@@ -973,7 +973,7 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):

973 def compute_logits(

974 self,

975 hidden_states: torch.Tensor,

976 - ) -> Optional[torch.Tensor]:

976 + ) -> torch.Tensor | None:

977 """Compute logits for next token prediction.

978

979 Args:

vllm/model_executor/parameter.py

+12 -13

@@ -1,9 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from collections.abc import Hashable

4 +from collections.abc import Callable, Hashable

5 from fractions import Fraction

6 -from typing import Callable, Optional, Union

6 from weakref import WeakValueDictionary

7

8 import torch

@@ -36,7 +35,7 @@ class BasevLLMParameter(Parameter):

35 into the parameter when the provided weight loader is called.

36 """

37

39 - def __new__(cls, data: Optional[torch.Tensor], **kwargs):

38 + def __new__(cls, data: torch.Tensor | None, **kwargs):

39 return super().__new__(cls, data=data, requires_grad=False)

40

41 def __init__(self, data: torch.Tensor, weight_loader: Callable):

@@ -109,7 +108,7 @@ class BasevLLMParameter(Parameter):

108 def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):

109 self._assert_and_load(loaded_weight)

110

112 - def _shard_id_as_int(self, shard_id: Union[str, int]) -> int:

111 + def _shard_id_as_int(self, shard_id: str | int) -> int:

112 if isinstance(shard_id, int):

113 return shard_id

114

@@ -290,7 +289,7 @@ class PerTensorScaleParameter(BasevLLMParameter):

289 super().load_row_parallel_weight(*args, **kwargs)

290

291 def _load_into_shard_id(

293 - self, loaded_weight: torch.Tensor, shard_id: Union[str, int], **kwargs

292 + self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs

293 ):

294 """

295 Slice the parameter data based on the shard id for

@@ -320,10 +319,10 @@ class PackedColumnParameter(_ColumnvLLMParameter):

319

320 def __init__(

321 self,

323 - packed_factor: Union[int, Fraction],

322 + packed_factor: int | Fraction,

323 packed_dim: int,

325 - marlin_tile_size: Optional[int] = None,

326 - bitblas_tile_size: Optional[int] = None,

324 + marlin_tile_size: int | None = None,

325 + bitblas_tile_size: int | None = None,

326 **kwargs,

327 ):

328 self._packed_factor = packed_factor

@@ -371,10 +370,10 @@ class PackedvLLMParameter(ModelWeightParameter):

370

371 def __init__(

372 self,

374 - packed_factor: Union[int, Fraction],

373 + packed_factor: int | Fraction,

374 packed_dim: int,

376 - marlin_tile_size: Optional[int] = None,

377 - bitblas_tile_size: Optional[int] = None,

375 + marlin_tile_size: int | None = None,

376 + bitblas_tile_size: int | None = None,

377 **kwargs,

378 ):

379 self._packed_factor = packed_factor

@@ -437,7 +436,7 @@ class SharedWeightParameter(BasevLLMParameter):

436 local_tensors: set[torch.Tensor]

437

438 # dictionary mapping partition indices to associated parameters

440 - partitions: dict[int, Union[ModelWeightParameter, Parameter]]

439 + partitions: dict[int, ModelWeightParameter | Parameter]

440

441 def __new__(cls, **kwargs):

442 return super().__new__(cls, data=None, **kwargs)

@@ -547,7 +546,7 @@ class SharedWeightParameter(BasevLLMParameter):

546 self,

547 param: BasevLLMParameter,

548 loaded_weight: torch.Tensor,

550 - loaded_weight_shard_id: Optional[Union[str, int]],

549 + loaded_weight_shard_id: str | int | None,

550 ):

551 raise ValueError(

552 "When loading partition weights of "

vllm/model_executor/utils.py

+2 -2

@@ -3,7 +3,7 @@

3 """Utils for model executor."""

4

5 import copy

6 -from typing import Any, Optional

6 +from typing import Any

7

8 import torch

9

@@ -16,7 +16,7 @@ def set_random_seed(seed: int) -> None:

16

17 def set_weight_attrs(

18 weight: torch.Tensor,

19 - weight_attrs: Optional[dict[str, Any]],

19 + weight_attrs: dict[str, Any] | None,

20 ):

21 """Set attributes on a weight tensor.

22

vllm/multimodal/audio.py

+2 -2

@@ -3,7 +3,7 @@

3 import base64

4 from io import BytesIO

5 from pathlib import Path

6 -from typing import Literal, Optional

6 +from typing import Literal

7

8 import numpy as np

9 import numpy.typing as npt

@@ -53,7 +53,7 @@ class AudioResampler:

53

54 def __init__(

55 self,

56 - target_sr: Optional[float] = None,

56 + target_sr: float | None = None,

57 method: Literal["librosa", "scipy"] = "librosa",

58 ):

59 self.target_sr = target_sr

vllm/multimodal/cache.py

+20 -20

@@ -5,10 +5,10 @@ import sys

5 from abc import ABC, abstractmethod

6 from collections.abc import Mapping, Sequence

7 from multiprocessing.synchronize import Lock as LockType

8 -from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union, cast

8 +from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar, cast

9

10 import torch

11 -from typing_extensions import TypeAlias, override

11 +from typing_extensions import override

12

13 from vllm.distributed.device_communicators.shm_object_storage import (

14 MsgpackSerde,

@@ -85,14 +85,14 @@ class MultiModalProcessorCacheItemMetadata:

85 self.prompt_updates = prompt_updates

86

87

88 -MultiModalCacheValue = Union[

89 - MultiModalProcessorCacheItem,

90 - MultiModalProcessorCacheItemMetadata,

91 - MultiModalKwargsItems,

92 - MultiModalKwargsItem,

93 - MultiModalKwargs,

94 - Mapping[str, NestedTensors],

95 -]

88 +MultiModalCacheValue: TypeAlias = (

89 + MultiModalProcessorCacheItem

90 + | MultiModalProcessorCacheItemMetadata

91 + | MultiModalKwargsItems

92 + | MultiModalKwargsItem

93 + | MultiModalKwargs

94 + | Mapping[str, NestedTensors]

95 +)

96

97 _V = TypeVar("_V", bound=MultiModalCacheValue)

98

@@ -256,13 +256,13 @@ class BaseMultiModalCache(ABC, Generic[_I, _O]):

256 raise NotImplementedError

257

258

259 -MultiModalProcessorCacheInItem: TypeAlias = Optional[

260 - tuple[MultiModalKwargsItem, Sequence["ResolvedPromptUpdate"]]

261 -]

259 +MultiModalProcessorCacheInItem: TypeAlias = (

260 + tuple[MultiModalKwargsItem, Sequence["ResolvedPromptUpdate"]] | None

261 +)

262

263

264 MultiModalProcessorCacheOutItem: TypeAlias = tuple[

265 - Optional[MultiModalKwargsItem], Sequence["ResolvedPromptUpdate"]

265 + MultiModalKwargsItem | None, Sequence["ResolvedPromptUpdate"]

266 ]

267

268

@@ -575,7 +575,7 @@ def _enable_mm_input_shm_cache(vllm_config: "VllmConfig") -> bool:

575 def processor_cache_from_config(

576 vllm_config: "VllmConfig",

577 mm_registry: "MultiModalRegistry",

578 -) -> Optional[BaseMultiModalProcessorCache]:

578 +) -> BaseMultiModalProcessorCache | None:

579 """Return a `BaseMultiModalProcessorCache`, if enabled."""

580 model_config = vllm_config.model_config

581

@@ -602,7 +602,7 @@ def processor_only_cache_from_config(

602

603

604 class BaseMultiModalReceiverCache(

605 - BaseMultiModalCache[Optional[MultiModalKwargsItem], MultiModalKwargsItem]

605 + BaseMultiModalCache[MultiModalKwargsItem | None, MultiModalKwargsItem]

606 ):

607 """The required interface for caches on P1."""

608

@@ -640,7 +640,7 @@ class MultiModalReceiverCache(BaseMultiModalReceiverCache):

640 @override

641 def get_and_update_item(

642 self,

643 - mm_item: Optional[MultiModalKwargsItem],

643 + mm_item: MultiModalKwargsItem | None,

644 mm_hash: str,

645 ) -> MultiModalKwargsItem:

646 if (cached_item := self._cache.get(mm_hash)) is not None:

@@ -692,7 +692,7 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):

692 @override

693 def get_and_update_item(

694 self,

695 - mm_item: Optional[MultiModalKwargsItem],

695 + mm_item: MultiModalKwargsItem | None,

696 mm_hash: str,

697 ) -> MultiModalKwargsItem:

698 assert mm_item is not None, f"Expected an address item for {mm_hash=}"

@@ -711,7 +711,7 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):

711 def engine_receiver_cache_from_config(

712 vllm_config: "VllmConfig",

713 mm_registry: "MultiModalRegistry",

714 -) -> Optional[BaseMultiModalReceiverCache]:

714 +) -> BaseMultiModalReceiverCache | None:

715 """

716 This is used in the engine process.

717 Return a `BaseMultiModalReceiverCache` only when IPC caching is enabled and

@@ -735,7 +735,7 @@ def worker_receiver_cache_from_config(

735 vllm_config: "VllmConfig",

736 mm_registry: "MultiModalRegistry",

737 shared_worker_lock: LockType,

738 -) -> Optional[BaseMultiModalReceiverCache]:

738 +) -> BaseMultiModalReceiverCache | None:

739 """

740 This is used in the worker process.

741 Return a `BaseMultiModalReceiverCache` only when IPC caching is enabled and

vllm/multimodal/evs.py

+1 -2

@@ -9,7 +9,6 @@

9 # license agreement from NVIDIA CORPORATION is strictly prohibited.

10

11 import typing

12 -from typing import Union

12

13 import torch

14

@@ -38,7 +37,7 @@ def compute_retained_tokens_count(

37

38 def compute_retention_mask(

39 video_embeds: torch.Tensor,

41 - video_size_thw: Union[torch.LongTensor, tuple[int, int, int]],

40 + video_size_thw: torch.LongTensor | tuple[int, int, int],

41 spatial_merge_size: int,

42 q: float,

43 ) -> torch.Tensor:

vllm/multimodal/hasher.py

+2 -3

@@ -4,7 +4,6 @@

4 import pickle

5 import uuid

6 from collections.abc import Iterable

7 -from typing import Union

7

8 import numpy as np

9 import torch

@@ -18,7 +17,7 @@ logger = init_logger(__name__)

17

18 class MultiModalHasher:

19 @classmethod

21 - def serialize_item(cls, obj: object) -> Iterable[Union[bytes, memoryview]]:

20 + def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]:

21 # Simple cases

22 if isinstance(obj, (bytes, memoryview)):

23 return (obj,)

@@ -84,7 +83,7 @@ class MultiModalHasher:

83 cls,

84 key: str,

85 obj: object,

87 - ) -> Iterable[Union[bytes, memoryview]]:

86 + ) -> Iterable[bytes | memoryview]:

87 # Recursive cases

88 if isinstance(obj, (list, tuple)):

89 for i, elem in enumerate(obj):

vllm/multimodal/image.py

+1 -2

@@ -3,7 +3,6 @@

3

4 from io import BytesIO

5 from pathlib import Path

6 -from typing import Union

6

7 import pybase64

8 import torch

@@ -26,7 +25,7 @@ def rescale_image_size(

25

26 def rgba_to_rgb(

27 image: Image.Image,

29 - background_color: Union[tuple[int, int, int], list[int]] = (255, 255, 255),

28 + background_color: tuple[int, int, int] | list[int] = (255, 255, 255),

29 ) -> Image.Image:

30 """Convert an RGBA image to RGB with filled background color."""

31 assert image.mode == "RGBA"

vllm/multimodal/inputs.py

+21 -11

@@ -7,10 +7,20 @@ from collections.abc import Mapping, Sequence

7 from dataclasses import dataclass

8 from functools import partial

9 from itertools import accumulate

10 -from typing import TYPE_CHECKING, Any, Literal, Optional, TypedDict, Union, cast, final

10 +from typing import (

11 + TYPE_CHECKING,

12 + Any,

13 + Literal,

14 + Optional,

15 + TypeAlias,

16 + TypedDict,

17 + Union,

18 + cast,

19 + final,

20 +)

21

22 import numpy as np

13 -from typing_extensions import NotRequired, TypeAlias, TypeVar, deprecated

23 +from typing_extensions import NotRequired, TypeVar, deprecated

24

25 from vllm.utils import LazyLoader, full_groupby, is_list_of

26 from vllm.utils.jsontree import json_map_leaves

@@ -85,7 +95,7 @@ which are treated as audio embeddings;

95 these are directly passed to the model without HF processing.

96 """

97

88 -ModalityData: TypeAlias = Union[_T, list[Optional[_T]], None]

98 +ModalityData: TypeAlias = _T | list[_T | None] | None

99 """

100 Either a single data item, or a list of data items. Can only be None if UUID

101 is provided.

@@ -117,7 +127,7 @@ The built-in modalities are defined by

127 [`MultiModalDataBuiltins`][vllm.multimodal.inputs.MultiModalDataBuiltins].

128 """

129

120 -MultiModalUUIDDict: TypeAlias = Mapping[str, Union[list[Optional[str]], str]]

130 +MultiModalUUIDDict: TypeAlias = Mapping[str, list[str | None] | str]

131 """

132 A dictionary containing user-provided UUIDs for items in each modality.

133 If a UUID for an item is not provided, its entry will be `None` and

@@ -412,7 +422,7 @@ class MultiModalFlatField(BaseMultiModalField):

422 [`MultiModalFieldConfig.flat_from_sizes`][vllm.multimodal.inputs.MultiModalFieldConfig.flat_from_sizes]

423 """

424

415 - slices: Union[Sequence[slice], Sequence[Sequence[slice]]]

425 + slices: Sequence[slice] | Sequence[Sequence[slice]]

426 dim: int = 0

427

428 def build_elems(

@@ -524,7 +534,7 @@ class MultiModalFieldConfig:

534 @staticmethod

535 def flat(

536 modality: str,

527 - slices: Union[Sequence[slice], Sequence[Sequence[slice]]],

537 + slices: Sequence[slice] | Sequence[Sequence[slice]],

538 dim: int = 0,

539 ):

540 """

@@ -729,7 +739,7 @@ class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]):

739 _I = TypeVar(

740 "_I",

741 MultiModalKwargsItem,

732 - Optional[MultiModalKwargsItem],

742 + MultiModalKwargsItem | None,

743 default=MultiModalKwargsItem,

744 )

745

@@ -818,10 +828,10 @@ class MultiModalKwargsItems(UserDict[str, Sequence[_I]]):

828 )

829

830

821 -MultiModalKwargsOptionalItems: TypeAlias = Union[

822 - MultiModalKwargsItems[MultiModalKwargsItem],

823 - MultiModalKwargsItems[Optional[MultiModalKwargsItem]],

824 -]

831 +MultiModalKwargsOptionalItems: TypeAlias = (

832 + MultiModalKwargsItems[MultiModalKwargsItem]

833 + | MultiModalKwargsItems[MultiModalKwargsItem | None]

834 +)

835

836

837 class MultiModalKwargs(UserDict[str, NestedTensors]):

vllm/multimodal/parse.py

+22 -24

@@ -10,14 +10,14 @@ from typing import (

10 Generic,

11 Literal,

12 NamedTuple,

13 - Optional,

13 + TypeAlias,

14 + TypeGuard,

15 TypeVar,

15 - Union,

16 )

17

18 import numpy as np

19 import torch

20 -from typing_extensions import TypeAlias, TypeGuard, assert_never

20 +from typing_extensions import assert_never

21

22 from vllm.utils import LazyLoader, is_list_of

23

@@ -111,7 +111,7 @@ class ProcessorBatchItems(ModalityDataItems[Sequence[_T], _T]):

111

112

113 class EmbeddingItems(

114 - ModalityDataItems[Union[torch.Tensor, list[torch.Tensor]], torch.Tensor]

114 + ModalityDataItems[torch.Tensor | list[torch.Tensor], torch.Tensor]

115 ):

116 """

117 Base class for data items that are expressed as a batched embedding tensor,

@@ -195,7 +195,7 @@ class DictEmbeddingItems(

195

196

197 class AudioProcessorItems(ProcessorBatchItems[HfAudioItem]):

198 - def __init__(self, data: Optional[Sequence[HfAudioItem]]) -> None:

198 + def __init__(self, data: Sequence[HfAudioItem] | None) -> None:

199 if data is None:

200 data = [None]

201 super().__init__(data, "audio")

@@ -206,7 +206,7 @@ class AudioProcessorItems(ProcessorBatchItems[HfAudioItem]):

206

207

208 class AudioEmbeddingItems(EmbeddingItems):

209 - def __init__(self, data: Union[torch.Tensor, list[torch.Tensor]]) -> None:

209 + def __init__(self, data: torch.Tensor | list[torch.Tensor]) -> None:

210 super().__init__(data, "audio")

211

212

@@ -216,7 +216,7 @@ class ImageSize(NamedTuple):

216

217

218 class ImageProcessorItems(ProcessorBatchItems[HfImageItem]):

219 - def __init__(self, data: Optional[Sequence[HfImageItem]]) -> None:

219 + def __init__(self, data: Sequence[HfImageItem] | None) -> None:

220 if data is None:

221 data = [None]

222 super().__init__(data, "image")

@@ -234,17 +234,15 @@ class ImageProcessorItems(ProcessorBatchItems[HfImageItem]):

234

235

236 class ImageEmbeddingItems(EmbeddingItems):

237 - def __init__(self, data: Union[torch.Tensor, list[torch.Tensor]]) -> None:

237 + def __init__(self, data: torch.Tensor | list[torch.Tensor]) -> None:

238 super().__init__(data, "image")

239

240

241 class VideoProcessorItems(ProcessorBatchItems[HfVideoItem]):

242 def __init__(

243 self,

244 - data: Optional[Sequence[HfVideoItem]],

245 - metadata: Optional[

246 - Union[dict[str, Any], list[Optional[dict[str, Any]]]]

247 - ] = None,

244 + data: Sequence[HfVideoItem] | None,

245 + metadata: dict[str, Any] | list[dict[str, Any] | None] | None = None,

246 ) -> None:

247 if data is None:

248 data = [None]

@@ -267,7 +265,7 @@ class VideoProcessorItems(ProcessorBatchItems[HfVideoItem]):

265

266

267 class VideoEmbeddingItems(EmbeddingItems):

270 - def __init__(self, data: Union[torch.Tensor, list[torch.Tensor]]) -> None:

268 + def __init__(self, data: torch.Tensor | list[torch.Tensor]) -> None:

269 super().__init__(data, "video")

270

271

@@ -306,7 +304,7 @@ class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]):

304 def get_items(

305 self,

306 modality: str,

309 - typ: Union[type[_D], tuple[type[_D], ...]],

307 + typ: type[_D] | tuple[type[_D], ...],

308 ) -> _D:

309 """

310 Get the data items belonging to a modality,

@@ -331,7 +329,7 @@ class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]):

329

330

331 ModalityDataParser: TypeAlias = Callable[

334 - [ModalityData[Any]], Optional[ModalityDataItems[Any, Any]]

332 + [ModalityData[Any]], ModalityDataItems[Any, Any] | None

333 ]

334

335

@@ -348,7 +346,7 @@ class MultiModalDataParser:

346 def __init__(

347 self,

348 *,

351 - target_sr: Optional[float] = None,

349 + target_sr: float | None = None,

350 audio_resample_method: Literal["librosa", "scipy"] = "librosa",

351 video_needs_metadata: bool = False,

352 ) -> None:

@@ -362,7 +360,7 @@ class MultiModalDataParser:

360

361 def _is_embeddings(

362 self, data: object

365 - ) -> TypeGuard[Union[torch.Tensor, list[torch.Tensor]]]:

363 + ) -> TypeGuard[torch.Tensor | list[torch.Tensor]]:

364 if isinstance(data, torch.Tensor):

365 return data.ndim == 3

366 if is_list_of(data, torch.Tensor):

@@ -381,7 +379,7 @@ class MultiModalDataParser:

379 def _get_audio_with_sr(

380 self,

381 audio: AudioItem,

384 - ) -> tuple[np.ndarray, Optional[float]]:

382 + ) -> tuple[np.ndarray, float | None]:

383 if isinstance(audio, tuple):

384 return audio

385 if isinstance(audio, list):

@@ -396,7 +394,7 @@ class MultiModalDataParser:

394 def _get_video_with_metadata(

395 self,

396 video: VideoItem,

399 - ) -> tuple[np.ndarray, Optional[dict[str, Any]]]:

397 + ) -> tuple[np.ndarray, dict[str, Any] | None]:

398 if isinstance(video, tuple):

399 return video

400 if isinstance(video, list):

@@ -411,7 +409,7 @@ class MultiModalDataParser:

409 def _parse_audio_data(

410 self,

411 data: ModalityData[AudioItem],

414 - ) -> Optional[ModalityDataItems[Any, Any]]:

412 + ) -> ModalityDataItems[Any, Any] | None:

413 if data is None:

414 return AudioProcessorItems(None)

415

@@ -451,7 +449,7 @@ class MultiModalDataParser:

449 def _parse_image_data(

450 self,

451 data: ModalityData[ImageItem],

454 - ) -> Optional[ModalityDataItems[Any, Any]]:

452 + ) -> ModalityDataItems[Any, Any] | None:

453 if data is None:

454 return ImageProcessorItems(None)

455

@@ -477,7 +475,7 @@ class MultiModalDataParser:

475 def _parse_video_data(

476 self,

477 data: ModalityData[VideoItem],

480 - ) -> Optional[ModalityDataItems[Any, Any]]:

478 + ) -> ModalityDataItems[Any, Any] | None:

479 if data is None:

480 return VideoProcessorItems(None)

481

@@ -500,8 +498,8 @@ class MultiModalDataParser:

498 else:

499 data_items = data

500

503 - new_videos = list[tuple[np.ndarray, Optional[dict[str, Any]]]]()

504 - metadata_lst: list[Optional[dict[str, Any]]] = []

501 + new_videos = list[tuple[np.ndarray, dict[str, Any] | None]]()

502 + metadata_lst: list[dict[str, Any] | None] = []

503 for data_item in data_items:

504 video, metadata = self._get_video_with_metadata(data_item)

505 if self.video_needs_metadata:

vllm/multimodal/processing.py

+69 -62

@@ -12,9 +12,8 @@ from typing import (

12 Any,

13 Generic,

14 NamedTuple,

15 - Optional,

15 Protocol,

17 - Union,

16 + TypeAlias,

17 cast,

18 overload,

19 )

@@ -57,12 +56,20 @@ if TYPE_CHECKING:

56

57 from .cache import BaseMultiModalProcessorCache

58 from .profiling import BaseDummyInputsBuilder

59 +else:

60 + PretrainedConfig = object

61 + BatchFeature = object

62 + ProcessorMixin = object

63 +

64 + ModelConfig = object

65 +

66 + BaseMultiModalProcessorCache = object

67

68 logger = init_logger(__name__)

69

70 _S = TypeVar("_S", str, list[int])

71

65 -PromptSeq = Union[str, list[int]]

72 +PromptSeq: TypeAlias = str | list[int]

73 """A token sequence (list of token IDs) or text."""

74

75

@@ -71,7 +78,7 @@ def _cached_encode(

78 tokenizer: AnyTokenizer,

79 text: str,

80 *,

74 - add_special_tokens: Optional[bool] = None,

81 + add_special_tokens: bool | None = None,

82 ) -> list[int]:

83 return encode_tokens(tokenizer, text, add_special_tokens=add_special_tokens)

84

@@ -81,7 +88,7 @@ def _cached_decode(

88 tokenizer: AnyTokenizer,

89 token_ids: tuple[int, ...],

90 *,

84 - skip_special_tokens: Optional[bool] = None,

91 + skip_special_tokens: bool | None = None,

92 ) -> str:

93 return decode_tokens(

94 tokenizer, list(token_ids), skip_special_tokens=skip_special_tokens

@@ -108,7 +115,7 @@ class _GetMatchIndex(Protocol):

115 tokenizer: AnyTokenizer,

116 prompt: PromptSeq,

117 start_idx: int = 0,

111 - ) -> Optional[int]: ...

118 + ) -> int | None: ...

119

120

121 @dataclass

@@ -138,7 +145,7 @@ class PromptIndexTargets:

145 tokenizer: AnyTokenizer,

146 prompt: PromptSeq,

147 start_idx: int = 0,

141 - ) -> Optional[int]:

148 + ) -> int | None:

149 if start_idx != 0:

150 return None

151

@@ -168,12 +175,12 @@ class PromptIndexTargets:

175 return PromptIndex(lambda tokenizer, prompt, start_idx=0: len(prompt))

176

177

171 -UpdateTarget = Union[PromptSeq, PromptIndex]

178 +UpdateTarget: TypeAlias = PromptSeq | PromptIndex

179 """

180 The token sequence or text to update.

181 """

182

176 -PromptUpdateTarget = Union[Callable[[int], UpdateTarget], UpdateTarget]

183 +PromptUpdateTarget: TypeAlias = Callable[[int], UpdateTarget] | UpdateTarget

184 """

185 Given the index of the processed item within

186 [`modality`][vllm.multimodal.processing.PromptUpdate.modality],

@@ -191,7 +198,7 @@ class PromptUpdateDetails(Generic[_S]):

198 full: _S

199 """The full content."""

200

194 - is_embed: Optional[Callable[[AnyTokenizer, PromptSeq], torch.Tensor]] = None

201 + is_embed: Callable[[AnyTokenizer, PromptSeq], torch.Tensor] | None = None

202 """

203 Given [`full`][vllm.multimodal.processing.PromptUpdateDetails.full],

204 return a boolean mask of shape `(len(full),)` indicating which positions

@@ -236,7 +243,7 @@ class PromptUpdateDetails(Generic[_S]):

243 return PromptUpdateDetails(full=seq, is_embed=is_embed)

244

245

239 -PromptUpdateInfo = Union[PromptSeq, PromptUpdateDetails]

246 +PromptUpdateInfo: TypeAlias = PromptSeq | PromptUpdateDetails

247 """

248 The token sequence or text that are part of the update.

249

@@ -245,7 +252,7 @@ use [`PromptUpdateDetails`][vllm.multimodal.processing.PromptUpdateDetails] to

252 specify which part.

253 """

254

248 -PromptUpdateContent = Union[Callable[[int], PromptUpdateInfo], PromptUpdateInfo]

255 +PromptUpdateContent: TypeAlias = Callable[[int], PromptUpdateInfo] | PromptUpdateInfo

256 """

257 Given the index of the processed item within

258 [`modality`][vllm.multimodal.processing.PromptUpdate.modality],

@@ -472,7 +479,7 @@ class _HasModalityProp(Protocol):

479 def modality(self) -> str: ...

480

481

475 -_M = TypeVar("_M", bound=Union[_HasModalityAttr, _HasModalityProp])

482 +_M = TypeVar("_M", bound=_HasModalityAttr | _HasModalityProp)

483

484

485 def full_groupby_modality(values: Iterable[_M]) -> ItemsView[str, list[_M]]:

@@ -554,7 +561,7 @@ class ResolvedPromptUpdate:

561

562 def iter_matches(

563 self,

557 - prompt: Union[list[int], str],

564 + prompt: list[int] | str,

565 tokenizer: AnyTokenizer,

566 *,

567 start_idx: int = 0,

@@ -642,7 +649,7 @@ class PlaceholderFeaturesInfo:

649 item_idx: int

650 start_idx: int

651 tokens: list[int]

645 - is_embed: Optional[torch.Tensor]

652 + is_embed: torch.Tensor | None

653

654 @property

655 def length(self) -> int:

@@ -668,8 +675,8 @@ def _find_matches(

675 *,

676 prev_end_idx: int = 0,

677 current_result: "MultiModalPromptUpdatesApplyResult",

671 -) -> tuple[Optional[UpdateMode], list[_MatchToApply]]:

672 - mode: Optional[UpdateMode] = None

678 +) -> tuple[UpdateMode | None, list[_MatchToApply]]:

679 + mode: UpdateMode | None = None

680 mm_matches = dict[tuple[str, int], tuple[PromptTargetMatch, int]]()

681

682 for modality, modality_updates in mm_prompt_updates.items():

@@ -723,7 +730,7 @@ def _apply_matches(

730 ) -> tuple[list[_S], "MultiModalPromptUpdatesApplyResult"]:

731 prompt_len = len(prompt)

732

726 - out_seqs = list[Union[str, list[int]]]()

733 + out_seqs = list[str | list[int]]()

734 out_result: MultiModalPromptUpdatesApplyResult = {

735 m: [None] * len(items) for m, items in mm_prompt_updates.items()

736 }

@@ -880,8 +887,8 @@ def find_mm_placeholders(

887

888

889 _T = TypeVar("_T")

883 -_C = TypeVar("_C", bound="PretrainedConfig", default="PretrainedConfig")

884 -_P = TypeVar("_P", bound="ProcessorMixin", default="ProcessorMixin")

890 +_C = TypeVar("_C", bound=PretrainedConfig, default=PretrainedConfig)

891 +_P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin)

892

893

894 @dataclass(frozen=True)

@@ -891,25 +898,25 @@ class InputProcessingContext:

898 modify the inputs.

899 """

900

894 - model_config: "ModelConfig"

901 + model_config: ModelConfig

902 """The configuration of the model."""

903

904 tokenizer: AnyTokenizer

905 """The tokenizer used to tokenize the inputs."""

906

907 @overload

901 - def get_hf_config(self, /) -> "PretrainedConfig": ...

908 + def get_hf_config(self, /) -> PretrainedConfig: ...

909

910 @overload

911 def get_hf_config(

912 self,

906 - typ: Union[type[_C], tuple[type[_C], ...]],

913 + typ: type[_C] | tuple[type[_C], ...],

914 /,

915 ) -> _C: ...

916

917 def get_hf_config(

918 self,

912 - typ: Optional[Union[type[Any], tuple[type[Any], ...]]] = None,

919 + typ: type[Any] | tuple[type[Any], ...] | None = None,

920 /,

921 ) -> Any:

922 """

@@ -955,19 +962,19 @@ class InputProcessingContext:

962 return mm_config

963

964 @overload

958 - def get_hf_processor(self, /, **kwargs: object) -> "ProcessorMixin": ...

965 + def get_hf_processor(self, /, **kwargs: object) -> ProcessorMixin: ...

966

967 @overload

968 def get_hf_processor(

969 self,

963 - typ: Union[type[_P], tuple[type[_P], ...]],

970 + typ: type[_P] | tuple[type[_P], ...],

971 /,

972 **kwargs: object,

973 ) -> _P: ...

974

975 def get_hf_processor(

976 self,

970 - typ: Optional[Union[type[Any], tuple[type[Any], ...]]] = None,

977 + typ: type[Any] | tuple[type[Any], ...] | None = None,

978 /,

979 **kwargs: object,

980 ) -> Any:

@@ -1026,13 +1033,13 @@ class InputProcessingContext:

1033

1034 def call_hf_processor(

1035 self,

1029 - hf_processor: "ProcessorMixin",

1036 + hf_processor: ProcessorMixin,

1037 data: Mapping[str, object],

1038 kwargs: Mapping[str, object] = {},

1039 *,

1040 num_tries: int = 1,

1041 max_tries: int = 5,

1035 - ) -> Union["BatchFeature", JSONTree]:

1042 + ) -> BatchFeature | JSONTree:

1043 """

1044 Call `hf_processor` on the prompt `data`

1045 (text, image, audio...) with configurable options `kwargs`.

@@ -1113,10 +1120,10 @@ class BaseProcessingInfo:

1120 def get_tokenizer(self) -> AnyTokenizer:

1121 return self.ctx.tokenizer

1122

1116 - def get_hf_config(self) -> "PretrainedConfig":

1123 + def get_hf_config(self) -> PretrainedConfig:

1124 return self.ctx.get_hf_config()

1125

1119 - def get_hf_processor(self, **kwargs: object) -> "ProcessorMixin":

1126 + def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:

1127 """

1128 Subclasses can override this method to handle

1129 specific kwargs from model config or user inputs.

@@ -1124,7 +1131,7 @@ class BaseProcessingInfo:

1131 return self.ctx.get_hf_processor(**kwargs)

1132

1133 @abstractmethod

1127 - def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

1134 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

1135 """

1136 Return the maximum supported number of items for each modality.

1137

@@ -1156,7 +1163,7 @@ class BaseProcessingInfo:

1163 self,

1164 seq_len: int,

1165 mm_counts: Mapping[str, int],

1159 - ) -> Optional[Mapping[str, int]]:

1166 + ) -> Mapping[str, int] | None:

1167 """

1168 Return the maximum number of tokens per item of for each modality.

1169

@@ -1193,7 +1200,7 @@ A collection of prompt updates with a similar structure as

1200 [`MultiModalKwargsItems`][vllm.multimodal.inputs.MultiModalKwargsItems].

1201 """

1202

1196 -MultiModalPromptUpdatesApplyResult = Mapping[str, list[Optional[int]]]

1203 +MultiModalPromptUpdatesApplyResult = Mapping[str, list[int | None]]

1204 """

1205 For an item `MultiModalPromptUpdates[k][i]`,

1206 `MultiModalPromptUpdatesApplyResult[k][i]` represents the index of the

@@ -1220,7 +1227,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1227 info: _I,

1228 dummy_inputs: "BaseDummyInputsBuilder[_I]",

1229 *,

1223 - cache: Optional["BaseMultiModalProcessorCache"] = None,

1230 + cache: BaseMultiModalProcessorCache | None = None,

1231 ) -> None:

1232 super().__init__()

1233

@@ -1248,7 +1255,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1255 mm_data: MultiModalDataDict,

1256 hf_processor_mm_kwargs: Mapping[str, object],

1257 *,

1251 - mm_uuids: Optional[MultiModalUUIDDict] = None,

1258 + mm_uuids: MultiModalUUIDDict | None = None,

1259 ) -> MultiModalInputs:

1260 return self.apply(prompt, mm_data, hf_processor_mm_kwargs, mm_uuids=mm_uuids)

1261

@@ -1305,7 +1312,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1312 @abstractmethod

1313 def _get_mm_fields_config(

1314 self,

1308 - hf_inputs: "BatchFeature",

1315 + hf_inputs: BatchFeature,

1316 hf_processor_mm_kwargs: Mapping[str, object],

1317 ) -> Mapping[str, MultiModalFieldConfig]:

1318 """Given the HF-processed data, output the metadata of each field."""

@@ -1411,7 +1418,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1418 mm_data: Mapping[str, object],

1419 mm_kwargs: Mapping[str, object],

1420 tok_kwargs: Mapping[str, object],

1414 - ) -> "BatchFeature":

1421 + ) -> BatchFeature:

1422 """

1423 Call the HF processor on the prompt text and

1424 associated multi-modal data.

@@ -1447,7 +1454,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1454 mm_items: MultiModalDataItems,

1455 hf_processor_mm_kwargs: Mapping[str, object],

1456 tokenization_kwargs: Mapping[str, object],

1450 - ) -> tuple[list[int], "BatchFeature", bool]:

1457 + ) -> tuple[list[int], BatchFeature, bool]:

1458 """

1459 Apply the HF processor on the prompt text and multi-modal data

1460 together.

@@ -1518,7 +1525,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1525 mm_items: MultiModalDataItems,

1526 hf_processor_mm_kwargs: Mapping[str, object],

1527 tokenization_kwargs: Mapping[str, object],

1521 - ) -> "BatchFeature":

1528 + ) -> BatchFeature:

1529 """

1530 Apply the HF processor on the multi-modal data only.

1531

@@ -1540,13 +1547,13 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1547

1548 def _apply_hf_processor_main(

1549 self,

1543 - prompt: Union[str, list[int]],

1550 + prompt: str | list[int],

1551 mm_items: MultiModalDataItems,

1552 hf_processor_mm_kwargs: Mapping[str, object],

1553 tokenization_kwargs: Mapping[str, object],

1554 *,

1555 enable_hf_prompt_update: bool,

1549 - ) -> tuple[list[int], "BatchFeature", bool]:

1556 + ) -> tuple[list[int], BatchFeature, bool]:

1557 """

1558 Apply the HF processor on the prompt text and multi-modal data.

1559

@@ -1585,7 +1592,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1592 hf_processor_mm_kwargs: Mapping[str, object],

1593 tokenization_kwargs: Mapping[str, object],

1594 *,

1588 - mm_uuids: Optional[MultiModalUUIDDict] = None,

1595 + mm_uuids: MultiModalUUIDDict | None = None,

1596 ) -> MultiModalHashes:

1597 """Create MM hashes to be returned.

1598

@@ -1647,7 +1654,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1654

1655 def _get_cache_missing_items(

1656 self,

1650 - cache: "BaseMultiModalProcessorCache",

1657 + cache: BaseMultiModalProcessorCache,

1658 mm_data_items: MultiModalDataItems,

1659 mm_hashes: MultiModalHashes,

1660 ) -> MultiModalDataItems:

@@ -1692,7 +1699,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1699

1700 def _merge_mm_kwargs(

1701 self,

1695 - cache: "BaseMultiModalProcessorCache",

1702 + cache: BaseMultiModalProcessorCache,

1703 mm_hashes: MultiModalHashes,

1704 mm_missing_kwargs: MultiModalKwargsItems,

1705 mm_missing_prompt_updates: MultiModalPromptUpdates,

@@ -1705,7 +1712,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1712

1713 mm_missing_next_idx = defaultdict[str, int](lambda: 0)

1714

1708 - merged_kwargs = defaultdict[str, list[Optional[MultiModalKwargsItem]]](list)

1715 + merged_kwargs = defaultdict[str, list[MultiModalKwargsItem | None]](list)

1716 merged_prompt_updates = defaultdict[str, list[Sequence[ResolvedPromptUpdate]]](

1717 list

1718 )

@@ -1714,7 +1721,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1721 missing_prompt_updates = mm_missing_prompt_updates.get(modality, [])

1722

1723 for item_idx, item_hash in enumerate(hashes):

1717 - kwargs: Optional[MultiModalKwargsItem]

1724 + kwargs: MultiModalKwargsItem | None

1725 if not mm_is_cached[modality][item_idx]:

1726 missing_next_idx = mm_missing_next_idx[modality]

1727 kwargs = missing_kwargs[missing_next_idx]

@@ -1743,12 +1750,12 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1750

1751 def _apply_hf_processor(

1752 self,

1746 - prompt: Union[str, list[int]],

1753 + prompt: str | list[int],

1754 mm_data_items: MultiModalDataItems,

1755 hf_processor_mm_kwargs: Mapping[str, object],

1756 tokenization_kwargs: Mapping[str, object],

1757 *,

1751 - mm_uuids: Optional[MultiModalUUIDDict] = None,

1758 + mm_uuids: MultiModalUUIDDict | None = None,

1759 ) -> tuple[list[int], MultiModalProcessingInfo, bool]:

1760 (

1761 prompt_ids,

@@ -1791,12 +1798,12 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1798

1799 def _cached_apply_hf_processor(

1800 self,

1794 - prompt: Union[str, list[int]],

1801 + prompt: str | list[int],

1802 mm_data_items: MultiModalDataItems,

1803 hf_processor_mm_kwargs: Mapping[str, object],

1804 tokenization_kwargs: Mapping[str, object],

1805 *,

1799 - mm_uuids: Optional[MultiModalUUIDDict] = None,

1806 + mm_uuids: MultiModalUUIDDict | None = None,

1807 ) -> tuple[list[int], MultiModalProcessingInfo, bool]:

1808 """

1809 Apply the HF processor on the full prompt text,

@@ -2026,12 +2033,12 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

2033

2034 def apply(

2035 self,

2029 - prompt: Union[str, list[int]],

2036 + prompt: str | list[int],

2037 mm_data: MultiModalDataDict,

2038 hf_processor_mm_kwargs: Mapping[str, object],

2032 - tokenization_kwargs: Optional[Mapping[str, object]] = None,

2039 + tokenization_kwargs: Mapping[str, object] | None = None,

2040 *,

2034 - mm_uuids: Optional[MultiModalUUIDDict] = None,

2041 + mm_uuids: MultiModalUUIDDict | None = None,

2042 ) -> MultiModalInputs:

2043 """

2044 Process multi-modal inputs to be used in vLLM.

@@ -2090,9 +2097,9 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):

2097 @abstractmethod

2098 def create_encoder_prompt(

2099 self,

2093 - prompt: Union[str, list[int]],

2100 + prompt: str | list[int],

2101 mm_data: MultiModalDataDict,

2095 - ) -> Union[str, list[int]]:

2102 + ) -> str | list[int]:

2103 """

2104 Create input prompt for the encoder. HF processor will be applied on

2105 this prompt during profiling and generation.

@@ -2105,15 +2112,15 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):

2112

2113 def create_decoder_prompt(

2114 self,

2108 - prompt: Union[str, list[int]],

2115 + prompt: str | list[int],

2116 mm_data: MultiModalDataDict,

2110 - ) -> Union[str, list[int]]:

2117 + ) -> str | list[int]:

2118 """Create input prompt for the decoder."""

2119 return prompt

2120

2121 def _get_enc_dec_inputs(

2122 self,

2116 - prompt: Union[str, list[int]],

2123 + prompt: str | list[int],

2124 mm_data: MultiModalDataDict,

2125 encoder_inputs: MultiModalInputs,

2126 ):

@@ -2135,12 +2142,12 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):

2142

2143 def apply(

2144 self,

2138 - prompt: Union[str, list[int]],

2145 + prompt: str | list[int],

2146 mm_data: MultiModalDataDict,

2147 hf_processor_mm_kwargs: Mapping[str, object],

2141 - tokenization_kwargs: Optional[Mapping[str, object]] = None,

2148 + tokenization_kwargs: Mapping[str, object] | None = None,

2149 *,

2143 - mm_uuids: Optional[MultiModalUUIDDict] = None,

2150 + mm_uuids: MultiModalUUIDDict | None = None,

2151 ) -> MultiModalEncDecInputs:

2152 """

2153 Process multi-modal inputs to be used in vLLM.

vllm/multimodal/profiling.py

+15 -15

@@ -3,7 +3,7 @@

3 from abc import ABC, abstractmethod

4 from collections.abc import Mapping

5 from dataclasses import dataclass, field

6 -from typing import Generic, NamedTuple, Optional, TypeVar, Union, cast

6 +from typing import Generic, NamedTuple, TypeVar, cast

7

8 import numpy as np

9 import numpy.typing as npt

@@ -41,7 +41,7 @@ class ProcessorInputs:

41 [`vllm.multimodal.processing.BaseMultiModalProcessor.apply`][].

42 """

43

44 - prompt: Union[str, list[int]]

44 + prompt: str | list[int]

45 mm_data: MultiModalDataDict

46 hf_processor_mm_kwargs: Mapping[str, object] = field(default_factory=dict)

47 tokenization_kwargs: Mapping[str, object] = field(default_factory=dict)

@@ -87,7 +87,7 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]):

87 self,

88 seq_len: int,

89 mm_counts: Mapping[str, int],

90 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

90 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

91 ) -> MultiModalDataDict:

92 """

93 Build the multimodal input which, after processing, results in

@@ -107,7 +107,7 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]):

107 self,

108 seq_len: int,

109 mm_counts: Mapping[str, int],

110 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

110 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

111 ) -> ProcessorInputs:

112 """

113 Build the input which, after processing, results in

@@ -136,7 +136,7 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]):

136 *,

137 length: int,

138 num_audios: int,

139 - overrides: Optional[AudioDummyOptions] = None,

139 + overrides: AudioDummyOptions | None = None,

140 ) -> list[npt.NDArray]:

141 if num_audios == 0:

142 return []

@@ -158,7 +158,7 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]):

158 width: int,

159 height: int,

160 num_images: int,

161 - overrides: Optional[ImageDummyOptions] = None,

161 + overrides: ImageDummyOptions | None = None,

162 ) -> list[Image.Image]:

163 if num_images == 0:

164 return []

@@ -191,7 +191,7 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]):

191 height: int,

192 num_frames: int,

193 num_videos: int,

194 - overrides: Optional[VideoDummyOptions] = None,

194 + overrides: VideoDummyOptions | None = None,

195 ) -> list[npt.NDArray]:

196 if num_videos == 0:

197 return []

@@ -254,8 +254,8 @@ class MultiModalProfiler(Generic[_I]):

254 def _get_dummy_mm_inputs(

255 self,

256 seq_len: int,

257 - mm_counts: Optional[Mapping[str, int]] = None,

258 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

257 + mm_counts: Mapping[str, int] | None = None,

258 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

259 ) -> MultiModalInputs:

260 if mm_counts is None:

261 mm_counts = self.get_mm_limits()

@@ -290,8 +290,8 @@ class MultiModalProfiler(Generic[_I]):

290 def get_encoder_dummy_data(

291 self,

292 seq_len: int,

293 - mm_counts: Optional[Mapping[str, int]] = None,

294 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

293 + mm_counts: Mapping[str, int] | None = None,

294 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

295 ) -> DummyEncoderData:

296 mm_inputs = self._get_dummy_mm_inputs(seq_len, mm_counts, mm_options)

297 mm_inputs = cast(MultiModalEncDecInputs, mm_inputs)

@@ -324,8 +324,8 @@ class MultiModalProfiler(Generic[_I]):

324 def get_decoder_dummy_data(

325 self,

326 seq_len: int,

327 - mm_counts: Optional[Mapping[str, int]] = None,

328 - mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,

327 + mm_counts: Mapping[str, int] | None = None,

328 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

329 ) -> DummyDecoderData:

330 mm_inputs = self._get_dummy_mm_inputs(seq_len, mm_counts, mm_options)

331

@@ -344,7 +344,7 @@ class MultiModalProfiler(Generic[_I]):

344 def _get_mm_max_tokens(

345 self,

346 seq_len: int,

347 - mm_counts: Optional[Mapping[str, int]] = None,

347 + mm_counts: Mapping[str, int] | None = None,

348 mm_embeddings_only: bool = True,

349 ) -> Mapping[str, int]:

350 if mm_counts is None:

@@ -363,7 +363,7 @@ class MultiModalProfiler(Generic[_I]):

363 def get_mm_max_contiguous_tokens(

364 self,

365 seq_len: int,

366 - mm_counts: Optional[Mapping[str, int]] = None,

366 + mm_counts: Mapping[str, int] | None = None,

367 ):

368 """

369 Returns the maximum length of the multimodal (image placeholders+text)

vllm/multimodal/registry.py

+15 -15

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Mapping

4 from dataclasses import dataclass

5 -from typing import TYPE_CHECKING, Generic, Optional, Protocol, TypeVar

5 +from typing import TYPE_CHECKING, Generic, Protocol, TypeVar

6

7 import torch.nn as nn

8

@@ -69,7 +69,7 @@ class MultiModalProcessorFactory(Protocol[_I]): # type: ignore[misc]

69 info: _I,

70 dummy_inputs: BaseDummyInputsBuilder[_I],

71 *,

72 - cache: Optional[BaseMultiModalProcessorCache] = None,

72 + cache: BaseMultiModalProcessorCache | None = None,

73 ) -> BaseMultiModalProcessor[_I]: ...

74

75

@@ -83,7 +83,7 @@ class _ProcessorFactories(Generic[_I]):

83 self,

84 ctx: InputProcessingContext,

85 *,

86 - cache: Optional[BaseMultiModalProcessorCache] = None,

86 + cache: BaseMultiModalProcessorCache | None = None,

87 ):

88 info = self.info(ctx)

89 dummy_inputs_builder = self.dummy_inputs(info)

@@ -101,7 +101,7 @@ class MultiModalRegistry:

101 def _extract_mm_options(

102 self,

103 model_config: "ModelConfig",

104 - ) -> Optional[Mapping[str, BaseDummyOptions]]:

104 + ) -> Mapping[str, BaseDummyOptions] | None:

105 """

106 Extract multimodal dummy options from model config.

107

@@ -151,7 +151,7 @@ class MultiModalRegistry:

151 self,

152 model_config: "ModelConfig",

153 *,

154 - cache: Optional[BaseMultiModalProcessorCache] = None,

154 + cache: BaseMultiModalProcessorCache | None = None,

155 ) -> Mapping[str, int]:

156 """

157 Get the maximum number of tokens per data item from each modality based

@@ -175,7 +175,7 @@ class MultiModalRegistry:

175 self,

176 model_config: "ModelConfig",

177 *,

178 - cache: Optional[BaseMultiModalProcessorCache] = None,

178 + cache: BaseMultiModalProcessorCache | None = None,

179 ) -> Mapping[str, int]:

180 """

181 Get the maximum number of tokens per data item from each modality based

@@ -202,7 +202,7 @@ class MultiModalRegistry:

202 self,

203 model_config: "ModelConfig",

204 *,

205 - cache: Optional[BaseMultiModalProcessorCache] = None,

205 + cache: BaseMultiModalProcessorCache | None = None,

206 ) -> Mapping[str, int]:

207 """

208 Get the maximum number of multi-modal input instances for each modality

@@ -259,7 +259,7 @@ class MultiModalRegistry:

259 def _create_processing_ctx(

260 self,

261 model_config: "ModelConfig",

262 - tokenizer: Optional[AnyTokenizer] = None,

262 + tokenizer: AnyTokenizer | None = None,

263 ) -> InputProcessingContext:

264 if tokenizer is None and not model_config.skip_tokenizer_init:

265 tokenizer = cached_tokenizer_from_config(model_config)

@@ -269,7 +269,7 @@ class MultiModalRegistry:

269 self,

270 model_config: "ModelConfig",

271 *,

272 - tokenizer: Optional[AnyTokenizer] = None,

272 + tokenizer: AnyTokenizer | None = None,

273 ) -> BaseProcessingInfo:

274 model_cls = self._get_model_cls(model_config)

275 factories = self._processor_factories[model_cls]

@@ -280,8 +280,8 @@ class MultiModalRegistry:

280 self,

281 model_config: "ModelConfig",

282 *,

283 - tokenizer: Optional[AnyTokenizer] = None,

284 - cache: Optional[BaseMultiModalProcessorCache] = None,

283 + tokenizer: AnyTokenizer | None = None,

284 + cache: BaseMultiModalProcessorCache | None = None,

285 ) -> BaseMultiModalProcessor[BaseProcessingInfo]:

286 """

287 Create a multi-modal processor for a specific model and tokenizer.

@@ -300,9 +300,9 @@ class MultiModalRegistry:

300 self,

301 model_config: "ModelConfig",

302 seq_len: int,

303 - mm_counts: Optional[Mapping[str, int]] = None,

303 + mm_counts: Mapping[str, int] | None = None,

304 *,

305 - cache: Optional[BaseMultiModalProcessorCache] = None,

305 + cache: BaseMultiModalProcessorCache | None = None,

306 ) -> DummyDecoderData:

307 """

308 Create dummy data for profiling the memory usage of a model.

@@ -333,9 +333,9 @@ class MultiModalRegistry:

333 self,

334 model_config: "ModelConfig",

335 seq_len: int,

336 - mm_counts: Optional[Mapping[str, int]] = None,

336 + mm_counts: Mapping[str, int] | None = None,

337 *,

338 - cache: Optional[BaseMultiModalProcessorCache] = None,

338 + cache: BaseMultiModalProcessorCache | None = None,

339 ) -> DummyEncoderData:

340 """

341 Create dummy data for profiling the memory usage of a model.

vllm/multimodal/utils.py

+12 -12

@@ -7,7 +7,7 @@ from collections.abc import Iterable

7 from concurrent.futures import ThreadPoolExecutor

8 from itertools import groupby

9 from pathlib import Path

10 -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union

10 +from typing import TYPE_CHECKING, Any, TypeVar

11 from urllib.parse import ParseResult, urlparse

12 from urllib.request import url2pathname

13

@@ -49,11 +49,11 @@ atexit.register(global_thread_pool.shutdown)

49 class MediaConnector:

50 def __init__(

51 self,

52 - media_io_kwargs: Optional[dict[str, dict[str, Any]]] = None,

52 + media_io_kwargs: dict[str, dict[str, Any]] | None = None,

53 connection: HTTPConnection = global_http_connection,

54 *,

55 allowed_local_media_path: str = "",

56 - allowed_media_domains: Optional[list[str]] = None,

56 + allowed_media_domains: list[str] | None = None,

57 ) -> None:

58 """

59 Args:

@@ -143,7 +143,7 @@ class MediaConnector:

143 url: str,

144 media_io: MediaIO[_M],

145 *,

146 - fetch_timeout: Optional[int] = None,

146 + fetch_timeout: int | None = None,

147 ) -> _M: # type: ignore[type-var]

148 url_spec = urlparse(url)

149

@@ -173,7 +173,7 @@ class MediaConnector:

173 url: str,

174 media_io: MediaIO[_M],

175 *,

176 - fetch_timeout: Optional[int] = None,

176 + fetch_timeout: int | None = None,

177 ) -> _M:

178 url_spec = urlparse(url)

179 loop = asyncio.get_running_loop()

@@ -207,7 +207,7 @@ class MediaConnector:

207 def fetch_audio(

208 self,

209 audio_url: str,

210 - ) -> tuple[np.ndarray, Union[int, float]]:

210 + ) -> tuple[np.ndarray, int | float]:

211 """

212 Load audio from a URL.

213 """

@@ -222,7 +222,7 @@ class MediaConnector:

222 async def fetch_audio_async(

223 self,

224 audio_url: str,

225 - ) -> tuple[np.ndarray, Union[int, float]]:

225 + ) -> tuple[np.ndarray, int | float]:

226 """

227 Asynchronously fetch audio from a URL.

228 """

@@ -396,7 +396,7 @@ def group_mm_kwargs_by_modality(

396 *,

397 device: torch.types.Device = None,

398 pin_memory: bool = False,

399 - merge_by_field_config: Optional[bool] = None,

399 + merge_by_field_config: bool | None = None,

400 ) -> Iterable[tuple[str, int, BatchedTensorInputs]]:

401 """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same

402 modality together into the same `MultiModalKwargs` instance.

@@ -452,8 +452,8 @@ def group_mm_kwargs_by_modality(

452

453 def fetch_audio(

454 audio_url: str,

455 - audio_io_kwargs: Optional[dict[str, Any]] = None,

456 -) -> tuple[np.ndarray, Union[int, float]]:

455 + audio_io_kwargs: dict[str, Any] | None = None,

456 +) -> tuple[np.ndarray, int | float]:

457 """

458 Args:

459 audio_url: URL of the audio file to fetch.

@@ -466,7 +466,7 @@ def fetch_audio(

466

467 def fetch_image(

468 image_url: str,

469 - image_io_kwargs: Optional[dict[str, Any]] = None,

469 + image_io_kwargs: dict[str, Any] | None = None,

470 ) -> Image.Image:

471 """

472 Args:

@@ -480,7 +480,7 @@ def fetch_image(

480

481 def fetch_video(

482 video_url: str,

483 - video_io_kwargs: Optional[dict[str, Any]] = None,

483 + video_io_kwargs: dict[str, Any] | None = None,

484 ) -> tuple[npt.NDArray, dict[str, Any]]:

485 """

486 Args:

vllm/multimodal/video.py

+2 -2

@@ -6,7 +6,7 @@ from abc import abstractmethod

6 from functools import partial

7 from io import BytesIO

8 from pathlib import Path

9 -from typing import Any, Union

9 +from typing import Any

10

11 import numpy as np

12 import numpy.typing as npt

@@ -192,7 +192,7 @@ class OpenCVDynamicVideoBackend(OpenCVVideoBackend):

192

193 # Refer to:

194 # https://github.com/huggingface/transformers/blob/v4.55.4/src/transformers/models/glm4v/video_processing_glm4v.py#L103-L140

195 - frame_indices: Union[range, list[int]]

195 + frame_indices: range | list[int]

196 if duration <= max_duration:

197 n = int(math.floor(duration * fps))

198 frame_indices = sorted(

vllm/outputs.py

+16 -16

@@ -4,7 +4,7 @@

4 from collections.abc import MutableSequence

5 from collections.abc import Sequence as GenericSequence

6 from dataclasses import dataclass

7 -from typing import Any, Generic, Optional, Union

7 +from typing import Any, Generic

8

9 import torch

10 from typing_extensions import TypeVar

@@ -41,11 +41,11 @@ class CompletionOutput:

41 index: int

42 text: str

43 token_ids: GenericSequence[int]

44 - cumulative_logprob: Optional[float]

45 - logprobs: Optional[SampleLogprobs]

46 - finish_reason: Optional[str] = None

47 - stop_reason: Union[int, str, None] = None

48 - lora_request: Optional[LoRARequest] = None

44 + cumulative_logprob: float | None

45 + logprobs: SampleLogprobs | None

46 + finish_reason: str | None = None

47 + stop_reason: int | str | None = None

48 + lora_request: LoRARequest | None = None

49

50 def finished(self) -> bool:

51 return self.finish_reason is not None

@@ -108,19 +108,19 @@ class RequestOutput:

108 def __init__(

109 self,

110 request_id: str,

111 - prompt: Optional[str],

112 - prompt_token_ids: Optional[list[int]],

113 - prompt_logprobs: Optional[PromptLogprobs],

111 + prompt: str | None,

112 + prompt_token_ids: list[int] | None,

113 + prompt_logprobs: PromptLogprobs | None,

114 outputs: list[CompletionOutput],

115 finished: bool,

116 - metrics: Optional[Union[RequestMetrics, RequestStateStats]] = None,

117 - lora_request: Optional[LoRARequest] = None,

118 - encoder_prompt: Optional[str] = None,

119 - encoder_prompt_token_ids: Optional[list[int]] = None,

120 - num_cached_tokens: Optional[int] = None,

116 + metrics: RequestMetrics | RequestStateStats | None = None,

117 + lora_request: LoRARequest | None = None,

118 + encoder_prompt: str | None = None,

119 + encoder_prompt_token_ids: list[int] | None = None,

120 + num_cached_tokens: int | None = None,

121 *,

122 - multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None,

123 - kv_transfer_params: Optional[dict[str, Any]] = None,

122 + multi_modal_placeholders: MultiModalPlaceholderDict | None = None,

123 + kv_transfer_params: dict[str, Any] | None = None,

124 # Forward compatibility, code that uses args added in new release can

125 # still run with older versions of vLLM without breaking.

126 **kwargs: Any,

vllm/platforms/__init__.py

+6 -6

@@ -3,7 +3,7 @@

3 import logging

4 import traceback

5 from itertools import chain

6 -from typing import TYPE_CHECKING, Optional

6 +from typing import TYPE_CHECKING

7

8 from vllm import envs

9 from vllm.plugins import load_plugins_by_group

@@ -31,7 +31,7 @@ def vllm_version_matches_substr(substr: str) -> bool:

31 return substr in vllm_version

32

33

34 -def tpu_platform_plugin() -> Optional[str]:

34 +def tpu_platform_plugin() -> str | None:

35 logger.debug("Checking if TPU platform is available.")

36

37 # Check for Pathways TPU proxy

@@ -55,7 +55,7 @@ def tpu_platform_plugin() -> Optional[str]:

55 return None

56

57

58 -def cuda_platform_plugin() -> Optional[str]:

58 +def cuda_platform_plugin() -> str | None:

59 is_cuda = False

60 logger.debug("Checking if CUDA platform is available.")

61 try:

@@ -106,7 +106,7 @@ def cuda_platform_plugin() -> Optional[str]:

106 return "vllm.platforms.cuda.CudaPlatform" if is_cuda else None

107

108

109 -def rocm_platform_plugin() -> Optional[str]:

109 +def rocm_platform_plugin() -> str | None:

110 is_rocm = False

111 logger.debug("Checking if ROCm platform is available.")

112 try:

@@ -127,7 +127,7 @@ def rocm_platform_plugin() -> Optional[str]:

127 return "vllm.platforms.rocm.RocmPlatform" if is_rocm else None

128

129

130 -def xpu_platform_plugin() -> Optional[str]:

130 +def xpu_platform_plugin() -> str | None:

131 is_xpu = False

132 logger.debug("Checking if XPU platform is available.")

133 try:

@@ -154,7 +154,7 @@ def xpu_platform_plugin() -> Optional[str]:

154 return "vllm.platforms.xpu.XPUPlatform" if is_xpu else None

155

156

157 -def cpu_platform_plugin() -> Optional[str]:

157 +def cpu_platform_plugin() -> str | None:

158 is_cpu = False

159 logger.debug("Checking if CPU platform is available.")

160 try:

vllm/platforms/cpu.py

+3 -3

@@ -4,13 +4,13 @@

4 import json

5 import os

6 import platform

7 -import re

7 import subprocess

8 import sys

9 from dataclasses import dataclass

10 from importlib.util import find_spec

12 -from typing import TYPE_CHECKING, Optional

11 +from typing import TYPE_CHECKING

12

13 +import regex as re

14 import torch

15

16 from vllm.logger import init_logger

@@ -128,7 +128,7 @@ class CpuPlatform(Platform):

128 selected_backend: "_Backend",

129 head_size: int,

130 dtype: torch.dtype,

131 - kv_cache_dtype: Optional[str],

131 + kv_cache_dtype: str | None,

132 block_size: int,

133 use_v1: bool,

134 use_mla: bool,

vllm/platforms/cuda.py

+6 -5

@@ -5,9 +5,10 @@ pynvml. However, it should not initialize cuda context.

5 """

6

7 import os

8 +from collections.abc import Callable

9 from datetime import timedelta

10 from functools import cache, wraps

10 -from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union

11 +from typing import TYPE_CHECKING, TypeVar

12

13 import torch

14 from torch.distributed import PrefixStore, ProcessGroup

@@ -85,7 +86,7 @@ class CudaPlatformBase(Platform):

86 _ = torch.zeros(1, device=device)

87

88 @classmethod

88 - def get_device_capability(cls, device_id: int = 0) -> Optional[DeviceCapability]:

89 + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:

90 raise NotImplementedError

91

92 @classmethod

@@ -210,7 +211,7 @@ class CudaPlatformBase(Platform):

211

212 @classmethod

213 def get_current_memory_usage(

213 - cls, device: Optional[torch.types.Device] = None

214 + cls, device: torch.types.Device | None = None

215 ) -> float:

216 torch.cuda.empty_cache()

217 torch.cuda.reset_peak_memory_stats(device)

@@ -594,7 +595,7 @@ class NvmlCudaPlatform(CudaPlatformBase):

595 @classmethod

596 @cache

597 @with_nvml_context

597 - def get_device_capability(cls, device_id: int = 0) -> Optional[DeviceCapability]:

598 + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:

599 try:

600 physical_device_id = cls.device_id_to_physical_device_id(device_id)

601 handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)

@@ -607,7 +608,7 @@ class NvmlCudaPlatform(CudaPlatformBase):

608 @with_nvml_context

609 def has_device_capability(

610 cls,

610 - capability: Union[tuple[int, int], int],

611 + capability: tuple[int, int] | int,

612 device_id: int = 0,

613 ) -> bool:

614 try:

vllm/platforms/interface.py

+20 -22

@@ -8,7 +8,7 @@ import random

8 import sys

9 from datetime import timedelta

10 from platform import uname

11 -from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union

11 +from typing import TYPE_CHECKING, Any, NamedTuple

12

13 import numpy as np

14 import torch

@@ -20,18 +20,16 @@ from vllm.logger import init_logger

20 if TYPE_CHECKING:

21 from vllm.attention.backends.registry import _Backend

22 from vllm.config import ModelConfig, VllmConfig

23 - from vllm.lora.request import LoRARequest

23 from vllm.pooling_params import PoolingParams

24 from vllm.sampling_params import SamplingParams

25 from vllm.utils import FlexibleArgumentParser

26 else:

28 - _Backend = None

29 - ModelConfig = None

30 - VllmConfig = None

31 - LoRARequest = None

32 - PoolingParams = None

33 - SamplingParams = None

34 - FlexibleArgumentParser = None

27 + _Backend = object

28 + ModelConfig = object

29 + VllmConfig = object

30 + PoolingParams = object

31 + SamplingParams = object

32 + FlexibleArgumentParser = object

33

34 logger = init_logger(__name__)

35

@@ -113,7 +111,7 @@ class Platform:

111

112 additional_env_vars: list[str] = []

113

116 - _global_graph_pool: Optional[Any] = None

114 + _global_graph_pool: Any | None = None

115

116 @property

117 def supported_dtypes(self) -> list[torch.dtype]:

@@ -180,7 +178,7 @@ class Platform:

178 import vllm._moe_C # noqa: F401

179

180 @classmethod

183 - def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> "_Backend":

181 + def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> _Backend:

182 from vllm.attention.backends.registry import _Backend

183

184 return _Backend.TORCH_SDPA

@@ -188,10 +186,10 @@ class Platform:

186 @classmethod

187 def get_attn_backend_cls(

188 cls,

191 - selected_backend: "_Backend",

189 + selected_backend: _Backend,

190 head_size: int,

191 dtype: torch.dtype,

194 - kv_cache_dtype: Optional[str],

192 + kv_cache_dtype: str | None,

193 block_size: int,

194 use_v1: bool,

195 use_mla: bool,

@@ -205,14 +203,14 @@ class Platform:

203 def get_device_capability(

204 cls,

205 device_id: int = 0,

208 - ) -> Optional[DeviceCapability]:

206 + ) -> DeviceCapability | None:

207 """Stateless version of [torch.cuda.get_device_capability][]."""

208 return None

209

210 @classmethod

211 def has_device_capability(

212 cls,

215 - capability: Union[tuple[int, int], int],

213 + capability: tuple[int, int] | int,

214 device_id: int = 0,

215 ) -> bool:

216 """

@@ -236,7 +234,7 @@ class Platform:

234 @classmethod

235 def is_device_capability(

236 cls,

239 - capability: Union[tuple[int, int], int],

237 + capability: tuple[int, int] | int,

238 device_id: int = 0,

239 ) -> bool:

240 """

@@ -283,7 +281,7 @@ class Platform:

281 return torch.inference_mode(mode=True)

282

283 @classmethod

286 - def seed_everything(cls, seed: Optional[int] = None) -> None:

284 + def seed_everything(cls, seed: int | None = None) -> None:

285 """

286 Set the seed of each random module.

287 `torch.manual_seed` will set seed on all devices.

@@ -304,7 +302,7 @@ class Platform:

302

303 @classmethod

304 def pre_register_and_update(

307 - cls, parser: Optional[FlexibleArgumentParser] = None

305 + cls, parser: FlexibleArgumentParser | None = None

306 ) -> None:

307 """

308 Do some pre-registration or update action for the current platform.

@@ -389,7 +387,7 @@ class Platform:

387

388 @classmethod

389 def get_current_memory_usage(

392 - cls, device: Optional[torch.types.Device] = None

390 + cls, device: torch.types.Device | None = None

391 ) -> float:

392 """

393 Return the memory usage in bytes.

@@ -501,7 +499,7 @@ class Platform:

499 def validate_request(

500 cls,

501 prompt: PromptType,

504 - params: Union[SamplingParams, PoolingParams],

502 + params: SamplingParams | PoolingParams,

503 processed_inputs: ProcessorInputs,

504 ) -> None:

505 """Raises if this request is unsupported on this platform"""

@@ -557,7 +555,7 @@ class Platform:

555

556 @classmethod

557 def is_kv_cache_dtype_supported(

560 - cls, kv_cache_dtype: str, model_config: "ModelConfig"

558 + cls, kv_cache_dtype: str, model_config: ModelConfig

559 ) -> bool:

560 """

561 Returns if the kv_cache_dtype is supported by the current platform.

@@ -617,7 +615,7 @@ class Platform:

615 return {}

616

617 @classmethod

620 - def get_nixl_memory_type(cls) -> Optional[str]:

618 + def get_nixl_memory_type(cls) -> str | None:

619 """

620 Returns the nixl memory type for the current platform.

621 """

vllm/platforms/rocm.py

+5 -5

@@ -4,7 +4,7 @@

4 import os

5 from datetime import timedelta

6 from functools import cache, lru_cache, wraps

7 -from typing import TYPE_CHECKING, Optional

7 +from typing import TYPE_CHECKING

8

9 import torch

10 from torch.distributed import PrefixStore, ProcessGroup

@@ -140,8 +140,8 @@ def use_rocm_custom_paged_attention(

140 max_seq_len: int,

141 sliding_window: int,

142 kv_cache_dtype: str,

143 - alibi_slopes: Optional[torch.Tensor] = None,

144 - sinks: Optional[torch.Tensor] = None,

143 + alibi_slopes: torch.Tensor | None = None,

144 + sinks: torch.Tensor | None = None,

145 ) -> bool:

146 GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName

147 ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])

@@ -320,7 +320,7 @@ class RocmPlatform(Platform):

320

321 @classmethod

322 @lru_cache(maxsize=8)

323 - def get_device_capability(cls, device_id: int = 0) -> Optional[DeviceCapability]:

323 + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:

324 major, minor = torch.cuda.get_device_capability(device_id)

325 return DeviceCapability(major=major, minor=minor)

326

@@ -420,7 +420,7 @@ class RocmPlatform(Platform):

420

421 @classmethod

422 def get_current_memory_usage(

423 - cls, device: Optional[torch.types.Device] = None

423 + cls, device: torch.types.Device | None = None

424 ) -> float:

425 torch.cuda.reset_peak_memory_stats(device)

426 return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(device)[0]

vllm/platforms/tpu.py

+3 -3

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import contextlib

5 -from typing import TYPE_CHECKING, Optional, Union, cast

5 +from typing import TYPE_CHECKING, cast

6

7 import torch

8 from tpu_info import device

@@ -57,7 +57,7 @@ class TpuPlatform(Platform):

57 selected_backend: "_Backend",

58 head_size: int,

59 dtype: torch.dtype,

60 - kv_cache_dtype: Optional[str],

60 + kv_cache_dtype: str | None,

61 block_size: int,

62 use_v1: bool,

63 use_mla: bool,

@@ -211,7 +211,7 @@ class TpuPlatform(Platform):

211 def validate_request(

212 cls,

213 prompt: PromptType,

214 - params: Union[SamplingParams, PoolingParams],

214 + params: SamplingParams | PoolingParams,

215 processed_inputs: ProcessorInputs,

216 ) -> None:

217 """Raises if this request is unsupported on this platform"""

vllm/platforms/xpu.py

+4 -4

@@ -3,7 +3,7 @@

3

4 import contextlib

5 import os

6 -from typing import TYPE_CHECKING, Optional

6 +from typing import TYPE_CHECKING

7

8 import torch

9

@@ -47,7 +47,7 @@ class XPUPlatform(Platform):

47 selected_backend: "_Backend",

48 head_size: int,

49 dtype: torch.dtype,

50 - kv_cache_dtype: Optional[str],

50 + kv_cache_dtype: str | None,

51 block_size: int,

52 use_v1: bool,

53 use_mla: bool,

@@ -113,7 +113,7 @@ class XPUPlatform(Platform):

113 def get_device_capability(

114 cls,

115 device_id: int = 0,

116 - ) -> Optional[DeviceCapability]:

116 + ) -> DeviceCapability | None:

117 # capacity format differs from cuda's and will cause unexpected

118 # failure, so use None directly

119 return None

@@ -213,7 +213,7 @@ class XPUPlatform(Platform):

213

214 @classmethod

215 def get_current_memory_usage(

216 - cls, device: Optional[torch.types.Device] = None

216 + cls, device: torch.types.Device | None = None

217 ) -> float:

218 torch.xpu.reset_peak_memory_stats(device)

219 return torch.xpu.max_memory_allocated(device)

vllm/plugins/__init__.py

+2 -1

@@ -2,7 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import logging

5 -from typing import Any, Callable

5 +from collections.abc import Callable

6 +from typing import Any

7

8 import vllm.envs as envs

9

vllm/plugins/io_processors/__init__.py

-2

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import logging

5

6 from vllm.config import VllmConfig

vllm/plugins/io_processors/interface.py

+7 -7

@@ -3,7 +3,7 @@

3

4 from abc import ABC, abstractmethod

5 from collections.abc import AsyncGenerator, Sequence

6 -from typing import Any, Generic, Optional, TypeVar, Union

6 +from typing import Any, Generic, TypeVar

7

8 from vllm.config import VllmConfig

9 from vllm.entrypoints.openai.protocol import IOProcessorResponse

@@ -22,24 +22,24 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):

22 def pre_process(

23 self,

24 prompt: IOProcessorInput,

25 - request_id: Optional[str] = None,

25 + request_id: str | None = None,

26 **kwargs,

27 - ) -> Union[PromptType, Sequence[PromptType]]:

27 + ) -> PromptType | Sequence[PromptType]:

28 raise NotImplementedError

29

30 async def pre_process_async(

31 self,

32 prompt: IOProcessorInput,

33 - request_id: Optional[str] = None,

33 + request_id: str | None = None,

34 **kwargs,

35 - ) -> Union[PromptType, Sequence[PromptType]]:

35 + ) -> PromptType | Sequence[PromptType]:

36 return self.pre_process(prompt, request_id, **kwargs)

37

38 @abstractmethod

39 def post_process(

40 self,

41 model_output: Sequence[PoolingRequestOutput],

42 - request_id: Optional[str] = None,

42 + request_id: str | None = None,

43 **kwargs,

44 ) -> IOProcessorOutput:

45 raise NotImplementedError

@@ -47,7 +47,7 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):

47 async def post_process_async(

48 self,

49 model_output: AsyncGenerator[tuple[int, PoolingRequestOutput]],

50 - request_id: Optional[str] = None,

50 + request_id: str | None = None,

51 **kwargs,

52 ) -> IOProcessorOutput:

53 # We cannot guarantee outputs are returned in the same order they were

vllm/plugins/lora_resolvers/filesystem_resolver.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import json

4 import os

5 -from typing import Optional

5

6 import vllm.envs as envs

7 from vllm.lora.request import LoRARequest

@@ -15,7 +14,7 @@ class FilesystemResolver(LoRAResolver):

14

15 async def resolve_lora(

16 self, base_model_name: str, lora_name: str

18 - ) -> Optional[LoRARequest]:

17 + ) -> LoRARequest | None:

18 lora_path = os.path.join(self.lora_cache_dir, lora_name)

19 if os.path.exists(lora_path):

20 adapter_config_path = os.path.join(

vllm/pooling_params.py

+9 -9

@@ -34,32 +34,32 @@ class PoolingParams(

34 """

35

36 # --8<-- [start:common-pooling-params]

37 - truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=-1)]] = None

37 + truncate_prompt_tokens: Annotated[int, msgspec.Meta(ge=-1)] | None = None

38 # --8<-- [end:common-pooling-params]

39

40 ## for embeddings models

41 # --8<-- [start:embedding-pooling-params]

42 - dimensions: Optional[int] = None

43 - normalize: Optional[bool] = None

42 + dimensions: int | None = None

43 + normalize: bool | None = None

44 # --8<-- [end:embedding-pooling-params]

45

46 ## for classification, scoring and rerank

47 # --8<-- [start:classification-pooling-params]

48 - activation: Optional[bool] = None

48 + activation: bool | None = None

49 # --8<-- [end:classification-pooling-params]

50

51 ## for reward models

52 - softmax: Optional[bool] = None

53 - step_tag_id: Optional[int] = None

54 - returned_token_ids: Optional[list[int]] = None

52 + softmax: bool | None = None

53 + step_tag_id: int | None = None

54 + returned_token_ids: list[int] | None = None

55

56 - task: Optional[PoolingTask] = None

56 + task: PoolingTask | None = None

57 """Internal use only."""

58

59 requires_token_ids: bool = False

60 """Internal use only."""

61

62 - extra_kwargs: Optional[dict[str, Any]] = None

62 + extra_kwargs: dict[str, Any] | None = None

63 """Internal use only."""

64

65 output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY

vllm/profiler/layerwise_profile.py

+10 -9

@@ -3,8 +3,9 @@

3

4 import copy

5 from collections import defaultdict

6 +from collections.abc import Callable

7 from dataclasses import asdict, dataclass, field

7 -from typing import Any, Callable, Optional, TypeAlias, Union

8 +from typing import Any, Optional, TypeAlias

9

10 import pandas as pd

11 from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult

@@ -62,14 +63,14 @@ class ModelStatsEntry:

63 trace: str

64

65

65 -StatsEntry: TypeAlias = Union[ModelStatsEntry, SummaryStatsEntry]

66 +StatsEntry: TypeAlias = ModelStatsEntry | SummaryStatsEntry

67

68

69 @dataclass

70 class _StatsTreeNode:

71 entry: StatsEntry

72 children: list[StatsEntry]

72 - parent: Optional[StatsEntry]

73 + parent: StatsEntry | None

74

75

76 @dataclass

@@ -82,7 +83,7 @@ class LayerwiseProfileResults(profile):

83 _summary_stats_tree: list[_StatsTreeNode] = field(init=False)

84

85 # profile metadata

85 - num_running_seqs: Optional[int] = None

86 + num_running_seqs: int | None = None

87

88 def __post_init__(self):

89 self._build_correlation_map()

@@ -150,7 +151,7 @@ class LayerwiseProfileResults(profile):

151 @staticmethod

152 def _indent_row_names_based_on_depth(

153 depths_rows: list[tuple[int, StatsEntry]],

153 - indent_style: Union[Callable[[int], str], str] = " ",

154 + indent_style: Callable[[int], str] | str = " ",

155 ):

156 indented_rows = []

157 for depth, row in depths_rows:

@@ -171,7 +172,7 @@ class LayerwiseProfileResults(profile):

172 event_tree = self._kineto_results.experimental_event_tree()

173

174 def _df_traversal(

174 - event: _ProfilerEvent, curr_node: Optional[_ModuleTreeNode] = None

175 + event: _ProfilerEvent, curr_node: _ModuleTreeNode | None = None

176 ):

177 # For the tensor parallel case for now only look at task 1

178 if event.start_tid != 1:

@@ -242,7 +243,7 @@ class LayerwiseProfileResults(profile):

243

244 def build_summary_stats_tree_df(

245 node: _ModuleTreeNode,

245 - parent: Optional[_StatsTreeNode] = None,

246 + parent: _StatsTreeNode | None = None,

247 summary_trace: tuple[str] = (),

248 ):

249 if event_has_module(node.event):

@@ -287,7 +288,7 @@ class LayerwiseProfileResults(profile):

288 self._summary_stats_tree.append(build_summary_stats_tree_df(root))

289

290 def build_model_stats_tree_df(

290 - node: _ModuleTreeNode, parent: Optional[_StatsTreeNode] = None

291 + node: _ModuleTreeNode, parent: _StatsTreeNode | None = None

292 ):

293 if event_has_module(

294 node.event,

@@ -357,7 +358,7 @@ class LayerwiseProfileResults(profile):

358

359

360 class layerwise_profile(profile):

360 - def __init__(self, num_running_seqs: Optional[int] = None):

361 + def __init__(self, num_running_seqs: int | None = None):

362 """

363 layerwise profile constructor.

364

vllm/profiler/utils.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import dataclasses

5 -from typing import Callable, Union

5 +from collections.abc import Callable

6

7 from torch._C._profiler import _EventType, _ProfilerEvent, _TensorMetadata

8

@@ -78,7 +78,7 @@ class TablePrinter:

78

79

80 def indent_string(

81 - string: str, indent: int, indent_style: Union[Callable[[int], str], str] = " "

81 + string: str, indent: int, indent_style: Callable[[int], str] | str = " "

82 ) -> str:

83 if indent:

84 if isinstance(indent_style, str):

vllm/ray/ray_env.py

+3 -4

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import json

4 import os

5 -from typing import Optional

5

6 import vllm.envs as envs

7 from vllm.logger import init_logger

@@ -32,9 +31,9 @@ except json.JSONDecodeError:

31

32

33 def get_env_vars_to_copy(

35 - exclude_vars: Optional[set[str]] = None,

36 - additional_vars: Optional[set[str]] = None,

37 - destination: Optional[str] = None,

34 + exclude_vars: set[str] | None = None,

35 + additional_vars: set[str] | None = None,

36 + destination: str | None = None,

37 ) -> set[str]:

38 """

39 Get the environment variables to copy to downstream Ray actors.

vllm/reasoning/abs_reasoning_parsers.py

+8 -10

@@ -1,13 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import os

5 from abc import abstractmethod

8 -from collections.abc import Sequence

6 +from collections.abc import Callable, Sequence

7 from functools import cached_property

10 -from typing import TYPE_CHECKING, Any, Callable, Union

8 +from typing import TYPE_CHECKING, Any

9

10 from vllm.logger import init_logger

11 from vllm.utils import import_from_path, is_list_of

@@ -78,7 +76,7 @@ class ReasoningParser:

76 def extract_reasoning_content(

77 self,

78 model_output: str,

81 - request: Union[ChatCompletionRequest, ResponsesRequest],

79 + request: ChatCompletionRequest | ResponsesRequest,

80 ) -> tuple[str | None, str | None]:

81 """

82 Extract reasoning content from a complete model-generated string.

@@ -107,7 +105,7 @@ class ReasoningParser:

105 previous_token_ids: Sequence[int],

106 current_token_ids: Sequence[int],

107 delta_token_ids: Sequence[int],

110 - ) -> Union[DeltaMessage, None]:

108 + ) -> DeltaMessage | None:

109 """

110 Instance method that should be implemented for extracting reasoning

111 from an incomplete response; for use when handling reasoning calls and

@@ -136,7 +134,7 @@ class ReasoningParserManager:

134 def _register_module(

135 cls,

136 module: type,

139 - module_name: Union[str, list[str]] | None = None,

137 + module_name: str | list[str] | None = None,

138 force: bool = True,

139 ) -> None:

140 if not issubclass(module, ReasoningParser):

@@ -158,10 +156,10 @@ class ReasoningParserManager:

156 @classmethod

157 def register_module(

158 cls,

161 - name: Union[str, list[str]] | None = None,

159 + name: str | list[str] | None = None,

160 force: bool = True,

163 - module: Union[type, None] = None,

164 - ) -> Union[type, Callable]:

161 + module: type | None = None,

162 + ) -> type | Callable:

163 """

164 Register module with the given name or name list. it can be used as a

165 decoder(with module as None) or normal function(with module as not

vllm/reasoning/basic_parsers.py

+3 -4

@@ -3,7 +3,6 @@

3

4 from abc import abstractmethod

5 from collections.abc import Sequence

6 -from typing import Optional, Union

6

7 from vllm.entrypoints.openai.protocol import (

8 ChatCompletionRequest,

@@ -79,7 +78,7 @@ class BaseThinkingReasoningParser(ReasoningParser):

78 previous_token_ids: Sequence[int],

79 current_token_ids: Sequence[int],

80 delta_token_ids: Sequence[int],

82 - ) -> Union[DeltaMessage, None]:

81 + ) -> DeltaMessage | None:

82 """

83 Extract reasoning content from a delta message.

84 Handles streaming output where previous + delta = current.

@@ -135,8 +134,8 @@ class BaseThinkingReasoningParser(ReasoningParser):

134 return DeltaMessage(content=delta_text)

135

136 def extract_reasoning_content(

138 - self, model_output: str, request: Union[ChatCompletionRequest, ResponsesRequest]

139 - ) -> tuple[Optional[str], Optional[str]]:

137 + self, model_output: str, request: ChatCompletionRequest | ResponsesRequest

138 + ) -> tuple[str | None, str | None]:

139 """

140 Extract reasoning content from the model output.

141

vllm/reasoning/deepseek_r1_reasoning_parser.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Sequence

5 -from typing import Union

5

6 from vllm.entrypoints.openai.protocol import DeltaMessage

7 from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager

@@ -36,7 +35,7 @@ class DeepSeekR1ReasoningParser(BaseThinkingReasoningParser):

35 previous_token_ids: Sequence[int],

36 current_token_ids: Sequence[int],

37 delta_token_ids: Sequence[int],

39 - ) -> Union[DeltaMessage, None]:

38 + ) -> DeltaMessage | None:

39 ret = super().extract_reasoning_content_streaming(

40 previous_text,

41 current_text,

vllm/reasoning/glm4_moe_reasoning_parser.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Sequence

5 -from typing import Optional, Union

5

6 from transformers import PreTrainedTokenizerBase

7

@@ -80,7 +79,7 @@ class Glm4MoeModelReasoningParser(ReasoningParser):

79 previous_token_ids: Sequence[int],

80 current_token_ids: Sequence[int],

81 delta_token_ids: Sequence[int],

83 - ) -> Union[DeltaMessage, None]:

82 + ) -> DeltaMessage | None:

83 """

84 Extract reasoning content from a delta message.

85 Handles streaming output where previous + delta = current.

@@ -137,7 +136,7 @@ class Glm4MoeModelReasoningParser(ReasoningParser):

136

137 def extract_reasoning_content(

138 self, model_output: str, request: ChatCompletionRequest

140 - ) -> tuple[Optional[str], Optional[str]]:

139 + ) -> tuple[str | None, str | None]:

140 """

141 Extract reasoning content from the model output.

142

vllm/reasoning/gptoss_reasoning_parser.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Sequence

5 -from typing import Optional, Union

5

6 from transformers import PreTrainedTokenizerBase

7

@@ -53,7 +52,7 @@ class GptOssReasoningParser(ReasoningParser):

52 previous_token_ids: Sequence[int],

53 current_token_ids: Sequence[int],

54 delta_token_ids: Sequence[int],

56 - ) -> Union[DeltaMessage, None]:

55 + ) -> DeltaMessage | None:

56 prev_reasoning, prev_content, _ = parse_chat_output(list(previous_token_ids))

57 cur_reasoning, cur_content, _ = parse_chat_output(list(current_token_ids))

58 reasoning_delta = None

@@ -78,7 +77,7 @@ class GptOssReasoningParser(ReasoningParser):

77 self,

78 model_output: str,

79 request: ChatCompletionRequest,

81 - ) -> tuple[Optional[str], Optional[str]]:

80 + ) -> tuple[str | None, str | None]:

81 raise NotImplementedError(

82 "gpt-oss has a special branch for parsing reasoning in non-streaming mode. This method shouldn't be used." # noqa: E501

83 )

vllm/reasoning/granite_reasoning_parser.py

+3 -4

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Sequence

5 -from typing import Optional, Union

5

6 import regex as re

7 from transformers import PreTrainedTokenizerBase

@@ -53,7 +52,7 @@ class GraniteReasoningParser(ReasoningParser):

52

53 def extract_reasoning_content(

54 self, model_output: str, request: ChatCompletionRequest

56 - ) -> tuple[Optional[str], Optional[str]]:

55 + ) -> tuple[str | None, str | None]:

56 """Extract the reasoning content & content sections, respectively.

57 If the sequence doesn't match what we expect, i.e., the model generates

58 something else, all content is considered non-reasoning content.

@@ -82,7 +81,7 @@ class GraniteReasoningParser(ReasoningParser):

81 previous_token_ids: Sequence[int],

82 current_token_ids: Sequence[int],

83 delta_token_ids: Sequence[int],

85 - ) -> Union[DeltaMessage, None]:

84 + ) -> DeltaMessage | None:

85 """Extract the reasoning content / content emitted by granite models;

86 If the sequence doesn't match what we expect, i.e., the model generates

87 something else, all content is considered non-reasoning content.

@@ -322,7 +321,7 @@ class GraniteReasoningParser(ReasoningParser):

321

322 def _get_content_sections(

323 self, current_text: str

325 - ) -> tuple[Optional[str], Optional[int], Optional[str]]:

324 + ) -> tuple[str | None, int | None, str | None]:

325 """Parse the text to extract the reasoning content / content

326 if we have them.

327

vllm/reasoning/hunyuan_a13b_reasoning_parser.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Sequence

5 -from typing import Optional, Union

5

6 import regex as re

7 from transformers import PreTrainedTokenizerBase

@@ -90,7 +89,7 @@ class HunyuanA13BReasoningParser(ReasoningParser):

89

90 def extract_reasoning_content(

91 self, model_output: str, request: ChatCompletionRequest

93 - ) -> tuple[Optional[str], Optional[str]]:

92 + ) -> tuple[str | None, str | None]:

93 """Extract the reasoning content & content sections, respectively.

94 If the sequence doesn't match what we expect, i.e., the model generates

95 something else, all content is considered non-reasoning content.

@@ -150,7 +149,7 @@ class HunyuanA13BReasoningParser(ReasoningParser):

149 previous_token_ids: Sequence[int],

150 current_token_ids: Sequence[int],

151 delta_token_ids: Sequence[int],

153 - ) -> Union[DeltaMessage, None]:

152 + ) -> DeltaMessage | None:

153 """Extract content using token ID sequence state machine"""

154 # Define sequences

155 think_start_sequence = self.think_start_ids

vllm/reasoning/olmo3_reasoning_parser.py

+8 -8

@@ -4,7 +4,7 @@

4 import dataclasses as dt

5 import enum

6 from collections.abc import Sequence

7 -from typing import TYPE_CHECKING, Optional, Union

7 +from typing import TYPE_CHECKING

8

9 import regex as re

10

@@ -36,7 +36,7 @@ class Indices:

36 return self.end - self.start

37

38

39 -def string_overlap(a: str, b: str) -> tuple[Optional[Indices], Optional[Indices]]:

39 +def string_overlap(a: str, b: str) -> tuple[Indices | None, Indices | None]:

40 """

41 Find the longest overlap where the end of string a matches the start

42 of string b.

@@ -90,7 +90,7 @@ class Olmo3ReasoningBuffer:

90 # is when we switch to content state.

91 state: Olmo3ReasoningState = Olmo3ReasoningState.REASONING

92

93 - def process_buffer(self) -> Optional[DeltaMessage]:

93 + def process_buffer(self) -> DeltaMessage | None:

94 start_think_idx = self.buffer.find(self.think_start)

95

96 if start_think_idx >= 0:

@@ -142,12 +142,12 @@ class Olmo3ReasoningBuffer:

142 # is the length of the text buffer

143 return len(self.buffer)

144

145 - def add_text(self, delta_text: str) -> Optional[DeltaMessage]:

145 + def add_text(self, delta_text: str) -> DeltaMessage | None:

146 # we start by adding the delta text to the buffer

147 self.buffer += delta_text

148

149 # setting this to empty before starting

150 - delta_message: Optional[DeltaMessage] = None

150 + delta_message: DeltaMessage | None = None

151

152 # we start by computing the overlap between the delta_text

153 # and start/end of think tokens.

@@ -254,8 +254,8 @@ class Olmo3ReasoningParser(ReasoningParser):

254 def extract_reasoning_content(

255 self,

256 model_output: str,

257 - request: Union[ChatCompletionRequest, ResponsesRequest],

258 - ) -> tuple[Optional[str], Optional[str]]:

257 + request: ChatCompletionRequest | ResponsesRequest,

258 + ) -> tuple[str | None, str | None]:

259 """Extract the reasoning content & content sections, respectively.

260 If the sequence doesn't match what we expect, i.e., the model generates

261 something else, all content is considered non-reasoning content.

@@ -287,7 +287,7 @@ class Olmo3ReasoningParser(ReasoningParser):

287 previous_token_ids: Sequence[int],

288 current_token_ids: Sequence[int],

289 delta_token_ids: Sequence[int],

290 - ) -> Union[DeltaMessage, None]:

290 + ) -> DeltaMessage | None:

291 """Extract content using token ID sequence state machine"""

292

293 delta_message = self.buffer.add_text(delta_text)

vllm/reasoning/qwen3_reasoning_parser.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional, Union

4

5 from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ResponsesRequest

6 from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager

@@ -31,8 +30,8 @@ class Qwen3ReasoningParser(BaseThinkingReasoningParser):

30 return "</think>"

31

32 def extract_reasoning_content(

34 - self, model_output: str, request: Union[ChatCompletionRequest, ResponsesRequest]

35 - ) -> tuple[Optional[str], Optional[str]]:

33 + self, model_output: str, request: ChatCompletionRequest | ResponsesRequest

34 + ) -> tuple[str | None, str | None]:

35 """

36 Extract reasoning content from the model output.

37

vllm/reasoning/step3_reasoning_parser.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Sequence

5 -from typing import Optional, Union

5

6 import regex as re

7 from transformers import PreTrainedTokenizerBase

@@ -50,7 +49,7 @@ class Step3ReasoningParser(ReasoningParser):

49 previous_token_ids: Sequence[int],

50 current_token_ids: Sequence[int],

51 delta_token_ids: Sequence[int],

53 - ) -> Union[DeltaMessage, None]:

52 + ) -> DeltaMessage | None:

53 """

54 Extract reasoning content from a delta message.

55 Handles streaming output where previous + delta = current.

@@ -81,7 +80,7 @@ class Step3ReasoningParser(ReasoningParser):

80

81 def extract_reasoning_content(

82 self, model_output: str, request: ChatCompletionRequest

84 - ) -> tuple[Optional[str], Optional[str]]:

83 + ) -> tuple[str | None, str | None]:

84 # Check if the model output contains the </think> token

85 if self.think_end_token not in model_output:

86 # If no </think> token, everything is reasoning content

vllm/sampling_params.py

+51 -51

@@ -7,7 +7,7 @@ import warnings

7 from dataclasses import field

8 from enum import Enum, IntEnum

9 from functools import cached_property

10 -from typing import Annotated, Any, Optional, Union

10 +from typing import Annotated, Any

11

12 import msgspec

13 from pydantic.dataclasses import dataclass

@@ -32,19 +32,19 @@ class SamplingType(IntEnum):

32 @dataclass

33 class StructuredOutputsParams:

34 # One of these fields will be used to build a logit processor.

35 - json: Optional[Union[str, dict]] = None

36 - regex: Optional[str] = None

37 - choice: Optional[list[str]] = None

38 - grammar: Optional[str] = None

39 - json_object: Optional[bool] = None

35 + json: str | dict | None = None

36 + regex: str | None = None

37 + choice: list[str] | None = None

38 + grammar: str | None = None

39 + json_object: bool | None = None

40 # These are other options that can be set.

41 disable_fallback: bool = False

42 disable_any_whitespace: bool = False

43 disable_additional_properties: bool = False

44 - whitespace_pattern: Optional[str] = None

45 - structural_tag: Optional[str] = None

44 + whitespace_pattern: str | None = None

45 + structural_tag: str | None = None

46

47 - _backend: Optional[str] = field(default=None, init=False)

47 + _backend: str | None = field(default=None, init=False)

48 """CAUTION: Should only be set by Processor._validate_structured_output"""

49 _backend_was_auto: bool = field(default=False, init=False)

50 """CAUTION: Should only be set by Processor._validate_structured_output"""

@@ -110,12 +110,12 @@ class SamplingParams(

110 are generated and streamed cumulatively per request. To see all `n`

111 outputs upon completion, use `output_kind=RequestOutputKind.FINAL_ONLY`

112 in `SamplingParams`."""

113 - best_of: Optional[int] = None

113 + best_of: int | None = None

114 """Number of output sequences that are generated from the prompt. From

115 these `best_of` sequences, the top `n` sequences are returned. `best_of`

116 must be greater than or equal to `n`. By default, `best_of` is set to `n`.

117 Warning, this is only supported in V0."""

118 - _real_n: Optional[int] = None

118 + _real_n: int | None = None

119 presence_penalty: float = 0.0

120 """Penalizes new tokens based on whether they appear in the generated text

121 so far. Values > 0 encourage the model to use new tokens, while values < 0

@@ -142,24 +142,24 @@ class SamplingParams(

142 """Represents the minimum probability for a token to be considered,

143 relative to the probability of the most likely token. Must be in [0, 1].

144 Set to 0 to disable this."""

145 - seed: Optional[int] = None

145 + seed: int | None = None

146 """Random seed to use for the generation."""

147 - stop: Optional[Union[str, list[str]]] = None

147 + stop: str | list[str] | None = None

148 """String(s) that stop the generation when they are generated. The returned

149 output will not contain the stop strings."""

150 - stop_token_ids: Optional[list[int]] = None

150 + stop_token_ids: list[int] | None = None

151 """Token IDs that stop the generation when they are generated. The returned

152 output will contain the stop tokens unless the stop tokens are special

153 tokens."""

154 ignore_eos: bool = False

155 """Whether to ignore the EOS token and continue generating

156 tokens after the EOS token is generated."""

157 - max_tokens: Optional[int] = 16

157 + max_tokens: int | None = 16

158 """Maximum number of tokens to generate per output sequence."""

159 min_tokens: int = 0

160 """Minimum number of tokens to generate per output sequence before EOS or

161 `stop_token_ids` can be generated"""

162 - logprobs: Optional[int] = None

162 + logprobs: int | None = None

163 """Number of log probabilities to return per output token. When set to

164 `None`, no probability is returned. If set to a non-`None` value, the

165 result includes the log probabilities of the specified number of most

@@ -167,7 +167,7 @@ class SamplingParams(

167 follows the OpenAI API: The API will always return the log probability of

168 the sampled token, so there may be up to `logprobs+1` elements in the

169 response. When set to -1, return all `vocab_size` log probabilities."""

170 - prompt_logprobs: Optional[int] = None

170 + prompt_logprobs: int | None = None

171 """Number of log probabilities to return per prompt token.

172 When set to -1, return all `vocab_size` log probabilities."""

173 # NOTE: This parameter is only exposed at the engine level for now.

@@ -179,14 +179,14 @@ class SamplingParams(

179 """Whether to skip special tokens in the output."""

180 spaces_between_special_tokens: bool = True

181 """Whether to add spaces between special tokens in the output."""

182 - # Optional[list[LogitsProcessor]] type. We use Any here because

183 - # Optional[list[LogitsProcessor]] type is not supported by msgspec.

184 - logits_processors: Optional[Any] = None

182 + # `list[LogitsProcessor] | None` type. We use Any here because

183 + # `list[LogitsProcessor] | None` type is not supported by msgspec.

184 + logits_processors: Any | None = None

185 """Functions that modify logits based on previously generated tokens, and

186 optionally prompt tokens as a first argument."""

187 include_stop_str_in_output: bool = False

188 """Whether to include the stop strings in output text."""

189 - truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=-1)]] = None

189 + truncate_prompt_tokens: Annotated[int, msgspec.Meta(ge=-1)] | None = None

190 """If set to -1, will use the truncation size supported by the model. If

191 set to an integer k, will use only the last k tokens from the prompt

192 (i.e., left truncation). If set to `None`, truncation is disabled."""

@@ -198,60 +198,60 @@ class SamplingParams(

198 _all_stop_token_ids: set[int] = msgspec.field(default_factory=set)

199

200 # Fields used to construct logits processors

201 - structured_outputs: Optional[StructuredOutputsParams] = None

201 + structured_outputs: StructuredOutputsParams | None = None

202 """Parameters for configuring structured outputs."""

203 - guided_decoding: Optional[GuidedDecodingParams] = None

203 + guided_decoding: GuidedDecodingParams | None = None

204 """Deprecated alias for structured_outputs."""

205 - logit_bias: Optional[dict[int, float]] = None

205 + logit_bias: dict[int, float] | None = None

206 """If provided, the engine will construct a logits processor that applies

207 these logit biases."""

208 - allowed_token_ids: Optional[list[int]] = None

208 + allowed_token_ids: list[int] | None = None

209 """If provided, the engine will construct a logits processor which only

210 retains scores for the given token ids."""

211 - extra_args: Optional[dict[str, Any]] = None

211 + extra_args: dict[str, Any] | None = None

212 """Arbitrary additional args, that can be used by custom sampling

213 implementations, plugins, etc. Not used by any in-tree sampling

214 implementations."""

215

216 # Fields used for bad words

217 - bad_words: Optional[list[str]] = None

217 + bad_words: list[str] | None = None

218 """Words that are not allowed to be generated. More precisely, only the

219 last token of a corresponding token sequence is not allowed when the next

220 generated token can complete the sequence."""

221 - _bad_words_token_ids: Optional[list[list[int]]] = None

221 + _bad_words_token_ids: list[list[int]] | None = None

222

223 @staticmethod

224 def from_optional(

225 - n: Optional[int] = 1,

226 - best_of: Optional[int] = None,

227 - presence_penalty: Optional[float] = 0.0,

228 - frequency_penalty: Optional[float] = 0.0,

229 - repetition_penalty: Optional[float] = 1.0,

230 - temperature: Optional[float] = 1.0,

231 - top_p: Optional[float] = 1.0,

225 + n: int | None = 1,

226 + best_of: int | None = None,

227 + presence_penalty: float | None = 0.0,

228 + frequency_penalty: float | None = 0.0,

229 + repetition_penalty: float | None = 1.0,

230 + temperature: float | None = 1.0,

231 + top_p: float | None = 1.0,

232 top_k: int = 0,

233 min_p: float = 0.0,

234 - seed: Optional[int] = None,

235 - stop: Optional[Union[str, list[str]]] = None,

236 - stop_token_ids: Optional[list[int]] = None,

237 - bad_words: Optional[list[str]] = None,

234 + seed: int | None = None,

235 + stop: str | list[str] | None = None,

236 + stop_token_ids: list[int] | None = None,

237 + bad_words: list[str] | None = None,

238 include_stop_str_in_output: bool = False,

239 ignore_eos: bool = False,

240 - max_tokens: Optional[int] = 16,

240 + max_tokens: int | None = 16,

241 min_tokens: int = 0,

242 - logprobs: Optional[int] = None,

243 - prompt_logprobs: Optional[int] = None,

242 + logprobs: int | None = None,

243 + prompt_logprobs: int | None = None,

244 detokenize: bool = True,

245 skip_special_tokens: bool = True,

246 spaces_between_special_tokens: bool = True,

247 - logits_processors: Optional[list[LogitsProcessor]] = None,

248 - truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=-1)]] = None,

247 + logits_processors: list[LogitsProcessor] | None = None,

248 + truncate_prompt_tokens: Annotated[int, msgspec.Meta(ge=-1)] | None = None,

249 output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE,

250 - structured_outputs: Optional[StructuredOutputsParams] = None,

251 - guided_decoding: Optional[GuidedDecodingParams] = None,

252 - logit_bias: Optional[Union[dict[int, float], dict[str, float]]] = None,

253 - allowed_token_ids: Optional[list[int]] = None,

254 - extra_args: Optional[dict[str, Any]] = None,

250 + structured_outputs: StructuredOutputsParams | None = None,

251 + guided_decoding: GuidedDecodingParams | None = None,

252 + logit_bias: dict[int, float] | dict[str, float] | None = None,

253 + allowed_token_ids: list[int] | None = None,

254 + extra_args: dict[str, Any] | None = None,

255 ) -> "SamplingParams":

256 if logit_bias is not None:

257 # Convert token_id to integer

@@ -483,7 +483,7 @@ class SamplingParams(

483 def update_from_generation_config(

484 self,

485 generation_config: dict[str, Any],

486 - model_eos_token_id: Optional[int] = None,

486 + model_eos_token_id: int | None = None,

487 ) -> None:

488 """Update if there are non-default values from generation_config"""

489

@@ -559,7 +559,7 @@ class SamplingParams(

559 return self._all_stop_token_ids

560

561 @property

562 - def bad_words_token_ids(self) -> Optional[list[list[int]]]:

562 + def bad_words_token_ids(self) -> list[list[int]] | None:

563 # For internal use only. Backward compatibility not guaranteed

564 return self._bad_words_token_ids

565

vllm/scalar_type.py

+6 -7

@@ -5,7 +5,6 @@ import functools

5 import struct

6 from dataclasses import dataclass

7 from enum import Enum

8 -from typing import Optional, Union

8

9 _SCALAR_TYPES_ID_MAP = {}

10

@@ -105,7 +104,7 @@ class ScalarType:

104 double_raw = self._floating_point_max_int()

105 return struct.unpack("!d", struct.pack("!Q", double_raw))[0]

106

108 - def _raw_max(self) -> Union[int, float]:

107 + def _raw_max(self) -> int | float:

108 if self.is_floating_point():

109 return self._floating_point_max()

110 else:

@@ -114,7 +113,7 @@ class ScalarType:

113 )

114 return (1 << self.mantissa) - 1

115

117 - def _raw_min(self) -> Union[int, float]:

116 + def _raw_min(self) -> int | float:

117 if self.is_floating_point():

118 assert self.is_signed(), (

119 "We currently assume all floating point types are signed"

@@ -168,14 +167,14 @@ class ScalarType:

167 def size_bits(self) -> int:

168 return self.exponent + self.mantissa + int(self.signed)

169

171 - def min(self) -> Union[int, float]:

170 + def min(self) -> int | float:

171 """

172 Min representable value for this scalar type.

173 (accounting for bias if there is one)

174 """

175 return self._raw_min() - self.bias

176

178 - def max(self) -> Union[int, float]:

177 + def max(self) -> int | float:

178 """

179 Max representable value for this scalar type.

180 (accounting for bias if there is one)

@@ -265,14 +264,14 @@ class ScalarType:

264 #

265

266 @classmethod

268 - def int_(cls, size_bits: int, bias: Optional[int]) -> "ScalarType":

267 + def int_(cls, size_bits: int, bias: int | None) -> "ScalarType":

268 "Create a signed integer scalar type (size_bits includes sign-bit)."

269 ret = cls(0, size_bits - 1, True, bias if bias else 0)

270 ret.id # noqa B018: make sure the id is cached

271 return ret

272

273 @classmethod

275 - def uint(cls, size_bits: int, bias: Optional[int]) -> "ScalarType":

274 + def uint(cls, size_bits: int, bias: int | None) -> "ScalarType":

275 """Create an unsigned integer scalar type."""

276 ret = cls(0, size_bits, False, bias if bias else 0)

277 ret.id # noqa B018: make sure the id is cached

vllm/sequence.py

+10 -10

@@ -3,7 +3,7 @@

3 """Sequence and its related classes."""

4

5 from dataclasses import dataclass

6 -from typing import TYPE_CHECKING, Any, Optional, Union

6 +from typing import TYPE_CHECKING, Any

7

8 import msgspec

9 import torch

@@ -39,13 +39,13 @@ class RequestMetrics:

39

40 arrival_time: float

41 last_token_time: float

42 - first_scheduled_time: Optional[float]

43 - first_token_time: Optional[float]

44 - time_in_queue: Optional[float]

45 - finished_time: Optional[float] = None

46 - scheduler_time: Optional[float] = None

47 - model_forward_time: Optional[float] = None

48 - model_execute_time: Optional[float] = None

42 + first_scheduled_time: float | None

43 + first_token_time: float | None

44 + time_in_queue: float | None

45 + finished_time: float | None = None

46 + scheduler_time: float | None = None

47 + model_forward_time: float | None = None

48 + model_execute_time: float | None = None

49

50

51 # cannot use msgspec.Struct here because Dynamo does not support it

@@ -59,7 +59,7 @@ class IntermediateTensors:

59 """

60

61 tensors: dict[str, torch.Tensor]

62 - kv_connector_output: Optional[KVConnectorOutput]

62 + kv_connector_output: KVConnectorOutput | None

63

64 def __init__(self, tensors):

65 # manually define this function, so that

@@ -68,7 +68,7 @@ class IntermediateTensors:

68 # a string, and we will lose the information about the source file.

69 self.tensors = tensors

70

71 - def __getitem__(self, key: Union[str, slice]):

71 + def __getitem__(self, key: str | slice):

72 if isinstance(key, str):

73 return self.tensors[key]

74 elif isinstance(key, slice):

vllm/tracing.py

+3 -4

@@ -3,7 +3,6 @@

3

4 import os

5 from collections.abc import Mapping

6 -from typing import Optional

6

7 from vllm.logger import init_logger

8 from vllm.utils import run_once

@@ -13,7 +12,7 @@ TRACE_HEADERS = ["traceparent", "tracestate"]

12 logger = init_logger(__name__)

13

14 _is_otel_imported = False

16 -otel_import_error_traceback: Optional[str] = None

15 +otel_import_error_traceback: str | None = None

16 try:

17 from opentelemetry.context.context import Context

18 from opentelemetry.sdk.environment_variables import (

@@ -55,7 +54,7 @@ def is_otel_available() -> bool:

54

55 def init_tracer(

56 instrumenting_module_name: str, otlp_traces_endpoint: str

58 -) -> Optional[Tracer]:

57 +) -> Tracer | None:

58 if not is_otel_available():

59 raise ValueError(

60 "OpenTelemetry is not available. Unable to initialize "

@@ -88,7 +87,7 @@ def get_span_exporter(endpoint):

87 return OTLPSpanExporter(endpoint=endpoint)

88

89

91 -def extract_trace_context(headers: Optional[Mapping[str, str]]) -> Optional[Context]:

90 +def extract_trace_context(headers: Mapping[str, str] | None) -> Context | None:

91 if is_otel_available():

92 headers = headers or {}

93 return TraceContextTextMapPropagator().extract(headers)

vllm/transformers_utils/chat_templates/registry.py

+6 -5

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Callable

4 from pathlib import Path

4 -from typing import Callable, Optional, Union

5 +from typing import TypeAlias

6

7 from vllm.logger import init_logger

8

@@ -9,17 +10,17 @@ logger = init_logger(__file__)

10

11 CHAT_TEMPLATES_DIR = Path(__file__).parent

12

12 -ChatTemplatePath = Union[Path, Callable[[str], Optional[Path]]]

13 +ChatTemplatePath: TypeAlias = Path | Callable[[str], Path | None]

14

15

15 -def _get_qwen_chat_template_fallback(tokenizer_name_or_path: str) -> Optional[Path]:

16 +def _get_qwen_chat_template_fallback(tokenizer_name_or_path: str) -> Path | None:

17 if tokenizer_name_or_path.endswith("-Chat"):

18 return CHAT_TEMPLATES_DIR / "template_chatml.jinja"

19

20 return CHAT_TEMPLATES_DIR / "template_basic.jinja"

21

22

22 -def _get_minicpmv_chat_template_fallback(tokenizer_name_or_path: str) -> Optional[Path]:

23 +def _get_minicpmv_chat_template_fallback(tokenizer_name_or_path: str) -> Path | None:

24 # MiniCPM-V-4.5 version uses a dedicated template

25 if "4.5" in tokenizer_name_or_path or "4_5" in tokenizer_name_or_path:

26 return CHAT_TEMPLATES_DIR / "template_minicpmv45.jinja"

@@ -58,7 +59,7 @@ def register_chat_template_fallback_path(

59 def get_chat_template_fallback_path(

60 model_type: str,

61 tokenizer_name_or_path: str,

61 -) -> Optional[Path]:

62 +) -> Path | None:

63 chat_template = _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK.get(model_type)

64 if callable(chat_template):

65 chat_template = chat_template(tokenizer_name_or_path)

vllm/transformers_utils/config.py

+47 -46

@@ -4,10 +4,11 @@

4 import json

5 import os

6 import time

7 +from collections.abc import Callable

8 from dataclasses import asdict

9 from functools import cache, partial

10 from pathlib import Path

10 -from typing import Any, Callable, Literal, Optional, TypeVar, Union

11 +from typing import Any, Literal, TypeVar

12

13 import huggingface_hub

14 from huggingface_hub import (

@@ -47,7 +48,7 @@ MISTRAL_CONFIG_NAME = "params.json"

48 logger = init_logger(__name__)

49

50

50 -def _get_hf_token() -> Optional[str]:

51 +def _get_hf_token() -> str | None:

52 """

53 Get the HuggingFace token from environment variable.

54

@@ -108,10 +109,10 @@ _AUTO_CONFIG_KWARGS_OVERRIDES: dict[str, dict[str, Any]] = {

109 class HFConfigParser(ConfigParserBase):

110 def parse(

111 self,

111 - model: Union[str, Path],

112 + model: str | Path,

113 trust_remote_code: bool,

113 - revision: Optional[str] = None,

114 - code_revision: Optional[str] = None,

114 + revision: str | None = None,

115 + code_revision: str | None = None,

116 **kwargs,

117 ) -> tuple[dict, PretrainedConfig]:

118 kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE

@@ -173,10 +174,10 @@ class HFConfigParser(ConfigParserBase):

174 class MistralConfigParser(ConfigParserBase):

175 def parse(

176 self,

176 - model: Union[str, Path],

177 + model: str | Path,

178 trust_remote_code: bool,

178 - revision: Optional[str] = None,

179 - code_revision: Optional[str] = None,

179 + revision: str | None = None,

180 + code_revision: str | None = None,

181 **kwargs,

182 ) -> tuple[dict, PretrainedConfig]:

183 # This function loads a params.json config which

@@ -247,8 +248,8 @@ def register_config_parser(config_format: str):

248 ... self,

249 ... model: Union[str, Path],

250 ... trust_remote_code: bool,

250 - ... revision: Optional[str] = None,

251 - ... code_revision: Optional[str] = None,

251 + ... revision: str | None = None,

252 + ... code_revision: str | None = None,

253 ... **kwargs,

254 ... ) -> tuple[dict, PretrainedConfig]:

255 ... raise NotImplementedError

@@ -310,9 +311,9 @@ def with_retry(

311 def list_repo_files(

312 repo_id: str,

313 *,

313 - revision: Optional[str] = None,

314 - repo_type: Optional[str] = None,

315 - token: Union[str, bool, None] = None,

314 + revision: str | None = None,

315 + repo_type: str | None = None,

316 + token: str | bool | None = None,

317 ) -> list[str]:

318 def lookup_files() -> list[str]:

319 # directly list files if model is local

@@ -348,9 +349,9 @@ def file_exists(

349 repo_id: str,

350 file_name: str,

351 *,

351 - repo_type: Optional[str] = None,

352 - revision: Optional[str] = None,

353 - token: Union[str, bool, None] = None,

352 + repo_type: str | None = None,

353 + revision: str | None = None,

354 + token: str | bool | None = None,

355 ) -> bool:

356 file_list = list_repo_files(

357 repo_id, repo_type=repo_type, revision=revision, token=token

@@ -360,7 +361,7 @@ def file_exists(

361

362 # In offline mode the result can be a false negative

363 def file_or_path_exists(

363 - model: Union[str, Path], config_name: str, revision: Optional[str]

364 + model: str | Path, config_name: str, revision: str | None

365 ) -> bool:

366 if (local_path := Path(model)).exists():

367 return (local_path / config_name).is_file()

@@ -493,10 +494,10 @@ def maybe_override_with_speculators(

494 model: str,

495 tokenizer: str,

496 trust_remote_code: bool,

496 - revision: Optional[str] = None,

497 - vllm_speculative_config: Optional[dict[str, Any]] = None,

497 + revision: str | None = None,

498 + vllm_speculative_config: dict[str, Any] | None = None,

499 **kwargs,

499 -) -> tuple[str, str, Optional[dict[str, Any]]]:

500 +) -> tuple[str, str, dict[str, Any] | None]:

501 """

502 Resolve model configuration when speculators are detected.

503

@@ -551,13 +552,13 @@ def maybe_override_with_speculators(

552

553

554 def get_config(

554 - model: Union[str, Path],

555 + model: str | Path,

556 trust_remote_code: bool,

556 - revision: Optional[str] = None,

557 - code_revision: Optional[str] = None,

558 - config_format: Union[str, ConfigFormat] = "auto",

559 - hf_overrides_kw: Optional[dict[str, Any]] = None,

560 - hf_overrides_fn: Optional[Callable[[PretrainedConfig], PretrainedConfig]] = None,

557 + revision: str | None = None,

558 + code_revision: str | None = None,

559 + config_format: str | ConfigFormat = "auto",

560 + hf_overrides_kw: dict[str, Any] | None = None,

561 + hf_overrides_fn: Callable[[PretrainedConfig], PretrainedConfig] | None = None,

562 **kwargs,

563 ) -> PretrainedConfig:

564 # Separate model folder from file path for GGUF models

@@ -669,8 +670,8 @@ def get_config(

670

671

672 def try_get_local_file(

672 - model: Union[str, Path], file_name: str, revision: Optional[str] = "main"

673 -) -> Optional[Path]:

673 + model: str | Path, file_name: str, revision: str | None = "main"

674 +) -> Path | None:

675 file_path = Path(model) / file_name

676 if file_path.is_file():

677 return file_path

@@ -687,7 +688,7 @@ def try_get_local_file(

688

689

690 def get_hf_file_to_dict(

690 - file_name: str, model: Union[str, Path], revision: Optional[str] = "main"

691 + file_name: str, model: str | Path, revision: str | None = "main"

692 ):

693 """

694 Downloads a file from the Hugging Face Hub and returns

@@ -735,7 +736,7 @@ def get_hf_file_to_dict(

736

737

738 @cache

738 -def get_pooling_config(model: str, revision: Optional[str] = "main") -> Optional[dict]:

739 +def get_pooling_config(model: str, revision: str | None = "main") -> dict | None:

740 """

741 This function gets the pooling and normalize

742 config from the model - only applies to

@@ -799,7 +800,7 @@ def get_pooling_config(model: str, revision: Optional[str] = "main") -> Optional

800 return None

801

802

802 -def get_pooling_config_name(pooling_name: str) -> Union[str, None]:

803 +def get_pooling_config_name(pooling_name: str) -> str | None:

804 if "pooling_mode_" in pooling_name:

805 pooling_name = pooling_name.replace("pooling_mode_", "")

806

@@ -820,7 +821,7 @@ def get_pooling_config_name(pooling_name: str) -> Union[str, None]:

821

822 @cache

823 def get_sentence_transformer_tokenizer_config(

823 - model: Union[str, Path], revision: Optional[str] = "main"

824 + model: str | Path, revision: str | None = "main"

825 ):

826 """

827 Returns the tokenization configuration dictionary for a

@@ -958,9 +959,9 @@ def maybe_register_config_serialize_by_value() -> None:

959

960

961 def get_hf_image_processor_config(

961 - model: Union[str, Path],

962 - hf_token: Optional[Union[bool, str]] = None,

963 - revision: Optional[str] = None,

962 + model: str | Path,

963 + hf_token: bool | str | None = None,

964 + revision: str | None = None,

965 **kwargs,

966 ) -> dict[str, Any]:

967 # ModelScope does not provide an interface for image_processor

@@ -992,9 +993,9 @@ def get_hf_text_config(config: PretrainedConfig):

993 def try_get_generation_config(

994 model: str,

995 trust_remote_code: bool,

995 - revision: Optional[str] = None,

996 - config_format: Union[str, ConfigFormat] = "auto",

997 -) -> Optional[GenerationConfig]:

996 + revision: str | None = None,

997 + config_format: str | ConfigFormat = "auto",

998 +) -> GenerationConfig | None:

999 try:

1000 return GenerationConfig.from_pretrained(

1001 model,

@@ -1016,7 +1017,7 @@ def try_get_generation_config(

1017 def try_get_safetensors_metadata(

1018 model: str,

1019 *,

1019 - revision: Optional[str] = None,

1020 + revision: str | None = None,

1021 ):

1022 get_safetensors_metadata_partial = partial(

1023 get_safetensors_metadata,

@@ -1034,10 +1035,10 @@ def try_get_safetensors_metadata(

1035

1036

1037 def try_get_tokenizer_config(

1037 - pretrained_model_name_or_path: Union[str, os.PathLike],

1038 + pretrained_model_name_or_path: str | os.PathLike,

1039 trust_remote_code: bool,

1039 - revision: Optional[str] = None,

1040 -) -> Optional[dict[str, Any]]:

1040 + revision: str | None = None,

1041 +) -> dict[str, Any] | None:

1042 try:

1043 return get_tokenizer_config(

1044 pretrained_model_name_or_path,

@@ -1051,7 +1052,7 @@ def try_get_tokenizer_config(

1052 def get_safetensors_params_metadata(

1053 model: str,

1054 *,

1054 - revision: Optional[str] = None,

1055 + revision: str | None = None,

1056 ) -> dict[str, Any]:

1057 """

1058 Get the safetensors metadata for remote model repository.

@@ -1112,7 +1113,7 @@ def _maybe_retrieve_max_pos_from_hf(model, revision, **kwargs) -> int:

1113 return max_position_embeddings

1114

1115

1115 -def get_model_path(model: Union[str, Path], revision: Optional[str] = None):

1116 +def get_model_path(model: str | Path, revision: str | None = None):

1117 if os.path.exists(model):

1118 return model

1119 assert huggingface_hub.constants.HF_HUB_OFFLINE

@@ -1132,8 +1133,8 @@ def get_model_path(model: Union[str, Path], revision: Optional[str] = None):

1133

1134

1135 def get_hf_file_bytes(

1135 - file_name: str, model: Union[str, Path], revision: Optional[str] = "main"

1136 -) -> Optional[bytes]:

1136 + file_name: str, model: str | Path, revision: str | None = "main"

1137 +) -> bytes | None:

1138 """Get file contents from HuggingFace repository as bytes."""

1139 file_path = try_get_local_file(model=model, file_name=file_name, revision=revision)

1140

vllm/transformers_utils/config_parser_base.py

+3 -4

@@ -3,7 +3,6 @@

3

4 from abc import ABC, abstractmethod

5 from pathlib import Path

6 -from typing import Optional, Union

6

7 from transformers import PretrainedConfig

8

@@ -12,10 +11,10 @@ class ConfigParserBase(ABC):

11 @abstractmethod

12 def parse(

13 self,

15 - model: Union[str, Path],

14 + model: str | Path,

15 trust_remote_code: bool,

17 - revision: Optional[str] = None,

18 - code_revision: Optional[str] = None,

16 + revision: str | None = None,

17 + code_revision: str | None = None,

18 **kwargs,

19 ) -> tuple[dict, PretrainedConfig]:

20 raise NotImplementedError

vllm/transformers_utils/configs/dotsocr.py

+2 -2

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Any, Optional

3 +from typing import Any

4

5 from transformers.configuration_utils import PretrainedConfig

6 from transformers.models.qwen2 import Qwen2Config

@@ -57,7 +57,7 @@ class DotsOCRConfig(Qwen2Config):

57 self,

58 image_token_id=151665,

59 video_token_id=151656,

60 - vision_config: Optional[dict] = None,

60 + vision_config: dict | None = None,

61 *args,

62 **kwargs,

63 ):

vllm/transformers_utils/configs/eagle.py

+5 -6

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import os

5 -from typing import Optional, Union

5

6 from transformers import AutoConfig, PretrainedConfig

7

@@ -14,12 +13,12 @@ class EAGLEConfig(PretrainedConfig):

13

14 def __init__(

15 self,

17 - model: Union[PretrainedConfig, dict, None] = None,

18 - truncated_vocab_size: Optional[int] = None,

19 - method: Optional[str] = "eagle",

16 + model: PretrainedConfig | dict | None = None,

17 + truncated_vocab_size: int | None = None,

18 + method: str | None = "eagle",

19 **kwargs,

20 ):

22 - model_config: Union[PretrainedConfig, DeepseekV2Config, None]

21 + model_config: PretrainedConfig | DeepseekV2Config | None

22 if isinstance(model, dict):

23 archs = model.get("architectures", [])

24 target_archs = ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]

@@ -84,7 +83,7 @@ class EAGLEConfig(PretrainedConfig):

83 @classmethod

84 def from_pretrained(

85 cls,

87 - pretrained_model_name_or_path: Union[str, os.PathLike],

86 + pretrained_model_name_or_path: str | os.PathLike,

87 **kwargs,

88 ) -> "EAGLEConfig":

89 config_dict, kwargs = cls.get_config_dict(

vllm/transformers_utils/configs/kimi_vl.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 # Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py

4 -from typing import Optional, Union

4

5 from transformers.configuration_utils import PretrainedConfig

6

@@ -14,8 +13,8 @@ class KimiVLConfig(PretrainedConfig):

13

14 def __init__(

15 self,

17 - vision_config: Optional[Union[dict, MoonViTConfig]] = None,

18 - text_config: Optional[Union[dict, DeepseekV2Config]] = None,

16 + vision_config: dict | MoonViTConfig | None = None,

17 + text_config: dict | DeepseekV2Config | None = None,

18 ignore_index: int = -100,

19 media_placeholder_token_id: int = 163605,

20 pad_token_id: int = 0,

vllm/transformers_utils/configs/lfm2_moe.py

+1 -2

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Optional

3

4 from transformers.configuration_utils import PretrainedConfig

5

@@ -115,7 +114,7 @@ class Lfm2MoeConfig(PretrainedConfig):

114 use_expert_bias: bool = True,

115 routed_scaling_factor: float = 1.0,

116 norm_topk_prob: bool = True,

118 - layer_types: Optional[list[str]] = None,

117 + layer_types: list[str] | None = None,

118 **kwargs,

119 ):

120 self.vocab_size = vocab_size

vllm/transformers_utils/configs/medusa.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import os

5 -from typing import Optional, Union

5

6 from transformers import PretrainedConfig

7

@@ -18,7 +17,7 @@ class MedusaConfig(PretrainedConfig):

17 num_hidden_layers: int = 1,

18 max_paths: int = 64,

19 topk: int = 10,

21 - truncated_vocab_size: Optional[int] = None,

20 + truncated_vocab_size: int | None = None,

21 **kwargs,

22 ):

23 self.hidden_size = hidden_size

@@ -39,7 +38,7 @@ class MedusaConfig(PretrainedConfig):

38 @classmethod

39 def from_pretrained(

40 cls,

42 - pretrained_model_name_or_path: Union[str, os.PathLike],

41 + pretrained_model_name_or_path: str | os.PathLike,

42 **kwargs,

43 ) -> "MedusaConfig":

44 config_dict, kwargs = cls.get_config_dict(

vllm/transformers_utils/configs/midashenglm.py

+6 -7

@@ -21,7 +21,6 @@

21 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 -from typing import Optional, Union

24

25 from transformers import PretrainedConfig

26 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (

@@ -36,15 +35,15 @@ class DashengConfig(PretrainedConfig):

35 self,

36 embed_dim: int = 768,

37 outputdim: int = 527,

39 - patch_size: Union[int, tuple[int, int]] = 16,

40 - patch_stride: Union[int, tuple[int, int]] = 16,

38 + patch_size: int | tuple[int, int] = 16,

39 + patch_stride: int | tuple[int, int] = 16,

40 input_channels: int = 1,

41 target_length: int = 1012,

42 depth: int = 12,

43 num_heads: int = 12,

44 mlp_ratio: float = 4.0,

45 qkv_bias: bool = True,

47 - init_values: Optional[float] = None,

46 + init_values: float | None = None,

47 drop_rate: float = 0.0,

48 attn_drop_rate: float = 0.0,

49 f_min: float = 0.0,

@@ -86,10 +85,10 @@ class MiDashengLMConfig(PretrainedConfig):

85

86 def __init__(

87 self,

89 - audio_encoder_config: Optional[dict] = None,

88 + audio_encoder_config: dict | None = None,

89 subsample_factor: int = 5,

91 - text_config: Optional[dict] = None,

92 - audio_token_id: Optional[int] = None,

90 + text_config: dict | None = None,

91 + audio_token_id: int | None = None,

92 **kwargs,

93 ):

94 self.audio_encoder_config = DashengConfig(**(audio_encoder_config or {}))

vllm/transformers_utils/configs/mlp_speculator.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 from transformers import PretrainedConfig

6

@@ -19,7 +18,7 @@ class MLPSpeculatorConfig(PretrainedConfig):

18 emb_dim: int = 4096,

19 inner_dim: int = 0,

20 n_predict: int = 3,

22 - top_k_tokens_per_head: Optional[list[int]] = None,

21 + top_k_tokens_per_head: list[int] | None = None,

22 n_candidates: int = 5,

23 tie_weights: bool = False,

24 scale_input: bool = False,

vllm/transformers_utils/configs/ovis.py

+4 -4

@@ -5,7 +5,7 @@

5 # adapted from https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_aimv2.py

6 # and https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_ovis.py

7 # Ovis Config with AimV2 config registration removed for Transformers compatibility

8 -from typing import Any, Optional, Union

8 +from typing import Any

9

10 from transformers import AutoConfig, PretrainedConfig

11

@@ -76,7 +76,7 @@ class BaseVisualTokenizerConfig(PretrainedConfig):

76 tau=1.0,

77 depths=None,

78 drop_cls_token=False,

79 - backbone_config: Optional[Union[PretrainedConfig, dict]] = None,

79 + backbone_config: PretrainedConfig | dict | None = None,

80 hidden_stride: int = 1,

81 **kwargs,

82 ):

@@ -142,8 +142,8 @@ class OvisConfig(PretrainedConfig):

142

143 def __init__(

144 self,

145 - llm_config: Optional[Union[PretrainedConfig, dict]] = None,

146 - visual_tokenizer_config: Optional[Union[PretrainedConfig, dict]] = None,

145 + llm_config: PretrainedConfig | dict | None = None,

146 + visual_tokenizer_config: PretrainedConfig | dict | None = None,

147 multimodal_max_length=8192,

148 hidden_size=None,

149 conversation_formatter_class=None,

vllm/transformers_utils/configs/radio.py

+3 -5

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Radio vision model configuration"""

4

5 -from typing import Optional, Union

6 -

5 from transformers.configuration_utils import PretrainedConfig

6 from transformers.utils import logging

7

@@ -60,9 +58,9 @@ class RadioConfig(PretrainedConfig):

58 initializer_factor: float = 1.0,

59 hidden_act: str = "gelu",

60 max_img_size: int = 2048,

63 - norm_mean: Union[tuple[float, float, float], list] = OPENAI_CLIP_MEAN,

64 - norm_std: Union[tuple[float, float, float], list] = OPENAI_CLIP_STD,

65 - reg_tokens: Optional[int] = None,

61 + norm_mean: tuple[float, float, float] | list = OPENAI_CLIP_MEAN,

62 + norm_std: tuple[float, float, float] | list = OPENAI_CLIP_STD,

63 + reg_tokens: int | None = None,

64 **kwargs,

65 ):

66 self.model_name = model_name

vllm/transformers_utils/configs/speculators/base.py

+2 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import os

4 -from typing import Any, Union

4 +from typing import Any

5

6 from transformers import PretrainedConfig

7

@@ -18,7 +18,7 @@ class SpeculatorsConfig(PretrainedConfig):

18 @classmethod

19 def from_pretrained(

20 cls,

21 - pretrained_model_name_or_path: Union[str, os.PathLike],

21 + pretrained_model_name_or_path: str | os.PathLike,

22 **kwargs,

23 ) -> "SpeculatorsConfig":

24 """Load speculators Eagle config and convert to vLLM format."""

vllm/transformers_utils/configs/step3_vl.py

+4 -4

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Any, Optional, Union

3 +from typing import Any

4

5 from transformers.configuration_utils import PretrainedConfig

6

@@ -53,7 +53,7 @@ class Step3TextConfig(PretrainedConfig):

53 moe_num_experts: int = 48,

54 moe_top_k: int = 3,

55 rope_theta: float = 500000,

56 - rope_scaling: Optional[dict[str, Any]] = None,

56 + rope_scaling: dict[str, Any] | None = None,

57 max_position_embedding: int = 65536,

58 share_expert_dim: int = 5120,

59 share_q_dim: int = 2048,

@@ -147,8 +147,8 @@ class Step3VLConfig(PretrainedConfig):

147

148 def __init__(

149 self,

150 - vision_config: Optional[Union[dict, Step3VisionEncoderConfig]] = None,

151 - text_config: Optional[Union[dict, Step3TextConfig]] = None,

150 + vision_config: dict | Step3VisionEncoderConfig | None = None,

151 + text_config: dict | Step3TextConfig | None = None,

152 understand_projector_stride: int = 1,

153 projector_bias: bool = True,

154 image_token_id: int = 128001,

vllm/transformers_utils/configs/ultravox.py

+5 -5

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 # Adapted from https://github.com/fixie-ai/ultravox/blob/ecd58c4041030bae2ad15aa6bcf04ab43199ea02/ultravox/model/ultravox_config.py

5 -from typing import Any, Optional

5 +from typing import Any

6

7 import transformers

8

@@ -50,10 +50,10 @@ class UltravoxConfig(transformers.PretrainedConfig):

50

51 def __init__(

52 self,

53 - audio_config: Optional[dict[str, Any]] = None,

54 - text_config: Optional[dict[str, Any]] = None,

55 - audio_model_id: Optional[str] = None,

56 - text_model_id: Optional[str] = None,

53 + audio_config: dict[str, Any] | None = None,

54 + text_config: dict[str, Any] | None = None,

55 + audio_model_id: str | None = None,

56 + text_model_id: str | None = None,

57 ignore_index: int = -100,

58 audio_token_index: int = 32000,

59 hidden_size: int = 4096,

vllm/transformers_utils/detokenizer_utils.py

+2 -3

@@ -1,12 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 from .tokenizer import AnyTokenizer

6

7

9 -def _replace_none_with_empty(tokens: list[Optional[str]]):

8 +def _replace_none_with_empty(tokens: list[str | None]):

9 for i, token in enumerate(tokens):

10 if token is None:

11 tokens[i] = ""

@@ -111,7 +110,7 @@ def convert_ids_list_to_tokens(

110 def detokenize_incrementally(

111 tokenizer: AnyTokenizer,

112 all_input_ids: list[int],

114 - prev_tokens: Optional[list[str]],

113 + prev_tokens: list[str] | None,

114 prefix_offset: int,

115 read_offset: int,

116 skip_special_tokens: bool = False,

vllm/transformers_utils/dynamic_module.py

+8 -9

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import os

4 -from typing import Optional, Union

4

5 from transformers.dynamic_module_utils import get_class_from_dynamic_module

6

@@ -14,18 +13,18 @@ logger = init_logger(__name__)

13 def try_get_class_from_dynamic_module(

14 class_reference: str,

15 pretrained_model_name_or_path: str,

17 - cache_dir: Optional[Union[str, os.PathLike]] = None,

16 + cache_dir: str | os.PathLike | None = None,

17 force_download: bool = False,

19 - resume_download: Optional[bool] = None,

20 - proxies: Optional[dict[str, str]] = None,

21 - token: Optional[Union[bool, str]] = None,

22 - revision: Optional[str] = None,

18 + resume_download: bool | None = None,

19 + proxies: dict[str, str] | None = None,

20 + token: bool | str | None = None,

21 + revision: str | None = None,

22 local_files_only: bool = False,

24 - repo_type: Optional[str] = None,

25 - code_revision: Optional[str] = None,

23 + repo_type: str | None = None,

24 + code_revision: str | None = None,

25 warn_on_fail: bool = True,

26 **kwargs,

28 -) -> Optional[type]:

27 +) -> type | None:

28 """

29 As `transformers.dynamic_module_utils.get_class_from_dynamic_module`,

30 but ignoring any errors.

vllm/transformers_utils/processor.py

+11 -11

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from functools import lru_cache

5 -from typing import TYPE_CHECKING, Any, Optional, Union, cast

5 +from typing import TYPE_CHECKING, Any, cast

6

7 from transformers import (

8 AutoFeatureExtractor,

@@ -45,7 +45,7 @@ class HashableList(list):

45 return hash(tuple(self))

46

47

48 -def _get_processor_factory_fn(processor_cls: Union[type, tuple[type, ...]]):

48 +def _get_processor_factory_fn(processor_cls: type | tuple[type, ...]):

49 if isinstance(processor_cls, tuple) or processor_cls == ProcessorMixin:

50 return AutoProcessor.from_pretrained

51 if hasattr(processor_cls, "from_pretrained"):

@@ -56,7 +56,7 @@ def _get_processor_factory_fn(processor_cls: Union[type, tuple[type, ...]]):

56

57 def _merge_mm_kwargs(

58 model_config: "ModelConfig",

59 - processor_cls: Union[type, tuple[type, ...]],

59 + processor_cls: type | tuple[type, ...],

60 /,

61 **kwargs,

62 ):

@@ -86,9 +86,9 @@ def _merge_mm_kwargs(

86 def get_processor(

87 processor_name: str,

88 *args: Any,

89 - revision: Optional[str] = None,

89 + revision: str | None = None,

90 trust_remote_code: bool = False,

91 - processor_cls: Union[type[_P], tuple[type[_P], ...]] = ProcessorMixin,

91 + processor_cls: type[_P] | tuple[type[_P], ...] = ProcessorMixin,

92 **kwargs: Any,

93 ) -> _P:

94 """Load a processor for the given model name via HuggingFace."""

@@ -146,7 +146,7 @@ cached_get_processor = lru_cache(get_processor)

146

147 def cached_processor_from_config(

148 model_config: "ModelConfig",

149 - processor_cls: Union[type[_P], tuple[type[_P], ...]] = ProcessorMixin,

149 + processor_cls: type[_P] | tuple[type[_P], ...] = ProcessorMixin,

150 **kwargs: Any,

151 ) -> _P:

152 return cached_get_processor(

@@ -161,7 +161,7 @@ def cached_processor_from_config(

161 def get_feature_extractor(

162 processor_name: str,

163 *args: Any,

164 - revision: Optional[str] = None,

164 + revision: str | None = None,

165 trust_remote_code: bool = False,

166 **kwargs: Any,

167 ):

@@ -211,7 +211,7 @@ def cached_feature_extractor_from_config(

211 def get_image_processor(

212 processor_name: str,

213 *args: Any,

214 - revision: Optional[str] = None,

214 + revision: str | None = None,

215 trust_remote_code: bool = False,

216 **kwargs: Any,

217 ):

@@ -261,9 +261,9 @@ def cached_image_processor_from_config(

261 def get_video_processor(

262 processor_name: str,

263 *args: Any,

264 - revision: Optional[str] = None,

264 + revision: str | None = None,

265 trust_remote_code: bool = False,

266 - processor_cls_overrides: Optional[type[_V]] = None,

266 + processor_cls_overrides: type[_V] | None = None,

267 **kwargs: Any,

268 ):

269 """Load a video processor for the given model name via HuggingFace."""

@@ -300,7 +300,7 @@ cached_get_video_processor = lru_cache(get_video_processor)

300

301 def cached_video_processor_from_config(

302 model_config: "ModelConfig",

303 - processor_cls: Optional[type[_V]] = None,

303 + processor_cls: type[_V] | None = None,

304 **kwargs: Any,

305 ):

306 return cached_get_video_processor(

vllm/transformers_utils/processors/ovis.py

+4 -4

@@ -23,7 +23,6 @@

23 # See the License for the specific language governing permissions and

24 # limitations under the License.

25 from functools import cached_property

26 -from typing import Union

26

27 import PIL

28 import torch

@@ -104,9 +103,10 @@ class OvisProcessor(ProcessorMixin):

103 def __call__(

104 self,

105 images: ImageInput = None,

107 - text: Union[

108 - TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]

109 - ] = None,

106 + text: TextInput

107 + | PreTokenizedInput

108 + | list[TextInput]

109 + | list[PreTokenizedInput] = None,

110 **kwargs: Unpack[OvisProcessorKwargs],

111 ) -> BatchFeature:

112 """

vllm/transformers_utils/processors/ovis2_5.py

+9 -9

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 from functools import cached_property

5 -from typing import Optional, Union

5

6 import numpy as np

7 import PIL

@@ -99,10 +98,11 @@ class Ovis2_5Processor(ProcessorMixin):

98 def __call__(

99 self,

100 images: ImageInput = None,

102 - videos: Union[np.ndarray, list[ImageInput]] = None,

103 - text: Union[

104 - TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]

105 - ] = None,

101 + videos: np.ndarray | list[ImageInput] = None,

102 + text: TextInput

103 + | PreTokenizedInput

104 + | list[TextInput]

105 + | list[PreTokenizedInput] = None,

106 **kwargs: Unpack[Ovis2_5ProcessorKwargs],

107 ) -> BatchFeature:

108 """

@@ -376,12 +376,12 @@ class Ovis2_5Processor(ProcessorMixin):

376

377 def preprocess_multidata(

378 self,

379 - images: Optional[Union[PIL.Image.Image, list[PIL.Image.Image]]] = None,

380 - video: Optional[Union[list[PIL.Image.Image], np.ndarray]] = None,

381 - convert_to_rgb: Optional[bool] = True,

379 + images: PIL.Image.Image | list[PIL.Image.Image] | None = None,

380 + video: list[PIL.Image.Image] | np.ndarray | None = None,

381 + convert_to_rgb: bool | None = True,

382 min_pixels: int = MIN_PIXELS,

383 max_pixels: int = MAX_PIXELS,

384 - return_tensors: Optional[str] = "pt",

384 + return_tensors: str | None = "pt",

385 ):

386 is_video = False

387 if images is not None:

vllm/transformers_utils/runai_utils.py

+2 -3

@@ -5,7 +5,6 @@ import hashlib

5 import os

6 import shutil

7 import signal

8 -from typing import Optional

8

9 from vllm import envs

10 from vllm.assets.base import get_cache_dir

@@ -88,8 +87,8 @@ class ObjectStorageModel:

87 def pull_files(

88 self,

89 model_path: str = "",

91 - allow_pattern: Optional[list[str]] = None,

92 - ignore_pattern: Optional[list[str]] = None,

90 + allow_pattern: list[str] | None = None,

91 + ignore_pattern: list[str] | None = None,

92 ) -> None:

93 """

94 Pull files from object storage into the temporary directory.

vllm/transformers_utils/s3_utils.py

+3 -3

@@ -34,7 +34,7 @@ def _filter_ignore(paths: list[str], patterns: list[str]) -> list[str]:

34 def glob(

35 s3: Optional["BaseClient"] = None,

36 path: str = "",

37 - allow_pattern: Optional[list[str]] = None,

37 + allow_pattern: list[str] | None = None,

38 ) -> list[str]:

39 """

40 List full file names from S3 path and filter by allow pattern.

@@ -58,8 +58,8 @@ def glob(

58 def list_files(

59 s3: "BaseClient",

60 path: str,

61 - allow_pattern: Optional[list[str]] = None,

62 - ignore_pattern: Optional[list[str]] = None,

61 + allow_pattern: list[str] | None = None,

62 + ignore_pattern: list[str] | None = None,

63 ) -> tuple[str, str, list[str]]:

64 """

65 List files from S3 path and filter by pattern.

vllm/transformers_utils/tokenizer.py

+9 -9

@@ -7,7 +7,7 @@ import os

7 import warnings

8 from functools import lru_cache

9 from pathlib import Path

10 -from typing import TYPE_CHECKING, Any, Optional, Union

10 +from typing import TYPE_CHECKING, Any, TypeAlias

11

12 import huggingface_hub

13 from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast

@@ -30,14 +30,14 @@ else:

30

31 logger = init_logger(__name__)

32

33 -AnyTokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast, TokenizerBase]

33 +AnyTokenizer: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast | TokenizerBase

34

35

36 def decode_tokens(

37 tokenizer: AnyTokenizer,

38 token_ids: list[int],

39 *,

40 - skip_special_tokens: Optional[bool] = None,

40 + skip_special_tokens: bool | None = None,

41 ) -> str:

42 """

43 Backend-agnostic equivalent of HF's

@@ -56,9 +56,9 @@ def encode_tokens(

56 tokenizer: AnyTokenizer,

57 text: str,

58 *,

59 - truncation: Optional[bool] = None,

60 - max_length: Optional[int] = None,

61 - add_special_tokens: Optional[bool] = None,

59 + truncation: bool | None = None,

60 + max_length: int | None = None,

61 + add_special_tokens: bool | None = None,

62 ) -> list[int]:

63 """

64 Backend-agnostic equivalent of HF's

@@ -137,12 +137,12 @@ def get_cached_tokenizer(tokenizer: AnyTokenizer) -> AnyTokenizer:

137

138

139 def get_tokenizer(

140 - tokenizer_name: Union[str, Path],

140 + tokenizer_name: str | Path,

141 *args,

142 tokenizer_mode: str = "auto",

143 trust_remote_code: bool = False,

144 - revision: Optional[str] = None,

145 - download_dir: Optional[str] = None,

144 + revision: str | None = None,

145 + download_dir: str | None = None,

146 **kwargs,

147 ) -> AnyTokenizer:

148 """Gets a tokenizer for the given model name via HuggingFace or ModelScope."""

vllm/transformers_utils/tokenizer_base.py

+10 -12

@@ -3,7 +3,7 @@

3

4 import importlib

5 from abc import ABC, abstractmethod

6 -from typing import TYPE_CHECKING, Any, Optional, Union

6 +from typing import TYPE_CHECKING, Any

7

8 if TYPE_CHECKING:

9 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

@@ -71,11 +71,11 @@ class TokenizerBase(ABC):

71 @abstractmethod

72 def __call__(

73 self,

74 - text: Union[str, list[str], list[int]],

75 - text_pair: Optional[str] = None,

74 + text: str | list[str] | list[int],

75 + text_pair: str | None = None,

76 add_special_tokens: bool = False,

77 truncation: bool = False,

78 - max_length: Optional[int] = None,

78 + max_length: int | None = None,

79 ):

80 raise NotImplementedError()

81

@@ -92,7 +92,7 @@ class TokenizerBase(ABC):

92 self,

93 text: str,

94 truncation: bool = False,

95 - max_length: Optional[int] = None,

95 + max_length: int | None = None,

96 ) -> list[int]:

97 raise NotImplementedError()

98

@@ -100,9 +100,9 @@ class TokenizerBase(ABC):

100 def encode(

101 self,

102 text: str,

103 - truncation: Optional[bool] = None,

104 - max_length: Optional[int] = None,

105 - add_special_tokens: Optional[bool] = None,

103 + truncation: bool | None = None,

104 + max_length: int | None = None,

105 + add_special_tokens: bool | None = None,

106 ) -> list[int]:

107 raise NotImplementedError()

108

@@ -110,7 +110,7 @@ class TokenizerBase(ABC):

110 def apply_chat_template(

111 self,

112 messages: list["ChatCompletionMessageParam"],

113 - tools: Optional[list[dict[str, Any]]] = None,

113 + tools: list[dict[str, Any]] | None = None,

114 **kwargs,

115 ) -> list[int]:

116 raise NotImplementedError()

@@ -120,9 +120,7 @@ class TokenizerBase(ABC):

120 raise NotImplementedError()

121

122 @abstractmethod

123 - def decode(

124 - self, ids: Union[list[int], int], skip_special_tokens: bool = True

125 - ) -> str:

123 + def decode(self, ids: list[int] | int, skip_special_tokens: bool = True) -> str:

124 raise NotImplementedError()

125

126 @abstractmethod

vllm/transformers_utils/tokenizers/mistral.py

+14 -16

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import TYPE_CHECKING, Any, Optional, Union, cast

4 +from typing import TYPE_CHECKING, Any, cast

5

6 from vllm.logger import init_logger

7 from vllm.transformers_utils.tokenizer_base import TokenizerBase

@@ -90,10 +90,10 @@ def truncate_tool_call_ids(request: "MistralChatCompletionRequest"):

90

91 def _prepare_apply_chat_template_tools_and_messages(

92 messages: list["ChatCompletionMessageParam"],

93 - tools: Optional[list[dict[str, Any]]] = None,

93 + tools: list[dict[str, Any]] | None = None,

94 continue_final_message: bool = False,

95 add_generation_prompt: bool = False,

96 -) -> tuple[list["ChatCompletionMessageParam"], Optional[list[dict[str, Any]]]]:

96 +) -> tuple[list["ChatCompletionMessageParam"], list[dict[str, Any]] | None]:

97 if add_generation_prompt and continue_final_message:

98 raise ValueError(

99 "Cannot set both `add_generation_prompt` and "

@@ -144,7 +144,7 @@ def validate_request_params(request: "ChatCompletionRequest"):

144 raise ValueError("chat_template is not supported for Mistral tokenizers.")

145

146

147 -def _tekken_token_to_id(tokenizer: "Tekkenizer", t: Union[str, bytes]) -> int:

147 +def _tekken_token_to_id(tokenizer: "Tekkenizer", t: str | bytes) -> int:

148 from mistral_common.tokens.tokenizers.tekken import Tekkenizer

149

150 assert isinstance(tokenizer, Tekkenizer), type(tokenizer)

@@ -197,7 +197,7 @@ class MistralTokenizer(TokenizerBase):

197

198 @classmethod

199 def from_pretrained(

200 - cls, path_or_repo_id: str, *, revision: Optional[str] = None

200 + cls, path_or_repo_id: str, *, revision: str | None = None

201 ) -> "MistralTokenizer":

202 from transformers.tokenization_mistral_common import (

203 MistralCommonTokenizer as TransformersMistralTokenizer,

@@ -298,11 +298,11 @@ class MistralTokenizer(TokenizerBase):

298

299 def __call__(

300 self,

301 - text: Union[str, list[str], list[int]],

302 - text_pair: Optional[str] = None,

301 + text: str | list[str] | list[int],

302 + text_pair: str | None = None,

303 add_special_tokens: bool = False,

304 truncation: bool = False,

305 - max_length: Optional[int] = None,

305 + max_length: int | None = None,

306 ):

307 return self.transformers_tokenizer(

308 text=text,

@@ -327,7 +327,7 @@ class MistralTokenizer(TokenizerBase):

327 self,

328 text: str,

329 truncation: bool = False,

330 - max_length: Optional[int] = None,

330 + max_length: int | None = None,

331 ) -> list[int]:

332 # Mistral Tokenizers should not add special tokens

333 return self.transformers_tokenizer.encode(

@@ -337,9 +337,9 @@ class MistralTokenizer(TokenizerBase):

337 def encode(

338 self,

339 text: str,

340 - truncation: Optional[bool] = None,

341 - max_length: Optional[int] = None,

342 - add_special_tokens: Optional[bool] = None,

340 + truncation: bool | None = None,

341 + max_length: int | None = None,

342 + add_special_tokens: bool | None = None,

343 ) -> list[int]:

344 if add_special_tokens is not None:

345 return self.transformers_tokenizer.encode(

@@ -359,7 +359,7 @@ class MistralTokenizer(TokenizerBase):

359 def apply_chat_template(

360 self,

361 messages: list["ChatCompletionMessageParam"],

362 - tools: Optional[list[dict[str, Any]]] = None,

362 + tools: list[dict[str, Any]] | None = None,

363 **kwargs,

364 ) -> list[int]:

365 add_generation_prompt = kwargs.pop("add_generation_prompt", False)

@@ -384,9 +384,7 @@ class MistralTokenizer(TokenizerBase):

384 return_dict=False,

385 )

386

387 - def decode(

388 - self, ids: Union[list[int], int], skip_special_tokens: bool = True

389 - ) -> str:

387 + def decode(self, ids: list[int] | int, skip_special_tokens: bool = True) -> str:

388 return self.transformers_tokenizer.decode(

389 ids, skip_special_tokens=skip_special_tokens

390 )

vllm/transformers_utils/utils.py

+7 -7

@@ -6,7 +6,7 @@ import struct

6 from functools import cache

7 from os import PathLike

8 from pathlib import Path

9 -from typing import Any, Optional, Union

9 +from typing import Any

10

11 from vllm.envs import VLLM_MODEL_REDIRECT_PATH

12 from vllm.logger import init_logger

@@ -18,7 +18,7 @@ def is_s3(model_or_path: str) -> bool:

18 return model_or_path.lower().startswith("s3://")

19

20

21 -def check_gguf_file(model: Union[str, PathLike]) -> bool:

21 +def check_gguf_file(model: str | PathLike) -> bool:

22 """Check if the file is a GGUF model."""

23 model = Path(model)

24 if not model.is_file():

@@ -38,8 +38,8 @@ def check_gguf_file(model: Union[str, PathLike]) -> bool:

38

39 def modelscope_list_repo_files(

40 repo_id: str,

41 - revision: Optional[str] = None,

42 - token: Union[str, bool, None] = None,

41 + revision: str | None = None,

42 + token: str | bool | None = None,

43 ) -> list[str]:

44 """List files in a modelscope repo."""

45 from modelscope.hub.api import HubApi

@@ -57,7 +57,7 @@ def modelscope_list_repo_files(

57 return files

58

59

60 -def _maybe_json_dict(path: Union[str, PathLike]) -> dict[str, str]:

60 +def _maybe_json_dict(path: str | PathLike) -> dict[str, str]:

61 with open(path) as f:

62 try:

63 return json.loads(f.read())

@@ -65,7 +65,7 @@ def _maybe_json_dict(path: Union[str, PathLike]) -> dict[str, str]:

65 return dict[str, str]()

66

67

68 -def _maybe_space_split_dict(path: Union[str, PathLike]) -> dict[str, str]:

68 +def _maybe_space_split_dict(path: str | PathLike) -> dict[str, str]:

69 parsed_dict = dict[str, str]()

70 with open(path) as f:

71 for line in f.readlines():

@@ -104,7 +104,7 @@ def maybe_model_redirect(model: str) -> str:

104 return model

105

106

107 -def parse_safetensors_file_metadata(path: Union[str, PathLike]) -> dict[str, Any]:

107 +def parse_safetensors_file_metadata(path: str | PathLike) -> dict[str, Any]:

108 with open(path, "rb") as f:

109 length_of_metadata = struct.unpack("<Q", f.read(8))[0]

110 metadata = json.loads(f.read(length_of_metadata).decode("utf-8"))

vllm/usage/usage_lib.py

+21 -21

@@ -10,7 +10,7 @@ import time

10 from enum import Enum

11 from pathlib import Path

12 from threading import Thread

13 -from typing import Any, Optional, Union

13 +from typing import Any

14 from uuid import uuid4

15

16 import cpuinfo

@@ -32,7 +32,7 @@ _USAGE_STATS_DO_NOT_TRACK_PATH = os.path.join(_config_home, "do_not_track")

32 _USAGE_STATS_ENABLED = None

33 _USAGE_STATS_SERVER = envs.VLLM_USAGE_STATS_SERVER

34

35 -_GLOBAL_RUNTIME_DATA = dict[str, Union[str, int, bool]]()

35 +_GLOBAL_RUNTIME_DATA = dict[str, str | int | bool]()

36

37 _USAGE_ENV_VARS_TO_COLLECT = [

38 "VLLM_USE_MODELSCOPE",

@@ -46,7 +46,7 @@ _USAGE_ENV_VARS_TO_COLLECT = [

46 ]

47

48

49 -def set_runtime_usage_data(key: str, value: Union[str, int, bool]) -> None:

49 +def set_runtime_usage_data(key: str, value: str | int | bool) -> None:

50 """Set global usage data that will be sent with every usage heartbeat."""

51 _GLOBAL_RUNTIME_DATA[key] = value

52

@@ -131,33 +131,33 @@ class UsageMessage:

131 self.uuid = str(uuid4())

132

133 # Environment Information

134 - self.provider: Optional[str] = None

135 - self.num_cpu: Optional[int] = None

136 - self.cpu_type: Optional[str] = None

137 - self.cpu_family_model_stepping: Optional[str] = None

138 - self.total_memory: Optional[int] = None

139 - self.architecture: Optional[str] = None

140 - self.platform: Optional[str] = None

141 - self.cuda_runtime: Optional[str] = None

142 - self.gpu_count: Optional[int] = None

143 - self.gpu_type: Optional[str] = None

144 - self.gpu_memory_per_device: Optional[int] = None

145 - self.env_var_json: Optional[str] = None

134 + self.provider: str | None = None

135 + self.num_cpu: int | None = None

136 + self.cpu_type: str | None = None

137 + self.cpu_family_model_stepping: str | None = None

138 + self.total_memory: int | None = None

139 + self.architecture: str | None = None

140 + self.platform: str | None = None

141 + self.cuda_runtime: str | None = None

142 + self.gpu_count: int | None = None

143 + self.gpu_type: str | None = None

144 + self.gpu_memory_per_device: int | None = None

145 + self.env_var_json: str | None = None

146

147 # vLLM Information

148 - self.model_architecture: Optional[str] = None

149 - self.vllm_version: Optional[str] = None

150 - self.context: Optional[str] = None

148 + self.model_architecture: str | None = None

149 + self.vllm_version: str | None = None

150 + self.context: str | None = None

151

152 # Metadata

153 - self.log_time: Optional[int] = None

154 - self.source: Optional[str] = None

153 + self.log_time: int | None = None

154 + self.source: str | None = None

155

156 def report_usage(

157 self,

158 model_architecture: str,

159 usage_context: UsageContext,

160 - extra_kvs: Optional[dict[str, Any]] = None,

160 + extra_kvs: dict[str, Any] | None = None,

161 ) -> None:

162 t = Thread(

163 target=self._report_usage_worker,

vllm/utils/__init__.py

+33 -31

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import asyncio

5 import concurrent

6 import contextlib

@@ -46,6 +44,7 @@ from collections import UserDict, defaultdict

44 from collections.abc import (

45 AsyncGenerator,

46 Awaitable,

47 + Callable,

48 Collection,

49 Generator,

50 Hashable,

@@ -62,12 +61,10 @@ from pathlib import Path

61 from typing import (

62 TYPE_CHECKING,

63 Any,

65 - Callable,

64 Generic,

65 Literal,

66 TextIO,

67 TypeVar,

70 - Union,

68 )

69 from urllib.parse import urlparse

70 from uuid import uuid4

@@ -99,6 +96,12 @@ if TYPE_CHECKING:

96

97 from vllm.config import ModelConfig, VllmConfig

98 from vllm.sequence import IntermediateTensors

99 +else:

100 + Namespace = object

101 +

102 + ModelConfig = object

103 + VllmConfig = object

104 + IntermediateTensors = object

105

106 logger = init_logger(__name__)

107

@@ -245,9 +248,7 @@ class AsyncMicrobatchTokenizer:

248 self._queues: dict[

249 tuple,

250 asyncio.Queue[

248 - Union[

249 - tuple[str, dict, asyncio.Future], tuple[list[int], asyncio.Future]

250 - ]

251 + tuple[str, dict, asyncio.Future] | tuple[list[int], asyncio.Future]

252 ],

253 ] = {}

254 self._batcher_tasks: list[asyncio.Task] = []

@@ -274,7 +275,7 @@ class AsyncMicrobatchTokenizer:

275 def _get_queue(

276 self, loop: asyncio.AbstractEventLoop, key: tuple

277 ) -> asyncio.Queue[

277 - Union[tuple[str, dict, asyncio.Future], tuple[list[int], asyncio.Future]]

278 + tuple[str, dict, asyncio.Future] | tuple[list[int], asyncio.Future]

279 ]:

280 """Get the request queue for the given operation key, creating a new

281 queue and batcher task if needed."""

@@ -431,7 +432,7 @@ def cancel_task_threadsafe(task: Task):

432 run_in_loop(task.get_loop(), task.cancel)

433

434

434 -def close_sockets(sockets: Sequence[Union[zmq.Socket, zmq.asyncio.Socket]]):

435 +def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]):

436 for sock in sockets:

437 if sock is not None:

438 sock.close(linger=0)

@@ -767,8 +768,8 @@ def _generate_random_fp8(

768

769

770 def get_kv_cache_torch_dtype(

770 - cache_dtype: Union[str, torch.dtype] | None,

771 - model_dtype: Union[str, torch.dtype] | None = None,

771 + cache_dtype: str | torch.dtype | None,

772 + model_dtype: str | torch.dtype | None = None,

773 ) -> torch.dtype:

774 if isinstance(cache_dtype, str):

775 if cache_dtype == "auto":

@@ -795,8 +796,8 @@ def create_kv_caches_with_random_flash(

796 num_layers: int,

797 num_heads: int,

798 head_size: int,

798 - cache_dtype: Union[str, torch.dtype] | None,

799 - model_dtype: Union[str, torch.dtype] | None = None,

799 + cache_dtype: str | torch.dtype | None,

800 + model_dtype: str | torch.dtype | None = None,

801 seed: int | None = None,

802 device: str | None = "cuda",

803 cache_layout: str | None = "NHD",

@@ -837,8 +838,8 @@ def create_kv_caches_with_random(

838 num_layers: int,

839 num_heads: int,

840 head_size: int,

840 - cache_dtype: Union[str, torch.dtype] | None,

841 - model_dtype: Union[str, torch.dtype] | None = None,

841 + cache_dtype: str | torch.dtype | None,

842 + model_dtype: str | torch.dtype | None = None,

843 seed: int | None = None,

844 device: str | None = "cuda",

845 ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:

@@ -952,7 +953,7 @@ def make_tensor_with_pad(

953 dtype: torch.dtype,

954 *,

955 max_len: int | None = None,

955 - device: Union[str, torch.device] | None = None,

956 + device: str | torch.device | None = None,

957 pin_memory: bool = False,

958 ) -> torch.Tensor:

959 """

@@ -974,7 +975,7 @@ def make_tensor_with_pad(

975 def async_tensor_h2d(

976 data: list,

977 dtype: torch.dtype,

977 - target_device: Union[str, torch.device],

978 + target_device: str | torch.device,

979 pin_memory: bool,

980 ) -> torch.Tensor:

981 """Asynchronously create a tensor and copy it from host to device."""

@@ -1041,7 +1042,7 @@ def as_list(maybe_list: Iterable[T]) -> list[T]:

1042 return maybe_list if isinstance(maybe_list, list) else list(maybe_list)

1043

1044

1044 -def as_iter(obj: Union[T, Iterable[T]]) -> Iterable[T]:

1045 +def as_iter(obj: T | Iterable[T]) -> Iterable[T]:

1046 if isinstance(obj, str) or not isinstance(obj, Iterable):

1047 return [obj] # type: ignore[list-item]

1048 return obj

@@ -1050,7 +1051,7 @@ def as_iter(obj: Union[T, Iterable[T]]) -> Iterable[T]:

1051 # `collections` helpers

1052 def is_list_of(

1053 value: object,

1053 - typ: Union[type[T], tuple[type[T], ...]],

1054 + typ: type[T] | tuple[type[T], ...],

1055 *,

1056 check: Literal["first", "all"] = "first",

1057 ) -> TypeIs[list[T]]:

@@ -1266,7 +1267,7 @@ F = TypeVar("F", bound=Callable[..., Any])

1267

1268 def deprecate_args(

1269 start_index: int,

1269 - is_deprecated: Union[bool, Callable[[], bool]] = True,

1270 + is_deprecated: bool | Callable[[], bool] = True,

1271 additional_message: str | None = None,

1272 ) -> Callable[[F], F]:

1273 if not callable(is_deprecated):

@@ -1306,7 +1307,7 @@ def deprecate_args(

1307

1308 def deprecate_kwargs(

1309 *kws: str,

1309 - is_deprecated: Union[bool, Callable[[], bool]] = True,

1310 + is_deprecated: bool | Callable[[], bool] = True,

1311 additional_message: str | None = None,

1312 ) -> Callable[[F], F]:

1313 deprecated_kws = set(kws)

@@ -1892,7 +1893,7 @@ class FlexibleArgumentParser(ArgumentParser):

1893 # only expecting a flat dictionary of atomic types

1894 processed_args: list[str] = []

1895

1895 - config: dict[str, Union[int, str]] = {}

1896 + config: dict[str, int | str] = {}

1897 try:

1898 with open(file_path) as config_file:

1899 config = yaml.safe_load(config_file)

@@ -2149,10 +2150,11 @@ def weak_ref_tensor(tensor: Any) -> Any:

2150

2151

2152 def weak_ref_tensors(

2152 - tensors: Union[

2153 - torch.Tensor, list[torch.Tensor], tuple[torch.Tensor], IntermediateTensors

2154 - ],

2155 -) -> Union[torch.Tensor, list[Any], tuple[Any], Any]:

2153 + tensors: torch.Tensor

2154 + | list[torch.Tensor]

2155 + | tuple[torch.Tensor]

2156 + | IntermediateTensors,

2157 +) -> torch.Tensor | list[Any] | tuple[Any] | Any:

2158 """

2159 Convenience function to create weak references to tensors,

2160 for single tensor, list of tensors or tuple of tensors.

@@ -2183,7 +2185,7 @@ def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:

2185 return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)

2186

2187

2186 -def import_from_path(module_name: str, file_path: Union[str, os.PathLike]):

2188 +def import_from_path(module_name: str, file_path: str | os.PathLike):

2189 """

2190 Import a Python file according to its file path.

2191

@@ -2584,7 +2586,7 @@ class MemorySnapshot:

2586 self.non_torch_memory = self.cuda_memory - self.torch_memory

2587 self.timestamp = time.time()

2588

2587 - def __sub__(self, other: MemorySnapshot) -> MemorySnapshot:

2589 + def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":

2590 return MemorySnapshot(

2591 torch_peak=self.torch_peak - other.torch_peak,

2592 free_memory=self.free_memory - other.free_memory,

@@ -2778,13 +2780,13 @@ def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:

2780

2781 # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501

2782 def make_zmq_socket(

2781 - ctx: Union[zmq.asyncio.Context, zmq.Context], # type: ignore[name-defined]

2783 + ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined]

2784 path: str,

2785 socket_type: Any,

2786 bind: bool | None = None,

2787 identity: bytes | None = None,

2788 linger: int | None = None,

2787 -) -> Union[zmq.Socket, zmq.asyncio.Socket]: # type: ignore[name-defined]

2789 +) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined]

2790 """Make a ZMQ socket with the proper bind/connect semantics."""

2791

2792 mem = psutil.virtual_memory()

@@ -2950,7 +2952,7 @@ def bind_kv_cache(

2952

2953 def run_method(

2954 obj: Any,

2953 - method: Union[str, bytes, Callable],

2955 + method: str | bytes | Callable,

2956 args: tuple[Any],

2957 kwargs: dict[str, Any],

2958 ) -> Any:

vllm/utils/cache.py

+9 -15

@@ -1,11 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from __future__ import annotations

4 -

3 from collections import UserDict

6 -from collections.abc import Hashable, Iterator, KeysView, Mapping

4 +from collections.abc import Callable, Hashable, Iterator, KeysView, Mapping

5 from types import MappingProxyType

8 -from typing import Callable, Generic, NamedTuple, TypeVar, Union, cast, overload

6 +from typing import Generic, NamedTuple, TypeVar, cast, overload

7

8 import cachetools

9

@@ -43,7 +41,7 @@ class CacheInfo(NamedTuple):

41

42 return self.hits / self.total

43

46 - def __sub__(self, other: CacheInfo):

44 + def __sub__(self, other: "CacheInfo"):

45 return CacheInfo(

46 hits=self.hits - other.hits,

47 total=self.total - other.total,

@@ -129,12 +127,10 @@ class LRUCache(cachetools.LRUCache[_K, _V], Generic[_K, _V]):

127 def get(self, key: _K, /) -> _V | None: ...

128

129 @overload

132 - def get(self, key: _K, /, default: Union[_V, _T]) -> Union[_V, _T]: ...

130 + def get(self, key: _K, /, default: _V | _T) -> _V | _T: ...

131

134 - def get(

135 - self, key: _K, /, default: Union[_V, _T] | None = None

136 - ) -> Union[_V, _T] | None:

137 - value: Union[_V, _T] | None

132 + def get(self, key: _K, /, default: _V | _T | None = None) -> _V | _T | None:

133 + value: _V | _T | None

134 if key in self:

135 value = self.__getitem__(key, update_info=False) # type: ignore[call-arg]

136

@@ -149,12 +145,10 @@ class LRUCache(cachetools.LRUCache[_K, _V], Generic[_K, _V]):

145 def pop(self, key: _K) -> _V: ...

146

147 @overload

152 - def pop(self, key: _K, default: Union[_V, _T]) -> Union[_V, _T]: ...

148 + def pop(self, key: _K, default: _V | _T) -> _V | _T: ...

149

154 - def pop(

155 - self, key: _K, default: Union[_V, _T] | None = None

156 - ) -> Union[_V, _T] | None:

157 - value: Union[_V, _T] | None

150 + def pop(self, key: _K, default: _V | _T | None = None) -> _V | _T | None:

151 + value: _V | _T | None

152 if key not in self:

153 return default

154

vllm/utils/deep_gemm.py

+2 -3

@@ -5,12 +5,11 @@

5 Users of vLLM should always import **only** these wrappers.

6 """

7

8 -from __future__ import annotations

9 -

8 import functools

9 import importlib

10 import os

13 -from typing import Any, Callable, NoReturn

11 +from collections.abc import Callable

12 +from typing import Any, NoReturn

13

14 import torch

15

vllm/utils/flashinfer.py

+2 -3

@@ -5,15 +5,14 @@

5 Users of vLLM should always import **only** these wrappers.

6 """

7

8 -from __future__ import annotations

9 -

8 import contextlib

9 import functools

10 import importlib

11 import importlib.util

12 import os

13 import shutil

16 -from typing import Any, Callable, NoReturn

14 +from collections.abc import Callable

15 +from typing import Any, NoReturn

16

17 import requests

18 import torch

vllm/utils/gc_utils.py

+2 -2

@@ -5,7 +5,7 @@ import json

5 import time

6 from collections import Counter

7 from contextlib import suppress

8 -from typing import Any, Optional

8 +from typing import Any

9

10 from vllm.envs import VLLM_GC_DEBUG

11 from vllm.logger import init_logger

@@ -21,7 +21,7 @@ class GCDebugConfig:

21 - '{"top_objects":5}': enable GC debugger with top 5 collected objects

22 """

23

24 - def __init__(self, gc_debug_conf: Optional[str] = None) -> None:

24 + def __init__(self, gc_debug_conf: str | None = None) -> None:

25 self.enabled: bool = False

26 self.top_objects: int = -1

27

vllm/utils/jsontree.py

+27 -30

@@ -2,9 +2,9 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Helper functions to work with nested JSON structures."""

4

5 -from collections.abc import Iterable

5 +from collections.abc import Callable, Iterable

6 from functools import reduce

7 -from typing import TYPE_CHECKING, Callable, TypeVar, Union, cast, overload

7 +from typing import TYPE_CHECKING, TypeAlias, TypeVar, cast, overload

8

9 if TYPE_CHECKING:

10 import torch

@@ -14,23 +14,20 @@ if TYPE_CHECKING:

14 _T = TypeVar("_T")

15 _U = TypeVar("_U")

16

17 -JSONTree = Union[

18 - dict[str, "JSONTree[_T]"],

19 - list["JSONTree[_T]"],

20 - tuple["JSONTree[_T]", ...],

21 - _T,

22 -]

17 +JSONTree: TypeAlias = (

18 + dict[str, "JSONTree[_T]"] | list["JSONTree[_T]"] | tuple["JSONTree[_T]", ...] | _T

19 +)

20 """A nested JSON structure where the leaves need not be JSON-serializable."""

21

25 -_JSONTree = Union[

26 - dict[str, "JSONTree[_T]"],

27 - list["JSONTree[_T]"],

28 - tuple["JSONTree[_T]", ...],

29 - dict[str, _T],

30 - list[_T],

31 - tuple[_T, ...],

32 - _T,

33 -]

22 +_JSONTree: TypeAlias = (

23 + dict[str, "JSONTree[_T]"]

24 + | list["JSONTree[_T]"]

25 + | tuple["JSONTree[_T]", ...]

26 + | dict[str, _T]

27 + | list[_T]

28 + | tuple[_T, ...]

29 + | _T

30 +)

31 """

32 Same as `JSONTree` but with additional `Union` members to satisfy overloads.

33 """

@@ -58,22 +55,22 @@ def json_map_leaves(

55 @overload

56 def json_map_leaves(

57 func: Callable[[_T], _U],

61 - value: Union[_T, dict[str, _T]],

62 -) -> Union[_U, dict[str, _U]]: ...

58 + value: _T | dict[str, _T],

59 +) -> _U | dict[str, _U]: ...

60

61

62 @overload

63 def json_map_leaves(

64 func: Callable[[_T], _U],

68 - value: Union[_T, list[_T]],

69 -) -> Union[_U, list[_U]]: ...

65 + value: _T | list[_T],

66 +) -> _U | list[_U]: ...

67

68

69 @overload

70 def json_map_leaves(

71 func: Callable[[_T], _U],

75 - value: Union[_T, tuple[_T, ...]],

76 -) -> Union[_U, tuple[_U, ...]]: ...

72 + value: _T | tuple[_T, ...],

73 +) -> _U | tuple[_U, ...]: ...

74

75

76 @overload

@@ -85,8 +82,8 @@ def json_map_leaves(

82

83 def json_map_leaves(

84 func: Callable[[_T], _U],

88 - value: Union["BatchedTensorInputs", _JSONTree[_T]],

89 -) -> Union["BatchedTensorInputs", _JSONTree[_U]]:

85 + value: "BatchedTensorInputs" | _JSONTree[_T],

86 +) -> "BatchedTensorInputs" | _JSONTree[_U]:

87 """Apply a function to each leaf in a nested JSON structure."""

88 if isinstance(value, dict):

89 return {

@@ -104,7 +101,7 @@ def json_map_leaves(

101 @overload

102 def json_reduce_leaves(

103 func: Callable[[_T, _T], _T],

107 - value: Union[_T, dict[str, _T]],

104 + value: _T | dict[str, _T],

105 /,

106 ) -> _T: ...

107

@@ -112,7 +109,7 @@ def json_reduce_leaves(

109 @overload

110 def json_reduce_leaves(

111 func: Callable[[_T, _T], _T],

115 - value: Union[_T, list[_T]],

112 + value: _T | list[_T],

113 /,

114 ) -> _T: ...

115

@@ -120,7 +117,7 @@ def json_reduce_leaves(

117 @overload

118 def json_reduce_leaves(

119 func: Callable[[_T, _T], _T],

123 - value: Union[_T, tuple[_T, ...]],

120 + value: _T | tuple[_T, ...],

121 /,

122 ) -> _T: ...

123

@@ -143,11 +140,11 @@ def json_reduce_leaves(

140

141

142 def json_reduce_leaves(

146 - func: Callable[..., Union[_T, _U]],

143 + func: Callable[..., _T | _U],

144 value: _JSONTree[_T],

145 initial: _U = cast(_U, ...), # noqa: B008

146 /,

150 -) -> Union[_T, _U]:

147 +) -> _T | _U:

148 """

149 Apply a function of two arguments cumulatively to each leaf in a

150 nested JSON structure, from left to right, so as to reduce the

vllm/utils/tensor_schema.py

+12 -10

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Annotated, Any, Optional, Union, get_args, get_origin, get_type_hints

3 +from types import UnionType

4 +from typing import Annotated, Any, Union, get_args, get_origin, get_type_hints

5

6 import torch

7

@@ -12,16 +13,16 @@ logger = init_logger(__name__)

13 class TensorShape:

14 def __init__(

15 self,

15 - *dims: Union[int, str],

16 - dynamic_dims: Optional[set[str]] = None,

16 + *dims: int | str,

17 + dynamic_dims: set[str] | None = None,

18 ) -> None:

19 super().__init__()

20

21 self.dims = dims

22 self.dynamic_dims = dynamic_dims if dynamic_dims else set()

23

23 - def resolve(self, **bindings: int) -> tuple[Union[int, str], ...]:

24 - resolved = list[Union[int, str]]()

24 + def resolve(self, **bindings: int) -> tuple[int | str, ...]:

25 + resolved = list[int | str]()

26 for dim in self.dims:

27 if isinstance(dim, str) and dim in bindings:

28 resolved.append(bindings[dim])

@@ -48,7 +49,7 @@ class TensorSchema:

49 self,

50 *,

51 validate: bool = True,

51 - resolve_bindings: Optional[dict[str, int]] = None,

52 + resolve_bindings: dict[str, int] | None = None,

53 **kwargs: Any,

54 ) -> None:

55 super().__init__()

@@ -71,7 +72,7 @@ class TensorSchema:

72 self,

73 actual: tuple[int, ...],

74 reference: tuple[int, ...],

74 - expected_shape: tuple[Union[int, str], ...],

75 + expected_shape: tuple[int | str, ...],

76 dynamic_dims: set[str],

77 ) -> bool:

78 if len(actual) != len(reference) or len(actual) > len(expected_shape):

@@ -100,7 +101,7 @@ class TensorSchema:

101 self,

102 value: object,

103 field_name: str,

103 - expected_shape: tuple[Union[int, str], ...],

104 + expected_shape: tuple[int | str, ...],

105 dynamic_dims: set[str],

106 leading_idxs: tuple[int, ...] = (),

107 ) -> tuple[int, ...]:

@@ -154,7 +155,7 @@ class TensorSchema:

155 def _validate_tensor_shape_expected(

156 self,

157 actual_shape: tuple[int, ...],

157 - expected_shape: tuple[Union[int, str], ...],

158 + expected_shape: tuple[int | str, ...],

159 field_name: str,

160 shape_env: dict[str, int],

161 dynamic_dims: set[str],

@@ -209,7 +210,8 @@ class TensorSchema:

210 actual_type = args[0]

211

212 # Check arg was provided as Union

212 - if get_origin(actual_type) is Union:

213 + if get_origin(actual_type) in {Union, UnionType}:

214 + # Union for Union[X, Y] and UnionType for X | Y

215 args = get_args(actual_type)

216 # Skip validation when Union contains None

217 if type(None) in args:

vllm/v1/attention/backends/cpu_attn.py

+29 -29

@@ -110,7 +110,7 @@ class TorchSDPAMetadata(AttentionMetadata):

110 """Metadata for PagedAttention."""

111 # (batch_size,). The length of sequences (entire tokens seen so far) per

112 # sequence.

113 - decode_seq_lens_tensor: Optional[torch.Tensor]

113 + decode_seq_lens_tensor: torch.Tensor | None

114 # Maximum sequence length in the batch. 0 if it is prefill-only batch.

115 decode_max_seq_len: int

116 # (batch_size, max_blocks_per_seq).

@@ -119,39 +119,39 @@ class TorchSDPAMetadata(AttentionMetadata):

119 # in the kv cache. Each block can contain up to block_size tokens.

120 # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph

121 # captured.

122 - decode_block_tables: Optional[torch.Tensor]

122 + decode_block_tables: torch.Tensor | None

123 """Metadata for TorchSDPABackend.

124 """

125 # Currently, input sequences can only contain all prompts

126 # or all decoding. True if all sequences are prompts.

127 chunked_prefill: bool

128 - seq_lens: Optional[list[int]] = None # For non-chunked prefill

128 + seq_lens: list[int] | None = None # For non-chunked prefill

129

130 # For chunked prefill only

131 - max_query_len: Optional[int] = None

132 - prefill_max_seq_len: Optional[int] = None

133 - prefill_query_start_loc: Optional[torch.Tensor] = None

134 - prefill_seq_start_loc: Optional[torch.Tensor] = None

135 - prefill_block_tables: Optional[torch.Tensor] = None

131 + max_query_len: int | None = None

132 + prefill_max_seq_len: int | None = None

133 + prefill_query_start_loc: torch.Tensor | None = None

134 + prefill_seq_start_loc: torch.Tensor | None = None

135 + prefill_block_tables: torch.Tensor | None = None

136

137 # For V1 logits index only

138 - query_start_loc: Optional[torch.Tensor] = None

138 + query_start_loc: torch.Tensor | None = None

139

140 # Begin encoder attn & enc/dec cross-attn fields...

141 # Encoder sequence lengths representation

142 - encoder_seq_lens: Optional[list[int]] = None

143 - encoder_seq_lens_tensor: Optional[torch.Tensor] = None

142 + encoder_seq_lens: list[int] | None = None

143 + encoder_seq_lens_tensor: torch.Tensor | None = None

144

145 # Maximum sequence length among encoder sequences

146 - max_encoder_seq_len: Optional[int] = None

146 + max_encoder_seq_len: int | None = None

147

148 # Number of tokens input to encoder

149 - num_encoder_tokens: Optional[int] = None

149 + num_encoder_tokens: int | None = None

150

151 # Cross-attention memory-mapping data structures: slot mapping

152 # and block tables

153 - cross_slot_mapping: Optional[torch.Tensor] = None

154 - cross_block_tables: Optional[torch.Tensor] = None

153 + cross_slot_mapping: torch.Tensor | None = None

154 + cross_block_tables: torch.Tensor | None = None

155

156 def __post_init__(self):

157 # Set during the execution of the first attention op.

@@ -159,9 +159,9 @@ class TorchSDPAMetadata(AttentionMetadata):

159 # when alibi slopes is used. It is because of the limitation

160 # from xformer API.

161 # will not appear in the __repr__ and __init__

162 - self.attn_bias: Optional[list[torch.Tensor]] = None

163 - self.encoder_attn_bias: Optional[list[torch.Tensor]] = None

164 - self.cross_attn_bias: Optional[list[torch.Tensor]] = None

162 + self.attn_bias: list[torch.Tensor] | None = None

163 + self.encoder_attn_bias: list[torch.Tensor] | None = None

164 + self.cross_attn_bias: list[torch.Tensor] | None = None

165

166 @property

167 def is_all_encoder_attn_metadata_set(self):

@@ -237,7 +237,7 @@ class TorchSDPAMetadata(AttentionMetadata):

237 def get_attn_bias(

238 self,

239 attn_type: str,

240 - ) -> Optional[list[torch.Tensor]]:

240 + ) -> list[torch.Tensor] | None:

241 """

242 Extract appropriate attention bias from attention metadata

243 according to attention type.

@@ -439,12 +439,12 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):

439 head_size: int,

440 scale: float,

441 num_kv_heads: int,

442 - alibi_slopes: Optional[list[float]],

443 - sliding_window: Optional[int],

442 + alibi_slopes: list[float] | None,

443 + sliding_window: int | None,

444 kv_cache_dtype: str,

445 - logits_soft_cap: Optional[float] = None,

445 + logits_soft_cap: float | None = None,

446 attn_type: str = AttentionType.DECODER,

447 - kv_sharing_target_layer_name: Optional[str] = None,

447 + kv_sharing_target_layer_name: str | None = None,

448 ) -> None:

449 if kv_sharing_target_layer_name is not None:

450 raise NotImplementedError("KV sharing is not supported in V0.")

@@ -484,9 +484,9 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):

484 value: torch.Tensor,

485 kv_cache: torch.Tensor,

486 attn_metadata: TorchSDPAMetadata, # type: ignore

487 - output: Optional[torch.Tensor] = None,

488 - output_scale: Optional[torch.Tensor] = None,

489 - output_block_scale: Optional[torch.Tensor] = None,

487 + output: torch.Tensor | None = None,

488 + output_scale: torch.Tensor | None = None,

489 + output_block_scale: torch.Tensor | None = None,

490 ) -> torch.Tensor:

491 """Forward pass with torch SDPA and PagedAttention.

492

@@ -737,7 +737,7 @@ def _make_alibi_bias(

737

738 def _make_sliding_window_bias(

739 seq_lens: list[int],

740 - window_size: Optional[int],

740 + window_size: int | None,

741 dtype: torch.dtype,

742 ) -> list[torch.Tensor]:

743 attn_biases: list[torch.Tensor] = []

@@ -824,7 +824,7 @@ class _PagedAttention:

824 kv_cache_dtype: str,

825 num_kv_heads: int,

826 scale: float,

827 - alibi_slopes: Optional[torch.Tensor],

827 + alibi_slopes: torch.Tensor | None,

828 k_scale: torch.Tensor,

829 v_scale: torch.Tensor,

830 *args,

@@ -907,7 +907,7 @@ class _IPEXPagedAttention(_PagedAttention):

907 kv_cache_dtype: str,

908 num_kv_heads: int,

909 scale: float,

910 - alibi_slopes: Optional[torch.Tensor],

910 + alibi_slopes: torch.Tensor | None,

911 k_scale: torch.Tensor,

912 v_scale: torch.Tensor,

913 *args,

vllm/v1/attention/backends/flash_attn.py

+24 -25

@@ -3,7 +3,6 @@

3 """Attention layer with FlashAttention."""

4

5 from dataclasses import dataclass

6 -from typing import Optional, Union

6

7 import numpy as np

8 import torch

@@ -59,7 +58,7 @@ class FlashAttentionBackend(AttentionBackend):

58 return [32, 64, 96, 128, 160, 192, 224, 256]

59

60 @staticmethod

62 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

61 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

62 return [MultipleOf(16)]

63

64 @classmethod

@@ -144,13 +143,13 @@ class FlashAttentionMetadata:

143 # For cascade attention.

144 use_cascade: bool

145 common_prefix_len: int

147 - cu_prefix_query_lens: Optional[torch.Tensor]

148 - prefix_kv_lens: Optional[torch.Tensor]

149 - suffix_kv_lens: Optional[torch.Tensor]

146 + cu_prefix_query_lens: torch.Tensor | None

147 + prefix_kv_lens: torch.Tensor | None

148 + suffix_kv_lens: torch.Tensor | None

149

150 # Optional aot scheduling

152 - scheduler_metadata: Optional[torch.Tensor] = None

153 - prefix_scheduler_metadata: Optional[torch.Tensor] = None

151 + scheduler_metadata: torch.Tensor | None = None

152 + prefix_scheduler_metadata: torch.Tensor | None = None

153 max_num_splits: int = 0

154

155 causal: bool = True

@@ -158,9 +157,9 @@ class FlashAttentionMetadata:

157

158 def _get_sliding_window_configs(

159 vllm_config: VllmConfig,

161 -) -> set[Optional[tuple[int, int]]]:

160 +) -> set[tuple[int, int] | None]:

161 """Get the set of all sliding window configs used in the model."""

163 - sliding_window_configs: set[Optional[tuple[int, int]]] = set()

162 + sliding_window_configs: set[tuple[int, int] | None] = set()

163 layers = get_layers_from_vllm_config(vllm_config, Attention)

164 for layer in layers.values():

165 assert isinstance(layer.impl, FlashAttentionImpl)

@@ -242,7 +241,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

241

242 # Sliding window size to be used with the AOT scheduler will be

243 # populated on first build() call.

245 - self.aot_sliding_window: Optional[tuple[int, int]] = None

244 + self.aot_sliding_window: tuple[int, int] | None = None

245

246 def build(

247 self,

@@ -403,13 +402,13 @@ class FlashAttentionImpl(AttentionImpl):

402 head_size: int,

403 scale: float,

404 num_kv_heads: int,

406 - alibi_slopes: Optional[list[float]],

407 - sliding_window: Optional[int],

405 + alibi_slopes: list[float] | None,

406 + sliding_window: int | None,

407 kv_cache_dtype: str,

409 - logits_soft_cap: Optional[float] = None,

408 + logits_soft_cap: float | None = None,

409 attn_type: AttentionType = AttentionType.DECODER,

411 - kv_sharing_target_layer_name: Optional[str] = None,

412 - sinks: Optional[torch.Tensor] = None,

410 + kv_sharing_target_layer_name: str | None = None,

411 + sinks: torch.Tensor | None = None,

412 ) -> None:

413 self.num_heads = num_heads

414 self.head_size = head_size

@@ -460,9 +459,9 @@ class FlashAttentionImpl(AttentionImpl):

459 value: torch.Tensor,

460 kv_cache: torch.Tensor,

461 attn_metadata: FlashAttentionMetadata,

463 - output: Optional[torch.Tensor] = None,

464 - output_scale: Optional[torch.Tensor] = None,

465 - output_block_scale: Optional[torch.Tensor] = None,

462 + output: torch.Tensor | None = None,

463 + output_scale: torch.Tensor | None = None,

464 + output_block_scale: torch.Tensor | None = None,

465 ) -> torch.Tensor:

466 """Forward pass with FlashAttention.

467

@@ -762,18 +761,18 @@ def cascade_attention(

761 suffix_kv_lens: torch.Tensor,

762 max_kv_len: int,

763 softmax_scale: float,

765 - alibi_slopes: Optional[torch.Tensor],

764 + alibi_slopes: torch.Tensor | None,

765 sliding_window: tuple[int, int],

766 logits_soft_cap: float,

767 block_table: torch.Tensor,

768 common_prefix_len: int,

769 fa_version: int,

771 - prefix_scheduler_metadata: Optional[torch.Tensor] = None,

772 - suffix_scheduler_metadata: Optional[torch.Tensor] = None,

773 - q_descale: Optional[torch.Tensor] = None,

774 - k_descale: Optional[torch.Tensor] = None,

775 - v_descale: Optional[torch.Tensor] = None,

776 - s_aux: Optional[torch.Tensor] = None,

770 + prefix_scheduler_metadata: torch.Tensor | None = None,

771 + suffix_scheduler_metadata: torch.Tensor | None = None,

772 + q_descale: torch.Tensor | None = None,

773 + k_descale: torch.Tensor | None = None,

774 + v_descale: torch.Tensor | None = None,

775 + s_aux: torch.Tensor | None = None,

776 ) -> torch.Tensor:

777 assert alibi_slopes is None, "Cascade attention does not support ALiBi."

778 # TODO: Support sliding window.

vllm/v1/attention/backends/flashinfer.py

+8 -10

@@ -2,10 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Attention layer with FlashInfer."""

4

5 -from __future__ import annotations

6 -

5 from dataclasses import dataclass

8 -from typing import ClassVar, Union

6 +from typing import ClassVar

7

8 import numpy as np

9 import torch

@@ -167,7 +165,7 @@ class FlashInferBackend(AttentionBackend):

165 return [64, 128, 256]

166

167 @staticmethod

170 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

168 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

169 # Note: Not sure for all platforms,

170 # but on Blackwell, only support a page size of

171 # 16, 32, 64

@@ -190,15 +188,15 @@ class FlashInferBackend(AttentionBackend):

188 return "FLASHINFER"

189

190 @staticmethod

193 - def get_impl_cls() -> type[FlashInferImpl]:

191 + def get_impl_cls() -> type["FlashInferImpl"]:

192 return FlashInferImpl

193

194 @staticmethod

197 - def get_metadata_cls() -> type[FlashInferMetadata]:

195 + def get_metadata_cls() -> type["FlashInferMetadata"]:

196 return FlashInferMetadata

197

198 @staticmethod

201 - def get_builder_cls() -> type[FlashInferMetadataBuilder]:

199 + def get_builder_cls() -> type["FlashInferMetadataBuilder"]:

200 return FlashInferMetadataBuilder

201

202 @staticmethod

@@ -1116,9 +1114,9 @@ def fast_plan_decode(

1114 pos_encoding_mode: str = "NONE",

1115 window_left: int = -1,

1116 logits_soft_cap: float | None = None,

1119 - q_data_type: Union[str, torch.dtype] | None = "float16",

1120 - kv_data_type: Union[str, torch.dtype] | None = None,

1121 - data_type: Union[str, torch.dtype] | None = None,

1117 + q_data_type: str | torch.dtype | None = "float16",

1118 + kv_data_type: str | torch.dtype | None = None,

1119 + data_type: str | torch.dtype | None = None,

1120 sm_scale: float | None = None,

1121 rope_scale: float | None = None,

1122 rope_theta: float | None = None,

vllm/v1/attention/backends/flex_attention.py

+21 -22

@@ -3,7 +3,6 @@

3 """Attention layer with FlexAttention."""

4

5 from dataclasses import dataclass

6 -from typing import Optional, Union

6

7 import torch

8 import torch._dynamo.decorators

@@ -282,9 +281,9 @@ class FlexAttentionMetadata:

281

282 use_cascade: bool

283 common_prefix_len: int

285 - cu_prefix_query_lens: Optional[torch.Tensor]

286 - prefix_kv_lens: Optional[torch.Tensor]

287 - suffix_kv_lens: Optional[torch.Tensor]

284 + cu_prefix_query_lens: torch.Tensor | None

285 + prefix_kv_lens: torch.Tensor | None

286 + suffix_kv_lens: torch.Tensor | None

287

288 # Block info

289 total_cache_tokens: int

@@ -300,15 +299,15 @@ class FlexAttentionMetadata:

299

300 # Flex Metadata

301 num_blocks = 0

303 - block_mask: Optional[BlockMask] = None

304 - score_mod: Optional[_score_mod_signature] = None

302 + block_mask: BlockMask | None = None

303 + score_mod: _score_mod_signature | None = None

304 logical_mask_mod: _mask_mod_signature = causal_mask_mod

306 - doc_ids: Optional[torch.Tensor] = None

305 + doc_ids: torch.Tensor | None = None

306 direct_build: bool = True

307 q_block_size: int = 16

308 kv_block_size: int = 16

310 - transformed_score_mod: Optional[_score_mod_signature] = None

311 - sliding_window: Optional[int] = None

309 + transformed_score_mod: _score_mod_signature | None = None

310 + sliding_window: int | None = None

311

312 def _convert_physical_to_logical(

313 self,

@@ -443,7 +442,7 @@ class FlexAttentionMetadata:

442 mask_mod = and_masks(mask_mod, sliding_window_mask_mod)

443 return mask_mod

444

446 - def get_transformed_score_mod(self) -> Optional[_score_mod_signature]:

445 + def get_transformed_score_mod(self) -> _score_mod_signature | None:

446 """Creates the transformed score_mod function for FlexAttention.

447

448 This function wraps the user's score_mod to handle physical-to-logical

@@ -669,9 +668,9 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

668

669

670 class FlexAttentionImpl(AttentionImpl):

672 - sliding_window: Optional[int]

673 - alibi_slopes: Optional[torch.Tensor]

674 - logits_soft_cap: Optional[float]

671 + sliding_window: int | None

672 + alibi_slopes: torch.Tensor | None

673 + logits_soft_cap: float | None

674

675 def __init__(

676 self,

@@ -679,12 +678,12 @@ class FlexAttentionImpl(AttentionImpl):

678 head_size: int,

679 scale: float,

680 num_kv_heads: int,

682 - alibi_slopes: Optional[list[float]],

683 - sliding_window: Optional[int],

681 + alibi_slopes: list[float] | None,

682 + sliding_window: int | None,

683 kv_cache_dtype: str,

685 - logits_soft_cap: Optional[float] = None,

684 + logits_soft_cap: float | None = None,

685 attn_type: AttentionType = AttentionType.DECODER,

687 - kv_sharing_target_layer_name: Optional[str] = None,

686 + kv_sharing_target_layer_name: str | None = None,

687 **kwargs,

688 ) -> None:

689 self.num_heads = num_heads

@@ -742,9 +741,9 @@ class FlexAttentionImpl(AttentionImpl):

741 value: torch.Tensor,

742 kv_cache: torch.Tensor,

743 attn_metadata: FlexAttentionMetadata,

745 - output: Optional[torch.Tensor] = None,

746 - output_scale: Optional[torch.Tensor] = None,

747 - output_block_scale: Optional[torch.Tensor] = None,

744 + output: torch.Tensor | None = None,

745 + output_scale: torch.Tensor | None = None,

746 + output_block_scale: torch.Tensor | None = None,

747 ) -> torch.Tensor:

748 """Forward pass with FLexAttention.

749

@@ -860,8 +859,8 @@ class FlexAttentionImpl(AttentionImpl):

859

860 def get_kernel_options(

861 query, block_m, block_n, use_direct_build: bool

863 -) -> dict[str, Union[int, bool]]:

864 - kernel_options: dict[str, Union[int, bool]] = {

862 +) -> dict[str, int | bool]:

863 + kernel_options: dict[str, int | bool] = {

864 "FORCE_USE_FLEX_ATTENTION": True,

865 }

866 if vllm_kernel_override_batch_invariant():

vllm/v1/attention/backends/gdn_attn.py

+13 -16

@@ -3,7 +3,6 @@

3 """Backend for GatedDeltaNet attention."""

4

5 from dataclasses import dataclass

6 -from typing import Optional

6

7 import torch

8

@@ -36,29 +35,27 @@ class GDNAttentionMetadata:

35 num_spec_decode_tokens: int

36 num_actual_tokens: int

37

39 - has_initial_state: Optional[torch.Tensor] = None

38 + has_initial_state: torch.Tensor | None = None

39

41 - spec_query_start_loc: Optional[torch.Tensor] = (

42 - None # shape: [num_spec_decodes + 1,]

43 - )

44 - non_spec_query_start_loc: Optional[torch.Tensor] = (

40 + spec_query_start_loc: torch.Tensor | None = None # shape: [num_spec_decodes + 1,]

41 + non_spec_query_start_loc: torch.Tensor | None = (

42 None # shape: [batch - num_spec_decodes + 1,]

43 )

44

48 - spec_state_indices_tensor: Optional[torch.Tensor] = None # shape: [batch, num_spec]

49 - non_spec_state_indices_tensor: Optional[torch.Tensor] = (

45 + spec_state_indices_tensor: torch.Tensor | None = None # shape: [batch, num_spec]

46 + non_spec_state_indices_tensor: torch.Tensor | None = (

47 None # shape: [batch - num_spec_decodes,]

48 )

52 - spec_sequence_masks: Optional[torch.Tensor] = None # shape: [batch,]

53 - spec_token_masks: Optional[torch.Tensor] = (

49 + spec_sequence_masks: torch.Tensor | None = None # shape: [batch,]

50 + spec_token_masks: torch.Tensor | None = (

51 None # shape: [num_prefill_tokens + num_decode_tokens,]

52 )

56 - num_accepted_tokens: Optional[torch.Tensor] = None # shape: [batch,]

53 + num_accepted_tokens: torch.Tensor | None = None # shape: [batch,]

54

55 # The following attributes are for triton implementation of causal_conv1d

59 - nums_dict: Optional[dict] = None

60 - batch_ptr: Optional[torch.Tensor] = None

61 - token_chunk_offset_ptr: Optional[torch.Tensor] = None

56 + nums_dict: dict | None = None

57 + batch_ptr: torch.Tensor | None = None

58 + token_chunk_offset_ptr: torch.Tensor | None = None

59

60

61 class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]):

@@ -133,8 +130,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

130 self,

131 common_prefix_len: int,

132 common_attn_metadata: CommonAttentionMetadata,

136 - num_accepted_tokens: Optional[torch.Tensor] = None,

137 - num_decode_draft_tokens_cpu: Optional[torch.Tensor] = None,

133 + num_accepted_tokens: torch.Tensor | None = None,

134 + num_decode_draft_tokens_cpu: torch.Tensor | None = None,

135 fast_build: bool = False,

136 ) -> GDNAttentionMetadata:

137 m = common_attn_metadata

vllm/v1/attention/backends/mamba1_attn.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import torch

7

@@ -26,7 +25,7 @@ class Mamba1AttentionMetadata:

25 query_start_loc: torch.Tensor

26 context_lens_tensor: torch.Tensor

27 state_indices_tensor: torch.Tensor

29 - has_initial_states: Optional[torch.Tensor]

28 + has_initial_states: torch.Tensor | None

29 num_prefills: int

30 num_prefill_tokens: int

31 num_decodes: int

vllm/v1/attention/backends/mamba2_attn.py

+7 -8

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import itertools

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import torch

7

@@ -108,18 +107,18 @@ class Mamba2AttentionMetadata:

107

108 # The following tensors only contain prefill requests and will be None if

109 # the batch has no prefill request.

111 - has_initial_states_p: Optional[torch.Tensor]

112 - seq_idx_p: Optional[torch.Tensor]

110 + has_initial_states_p: torch.Tensor | None

111 + seq_idx_p: torch.Tensor | None

112

113 # cu_chunk_seqlen_p is a tensor of shape (nchunks+1,) that contains, for

114 # each chunk, its offests into the varlen sequence dimension. It is defined

115 # such that the i-th chunk contains tokens from cu_chunk_seqlen_p[i] to

116 # cu_chunk_seqlen_p[i+1].

118 - cu_chunk_seqlen_p: Optional[torch.Tensor]

117 + cu_chunk_seqlen_p: torch.Tensor | None

118

119 # last_chunk_indices_p is a tensor of shape (batch,) that contains the

120 # index of the last chunk for every sequence in the (prefill) batch.

122 - last_chunk_indices_p: Optional[torch.Tensor]

121 + last_chunk_indices_p: torch.Tensor | None

122

123 state_indices_tensor: torch.Tensor # shape: [batch,]

124 block_idx_last_scheduled_token: torch.Tensor # shape: [batch,]

@@ -128,9 +127,9 @@ class Mamba2AttentionMetadata:

127 num_computed_tokens_p: torch.Tensor # shape: [batch,]

128

129 # The following attributes are for triton implementation of causal_conv1d

131 - nums_dict: Optional[dict] = None

132 - batch_ptr: Optional[torch.Tensor] = None

133 - token_chunk_offset_ptr: Optional[torch.Tensor] = None

130 + nums_dict: dict | None = None

131 + batch_ptr: torch.Tensor | None = None

132 + token_chunk_offset_ptr: torch.Tensor | None = None

133

134

135 class Mamba2AttentionMetadataBuilder(

vllm/v1/attention/backends/mla/common.py

+37 -32

@@ -190,7 +190,7 @@ return curr_o @ W_O

190 import functools

191 from abc import abstractmethod

192 from dataclasses import dataclass, field

193 -from typing import ClassVar, Generic, Optional, TypeVar, Union

193 +from typing import ClassVar, Generic, TypeVar

194

195 import torch

196 from tqdm import tqdm

@@ -243,6 +243,8 @@ try:

243

244 flashinfer_available = True

245 except ImportError:

246 + BatchPrefillWithRaggedKVCacheWrapper = object

247 +

248 flashinfer_available = False

249

250

@@ -337,22 +339,22 @@ class MLACommonPrefillMetadata:

339 workspace: torch.Tensor

340

341 # for mla DCP

340 - cp_chunk_seq_lens: Optional[list[list[int]]] = None

341 - origin_context_lens: Optional[list[int]] = None

342 - cp_cu_seq_lens: Optional[torch.Tensor] = None

343 - chunk_size: Optional[int] = None

344 - cu_seq_lens_lst: Optional[list[list[int]]] = None

342 + cp_chunk_seq_lens: list[list[int]] | None = None

343 + origin_context_lens: list[int] | None = None

344 + cp_cu_seq_lens: torch.Tensor | None = None

345 + chunk_size: int | None = None

346 + cu_seq_lens_lst: list[list[int]] | None = None

347

348 block_table: torch.Tensor

349 query_start_loc: torch.Tensor

350 max_query_len: int

349 - chunked_context: Optional[ChunkedContextMetadata] = None

351 + chunked_context: ChunkedContextMetadata | None = None

352

353

354 @dataclass

355 class FlashInferPrefillMetadata(MLACommonPrefillMetadata):

354 - prefill_main: Optional["BatchPrefillWithRaggedKVCacheWrapper"] = None

355 - prefill_chunks: list["BatchPrefillWithRaggedKVCacheWrapper"] = field(

356 + prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None

357 + prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = field(

358 default_factory=list

359 )

360

@@ -362,15 +364,15 @@ class CudnnPrefillMetadata(MLACommonPrefillMetadata):

364 class ChunkedContextMetadata(MLACommonPrefillMetadata.ChunkedContextMetadata):

365 seq_lens: torch.Tensor

366

365 - query_seq_lens: Optional[torch.Tensor] = None

366 - cudnn_workspace: Optional[torch.Tensor] = None

367 + query_seq_lens: torch.Tensor | None = None

368 + cudnn_workspace: torch.Tensor | None = None

369

370

371 @dataclass

372 class MLACommonDecodeMetadata:

373 block_table: torch.Tensor

374 seq_lens: torch.Tensor

373 - dcp_tot_seq_lens: Optional[torch.Tensor]

375 + dcp_tot_seq_lens: torch.Tensor | None

376

377

378 D = TypeVar("D", bound=MLACommonDecodeMetadata)

@@ -407,12 +409,15 @@ class MLACommonMetadata(Generic[D]):

409 num_prefills: int

410

411 # The dimension of the attention heads

410 - head_dim: Optional[int] = None

412 + head_dim: int | None = None

413

412 - decode: Optional[D] = None

413 - prefill: Optional[

414 - Union[MLACommonPrefillMetadata, FlashInferPrefillMetadata, CudnnPrefillMetadata]

415 - ] = None

414 + decode: D | None = None

415 + prefill: (

416 + MLACommonPrefillMetadata

417 + | FlashInferPrefillMetadata

418 + | CudnnPrefillMetadata

419 + | None

420 + ) = None

421

422 def __post_init__(self):

423 if self.head_dim is not None:

@@ -508,7 +513,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

513 layer_names: list[str],

514 vllm_config: VllmConfig,

515 device: torch.device,

511 - metadata_cls: Optional[type[M]] = None,

516 + metadata_cls: type[M] | None = None,

517 ):

518 self.metadata_cls = (

519 metadata_cls if metadata_cls is not None else MLACommonMetadata

@@ -580,7 +585,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

585 FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device=device

586 )

587

583 - self._fi_prefill_main: Optional[BatchPrefillWithRaggedKVCacheWrapper] = None

588 + self._fi_prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None

589 self._fi_prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []

590

591 self._global_hyperparameters = infer_global_hyperparameters(

@@ -683,7 +688,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

688 query_start_loc_cpu: torch.Tensor,

689 query_start_loc_device: torch.Tensor,

690 num_decode_tokens: int,

686 - dcp_tot_seq_lens_device: Optional[torch.Tensor],

691 + dcp_tot_seq_lens_device: torch.Tensor | None,

692 ) -> MLACommonDecodeMetadata:

693 return MLACommonDecodeMetadata(

694 block_table=block_table_tensor,

@@ -1023,14 +1028,14 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

1028 head_size: int,

1029 scale: float,

1030 num_kv_heads: int,

1026 - alibi_slopes: Optional[list[float]],

1027 - sliding_window: Optional[int],

1031 + alibi_slopes: list[float] | None,

1032 + sliding_window: int | None,

1033 kv_cache_dtype: str,

1029 - logits_soft_cap: Optional[float],

1034 + logits_soft_cap: float | None,

1035 attn_type: str,

1031 - kv_sharing_target_layer_name: Optional[str],

1036 + kv_sharing_target_layer_name: str | None,

1037 # MLA Specific Arguments

1033 - q_lora_rank: Optional[int],

1038 + q_lora_rank: int | None,

1039 kv_lora_rank: int,

1040 qk_nope_head_dim: int,

1041 qk_rope_head_dim: int,

@@ -1038,7 +1043,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

1043 v_head_dim: int,

1044 kv_b_proj: ColumnParallelLinear,

1045 indexer=None,

1041 - q_pad_num_heads: Optional[int] = None,

1046 + q_pad_num_heads: int | None = None,

1047 ) -> None:

1048 if kv_sharing_target_layer_name is not None:

1049 raise NotImplementedError("KV sharing is not supported for MLA")

@@ -1226,7 +1231,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1231 and current_platform.get_device_capability()[0] == 9

1232 )

1233

1229 - self.dcp_world_size: Optional[int] = None

1234 + self.dcp_world_size: int | None = None

1235

1236 self.chunked_prefill_workspace_size = (

1237 MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(

@@ -1710,11 +1715,11 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1715 @abstractmethod

1716 def _forward_decode(

1717 self,

1713 - q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

1718 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

1719 kv_c_and_k_pe_cache: torch.Tensor,

1720 attn_metadata: M,

1721 layer: AttentionLayer,

1717 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

1722 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

1723 raise NotImplementedError

1724

1725 def forward(

@@ -1725,9 +1730,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1730 k_pe: torch.Tensor, # value in unified attn

1731 kv_cache: torch.Tensor,

1732 attn_metadata: M,

1728 - output: Optional[torch.Tensor] = None,

1729 - output_scale: Optional[torch.Tensor] = None,

1730 - output_block_scale: Optional[torch.Tensor] = None,

1733 + output: torch.Tensor | None = None,

1734 + output_scale: torch.Tensor | None = None,

1735 + output_block_scale: torch.Tensor | None = None,

1736 ) -> torch.Tensor:

1737 assert output is not None, "Output tensor must be provided."

1738

vllm/v1/attention/backends/mla/cutlass_mla.py

+8 -8

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import os

5 -from typing import ClassVar, Optional, Union

5 +from typing import ClassVar

6

7 import torch

8

@@ -46,7 +46,7 @@ class CutlassMLABackend(MLACommonBackend):

46 return CutlassMLAMetadataBuilder

47

48 @staticmethod

49 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

49 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

50 return [128]

51

52

@@ -95,12 +95,12 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

95 head_size: int,

96 scale: float,

97 num_kv_heads: int,

98 - alibi_slopes: Optional[list[float]],

99 - sliding_window: Optional[int],

98 + alibi_slopes: list[float] | None,

99 + sliding_window: int | None,

100 kv_cache_dtype: str,

101 - logits_soft_cap: Optional[float],

101 + logits_soft_cap: float | None,

102 attn_type: str,

103 - kv_sharing_target_layer_name: Optional[str],

103 + kv_sharing_target_layer_name: str | None,

104 # MLA Specific Arguments

105 **mla_args,

106 ) -> None:

@@ -232,11 +232,11 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

232

233 def _forward_decode(

234 self,

235 - q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

235 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

236 kv_c_and_k_pe_cache: torch.Tensor,

237 attn_metadata: MLACommonMetadata,

238 layer: AttentionLayer,

239 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

239 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

240 assert kv_c_and_k_pe_cache.numel() > 0

241 assert attn_metadata.decode is not None

242

vllm/v1/attention/backends/mla/flashattn_mla.py

+9 -9

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import ClassVar, Optional, Union

5 +from typing import ClassVar

6

7 import torch

8

@@ -55,7 +55,7 @@ class FlashAttnMLADecodeMetadata(MLACommonDecodeMetadata):

55 query_start_loc: torch.Tensor

56 max_query_len: int

57 max_seq_len: int

58 - scheduler_metadata: Optional[torch.Tensor] = None

58 + scheduler_metadata: torch.Tensor | None = None

59 max_num_splits: int = 0

60

61

@@ -135,7 +135,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

135 query_start_loc_cpu: torch.Tensor,

136 query_start_loc_device: torch.Tensor,

137 num_decode_tokens: int,

138 - dcp_tot_seq_lens_device: Optional[torch.Tensor],

138 + dcp_tot_seq_lens_device: torch.Tensor | None,

139 ) -> FlashAttnMLADecodeMetadata:

140 query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

141 max_query_len = query_lens_cpu.max().item()

@@ -195,12 +195,12 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):

195 head_size: int,

196 scale: float,

197 num_kv_heads: int,

198 - alibi_slopes: Optional[list[float]],

199 - sliding_window: Optional[int],

198 + alibi_slopes: list[float] | None,

199 + sliding_window: int | None,

200 kv_cache_dtype: str,

201 - logits_soft_cap: Optional[float],

201 + logits_soft_cap: float | None,

202 attn_type: str,

203 - kv_sharing_target_layer_name: Optional[str],

203 + kv_sharing_target_layer_name: str | None,

204 # MLA Specific Arguments

205 **mla_args,

206 ) -> None:

@@ -242,11 +242,11 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):

242

243 def _forward_decode(

244 self,

245 - q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

245 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

246 kv_c_and_k_pe_cache: torch.Tensor,

247 attn_metadata: FlashAttnMLAMetadata,

248 layer: AttentionLayer,

249 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

249 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

250 assert kv_c_and_k_pe_cache.numel() > 0

251 assert attn_metadata.decode is not None

252

vllm/v1/attention/backends/mla/flashinfer_mla.py

+9 -9

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import ClassVar, Optional, Union

4 +from typing import ClassVar

5

6 import torch

7 from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla

@@ -57,12 +57,12 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

57 head_size: int,

58 scale: float,

59 num_kv_heads: int,

60 - alibi_slopes: Optional[list[float]],

61 - sliding_window: Optional[int],

60 + alibi_slopes: list[float] | None,

61 + sliding_window: int | None,

62 kv_cache_dtype: str,

63 - logits_soft_cap: Optional[float],

63 + logits_soft_cap: float | None,

64 attn_type: str,

65 - kv_sharing_target_layer_name: Optional[str],

65 + kv_sharing_target_layer_name: str | None,

66 # MLA Specific Arguments

67 **mla_args,

68 ) -> None:

@@ -96,16 +96,16 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

96 )

97

98 self._workspace_buffer = g_fi_workspace

99 - self.bmm1_scale: Optional[float] = None

100 - self.bmm2_scale: Optional[float] = None

99 + self.bmm1_scale: float | None = None

100 + self.bmm2_scale: float | None = None

101

102 def _forward_decode(

103 self,

104 - q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

104 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

105 kv_c_and_k_pe_cache: torch.Tensor,

106 attn_metadata: MLACommonMetadata,

107 layer: AttentionLayer,

108 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

108 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

109 assert kv_c_and_k_pe_cache.numel() > 0

110 assert attn_metadata.decode is not None

111

vllm/v1/attention/backends/mla/flashmla.py

+9 -9

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import ClassVar, Optional, Union

5 +from typing import ClassVar

6

7 import torch

8

@@ -45,7 +45,7 @@ class FlashMLABackend(MLACommonBackend):

45 return FlashMLAImpl

46

47 @staticmethod

48 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

48 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

49 return [64]

50

51

@@ -106,7 +106,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

106 query_start_loc_cpu: torch.Tensor,

107 query_start_loc_device: torch.Tensor,

108 num_decode_tokens: int,

109 - dcp_tot_seq_lens_device: Optional[torch.Tensor],

109 + dcp_tot_seq_lens_device: torch.Tensor | None,

110 ) -> FlashMLADecodeMetadata:

111 tile_scheduler_metadata, num_splits = get_mla_metadata(

112 seq_lens_device,

@@ -160,12 +160,12 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

160 head_size: int,

161 scale: float,

162 num_kv_heads: int,

163 - alibi_slopes: Optional[list[float]],

164 - sliding_window: Optional[int],

163 + alibi_slopes: list[float] | None,

164 + sliding_window: int | None,

165 kv_cache_dtype: str,

166 - logits_soft_cap: Optional[float],

166 + logits_soft_cap: float | None,

167 attn_type: str,

168 - kv_sharing_target_layer_name: Optional[str],

168 + kv_sharing_target_layer_name: str | None,

169 # MLA Specific Arguments

170 **mla_args,

171 ) -> None:

@@ -203,11 +203,11 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

203

204 def _forward_decode(

205 self,

206 - q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

206 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

207 kv_c_and_k_pe_cache: torch.Tensor,

208 attn_metadata: FlashMLAMetadata,

209 layer: AttentionLayer,

210 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

210 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

211 # TODO: (zyongye) decode function for mla here

212 assert kv_c_and_k_pe_cache.numel() > 0

213 assert attn_metadata.decode is not None

vllm/v1/attention/backends/mla/flashmla_sparse.py

+10 -10

@@ -110,12 +110,12 @@ class FlashMLASparseMetadata:

110

111 @dataclass

112 class FP8KernelMetadata:

113 - scheduler_metadata: Optional[torch.Tensor]

113 + scheduler_metadata: torch.Tensor | None

114 num_splits: torch.Tensor

115 dummy_block_table: torch.Tensor

116 cache_lens: torch.Tensor

117

118 - fp8_extra_metadata: Optional[FP8KernelMetadata] = None

118 + fp8_extra_metadata: FP8KernelMetadata | None = None

119

120

121 @triton.jit

@@ -373,14 +373,14 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

373 head_size: int,

374 scale: float,

375 num_kv_heads: int,

376 - alibi_slopes: Optional[list[float]],

377 - sliding_window: Optional[int],

376 + alibi_slopes: list[float] | None,

377 + sliding_window: int | None,

378 kv_cache_dtype: str,

379 - logits_soft_cap: Optional[float],

379 + logits_soft_cap: float | None,

380 attn_type: str,

381 - kv_sharing_target_layer_name: Optional[str],

381 + kv_sharing_target_layer_name: str | None,

382 # MLA Specific Arguments

383 - topk_indice_buffer: Optional[torch.Tensor] = None,

383 + topk_indice_buffer: torch.Tensor | None = None,

384 indexer: Optional["Indexer"] = None,

385 **mla_args,

386 ) -> None:

@@ -466,9 +466,9 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

466 k_pe: torch.Tensor, # value in unified attn

467 kv_cache: torch.Tensor,

468 attn_metadata: FlashMLASparseMetadata,

469 - output: Optional[torch.Tensor] = None,

470 - output_scale: Optional[torch.Tensor] = None,

471 - output_block_scale: Optional[torch.Tensor] = None,

469 + output: torch.Tensor | None = None,

470 + output_scale: torch.Tensor | None = None,

471 + output_block_scale: torch.Tensor | None = None,

472 ) -> torch.Tensor:

473 # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use

474 # MQA 576/512 approach for both prefill and decode

vllm/v1/attention/backends/mla/indexer.py

+4 -4

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 -from typing import ClassVar, Optional, Union

4 +from typing import ClassVar

5

6 import torch

7

@@ -52,7 +52,7 @@ class DeepseekV32IndexerBackend(AttentionBackend):

52 return (0, 1, 2)

53

54 @classmethod

55 - def get_supported_kernel_block_size(cls) -> list[Union[int, MultipleOf]]:

55 + def get_supported_kernel_block_size(cls) -> list[int | MultipleOf]:

56 return [64]

57

58

@@ -105,8 +105,8 @@ class DeepseekV32IndexerMetadata:

105 num_prefills: int

106 num_prefill_tokens: int

107

108 - decode: Optional[DeepSeekV32IndexerDecodeMetadata] = None

109 - prefill: Optional[DeepseekV32IndexerPrefillMetadata] = None

108 + decode: DeepSeekV32IndexerDecodeMetadata | None = None

109 + prefill: DeepseekV32IndexerPrefillMetadata | None = None

110

111

112 # TODO (zyongye) optimize this, this is now vibe coded

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+12 -12

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import ClassVar, Optional, Union

5 +from typing import ClassVar

6

7 import torch

8

@@ -47,14 +47,14 @@ class AiterMLABackend(MLACommonBackend):

47 @dataclass

48 class AiterMLADecodeMetadata(MLACommonDecodeMetadata):

49 # The indptr of the paged kv cache, shape: [batch_size + 1]

50 - paged_kv_indptr: Optional[torch.Tensor] = None

50 + paged_kv_indptr: torch.Tensor | None = None

51 # The page indices of the paged kv cache

52 - paged_kv_indices: Optional[torch.Tensor] = None

52 + paged_kv_indices: torch.Tensor | None = None

53 # The number of entries in the last page of each request in

54 # the paged kv cache, shape: [batch_size]

55 - paged_kv_last_page_len: Optional[torch.Tensor] = None

55 + paged_kv_last_page_len: torch.Tensor | None = None

56 # The query indptr, shape : [num_decode + 1]

57 - qo_indptr: Optional[torch.Tensor] = None

57 + qo_indptr: torch.Tensor | None = None

58

59

60 class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):

@@ -116,7 +116,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

116 query_start_loc_cpu: torch.Tensor,

117 query_start_loc_device: torch.Tensor,

118 num_decode_tokens: int,

119 - dcp_tot_seq_lens_device: Optional[torch.Tensor],

119 + dcp_tot_seq_lens_device: torch.Tensor | None,

120 ) -> AiterMLADecodeMetadata:

121 page_size = self.kv_cache_spec.block_size

122 block_table_bounds = (seq_lens_device + page_size - 1) // page_size

@@ -188,12 +188,12 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

188 head_size: int,

189 scale: float,

190 num_kv_heads: int,

191 - alibi_slopes: Optional[list[float]],

192 - sliding_window: Optional[int],

191 + alibi_slopes: list[float] | None,

192 + sliding_window: int | None,

193 kv_cache_dtype: str,

194 - logits_soft_cap: Optional[float],

194 + logits_soft_cap: float | None,

195 attn_type: str,

196 - kv_sharing_target_layer_name: Optional[str],

196 + kv_sharing_target_layer_name: str | None,

197 # MLA Specific Arguments

198 **mla_args,

199 ) -> None:

@@ -242,11 +242,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

242

243 def _forward_decode(

244 self,

245 - q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

245 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

246 kv_c_and_k_pe_cache: torch.Tensor,

247 attn_metadata: AiterMLAMetadata,

248 layer: AttentionLayer,

249 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

249 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

250 assert kv_c_and_k_pe_cache.numel() > 0

251 assert attn_metadata.decode is not None

252

vllm/v1/attention/backends/mla/triton_mla.py

+6 -7

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional, Union

4

5 import torch

6

@@ -44,12 +43,12 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

43 head_size: int,

44 scale: float,

45 num_kv_heads: int,

47 - alibi_slopes: Optional[list[float]],

48 - sliding_window: Optional[int],

46 + alibi_slopes: list[float] | None,

47 + sliding_window: int | None,

48 kv_cache_dtype: str,

50 - logits_soft_cap: Optional[float],

49 + logits_soft_cap: float | None,

50 attn_type: str,

52 - kv_sharing_target_layer_name: Optional[str],

51 + kv_sharing_target_layer_name: str | None,

52 # MLA Specific Arguments

53 **mla_args,

54 ) -> None:

@@ -138,11 +137,11 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

137

138 def _forward_decode(

139 self,

141 - q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

140 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

141 kv_c_and_k_pe_cache: torch.Tensor,

142 attn_metadata: MLACommonMetadata,

143 layer: AttentionLayer,

145 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

144 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

145 assert kv_c_and_k_pe_cache.numel() > 0

146 assert attn_metadata.decode is not None

147

vllm/v1/attention/backends/pallas.py

+8 -9

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import torch

7

@@ -201,12 +200,12 @@ class PallasAttentionBackendImpl(AttentionImpl):

200 head_size: int,

201 scale: float,

202 num_kv_heads: int,

204 - alibi_slopes: Optional[list[float]],

205 - sliding_window: Optional[int],

203 + alibi_slopes: list[float] | None,

204 + sliding_window: int | None,

205 kv_cache_dtype: str,

207 - logits_soft_cap: Optional[float] = None,

206 + logits_soft_cap: float | None = None,

207 attn_type: str = AttentionType.DECODER,

209 - kv_sharing_target_layer_name: Optional[int] = None,

208 + kv_sharing_target_layer_name: int | None = None,

209 ) -> None:

210 self.num_heads = num_heads

211 self.head_size = head_size

@@ -242,9 +241,9 @@ class PallasAttentionBackendImpl(AttentionImpl):

241 value: torch.Tensor,

242 kv_cache: torch.Tensor,

243 attn_metadata: PallasMetadata,

245 - output: Optional[torch.Tensor] = None,

246 - output_scale: Optional[torch.Tensor] = None,

247 - output_block_scale: Optional[torch.Tensor] = None,

244 + output: torch.Tensor | None = None,

245 + output_scale: torch.Tensor | None = None,

246 + output_block_scale: torch.Tensor | None = None,

247 ) -> torch.Tensor:

248 """Forward pass with Pallas attention.

249

@@ -342,7 +341,7 @@ def write_to_kv_cache(

341 slot_mapping: torch.Tensor,

342 num_slices_per_kv_cache_update_block: int,

343 num_kv_update_slices: torch.Tensor,

345 - kv_cache_quantized_dtype: Optional[torch.dtype] = None,

344 + kv_cache_quantized_dtype: torch.dtype | None = None,

345 k_scale: float = 1.0,

346 v_scale: float = 1.0,

347 ) -> None:

vllm/v1/attention/backends/rocm_aiter_fa.py

+14 -15

@@ -3,7 +3,6 @@

3 """Attention layer with AiterFlashAttention."""

4

5 from dataclasses import dataclass

6 -from typing import Optional, Union

6

7 import torch

8

@@ -160,8 +159,8 @@ if current_platform.is_rocm():

159 max_seqlen_q: int,

160 max_seqlen_k: int,

161 softmax_scale: float,

163 - window_size: Optional[list[int]], # -1 means infinite context window

164 - alibi_slopes: Optional[list[float]],

162 + window_size: list[int] | None, # -1 means infinite context window

163 + alibi_slopes: list[float] | None,

164 block_table: torch.Tensor,

165 k_scale: torch.Tensor,

166 v_scale: torch.Tensor,

@@ -209,8 +208,8 @@ if current_platform.is_rocm():

208 max_seqlen_q: int,

209 max_seqlen_k: int,

210 softmax_scale: float,

212 - window_size: Optional[list[int]], # -1 means infinite context window

213 - alibi_slopes: Optional[list[float]],

211 + window_size: list[int] | None, # -1 means infinite context window

212 + alibi_slopes: list[float] | None,

213 block_table: torch.Tensor,

214 k_scale: torch.Tensor,

215 v_scale: torch.Tensor,

@@ -249,7 +248,7 @@ class AiterFlashAttentionMetadata:

248 seq_lens: torch.Tensor

249 slot_mapping: torch.Tensor

250 block_table: torch.Tensor

252 - cu_seq_lens: Optional[torch.Tensor]

251 + cu_seq_lens: torch.Tensor | None

252

253 # For cascade attention.

254 use_cascade: bool

@@ -283,7 +282,7 @@ class AiterFlashAttentionMetadataBuilder(

282 self.block_size = kv_cache_spec.block_size

283 # Sliding window size to be used with the AOT scheduler will be

284 # populated on first build() call.

286 - self.aot_sliding_window: Optional[tuple[int, int]] = None

285 + self.aot_sliding_window: tuple[int, int] | None = None

286 self.total_tokens: int = 0

287

288 def build_for_cudagraph_capture(

@@ -361,7 +360,7 @@ class AiterFlashAttentionBackend(AttentionBackend):

360 return [64, 128, 256]

361

362 @staticmethod

364 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

363 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

364 return [MultipleOf(16)]

365

366 @classmethod

@@ -412,12 +411,12 @@ class AiterFlashAttentionImpl(AttentionImpl):

411 head_size: int,

412 scale: float,

413 num_kv_heads: int,

415 - alibi_slopes: Optional[list[float]],

416 - sliding_window: Optional[int],

414 + alibi_slopes: list[float] | None,

415 + sliding_window: int | None,

416 kv_cache_dtype: str,

418 - logits_soft_cap: Optional[float] = None,

417 + logits_soft_cap: float | None = None,

418 attn_type: AttentionType = AttentionType.DECODER,

420 - kv_sharing_target_layer_name: Optional[int] = None,

419 + kv_sharing_target_layer_name: int | None = None,

420 ) -> None:

421 self.num_heads = num_heads

422 self.head_size = head_size

@@ -458,9 +457,9 @@ class AiterFlashAttentionImpl(AttentionImpl):

457 value: torch.Tensor,

458 kv_cache: torch.Tensor,

459 attn_metadata: AiterFlashAttentionMetadata,

461 - output: Optional[torch.Tensor] = None,

462 - output_scale: Optional[torch.Tensor] = None,

463 - output_block_scale: Optional[torch.Tensor] = None,

460 + output: torch.Tensor | None = None,

461 + output_scale: torch.Tensor | None = None,

462 + output_block_scale: torch.Tensor | None = None,

463 ) -> torch.Tensor:

464 """Forward pass with AiterFlashAttention.

465

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+8 -10

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Attention layer with PagedAttention and Triton prefix prefill."""

4

5 -from typing import Optional

6 -

5 import torch

6

7 from vllm import _custom_ops as ops

@@ -70,13 +68,13 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

68 head_size: int,

69 scale: float,

70 num_kv_heads: int,

73 - alibi_slopes: Optional[list[float]],

74 - sliding_window: Optional[int],

71 + alibi_slopes: list[float] | None,

72 + sliding_window: int | None,

73 kv_cache_dtype: str,

76 - logits_soft_cap: Optional[float] = None,

74 + logits_soft_cap: float | None = None,

75 attn_type: AttentionType = AttentionType.DECODER,

78 - kv_sharing_target_layer_name: Optional[int] = None,

79 - sinks: Optional[torch.Tensor] = None,

76 + kv_sharing_target_layer_name: int | None = None,

77 + sinks: torch.Tensor | None = None,

78 ) -> None:

79 super().__init__(

80 num_heads,

@@ -106,9 +104,9 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

104 value: torch.Tensor,

105 kv_cache: torch.Tensor,

106 attn_metadata: FlashAttentionMetadata,

109 - output: Optional[torch.Tensor] = None,

110 - output_scale: Optional[torch.Tensor] = None,

111 - output_block_scale: Optional[torch.Tensor] = None,

107 + output: torch.Tensor | None = None,

108 + output_scale: torch.Tensor | None = None,

109 + output_block_scale: torch.Tensor | None = None,

110 ) -> torch.Tensor:

111 """Forward pass with FlashAttention.

112

vllm/v1/attention/backends/rocm_attn.py

+14 -14

@@ -3,7 +3,7 @@

3 """Attention layer with PagedAttention and Triton prefix prefill."""

4

5 from dataclasses import dataclass

6 -from typing import ClassVar, Optional

6 +from typing import ClassVar

7

8 import torch

9

@@ -54,13 +54,13 @@ class RocmAttentionMetadata:

54 # For cascade attention.

55 use_cascade: bool

56 common_prefix_len: int

57 - cu_prefix_query_lens: Optional[torch.Tensor]

58 - prefix_kv_lens: Optional[torch.Tensor]

59 - suffix_kv_lens: Optional[torch.Tensor]

57 + cu_prefix_query_lens: torch.Tensor | None

58 + prefix_kv_lens: torch.Tensor | None

59 + suffix_kv_lens: torch.Tensor | None

60

61 # Optional aot scheduling

62 - scheduler_metadata: Optional[torch.Tensor] = None

63 - prefix_scheduler_metadata: Optional[torch.Tensor] = None

62 + scheduler_metadata: torch.Tensor | None = None

63 + prefix_scheduler_metadata: torch.Tensor | None = None

64

65

66 class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadata]):

@@ -217,13 +217,13 @@ class RocmAttentionImpl(AttentionImpl):

217 head_size: int,

218 scale: float,

219 num_kv_heads: int,

220 - alibi_slopes: Optional[list[float]],

221 - sliding_window: Optional[int],

220 + alibi_slopes: list[float] | None,

221 + sliding_window: int | None,

222 kv_cache_dtype: str,

223 - logits_soft_cap: Optional[float] = None,

223 + logits_soft_cap: float | None = None,

224 attn_type: AttentionType = AttentionType.DECODER,

225 - kv_sharing_target_layer_name: Optional[int] = None,

226 - sinks: Optional[torch.Tensor] = None,

225 + kv_sharing_target_layer_name: int | None = None,

226 + sinks: torch.Tensor | None = None,

227 ) -> None:

228 self.num_heads = num_heads

229 self.head_size = head_size

@@ -273,9 +273,9 @@ class RocmAttentionImpl(AttentionImpl):

273 value: torch.Tensor,

274 kv_cache: torch.Tensor,

275 attn_metadata: FlashAttentionMetadata,

276 - output: Optional[torch.Tensor] = None,

277 - output_scale: Optional[torch.Tensor] = None,

278 - output_block_scale: Optional[torch.Tensor] = None,

276 + output: torch.Tensor | None = None,

277 + output_scale: torch.Tensor | None = None,

278 + output_block_scale: torch.Tensor | None = None,

279 ) -> torch.Tensor:

280 """Forward pass with FlashAttention.

281

vllm/v1/attention/backends/short_conv_attn.py

+4 -5

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 -from typing import Optional

4

5 import torch

6

@@ -30,12 +29,12 @@ class ShortConvAttentionMetadata:

29

30 query_start_loc: torch.Tensor

31 state_indices_tensor: torch.Tensor

33 - has_initial_states_p: Optional[torch.Tensor]

32 + has_initial_states_p: torch.Tensor | None

33

34 # For causal_conv1d

36 - nums_dict: Optional[dict] = None

37 - batch_ptr: Optional[torch.Tensor] = None

38 - token_chunk_offset_ptr: Optional[torch.Tensor] = None

35 + nums_dict: dict | None = None

36 + batch_ptr: torch.Tensor | None = None

37 + token_chunk_offset_ptr: torch.Tensor | None = None

38

39

40 class ShortConvAttentionMetadataBuilder(

vllm/v1/attention/backends/tree_attn.py

+12 -12

@@ -4,7 +4,7 @@

4

5 import ast

6 from dataclasses import dataclass

7 -from typing import Optional, Union

7 +from typing import Optional

8

9 import torch

10

@@ -41,7 +41,7 @@ class TreeAttentionBackend(AttentionBackend):

41 return [32, 64, 96, 128, 160, 192, 224, 256]

42

43 @staticmethod

44 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

44 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

45 return [MultipleOf(16)]

46

47 @classmethod

@@ -104,7 +104,7 @@ class TreeAttentionMetadata:

104 num_prefills: int = 0

105 num_decodes: int = 0

106

107 - tree_attn_bias: Optional[torch.Tensor] = None

107 + tree_attn_bias: torch.Tensor | None = None

108

109 # Cached Prefill/decode metadata.

110 _cached_prefill_metadata: Optional["TreeAttentionMetadata"] = None

@@ -267,8 +267,8 @@ def _get_depth_counts(sorted_tree_choices: list[tuple[int, ...]]) -> list[int]:

267 def _prepare_tree_attn_bias(

268 sorted_tree_choices: list[tuple[int, ...]],

269 depth_counts: list[int],

270 - dtype: Optional[torch.dtype],

271 - device: Optional[torch.device],

270 + dtype: torch.dtype | None,

271 + device: torch.device | None,

272 ) -> torch.Tensor:

273 # +1 comes from the additional root node.

274 tree_len = len(sorted_tree_choices) + 1

@@ -310,12 +310,12 @@ class TreeAttentionImpl(AttentionImpl):

310 head_size: int,

311 scale: float,

312 num_kv_heads: int,

313 - alibi_slopes: Optional[list[float]],

314 - sliding_window: Optional[int],

313 + alibi_slopes: list[float] | None,

314 + sliding_window: int | None,

315 kv_cache_dtype: str,

316 - logits_soft_cap: Optional[float] = None,

316 + logits_soft_cap: float | None = None,

317 attn_type: AttentionType = AttentionType.DECODER,

318 - kv_sharing_target_layer_name: Optional[str] = None,

318 + kv_sharing_target_layer_name: str | None = None,

319 ) -> None:

320 self.num_heads = num_heads

321 self.head_size = head_size

@@ -354,9 +354,9 @@ class TreeAttentionImpl(AttentionImpl):

354 value: torch.Tensor,

355 kv_cache: torch.Tensor,

356 attn_metadata: TreeAttentionMetadata,

357 - output: Optional[torch.Tensor] = None,

358 - output_scale: Optional[torch.Tensor] = None,

359 - output_block_scale: Optional[torch.Tensor] = None,

357 + output: torch.Tensor | None = None,

358 + output_scale: torch.Tensor | None = None,

359 + output_block_scale: torch.Tensor | None = None,

360 ) -> torch.Tensor:

361 """Forward pass with TreeAttention.

362

vllm/v1/attention/backends/triton_attn.py

+15 -15

@@ -3,7 +3,7 @@

3 """High-Performance Triton-only Attention layer."""

4

5 from dataclasses import dataclass

6 -from typing import ClassVar, Optional, Union

6 +from typing import ClassVar

7

8 import torch

9

@@ -61,13 +61,13 @@ class TritonAttentionMetadata:

61 # For cascade attention.

62 use_cascade: bool

63 common_prefix_len: int

64 - cu_prefix_query_lens: Optional[torch.Tensor]

65 - prefix_kv_lens: Optional[torch.Tensor]

66 - suffix_kv_lens: Optional[torch.Tensor]

64 + cu_prefix_query_lens: torch.Tensor | None

65 + prefix_kv_lens: torch.Tensor | None

66 + suffix_kv_lens: torch.Tensor | None

67

68 # Optional aot scheduling

69 - scheduler_metadata: Optional[torch.Tensor] = None

70 - prefix_scheduler_metadata: Optional[torch.Tensor] = None

69 + scheduler_metadata: torch.Tensor | None = None

70 + prefix_scheduler_metadata: torch.Tensor | None = None

71

72

73 class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):

@@ -159,7 +159,7 @@ class TritonAttentionBackend(AttentionBackend):

159 return [torch.float16, torch.bfloat16, torch.float32]

160

161 @staticmethod

162 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

162 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

163 return [MultipleOf(16)]

164

165 @classmethod

@@ -216,13 +216,13 @@ class TritonAttentionImpl(AttentionImpl):

216 head_size: int,

217 scale: float,

218 num_kv_heads: int,

219 - alibi_slopes: Optional[list[float]],

220 - sliding_window: Optional[int],

219 + alibi_slopes: list[float] | None,

220 + sliding_window: int | None,

221 kv_cache_dtype: str,

222 - logits_soft_cap: Optional[float] = None,

222 + logits_soft_cap: float | None = None,

223 attn_type: AttentionType = AttentionType.DECODER,

224 - kv_sharing_target_layer_name: Optional[int] = None,

225 - sinks: Optional[torch.Tensor] = None,

224 + kv_sharing_target_layer_name: int | None = None,

225 + sinks: torch.Tensor | None = None,

226 ) -> None:

227 self.num_heads = num_heads

228 self.head_size = head_size

@@ -272,9 +272,9 @@ class TritonAttentionImpl(AttentionImpl):

272 value: torch.Tensor,

273 kv_cache: torch.Tensor,

274 attn_metadata: TritonAttentionMetadata,

275 - output: Optional[torch.Tensor] = None,

276 - output_scale: Optional[torch.Tensor] = None,

277 - output_block_scale: Optional[torch.Tensor] = None,

275 + output: torch.Tensor | None = None,

276 + output_scale: torch.Tensor | None = None,

277 + output_block_scale: torch.Tensor | None = None,

278 ) -> torch.Tensor:

279 """Forward pass with Paged Attention impl. in Triton.

280

vllm/v1/attention/backends/utils.py

+10 -12

@@ -11,10 +11,8 @@ from typing import (

11 ClassVar,

12 Generic,

13 Literal,

14 - Optional,

14 Protocol,

15 TypeVar,

17 - Union,

16 get_args,

17 )

18

@@ -42,7 +40,7 @@ from vllm.v1.worker.ubatch_utils import UBatchSlice

40

41 logger = init_logger(__name__)

42 KVCacheLayoutType = Literal["NHD", "HND"]

45 -_KV_CACHE_LAYOUT_OVERRIDE: Union[KVCacheLayoutType, None] = None

43 +_KV_CACHE_LAYOUT_OVERRIDE: KVCacheLayoutType | None = None

44

45 PAD_SLOT_ID = -1

46

@@ -87,13 +85,13 @@ class CommonAttentionMetadata:

85 causal: bool = True

86

87 # Needed by FastPrefillAttentionBuilder

90 - logits_indices_padded: Optional[torch.Tensor] = None

91 - num_logits_indices: Optional[int] = None

88 + logits_indices_padded: torch.Tensor | None = None

89 + num_logits_indices: int | None = None

90

91 # Needed by CrossAttentionBuilder

94 - encoder_seq_lens: Optional[np.ndarray] = None

92 + encoder_seq_lens: np.ndarray | None = None

93

96 - dcp_local_seq_lens: Optional[torch.Tensor] = None

94 + dcp_local_seq_lens: torch.Tensor | None = None

95 """Sequence lengths of the local rank in decode context parallelism world"""

96

97

@@ -250,7 +248,7 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

248 # Does this backend/builder reorder the batch?

249 # If not, set this to None. Otherwise set it to the query

250 # length that will be pulled into the front of the batch.

253 - reorder_batch_threshold: Optional[int] = None

251 + reorder_batch_threshold: int | None = None

252

253 @abstractmethod

254 def __init__(

@@ -395,12 +393,12 @@ class PerLayerParameters:

393 """

394

395 window_left: int

398 - logits_soft_cap: Optional[float]

396 + logits_soft_cap: float | None

397 sm_scale: float

398 has_sinks: bool = False

399 # has same params for all layers

402 - has_same_window_lefts: Optional[bool] = field(default=None, compare=False)

403 - has_same_all_params: Optional[bool] = field(default=None, compare=False)

400 + has_same_window_lefts: bool | None = field(default=None, compare=False)

401 + has_same_all_params: bool | None = field(default=None, compare=False)

402

403

404 def get_per_layer_parameters(

@@ -873,7 +871,7 @@ def reshape_attn_output_for_spec_decode(attn_output: torch.Tensor) -> torch.Tens

871

872

873 KV_SHARING_FAST_PREFILL_METADATA_FIELDS = [

876 - ("logits_indices_padded", Optional[torch.Tensor], None),

874 + ("logits_indices_padded", torch.Tensor | None, None),

875 ("num_logits_indices", int, 0),

876 ]

877

vllm/v1/attention/backends/xformers.py

+9 -9

@@ -3,7 +3,7 @@

3 """Attention layer with XFormersAttention."""

4

5 from dataclasses import dataclass

6 -from typing import Optional, Union

6 +from typing import Optional

7

8 import torch

9

@@ -82,7 +82,7 @@ class XFormersAttentionBackend(AttentionBackend):

82 ]

83

84 @staticmethod

85 - def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:

85 + def get_supported_kernel_block_size() -> list[int | MultipleOf]:

86 return [MultipleOf(16)]

87

88 @classmethod

@@ -280,12 +280,12 @@ class XFormersAttentionImpl(AttentionImpl):

280 head_size: int,

281 scale: float,

282 num_kv_heads: int,

283 - alibi_slopes: Optional[list[float]],

284 - sliding_window: Optional[int],

283 + alibi_slopes: list[float] | None,

284 + sliding_window: int | None,

285 kv_cache_dtype: str,

286 - logits_soft_cap: Optional[float] = None,

286 + logits_soft_cap: float | None = None,

287 attn_type: AttentionType = AttentionType.DECODER,

288 - kv_sharing_target_layer_name: Optional[str] = None,

288 + kv_sharing_target_layer_name: str | None = None,

289 ) -> None:

290 if kv_sharing_target_layer_name is not None:

291 raise NotImplementedError("KV sharing is not supported in V0.")

@@ -328,9 +328,9 @@ class XFormersAttentionImpl(AttentionImpl):

328 value: torch.Tensor,

329 kv_cache: torch.Tensor,

330 attn_metadata: XFormersAttentionMetadata,

331 - output: Optional[torch.Tensor] = None,

332 - output_scale: Optional[torch.Tensor] = None,

333 - output_block_scale: Optional[torch.Tensor] = None,

331 + output: torch.Tensor | None = None,

332 + output_scale: torch.Tensor | None = None,

333 + output_block_scale: torch.Tensor | None = None,

334 ) -> torch.Tensor:

335 """Forward pass with XFormers.

336

vllm/v1/core/block_pool.py

+7 -7

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable

4 -from typing import Any, Optional, Union

4 +from typing import Any

5

6 from vllm.distributed.kv_events import (

7 MEDIUM_GPU,

@@ -51,10 +51,10 @@ class BlockHashToBlockMap:

51

52 def __init__(self):

53 self._cache: dict[

54 - BlockHashWithGroupId, Union[KVCacheBlock, dict[int, KVCacheBlock]]

54 + BlockHashWithGroupId, KVCacheBlock | dict[int, KVCacheBlock]

55 ] = {}

56

57 - def get_one_block(self, key: BlockHashWithGroupId) -> Optional[KVCacheBlock]:

57 + def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None:

58 """

59 Gets any block with the given block hash key.

60 """

@@ -85,7 +85,7 @@ class BlockHashToBlockMap:

85 else:

86 self._unexpected_blocks_type(blocks)

87

88 - def pop(self, key: BlockHashWithGroupId, block_id: int) -> Optional[KVCacheBlock]:

88 + def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None:

89 """

90 Checks if block_hash exists and pop block_id from the cache

91 """

@@ -168,7 +168,7 @@ class BlockPool:

168

169 def get_cached_block(

170 self, block_hash: BlockHash, kv_cache_group_ids: list[int]

171 - ) -> Optional[list[KVCacheBlock]]:

171 + ) -> list[KVCacheBlock] | None:

172 """Get the cached block by the block hash for each group in

173 `kv_cache_group_ids`, or None if cache miss for any group.

174 If there are duplicated blocks, we return the first block in the cache.

@@ -225,7 +225,7 @@ class BlockPool:

225 assert len(request.block_hashes) >= num_full_blocks

226 new_block_hashes = request.block_hashes[num_cached_blocks:]

227

228 - new_hashes: Optional[list[ExternalBlockHash]] = (

228 + new_hashes: list[ExternalBlockHash] | None = (

229 [] if self.enable_kv_cache_events else None

230 )

231 for i, blk in enumerate(new_full_blocks):

@@ -243,7 +243,7 @@ class BlockPool:

243

244 if self.enable_kv_cache_events:

245 if num_cached_blocks == 0:

246 - parent_block_hash: Optional[ExternalBlockHash] = None

246 + parent_block_hash: ExternalBlockHash | None = None

247 else:

248 parent_block = blocks[num_cached_blocks - 1]

249 assert parent_block.block_hash is not None

vllm/v1/core/kv_cache_coordinator.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from abc import ABC, abstractmethod

4 -from typing import Optional

4

5 from vllm.v1.core.block_pool import BlockPool

6 from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

@@ -320,8 +319,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

319 one of them is full attention. Then, split the kv cache groups into full

320 attention groups and other groups.

321 """

323 - full_attention_spec: Optional[FullAttentionSpec] = None

324 - other_spec: Optional[KVCacheSpec] = None

322 + full_attention_spec: FullAttentionSpec | None = None

323 + other_spec: KVCacheSpec | None = None

324 self.full_attention_group_ids: list[int] = []

325 self.other_group_ids: list[int] = []

326 for i, g in enumerate(self.kv_cache_config.kv_cache_groups):

vllm/v1/core/kv_cache_manager.py

+7 -7

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import Literal, Optional, overload

5 +from typing import Literal, overload

6

7 from vllm.distributed.kv_events import KVCacheEvent

8 from vllm.logger import init_logger

@@ -49,12 +49,12 @@ class KVCacheBlocks:

49 def get_block_ids(

50 self,

51 allow_none: Literal[True] = True,

52 - ) -> Optional[tuple[list[int], ...]]: ...

52 + ) -> tuple[list[int], ...] | None: ...

53

54 def get_block_ids(

55 self,

56 allow_none: bool = False,

57 - ) -> Optional[tuple[list[int], ...]]:

57 + ) -> tuple[list[int], ...] | None:

58 """

59 Converts the KVCacheBlocks instance to block_ids.

60

@@ -97,7 +97,7 @@ class KVCacheManager:

97 # FIXME: make prefix cache stats conditional on log_stats

98 self.prefix_cache_stats = PrefixCacheStats() if log_stats else None

99

100 - self.block_size: Optional[int] = None

100 + self.block_size: int | None = None

101 if self.enable_caching:

102 assert (

103 len(

@@ -140,7 +140,7 @@ class KVCacheManager:

140 """

141 return self.block_pool.get_usage()

142

143 - def make_prefix_cache_stats(self) -> Optional[PrefixCacheStats]:

143 + def make_prefix_cache_stats(self) -> PrefixCacheStats | None:

144 """Get (and reset) the prefix cache stats.

145

146 Returns:

@@ -205,11 +205,11 @@ class KVCacheManager:

205 request: Request,

206 num_new_tokens: int,

207 num_new_computed_tokens: int = 0,

208 - new_computed_blocks: Optional[KVCacheBlocks] = None,

208 + new_computed_blocks: KVCacheBlocks | None = None,

209 num_lookahead_tokens: int = 0,

210 delay_cache_blocks: bool = False,

211 num_encoder_tokens: int = 0,

212 - ) -> Optional[KVCacheBlocks]:

212 + ) -> KVCacheBlocks | None:

213 """Add slots for a request with new tokens to append.

214

215 Args:

vllm/v1/core/kv_cache_utils.py

+10 -10

@@ -5,9 +5,9 @@

5 import copy

6 import os

7 from collections import defaultdict

8 -from collections.abc import Iterable, Sequence

8 +from collections.abc import Callable, Iterable, Sequence

9 from dataclasses import dataclass

10 -from typing import Any, Callable, NewType, Optional, Union

10 +from typing import Any, NewType, TypeAlias

11

12 from vllm import envs

13 from vllm.config import VllmConfig

@@ -38,7 +38,7 @@ BlockHashWithGroupId = NewType("BlockHashWithGroupId", bytes)

38 # ExternalBlockHash is used for reproducible prefix-cache block hashing.

39 # It's a union of ``bytes`` and ``int`` to keep backward compatibility

40 # after we default block hashing to use sha256 bytes.

41 -ExternalBlockHash = Union[bytes, int]

41 +ExternalBlockHash: TypeAlias = bytes | int

42

43

44 def make_block_hash_with_group_id(

@@ -110,18 +110,18 @@ class KVCacheBlock:

110 ref_cnt: int = 0

111 # The hash key (block hash + group id) of the block, only available

112 # when the block is full and cached.

113 - _block_hash: Optional[BlockHashWithGroupId] = None

113 + _block_hash: BlockHashWithGroupId | None = None

114

115 # Used to construct a doubly linked list for free blocks.

116 # These two attributes should only be manipulated by FreeKVCacheBlockQueue.

117 - prev_free_block: Optional["KVCacheBlock"] = None

118 - next_free_block: Optional["KVCacheBlock"] = None

117 + prev_free_block: "KVCacheBlock | None" = None

118 + next_free_block: "KVCacheBlock | None" = None

119

120 # Whether the block is a null block that should never be cached.

121 is_null: bool = False

122

123 @property

124 - def block_hash(self) -> Optional[BlockHashWithGroupId]:

124 + def block_hash(self) -> BlockHashWithGroupId | None:

125 return self._block_hash

126

127 @block_hash.setter

@@ -461,7 +461,7 @@ def _gen_lora_extra_hash_keys(request: Request) -> list[int]:

461

462 def generate_block_hash_extra_keys(

463 request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int

464 -) -> tuple[Optional[tuple[Any, ...]], int]:

464 +) -> tuple[tuple[Any, ...] | None, int]:

465 """Generate extra keys for the block hash. The extra keys can come from

466 the multi-modal inputs and request specific metadata (e.g., LoRA ID).

467

@@ -493,9 +493,9 @@ def generate_block_hash_extra_keys(

493

494 def hash_block_tokens(

495 hash_function: Callable[[Any], bytes],

496 - parent_block_hash: Optional[BlockHash],

496 + parent_block_hash: BlockHash | None,

497 curr_block_token_ids: Sequence[int],

498 - extra_keys: Optional[tuple[Any, ...]] = None,

498 + extra_keys: tuple[Any, ...] | None = None,

499 ) -> BlockHash:

500 """Computes a hash value corresponding to the contents of a block and

501 the contents of the preceding block(s). The hash value is used for

vllm/v1/core/sched/async_scheduler.py

-2

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 from vllm.logger import init_logger

5 from vllm.v1.core.sched.output import SchedulerOutput

6 from vllm.v1.core.sched.scheduler import Scheduler

vllm/v1/core/sched/interface.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from abc import ABC, abstractmethod

4 from collections.abc import Iterable

5 -from typing import TYPE_CHECKING, Optional, Union

5 +from typing import TYPE_CHECKING, Optional

6

7 if TYPE_CHECKING:

8 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1

@@ -80,7 +80,7 @@ class SchedulerInterface(ABC):

80 @abstractmethod

81 def finish_requests(

82 self,

83 - request_ids: Union[str, Iterable[str]],

83 + request_ids: str | Iterable[str],

84 finished_status: "RequestStatus",

85 ) -> None:

86 """Finish the requests in the scheduler's internal queue. If the request

vllm/v1/core/sched/output.py

+11 -6

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 from dataclasses import dataclass

5 from typing import TYPE_CHECKING

6

@@ -19,6 +17,13 @@ if TYPE_CHECKING:

17 from vllm.pooling_params import PoolingParams

18 from vllm.sampling_params import SamplingParams

19 from vllm.v1.request import Request

20 +else:

21 + KVConnectorMetadata = object

22 + LoRARequest = object

23 + MultiModalFeatureSpec = object

24 + PoolingParams = object

25 + SamplingParams = object

26 + Request = object

27

28

29 @bc_linter_include

@@ -32,14 +37,14 @@ class NewRequestData:

37 block_ids: tuple[list[int], ...]

38 num_computed_tokens: int

39 lora_request: LoRARequest | None

35 - prompt_embeds: torch.Tensor | None = None

40 + prompt_embeds: "torch.Tensor | None" = None

41

42 @classmethod

43 def from_request(

44 cls,

45 request: Request,

46 block_ids: tuple[list[int], ...],

42 - ) -> NewRequestData:

47 + ) -> "NewRequestData":

48 return cls(

49 req_id=request.request_id,

50 prompt_token_ids=request.prompt_token_ids,

@@ -110,7 +115,7 @@ class CachedRequestData:

115 return len(self.req_ids)

116

117 @classmethod

113 - def make_empty(cls) -> CachedRequestData:

118 + def make_empty(cls) -> "CachedRequestData":

119 return cls(

120 req_ids=[],

121 resumed_from_preemption=[],

@@ -164,7 +169,7 @@ class SchedulerOutput:

169 # for filling the next token bitmask

170 structured_output_request_ids: dict[str, int]

171 # the bitmask for the whole batch

167 - grammar_bitmask: npt.NDArray[np.int32] | None

172 + grammar_bitmask: "npt.NDArray[np.int32] | None"

173

174 # KV Cache Connector metadata.

175 kv_connector_metadata: KVConnectorMetadata | None = None

vllm/v1/core/sched/request_queue.py

+1 -3

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import heapq

5 from abc import ABC, abstractmethod

6 from collections import deque

@@ -43,7 +41,7 @@ class RequestQueue(ABC):

41 pass

42

43 @abstractmethod

46 - def prepend_requests(self, requests: RequestQueue) -> None:

44 + def prepend_requests(self, requests: "RequestQueue") -> None:

45 """Prepend all requests from another queue to the front of this

46 queue."""

47 pass

vllm/v1/core/sched/scheduler.py

+2 -4

@@ -1,13 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import itertools

5 import time

6 from collections import defaultdict

7 from collections.abc import Iterable

10 -from typing import Any, Union

8 +from typing import Any

9

10 from vllm.config import VllmConfig

11 from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch

@@ -1168,7 +1166,7 @@ class Scheduler(SchedulerInterface):

1166

1167 def finish_requests(

1168 self,

1171 - request_ids: Union[str, Iterable[str]],

1169 + request_ids: str | Iterable[str],

1170 finished_status: RequestStatus,

1171 ) -> None:

1172 """Handles the finish signal from outside the scheduler.

vllm/v1/core/sched/utils.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import contextlib

4 -from typing import Optional

4

5 import torch

6

@@ -41,7 +40,7 @@ def remove_all(lst: list, items_to_remove: set) -> list:

40

41

42 def check_stop(

44 - request: Request, max_model_len: int, pooler_output: Optional[torch.Tensor] = None

43 + request: Request, max_model_len: int, pooler_output: torch.Tensor | None = None

44 ) -> bool:

45 if (

46 request.num_tokens >= max_model_len

vllm/v1/cudagraph_dispatcher.py

+1 -2

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Optional

3

4 from vllm.config import CUDAGraphMode, VllmConfig

5 from vllm.forward_context import BatchDescriptor

@@ -102,7 +101,7 @@ class CudagraphDispatcher:

101

102 def dispatch(

103 self, batch_descriptor: BatchDescriptor, use_cascade_attn: bool = False

105 - ) -> tuple[CUDAGraphMode, Optional[BatchDescriptor]]:

104 + ) -> tuple[CUDAGraphMode, BatchDescriptor | None]:

105 """

106 Given conditions(e.g.,batch descriptor and if using cascade attention),

107 dispatch to a cudagraph runtime mode and the valid batch descriptor.

vllm/v1/engine/__init__.py

+27 -27

@@ -4,7 +4,7 @@

4 import enum

5 import time

6 from collections.abc import Mapping

7 -from typing import Any, Optional, Union

7 +from typing import Any

8

9 import msgspec

10 import torch

@@ -48,16 +48,16 @@ class EngineCoreRequest(

48 gc=False,

49 ): # type: ignore[call-arg]

50 request_id: str

51 - prompt_token_ids: Optional[list[int]]

52 - mm_features: Optional[list[MultiModalFeatureSpec]]

53 - sampling_params: Optional[SamplingParams]

54 - pooling_params: Optional[PoolingParams]

55 - eos_token_id: Optional[int]

51 + prompt_token_ids: list[int] | None

52 + mm_features: list[MultiModalFeatureSpec] | None

53 + sampling_params: SamplingParams | None

54 + pooling_params: PoolingParams | None

55 + eos_token_id: int | None

56 arrival_time: float

57 - lora_request: Optional[LoRARequest]

58 - cache_salt: Optional[str]

59 - data_parallel_rank: Optional[int]

60 - prompt_embeds: Optional[torch.Tensor] = None

57 + lora_request: LoRARequest | None

58 + cache_salt: str | None

59 + data_parallel_rank: int | None

60 + prompt_embeds: torch.Tensor | None = None

61

62 # Index of the client, used to ensure outputs are sent back to the same

63 # client for this request when scaling out the front-end.

@@ -69,7 +69,7 @@ class EngineCoreRequest(

69 current_wave: int = 0

70 priority: int = 0

71

72 - trace_headers: Optional[Mapping[str, str]] = None

72 + trace_headers: Mapping[str, str] | None = None

73

74

75 class EngineCoreEventType(enum.IntEnum):

@@ -93,7 +93,7 @@ class EngineCoreEvent(msgspec.Struct):

93

94 @classmethod

95 def new_event(

96 - cls, event_type: EngineCoreEventType, timestamp: Optional[float] = None

96 + cls, event_type: EngineCoreEventType, timestamp: float | None = None

97 ) -> "EngineCoreEvent":

98 timestamp = time.monotonic() if timestamp is None else timestamp

99 return cls(event_type, timestamp)

@@ -108,17 +108,17 @@ class EngineCoreOutput(

108 request_id: str

109 new_token_ids: list[int]

110

111 - new_logprobs: Optional[LogprobsLists] = None

112 - new_prompt_logprobs_tensors: Optional[LogprobsTensors] = None

111 + new_logprobs: LogprobsLists | None = None

112 + new_prompt_logprobs_tensors: LogprobsTensors | None = None

113

114 - pooling_output: Optional[torch.Tensor] = None

114 + pooling_output: torch.Tensor | None = None

115

116 - finish_reason: Optional[FinishReason] = None

117 - stop_reason: Union[int, str, None] = None

118 - events: Optional[list[EngineCoreEvent]] = None

119 - kv_transfer_params: Optional[dict[str, Any]] = None

116 + finish_reason: FinishReason | None = None

117 + stop_reason: int | str | None = None

118 + events: list[EngineCoreEvent] | None = None

119 + kv_transfer_params: dict[str, Any] | None = None

120

121 - trace_headers: Optional[Mapping[str, str]] = None

121 + trace_headers: Mapping[str, str] | None = None

122 # The number of tokens with prefix cache hits.

123 num_cached_tokens: int = 0

124

@@ -142,8 +142,8 @@ class UtilityOutput(

142 call_id: int

143

144 # Non-None implies the call failed, result should be None.

145 - failure_message: Optional[str] = None

146 - result: Optional[UtilityResult] = None

145 + failure_message: str | None = None

146 + result: UtilityResult | None = None

147

148

149 class EngineCoreOutputs(

@@ -159,18 +159,18 @@ class EngineCoreOutputs(

159

160 # [num_reqs]

161 outputs: list[EngineCoreOutput] = []

162 - scheduler_stats: Optional[SchedulerStats] = None

162 + scheduler_stats: SchedulerStats | None = None

163 timestamp: float = 0.0

164

165 - utility_output: Optional[UtilityOutput] = None

166 - finished_requests: Optional[set[str]] = None

165 + utility_output: UtilityOutput | None = None

166 + finished_requests: set[str] | None = None

167

168 # In DP case, used to signal that the current wave of requests

169 # has finished and the engines are paused.

170 - wave_complete: Optional[int] = None

170 + wave_complete: int | None = None

171 # In DP case, used to signal that a request was received for an

172 # "old" wave, so the next wave needs to be started in other engines.

173 - start_wave: Optional[int] = None

173 + start_wave: int | None = None

174

175 def __post_init__(self):

176 if self.timestamp == 0.0:

vllm/v1/engine/async_llm.py

+35 -35

@@ -6,7 +6,7 @@ import socket

6 import time

7 from collections.abc import AsyncGenerator, Iterable, Mapping

8 from copy import copy

9 -from typing import Any, Optional, Union

9 +from typing import Any

10

11 import numpy as np

12 import torch

@@ -56,8 +56,8 @@ class AsyncLLM(EngineClient):

56 use_cached_outputs: bool = False,

57 log_requests: bool = True,

58 start_engine_loop: bool = True,

59 - stat_loggers: Optional[list[StatLoggerFactory]] = None,

60 - client_addresses: Optional[dict[str, str]] = None,

59 + stat_loggers: list[StatLoggerFactory] | None = None,

60 + client_addresses: dict[str, str] | None = None,

61 client_count: int = 1,

62 client_index: int = 0,

63 ) -> None:

@@ -136,7 +136,7 @@ class AsyncLLM(EngineClient):

136 )

137

138 # Loggers.

139 - self.logger_manager: Optional[StatLoggerManager] = None

139 + self.logger_manager: StatLoggerManager | None = None

140 if self.log_stats:

141 self.logger_manager = StatLoggerManager(

142 vllm_config=vllm_config,

@@ -147,7 +147,7 @@ class AsyncLLM(EngineClient):

147 )

148 self.logger_manager.log_engine_initialized()

149

150 - self.output_handler: Optional[asyncio.Task] = None

150 + self.output_handler: asyncio.Task | None = None

151 try:

152 # Start output handler eagerly if we are in the asyncio eventloop.

153 asyncio.get_running_loop()

@@ -185,10 +185,10 @@ class AsyncLLM(EngineClient):

185 vllm_config: VllmConfig,

186 start_engine_loop: bool = True,

187 usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,

188 - stat_loggers: Optional[list[StatLoggerFactory]] = None,

188 + stat_loggers: list[StatLoggerFactory] | None = None,

189 enable_log_requests: bool = False,

190 disable_log_stats: bool = False,

191 - client_addresses: Optional[dict[str, str]] = None,

191 + client_addresses: dict[str, str] | None = None,

192 client_count: int = 1,

193 client_index: int = 0,

194 disable_log_requests: bool = True, # Deprecated, will be removed

@@ -221,7 +221,7 @@ class AsyncLLM(EngineClient):

221 engine_args: AsyncEngineArgs,

222 start_engine_loop: bool = True,

223 usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,

224 - stat_loggers: Optional[list[StatLoggerFactory]] = None,

224 + stat_loggers: list[StatLoggerFactory] | None = None,

225 ) -> "AsyncLLM":

226 """Create an AsyncLLM from the EngineArgs."""

227

@@ -259,15 +259,15 @@ class AsyncLLM(EngineClient):

259 async def add_request(

260 self,

261 request_id: str,

262 - prompt: Union[EngineCoreRequest, PromptType],

263 - params: Union[SamplingParams, PoolingParams],

264 - arrival_time: Optional[float] = None,

265 - lora_request: Optional[LoRARequest] = None,

266 - tokenization_kwargs: Optional[dict[str, Any]] = None,

267 - trace_headers: Optional[Mapping[str, str]] = None,

262 + prompt: EngineCoreRequest | PromptType,

263 + params: SamplingParams | PoolingParams,

264 + arrival_time: float | None = None,

265 + lora_request: LoRARequest | None = None,

266 + tokenization_kwargs: dict[str, Any] | None = None,

267 + trace_headers: Mapping[str, str] | None = None,

268 priority: int = 0,

269 - data_parallel_rank: Optional[int] = None,

270 - prompt_text: Optional[str] = None,

269 + data_parallel_rank: int | None = None,

270 + prompt_text: str | None = None,

271 ) -> RequestOutputCollector:

272 """Add new request to the AsyncLLM."""

273

@@ -325,8 +325,8 @@ class AsyncLLM(EngineClient):

325 async def _add_request(

326 self,

327 request: EngineCoreRequest,

328 - prompt: Optional[str],

329 - parent_req: Optional[ParentRequest],

328 + prompt: str | None,

329 + parent_req: ParentRequest | None,

330 index: int,

331 queue: RequestOutputCollector,

332 ):

@@ -346,16 +346,16 @@ class AsyncLLM(EngineClient):

346 # re-multiplexed in the API server anyhow.

347 async def generate(

348 self,

349 - prompt: Union[EngineCoreRequest, PromptType],

349 + prompt: EngineCoreRequest | PromptType,

350 sampling_params: SamplingParams,

351 request_id: str,

352 *,

353 - prompt_text: Optional[str] = None,

354 - lora_request: Optional[LoRARequest] = None,

355 - tokenization_kwargs: Optional[dict[str, Any]] = None,

356 - trace_headers: Optional[Mapping[str, str]] = None,

353 + prompt_text: str | None = None,

354 + lora_request: LoRARequest | None = None,

355 + tokenization_kwargs: dict[str, Any] | None = None,

356 + trace_headers: Mapping[str, str] | None = None,

357 priority: int = 0,

358 - data_parallel_rank: Optional[int] = None,

358 + data_parallel_rank: int | None = None,

359 ) -> AsyncGenerator[RequestOutput, None]:

360 """

361 Main function called by the API server to kick off a request

@@ -520,7 +520,7 @@ class AsyncLLM(EngineClient):

520

521 self.output_handler = asyncio.create_task(output_handler())

522

523 - async def abort(self, request_id: Union[str, Iterable[str]]) -> None:

523 + async def abort(self, request_id: str | Iterable[str]) -> None:

524 """Abort RequestId in OutputProcessor and EngineCore."""

525

526 request_ids = (

@@ -537,11 +537,11 @@ class AsyncLLM(EngineClient):

537 prompt: PromptType,

538 pooling_params: PoolingParams,

539 request_id: str,

540 - lora_request: Optional[LoRARequest] = None,

541 - trace_headers: Optional[Mapping[str, str]] = None,

540 + lora_request: LoRARequest | None = None,

541 + trace_headers: Mapping[str, str] | None = None,

542 priority: int = 0,

543 - truncate_prompt_tokens: Optional[int] = None,

544 - tokenization_kwargs: Optional[dict[str, Any]] = None,

543 + truncate_prompt_tokens: int | None = None,

544 + tokenization_kwargs: dict[str, Any] | None = None,

545 ) -> AsyncGenerator[PoolingRequestOutput, None]:

546 """

547 Main function called by the API server to kick off a request

@@ -622,11 +622,11 @@ class AsyncLLM(EngineClient):

622 raise EngineGenerateError() from e

623

624 @property

625 - def tokenizer(self) -> Optional[AnyTokenizer]:

625 + def tokenizer(self) -> AnyTokenizer | None:

626 return self.processor.tokenizer

627

628 @tokenizer.setter

629 - def tokenizer(self, tokenizer: Optional[AnyTokenizer]) -> None:

629 + def tokenizer(self, tokenizer: AnyTokenizer | None) -> None:

630 self.processor.tokenizer = tokenizer

631

632 async def get_tokenizer(self) -> AnyTokenizer:

@@ -665,7 +665,7 @@ class AsyncLLM(EngineClient):

665 self.processor.clear_mm_cache()

666 await self.engine_core.reset_mm_cache_async()

667

668 - async def reset_prefix_cache(self, device: Optional[Device] = None) -> None:

668 + async def reset_prefix_cache(self, device: Device | None = None) -> None:

669 if device == Device.CPU:

670 raise ValueError("Not supported on CPU.")

671 await self.engine_core.reset_prefix_cache_async()

@@ -674,7 +674,7 @@ class AsyncLLM(EngineClient):

674 await self.reset_prefix_cache()

675 await self.engine_core.sleep_async(level)

676

677 - async def wake_up(self, tags: Optional[list[str]] = None) -> None:

677 + async def wake_up(self, tags: list[str] | None = None) -> None:

678 await self.engine_core.wake_up_async(tags)

679

680 async def is_sleeping(self) -> bool:

@@ -699,9 +699,9 @@ class AsyncLLM(EngineClient):

699 async def collective_rpc(

700 self,

701 method: str,

702 - timeout: Optional[float] = None,

702 + timeout: float | None = None,

703 args: tuple = (),

704 - kwargs: Optional[dict] = None,

704 + kwargs: dict | None = None,

705 ):

706 """

707 Perform a collective RPC call to the given path.

vllm/v1/engine/coordinator.py

+2 -3

@@ -4,7 +4,6 @@ import copy

4 import multiprocessing

5 import time

6 import weakref

7 -from typing import Optional

7

8 import msgspec.msgpack

9 import zmq

@@ -155,7 +154,7 @@ class DPCoordinatorProc:

154 stats_changed = False

155 last_stats_step = -1

156 last_stats_wave = -1

158 - last_step_counts: Optional[list[list[int]]] = None

157 + last_step_counts: list[list[int]] | None = None

158

159 with (

160 make_zmq_socket(

@@ -360,7 +359,7 @@ class DPCoordinatorProc:

359

360 @staticmethod

361 def _send_start_wave(

363 - socket: zmq.Socket, wave: int, exclude_engine_index: Optional[int]

362 + socket: zmq.Socket, wave: int, exclude_engine_index: int | None

363 ):

364 """Broadcast the START_DP_WAVE message to all the engines.

365 It includes the current wave number and index of engine which

vllm/v1/engine/core.py

+24 -24

@@ -7,12 +7,12 @@ import signal

7 import threading

8 import time

9 from collections import deque

10 -from collections.abc import Generator

10 +from collections.abc import Callable, Generator

11 from concurrent.futures import Future

12 from contextlib import ExitStack, contextmanager

13 from inspect import isclass, signature

14 from logging import DEBUG

15 -from typing import Any, Callable, Optional, TypeVar, Union

15 +from typing import Any, TypeVar

16

17 import msgspec

18 import zmq

@@ -83,7 +83,7 @@ class EngineCore:

83 vllm_config: VllmConfig,

84 executor_class: type[Executor],

85 log_stats: bool,

86 - executor_fail_callback: Optional[Callable] = None,

86 + executor_fail_callback: Callable | None = None,

87 ):

88 # plugins need to be loaded at the engine/scheduler level too

89 from vllm.plugins import load_general_plugins

@@ -171,14 +171,14 @@ class EngineCore:

171 # schedule and execute batches, and is required by pipeline parallelism

172 # to eliminate pipeline bubbles.

173 self.batch_queue_size = self.model_executor.max_concurrent_batches

174 - self.batch_queue: Optional[

175 - deque[tuple[Future[ModelRunnerOutput], SchedulerOutput]]

176 - ] = None

174 + self.batch_queue: (

175 + deque[tuple[Future[ModelRunnerOutput], SchedulerOutput]] | None

176 + ) = None

177 if self.batch_queue_size > 1:

178 logger.info("Batch queue is enabled with size %d", self.batch_queue_size)

179 self.batch_queue = deque(maxlen=self.batch_queue_size)

180

181 - self.request_block_hasher: Optional[Callable[[Request], list[BlockHash]]] = None

181 + self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None

182 if (

183 self.vllm_config.cache_config.enable_prefix_caching

184 or self.scheduler.get_kv_connector() is not None

@@ -337,7 +337,7 @@ class EngineCore:

337

338 def step_with_batch_queue(

339 self,

340 - ) -> tuple[Optional[dict[int, EngineCoreOutputs]], bool]:

340 + ) -> tuple[dict[int, EngineCoreOutputs] | None, bool]:

341 """Schedule and execute batches with the batch queue.

342 Note that if nothing to output in this step, None is returned.

343

@@ -424,7 +424,7 @@ class EngineCore:

424 def sleep(self, level: int = 1):

425 self.model_executor.sleep(level)

426

427 - def wake_up(self, tags: Optional[list[str]] = None):

427 + def wake_up(self, tags: list[str] | None = None):

428 self.model_executor.wake_up(tags)

429

430 def is_sleeping(self) -> bool:

@@ -448,8 +448,8 @@ class EngineCore:

448 def save_sharded_state(

449 self,

450 path: str,

451 - pattern: Optional[str] = None,

452 - max_size: Optional[int] = None,

451 + pattern: str | None = None,

452 + max_size: int | None = None,

453 ) -> None:

454 self.model_executor.save_sharded_state(

455 path=path, pattern=pattern, max_size=max_size

@@ -457,10 +457,10 @@ class EngineCore:

457

458 def collective_rpc(

459 self,

460 - method: Union[str, Callable[..., _R]],

461 - timeout: Optional[float] = None,

460 + method: str | Callable[..., _R],

461 + timeout: float | None = None,

462 args: tuple = (),

463 - kwargs: Optional[dict[str, Any]] = None,

463 + kwargs: dict[str, Any] | None = None,

464 ) -> list[_R]:

465 return self.model_executor.collective_rpc(method, timeout, args, kwargs)

466

@@ -509,11 +509,11 @@ class EngineCoreProc(EngineCore):

509 handshake_address: str,

510 executor_class: type[Executor],

511 log_stats: bool,

512 - client_handshake_address: Optional[str] = None,

512 + client_handshake_address: str | None = None,

513 engine_index: int = 0,

514 ):

515 self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]()

516 - self.output_queue = queue.Queue[Union[tuple[int, EngineCoreOutputs], bytes]]()

516 + self.output_queue = queue.Queue[tuple[int, EngineCoreOutputs] | bytes]()

517 executor_fail_callback = lambda: self.input_queue.put_nowait(

518 (EngineCoreRequestType.EXECUTOR_FAILED, b"")

519 )

@@ -606,7 +606,7 @@ class EngineCoreProc(EngineCore):

606 identity: bytes,

607 local_client: bool,

608 vllm_config: VllmConfig,

609 - client_handshake_address: Optional[str],

609 + client_handshake_address: str | None,

610 ) -> Generator[EngineZmqAddresses, None, None]:

611 """

612 Perform startup handshakes.

@@ -667,7 +667,7 @@ class EngineCoreProc(EngineCore):

667 local_client: bool,

668 headless: bool,

669 vllm_config: VllmConfig,

670 - parallel_config_to_update: Optional[ParallelConfig] = None,

670 + parallel_config_to_update: ParallelConfig | None = None,

671 ) -> Generator[EngineZmqAddresses, None, None]:

672 with make_zmq_socket(

673 ctx,

@@ -710,7 +710,7 @@ class EngineCoreProc(EngineCore):

710 handshake_socket: zmq.Socket,

711 local_client: bool,

712 headless: bool,

713 - parallel_config: Optional[ParallelConfig] = None,

713 + parallel_config: ParallelConfig | None = None,

714 ) -> EngineZmqAddresses:

715 # Send registration message.

716 handshake_socket.send(

@@ -765,7 +765,7 @@ class EngineCoreProc(EngineCore):

765 signal.signal(signal.SIGTERM, signal_handler)

766 signal.signal(signal.SIGINT, signal_handler)

767

768 - engine_core: Optional[EngineCoreProc] = None

768 + engine_core: EngineCoreProc | None = None

769 try:

770 parallel_config: ParallelConfig = kwargs["vllm_config"].parallel_config

771 if parallel_config.data_parallel_size > 1 or dp_rank > 0:

@@ -911,7 +911,7 @@ class EngineCoreProc(EngineCore):

911 def process_input_sockets(

912 self,

913 input_addresses: list[str],

914 - coord_input_address: Optional[str],

914 + coord_input_address: str | None,

915 identity: bytes,

916 ready_event: threading.Event,

917 ):

@@ -980,7 +980,7 @@ class EngineCoreProc(EngineCore):

980 def process_output_sockets(

981 self,

982 output_paths: list[str],

983 - coord_output_path: Optional[str],

983 + coord_output_path: str | None,

984 engine_index: int,

985 ):

986 """Output socket IO thread."""

@@ -1059,7 +1059,7 @@ class DPEngineCoreProc(EngineCoreProc):

1059 handshake_address: str,

1060 executor_class: type[Executor],

1061 log_stats: bool,

1062 - client_handshake_address: Optional[str] = None,

1062 + client_handshake_address: str | None = None,

1063 ):

1064 # Counts forward-passes of the model so that we can synchronize

1065 # finished with DP peers every N steps.

@@ -1332,7 +1332,7 @@ class DPEngineCoreActor(DPEngineCoreProc):

1332 identity: bytes,

1333 local_client: bool,

1334 vllm_config: VllmConfig,

1335 - client_handshake_address: Optional[str],

1335 + client_handshake_address: str | None,

1336 ):

1337 """

1338 For Ray, we don't need to actually perform handshake.

vllm/v1/engine/core_client.py

+50 -52

@@ -9,11 +9,11 @@ import uuid

9 import weakref

10 from abc import ABC, abstractmethod

11 from collections import defaultdict, deque

12 -from collections.abc import Awaitable, Sequence

12 +from collections.abc import Awaitable, Callable, Sequence

13 from concurrent.futures import Future

14 from dataclasses import dataclass

15 from threading import Thread

16 -from typing import Any, Callable, Optional, TypeVar, Union

16 +from typing import Any, TypeAlias, TypeVar

17

18 import msgspec.msgpack

19 import zmq

@@ -51,7 +51,7 @@ from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr

51

52 logger = init_logger(__name__)

53

54 -AnyFuture = Union[asyncio.Future[Any], Future[Any]]

54 +AnyFuture: TypeAlias = asyncio.Future[Any] | Future[Any]

55

56 _R = TypeVar("_R") # Return type for collective_rpc

57

@@ -99,7 +99,7 @@ class EngineCoreClient(ABC):

99 vllm_config: VllmConfig,

100 executor_class: type[Executor],

101 log_stats: bool,

102 - client_addresses: Optional[dict[str, str]] = None,

102 + client_addresses: dict[str, str] | None = None,

103 client_count: int = 1,

104 client_index: int = 0,

105 ) -> "MPClient":

@@ -144,7 +144,7 @@ class EngineCoreClient(ABC):

144 def sleep(self, level: int = 1) -> None:

145 raise NotImplementedError

146

147 - def wake_up(self, tags: Optional[list[str]] = None) -> None:

147 + def wake_up(self, tags: list[str] | None = None) -> None:

148 raise NotImplementedError

149

150 def is_sleeping(self) -> bool:

@@ -172,16 +172,16 @@ class EngineCoreClient(ABC):

172 raise NotImplementedError

173

174 def save_sharded_state(

175 - self, path: str, pattern: Optional[str] = None, max_size: Optional[int] = None

175 + self, path: str, pattern: str | None = None, max_size: int | None = None

176 ) -> None:

177 raise NotImplementedError

178

179 def collective_rpc(

180 self,

181 - method: Union[str, Callable[..., _R]],

182 - timeout: Optional[float] = None,

181 + method: str | Callable[..., _R],

182 + timeout: float | None = None,

183 args: tuple = (),

184 - kwargs: Optional[dict[str, Any]] = None,

184 + kwargs: dict[str, Any] | None = None,

185 ) -> list[_R]:

186 raise NotImplementedError

187

@@ -214,7 +214,7 @@ class EngineCoreClient(ABC):

214 async def sleep_async(self, level: int = 1) -> None:

215 raise NotImplementedError

216

217 - async def wake_up_async(self, tags: Optional[list[str]] = None) -> None:

217 + async def wake_up_async(self, tags: list[str] | None = None) -> None:

218 raise NotImplementedError

219

220 async def is_sleeping_async(self) -> bool:

@@ -236,16 +236,16 @@ class EngineCoreClient(ABC):

236 raise NotImplementedError

237

238 async def save_sharded_state_async(

239 - self, path: str, pattern: Optional[str] = None, max_size: Optional[int] = None

239 + self, path: str, pattern: str | None = None, max_size: int | None = None

240 ) -> None:

241 raise NotImplementedError

242

243 async def collective_rpc_async(

244 self,

245 - method: Union[str, Callable[..., _R]],

246 - timeout: Optional[float] = None,

245 + method: str | Callable[..., _R],

246 + timeout: float | None = None,

247 args: tuple = (),

248 - kwargs: Optional[dict[str, Any]] = None,

248 + kwargs: dict[str, Any] | None = None,

249 ) -> list[_R]:

250 raise NotImplementedError

251

@@ -293,7 +293,7 @@ class InprocClient(EngineCoreClient):

293 def sleep(self, level: int = 1) -> None:

294 self.engine_core.sleep(level)

295

296 - def wake_up(self, tags: Optional[list[str]] = None) -> None:

296 + def wake_up(self, tags: list[str] | None = None) -> None:

297 self.engine_core.wake_up(tags)

298

299 def is_sleeping(self) -> bool:

@@ -315,16 +315,16 @@ class InprocClient(EngineCoreClient):

315 return self.engine_core.pin_lora(lora_id)

316

317 def save_sharded_state(

318 - self, path: str, pattern: Optional[str] = None, max_size: Optional[int] = None

318 + self, path: str, pattern: str | None = None, max_size: int | None = None

319 ) -> None:

320 self.engine_core.save_sharded_state(path, pattern, max_size)

321

322 def collective_rpc(

323 self,

324 - method: Union[str, Callable[..., _R]],

325 - timeout: Optional[float] = None,

324 + method: str | Callable[..., _R],

325 + timeout: float | None = None,

326 args: tuple = (),

327 - kwargs: Optional[dict[str, Any]] = None,

327 + kwargs: dict[str, Any] | None = None,

328 ) -> list[_R]:

329 return self.engine_core.collective_rpc(method, timeout, args, kwargs)

330

@@ -340,18 +340,16 @@ class BackgroundResources:

340 ctx: zmq.Context

341 # If CoreEngineProcManager, it manages local engines;

342 # if CoreEngineActorManager, it manages all engines.

343 - engine_manager: Optional[Union[CoreEngineProcManager, CoreEngineActorManager]] = (

344 - None

345 - )

346 - coordinator: Optional[DPCoordinator] = None

347 - output_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

348 - input_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None

349 - first_req_send_socket: Optional[zmq.asyncio.Socket] = None

350 - first_req_rcv_socket: Optional[zmq.asyncio.Socket] = None

351 - stats_update_socket: Optional[zmq.asyncio.Socket] = None

352 - output_queue_task: Optional[asyncio.Task] = None

353 - stats_update_task: Optional[asyncio.Task] = None

354 - shutdown_path: Optional[str] = None

343 + engine_manager: CoreEngineProcManager | CoreEngineActorManager | None = None

344 + coordinator: DPCoordinator | None = None

345 + output_socket: zmq.Socket | zmq.asyncio.Socket | None = None

346 + input_socket: zmq.Socket | zmq.asyncio.Socket | None = None

347 + first_req_send_socket: zmq.asyncio.Socket | None = None

348 + first_req_rcv_socket: zmq.asyncio.Socket | None = None

349 + stats_update_socket: zmq.asyncio.Socket | None = None

350 + output_queue_task: asyncio.Task | None = None

351 + stats_update_task: asyncio.Task | None = None

352 + shutdown_path: str | None = None

353

354 # Set if any of the engines are dead. Here so that the output

355 # processing threads can access it without holding a ref to the client.

@@ -438,7 +436,7 @@ class MPClient(EngineCoreClient):

436 vllm_config: VllmConfig,

437 executor_class: type[Executor],

438 log_stats: bool,

441 - client_addresses: Optional[dict[str, str]] = None,

439 + client_addresses: dict[str, str] | None = None,

440 ):

441 self.vllm_config = vllm_config

442 # Serialization setup.

@@ -459,7 +457,7 @@ class MPClient(EngineCoreClient):

457 # State used for data parallel.

458 self.engines_running = False

459

462 - self.stats_update_address: Optional[str] = None

460 + self.stats_update_address: str | None = None

461 if client_addresses:

462 # Engines are managed externally to this client.

463 input_address = client_addresses["input_address"]

@@ -646,7 +644,7 @@ class SyncMPClient(MPClient):

644 )

645

646 self.is_dp = self.vllm_config.parallel_config.data_parallel_size > 1

649 - self.outputs_queue = queue.Queue[Union[EngineCoreOutputs, Exception]]()

647 + self.outputs_queue = queue.Queue[EngineCoreOutputs | Exception]()

648

649 # Ensure that the outputs socket processing thread does not have

650 # a ref to the client which prevents gc.

@@ -770,7 +768,7 @@ class SyncMPClient(MPClient):

768 def sleep(self, level: int = 1) -> None:

769 self.call_utility("sleep", level)

770

773 - def wake_up(self, tags: Optional[list[str]] = None) -> None:

771 + def wake_up(self, tags: list[str] | None = None) -> None:

772 self.call_utility("wake_up", tags)

773

774 def is_sleeping(self) -> bool:

@@ -781,15 +779,15 @@ class SyncMPClient(MPClient):

779

780 def collective_rpc(

781 self,

784 - method: Union[str, Callable[..., _R]],

785 - timeout: Optional[float] = None,

782 + method: str | Callable[..., _R],

783 + timeout: float | None = None,

784 args: tuple = (),

787 - kwargs: Optional[dict[str, Any]] = None,

785 + kwargs: dict[str, Any] | None = None,

786 ) -> list[_R]:

787 return self.call_utility("collective_rpc", method, timeout, args, kwargs)

788

789 def save_sharded_state(

792 - self, path: str, pattern: Optional[str] = None, max_size: Optional[int] = None

790 + self, path: str, pattern: str | None = None, max_size: int | None = None

791 ) -> None:

792 self.call_utility("save_sharded_state", path, pattern, max_size)

793

@@ -802,7 +800,7 @@ class AsyncMPClient(MPClient):

800 vllm_config: VllmConfig,

801 executor_class: type[Executor],

802 log_stats: bool,

805 - client_addresses: Optional[dict[str, str]] = None,

803 + client_addresses: dict[str, str] | None = None,

804 client_count: int = 1,

805 client_index: int = 0,

806 ):

@@ -816,7 +814,7 @@ class AsyncMPClient(MPClient):

814

815 self.client_count = client_count

816 self.client_index = client_index

819 - self.outputs_queue = asyncio.Queue[Union[EngineCoreOutputs, Exception]]()

817 + self.outputs_queue = asyncio.Queue[EngineCoreOutputs | Exception]()

818 try:

819 # If we are running in an asyncio event loop, start the queue task.

820 # Otherwise, it will be started lazily. If it is not started here,

@@ -837,9 +835,9 @@ class AsyncMPClient(MPClient):

835 decoder = self.decoder

836 utility_results = self.utility_results

837 outputs_queue = self.outputs_queue

840 - output_handler: Optional[

841 - Callable[[AsyncMPClient, EngineCoreOutputs], Awaitable[None]]

842 - ] = getattr(self.__class__, "process_engine_outputs", None)

838 + output_handler: (

839 + Callable[[AsyncMPClient, EngineCoreOutputs], Awaitable[None]] | None

840 + ) = getattr(self.__class__, "process_engine_outputs", None)

841 _self_ref = weakref.ref(self) if output_handler else None

842 output_socket = resources.output_socket

843 assert output_socket is not None

@@ -888,7 +886,7 @@ class AsyncMPClient(MPClient):

886 self,

887 request_type: EngineCoreRequestType,

888 request: Any,

891 - engine: Optional[EngineIdentity] = None,

889 + engine: EngineIdentity | None = None,

890 ) -> Awaitable[Any]:

891 if engine is None:

892 engine = self.core_engine

@@ -962,7 +960,7 @@ class AsyncMPClient(MPClient):

960 async def sleep_async(self, level: int = 1) -> None:

961 await self.call_utility_async("sleep", level)

962

965 - async def wake_up_async(self, tags: Optional[list[str]] = None) -> None:

963 + async def wake_up_async(self, tags: list[str] | None = None) -> None:

964 await self.call_utility_async("wake_up", tags)

965

966 async def is_sleeping_async(self) -> bool:

@@ -984,16 +982,16 @@ class AsyncMPClient(MPClient):

982 return await self.call_utility_async("pin_lora", lora_id)

983

984 async def save_sharded_state_async(

987 - self, path: str, pattern: Optional[str] = None, max_size: Optional[int] = None

985 + self, path: str, pattern: str | None = None, max_size: int | None = None

986 ) -> None:

987 await self.call_utility_async("save_sharded_state", path, pattern, max_size)

988

989 async def collective_rpc_async(

990 self,

993 - method: Union[str, Callable[..., _R]],

994 - timeout: Optional[float] = None,

991 + method: str | Callable[..., _R],

992 + timeout: float | None = None,

993 args: tuple = (),

996 - kwargs: Optional[dict[str, Any]] = None,

994 + kwargs: dict[str, Any] | None = None,

995 ) -> list[_R]:

996 return await self.call_utility_async(

997 "collective_rpc", method, timeout, args, kwargs

@@ -1009,7 +1007,7 @@ class DPAsyncMPClient(AsyncMPClient):

1007 vllm_config: VllmConfig,

1008 executor_class: type[Executor],

1009 log_stats: bool,

1012 - client_addresses: Optional[dict[str, str]] = None,

1010 + client_addresses: dict[str, str] | None = None,

1011 client_count: int = 1,

1012 client_index: int = 0,

1013 ):

@@ -1166,7 +1164,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1164 vllm_config: VllmConfig,

1165 executor_class: type[Executor],

1166 log_stats: bool,

1169 - client_addresses: Optional[dict[str, str]] = None,

1167 + client_addresses: dict[str, str] | None = None,

1168 client_count: int = 1,

1169 client_index: int = 0,

1170 ):

vllm/v1/engine/detokenizer.py

+5 -6

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from abc import ABC, abstractmethod

4 -from typing import Optional

4

5 import tokenizers

6 from packaging import version

@@ -36,7 +35,7 @@ class IncrementalDetokenizer:

35 def output_token_ids(self) -> list[int]:

36 return self.token_ids

37

39 - def update(self, new_token_ids: list[int], stop_terminated: bool) -> Optional[str]:

38 + def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None:

39 self.token_ids.extend(new_token_ids)

40 return None

41

@@ -46,7 +45,7 @@ class IncrementalDetokenizer:

45 @classmethod

46 def from_new_request(

47 cls,

49 - tokenizer: Optional[AnyTokenizer],

48 + tokenizer: AnyTokenizer | None,

49 request: EngineCoreRequest,

50 ) -> "IncrementalDetokenizer":

51 assert request.sampling_params is not None

@@ -85,7 +84,7 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC):

84 # Generation data

85 self.output_text = ""

86

88 - def update(self, new_token_ids: list[int], stop_terminated: bool) -> Optional[str]:

87 + def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None:

88 """

89 Update RequestState for the request_id by:

90 1) Detokenize the new token ids incrementally.

@@ -224,7 +223,7 @@ class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):

223

224 return token or ""

225

227 - def _protected_step(self, next_token_id: int) -> Optional[str]:

226 + def _protected_step(self, next_token_id: int) -> str | None:

227 try:

228 token = self.stream.step(self.tokenizer, next_token_id)

229 except (OverflowError, TypeError):

@@ -312,7 +311,7 @@ def check_stop_strings(

311 new_char_count: int,

312 stop: list[str],

313 include_in_output: bool,

315 -) -> Optional[tuple[str, int]]:

314 +) -> tuple[str, int] | None:

315 """Check if any stop strings are matched and truncate sequence

316 output text accordingly.

317

vllm/v1/engine/llm_engine.py

+21 -21

@@ -2,9 +2,9 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import time

5 -from collections.abc import Mapping

5 +from collections.abc import Callable, Mapping

6 from copy import copy

7 -from typing import Any, Callable, Optional, Union

7 +from typing import Any

8

9 import torch.nn as nn

10 from typing_extensions import TypeVar

@@ -52,7 +52,7 @@ class LLMEngine:

52 executor_class: type[Executor],

53 log_stats: bool,

54 usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,

55 - stat_loggers: Optional[list[StatLoggerFactory]] = None,

55 + stat_loggers: list[StatLoggerFactory] | None = None,

56 mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,

57 use_cached_outputs: bool = False,

58 multiprocess_mode: bool = False,

@@ -126,7 +126,7 @@ class LLMEngine:

126 log_stats=self.log_stats,

127 )

128

129 - self.logger_manager: Optional[StatLoggerManager] = None

129 + self.logger_manager: StatLoggerManager | None = None

130 if self.log_stats:

131 self.logger_manager = StatLoggerManager(

132 vllm_config=vllm_config,

@@ -152,7 +152,7 @@ class LLMEngine:

152 cls,

153 vllm_config: VllmConfig,

154 usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,

155 - stat_loggers: Optional[list[StatLoggerFactory]] = None,

155 + stat_loggers: list[StatLoggerFactory] | None = None,

156 disable_log_stats: bool = False,

157 ) -> "LLMEngine":

158 return cls(

@@ -169,7 +169,7 @@ class LLMEngine:

169 cls,

170 engine_args: EngineArgs,

171 usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,

172 - stat_loggers: Optional[list[StatLoggerFactory]] = None,

172 + stat_loggers: list[StatLoggerFactory] | None = None,

173 enable_multiprocessing: bool = False,

174 ) -> "LLMEngine":

175 """Creates an LLM engine from the engine arguments."""

@@ -225,14 +225,14 @@ class LLMEngine:

225 def add_request(

226 self,

227 request_id: str,

228 - prompt: Union[EngineCoreRequest, PromptType],

229 - params: Union[SamplingParams, PoolingParams],

230 - arrival_time: Optional[float] = None,

231 - lora_request: Optional[LoRARequest] = None,

232 - tokenization_kwargs: Optional[dict[str, Any]] = None,

233 - trace_headers: Optional[Mapping[str, str]] = None,

228 + prompt: EngineCoreRequest | PromptType,

229 + params: SamplingParams | PoolingParams,

230 + arrival_time: float | None = None,

231 + lora_request: LoRARequest | None = None,

232 + tokenization_kwargs: dict[str, Any] | None = None,

233 + trace_headers: Mapping[str, str] | None = None,

234 priority: int = 0,

235 - prompt_text: Optional[str] = None,

235 + prompt_text: str | None = None,

236 ) -> None:

237 # Validate the request_id type.

238 if not isinstance(request_id, str):

@@ -283,7 +283,7 @@ class LLMEngine:

283 # Add the request to EngineCore.

284 self.engine_core.add_request(child_request)

285

286 - def step(self) -> Union[list[RequestOutput], list[PoolingRequestOutput]]:

286 + def step(self) -> list[RequestOutput] | list[PoolingRequestOutput]:

287 if self.should_execute_dummy_batch:

288 self.should_execute_dummy_batch = False

289 self.engine_core.execute_dummy_batch()

@@ -326,13 +326,13 @@ class LLMEngine:

326 self.processor.clear_mm_cache()

327 self.engine_core.reset_mm_cache()

328

329 - def reset_prefix_cache(self, device: Optional[Device] = None):

329 + def reset_prefix_cache(self, device: Device | None = None):

330 self.engine_core.reset_prefix_cache()

331

332 def sleep(self, level: int = 1):

333 self.engine_core.sleep(level)

334

335 - def wake_up(self, tags: Optional[list[str]] = None):

335 + def wake_up(self, tags: list[str] | None = None):

336 self.engine_core.wake_up(tags)

337

338 def is_sleeping(self) -> bool:

@@ -343,11 +343,11 @@ class LLMEngine:

343 return get_metrics_snapshot()

344

345 @property

346 - def tokenizer(self) -> Optional[AnyTokenizer]:

346 + def tokenizer(self) -> AnyTokenizer | None:

347 return self.processor.tokenizer

348

349 @tokenizer.setter

350 - def tokenizer(self, tokenizer: Optional[AnyTokenizer]) -> None:

350 + def tokenizer(self, tokenizer: AnyTokenizer | None) -> None:

351 self.processor.tokenizer = tokenizer

352

353 def get_tokenizer(self) -> AnyTokenizer:

@@ -390,10 +390,10 @@ class LLMEngine:

390

391 def collective_rpc(

392 self,

393 - method: Union[str, Callable[[WorkerBase], _R]],

394 - timeout: Optional[float] = None,

393 + method: str | Callable[[WorkerBase], _R],

394 + timeout: float | None = None,

395 args: tuple = (),

396 - kwargs: Optional[dict[str, Any]] = None,

396 + kwargs: dict[str, Any] | None = None,

397 ) -> list[_R]:

398 return self.engine_core.collective_rpc(method, timeout, args, kwargs)

399

vllm/v1/engine/logprobs.py

+9 -10

@@ -4,7 +4,6 @@

4 import itertools

5 from collections.abc import Iterable

6 from dataclasses import dataclass

7 -from typing import Optional

7

8 from vllm.logger import init_logger

9 from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs

@@ -24,19 +23,19 @@ NONES = itertools.repeat(None)

23 class LogprobsProcessor:

24 # Tokenizer for this request,

25 # None if detokenization is disabled.

27 - tokenizer: Optional[AnyTokenizer]

26 + tokenizer: AnyTokenizer | None

27

28 # Logprobs for this request

30 - logprobs: Optional[SampleLogprobs]

31 - prompt_logprobs: Optional[PromptLogprobs]

32 - cumulative_logprob: Optional[float]

33 - num_logprobs: Optional[int]

34 - num_prompt_logprobs: Optional[int]

29 + logprobs: SampleLogprobs | None

30 + prompt_logprobs: PromptLogprobs | None

31 + cumulative_logprob: float | None

32 + num_logprobs: int | None

33 + num_prompt_logprobs: int | None

34

35 @classmethod

36 def from_new_request(

37 cls,

39 - tokenizer: Optional[AnyTokenizer],

38 + tokenizer: AnyTokenizer | None,

39 request: EngineCoreRequest,

40 ) -> "LogprobsProcessor":

41 assert request.sampling_params is not None

@@ -148,7 +147,7 @@ class LogprobsProcessor:

147 )

148 )

149

151 - def pop_prompt_logprobs(self) -> Optional[PromptLogprobs]:

150 + def pop_prompt_logprobs(self) -> PromptLogprobs | None:

151 """Pop and return all request prompt logprobs

152

153 The logprobs processor aggregates prompt chunk logprobs

@@ -171,7 +170,7 @@ class LogprobsProcessor:

170 def _make_logprob_dict(

171 logprobs: list[float],

172 logprob_token_ids: list[int],

174 - decoded_tokens: Iterable[Optional[str]],

173 + decoded_tokens: Iterable[str | None],

174 rank: int,

175 num_logprobs: int,

176 ) -> dict[int, Logprob]:

vllm/v1/engine/output_processor.py

+43 -47

@@ -4,7 +4,7 @@

4 import asyncio

5 from collections.abc import Iterable

6 from dataclasses import dataclass

7 -from typing import Any, Optional, Union, cast

7 +from typing import Any, cast

8

9 import torch

10

@@ -36,14 +36,10 @@ class RequestOutputCollector:

36

37 def __init__(self, output_kind: RequestOutputKind):

38 self.aggregate = output_kind == RequestOutputKind.DELTA

39 - self.output: Optional[Union[RequestOutput, PoolingRequestOutput, Exception]] = (

40 - None

41 - )

39 + self.output: RequestOutput | PoolingRequestOutput | Exception | None = None

40 self.ready = asyncio.Event()

41

44 - def put(

45 - self, output: Union[RequestOutput, PoolingRequestOutput, Exception]

46 - ) -> None:

42 + def put(self, output: RequestOutput | PoolingRequestOutput | Exception) -> None:

43 """Non-blocking put operation."""

44 if self.output is None or isinstance(output, Exception):

45 self.output = output

@@ -53,7 +49,7 @@ class RequestOutputCollector:

49 # (if n > 1) do not override each other.

50 self.output.add(output, aggregate=self.aggregate)

51

56 - async def get(self) -> Union[RequestOutput, PoolingRequestOutput]:

52 + async def get(self) -> RequestOutput | PoolingRequestOutput:

53 """Get operation blocks on put event."""

54 while (output := self.output) is None:

55 await self.ready.wait()

@@ -63,7 +59,7 @@ class RequestOutputCollector:

59 raise output

60 return output

61

66 - def get_nowait(self) -> Optional[Union[RequestOutput, PoolingRequestOutput]]:

62 + def get_nowait(self) -> RequestOutput | PoolingRequestOutput | None:

63 """Non-blocking get operation."""

64 output = self.output

65 if output is not None:

@@ -76,7 +72,7 @@ class RequestOutputCollector:

72

73 @dataclass

74 class OutputProcessorOutput:

79 - request_outputs: list[Union[RequestOutput, PoolingRequestOutput]]

75 + request_outputs: list[RequestOutput | PoolingRequestOutput]

76 reqs_to_abort: list[str]

77

78

@@ -84,22 +80,22 @@ class RequestState:

80 def __init__(

81 self,

82 request_id: str,

87 - parent_req: Optional[ParentRequest],

83 + parent_req: ParentRequest | None,

84 request_index: int,

89 - lora_name: Optional[str],

85 + lora_name: str | None,

86 output_kind: RequestOutputKind,

91 - prompt: Optional[str],

92 - prompt_token_ids: Optional[list[int]],

93 - prompt_embeds: Optional[torch.Tensor],

94 - logprobs_processor: Optional[LogprobsProcessor],

95 - detokenizer: Optional[IncrementalDetokenizer],

96 - max_tokens_param: Optional[int],

87 + prompt: str | None,

88 + prompt_token_ids: list[int] | None,

89 + prompt_embeds: torch.Tensor | None,

90 + logprobs_processor: LogprobsProcessor | None,

91 + detokenizer: IncrementalDetokenizer | None,

92 + max_tokens_param: int | None,

93 arrival_time: float,

98 - queue: Optional[RequestOutputCollector],

94 + queue: RequestOutputCollector | None,

95 log_stats: bool,

100 - top_p: Optional[float] = None,

101 - n: Optional[int] = None,

102 - temperature: Optional[float] = None,

96 + top_p: float | None = None,

97 + n: int | None = None,

98 + temperature: float | None = None,

99 ):

100 self.request_id = request_id

101 self.parent_req = parent_req

@@ -129,10 +125,10 @@ class RequestState:

125 cls,

126 tokenizer: AnyTokenizer,

127 request: EngineCoreRequest,

132 - prompt: Optional[str],

133 - parent_req: Optional[ParentRequest],

128 + prompt: str | None,

129 + parent_req: ParentRequest | None,

130 request_index: int,

135 - queue: Optional[RequestOutputCollector],

131 + queue: RequestOutputCollector | None,

132 log_stats: bool,

133 ) -> "RequestState":

134 if sampling_params := request.sampling_params:

@@ -186,11 +182,11 @@ class RequestState:

182 def make_request_output(

183 self,

184 new_token_ids: list[int],

189 - pooling_output: Optional[torch.Tensor],

190 - finish_reason: Optional[FinishReason],

191 - stop_reason: Union[int, str, None],

192 - kv_transfer_params: Optional[dict[str, Any]] = None,

193 - ) -> Optional[Union[RequestOutput, PoolingRequestOutput]]:

185 + pooling_output: torch.Tensor | None,

186 + finish_reason: FinishReason | None,

187 + stop_reason: int | str | None,

188 + kv_transfer_params: dict[str, Any] | None = None,

189 + ) -> RequestOutput | PoolingRequestOutput | None:

190 finished = finish_reason is not None

191 final_only = self.output_kind == RequestOutputKind.FINAL_ONLY

192

@@ -222,10 +218,10 @@ class RequestState:

218 def _new_request_output(

219 self,

220 request_id: str,

225 - outputs: Union[list[CompletionOutput], list[PoolingOutput]],

221 + outputs: list[CompletionOutput] | list[PoolingOutput],

222 finished: bool,

227 - kv_transfer_params: Optional[dict[str, Any]] = None,

228 - ) -> Union[RequestOutput, PoolingRequestOutput]:

223 + kv_transfer_params: dict[str, Any] | None = None,

224 + ) -> RequestOutput | PoolingRequestOutput:

225 first_output = outputs[0]

226 if isinstance(first_output, PoolingOutput):

227 assert len(outputs) == 1

@@ -264,8 +260,8 @@ class RequestState:

260 def _new_completion_output(

261 self,

262 token_ids: list[int],

267 - finish_reason: Optional[FinishReason],

268 - stop_reason: Union[int, str, None],

263 + finish_reason: FinishReason | None,

264 + stop_reason: int | str | None,

265 ) -> CompletionOutput:

266 assert self.detokenizer is not None

267 assert self.logprobs_processor is not None

@@ -308,7 +304,7 @@ class OutputProcessor:

304 self.request_states: dict[str, RequestState] = {}

305 self.parent_requests: dict[str, ParentRequest] = {}

306 self.lora_states = LoRARequestStates()

311 - self.tracer: Optional[Tracer] = None

307 + self.tracer: Tracer | None = None

308

309 def get_num_unfinished_requests(self):

310 return len(self.request_states)

@@ -360,10 +356,10 @@ class OutputProcessor:

356 def add_request(

357 self,

358 request: EngineCoreRequest,

363 - prompt: Optional[str],

364 - parent_req: Optional[ParentRequest] = None,

359 + prompt: str | None,

360 + parent_req: ParentRequest | None = None,

361 request_index: int = 0,

366 - queue: Optional[RequestOutputCollector] = None,

362 + queue: RequestOutputCollector | None = None,

363 ) -> None:

364 request_id = request.request_id

365 if request_id in self.request_states:

@@ -386,8 +382,8 @@ class OutputProcessor:

382 def process_outputs(

383 self,

384 engine_core_outputs: list[EngineCoreOutput],

389 - engine_core_timestamp: Optional[float] = None,

390 - iteration_stats: Optional[IterationStats] = None,

385 + engine_core_timestamp: float | None = None,

386 + iteration_stats: IterationStats | None = None,

387 ) -> OutputProcessorOutput:

388 """

389 Process the EngineCoreOutputs:

@@ -411,7 +407,7 @@ class OutputProcessor:

407 within the loop below.

408 """

409

414 - request_outputs: Union[list[RequestOutput], list[PoolingRequestOutput]] = []

410 + request_outputs: list[RequestOutput] | list[PoolingRequestOutput] = []

411 reqs_to_abort: list[str] = []

412 for engine_core_output in engine_core_outputs:

413 req_id = engine_core_output.request_id

@@ -492,7 +488,7 @@ class OutputProcessor:

488 self,

489 engine_core_output: EngineCoreOutput,

490 req_state: RequestState,

495 - iteration_stats: Optional[IterationStats],

491 + iteration_stats: IterationStats | None,

492 ) -> None:

493 assert req_state.stats is not None

494 assert iteration_stats is not None

@@ -555,8 +551,8 @@ class OutputProcessor:

551 self,

552 req_state: RequestState,

553 engine_core_output: EngineCoreOutput,

558 - engine_core_timestamp: Optional[float],

559 - iteration_stats: Optional[IterationStats],

554 + engine_core_timestamp: float | None,

555 + iteration_stats: IterationStats | None,

556 ):

557 if iteration_stats is None:

558 return

@@ -577,8 +573,8 @@ class OutputProcessor:

573 def _update_stats_from_finished(

574 self,

575 req_state: RequestState,

580 - finish_reason: Optional[FinishReason],

581 - iteration_stats: Optional[IterationStats],

576 + finish_reason: FinishReason | None,

577 + iteration_stats: IterationStats | None,

578 ):

579 if iteration_stats is None:

580 return

vllm/v1/engine/parallel_sampling.py

+1 -1

@@ -29,7 +29,7 @@ class ParentRequest:

29 max_num_generation_tokens: int

30

31 # To efficiently obtain child sampling params

32 - cached_child_sampling_params: Optional[SamplingParams]

32 + cached_child_sampling_params: SamplingParams | None

33

34 def __init__(self, request_id: str, sampling_params: SamplingParams) -> None:

35 self.request_id = request_id

vllm/v1/engine/processor.py

+17 -17

@@ -3,7 +3,7 @@

3

4 import time

5 from collections.abc import Mapping

6 -from typing import Any, Literal, Optional, Union

6 +from typing import Any, Literal

7

8 from vllm.config import VllmConfig

9 from vllm.inputs import ProcessorInputs, PromptType, SingletonInputs

@@ -38,7 +38,7 @@ class Processor:

38 def __init__(

39 self,

40 vllm_config: VllmConfig,

41 - tokenizer: Optional[AnyTokenizer],

41 + tokenizer: AnyTokenizer | None,

42 mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,

43 ) -> None:

44 self.vllm_config = vllm_config

@@ -60,11 +60,11 @@ class Processor:

60 )

61

62 @property

63 - def tokenizer(self) -> Optional[AnyTokenizer]:

63 + def tokenizer(self) -> AnyTokenizer | None:

64 return self.input_preprocessor.tokenizer

65

66 @tokenizer.setter

67 - def tokenizer(self, tokenizer: Optional[AnyTokenizer]) -> None:

67 + def tokenizer(self, tokenizer: AnyTokenizer | None) -> None:

68 self.input_preprocessor.tokenizer = tokenizer

69

70 def _validate_logprobs(

@@ -152,7 +152,7 @@ class Processor:

152

153 def _validate_params(

154 self,

155 - params: Union[SamplingParams, PoolingParams],

155 + params: SamplingParams | PoolingParams,

156 ):

157 """

158 Validate supported SamplingParam.

@@ -174,7 +174,7 @@ class Processor:

174 auto-hashed downstream.

175 """

176

177 - def _validate_single_prompt(single_prompt: Union[dict, str]) -> None:

177 + def _validate_single_prompt(single_prompt: dict | str) -> None:

178 if not isinstance(single_prompt, dict):

179 return

180 mm_data = single_prompt.get("multi_modal_data")

@@ -214,7 +214,7 @@ class Processor:

214 else:

215 _validate_single_prompt(prompt) # type: ignore[arg-type]

216

217 - def _validate_lora(self, lora_request: Optional[LoRARequest]) -> None:

217 + def _validate_lora(self, lora_request: LoRARequest | None) -> None:

218 if lora_request is None:

219 return

220

@@ -309,7 +309,7 @@ class Processor:

309 self,

310 request_id: str,

311 prompt: PromptType,

312 - ) -> Optional[MultiModalUUIDDict]:

312 + ) -> MultiModalUUIDDict | None:

313 """Build per-item multimodal hash overrides when enabled. In this case,

314 multimodal data items are identified by their request id, modality and

315 index rather than their content.

@@ -342,13 +342,13 @@ class Processor:

342 self,

343 request_id: str,

344 prompt: PromptType,

345 - params: Union[SamplingParams, PoolingParams],

346 - arrival_time: Optional[float] = None,

347 - lora_request: Optional[LoRARequest] = None,

348 - tokenization_kwargs: Optional[dict[str, Any]] = None,

349 - trace_headers: Optional[Mapping[str, str]] = None,

345 + params: SamplingParams | PoolingParams,

346 + arrival_time: float | None = None,

347 + lora_request: LoRARequest | None = None,

348 + tokenization_kwargs: dict[str, Any] | None = None,

349 + trace_headers: Mapping[str, str] | None = None,

350 priority: int = 0,

351 - data_parallel_rank: Optional[int] = None,

351 + data_parallel_rank: int | None = None,

352 ) -> EngineCoreRequest:

353 self._validate_lora(lora_request)

354 self._validate_params(params)

@@ -445,7 +445,7 @@ class Processor:

445 pooling_params = params.clone()

446

447 # Multimodal related.

448 - mm_features: Optional[list[MultiModalFeatureSpec]] = None

448 + mm_features: list[MultiModalFeatureSpec] | None = None

449

450 if decoder_inputs["type"] == "multimodal":

451 decoder_mm_inputs = decoder_inputs["mm_kwargs"]

@@ -485,7 +485,7 @@ class Processor:

485 )

486

487 def _validate_model_inputs(

488 - self, encoder_inputs: Optional[SingletonInputs], decoder_inputs: SingletonInputs

488 + self, encoder_inputs: SingletonInputs | None, decoder_inputs: SingletonInputs

489 ):

490 if encoder_inputs is not None:

491 self._validate_model_input(encoder_inputs, prompt_type="encoder")

@@ -574,7 +574,7 @@ class Processor:

574 # check that chunked prefill does not truncate them

575 # max_batch_len = self.scheduler_config.max_num_batched_tokens

576

577 - def stat_mm_cache(self) -> Optional[MultiModalCacheStats]:

577 + def stat_mm_cache(self) -> MultiModalCacheStats | None:

578 return self.input_preprocessor.stat_mm_cache()

579

580 def clear_mm_cache(self) -> None:

vllm/v1/engine/utils.py

+14 -14

@@ -4,12 +4,12 @@

4 import contextlib

5 import os

6 import weakref

7 -from collections.abc import Iterator

7 +from collections.abc import Callable, Iterator

8 from dataclasses import dataclass

9 from enum import Enum, auto

10 from multiprocessing import Process, connection

11 from multiprocessing.process import BaseProcess

12 -from typing import TYPE_CHECKING, Callable, Optional, Union

12 +from typing import TYPE_CHECKING

13 from unittest.mock import patch

14

15 import msgspec

@@ -56,13 +56,13 @@ class EngineZmqAddresses:

56 # ZMQ output socket addresses for each front-end client (responses)

57 outputs: list[str]

58 # ZMQ input socket address of DP coordinator if applicable

59 - coordinator_input: Optional[str] = None

59 + coordinator_input: str | None = None

60 # ZMQ output socket address of DP coordinator if applicable

61 - coordinator_output: Optional[str] = None

61 + coordinator_output: str | None = None

62 # ZMQ socket for front-end to connect to DP coordinator.

63 # Not used by engine, just relayed to front-end in handshake response.

64 # Only required for external DP LB case.

65 - frontend_stats_publish_address: Optional[str] = None

65 + frontend_stats_publish_address: str | None = None

66

67

68 @dataclass

@@ -73,8 +73,8 @@ class EngineHandshakeMetadata:

73 """

74

75 addresses: EngineZmqAddresses

76 - parallel_config: dict[str, Union[int, str, list[int]]]

77 - parallel_config_hash: Optional[str] = None

76 + parallel_config: dict[str, int | str | list[int]]

77 + parallel_config_hash: str | None = None

78

79

80 class CoreEngineProcManager:

@@ -94,7 +94,7 @@ class CoreEngineProcManager:

94 handshake_address: str,

95 executor_class: type[Executor],

96 log_stats: bool,

97 - client_handshake_address: Optional[str] = None,

97 + client_handshake_address: str | None = None,

98 ):

99 context = get_mp_context()

100 common_kwargs = {

@@ -221,8 +221,8 @@ class CoreEngineActorManager:

221 addresses: EngineZmqAddresses,

222 executor_class: type[Executor],

223 log_stats: bool,

224 - placement_groups: Optional[list["PlacementGroup"]] = None,

225 - local_dp_ranks: Optional[list[int]] = None,

224 + placement_groups: list["PlacementGroup"] | None = None,

225 + local_dp_ranks: list[int] | None = None,

226 ):

227 import copy

228

@@ -675,8 +675,8 @@ def launch_core_engines(

675 num_api_servers: int = 1,

676 ) -> Iterator[

677 tuple[

678 - Optional[Union[CoreEngineProcManager, CoreEngineActorManager]],

679 - Optional[DPCoordinator],

678 + CoreEngineProcManager | CoreEngineActorManager | None,

679 + DPCoordinator | None,

680 EngineZmqAddresses,

681 ]

682 ]:

@@ -829,8 +829,8 @@ def wait_for_engine_startup(

829 core_engines: list[CoreEngine],

830 parallel_config: ParallelConfig,

831 cache_config: CacheConfig,

832 - proc_manager: Optional[CoreEngineProcManager],

833 - coord_process: Optional[Process],

832 + proc_manager: CoreEngineProcManager | None,

833 + coord_process: Process | None,

834 ):

835 # Wait for engine core process(es) to send ready messages.

836 local_count = parallel_config.data_parallel_size_local

vllm/v1/executor/abstract.py

+7 -6

@@ -1,8 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from collections.abc import Callable

5 from concurrent.futures import Future

5 -from typing import Any, Callable, Optional, Union

6 +from typing import Any

7

8 import torch

9 import torch.distributed as dist

@@ -91,10 +92,10 @@ class Executor(ExecutorBase):

92

93 def collective_rpc(

94 self,

94 - method: Union[str, Callable],

95 - timeout: Optional[float] = None,

95 + method: str | Callable,

96 + timeout: float | None = None,

97 args: tuple = (),

97 - kwargs: Optional[dict] = None,

98 + kwargs: dict | None = None,

99 non_block: bool = False,

100 ) -> list[Any]:

101 raise NotImplementedError

@@ -103,7 +104,7 @@ class Executor(ExecutorBase):

104 self,

105 scheduler_output: SchedulerOutput,

106 non_block: bool = False,

106 - ) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:

107 + ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:

108 output = self.collective_rpc(

109 "execute_model", args=(scheduler_output,), non_block=non_block

110 )

@@ -112,7 +113,7 @@ class Executor(ExecutorBase):

113 def execute_dummy_batch(self) -> None:

114 self.collective_rpc("execute_dummy_batch")

115

115 - def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

116 + def take_draft_token_ids(self) -> DraftTokenIds | None:

117 output = self.collective_rpc("take_draft_token_ids")

118 return output[0]

119

vllm/v1/executor/multiproc_executor.py

+16 -15

@@ -9,6 +9,7 @@ import threading

9 import time

10 import traceback

11 import weakref

12 +from collections.abc import Callable

13 from concurrent.futures import Future, ThreadPoolExecutor

14 from dataclasses import dataclass

15 from enum import Enum, auto

@@ -17,7 +18,7 @@ from multiprocessing.connection import Connection

18 from multiprocessing.process import BaseProcess

19 from multiprocessing.synchronize import Lock as LockType

20 from threading import Thread

20 -from typing import Any, Callable, Optional, Union, cast

21 +from typing import Any, cast

22

23 import cloudpickle

24 import torch

@@ -59,8 +60,8 @@ class MultiprocExecutor(Executor):

60 self._finalizer = weakref.finalize(self, self.shutdown)

61 self.is_failed = False

62 self.shutdown_event = threading.Event()

62 - self.failure_callback: Optional[FailureCallback] = None

63 - self.io_thread_pool: Optional[ThreadPoolExecutor] = None

63 + self.failure_callback: FailureCallback | None = None

64 + self.io_thread_pool: ThreadPoolExecutor | None = None

65

66 self.world_size = self.parallel_config.world_size

67 tensor_parallel_size = self.parallel_config.tensor_parallel_size

@@ -179,7 +180,7 @@ class MultiprocExecutor(Executor):

180 self,

181 scheduler_output: SchedulerOutput,

182 non_block: bool = False,

182 - ) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:

183 + ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:

184 if not self.has_connector:

185 # get output only from a single worker (output_rank)

186 (output,) = self.collective_rpc(

@@ -207,7 +208,7 @@ class MultiprocExecutor(Executor):

208 def execute_dummy_batch(self) -> None:

209 self.collective_rpc("execute_dummy_batch", unique_reply_rank=self.output_rank)

210

210 - def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

211 + def take_draft_token_ids(self) -> DraftTokenIds | None:

212 # OPTIMIZATION: Get output only from a single worker (output_rank)

213 outputs = self.collective_rpc(

214 "take_draft_token_ids", unique_reply_rank=self.output_rank

@@ -216,12 +217,12 @@ class MultiprocExecutor(Executor):

217

218 def collective_rpc(

219 self,

219 - method: Union[str, Callable],

220 - timeout: Optional[float] = None,

220 + method: str | Callable,

221 + timeout: float | None = None,

222 args: tuple = (),

222 - kwargs: Optional[dict] = None,

223 + kwargs: dict | None = None,

224 non_block: bool = False,

224 - unique_reply_rank: Optional[int] = None,

225 + unique_reply_rank: int | None = None,

226 ) -> list[Any]:

227 if self.is_failed:

228 raise RuntimeError("Executor failed.")

@@ -252,8 +253,8 @@ class MultiprocExecutor(Executor):

253

254 def get_response(

255 w: WorkerProcHandle,

255 - dequeue_timeout: Optional[float] = None,

256 - cancel_event: Optional[threading.Event] = None,

256 + dequeue_timeout: float | None = None,

257 + cancel_event: threading.Event | None = None,

258 ):

259 status, result = w.worker_response_mq.dequeue(

260 timeout=dequeue_timeout, cancel=cancel_event

@@ -370,7 +371,7 @@ class UnreadyWorkerProcHandle:

371 proc: BaseProcess

372 rank: int

373 ready_pipe: Connection

373 - death_writer: Optional[Connection] = None

374 + death_writer: Connection | None = None

375

376

377 @dataclass

@@ -378,7 +379,7 @@ class WorkerProcHandle:

379 proc: BaseProcess

380 rank: int

381 worker_response_mq: MessageQueue # The worker process writes to this MQ

381 - death_writer: Optional[Connection] = None

382 + death_writer: Connection | None = None

383

384 @classmethod

385 def from_unready_handle(

@@ -505,7 +506,7 @@ class WorkerProc:

506 )

507

508 pipes = {handle.ready_pipe: handle for handle in unready_proc_handles}

508 - ready_proc_handles: list[Optional[WorkerProcHandle]] = [None] * len(

509 + ready_proc_handles: list[WorkerProcHandle | None] = [None] * len(

510 unready_proc_handles

511 )

512 while pipes:

@@ -674,7 +675,7 @@ class WorkerProc:

675 output = self.async_output_queue.get()

676 self.enqueue_output(output)

677

677 - def worker_busy_loop(self, cancel: Optional[threading.Event] = None):

678 + def worker_busy_loop(self, cancel: threading.Event | None = None):

679 """Main busy loop for Multiprocessing Workers"""

680 while True:

681 method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue(

vllm/v1/executor/ray_distributed_executor.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from concurrent.futures import Future

5 -from typing import Optional, Union

5

6 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator

7 from vllm.executor.ray_distributed_executor import ( # noqa

@@ -26,7 +25,7 @@ class FutureWrapper(Future):

25 the result() call. If not only the first worker's output is returned.

26 """

27

29 - def __init__(self, refs, aggregator: Optional[KVOutputAggregator] = None):

28 + def __init__(self, refs, aggregator: KVOutputAggregator | None = None):

29 super().__init__()

30 self.refs = refs

31 self.aggregator = aggregator

@@ -66,7 +65,7 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):

65 self,

66 scheduler_output: SchedulerOutput,

67 non_block: bool = False,

69 - ) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:

68 + ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:

69 """Execute the model on the Ray workers.

70

71 Args:

vllm/v1/kv_cache_interface.py

+6 -7

@@ -4,7 +4,6 @@

4 import copy

5 from dataclasses import dataclass, fields

6 from math import prod

7 -from typing import Optional

7

8 import torch

9 from typing_extensions import Self

@@ -74,8 +73,8 @@ class AttentionSpec(KVCacheSpec):

73

74 @dataclass(frozen=True)

75 class FullAttentionSpec(AttentionSpec):

77 - sliding_window: Optional[int] = None

78 - attention_chunk_size: Optional[int] = None

76 + sliding_window: int | None = None

77 + attention_chunk_size: int | None = None

78 """

79 When hybrid allocator is disabled and the model contains both full

80 attention layers and sliding window attention layers, sliding

@@ -96,7 +95,7 @@ class FullAttentionSpec(AttentionSpec):

95 return cdiv(max_model_len, self.block_size) * self.page_size_bytes

96

97 @classmethod

99 - def merge_window_sizes(cls, window_sizes: set[int]) -> Optional[int]:

98 + def merge_window_sizes(cls, window_sizes: set[int]) -> int | None:

99 if len(window_sizes) == 0:

100 return None

101 elif len(window_sizes) == 1:

@@ -154,7 +153,7 @@ class FullAttentionSpec(AttentionSpec):

153 @dataclass(frozen=True)

154 class MLAAttentionSpec(FullAttentionSpec):

155 # TODO(Lucas/Chen): less hacky way to do this

157 - cache_dtype_str: Optional[str] = None

156 + cache_dtype_str: str | None = None

157

158 @property

159 def page_size_bytes(self) -> int:

@@ -237,7 +236,7 @@ class SlidingWindowSpec(AttentionSpec):

236 class MambaSpec(KVCacheSpec):

237 shapes: tuple[tuple[int, ...], ...]

238 dtypes: tuple[torch.dtype]

240 - page_size_padded: Optional[int] = None

239 + page_size_padded: int | None = None

240 mamba_type: str = "mamba2"

241 num_speculative_blocks: int = 0

242

@@ -342,7 +341,7 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):

341 )

342

343 @classmethod

345 - def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Optional[Self]:

344 + def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None:

345 """

346 Return a SameTypeKVCacheSpecs object if all layers have the same type

347 of KV cache spec. Return None if not.

vllm/v1/kv_offload/abstract.py

+1 -2

@@ -30,7 +30,6 @@ The class provides the following primitives:

30 from abc import ABC, abstractmethod

31 from collections.abc import Iterable

32 from dataclasses import dataclass

33 -from typing import Optional

33

34 from vllm.v1.core.kv_cache_utils import BlockHash

35

@@ -122,7 +121,7 @@ class OffloadingManager(ABC):

121 @abstractmethod

122 def prepare_store(

123 self, block_hashes: Iterable[BlockHash]

125 - ) -> Optional[PrepareStoreOutput]:

124 + ) -> PrepareStoreOutput | None:

125 """

126 Prepare the given blocks to be offloaded.

127 The given blocks will be protected from eviction until

vllm/v1/kv_offload/cpu.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterator

4 -from typing import Optional

4

5 import torch

6

@@ -29,10 +28,10 @@ class CPUOffloadingSpec(OffloadingSpec):

28 self.num_cpu_blocks: int = num_cpu_blocks

29

30 # scheduler-side

32 - self._manager: Optional[OffloadingManager] = None

31 + self._manager: OffloadingManager | None = None

32

33 # worker-side

35 - self._handler: Optional[OffloadingHandler] = None

34 + self._handler: OffloadingHandler | None = None

35

36 def get_manager(self) -> OffloadingManager:

37 if not self._manager:

vllm/v1/kv_offload/factory.py

+2 -1

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import importlib

4 -from typing import TYPE_CHECKING, Callable

4 +from collections.abc import Callable

5 +from typing import TYPE_CHECKING

6

7 from vllm.logger import init_logger

8 from vllm.v1.kv_offload.spec import OffloadingSpec

vllm/v1/kv_offload/lru_manager.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections import OrderedDict

4 from collections.abc import Iterable

5 -from typing import Optional

5

6 from vllm.v1.core.kv_cache_utils import BlockHash

7 from vllm.v1.kv_offload.abstract import (

@@ -23,7 +22,7 @@ class LRUOffloadingManager(OffloadingManager):

22 self.backend: Backend = backend

23 # block_hash -> BlockStatus

24 self.blocks: OrderedDict[BlockHash, BlockStatus] = OrderedDict()

26 - self.events: Optional[list[OffloadingEvent]] = [] if enable_events else None

25 + self.events: list[OffloadingEvent] | None = [] if enable_events else None

26

27 def lookup(self, block_hashes: Iterable[BlockHash]) -> int:

28 hit_count = 0

@@ -57,7 +56,7 @@ class LRUOffloadingManager(OffloadingManager):

56

57 def prepare_store(

58 self, block_hashes: Iterable[BlockHash]

60 - ) -> Optional[PrepareStoreOutput]:

59 + ) -> PrepareStoreOutput | None:

60 # filter out blocks that are already stored

61 block_hashes_to_store = [

62 block_hash for block_hash in block_hashes if block_hash not in self.blocks

vllm/v1/metrics/loggers.py

+25 -30

@@ -4,9 +4,10 @@

4 import logging

5 import time

6 from abc import ABC, abstractmethod

7 -from typing import Callable, Optional, Union

7 +from collections.abc import Callable

8 +from typing import TypeAlias

9

9 -import prometheus_client

10 +from prometheus_client import Counter, Gauge, Histogram

11

12 from vllm.config import SupportsMetricsInfo, VllmConfig

13 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorLogging

@@ -40,9 +41,9 @@ class StatLoggerBase(ABC):

41 @abstractmethod

42 def record(

43 self,

43 - scheduler_stats: Optional[SchedulerStats],

44 - iteration_stats: Optional[IterationStats],

45 - mm_cache_stats: Optional[MultiModalCacheStats] = None,

44 + scheduler_stats: SchedulerStats | None,

45 + iteration_stats: IterationStats | None,

46 + mm_cache_stats: MultiModalCacheStats | None = None,

47 engine_idx: int = 0,

48 ): ...

49

@@ -93,9 +94,9 @@ class LoggingStatLogger(StatLoggerBase):

94

95 def record(

96 self,

96 - scheduler_stats: Optional[SchedulerStats],

97 - iteration_stats: Optional[IterationStats],

98 - mm_cache_stats: Optional[MultiModalCacheStats] = None,

97 + scheduler_stats: SchedulerStats | None,

98 + iteration_stats: IterationStats | None,

99 + mm_cache_stats: MultiModalCacheStats | None = None,

100 engine_idx: int = 0,

101 ):

102 """Log Stats to standard output."""

@@ -178,13 +179,13 @@ class LoggingStatLogger(StatLoggerBase):

179

180

181 class PrometheusStatLogger(StatLoggerBase):

181 - _gauge_cls = prometheus_client.Gauge

182 - _counter_cls = prometheus_client.Counter

183 - _histogram_cls = prometheus_client.Histogram

182 + _gauge_cls = Gauge

183 + _counter_cls = Counter

184 + _histogram_cls = Histogram

185 _spec_decoding_cls = SpecDecodingProm

186

187 def __init__(

187 - self, vllm_config: VllmConfig, engine_indexes: Optional[list[int]] = None

188 + self, vllm_config: VllmConfig, engine_indexes: list[int] | None = None

189 ):

190 if engine_indexes is None:

191 engine_indexes = [0]

@@ -368,9 +369,7 @@ class PrometheusStatLogger(StatLoggerBase):

369 counter_generation_tokens, engine_indexes, model_name

370 )

371

371 - self.counter_request_success: dict[

372 - FinishReason, dict[int, prometheus_client.Counter]

373 - ] = {}

372 + self.counter_request_success: dict[FinishReason, dict[int, Counter]] = {}

373 counter_request_success_base = self._counter_cls(

374 name="vllm:request_success",

375 documentation="Count of successfully processed requests.",

@@ -660,7 +659,7 @@ class PrometheusStatLogger(StatLoggerBase):

659

660 # TODO: This metric might be incorrect in case of using multiple

661 # api_server counts which uses prometheus mp.

663 - self.gauge_lora_info: Optional[prometheus_client.Gauge] = None

662 + self.gauge_lora_info: Gauge | None = None

663 if vllm_config.lora_config is not None:

664 if len(self.engine_indexes) > 1:

665 raise NotImplementedError("LoRA in DP mode is not supported yet.")

@@ -705,9 +704,9 @@ class PrometheusStatLogger(StatLoggerBase):

704

705 def record(

706 self,

708 - scheduler_stats: Optional[SchedulerStats],

709 - iteration_stats: Optional[IterationStats],

710 - mm_cache_stats: Optional[MultiModalCacheStats] = None,

707 + scheduler_stats: SchedulerStats | None,

708 + iteration_stats: IterationStats | None,

709 + mm_cache_stats: MultiModalCacheStats | None = None,

710 engine_idx: int = 0,

711 ):

712 """Log to prometheus."""

@@ -826,11 +825,7 @@ class PrometheusStatLogger(StatLoggerBase):

825 self.log_metrics_info("cache_config", self.vllm_config.cache_config)

826

827

829 -PromMetric = Union[

830 - prometheus_client.Gauge,

831 - prometheus_client.Counter,

832 - prometheus_client.Histogram,

833 -]

828 +PromMetric: TypeAlias = Gauge | Counter | Histogram

829

830

831 def make_per_engine(

@@ -882,8 +877,8 @@ class StatLoggerManager:

877 def __init__(

878 self,

879 vllm_config: VllmConfig,

885 - engine_idxs: Optional[list[int]] = None,

886 - custom_stat_loggers: Optional[list[StatLoggerFactory]] = None,

880 + engine_idxs: list[int] | None = None,

881 + custom_stat_loggers: list[StatLoggerFactory] | None = None,

882 enable_default_loggers: bool = True,

883 client_count: int = 1,

884 ):

@@ -924,10 +919,10 @@ class StatLoggerManager:

919

920 def record(

921 self,

927 - scheduler_stats: Optional[SchedulerStats],

928 - iteration_stats: Optional[IterationStats],

929 - mm_cache_stats: Optional[MultiModalCacheStats] = None,

930 - engine_idx: Optional[int] = None,

922 + scheduler_stats: SchedulerStats | None,

923 + iteration_stats: IterationStats | None,

924 + mm_cache_stats: MultiModalCacheStats | None = None,

925 + engine_idx: int | None = None,

926 ):

927 if engine_idx is None:

928 engine_idx = 0

vllm/v1/metrics/prometheus.py

+1 -2

@@ -3,7 +3,6 @@

3

4 import os

5 import tempfile

6 -from typing import Optional

6

7 from prometheus_client import REGISTRY, CollectorRegistry, multiprocess

8

@@ -12,7 +11,7 @@ from vllm.logger import init_logger

11 logger = init_logger(__name__)

12

13 # Global temporary directory for prometheus multiprocessing

15 -_prometheus_multiproc_dir: Optional[tempfile.TemporaryDirectory] = None

14 +_prometheus_multiproc_dir: tempfile.TemporaryDirectory | None = None

15

16

17 def setup_multiprocess_prometheus():

vllm/v1/metrics/ray_wrappers.py

+11 -12

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import time

4 -from typing import Optional, Union

4

5 from vllm.v1.metrics.loggers import PrometheusStatLogger

6 from vllm.v1.spec_decode.metrics import SpecDecodingProm

@@ -63,9 +62,9 @@ class RayGaugeWrapper(RayPrometheusMetric):

62 def __init__(

63 self,

64 name: str,

66 - documentation: Optional[str] = "",

67 - labelnames: Optional[list[str]] = None,

68 - multiprocess_mode: Optional[str] = "",

65 + documentation: str | None = "",

66 + labelnames: list[str] | None = None,

67 + multiprocess_mode: str | None = "",

68 ):

69 # All Ray metrics are keyed by WorkerId, so multiprocess modes like

70 # "mostrecent", "all", "sum" do not apply. This logic can be manually

@@ -77,7 +76,7 @@ class RayGaugeWrapper(RayPrometheusMetric):

76 name=name, description=documentation, tag_keys=labelnames_tuple

77 )

78

80 - def set(self, value: Union[int, float]):

79 + def set(self, value: int | float):

80 return self.metric.set(value)

81

82 def set_to_current_time(self):

@@ -92,8 +91,8 @@ class RayCounterWrapper(RayPrometheusMetric):

91 def __init__(

92 self,

93 name: str,

95 - documentation: Optional[str] = "",

96 - labelnames: Optional[list[str]] = None,

94 + documentation: str | None = "",

95 + labelnames: list[str] | None = None,

96 ):

97 labelnames_tuple = tuple(labelnames) if labelnames else None

98 name = self._get_sanitized_opentelemetry_name(name)

@@ -101,7 +100,7 @@ class RayCounterWrapper(RayPrometheusMetric):

100 name=name, description=documentation, tag_keys=labelnames_tuple

101 )

102

104 - def inc(self, value: Union[int, float] = 1.0):

103 + def inc(self, value: int | float = 1.0):

104 if value == 0:

105 return

106 return self.metric.inc(value)

@@ -114,9 +113,9 @@ class RayHistogramWrapper(RayPrometheusMetric):

113 def __init__(

114 self,

115 name: str,

117 - documentation: Optional[str] = "",

118 - labelnames: Optional[list[str]] = None,

119 - buckets: Optional[list[float]] = None,

116 + documentation: str | None = "",

117 + labelnames: list[str] | None = None,

118 + buckets: list[float] | None = None,

119 ):

120 labelnames_tuple = tuple(labelnames) if labelnames else None

121 name = self._get_sanitized_opentelemetry_name(name)

@@ -128,7 +127,7 @@ class RayHistogramWrapper(RayPrometheusMetric):

127 boundaries=boundaries,

128 )

129

131 - def observe(self, value: Union[int, float]):

130 + def observe(self, value: int | float):

131 return self.metric.observe(value)

132

133

vllm/v1/metrics/reader.py

+1 -2

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 from prometheus_client import REGISTRY

7 from prometheus_client import Metric as PromMetric

@@ -144,7 +143,7 @@ def get_metrics_snapshot() -> list[Metric]:

143 return collected

144

145

147 -def _get_samples(metric: PromMetric, suffix: Optional[str] = None) -> list[Sample]:

146 +def _get_samples(metric: PromMetric, suffix: str | None = None) -> list[Sample]:

147 name = (metric.name + suffix) if suffix is not None else metric.name

148 return [s for s in metric.samples if s.name == name]

149

vllm/v1/metrics/stats.py

+11 -11

@@ -4,7 +4,7 @@

4 import time

5 from collections import deque

6 from dataclasses import dataclass, field

7 -from typing import TYPE_CHECKING, Any, Optional

7 +from typing import TYPE_CHECKING, Any

8

9 from vllm.v1.spec_decode.metrics import SpecDecodingStats

10

@@ -152,8 +152,8 @@ class SchedulerStats:

152

153 prefix_cache_stats: PrefixCacheStats = field(default_factory=PrefixCacheStats)

154

155 - spec_decoding_stats: Optional[SpecDecodingStats] = None

156 - kv_connector_stats: Optional[dict[str, Any]] = None

155 + spec_decoding_stats: SpecDecodingStats | None = None

156 + kv_connector_stats: dict[str, Any] | None = None

157

158 num_corrupted_reqs: int = 0

159

@@ -191,7 +191,7 @@ class FinishedRequestStats:

191 e2e_latency: float = 0.0

192 num_prompt_tokens: int = 0

193 num_generation_tokens: int = 0

194 - max_tokens_param: Optional[int] = None

194 + max_tokens_param: int | None = None

195 queued_time: float = 0.0

196 prefill_time: float = 0.0

197 inference_time: float = 0.0

@@ -230,7 +230,7 @@ class IterationStats:

230 is_prefilling: bool,

231 prompt_len: int,

232 req_stats: RequestStateStats,

233 - lora_stats: Optional[LoRAStats],

233 + lora_stats: LoRAStats | None,

234 ):

235 num_new_generation_tokens = len(output.new_token_ids)

236

@@ -265,7 +265,7 @@ class IterationStats:

265 events: list["EngineCoreEvent"],

266 is_prefilling: bool,

267 req_stats: RequestStateStats,

268 - lora_stats: Optional[LoRAStats],

268 + lora_stats: LoRAStats | None,

269 ):

270 # Avoid circular dependency

271 from vllm.v1.engine import EngineCoreEventType

@@ -287,7 +287,7 @@ class IterationStats:

287 self,

288 finish_reason: "FinishReason",

289 num_prompt_tokens: int,

290 - max_tokens_param: Optional[int],

290 + max_tokens_param: int | None,

291 req_stats: RequestStateStats,

292 ):

293 e2e_latency = self._time_since(req_stats.arrival_time)

@@ -335,7 +335,7 @@ class LoRARequestStates:

335 def __init__(self):

336 self.lora_name_to_stats: dict[str, LoRAStats] = {}

337

338 - def get_stats(self, req_state: "RequestState") -> Optional[LoRAStats]:

338 + def get_stats(self, req_state: "RequestState") -> LoRAStats | None:

339 if req_state.lora_name is None:

340 return None

341 if req_state.lora_name not in self.lora_name_to_stats:

@@ -362,20 +362,20 @@ class LoRARequestStates:

362 # Break the pattern for this lifecycle methods so we can

363 # call this from IterationStats.update_from_events()

364 @staticmethod

365 - def scheduled_request(lora_stats: Optional[LoRAStats], request_id: str):

365 + def scheduled_request(lora_stats: LoRAStats | None, request_id: str):

366 if lora_stats is None:

367 return

368 lora_stats.waiting_requests.remove(request_id)

369 lora_stats.running_requests.add(request_id)

370

371 @staticmethod

372 - def preempted_request(lora_stats: Optional[LoRAStats], request_id: str):

372 + def preempted_request(lora_stats: LoRAStats | None, request_id: str):

373 if lora_stats is None:

374 return

375 lora_stats.running_requests.remove(request_id)

376 lora_stats.waiting_requests.add(request_id)

377

378 - def update_iteration_stats(self, iteration_stats: Optional[IterationStats]):

378 + def update_iteration_stats(self, iteration_stats: IterationStats | None):

379 if iteration_stats is None:

380 return

381 for lora_name, stats in self.lora_name_to_stats.items():

vllm/v1/outputs.py

+13 -11

@@ -3,12 +3,14 @@

3

4 from abc import ABC, abstractmethod

5 from dataclasses import dataclass, field

6 -from typing import TYPE_CHECKING, NamedTuple, Optional, Union

6 +from typing import TYPE_CHECKING, NamedTuple

7

8 import torch

9

10 if TYPE_CHECKING:

11 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats

12 +else:

13 + KVConnectorStats = object

14

15

16 class LogprobsLists(NamedTuple):

@@ -64,7 +66,7 @@ class LogprobsTensors(NamedTuple):

66

67 # [num_reqs, <dynamic>]

68 # The shape of each element depends on the pooler used

67 -PoolerOutput = Union[torch.Tensor, list[torch.Tensor]]

69 +PoolerOutput = torch.Tensor | list[torch.Tensor]

70

71

72 @dataclass

@@ -74,15 +76,15 @@ class SamplerOutput:

76 # All requests are padded to max_num_generated_tokens.

77 # PLACEHOLDER_TOKEN_ID (-1 by default) is used for padding.

78 sampled_token_ids: torch.Tensor

77 - logprobs_tensors: Optional[LogprobsTensors]

79 + logprobs_tensors: LogprobsTensors | None

80

81

82 @dataclass

83 class KVConnectorOutput:

84 # [req_ids]

83 - finished_sending: Optional[set[str]] = None

84 - finished_recving: Optional[set[str]] = None

85 - kv_connector_stats: Optional["KVConnectorStats"] = None

85 + finished_sending: set[str] | None = None

86 + finished_recving: set[str] | None = None

87 + kv_connector_stats: KVConnectorStats | None = None

88 # IDs of externally computed KV blocks that failed to load.

89 # Requests referencing these blocks should be rescheduled to recompute them.

90 invalid_block_ids: set[int] = field(default_factory=set)

@@ -114,21 +116,21 @@ class ModelRunnerOutput:

116 # [num_reqs, max_num_logprobs + 1]

117 # [num_reqs, max_num_logprobs + 1]

118 # [num_reqs]

117 - logprobs: Optional[LogprobsLists]

119 + logprobs: LogprobsLists | None

120

121 # req_id -> (token_ids, logprobs, ranks)

122 # [prompt_len, num_prompt_logprobs]

123 # [prompt_len, num_prompt_logprobs]

124 # [prompt_len]

123 - prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]]

125 + prompt_logprobs_dict: dict[str, LogprobsTensors | None]

126

127 # [num_reqs, hidden_size]

126 - pooler_output: list[Optional[torch.Tensor]]

128 + pooler_output: list[torch.Tensor | None]

129

128 - kv_connector_output: Optional[KVConnectorOutput] = None

130 + kv_connector_output: KVConnectorOutput | None = None

131

132 # req_id -> num_nans_in_logits

131 - num_nans_in_logits: Optional[dict[str, int]] = None

133 + num_nans_in_logits: dict[str, int] | None = None

134

135

136 # ModelRunnerOutput wrapper for async scheduling.

vllm/v1/pool/metadata.py

+2 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 -from typing import Optional

4

5 import torch

6

@@ -37,9 +36,9 @@ class PoolingMetadata:

36 """Tensors for pooling."""

37

38 prompt_lens: torch.Tensor # CPU Tensor

40 - prompt_token_ids: Optional[torch.Tensor]

39 + prompt_token_ids: torch.Tensor | None

40 pooling_params: list[PoolingParams]

42 - pooling_cursor: Optional[PoolingCursor] = None

41 + pooling_cursor: PoolingCursor | None = None

42

43 def __getitem__(self, indices: slice):

44 return PoolingMetadata(

vllm/v1/request.py

+22 -22

@@ -3,9 +3,9 @@

3

4 import enum

5 import time

6 -from collections.abc import Mapping

6 +from collections.abc import Callable, Mapping

7 from functools import partial

8 -from typing import TYPE_CHECKING, Any, Callable, Optional, Union

8 +from typing import TYPE_CHECKING, Any, Optional

9

10 import torch

11

@@ -31,20 +31,20 @@ class Request:

31 def __init__(

32 self,

33 request_id: str,

34 - prompt_token_ids: Optional[list[int]],

35 - sampling_params: Optional[SamplingParams],

36 - pooling_params: Optional[PoolingParams],

37 - eos_token_id: Optional[int],

34 + prompt_token_ids: list[int] | None,

35 + sampling_params: SamplingParams | None,

36 + pooling_params: PoolingParams | None,

37 + eos_token_id: int | None,

38 client_index: int = 0,

39 - arrival_time: Optional[float] = None,

40 - prompt_embeds: Optional[torch.Tensor] = None,

41 - mm_features: Optional[list[MultiModalFeatureSpec]] = None,

39 + arrival_time: float | None = None,

40 + prompt_embeds: torch.Tensor | None = None,

41 + mm_features: list[MultiModalFeatureSpec] | None = None,

42 lora_request: Optional["LoRARequest"] = None,

43 structured_output_request: Optional["StructuredOutputRequest"] = None,

44 - cache_salt: Optional[str] = None,

44 + cache_salt: str | None = None,

45 priority: int = 0,

46 - trace_headers: Optional[Mapping[str, str]] = None,

47 - block_hasher: Optional[Callable[["Request"], list["BlockHash"]]] = None,

46 + trace_headers: Mapping[str, str] | None = None,

47 + block_hasher: Callable[["Request"], list["BlockHash"]] | None = None,

48 ) -> None:

49 self.request_id = request_id

50 self.client_index = client_index

@@ -60,10 +60,10 @@ class Request:

60 self.status = RequestStatus.WAITING

61 self.use_structured_output = False

62 self.events: list[EngineCoreEvent] = []

63 - self.stop_reason: Union[int, str, None] = None

63 + self.stop_reason: int | str | None = None

64

65 # P/D: Connector-specific KV transfer parameters.

66 - self.kv_transfer_params: Optional[dict[str, Any]] = None

66 + self.kv_transfer_params: dict[str, Any] | None = None

67

68 if pooling_params is not None:

69 # Pooling models.

@@ -97,7 +97,7 @@ class Request:

97 self.num_output_placeholders = 0 # Used in async scheduling.

98 self.spec_token_ids: list[int] = []

99 self.num_computed_tokens = 0

100 - self.cache_salt: Optional[str] = cache_salt

100 + self.cache_salt: str | None = cache_salt

101

102 # Multi-modal related

103 self.mm_features = mm_features or []

@@ -123,7 +123,7 @@ class Request:

123 self.num_preemptions = 0

124

125 self.block_hashes: list[BlockHash] = []

126 - self.get_hash_new_full_blocks: Optional[Callable[[], list[BlockHash]]] = None

126 + self.get_hash_new_full_blocks: Callable[[], list[BlockHash]] | None = None

127 if block_hasher is not None:

128 self.get_hash_new_full_blocks = partial(block_hasher, self)

129 self.block_hashes = self.get_hash_new_full_blocks()

@@ -132,7 +132,7 @@ class Request:

132 def from_engine_core_request(

133 cls,

134 request: EngineCoreRequest,

135 - block_hasher: Optional[Callable[["Request"], list["BlockHash"]]],

135 + block_hasher: Callable[["Request"], list["BlockHash"]] | None,

136 ) -> "Request":

137 return cls(

138 request_id=request.request_id,

@@ -158,7 +158,7 @@ class Request:

158

159 def append_output_token_ids(

160 self,

161 - token_ids: Union[int, list[int]],

161 + token_ids: int | list[int],

162 ) -> None:

163 if isinstance(token_ids, int):

164 self._output_token_ids.append(token_ids)

@@ -189,7 +189,7 @@ class Request:

189 def is_finished(self) -> bool:

190 return RequestStatus.is_finished(self.status)

191

192 - def get_finished_reason(self) -> Union[FinishReason, None]:

192 + def get_finished_reason(self) -> FinishReason | None:

193 return RequestStatus.get_finished_reason(self.status)

194

195 def get_num_encoder_tokens(self, input_id: int) -> int:

@@ -200,11 +200,11 @@ class Request:

200 def record_event(

201 self,

202 event_type: EngineCoreEventType,

203 - timestamp: Optional[float] = None,

203 + timestamp: float | None = None,

204 ) -> None:

205 self.events.append(EngineCoreEvent.new_event(event_type, timestamp))

206

207 - def take_events(self) -> Optional[list[EngineCoreEvent]]:

207 + def take_events(self) -> list[EngineCoreEvent] | None:

208 if not self.events:

209 return None

210 events, self.events = self.events, []

@@ -234,7 +234,7 @@ class RequestStatus(enum.IntEnum):

234 return status > RequestStatus.PREEMPTED

235

236 @staticmethod

237 - def get_finished_reason(status: "RequestStatus") -> Union[FinishReason, None]:

237 + def get_finished_reason(status: "RequestStatus") -> FinishReason | None:

238 return _FINISHED_REASON_MAP.get(status)

239

240

vllm/v1/sample/logits_processor/__init__.py

+8 -8

@@ -6,7 +6,7 @@ import itertools

6 from abc import abstractmethod

7 from collections.abc import Sequence

8 from functools import partial

9 -from typing import TYPE_CHECKING, Optional, Union

9 +from typing import TYPE_CHECKING

10

11 import torch

12

@@ -81,7 +81,7 @@ def _load_logitsprocs_plugins() -> list[type[LogitsProcessor]]:

81

82

83 def _load_logitsprocs_by_fqcns(

84 - logits_processors: Optional[Sequence[Union[str, type[LogitsProcessor]]]],

84 + logits_processors: Sequence[str | type[LogitsProcessor]] | None,

85 ) -> list[type[LogitsProcessor]]:

86 """Load logit processor types, identifying them by fully-qualified class

87 names (FQCNs).

@@ -146,7 +146,7 @@ def _load_logitsprocs_by_fqcns(

146

147

148 def _load_custom_logitsprocs(

149 - logits_processors: Optional[Sequence[Union[str, type[LogitsProcessor]]]],

149 + logits_processors: Sequence[str | type[LogitsProcessor]] | None,

150 ) -> list[type[LogitsProcessor]]:

151 """Load all custom logits processors.

152

@@ -176,7 +176,7 @@ def build_logitsprocs(

176 device: torch.device,

177 is_pin_memory: bool,

178 is_pooling_model: bool,

179 - custom_logitsprocs: Sequence[Union[str, type[LogitsProcessor]]] = (),

179 + custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = (),

180 ) -> LogitsProcessors:

181 if is_pooling_model:

182 if custom_logitsprocs:

@@ -249,7 +249,7 @@ class AdapterLogitsProcessor(LogitsProcessor):

249 def new_req_logits_processor(

250 self,

251 params: SamplingParams,

252 - ) -> Optional[RequestLogitsProcessor]:

252 + ) -> RequestLogitsProcessor | None:

253 """Consume request info; return a per-request logits processor.

254

255 Return None if logits processor does not need to be applied to request

@@ -267,9 +267,9 @@ class AdapterLogitsProcessor(LogitsProcessor):

267 def _new_state(

268 self,

269 params: SamplingParams,

270 - prompt_ids: Optional[list[int]],

270 + prompt_ids: list[int] | None,

271 output_ids: list[int],

272 - ) -> Optional[partial[torch.Tensor]]:

272 + ) -> partial[torch.Tensor] | None:

273 """Return state representation for new request

274

275 Returns None if logits processor is not applicable to request

@@ -292,7 +292,7 @@ class AdapterLogitsProcessor(LogitsProcessor):

292 return partial(req_lp, *args)

293 return None

294

295 - def update_state(self, batch_update: Optional[BatchUpdate]):

295 + def update_state(self, batch_update: BatchUpdate | None):

296 process_dict_updates(

297 self.req_info,

298 batch_update,

vllm/v1/sample/logits_processor/builtin.py

+9 -9

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from collections.abc import Sequence

4 -from typing import TYPE_CHECKING, Callable, Optional, TypeVar

3 +from collections.abc import Callable, Sequence

4 +from typing import TYPE_CHECKING, TypeVar

5

6 import torch

7

@@ -49,7 +49,7 @@ class MinPLogitsProcessor(LogitsProcessor):

49 def get_min_p_by_index(self, index: int) -> float:

50 return float(self.min_p_cpu[index])

51

52 - def update_state(self, batch_update: Optional[BatchUpdate]):

52 + def update_state(self, batch_update: BatchUpdate | None):

53 if not batch_update:

54 return

55

@@ -131,7 +131,7 @@ class LogitBiasLogitsProcessor(LogitsProcessor):

131 outcome of argmax in greedy sampling."""

132 return False

133

134 - def update_state(self, batch_update: Optional[BatchUpdate]):

134 + def update_state(self, batch_update: BatchUpdate | None):

135 needs_update = process_dict_updates(

136 self.biases, batch_update, lambda params, _, __: params.logit_bias or None

137 )

@@ -185,14 +185,14 @@ class MinTokensLogitsProcessor(LogitsProcessor):

185

186 @staticmethod

187 def add_request(

188 - params: SamplingParams, _: Optional[list[int]], output_tok_ids: list[int]

189 - ) -> Optional[tuple[int, Sequence[int], set[int]]]:

188 + params: SamplingParams, _: list[int] | None, output_tok_ids: list[int]

189 + ) -> tuple[int, Sequence[int], set[int]] | None:

190 min_tokens = params.min_tokens

191 if not min_tokens or len(output_tok_ids) >= min_tokens:

192 return None

193 return min_tokens, output_tok_ids, params.all_stop_token_ids

194

195 - def update_state(self, batch_update: Optional[BatchUpdate]):

195 + def update_state(self, batch_update: BatchUpdate | None):

196 needs_update = process_dict_updates(

197 self.min_toks, batch_update, self.add_request

198 )

@@ -235,8 +235,8 @@ class MinTokensLogitsProcessor(LogitsProcessor):

235

236 def process_dict_updates(

237 req_entries: dict[int, T],

238 - batch_update: Optional[BatchUpdate],

239 - new_state: Callable[[SamplingParams, Optional[list[int]], list[int]], Optional[T]],

238 + batch_update: BatchUpdate | None,

239 + new_state: Callable[[SamplingParams, list[int] | None, list[int]], T | None],

240 ) -> bool:

241 """Utility function to update dict state for sparse LogitsProcessors."""

242

vllm/v1/sample/logits_processor/interface.py

+1 -1

@@ -26,7 +26,7 @@ RemovedRequest = int

26

27 # (index, params, prompt_tok_ids, output_tok_ids) tuples for new

28 # requests added to the batch.

29 -AddedRequest = tuple[int, SamplingParams, Optional[list[int]], list[int]]

29 +AddedRequest = tuple[int, SamplingParams, list[int] | None, list[int]]

30

31 # (index 1, index 2, directionality) tuples representing

32 # one-way moves or two-way swaps of requests in batch

vllm/v1/sample/logits_processor/state.py

+8 -10

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterator

4 from itertools import chain

5 -from typing import TYPE_CHECKING, Optional

5 +from typing import TYPE_CHECKING

6

7 from vllm.v1.sample.logits_processor.interface import (

8 AddedRequest,

@@ -43,9 +43,9 @@ class BatchUpdateBuilder:

43

44 def __init__(

45 self,

46 - removed: Optional[list[RemovedRequest]] = None,

47 - added: Optional[list[AddedRequest]] = None,

48 - moved: Optional[list[MovedRequest]] = None,

46 + removed: list[RemovedRequest] | None = None,

47 + added: list[AddedRequest] | None = None,

48 + moved: list[MovedRequest] | None = None,

49 ) -> None:

50 self._removed = removed or []

51 self.added = added or []

@@ -92,14 +92,14 @@ class BatchUpdateBuilder:

92 def has_removed(self) -> bool:

93 return bool(self._removed)

94

95 - def peek_removed(self) -> Optional[int]:

95 + def peek_removed(self) -> int | None:

96 """Return lowest removed request index"""

97 if self.has_removed():

98 self._ensure_removed_sorted()

99 return self._removed[-1]

100 return None

101

102 - def pop_removed(self) -> Optional[int]:

102 + def pop_removed(self) -> int | None:

103 """Pop lowest removed request index"""

104 if self.has_removed():

105 self._ensure_removed_sorted()

@@ -116,7 +116,7 @@ class BatchUpdateBuilder:

116 self.batch_changed = False

117 return batch_changed

118

119 - def get_and_reset(self, batch_size: int) -> Optional[BatchUpdate]:

119 + def get_and_reset(self, batch_size: int) -> BatchUpdate | None:

120 """Generate a logitsprocs batch update data structure and reset

121 internal batch update builder state.

122

@@ -148,9 +148,7 @@ class BatchUpdateBuilder:

148 class LogitsProcessors:

149 """Encapsulates initialized logitsproc objects."""

150

151 - def __init__(

152 - self, logitsprocs: Optional[Iterator["LogitsProcessor"]] = None

153 - ) -> None:

151 + def __init__(self, logitsprocs: Iterator["LogitsProcessor"] | None = None) -> None:

152 self.argmax_invariant: list[LogitsProcessor] = []

153 self.non_argmax_invariant: list[LogitsProcessor] = []

154 if logitsprocs:

vllm/v1/sample/metadata.py

+7 -8

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from dataclasses import dataclass

5 -from typing import Optional

5

6 import torch

7

@@ -11,20 +10,20 @@ from vllm.v1.sample.logits_processor import LogitsProcessors

10

11 @dataclass

12 class SamplingMetadata:

14 - temperature: Optional[torch.Tensor]

13 + temperature: torch.Tensor | None

14 all_greedy: bool

15 all_random: bool

16

18 - top_p: Optional[torch.Tensor]

19 - top_k: Optional[torch.Tensor]

17 + top_p: torch.Tensor | None

18 + top_k: torch.Tensor | None

19

20 generators: dict[int, torch.Generator]

21

22 # None means no logprobs, 0 means sampled token logprobs only

24 - max_num_logprobs: Optional[int]

23 + max_num_logprobs: int | None

24

25 no_penalties: bool

27 - prompt_token_ids: Optional[torch.Tensor]

26 + prompt_token_ids: torch.Tensor | None

27 frequency_penalties: torch.Tensor

28 presence_penalties: torch.Tensor

29 repetition_penalties: torch.Tensor

@@ -33,7 +32,7 @@ class SamplingMetadata:

32

33 # `allowed_token_ids_mask` is a 2D bool tensor of shape (max batch size,

34 # vocab size).

36 - allowed_token_ids_mask: Optional[torch.Tensor]

35 + allowed_token_ids_mask: torch.Tensor | None

36

37 # req_index -> bad_words_token_ids

38 bad_words_token_ids: dict[int, list[list[int]]]

@@ -42,4 +41,4 @@ class SamplingMetadata:

41 logitsprocs: LogitsProcessors

42

43 # Speculative token ids

45 - spec_token_ids: Optional[list[list[int]]] = None

44 + spec_token_ids: list[list[int]] | None = None

vllm/v1/sample/ops/topk_topp_sampler.py

+13 -14

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import Optional

4

5 import torch

6 import torch.nn as nn

@@ -84,9 +83,9 @@ class TopKTopPSampler(nn.Module):

83 self,

84 logits: torch.Tensor,

85 generators: dict[int, torch.Generator],

87 - k: Optional[torch.Tensor],

88 - p: Optional[torch.Tensor],

89 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

86 + k: torch.Tensor | None,

87 + p: torch.Tensor | None,

88 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

89 """

90 PyTorch-native implementation of top-k and top-p sampling.

91

@@ -105,9 +104,9 @@ class TopKTopPSampler(nn.Module):

104 self,

105 logits: torch.Tensor,

106 generators: dict[int, torch.Generator],

108 - k: Optional[torch.Tensor],

109 - p: Optional[torch.Tensor],

110 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

107 + k: torch.Tensor | None,

108 + p: torch.Tensor | None,

109 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

110 """More optimized implementation for top-k and top-p sampling."""

111 # We prefer `random_sample` over `flashinfer_sample` when sorting is

112 # not needed. This is because `random_sample` does not require

@@ -132,9 +131,9 @@ class TopKTopPSampler(nn.Module):

131 self,

132 logits: torch.Tensor,

133 generators: dict[int, torch.Generator],

135 - k: Optional[torch.Tensor],

136 - p: Optional[torch.Tensor],

137 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

134 + k: torch.Tensor | None,

135 + p: torch.Tensor | None,

136 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

137 """

138 PyTorch-native implementation of top-k and top-p sampling for CPU.

139

@@ -170,8 +169,8 @@ class TopKTopPSampler(nn.Module):

169

170 def apply_top_k_top_p(

171 logits: torch.Tensor,

173 - k: Optional[torch.Tensor],

174 - p: Optional[torch.Tensor],

172 + k: torch.Tensor | None,

173 + p: torch.Tensor | None,

174 ) -> torch.Tensor:

175 """Apply top-k and top-p masks to the logits.

176

@@ -262,8 +261,8 @@ def random_sample(

261

262 def flashinfer_sample(

263 logits: torch.Tensor,

265 - k: Optional[torch.Tensor],

266 - p: Optional[torch.Tensor],

264 + k: torch.Tensor | None,

265 + p: torch.Tensor | None,

266 generators: dict[int, torch.Generator],

267 ) -> torch.Tensor:

268 """Sample from the logits using FlashInfer.

vllm/v1/sample/rejection_sampler.py

+4 -5

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Optional

3

4 import torch

5 import torch.nn as nn

@@ -49,7 +48,7 @@ class RejectionSampler(nn.Module):

48 self,

49 metadata: SpecDecodeMetadata,

50 # [num_tokens, vocab_size]

52 - draft_probs: Optional[torch.Tensor],

51 + draft_probs: torch.Tensor | None,

52 # [num_tokens, vocab_size]

53 target_logits: torch.Tensor,

54 # [batch_size, 1]

@@ -216,7 +215,7 @@ class RejectionSampler(nn.Module):

215 @staticmethod

216 def _combine_outputs_with_spec_tokens(

217 output_token_ids: list[list[int]],

219 - spec_token_ids: Optional[list[list[int]]] = None,

218 + spec_token_ids: list[list[int]] | None = None,

219 ) -> list[list[int]]:

220 if spec_token_ids is None:

221 return output_token_ids

@@ -240,7 +239,7 @@ def rejection_sample(

239 # [batch_size]

240 cu_num_draft_tokens: torch.Tensor,

241 # [num_tokens, vocab_size]

243 - draft_probs: Optional[torch.Tensor],

242 + draft_probs: torch.Tensor | None,

243 # [num_tokens, vocab_size]

244 target_probs: torch.Tensor,

245 # [batch_size, 1]

@@ -493,7 +492,7 @@ def sample_recovered_tokens(

492 # [num_tokens]

493 draft_token_ids: torch.Tensor,

494 # [num_tokens, vocab_size]

496 - draft_probs: Optional[torch.Tensor],

495 + draft_probs: torch.Tensor | None,

496 # [num_tokens, vocab_size]

497 target_probs: torch.Tensor,

498 sampling_metadata: SamplingMetadata,

vllm/v1/sample/sampler.py

+2 -4

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """A layer that samples the next tokens from the model's outputs."""

4

5 -from typing import Optional

6 -

5 import torch

6 import torch.nn as nn

7

@@ -140,7 +138,7 @@ class Sampler(nn.Module):

138 self,

139 logits: torch.Tensor,

140 sampling_metadata: SamplingMetadata,

143 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

141 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

142 """Sample logits based on sampling metadata.

143

144 The various logits processing functions called in this method

@@ -243,7 +241,7 @@ class Sampler(nn.Module):

241 @staticmethod

242 def _combine_outputs_with_spec_tokens(

243 output_token_ids: list[list[int]],

246 - spec_token_ids: Optional[list[list[int]]] = None,

244 + spec_token_ids: list[list[int]] | None = None,

245 ) -> list[list[int]]:

246 if spec_token_ids is None:

247 return output_token_ids

vllm/v1/sample/tpu/metadata.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass, field

4 -from typing import Optional

4

5 import torch

6

@@ -48,7 +47,7 @@ class TPUSupportedSamplingMetadata:

47

48 min_tokens = None # impl is not vectorized

49

51 - logit_bias: list[Optional[dict[int, float]]] = field(default_factory=lambda: list())

50 + logit_bias: list[dict[int, float] | None] = field(default_factory=lambda: list())

51

52 allowed_token_ids_mask = None

53 bad_words_token_ids = None

vllm/v1/sample/tpu/sampler.py

+2 -4

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Sampler layer implementing TPU supported operations."""

4

5 -from typing import Optional

6 -

5 import torch

6 import torch.nn as nn

7

@@ -166,8 +164,8 @@ class Sampler(nn.Module):

164

165 def apply_top_k_top_p(

166 logits: torch.Tensor,

169 - k: Optional[torch.Tensor],

170 - p: Optional[torch.Tensor],

167 + k: torch.Tensor | None,

168 + p: torch.Tensor | None,

169 ) -> torch.Tensor:

170 """

171 Apply top-k and top-p optimized for TPU.

vllm/v1/serial_utils.py

+10 -10

@@ -4,10 +4,10 @@

4 import dataclasses

5 import importlib

6 import pickle

7 -from collections.abc import Sequence

7 +from collections.abc import Callable, Sequence

8 from inspect import isclass

9 from types import FunctionType

10 -from typing import Any, Callable, Optional, Union

10 +from typing import Any, TypeAlias

11

12 import cloudpickle

13 import msgspec

@@ -47,7 +47,7 @@ MMF_CLASS_TO_FACTORY: dict[type[BaseMultiModalField], str] = {

47 MultiModalBatchedField: "batched",

48 }

49

50 -bytestr = Union[bytes, bytearray, memoryview, zmq.Frame]

50 +bytestr: TypeAlias = bytes | bytearray | memoryview | zmq.Frame

51

52

53 def _log_insecure_serialization_warning():

@@ -57,7 +57,7 @@ def _log_insecure_serialization_warning():

57 )

58

59

60 -def _typestr(val: Any) -> Optional[tuple[str, str]]:

60 +def _typestr(val: Any) -> tuple[str, str] | None:

61 if val is None:

62 return None

63 t = type(val)

@@ -111,14 +111,14 @@ class MsgpackEncoder:

111 via dedicated messages. Note that this is a per-tensor limit.

112 """

113

114 - def __init__(self, size_threshold: Optional[int] = None):

114 + def __init__(self, size_threshold: int | None = None):

115 if size_threshold is None:

116 size_threshold = envs.VLLM_MSGPACK_ZERO_COPY_THRESHOLD

117 self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)

118 # This is used as a local stash of buffers that we can then access from

119 # our custom `msgspec` hook, `enc_hook`. We don't have a way to

120 # pass custom data to the hook otherwise.

121 - self.aux_buffers: Optional[list[bytestr]] = None

121 + self.aux_buffers: list[bytestr] | None = None

122 self.size_threshold = size_threshold

123 if envs.VLLM_ALLOW_INSECURE_SERIALIZATION:

124 _log_insecure_serialization_warning()

@@ -195,7 +195,7 @@ class MsgpackEncoder:

195

196 def _encode_ndarray(

197 self, obj: np.ndarray

198 - ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:

198 + ) -> tuple[str, tuple[int, ...], int | memoryview]:

199 assert self.aux_buffers is not None

200 # If the array is non-contiguous, we need to copy it first

201 arr_data = obj.data if obj.flags.c_contiguous else obj.tobytes()

@@ -215,7 +215,7 @@ class MsgpackEncoder:

215

216 def _encode_tensor(

217 self, obj: torch.Tensor

218 - ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:

218 + ) -> tuple[str, tuple[int, ...], int | memoryview]:

219 assert self.aux_buffers is not None

220 # view the tensor as a contiguous 1D array of bytes

221 arr = obj.flatten().contiguous().view(torch.uint8).numpy()

@@ -280,7 +280,7 @@ class MsgpackDecoder:

280 not thread-safe when encoding tensors / numpy arrays.

281 """

282

283 - def __init__(self, t: Optional[Any] = None):

283 + def __init__(self, t: Any | None = None):

284 args = () if t is None else (t,)

285 self.decoder = msgpack.Decoder(

286 *args, ext_hook=self.ext_hook, dec_hook=self.dec_hook

@@ -289,7 +289,7 @@ class MsgpackDecoder:

289 if envs.VLLM_ALLOW_INSECURE_SERIALIZATION:

290 _log_insecure_serialization_warning()

291

292 - def decode(self, bufs: Union[bytestr, Sequence[bytestr]]) -> Any:

292 + def decode(self, bufs: bytestr | Sequence[bytestr]) -> Any:

293 if isinstance(bufs, bytestr): # type: ignore

294 return self.decoder.decode(bufs)

295

vllm/v1/spec_decode/eagle.py

+5 -6

@@ -3,7 +3,6 @@

3 import ast

4 from dataclasses import replace

5 from importlib.util import find_spec

6 -from typing import Optional

6

7 import numpy as np

8 import torch

@@ -79,8 +78,8 @@ class EagleProposer:

78 vllm_config.model_config

79 )

80

82 - self.attn_metadata_builder: Optional[AttentionMetadataBuilder] = None

83 - self.draft_indexer_metadata_builder: Optional[AttentionMetadataBuilder] = None

81 + self.attn_metadata_builder: AttentionMetadataBuilder | None = None

82 + self.draft_indexer_metadata_builder: AttentionMetadataBuilder | None = None

83 self.attn_layer_names: list[str] = []

84 self.indexer_layer_names: list[str] = []

85

@@ -149,7 +148,7 @@ class EagleProposer:

148 )

149

150 # Determine allowed attention backends once during initialization.

152 - self.allowed_attn_types: Optional[tuple] = None

151 + self.allowed_attn_types: tuple | None = None

152 if current_platform.is_rocm():

153 rocm_types = [TritonAttentionMetadata, FlashAttentionMetadata]

154 # vllm.v1.attention.backends.rocm_aiter_fa is an optional backend

@@ -207,10 +206,10 @@ class EagleProposer:

206 target_hidden_states: torch.Tensor,

207 # [batch_size]

208 next_token_ids: torch.Tensor,

210 - last_token_indices: Optional[torch.Tensor],

209 + last_token_indices: torch.Tensor | None,

210 common_attn_metadata: CommonAttentionMetadata,

211 sampling_metadata: SamplingMetadata,

213 - mm_embed_inputs: Optional[tuple[list[torch.Tensor], torch.Tensor]] = None,

212 + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

213 ) -> torch.Tensor:

214 num_tokens = target_token_ids.shape[0]

215 batch_size = next_token_ids.shape[0]

vllm/v1/spec_decode/metrics.py

+1 -2

@@ -3,7 +3,6 @@

3

4 import time

5 from dataclasses import dataclass, field

6 -from typing import Optional

6

7 import numpy as np

8 import prometheus_client

@@ -143,7 +142,7 @@ class SpecDecodingProm:

142

143 def __init__(

144 self,

146 - speculative_config: Optional[SpeculativeConfig],

145 + speculative_config: SpeculativeConfig | None,

146 labelnames: list[str],

147 per_engine_labelvalues: dict[int, list[str]],

148 ):

vllm/v1/structured_output/__init__.py

+4 -3

@@ -1,7 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from __future__ import annotations

4 -

3 import multiprocessing

4 from concurrent.futures import Future, ThreadPoolExecutor

5 from typing import TYPE_CHECKING

@@ -28,6 +26,9 @@ if TYPE_CHECKING:

26 else:

27 torch = LazyLoader("torch", globals(), "torch")

28

29 + ReasoningParser = object

30 + Request = object

31 +

32 logger = init_logger(__name__)

33

34

@@ -168,7 +169,7 @@ class StructuredOutputManager:

169 requests: dict[str, Request],

170 structured_output_request_ids: dict[str, int],

171 scheduled_spec_decode_tokens: dict[str, list[int]],

171 - ) -> npt.NDArray[np.int32] | None:

172 + ) -> "npt.NDArray[np.int32] | None":

173 # Prepare the structured output bitmask for this batch.

174 if not structured_output_request_ids:

175 return None

vllm/v1/structured_output/backend_guidance.py

+4 -6

@@ -1,13 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import copy

5 import json

6 import os

7 from dataclasses import dataclass

10 -from typing import TYPE_CHECKING, Any, Union

8 +from typing import TYPE_CHECKING, Any

9

10 import torch

11

@@ -47,7 +45,7 @@ def _walk_json_for_additional_properties(data: object):

45

46

47 def process_for_additional_properties(

50 - guide_json: Union[str, dict[str, Any]],

48 + guide_json: str | dict[str, Any],

49 ) -> dict[str, Any]:

50 if isinstance(guide_json, str):

51 guide_json_obj = json.loads(guide_json)

@@ -184,12 +182,12 @@ class GuidanceGrammar(StructuredOutputGrammar):

182

183 def serialize_guidance_grammar(

184 request_type: StructuredOutputOptions,

187 - grammar_spec: Union[str, dict[str, Any]],

185 + grammar_spec: str | dict[str, Any],

186 disable_any_whitespace: bool = False,

187 disable_additional_properties: bool = False,

188 ) -> str:

189 def _process_schema(

192 - grammar_spec: Union[str, dict[str, Any]],

190 + grammar_spec: str | dict[str, Any],

191 ) -> str:

192 if disable_additional_properties:

193 grammar_spec = process_for_additional_properties(grammar_spec)

vllm/v1/structured_output/backend_lm_format_enforcer.py

+1 -3

@@ -1,7 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from __future__ import annotations

4 -

3 import ast

4 import json

5 from dataclasses import dataclass, field

@@ -34,7 +32,7 @@ else:

32 @lru_cache

33 def _cached_build_vllm_token_enforcer_tokenizer_data(

34 tokenizer: PreTrainedTokenizerBase, vocab_size: int

37 -) -> lmfe_vllm.TokenEnforcerTokenizerData:

35 +) -> "lmfe_vllm.TokenEnforcerTokenizerData":

36 return lmfe_vllm.build_vllm_token_enforcer_tokenizer_data(

37 tokenizer, use_bitmask=True, vocab_size=vocab_size

38 )

vllm/v1/structured_output/backend_outlines.py

-2

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright 2025-present the Outlines developers

3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from __future__ import annotations

5 -

4 import ast

5 import importlib

6 import json

vllm/v1/structured_output/backend_types.py

+5 -4

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import enum

5 from abc import ABC, abstractmethod

6 from dataclasses import dataclass

@@ -13,6 +11,9 @@ if TYPE_CHECKING:

11

12 from vllm.config import VllmConfig

13 from vllm.transformers_utils.tokenizer import AnyTokenizer

14 +else:

15 + VllmConfig = object

16 + AnyTokenizer = object

17

18

19 class StructuredOutputOptions(enum.Enum):

@@ -69,7 +70,7 @@ class StructuredOutputGrammar(ABC):

70 """

71

72 @abstractmethod

72 - def fill_bitmask(self, bitmask: torch.Tensor, batch_index: int) -> None:

73 + def fill_bitmask(self, bitmask: "torch.Tensor", batch_index: int) -> None:

74 """

75 Fills the bitmask for a specific batch index.

76

@@ -119,7 +120,7 @@ class StructuredOutputBackend(ABC):

120 """

121

122 @abstractmethod

122 - def allocate_token_bitmask(self, max_num_seqs: int) -> torch.Tensor:

123 + def allocate_token_bitmask(self, max_num_seqs: int) -> "torch.Tensor":

124 """

125 Allocates a token bitmask for the specified maximum number of sequences.

126

vllm/v1/structured_output/backend_xgrammar.py

-2

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import json

5 from dataclasses import dataclass, field

6 from typing import TYPE_CHECKING, Any

vllm/v1/structured_output/request.py

+4 -10

@@ -1,13 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from __future__ import annotations

4 -

3 import dataclasses

4 import functools

5 import json

6 from concurrent.futures import Future

7 from concurrent.futures._base import TimeoutError

10 -from typing import Optional, Union, cast

8 +from typing import cast

9

10 from vllm.sampling_params import SamplingParams

11 from vllm.v1.structured_output.backend_types import (

@@ -20,9 +18,7 @@ from vllm.v1.structured_output.backend_types import (

18 @dataclasses.dataclass

19 class StructuredOutputRequest:

20 sampling_params: SamplingParams

23 - _grammar: Union[Future[StructuredOutputGrammar], StructuredOutputGrammar] | None = (

24 - None

25 - )

21 + _grammar: Future[StructuredOutputGrammar] | StructuredOutputGrammar | None = None

22 reasoning_ended: bool | None = None

23

24 def _check_grammar_completion(self) -> bool:

@@ -46,14 +42,12 @@ class StructuredOutputRequest:

42 def grammar(self) -> StructuredOutputGrammar | None:

43 completed = self._check_grammar_completion()

44 return (

49 - cast(Optional[StructuredOutputGrammar], self._grammar)

50 - if completed

51 - else None

45 + cast(StructuredOutputGrammar | None, self._grammar) if completed else None

46 )

47

48 @grammar.setter

49 def grammar(

56 - self, grammar: Union[StructuredOutputGrammar, Future[StructuredOutputGrammar]]

50 + self, grammar: StructuredOutputGrammar | Future[StructuredOutputGrammar]

51 ) -> None:

52 self._grammar = grammar

53

vllm/v1/structured_output/utils.py

+4 -3

@@ -1,8 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -from __future__ import annotations

5 -

3 import hashlib

4 import importlib.metadata

5 import os

@@ -37,6 +34,10 @@ else:

34 "transformers.models.gpt2.tokenization_gpt2",

35 )

36

37 + AnyTokenizer = object

38 + SchedulerOutput = object

39 + InputBatch = object

40 +

41 logger = init_logger(__name__)

42

43 CACHE = None

vllm/v1/utils.py

+10 -12

@@ -5,14 +5,13 @@ import contextlib

5 import multiprocessing

6 import time

7 import weakref

8 -from collections.abc import Sequence

8 +from collections.abc import Callable, Sequence

9 from contextlib import AbstractContextManager

10 from multiprocessing import connection

11 from multiprocessing.process import BaseProcess

12 from typing import (

13 TYPE_CHECKING,

14 Any,

15 - Callable,

15 Generic,

16 Optional,

17 TypeVar,

@@ -66,7 +65,7 @@ class ConstantList(Generic[T], Sequence):

65 def clear(self):

66 raise TypeError("Cannot clear a constant list")

67

69 - def index(self, item: T, start: int = 0, stop: Optional[int] = None) -> int:

68 + def index(self, item: T, start: int = 0, stop: int | None = None) -> int:

69 return self._x.index(item, start, stop if stop is not None else len(self._x))

70

71 @overload

@@ -75,7 +74,7 @@ class ConstantList(Generic[T], Sequence):

74 @overload

75 def __getitem__(self, s: slice, /) -> list[T]: ...

76

78 - def __getitem__(self, item: Union[int, slice]) -> Union[T, list[T]]:

77 + def __getitem__(self, item: int | slice) -> T | list[T]:

78 return self._x[item]

79

80 @overload

@@ -84,7 +83,7 @@ class ConstantList(Generic[T], Sequence):

83 @overload

84 def __setitem__(self, s: slice, value: T, /): ...

85

87 - def __setitem__(self, item: Union[int, slice], value: Union[T, list[T]]):

86 + def __setitem__(self, item: int | slice, value: T | list[T]):

87 raise TypeError("Cannot set item in a constant list")

88

89 def __delitem__(self, item):

@@ -108,7 +107,7 @@ class CpuGpuBuffer:

107

108 def __init__(

109 self,

111 - *size: Union[int, torch.SymInt],

110 + *size: int | torch.SymInt,

111 dtype: torch.dtype,

112 device: torch.device,

113 pin_memory: bool,

@@ -128,12 +127,12 @@ class CpuGpuBuffer:

127 )

128 self.np = self.cpu.numpy()

129

131 - def copy_to_gpu(self, n: Optional[int] = None) -> torch.Tensor:

130 + def copy_to_gpu(self, n: int | None = None) -> torch.Tensor:

131 if n is None:

132 return self.gpu.copy_(self.cpu, non_blocking=True)

133 return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)

134

136 - def copy_to_cpu(self, n: Optional[int] = None) -> torch.Tensor:

135 + def copy_to_cpu(self, n: int | None = None) -> torch.Tensor:

136 """NOTE: Because this method is non-blocking, explicit synchronization

137 is needed to ensure the data is copied to CPU."""

138 if n is None:

@@ -173,7 +172,7 @@ class APIServerProcessManager:

172 num_servers: int,

173 input_addresses: list[str],

174 output_addresses: list[str],

176 - stats_update_address: Optional[str] = None,

175 + stats_update_address: str | None = None,

176 ):

177 """Initialize and start API server worker processes.

178

@@ -227,9 +226,8 @@ class APIServerProcessManager:

226

227 def wait_for_completion_or_failure(

228 api_server_manager: APIServerProcessManager,

230 - engine_manager: Optional[

231 - Union["CoreEngineProcManager", "CoreEngineActorManager"]

232 - ] = None,

229 + engine_manager: Union["CoreEngineProcManager", "CoreEngineActorManager"]

230 + | None = None,

231 coordinator: Optional["DPCoordinator"] = None,

232 ) -> None:

233 """Wait for all processes to complete or detect if any fail.

vllm/v1/worker/block_table.py

+1 -2

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Union

3

4 import numpy as np

5 import torch

@@ -215,7 +214,7 @@ class BlockTable:

214 return self.block_table.np

215

216 def _make_buffer(

218 - self, *size: Union[int, torch.SymInt], dtype: torch.dtype

217 + self, *size: int | torch.SymInt, dtype: torch.dtype

218 ) -> CpuGpuBuffer:

219 return CpuGpuBuffer(

220 *size, dtype=dtype, device=self.device, pin_memory=self.pin_memory

vllm/v1/worker/cpu_model_runner.py

+2 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from contextlib import contextmanager

4 -from typing import TYPE_CHECKING, Any, Optional

4 +from typing import TYPE_CHECKING, Any

5

6 import torch

7 import torch.nn as nn

@@ -95,7 +95,7 @@ class CPUModelRunner(GPUModelRunner):

95 def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:

96 return sampled_token_ids.tolist()

97

98 - def get_dp_padding(self, num_tokens: int) -> tuple[int, Optional[torch.Tensor]]:

98 + def get_dp_padding(self, num_tokens: int) -> tuple[int, torch.Tensor | None]:

99 # Note: For CPU backend, dp padding is not required for now.

100 return 0, None

101

vllm/v1/worker/cpu_worker.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import os

4 import platform

5 -from typing import Callable, Optional

5 +from collections.abc import Callable

6

7 import torch

8

@@ -91,7 +91,7 @@ class CPUWorker(Worker):

91 logger.warning("sleep mode is not supported on CPU, ignore it.")

92 pass

93

94 - def wake_up(self, tags: Optional[list[str]] = None) -> None:

94 + def wake_up(self, tags: list[str] | None = None) -> None:

95 logger.warning("sleep mode is not supported on CPU, ignore it.")

96 pass

97

vllm/v1/worker/dp_utils.py

+5 -6

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Optional

3

4 import numpy as np

5 import torch

@@ -95,7 +94,7 @@ def _synchronize_dp_ranks(

94 should_attempt_ubatching: bool,

95 should_attempt_dp_padding: bool,

96 parallel_config: ParallelConfig,

98 -) -> tuple[bool, Optional[torch.Tensor]]:

97 +) -> tuple[bool, torch.Tensor | None]:

98 """

99 1. Decides if each DP rank is going to microbatch. Either all ranks

100 run with microbatching or none of them do.

@@ -156,10 +155,10 @@ def coordinate_batch_across_dp(

155 allow_microbatching: bool,

156 allow_dp_padding: bool,

157 parallel_config: ParallelConfig,

159 - num_tokens_padded: Optional[int] = None,

160 - uniform_decode: Optional[bool] = None,

161 - num_scheduled_tokens_per_request: Optional[np.ndarray] = None,

162 -) -> tuple[Optional[UBatchSlices], Optional[torch.Tensor]]:

158 + num_tokens_padded: int | None = None,

159 + uniform_decode: bool | None = None,

160 + num_scheduled_tokens_per_request: np.ndarray | None = None,

161 +) -> tuple[UBatchSlices | None, torch.Tensor | None]:

162 """

163 Coordinates amongst all DP ranks to determine if and how the full batch

164 should be split into microbatches.

vllm/v1/worker/gpu_input_batch.py

+22 -22

@@ -3,7 +3,7 @@

3 # Datastructures defining a GPU input batch

4

5 from dataclasses import dataclass

6 -from typing import Optional, cast

6 +from typing import cast

7

8 import numpy as np

9 import torch

@@ -29,21 +29,21 @@ from vllm.v1.worker.block_table import MultiGroupBlockTable

29 @dataclass

30 class CachedRequestState:

31 req_id: str

32 - prompt_token_ids: Optional[list[int]]

32 + prompt_token_ids: list[int] | None

33 mm_features: list[MultiModalFeatureSpec]

34 - sampling_params: Optional[SamplingParams]

35 - pooling_params: Optional[PoolingParams]

36 - generator: Optional[torch.Generator]

34 + sampling_params: SamplingParams | None

35 + pooling_params: PoolingParams | None

36 + generator: torch.Generator | None

37

38 block_ids: tuple[list[int], ...]

39 num_computed_tokens: int

40 output_token_ids: list[int]

41

42 - mrope_positions: Optional[torch.Tensor] = None

43 - mrope_position_delta: Optional[int] = None

42 + mrope_positions: torch.Tensor | None = None

43 + mrope_position_delta: int | None = None

44

45 - lora_request: Optional[LoRARequest] = None

46 - prompt_embeds: Optional[torch.Tensor] = None

45 + lora_request: LoRARequest | None = None

46 + prompt_embeds: torch.Tensor | None = None

47

48 def __post_init__(self):

49 self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

@@ -78,7 +78,7 @@ class InputBatch:

78 vocab_size: int,

79 block_sizes: list[int], # The block_size of each kv cache group

80 kernel_block_sizes: list[int],

81 - logitsprocs: Optional[LogitsProcessors] = None,

81 + logitsprocs: LogitsProcessors | None = None,

82 logitsprocs_need_output_token_ids: bool = False,

83 is_spec_decode: bool = False,

84 is_pooling_model: bool = False,

@@ -93,7 +93,7 @@ class InputBatch:

93 self.pin_memory = pin_memory

94 self.vocab_size = vocab_size

95

96 - self._req_ids: list[Optional[str]] = []

96 + self._req_ids: list[str | None] = []

97 self.req_id_to_index: dict[str, int] = {}

98

99 # TODO(woosuk): This buffer could be too large if max_model_len is big.

@@ -228,15 +228,15 @@ class InputBatch:

228 self.has_allowed_token_ids: set[str] = set()

229 # NOTE(lufang): In the mask tensor, if the corresponding token allowed,

230 # the value is False. Since we use masked_fill_ to set -inf.

231 - self.allowed_token_ids_mask: Optional[torch.Tensor] = None

232 - self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None

231 + self.allowed_token_ids_mask: torch.Tensor | None = None

232 + self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None

233

234 # req_index -> bad_words_token_ids

235 self.bad_words_token_ids: dict[int, list[list[int]]] = {}

236

237 self.logits_processing_needs_token_ids = np.zeros(max_num_reqs, dtype=bool)

238

239 - self.req_output_token_ids: list[Optional[list[int]]] = []

239 + self.req_output_token_ids: list[list[int] | None] = []

240

241 # Store provided logitsprocs. If none are provided, initialize empty

242 # data structure

@@ -244,7 +244,7 @@ class InputBatch:

244 self.logitsprocs_need_output_token_ids = logitsprocs_need_output_token_ids

245

246 # Store last speculative tokens for sampler.

247 - self.spec_token_ids: list[Optional[list[int]]] = []

247 + self.spec_token_ids: list[list[int] | None] = []

248

249 # This is updated each time the batch constituents change.

250 self.sampling_metadata = self._make_sampling_metadata()

@@ -252,13 +252,13 @@ class InputBatch:

252 self.pooling_params: dict[str, PoolingParams] = {}

253

254 # Cached reference to the GPU tensor of previously sampled tokens

255 - self.prev_sampled_token_ids: Optional[torch.Tensor] = None

256 - self.prev_req_id_to_index: Optional[dict[str, int]] = None

255 + self.prev_sampled_token_ids: torch.Tensor | None = None

256 + self.prev_req_id_to_index: dict[str, int] | None = None

257 # These are used to update output_token_ids with real sampled

258 # ids from prior step, if required by current sampling params

259 # (e.g. penalties).

260 - self.sampled_token_ids_cpu: Optional[torch.Tensor] = None

261 - self.async_copy_ready_event: Optional[torch.cuda.Event] = None

260 + self.sampled_token_ids_cpu: torch.Tensor | None = None

261 + self.async_copy_ready_event: torch.cuda.Event | None = None

262

263 @property

264 def req_ids(self) -> list[str]:

@@ -438,7 +438,7 @@ class InputBatch:

438

439 return req_index

440

441 - def remove_request(self, req_id: str) -> Optional[int]:

441 + def remove_request(self, req_id: str) -> int | None:

442 """This method must always be followed by a call to condense().

443

444 Args:

@@ -796,7 +796,7 @@ class InputBatch:

796 else []

797 )

798

799 - allowed_token_ids_mask: Optional[torch.Tensor] = None

799 + allowed_token_ids_mask: torch.Tensor | None = None

800 if not self.no_allowed_token_ids:

801 assert self.allowed_token_ids_mask is not None

802 copy_slice(

@@ -954,7 +954,7 @@ class InputBatch:

954 )

955

956 @property

957 - def max_num_logprobs(self) -> Optional[int]:

957 + def max_num_logprobs(self) -> int | None:

958 return max(self.num_logprobs.values()) if self.num_logprobs else None

959

960 @property

vllm/v1/worker/gpu_model_runner.py

+41 -42

@@ -8,14 +8,13 @@ from collections import defaultdict

8 from collections.abc import Iterator

9 from contextlib import contextmanager

10 from copy import deepcopy

11 -from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union, cast

11 +from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast

12

13 import numpy as np

14 import torch

15 import torch.distributed

16 import torch.nn as nn

17 from tqdm import tqdm

18 -from typing_extensions import TypeAlias

18

19 import vllm.envs as envs

20 from vllm.attention import Attention, AttentionType

@@ -162,7 +161,7 @@ logger = init_logger(__name__)

161

162 AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata]

163 # list when ubatching is enabled

165 -PerLayerAttnMetadata: TypeAlias = Union[list[AttnMetadataDict], AttnMetadataDict]

164 +PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict

165

166

167 # Wrapper for ModelRunnerOutput to support overlapped execution.

@@ -295,7 +294,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

294 # Sampler

295 self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)

296

298 - self.eplb_state: Optional[EplbState] = None

297 + self.eplb_state: EplbState | None = None

298 """

299 State of the expert parallelism load balancer.

300

@@ -447,14 +446,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

446

447 # CUDA event to synchronize use of reused CPU tensors between steps

448 # when async scheduling is enabled.

450 - self.prepare_inputs_event: Optional[torch.cuda.Event] = None

449 + self.prepare_inputs_event: torch.cuda.Event | None = None

450 if self.use_async_scheduling:

451 self.prepare_inputs_event = torch.cuda.Event()

452 # Start in a completed state.

453 self.prepare_inputs_event.record(torch.cuda.default_stream())

454

455 # None in the first PP rank. The rest are set after load_model.

457 - self.intermediate_tensors: Optional[IntermediateTensors] = None

456 + self.intermediate_tensors: IntermediateTensors | None = None

457

458 # OPTIMIZATION: Cache the tensors rather than creating them every step.

459 # Keep in int64 to avoid overflow with long context

@@ -495,7 +494,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

494 else None

495 )

496

498 - self.reorder_batch_threshold: Optional[int] = None

497 + self.reorder_batch_threshold: int | None = None

498

499 # Attention layers that are only in the KVCacheConfig of the runner

500 # (e.g., KV sharing, encoder-only attention), but not in the

@@ -503,7 +502,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

502 self.runner_only_attn_layers: set[str] = set()

503

504 # Cached outputs.

506 - self._draft_token_ids: Optional[Union[list[list[int]], torch.Tensor]] = None

505 + self._draft_token_ids: list[list[int]] | torch.Tensor | None = None

506 self.transfer_event = torch.cuda.Event()

507 self.sampled_token_ids_pinned_cpu = torch.empty(

508 (self.max_model_len, 1),

@@ -527,7 +526,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

526 return self.positions.gpu[num_tokens]

527

528 def _make_buffer(

530 - self, *size: Union[int, torch.SymInt], dtype: torch.dtype, numpy: bool = True

529 + self, *size: int | torch.SymInt, dtype: torch.dtype, numpy: bool = True

530 ) -> CpuGpuBuffer:

531 return CpuGpuBuffer(

532 *size,

@@ -928,7 +927,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

927 def _get_cumsum_and_arange(

928 self,

929 num_tokens: np.ndarray,

931 - cumsum_dtype: Optional[np.dtype] = None,

930 + cumsum_dtype: np.dtype | None = None,

931 ) -> tuple[np.ndarray, np.ndarray]:

932 """Get the cumulative sum and batched arange of the given array.

933 # E.g., [2, 5, 3] -> ([2, 7, 10], [0, 1, 0, 1, 2, 3, 4, 0, 1, 2])

@@ -1024,7 +1023,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1023 scheduler_output: "SchedulerOutput",

1024 kv_cache_spec: KVCacheSpec,

1025 num_reqs: int,

1027 - ) -> Optional[np.ndarray]:

1026 + ) -> np.ndarray | None:

1027 if not isinstance(kv_cache_spec, CrossAttentionSpec):

1028 return None

1029

@@ -1042,12 +1041,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1041 ) -> tuple[

1042 PerLayerAttnMetadata,

1043 torch.Tensor,

1045 - Optional[SpecDecodeMetadata],

1044 + SpecDecodeMetadata | None,

1045 np.ndarray,

1047 - Optional[CommonAttentionMetadata],

1046 + CommonAttentionMetadata | None,

1047 int,

1049 - Optional[UBatchSlices],

1050 - Optional[torch.Tensor],

1048 + UBatchSlices | None,

1049 + torch.Tensor | None,

1050 bool,

1051 ]:

1052 """

@@ -2048,7 +2047,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2047 )

2048 self._sync_device()

2049

2051 - pooler_output: list[Optional[torch.Tensor]] = []

2050 + pooler_output: list[torch.Tensor | None] = []

2051 for raw_output, seq_len, prompt_len in zip(

2052 raw_pooler_output, seq_lens_cpu, pooling_metadata.prompt_lens

2053 ):

@@ -2091,13 +2090,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2090 self,

2091 scheduler_output: "SchedulerOutput",

2092 num_input_tokens: int, # Padded

2094 - intermediate_tensors: Optional[IntermediateTensors] = None,

2093 + intermediate_tensors: IntermediateTensors | None = None,

2094 ) -> tuple[

2095 int,

2097 - Optional[torch.Tensor],

2098 - Optional[torch.Tensor],

2096 + torch.Tensor | None,

2097 + torch.Tensor | None,

2098 torch.Tensor,

2100 - Optional[IntermediateTensors],

2099 + IntermediateTensors | None,

2100 dict[str, Any],

2101 ]:

2102 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

@@ -2197,8 +2196,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2196

2197 def _sample(

2198 self,

2200 - logits: Optional[torch.Tensor],

2201 - spec_decode_metadata: Optional[SpecDecodeMetadata],

2199 + logits: torch.Tensor | None,

2200 + spec_decode_metadata: SpecDecodeMetadata | None,

2201 ) -> SamplerOutput:

2202 # Sample the next token and get logprobs if needed.

2203 sampling_metadata = self.input_batch.sampling_metadata

@@ -2243,14 +2242,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2242 self,

2243 scheduler_output: "SchedulerOutput",

2244 sampler_output: SamplerOutput,

2246 - logits: Optional[torch.Tensor],

2245 + logits: torch.Tensor | None,

2246 hidden_states: torch.Tensor,

2247 num_scheduled_tokens: int,

2248 ) -> tuple[

2249 dict[str, int],

2251 - Optional[LogprobsLists],

2250 + LogprobsLists | None,

2251 list[list[int]],

2253 - dict[str, Optional[LogprobsTensors]],

2252 + dict[str, LogprobsTensors | None],

2253 list[str],

2254 dict[str, int],

2255 list[int],

@@ -2377,10 +2376,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2376

2377 def _model_forward(

2378 self,

2380 - input_ids: Optional[torch.Tensor] = None,

2381 - positions: Optional[torch.Tensor] = None,

2382 - intermediate_tensors: Optional[IntermediateTensors] = None,

2383 - inputs_embeds: Optional[torch.Tensor] = None,

2379 + input_ids: torch.Tensor | None = None,

2380 + positions: torch.Tensor | None = None,

2381 + intermediate_tensors: IntermediateTensors | None = None,

2382 + inputs_embeds: torch.Tensor | None = None,

2383 **model_kwargs: dict[str, Any],

2384 ) -> Any:

2385 """Helper method to call the model forward pass.

@@ -2411,8 +2410,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2410 def execute_model(

2411 self,

2412 scheduler_output: "SchedulerOutput",

2414 - intermediate_tensors: Optional[IntermediateTensors] = None,

2415 - ) -> Union[ModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:

2413 + intermediate_tensors: IntermediateTensors | None = None,

2414 + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:

2415 with record_function_or_nullcontext("Preprocess"):

2416 with self.synchronize_input_prep():

2417 # Update persistent batch states.

@@ -2678,7 +2677,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2677

2678 return async_output

2679

2681 - def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

2680 + def take_draft_token_ids(self) -> DraftTokenIds | None:

2681 if self._draft_token_ids is None:

2682 return None

2683 req_ids = self.input_batch.req_ids

@@ -2692,14 +2691,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2691 def propose_draft_token_ids(

2692 self,

2693 scheduler_output: "SchedulerOutput",

2695 - sampled_token_ids: Union[torch.Tensor, list[list[int]]],

2694 + sampled_token_ids: torch.Tensor | list[list[int]],

2695 sampling_metadata: SamplingMetadata,

2696 hidden_states: torch.Tensor,

2697 sample_hidden_states: torch.Tensor,

2699 - aux_hidden_states: Optional[list[torch.Tensor]],

2700 - spec_decode_metadata: Optional[SpecDecodeMetadata],

2698 + aux_hidden_states: list[torch.Tensor] | None,

2699 + spec_decode_metadata: SpecDecodeMetadata | None,

2700 common_attn_metadata: CommonAttentionMetadata,

2702 - ) -> Union[list[list[int]], torch.Tensor]:

2701 + ) -> list[list[int]] | torch.Tensor:

2702 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

2703 if self.speculative_config.method == "ngram":

2704 assert isinstance(sampled_token_ids, list)

@@ -2961,7 +2960,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2960 self.model, self.vllm_config, CUDAGraphMode.NONE, self.device

2961 )

2962

2964 - def _get_eagle3_aux_layers_from_config(self) -> Optional[tuple[int, ...]]:

2963 + def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None:

2964 """Extract Eagle3 auxiliary layer indices from speculative config.

2965

2966 These indices specify which hidden states from the base model should

@@ -3007,13 +3006,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3006 self,

3007 hidden_states: torch.Tensor,

3008 num_scheduled_tokens: dict[str, int],

3010 - ) -> dict[str, Optional[LogprobsTensors]]:

3009 + ) -> dict[str, LogprobsTensors | None]:

3010 num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs

3011 if not num_prompt_logprobs_dict:

3012 return {}

3013

3014 in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu

3016 - prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]] = {}

3015 + prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}

3016

3017 # Since prompt logprobs are a rare feature, prioritize simple,

3018 # maintainable loop over optimal performance.

@@ -3107,7 +3106,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3106

3107 def _get_nans_in_logits(

3108 self,

3110 - logits: Optional[torch.Tensor],

3109 + logits: torch.Tensor | None,

3110 ) -> dict[str, int]:

3111 try:

3112 if logits is None:

@@ -3190,7 +3189,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3189 def _dummy_run(

3190 self,

3191 num_tokens: int,

3193 - cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,

3192 + cudagraph_runtime_mode: CUDAGraphMode | None = None,

3193 force_attention: bool = False,

3194 uniform_decode: bool = False,

3195 allow_microbatching: bool = True,

@@ -3294,7 +3293,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3293 dp_rank = self.parallel_config.data_parallel_rank

3294 num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])

3295

3297 - attn_metadata: Optional[PerLayerAttnMetadata] = None

3296 + attn_metadata: PerLayerAttnMetadata | None = None

3297

3298 # If force_attention is True, we always capture attention. Otherwise,

3299 # it only happens for cudagraph_runtime_mode=FULL.

vllm/v1/worker/gpu_ubatch_wrapper.py

+5 -4

@@ -2,8 +2,9 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import threading

5 +from collections.abc import Callable

6 from dataclasses import dataclass

6 -from typing import Any, Callable, Optional

7 +from typing import Any

8

9 import torch

10

@@ -32,8 +33,8 @@ class UbatchMetadata:

33 context: UBatchContext

34 input_ids: torch.Tensor

35 positions: torch.Tensor

35 - inputs_embeds: Optional[torch.Tensor]

36 - intermediate_tensors: Optional[IntermediateTensors]

36 + inputs_embeds: torch.Tensor | None

37 + intermediate_tensors: IntermediateTensors | None

38 num_tokens: int

39

40

@@ -41,7 +42,7 @@ class UbatchMetadata:

42 class CUDAGraphMetaData:

43 cudagraph: torch.cuda.CUDAGraph

44 ubatch_metadata: UbatchMetadata

44 - outputs: Optional[Any] = None

45 + outputs: Any | None = None

46

47

48 class SMControlContextManager:

vllm/v1/worker/gpu_worker.py

+9 -9

@@ -6,7 +6,7 @@ import copy

6 import gc

7 import os

8 from contextlib import AbstractContextManager, nullcontext

9 -from typing import TYPE_CHECKING, Any, Optional, Union

9 +from typing import TYPE_CHECKING, Any

10

11 import torch

12 import torch.distributed

@@ -132,7 +132,7 @@ class Worker(WorkerBase):

132 used_bytes / GiB_bytes,

133 )

134

135 - def wake_up(self, tags: Optional[list[str]] = None) -> None:

135 + def wake_up(self, tags: list[str] | None = None) -> None:

136 from vllm.device_allocator.cumem import CuMemAllocator

137

138 allocator = CuMemAllocator.get_instance()

@@ -456,7 +456,7 @@ class Worker(WorkerBase):

456 def execute_model(

457 self,

458 scheduler_output: "SchedulerOutput",

459 - ) -> Optional[Union[ModelRunnerOutput, AsyncModelRunnerOutput]]:

459 + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:

460 intermediate_tensors = None

461 forward_pass = scheduler_output.total_num_scheduled_tokens > 0

462 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

@@ -504,7 +504,7 @@ class Worker(WorkerBase):

504 output.kv_connector_output = kv_connector_output

505 return output

506

507 - def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

507 + def take_draft_token_ids(self) -> DraftTokenIds | None:

508 return self.model_runner.take_draft_token_ids()

509

510 def profile(self, is_start: bool = True):

@@ -565,7 +565,7 @@ class Worker(WorkerBase):

565 self,

566 old_ep_size: int,

567 new_ep_size: int,

568 - global_expert_load: Optional[torch.Tensor],

568 + global_expert_load: torch.Tensor | None,

569 ) -> None:

570 from vllm.distributed.parallel_state import get_ep_group

571

@@ -611,7 +611,7 @@ class Worker(WorkerBase):

611

612 def _reconfigure_moe(

613 self, old_ep_size: int, new_ep_size: int

614 - ) -> Optional[torch.Tensor]:

614 + ) -> torch.Tensor | None:

615 """

616 Reconfigure MoE modules with provided reconfig_request

617

@@ -730,8 +730,8 @@ class Worker(WorkerBase):

730 def save_sharded_state(

731 self,

732 path: str,

733 - pattern: Optional[str] = None,

734 - max_size: Optional[int] = None,

733 + pattern: str | None = None,

734 + max_size: int | None = None,

735 ) -> None:

736 from vllm.model_executor.model_loader import ShardedStateLoader

737

@@ -758,7 +758,7 @@ class Worker(WorkerBase):

758 def init_worker_distributed_environment(

759 vllm_config: VllmConfig,

760 rank: int,

761 - distributed_init_method: Optional[str] = None,

761 + distributed_init_method: str | None = None,

762 local_rank: int = -1,

763 backend: str = "nccl",

764 ) -> None:

vllm/v1/worker/kv_connector_model_runner_mixin.py

+3 -4

@@ -9,7 +9,6 @@ from collections.abc import Generator

9 from contextlib import AbstractContextManager, contextmanager, nullcontext

10 from typing import (

11 TYPE_CHECKING, # noqa: UP035

12 - Optional,

12 )

13

14 from vllm.config import VllmConfig

@@ -65,7 +64,7 @@ class KVConnectorModelRunnerMixin:

64 @staticmethod

65 def get_finished_kv_transfers(

66 scheduler_output: "SchedulerOutput",

68 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:

67 + ) -> tuple[set[str] | None, set[str] | None]:

68 if has_kv_transfer_group():

69 return get_kv_transfer_group().get_finished(

70 scheduler_output.finished_req_ids

@@ -95,7 +94,7 @@ class KVConnectorModelRunnerMixin:

94 @staticmethod

95 def maybe_get_kv_connector_output(

96 scheduler_output: "SchedulerOutput",

98 - ) -> AbstractContextManager[Optional[KVConnectorOutput]]:

97 + ) -> AbstractContextManager[KVConnectorOutput | None]:

98 return (

99 KVConnectorModelRunnerMixin._get_kv_connector_output(scheduler_output)

100 if has_kv_transfer_group()

@@ -139,7 +138,7 @@ class KVConnectorModelRunnerMixin:

138 kv_connector.clear_connector_metadata()

139

140 @staticmethod

142 - def get_kv_connector_stats() -> Optional[KVConnectorStats]:

141 + def get_kv_connector_stats() -> KVConnectorStats | None:

142 if has_kv_transfer_group():

143 return get_kv_transfer_group().get_kv_connector_stats()

144 return None

vllm/v1/worker/lora_model_runner_mixin.py

+5 -6

@@ -5,7 +5,6 @@ Define LoRA functionality mixin for model runners.

5 """

6

7 from contextlib import contextmanager

8 -from typing import Optional, Union

8

9 import numpy as np

10 import torch

@@ -21,7 +20,7 @@ from vllm.model_executor.models import supports_lora, supports_multimodal

20 from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch

21 from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch

22

24 -InputBatch = Union[TPUInputBatch, GPUInputBatch]

23 +InputBatch = TPUInputBatch | GPUInputBatch

24

25 logger = init_logger(__name__)

26

@@ -85,7 +84,7 @@ class LoRAModelRunnerMixin:

84

85 @contextmanager

86 def maybe_setup_dummy_loras(

88 - self, lora_config: Optional[LoRAConfig], remove_lora: bool = True

87 + self, lora_config: LoRAConfig | None, remove_lora: bool = True

88 ):

89 if lora_config is None:

90 yield

@@ -121,7 +120,7 @@ class LoRAModelRunnerMixin:

120

121 @contextmanager

122 def maybe_select_dummy_loras(

124 - self, lora_config: Optional[LoRAConfig], num_scheduled_tokens: np.ndarray

123 + self, lora_config: LoRAConfig | None, num_scheduled_tokens: np.ndarray

124 ):

125 if lora_config is None:

126 yield

@@ -158,7 +157,7 @@ class LoRAModelRunnerMixin:

157 @contextmanager

158 def maybe_dummy_run_with_lora(

159 self,

161 - lora_config: Optional[LoRAConfig],

160 + lora_config: LoRAConfig | None,

161 num_scheduled_tokens: np.ndarray,

162 remove_lora: bool = True,

163 ):

@@ -168,7 +167,7 @@ class LoRAModelRunnerMixin:

167 ):

168 yield

169

171 - def maybe_remove_all_loras(self, lora_config: Optional[LoRAConfig]):

170 + def maybe_remove_all_loras(self, lora_config: LoRAConfig | None):

171 if lora_config is None:

172 return

173 self.lora_manager.remove_all_adapters()

vllm/v1/worker/tpu_input_batch.py

+9 -9

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 # Datastructures defining a TPU input batch

4

5 -from typing import Optional, cast

5 +from typing import cast

6

7 import numpy as np

8 import torch

@@ -36,7 +36,7 @@ class InputBatch:

36 self.pin_memory = pin_memory

37 self.vocab_size = vocab_size

38

39 - self._req_ids: list[Optional[str]] = []

39 + self._req_ids: list[str | None] = []

40 self.req_id_to_index: dict[str, int] = {}

41

42 # TODO(woosuk): This buffer could be too large if max_model_len is big.

@@ -155,17 +155,17 @@ class InputBatch:

155 # To accumulate prompt logprobs tensor chunks across prefill steps.

156 self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

157

158 - self.logit_bias: list[Optional[dict[int, float]]] = [None] * max_num_reqs

158 + self.logit_bias: list[dict[int, float] | None] = [None] * max_num_reqs

159 self.has_allowed_token_ids: set[str] = set()

160 # NOTE(lufang): In the mask tensor, if the corresponding token allowed,

161 # the value is False. Since we use masked_fill_ to set -inf.

162 - self.allowed_token_ids_mask: Optional[torch.Tensor] = None

163 - self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None

162 + self.allowed_token_ids_mask: torch.Tensor | None = None

163 + self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None

164

165 # req_index -> bad_words_token_ids

166 self.bad_words_token_ids: dict[int, list[list[int]]] = {}

167

168 - self.req_output_token_ids: list[Optional[list[int]]] = []

168 + self.req_output_token_ids: list[list[int] | None] = []

169

170 @property

171 def req_ids(self) -> list[str]:

@@ -176,7 +176,7 @@ class InputBatch:

176 def add_request(

177 self,

178 request: "CachedRequestState",

179 - req_index: Optional[int] = None,

179 + req_index: int | None = None,

180 ) -> None:

181 if req_index is None:

182 req_index = self.num_reqs

@@ -296,7 +296,7 @@ class InputBatch:

296 # No LoRA

297 self.request_lora_mapping[req_index] = 0

298

299 - def remove_request(self, req_id: str) -> Optional[int]:

299 + def remove_request(self, req_id: str) -> int | None:

300 """This method must always be followed by a call to condense()."""

301

302 req_index = self.req_id_to_index.pop(req_id, None)

@@ -580,7 +580,7 @@ class InputBatch:

580 )

581

582 @property

583 - def max_num_logprobs(self) -> Optional[int]:

583 + def max_num_logprobs(self) -> int | None:

584 return max(self.num_logprobs.values()) if self.num_logprobs else None

585

586 @property

vllm/v1/worker/tpu_model_runner.py

+7 -7

@@ -3,7 +3,7 @@

3 import bisect

4 import gc

5 import time

6 -from typing import TYPE_CHECKING, Any, Optional, cast

6 +from typing import TYPE_CHECKING, Any, cast

7 from unittest.mock import patch

8

9 import numpy as np

@@ -140,7 +140,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

140 self,

141 vllm_config: VllmConfig,

142 device: torch.device,

143 - original_parallel_config: Optional[ParallelConfig] = None,

143 + original_parallel_config: ParallelConfig | None = None,

144 ):

145 self.vllm_config = vllm_config

146 self.model_config = vllm_config.model_config

@@ -1050,7 +1050,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1050 def _get_model_inputs(

1051 self,

1052 input_ids: torch.Tensor,

1053 - mm_embed_inputs: Optional[tuple[list[torch.Tensor], torch.Tensor]],

1053 + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,

1054 ):

1055 if self.supports_mm_inputs:

1056 mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

@@ -1076,7 +1076,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1076 def execute_model(

1077 self,

1078 scheduler_output: "SchedulerOutput",

1079 - intermediate_tensors: Optional[IntermediateTensors] = None,

1079 + intermediate_tensors: IntermediateTensors | None = None,

1080 ) -> ModelRunnerOutput:

1081 # Update cached state

1082 self._update_states(scheduler_output)

@@ -1220,7 +1220,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1220 ), "req_ids contains None"

1221 req_ids = cast(list[str], self.input_batch.req_ids[:num_reqs])

1222

1223 - prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]] = {}

1223 + prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}

1224 for req_id in self.input_batch.req_ids[:num_reqs]:

1225 prompt_logprobs_dict[req_id] = None

1226

@@ -2127,8 +2127,8 @@ def replace_set_lora(model):

2127 index: int,

2128 lora_a: torch.Tensor,

2129 lora_b: torch.Tensor,

2130 - embeddings_tensor: Optional[torch.Tensor],

2131 - bias: Optional[torch.Tensor] = None,

2130 + embeddings_tensor: torch.Tensor | None,

2131 + bias: torch.Tensor | None = None,

2132 ):

2133 # TODO: The integer index leads to a recompilation, but converting it

2134 # to a tensor doesn't seem to work anymore. This might be fixed with a

vllm/v1/worker/tpu_worker.py

+4 -3

@@ -3,7 +3,8 @@

3 """A TPU worker class."""

4

5 import os

6 -from typing import Any, Callable, Optional, TypeVar

6 +from collections.abc import Callable

7 +from typing import Any, TypeVar

8

9 import torch

10 import torch.distributed

@@ -257,7 +258,7 @@ class TPUWorker:

258 def execute_model(

259 self,

260 scheduler_output: "SchedulerOutput",

260 - ) -> Optional[ModelRunnerOutput]:

261 + ) -> ModelRunnerOutput | None:

262 output = self.model_runner.execute_model(scheduler_output)

263 # every worker's output is needed when kv_transfer_group is set up

264 return output if self.is_driver_worker or has_kv_transfer_group() else None

@@ -317,7 +318,7 @@ class TPUWorker:

318 self,

319 vllm_config: VllmConfig,

320 rank: int,

320 - distributed_init_method: Optional[str] = None,

321 + distributed_init_method: str | None = None,

322 local_rank: int = -1,

323 ) -> None:

324 """Initialize the distributed environment."""

vllm/v1/worker/ubatch_utils.py

+1 -1

@@ -1,9 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 +from typing import TypeAlias

5

6 import numpy as np

6 -from typing_extensions import TypeAlias

7

8 from vllm.config import ParallelConfig

9

vllm/v1/worker/utils.py

+5 -5

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections import defaultdict

4 from dataclasses import dataclass

5 -from typing import TYPE_CHECKING, Optional

5 +from typing import TYPE_CHECKING

6

7 import torch

8

@@ -193,7 +193,7 @@ def sanity_check_mm_encoder_outputs(

193

194 def scatter_mm_placeholders(

195 embeds: torch.Tensor,

196 - is_embed: Optional[torch.Tensor],

196 + is_embed: torch.Tensor | None,

197 ) -> torch.Tensor:

198 """

199 Scatter the multimodal embeddings into a contiguous tensor that represents

@@ -221,7 +221,7 @@ def scatter_mm_placeholders(

221

222 def gather_mm_placeholders(

223 placeholders: torch.Tensor,

224 - is_embed: Optional[torch.Tensor],

224 + is_embed: torch.Tensor | None,

225 ) -> torch.Tensor:

226 """

227 Reconstructs the embeddings from the placeholder tokens.

@@ -238,7 +238,7 @@ def gather_mm_placeholders(

238 def add_kv_sharing_layers_to_kv_cache_groups(

239 shared_kv_cache_layers: dict[str, str],

240 kv_cache_groups: list[KVCacheGroupSpec],

241 - runner_only_attn_layers: Optional[set[str]] = None,

241 + runner_only_attn_layers: set[str] | None = None,

242 ) -> None:

243 """

244 Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches`

@@ -270,7 +270,7 @@ def bind_kv_cache(

270 kv_caches: dict[str, torch.Tensor],

271 forward_context: dict[str, "Attention"],

272 runner_kv_caches: list[torch.Tensor],

273 - num_attn_module: Optional[int] = 1,

273 + num_attn_module: int | None = 1,

274 ) -> None:

275 """

276 Bind the allocated KV cache to both ModelRunner and forward context so

vllm/v1/worker/worker_base.py

+6 -4

@@ -1,10 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from __future__ import annotations

5 -

4 import os

7 -from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union

5 +from collections.abc import Callable

6 +from typing import TYPE_CHECKING, Any, TypeVar

7

8 import torch

9 import torch.nn as nn

@@ -26,6 +25,9 @@ from vllm.v1.kv_cache_interface import KVCacheSpec

25 if TYPE_CHECKING:

26 from vllm.v1.core.sched.output import SchedulerOutput

27 from vllm.v1.outputs import ModelRunnerOutput

28 +else:

29 + SchedulerOutput = object

30 + ModelRunnerOutput = object

31

32 logger = init_logger(__name__)

33

@@ -329,7 +331,7 @@ class WorkerWrapperBase:

331 # To make vLLM config available during device initialization

332 self.worker.init_device() # type: ignore

333

332 - def execute_method(self, method: Union[str, bytes], *args, **kwargs):

334 + def execute_method(self, method: str | bytes, *args, **kwargs):

335 try:

336 # method resolution order:

337 # if a method is defined in this class, it will be called directly.