From e60f799b4019b251269d0e83bb07cba8dd6c506d Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <1182563586@qq.com>
Date: Thu, 28 May 2026 13:51:33 +0800
Subject: [PATCH] Enable Kimi-K2.5 piecewise CUDA graph (#26382)

---
 python/sglang/srt/layers/layernorm.py |  1 +
 python/sglang/srt/models/kimi_k25.py  | 13 +++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py
index 6e4b7eee4..84d78c2f3 100644
--- a/python/sglang/srt/layers/layernorm.py
+++ b/python/sglang/srt/layers/layernorm.py
@@ -178,6 +178,7 @@ def _forward_with_allreduce_fusion(
                     residual=residual,
                     weight=weight,
                     eps=norm_module.variance_epsilon,
+                    max_token_num=max(x.shape[0], 2048),
                     use_attn_tp_group=use_attn_tp_group,
                 )
                 if fused_result[0] is not None:
diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py
index c1d134e8e..b4a264848 100644
--- a/python/sglang/srt/models/kimi_k25.py
+++ b/python/sglang/srt/models/kimi_k25.py
@@ -674,6 +674,19 @@ class KimiK25ForConditionalGeneration(nn.Module):
             self.vision_tower = self.vision_tower.to(dtype=target_dtype)
             self.mm_projector = self.mm_projector.to(dtype=target_dtype)
 
+    @property
+    def model(self):
+        # Alias .model to .language_model so this class satisfies the piecewise
+        # CUDA graph gate, which checks `hasattr(model, "model")`.
+        return self.language_model
+
+    def __setattr__(self, name, value):
+        # Skip redundant self.model.model assignment in runner to avoid duplicate
+        # nn.Module registration.
+        if name == "model":
+            return
+        super().__setattr__(name, value)
+
     def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
         device = self.vision_tower.device
         target_dtype = self.vision_tower.patch_embed.proj.weight.dtype