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

Fork 0

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

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

Synced Nov 15, 08:30 UTC

Add ability to use CUDAGraphs with use_inductor=False (#17345)

Signed-off-by: rzou <zou3519@gmail.com>

Richard Zou committed May 28, 2025 at 22:16 UTC 26b4fa45bead5d65d4e15bfaffaa52ac71bea270

5 files changed +51 -13

tests/compile/piecewise/test_simple.py

+10 -1

@@ -74,11 +74,12 @@ class SillyModel(nn.Module):

74 return x

77 -def test_simple_piecewise_compile():

77 +def _test_simple_piecewise_compile(*, use_inductor):

79 vllm_config = VllmConfig(compilation_config=CompilationConfig(

80 level=CompilationLevel.PIECEWISE,

81 use_cudagraph=True,

82 + use_inductor=use_inductor,

83 splitting_ops=["silly.attention"],

84 cudagraph_copy_inputs=True,

85 cudagraph_capture_sizes=[1, 2],

@@ -108,3 +109,11 @@ def test_simple_piecewise_compile():

109 output = model(input)

110 assert global_counter == 2

111 assert torch.allclose(output.cpu(), torch.tensor([3., 1.]))

112 +

113 +

114 +def test_simple_piecewise_compile_inductor():

115 + _test_simple_piecewise_compile(use_inductor=True)

116 +

117 +

118 +def test_simple_piecewise_compile_no_inductor():

119 + _test_simple_piecewise_compile(use_inductor=False)

tests/compile/piecewise/test_toy_llama.py

+33 -7

@@ -261,12 +261,14 @@ def tractable_computation(input_ids: torch.Tensor,

261 @torch.inference_mode

262 def run_model(llama_config,

263 use_compile: bool,

264 + use_inductor: bool,

265 split_attn: bool = False) -> torch.Tensor:

266

267 if use_compile:

268 compilation_config = CompilationConfig(

269 level=CompilationLevel.PIECEWISE,

270 use_cudagraph=True,

271 + use_inductor=use_inductor,

272 cudagraph_capture_sizes=[1, 2],

273 )

274 if split_attn:

@@ -304,7 +306,7 @@ def run_model(llama_config,

306 return output.cpu()

307

308

307 -def test_toy_llama():

309 +def _test_toy_llama(*, use_inductor):

310 # compare output with and without piecewise compilation

311

312 llama_config = LlamaConfig(hidden_size=128,

@@ -326,8 +328,14 @@ def test_toy_llama():

328 num_backend_compilations=0,

329 num_cudagraph_caputured=0,

330 ):

329 - outputs.append(run_model(llama_config, use_compile=False))

330 - run_model(tractable_config, use_compile=False)

331 + outputs.append(

332 + run_model(llama_config, use_inductor=False, use_compile=False))

333 + run_model(tractable_config, use_inductor=False, use_compile=False)

334 +

335 + if use_inductor:

336 + kwargs = {"num_inductor_compiles": 1, "num_eager_compiles": 0}

337 + else:

338 + kwargs = {"num_eager_compiles": 1, "num_inductor_compiles": 0}

339

340 with compilation_counter.expect(

341 num_graphs_seen=1, # one graph for the model

@@ -336,9 +344,13 @@ def test_toy_llama():

344 num_backend_compilations=1, # num_piecewise_capturable_graphs_seen

345 num_cudagraph_caputured=

346 2, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen

347 + **kwargs,

348 ):

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

341 - run_model(tractable_config, use_compile=True)

349 + outputs.append(

350 + run_model(llama_config,

351 + use_inductor=use_inductor,

352 + use_compile=True))

353 + run_model(tractable_config, use_inductor=use_inductor, use_compile=True)

354

355 with compilation_counter.expect(

356 num_graphs_seen=1, # one graph for the model

@@ -353,13 +365,27 @@ def test_toy_llama():

365 ), # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen

366 ):

367 outputs.append(

356 - run_model(llama_config, use_compile=True, split_attn=True))

357 - run_model(tractable_config, use_compile=True, split_attn=True)

368 + run_model(llama_config,

369 + use_inductor=use_inductor,

370 + use_compile=True,

371 + split_attn=True))

372 + run_model(tractable_config,

373 + use_inductor=use_inductor,

374 + use_compile=True,

375 + split_attn=True)

376

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

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

379

380

381 +def test_toy_llama_inductor():

382 + _test_toy_llama(use_inductor=True)

383 +

384 +

385 +def test_toy_no_inductor():

386 + _test_toy_llama(use_inductor=False)

387 +

388 +

389 @torch.inference_mode

390 def benchmark():

391 from triton.testing import do_bench

vllm/compilation/compiler_interface.py

@@ -12,6 +12,7 @@ import torch._inductor.compile_fx

12 import torch.fx as fx

14 import vllm.envs as envs

15 +from vllm.compilation.counter import compilation_counter

16 from vllm.config import VllmConfig

17 from vllm.utils import is_torch_equal_or_newer

@@ -175,6 +176,7 @@ class InductorStandaloneAdaptor(CompilerInterface):

176 runtime_shape: Optional[int] = None,

177 key: Optional[str] = None,

178 ) -> tuple[Optional[Callable], Optional[Any]]:

179 + compilation_counter.num_inductor_compiles += 1

180 current_config = {}

181 if compiler_config is not None:

182 current_config.update(compiler_config)

@@ -262,6 +264,7 @@ class InductorAdaptor(CompilerInterface):

264 runtime_shape: Optional[int] = None,

265 key: Optional[str] = None,

266 ) -> tuple[Optional[Callable], Optional[Any]]:

267 + compilation_counter.num_inductor_compiles += 1

268 from torch._inductor.compile_fx import compile_fx

269 current_config = {}

270 if compiler_config is not None:

@@ -528,6 +531,7 @@ class EagerAdaptor(CompilerInterface):

531 runtime_shape: Optional[int] = None,

532 key: Optional[str] = None,

533 ) -> tuple[Optional[Callable], Optional[Any]]:

534 + compilation_counter.num_eager_compiles += 1

535 # we don't need to compile the graph, just return the graph itself.

536 # It does not support caching, return None for the handle.

537 return graph, None

vllm/compilation/counter.py

@@ -15,6 +15,10 @@ class CompilationCounter:

15 num_piecewise_capturable_graphs_seen: int = 0

16 num_backend_compilations: int = 0

17 num_cudagraph_caputured: int = 0

18 + # InductorAdapter.compile calls

19 + num_inductor_compiles: int = 0

20 + # EagerAdapter.compile calls

21 + num_eager_compiles: int = 0

23 def clone(self) -> "CompilationCounter":

24 return copy.deepcopy(self)

vllm/config.py

-5

@@ -4315,15 +4315,10 @@ class VllmConfig:

4315 self.compilation_config.custom_ops.append("+rms_norm")

4316 if envs.VLLM_USE_V1 and self.model_config is not None and \

4317 not self.model_config.enforce_eager:

4318 - # NOTE(woosuk): Currently, we use inductor because the piecewise

4319 - # CUDA graphs do not work properly with the custom CUDA kernels.

4320 - # FIXME(woosuk): Disable inductor to reduce the compilation time

4321 - # and avoid any potential issues with the inductor.

4318 # FIXME(rob): Add function to set all of these.

4319 if not self.compilation_config.custom_ops:

4320 self.compilation_config.custom_ops = ["none"]

4321 self.compilation_config.use_cudagraph = True

4326 - self.compilation_config.use_inductor = True

4322 self.compilation_config.cudagraph_num_of_warmups = 1

4323 self.compilation_config.pass_config.enable_fusion = False

4324 self.compilation_config.pass_config.enable_noop = False