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

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] Fix tests for torch==2.9 inductor partition (#26116)

Signed-off-by: ProExpertProg <lgovedic@redhat.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com>

Luka Govedič committed Oct 14, 2025 at 19:55 UTC 2dcd12d3571b070432ad1cd321a67b840b4a34b2

8 files changed +138 -72

tests/compile/piecewise/test_full_cudagraph.py

+21 -8

@@ -11,6 +11,7 @@ from tests.v1.attention.utils import full_cg_backend_configs as backend_configs

11 from vllm import LLM, SamplingParams

12 from vllm.config import CompilationConfig

13 from vllm.platforms import current_platform

14 +from vllm.utils import is_torch_equal_or_newer

17 @contextlib.contextmanager

@@ -32,13 +33,13 @@ def temporary_environ(env_vars):

33 os.environ[k] = v

35 -test_params_full_cudagraph = []

36 +model_backends_full_cudagraph = []

38 # deepseek-ai/DeepSeek-V2-Lite with MLA

39 MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA"]

40 for mla_backend in MLA_backends:

40 - test_params_full_cudagraph.append(

41 - pytest.param(("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend]))

41 + model_backends_full_cudagraph.append(

42 + ("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend])

43 )

45 # Qwen/Qwen2-1.5B-Instruct with other backends

@@ -46,14 +47,18 @@ other_backend_configs = [

47 backend_configs[c] for c in backend_configs if c not in MLA_backends

48 ]

49 for backend_config in other_backend_configs:

49 - test_params_full_cudagraph.append(

50 - pytest.param(("Qwen/Qwen2-1.5B-Instruct", backend_config))

51 - )

50 + model_backends_full_cudagraph.append(("Qwen/Qwen2-1.5B-Instruct", backend_config))

53 @pytest.fixture(scope="class")

54 def llm_pair(request):

56 - model, backend_config = request.param

55 + model, backend_config, use_inductor_graph_partition = request.param

56 + backend_config.comp_config["use_inductor_graph_partition"] = (

57 + use_inductor_graph_partition

58 + )

59 +

60 + if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):

61 + pytest.skip("Inductor graph partition only supported in torch>=2.9")

63 # Dynamically skip test if GPU capability is not met

64 if (

@@ -104,7 +109,15 @@ def llm_pair(request):

109 )

110

111

107 -@pytest.mark.parametrize("llm_pair", test_params_full_cudagraph, indirect=True)

112 +@pytest.mark.parametrize(

113 + "llm_pair",

114 + [

115 + pytest.param((model, backend_config, use_inductor_graph_partition))

116 + for model, backend_config in model_backends_full_cudagraph

117 + for use_inductor_graph_partition in [True, False]

118 + ],

119 + indirect=True,

120 +)

121 class TestFullCUDAGraph:

122 """

123 Use a class such that an llm pair is constructed once for all

tests/compile/piecewise/test_multiple_graphs.py

+27 -11

@@ -5,6 +5,7 @@ Test (piecewise) compilation with a simple model where multiple submodules

5 are compiled and graph captured separately.

6 """

8 +import pytest

9 import torch

10 from torch import nn

@@ -190,7 +191,12 @@ def run_model(

191 return output.cpu()

192

193

193 -def test_multi_graph_piecewise_compile_outputs_equal():

194 +@pytest.mark.parametrize("use_inductor_graph_partition", [False, True])

195 +def test_multi_graph_piecewise_compile(use_inductor_graph_partition: bool):

196 + if use_inductor_graph_partition:

197 + # FIXME(luka/boyuan): this currently fails

198 + pytest.skip("Inductor graph partition not supported with multi-graph")

199 +

200 outputs = []

201

202 # piecewise compile

@@ -200,6 +206,7 @@ def test_multi_graph_piecewise_compile_outputs_equal():

206 use_cudagraph=True,

207 splitting_ops=["silly::attention"],

208 cudagraph_capture_sizes=[1, 2],

209 + use_inductor_graph_partition=use_inductor_graph_partition,

210 )

211 )

212 cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

@@ -220,16 +227,24 @@ def test_multi_graph_piecewise_compile_outputs_equal():

227 # static tensor addresses

228 inputs = torch.randn(BATCH_SIZE, MLP_SIZE).cuda()

229

223 - with compilation_counter.expect(

224 - num_graphs_seen=2, # two graphs for the model

225 - num_piecewise_graphs_seen=6,

230 + if use_inductor_graph_partition:

231 + # Splitting happens at Inductor lowering level,

232 + # total piecewise fx graphs is equal to total graphs

233 + num_piecewise_fx = 2

234 + num_piecewise_capturable_fx = 2

235 + else:

236 # attn_one, attn_two each has 3 piecewise graphs

237 # (pre attn, post attn, silly_attention) each

228 - num_piecewise_capturable_graphs_seen=4,

238 + num_piecewise_fx = 6

239 # attn_one, attn_two has pre attn and post attn each, total=4

230 - num_backend_compilations=4, # num_piecewise_capturable_graphs_seen

231 - num_cudagraph_captured=8,

232 - # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen

240 + num_piecewise_capturable_fx = 4

241 +

242 + with compilation_counter.expect(

243 + num_graphs_seen=2, # two graphs for the model

244 + num_piecewise_graphs_seen=num_piecewise_fx,

245 + num_piecewise_capturable_graphs_seen=num_piecewise_capturable_fx,

246 + num_backend_compilations=num_piecewise_capturable_fx,

247 + num_cudagraph_captured=8, # num_cudagraph_sizes * num_partitions

248 ):

249 outputs.append(run_model(vllm_config, model, inputs, cudagraph_runtime_mode))

250

@@ -268,6 +283,7 @@ def test_multi_graph_piecewise_compile_outputs_equal():

283 level=CompilationLevel.PIECEWISE,

284 use_cudagraph=False,

285 splitting_ops=["silly::attention"],

286 + use_inductor_graph_partition=use_inductor_graph_partition,

287 )

288 )

289 cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

@@ -286,9 +302,9 @@ def test_multi_graph_piecewise_compile_outputs_equal():

302

303 with compilation_counter.expect(

304 num_graphs_seen=2,

289 - num_piecewise_graphs_seen=6,

290 - num_piecewise_capturable_graphs_seen=4,

291 - num_backend_compilations=4,

305 + num_piecewise_graphs_seen=num_piecewise_fx,

306 + num_piecewise_capturable_graphs_seen=num_piecewise_capturable_fx,

307 + num_backend_compilations=num_piecewise_capturable_fx,

308 num_cudagraph_captured=0, # no cudagraph captured

309 ):

310 outputs.append(run_model(vllm_config, model, inputs, cudagraph_runtime_mode))

tests/compile/piecewise/test_toy_llama.py

+74 -43

@@ -9,6 +9,7 @@ if the config `tractable_init` is set to True. Otherwise, the weights are

9 initialized randomly with a fixed seed.

10 """

12 +from copy import deepcopy

13 from dataclasses import dataclass

14 from typing import Any

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

27 set_current_vllm_config,

28 )

29 from vllm.forward_context import BatchDescriptor, set_forward_context

30 +from vllm.utils import is_torch_equal_or_newer

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

33 from .. import silly_attention # noqa: F401

@@ -257,27 +259,13 @@ def tractable_computation(

259

260

261 @torch.inference_mode

260 -def run_model(

261 - llama_config, use_compile: bool, backend: str, split_attn: bool = False

262 -) -> torch.Tensor:

263 - if use_compile:

264 - compilation_config = CompilationConfig(

265 - level=CompilationLevel.PIECEWISE,

266 - use_cudagraph=True,

267 - backend=backend,

268 - cudagraph_capture_sizes=[1, 2],

269 - )

270 - if split_attn:

271 - compilation_config.splitting_ops = ["silly::attention"]

272 - cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

273 - else:

274 - compilation_config = CompilationConfig(

275 - level=CompilationLevel.NO_COMPILATION,

276 - )

277 - cudagraph_runtime_mode = CUDAGraphMode.NONE

262 +def run_model(llama_config, compile_config: CompilationConfig) -> torch.Tensor:

263 + # Start with a fresh copy to make sure there's no cache dir sharing

264 + compile_config = deepcopy(compile_config)

265 + cudagraph_runtime_mode = compile_config.cudagraph_mode

266

267 vllm_config = VllmConfig(

280 - compilation_config=compilation_config, additional_config=llama_config

268 + compilation_config=compile_config, additional_config=llama_config

269 )

270 with set_current_vllm_config(vllm_config):

271 model = (

@@ -338,8 +326,25 @@ def run_model(

326 return output.cpu()

327

328

341 -@pytest.mark.parametrize("backend", ["inductor", "eager"])

342 -def test_toy_llama(backend: str):

329 +@pytest.mark.parametrize(

330 + "backend, use_inductor_graph_partition",

331 + [

332 + ("eager", False), # No inductor

333 + ("inductor", False), # Inductor, Dynamo partition

334 + ("inductor", True), # Inductor, Inductor partition

335 + ],

336 +)

337 +def test_toy_llama(

338 + backend: str, use_inductor_graph_partition: bool, monkeypatch, tmp_path

339 +):

340 + # We disable the vLLM compile cache into a new tmp dir for 2 reasons:

341 + # 1. To make sure we can properly track the number of Inductor compilations.

342 + # 2. Inductor partitioning does not play nicely with Autograd cache (below)

343 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")

344 +

345 + if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):

346 + pytest.skip("Inductor graph partition only supported in torch>=2.9")

347 +

348 # compare output with and without piecewise compilation

349

350 llama_config = LlamaConfig(

@@ -350,6 +355,32 @@ def test_toy_llama(backend: str):

355 hidden_size=128, mlp_size=256, vocab_size=128, num_layers=2, tractable_init=True

356 )

357

358 + compile_config_no_compile = CompilationConfig(

359 + level=CompilationLevel.NO_COMPILATION,

360 + cudagraph_mode=CUDAGraphMode.NONE,

361 + backend="eager",

362 + )

363 +

364 + compile_config_no_split = CompilationConfig(

365 + level=CompilationLevel.PIECEWISE,

366 + use_inductor_graph_partition=use_inductor_graph_partition,

367 + cudagraph_mode=CUDAGraphMode.PIECEWISE,

368 + backend=backend,

369 + cudagraph_capture_sizes=[1, 2],

370 + )

371 +

372 + # FIXME(luka/boyuan): the graph from the previous test case

373 + # (no inductor partition) gets cached by AotAutograd so then the

374 + # compilation with inductor partitioning incorrectly loads an unpartitioned

375 + # graph and never partitions. I think this is a bug with custom inductor

376 + # partitioning but does not affect vLLM more generally as vLLM uses its own

377 + # cache (which takes inductor partitioning into account).

378 + if use_inductor_graph_partition:

379 + compile_config_no_split.inductor_compile_config["force_disable_caches"] = True

380 +

381 + compile_config_split = deepcopy(compile_config_no_split)

382 + compile_config_split.splitting_ops = ["silly::attention"]

383 +

384 outputs = []

385 with compilation_counter.expect(

386 num_graphs_seen=0,

@@ -358,8 +389,9 @@ def test_toy_llama(backend: str):

389 num_backend_compilations=0,

390 num_cudagraph_captured=0,

391 ):

361 - outputs.append(run_model(llama_config, backend="eager", use_compile=False))

362 - run_model(tractable_config, backend="eager", use_compile=False)

392 + outputs.append(run_model(llama_config, compile_config_no_compile))

393 +

394 + run_model(tractable_config, compile_config_no_compile)

395

396 if backend == "inductor":

397 kwargs = {"num_inductor_compiles": 1, "num_eager_compiles": 0}

@@ -367,35 +399,34 @@ def test_toy_llama(backend: str):

399 kwargs = {"num_eager_compiles": 1, "num_inductor_compiles": 0}

400

401 with compilation_counter.expect(

370 - # One graph for the model

371 - num_graphs_seen=1,

402 + num_graphs_seen=1, # one graph for the model

403 num_piecewise_graphs_seen=1,

404 num_piecewise_capturable_graphs_seen=1,

374 - # num_piecewise_capturable_graphs_seen

375 - num_backend_compilations=1,

376 - # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen

405 + num_backend_compilations=1, # num_piecewise_capturable_graphs_seen

406 num_cudagraph_captured=2,

407 **kwargs,

408 ):

380 - outputs.append(run_model(llama_config, backend=backend, use_compile=True))

381 - run_model(tractable_config, backend=backend, use_compile=True)

409 + outputs.append(run_model(llama_config, compile_config_no_split))

410 +

411 + run_model(tractable_config, compile_config_no_split)

412 +

413 + if use_inductor_graph_partition:

414 + num_piecewise_fx = 1

415 + num_piecewise_capturable_fx = 1

416 + else:

417 + num_piecewise_fx = 2 * llama_config.num_layers + 1

418 + num_piecewise_capturable_fx = 1 + llama_config.num_layers

419

420 with compilation_counter.expect(

421 num_graphs_seen=1, # one graph for the model

385 - num_piecewise_graphs_seen=2 * llama_config.num_layers + 1, # 2 * num_layers + 1

386 - num_piecewise_capturable_graphs_seen=1

387 - + llama_config.num_layers, # 1 + num_layers

388 - num_backend_compilations=1

389 - + llama_config.num_layers, # num_piecewise_capturable_graphs_seen

390 - num_cudagraph_captured=2

391 - * (

392 - 1 + llama_config.num_layers

393 - ), # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen

422 + num_piecewise_graphs_seen=num_piecewise_fx,

423 + num_piecewise_capturable_graphs_seen=num_piecewise_capturable_fx,

424 + num_backend_compilations=num_piecewise_capturable_fx,

425 + # num_cudagraph_sizes * num_partitions

426 + num_cudagraph_captured=2 * (1 + llama_config.num_layers),

427 ):

395 - outputs.append(

396 - run_model(llama_config, backend=backend, use_compile=True, split_attn=True)

397 - )

398 - run_model(tractable_config, backend=backend, use_compile=True, split_attn=True)

428 + outputs.append(run_model(llama_config, compile_config_split))

429 + run_model(tractable_config, compile_config_split)

430

431 for i in range(1, len(outputs)):

432 assert torch.allclose(outputs[0], outputs[i])

tests/compile/silly_attention.py

-1

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

62 mutates_args=["out"],

63 fake_impl=silly_attention_fake,

64 target_lib=silly_lib,

65 - tags=(torch._C.Tag.cudagraph_unsafe,),

65 )

tests/compile/test_decorator.py

@@ -73,6 +73,7 @@ def test_ignore_torch_compile_decorator():

73 use_cudagraph=True,

74 splitting_ops=["silly::attention"],

75 cudagraph_capture_sizes=[1, 2],

76 + use_inductor_graph_partition=False, # TODO test both?

77 )

78 )

79 cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

@@ -188,6 +189,7 @@ def test_conditional_compile_enable_if():

189 use_cudagraph=True,

190 splitting_ops=["silly::attention"],

191 cudagraph_capture_sizes=[1, 2],

192 + use_inductor_graph_partition=False, # TODO test both

193 ),

194 )

195 cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

@@ -220,6 +222,7 @@ def test_conditional_compile_enable_if():

222 use_cudagraph=True,

223 splitting_ops=["silly::attention"],

224 cudagraph_capture_sizes=[1, 2],

225 + use_inductor_graph_partition=False, # TODO test both?

226 ),

227 )

228

vllm/attention/layer.py

-6

@@ -38,10 +38,6 @@ from vllm.utils import GiB_bytes, direct_register_custom_op

39 logger = init_logger(__name__)

40 USE_XFORMERS_OPS = None

41 -try:

42 - tag_cudagraph_unsafe = (torch._C.Tag.cudagraph_unsafe,)

43 -except AttributeError:

44 - tag_cudagraph_unsafe = () # type: ignore[assignment]

43 def check_xformers_availability():

@@ -879,7 +875,6 @@ direct_register_custom_op(

875 op_name="unified_attention",

876 op_func=unified_attention,

877 fake_impl=unified_attention_fake,

882 - tags=tag_cudagraph_unsafe,

878 )

879

880

@@ -931,7 +926,6 @@ direct_register_custom_op(

926 op_func=unified_attention_with_output,

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

928 fake_impl=unified_attention_with_output_fake,

934 - tags=tag_cudagraph_unsafe,

929 )

930

931

vllm/compilation/partition_rules.py

+11 -2

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

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

4 import contextlib

5 +import logging

6 from typing import TYPE_CHECKING

8 from torch._library.utils import lookup_op

@@ -38,8 +39,16 @@ def resolve_defined_ops(op_names: list[str]) -> list["torch._ops.OpOverload"]:

39 resolved.append(lookup_op(op_name))

40 except Exception:

41 # Skip operators that don't exist (e.g., model-specific ops)

41 - logger.warning(

42 - "Failed to resolve operator for Inductor partition: %s", op_name

42 + # Do not warn for attention ops, warn for others

43 + # (most likely manually specified)

44 + from vllm.config import CompilationConfig

45 +

46 + logger.log(

47 + logging.DEBUG

48 + if op_name in CompilationConfig._attention_ops

49 + else logging.WARNING,

50 + "Failed to resolve operator for CUDAGraph partition: %s",

51 + op_name,

52 )

53 continue

vllm/config/compilation.py

+2 -1

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

201 (it sees a part of the graph). The backend can not be custom for compilation

202 level 3, i.e. the backend must be either eager or inductor. Furthermore,

203 compilation is only piecewise if splitting ops is set accordingly and

204 - use_inductor_cudagraphs_partition is off. Note that the default options for

204 + use_inductor_graph_partition is off. Note that the default options for

205 splitting ops are sufficient for piecewise compilation.

206 """

207 custom_ops: list[str] = field(default_factory=list)

@@ -431,6 +431,7 @@ class CompilationConfig:

431 factors.append(self.custom_ops)

432 factors.append(self.splitting_ops)

433 factors.append(self.use_inductor)

434 + factors.append(self.use_inductor_graph_partition)

435 factors.append(self.inductor_compile_config)

436 factors.append(self.inductor_passes)

437 factors.append(self.pass_config.uuid())