From 43435a2f8e3474582df5ac5bfdbe1bfccd134126 Mon Sep 17 00:00:00 2001
From: Chengze Fan <fancz2002@gmail.com>
Date: Fri, 26 Jun 2026 22:34:27 -0700
Subject: [PATCH] [mori] Add a combine-kwargs hook and use_external_inp_buf
 plumbing (#29097)

---
 .../srt/layers/moe/token_dispatcher/moriep.py    | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py
index 6f486194f..9d5ac4830 100644
--- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py
+++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py
@@ -210,6 +210,7 @@ def init_mori_op(
     dispatch_dtype=DispatchDtype.bf16,
     combine_dtype=CombineDtype.bf16,
     enable_sdma=False,
+    use_external_inp_buf=True,
 ):
 
     import mori
@@ -290,6 +291,7 @@ def init_mori_op(
         f"[MORI init] {world_size=} {rank=} {hidden_size=} {params_dtype=} "
         f"{num_max_dispatch_tokens_per_rank=} {num_local_experts=} "
         f"{router_topk=} {mode=} {dispatch_dtype=} {combine_dtype=} "
+        f"{use_external_inp_buf=} "
     )
 
     def check_mori_compatibility(kwargs: dict) -> None:
@@ -321,6 +323,7 @@ def init_mori_op(
         max_total_recv_tokens=get_int_env_var(
             "SGLANG_MORI_PREALLOC_MAX_RECV_TOKENS", 0
         ),
+        use_external_inp_buf=use_external_inp_buf,
         kernel_type=kernel_type,
         gpu_per_node=gpu_per_node,
         rdma_block_num=rdma_block_num,
@@ -389,6 +392,7 @@ class _MoriEPDispatcherImplBase:
         )
 
         self.enable_sdma = get_bool_env_var("MORI_ENABLE_SDMA", "false")
+        self.use_external_inp_buf = True
 
         self._mori_op = None
         self.dispatch_dtype = DispatchDtype.bf16
@@ -418,6 +422,7 @@ class _MoriEPDispatcherImplBase:
                 self.dispatch_dtype,
                 self.combine_dtype,
                 self.enable_sdma,
+                self.use_external_inp_buf,
             )
         return self._mori_op
 
@@ -512,6 +517,9 @@ class _MoriEPDispatcherImplBase:
         self.overlap_args = None
         self.meta_overlap_args = None
 
+    def _combine_kwargs(self, hidden_states: torch.Tensor) -> dict:
+        return {}
+
 
 class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
     def __init__(self, async_finish: bool, **kwargs):
@@ -773,7 +781,10 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
                     if self.enable_sdma
                     else self.mori_op.combine
                 )
-                combined_hidden_states = combine_fn(hidden_states, None, topk_ids)[0]
+                combine_kwargs = self._combine_kwargs(hidden_states)
+                combined_hidden_states = combine_fn(
+                    hidden_states, None, topk_ids, **combine_kwargs
+                )[0]
                 if self.enable_sdma:
                     self.mori_op.combine_recv()
 
@@ -786,8 +797,9 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
             combined_hidden_states.record_stream(comm_stream)
 
         else:
+            combine_kwargs = self._combine_kwargs(hidden_states)
             combined_hidden_states = self.mori_op.combine(
-                hidden_states, None, topk_ids
+                hidden_states, None, topk_ids, **combine_kwargs
             )[0]
 
         return combined_hidden_states, done_event