diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index b66c39551..06027edce 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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 = [ diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 85fe8653e..a2eb8f07b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/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) - mori.shmem.shmem_torch_process_group_init("mori") + try: + 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 - - 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] + num_token = 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: - # NOTE: aiter is able to handle token=0 case in UT. But for some reason it failed at e2e case. Root cause TBD. + 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. 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():