[AMD][TBO] Fix mori ep dual stream accuracy (#19888)

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -1035,7 +1035,8 @@ class MaybeTboDeepEPDispatcher(BaseDispatcher):
             ]
         elif get_moe_a2a_backend().is_mori():
             self._inners = [
-                MoriEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
+                MoriEPDispatcher(instance_id=i, **kwargs)
+                for i in range(num_inner_dispatchers)
             ]
         elif get_moe_a2a_backend().is_nixl():
             self._inners = [

python/sglang/srt/layers/moe/token_dispatcher/moriep.py

@@ -195,6 +195,7 @@ def init_mori_op(
     params_dtype,
     num_max_dispatch_tokens_per_rank,
     deepep_mode,
+    instance_id=0,
 ):
 
     import mori
@@ -204,9 +205,23 @@ def init_mori_op(
 
     gpu_per_node = 8 if world_size >= 8 else world_size
 
+    group_name = f"mori"
     cpu_group = group.cpu_group
-    torch._C._distributed_c10d._register_process_group("mori", cpu_group)
+    try:
-    mori.shmem.shmem_torch_process_group_init("mori")
+        torch._C._distributed_c10d._register_process_group(group_name, cpu_group)
+    except Exception as e:
+        if "already registered" in str(e):
+            logger.info(
+                f"[MORI init] The same process group is already "
+                f"registered. Ignoring [{str(e)}]"
+            )
+        else:
+            raise
+    else:
+        # If new group is newly registered then need to init mori shmem. However
+        # if the group is registered already then need to skip init mori shmem
+        # and reuse the previous one.
+        mori.shmem.shmem_torch_process_group_init(group_name)
 
     mode = EpMode.INTRA_NODE if world_size <= 8 else EpMode.INTER_NODE
     async_mode = deepep_mode.enable_low_latency()
@@ -214,7 +229,9 @@ def init_mori_op(
         mode = EpMode.LOW_LATENCY
 
     logger.info(
-        f"[MORI init] {world_size=} {rank=} {hidden_size=} {params_dtype=} {num_max_dispatch_tokens_per_rank=} {num_local_experts=} {router_topk=} {mode=}"
+        f"[MORI init] {world_size=} {rank=} {hidden_size=} {params_dtype=} "
+        f"{num_max_dispatch_tokens_per_rank=} {num_local_experts=} "
+        f"{router_topk=} {mode=}"
     )
 
     cfg = get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank)[mode]
@@ -234,7 +251,10 @@ def init_mori_op(
     if fp8_dispatch:
         scale_dim = hidden_size // 128
     elif fp4_dispatch:
-        # FP4 kernel still takes the original hidden size and do quantization internally, so hidden_dim is not reduced. The reason is that for FP4 quantization, we need to keep the original hidden size to calculate the quantization scale correctly. don't use packed hidden size for FP4 kernel.
+        # FP4 kernel still takes the original hidden size and do quantization
+        # internally, so hidden_dim is not reduced. The reason is that for FP4
+        # quantization, we need to keep the original hidden size to calculate
+        # the quantization scale correctly. Don't use packed hidden size for FP4 kernel.
         hidden_dim = hidden_size
         scale_dim = hidden_size // 32
         data_type = torch.float4_e2m1fn_x2
@@ -308,6 +328,7 @@ class _MoriEPDispatcherImplBase:
         hidden_size: int,
         params_dtype: torch.dtype,
         deepep_mode: DeepEPMode,
+        instance_id: int = 0,
     ):
         try:
             import mori  # noqa: F401
@@ -321,6 +342,7 @@ class _MoriEPDispatcherImplBase:
         self.hidden_size = hidden_size
         self.params_dtype = params_dtype
         self.deepep_mode = deepep_mode
+        self.instance_id = instance_id
 
         self.num_max_dispatch_tokens_per_rank = get_int_env_var(
             "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 4096
@@ -335,6 +357,7 @@ class _MoriEPDispatcherImplBase:
             self.params_dtype,
             self.num_max_dispatch_tokens_per_rank,
             self.deepep_mode,
+            self.instance_id,
         )
 
         self.quant_config: Optional[dict] = None
@@ -383,7 +406,6 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
 
         self.async_finish = async_finish
         self.quant_config = {}
-        # [kk TODO] need to support mxfp4 type
         self.fp8_quant_func = get_hip_quant(QuantType.per_1x128)
         self.fp4_quant_func = get_hip_quant(QuantType.per_1x32)
         self.enable_dual_stream = is_tbo_enabled()
@@ -406,18 +428,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
     ):
         topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
 
-        previous_event = self._capture_event_if_async() if self._comm_stream else None
+        num_token = hidden_states.shape[0]
-
-        return (hidden_states, topk_weights, topk_ids, previous_event)
-
-    def dispatch_b(
-        self,
-        hidden_states,
-        topk_weights,
-        topk_ids,
-        previous_event,
-    ):
-        num_tokens = hidden_states.shape[0]
         output_dtype = hidden_states.dtype
         scale = None
 
@@ -425,8 +436,9 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
 
         if fp8_dispatch:
             # FP8 quant
-            if num_tokens > 0:
+            if num_token > 0:
-                # NOTE: aiter is able to handle token=0 case in UT. But for some reason it failed at e2e case. Root cause TBD.
+                # NOTE: aiter is able to handle token=0 case in UT. But for some
+                # reason it failed at e2e case. Root cause TBD.
                 hidden_states, scale = self.fp8_quant_func(
                     hidden_states, quant_dtype=fp8_dtype
                 )
@@ -442,7 +454,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
 
         elif fp4_dispatch:
             # FP4 quant
-            if num_tokens > 0:
+            if num_token > 0:
                 hidden_states, scale = self.fp4_quant_func(hidden_states, shuffle=False)
             else:
                 hidden_states = torch.empty(
@@ -456,6 +468,27 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
                     device=hidden_states.device,
                 )
 
+        previous_event = self._capture_event_if_async() if self._comm_stream else None
+
+        return (
+            hidden_states,
+            topk_weights,
+            topk_ids,
+            scale,
+            output_dtype,
+            previous_event,
+        )
+
+    def dispatch_b(
+        self,
+        hidden_states,
+        topk_weights,
+        topk_ids,
+        scale,
+        output_dtype,
+        previous_event,
+    ):
+
         (
             packed_recv_hidden,
             recv_topk_weights,
@@ -650,7 +683,8 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
         if fp8_dispatch:
             # FP8 quant
             if num_tokens > 0:
-                # NOTE: aiter is able to handle token=0 case in UT. But for some reason it failed at e2e case. Root cause TBD.
+                # NOTE: aiter is able to handle token=0 case in UT. But for some
+                # reason it failed at e2e case. Root cause TBD.
                 hidden_states, scale = self.fp8_quant_func(
                     hidden_states, quant_dtype=fp8_dtype
                 )
@@ -818,6 +852,7 @@ class MoriEPDispatcher(BaseDispatcher):
         deepep_mode: DeepEPMode = DeepEPMode.AUTO,
         async_finish: bool = False,
         return_recv_hook: bool = False,
+        instance_id: int = 0,
     ):
         super().__init__()
 
@@ -832,6 +867,7 @@ class MoriEPDispatcher(BaseDispatcher):
             hidden_size=hidden_size,
             params_dtype=params_dtype,
             deepep_mode=deepep_mode,
+            instance_id=instance_id,
         )
 
         if self.deepep_mode.enable_low_latency():