xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-01-23 10:24:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] fix IMA issue in certain cases of the moe marlin kernel (#28619)

Browse Source 
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>
This commit is contained in:
Jinzhen Lin 2025-11-27 11:02:21 +08:00 committed by GitHub
parent 77740191de
commit a67dec7cba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
csrc/moe/marlin_moe_wna16/marlin_template.h
View File

@ -489,14 +489,16 @@ __global__ void Marlin(
#pragma unroll
for (int i = 0; i < 4; i++) {
int idx = tid4 * 4 + i;
idx = idx < block_num_valid_tokens ? idx : 0;
if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
sh_block_topk_weights[idx] = __hmul2(
global_scale, Dtype::num2num2(Dtype::float2num(
topk_weights_ptr[sh_block_sorted_ids[idx]])));
} else {
sh_block_topk_weights[idx] = Dtype::num2num2(
Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]));
if (idx < block_num_valid_tokens) {
if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
sh_block_topk_weights[idx] =
__hmul2(global_scale,
Dtype::num2num2(Dtype::float2num(
topk_weights_ptr[sh_block_sorted_ids[idx]])));
} else {
sh_block_topk_weights[idx] = Dtype::num2num2(
Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]));
}
}
}
}

1
vllm/model_executor/layers/fused_moe/shared_fused_moe.py
View File

@ -38,7 +38,6 @@ class SharedFusedMoE(FusedMoE):
# TODO(wentao): find the root cause and remove this condition
self.enable_eplb
or (self.moe_config.use_flashinfer_cutlass_kernels and self.dp_size > 1)
or self.use_marlin_kernels
)
and self._shared_experts is not None
)