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

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 9, 02:30 UTC

Reduce the kernel overhead when num of active loras is smaller than max loras. Multiple cuda graphs are captured for each num of active-loras. (#32005)

Reduce the kernel overhead when num of active loras is smaller than max loras. Multiple cuda graphs are captured for each num of active-loras. (#32005) Signed-off-by: Yu Gong <yu3.gong@gmail.com>

yugong333 committed Feb 2, 2026 at 09:30 UTC ffe1fc7a28841973135b981fb68ce515b409a236

15 files changed +323 -66

tests/lora/test_fused_moe_lora_kernel.py

+10

@@ -181,6 +181,10 @@ def use_fused_moe_lora_kernel(

181 expert_ids = expert_ids.view(max_loras, -1)

182 sorted_token_ids = sorted_token_ids.view(max_loras, -1)

183

184 + # num_active_loras is the number of active LoRAs

185 + # (max_loras + 1 to include no-lora case)

186 + num_active_loras = max_loras + 1

187 +

188 fused_moe_lora(

189 output,

190 hidden_states,

@@ -194,6 +198,7 @@ def use_fused_moe_lora_kernel(

198 max_lora_rank,

199 top_k_num,

200 lora_ids,

201 + num_active_loras,

202 adapter_enabled,

203 config["BLOCK_SIZE_M"],

204 config["BLOCK_SIZE_N"],

@@ -376,6 +381,10 @@ def use_fused_moe_lora_kernel_naive(

381 adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)

382 lora_ids = torch.arange(max_loras + 2, dtype=torch.int32)

383

384 + # num_active_loras is the number of active LoRAs

385 + # (max_loras + 1 to include no-lora case)

386 + num_active_loras = max_loras + 1

387 +

388 fused_moe_lora(

389 output,

390 hidden_states,

@@ -389,6 +398,7 @@ def use_fused_moe_lora_kernel_naive(

398 max_lora_rank,

399 top_k_num,

400 lora_ids,

401 + num_active_loras,

402 adapter_enabled,

403 config["BLOCK_SIZE_M"],

404 config["BLOCK_SIZE_N"],

tests/lora/test_punica_ops.py

+2 -2

@@ -161,7 +161,7 @@ def check_lora_shrink_kernel(

161 data.inputs_tensor,

162 data.lora_weights,

163 out_tensor,

164 - *lora_meta.meta_args(token_nums=token_nums),

164 + *lora_meta.meta_args(token_nums=token_nums, specialize_active_lora=False),

165 scaling,

166 )

167

@@ -234,7 +234,7 @@ def check_lora_expand_kernel(

234 data.inputs_tensor,

235 data.lora_weights,

236 out_tensor,

237 - *lora_meta.meta_args(token_nums=token_nums),

237 + *lora_meta.meta_args(token_nums=token_nums, specialize_active_lora=False),

238 offset_start=0,

239 add_inputs=add_inputs,

240 )

tests/v1/cudagraph/test_cudagraph_dispatch.py

+13 -2

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

17 SchedulerConfig,

18 VllmConfig,

19 )

20 +from vllm.config.lora import LoRAConfig

21 from vllm.forward_context import BatchDescriptor, set_forward_context

22 from vllm.platforms import current_platform

23 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

@@ -47,6 +48,12 @@ def _create_vllm_config(

48 mock_config.speculative_config = None # No speculative decoding

49 if not lora_config:

50 mock_config.lora_config = None

51 + else:

52 + # Create a real LoRAConfig with specialize_active_lora enabled

53 + mock_config.lora_config = LoRAConfig(

54 + max_loras=4,

55 + specialize_active_lora=True,

56 + )

57 # Mimic the behavior of VllmConfig.__post_init__()

58 if compilation_config.mode == CompilationMode.VLLM_COMPILE:

59 compilation_config.set_splitting_ops_for_v1(

@@ -106,15 +113,19 @@ class TestCudagraphDispatcher:

113 )

114

115 # Verify the key is initialized correctly

116 + # With LoRA specialization (max_loras=4, specialize_active_lora=True):

117 + # - lora_cases = [0, 1, 2, 4, 5] (no-lora + powers of 2 up to 4 + max_loras+1)

118 + # - capture_sizes = [1, 8]

119 + # - Total keys = 2 sizes × 5 lora_cases = 10

120 if cudagraph_mode_str in ["FULL_AND_PIECEWISE", "PIECEWISE"]:

121 assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == (

111 - 4 if lora_config else 2

122 + 10 if lora_config else 2

123 )

124 else:

125 assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == 0

126 if cudagraph_mode_str not in ["NONE", "PIECEWISE"]:

127 assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == (

117 - 4 if lora_config else 2

128 + 10 if lora_config else 2

129 )

130 else:

131 assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == 0

vllm/config/lora.py

@@ -60,6 +60,13 @@ class LoRAConfig:

60 of multimodal models will be enabled. This is an experimental feature and

61 currently only supports some MM models such as the Qwen VL series. The default

62 is False."""

63 + specialize_active_lora: bool = False

64 + """Whether to construct lora kernel grid by the number of active LoRA adapters.

65 + When set to True, separate cuda graphs will be captured for different counts

66 + of active LoRAs (powers of 2 up to max_loras), which can improve performance

67 + for variable LoRA usage patterns at the cost of increased startup time and

68 + memory usage. Only takes effect when cudagraph_specialize_lora is True.

69 + """

71 def compute_hash(self) -> str:

72 """

vllm/engine/arg_utils.py

@@ -485,6 +485,7 @@ class EngineArgs:

485 max_cpu_loras: int | None = LoRAConfig.max_cpu_loras

486 lora_dtype: str | torch.dtype | None = LoRAConfig.lora_dtype

487 enable_tower_connector_lora: bool = LoRAConfig.enable_tower_connector_lora

488 + specialize_active_lora: bool = LoRAConfig.specialize_active_lora

489

490 ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight

491 num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override

@@ -1026,6 +1027,9 @@ class EngineArgs:

1027 "--fully-sharded-loras", **lora_kwargs["fully_sharded_loras"]

1028 )

1029 lora_group.add_argument("--default-mm-loras", **lora_kwargs["default_mm_loras"])

1030 + lora_group.add_argument(

1031 + "--specialize-active-lora", **lora_kwargs["specialize_active_lora"]

1032 + )

1033

1034 # Observability arguments

1035 observability_kwargs = get_kwargs(ObservabilityConfig)

@@ -1657,6 +1661,7 @@ class EngineArgs:

1661 fully_sharded_loras=self.fully_sharded_loras,

1662 lora_dtype=self.lora_dtype,

1663 enable_tower_connector_lora=self.enable_tower_connector_lora,

1664 + specialize_active_lora=self.specialize_active_lora,

1665 max_cpu_loras=self.max_cpu_loras

1666 if self.max_cpu_loras and self.max_cpu_loras > 0

1667 else None,

vllm/forward_context.py

+13 -1

@@ -47,6 +47,14 @@ class BatchDescriptor(NamedTuple):

47 """

48 Whether this batch has active LoRA adapters.

49 """

50 + num_active_loras: int = 0

51 + """

52 + Number of distinct active LoRA adapters in this batch.

53 + When cudagraph_specialize_lora_count is enabled, separate CUDA graphs

54 + are captured for each num_active_loras value. This allows kernels

55 + (like fused_moe_lora) whose grid size depends on num_active_loras

56 + to be properly captured.

57 + """

59 def relax_for_mixed_batch_cudagraphs(self) -> "BatchDescriptor":

60 """

@@ -54,7 +62,11 @@ class BatchDescriptor(NamedTuple):

62 with PIECEWISE cudagraphs (or mixed prefill-decode FA cudagraphs).

63 """

64 return BatchDescriptor(

57 - self.num_tokens, num_reqs=None, uniform=False, has_lora=self.has_lora

65 + self.num_tokens,

66 + num_reqs=None,

67 + uniform=False,

68 + has_lora=self.has_lora,

69 + num_active_loras=self.num_active_loras,

70 )

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

+16 -4

@@ -95,7 +95,7 @@ def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):

97 def _adjust_kernel_inputs(

98 - max_loras: int,

98 + num_active_loras: int,

99 sorted_token_ids: torch.Tensor | None,

100 expert_ids: torch.Tensor,

101 ):

@@ -109,7 +109,7 @@ def _adjust_kernel_inputs(

109 else:

110 stride_tl = sorted_token_ids.stride(0)

111 stride_el = expert_ids.stride(0)

112 - grid_lora_dim = max_loras + 1

112 + grid_lora_dim = num_active_loras

113 return grid_lora_dim, stride_tl, stride_el

114

115

@@ -354,6 +354,7 @@ def _fused_moe_lora_shrink(

354 num_warps: int,

355 num_stages: int,

356 split_k: int,

357 + num_active_loras: int,

358 mul_routed_weight: bool = False,

359 use_gdc: bool = False,

360 ) -> None:

@@ -373,7 +374,7 @@ def _fused_moe_lora_shrink(

374 b_ptr = _get_ptr(lora_a_stacked, device)

375

376 grid_lora_dim, stride_tl, stride_el = _adjust_kernel_inputs(

376 - w1_lora_a_stacked.shape[0], sorted_token_ids, expert_ids

377 + num_active_loras, sorted_token_ids, expert_ids

378 )

379 grid = lambda META: (

380 split_k

@@ -457,6 +458,7 @@ def _fused_moe_lora_expand(

458 num_warps: int,

459 num_stages: int,

460 split_k: int,

461 + num_active_loras: int,

462 mul_routed_weight: bool = False,

463 offset: int = 0,

464 use_gdc: bool = False,

@@ -484,7 +486,7 @@ def _fused_moe_lora_expand(

486 }

487

488 grid_lora_dim, stride_tl, stride_el = _adjust_kernel_inputs(

487 - w1_lora_b_stacked.shape[0], sorted_token_ids, expert_ids

489 + num_active_loras, sorted_token_ids, expert_ids

490 )

491

492 grid = lambda META: (

@@ -557,6 +559,7 @@ def _fused_moe_lora(

559 max_lora_rank: int,

560 top_k_num: int,

561 lora_ids: torch.Tensor,

562 + num_active_loras: int,

563 adapter_enabled: torch.Tensor,

564 shrink_block_size_m: int,

565 shrink_block_size_n: int,

@@ -648,6 +651,7 @@ def _fused_moe_lora(

651 shrink_num_warps,

652 shrink_num_stages,

653 shrink_split_k,

654 + num_active_loras,

655 mul_routed_weight,

656 use_gdc=use_gdc,

657 )

@@ -695,6 +699,7 @@ def _fused_moe_lora(

699 expand_num_warps,

700 expand_num_stages,

701 expand_split_k,

702 + num_active_loras,

703 mul_routed_weight,

704 offset,

705 use_gdc=use_gdc,

@@ -714,6 +719,7 @@ def _fused_moe_lora_fake(

719 max_lora_rank: int,

720 top_k_num: int,

721 lora_ids: torch.Tensor,

722 + num_active_loras: int,

723 adapter_enabled: torch.Tensor,

724 shrink_block_size_m: int,

725 shrink_block_size_n: int,

@@ -730,6 +736,8 @@ def _fused_moe_lora_fake(

736 expand_num_stages: int,

737 expand_split_k: int,

738 mul_routed_weight: bool = False,

739 + fully_sharded: bool = False,

740 + offset: int = 0,

741 ) -> None:

742 return

743

@@ -761,6 +769,7 @@ def _fused_moe_lora_shrink_fake(

769 num_warps: int,

770 num_stages: int,

771 split_k: int,

772 + num_active_loras: int,

773 mul_routed_weight: bool = False,

774 use_gdc: bool = False,

775 ) -> None:

@@ -770,6 +779,7 @@ def _fused_moe_lora_shrink_fake(

779 def _fused_moe_lora_expand_fake(

780 output: torch.Tensor,

781 a_intermediate_cache1: torch.Tensor,

782 + b_intermediate_cache1: torch.Tensor,

783 lora_b_stacked: list[torch.Tensor],

784 topk_weights: torch.Tensor,

785 sorted_token_ids: torch.Tensor | None,

@@ -796,7 +806,9 @@ def _fused_moe_lora_expand_fake(

806 num_warps: int,

807 num_stages: int,

808 split_k: int,

809 + num_active_loras: int,

810 mul_routed_weight: bool = False,

811 + offset: int = 0,

812 use_gdc: bool = False,

813 ) -> None:

814 return

vllm/lora/ops/triton_ops/lora_expand_op.py

+3 -4

@@ -138,6 +138,7 @@ def _lora_expand(

138 lora_token_start_loc: torch.Tensor, # shape [max-loras + 2]

139 lora_ids: torch.Tensor, # shape [max-loras + 1]

140 no_lora_flag_cpu: torch.Tensor, # shape [1]

141 + num_active_loras: int, # number of active LoRAs (unused here, for API compat)

142 offset_start: int = 0,

143 add_inputs: bool = False,

144 ) -> None:

@@ -234,10 +235,7 @@ def _lora_expand(

235 grid = (

236 triton.cdiv(M, BLOCK_M) * triton.cdiv(MAX_N, BLOCK_N),

237 NUM_SLICES,

237 - # Each LoRA receives its own set of thread blocks for output

238 - # computation. If some LoRA doesn't have any tokens to process, its

239 - # thread blocks simply exit.

240 - MAX_LORAS,

238 + num_active_loras,

239 )

240 # We disable PDL temporarily because LoRA kernels are not launching back-to-back,

241 # making PDL invalid and affecting the kernel performance.

@@ -291,6 +289,7 @@ def _lora_expand_fake(

289 lora_token_start_loc: torch.Tensor,

290 lora_ids: torch.Tensor,

291 no_lora_flag_cpu: torch.Tensor,

292 + num_active_loras: int,

293 offset_start: int = 0,

294 add_inputs: bool = False,

295 ) -> None:

vllm/lora/ops/triton_ops/lora_kernel_metadata.py

+36 -3

@@ -4,7 +4,8 @@

4 LoRA kernels metadata preparation utilities.

5 """

7 -from dataclasses import dataclass

7 +import bisect

8 +from dataclasses import dataclass, field

10 import torch

@@ -28,9 +29,22 @@ class LoRAKernelMeta:

29 # to early exit from inside the lora_expand / lora_shrink torch operation.

30 no_lora_flag_cpu: torch.Tensor

32 + # Number of active LoRAs (unique non-(-1) values in token_lora_mapping)

33 + # Stored as a Python int to avoid GPU->CPU sync during forward pass

34 + num_active_loras: int = 0

35 +

36 + # Captured LoRA counts for cudagraph specialization (sorted list).

37 + # When specialize_active_lora is enabled, num_active_loras is rounded up

38 + # to the nearest value in this list to match cudagraph capture keys.

39 + # Empty list means no specialization (use actual count).

40 + captured_lora_counts: list[int] = field(default_factory=list)

41 +

42 @staticmethod

43 def make(

33 - max_loras: int, max_num_tokens: int, device: torch.device | str

44 + max_loras: int,

45 + max_num_tokens: int,

46 + device: torch.device | str,

47 + captured_lora_counts: list[int] | None = None,

48 ) -> "LoRAKernelMeta":

49 token_lora_mapping = torch.empty(

50 max_num_tokens, dtype=torch.int32, device=device

@@ -66,6 +80,9 @@ class LoRAKernelMeta:

80 num_tokens_per_lora=num_tokens_per_lora,

81 lora_token_start_loc=lora_token_start_loc,

82 no_lora_flag_cpu=no_lora_flag_cpu,

83 + captured_lora_counts=sorted(captured_lora_counts)

84 + if captured_lora_counts

85 + else [],

86 )

88 def _reset(self):

@@ -73,6 +90,8 @@ class LoRAKernelMeta:

90 self.num_tokens_per_lora.fill_(0)

91 self.lora_token_start_loc.fill_(0)

92 self.no_lora_flag_cpu.fill_(False)

93 + self.num_active_loras = 0

94 + self.captured_lora_counts = []

96 def prepare_tensors(self, token_lora_mapping: torch.Tensor) -> None:

97 """

@@ -118,6 +137,15 @@ class LoRAKernelMeta:

137 num_tokens_per_lora, non_blocking=True

138 )

139

140 + self.num_active_loras = lora_ids.size(0)

141 +

142 + # Round up num_active_loras to match cudagraph capture keys.

143 + # This ensures the kernel grid dimension matches the captured graph.

144 + if self.captured_lora_counts and self.num_active_loras > 0:

145 + idx = bisect.bisect_left(self.captured_lora_counts, self.num_active_loras)

146 + if idx < len(self.captured_lora_counts):

147 + self.num_active_loras = self.captured_lora_counts[idx]

148 +

149 # lora_token_start_loc

150 lora_token_start_loc = torch.cumsum(num_tokens_per_lora, dim=0)

151 self.lora_token_start_loc[1 : 1 + lora_token_start_loc.size(0)].copy_(

@@ -125,7 +153,9 @@ class LoRAKernelMeta:

153 )

154

155 def meta_args(

128 - self, token_nums: int

156 + self,

157 + token_nums: int,

158 + specialize_active_lora: bool,

159 ) -> tuple[

160 torch.Tensor,

161 torch.Tensor,

@@ -133,6 +163,7 @@ class LoRAKernelMeta:

163 torch.Tensor,

164 torch.Tensor,

165 torch.Tensor,

166 + int,

167 ]:

168 """

169 This function returns the kernel metadata required for the current

@@ -144,6 +175,7 @@ class LoRAKernelMeta:

175 token_nums (int): Number of input tokens in the current forward

176 pass of the kernel.

177 """

178 + max_loras = self.active_lora_ids.size(0) - 1

179 return (

180 self.token_lora_mapping[:token_nums],

181 self.token_indices_sorted_by_lora_ids[:token_nums],

@@ -151,4 +183,5 @@ class LoRAKernelMeta:

183 self.lora_token_start_loc,

184 self.active_lora_ids,

185 self.no_lora_flag_cpu,

186 + self.num_active_loras if specialize_active_lora else max_loras + 1,

187 )

vllm/lora/ops/triton_ops/lora_shrink_op.py

+3 -4

@@ -134,6 +134,7 @@ def _lora_shrink(

134 lora_token_start_loc: torch.Tensor, # shape [max-loras + 2]

135 lora_ids: torch.Tensor, # shape [max-loras + 1]

136 no_lora_flag_cpu: torch.Tensor, # shape [1]

137 + num_active_loras: int, # number of active LoRAs (unused here, for API compat)

138 scaling: float,

139 ) -> None:

140 """

@@ -214,10 +215,7 @@ def _lora_shrink(

215 grid = (

216 SPLIT_K * triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),

217 NUM_SLICES,

217 - # Each LoRA receives its own set of thread blocks for output

218 - # computation. If some LoRA doesn't have any tokens to process, its

219 - # thread blocks exit early.

220 - MAX_LORAS,

218 + num_active_loras,

219 )

220 # We disable PDL temporarily because LoRA kernels are not launching back-to-back,

221 # making PDL invalid and affecting the kernel performance.

@@ -269,6 +267,7 @@ def _lora_shrink_fake(

267 lora_token_start_loc: torch.Tensor,

268 lora_ids: torch.Tensor,

269 no_lora_flag_cpu: torch.Tensor,

270 + num_active_loras: int,

271 scaling: float,

272 ) -> None:

273 return

vllm/lora/punica_wrapper/punica_gpu.py

+38 -10

@@ -12,6 +12,7 @@ from typing import final

12 import torch

14 from vllm.lora.layers import LoRAMapping

15 +from vllm.lora.utils import get_captured_lora_counts

16 from vllm.triton_utils import HAS_TRITON, triton

17 from vllm.utils.math_utils import round_up

@@ -48,8 +49,16 @@ class PunicaWrapperGPU(PunicaWrapperBase):

49 self.lora_config = kwargs["lora_config"]

50 self.max_loras = self.lora_config.max_loras

52 + # Compute captured LoRA counts for cudagraph specialization.

53 + captured_lora_counts = get_captured_lora_counts(

54 + self.max_loras, self.lora_config.specialize_active_lora

55 + )

56 +

57 self.token_mapping_meta = LoRAKernelMeta.make(

52 - self.max_loras, max_num_batched_tokens, device=device

58 + self.max_loras,

59 + max_num_batched_tokens,

60 + device=device,

61 + captured_lora_counts=captured_lora_counts,

62 )

64 # When speculative decoding is enabled, max_num_samples is

@@ -57,7 +66,10 @@ class PunicaWrapperGPU(PunicaWrapperBase):

66 # This line can be optimized by replacing max_num_batched_tokens

67 # to max_batches * (num_speculative_decoding_tokens + 1).

68 self.prompt_mapping_meta = LoRAKernelMeta.make(

60 - self.max_loras, max_num_batched_tokens, device=device

69 + self.max_loras,

70 + max_num_batched_tokens,

71 + device=device,

72 + captured_lora_counts=captured_lora_counts,

73 )

75 def update_metadata(

@@ -102,7 +114,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):

114 x,

115 lora_a_stacked,

116 y,

105 - *self.token_mapping_meta.meta_args(x.size(0)),

117 + *self.token_mapping_meta.meta_args(

118 + x.size(0), self.lora_config.specialize_active_lora

119 + ),

120 scale,

121 )

122

@@ -143,7 +157,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):

157 x,

158 lora_b_stacked,

159 y,

146 - *self.token_mapping_meta.meta_args(num_tokens),

160 + *self.token_mapping_meta.meta_args(

161 + num_tokens, self.lora_config.specialize_active_lora

162 + ),

163 offset_start=offset_start,

164 add_inputs=True,

165 )

@@ -175,7 +191,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):

191 x.unsqueeze(dim=0),

192 (lora_b_stacked,),

193 y,

178 - *self.token_mapping_meta.meta_args(x.size(0)),

194 + *self.token_mapping_meta.meta_args(

195 + x.size(0), self.lora_config.specialize_active_lora

196 + ),

197 offset_start=0,

198 add_inputs=add_inputs,

199 )

@@ -287,7 +305,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):

305 x,

306 [lora_a_stacked],

307 buffer.unsqueeze(dim=0),

290 - *self.prompt_mapping_meta.meta_args(x.size(0)),

308 + *self.prompt_mapping_meta.meta_args(

309 + x.size(0), self.lora_config.specialize_active_lora

310 + ),

311 scale,

312 )

313

@@ -295,7 +315,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):

315 buffer.unsqueeze(dim=0),

316 [lora_b_stacked],

317 y,

298 - *self.prompt_mapping_meta.meta_args(buffer.size(0)),

318 + *self.prompt_mapping_meta.meta_args(

319 + buffer.size(0), self.lora_config.specialize_active_lora

320 + ),

321 add_inputs=True,

322 )

323 y = y.view_as(y_org)

@@ -316,8 +338,10 @@ class PunicaWrapperGPU(PunicaWrapperBase):

338 Aligns tokens and experts into block-sized chunks for LoRA-based

339 mixture-of-experts (MoE) execution.

340 """

319 - (token_lora_mapping, _, _, _, lora_ids, _) = self.token_mapping_meta.meta_args(

320 - num_tokens

341 + (token_lora_mapping, _, _, _, lora_ids, _, _) = (

342 + self.token_mapping_meta.meta_args(

343 + num_tokens, self.lora_config.specialize_active_lora

344 + )

345 )

346 if naive_block_assignment:

347 expert_ids = topk_ids.reshape(-1)

@@ -392,7 +416,10 @@ class PunicaWrapperGPU(PunicaWrapperBase):

416 _,

417 lora_ids,

418 _,

395 - ) = self.token_mapping_meta.meta_args(x.size(0))

419 + num_active_loras,

420 + ) = self.token_mapping_meta.meta_args(

421 + x.size(0), self.lora_config.specialize_active_lora

422 + )

423 if token_lora_mapping is None:

424 token_lora_mapping = token_lora_mapping_meta

425 fused_moe_lora(

@@ -408,6 +435,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

435 max_lora_rank,

436 top_k_num,

437 lora_ids,

438 + num_active_loras,

439 adapter_enabled,

440 shrink_config.get("BLOCK_SIZE_M", 64),

441 shrink_config.get("BLOCK_SIZE_N", 64),

vllm/lora/utils.py

+19

@@ -44,6 +44,25 @@ if TYPE_CHECKING:

45 logger = init_logger(__name__)

47 +

48 +def get_captured_lora_counts(max_loras: int, specialize: bool) -> list[int]:

49 + """

50 + Returns num_active_loras values for cudagraph capture.

51 +

52 + When specialize=True: powers of 2 up to max_loras, plus max_loras + 1.

53 + When specialize=False: just [max_loras + 1].

54 +

55 + This is the single source of truth for LoRA capture cases, used by both

56 + CudagraphDispatcher and PunicaWrapperGPU.

57 + """

58 + if not specialize:

59 + return [max_loras + 1]

60 +

61 + return [

62 + n for n in range(1, max_loras + 2) if (n & (n - 1)) == 0 or n == max_loras + 1

63 + ]

64 +

65 +

66 _GLOBAL_LORA_ID = 0

vllm/v1/cudagraph_dispatcher.py

+65 -14

@@ -5,6 +5,7 @@ from itertools import product

5 from vllm.config import CUDAGraphMode, VllmConfig

6 from vllm.forward_context import BatchDescriptor

7 from vllm.logger import init_logger

8 +from vllm.lora.utils import get_captured_lora_counts

10 logger = init_logger(__name__)

@@ -57,6 +58,11 @@ class CudagraphDispatcher:

58 )

60 self.keys_initialized = False

61 + self.specialize_lora_count = (

62 + self.vllm_config.lora_config.specialize_active_lora

63 + if self.vllm_config.lora_config is not None

64 + else False

65 + )

66 # Default cudagraph_mode to NONE until initialize_cudagraph_keys is called

67 self.cudagraph_mode = CUDAGraphMode.NONE

@@ -92,8 +98,33 @@ class CudagraphDispatcher:

98 "Use values from cudagraph_capture_sizes."

99 )

100

101 + def _get_lora_cases(self) -> list[int]:

102 + """

103 + Returns list of has_lora values for CUDA graph capture.

104 + This is the single source of truth for LoRA capture cases.

105 + """

106 + lora_config = self.vllm_config.lora_config

107 + if lora_config is None:

108 + # No LoRA configured - single case with no LoRA

109 + return [0]

110 +

111 + # LoRA is enabled - capture graphs based on cudagraph_specialize_lora

112 + if self.compilation_config.cudagraph_specialize_lora:

113 + captured_counts = get_captured_lora_counts(

114 + lora_config.max_loras, self.specialize_lora_count

115 + )

116 + # Specialize: capture separate graphs for with and without LoRA

117 + return [0] + captured_counts

118 + else:

119 + # No specialization: only capture graphs with LoRA active

120 + return [lora_config.max_loras + 1]

121 +

122 def _create_padded_batch_descriptor(

96 - self, num_tokens: int, uniform_decode: bool, has_lora: bool

123 + self,

124 + num_tokens: int,

125 + uniform_decode: bool,

126 + has_lora: bool,

127 + num_active_loras: int = 0,

128 ) -> BatchDescriptor:

129 max_num_seqs = self.vllm_config.scheduler_config.max_num_seqs

130 uniform_decode_query_len = self.uniform_decode_query_len

@@ -111,6 +142,7 @@ class CudagraphDispatcher:

142 num_reqs=num_reqs,

143 uniform=uniform_decode,

144 has_lora=has_lora,

145 + num_active_loras=num_active_loras,

146 )

147

148 def add_cudagraph_key(

@@ -135,26 +167,23 @@ class CudagraphDispatcher:

167

168 self._compute_bs_to_padded_graph_size()

169

138 - # LoRA activation cases to specialize the cuda graphs on

139 - if self.vllm_config.lora_config:

140 - if self.compilation_config.cudagraph_specialize_lora:

141 - lora_cases = [True, False]

142 - else:

143 - lora_cases = [True]

144 - else:

145 - lora_cases = [False]

170 + # Get LoRA cases to capture

171 + lora_cases = self._get_lora_cases()

172 + self.captured_lora_counts = [

173 + lora_count for lora_count in lora_cases if lora_count

174 + ]

175

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

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

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

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

151 - for bs, has_lora in product(

180 + for bs, num_active_loras in product(

181 self.compilation_config.cudagraph_capture_sizes, lora_cases

182 ):

183 self.add_cudagraph_key(

184 cudagraph_mode.mixed_mode(),

185 self._create_padded_batch_descriptor(

157 - bs, False, has_lora

186 + bs, False, num_active_loras > 0, num_active_loras

187 ).relax_for_mixed_batch_cudagraphs(),

188 )

189

@@ -173,10 +202,14 @@ class CudagraphDispatcher:

202 for x in self.compilation_config.cudagraph_capture_sizes

203 if x <= max_num_tokens and x >= uniform_decode_query_len

204 ]

176 - for bs, has_lora in product(cudagraph_capture_sizes_for_decode, lora_cases):

205 + for bs, num_active_loras in product(

206 + cudagraph_capture_sizes_for_decode, lora_cases

207 + ):

208 self.add_cudagraph_key(

209 CUDAGraphMode.FULL,

179 - self._create_padded_batch_descriptor(bs, True, has_lora),

210 + self._create_padded_batch_descriptor(

211 + bs, True, num_active_loras > 0, num_active_loras

212 + ),

213 )

214

215 self.keys_initialized = True

@@ -187,6 +220,7 @@ class CudagraphDispatcher:

220 uniform_decode: bool = False,

221 has_lora: bool = False,

222 disable_full: bool = False,

223 + num_active_loras: int = 0,

224 ) -> tuple[CUDAGraphMode, BatchDescriptor]:

225 """

226 Given conditions(e.g.,batch descriptor and if using piecewise only),

@@ -202,6 +236,7 @@ class CudagraphDispatcher:

236 disable_full: If True, skip FULL cudagraph checks and

237 return PIECEWISE or NONE only. (can be used for features like

238 cascade attention that are not supported by full cudagraphs)

239 + num_active_loras: Number of distinct active LoRA adapters.

240 """

241 if (

242 not self.keys_initialized

@@ -210,8 +245,24 @@ class CudagraphDispatcher:

245 ):

246 return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)

247

248 + effective_num_active_loras = num_active_loras

249 + if has_lora and num_active_loras > 0:

250 + if self.specialize_lora_count:

251 + # Find the smallest captured `num_active_loras` that is >= the current

252 + # `num_active_loras`. This is because we only capture graphs for

253 + # a subset of possible `num_active_loras` values (powers of 2).

254 + import bisect

255 +

256 + idx = bisect.bisect_left(self.captured_lora_counts, num_active_loras)

257 + if idx < len(self.captured_lora_counts):

258 + effective_num_active_loras = self.captured_lora_counts[idx]

259 + else:

260 + # When not specializing, graphs are captured only with max_loras + 1,

261 + # so we must use max_loras + 1 for dispatch to find a matching graph.

262 + effective_num_active_loras = self.vllm_config.lora_config.max_loras + 1

263 +

264 batch_desc = self._create_padded_batch_descriptor(

214 - num_tokens, uniform_decode, has_lora

265 + num_tokens, uniform_decode, has_lora, effective_num_active_loras

266 )

267 relaxed_batch_desc = batch_desc.relax_for_mixed_batch_cudagraphs()

268

vllm/v1/worker/gpu_model_runner.py

+23 -12

@@ -3082,6 +3082,7 @@ class GPUModelRunner(

3082 # be improved in model runner v2)

3083 force_uniform_decode: bool | None = None,

3084 force_has_lora: bool | None = None,

3085 + force_num_active_loras: int | None = None,

3086 num_encoder_reqs: int = 0,

3087 ) -> tuple[

3088 CUDAGraphMode,

@@ -3103,11 +3104,13 @@ class GPUModelRunner(

3104 self.model_config.is_encoder_decoder and num_encoder_reqs > 0

3105 )

3106

3106 - has_lora = (

3107 - len(self.input_batch.lora_id_to_lora_request) > 0

3108 - if force_has_lora is None

3109 - else force_has_lora

3107 + # Compute LoRA state for cudagraph dispatch

3108 + num_active_loras = (

3109 + force_num_active_loras

3110 + if force_num_active_loras is not None

3111 + else len(self.input_batch.lora_id_to_lora_request)

3112 )

3113 + has_lora = num_active_loras > 0 if force_has_lora is None else force_has_lora

3114

3115 num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens)

3116 dispatch_cudagraph = (

@@ -3116,6 +3119,7 @@ class GPUModelRunner(

3119 has_lora=has_lora,

3120 uniform_decode=uniform_decode,

3121 disable_full=disable_full,

3122 + num_active_loras=num_active_loras,

3123 )

3124 if not force_eager

3125 else (CUDAGraphMode.NONE, BatchDescriptor(num_tokens_padded))

@@ -4606,8 +4610,8 @@ class GPUModelRunner(

4610 is_profile: bool = False,

4611 create_mixed_batch: bool = False,

4612 remove_lora: bool = True,

4609 - activate_lora: bool = False,

4613 is_graph_capturing: bool = False,

4614 + num_active_loras: int = 0,

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

4616 """

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

@@ -4630,7 +4634,8 @@ class GPUModelRunner(

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

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

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

4633 - activate_lora: If False, dummy_run is performed without LoRAs.

4637 + num_active_loras: Number of distinct active LoRAs to capture for.

4638 + LoRA is activated when num_active_loras > 0.

4639 """

4640 mm_config = self.vllm_config.model_config.multimodal_config

4641 if mm_config and mm_config.mm_encoder_only:

@@ -4712,7 +4717,10 @@ class GPUModelRunner(

4717 # `force_has_lora` is used for cudagraph capture; because LoRA is

4718 # activated later in the context manager, but we need to know the

4719 # LoRA state when determining the batch descriptor for capture

4715 - force_has_lora=activate_lora,

4720 + force_has_lora=num_active_loras > 0,

4721 + # `force_num_active_loras` is used for cudagraph capture; because we

4722 + # need to capture graphs for specific num_active_loras counts

4723 + force_num_active_loras=num_active_loras,

4724 )

4725 )

4726

@@ -4782,8 +4790,8 @@ class GPUModelRunner(

4790 self.lora_config,

4791 num_scheduled_tokens,

4792 num_sampled_tokens,

4785 - activate_lora,

4793 remove_lora,

4794 + num_active_loras,

4795 ):

4796 # Make sure padding doesn't exceed max_num_tokens

4797 assert num_tokens_padded <= self.max_num_tokens

@@ -4884,7 +4892,10 @@ class GPUModelRunner(

4892 # lora cases when cudagraph_specialize_lora is enabled. This is a

4893 # short term mitigation for issue mentioned in

4894 # https://github.com/vllm-project/vllm/issues/28334

4887 - if self.compilation_config.cudagraph_specialize_lora and activate_lora:

4895 + if (

4896 + self.compilation_config.cudagraph_specialize_lora

4897 + and num_active_loras > 0

4898 + ):

4899 use_cudagraphs = False

4900

4901 self.drafter.dummy_run(

@@ -5259,7 +5270,7 @@ class GPUModelRunner(

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

5271 for batch_desc in batch_descriptors:

5272 num_tokens = batch_desc.num_tokens

5262 - activate_lora = batch_desc.has_lora

5273 + num_active_loras = batch_desc.num_active_loras

5274

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

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

@@ -5286,7 +5297,7 @@ class GPUModelRunner(

5297 num_tokens,

5298 cudagraph_runtime_mode=CUDAGraphMode.NONE,

5299 allow_microbatching=allow_microbatching,

5289 - activate_lora=activate_lora,

5300 + num_active_loras=num_active_loras,

5301 )

5302

5303 # Capture run

@@ -5294,7 +5305,7 @@ class GPUModelRunner(

5305 num_tokens,

5306 cudagraph_runtime_mode=cudagraph_runtime_mode,

5307 allow_microbatching=allow_microbatching,

5297 - activate_lora=activate_lora,

5308 + num_active_loras=num_active_loras,

5309 is_graph_capturing=True,

5310 )

5311 self.maybe_remove_all_loras(self.lora_config)

vllm/v1/worker/lora_model_runner_mixin.py

+70 -10

@@ -133,11 +133,23 @@ class LoRAModelRunnerMixin:

133 num_scheduled_tokens: np.ndarray,

134 mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,

135 num_sampled_tokens: np.ndarray | None = None,

136 - activate_lora: bool = True,

136 + num_active_loras: int = 0,

137 ):

138 + """

139 + Context manager to select dummy LoRAs for capture/warmup.

140 +

141 + Args:

142 + lora_config: LoRA configuration, or None if LoRA is disabled.

143 + num_scheduled_tokens: Array of scheduled token counts per request.

144 + num_sampled_tokens: Array of sampled token counts per request.

145 + num_active_loras: Number of distinct active LoRAs to use.

146 + - 0: No LoRA active (set up zero mappings).

147 + - >0: Use exactly this many distinct LoRAs.

148 + """

149 if num_sampled_tokens is None:

150 num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32)

151

152 + # Skip LoRA setup entirely only if no LoRA config

153 if lora_config is None:

154 yield

155 else:

@@ -145,15 +157,52 @@ class LoRAModelRunnerMixin:

157 assert self.lora_manager is not None, "LoRA is not enabled"

158

159 num_reqs = len(num_scheduled_tokens)

148 - num_loras = lora_config.max_loras

160 + max_loras = lora_config.max_loras

161 +

162 + # Determine how many distinct LoRAs to use and whether to include

163 + # no-LoRA tokens (-1 entries).

164 + # When num_active_loras > max_loras (e.g., max_loras + 1), we need

165 + # to include -1 entries to simulate batches with both LoRA and

166 + # no-LoRA tokens. This ensures prepare_tensors computes the correct

167 + # num_active_loras that matches the cudagraph capture key.

168 + if num_active_loras == 0:

169 + # No LoRA active - use 0 mappings like the original code

170 + effective_num_loras = 0

171 + include_no_lora = False

172 + elif num_active_loras > max_loras:

173 + # num_active_loras > max_loras means we want max_loras adapters

174 + # PLUS no-LoRA tokens (-1). This is the max_loras + 1 case.

175 + effective_num_loras = max_loras

176 + include_no_lora = True

177 + else:

178 + # Specific number of active LoRAs requested

179 + effective_num_loras = min(num_active_loras, max_loras)

180 + include_no_lora = False

181

182 # Make prompt lora mapping

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

152 - if activate_lora:

153 - prompt_lora_mapping = (

154 - np.arange(num_reqs, dtype=np.int32) % num_loras

155 - ) + 1

184 + # LoRA IDs are 1-indexed (1 to max_loras) as required by LoRARequest.

185 + # convert_mapping() will convert these to 0-indexed slot indices.

186 + if effective_num_loras > 0:

187 + if include_no_lora:

188 + # Include -1 (no-LoRA) entries by cycling through

189 + # -1, 1, 2, ..., effective_num_loras

190 + # This ensures prepare_tensors sees both LoRA and no-LoRA

191 + # tokens, computing num_active_loras = effective_num_loras+1

192 + cycle_values = np.array(

193 + list(range(1, effective_num_loras + 1)),

194 + dtype=np.int32,

195 + )

196 + prompt_lora_mapping = cycle_values[

197 + np.arange(num_reqs, dtype=np.int32) % len(cycle_values)

198 + ]

199 + else:

200 + # Use 1 to effective_num_loras (1-indexed lora IDs)

201 + prompt_lora_mapping = (

202 + np.arange(num_reqs, dtype=np.int32) % effective_num_loras

203 + ) + 1

204 else:

205 + # No LoRA active - use 0 for all tokens (original behavior)

206 prompt_lora_mapping = np.zeros(num_reqs, dtype=np.int32)

207

208 # Make sample lora mapping

@@ -162,14 +211,14 @@ class LoRAModelRunnerMixin:

211 # Make token lora mapping

212 token_lora_mapping = np.repeat(prompt_lora_mapping, num_scheduled_tokens)

213

165 - # Make dummy lora requests

214 + # Make dummy lora requests (only for the active LoRAs)

215 lora_requests: set[LoRARequest] = {

216 LoRARequest(

217 lora_name=f"warmup_{lora_id}",

218 lora_int_id=lora_id,

219 lora_path="/not/a/real/path",

220 )

172 - for lora_id in range(1, num_loras + 1)

221 + for lora_id in range(1, effective_num_loras + 1)

222 }

223

224 self._set_active_loras(

@@ -187,10 +236,21 @@ class LoRAModelRunnerMixin:

236 lora_config: LoRAConfig | None,

237 num_scheduled_tokens: np.ndarray,

238 num_sampled_tokens: np.ndarray,

190 - activate_lora: bool = True,

239 remove_lora: bool = True,

240 + num_active_loras: int = 0,

241 mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,

242 ):

243 + """

244 + Context manager for dummy runs with LoRA.

245 +

246 + Args:

247 + lora_config: LoRA configuration.

248 + num_scheduled_tokens: Array of scheduled token counts per request.

249 + num_sampled_tokens: Array of sampled token counts per request.

250 + remove_lora: Whether to remove LoRAs after the context exits.

251 + num_active_loras: Number of distinct active LoRAs to use.

252 + LoRA is activated when num_active_loras > 0.

253 + """

254 with (

255 self.maybe_setup_dummy_loras(lora_config, remove_lora),

256 self.maybe_select_dummy_loras(

@@ -198,7 +258,7 @@ class LoRAModelRunnerMixin:

258 num_scheduled_tokens,

259 mapping_type,

260 num_sampled_tokens,

201 - activate_lora,

261 + num_active_loras,

262 ),

263 ):

264 yield