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, 16:00 UTC
Using max_loras + 1 to construct grid in fused_moe_lora (#32277)
Signed-off-by: Yu Gong <yu3.gong@gmail.com>
yugong333 committed
Jan 24, 2026 at 09:39 UTC
d4dbb7af632e16e611f496bbbafbc0c276e48ac8
1 file changed
+8
-3
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+8
-3
@@ -104,7 +104,10 @@ def _fused_moe_lora_kernel(
104
if moe_enabled == 0:
105
# Early exit for the no moe lora case.
106
return
107
- max_loras = tl.num_programs(axis=2)
107
+ # The grid size on axis 2 is (max_loras + 1) to handle the no-lora case
108
+ # (lora_id == -1), but sorted_token_ids and expert_ids are allocated with
109
+ # shape (max_loras, ...). Use (num_programs - 1) for correct bounds checking.
110
+ max_loras = tl.num_programs(axis=2) - 1
111
grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
112
113
# calculate pid_m,pid_n
@@ -255,7 +258,8 @@ def _fused_moe_lora_shrink(
258
* triton.cdiv(EM, META["BLOCK_SIZE_M"])
259
* triton.cdiv(N, META["BLOCK_SIZE_N"]),
260
len(lora_a_stacked),
258
- lora_a_stacked[0].shape[0],
261
+ ## max_loras + 1 to handle the no-lora case (lora_id == -1)
262
+ lora_a_stacked[0].shape[0] + 1,
263
)
264
_fused_moe_lora_kernel[grid](
265
qcurr_hidden_states,
@@ -355,7 +359,8 @@ def _fused_moe_lora_expand(
359
grid = lambda META: (
360
triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
361
len(lora_b_stacked),
358
- lora_b_stacked[0].shape[0],
362
+ ## max_loras + 1 to handle the no-lora case (lora_id == -1)
363
+ lora_b_stacked[0].shape[0] + 1,
364
)
365
_fused_moe_lora_kernel[grid](
366
a_intermediate_cache1,