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

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

[LoRA] LoRA cuda graph specialization (#25914)

Signed-off-by: Andy Lo <andy@mistral.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Andy Lo committed Oct 20, 2025 at 05:21 UTC b63f2143f8db6ef18c34eb25237fd06c5ad1c450

9 files changed +122 -34

tests/lora/test_chatglm3_tp.py

+15 -6

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

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

4 import vllm

5 +import vllm.config

6 from vllm.lora.request import LoRARequest

8 from ..utils import create_new_process_for_each_test, multi_gpu_test

@@ -53,9 +54,10 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:

54 def test_chatglm3_lora(chatglm3_lora_files):

55 llm = vllm.LLM(

56 MODEL_PATH,

56 - max_model_len=1024,

57 + max_model_len=512,

58 enable_lora=True,

58 - max_loras=4,

59 + max_loras=2,

60 + max_num_seqs=16,

61 max_lora_rank=64,

62 trust_remote_code=True,

63 )

@@ -72,13 +74,17 @@ def test_chatglm3_lora(chatglm3_lora_files):

74 def test_chatglm3_lora_tp4(chatglm3_lora_files):

75 llm = vllm.LLM(

76 MODEL_PATH,

75 - max_model_len=1024,

77 + max_model_len=512,

78 enable_lora=True,

77 - max_loras=4,

79 + max_loras=2,

80 max_lora_rank=64,

81 + max_num_seqs=16,

82 tensor_parallel_size=4,

83 trust_remote_code=True,

84 fully_sharded_loras=False,

85 + compilation_config=vllm.config.CompilationConfig( # Avoid OOM

86 + cudagraph_specialize_lora=False,

87 + ),

88 )

90 output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)

@@ -96,14 +102,17 @@ def test_chatglm3_lora_tp4_fully_sharded_loras(chatglm3_lora_files):

102 # more GPU memory causing vLLM to OOM

103 llm = vllm.LLM(

104 MODEL_PATH,

99 - max_model_len=1024,

105 + max_model_len=512,

106 enable_lora=True,

101 - max_loras=4,

107 + max_loras=2,

108 max_lora_rank=64,

109 tensor_parallel_size=4,

110 trust_remote_code=True,

111 fully_sharded_loras=True,

112 gpu_memory_utilization=0.85,

113 + compilation_config=vllm.config.CompilationConfig( # Avoid OOM

114 + cudagraph_specialize_lora=False,

115 + ),

116 )

117 output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)

118 for i in range(len(EXPECTED_LORA_OUTPUT)):

tests/lora/test_llama_tp.py

+8 -1

@@ -3,7 +3,10 @@

3 import subprocess

4 import sys

6 +import pytest

7 +

8 import vllm

9 +import vllm.config

10 from vllm import LLM

11 from vllm.lora.request import LoRARequest

12 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

@@ -100,7 +103,8 @@ def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None =

103

104

105 @create_new_process_for_each_test()

103 -def test_llama_lora(sql_lora_files):

106 +@pytest.mark.parametrize("cudagraph_specialize_lora", [True, False])

107 +def test_llama_lora(sql_lora_files, cudagraph_specialize_lora: bool):

108 llm = vllm.LLM(

109 MODEL_PATH,

110 tokenizer=sql_lora_files,

@@ -108,6 +112,9 @@ def test_llama_lora(sql_lora_files):

112 # also test odd max_num_seqs

113 max_num_seqs=13,

114 max_loras=4,

115 + compilation_config=vllm.config.CompilationConfig(

116 + cudagraph_specialize_lora=cudagraph_specialize_lora,

117 + ),

118 )

119 generate_and_test(llm, sql_lora_files)

120

vllm/config/compilation.py

@@ -366,6 +366,14 @@ class CompilationConfig:

366 minor release, i.e. v0.11.0 or v1.0.0. Please use cudagraph_mode=

367 FULL_AND_PIECEWISE instead.

368 """

369 + cudagraph_specialize_lora: bool = True

370 + """Whether to create separate cuda graphs for cases with and without active

371 + LoRA adapters. When set to False, the LoRA-enabled cuda graph will be used

372 + for all cases, incurring the overhead of running LoRA ops even when no

373 + adapters are active. Setting this to True will remove this overhead at the

374 + cost of increased startup time and slightly higher memory usage.

375 + When `enable_lora` is False, this option has no effect.

376 + """

377

378 use_inductor_graph_partition: bool = False

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

vllm/forward_context.py

+7 -1

@@ -40,13 +40,19 @@ class BatchDescriptor(NamedTuple):

40 False can also be used for an uniform decode batch to dispatch to the

41 cudagraph supporting non-uniform batches.

42 """

43 + has_lora: bool = False

44 + """

45 + Whether this batch has active LoRA adapters.

46 + """

48 @property

49 def non_uniform(self) -> "BatchDescriptor":

50 """

51 Return a non-uniform version of current batch descriptor.

52 """

49 - return BatchDescriptor(self.num_tokens, uniform_decode=False)

53 + return BatchDescriptor(

54 + self.num_tokens, uniform_decode=False, has_lora=self.has_lora

55 + )

58 def _compute_sp_num_tokens(

vllm/lora/ops/triton_ops/lora_shrink_op.py

@@ -169,6 +169,8 @@ def _lora_shrink(

169 assert lora_ids.size(0) == num_tokens_per_lora.size(0)

170 assert lora_token_start_loc.size(0) == lora_ids.size(0) + 1

171

172 + output_tensor.zero_()

173 +

174 (lora_ptr_tensor, lora_strides_d0, lora_strides_d1, lora_strides_d2) = (

175 _get_lora_a_ptr(lora_a_weights, inputs.device)

176 )

vllm/lora/punica_wrapper/punica_gpu.py

+21 -13

@@ -205,15 +205,18 @@ class PunicaWrapperGPU(PunicaWrapperBase):

205

206 assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)

207

208 - if buffer is None:

209 - r = lora_b_stacked[0].size(-1)

210 - # We set the buffer to be float32 by default, refer to:

211 - # https://github.com/triton-lang/triton/issues/1387

212 - buffer = torch.zeros( # type: ignore

213 - (len(output_slices), x.size(0), r),

214 - dtype=torch.float32,

215 - device=x.device,

216 - )

208 + assert buffer is None, (

209 + "To minimize overhead, the buffer should be created by "

210 + ".add_lora_linear() instead of being passed in."

211 + )

212 + r = lora_b_stacked[0].size(-1)

213 + # We set the buffer to be float32 by default, refer to:

214 + # https://github.com/triton-lang/triton/issues/1387

215 + # Note: buffer is zeroed inside the shrink op

216 + buffer = torch.empty(

217 + (len(output_slices), x.size(0), r), dtype=torch.float32, device=x.device

218 + )

219 +

220 self.add_shrink(

221 buffer, # type: ignore

222 x,

@@ -260,10 +263,15 @@ class PunicaWrapperGPU(PunicaWrapperBase):

263 y = y.view(-1, y.shape[-1])

264 x = x.view(-1, x.shape[-1])

265 r = lora_b_stacked.size(-1)

263 - if buffer is None:

264 - # We set the buffer to be float32 by default, refer to:

265 - # https://github.com/triton-lang/triton/issues/1387

266 - buffer = torch.zeros((x.size(0), r), dtype=torch.float32, device=x.device)

266 +

267 + assert buffer is None, (

268 + "To minimize overhead, the buffer should be created by "

269 + ".add_lora_linear() instead of being passed in."

270 + )

271 + # We set the buffer to be float32 by default, refer to:

272 + # https://github.com/triton-lang/triton/issues/1387

273 + # Note: buffer is zeroed inside the shrink op

274 + buffer = torch.empty((x.size(0), r), dtype=torch.float32, device=x.device)

275

276 lora_shrink(

277 x,

vllm/v1/cudagraph_dispatcher.py

+20 -4

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from itertools import product

5 from vllm.config import CUDAGraphMode, VllmConfig

6 from vllm.forward_context import BatchDescriptor

@@ -67,14 +68,27 @@ class CudagraphDispatcher:

68 ):

69 # This should be called only after attention backend is initialized.

71 + # LoRA activation cases to specialize the cuda graphs on

72 + if self.vllm_config.lora_config:

73 + if self.compilation_config.cudagraph_specialize_lora:

74 + lora_cases = [True, False]

75 + else:

76 + lora_cases = [True]

77 + else:

78 + lora_cases = [False]

79 +

80 # Note: we create all valid keys for cudagraph here but do not

81 # guarantee all keys would be used. For example, if we allow lazy

82 # capturing in future PR, some keys may never be triggered.

83 if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:

74 - for bs in self.compilation_config.cudagraph_capture_sizes:

84 + for bs, has_lora in product(

85 + self.compilation_config.cudagraph_capture_sizes, lora_cases

86 + ):

87 self.add_cudagraph_key(

88 cudagraph_mode.mixed_mode(),

77 - BatchDescriptor(num_tokens=bs, uniform_decode=False),

89 + BatchDescriptor(

90 + num_tokens=bs, uniform_decode=False, has_lora=has_lora

91 + ),

92 )

94 # if decode cudagraph mode is FULL, and we don't already have mixed

@@ -92,10 +106,12 @@ class CudagraphDispatcher:

106 for x in self.compilation_config.cudagraph_capture_sizes

107 if x <= max_num_tokens and x >= uniform_decode_query_len

108 ]

95 - for bs in cudagraph_capture_sizes_for_decode:

109 + for bs, has_lora in product(cudagraph_capture_sizes_for_decode, lora_cases):

110 self.add_cudagraph_key(

111 CUDAGraphMode.FULL,

98 - BatchDescriptor(num_tokens=bs, uniform_decode=True),

112 + BatchDescriptor(

113 + num_tokens=bs, uniform_decode=True, has_lora=has_lora

114 + ),

115 )

116 self.keys_initialized = True

117

vllm/v1/worker/gpu_model_runner.py

+27 -6

@@ -8,6 +8,7 @@ from collections import defaultdict

8 from collections.abc import Iterator

9 from contextlib import contextmanager

10 from copy import deepcopy

11 +from itertools import product

12 from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast

14 import numpy as np

@@ -2469,7 +2470,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2470 num_scheduled_tokens == self.input_batch.num_reqs * max_query_len

2471 )

2472 batch_descriptor = BatchDescriptor(

2472 - num_tokens=num_input_tokens, uniform_decode=uniform_decode

2473 + num_tokens=num_input_tokens,

2474 + uniform_decode=uniform_decode,

2475 + has_lora=len(self.input_batch.lora_id_to_lora_request) > 0,

2476 )

2477 cudagraph_runtime_mode, batch_descriptor = (

2478 self.cudagraph_dispatcher.dispatch(batch_descriptor, use_cascade_attn)

@@ -3193,6 +3196,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3196 is_profile: bool = False,

3197 create_mixed_batch: bool = False,

3198 remove_lora: bool = True,

3199 + activate_lora: bool = False,

3200 ) -> tuple[torch.Tensor, torch.Tensor]:

3201 """

3202 Run a dummy forward pass to warm up/profile run or capture the

@@ -3215,6 +3219,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3219 create_mixed_batch: If True, create a mixed batch with both decode

3220 (1 token) and prefill (multiple tokens) requests.

3221 remove_lora: If False, dummy LoRAs are not destroyed after the run

3222 + activate_lora: If False, dummy_run is performed without LoRAs.

3223 """

3224 assert (

3225 cudagraph_runtime_mode is None

@@ -3364,7 +3369,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3369 attn_metadata[layer_name] = attn_metadata_i

3370

3371 with self.maybe_dummy_run_with_lora(

3367 - self.lora_config, num_scheduled_tokens, remove_lora

3372 + self.lora_config, num_scheduled_tokens, activate_lora, remove_lora

3373 ):

3374 # Make sure padding doesn't exceed max_num_tokens

3375 assert num_tokens_after_padding <= self.max_num_tokens

@@ -3411,6 +3416,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3416 BatchDescriptor(

3417 num_tokens=num_tokens_after_padding,

3418 uniform_decode=uniform_decode,

3419 + has_lora=activate_lora and self.lora_config is not None,

3420 )

3421 )

3422 if not is_profile

@@ -3769,10 +3775,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3775 start_free_gpu_memory = torch.cuda.mem_get_info()[0]

3776 cudagraph_mode = self.compilation_config.cudagraph_mode

3777 assert cudagraph_mode is not None

3778 +

3779 + if self.lora_config:

3780 + if self.compilation_config.cudagraph_specialize_lora:

3781 + lora_cases = [True, False]

3782 + else:

3783 + lora_cases = [True]

3784 + else:

3785 + lora_cases = [False]

3786 +

3787 if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:

3788 cudagraph_runtime_mode = cudagraph_mode.mixed_mode()

3789

3775 - compilation_cases = list(reversed(self.cudagraph_batch_sizes))

3790 + compilation_cases = list(

3791 + product(reversed(self.cudagraph_batch_sizes), lora_cases)

3792 + )

3793 self._capture_cudagraphs(

3794 compilation_cases,

3795 cudagraph_runtime_mode=cudagraph_runtime_mode,

@@ -3793,7 +3810,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3810 for x in self.cudagraph_batch_sizes

3811 if max_num_tokens >= x >= self.uniform_decode_query_len

3812 ]

3796 - compilation_cases_decode = list(reversed(decode_cudagraph_batch_sizes))

3813 + compilation_cases_decode = list(

3814 + product(reversed(decode_cudagraph_batch_sizes), lora_cases)

3815 + )

3816 self._capture_cudagraphs(

3817 compilation_cases=compilation_cases_decode,

3818 cudagraph_runtime_mode=CUDAGraphMode.FULL,

@@ -3823,7 +3842,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3842

3843 def _capture_cudagraphs(

3844 self,

3826 - compilation_cases: list[int],

3845 + compilation_cases: list[tuple[int, bool]],

3846 cudagraph_runtime_mode: CUDAGraphMode,

3847 uniform_decode: bool,

3848 ):

@@ -3844,7 +3863,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3863 )

3864

3865 # We skip EPLB here since we don't want to record dummy metrics

3847 - for num_tokens in compilation_cases:

3866 + for num_tokens, activate_lora in compilation_cases:

3867 # We currently only capture ubatched graphs when its a FULL

3868 # cudagraph, a uniform decode batch, and the number of tokens

3869 # is above the threshold. Otherwise we just capture a non-ubatched

@@ -3875,6 +3894,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3894 allow_microbatching=allow_microbatching,

3895 skip_eplb=True,

3896 remove_lora=False,

3897 + activate_lora=activate_lora,

3898 )

3899 self._dummy_run(

3900 num_tokens,

@@ -3883,6 +3903,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3903 allow_microbatching=allow_microbatching,

3904 skip_eplb=True,

3905 remove_lora=False,

3906 + activate_lora=activate_lora,

3907 )

3908 self.maybe_remove_all_loras(self.lora_config)

3909

vllm/v1/worker/lora_model_runner_mixin.py

+14 -3

@@ -120,7 +120,10 @@ class LoRAModelRunnerMixin:

120

121 @contextmanager

122 def maybe_select_dummy_loras(

123 - self, lora_config: LoRAConfig | None, num_scheduled_tokens: np.ndarray

123 + self,

124 + lora_config: LoRAConfig | None,

125 + num_scheduled_tokens: np.ndarray,

126 + activate_lora: bool = True,

127 ):

128 if lora_config is None:

129 yield

@@ -133,7 +136,12 @@ class LoRAModelRunnerMixin:

136

137 # Make prompt lora mapping

138 # Assign LoRA IDs cyclically to simulate a worst-case scenario.

136 - prompt_lora_mapping = (np.arange(num_reqs, dtype=np.int32) % num_loras) + 1

139 + if activate_lora:

140 + prompt_lora_mapping = (

141 + np.arange(num_reqs, dtype=np.int32) % num_loras

142 + ) + 1

143 + else:

144 + prompt_lora_mapping = np.zeros(num_reqs, dtype=np.int32)

145

146 # Make token lora mapping

147 token_lora_mapping = np.repeat(prompt_lora_mapping, num_scheduled_tokens)

@@ -159,11 +167,14 @@ class LoRAModelRunnerMixin:

167 self,

168 lora_config: LoRAConfig | None,

169 num_scheduled_tokens: np.ndarray,

170 + activate_lora: bool = True,

171 remove_lora: bool = True,

172 ):

173 with (

174 self.maybe_setup_dummy_loras(lora_config, remove_lora),

166 - self.maybe_select_dummy_loras(lora_config, num_scheduled_tokens),

175 + self.maybe_select_dummy_loras(

176 + lora_config, num_scheduled_tokens, activate_lora

177 + ),

178 ):

179 yield

180