@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4b2ed7926a

Fork 0

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

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

Synced Nov 15, 11:00 UTC

Improve configs - the rest! (#17562)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed May 9, 2025 at 23:18 UTC 4b2ed7926a1d93d4189ac112209f2e34cd80846a

14 files changed +461 -345

tests/compile/test_full_graph.py

+1 -4

@@ -9,7 +9,7 @@ import torch

10 from tests.quantization.utils import is_quant_method_supported

11 from vllm import LLM, SamplingParams

12 -from vllm.config import CompilationConfig, CompilationLevel

12 +from vllm.config import CompilationConfig, CompilationLevel, PassConfig

13 from vllm.platforms import current_platform

15 from ..utils import create_new_process_for_each_test

@@ -95,9 +95,6 @@ def test_full_graph(

95 run_model(optimization_level, model, model_kwargs)

98 -PassConfig = CompilationConfig.PassConfig

99 -

100 -

98 # TODO(luka) add other supported compilation config scenarios here

99 @pytest.mark.parametrize(

100 "compilation_config, model_info",

tests/compile/test_functionalization.py

+3 -4

@@ -11,7 +11,7 @@ from vllm.compilation.fusion import (FUSED_OPS, FusionPass, QuantKey,

11 kFp8DynamicTokenSym, kFp8StaticTensorSym)

12 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func

13 from vllm.compilation.noop_elimination import NoOpEliminationPass

14 -from vllm.config import CompilationConfig, VllmConfig

14 +from vllm.config import CompilationConfig, PassConfig, VllmConfig

16 from .backend import TestBackend

@@ -53,9 +53,8 @@ def test_fix_functionalization(model: str, quant_key: QuantKey,

53 torch.set_default_device("cuda")

55 vllm_config = VllmConfig()

56 - vllm_config.compilation_config = CompilationConfig(pass_config= \

57 - CompilationConfig.PassConfig(enable_fusion=do_fusion,

58 - enable_noop=True))

56 + vllm_config.compilation_config = CompilationConfig(

57 + pass_config=PassConfig(enable_fusion=do_fusion, enable_noop=True))

58 noop_pass = NoOpEliminationPass(vllm_config)

59 fusion_pass = FusionPass.instance(vllm_config)

60 act_quant_fusion_pass = ActivationQuantFusionPass(vllm_config)

tests/compile/test_fusion.py

+3 -3

@@ -9,7 +9,8 @@ from vllm.compilation.fusion import (FUSED_OPS, QUANT_OPS, FusedRMSQuantKey,

9 FusionPass, QuantKey)

10 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe

11 from vllm.compilation.noop_elimination import NoOpEliminationPass

12 -from vllm.config import CompilationConfig, CompilationLevel, VllmConfig

12 +from vllm.config import (CompilationConfig, CompilationLevel, PassConfig,

13 + VllmConfig)

14 from vllm.model_executor.layers.layernorm import RMSNorm

15 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

16 CUTLASS_FP8_SUPPORTED, Fp8LinearOp, maybe_create_device_identity)

@@ -78,8 +79,7 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,

79 vllm_config = VllmConfig(compilation_config=CompilationConfig(

80 level=CompilationLevel.PIECEWISE, custom_ops=["+rms_norm"]))

81 vllm_config.compilation_config.pass_config = \

81 - CompilationConfig.PassConfig(enable_fusion=True,

82 - enable_noop=True)

82 + PassConfig(enable_fusion=True, enable_noop=True)

83 with vllm.config.set_current_vllm_config(vllm_config):

84 # Reshape pass is needed for the fusion pass to work

85 noop_pass = NoOpEliminationPass(vllm_config)

tests/compile/test_sequence_parallelism.py

+3 -4

@@ -10,7 +10,7 @@ from vllm.compilation.fx_utils import (find_auto_fn, find_auto_fn_maybe,

10 find_specified_fn_maybe, is_func)

11 from vllm.compilation.sequence_parallelism import SequenceParallelismPass

12 from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,

13 - VllmConfig)

13 + PassConfig, VllmConfig)

14 from vllm.distributed import tensor_model_parallel_all_reduce

15 from vllm.distributed.parallel_state import (init_distributed_environment,

16 initialize_model_parallel)

@@ -126,9 +126,8 @@ def sequence_parallelism_pass_on_test_model(local_rank: int, world_size: int,

126

127 # configure vllm config for SequenceParallelismPass

128 vllm_config = VllmConfig()

129 - vllm_config.compilation_config = CompilationConfig(

130 - pass_config=CompilationConfig.PassConfig(

131 - enable_sequence_parallelism=True, ), )

129 + vllm_config.compilation_config = CompilationConfig(pass_config=PassConfig(

130 + enable_sequence_parallelism=True))

131 vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))

132

133 # this is a fake model name to construct the model config

tests/compile/test_silu_mul_quant_fusion.py

+2 -3

@@ -6,7 +6,7 @@ import vllm.envs as envs

6 from vllm._custom_ops import scaled_fp8_quant

7 from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass

8 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe

9 -from vllm.config import CompilationConfig, VllmConfig

9 +from vllm.config import CompilationConfig, PassConfig, VllmConfig

10 from vllm.model_executor.layers.activation import SiluAndMul

12 from .backend import TestBackend

@@ -36,8 +36,7 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size):

36 # Reshape pass is needed for the fusion pass to work

37 config = VllmConfig()

38 config.compilation_config = CompilationConfig(

39 - pass_config=CompilationConfig.PassConfig(enable_fusion=True,

40 - enable_reshape=True))

39 + pass_config=PassConfig(enable_fusion=True, enable_reshape=True))

40 fusion_pass = ActivationQuantFusionPass(config)

42 backend = TestBackend(fusion_pass)

tests/distributed/test_sequence_parallel.py

+2 -2

@@ -206,7 +206,7 @@ def _compare_sp(

206 'compile_sizes': [4, 8],

207 'splitting_ops': [],

208 'pass_config': {

209 - 'enable_sequence_parallism': sp_enabled,

209 + 'enable_sequence_parallelism': sp_enabled,

210 'enable_noop': True,

211 'enable_fusion': True,

212 },

@@ -223,7 +223,7 @@ def _compare_sp(

223 "--distributed-executor-backend",

224 distributed_backend,

225 "--compilation_config",

226 - str(compilation_config),

226 + json.dumps(compilation_config),

227 ]

228

229 tp_env = {

tests/engine/test_arg_utils.py

+61 -15

@@ -8,21 +8,18 @@ from typing import Literal, Optional

9 import pytest

11 -from vllm.config import config

11 +from vllm.config import CompilationConfig, config

12 from vllm.engine.arg_utils import (EngineArgs, contains_type, get_kwargs,

13 get_type, is_not_builtin, is_type,

14 literal_to_kwargs, nullable_kvs,

15 - optional_type)

15 + optional_type, parse_type)

16 from vllm.utils import FlexibleArgumentParser

19 @pytest.mark.parametrize(("type", "value", "expected"), [

20 (int, "42", 42),

21 - (int, "None", None),

21 (float, "3.14", 3.14),

23 - (float, "None", None),

22 (str, "Hello World!", "Hello World!"),

25 - (str, "None", None),

23 (json.loads, '{"foo":1,"bar":2}', {

24 "foo": 1,

25 "bar": 2

@@ -31,15 +28,20 @@ from vllm.utils import FlexibleArgumentParser

28 "foo": 1,

29 "bar": 2

30 }),

34 - (json.loads, "None", None),

31 ])

36 -def test_optional_type(type, value, expected):

37 - optional_type_func = optional_type(type)

32 +def test_parse_type(type, value, expected):

33 + parse_type_func = parse_type(type)

34 context = nullcontext()

35 if value == "foo=1,bar=2":

36 context = pytest.warns(DeprecationWarning)

37 with context:

42 - assert optional_type_func(value) == expected

38 + assert parse_type_func(value) == expected

39 +

40 +

41 +def test_optional_type():

42 + optional_type_func = optional_type(int)

43 + assert optional_type_func("None") is None

44 + assert optional_type_func("42") == 42

47 @pytest.mark.parametrize(("type_hint", "type", "expected"), [

@@ -89,7 +91,40 @@ def test_literal_to_kwargs(type_hints, expected):

92 @config

93 @dataclass

92 -class DummyConfigClass:

94 +class NestedConfig:

95 + field: int = 1

96 + """field"""

97 +

98 +

99 +@config

100 +@dataclass

101 +class FromCliConfig1:

102 + field: int = 1

103 + """field"""

104 +

105 + @classmethod

106 + def from_cli(cls, cli_value: str):

107 + inst = cls(**json.loads(cli_value))

108 + inst.field += 1

109 + return inst

110 +

111 +

112 +@config

113 +@dataclass

114 +class FromCliConfig2:

115 + field: int = 1

116 + """field"""

117 +

118 + @classmethod

119 + def from_cli(cls, cli_value: str):

120 + inst = cls(**json.loads(cli_value))

121 + inst.field += 2

122 + return inst

123 +

124 +

125 +@config

126 +@dataclass

127 +class DummyConfig:

128 regular_bool: bool = True

129 """Regular bool with default True"""

130 optional_bool: Optional[bool] = None

@@ -108,18 +143,24 @@ class DummyConfigClass:

143 """Literal of literals with default 1"""

144 json_tip: dict = field(default_factory=dict)

145 """Dict which will be JSON in CLI"""

146 + nested_config: NestedConfig = field(default_factory=NestedConfig)

147 + """Nested config"""

148 + from_cli_config1: FromCliConfig1 = field(default_factory=FromCliConfig1)

149 + """Config with from_cli method"""

150 + from_cli_config2: FromCliConfig2 = field(default_factory=FromCliConfig2)

151 + """Different config with from_cli method"""

152

153

154 @pytest.mark.parametrize(("type_hint", "expected"), [

155 (int, False),

115 - (DummyConfigClass, True),

156 + (DummyConfig, True),

157 ])

158 def test_is_not_builtin(type_hint, expected):

159 assert is_not_builtin(type_hint) == expected

160

161

162 def test_get_kwargs():

122 - kwargs = get_kwargs(DummyConfigClass)

163 + kwargs = get_kwargs(DummyConfig)

164 print(kwargs)

165

166 # bools should not have their type set

@@ -142,6 +183,11 @@ def test_get_kwargs():

183 # dict should have json tip in help

184 json_tip = "\n\nShould be a valid JSON string."

185 assert kwargs["json_tip"]["help"].endswith(json_tip)

186 + # nested config should should construct the nested config

187 + assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2)

188 + # from_cli configs should be constructed with the correct method

189 + assert kwargs["from_cli_config1"]["type"]('{"field": 2}').field == 3

190 + assert kwargs["from_cli_config2"]["type"]('{"field": 2}').field == 4

191

192

193 @pytest.mark.parametrize(("arg", "expected"), [

@@ -177,7 +223,7 @@ def test_compilation_config():

223

224 # default value

225 args = parser.parse_args([])

180 - assert args.compilation_config is None

226 + assert args.compilation_config == CompilationConfig()

227

228 # set to O3

229 args = parser.parse_args(["-O3"])

@@ -194,7 +240,7 @@ def test_compilation_config():

240 # set to string form of a dict

241 args = parser.parse_args([

242 "--compilation-config",

197 - "{'level': 3, 'cudagraph_capture_sizes': [1, 2, 4, 8]}",

243 + '{"level": 3, "cudagraph_capture_sizes": [1, 2, 4, 8]}',

244 ])

245 assert (args.compilation_config.level == 3 and

246 args.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8])

@@ -202,7 +248,7 @@ def test_compilation_config():

248 # set to string form of a dict

249 args = parser.parse_args([

250 "--compilation-config="

205 - "{'level': 3, 'cudagraph_capture_sizes': [1, 2, 4, 8]}",

251 + '{"level": 3, "cudagraph_capture_sizes": [1, 2, 4, 8]}',

252 ])

253 assert (args.compilation_config.level == 3 and

254 args.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8])

vllm/compilation/vllm_inductor_pass.py

+2 -5

@@ -4,7 +4,7 @@ import time

5 import torch

7 -from vllm.config import CompilationConfig, VllmConfig

7 +from vllm.config import PassConfig, VllmConfig

8 # yapf: disable

9 from vllm.distributed import get_tensor_model_parallel_rank as get_tp_rank

10 from vllm.distributed import (

@@ -56,10 +56,7 @@ class VllmInductorPass(InductorPass):

57 class PrinterInductorPass(VllmInductorPass):

59 - def __init__(self,

60 - name: str,

61 - config: CompilationConfig.PassConfig,

62 - always=False):

59 + def __init__(self, name: str, config: PassConfig, always=False):

60 super().__init__(config)

61 self.name = name

62 self.always = always

vllm/config.py

+303 -220

@@ -11,8 +11,8 @@ import textwrap

11 import warnings

12 from collections import Counter

13 from contextlib import contextmanager

14 -from dataclasses import (MISSING, dataclass, field, fields, is_dataclass,

15 - replace)

14 +from dataclasses import (MISSING, Field, asdict, dataclass, field, fields,

15 + is_dataclass, replace)

16 from functools import cached_property

17 from importlib.util import find_spec

18 from pathlib import Path

@@ -20,7 +20,6 @@ from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Literal, Optional,

20 Protocol, TypeVar, Union, cast, get_args, get_origin)

22 import torch

23 -from pydantic import BaseModel, Field, PrivateAttr

23 from torch.distributed import ProcessGroup, ReduceOp

24 from transformers import PretrainedConfig

25 from typing_extensions import deprecated

@@ -57,7 +56,7 @@ if TYPE_CHECKING:

57 ConfigType = type[DataclassInstance]

58 else:

60 - QuantizationConfig = None

59 + QuantizationConfig = Any

60 ConfigType = type

62 logger = init_logger(__name__)

@@ -169,6 +168,12 @@ def config(cls: ConfigT) -> ConfigT:

168 """

169 A decorator that ensures all fields in a dataclass have default values

170 and that each field has a docstring.

171 +

172 + If a `ConfigT` is used as a CLI argument itself, the default value provided

173 + by `get_kwargs` will be the result parsing a JSON string as the kwargs

174 + (i.e. `ConfigT(**json.loads(cli_arg))`). However, if a particular `ConfigT`

175 + requires custom construction from CLI (i.e. `CompilationConfig`), it can

176 + have a `from_cli` method, which will be called instead.

177 """

178 if not is_dataclass(cls):

179 raise TypeError("The decorated class must be a dataclass.")

@@ -202,7 +207,7 @@ def get_field(cls: ConfigType, name: str) -> Field:

207 cls_fields = {f.name: f for f in fields(cls)}

208 if name not in cls_fields:

209 raise ValueError(f"Field '{name}' not found in {cls.__name__}.")

205 - named_field: Field = cls_fields.get(name)

210 + named_field: Field = cls_fields[name]

211 if (default_factory := named_field.default_factory) is not MISSING:

212 return field(default_factory=default_factory)

213 if (default := named_field.default) is not MISSING:

@@ -211,6 +216,10 @@ def get_field(cls: ConfigType, name: str) -> Field:

216 f"{cls.__name__}.{name} must have a default value or default factory.")

217

218

219 +def is_init_field(cls: ConfigType, name: str) -> bool:

220 + return next(f for f in fields(cls) if f.name == name).init

221 +

222 +

223 TokenizerMode = Literal["auto", "slow", "mistral", "custom"]

224 ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]

225

@@ -2007,13 +2016,13 @@ class SchedulerConfig:

2016 def __post_init__(self) -> None:

2017 if self.max_model_len is None:

2018 self.max_model_len = 8192

2010 - logger.warning(

2019 + logger.warning_once(

2020 "max_model_len was is not set. Defaulting to arbitrary value "

2021 "of %d.", self.max_model_len)

2022

2023 if self.max_num_seqs is None:

2024 self.max_num_seqs = 128

2016 - logger.warning(

2025 + logger.warning_once(

2026 "max_num_seqs was is not set. Defaulting to arbitrary value "

2027 "of %d.", self.max_num_seqs)

2028

@@ -2840,8 +2849,8 @@ class PromptAdapterConfig:

2849 class MultiModalConfig:

2850 """Controls the behavior of multimodal models."""

2851

2843 - limit_per_prompt: dict[str, int] = get_field(ModelConfig,

2844 - "limit_mm_per_prompt")

2852 + limit_per_prompt: dict[str, int] = \

2853 + cast(dict[str, int], get_field(ModelConfig, "limit_mm_per_prompt"))

2854 """

2855 The maximum number of input items allowed per prompt for each modality.

2856 Defaults to 1 (V0) or 999 (V1) for each modality.

@@ -3415,41 +3424,49 @@ class ObservabilityConfig:

3424 self.collect_detailed_traces[0].split(","))

3425

3426

3418 -class KVTransferConfig(BaseModel):

3427 +KVProducer = Literal["kv_producer", "kv_both"]

3428 +KVConsumer = Literal["kv_consumer", "kv_both"]

3429 +KVRole = Literal[KVProducer, KVConsumer]

3430 +

3431 +

3432 +@config

3433 +@dataclass

3434 +class KVTransferConfig:

3435 """Configuration for distributed KV cache transfer."""

3436

3421 - # The KV connector for vLLM to transmit KV caches between vLLM instances.

3437 kv_connector: Optional[str] = None

3438 + """The KV connector for vLLM to transmit KV caches between vLLM instances.

3439 + """

3440

3424 - # The device used by kv connector to buffer the KV cache.

3425 - # Currently only support 'cuda'.

3441 kv_buffer_device: Optional[str] = "cuda"

3442 + """The device used by kv connector to buffer the KV cache.

3443 + Currently only support 'cuda'."""

3444

3428 - # The buffer size for TorchDistributedConnector. Measured in number of

3429 - # bytes. Recommended value: 1e9 (about 1GB).

3445 kv_buffer_size: float = 1e9

3446 + """The buffer size for TorchDistributedConnector. Measured in number of

3447 + bytes. Recommended value: 1e9 (about 1GB)."""

3448

3432 - # Whether this vLLM instance produces, consumes KV cache, or both. Choices

3433 - # are 'kv_producer', 'kv_consumer', and 'both'.

3434 - kv_role: Optional[str] = None

3449 + kv_role: Optional[KVRole] = None

3450 + """Whether this vLLM instance produces, consumes KV cache, or both. Choices

3451 + are 'kv_producer', 'kv_consumer', and 'both'."""

3452

3436 - # The rank of this vLLM instance in the KV cache transfer. Typical value:

3437 - # 0 for prefill instance, 1 for decode instance.

3438 - # Currently only 1P1D is supported.

3453 kv_rank: Optional[int] = None

3454 + """The rank of this vLLM instance in the KV cache transfer. Typical value:

3455 + 0 for prefill instance, 1 for decode instance.

3456 + Currently only 1P1D is supported."""

3457

3441 - # The number of parallel instances for KV cache transfer. For

3442 - # PyNcclConnector, this should be 2.

3458 kv_parallel_size: int = 1

3459 + """The number of parallel instances for KV cache transfer. For

3460 + PyNcclConnector, this should be 2."""

3461

3445 - # The KV connector ip, used to build distributed connection

3462 kv_ip: str = "127.0.0.1"

3463 + """The KV connector ip, used to build distributed connection."""

3464

3448 - # The KV connector port, used to build distributed connection

3465 kv_port: int = 14579

3466 + """The KV connector port, used to build distributed connection."""

3467

3451 - # any extra config that the connector may need

3452 - kv_connector_extra_config: dict[str, Any] = {}

3468 + kv_connector_extra_config: dict[str, Any] = field(default_factory=dict)

3469 + """any extra config that the connector may need."""

3470

3471 def compute_hash(self) -> str:

3472 """

@@ -3470,46 +3487,37 @@ class KVTransferConfig(BaseModel):

3487 usedforsecurity=False).hexdigest()

3488 return hash_str

3489

3473 - @classmethod

3474 - def from_cli(cls, cli_value: str) -> "KVTransferConfig":

3475 - """Parse the CLI value for the kv cache transfer config."""

3476 - return KVTransferConfig.model_validate_json(cli_value)

3477 -

3478 - def model_post_init(self, __context: Any) -> None:

3479 -

3480 - if self.kv_role is not None and self.kv_role not in [

3481 - "kv_producer", "kv_consumer", "kv_both"

3482 - ]:

3483 - raise ValueError(

3484 - f"Unsupported kv_role: {self.kv_role}. "

3485 - f"Supported roles are `kv_producer`, `kv_consumer`, "

3486 - f"and `kv_both`")

3490 + def __post_init__(self) -> None:

3491 + if self.kv_role is not None and self.kv_role not in get_args(KVRole):

3492 + raise ValueError(f"Unsupported kv_role: {self.kv_role}. "

3493 + f"Supported roles are {get_args(KVRole)}")

3494

3495 if self.kv_connector is not None and self.kv_role is None:

3496 raise ValueError("Please specify kv_disagg_role when kv_connector "

3490 - "is set, supported roles are `kv_producer`, "

3491 - "`kv_consumer`, and `kv_both`")

3497 + f"is set, supported roles are {get_args(KVRole)}")

3498

3499 @property

3500 def is_kv_transfer_instance(self) -> bool:

3501 return self.kv_connector is not None and \

3496 - self.kv_role in ["kv_producer", "kv_consumer", "kv_both"]

3502 + self.kv_role in get_args(KVRole)

3503

3504 @property

3505 def is_kv_producer(self) -> bool:

3506 return self.kv_connector is not None and \

3501 - self.kv_role in ["kv_producer", "kv_both"]

3507 + self.kv_role in get_args(KVProducer)

3508

3509 @property

3510 def is_kv_consumer(self) -> bool:

3511 return self.kv_connector is not None and \

3506 - self.kv_role in ["kv_consumer", "kv_both"]

3512 + self.kv_role in get_args(KVConsumer)

3513

3514 def get_from_extra_config(self, key, default) -> Any:

3515 return self.kv_connector_extra_config.get(key, default)

3516

3517

3512 -class KVEventsConfig(BaseModel):

3518 +@config

3519 +@dataclass

3520 +class KVEventsConfig:

3521 """Configuration for KV event publishing."""

3522

3523 enable_kv_cache_events: bool = False

@@ -3548,11 +3556,6 @@ class KVEventsConfig(BaseModel):

3556 this topic to receive events.

3557 """

3558

3551 - @classmethod

3552 - def from_cli(cls, cli_value: str) -> "KVEventsConfig":

3553 - """Parse the CLI value for the event publisher config."""

3554 - return KVEventsConfig.model_validate_json(cli_value)

3555 -

3559

3560 class CompilationLevel:

3561 # constants for the levels of the compilation process

@@ -3562,80 +3565,72 @@ class CompilationLevel:

3565 PIECEWISE = 3

3566

3567

3565 -class CompilationConfig(BaseModel):

3566 - """

3567 - Configuration for compilation.

3568 - It has three parts:

3568 +@config

3569 +@dataclass

3570 +class PassConfig:

3571 + """Configuration for custom Inductor passes.

3572 +

3573 + This is separate from general `CompilationConfig` so that inductor passes

3574 + don't all have access to full configuration - that would create a cycle as

3575 + the `PassManager` is set as a property of config."""

3576 +

3577 + dump_graph_stages: list[str] = field(default_factory=list)

3578 + """List of stages for which we want to dump the graph. Each pass defines

3579 + its own stages (before, after, maybe in-between)."""

3580 + dump_graph_dir: Path = Path(".")

3581 + """Directory to dump the graphs."""

3582 + # TODO(luka) better pass enabling system.

3583 + enable_fusion: bool = True

3584 + """Whether to enable the custom fusion pass."""

3585 + enable_noop: bool = True

3586 + """Whether to enable the custom no-op elimination pass."""

3587 + enable_sequence_parallelism: bool = False

3588 + """Whether to enable sequence parallelism."""

3589 +

3590 + def uuid(self):

3591 + """

3592 + Produces a hash unique to the pass configuration.

3593 + Any new fields that affect compilation should be added to the hash.

3594 + Do not include dump_graph_* in the hash - they don't affect

3595 + compilation.

3596 + """

3597 + include = {

3598 + "enable_fusion", "enable_noop", "enable_sequence_parallelism"

3599 + }

3600 + dict_ = {k: v for k, v in asdict(self).items() if k in include}

3601 + return InductorPass.hash_dict(dict_)

3602 +

3603 + def __post_init__(self) -> None:

3604 + if not self.enable_noop and self.enable_fusion:

3605 + logger.warning_once(

3606 + "Fusion enabled but reshape elimination disabled. "

3607 + "RMSNorm + quant (fp8) fusion might not work")

3608 +

3609 +

3610 +@config

3611 +@dataclass

3612 +class CompilationConfig:

3613 + """Configuration for compilation. It has three parts:

3614 +

3615 - Top-level Compilation control:

3570 - - level: the level of compilation.

3571 - - 0: no compilation.

3572 - - 1: dynamo as is.

3573 - - 2: dynamo once.

3574 - - 3: piecewise compilation.

3575 - - debug_dump_path: the path to dump the debug information.

3576 - - cache_dir: the directory to store the compiled graph, to

3577 - accelerate Inductor compilation. By default, it will use

3578 - model-related information to generate a cache directory.

3579 - - backend: the backend for compilation. It needs to be a string.

3580 - - "" (empty string): use the default backend.

3581 - - "eager"/"openxla"/...: use the specified backend registered in PyTorch.

3582 - - "full.module.name": a qualified name which can be used to import the backend function.

3583 - We use string to avoid serialization issues when using compilation in a distributed setting.

3584 - When the compilation level is 1 or 2, the backend is used for the compilation directly (it sees the whole graph).

3585 - When the compilation level is 3, the backend is used for the piecewise compilation (it sees a part of the graph).

3586 - - custom_ops: fine-grained control over which custom ops to enable/disable.

3587 - Use 'all' to enable all, 'none' to disable all.

3588 - Also specify a list of custom op names to enable (prefixed with a '+'),

3589 - or disable (prefixed with a '-').

3590 - Examples:

3591 - - 'all,-op1' to enable all except op1

3592 - - 'none,+op1,+op2' to enable only op1 and op2

3593 - By default, all custom ops are enabled when running without Inductor

3594 - and disabled when running with Inductor (compile_level >= Inductor).

3595 - - splitting_ops: a list of ops to split the full graph into subgraphs, used in piecewise compilation.

3616 + - {attr}`level`

3617 + - {attr}`debug_dump_path`

3618 + - {attr}`cache_dir`

3619 + - {attr}`backend`

3620 + - {attr}`custom_ops`

3621 + - {attr}`splitting_ops`

3622 - CudaGraph capture:

3597 - - use_cudagraph: whether to use cudagraph inside compilation.

3598 - - False: cudagraph inside compilation is not used.

3599 - - True: cudagraph inside compilation is used. It requires

3600 - that all input buffers have fixed addresses, and all

3601 - splitting ops write their outputs to input buffers.

3602 - Note that this is orthogonal to the cudagraph capture logic

3603 - outside of compilation.

3604 - TODO: move outside cudagraph logic into compilation.

3605 - torch.compile will handle cudagraph capture logic in the future.

3606 - - cudagraph_capture_sizes: sizes to capture cudagraph.

3607 - - None (default): capture sizes are inferred from vllm config.

3608 - - list[int]: capture sizes are specified as given.

3609 - - cudagraph_num_of_warmups: number of warmup runs for cudagraph.

3610 - It means the first several runs will be treated as warmup runs.

3611 - Only after that, the execution will be recorded, and the recorded

3612 - cudagraph will be used for subsequent runs.

3613 - - cudagraph_copy_inputs: whether to copy input tensors for

3614 - cudagraph. If the caller can guarantee that the same input buffers

3615 - are always used, it can set this to False. Otherwise, it should

3616 - set this to True, and the compiler will copy the input to an

3617 - internally managed buffer. Default is False.

3618 - - full_cuda_graph: whether to use a full cuda graph for the entire forward

3619 - pass rather than splitting certain operations such as attention into subgraphs.

3620 - Thus this flag cannot be used together with splitting_ops. This may provide

3621 - performance benefits for smaller models.

3623 + - {attr}`use_cudagraph`

3624 + - {attr}`cudagraph_capture_sizes`

3625 + - {attr}`cudagraph_num_of_warmups`

3626 + - {attr}`cudagraph_copy_inputs`

3627 + - {attr}`full_cuda_graph`

3628 - Inductor compilation:

3623 - - use_inductor: whether to use inductor compilation.

3624 - - False: inductor compilation is not used. graph runs in eager.

3625 - - True: inductor compilation is used. one graph for symbolic shape

3626 - is compiled. In addition, compile for compile_sizes,

3627 - using configurations in inductor_compile_config.

3628 - - compile_sizes: sizes to compile for inductor. In addition

3629 - to integers, it also supports "cudagraph_capture_sizes" to

3630 - specify the sizes for cudagraph capture.

3631 - - inductor_compile_config: additional configurations for inductor.

3632 - - None: use default configurations.

3633 - - inductor_passes: additional passes for inductor. It is a dictionary

3634 - from pass name to pass function qualified name. We use function

3635 - name because the config uses json format. If we pass the config

3636 - from Python, functions can also be passed directly via Python object

3637 - constructor, e.g. `CompilationConfig(inductor_passes={"a": func})`

3638 - - custom inductor passes: see PassConfig for more details

3629 + - {attr}`use_inductor`

3630 + - {attr}`compile_sizes`

3631 + - {attr}`inductor_compile_config`

3632 + - {attr}`inductor_passes`

3633 + - custom inductor passes

3634

3635 Why we have different sizes for cudagraph and inductor:

3636 - cudagraph: a cudagraph captured for a specific size can only be used

@@ -3646,83 +3641,135 @@ class CompilationConfig(BaseModel):

3641 static shapes. However, we find the general shape compilation is

3642 sufficient for most cases. It might be beneficial to compile for

3643 certain small batchsizes, where inductor is good at optimizing.

3649 - """ # noqa

3644 + """

3645 + # Top-level Compilation control

3646 level: int = 0

3647 + """The level of compilation:

3648 +

3649 + - 0: no compilation.

3650 + - 1: dynamo as is.

3651 + - 2: dynamo once.

3652 + - 3: piecewise compilation."""

3653 debug_dump_path: str = ""

3654 + """The path to dump the debug information."""

3655 cache_dir: str = ""

3656 + """The directory to store the compiled graph, to accelerate Inductor

3657 + compilation. By default, it will use model-related information to generate

3658 + a cache directory."""

3659 backend: str = ""

3654 - custom_ops: list[str] = Field(default_factory=list)

3655 - splitting_ops: list[str] = Field(default=None) # type: ignore

3656 -

3660 + """The backend for compilation. It needs to be a string:

3661 +

3662 + - "" (empty string): use the default backend.

3663 + - "eager"/"openxla"/...: use the specified backend registered in PyTorch.

3664 + - "full.module.name": a qualified name which can be used to import the

3665 +

3666 + backend function.

3667 + We use string to avoid serialization issues when using compilation in a

3668 + distributed setting. When the compilation level is 1 or 2, the backend is

3669 + used for the compilation directly (it sees the whole graph). When the

3670 + compilation level is 3, the backend is used for the piecewise compilation

3671 + (it sees a part of the graph)."""

3672 + custom_ops: list[str] = field(default_factory=list)

3673 + """Fine-grained control over which custom ops to enable/disable. Use 'all'

3674 + to enable all, 'none' to disable all. Also specify a list of custom op

3675 + names to enable (prefixed with a '+'), or disable (prefixed with a '-').

3676 + Examples:

3677 +

3678 + - 'all,-op1' to enable all except op1

3679 + - 'none,+op1,+op2' to enable only op1 and op2

3680 +

3681 + By default, all custom ops are enabled when running without Inductor and

3682 + disabled when running with Inductor (compile_level >= Inductor)."""

3683 + splitting_ops: list[str] = field(default_factory=list)

3684 + """A list of ops to split the full graph into subgraphs, used in piecewise

3685 + compilation."""

3686 +

3687 + # Inductor capture

3688 use_inductor: bool = True

3658 - compile_sizes: Optional[list[Union[int, str]]] = Field(default=None)

3659 - inductor_compile_config: dict = Field(default_factory=dict)

3660 - inductor_passes: dict[str, str] = Field(default_factory=dict)

3661 -

3689 + """Whether to use inductor compilation:

3690 +

3691 + - False: inductor compilation is not used. graph runs in eager.

3692 + - True: inductor compilation is used. one graph for symbolic shape

3693 + is compiled. In addition, compile for compile_sizes,

3694 + using configurations in inductor_compile_config."""

3695 + compile_sizes: Optional[list[Union[int, str]]] = None

3696 + """Sizes to compile for inductor. In addition

3697 + to integers, it also supports "cudagraph_capture_sizes" to

3698 + specify the sizes for cudagraph capture."""

3699 + inductor_compile_config: dict = field(default_factory=dict)

3700 + """Additional configurations for inductor.

3701 + - None: use default configurations."""

3702 + inductor_passes: dict[str, str] = field(default_factory=dict)

3703 + """Additional passes for inductor. It is a dictionary

3704 + from pass name to pass function qualified name. We use function

3705 + name because the config uses JSON format. If we pass the config

3706 + from Python, functions can also be passed directly via Python object

3707 + constructor, e.g. `CompilationConfig(inductor_passes={"a": func})`."""

3708 +

3709 + # CudaGraph compilation

3710 use_cudagraph: bool = False

3711 + """Whether to use cudagraph inside compilation.

3712 + - False: cudagraph inside compilation is not used.

3713 + - True: cudagraph inside compilation is used. It requires

3714 + that all input buffers have fixed addresses, and all

3715 + splitting ops write their outputs to input buffers.

3716 + Note that this is orthogonal to the cudagraph capture logic

3717 + outside of compilation.

3718 + TODO: move outside cudagraph logic into compilation.

3719 + torch.compile will handle cudagraph capture logic in the future."""

3720 cudagraph_num_of_warmups: int = 0

3721 + """Number of warmup runs for cudagraph.

3722 + It means the first several runs will be treated as warmup runs.

3723 + Only after that, the execution will be recorded, and the recorded

3724 + cudagraph will be used for subsequent runs."""

3725 cudagraph_capture_sizes: Optional[list[int]] = None

3726 + """Sizes to capture cudagraph.

3727 + - None (default): capture sizes are inferred from vllm config.

3728 + - list[int]: capture sizes are specified as given."""

3729 cudagraph_copy_inputs: bool = False

3730 + """Whether to copy input tensors for

3731 + cudagraph. If the caller can guarantee that the same input buffers

3732 + are always used, it can set this to False. Otherwise, it should

3733 + set this to True, and the compiler will copy the input to an

3734 + internally managed buffer. Default is False."""

3735 full_cuda_graph: bool = False

3667 -

3668 - class PassConfig(BaseModel):

3669 - """

3670 - Configuration for custom Inductor passes.

3671 - This is separate from general CompilationConfig so that inductor passes

3672 - don't all have access to full configuration - that would create a cycle

3673 - as the PassManager is set as a property of config.

3674 - - dump_graph_stages: list of stages for which we want to dump the graph.

3675 - Each pass defines its own stages (before, after, maybe in-between).

3676 - - dump_graph_dir: directory to dump the graphs. Default is .

3677 - - enable_fusion: whether to enable the custom fusion pass.

3678 - - enable_noop: whether to enable the custom no-op elimination pass.

3679 - TODO(luka) better pass enabling system.

3680 - - enable_sequence_parallelism: whether to enable sequence parallelism.

3681 - """

3682 - dump_graph_stages: list[str] = Field(default_factory=list)

3683 - dump_graph_dir: Path = Field(default=Path("."))

3684 - enable_fusion: bool = True

3685 - enable_noop: bool = True

3686 - enable_sequence_parallelism: bool = False

3687 -

3688 - def uuid(self):

3689 - """

3690 - Produces a hash unique to the pass configuration.

3691 - Any new fields that affect compilation should be added to the hash.

3692 - Do not include dump_graph_* in the hash - they don't affect

3693 - compilation.

3694 - """

3695 - dict_ = self.model_dump(include={"enable_fusion", "enable_noop", \

3696 - "enable_sequence_parallelism"})

3697 - return InductorPass.hash_dict(dict_)

3698 -

3699 - def model_post_init(self, __context: Any) -> None:

3700 - if not self.enable_noop and self.enable_fusion:

3701 - logger.warning_once(

3702 - "Fusion enabled but reshape elimination disabled. "

3703 - "RMSNorm + quant (fp8) fusion might not work")

3704 -

3705 - pass_config: PassConfig = Field(default_factory=PassConfig)

3706 -

3707 - # not configurable, computed after init

3708 - max_capture_size: int = PrivateAttr

3709 - local_cache_dir: str = PrivateAttr # local cache dir for each rank

3710 - # optimization:

3711 - # Intuitively, bs_to_padded_graph_size should be dict[int, int].

3712 - # since we know all keys are in a range [0, max_capture_size],

3713 - # we can optimize it to list[int] for better lookup performance.

3714 - bs_to_padded_graph_size: list[int] = PrivateAttr

3736 + """whether to use a full cuda graph for the entire forward pass rather than

3737 + splitting certain operations such as attention into subgraphs. Thus this

3738 + flag cannot be used together with splitting_ops. This may provide

3739 + performance benefits for smaller models."""

3740 +

3741 + pass_config: PassConfig = field(default_factory=PassConfig)

3742 + """Custom inductor passes, see PassConfig for more details"""

3743 +

3744 + max_capture_size: int = field(default=None, init=False) # type: ignore

3745 + """not configurable, computed after init"""

3746 + local_cache_dir: str = field(default=None, init=False) # type: ignore

3747 + """local cache dir for each rank"""

3748 + bs_to_padded_graph_size: list[int] = field(

3749 + default=None, # type: ignore

3750 + init=False)

3751 + """optimization:

3752 + Intuitively, bs_to_padded_graph_size should be dict[int, int].

3753 + since we know all keys are in a range [0, max_capture_size],

3754 + we can optimize it to list[int] for better lookup performance."""

3755

3756 # keep track of enabled and disabled custom ops

3717 - enabled_custom_ops: Counter[str] = PrivateAttr

3718 - disabled_custom_ops: Counter[str] = PrivateAttr

3719 - traced_files: set[str] = PrivateAttr

3720 - compilation_time: float = PrivateAttr

3721 -

3722 - # Per-model forward context

3723 - # Map from layer name to layer objects that need to be accessed outside

3724 - # model code, e.g., Attention, FusedMOE when dp_size>1.

3725 - static_forward_context: dict[str, Any] = PrivateAttr

3757 + enabled_custom_ops: Counter[str] = field(default_factory=Counter,

3758 + init=False)

3759 + """custom ops that are enabled"""

3760 + disabled_custom_ops: Counter[str] = field(default_factory=Counter,

3761 + init=False)

3762 + """custom ops that are disabled"""

3763 + traced_files: set[str] = field(default_factory=set, init=False)

3764 + """files that are traced for compilation"""

3765 + compilation_time: float = field(default=0.0, init=False)

3766 + """time taken for compilation"""

3767 +

3768 + static_forward_context: dict[str, Any] = field(default_factory=dict,

3769 + init=False)

3770 + """Per-model forward context

3771 + Map from layer name to layer objects that need to be accessed outside

3772 + model code, e.g., Attention, FusedMOE when dp_size>1."""

3773

3774 def compute_hash(self) -> str:

3775 """

@@ -3757,7 +3804,17 @@ class CompilationConfig(BaseModel):

3804 "pass_config",

3805 "traced_files",

3806 }

3760 - return self.model_dump_json(exclude=exclude, exclude_unset=True)

3807 + include = dict()

3808 + for k, v in asdict(self).items():

3809 + if k in exclude:

3810 + continue

3811 + f = get_field(CompilationConfig, k)

3812 + if (d := f.default) is not MISSING and d == v:

3813 + continue

3814 + if (df := f.default_factory) is not MISSING and df() == v:

3815 + continue

3816 + include[k] = v

3817 + return json.dumps(include)

3818

3819 __str__ = __repr__

3820

@@ -3766,12 +3823,9 @@ class CompilationConfig(BaseModel):

3823 """Parse the CLI value for the compilation config."""

3824 if cli_value in ["0", "1", "2", "3"]:

3825 return cls(level=int(cli_value))

3769 - # do not use `eval`, it is dangerous and can execute arbitrary code

3770 - dict_value = ast.literal_eval(cli_value)

3771 - return CompilationConfig.model_validate(dict_value)

3772 -

3773 - def model_post_init(self, __context: Any) -> None:

3826 + return cls(**json.loads(cli_value))

3827

3828 + def __post_init__(self) -> None:

3829 count_none = self.custom_ops.count("none")

3830 count_all = self.custom_ops.count("all")

3831 assert count_none + count_all <= 1, "Can only specify 'none' or 'all'"

@@ -3789,9 +3843,6 @@ class CompilationConfig(BaseModel):

3843 if KEY not in self.inductor_compile_config:

3844 self.inductor_compile_config[KEY] = False

3845

3792 - if self.splitting_ops is None:

3793 - self.splitting_ops = []

3794 -

3846 for k, v in self.inductor_passes.items():

3847 if not isinstance(v, str):

3848 assert callable(v), (

@@ -3808,11 +3859,8 @@ class CompilationConfig(BaseModel):

3859 self.inductor_compile_config[k] = func if isinstance(

3860 func, InductorPass) else CallableInductorPass(func)

3861

3811 - self.enabled_custom_ops = Counter()

3812 - self.disabled_custom_ops = Counter()

3813 - self.traced_files = set()

3814 - self.static_forward_context = {}

3815 - self.compilation_time = 0.0

3862 + if isinstance(self.pass_config, dict):

3863 + self.pass_config = PassConfig(**self.pass_config)

3864

3865 def init_backend(self, vllm_config: "VllmConfig") -> Union[str, Callable]:

3866 if self.level == CompilationLevel.NO_COMPILATION:

@@ -3899,39 +3947,67 @@ class CompilationConfig(BaseModel):

3947 ]

3948

3949

3950 +@config

3951 @dataclass

3952 class VllmConfig:

3953 """Dataclass which contains all vllm-related configuration. This

3954 simplifies passing around the distinct configurations in the codebase.

3955 """

3956

3908 - model_config: ModelConfig = field(default=None, init=True) # type: ignore

3909 - cache_config: CacheConfig = field(default=None, init=True) # type: ignore

3910 - parallel_config: ParallelConfig = field(default_factory=ParallelConfig,

3911 - init=True)

3912 - scheduler_config: SchedulerConfig = field(default_factory=SchedulerConfig,

3913 - init=True)

3914 - device_config: DeviceConfig = field(default=None,

3915 - init=True) # type: ignore

3916 - load_config: LoadConfig = field(default=None, init=True) # type: ignore

3957 + model_config: ModelConfig = field(default_factory=ModelConfig)

3958 + """Model configuration."""

3959 + cache_config: CacheConfig = field(default_factory=CacheConfig)

3960 + """Cache configuration."""

3961 + parallel_config: ParallelConfig = field(default_factory=ParallelConfig)

3962 + """Parallel configuration."""

3963 + scheduler_config: SchedulerConfig = field(default_factory=SchedulerConfig)

3964 + """Scheduler configuration."""

3965 + device_config: DeviceConfig = field(default_factory=DeviceConfig)

3966 + """Device configuration."""

3967 + load_config: LoadConfig = field(default_factory=LoadConfig)

3968 + """Load configuration."""

3969 lora_config: Optional[LoRAConfig] = None

3918 - speculative_config: SpeculativeConfig = field(default=None,

3919 - init=True) # type: ignore

3970 + """LoRA configuration."""

3971 + speculative_config: Optional[SpeculativeConfig] = None

3972 + """Speculative decoding configuration."""

3973 decoding_config: Optional[DecodingConfig] = None

3974 + """Decoding configuration."""

3975 observability_config: Optional[ObservabilityConfig] = None

3976 + """Observability configuration."""

3977 prompt_adapter_config: Optional[PromptAdapterConfig] = None

3978 + """Prompt adapter configuration."""

3979 quant_config: Optional[QuantizationConfig] = None

3924 - compilation_config: CompilationConfig = field(default=None,

3925 - init=True) # type: ignore

3926 - kv_transfer_config: KVTransferConfig = field(default=None,

3927 - init=True) # type: ignore

3980 + """Quantization configuration."""

3981 + compilation_config: CompilationConfig = field(

3982 + default_factory=CompilationConfig)

3983 + """`torch.compile` configuration for the model.

3984 +

3985 + When it is a number (0, 1, 2, 3), it will be interpreted as the

3986 + optimization level.

3987 +

3988 + NOTE: level 0 is the default level without any optimization. level 1 and 2

3989 + are for internal testing only. level 3 is the recommended level for

3990 + production.

3991 +

3992 + Following the convention of traditional compilers, using `-O` without space

3993 + is also supported. `-O3` is equivalent to `-O 3`.

3994 +

3995 + You can specify the full compilation config like so:

3996 + `{"level": 3, "cudagraph_capture_sizes": [1, 2, 4, 8]}`

3997 + """

3998 + kv_transfer_config: Optional[KVTransferConfig] = None

3999 + """The configurations for distributed KV cache transfer."""

4000 kv_events_config: Optional[KVEventsConfig] = None

4001 + """The configurations for event publishing."""

4002 # some opaque config, only used to provide additional information

4003 # for the hash computation, mainly used for testing, debugging or out of

4004 # tree config registration.

3932 - additional_config: SupportsHash = field(default=None,

3933 - init=True) # type: ignore

4005 + additional_config: Union[dict, SupportsHash] = field(default_factory=dict)

4006 + """Additional config for specified platform. Different platforms may

4007 + support different configs. Make sure the configs are valid for the platform

4008 + you are using. Contents must be hashable."""

4009 instance_id: str = ""

4010 + """The ID of the vLLM instance."""

4011

4012 def compute_hash(self) -> str:

4013 """

@@ -4012,7 +4088,14 @@ class VllmConfig:

4088 else:

4089 vllm_factors.append("None")

4090 if self.additional_config:

4015 - vllm_factors.append(self.additional_config.compute_hash())

4091 + if isinstance(additional_config := self.additional_config, dict):

4092 + additional_config_hash = hashlib.md5(

4093 + json.dumps(additional_config, sort_keys=True).encode(),

4094 + usedforsecurity=False,

4095 + ).hexdigest()

4096 + else:

4097 + additional_config_hash = additional_config.compute_hash()

4098 + vllm_factors.append(additional_config_hash)

4099 else:

4100 vllm_factors.append("None")

4101 factors.append(vllm_factors)

vllm/distributed/kv_events.py

+2 -1

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

5 import time

6 from abc import ABC, abstractmethod

7 from collections import deque

8 +from dataclasses import asdict

9 from itertools import count

10 from queue import Queue

11 from typing import Any, Callable, Optional, Union

@@ -284,7 +285,7 @@ class EventPublisherFactory:

285 if not config:

286 return NullEventPublisher()

287

287 - config_dict = config.model_dump()

288 + config_dict = asdict(config)

289

290 kind = config_dict.pop("publisher", "null")

291 config_dict.pop("enable_kv_cache_events")

vllm/engine/arg_utils.py

+56 -75

@@ -7,10 +7,10 @@ import json

7 import re

8 import threading

9 import warnings

10 -from dataclasses import MISSING, dataclass, fields

10 +from dataclasses import MISSING, dataclass, fields, is_dataclass

11 from itertools import permutations

12 -from typing import (Any, Callable, Dict, List, Literal, Optional, Type,

13 - TypeVar, Union, cast, get_args, get_origin)

12 +from typing import (Annotated, Any, Callable, Dict, List, Literal, Optional,

13 + Type, TypeVar, Union, cast, get_args, get_origin)

15 import torch

16 from typing_extensions import TypeIs, deprecated

@@ -36,7 +36,8 @@ from vllm.reasoning import ReasoningParserManager

36 from vllm.test_utils import MODEL_WEIGHTS_S3_BUCKET, MODELS_ON_S3

37 from vllm.transformers_utils.utils import check_gguf_file

38 from vllm.usage.usage_lib import UsageContext

39 -from vllm.utils import FlexibleArgumentParser, GiB_bytes, is_in_ray_actor

39 +from vllm.utils import (FlexibleArgumentParser, GiB_bytes, is_in_doc_build,

40 + is_in_ray_actor)

42 # yapf: enable

@@ -48,12 +49,9 @@ TypeHint = Union[type[Any], object]

49 TypeHintT = Union[type[T], object]

51 -def optional_type(

52 - return_type: Callable[[str], T]) -> Callable[[str], Optional[T]]:

52 +def parse_type(return_type: Callable[[str], T]) -> Callable[[str], T]:

54 - def _optional_type(val: str) -> Optional[T]:

55 - if val == "" or val == "None":

56 - return None

54 + def _parse_type(val: str) -> T:

55 try:

56 if return_type is json.loads and not re.match("^{.*}$", val):

57 return cast(T, nullable_kvs(val))

@@ -62,14 +60,24 @@ def optional_type(

60 raise argparse.ArgumentTypeError(

61 f"Value {val} cannot be converted to {return_type}.") from e

63 + return _parse_type

64 +

65 +

66 +def optional_type(

67 + return_type: Callable[[str], T]) -> Callable[[str], Optional[T]]:

68 +

69 + def _optional_type(val: str) -> Optional[T]:

70 + if val == "" or val == "None":

71 + return None

72 + return parse_type(return_type)(val)

73 +

74 return _optional_type

77 def union_dict_and_str(val: str) -> Optional[Union[str, dict[str, str]]]:

78 if not re.match("^{.*}$", val):

79 return str(val)

71 - else:

72 - return optional_type(json.loads)(val)

80 + return optional_type(json.loads)(val)

83 @deprecated(

@@ -144,10 +152,25 @@ def get_kwargs(cls: ConfigType) -> dict[str, Any]:

152 cls_docs = get_attr_docs(cls)

153 kwargs = {}

154 for field in fields(cls):

155 + # Get the set of possible types for the field

156 + type_hints: set[TypeHint] = set()

157 + if get_origin(field.type) in {Union, Annotated}:

158 + type_hints.update(get_args(field.type))

159 + else:

160 + type_hints.add(field.type)

161 +

162 + # If the field is a dataclass, we can use the model_validate_json

163 + generator = (th for th in type_hints if is_dataclass(th))

164 + dataclass_cls = next(generator, None)

165 +

166 # Get the default value of the field

148 - default = field.default

149 - if field.default_factory is not MISSING:

150 - default = field.default_factory()

167 + if field.default is not MISSING:

168 + default = field.default

169 + elif field.default_factory is not MISSING:

170 + if is_dataclass(field.default_factory) and is_in_doc_build():

171 + default = {}

172 + else:

173 + default = field.default_factory()

174

175 # Get the help text for the field

176 name = field.name

@@ -158,16 +181,17 @@ def get_kwargs(cls: ConfigType) -> dict[str, Any]:

181 # Initialise the kwargs dictionary for the field

182 kwargs[name] = {"default": default, "help": help}

183

161 - # Get the set of possible types for the field

162 - type_hints: set[TypeHint] = set()

163 - if get_origin(field.type) is Union:

164 - type_hints.update(get_args(field.type))

165 - else:

166 - type_hints.add(field.type)

167 -

184 # Set other kwargs based on the type hints

185 json_tip = "\n\nShould be a valid JSON string."

170 - if contains_type(type_hints, bool):

186 + if dataclass_cls is not None:

187 + dataclass_init = lambda x, f=dataclass_cls: f(**json.loads(x))

188 + # Special case for configs with a from_cli method

189 + if hasattr(dataclass_cls, "from_cli"):

190 + from_cli = dataclass_cls.from_cli

191 + dataclass_init = lambda x, f=from_cli: f(x)

192 + kwargs[name]["type"] = dataclass_init

193 + kwargs[name]["help"] += json_tip

194 + elif contains_type(type_hints, bool):

195 # Creates --no-<name> and --<name> flags

196 kwargs[name]["action"] = argparse.BooleanOptionalAction

197 elif contains_type(type_hints, Literal):

@@ -202,7 +226,7 @@ def get_kwargs(cls: ConfigType) -> dict[str, Any]:

226 kwargs[name]["type"] = union_dict_and_str

227 elif contains_type(type_hints, dict):

228 # Dict arguments will always be optional

205 - kwargs[name]["type"] = optional_type(json.loads)

229 + kwargs[name]["type"] = parse_type(json.loads)

230 kwargs[name]["help"] += json_tip

231 elif (contains_type(type_hints, str)

232 or any(is_not_builtin(th) for th in type_hints)):

@@ -771,63 +795,20 @@ class EngineArgs:

795 scheduler_group.add_argument("--scheduler-cls",

796 **scheduler_kwargs["scheduler_cls"])

797

774 - # Compilation arguments

775 - # compilation_kwargs = get_kwargs(CompilationConfig)

776 - compilation_group = parser.add_argument_group(

777 - title="CompilationConfig",

778 - description=CompilationConfig.__doc__,

779 - )

780 - compilation_group.add_argument(

781 - "--compilation-config",

782 - "-O",

783 - type=CompilationConfig.from_cli,

784 - default=None,

785 - help="torch.compile configuration for the model. "

786 - "When it is a number (0, 1, 2, 3), it will be "

787 - "interpreted as the optimization level.\n"

788 - "NOTE: level 0 is the default level without "

789 - "any optimization. level 1 and 2 are for internal "

790 - "testing only. level 3 is the recommended level "

791 - "for production.\n"

792 - "To specify the full compilation config, "

793 - "use a JSON string, e.g. ``{\"level\": 3, "

794 - "\"cudagraph_capture_sizes\": [1, 2, 4, 8]}``\n"

795 - "Following the convention of traditional "

796 - "compilers, using ``-O`` without space is also "

797 - "supported. ``-O3`` is equivalent to ``-O 3``.")

798 -

799 - # KVTransfer arguments

800 - # kv_transfer_kwargs = get_kwargs(KVTransferConfig)

801 - kv_transfer_group = parser.add_argument_group(

802 - title="KVTransferConfig",

803 - description=KVTransferConfig.__doc__,

804 - )

805 - kv_transfer_group.add_argument(

806 - "--kv-transfer-config",

807 - type=KVTransferConfig.from_cli,

808 - default=None,

809 - help="The configurations for distributed KV cache "

810 - "transfer. Should be a JSON string.")

811 - kv_transfer_group.add_argument(

812 - '--kv-events-config',

813 - type=KVEventsConfig.from_cli,

814 - default=None,

815 - help='The configurations for event publishing.')

816 -

798 # vLLM arguments

818 - # vllm_kwargs = get_kwargs(VllmConfig)

799 + vllm_kwargs = get_kwargs(VllmConfig)

800 vllm_group = parser.add_argument_group(

801 title="VllmConfig",

802 description=VllmConfig.__doc__,

803 )

823 - vllm_group.add_argument(

824 - "--additional-config",

825 - type=json.loads,

826 - default=None,

827 - help="Additional config for specified platform in JSON format. "

828 - "Different platforms may support different configs. Make sure the "

829 - "configs are valid for the platform you are using. The input format"

830 - " is like '{\"config_key\":\"config_value\"}'")

804 + vllm_group.add_argument("--kv-transfer-config",

805 + **vllm_kwargs["kv_transfer_config"])

806 + vllm_group.add_argument('--kv-events-config',

807 + **vllm_kwargs["kv_events_config"])

808 + vllm_group.add_argument("--compilation-config", "-O",

809 + **vllm_kwargs["compilation_config"])

810 + vllm_group.add_argument("--additional-config",

811 + **vllm_kwargs["additional_config"])

812

813 # Other arguments

814 parser.add_argument('--use-v2-block-manager',

vllm/entrypoints/llm.py

+9 -4

@@ -13,7 +13,8 @@ from typing_extensions import TypeVar, deprecated

14 from vllm.beam_search import (BeamSearchInstance, BeamSearchOutput,

15 BeamSearchSequence, get_beam_search_score)

16 -from vllm.config import CompilationConfig, ModelDType, TokenizerMode

16 +from vllm.config import (CompilationConfig, ModelDType, TokenizerMode,

17 + is_init_field)

18 from vllm.engine.arg_utils import (EngineArgs, HfOverrides, PoolerConfig,

19 TaskOption)

20 from vllm.engine.llm_engine import LLMEngine

@@ -204,9 +205,13 @@ class LLM:

205 kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)

206

207 if compilation_config is not None:

207 - if isinstance(compilation_config, (int, dict)):

208 - compilation_config_instance = CompilationConfig.from_cli(

209 - str(compilation_config))

208 + if isinstance(compilation_config, int):

209 + compilation_config_instance = CompilationConfig(

210 + level=compilation_config)

211 + elif isinstance(compilation_config, dict):

212 + predicate = lambda x: is_init_field(CompilationConfig, x[0])

213 + compilation_config_instance = CompilationConfig(

214 + **dict(filter(predicate, compilation_config.items())))

215 else:

216 compilation_config_instance = compilation_config

217 else:

vllm/platforms/tpu.py

+6 -5

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from typing import TYPE_CHECKING, Optional, Tuple, Union, cast

5 import torch

6 from tpu_info import device

@@ -13,9 +13,10 @@ from vllm.sampling_params import SamplingParams, SamplingType

13 from .interface import Platform, PlatformEnum, _Backend

15 if TYPE_CHECKING:

16 - from vllm.config import ModelConfig, VllmConfig

16 + from vllm.config import BlockSize, ModelConfig, VllmConfig

17 from vllm.pooling_params import PoolingParams

18 else:

19 + BlockSize = None

20 ModelConfig = None

21 VllmConfig = None

22 PoolingParams = None

@@ -94,7 +95,7 @@ class TpuPlatform(Platform):

95 cache_config = vllm_config.cache_config

96 # For v0, the default block size is 16.

97 if cache_config and cache_config.block_size is None:

97 - cache_config.block_size = 16

98 + cache_config.block_size = cast(BlockSize, 16)

99 compilation_config = vllm_config.compilation_config

100

101 # TPU only supports DYNAMO_ONCE compilation level

@@ -118,7 +119,7 @@ class TpuPlatform(Platform):

119 from vllm.v1.attention.backends.pallas import (

120 PallasAttentionBackend)

121 cache_config.block_size = PallasAttentionBackend.get_page_size(

121 - vllm_config)

122 + vllm_config) # type: ignore[assignment]

123 min_page_size = PallasAttentionBackend.get_min_page_size(

124 vllm_config)

125 if min_page_size > cache_config.block_size:

@@ -128,7 +129,7 @@ class TpuPlatform(Platform):

129 cache_config.block_size,

130 min_page_size,

131 )

131 - cache_config.block_size = min_page_size

132 + cache_config.block_size = min_page_size # type: ignore[assignment]

133

134 parallel_config = vllm_config.parallel_config

135 scheduler_config = vllm_config.scheduler_config

vllm/utils.py

@@ -1820,6 +1820,14 @@ def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:

1820 return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)

1821

1822

1823 +def is_in_doc_build() -> bool:

1824 + try:

1825 + from sphinx.ext.autodoc.mock import _MockModule

1826 + return isinstance(zmq, _MockModule)

1827 + except ModuleNotFoundError:

1828 + return False

1829 +

1830 +

1831 def import_from_path(module_name: str, file_path: Union[str, os.PathLike]):

1832 """

1833 Import a Python file according to its file path.