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, 17:06 UTC
[Bugfix] Fix FusedMoE LoRA kernel offs_token out of bound value (#32279)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Jan 23, 2026 at 17:41 UTC
ecc3dd66cceebcd9fc48216ae3852a0e994fbdc1
1 file changed
+4
-2
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+4
-2
@@ -139,7 +139,9 @@ def _fused_moe_lora_kernel(
139
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
140
token_ind = stride_tl * lora_id + offs_token_id
141
offs_token = tl.load(
142
- sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0
142
+ sorted_token_ids_ptr + token_ind,
143
+ mask=token_ind < max_loras * stride_tl,
144
+ other=num_valid_tokens,
145
)
146
token_mask = offs_token < num_valid_tokens
147
@@ -185,7 +187,7 @@ def _fused_moe_lora_kernel(
187
b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
188
189
if MUL_ROUTED_WEIGHT:
188
- moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
190
+ moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
191
accumulator = accumulator * moe_weight[:, None]
192
accumulator = accumulator.to(c_ptr.dtype.element_ty)
193
# Write back the block of the output