From 56d12b4aea9a8d44609f013a11d837baf6bc7218 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Thu, 18 Dec 2025 16:58:23 +0800
Subject: [PATCH] Fix warp illegal instruction in kimi k2 thinking PCG (#15306)

---
 sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu b/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu
index f674e010c..7f719e52e 100644
--- a/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu
+++ b/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu
@@ -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();