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 15, 11:00 UTC
[Bugfix] fix moe marlin `topk_weight` loading (#18080)
Co-authored-by: mgoin <mgoin64@gmail.com>
Jinzhen Lin committed
May 14, 2025 at 14:31 UTC
d4154c35a241077d27418940f2553003c58dd903
1 file changed
+7
-7
csrc/moe/marlin_moe_wna16/marlin_template.h
+7
-7
@@ -473,15 +473,15 @@ __global__ void Marlin(
473
if (mul_topk_weights) {
474
#pragma unroll
475
for (int i = 0; i < 4; i++) {
476
+ int idx = tid4 * 4 + i;
477
+ idx = idx < block_num_valid_tokens ? idx : 0;
478
if constexpr (w_type == vllm::kFE2M1f) {
477
- sh_block_topk_weights[tid4 * 4 + i] = __hmul2(
478
- global_scale,
479
- Dtype::num2num2(Dtype::float2num(
480
- topk_weights_ptr[sh_block_sorted_ids[tid4 * 4 + i]])));
479
+ sh_block_topk_weights[idx] = __hmul2(
480
+ global_scale, Dtype::num2num2(Dtype::float2num(
481
+ topk_weights_ptr[sh_block_sorted_ids[idx]])));
482
} else {
482
- sh_block_topk_weights[tid4 * 4 + i] =
483
- Dtype::num2num2(Dtype::float2num(
484
- topk_weights_ptr[sh_block_sorted_ids[tid4 * 4 + i]]));
483
+ sh_block_topk_weights[idx] = Dtype::num2num2(
484
+ Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]));
485
}
486
}
487
}