From bda01d24357b0300e75e404b71195bf997c28108 Mon Sep 17 00:00:00 2001 From: JoeLee314 Date: Sat, 16 May 2026 13:46:44 +0800 Subject: [PATCH] [Disagg] Fix MegaMoE topk_ids dtype mismatch and FakeKVManager missing kv_args (#25380) Co-authored-by: JoeLee314 --- python/sglang/srt/disaggregation/fake/conn.py | 1 + python/sglang/srt/layers/moe/mega_moe.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py index d59641c3c..e44cbb7b3 100644 --- a/python/sglang/srt/disaggregation/fake/conn.py +++ b/python/sglang/srt/disaggregation/fake/conn.py @@ -28,6 +28,7 @@ class FakeKVManager(BaseKVManager): is_mla_backend: Optional[bool] = False, ): super().__init__(args, disaggregation_mode, server_args, is_mla_backend) + self.kv_args = args self.req_to_decode_prefix_len = {} def register_to_bootstrap(self): diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 93f4d9a15..1c13f7be9 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -206,8 +206,8 @@ def _run_mega_routed( ) if num_tokens > 0: - topk_ids_in = topk_ids - topk_weights_in = topk_weights + topk_ids_in = topk_ids.to(torch.int32) + topk_weights_in = topk_weights.to(torch.float32) else: topk_ids_in = hidden_states.new_empty((0, top_k), dtype=torch.int32) topk_weights_in = hidden_states.new_empty((0, top_k), dtype=torch.float32)