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

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

[torch.compile] CUDAGraph Inductor partition integration (#24281)

Signed-off-by: Boyuan Feng <boyuan@meta.com> Signed-off-by: Boyuan Feng <fby.1994@gmail.com> Signed-off-by: boyuanfeng <boyuan@meta.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Boyuan Feng committed Sep 19, 2025 at 18:02 UTC 8945b001db3202f882108e50d16b6f9c5e6f01ed

9 files changed +280 -32

tests/compile/piecewise/test_simple.py

+60 -11

@@ -15,6 +15,7 @@ from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,

15 VllmConfig, set_current_vllm_config)

16 from vllm.envs import VLLM_USE_V1

17 from vllm.forward_context import BatchDescriptor, set_forward_context

18 +from vllm.utils import is_torch_equal_or_newer

20 # This import automatically registers `torch.ops.silly.attention`

21 from ..silly_attention import get_global_counter, reset_global_counter

@@ -50,16 +51,21 @@ class SillyModel(nn.Module):

51 return x

53 -@pytest.mark.parametrize("use_inductor", [True, False])

54 -@torch.inference_mode()

55 -def test_simple_piecewise_compile(use_inductor):

56 - assert VLLM_USE_V1

57 -

54 +def _run_simple_model(

55 + splitting_ops,

56 + use_inductor_graph_partition,

57 + use_inductor,

58 + expected_num_piecewise_graphs_seen,

59 + expected_num_piecewise_capturable_graphs_seen,

60 + expected_num_backend_compilations,

61 + expected_num_cudagraph_captured,

62 +):

63 vllm_config = VllmConfig(compilation_config=CompilationConfig(

64 level=CompilationLevel.PIECEWISE,

65 use_cudagraph=True,

66 use_inductor=use_inductor,

62 - splitting_ops=["silly.attention"],

67 + splitting_ops=splitting_ops,

68 + use_inductor_graph_partition=use_inductor_graph_partition,

69 cudagraph_copy_inputs=True,

70 cudagraph_capture_sizes=[1, 2],

71 ))

@@ -70,11 +76,11 @@ def test_simple_piecewise_compile(use_inductor):

77 with compilation_counter.expect(

78 num_graphs_seen=1, # one graph for the model

73 - num_piecewise_graphs_seen=5, # 2 * num_layers + 1

74 - num_piecewise_capturable_graphs_seen=3, # 1 + num_layers

75 - num_backend_compilations=3, # num_piecewise_capturable_graphs_seen

76 - num_cudagraph_captured=

77 - 6, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen

79 + num_piecewise_graphs_seen=expected_num_piecewise_graphs_seen,

80 + num_piecewise_capturable_graphs_seen=

81 + expected_num_piecewise_capturable_graphs_seen,

82 + num_backend_compilations=expected_num_backend_compilations,

83 + num_cudagraph_captured=expected_num_cudagraph_captured,

84 ), set_forward_context(None,

85 vllm_config=vllm_config): # background context

86 # warm up with background context

@@ -104,3 +110,46 @@ def test_simple_piecewise_compile(use_inductor):

110 output = model(input)

111 assert get_global_counter() == 2

112 assert torch.allclose(output.cpu(), torch.tensor([19.0, 19.0]))

113 +

114 +

115 +@pytest.mark.parametrize("use_inductor", [True, False])

116 +@torch.inference_mode()

117 +def test_simple_piecewise_compile(use_inductor):

118 + assert VLLM_USE_V1

119 + _run_simple_model(

120 + splitting_ops=["silly.attention"],

121 + use_inductor_graph_partition=False,

122 + use_inductor=use_inductor,

123 + expected_num_piecewise_graphs_seen=5, # 2 * num_layers + 1

124 + expected_num_piecewise_capturable_graphs_seen=3, # 1 + num_layers

125 + expected_num_backend_compilations=

126 + 3, # num_piecewise_capturable_graphs_seen

127 + expected_num_cudagraph_captured=

128 + 6, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen

129 + )

130 +

131 +

132 +@torch.inference_mode()

133 +@pytest.mark.parametrize("splitting_ops", [["silly.attention"], []])

134 +def test_simple_inductor_graph_partition(splitting_ops):

135 + assert VLLM_USE_V1

136 + if not is_torch_equal_or_newer("2.9.0.dev"):

137 + pytest.skip("inductor graph partition is only available "

138 + "in PyTorch 2.9+")

139 +

140 + _run_simple_model(

141 + # inductor graph partition automatically resets splitting_ops

142 + # to be an empty list

143 + splitting_ops=splitting_ops,

144 + use_inductor_graph_partition=True,

145 + use_inductor=True,

146 + expected_num_piecewise_graphs_seen=

147 + 1, # since not splitting at fx graph level

148 + expected_num_piecewise_capturable_graphs_seen=

149 + 1, # since not splitting at fx graph level

150 + expected_num_backend_compilations=

151 + 1, # since not splitting at fx graph level

152 + expected_num_cudagraph_captured=

153 + 6, # inductor graph partition still captures 6

154 + # graph, same as fx graph partition.

155 + )

tests/compile/silly_attention.py

@@ -60,4 +60,5 @@ direct_register_custom_op(

60 mutates_args=["out"],

61 fake_impl=silly_attention_fake,

62 target_lib=silly_lib,

63 + tags=(torch._C.Tag.cudagraph_unsafe, ),

64 )

tests/compile/test_full_graph.py

+58 -1

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

4 from __future__ import annotations

6 +import logging

7 import tempfile

8 from typing import Any, Optional, Union

@@ -10,9 +11,13 @@ import pytest

11 import torch

13 from tests.quantization.utils import is_quant_method_supported

14 +from tests.v1.attention.utils import _Backend

15 from vllm import LLM, SamplingParams

14 -from vllm.config import CompilationConfig, CompilationLevel, PassConfig

16 +from vllm.attention.selector import global_force_attn_backend_context_manager

17 +from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,

18 + PassConfig)

19 from vllm.platforms import current_platform

20 +from vllm.utils import is_torch_equal_or_newer

22 from ..utils import create_new_process_for_each_test

@@ -105,6 +110,18 @@ def test_full_graph(

110 (CompilationConfig(level=CompilationLevel.PIECEWISE,

111 debug_dump_path=tempfile.gettempdir()),

112 ("facebook/opt-125m", {})),

113 + ] + [

114 + # graph inductor partition

115 + (

116 + CompilationConfig(

117 + level=CompilationLevel.PIECEWISE,

118 + # inductor graph partition uses

119 + # torch._C.Tag.cudagraph_unsafe to specify splitting ops

120 + use_inductor_graph_partition=True,

121 + cudagraph_mode=CUDAGraphMode.PIECEWISE,

122 + compile_sizes=[1, 2]),

123 + model) for model in models_list(all=False)

124 + if is_torch_equal_or_newer("2.9.0.dev")

125 ])

126 # only test some of the models

127 @create_new_process_for_each_test()

@@ -112,11 +129,51 @@ def test_custom_compile_config(

129 compilation_config: CompilationConfig,

130 model_info: tuple[str, dict[str, Any]],

131 ):

132 + if (compilation_config.use_inductor_graph_partition

133 + and not is_torch_equal_or_newer("2.9.0.dev")):

134 + pytest.skip("inductor graph partition is only available "

135 + "in PyTorch 2.9+")

136 +

137 model, model_kwargs = model_info

138 print(f"MODEL={model}")

139 run_model(compilation_config, model, model_kwargs)

140

141

142 +def test_inductor_graph_partition_attn_fusion(caplog_vllm):

143 + if not is_torch_equal_or_newer("2.9.0.dev"):

144 + pytest.skip("inductor graph partition is only available "

145 + "in PyTorch 2.9+")

146 +

147 + model = "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"

148 + compilation_config = CompilationConfig(

149 + level=CompilationLevel.PIECEWISE,

150 + use_inductor_graph_partition=True,

151 + cudagraph_mode=CUDAGraphMode.PIECEWISE,

152 + custom_ops=["+quant_fp8"],

153 + pass_config=PassConfig(enable_attn_fusion=True, enable_noop=True),

154 + )

155 + model_kwargs = {

156 + "kv_cache_dtype": "fp8",

157 + "max_model_len": 1024,

158 + }

159 + with caplog_vllm.at_level(

160 + logging.DEBUG), global_force_attn_backend_context_manager(

161 + _Backend.FLASHINFER):

162 + run_model(compilation_config, model, model_kwargs)

163 +

164 + try:

165 + assert ("Fused quantization onto 48 attention nodes"

166 + in caplog_vllm.text), caplog_vllm.text

167 + except AssertionError:

168 + # Note: this message is only triggered when the compilation goes

169 + # through the custom pass. Due to multiple layers of cache on

170 + # PyTorch side, the compilation of a graph may be cached such

171 + # that custom pass directly goes through cache. In this case,

172 + # we go through this branch and assert that the pass is not

173 + # triggered.

174 + assert "Fused quantization" not in caplog_vllm.text

175 +

176 +

177 def run_model(compile_config: Union[int, CompilationConfig], model: str,

178 model_kwargs: dict[str, Any]):

179 prompts = [

tests/compile/test_fusion_attn.py

+15 -1

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

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

28 Fp8LinearOp)

29 from vllm.platforms import current_platform

30 +from vllm.utils import is_torch_equal_or_newer

31 from vllm.v1.kv_cache_interface import AttentionSpec

33 FP8_DTYPE = current_platform.fp8_dtype()

@@ -339,6 +340,10 @@ else:

340 @pytest.mark.parametrize(

341 "split_attention",

342 [False, True] if current_platform.is_rocm() else [False])

343 +# TODO(boyuan): test inductor graph partition on rocm

344 +@pytest.mark.parametrize(

345 + "use_inductor_graph_partition",

346 + [False] if current_platform.is_rocm() else [False, True])

347 @pytest.mark.skipif(not current_platform.is_cuda_alike(),

348 reason="Only test ROCm or CUDA")

349 @pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")

@@ -352,9 +357,15 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,

357 dtype: torch.dtype, model_name: str,

358 model_class: type[AttentionQuantPatternModel],

359 backend: _Backend, split_attention: bool,

355 - monkeypatch, dist_init):

360 + use_inductor_graph_partition: bool,

361 + monkeypatch, dist_init, caplog_vllm):

362 """Test AttentionStaticQuantPattern fusion pass"""

363

364 + if use_inductor_graph_partition and not is_torch_equal_or_newer(

365 + "2.9.0.dev"):

366 + pytest.skip("inductor graph partition is only available "

367 + "in PyTorch 2.9+")

368 +

369 monkeypatch.setenv("VLLM_USE_V1", "1")

370 if split_attention:

371 monkeypatch.setenv("VLLM_V1_USE_PREFILL_DECODE_ATTENTION", "1")

@@ -372,6 +383,7 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,

383 compilation_config=CompilationConfig(

384 level=CompilationLevel.PIECEWISE,

385 custom_ops=["+quant_fp8"],

386 + use_inductor_graph_partition=use_inductor_graph_partition,

387 ),

388 cache_config=CacheConfig(cache_dtype="fp8"))

389

@@ -444,6 +456,7 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,

456 backend=test_backend,

457 fullgraph=True)

458 assert model_compiled.attn._o_scale_float is None

459 +

460 result_fused_1 = model_compiled(q, k, v)

461

462 if backend == _Backend.FLASHINFER:

@@ -453,6 +466,7 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,

466 # _o_scale_float

467 assert model_compiled.attn._o_scale_float is not None

468 result_fused_2 = model_compiled(q, k, v)

469 +

470 assert model_compiled.attn._o_scale_float is not None

471

472 torch.testing.assert_close(result_unfused,

vllm/attention/layer.py

@@ -577,6 +577,7 @@ direct_register_custom_op(

577 mutates_args=[],

578 fake_impl=unified_attention_fake,

579 dispatch_key=current_platform.dispatch_key,

580 + tags=(torch._C.Tag.cudagraph_unsafe, ),

581 )

582

583

@@ -627,4 +628,5 @@ direct_register_custom_op(

628 mutates_args=["output", "output_block_scale"],

629 fake_impl=unified_attention_with_output_fake,

630 dispatch_key=current_platform.dispatch_key,

631 + tags=(torch._C.Tag.cudagraph_unsafe, ),

632 )

vllm/compilation/backends.py

+8 -2

@@ -329,6 +329,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):

329 i for i, x in enumerate(args) if isinstance(x, torch.SymInt)

330 ]

331 global compilation_start_time

332 +

333 compiled_graph_for_dynamic_shape = self.vllm_backend.\

334 compiler_manager.compile(

335 submod,

@@ -339,7 +340,6 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):

340 num_graphs=len(self.compile_submod_names),

341 runtime_shape=None)

342 # Lazy import here to avoid circular import

342 - from .cuda_graph import CUDAGraphOptions

343 from .cuda_piecewise_backend import PiecewiseBackend

344

345 piecewise_backend = PiecewiseBackend(

@@ -347,7 +347,13 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):

347 len(self.compile_submod_names), sym_shape_indices,

348 compiled_graph_for_dynamic_shape, self.vllm_backend)

349

350 - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:

350 + if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

351 + and

352 + not self.compilation_config.use_inductor_graph_partition):

353 + # We're using Dynamo-based piecewise splitting, so we wrap

354 + # the whole subgraph with a static graph wrapper.

355 + from .cuda_graph import CUDAGraphOptions

356 +

357 # resolve the static graph wrapper class (e.g. CUDAGraphWrapper

358 # class) as platform dependent.

359 static_graph_wrapper_class = resolve_obj_by_qualname(

vllm/compilation/decorators.py

+55 -2

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

1 # SPDX-License-Identifier: Apache-2.0

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

4 +import contextlib

5 import inspect

6 from typing import Callable, Optional, TypeVar, Union, overload

7 from unittest.mock import patch

@@ -14,7 +15,7 @@ from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher

15 from vllm.config import CompilationLevel, VllmConfig

16 from vllm.logger import init_logger

17 from vllm.sequence import IntermediateTensors

17 -from vllm.utils import supports_dynamo

18 +from vllm.utils import resolve_obj_by_qualname, supports_dynamo

20 from .monitor import start_monitoring_torch_compile

@@ -301,8 +302,11 @@ def _support_torch_compile(

302

303 with patch.object(InliningInstructionTranslator, 'inline_call',

304 patched_inline_call), torch._dynamo.config.patch(

304 - **dynamo_config_patches):

305 + **dynamo_config_patches

306 + ), maybe_use_cudagraph_partition_wrapper(

307 + self.vllm_config):

308 output = self.compiled_callable(*args, **kwargs)

309 +

310 return output

311

312 # usually, capturing the model once is enough, and then we can

@@ -314,3 +318,52 @@ def _support_torch_compile(

318

319 cls.__call__ = __call__

320 return cls

321 +

322 +

323 +@contextlib.contextmanager

324 +def maybe_use_cudagraph_partition_wrapper(vllm_config: VllmConfig):

325 + """

326 + Context manager to set/unset customized cudagraph partition wrappers.

327 +

328 + If we're using Inductor-based graph partitioning, we currently have the

329 + whole `fx.Graph` before Inductor lowering and and the piecewise

330 + splitting happens after all graph passes and fusions. Here, we add

331 + a custom hook for Inductor to wrap each partition with our static

332 + graph wrapper class to maintain more control over static graph

333 + capture and replay.

334 + """

335 + from vllm.config import CUDAGraphMode

336 +

337 + compilation_config = vllm_config.compilation_config

338 + if (compilation_config.cudagraph_mode != CUDAGraphMode.NONE

339 + and compilation_config.use_inductor_graph_partition):

340 + from torch._inductor.utils import CUDAGraphWrapperMetadata

341 +

342 + from vllm.compilation.cuda_graph import CUDAGraphOptions

343 + from vllm.platforms import current_platform

344 +

345 + static_graph_wrapper_class = resolve_obj_by_qualname(

346 + current_platform.get_static_graph_wrapper_cls())

347 +

348 + def customized_cudagraph_wrapper(f,

349 + metadata: CUDAGraphWrapperMetadata):

350 + partition_id = metadata.partition_index

351 + num_partitions = metadata.num_partitions

352 + return static_graph_wrapper_class(

353 + runnable=f,

354 + vllm_config=vllm_config,

355 + runtime_mode=CUDAGraphMode.PIECEWISE,

356 + cudagraph_options=CUDAGraphOptions(

357 + debug_log_enable=partition_id == 0,

358 + gc_disable=partition_id != 0,

359 + weak_ref_output=partition_id == num_partitions - 1,

360 + ))

361 +

362 + torch._inductor.utils.set_customized_partition_wrappers(

363 + customized_cudagraph_wrapper)

364 +

365 + yield

366 +

367 + if (compilation_config.cudagraph_mode != CUDAGraphMode.NONE

368 + and compilation_config.use_inductor_graph_partition):

369 + torch._inductor.utils.set_customized_partition_wrappers(None)

vllm/config/compilation.py

+73 -11

@@ -299,6 +299,26 @@ class CompilationConfig:

299 minor release, i.e. v0.11.0 or v1.0.0. Please use cudagraph_mode instead.

300 """

301

302 + use_inductor_graph_partition: bool = False

303 + """Use inductor graph partition to split the graph at cudagraph_unsafe ops.

304 + This partition happens at inductor codegen time after all passes and fusions

305 + are finished. It generates a single `call` function which wraps

306 + cudagraph-safe ops into partition functions and leave cudagraph-unsafe ops

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

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

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

310 + register the custom op.

311 +

312 + This config supports both full cudagraph and piecewise cudagraph without

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

314 + to each partition. For N+1 partitions, there would be N+1

315 + CUDAGraph wrapper instances.

316 +

317 + For full CUDAGraph, we always apply a single CUDAGraph wrapper outside the

318 + inductor `call` function in the model runner. The top-level full cudagraph

319 + capture ignores all partitioning.

320 + """

321 +

322 pass_config: PassConfig = field(default_factory=PassConfig)

323 """Custom inductor passes, see PassConfig for more details"""

324

@@ -461,6 +481,12 @@ class CompilationConfig:

481 "since full_cuda_graph is deprecated.")

482 self.cudagraph_mode = CUDAGraphMode.FULL

483

484 + if (self.use_inductor_graph_partition

485 + and not is_torch_equal_or_newer("2.9.0.dev")):

486 + raise ValueError("use_inductor_graph_partition is only "

487 + "supported with torch>=2.9.0.dev. Set "

488 + "use_inductor_graph_partition=False instead.")

489 +

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

491 if self.level == CompilationLevel.NO_COMPILATION:

492 raise ValueError("No compilation level is set.")

@@ -540,19 +566,36 @@ class CompilationConfig:

566 "set_splitting_ops_for_v1 should only be called when "

567 "level is CompilationLevel.PIECEWISE")

568

569 + use_inductor_graph_partition_msg = (

570 + "When use_inductor_graph_partition=True, splitting_ops "

571 + "are ignored and set to an empty list. Instead, "

572 + "\"tags=(torch._C.Tag.cudagraph_unsafe, ),\" is "

573 + "used to annotate custom ops for graph partition.")

574 +

575 if self.splitting_ops is None:

544 - # NOTE: When using full cudagraph, instead of setting an empty

545 - # list and capture the full cudagraph inside the flattened fx

546 - # graph, we keep the piecewise fx graph structure but capture the

547 - # full cudagraph outside the fx graph. This reduces some cpu

548 - # overhead when the runtime batch_size is not cudagraph captured.

549 - # see https://github.com/vllm-project/vllm/pull/20059 for details.

550 - # make a copy to avoid mutating the class-level list via reference.

551 - self.splitting_ops = list(self._attention_ops)

576 + if self.use_inductor_graph_partition:

577 + # When using inductor graph partition, we set splitting_ops

578 + # to be empty and rely on torch._C.Tag.cudagraph_unsafe to

579 + # annotate custom ops as splitting ops.

580 + logger.warning_once(use_inductor_graph_partition_msg)

581 + self.splitting_ops = []

582 + else:

583 + # NOTE: When using full cudagraph, instead of setting an empty

584 + # list and capture the full cudagraph inside the flattened fx

585 + # graph, we keep the piecewise fx graph structure but capture

586 + # the full cudagraph outside the fx graph. This reduces some

587 + # cpu overhead when the runtime batch_size is not cudagraph

588 + # captured. see https://github.com/vllm-project/vllm/pull/20059

589 + # for details. make a copy to avoid mutating the class-level

590 + # list via reference.

591 + self.splitting_ops = list(self._attention_ops)

592 elif len(self.splitting_ops) == 0:

553 - logger.warning_once("Using piecewise compilation with empty "

554 - "splitting_ops.")

555 - if self.cudagraph_mode == CUDAGraphMode.PIECEWISE:

593 + logger.warning_once(

594 + "Using piecewise compilation with empty "

595 + "splitting_ops and use_inductor_graph_partition"

596 + f"={self.use_inductor_graph_partition}.")

597 + if (self.cudagraph_mode == CUDAGraphMode.PIECEWISE

598 + and not self.use_inductor_graph_partition):

599 logger.warning_once(

600 "When compilation level is piecewise with empty "

601 "splitting_ops, PIECEWISE cudagraph_mode will be "

@@ -562,7 +605,26 @@ class CompilationConfig:

605 "any problems.")

606 self.cudagraph_mode = CUDAGraphMode.FULL

607 self.splitting_ops = []

608 + elif self.use_inductor_graph_partition:

609 + logger.warning_once(use_inductor_graph_partition_msg)

610 + self.splitting_ops = []

611

612 def splitting_ops_contain_attention(self) -> bool:

613 return self.splitting_ops is not None and all(

614 op in self.splitting_ops for op in self._attention_ops)

615 +

616 + def is_attention_compiled_piecewise(self) -> bool:

617 + use_fx_graph_piecewise_compilation = (

618 + self.level == CompilationLevel.PIECEWISE

619 + and self.splitting_ops_contain_attention())

620 +

621 + inductor_used = (self.level == CompilationLevel.PIECEWISE

622 + and self.use_inductor) or (

623 + self.level >= CompilationLevel.DYNAMO_AS_IS

624 + and self.backend == "inductor")

625 + use_inductor_piecewise_compilation = (

626 + inductor_used and self.use_inductor_graph_partition

627 + and not self.splitting_ops_contain_attention())

628 +

629 + return use_fx_graph_piecewise_compilation or \

630 + use_inductor_piecewise_compilation

vllm/v1/cudagraph_dispatcher.py

+8 -4

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

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

3 from typing import Optional

5 -from vllm.config import CompilationLevel, CUDAGraphMode, VllmConfig

5 +from vllm.config import CUDAGraphMode, VllmConfig

6 from vllm.forward_context import BatchDescriptor

7 from vllm.logger import init_logger

@@ -39,11 +39,15 @@ class CudagraphDispatcher:

39 CUDAGraphMode.FULL: set(),

40 }

42 - assert not self.cudagraph_mode.requires_piecewise_compilation() or \

43 - (self.compilation_config.level == CompilationLevel.PIECEWISE and

44 - self.compilation_config.splitting_ops_contain_attention()), \

42 + not_use_piecewise_compilation = (

43 + not self.cudagraph_mode.requires_piecewise_compilation())

44 +

45 + assert not_use_piecewise_compilation or \

46 + self.compilation_config.is_attention_compiled_piecewise(), \

47 "Compilation level should be CompilationLevel.PIECEWISE when "\

48 "cudagraph_mode piecewise cudagraphs is used, "\

49 + "and attention should be in splitting_ops or "\

50 + "inductor splitting should be used. " \

51 f"cudagraph_mode={self.cudagraph_mode}, "\

52 f"compilation_level={self.compilation_config.level}, "\

53 f"splitting_ops={self.compilation_config.splitting_ops}"