From 85ffce30af54fff826d6d969ccab5dcd48778250 Mon Sep 17 00:00:00 2001 From: b8zhong Date: Fri, 21 Nov 2025 09:39:59 -0800 Subject: [PATCH] [Piecewise CUDA Graph] Support Kimi-K2 (non-Thinking) (#13466) Co-authored-by: Brayden Zhong Co-authored-by: luoyuan.luo --- python/sglang/srt/layers/moe/topk.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 84396329d..066037db9 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -74,6 +74,29 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda: from sgl_kernel import kimi_k2_moe_fused_gate, moe_fused_gate + @torch.library.register_fake("sgl_kernel::kimi_k2_moe_fused_gate") + def _kimi_k2_moe_fused_gate( + input_tensor, + bias, + topk, + renormalize, + routed_scaling_factor, + apply_routed_scaling_factor_on_output, + ): + num_rows = input_tensor.shape[0] + topk_weights = input_tensor.new_empty( + num_rows, + topk, + dtype=torch.float32, + ) + topk_ids = input_tensor.new_empty( + num_rows, + topk, + dtype=torch.int32, + ) + return topk_weights, topk_ids + + if _is_cuda or _is_hip: from sgl_kernel import topk_softmax if _use_aiter: