[mori] Add a combine-kwargs hook and use_external_inp_buf plumbing (#29097)

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