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

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

[Fix Bug]`num_active_loras` always equals to zero (#34119)

Signed-off-by: Runkai Tao <rt572@physics.rutgers.edu> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Runkai Tao committed Mar 2, 2026 at 10:17 UTC ada4f4fadd20372b1bf349961a1e442b2d07c53d

7 files changed +58 -29

tests/lora/test_fused_moe_lora_kernel.py

+4 -2

@@ -187,7 +187,8 @@ def use_fused_moe_lora_kernel(

187

188 # num_active_loras is the number of active LoRAs

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

190 - num_active_loras = max_loras + 1

190 + # Stored as CPU tensor to match the kernel API (torch.compile compatibility)

191 + num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")

192

193 fused_moe_lora(

194 output,

@@ -399,7 +400,8 @@ def use_fused_moe_lora_kernel_naive(

400

401 # num_active_loras is the number of active LoRAs

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

402 - num_active_loras = max_loras + 1

403 + # Stored as CPU tensor to match the kernel API (torch.compile compatibility)

404 + num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")

405

406 fused_moe_lora(

407 output,

tests/lora/test_gptoss_tp.py

+6 -1

@@ -70,8 +70,12 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:

72 @pytest.mark.parametrize("mxfp4_use_marlin", [True, False])

73 +@pytest.mark.parametrize("specialize_active_lora", [True, False])

74 def test_gpt_oss_lora(

74 - monkeypatch: pytest.MonkeyPatch, gptoss20b_lora_files, mxfp4_use_marlin

75 + monkeypatch: pytest.MonkeyPatch,

76 + gptoss20b_lora_files,

77 + mxfp4_use_marlin,

78 + specialize_active_lora,

79 ):

80 with monkeypatch.context() as m:

81 m.setenv("VLLM_MXFP4_USE_MARLIN", "1" if mxfp4_use_marlin else "0")

@@ -83,6 +87,7 @@ def test_gpt_oss_lora(

87 max_lora_rank=8,

88 max_num_seqs=2,

89 max_num_batched_tokens=2048,

90 + specialize_active_lora=specialize_active_lora,

91 compilation_config=vllm.config.CompilationConfig( # Avoid OOM

92 cudagraph_specialize_lora=False,

93 ),

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

+8 -8

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

127

128

129 def _adjust_kernel_inputs(

130 - num_active_loras: int,

130 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

131 sorted_token_ids: torch.Tensor | None,

132 expert_ids: torch.Tensor,

133 ):

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

141 else:

142 stride_tl = sorted_token_ids.stride(0)

143 stride_el = expert_ids.stride(0)

144 - grid_lora_dim = num_active_loras

144 + grid_lora_dim = num_active_loras.item()

145 return grid_lora_dim, stride_tl, stride_el

146

147

@@ -444,7 +444,7 @@ def _fused_moe_lora_shrink(

444 num_warps: int,

445 num_stages: int,

446 split_k: int,

447 - num_active_loras: int,

447 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

448 mul_routed_weight: bool = False,

449 use_gdc: bool = False,

450 use_tma: bool = False,

@@ -562,7 +562,7 @@ def _fused_moe_lora_expand(

562 num_warps: int,

563 num_stages: int,

564 split_k: int,

565 - num_active_loras: int,

565 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

566 mul_routed_weight: bool = False,

567 offset: int = 0,

568 use_gdc: bool = False,

@@ -683,7 +683,7 @@ def _fused_moe_lora(

683 max_lora_rank: int,

684 top_k_num: int,

685 lora_ids: torch.Tensor,

686 - num_active_loras: int,

686 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

687 adapter_enabled: torch.Tensor,

688 shrink_block_size_m: int,

689 shrink_block_size_n: int,

@@ -871,7 +871,7 @@ def _fused_moe_lora_fake(

871 max_lora_rank: int,

872 top_k_num: int,

873 lora_ids: torch.Tensor,

874 - num_active_loras: int,

874 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

875 adapter_enabled: torch.Tensor,

876 shrink_block_size_m: int,

877 shrink_block_size_n: int,

@@ -921,7 +921,7 @@ def _fused_moe_lora_shrink_fake(

921 num_warps: int,

922 num_stages: int,

923 split_k: int,

924 - num_active_loras: int,

924 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

925 mul_routed_weight: bool = False,

926 use_gdc: bool = False,

927 use_tma: bool = False,

@@ -958,7 +958,7 @@ def _fused_moe_lora_expand_fake(

958 num_warps: int,

959 num_stages: int,

960 split_k: int,

961 - num_active_loras: int,

961 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

962 mul_routed_weight: bool = False,

963 offset: int = 0,

964 use_gdc: bool = False,

vllm/lora/ops/triton_ops/lora_expand_op.py

+3 -3

@@ -138,7 +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)

141 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

142 offset_start: int = 0,

143 add_inputs: bool = False,

144 ) -> None:

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

235 grid = (

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

237 NUM_SLICES,

238 - num_active_loras,

238 + num_active_loras.item(),

239 )

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

241 # making PDL invalid and affecting the kernel performance.

@@ -289,7 +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,

292 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

293 offset_start: int = 0,

294 add_inputs: bool = False,

295 ) -> None:

vllm/lora/ops/triton_ops/lora_kernel_metadata.py

+30 -12

@@ -29,9 +29,16 @@ 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

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

33 + # Stored as a CPU tensor (not a Python int) so that torch.compile treats

34 + # it as a dynamic value rather than baking it as a constant at trace time.

35 + # This follows the same pattern as no_lora_flag_cpu above.

36 + num_active_loras_cpu: torch.Tensor

37 +

38 + # Default num_active_loras value (max_loras + 1) as a CPU tensor,

39 + # used when specialize_active_lora is False to avoid allocating a

40 + # new tensor on every meta_args() call.

41 + default_num_active_loras_cpu: torch.Tensor

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

44 # When specialize_active_lora is enabled, num_active_loras is rounded up

@@ -73,6 +80,11 @@ class LoRAKernelMeta:

81 no_lora_flag_cpu = torch.tensor([False], dtype=torch.bool, device="cpu")

83 + num_active_loras_cpu = torch.tensor([0], dtype=torch.int32, device="cpu")

84 + default_num_active_loras_cpu = torch.tensor(

85 + [max_loras + 1], dtype=torch.int32, device="cpu"

86 + )

87 +

88 return LoRAKernelMeta(

89 token_lora_mapping=token_lora_mapping,

90 token_indices_sorted_by_lora_ids=token_indices_sorted_by_lora_ids,

@@ -80,6 +92,8 @@ class LoRAKernelMeta:

92 num_tokens_per_lora=num_tokens_per_lora,

93 lora_token_start_loc=lora_token_start_loc,

94 no_lora_flag_cpu=no_lora_flag_cpu,

95 + num_active_loras_cpu=num_active_loras_cpu,

96 + default_num_active_loras_cpu=default_num_active_loras_cpu,

97 captured_lora_counts=sorted(captured_lora_counts)

98 if captured_lora_counts

99 else [],

@@ -90,8 +104,7 @@ class LoRAKernelMeta:

104 self.num_tokens_per_lora.fill_(0)

105 self.lora_token_start_loc.fill_(0)

106 self.no_lora_flag_cpu.fill_(False)

93 - self.num_active_loras = 0

94 - self.captured_lora_counts = []

107 + self.num_active_loras_cpu.fill_(0)

108

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

110 """

@@ -137,14 +150,16 @@ class LoRAKernelMeta:

150 num_tokens_per_lora, non_blocking=True

151 )

152

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

153 + num_active_loras = lora_ids.size(0)

154

155 # Round up num_active_loras to match cudagraph capture keys.

156 # 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)

157 + if self.captured_lora_counts and num_active_loras > 0:

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

159 if idx < len(self.captured_lora_counts):

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

160 + num_active_loras = self.captured_lora_counts[idx]

161 +

162 + self.num_active_loras_cpu[0] = num_active_loras

163

164 # lora_token_start_loc

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

@@ -163,7 +178,7 @@ class LoRAKernelMeta:

178 torch.Tensor,

179 torch.Tensor,

180 torch.Tensor,

166 - int,

181 + torch.Tensor,

182 ]:

183 """

184 This function returns the kernel metadata required for the current

@@ -175,7 +190,10 @@ class LoRAKernelMeta:

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

191 pass of the kernel.

192 """

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

193 + if specialize_active_lora:

194 + num_active_loras = self.num_active_loras_cpu

195 + else:

196 + num_active_loras = self.default_num_active_loras_cpu

197 return (

198 self.token_lora_mapping[:token_nums],

199 self.token_indices_sorted_by_lora_ids[:token_nums],

@@ -183,5 +201,5 @@ class LoRAKernelMeta:

201 self.lora_token_start_loc,

202 self.active_lora_ids,

203 self.no_lora_flag_cpu,

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

204 + num_active_loras,

205 )

vllm/lora/ops/triton_ops/lora_shrink_op.py

+6 -3

@@ -134,7 +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)

137 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

138 scaling: float,

139 ) -> None:

140 """

@@ -157,6 +157,9 @@ def _lora_shrink(

157 lora_ids (torch.Tensor): LoRA ids to process.

158 no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates

159 if there are any requests that require LoRA.

160 + num_active_loras (torch.Tensor): A CPU tensor of size 1, containing the

161 + number of active LoRAs. Stored as a tensor (not int) so

162 + torch.compile treats it as dynamic rather than a constant.

163 scaling (float): Scaling factor.

164 """

165

@@ -215,7 +218,7 @@ def _lora_shrink(

218 grid = (

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

220 NUM_SLICES,

218 - num_active_loras,

221 + num_active_loras.item(),

222 )

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

224 # making PDL invalid and affecting the kernel performance.

@@ -267,7 +270,7 @@ def _lora_shrink_fake(

270 lora_token_start_loc: torch.Tensor,

271 lora_ids: torch.Tensor,

272 no_lora_flag_cpu: torch.Tensor,

270 - num_active_loras: int,

273 + num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

274 scaling: float,

275 ) -> None:

276 return

vllm/v1/worker/gpu_model_runner.py

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

5379 # if we want to warm up attention or not. This is

5380 # different from the case where `FULL` implies capture

5381 # attention while `PIECEWISE` implies no attention.

5382 +

5383 dummy_run(

5384 num_tokens,

5385 cudagraph_runtime_mode=CUDAGraphMode.NONE,