minke.yu/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
sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu
+9 -1
View File
@@ -126,6 +126,9 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token(
if (expert_id >= 0 && expert_id < NUM_EXPERTS) { if (expert_id >= 0 && expert_id < NUM_EXPERTS) {
output_ptr[row_idx * topk + k] = shared_original_scores[expert_id]; output_ptr[row_idx * topk + k] = shared_original_scores[expert_id];
indices_ptr[row_idx * topk + k] = 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; int64_t output_idx = row_idx * topk + k;
if (max_expert != -1) {
output_ptr[output_idx] = warp_original_scores[max_expert]; output_ptr[output_idx] = warp_original_scores[max_expert];
indices_ptr[output_idx] = max_expert; indices_ptr[output_idx] = max_expert;
warp_scores[max_expert] = -FLT_MAX; warp_scores[max_expert] = -FLT_MAX;
} else {
output_ptr[output_idx] = 0.0f;
indices_ptr[output_idx] = 0;
}
} }
__syncwarp(); __syncwarp();