chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix warp illegal instruction in kimi k2 thinking PCG (#15306)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-12-18 16:58:23 +08:00
committed by GitHub
parent 374ad4cce0
commit 56d12b4aea
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu
View File

@@ -126,6 +126,9 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token(
if (expert_id >= 0 && expert_id < NUM_EXPERTS) {
output_ptr[row_idx * topk + k] = shared_original_scores[expert_id];
indices_ptr[row_idx * topk + k] = expert_id;
} else {
output_ptr[row_idx * topk + k] = 0.0f;
indices_ptr[row_idx * topk + k] = 0;
}
}
@@ -219,11 +222,16 @@ __global__ void kimi_k2_moe_fused_gate_kernel(
}
}
if (lane_id == 0 && max_expert != -1) {
if (lane_id == 0) {
int64_t output_idx = row_idx * topk + k;
output_ptr[output_idx] = warp_original_scores[max_expert];
indices_ptr[output_idx] = max_expert;
warp_scores[max_expert] = -FLT_MAX;
if (max_expert != -1) {
output_ptr[output_idx] = warp_original_scores[max_expert];
indices_ptr[output_idx] = max_expert;
warp_scores[max_expert] = -FLT_MAX;
} else {
output_ptr[output_idx] = 0.0f;
indices_ptr[output_idx] = 0;
}
}
__syncwarp();