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

Fork 0

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

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

Synced Nov 14, 05:30 UTC

Extract `CompilationConfig` from `config.py` (#22524)

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

Harry Mellor committed Aug 9, 2025 at 00:34 UTC e3edc0a7a8f015b938d5cd77a44638dde28ab3a9

5 files changed +467 -480

tests/engine/test_arg_utils.py

-33

@@ -93,32 +93,6 @@ class NestedConfig:

93 """field"""

96 -@config

97 -@dataclass

98 -class FromCliConfig1:

99 - field: int = 1

100 - """field"""

101 -

102 - @classmethod

103 - def from_cli(cls, cli_value: str):

104 - inst = cls(**json.loads(cli_value))

105 - inst.field += 1

106 - return inst

107 -

108 -

109 -@config

110 -@dataclass

111 -class FromCliConfig2:

112 - field: int = 1

113 - """field"""

114 -

115 - @classmethod

116 - def from_cli(cls, cli_value: str):

117 - inst = cls(**json.loads(cli_value))

118 - inst.field += 2

119 - return inst

120 -

121 -

96 @config

97 @dataclass

98 class DummyConfig:

@@ -144,10 +118,6 @@ class DummyConfig:

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

119 nested_config: NestedConfig = field(default_factory=NestedConfig)

120 """Nested config"""

147 - from_cli_config1: FromCliConfig1 = field(default_factory=FromCliConfig1)

148 - """Config with from_cli method"""

149 - from_cli_config2: FromCliConfig2 = field(default_factory=FromCliConfig2)

150 - """Different config with from_cli method"""

121

122

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

@@ -199,9 +169,6 @@ def test_get_kwargs():

169 assert json_tip in kwargs["json_tip"]["help"]

170 # nested config should should construct the nested config

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

202 - # from_cli configs should be constructed with the correct method

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

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

172

173

174 @pytest.mark.parametrize(

vllm/config/__init__.py renamed

+7 -442

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

1 # SPDX-License-Identifier: Apache-2.0

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

4 +# ruff: noqa: F401

5 import ast

6 import copy

7 import enum

@@ -10,11 +11,9 @@ import json

11 import textwrap

12 import uuid

13 import warnings

13 -from collections import Counter

14 from collections.abc import Mapping

15 from contextlib import contextmanager

16 -from dataclasses import (MISSING, Field, asdict, field, fields, is_dataclass,

17 - replace)

16 +from dataclasses import MISSING, Field, field, fields, is_dataclass, replace

17 from functools import cached_property, lru_cache

18 from importlib.util import find_spec

19 from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Literal, Optional,

@@ -22,7 +21,7 @@ from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Literal, Optional,

22 import regex as re

23 import torch

25 -from pydantic import (ConfigDict, SkipValidation, TypeAdapter, field_validator,

24 +from pydantic import (ConfigDict, SkipValidation, field_validator,

25 model_validator)

26 from pydantic.dataclasses import dataclass

27 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE

@@ -31,7 +30,9 @@ from typing_extensions import Self, assert_never, runtime_checkable

31 import vllm.envs as envs

32 from vllm import version

34 -from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass

33 +from vllm.config.compilation import (CompilationConfig, CompilationLevel,

34 + PassConfig)

35 +from vllm.config.utils import ConfigType, config

36 from vllm.logger import init_logger

37 from vllm.model_executor.layers.quantization import QuantizationMethods

38 from vllm.platforms import current_platform

@@ -50,8 +51,7 @@ from vllm.utils import (DEFAULT_MAX_NUM_BATCHED_TOKENS,

51 POOLING_MODEL_MAX_NUM_BATCHED_TOKENS, GiB_bytes,

52 LayerBlockType, LazyLoader, common_broadcastable_dtype,

53 cuda_device_count_stateless, get_cpu_memory,

53 - get_open_port, is_torch_equal_or_newer, random_uuid,

54 - resolve_obj_by_qualname)

54 + get_open_port, random_uuid)

56 # yapf: enable

@@ -70,7 +70,6 @@ if TYPE_CHECKING:

70 from vllm.model_executor.model_loader import LoadFormats

71 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

73 - ConfigType = type[DataclassInstance]

73 HfOverrides = Union[dict, Callable[[type], type]]

74 else:

75 DataclassInstance = Any

@@ -83,7 +82,6 @@ else:

82 BaseModelLoader = Any

83 LoadFormats = Any

84 TensorizerConfig = Any

86 - ConfigType = type

85 HfOverrides = Union[dict[str, Any], Callable[[type], type]]

87 me_quant = LazyLoader("model_executor", globals(),

@@ -93,7 +91,6 @@ else:

92 logger = init_logger(__name__)

93 DataclassInstanceT = TypeVar("DataclassInstanceT", bound=DataclassInstance)

96 -ConfigT = TypeVar("ConfigT", bound=ConfigType)

95 TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",

96 "score", "reward", "transcription", "draft"]

@@ -234,23 +231,6 @@ def get_attr_docs(cls: type[Any]) -> dict[str, str]:

231 return out

232

233

237 -def config(cls: ConfigT) -> ConfigT:

238 - """

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

240 - and that each field has a docstring.

241 -

242 - If a `ConfigT` is used as a CLI argument itself, the default value provided

243 - by `get_kwargs` will be the result parsing a JSON string as the kwargs

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

245 - requires custom construction from CLI (i.e. `CompilationConfig`), it can

246 - have a `from_cli` method, which will be called instead.

247 -

248 - Config validation is performed by the tools/validate_config.py

249 - script, which is invoked during the pre-commit checks.

250 - """

251 - return cls

252 -

253 -

234 def get_field(cls: ConfigType, name: str) -> Field:

235 """Get the default factory field of a dataclass by name. Used for getting

236 default factory fields in `EngineArgs`."""

@@ -4154,421 +4134,6 @@ class KVEventsConfig:

4134 """

4135

4136

4157 -class CompilationLevel:

4158 - # constants for the levels of the compilation process

4159 - NO_COMPILATION = 0

4160 - DYNAMO_AS_IS = 1

4161 - DYNAMO_ONCE = 2

4162 - PIECEWISE = 3

4163 -

4164 -

4165 -@config

4166 -@dataclass

4167 -class PassConfig:

4168 - """Configuration for custom Inductor passes.

4169 -

4170 - This is separate from general `CompilationConfig` so that inductor passes

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

4172 - the `PassManager` is set as a property of config."""

4173 -

4174 - enable_fusion: bool = field(default_factory=lambda: not envs.VLLM_USE_V1)

4175 - """Whether to enable the custom fusion (RMSNorm/SiluMul+quant) pass."""

4176 - enable_attn_fusion: bool = False

4177 - """Whether to enable the custom attention+quant fusion pass."""

4178 - enable_noop: bool = field(default_factory=lambda: not envs.VLLM_USE_V1)

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

4180 - enable_sequence_parallelism: bool = False

4181 - """Whether to enable sequence parallelism."""

4182 - enable_async_tp: bool = False

4183 - """Whether to enable async TP."""

4184 - enable_fi_allreduce_fusion: bool = False

4185 - """Whether to enable flashinfer allreduce fusion."""

4186 - fi_allreduce_fusion_max_token_num: int = 16384

4187 - """Max number of tokens to used in flashinfer allreduce fusion."""

4188 -

4189 - # TODO(luka) better pass enabling system.

4190 -

4191 - def uuid(self):

4192 - """

4193 - Produces a hash unique to the pass configuration.

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

4195 - Any future fields that don't affect compilation should be excluded.

4196 - """

4197 - return InductorPass.hash_dict(asdict(self))

4198 -

4199 - def __post_init__(self) -> None:

4200 - if not self.enable_noop:

4201 - if self.enable_fusion:

4202 - logger.warning_once(

4203 - "Fusion enabled but reshape elimination disabled. "

4204 - "RMSNorm/SiluMul + quant (fp8) fusion might not work")

4205 - if self.enable_attn_fusion:

4206 - logger.warning_once(

4207 - "Fusion enabled but reshape elimination disabled. "

4208 - "Attention + quant (fp8) fusion might not work")

4209 -

4210 -

4211 -@config

4212 -@dataclass

4213 -class CompilationConfig:

4214 - """Configuration for compilation. It has three parts:

4215 -

4216 - - Top-level Compilation control:

4217 - - [`level`][vllm.config.CompilationConfig.level]

4218 - - [`debug_dump_path`][vllm.config.CompilationConfig.debug_dump_path]

4219 - - [`cache_dir`][vllm.config.CompilationConfig.cache_dir]

4220 - - [`backend`][vllm.config.CompilationConfig.backend]

4221 - - [`custom_ops`][vllm.config.CompilationConfig.custom_ops]

4222 - - [`splitting_ops`][vllm.config.CompilationConfig.splitting_ops]

4223 - - CudaGraph capture:

4224 - - [`use_cudagraph`][vllm.config.CompilationConfig.use_cudagraph]

4225 - - [`cudagraph_capture_sizes`]

4226 - [vllm.config.CompilationConfig.cudagraph_capture_sizes]

4227 - - [`cudagraph_num_of_warmups`]

4228 - [vllm.config.CompilationConfig.cudagraph_num_of_warmups]

4229 - - [`cudagraph_copy_inputs`]

4230 - [vllm.config.CompilationConfig.cudagraph_copy_inputs]

4231 - - [`full_cuda_graph`][vllm.config.CompilationConfig.full_cuda_graph]

4232 - - Inductor compilation:

4233 - - [`use_inductor`][vllm.config.CompilationConfig.use_inductor]

4234 - - [`compile_sizes`][vllm.config.CompilationConfig.compile_sizes]

4235 - - [`inductor_compile_config`]

4236 - [vllm.config.CompilationConfig.inductor_compile_config]

4237 - - [`inductor_passes`][vllm.config.CompilationConfig.inductor_passes]

4238 - - custom inductor passes

4239 -

4240 - Why we have different sizes for cudagraph and inductor:

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

4242 - for the same size. We need to capture all the sizes we want to use.

4243 - - inductor: a graph compiled by inductor for a general shape can be used

4244 - for different sizes. Inductor can also compile for specific sizes,

4245 - where it can have more information to optimize the graph with fully

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

4247 - sufficient for most cases. It might be beneficial to compile for

4248 - certain small batchsizes, where inductor is good at optimizing.

4249 - """

4250 - # Top-level Compilation control

4251 - level: Optional[int] = None

4252 - """The level of compilation:

4253 -

4254 - - None: If None, we will select the default compilation level.

4255 - For V1 engine this is 3, for V0 engine this is 0.

4256 - - 0: no compilation.

4257 - - 1: dynamo as is.

4258 - - 2: dynamo once.

4259 - - 3: piecewise compilation."""

4260 - debug_dump_path: str = ""

4261 - """The path to dump the debug information."""

4262 - cache_dir: str = ""

4263 - """The directory to store the compiled graph, to accelerate Inductor

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

4265 - a cache directory."""

4266 - backend: str = ""

4267 - """The backend for compilation. It needs to be a string:

4268 -

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

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

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

4272 -

4273 - backend function.

4274 - We use string to avoid serialization issues when using compilation in a

4275 - distributed setting. When the compilation level is 1 or 2, the backend is

4276 - used for the compilation directly (it sees the whole graph). When the

4277 - compilation level is 3, the backend is used for the piecewise compilation

4278 - (it sees a part of the graph)."""

4279 - custom_ops: list[str] = field(default_factory=list)

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

4281 - to enable all, 'none' to disable all. Also specify a list of custom op

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

4283 - Examples:

4284 -

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

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

4287 -

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

4289 - disabled when running with Inductor: level>=PIECEWISE and use_inductor=True.

4290 - Inductor generates (fused) Triton kernels for disabled custom ops."""

4291 - splitting_ops: list[str] = field(default_factory=list)

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

4293 - compilation."""

4294 -

4295 - # Inductor capture

4296 - use_inductor: bool = True

4297 - """Whether to use inductor compilation:

4298 -

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

4300 - (custom_ops enabled by default).

4301 - - True: inductor compilation is used (custom_ops disabled by default).

4302 - One graph for symbolic shape and one graph per size in compile_sizes

4303 - are compiled using configurations in inductor_compile_config.

4304 -

4305 - This setting is ignored if level<PIECEWISE."""

4306 - compile_sizes: Optional[list[Union[int, str]]] = None

4307 - """Sizes to compile for inductor. In addition

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

4309 - specify the sizes for cudagraph capture."""

4310 - inductor_compile_config: dict = field(default_factory=dict)

4311 - """Additional configurations for inductor.

4312 - - None: use default configurations."""

4313 - inductor_passes: dict[str, str] = field(default_factory=dict)

4314 - """Additional passes for inductor. It is a dictionary

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

4316 - name because the config uses JSON format. If we pass the config

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

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

4319 -

4320 - # CudaGraph compilation

4321 - use_cudagraph: bool = field(default_factory=lambda: envs.VLLM_USE_V1)

4322 - """Whether to use cudagraph inside compilation.

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

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

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

4326 - splitting ops write their outputs to input buffers.

4327 - In the vLLM V1 Engine, this flag only applies for

4328 - CompilationLevel.PIECEWISE (aka -O3).

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

4330 - outside of compilation.

4331 - TODO: move outside cudagraph logic into compilation.

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

4333 - cudagraph_num_of_warmups: int = 0

4334 - """Number of warmup runs for cudagraph.

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

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

4337 - cudagraph will be used for subsequent runs."""

4338 - cudagraph_capture_sizes: Optional[list[int]] = None

4339 - """Sizes to capture cudagraph.

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

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

4342 - cudagraph_copy_inputs: bool = False

4343 - """Whether to copy input tensors for

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

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

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

4347 - internally managed buffer. Default is False."""

4348 - full_cuda_graph: bool = False

4349 - """whether to use a full cuda graph for the entire forward pass rather than

4350 - splitting certain operations such as attention into subgraphs. Thus this

4351 - flag cannot be used together with splitting_ops. This may provide

4352 - performance benefits for smaller models."""

4353 -

4354 - pass_config: PassConfig = field(default_factory=PassConfig)

4355 - """Custom inductor passes, see PassConfig for more details"""

4356 -

4357 - max_capture_size: int = field(default=None, init=False) # type: ignore

4358 - """not configurable, computed after init"""

4359 - local_cache_dir: str = field(default=None, init=False) # type: ignore

4360 - """local cache dir for each rank"""

4361 - bs_to_padded_graph_size: list[int] = field(

4362 - default=None, # type: ignore

4363 - init=False)

4364 - """optimization:

4365 - Intuitively, bs_to_padded_graph_size should be dict[int, int].

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

4367 - we can optimize it to list[int] for better lookup performance."""

4368 -

4369 - # keep track of enabled and disabled custom ops

4370 - enabled_custom_ops: Counter[str] = field(default_factory=Counter,

4371 - init=False)

4372 - """custom ops that are enabled"""

4373 - disabled_custom_ops: Counter[str] = field(default_factory=Counter,

4374 - init=False)

4375 - """custom ops that are disabled"""

4376 - traced_files: set[str] = field(default_factory=set, init=False)

4377 - """files that are traced for compilation"""

4378 - compilation_time: float = field(default=0.0, init=False)

4379 - """time taken for compilation"""

4380 -

4381 - static_forward_context: dict[str, Any] = field(default_factory=dict,

4382 - init=False)

4383 - """Per-model forward context

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

4385 - model code, e.g., Attention, FusedMOE when dp_size>1."""

4386 -

4387 - def compute_hash(self) -> str:

4388 - """

4389 - WARNING: Whenever a new field is added to this config,

4390 - ensure that it is included in the factors list if

4391 - it affects the computation graph.

4392 -

4393 - Provide a hash that uniquely identifies all the configs

4394 - that affect the structure of the computation

4395 - graph from input ids/embeddings to the final hidden states,

4396 - excluding anything before input ids/embeddings and after

4397 - the final hidden states.

4398 - """

4399 - factors: list[Any] = []

4400 - factors.append(self.level)

4401 - factors.append(self.backend)

4402 - factors.append(self.custom_ops)

4403 - factors.append(self.splitting_ops)

4404 - factors.append(self.use_inductor)

4405 - factors.append(self.inductor_compile_config)

4406 - factors.append(self.inductor_passes)

4407 - factors.append(self.pass_config.uuid())

4408 - return hashlib.sha256(str(factors).encode()).hexdigest()

4409 -

4410 - def __repr__(self) -> str:

4411 - exclude = {

4412 - "static_forward_context": True,

4413 - "enabled_custom_ops": True,

4414 - "disabled_custom_ops": True,

4415 - "compilation_time": True,

4416 - "bs_to_padded_graph_size": True,

4417 - "traced_files": True,

4418 - "inductor_compile_config": {

4419 - "post_grad_custom_post_pass": True,

4420 - },

4421 - }

4422 -

4423 - # exclude default attr in pass_config

4424 - pass_config_exclude = {}

4425 - for attr, default_val in vars(PassConfig()).items():

4426 - if getattr(self.pass_config, attr) == default_val:

4427 - pass_config_exclude[attr] = True

4428 - if pass_config_exclude:

4429 - exclude["pass_config"] = pass_config_exclude

4430 -

4431 - # The cast to string is necessary because Pydantic is mocked in docs

4432 - # builds and sphinx-argparse doesn't know the return type of decode()

4433 - return str(

4434 - TypeAdapter(CompilationConfig).dump_json(

4435 - self,

4436 - exclude=exclude, # type: ignore[arg-type]

4437 - exclude_unset=True).decode())

4438 -

4439 - __str__ = __repr__

4440 -

4441 - @classmethod

4442 - def from_cli(cls, cli_value: str) -> "CompilationConfig":

4443 - """Parse the CLI value for the compilation config.

4444 - -O1, -O2, -O3, etc. is handled in FlexibleArgumentParser.

4445 - """

4446 - return TypeAdapter(CompilationConfig).validate_json(cli_value)

4447 -

4448 - def __post_init__(self) -> None:

4449 - count_none = self.custom_ops.count("none")

4450 - count_all = self.custom_ops.count("all")

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

4452 -

4453 - # TODO(zou3519/luka): There are 2 issues with auto-functionalization V2:

4454 - # 1. A bug in PyTorch, fixed in 2.7:

4455 - # https://github.com/pytorch/pytorch/issues/147924

4456 - # 2. Custom passes (fusion) rely on auto-functionalization V1 and don't

4457 - # work with V2. Addressing this will take extra engineering effort

4458 - # and it is not yet a priority. RFC here:

4459 - # https://github.com/vllm-project/vllm/issues/14703

4460 -

4461 - if is_torch_equal_or_newer("2.6"):

4462 - KEY = 'enable_auto_functionalized_v2'

4463 - if KEY not in self.inductor_compile_config:

4464 - self.inductor_compile_config[KEY] = False

4465 -

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

4467 - if not isinstance(v, str):

4468 - assert callable(v), (

4469 - f"pass {k} should be callable or a qualified name")

4470 - self.inductor_compile_config[k] = v if isinstance(

4471 - v, InductorPass) else CallableInductorPass(v)

4472 - continue

4473 -

4474 - # resolve function from qualified name

4475 - names = v.split(".")

4476 - module = ".".join(names[:-1])

4477 - func_name = names[-1]

4478 - func = __import__(module).__dict__[func_name]

4479 - self.inductor_compile_config[k] = func if isinstance(

4480 - func, InductorPass) else CallableInductorPass(func)

4481 -

4482 - if isinstance(self.pass_config, dict):

4483 - self.pass_config = PassConfig(**self.pass_config)

4484 -

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

4486 - if self.level == CompilationLevel.NO_COMPILATION:

4487 - raise ValueError("No compilation level is set.")

4488 -

4489 - from torch._dynamo.backends.registry import list_backends

4490 - torch_backends = list_backends(exclude_tags=tuple())

4491 - if self.level in [

4492 - CompilationLevel.DYNAMO_AS_IS, CompilationLevel.DYNAMO_ONCE

4493 - ]:

4494 - if self.backend == "":

4495 - return "eager"

4496 - if self.backend in torch_backends:

4497 - return self.backend

4498 - return resolve_obj_by_qualname(self.backend)

4499 -

4500 - # TODO: pass user-specified backend to piecewise compilation

4501 - # merge with the config use_inductor

4502 - assert self.level == CompilationLevel.PIECEWISE

4503 -

4504 - from vllm.compilation.backends import VllmBackend

4505 - return VllmBackend(vllm_config)

4506 -

4507 - def init_with_cudagraph_sizes(self,

4508 - cudagraph_capture_sizes: list[int]) -> None:

4509 - """To complete the initialization of config,

4510 - we need to know the cudagraph sizes."""

4511 -

4512 - if self.cudagraph_capture_sizes is None:

4513 - self.cudagraph_capture_sizes = cudagraph_capture_sizes

4514 - else:

4515 - # de-duplicate the sizes provided by the config

4516 - dedup_sizes = list(set(self.cudagraph_capture_sizes))

4517 - if len(dedup_sizes) < len(self.cudagraph_capture_sizes):

4518 - logger.info(("cudagraph sizes specified by model runner"

4519 - " %s is overridden by config %s"),

4520 - cudagraph_capture_sizes, dedup_sizes)

4521 - self.cudagraph_capture_sizes = dedup_sizes

4522 -

4523 - computed_compile_sizes = []

4524 - if self.compile_sizes is not None:

4525 - # de-duplicate the sizes provided by the config

4526 - self.compile_sizes = list(set(self.compile_sizes))

4527 - for x in self.compile_sizes:

4528 - if isinstance(x, str):

4529 - assert x == "cudagraph_capture_sizes", \

4530 - "Unrecognized size type in compile_sizes, " \

4531 - f"expect 'cudagraph_capture_sizes', got {x}"

4532 - computed_compile_sizes.extend(self.cudagraph_capture_sizes)

4533 - else:

4534 - assert isinstance(x, int)

4535 - computed_compile_sizes.append(x)

4536 - self.compile_sizes = computed_compile_sizes # type: ignore

4537 -

4538 - # sort to make sure cudagraph capture sizes are in descending order

4539 - self.cudagraph_capture_sizes.sort(reverse=True)

4540 - self.max_capture_size = self.cudagraph_capture_sizes[

4541 - 0] if self.cudagraph_capture_sizes else 0

4542 -

4543 - # pre-compute the mapping from batch size to padded graph size

4544 - self.bs_to_padded_graph_size = [

4545 - 0 for i in range(self.max_capture_size + 1)

4546 - ]

4547 - for end, start in zip(self.cudagraph_capture_sizes,

4548 - self.cudagraph_capture_sizes[1:] + [0]):

4549 - for bs in range(start, end):

4550 - if bs == start:

4551 - self.bs_to_padded_graph_size[bs] = start

4552 - else:

4553 - self.bs_to_padded_graph_size[bs] = end

4554 - self.bs_to_padded_graph_size[

4555 - self.max_capture_size] = self.max_capture_size

4556 -

4557 - def set_splitting_ops_for_v1(self):

4558 - # NOTE: this function needs to be called

4559 - if self.splitting_ops and self.full_cuda_graph:

4560 - raise ValueError("full_cuda_graph cannot be used together with "

4561 - "splitting_ops, as Full CUDA graph will override "

4562 - f"the splitting_ops: {self.splitting_ops}")

4563 -

4564 - if not self.splitting_ops:

4565 - self.splitting_ops = [] if self.full_cuda_graph else [

4566 - "vllm.unified_attention",

4567 - "vllm.unified_attention_with_output",

4568 - "vllm.mamba_mixer2",

4569 - ]

4570 -

4571 -

4137 @config

4138 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))

4139 class VllmConfig:

vllm/config/compilation.py new

+428

@@ -0,0 +1,428 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +import hashlib

5 +from collections import Counter

6 +from dataclasses import asdict, field

7 +from typing import TYPE_CHECKING, Any, Callable, Optional, Union

8 +

9 +from pydantic import TypeAdapter

10 +from pydantic.dataclasses import dataclass

11 +

12 +import vllm.envs as envs

13 +from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass

14 +from vllm.config.utils import config

15 +from vllm.logger import init_logger

16 +from vllm.utils import is_torch_equal_or_newer, resolve_obj_by_qualname

17 +

18 +if TYPE_CHECKING:

19 + from vllm.config.config import VllmConfig

20 +else:

21 + VllmConfig = object

22 +

23 +logger = init_logger(__name__)

24 +

25 +

26 +class CompilationLevel:

27 + # constants for the levels of the compilation process

28 + NO_COMPILATION = 0

29 + DYNAMO_AS_IS = 1

30 + DYNAMO_ONCE = 2

31 + PIECEWISE = 3

32 +

33 +

34 +@config

35 +@dataclass

36 +class PassConfig:

37 + """Configuration for custom Inductor passes.

38 +

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

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

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

42 +

43 + enable_fusion: bool = field(default_factory=lambda: not envs.VLLM_USE_V1)

44 + """Whether to enable the custom fusion (RMSNorm/SiluMul+quant) pass."""

45 + enable_attn_fusion: bool = False

46 + """Whether to enable the custom attention+quant fusion pass."""

47 + enable_noop: bool = field(default_factory=lambda: not envs.VLLM_USE_V1)

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

49 + enable_sequence_parallelism: bool = False

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

51 + enable_async_tp: bool = False

52 + """Whether to enable async TP."""

53 + enable_fi_allreduce_fusion: bool = False

54 + """Whether to enable flashinfer allreduce fusion."""

55 + fi_allreduce_fusion_max_token_num: int = 16384

56 + """Max number of tokens to used in flashinfer allreduce fusion."""

57 +

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

59 +

60 + def uuid(self):

61 + """

62 + Produces a hash unique to the pass configuration.

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

64 + Any future fields that don't affect compilation should be excluded.

65 + """

66 + return InductorPass.hash_dict(asdict(self))

67 +

68 + def __post_init__(self) -> None:

69 + if not self.enable_noop:

70 + if self.enable_fusion:

71 + logger.warning_once(

72 + "Fusion enabled but reshape elimination disabled. "

73 + "RMSNorm/SiluMul + quant (fp8) fusion might not work")

74 + if self.enable_attn_fusion:

75 + logger.warning_once(

76 + "Fusion enabled but reshape elimination disabled. "

77 + "Attention + quant (fp8) fusion might not work")

78 +

79 +

80 +@config

81 +@dataclass

82 +class CompilationConfig:

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

84 +

85 + - Top-level Compilation control:

86 + - [`level`][vllm.config.CompilationConfig.level]

87 + - [`debug_dump_path`][vllm.config.CompilationConfig.debug_dump_path]

88 + - [`cache_dir`][vllm.config.CompilationConfig.cache_dir]

89 + - [`backend`][vllm.config.CompilationConfig.backend]

90 + - [`custom_ops`][vllm.config.CompilationConfig.custom_ops]

91 + - [`splitting_ops`][vllm.config.CompilationConfig.splitting_ops]

92 + - CudaGraph capture:

93 + - [`use_cudagraph`][vllm.config.CompilationConfig.use_cudagraph]

94 + - [`cudagraph_capture_sizes`]

95 + [vllm.config.CompilationConfig.cudagraph_capture_sizes]

96 + - [`cudagraph_num_of_warmups`]

97 + [vllm.config.CompilationConfig.cudagraph_num_of_warmups]

98 + - [`cudagraph_copy_inputs`]

99 + [vllm.config.CompilationConfig.cudagraph_copy_inputs]

100 + - [`full_cuda_graph`][vllm.config.CompilationConfig.full_cuda_graph]

101 + - Inductor compilation:

102 + - [`use_inductor`][vllm.config.CompilationConfig.use_inductor]

103 + - [`compile_sizes`][vllm.config.CompilationConfig.compile_sizes]

104 + - [`inductor_compile_config`]

105 + [vllm.config.CompilationConfig.inductor_compile_config]

106 + - [`inductor_passes`][vllm.config.CompilationConfig.inductor_passes]

107 + - custom inductor passes

108 +

109 + Why we have different sizes for cudagraph and inductor:

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

111 + for the same size. We need to capture all the sizes we want to use.

112 + - inductor: a graph compiled by inductor for a general shape can be used

113 + for different sizes. Inductor can also compile for specific sizes,

114 + where it can have more information to optimize the graph with fully

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

116 + sufficient for most cases. It might be beneficial to compile for

117 + certain small batchsizes, where inductor is good at optimizing.

118 + """

119 + # Top-level Compilation control

120 + level: Optional[int] = None

121 + """The level of compilation:

122 +

123 + - None: If None, we will select the default compilation level.

124 + For V1 engine this is 3, for V0 engine this is 0.

125 + - 0: no compilation.

126 + - 1: dynamo as is.

127 + - 2: dynamo once.

128 + - 3: piecewise compilation."""

129 + debug_dump_path: str = ""

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

131 + cache_dir: str = ""

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

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

134 + a cache directory."""

135 + backend: str = ""

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

137 +

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

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

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

141 +

142 + backend function.

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

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

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

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

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

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

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

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

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

152 + Examples:

153 +

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

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

156 +

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

158 + disabled when running with Inductor: level>=PIECEWISE and use_inductor=True.

159 + Inductor generates (fused) Triton kernels for disabled custom ops."""

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

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

162 + compilation."""

163 +

164 + # Inductor capture

165 + use_inductor: bool = True

166 + """Whether to use inductor compilation:

167 +

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

169 + (custom_ops enabled by default).

170 + - True: inductor compilation is used (custom_ops disabled by default).

171 + One graph for symbolic shape and one graph per size in compile_sizes

172 + are compiled using configurations in inductor_compile_config.

173 +

174 + This setting is ignored if level<PIECEWISE."""

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

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

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

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

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

180 + """Additional configurations for inductor.

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

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

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

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

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

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

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

188 +

189 + # CudaGraph compilation

190 + use_cudagraph: bool = field(default_factory=lambda: envs.VLLM_USE_V1)

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

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

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

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

195 + splitting ops write their outputs to input buffers.

196 + In the vLLM V1 Engine, this flag only applies for

197 + CompilationLevel.PIECEWISE (aka -O3).

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

199 + outside of compilation.

200 + TODO: move outside cudagraph logic into compilation.

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

202 + cudagraph_num_of_warmups: int = 0

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

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

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

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

207 + cudagraph_capture_sizes: Optional[list[int]] = None

208 + """Sizes to capture cudagraph.

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

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

211 + cudagraph_copy_inputs: bool = False

212 + """Whether to copy input tensors for

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

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

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

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

217 + full_cuda_graph: bool = False

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

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

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

221 + performance benefits for smaller models."""

222 +

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

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

225 +

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

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

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

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

230 + bs_to_padded_graph_size: list[int] = field(

231 + default=None, # type: ignore

232 + init=False)

233 + """optimization:

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

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

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

237 +

238 + # keep track of enabled and disabled custom ops

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

240 + init=False)

241 + """custom ops that are enabled"""

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

243 + init=False)

244 + """custom ops that are disabled"""

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

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

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

248 + """time taken for compilation"""

249 +

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

251 + init=False)

252 + """Per-model forward context

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

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

255 +

256 + def compute_hash(self) -> str:

257 + """

258 + WARNING: Whenever a new field is added to this config,

259 + ensure that it is included in the factors list if

260 + it affects the computation graph.

261 +

262 + Provide a hash that uniquely identifies all the configs

263 + that affect the structure of the computation

264 + graph from input ids/embeddings to the final hidden states,

265 + excluding anything before input ids/embeddings and after

266 + the final hidden states.

267 + """

268 + factors: list[Any] = []

269 + factors.append(self.level)

270 + factors.append(self.backend)

271 + factors.append(self.custom_ops)

272 + factors.append(self.splitting_ops)

273 + factors.append(self.use_inductor)

274 + factors.append(self.inductor_compile_config)

275 + factors.append(self.inductor_passes)

276 + factors.append(self.pass_config.uuid())

277 + return hashlib.sha256(str(factors).encode()).hexdigest()

278 +

279 + def __repr__(self) -> str:

280 + exclude = {

281 + "static_forward_context": True,

282 + "enabled_custom_ops": True,

283 + "disabled_custom_ops": True,

284 + "compilation_time": True,

285 + "bs_to_padded_graph_size": True,

286 + "traced_files": True,

287 + "inductor_compile_config": {

288 + "post_grad_custom_post_pass": True,

289 + },

290 + }

291 +

292 + # exclude default attr in pass_config

293 + pass_config_exclude = {}

294 + for attr, default_val in vars(PassConfig()).items():

295 + if getattr(self.pass_config, attr) == default_val:

296 + pass_config_exclude[attr] = True

297 + if pass_config_exclude:

298 + exclude["pass_config"] = pass_config_exclude

299 +

300 + return TypeAdapter(CompilationConfig).dump_json(

301 + self,

302 + exclude=exclude, # type: ignore[arg-type]

303 + exclude_unset=True).decode()

304 +

305 + __str__ = __repr__

306 +

307 + def __post_init__(self) -> None:

308 + count_none = self.custom_ops.count("none")

309 + count_all = self.custom_ops.count("all")

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

311 +

312 + # TODO(zou3519/luka): There are 2 issues with auto-functionalization V2:

313 + # 1. A bug in PyTorch, fixed in 2.7:

314 + # https://github.com/pytorch/pytorch/issues/147924

315 + # 2. Custom passes (fusion) rely on auto-functionalization V1 and don't

316 + # work with V2. Addressing this will take extra engineering effort

317 + # and it is not yet a priority. RFC here:

318 + # https://github.com/vllm-project/vllm/issues/14703

319 +

320 + if is_torch_equal_or_newer("2.6"):

321 + KEY = 'enable_auto_functionalized_v2'

322 + if KEY not in self.inductor_compile_config:

323 + self.inductor_compile_config[KEY] = False

324 +

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

326 + if not isinstance(v, str):

327 + assert callable(v), (

328 + f"pass {k} should be callable or a qualified name")

329 + self.inductor_compile_config[k] = v if isinstance(

330 + v, InductorPass) else CallableInductorPass(v)

331 + continue

332 +

333 + # resolve function from qualified name

334 + names = v.split(".")

335 + module = ".".join(names[:-1])

336 + func_name = names[-1]

337 + func = __import__(module).__dict__[func_name]

338 + self.inductor_compile_config[k] = func if isinstance(

339 + func, InductorPass) else CallableInductorPass(func)

340 +

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

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

343 +

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

345 + if self.level == CompilationLevel.NO_COMPILATION:

346 + raise ValueError("No compilation level is set.")

347 +

348 + from torch._dynamo.backends.registry import list_backends

349 + torch_backends = list_backends(exclude_tags=tuple())

350 + if self.level in [

351 + CompilationLevel.DYNAMO_AS_IS, CompilationLevel.DYNAMO_ONCE

352 + ]:

353 + if self.backend == "":

354 + return "eager"

355 + if self.backend in torch_backends:

356 + return self.backend

357 + return resolve_obj_by_qualname(self.backend)

358 +

359 + # TODO: pass user-specified backend to piecewise compilation

360 + # merge with the config use_inductor

361 + assert self.level == CompilationLevel.PIECEWISE

362 +

363 + from vllm.compilation.backends import VllmBackend

364 + return VllmBackend(vllm_config)

365 +

366 + def init_with_cudagraph_sizes(self,

367 + cudagraph_capture_sizes: list[int]) -> None:

368 + """To complete the initialization of config,

369 + we need to know the cudagraph sizes."""

370 +

371 + if self.cudagraph_capture_sizes is None:

372 + self.cudagraph_capture_sizes = cudagraph_capture_sizes

373 + else:

374 + # de-duplicate the sizes provided by the config

375 + dedup_sizes = list(set(self.cudagraph_capture_sizes))

376 + if len(dedup_sizes) < len(self.cudagraph_capture_sizes):

377 + logger.info(("cudagraph sizes specified by model runner"

378 + " %s is overridden by config %s"),

379 + cudagraph_capture_sizes, dedup_sizes)

380 + self.cudagraph_capture_sizes = dedup_sizes

381 +

382 + computed_compile_sizes = []

383 + if self.compile_sizes is not None:

384 + # de-duplicate the sizes provided by the config

385 + self.compile_sizes = list(set(self.compile_sizes))

386 + for x in self.compile_sizes:

387 + if isinstance(x, str):

388 + assert x == "cudagraph_capture_sizes", \

389 + "Unrecognized size type in compile_sizes, " \

390 + f"expect 'cudagraph_capture_sizes', got {x}"

391 + computed_compile_sizes.extend(self.cudagraph_capture_sizes)

392 + else:

393 + assert isinstance(x, int)

394 + computed_compile_sizes.append(x)

395 + self.compile_sizes = computed_compile_sizes # type: ignore

396 +

397 + # sort to make sure cudagraph capture sizes are in descending order

398 + self.cudagraph_capture_sizes.sort(reverse=True)

399 + self.max_capture_size = self.cudagraph_capture_sizes[

400 + 0] if self.cudagraph_capture_sizes else 0

401 +

402 + # pre-compute the mapping from batch size to padded graph size

403 + self.bs_to_padded_graph_size = [

404 + 0 for i in range(self.max_capture_size + 1)

405 + ]

406 + for end, start in zip(self.cudagraph_capture_sizes,

407 + self.cudagraph_capture_sizes[1:] + [0]):

408 + for bs in range(start, end):

409 + if bs == start:

410 + self.bs_to_padded_graph_size[bs] = start

411 + else:

412 + self.bs_to_padded_graph_size[bs] = end

413 + self.bs_to_padded_graph_size[

414 + self.max_capture_size] = self.max_capture_size

415 +

416 + def set_splitting_ops_for_v1(self):

417 + # NOTE: this function needs to be called

418 + if self.splitting_ops and self.full_cuda_graph:

419 + raise ValueError("full_cuda_graph cannot be used together with "

420 + "splitting_ops, as Full CUDA graph will override "

421 + f"the splitting_ops: {self.splitting_ops}")

422 +

423 + if not self.splitting_ops:

424 + self.splitting_ops = [] if self.full_cuda_graph else [

425 + "vllm.unified_attention",

426 + "vllm.unified_attention_with_output",

427 + "vllm.mamba_mixer2",

428 + ]

vllm/config/utils.py new

+29

@@ -0,0 +1,29 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +from typing import TYPE_CHECKING, TypeVar

5 +

6 +if TYPE_CHECKING:

7 + from _typeshed import DataclassInstance

8 +

9 + ConfigType = type[DataclassInstance]

10 +else:

11 + ConfigType = type

12 +

13 +ConfigT = TypeVar("ConfigT", bound=ConfigType)

14 +

15 +

16 +def config(cls: ConfigT) -> ConfigT:

17 + """

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

19 + and that each field has a docstring.

20 +

21 + If a `ConfigT` is used as a CLI argument itself, the `type` keyword argument

22 + provided by `get_kwargs` will be

23 + `pydantic.TypeAdapter(ConfigT).validate_json(cli_arg)` which treats the

24 + `cli_arg` as a JSON string which gets validated by `pydantic`.

25 +

26 + Config validation is performed by the tools/validate_config.py

27 + script, which is invoked during the pre-commit checks.

28 + """

29 + return cls

vllm/engine/arg_utils.py

+3 -5

@@ -193,8 +193,6 @@ Additionally, list elements can be passed individually using `+`:

193

194 def parse_dataclass(val: str, cls=dataclass_cls) -> Any:

195 try:

196 - if hasattr(cls, "from_cli"):

197 - return cls.from_cli(val)

196 return TypeAdapter(cls).validate_json(val)

197 except ValidationError as e:

198 raise argparse.ArgumentTypeError(repr(e)) from e

@@ -455,9 +453,9 @@ class EngineArgs:

453 # support `EngineArgs(compilation_config={...})`

454 # without having to manually construct a

455 # CompilationConfig object

458 - if isinstance(self.compilation_config, (int, dict)):

459 - self.compilation_config = CompilationConfig.from_cli(

460 - str(self.compilation_config))

456 + if isinstance(self.compilation_config, dict):

457 + self.compilation_config = CompilationConfig(

458 + **self.compilation_config)

459 # Setup plugins

460 from vllm.plugins import load_general_plugins

461 load_general_plugins()