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 12, 12:15 UTC
[Bugfix] fix IMA issue in certain cases of the moe marlin kernel (#28619)
Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: youkaichao <youkaichao@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Jinzhen Lin committed
Nov 27, 2025 at 11:02 UTC
a67dec7cba6239022b5b713845e29d9cbb294ec7
2 files changed
+10
-9
csrc/moe/marlin_moe_wna16/marlin_template.h
+10
-8
@@ -489,14 +489,16 @@ __global__ void Marlin(
489
#pragma unroll
490
for (int i = 0; i < 4; i++) {
491
int idx = tid4 * 4 + i;
492
- idx = idx < block_num_valid_tokens ? idx : 0;
493
- if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
494
- sh_block_topk_weights[idx] = __hmul2(
495
- global_scale, Dtype::num2num2(Dtype::float2num(
496
- topk_weights_ptr[sh_block_sorted_ids[idx]])));
497
- } else {
498
- sh_block_topk_weights[idx] = Dtype::num2num2(
499
- Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]));
492
+ if (idx < block_num_valid_tokens) {
493
+ if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
494
+ sh_block_topk_weights[idx] =
495
+ __hmul2(global_scale,
496
+ Dtype::num2num2(Dtype::float2num(
497
+ topk_weights_ptr[sh_block_sorted_ids[idx]])));
498
+ } else {
499
+ sh_block_topk_weights[idx] = Dtype::num2num2(
500
+ Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]));
501
+ }
502
}
503
}
504
}
vllm/model_executor/layers/fused_moe/shared_fused_moe.py
-1
@@ -38,7 +38,6 @@ class SharedFusedMoE(FusedMoE):
38
# TODO(wentao): find the root cause and remove this condition
39
self.enable_eplb
40
or (self.moe_config.use_flashinfer_cutlass_kernels and self.dp_size > 1)
41
- or self.use_marlin_kernels
41
)
42
and self._shared_experts is not None
43
)