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

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 13, 10:45 UTC

[Frontend] Change CompilationMode to a proper Enum (#28165)

Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>

Yanan Cao committed Nov 11, 2025 at 16:46 UTC 48c879369f83ab1ab281a4bfe97f9a54790715d1

6 files changed +108 -23

tests/compile/test_basic_correctness.py

+4 -2

@@ -127,7 +127,9 @@ def test_compile_correctness(

127 CompilationMode.VLLM_COMPILE,

128 ]:

129 for mode in [CompilationMode.NONE, comp_mode]:

130 - all_args.append(final_args + [f"-O.mode={mode}", "-O.backend=inductor"])

130 + all_args.append(

131 + final_args + [f"-O.mode={mode.name}", "-O.backend=inductor"]

132 + )

133

134 # inductor will change the output, so we only compare if the output

135 # is close, not exactly the same.

@@ -146,7 +148,7 @@ def test_compile_correctness(

148 CompilationMode.DYNAMO_TRACE_ONCE,

149 CompilationMode.VLLM_COMPILE,

150 ]:

149 - all_args.append(final_args + [f"-O.mode={mode}", "-O.backend=eager"])

151 + all_args.append(final_args + [f"-O.mode={mode.name}", "-O.backend=eager"])

152 all_envs.append({})

153 all_envs.append({})

154

tests/utils_/test_argparse_utils.py

+60

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

8 import pytest

9 import yaml

10 from transformers import AutoTokenizer

11 +from pydantic import ValidationError

13 from vllm.transformers_utils.detokenizer_utils import convert_ids_list_to_tokens

@@ -376,6 +377,65 @@ def test_load_config_file(tmp_path):

377 os.remove(str(config_file_path))

378

379

380 +def test_compilation_mode_string_values(parser):

381 + """Test that -O.mode accepts both integer and string mode values."""

382 + args = parser.parse_args(["-O.mode", "0"])

383 + assert args.compilation_config == {"mode": 0}

384 +

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

386 + assert args.compilation_config == {"mode": 3}

387 +

388 + args = parser.parse_args(["-O.mode=NONE"])

389 + assert args.compilation_config == {"mode": "NONE"}

390 +

391 + args = parser.parse_args(["-O.mode", "STOCK_TORCH_COMPILE"])

392 + assert args.compilation_config == {"mode": "STOCK_TORCH_COMPILE"}

393 +

394 + args = parser.parse_args(["-O.mode=DYNAMO_TRACE_ONCE"])

395 + assert args.compilation_config == {"mode": "DYNAMO_TRACE_ONCE"}

396 +

397 + args = parser.parse_args(["-O.mode", "VLLM_COMPILE"])

398 + assert args.compilation_config == {"mode": "VLLM_COMPILE"}

399 +

400 + args = parser.parse_args(["-O.mode=none"])

401 + assert args.compilation_config == {"mode": "none"}

402 +

403 + args = parser.parse_args(["-O.mode=vllm_compile"])

404 + assert args.compilation_config == {"mode": "vllm_compile"}

405 +

406 +

407 +def test_compilation_config_mode_validator():

408 + """Test that CompilationConfig.mode field validator converts strings to integers."""

409 + from vllm.config.compilation import CompilationConfig, CompilationMode

410 +

411 + config = CompilationConfig(mode=0)

412 + assert config.mode == CompilationMode.NONE

413 +

414 + config = CompilationConfig(mode=3)

415 + assert config.mode == CompilationMode.VLLM_COMPILE

416 +

417 + config = CompilationConfig(mode="NONE")

418 + assert config.mode == CompilationMode.NONE

419 +

420 + config = CompilationConfig(mode="STOCK_TORCH_COMPILE")

421 + assert config.mode == CompilationMode.STOCK_TORCH_COMPILE

422 +

423 + config = CompilationConfig(mode="DYNAMO_TRACE_ONCE")

424 + assert config.mode == CompilationMode.DYNAMO_TRACE_ONCE

425 +

426 + config = CompilationConfig(mode="VLLM_COMPILE")

427 + assert config.mode == CompilationMode.VLLM_COMPILE

428 +

429 + config = CompilationConfig(mode="none")

430 + assert config.mode == CompilationMode.NONE

431 +

432 + config = CompilationConfig(mode="vllm_compile")

433 + assert config.mode == CompilationMode.VLLM_COMPILE

434 +

435 + with pytest.raises(ValidationError, match="Invalid compilation mode"):

436 + CompilationConfig(mode="INVALID_MODE")

437 +

438 +

439 def test_flat_product():

440 # Check regular itertools.product behavior

441 result1 = list(flat_product([1, 2, 3], ["a", "b"]))

vllm/compilation/wrapper.py

+3 -1

@@ -31,7 +31,9 @@ class TorchCompileWrapperWithCustomDispatcher:

31 """

33 def __init__(

34 - self, compiled_callable: Callable | None = None, compilation_mode: int = 0

34 + self,

35 + compiled_callable: Callable | None = None,

36 + compilation_mode: CompilationMode = CompilationMode.NONE,

37 ):

38 vllm_config = get_current_vllm_config()

39 self.vllm_config = vllm_config

vllm/config/compilation.py

+36 -15

@@ -28,7 +28,7 @@ else:

28 logger = init_logger(__name__)

31 -class CompilationMode:

31 +class CompilationMode(enum.IntEnum):

32 """The compilation approach used for torch.compile-based compilation of the

33 model."""

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

115 """The threshold of the communicated tensor sizes under which

116 vllm should use flashinfer fused allreduce. Specified as a

117 float in MB.

118 - Unspecified will fallback to default values

118 + Unspecified will fallback to default values

119 which are compute capability and world size dependent.

120 FI_ALLREDUCE_FUSION_MAX_SIZE_MB = {

121 90: {

@@ -244,7 +244,7 @@ class CompilationConfig:

244 Please use mode. Currently all levels are mapped to mode.

245 """

246 # Top-level Compilation control

247 - mode: int | None = None

247 + mode: CompilationMode | None = None

248 """The compilation approach used for torch.compile-based compilation of the

249 model.

250

@@ -377,23 +377,23 @@ class CompilationConfig:

377 FULL mode: Capture full cudagraph for all batches. Can be good for small

378 models or workloads with small prompts; not supported by many backends.

379 Generally for performance FULL_AND_PIECEWISE is better.

380 -

380 +

381 FULL_DECODE_ONLY mode: Capture full cudagraph for decode batches only.

382 Mixed prefill-decode batches are run without cudagraphs. Can be good for

383 decode instances in a P/D setup where prefill is not as important so we

384 can save some memory.

385 -

385 +

386 FULL_AND_PIECEWISE mode: Capture full cudagraph for decode batches and

387 piecewise cudagraph for prefill and mixed prefill-decode batches.

388 This is the most performant mode for most models and is the default.

389

390 Currently, the cudagraph mode is only used for the v1 engine.

391 - Note that the cudagraph logic is generally orthogonal to the

392 - compilation logic. While piecewise cudagraphs require piecewise

391 + Note that the cudagraph logic is generally orthogonal to the

392 + compilation logic. While piecewise cudagraphs require piecewise

393 compilation (mode=VLLM_COMPILE and non-empty splitting_ops), full

394 cudagraphs are supported with and without compilation.

395 -

396 - Warning: This flag is new and subject to change in addition

395 +

396 + Warning: This flag is new and subject to change in addition

397 more modes may be added.

398 """

399 use_cudagraph: bool = True

@@ -422,7 +422,7 @@ class CompilationConfig:

422 cudagraph. If the caller can guarantee that the same input buffers

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

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

425 - internally managed buffer. Default is False.

425 + internally managed buffer. Default is False.

426 Note that this flag is only effective when cudagraph_mode is PIECEWISE.

427 """

428 full_cuda_graph: bool | None = False

@@ -451,7 +451,7 @@ class CompilationConfig:

451 outside the partition functions. For a graph with N cudagraph-unsafe ops

452 (e.g., Attention), there would be N+1 partitions. To mark an op as

453 cudagraph unsafe, we can add `tags=(torch._C.Tag.cudagraph_unsafe)` when

454 - register the custom op.

454 + register the custom op.

455

456 This config supports both full cudagraph and piecewise cudagraph without

457 compiling twice. For piecewise cudagraph, it applies vLLM CUDAGraph wrapper

@@ -468,8 +468,8 @@ class CompilationConfig:

468

469 max_cudagraph_capture_size: int | None = field(default=None)

470 """The maximum cudagraph capture size.

471 -

472 - If cudagraph_capture_sizes is specified, this will be set to the largest

471 +

472 + If cudagraph_capture_sizes is specified, this will be set to the largest

473 size in that list (or checked for consistency if specified). If

474 cudagraph_capture_sizes is not specified, the list of sizes is generated

475 automatically following the pattern:

@@ -478,7 +478,7 @@ class CompilationConfig:

478 range(256, max_cudagraph_capture_size + 1, 16))

479

480 If not specified, max_cudagraph_capture_size is set to min(max_num_seqs*2,

481 - 512) by default. This voids OOM in tight memory scenarios with small

481 + 512) by default. This voids OOM in tight memory scenarios with small

482 max_num_seqs, and prevents capture of many large graphs (>512) that would

483 greatly increase startup time with limited performance benefit.

484 """

@@ -579,6 +579,27 @@ class CompilationConfig:

579

580 __str__ = __repr__

581

582 + @field_validator("mode", mode="before")

583 + @classmethod

584 + def validate_mode_before(cls, value: Any) -> Any:

585 + """

586 + Enable parsing the `mode` field from string mode names.

587 + Accepts both integers (0-3) and string names, like NONE, STOCK_TORCH_COMPILE,

588 + DYNAMO_TRACE_ONCE, VLLM_COMPILE.

589 + """

590 + if isinstance(value, str):

591 + # Convert string mode name to integer value

592 + mode_name = value.upper()

593 +

594 + if mode_name not in CompilationMode.__members__:

595 + raise ValueError(

596 + f"Invalid compilation mode: {value}. "

597 + f"Valid modes are: {', '.join(CompilationMode.__members__.keys())}"

598 + )

599 +

600 + return CompilationMode[mode_name]

601 + return value

602 +

603 @field_validator("cudagraph_mode", mode="before")

604 @classmethod

605 def validate_cudagraph_mode_before(cls, value: Any) -> Any:

@@ -904,7 +925,7 @@ class CompilationConfig:

925 return self.mode == CompilationMode.VLLM_COMPILE

926

927 # Inductor partition case

907 - return self.backend == "inductor" and self.mode > CompilationMode.NONE

928 + return self.backend == "inductor" and self.mode != CompilationMode.NONE

929

930 def custom_op_log_check(self):

931 """

vllm/config/vllm.py

+1 -4

@@ -422,16 +422,13 @@ class VllmConfig:

422 self.compilation_config.mode = CompilationMode.VLLM_COMPILE

423 else:

424 self.compilation_config.mode = CompilationMode.NONE

425 - else:

426 - assert self.compilation_config.mode >= CompilationMode.NONE

427 - assert self.compilation_config.mode <= CompilationMode.VLLM_COMPILE

425

426 # If user does not set custom ops via none or all set it here based on

427 # compilation mode and backend.

428 if all(s not in self.compilation_config.custom_ops for s in ("all", "none")):

429 if (

430 self.compilation_config.backend == "inductor"

434 - and self.compilation_config.mode > CompilationMode.NONE

431 + and self.compilation_config.mode != CompilationMode.NONE

432 ):

433 self.compilation_config.custom_ops.append("none")

434 else:

vllm/entrypoints/llm.py

+4 -1

@@ -23,6 +23,7 @@ from vllm.config import (

23 StructuredOutputsConfig,

24 is_init_field,

25 )

26 +from vllm.config.compilation import CompilationMode

27 from vllm.config.model import (

28 ConvertOption,

29 HfOverrides,

@@ -259,7 +260,9 @@ class LLM:

260

261 if compilation_config is not None:

262 if isinstance(compilation_config, int):

262 - compilation_config_instance = CompilationConfig(mode=compilation_config)

263 + compilation_config_instance = CompilationConfig(

264 + mode=CompilationMode(compilation_config)

265 + )

266 elif isinstance(compilation_config, dict):

267 compilation_config_instance = CompilationConfig(

268 **{